| """Hugging Face Inference Endpoints handler for the Kámárí CNN age model. |
| |
| A cheap, platform-diverse fallback for the age model: host `cnn-age-v0` on a small CPU |
| Inference Endpoint (scale-to-zero) and point the gateway's third CNN fallback at it. Mirrors |
| the Modal serve logic: detect and crop the largest face, then estimate age, under-18 |
| probability, uncertainty, and a sharpness-based quality. |
| |
| Send the raw image bytes with `Content-Type: image/jpeg`. Returns: |
| {estimated_age, p_under_18, uncertainty, face_quality, faces_detected, model_version} |
| """ |
| import base64 |
| import io |
| import os |
|
|
| import numpy as np |
| import timm |
| import torch |
| import torch.nn as nn |
| from PIL import Image |
|
|
| MEAN = [0.485, 0.456, 0.406] |
| STD = [0.229, 0.224, 0.225] |
| MODEL_VERSION = "cnn_v0.1.0-hf" |
| BACKBONE = os.environ.get("CNN_BACKBONE", "tf_efficientnetv2_s") |
|
|
|
|
| class AgeNet(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.backbone = timm.create_model(BACKBONE, pretrained=False, num_classes=0) |
| d = self.backbone.num_features |
| self.age = nn.Linear(d, 1) |
| self.minor = nn.Linear(d, 1) |
| self.logvar = nn.Linear(d, 1) |
|
|
| def forward(self, x): |
| f = self.backbone(x) |
| return self.age(f).squeeze(1), self.minor(f).squeeze(1), self.logvar(f).squeeze(1) |
|
|
|
|
| def _to_bytes(data) -> bytes: |
| """Accept the many shapes HF may hand a custom handler: raw bytes, a PIL image, |
| a base64 / data-URL string, or a dict wrapping any of those under `inputs`.""" |
| if isinstance(data, dict): |
| data = data.get("inputs", data) |
| if isinstance(data, (bytes, bytearray)): |
| return bytes(data) |
| if isinstance(data, str): |
| if data.startswith("data:"): |
| data = data.split(",", 1)[1] |
| return base64.b64decode(data) |
| buf = io.BytesIO() |
| data.convert("RGB").save(buf, format="JPEG") |
| return buf.getvalue() |
|
|
|
|
| class EndpointHandler: |
| def __init__(self, path: str = ""): |
| import cv2 |
| model = AgeNet() |
| model.load_state_dict( |
| torch.load(os.path.join(path, "best.pt"), map_location="cpu", weights_only=True)) |
| model.eval() |
| torch.set_num_threads(max(1, os.cpu_count() or 2)) |
| self.model = model |
| self.cv2 = cv2 |
| self.cascade = cv2.CascadeClassifier( |
| cv2.data.haarcascades + "haarcascade_frontalface_default.xml") |
|
|
| def __call__(self, data): |
| try: |
| return self._infer(_to_bytes(data)) |
| except Exception: |
| return {"estimated_age": 0.0, "p_under_18": 0.5, "uncertainty": 1.0, |
| "face_quality": 0.0, "faces_detected": 0, "model_version": MODEL_VERSION} |
|
|
| def _infer(self, image_bytes: bytes) -> dict: |
| cv2 = self.cv2 |
| try: |
| img = Image.open(io.BytesIO(image_bytes)).convert("RGB") |
| except Exception: |
| return {"estimated_age": 0.0, "p_under_18": 0.5, "uncertainty": 1.0, |
| "face_quality": 0.0, "faces_detected": 0, "model_version": MODEL_VERSION} |
| w, h = img.size |
| gray = cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2GRAY) |
| faces = self.cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, |
| minSize=(48, 48)) |
| if len(faces) == 0: |
| return {"estimated_age": 0.0, "p_under_18": 0.5, "uncertainty": 1.0, |
| "face_quality": 0.0, "faces_detected": 0, "model_version": MODEL_VERSION} |
|
|
| fx, fy, fw, fh = (int(v) for v in max(faces, key=lambda f: f[2] * f[3])) |
| mx, my = int(fw * 0.3), int(fh * 0.3) |
| x0, y0 = max(0, fx - mx), max(0, fy - my) |
| x1, y1 = min(w, fx + fw + mx), min(h, fy + fh + my) |
| face = img.crop((x0, y0, x1, y1)).resize((224, 224)) |
|
|
| arr = np.asarray(face, np.float32) / 255.0 |
| sharp = float(cv2.Laplacian(cv2.cvtColor(np.asarray(face), cv2.COLOR_RGB2GRAY), |
| cv2.CV_64F).var()) |
| quality = float(np.clip(0.55 + sharp / 800.0, 0.0, 1.0)) |
|
|
| x = ((arr - MEAN) / STD).transpose(2, 0, 1)[None].astype(np.float32) |
| with torch.no_grad(): |
| age, minor_logit, logvar = self.model(torch.from_numpy(x)) |
| p_under = 1.0 / (1.0 + float(np.exp(-float(minor_logit[0])))) |
| uncertainty = float(np.clip(float(np.exp(0.5 * float(logvar[0]))) / 10.0, 0, 1)) |
| return { |
| "estimated_age": round(float(age[0]), 1), |
| "p_under_18": round(p_under, 3), |
| "uncertainty": round(uncertainty, 3), |
| "face_quality": round(quality, 3), |
| "faces_detected": int(len(faces)), |
| "model_version": MODEL_VERSION, |
| } |
|
|
