""" Perceived-demographic prediction (age, gender, race) reproducing the exact preprocessing and decision rule the models were trained and validated with. from predict import DemographicPredictor p = DemographicPredictor() p.predict("advert.jpg", domain="real") # photographs p.predict("generated.png", domain="ai") # anything a generator made 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 5. Square resize to the checkpoint's own size, ImageNet normalisation 6. Race: temperature-scale the logits, add the per-domain offsets, argmax 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. Step 6 is not optional either, and `domain` is not a convenience flag: the classifier behaves differently on photographs and on generated images, so the offsets were fitted separately on each. Passing the wrong one shifts the race distribution by several points. See calibration.py. 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 from calibration import RACE4, calibrated_probs ImageFile.LOAD_TRUNCATED_IMAGES = True # -------------------------------------------------------------------------- # Configuration # -------------------------------------------------------------------------- REPO = os.environ.get("DEMOG_REPO", "TimmaJ/age-gender-race-prediction") SUBFOLDERS = {"age": "age", "gender": "gender", "race": "race"} GENDER_LABELS = ["male", "female"] RACE_LABELS = RACE4 # White, Black, Asian, Other # Fitted on the 1,000 human-labelled real ads and applied before the offsets. # Monotonic on its own, but it scales the logits against the fixed offsets, so # it is part of the decision rule rather than a cosmetic rescaling. RACE_TEMPERATURE = 0.909039 # Half-open bands: lo <= age < hi + 1. AGE_BANDS = [("0-17", 0, 17), ("18-24", 18, 24), ("25-34", 25, 34), ("35-44", 35, 44), ("45-54", 45, 54), ("55-64", 55, 64), ("65+", 65, 200)] 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) for name, lo, hi in AGE_BANDS: if lo <= a < hi + 1: return name return AGE_BANDS[-1][0] if a >= AGE_BANDS[-1][1] else AGE_BANDS[0][0] @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, or a local path device : "cuda" | "cpu" detect : run face detection. Set False only if inputs are already padded face crops in the training convention. domain : "real" for photographs, "ai" for generated imagery. Selects the race offsets; can also be overridden per call. """ def __init__(self, repo: str = REPO, device: Optional[str] = None, detect: bool = True, use_yolo_fallback: bool = True, domain: str = "real"): from transformers import (AutoImageProcessor as _Proc, AutoModelForImageClassification as _Model) self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") self.detect = detect self.domain = domain self.cropper = FaceCropper(self.device, use_yolo_fallback) if detect else None local = os.path.isdir(repo) self.proc, self.models = {}, {} for task, sub in SUBFOLDERS.items(): src = os.path.join(repo, sub) if local else repo kw = {} if local else {"subfolder": sub} self.proc[task] = _Proc.from_pretrained(src, **kw) self.models[task] = _Model.from_pretrained(src, **kw).to(self.device).eval() # -- internals --------------------------------------------------------- @torch.no_grad() def _forward(self, pils: list[Image.Image], domain: str) -> 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) gp = torch.softmax(self.models["gender"](**inp).logits, -1).cpu().numpy() for i, p in enumerate(gp): out[i]["gender"] = GENDER_LABELS[int(p.argmax())] out[i]["p_female"] = round(float(p[1]), 5) inp = self.proc["race"](pils, return_tensors="pt").to(self.device) logits = self.models["race"](**inp).logits.double().cpu().numpy() probs = torch.softmax(torch.tensor(logits) / RACE_TEMPERATURE, dim=1).numpy() cal = calibrated_probs(probs, domain=domain) for i in range(len(pils)): out[i]["race"] = RACE_LABELS[int(cal[i].argmax())] out[i]["race_probs"] = {c: round(float(probs[i, j]), 5) for j, c in enumerate(RACE_LABELS)} out[i]["domain"] = domain 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, domain: Optional[str] = None) -> dict: """image: path, file object, or PIL.Image.""" crop, fb = self._prepare(image) res = self._forward([crop], domain or self.domain)[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, domain: Optional[str] = None) -> 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] preds = self._forward([c for c, _ in prepared], domain or self.domain) 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)}")