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