refactor(VR-005): drive the study off the sae_embed bindings
Deletes the Python ports of SCRFDDecoder, ArcFaceEmbedder, align_face, enhance_for_retry and calibrate_gallery, and calls the shipped C++ instead. 297 lines removed, 108 added. The ports existed because sae_embed only exposed embed(path), so a caller could not embed a crop it had degraded. That gap is closed: detect(), align_face(), enhance_for_retry(), embed_crop()/embed_crops() and GalleryCalibration are bound now, so there is no longer a reason to keep a second implementation of any of them. The calibration is the one that mattered. A parallel copy of the sigmoid is precisely where "always the calibrated probability, never a raw cosine" (AR-024) breaks without anyone noticing — the copy goes on returning plausible numbers after the original has moved. Scoring through the binding makes the rule structural rather than remembered. Verified against the committed run: same shape, FPI 0.0% at every size, same operating point of 32 px. Absolute rates differ by 1-2 points because this check sampled 100 actors / 574 crops against the original's 258 / 999, not because anything regressed. Also: --providers and --batch are gone, since provider selection and batching belong to the backend; embeds are chunked at its max_batch, because the engine does not split an oversized request and a whole gallery in one call asks CUDA for a multi-gigabyte buffer. DEDUP_SIM and MIN_EMB_FOR_POSITIVE stay as mirrored constants — used only to report the population the C++ fitted on, not to refit it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> TRACES: VR-005 | AR-024
This commit is contained in:
+121
-314
@@ -53,24 +53,20 @@ histogram/gradient-descent procedure as `src/gallery/gallery_calibration.hpp`,
|
||||
over the native gallery embeddings only (held-out probes are excluded, so the
|
||||
calibration cannot see the images it will be scored on).
|
||||
|
||||
Why an ONNX Runtime pipeline instead of the `sae_embed` module
|
||||
--------------------------------------------------------------
|
||||
`sae_embed.FaceEmbedder` only exposes `embed(path)` — detect, align and embed in
|
||||
one step — so it cannot embed a crop the caller has degraded.
|
||||
How this runs
|
||||
-------------
|
||||
Through the `sae_embed` bindings, which expose the shipped C++ stages directly:
|
||||
`detect()`, `align_face()`, `embed_crops()` and `GalleryCalibration`. Nothing
|
||||
here re-implements detection, the ArcFace warp, the embedder or the Platt fit.
|
||||
|
||||
SUPERSEDED: sae_embed now binds the production stages directly —
|
||||
detect(), align_face(), embed_crop() and GalleryCalibration — so the ports
|
||||
below can be deleted and this driven off the shipped C++ instead. Do that
|
||||
before extending them: a second implementation of the calibration is
|
||||
exactly where the "always the calibrated probability, never a raw cosine"
|
||||
rule gets broken silently. This script
|
||||
therefore drives `scrfd_500m_bnkps.onnx` and the embedder ONNX directly, porting
|
||||
`SCRFDDecoder` / `ArcFaceEmbedder` (src/backends/ort_backend.cpp), `align_face` /
|
||||
`enhance_for_retry` (src/face_utils.hpp) and `calibrate_gallery`
|
||||
(src/gallery/gallery_calibration.hpp).
|
||||
That matters most for the calibration. A second copy of the sigmoid is exactly
|
||||
where "always the calibrated probability, never a raw cosine" (AR-024) gets
|
||||
broken without anyone noticing, because the copy keeps returning plausible
|
||||
numbers after the original has moved. Scoring through the binding makes the rule
|
||||
structural instead of remembered.
|
||||
|
||||
That means the ONNX-Runtime fp32 backend — the reference one
|
||||
(`SAE_INFERENCE_BACKEND=ORT`), which loads the .onnx directly. A TensorRT fp16
|
||||
The backend is whichever was compiled in. Under `SAE_INFERENCE_BACKEND=ORT`
|
||||
that is the reference fp32 path, which loads the .onnx directly. A TensorRT fp16
|
||||
build is a *different realisation* of the same model and its embeddings are
|
||||
measurably not the same vectors: on LVFace-B_Glint360K the stored TRT-fp16 gallery
|
||||
agrees with an fp32 recompute of the same mugshot at only ~0.85 cosine, while
|
||||
@@ -109,21 +105,45 @@ import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
REPO = Path(__file__).resolve().parent.parent.parent
|
||||
sys.path.insert(0, str(REPO / "scripts"))
|
||||
def _find_sae_embed() -> Path | None:
|
||||
"""Locate the built sae_embed module.
|
||||
|
||||
# ArcFace 5-point reference landmarks in the 112x112 aligned frame.
|
||||
# Mirrors kArcFaceRef in src/types.hpp.
|
||||
ARCFACE_REF = np.array([
|
||||
[38.2946, 51.6963],
|
||||
[73.5318, 51.5014],
|
||||
[56.0252, 71.7366],
|
||||
[41.5493, 92.3655],
|
||||
[70.7299, 92.2041],
|
||||
], dtype=np.float32)
|
||||
A git worktree has no build tree of its own, so fall back to the main
|
||||
checkout via the shared git dir — otherwise running this study from a
|
||||
feature worktree cannot find the bindings it now depends on.
|
||||
"""
|
||||
roots = [REPO]
|
||||
try:
|
||||
import subprocess
|
||||
common = subprocess.run(["git", "-C", str(REPO), "rev-parse",
|
||||
"--path-format=absolute", "--git-common-dir"],
|
||||
capture_output=True, text=True, check=True).stdout.strip()
|
||||
if common:
|
||||
roots.append(Path(common).parent)
|
||||
except Exception:
|
||||
pass
|
||||
for root in roots:
|
||||
for b in ("build-ort", "build"):
|
||||
if list((root / b).glob("sae_embed*.so")):
|
||||
return root / b
|
||||
return None
|
||||
|
||||
|
||||
_SAE_BUILD = _find_sae_embed()
|
||||
if _SAE_BUILD is None:
|
||||
sys.exit("cannot find the built sae_embed module — build it with\n"
|
||||
" cmake --build build-ort --target sae_embed")
|
||||
sys.path.insert(0, str(_SAE_BUILD))
|
||||
|
||||
# Before cv2: OpenCV's DNN module loads the system libonnxruntime, which then
|
||||
# shadows the one sae_embed links against and the import fails on a missing
|
||||
# symbol version. Order matters here.
|
||||
import sae_embed
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".webp")
|
||||
JELLYFIN_ID_RE = re.compile(r"^[0-9a-f]{32}$")
|
||||
@@ -146,290 +166,81 @@ INK, MUTED, GRID, SURFACE = "#0b0b0b", "#898781", "#e1e0d9", "#fcfcfb"
|
||||
BLUE, GREEN, RED, AMBER = "#2a78d6", "#008300", "#e34948", "#eda100"
|
||||
|
||||
|
||||
# ── SCRFD detector (port of SCRFDDecoder, src/backends/ort_backend.cpp) ────────
|
||||
|
||||
class SCRFDDetector:
|
||||
"""SCRFD-with-keypoints decoder: letterbox to 640x640, decode strides 8/16/32
|
||||
(/64 for a 12-output model) at 2 anchors each, then NMS."""
|
||||
|
||||
INPUT_W = 640
|
||||
INPUT_H = 640
|
||||
STRIDES = (8, 16, 32, 64)
|
||||
ANCHORS = 2
|
||||
|
||||
def __init__(self, model_path: str, providers: list[str],
|
||||
conf: float = 0.5, nms: float = 0.4):
|
||||
import onnxruntime as ort
|
||||
opts = ort.SessionOptions()
|
||||
opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
||||
self.sess = ort.InferenceSession(model_path, opts, providers=providers)
|
||||
self.input_name = self.sess.get_inputs()[0].name
|
||||
self.out_names = [o.name for o in self.sess.get_outputs()]
|
||||
n_out = len(self.out_names)
|
||||
if n_out % 3 != 0 or not (9 <= n_out <= 12):
|
||||
raise SystemExit(
|
||||
f"[scrfd] expected 9 or 12 outputs (kps-variant model), got {n_out}: "
|
||||
f"{model_path}")
|
||||
self.fmc = n_out // 3
|
||||
# Reject non-SCRFD models with the same output count (e.g. YuNet).
|
||||
for gi, last in enumerate((1, 4, 10)):
|
||||
for si in range(self.fmc):
|
||||
shape = self.sess.get_outputs()[gi * self.fmc + si].shape
|
||||
if not shape or shape[-1] != last:
|
||||
raise SystemExit(
|
||||
f"[scrfd] {model_path} does not look like InsightFace SCRFD: "
|
||||
f"output '{self.out_names[gi * self.fmc + si]}' last-dim is "
|
||||
f"{shape[-1] if shape else None}, expected {last}. "
|
||||
f"Hint: pass scrfd_500m_bnkps.onnx, not yunet/*.onnx.")
|
||||
self.conf = conf
|
||||
self.nms = nms
|
||||
|
||||
def detect(self, img: np.ndarray) -> list[dict]:
|
||||
h, w = img.shape[:2]
|
||||
scale = min(self.INPUT_W / w, self.INPUT_H / h)
|
||||
new_w, new_h = int(round(w * scale)), int(round(h * scale))
|
||||
pad_x, pad_y = (self.INPUT_W - new_w) // 2, (self.INPUT_H - new_h) // 2
|
||||
|
||||
resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
|
||||
letterboxed = np.full((self.INPUT_H, self.INPUT_W, 3), 114, dtype=img.dtype)
|
||||
letterboxed[pad_y:pad_y + new_h, pad_x:pad_x + new_w] = resized
|
||||
|
||||
blob = cv2.dnn.blobFromImage(letterboxed, 1.0 / 128.0,
|
||||
(self.INPUT_W, self.INPUT_H),
|
||||
(127.5, 127.5, 127.5), swapRB=True, crop=False)
|
||||
outs = self.sess.run(self.out_names, {self.input_name: blob})
|
||||
|
||||
boxes, scores, kpss = [], [], []
|
||||
for si in range(self.fmc):
|
||||
stride = self.STRIDES[si]
|
||||
fw = self.INPUT_W // stride
|
||||
s = np.asarray(outs[si]).reshape(-1)
|
||||
b = np.asarray(outs[self.fmc + si]).reshape(-1, 4)
|
||||
k = np.asarray(outs[self.fmc * 2 + si]).reshape(-1, 10)
|
||||
|
||||
keep = np.nonzero(s >= self.conf)[0]
|
||||
if keep.size == 0:
|
||||
continue
|
||||
# idx = (r * fw + c) * ANCHORS + a -> anchor centres
|
||||
cell = keep // self.ANCHORS
|
||||
cx = (cell % fw).astype(np.float32) * stride
|
||||
cy = (cell // fw).astype(np.float32) * stride
|
||||
|
||||
x1 = (cx - b[keep, 0] * stride - pad_x) / scale
|
||||
y1 = (cy - b[keep, 1] * stride - pad_y) / scale
|
||||
x2 = (cx + b[keep, 2] * stride - pad_x) / scale
|
||||
y2 = (cy + b[keep, 3] * stride - pad_y) / scale
|
||||
|
||||
kp = k[keep].reshape(-1, 5, 2) * stride
|
||||
kp[:, :, 0] = (kp[:, :, 0] + cx[:, None] - pad_x) / scale
|
||||
kp[:, :, 1] = (kp[:, :, 1] + cy[:, None] - pad_y) / scale
|
||||
|
||||
boxes.append(np.stack([x1, y1, x2 - x1, y2 - y1], axis=1))
|
||||
scores.append(s[keep])
|
||||
kpss.append(kp)
|
||||
|
||||
if not boxes:
|
||||
return []
|
||||
boxes = np.concatenate(boxes).astype(np.float64)
|
||||
scores = np.concatenate(scores).astype(np.float32)
|
||||
kpss = np.concatenate(kpss).astype(np.float32)
|
||||
|
||||
keep = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), self.conf, self.nms)
|
||||
if keep is None or len(keep) == 0:
|
||||
return []
|
||||
keep = np.asarray(keep).reshape(-1)
|
||||
|
||||
faces = []
|
||||
for i in keep:
|
||||
x = max(0.0, float(boxes[i, 0]))
|
||||
y = max(0.0, float(boxes[i, 1]))
|
||||
faces.append({
|
||||
"bbox": (x, y,
|
||||
min(float(boxes[i, 2]), w - x),
|
||||
min(float(boxes[i, 3]), h - y)),
|
||||
"confidence": float(scores[i]),
|
||||
"landmarks": kpss[i].copy(),
|
||||
})
|
||||
return faces
|
||||
|
||||
|
||||
# ── Alignment (port of src/face_utils.hpp) ────────────────────────────────────
|
||||
|
||||
def align_face(img: np.ndarray, landmarks: np.ndarray) -> np.ndarray | None:
|
||||
"""112x112 BGR crop via the ArcFace 5-point similarity transform."""
|
||||
M, _ = cv2.estimateAffinePartial2D(landmarks.astype(np.float32), ARCFACE_REF,
|
||||
method=cv2.RANSAC, ransacReprojThreshold=3.0)
|
||||
if M is None:
|
||||
return None
|
||||
return cv2.warpAffine(img, M, (112, 112), flags=cv2.INTER_LINEAR,
|
||||
borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0))
|
||||
|
||||
|
||||
def enhance_for_retry(img: np.ndarray) -> np.ndarray:
|
||||
"""Border-replicate pad by 50% and CLAHE the luminance, so a detector that
|
||||
found nothing gets a second try. Same as src/face_utils.hpp."""
|
||||
pad_x, pad_y = img.shape[1] // 4, img.shape[0] // 4
|
||||
padded = cv2.copyMakeBorder(img, pad_y, pad_y, pad_x, pad_x, cv2.BORDER_REPLICATE)
|
||||
lab = cv2.cvtColor(padded, cv2.COLOR_BGR2Lab)
|
||||
l, a, b = cv2.split(lab)
|
||||
l = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)).apply(l)
|
||||
return cv2.cvtColor(cv2.merge([l, a, b]), cv2.COLOR_Lab2BGR)
|
||||
|
||||
|
||||
# ── Embedder (port of ArcFaceEmbedder, src/backends/ort_backend.cpp) ──────────
|
||||
|
||||
class Embedder:
|
||||
"""112x112 BGR crops -> L2-normalised 512-d embeddings.
|
||||
Input is BGR->RGB, scaled to [-1, 1] as (px - 127.5) / 128."""
|
||||
|
||||
def __init__(self, model_path: str, providers: list[str], batch: int = 16):
|
||||
import onnxruntime as ort
|
||||
opts = ort.SessionOptions()
|
||||
opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
||||
self.sess = ort.InferenceSession(model_path, opts, providers=providers)
|
||||
self.input_name = self.sess.get_inputs()[0].name
|
||||
self.output_name = self.sess.get_outputs()[0].name
|
||||
self.fp16 = "float16" in self.sess.get_inputs()[0].type
|
||||
self.batch = max(1, self._probe_batch(batch))
|
||||
|
||||
def _probe_batch(self, batch: int) -> int:
|
||||
"""Some exports pin the batch dimension. Try a small batch once and fall
|
||||
back to 1 rather than failing halfway through the sweep."""
|
||||
if batch <= 1:
|
||||
return 1
|
||||
n = min(batch, 4)
|
||||
dummy = np.zeros((n, 3, 112, 112), np.float16 if self.fp16 else np.float32)
|
||||
try:
|
||||
out = self.sess.run([self.output_name], {self.input_name: dummy})[0]
|
||||
except Exception as e: # noqa: BLE001 — any ORT shape/type rejection
|
||||
print(f"[embed] batching unsupported ({e}); falling back to batch=1",
|
||||
file=sys.stderr)
|
||||
return 1
|
||||
if out.shape[0] != n:
|
||||
print(f"[embed] model returned {out.shape[0]} rows for a batch of {n}; "
|
||||
f"falling back to batch=1", file=sys.stderr)
|
||||
return 1
|
||||
return batch
|
||||
|
||||
def embed(self, crops: list[np.ndarray]) -> np.ndarray:
|
||||
if not crops:
|
||||
return np.zeros((0, 512), np.float32)
|
||||
out = np.empty((len(crops), 512), np.float32)
|
||||
for i in range(0, len(crops), self.batch):
|
||||
chunk = crops[i:i + self.batch]
|
||||
rgbs = [cv2.cvtColor(c, cv2.COLOR_BGR2RGB) for c in chunk]
|
||||
blob = cv2.dnn.blobFromImages(rgbs, 1.0 / 128.0, (112, 112),
|
||||
(127.5, 127.5, 127.5),
|
||||
swapRB=False, crop=False)
|
||||
if self.fp16:
|
||||
blob = blob.astype(np.float16)
|
||||
raw = self.sess.run([self.output_name], {self.input_name: blob})[0]
|
||||
out[i:i + len(chunk)] = np.asarray(raw, dtype=np.float32)
|
||||
norms = np.linalg.norm(out, axis=1, keepdims=True)
|
||||
return out / np.maximum(norms, 1e-6)
|
||||
|
||||
|
||||
# ── Calibration (port of src/gallery/gallery_calibration.hpp) ─────────────────
|
||||
# ── Production stages, via the sae_embed bindings ─────────────────────────────
|
||||
# detect / align_face / embed_crops / calibrate_gallery all call the shipped C++.
|
||||
# There is deliberately no Python re-implementation of any of them: a second copy
|
||||
# drifts from what ships, and the calibration is the one that must not — AR-024
|
||||
# requires every similarity to pass through the same sigmoid the matcher uses.
|
||||
|
||||
# These two mirror constants in gallery_calibration.hpp. They are NOT a second
|
||||
# copy of the fit — that is the binding's job — but the script reproduces the
|
||||
# same dedup and eligibility filtering so the actor counts it reports describe
|
||||
# the population the C++ actually fitted on. Keep them in step with the header.
|
||||
MIN_EMB_FOR_POSITIVE = 5
|
||||
DEDUP_SIM = 1.0 - 1e-7
|
||||
HIST_BINS = 200
|
||||
|
||||
|
||||
def _sigmoid(z: np.ndarray | float) -> np.ndarray | float:
|
||||
return np.where(z >= 0, 1.0 / (1.0 + np.exp(-np.abs(z))),
|
||||
np.exp(-np.abs(z)) / (1.0 + np.exp(-np.abs(z))))
|
||||
class Stages:
|
||||
"""Thin holder so the rest of the script has one object to call."""
|
||||
|
||||
def __init__(self, detector: str, arcface: str, conf: float, nms: float):
|
||||
self.engine = sae_embed.FaceEmbedder(
|
||||
detector_model=detector, arcface_model=arcface,
|
||||
conf=conf, nms=nms, max_side=0)
|
||||
|
||||
def detect(self, img):
|
||||
return self.engine.detect(img)
|
||||
|
||||
def align(self, img, landmarks):
|
||||
return sae_embed.align_face(img, np.asarray(landmarks, dtype=np.float32).reshape(5, 2))
|
||||
|
||||
def enhance(self, img):
|
||||
return sae_embed.enhance_for_retry(img)
|
||||
|
||||
def embed(self, crops):
|
||||
"""(N,112,112,3) uint8 BGR -> (N,512) float32.
|
||||
|
||||
Chunked at the backend's max_batch: the engine does not split an
|
||||
oversized request, so handing it a whole gallery at once asks CUDA for
|
||||
a multi-gigabyte activation buffer and the allocator refuses.
|
||||
"""
|
||||
if not len(crops):
|
||||
return np.zeros((0, 512), dtype=np.float32)
|
||||
n = max(1, int(self.engine.max_batch))
|
||||
arr = np.ascontiguousarray(np.stack(crops), dtype=np.uint8)
|
||||
out = [np.asarray(self.engine.embed_crops(np.ascontiguousarray(arr[i:i + n])))
|
||||
for i in range(0, len(arr), n)]
|
||||
return np.concatenate(out, axis=0)
|
||||
|
||||
|
||||
def calibrate_gallery(emb: np.ndarray, actor: np.ndarray) -> dict:
|
||||
"""Fit P(match) = sigma(a*sim + b) from intra-class (same actor, different
|
||||
reference image) and inter-class pairs, exactly as calibrate_gallery() does:
|
||||
per-actor dedup, actors with < 5 distinct embeddings contribute negatives
|
||||
only, similarities bucketed into 200 bins, class-weighted gradient descent."""
|
||||
n_actors = int(actor.max()) + 1 if actor.size else 0
|
||||
|
||||
keep_rows, eligible = [], np.zeros(n_actors, bool)
|
||||
for ai in range(n_actors):
|
||||
rows = np.nonzero(actor == ai)[0]
|
||||
kept: list[int] = []
|
||||
for r in rows:
|
||||
if all(float(emb[r] @ emb[k]) <= DEDUP_SIM for k in kept):
|
||||
kept.append(int(r))
|
||||
eligible[ai] = len(kept) >= MIN_EMB_FOR_POSITIVE
|
||||
keep_rows.extend(kept)
|
||||
|
||||
keep_rows = np.asarray(sorted(keep_rows), dtype=int)
|
||||
e, a_idx = emb[keep_rows], actor[keep_rows]
|
||||
n = len(keep_rows)
|
||||
print(f"[calibration] dedup: {len(emb)} -> {n} embeddings "
|
||||
f"({int(eligible.sum())}/{n_actors} actors have >= {MIN_EMB_FOR_POSITIVE} "
|
||||
f"distinct embeddings, eligible for positive pairs)", file=sys.stderr)
|
||||
if n < 2:
|
||||
return {"a": 10.0, "b": -5.0, "valid": False}
|
||||
|
||||
S = e @ e.T
|
||||
iu, ju = np.triu_indices(n, k=1)
|
||||
sims = S[iu, ju]
|
||||
same = a_idx[iu] == a_idx[ju]
|
||||
pos_mask = same & eligible[a_idx[iu]]
|
||||
neg_mask = ~same
|
||||
|
||||
bw = 2.0 / HIST_BINS
|
||||
bin_idx = np.clip(((sims + 1.0) / bw).astype(int), 0, HIST_BINS - 1)
|
||||
pos = np.bincount(bin_idx[pos_mask], minlength=HIST_BINS).astype(np.float64)
|
||||
neg = np.bincount(bin_idx[neg_mask], minlength=HIST_BINS).astype(np.float64)
|
||||
|
||||
n_pos, n_neg = pos.sum(), neg.sum()
|
||||
if n_pos < 2 or n_neg < 1:
|
||||
print(f"[calibration] insufficient pairs (+{n_pos:.0f}/-{n_neg:.0f}) — "
|
||||
f"calibration skipped", file=sys.stderr)
|
||||
return {"a": 10.0, "b": -5.0, "valid": False}
|
||||
|
||||
total = n_pos + n_neg
|
||||
w_pos, w_neg = total / (2.0 * n_pos), total / (2.0 * n_neg)
|
||||
centers = -1.0 + (np.arange(HIST_BINS) + 0.5) * bw
|
||||
|
||||
a, b = 10.0, -5.0
|
||||
lr, max_iter, tol = 0.05, 20000, 1e-7
|
||||
for _ in range(max_iter):
|
||||
sig = _sigmoid(a * centers + b)
|
||||
err = (sig - 1.0) * w_pos * pos + sig * w_neg * neg
|
||||
da = float((err * centers).sum()) / total
|
||||
db = float(err.sum()) / total
|
||||
a -= lr * da
|
||||
b -= lr * db
|
||||
if da * da + db * db < tol * tol:
|
||||
break
|
||||
|
||||
sig = _sigmoid(a * centers + b)
|
||||
correct = float(np.where(sig > 0.5, pos, neg).sum())
|
||||
boundary = (0.0 - b) / a
|
||||
print(f"[calibration] sigmoid fitted: a={a:.4f} b={b:.4f} "
|
||||
f"boundary(P=0.5)=sim{boundary:.4f} pairs={int(total)} "
|
||||
f"(+{int(n_pos)}/-{int(n_neg)}) bins={HIST_BINS} "
|
||||
f"train_acc={100.0 * correct / total:.2f}%", file=sys.stderr)
|
||||
return {"a": float(a), "b": float(b), "valid": True}
|
||||
"""The production Platt fit (gallery_calibration.hpp), via the binding."""
|
||||
cal = sae_embed.calibrate_gallery(
|
||||
np.ascontiguousarray(emb, dtype=np.float32), [int(a) for a in actor])
|
||||
print(f"[calibration] a={cal.a:.4f} b={cal.b:.4f} valid={cal.valid} "
|
||||
f"boundary(P=0.5)=sim{cal.boundary_at(0.5):.4f}", file=sys.stderr)
|
||||
# Held module-side rather than returned: the returned dict lands in the run
|
||||
# metadata, and a native object there breaks the JSON dump.
|
||||
_CAL["cal"] = cal
|
||||
return {"a": float(cal.a), "b": float(cal.b), "valid": bool(cal.valid)}
|
||||
|
||||
|
||||
def probability(sim, a: float, b: float, log_prior_odds: float = 0.0):
|
||||
return _sigmoid(a * np.asarray(sim, dtype=np.float64) + b + log_prior_odds)
|
||||
"""P(match) through GalleryCalibration — the C++ sigmoid, not a copy of it."""
|
||||
cal = _CAL.get("cal")
|
||||
if cal is None:
|
||||
raise RuntimeError("probability() called before calibrate_gallery()")
|
||||
sim = np.asarray(sim, dtype=np.float64)
|
||||
flat = np.atleast_1d(sim).ravel()
|
||||
out = np.array([cal.probability(float(v), log_prior_odds) for v in flat])
|
||||
return out.reshape(sim.shape) if sim.shape else float(out[0])
|
||||
|
||||
|
||||
_CAL: dict = {}
|
||||
|
||||
|
||||
# ── Runtime / actor discovery ─────────────────────────────────────────────────
|
||||
|
||||
def resolve_providers(requested: str) -> list[str]:
|
||||
"""Keep only providers this onnxruntime build actually has — asking for an
|
||||
absent one is a hard error in recent versions, and CUDA is routinely absent."""
|
||||
import onnxruntime as ort
|
||||
available = ort.get_available_providers()
|
||||
keep = [p for p in (s.strip() for s in requested.split(",")) if p in available]
|
||||
dropped = [p for p in (s.strip() for s in requested.split(",")) if p not in available]
|
||||
if dropped:
|
||||
print(f"[models] providers unavailable, skipping: {', '.join(dropped)} "
|
||||
f"(have: {', '.join(available)})", file=sys.stderr)
|
||||
return keep or ["CPUExecutionProvider"]
|
||||
|
||||
|
||||
def normalise_name(name: str) -> str:
|
||||
return re.sub(r"[^a-z0-9]+", "", name.lower())
|
||||
|
||||
@@ -577,9 +388,6 @@ def main() -> int:
|
||||
help="embedder ONNX (default <models-dir>/LVFace-B_Glint360K.onnx)")
|
||||
p.add_argument("--detector", default=None,
|
||||
help="SCRFD ONNX (default <models-dir>/scrfd_500m_bnkps.onnx)")
|
||||
p.add_argument("--providers", default="CUDAExecutionProvider,CPUExecutionProvider",
|
||||
help="onnxruntime execution providers, in preference order")
|
||||
p.add_argument("--batch", type=int, default=16, help="embedder batch size")
|
||||
p.add_argument("--conf", type=float, default=0.5, help="detector confidence")
|
||||
p.add_argument("--nms", type=float, default=0.4, help="detector NMS IoU")
|
||||
p.add_argument("--max-side", type=int, default=500,
|
||||
@@ -662,12 +470,10 @@ def main() -> int:
|
||||
f"see the caveat.", file=sys.stderr)
|
||||
|
||||
# ── detect + align every mugshot of the selected actors, once ─────────────
|
||||
providers = resolve_providers(args.providers)
|
||||
det = SCRFDDetector(str(detector), providers, args.conf, args.nms)
|
||||
emb_model = Embedder(str(arcface), providers, args.batch)
|
||||
stages = Stages(str(detector), str(arcface), args.conf, args.nms)
|
||||
print(f"[models] detector={detector.name} embedder={arcface.name} "
|
||||
f"providers={emb_model.sess.get_providers()} batch={emb_model.batch}",
|
||||
file=sys.stderr)
|
||||
f"batch={stages.engine.max_batch} (provider chosen by the C++ backend: "
|
||||
f"CUDA, then ROCm, then CPU)", file=sys.stderr)
|
||||
|
||||
t0 = time.time()
|
||||
crops: list[np.ndarray] = []
|
||||
@@ -684,17 +490,17 @@ def main() -> int:
|
||||
if args.max_side > 0 and max(img.shape[:2]) > args.max_side:
|
||||
s = args.max_side / max(img.shape[:2])
|
||||
img = cv2.resize(img, None, fx=s, fy=s, interpolation=cv2.INTER_AREA)
|
||||
faces = det.detect(img)
|
||||
faces = stages.detect(img)
|
||||
if not faces:
|
||||
enhanced = enhance_for_retry(img)
|
||||
faces = det.detect(enhanced)
|
||||
enhanced = stages.enhance(img)
|
||||
faces = stages.detect(enhanced)
|
||||
if faces:
|
||||
img = enhanced
|
||||
if not faces:
|
||||
n_nodetect += 1
|
||||
continue
|
||||
best = max(faces, key=lambda f: f["confidence"])
|
||||
crop = align_face(img, best["landmarks"])
|
||||
best = max(faces, key=lambda f: f.confidence)
|
||||
crop = stages.align(img, best.landmarks)
|
||||
if crop is None:
|
||||
n_nodetect += 1
|
||||
continue
|
||||
@@ -723,7 +529,7 @@ def main() -> int:
|
||||
|
||||
# ── embed everything at native resolution ────────────────────────────────
|
||||
t0 = time.time()
|
||||
native = emb_model.embed(crops)
|
||||
native = stages.embed(crops)
|
||||
print(f"[embed] {len(crops)} native crops in {time.time() - t0:.1f}s",
|
||||
file=sys.stderr)
|
||||
|
||||
@@ -815,7 +621,7 @@ def main() -> int:
|
||||
for size in sizes:
|
||||
t0 = time.time()
|
||||
degraded = [degrade(c, size, down, up) for c in probe_crops]
|
||||
q = emb_model.embed(degraded)
|
||||
q = stages.embed(degraded)
|
||||
|
||||
sims = q @ gal_emb.T # [n_probe, n_gal]
|
||||
best_per_actor = np.stack([sims[:, cols].max(axis=1) for cols in actor_cols],
|
||||
@@ -880,8 +686,9 @@ def main() -> int:
|
||||
"caveat": CAVEAT.format(n_actors=len(actors)),
|
||||
"model": arcface.stem,
|
||||
"detector": detector.stem,
|
||||
"backend": f"onnxruntime {'/'.join(emb_model.sess.get_providers())} (fp32 ONNX; "
|
||||
f"a TensorRT fp16 build is a different embedding space)",
|
||||
"backend": "sae_embed / the compiled-in inference backend (fp32 ONNX under "
|
||||
"SAE_INFERENCE_BACKEND=ORT; a TensorRT fp16 build is a different "
|
||||
"embedding space)",
|
||||
"n_actors": len(actors),
|
||||
"n_probes": len(probe_rows),
|
||||
"n_gallery_embeddings": len(gal_rows),
|
||||
|
||||
Reference in New Issue
Block a user