predict.py

"""
Perceived-demographic prediction (age, gender, race) reproducing the exact
preprocessing the models were trained and validated with.

    from predict import DemographicPredictor
    p = DemographicPredictor()
    p.predict("advert.jpg")

Pipeline, in order:
    1. InsightFace buffalo_l detection; keep the largest face, det_score >= 0.10
    2. YOLOv11n-face fallback (imgsz 1280, conf >= 0.10) when InsightFace finds nothing
    3. Expand the box by FACE_MARGIN (0.35) of the box size on every side, clamped
    4. If neither detector fires, use the whole image

Step 3 is not cosmetic. Tightening the margin from 0.35 to 0.25 costs ~2.8 points
of race accuracy and ~5.9 points of macro-F1 on our validation set.

Licence: MIT (code). Weights are non-commercial research use only.
"""

from __future__ import annotations

import os
from dataclasses import dataclass
from typing import Iterable, Optional

import numpy as np
import torch
from PIL import Image, ImageFile

ImageFile.LOAD_TRUNCATED_IMAGES = True

# --------------------------------------------------------------------------
# Configuration
# --------------------------------------------------------------------------

REPO = os.environ.get("DEMOG_REPO", "Tijmen/age-gender-race-prediction")

SUBFOLDERS = {"age": "age", "gender": "gender", "race": "race"}

GENDER_LABELS = ["Male", "Female"]
RACE_LABELS = [
    "White", "Black", "Indian", "East Asian",
    "Southeast Asian", "Middle Eastern", "Latino_Hispanic",
]

# Middle Eastern -> White, matching the US Census RACHSING recode.
RACE_TO_FOUR = {
    "White": "White",
    "Black": "Black",
    "Latino_Hispanic": "Hispanic",
    "Indian": "Asian",
    "East Asian": "Asian",
    "Southeast Asian": "Asian",
    "Middle Eastern": "White",
}

AGE_BANDS = ["0-16", "16-24", "25-34", "35-44", "45-54", "55-64", "65+"]

FACE_MARGIN = 0.35
MIN_INSIGHTFACE_SCORE = 0.10
MIN_YOLO_CONF = 0.10
YOLO_IMGSZ = 1280
YOLO_REPO = "AdamCodd/YOLOv11n-face-detection"


def age_to_band(age: Optional[float]) -> Optional[str]:
    if age is None or (isinstance(age, float) and np.isnan(age)):
        return None
    a = float(age)
    if a < 16:  return "0-16"
    if a <= 24: return "16-24"
    if a <= 34: return "25-34"
    if a <= 44: return "35-44"
    if a <= 54: return "45-54"
    if a <= 64: return "55-64"
    return "65+"


@dataclass
class FaceBox:
    box: Optional[tuple]     # (x1, y1, x2, y2) or None
    detector: str            # "insightface" | "yolo" | "none"
    score: Optional[float]


# --------------------------------------------------------------------------
# Detection
# --------------------------------------------------------------------------

class FaceCropper:
    """InsightFace primary, YOLOv11n-face fallback, whole image as last resort."""

    def __init__(self, device: str = "cuda", use_yolo_fallback: bool = True):
        self.device = device
        self.use_yolo_fallback = use_yolo_fallback
        self._insight = None
        self._yolo = None

    @property
    def insight(self):
        if self._insight is None:
            from insightface.app import FaceAnalysis
            providers = (["CUDAExecutionProvider", "CPUExecutionProvider"]
                         if self.device.startswith("cuda") else ["CPUExecutionProvider"])
            app = FaceAnalysis(name="buffalo_l", allowed_modules=["detection"],
                               providers=providers)
            app.prepare(ctx_id=0 if self.device.startswith("cuda") else -1,
                        det_size=(640, 640))
            self._insight = app
        return self._insight

    @property
    def yolo(self):
        if self._yolo is None and self.use_yolo_fallback:
            from huggingface_hub import hf_hub_download
            from ultralytics import YOLO
            path = hf_hub_download(repo_id=YOLO_REPO, filename="model.pt")
            # torch >= 2.6 defaults weights_only=True; ultralytics does not pass it.
            orig = torch.load
            def _compat(*a, **k):
                k.setdefault("weights_only", False)
                return orig(*a, **k)
            torch.load = _compat
            try:
                self._yolo = YOLO(path)
            finally:
                torch.load = orig
        return self._yolo

    def detect(self, img_rgb: np.ndarray) -> FaceBox:
        faces = self.insight.get(img_rgb)
        if faces:
            best = max(faces, key=lambda f: float(getattr(f, "det_score", 0.0)))
            score = float(getattr(best, "det_score", 0.0))
            if score >= MIN_INSIGHTFACE_SCORE:
                x1, y1, x2, y2 = [float(v) for v in best.bbox]
                return FaceBox((x1, y1, x2, y2), "insightface", score)

        if self.use_yolo_fallback and self.yolo is not None:
            res = self.yolo.predict(img_rgb, verbose=False, conf=0.01,
                                    imgsz=YOLO_IMGSZ, max_det=50)[0]
            if res.boxes is not None and len(res.boxes):
                xyxy = res.boxes.xyxy.cpu().numpy()
                conf = res.boxes.conf.cpu().numpy()
                i = int(np.argmax(conf))
                if float(conf[i]) >= MIN_YOLO_CONF:
                    return FaceBox(tuple(float(v) for v in xyxy[i]), "yolo", float(conf[i]))

        return FaceBox(None, "none", None)

    @staticmethod
    def crop(img: Image.Image, fb: FaceBox, margin: float = FACE_MARGIN) -> Image.Image:
        if fb.box is None:
            return img
        w, h = img.size
        x1, y1, x2, y2 = fb.box
        mx, my = (x2 - x1) * margin, (y2 - y1) * margin
        nx1 = max(0, int(round(x1 - mx)));  ny1 = max(0, int(round(y1 - my)))
        nx2 = min(w, int(round(x2 + mx)));  ny2 = min(h, int(round(y2 + my)))
        if nx2 <= nx1 + 1 or ny2 <= ny1 + 1:
            return img
        return img.crop((nx1, ny1, nx2, ny2))


# --------------------------------------------------------------------------
# Prediction
# --------------------------------------------------------------------------

class DemographicPredictor:
    """
    Parameters
    ----------
    repo    : Hugging Face repo id holding the three checkpoints
    device  : "cuda" | "cpu"
    detect  : run face detection. Set False only if inputs are already
              padded face crops in the training convention.
    """

    def __init__(self, repo: str = REPO, device: Optional[str] = None,
                 detect: bool = True, use_yolo_fallback: bool = True):
        try:
            from transformers import (AutoImageProcessor as _Proc,
                                      AutoModelForImageClassification as _Model)
        except Exception:
            # some environments have a broken transformers.models.auto lazy import
            from transformers import (ConvNextImageProcessor as _Proc,
                                      ConvNextForImageClassification as _Model)

        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.detect = detect
        self.cropper = FaceCropper(self.device, use_yolo_fallback) if detect else None

        self.proc, self.models = {}, {}
        for task, sub in SUBFOLDERS.items():
            self.proc[task] = _Proc.from_pretrained(repo, subfolder=sub)
            self.models[task] = (_Model
                                 .from_pretrained(repo, subfolder=sub)
                                 .to(self.device).eval())

    # -- internals ---------------------------------------------------------

    @torch.no_grad()
    def _forward(self, pils: list[Image.Image]) -> list[dict]:
        out = [{} for _ in pils]

        inp = self.proc["age"](pils, return_tensors="pt").to(self.device)
        ages = self.models["age"](**inp).logits.squeeze(-1).cpu().numpy().reshape(-1)
        for i, a in enumerate(ages):
            a = float(np.clip(a, 0, 120))
            out[i]["age"] = round(a, 1)
            out[i]["age_group"] = age_to_band(a)

        inp = self.proc["gender"](pils, return_tensors="pt").to(self.device)
        gidx = self.models["gender"](**inp).logits.argmax(-1).cpu().numpy()
        for i, g in enumerate(gidx):
            out[i]["gender"] = GENDER_LABELS[int(g)]

        inp = self.proc["race"](pils, return_tensors="pt").to(self.device)
        ridx = self.models["race"](**inp).logits.argmax(-1).cpu().numpy()
        for i, r in enumerate(ridx):
            label = RACE_LABELS[int(r)]
            out[i]["race"] = label
            out[i]["race_four"] = RACE_TO_FOUR[label]

        return out

    def _load(self, image) -> Image.Image:
        if isinstance(image, Image.Image):
            return image.convert("RGB")
        return Image.open(image).convert("RGB")

    def _prepare(self, image):
        img = self._load(image)
        if not self.detect:
            return img, FaceBox(None, "disabled", None)
        fb = self.cropper.detect(np.asarray(img))
        return self.cropper.crop(img, fb), fb

    # -- public ------------------------------------------------------------

    def predict(self, image) -> dict:
        """image: path, file object, or PIL.Image."""
        crop, fb = self._prepare(image)
        res = self._forward([crop])[0]
        res["face_found"] = fb.box is not None
        res["detector"] = fb.detector
        return res

    def predict_batch(self, images: Iterable, batch_size: int = 32) -> list[dict]:
        images = list(images)
        results: list[dict] = []
        for start in range(0, len(images), batch_size):
            chunk = images[start:start + batch_size]
            prepared = [self._prepare(im) for im in chunk]
            crops = [c for c, _ in prepared]
            preds = self._forward(crops)
            for pred, (_, fb) in zip(preds, prepared):
                pred["face_found"] = fb.box is not None
                pred["detector"] = fb.detector
            results.extend(preds)
        return results


if __name__ == "__main__":
    import sys, json
    if len(sys.argv) < 2:
        print("usage: python predict.py IMAGE [IMAGE ...]")
        raise SystemExit(1)
    p = DemographicPredictor()
    for path, res in zip(sys.argv[1:], p.predict_batch(sys.argv[1:])):
        print(f"{path}: {json.dumps(res)}")