210 lines
7.5 KiB
Python
210 lines
7.5 KiB
Python
"""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()
|