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