File size: 4,774 Bytes
93e4260 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | """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() # already a PIL image
data.convert("RGB").save(buf, format="JPEG")
return buf.getvalue()
class EndpointHandler:
def __init__(self, path: str = ""):
import cv2 # noqa: F401 - imported lazily so the model loads even if cv2 is slow
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,
}
|