# ───────────────────────────────────────────────────────────────────────────── # face_age_filter.py — age-classification pre-filter (no face detection deps). # # This rewrite drops facenet-pytorch / MTCNN entirely — they pull torchvision # which collides with Colab's current Pillow (the classic "_util.is_directory" # ImportError). The project CLAUDE.md flags this exact failure mode. # # Strategy: # - The age classifier (HF nateraw/vit-age-classifier) runs on PIL images # directly. For datasets where the image IS a centered face (FFHQ) or # where face bbox coords are provided (IMDB has `rect` in its CSV), no # face detector is needed. # - For deepfashion (face position unknown, possibly cropped out) we'll add # a lightweight detector later — a separate concern. # # Threshold logic unchanged from the previous draft: # reject if expected age < 24 OR P(0-2)+P(3-9)+P(10-19) > 0.20 # # Paste this cell ONCE per Colab session, after super_dataset_lib.py. # ───────────────────────────────────────────────────────────────────────────── # ═════════════════════════════════════════════════════════════════════════════ # 1. ENSURE DEPS (no force-upgrades — Colab's stock transformers/torch/PIL # are kept untouched to avoid the torchvision↔Pillow ImportError chain # documented in the project CLAUDE.md). # ═════════════════════════════════════════════════════════════════════════════ import importlib, subprocess, sys def _ensure(pkg_spec: str, import_name: str | None = None): name = import_name or pkg_spec.split(">=")[0].split("==")[0].split("<")[0] try: importlib.import_module(name) except ImportError: print(f" installing missing dep: {pkg_spec}") subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", pkg_spec]) _ensure("transformers") _ensure("torch") print("face_age_filter deps OK (no force-upgrades).") # ═════════════════════════════════════════════════════════════════════════════ # 2. IMPORTS + MODEL CONFIG # ═════════════════════════════════════════════════════════════════════════════ from dataclasses import dataclass from typing import Optional import numpy as np import torch from PIL import Image as _PILImage from transformers import AutoImageProcessor, AutoModelForImageClassification AGE_MODEL_ID = "nateraw/vit-age-classifier" AGE_THRESHOLD = 24.0 MINOR_MASS_MAX = 0.20 # Device selection. # DEVICE_OVERRIDE = None → auto-detect, GPU-test-then-fallback (default) # DEVICE_OVERRIDE = "cuda" → force GPU (ignore warnings) # DEVICE_OVERRIDE = "cpu" → force CPU (~10× slower but always works) # # Auto-detect catches the case where the installed PyTorch's bundled CUDA # kernels don't include your GPU's compute capability (e.g. stock Colab torch # topping out at sm_90 vs an RTX 6000 Blackwell at sm_120). We detect by # running a tiny model forward; if it crashes, fall back to CPU. DEVICE_OVERRIDE = None def _select_device() -> str: if DEVICE_OVERRIDE in ("cpu", "cuda"): return DEVICE_OVERRIDE if not torch.cuda.is_available(): return "cpu" # Check that the GPU's capability is in torch's compiled-for list. try: cap_major, cap_minor = torch.cuda.get_device_capability(0) my_sm = f"sm_{cap_major}{cap_minor}" # Some torch builds expose get_arch_list, some don't. arch_list = getattr(torch.cuda, "get_arch_list", lambda: [])() # arch_list entries look like "sm_80" / "compute_80"; normalize. compiled_sm = {a.replace("compute_", "sm_") for a in arch_list} if compiled_sm and my_sm not in compiled_sm: print(f" GPU is {my_sm} but PyTorch was compiled for {sorted(compiled_sm)}.") print(f" Trying GPU anyway — if forward fails we'll fall back to CPU.") except Exception: pass return "cuda" DEVICE = _select_device() # Bucket → midpoint mapping. Multiplied by per-bucket probability to get a # continuous expected age estimate. AGE_BUCKETS = [ ("0-2", 1.0), ("3-9", 6.0), ("10-19", 14.0), ("20-29", 24.0), ("30-39", 34.0), ("40-49", 44.0), ("50-59", 54.0), ("60-69", 64.0), ("more than 70", 75.0), ] MINOR_BUCKETS = {"0-2", "3-9", "10-19"} # ═════════════════════════════════════════════════════════════════════════════ # 3. MODEL LOAD (singleton) # ═════════════════════════════════════════════════════════════════════════════ print(f"Loading age classifier {AGE_MODEL_ID} ({DEVICE}) …") # Fast (Rust-backed) image preprocessing is the default in current transformers; # passing use_fast= now deprecation-warns, so we pass nothing. _AGE_PROCESSOR = AutoImageProcessor.from_pretrained(AGE_MODEL_ID) _AGE_MODEL = AutoModelForImageClassification.from_pretrained(AGE_MODEL_ID).to(DEVICE).eval() _MODEL_LABELS = [_AGE_MODEL.config.id2label[i] for i in range(_AGE_MODEL.config.num_labels)] _LABEL_TO_MIDPOINT = dict(AGE_BUCKETS) _missing = [lbl for lbl, _ in AGE_BUCKETS if lbl not in _MODEL_LABELS] if _missing: print(f" WARNING — model labels don't include AGE_BUCKETS entries: {_missing}") print(f" model labels: {_MODEL_LABELS}") # GPU smoke test: run a tiny zero-tensor forward to confirm the GPU kernels # actually execute on this device. If PyTorch was compiled without our SM # version (Blackwell sm_120 on stock Colab torch) this fails immediately # rather than crashing mid-ingest. if DEVICE == "cuda": try: with torch.no_grad(): _test_in = torch.zeros(1, 3, 224, 224, device=DEVICE) _ = _AGE_MODEL(_test_in) print(f" GPU smoke test passed. VRAM: {torch.cuda.memory_allocated()/1024**3:.2f} GB") except RuntimeError as e: msg = str(e).splitlines()[0] print(f" GPU smoke test FAILED ({msg!r}) — falling back to CPU.") DEVICE = "cpu" _AGE_MODEL = _AGE_MODEL.to(DEVICE) print(f" age model relocated to CPU.") else: print(f" running on CPU (slower, but compatible).") # ═════════════════════════════════════════════════════════════════════════════ # 4. RESULT TYPE # ═════════════════════════════════════════════════════════════════════════════ @dataclass class FaceCheckResult: """Outcome of running the age filter on ONE image.""" decision: str # "pass" | "fail" expected_age: float # continuous age estimate minor_mass: float # P(0-2)+P(3-9)+P(10-19) most_likely_bucket: str # argmax bucket label most_likely_prob: float # probability of argmax bucket reasons: list # human-readable reasons for fail def to_audit(self) -> dict: return { "decision": self.decision, "expected_age": round(self.expected_age, 1), "minor_mass": round(self.minor_mass, 3), "most_likely": f"{self.most_likely_bucket} ({self.most_likely_prob:.2f})", "reasons": self.reasons, } # ═════════════════════════════════════════════════════════════════════════════ # 5. FaceAgeFilter — age-classifier-only variant # ═════════════════════════════════════════════════════════════════════════════ class FaceAgeFilter: """Runs the age classifier over images (or pre-cropped face regions). Entry points: .check_one(pil, bbox=None) — single image (optional face bbox crop) .check_batch(pils, bboxes=None) — N images, batched on GPU `bbox` (if provided) is an (x1, y1, x2, y2) tuple in pixel coords — the image is cropped to that region before classification. Useful for IMDB where the CSV provides face bbox coords. For FFHQ leave bbox=None and the whole image is classified (each FFHQ image is a centered face crop). decision_mode controls how strict the reject rule is: "strict" — fail if expected_age < age_threshold OR minor_mass > minor_mass_max (catches every borderline; gives ~30-40% reject rate on FFHQ) "balanced" — fail only if most_likely bucket is a minor bucket OR minor_mass > 0.40 (single-bucket-argmax + relaxed mass; ~10-20% reject rate) "loose" — fail only if most_likely bucket is a minor bucket (most permissive; only rejects model-confident minors) """ def __init__(self, age_threshold: float = AGE_THRESHOLD, minor_mass_max: float = MINOR_MASS_MAX, decision_mode: str = "strict", # "strict" | "balanced" | "loose" batch_size: int = 32): assert decision_mode in ("strict", "balanced", "loose") self.age_threshold = age_threshold self.minor_mass_max = minor_mass_max self.decision_mode = decision_mode self.batch_size = batch_size # ── core ──────────────────────────────────────────────────────────────── def _prep_one(self, img: _PILImage.Image, bbox: Optional[tuple] = None) -> _PILImage.Image: if img.mode != "RGB": img = img.convert("RGB") if bbox is not None: x1, y1, x2, y2 = bbox W, H = img.size x1, y1 = max(0, int(x1)), max(0, int(y1)) x2, y2 = min(W, int(x2)), min(H, int(y2)) if x2 > x1 and y2 > y1: img = img.crop((x1, y1, x2, y2)) return img def _classify_batch(self, crops: list) -> tuple: """Returns (expected_ages, minor_masses, most_likely_buckets, most_likely_probs) per crop. Each is a list aligned with `crops`.""" if not crops: return [], [], [], [] inputs = _AGE_PROCESSOR(images=crops, return_tensors="pt").to(DEVICE) with torch.no_grad(): logits = _AGE_MODEL(**inputs).logits probs = torch.softmax(logits, dim=-1).cpu().numpy() expected_ages, minor_masses = [], [] most_likely_buckets, most_likely_probs = [], [] for row in probs: exp_age, minor_mass = 0.0, 0.0 for i, label in enumerate(_MODEL_LABELS): p = float(row[i]) exp_age += p * _LABEL_TO_MIDPOINT.get(label, 0.0) if label in MINOR_BUCKETS: minor_mass += p expected_ages.append(exp_age) minor_masses.append(minor_mass) mli = int(row.argmax()) most_likely_buckets.append(_MODEL_LABELS[mli]) most_likely_probs.append(float(row[mli])) return expected_ages, minor_masses, most_likely_buckets, most_likely_probs def _decide(self, exp_age: float, minor_mass: float, most_likely_bucket: str, most_likely_prob: float) -> tuple: reasons = [] mode = self.decision_mode if mode == "strict": if exp_age < self.age_threshold: reasons.append(f"expected_age={exp_age:.1f} < {self.age_threshold}") if minor_mass > self.minor_mass_max: reasons.append(f"minor_mass={minor_mass:.2f} > {self.minor_mass_max}") elif mode == "balanced": if most_likely_bucket in MINOR_BUCKETS: reasons.append(f"most_likely={most_likely_bucket} ({most_likely_prob:.2f}) is minor bucket") elif minor_mass > 0.40: reasons.append(f"minor_mass={minor_mass:.2f} > 0.40") elif mode == "loose": if most_likely_bucket in MINOR_BUCKETS: reasons.append(f"most_likely={most_likely_bucket} ({most_likely_prob:.2f}) is minor bucket") return (("fail", reasons) if reasons else ("pass", [])) # ── public ────────────────────────────────────────────────────────────── def check_one(self, img: _PILImage.Image, bbox: Optional[tuple] = None) -> FaceCheckResult: prepped = self._prep_one(img, bbox) ea, mm, mlb, mlp = self._classify_batch([prepped]) decision, reasons = self._decide(ea[0], mm[0], mlb[0], mlp[0]) return FaceCheckResult( decision=decision, expected_age=ea[0], minor_mass=mm[0], most_likely_bucket=mlb[0], most_likely_prob=mlp[0], reasons=reasons, ) def check_batch(self, images: list, bboxes: Optional[list] = None) -> list: """Process N images. `bboxes`, if given, must have same length as `images` (use None for items where no crop should happen).""" if not images: return [] if bboxes is None: bboxes = [None] * len(images) assert len(bboxes) == len(images), "bboxes and images must align" prepped = [self._prep_one(im, bb) for im, bb in zip(images, bboxes)] all_exp, all_mm, all_mlb, all_mlp = [], [], [], [] bs = self.batch_size for start in range(0, len(prepped), bs): ea, mm, mlb, mlp = self._classify_batch(prepped[start:start + bs]) all_exp.extend(ea); all_mm.extend(mm) all_mlb.extend(mlb); all_mlp.extend(mlp) results = [] for ea, mm, mlb, mlp in zip(all_exp, all_mm, all_mlb, all_mlp): decision, reasons = self._decide(ea, mm, mlb, mlp) results.append(FaceCheckResult( decision=decision, expected_age=ea, minor_mass=mm, most_likely_bucket=mlb, most_likely_prob=mlp, reasons=reasons, )) return results print(f"face_age_filter loaded. threshold={AGE_THRESHOLD}, " f"minor_mass_max={MINOR_MASS_MAX}, batch={32}")