diff --git a/scripts/validation/min_face_size.py b/scripts/validation/min_face_size.py index 5e62437..8c17538 100644 --- a/scripts/validation/min_face_size.py +++ b/scripts/validation/min_face_size.py @@ -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 /LVFace-B_Glint360K.onnx)") p.add_argument("--detector", default=None, help="SCRFD ONNX (default /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),