diff --git a/filter_text.py b/filter_text.py new file mode 100644 index 0000000..d26b498 --- /dev/null +++ b/filter_text.py @@ -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 --gpu 1 \ + --write-report text_report.csv # dry run (default) + python filter_text.py -I --gpu 1 \ + --write-report text_report.csv \ + --move --text-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: " + "/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()