diff --git a/README.md b/README.md index 4b1ff9d..bea1fcc 100644 --- a/README.md +++ b/README.md @@ -76,6 +76,17 @@ context leaks onto the wrong card: AGPL-3.0 too. Fine for personal/research dataset building; if this ever ships in a product, revisit. +## QC: auditing an existing collection + +`audit_stills.py` runs the same detector over a directory of stills and +reports how many contain no face at all — e.g. measuring the damage the +old person-proxy pipeline did: + +```bash +python audit_stills.py -I ~/Documents/image-training-datasets/gabriel-dropout \ + --gpu 1 --write-report old_run_report.csv +``` + ## Verification Run `--preview-dir` and eyeball the `_ann.png` files: every saved frame diff --git a/audit_stills.py b/audit_stills.py new file mode 100644 index 0000000..7908cd0 --- /dev/null +++ b/audit_stills.py @@ -0,0 +1,107 @@ +#!/usr/bin/env python3 +"""Audit an existing still collection with the anime-face detector: +how many saved images actually contain a qualifying face? + +Useful for measuring how much junk the old person-proxy pipeline left +behind, and as a QC pass over any facesaver output. + + python audit_stills.py -I --gpu 1 [--write-report report.csv] +""" + +import argparse +import csv +import os +import sys + + +def parse_args(): + p = argparse.ArgumentParser( + description=(__doc__ or "Audit stills for anime faces.").split("\n")[0]) + p.add_argument("-I", "--input-dir", required=True, + help="Directory of still images to audit.") + p.add_argument("--model", + default=os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "models", "yolov8x6_animeface.pt")) + p.add_argument("--gpu", default="1", + help="GPU index, 'cpu', or 'auto'.") + p.add_argument("--conf", type=float, default=0.3) + p.add_argument("--imgsz", type=int, default=1280) + p.add_argument("--batch", type=int, default=8) + p.add_argument("--write-report", default=None, + help="Optional CSV: filename, n_faces, max_w, max_h, " + "ok (1/0).") + return p.parse_args() + + +def main(): + args = parse_args() + if args.gpu.lower() == "cpu": + os.environ["CUDA_VISIBLE_DEVICES"] = "-1" + device = "cpu" + elif args.gpu.lower() != "auto": + os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu + device = "cuda" + else: + device = None + + import cv2 + import numpy as np + from ultralytics import YOLO + + exts = (".png", ".jpg", ".jpeg", ".webp", ".bmp") + files = sorted(f for f in os.listdir(args.input_dir) + if f.lower().endswith(exts)) + if not files: + print("No images found.") + sys.exit(1) + print(f"Auditing {len(files)} images with " + f"{os.path.basename(args.model)} on {device or 'auto'}") + + model = YOLO(args.model) + + rows = [] + no_face = 0 + # stream in chunks so 1000s of 1080p pngs don't eat RAM + for i in range(0, len(files), args.batch): + chunk = files[i:i + args.batch] + imgs, kept = [], [] + for f in chunk: + img = cv2.imread(os.path.join(args.input_dir, f)) + if img is None: + print(f" WARN unreadable: {f}") + continue + imgs.append(img) + kept.append(f) + results = model.predict(imgs, conf=args.conf, imgsz=args.imgsz, + device=device, verbose=False) + for f, res in zip(kept, results): + n, mw, mh = 0, 0, 0 + if res.boxes is not None and len(res.boxes): + xyxy = res.boxes.xyxy.cpu().numpy() + n = len(xyxy) + for (x1, y1, x2, y2) in xyxy: + mw = max(mw, int(x2 - x1)) + mh = max(mh, int(y2 - y1)) + ok = 1 if n > 0 else 0 + if not ok: + no_face += 1 + print(f" NO FACE: {f}") + rows.append((f, n, mw, mh, ok)) + done = min(i + args.batch, len(files)) + print(f" {done}/{len(files)}", flush=True) + + print() + print(f"Total: {len(rows)} with face: {len(rows) - no_face} " + f"NO face: {no_face} ({100.0 * no_face / 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_faces", "max_w", "max_h", "ok"]) + w.writerows(rows) + print(f"Report: {args.write_report}") + + +if __name__ == "__main__": + main()