dtourolle
f891e579c5
The parser reads a tag up to end of line, so `# TRACES: GR-004 | SR-001 — prose` swallowed the prose into the tag and the row went unmatched. Splitting the comment leaves the tag greppable by the same pattern as the code tags and the commit trailers, which is the point of the house format. Mechanical throughout; no logic touched. The regenerated report reflects this session's new tags: 137 -> 148 found, and one more tagged-but-unexecuted, which is the SuperHero accuracy assertion that is documented but not yet a test.
73 lines
3.1 KiB
Python
73 lines
3.1 KiB
Python
"""Shared loader for the sae_embed nanobind module (SCRFD + ArcFace).
|
|
|
|
sae_embed.FaceEmbedder loads both ONNX sessions once and exposes an
|
|
embed(path) -> FaceResult method, avoiding the per-process model reload cost
|
|
of spawning the embed_faces CLI binary for every image.
|
|
|
|
resolve_arcface() exposes the same default-resolution logic load_embedder uses,
|
|
so a caller can stamp the gallery it is about to write with the model that
|
|
actually produced its embeddings (GR-004) — the resolved path, not the CLI
|
|
argument, which is often None.
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
from pathlib import Path
|
|
|
|
DEFAULT_ARCFACE = "arcface_w600k_r50.onnx"
|
|
|
|
|
|
def resolve_arcface(models_dir: str, arcface: str | None = None) -> str:
|
|
"""The ArcFace/LVFace ONNX path load_embedder would use for these arguments.
|
|
|
|
TRACES: GR-004 | SR-001
|
|
|
|
Single source of truth for "which model is this", so the stamp written into
|
|
a gallery can never drift from the model loaded."""
|
|
return arcface if arcface else str(Path(models_dir) / DEFAULT_ARCFACE)
|
|
|
|
|
|
def load_embedder(build_dir: str, models_dir: str, arcface: str | None = None,
|
|
conf: float = 0.5, nms: float = 0.4, max_side: int = 500):
|
|
"""Import sae_embed from build_dir and construct a FaceEmbedder.
|
|
|
|
Exits with a clear error if the module or models are missing — there is
|
|
no subprocess fallback.
|
|
"""
|
|
build_path = Path(build_dir).resolve()
|
|
sys.path.insert(0, str(build_path))
|
|
try:
|
|
import sae_embed
|
|
except ImportError as e:
|
|
sys.exit(
|
|
f"sae_embed module not found in {build_path}: {e}\n"
|
|
f"Build it first: cmake --build {build_dir} --target sae_embed"
|
|
)
|
|
|
|
models_path = Path(models_dir)
|
|
detector_path = str(models_path / "scrfd_500m_bnkps.onnx")
|
|
arcface_path = resolve_arcface(models_dir, arcface)
|
|
for model, name in [(detector_path, "SCRFD"), (arcface_path, "ArcFace")]:
|
|
if not Path(model).is_file():
|
|
sys.exit(f"{name} model not found: {model}\nRun: bash scripts/download_models.sh")
|
|
|
|
# A TRT-backend build cannot load .onnx; it needs pre-built engines from
|
|
# scripts/build_trt_engines.sh.
|
|
#
|
|
# These are passed only on request. The old comment here claimed they were
|
|
# "ignored by ORT" — they are not. The ORT backend treats an engine path as
|
|
# an instruction and raises, which is the right behaviour (silently ignoring
|
|
# a requested engine would be worse), but it meant that merely HAVING a
|
|
# populated trt_cache/ broke every ORT gallery build in the repo, with an
|
|
# error naming a flag the caller never set.
|
|
use_engines = os.environ.get("SAE_USE_TRT_ENGINES", "") not in ("", "0", "false")
|
|
trt = Path(models_path).parent / "trt_cache"
|
|
det_engine = trt / "scrfd.scrfd_500m_bnkps.640.fp16.engine"
|
|
arc_engine = trt / f"arcface.{Path(arcface_path).stem}.b4.fp16.engine"
|
|
|
|
return sae_embed.FaceEmbedder(
|
|
detector_path, arcface_path, conf, nms, max_side,
|
|
str(det_engine) if (use_engines and det_engine.is_file()) else "",
|
|
str(arc_engine) if (use_engines and arc_engine.is_file()) else "",
|
|
)
|