dtourolle/scene-actor-extraction
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(VR-005): drive the study off the sae_embed bindings

Browse Source 
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:
dtourolle
2026-07-31 15:51:50 +02:00
co-authored by Claude Opus 5
parent 042e424961
commit 66c9ca0a0c
1 changed files with 121 additions and 314 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/validation/min_face_size.py
+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 over the native gallery embeddings only (held-out probes are excluded, so the
calibration cannot see the images it will be scored on). calibration cannot see the images it will be scored on).
Why an ONNX Runtime pipeline instead of the `sae_embed` module How this runs
-------------------------------------------------------------- -------------
`sae_embed.FaceEmbedder` only exposes `embed(path)` — detect, align and embed in Through the `sae_embed` bindings, which expose the shipped C++ stages directly:
one step — so it cannot embed a crop the caller has degraded. `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 — That matters most for the calibration. A second copy of the sigmoid is exactly
detect(), align_face(), embed_crop() and GalleryCalibration — so the ports where "always the calibrated probability, never a raw cosine" (AR-024) gets
below can be deleted and this driven off the shipped C++ instead. Do that broken without anyone noticing, because the copy keeps returning plausible
before extending them: a second implementation of the calibration is numbers after the original has moved. Scoring through the binding makes the rule
exactly where the "always the calibrated probability, never a raw cosine" structural instead of remembered.
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 means the ONNX-Runtime fp32 backend — the reference one The backend is whichever was compiled in. Under `SAE_INFERENCE_BACKEND=ORT`
(`SAE_INFERENCE_BACKEND=ORT`), which loads the .onnx directly. A TensorRT fp16 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 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 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 agrees with an fp32 recompute of the same mugshot at only ~0.85 cosine, while
@@ -109,21 +105,45 @@ import sys
import time import time
from pathlib import Path from pathlib import Path
import cv2
import numpy as np
REPO = Path(__file__).resolve().parent.parent.parent REPO = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(REPO / "scripts")) 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. A git worktree has no build tree of its own, so fall back to the main
# Mirrors kArcFaceRef in src/types.hpp. checkout via the shared git dir — otherwise running this study from a
ARCFACE_REF = np.array([ feature worktree cannot find the bindings it now depends on.
[38.2946, 51.6963], """
[73.5318, 51.5014], roots = [REPO]
[56.0252, 71.7366], try:
[41.5493, 92.3655], import subprocess
[70.7299, 92.2041], common = subprocess.run(["git", "-C", str(REPO), "rev-parse",
], dtype=np.float32) "--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") IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".webp")
JELLYFIN_ID_RE = re.compile(r"^[0-9a-f]{32}$") 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" BLUE, GREEN, RED, AMBER = "#2a78d6", "#008300", "#e34948", "#eda100"
# ── SCRFD detector (port of SCRFDDecoder, src/backends/ort_backend.cpp) ──────── # ── Production stages, via the sae_embed bindings ─────────────────────────────
# detect / align_face / embed_crops / calibrate_gallery all call the shipped C++.
class SCRFDDetector: # There is deliberately no Python re-implementation of any of them: a second copy
"""SCRFD-with-keypoints decoder: letterbox to 640x640, decode strides 8/16/32 # drifts from what ships, and the calibration is the one that must not — AR-024
(/64 for a 12-output model) at 2 anchors each, then NMS.""" # requires every similarity to pass through the same sigmoid the matcher uses.
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) ─────────────────
# 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 MIN_EMB_FOR_POSITIVE = 5
DEDUP_SIM = 1.0 - 1e-7 DEDUP_SIM = 1.0 - 1e-7
HIST_BINS = 200
def _sigmoid(z: np.ndarray | float) -> np.ndarray | float: class Stages:
return np.where(z >= 0, 1.0 / (1.0 + np.exp(-np.abs(z))), """Thin holder so the rest of the script has one object to call."""
np.exp(-np.abs(z)) / (1.0 + np.exp(-np.abs(z))))
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: def calibrate_gallery(emb: np.ndarray, actor: np.ndarray) -> dict:
"""Fit P(match) = sigma(a*sim + b) from intra-class (same actor, different """The production Platt fit (gallery_calibration.hpp), via the binding."""
reference image) and inter-class pairs, exactly as calibrate_gallery() does: cal = sae_embed.calibrate_gallery(
per-actor dedup, actors with < 5 distinct embeddings contribute negatives np.ascontiguousarray(emb, dtype=np.float32), [int(a) for a in actor])
only, similarities bucketed into 200 bins, class-weighted gradient descent.""" print(f"[calibration] a={cal.a:.4f} b={cal.b:.4f} valid={cal.valid} "
n_actors = int(actor.max()) + 1 if actor.size else 0 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
keep_rows, eligible = [], np.zeros(n_actors, bool) # metadata, and a native object there breaks the JSON dump.
for ai in range(n_actors): _CAL["cal"] = cal
rows = np.nonzero(actor == ai)[0] return {"a": float(cal.a), "b": float(cal.b), "valid": bool(cal.valid)}
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}
def probability(sim, a: float, b: float, log_prior_odds: float = 0.0): 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 ───────────────────────────────────────────────── # ── 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: def normalise_name(name: str) -> str:
return re.sub(r"[^a-z0-9]+", "", name.lower()) 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)") help="embedder ONNX (default <models-dir>/LVFace-B_Glint360K.onnx)")
p.add_argument("--detector", default=None, p.add_argument("--detector", default=None,
help="SCRFD ONNX (default <models-dir>/scrfd_500m_bnkps.onnx)") 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("--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("--nms", type=float, default=0.4, help="detector NMS IoU")
p.add_argument("--max-side", type=int, default=500, p.add_argument("--max-side", type=int, default=500,
@@ -662,12 +470,10 @@ def main() -> int:
f"see the caveat.", file=sys.stderr) f"see the caveat.", file=sys.stderr)
# ── detect + align every mugshot of the selected actors, once ───────────── # ── detect + align every mugshot of the selected actors, once ─────────────
providers = resolve_providers(args.providers) stages = Stages(str(detector), str(arcface), args.conf, args.nms)
det = SCRFDDetector(str(detector), providers, args.conf, args.nms)
emb_model = Embedder(str(arcface), providers, args.batch)
print(f"[models] detector={detector.name} embedder={arcface.name} " print(f"[models] detector={detector.name} embedder={arcface.name} "
f"providers={emb_model.sess.get_providers()} batch={emb_model.batch}", f"batch={stages.engine.max_batch} (provider chosen by the C++ backend: "
file=sys.stderr) f"CUDA, then ROCm, then CPU)", file=sys.stderr)
t0 = time.time() t0 = time.time()
crops: list[np.ndarray] = [] crops: list[np.ndarray] = []
@@ -684,17 +490,17 @@ def main() -> int:
if args.max_side > 0 and max(img.shape[:2]) > args.max_side: if args.max_side > 0 and max(img.shape[:2]) > args.max_side:
s = args.max_side / max(img.shape[:2]) s = args.max_side / max(img.shape[:2])
img = cv2.resize(img, None, fx=s, fy=s, interpolation=cv2.INTER_AREA) img = cv2.resize(img, None, fx=s, fy=s, interpolation=cv2.INTER_AREA)
faces = det.detect(img) faces = stages.detect(img)
if not faces: if not faces:
enhanced = enhance_for_retry(img) enhanced = stages.enhance(img)
faces = det.detect(enhanced) faces = stages.detect(enhanced)
if faces: if faces:
img = enhanced img = enhanced
if not faces: if not faces:
n_nodetect += 1 n_nodetect += 1
continue continue
best = max(faces, key=lambda f: f["confidence"]) best = max(faces, key=lambda f: f.confidence)
crop = align_face(img, best["landmarks"]) crop = stages.align(img, best.landmarks)
if crop is None: if crop is None:
n_nodetect += 1 n_nodetect += 1
continue continue
@@ -723,7 +529,7 @@ def main() -> int:
# ── embed everything at native resolution ──────────────────────────────── # ── embed everything at native resolution ────────────────────────────────
t0 = time.time() 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", print(f"[embed] {len(crops)} native crops in {time.time() - t0:.1f}s",
file=sys.stderr) file=sys.stderr)
@@ -815,7 +621,7 @@ def main() -> int:
for size in sizes: for size in sizes:
t0 = time.time() t0 = time.time()
degraded = [degrade(c, size, down, up) for c in probe_crops] 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] sims = q @ gal_emb.T # [n_probe, n_gal]
best_per_actor = np.stack([sims[:, cols].max(axis=1) for cols in actor_cols], 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)), "caveat": CAVEAT.format(n_actors=len(actors)),
"model": arcface.stem, "model": arcface.stem,
"detector": detector.stem, "detector": detector.stem,
"backend": f"onnxruntime {'/'.join(emb_model.sess.get_providers())} (fp32 ONNX; " "backend": "sae_embed / the compiled-in inference backend (fp32 ONNX under "
f"a TensorRT fp16 build is a different embedding space)", "SAE_INFERENCE_BACKEND=ORT; a TensorRT fp16 build is a different "
"embedding space)",
"n_actors": len(actors), "n_actors": len(actors),
"n_probes": len(probe_rows), "n_probes": len(probe_rows),
"n_gallery_embeddings": len(gal_rows), "n_gallery_embeddings": len(gal_rows),