face_age_filter.py

# ─────────────────────────────────────────────────────────────────────────────
# 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}")