#!/usr/bin/env python3 """ Pre-crop face images using MTCNN and save to a new directory. Runs face detection once over the dataset and saves cropped images to disk. Training configs can then point at the pre-cropped directory — no per-epoch MTCNN overhead during training. The output mirrors the source structure exactly: data/wiki/14/37591914.jpg -> cropped/classifier/wiki/14/37591914.jpg Resumable: already-cropped images are skipped by default. Usage: python tools/facecrop.py python tools/facecrop.py --data-dir data --output-dir cropped/classifier python tools/facecrop.py --sources wiki inpainting --device cpu python tools/facecrop.py --no-skip-existing # reprocess everything """ import argparse import sys import warnings from pathlib import Path # Suppress facenet_pytorch's torch.load FutureWarning — not fixable externally. warnings.filterwarnings("ignore", message=".*weights_only.*", category=FutureWarning) ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(ROOT)) SOURCES = ["wiki", "inpainting", "text2img", "insight"] _DETECTORS: dict[tuple[str, str], object] = {} def parse_args(): p = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter, ) p.add_argument("--data-dir", default="data", help="Source dataset root (default: data)") p.add_argument("--output-dir", default="cropped/classifier", help="Output root (default: cropped/classifier)") p.add_argument("--margin", type=float, default=0.6, help="Face box margin as fraction of box size (default: 0.6)") p.add_argument("--size", type=int, default=224, help="Output image size in px, square (default: 224)") p.add_argument("--device", default=None, help="'cpu' or 'cuda'. Default: auto-detect") p.add_argument("--sources", nargs="+", default=None, metavar="SOURCE", help=f"Only process these sources. Default: all ({', '.join(SOURCES)})") p.add_argument("--skip-existing", dest="skip_existing", action="store_true", default=True, help="Skip images already present in output-dir (default: on, resumable)") p.add_argument("--no-skip-existing", dest="skip_existing", action="store_false", help="Re-process all images even if already cropped") return p.parse_args() # ── crop helpers ────────────────────────────────────────────────────────────── def _crop_face(img, box, margin: float, size: int): from PIL import Image as PILImage x1, y1, x2, y2 = [float(v) for v in box] bw, bh = x2 - x1, y2 - y1 mx, my = bw * margin / 2, bh * margin / 2 x1 -= mx; y1 -= my; x2 += mx; y2 += my cx, cy = (x1 + x2) / 2, (y1 + y2) / 2 side = max(x2 - x1, y2 - y1) x1, y1 = cx - side / 2, cy - side / 2 x2, y2 = x1 + side, y1 + side w, h = img.size x1, y1 = max(0, x1), max(0, y1) x2, y2 = min(w, x2), min(h, y2) return img.crop((int(x1), int(y1), int(x2), int(y2))).resize((size, size), PILImage.BILINEAR) def _center_crop(img, size: int): from PIL import Image as PILImage w, h = img.size side = min(w, h) left, top = (w - side) // 2, (h - side) // 2 return img.crop((left, top, left + side, top + side)).resize((size, size), PILImage.BILINEAR) def _get_detectors(device: str): key_std = ("std", device) key_relaxed = ("relaxed", device) if key_std in _DETECTORS and key_relaxed in _DETECTORS: return _DETECTORS[key_std], _DETECTORS[key_relaxed] from facenet_pytorch import MTCNN detector = MTCNN( keep_all=False, select_largest=True, min_face_size=15, device=device, post_process=False, ) detector_relaxed = MTCNN( keep_all=False, select_largest=True, min_face_size=10, thresholds=[0.5, 0.6, 0.6], device=device, post_process=False, ) _DETECTORS[key_std] = detector _DETECTORS[key_relaxed] = detector_relaxed return detector, detector_relaxed class FaceCropper: """Reusable face cropper for notebooks/tools (not training pipeline).""" def __init__(self, margin: float = 0.6, size: int = 224, device: str | None = None): import torch self.margin = margin self.size = size self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") def __call__(self, img): from PIL import Image as PILImage detector, detector_relaxed = _get_detectors(self.device) boxes, _ = detector.detect(img) if boxes is not None and len(boxes) > 0: return _crop_face(img, boxes[0], self.margin, self.size) w, h = img.size img2x = img.resize((w * 2, h * 2), PILImage.BILINEAR) boxes2, _ = detector_relaxed.detect(img2x) if boxes2 is not None and len(boxes2) > 0: box_orig = [v / 2 for v in boxes2[0]] return _crop_face(img, box_orig, self.margin, self.size) return _center_crop(img, self.size) # ── main ────────────────────────────────────────────────────────────────────── def main(): args = parse_args() import torch from PIL import Image from tqdm import tqdm data_dir = Path(args.data_dir) output_dir = Path(args.output_dir) device = args.device or ("cuda" if torch.cuda.is_available() else "cpu") sources = args.sources or SOURCES if not data_dir.exists(): print(f"Error: data directory not found: {data_dir}") sys.exit(1) # Validate requested sources for src in sources: if not (data_dir / src).exists(): print(f"Error: source directory not found: {data_dir / src}") sys.exit(1) try: import facenet_pytorch # noqa: F401 except ImportError: print("Error: facenet_pytorch not installed.") print(" Run: pip install facenet-pytorch") sys.exit(1) print(f"Data dir: {data_dir.resolve()}") print(f"Output dir: {output_dir.resolve()}") print(f"Sources: {', '.join(sources)}") print(f"Device: {device}") print(f"Margin: {args.margin} | Size: {args.size}px") print(f"Skip exist: {args.skip_existing}") detector, detector_relaxed = _get_detectors(device) # Collect all image paths, grouped by source for per-source stats all_paths: list[Path] = [] for src in sources: for subdir in sorted((data_dir / src).iterdir()): if subdir.is_dir(): all_paths.extend(sorted(subdir.glob("*.jpg"))) print(f"\nTotal images: {len(all_paths):,}\n") n_processed = n_skipped = n_error = 0 # track per-source: detected / retry_detected / fallback src_stats: dict[str, dict] = {s: {"detected": 0, "retry": 0, "fallback": 0} for s in sources} for img_path in tqdm(all_paths, desc="Pre-cropping", unit="img"): rel = img_path.relative_to(data_dir) out_path = output_dir / rel src_name = img_path.parent.parent.name # data/wiki/14/file.jpg -> wiki if args.skip_existing and out_path.exists(): n_skipped += 1 continue out_path.parent.mkdir(parents=True, exist_ok=True) try: img = Image.open(img_path).convert("RGB") except Exception as exc: tqdm.write(f"[WARN] Cannot open {img_path.name}: {exc}") n_error += 1 continue cropped = None try: # Pass 1: detect on original image boxes, _ = detector.detect(img) if boxes is not None and len(boxes) > 0: cropped = _crop_face(img, boxes[0], args.margin, args.size) src_stats[src_name]["detected"] += 1 else: # Pass 2: upscale 2x and retry with relaxed thresholds w, h = img.size img2x = img.resize((w * 2, h * 2), Image.BILINEAR) boxes2, _ = detector_relaxed.detect(img2x) if boxes2 is not None and len(boxes2) > 0: # boxes are in upscaled coords — divide by 2 to get original coords box_orig = [v / 2 for v in boxes2[0]] cropped = _crop_face(img, box_orig, args.margin, args.size) src_stats[src_name]["retry"] += 1 else: cropped = _center_crop(img, args.size) src_stats[src_name]["fallback"] += 1 except Exception as exc: tqdm.write(f"[WARN] Detection failed for {img_path.name}: {exc}") cropped = _center_crop(img, args.size) src_stats[src_name]["fallback"] += 1 cropped.save(out_path, quality=95) n_processed += 1 total = n_processed + n_skipped n_detected = sum(s["detected"] for s in src_stats.values()) n_retry = sum(s["retry"] for s in src_stats.values()) n_fallback = sum(s["fallback"] for s in src_stats.values()) denom = max(n_processed, 1) print(f"\n{'─' * 55}") print(f" Total images : {total:>8,}") print(f" Processed : {n_processed:>8,}") print(f" Skipped (existed) : {n_skipped:>8,}") print(f" Errors : {n_error:>8,}") print(f" Pass-1 detected : {n_detected:>8,} ({n_detected / denom:.1%})") print(f" Pass-2 detected : {n_retry:>8,} ({n_retry / denom:.1%}) ← 2x upscale retry") print(f" Centre fallback : {n_fallback:>8,} ({n_fallback / denom:.1%})") print() print(f" {'Source':<12} {'pass-1':>8} {'pass-2':>8} {'fallback':>8} {'fallback%':>10}") print(f" {'─'*12} {'─'*8} {'─'*8} {'─'*8} {'─'*10}") for src in sources: s = src_stats[src] total_src = s["detected"] + s["retry"] + s["fallback"] fb_pct = s["fallback"] / max(total_src, 1) print(f" {src:<12} {s['detected']:>8,} {s['retry']:>8,} {s['fallback']:>8,} {fb_pct:>9.1%}") print(f"{'─' * 55}") print(f" Output: {output_dir.resolve()}") print() print("Next step — update your config:") print(f' "data_dir": "{output_dir}"') print(f' remove "face_crop": true (images are already cropped)') if __name__ == "__main__": main()