Add filter_text.py: flag/move text-heavy stills via EasyOCR detect pass
Dry-run default with CSV report; --move relocates to has-text dir (reversible). Flags on >=3 confident lines or one box >=3% frame area; 1280px downscale so big overlay cards hit and incidental small detail doesn't.
This commit is contained in:
+185
@@ -0,0 +1,185 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Flag (and optionally move aside) stills dominated by overlay text:
|
||||||
|
title cards, credit crawls, eyecatch cards, preview screens.
|
||||||
|
|
||||||
|
Uses EasyOCR's text detector+recognizer on a downscaled copy of each
|
||||||
|
image (large overlay text survives downscaling; incidental background
|
||||||
|
text like shop signs mostly doesn't — a feature here, since we want to
|
||||||
|
bin graphics cards, not sign-adjacent character shots).
|
||||||
|
|
||||||
|
An image is flagged when EITHER:
|
||||||
|
- it has at least --min-lines text lines with conf >= --min-conf, OR
|
||||||
|
- a single text box covers at least --min-big-frac of the frame area
|
||||||
|
(stylized logos OCR reads badly but whose footprint is huge).
|
||||||
|
|
||||||
|
Everything is written to a CSV first; --move relocates flagged files
|
||||||
|
into --text-dir instead of deleting, so a bad pass is one `mv * ..`
|
||||||
|
away from fixed.
|
||||||
|
|
||||||
|
python filter_text.py -I <stills_dir> --gpu 1 \
|
||||||
|
--write-report text_report.csv # dry run (default)
|
||||||
|
python filter_text.py -I <stills_dir> --gpu 1 \
|
||||||
|
--write-report text_report.csv \
|
||||||
|
--move --text-dir <stills_dir>/has-text # actually move
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import csv
|
||||||
|
import os
|
||||||
|
import shutil
|
||||||
|
import sys
|
||||||
|
|
||||||
|
|
||||||
|
def parse_args():
|
||||||
|
p = argparse.ArgumentParser(
|
||||||
|
description=(__doc__ or "").split("\n")[0])
|
||||||
|
p.add_argument("-I", "--input-dir", required=True)
|
||||||
|
p.add_argument("--gpu", default="1",
|
||||||
|
help="GPU index, 'cpu', or 'auto' (CUDA-pinned the "
|
||||||
|
"same way as anime_facesaver).")
|
||||||
|
p.add_argument("--langs", default="en,ja",
|
||||||
|
help="Comma-separated EasyOCR language codes.")
|
||||||
|
p.add_argument("--max-width", type=int, default=1280,
|
||||||
|
help="Downscale width before OCR (speed; big overlay "
|
||||||
|
"text survives, incidental small text dies).")
|
||||||
|
p.add_argument("--min-conf", type=float, default=0.3,
|
||||||
|
help="Recognition confidence to count a line as text.")
|
||||||
|
p.add_argument("--min-lines", type=int, default=3,
|
||||||
|
help="Confident text lines needed to flag.")
|
||||||
|
p.add_argument("--min-big-frac", type=float, default=0.03,
|
||||||
|
help="Single text box area fraction of frame that "
|
||||||
|
"flags on its own (big stylized titles).")
|
||||||
|
p.add_argument("--write-report", default=None,
|
||||||
|
help="CSV: filename,n_lines,max_area_frac,"
|
||||||
|
"flagged,sample_text.")
|
||||||
|
p.add_argument("--move", action="store_true",
|
||||||
|
help="Move flagged files into --text-dir (default off "
|
||||||
|
"= dry run).")
|
||||||
|
p.add_argument("--text-dir", default=None,
|
||||||
|
help="Destination for --move. Default: "
|
||||||
|
"<input-dir>/has-text")
|
||||||
|
return p.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
def polygon_area(pts):
|
||||||
|
# shoelace over the 4-point box
|
||||||
|
s = 0.0
|
||||||
|
n = len(pts)
|
||||||
|
for i in range(n):
|
||||||
|
x1, y1 = pts[i]
|
||||||
|
x2, y2 = pts[(i + 1) % n]
|
||||||
|
s += x1 * y2 - x2 * y1
|
||||||
|
return abs(s) / 2.0
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
args = parse_args()
|
||||||
|
|
||||||
|
# CUDA pin, locked before torch initializes (same trap as the
|
||||||
|
# main script: easyocr imports torch straight, no select_device,
|
||||||
|
# but we still warm-alloc to keep CVD honest).
|
||||||
|
device = None
|
||||||
|
if args.gpu.lower() == "cpu":
|
||||||
|
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
|
||||||
|
device = False
|
||||||
|
elif args.gpu.lower() != "auto":
|
||||||
|
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
|
||||||
|
import torch
|
||||||
|
if torch.cuda.is_available():
|
||||||
|
torch.zeros(1, device="cuda")
|
||||||
|
else:
|
||||||
|
device = False
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
import easyocr
|
||||||
|
|
||||||
|
exts = (".png", ".jpg", ".jpeg", ".webp", ".bmp")
|
||||||
|
skip_dirs = {os.path.abspath(args.text_dir) if args.text_dir else None,
|
||||||
|
os.path.abspath(args.input_dir) + "/has-text"}
|
||||||
|
files = []
|
||||||
|
for f in sorted(os.listdir(args.input_dir)):
|
||||||
|
full = os.path.abspath(os.path.join(args.input_dir, f))
|
||||||
|
if f.lower().endswith(exts) and os.path.dirname(full) not in skip_dirs:
|
||||||
|
files.append(f)
|
||||||
|
if not files:
|
||||||
|
print("No images found.")
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
langs = [l.strip() for l in args.langs.split(",") if l.strip()]
|
||||||
|
print(f"Init EasyOCR ({langs}) ...")
|
||||||
|
reader = easyocr.Reader(langs, gpu=(device is not False))
|
||||||
|
print(f"Auditing {len(files)} images "
|
||||||
|
f"(max_width={args.max_width}px)")
|
||||||
|
|
||||||
|
rows = []
|
||||||
|
flagged = 0
|
||||||
|
for i, f in enumerate(files):
|
||||||
|
img = cv2.imread(os.path.join(args.input_dir, f))
|
||||||
|
if img is None:
|
||||||
|
print(f" WARN unreadable: {f}")
|
||||||
|
continue
|
||||||
|
h, w = img.shape[:2]
|
||||||
|
if w > args.max_width:
|
||||||
|
img = cv2.resize(img, (args.max_width,
|
||||||
|
int(h * args.max_width / w)))
|
||||||
|
h2, w2 = img.shape[:2]
|
||||||
|
frame_area = float(w2 * h2)
|
||||||
|
|
||||||
|
try:
|
||||||
|
results = reader.readtext(img, detail=True, paragraph=False)
|
||||||
|
except Exception as e:
|
||||||
|
print(f" WARN OCR failed {f}: {e}")
|
||||||
|
results = []
|
||||||
|
|
||||||
|
n_lines = 0
|
||||||
|
max_frac = 0.0
|
||||||
|
samples = []
|
||||||
|
for box, text, conf in results:
|
||||||
|
pts = [(float(x), float(y)) for x, y in box]
|
||||||
|
frac = polygon_area(pts) / frame_area
|
||||||
|
max_frac = max(max_frac, frac)
|
||||||
|
if conf >= args.min_conf:
|
||||||
|
n_lines += 1
|
||||||
|
if len(samples) < 4:
|
||||||
|
samples.append(text)
|
||||||
|
|
||||||
|
is_flag = (n_lines >= args.min_lines) or (max_frac >= args.min_big_frac)
|
||||||
|
if is_flag:
|
||||||
|
flagged += 1
|
||||||
|
print(f" TEXT: {f} lines={n_lines} "
|
||||||
|
f"max_area={max_frac*100:.1f}% "
|
||||||
|
f"'{' | '.join(samples)[:60]}'")
|
||||||
|
rows.append([f, n_lines, round(max_frac, 4),
|
||||||
|
int(is_flag), " | ".join(samples)[:120]])
|
||||||
|
|
||||||
|
if (i + 1) % 100 == 0:
|
||||||
|
print(f" {i+1}/{len(files)} done, {flagged} flagged",
|
||||||
|
flush=True)
|
||||||
|
|
||||||
|
print(f"\nTotal: {len(rows)} flagged: {flagged} "
|
||||||
|
f"({100.0*flagged/max(1,len(rows)):.1f}%)")
|
||||||
|
|
||||||
|
if args.write_report:
|
||||||
|
with open(args.write_report, "w", newline="") as fh:
|
||||||
|
w = csv.writer(fh)
|
||||||
|
w.writerow(["filename", "n_lines", "max_area_frac",
|
||||||
|
"flagged", "sample_text"])
|
||||||
|
w.writerows(rows)
|
||||||
|
print(f"Report: {args.write_report}")
|
||||||
|
|
||||||
|
if args.move:
|
||||||
|
dest = args.text_dir or os.path.join(args.input_dir, "has-text")
|
||||||
|
os.makedirs(dest, exist_ok=True)
|
||||||
|
moved = 0
|
||||||
|
for r in rows:
|
||||||
|
if r[3] == 1:
|
||||||
|
src = os.path.join(args.input_dir, r[0])
|
||||||
|
if os.path.exists(src):
|
||||||
|
shutil.move(src, os.path.join(dest, r[0]))
|
||||||
|
moved += 1
|
||||||
|
print(f"Moved {moved} flagged files to {dest} "
|
||||||
|
f"(report lists exactly which).")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
Reference in New Issue
Block a user