Johnny Fernandes
263 lines
10 KiB
Python
263 lines
10 KiB
Python
#!/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()
|