Fotos: Fotowand je Verband + Gesichtssuche (lokales SCRFD/MobileFaceNet via onnxruntime)
This commit is contained in:
Nicolas Gruber
+209
@@ -0,0 +1,209 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user