"""Face matching: SCRFD detection + MobileFaceNet embeddings (onnxruntime, insightface buffalo_sc).""" import io import zipfile from pathlib import Path import numpy as np import onnxruntime as ort from PIL import Image DATA_DIR = Path(__file__).resolve().parent.parent / "data" MODEL_DIR = DATA_DIR / "models" ZIP_URL = "https://github.com/deepinsight/insightface/releases/download/v0.7/buffalo_sc.zip" DET_FILE = MODEL_DIR / "det_500m.onnx" REC_FILE = MODEL_DIR / "w600k_mbf.onnx" DET_SIZE = 640 DET_STRIDES = (8, 16, 32) NUM_ANCHORS = 2 DET_SCORE_THRESH = 0.40 DET_NMS_THRESH = 0.40 DET_INPUT_MEAN = 127.5 DET_INPUT_STD = 128.0 REC_SIZE = 112 REC_INPUT_MEAN = 127.5 REC_INPUT_STD = 127.5 # ArcFace 5-point reference (le, re, nose, lm, rm) in 112x112 REF5 = np.array([ [38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366], [41.5493, 92.3655], [70.7299, 92.2041], ], dtype=np.float64) def _nms(boxes, scores, thresh): x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3] areas = (x2 - x1 + 1) * (y2 - y1 + 1) order = scores.argsort()[::-1] keep = [] while order.size > 0: i = order[0] keep.append(i) rest = order[1:] if rest.size == 0: break xx1 = np.maximum(x1[i], x1[rest]) yy1 = np.maximum(y1[i], y1[rest]) xx2 = np.minimum(x2[i], x2[rest]) yy2 = np.minimum(y2[i], y2[rest]) inter = np.maximum(0.0, xx2 - xx1 + 1) * np.maximum(0.0, yy2 - yy1 + 1) iou = inter / (areas[i] + areas[rest] - inter) order = rest[iou <= thresh] return keep class FaceEngine: def __init__(self): self._det = None self._rec = None self._det_in = "" self._rec_in = "" self._lock = None def _load(self): if self._det is not None: return import threading self._lock = threading.Lock() with self._lock: if self._det is not None: return if not DET_FILE.exists() or not REC_FILE.exists(): MODEL_DIR.mkdir(parents=True, exist_ok=True) import time as _time import urllib.request req = urllib.request.Request(ZIP_URL, headers={"User-Agent": "Mozilla/5.0"}) for attempt in range(3): try: blob = urllib.request.urlopen(req, timeout=600).read() break except Exception: if attempt == 2: raise _time.sleep(2 * (attempt + 1)) with zipfile.ZipFile(io.BytesIO(blob)) as zf: zf.extractall(MODEL_DIR) opts = ort.SessionOptions() opts.intra_op_num_threads = 2 self._det = ort.InferenceSession(str(DET_FILE), opts, providers=["CPUExecutionProvider"]) self._rec = ort.InferenceSession(str(REC_FILE), opts, providers=["CPUExecutionProvider"]) self._det_in = self._det.get_inputs()[0].name self._rec_in = self._rec.get_inputs()[0].name def _detect(self, img: Image.Image): """SCRFD decode (insightface). Returns (box (x1,y1,x2,y2), kps (5,2)) in original coords, or None.""" h, w = img.height, img.width im_ratio = h / w if im_ratio > 1.0: new_h = DET_SIZE new_w = int(new_h / im_ratio) else: new_w = DET_SIZE new_h = int(new_w * im_ratio) det_scale = new_h / h det = np.zeros((DET_SIZE, DET_SIZE, 3), dtype=np.uint8) det[:new_h, :new_w] = np.asarray(img.convert("RGB").resize((new_w, new_h))) blob = ((det.astype(np.float32) - DET_INPUT_MEAN) / DET_INPUT_STD).transpose(2, 0, 1)[None] outs = self._det.run(None, {self._det_in: blob}) scores_list, bboxes_list, kpss_list = [], [], [] for idx, stride in enumerate(DET_STRIDES): scores = outs[idx][:, 0] keep = scores >= DET_SCORE_THRESH if not keep.any(): continue H = DET_SIZE // stride W = DET_SIZE // stride centers = (np.stack(np.mgrid[:H, :W][::-1], axis=-1).astype(np.float32) * stride).reshape(-1, 2) centers = np.stack([centers] * NUM_ANCHORS, axis=1).reshape(-1, 2) c = centers[keep] b = outs[idx + 3][keep] * stride boxes = np.stack([c[:, 0] - b[:, 0], c[:, 1] - b[:, 1], c[:, 0] + b[:, 2], c[:, 1] + b[:, 3]], axis=-1) k = outs[idx + 6][keep] * stride kpss = np.stack([c[:, 0:1] + k[:, 0::2], c[:, 1:2] + k[:, 1::2]], axis=-1).reshape(-1, 5, 2) scores_list.append(scores[keep]) bboxes_list.append(boxes / det_scale) kpss_list.append(kpss / det_scale) if not scores_list: return None scores = np.concatenate(scores_list) bboxes = np.concatenate(bboxes_list) kpss = np.concatenate(kpss_list) order = np.argsort(-scores, kind="stable") keep = _nms(bboxes[order], scores[order], DET_NMS_THRESH) if not keep: return None sel = [order[k] for k in keep] areas = (bboxes[sel, 2] - bboxes[sel, 0]) * (bboxes[sel, 3] - bboxes[sel, 1]) cx = (bboxes[sel, 0] + bboxes[sel, 2]) / 2 - w / 2 cy = (bboxes[sel, 1] + bboxes[sel, 3]) / 2 - h / 2 values = areas - 2.0 * (cx ** 2 + cy ** 2) best = sel[int(np.argmax(values))] return bboxes[best], kpss[best] @staticmethod def _align(img: Image.Image, pts) -> Image.Image: src = pts.astype(np.float64) dst = REF5 A = [] bvec = [] for (x, y), (X, Y) in zip(src, dst): A.append([x, -y, 1, 0]) A.append([y, x, 0, 1]) bvec.append(X) bvec.append(Y) sol, *_ = np.linalg.lstsq(np.array(A), np.array(bvec), rcond=None) a, b, c, d = sol s = a * a + b * b if s < 1e-6: return img ap = a / s bp = b / s cp = -(a * c + b * d) / s dp = -b / s ep = a / s fp = (b * c - a * d) / s return img.transform((REC_SIZE, REC_SIZE), Image.AFFINE, (ap, bp, cp, dp, ep, fp), Image.BICUBIC, fillcolor=(128, 128, 128)) def embed(self, img_bytes: bytes) -> np.ndarray | None: """512-d L2-normalized embedding of the most prominent face, or None.""" self._load() try: img = Image.open(io.BytesIO(img_bytes)) img.load() except Exception: return None det = self._detect(img) if not det: return None box, pts = det x1, y1, x2, y2 = [int(v) for v in box] x1 = max(0, x1) y1 = max(0, y1) x2 = min(img.width, x2) y2 = min(img.height, y2) if x2 - x1 < 20 or y2 - y1 < 20: return None if pts is not None: face = self._align(img, pts) else: face = img.crop((x1, y1, x2, y2)).resize((REC_SIZE, REC_SIZE)) arr = np.asarray(face.convert("RGB"), dtype=np.float32) arr = (arr - REC_INPUT_MEAN) / REC_INPUT_STD arr = arr.transpose(2, 0, 1)[None] emb = self._rec.run(None, {self._rec_in: arr})[0][0] n = np.linalg.norm(emb) if n < 1e-6: return None return (emb / n).astype(np.float32) @staticmethod def similarity(a: np.ndarray, b: np.ndarray) -> float: return float(np.dot(a, b)) engine = FaceEngine()