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:
2026-07-31 15:51:50 +02:00
co-authored by Claude Opus 5
parent 042e424961
commit 66c9ca0a0c
+121 -314
View File
@@ -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),