feat(faces): real ONNX face analyzer (SCRFD + ArcFace), opt-in
Implements the last Phase 2 piece: a working face detector/embedder behind
the new `faces-onnx` cargo feature (mirrors how `plugins` gates wasmtime).
Inert by default — the default build is unchanged and ships the no-op
analyzer.
Pipeline (InsightFace/immich pattern): SCRFD detection with 5-point
landmarks → least-squares similarity alignment to the canonical 112×112
template → ArcFace embedding → L2-normalized 512-d vector.
- face_geometry.rs (always compiled, unit-tested): SCRFD anchor/distance
decode, NMS, the closed-form (complex-number) similarity transform,
bilinear affine warp, NCHW normalization, L2-norm, Laplacian sharpness.
11 unit tests cover the error-prone math with no model needed.
- onnx_face_analyzer.rs (feature `faces-onnx`): wires the geometry to ONNX
Runtime via `ort` (load-dynamic, so libonnxruntime is dlopen'd at runtime
and the crate builds without it). Inference runs on spawn_blocking; each
session is serialized behind a Mutex. Loads via `ort::init_from` (fallible)
not ORT's lazy loader, which would panic under `panic = "abort"`.
- config: FacesConfig + OXICLOUD_FACES_{ORT_DYLIB,DETECTOR_MODEL,
EMBEDDER_MODEL,DET_SIZE,DET_THRESHOLD,NMS_THRESHOLD,INTRA_THREADS}.
- di: build_face_analyzer() loads the real analyzer when the feature is
compiled in and runtime+models are configured; any missing piece or load
failure degrades to the no-op analyzer (logged) so startup never fails.
- ort/ndarray added as optional deps; example.env documents the setup.
Models and the ONNX Runtime dylib are operator-provided at runtime and are
never committed. Cannot be exercised in CI (no models/dylib); the geometry
is unit-tested and the ONNX seam is isolated.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01JW6ghFMDtnRYuYNzZhb47M
This commit is contained in:
Claude
Generated
+80
@@ -3654,6 +3654,16 @@ version = "0.2.186"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
|
||||
|
||||
[[package]]
|
||||
name = "libloading"
|
||||
version = "0.9.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "754ca22de805bb5744484a5b151a9e1a8e837d5dc232c2d7d8c2e3492edc8b60"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "libm"
|
||||
version = "0.2.16"
|
||||
@@ -3866,6 +3876,16 @@ version = "0.8.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3"
|
||||
|
||||
[[package]]
|
||||
name = "matrixmultiply"
|
||||
version = "0.3.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08"
|
||||
dependencies = [
|
||||
"autocfg",
|
||||
"rawpointer",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "maybe-owned"
|
||||
version = "0.3.4"
|
||||
@@ -4071,6 +4091,21 @@ version = "0.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "13d2233c9842d08cfe13f9eac96e207ca6a2ea10b80259ebe8ad0268be27d2af"
|
||||
|
||||
[[package]]
|
||||
name = "ndarray"
|
||||
version = "0.17.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d"
|
||||
dependencies = [
|
||||
"matrixmultiply",
|
||||
"num-complex",
|
||||
"num-integer",
|
||||
"num-traits",
|
||||
"portable-atomic",
|
||||
"portable-atomic-util",
|
||||
"rawpointer",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nom"
|
||||
version = "7.1.3"
|
||||
@@ -4157,6 +4192,15 @@ dependencies = [
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-complex"
|
||||
version = "0.4.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
|
||||
dependencies = [
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-conv"
|
||||
version = "0.2.1"
|
||||
@@ -4238,6 +4282,25 @@ dependencies = [
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ort"
|
||||
version = "2.0.0-rc.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d7de3af33d24a745ffb8fab904b13478438d1cd52868e6f17735ef6e1f8bf133"
|
||||
dependencies = [
|
||||
"libloading",
|
||||
"ndarray",
|
||||
"ort-sys",
|
||||
"smallvec",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ort-sys"
|
||||
version = "2.0.0-rc.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d7b497d21a8b6fbb4b5a544f8fadb77e801a09ae0add9e411d31c6f89e3c1e90"
|
||||
|
||||
[[package]]
|
||||
name = "outref"
|
||||
version = "0.1.0"
|
||||
@@ -4676,7 +4739,9 @@ dependencies = [
|
||||
"mockall",
|
||||
"moka",
|
||||
"mp3-duration",
|
||||
"ndarray",
|
||||
"nom-exif",
|
||||
"ort",
|
||||
"oxc_allocator",
|
||||
"oxc_codegen",
|
||||
"oxc_minifier",
|
||||
@@ -5074,6 +5139,15 @@ version = "1.13.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49"
|
||||
|
||||
[[package]]
|
||||
name = "portable-atomic-util"
|
||||
version = "0.2.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618"
|
||||
dependencies = [
|
||||
"portable-atomic",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "postcard"
|
||||
version = "1.1.3"
|
||||
@@ -5496,6 +5570,12 @@ version = "1.7.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "973443cf09a9c8656b574a866ab68dfa19f0867d0340648c7d2f6a71b8a8ea68"
|
||||
|
||||
[[package]]
|
||||
name = "rawpointer"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
|
||||
|
||||
[[package]]
|
||||
name = "rayon"
|
||||
version = "1.12.0"
|
||||
|
||||
@@ -81,6 +81,8 @@ nom-exif = "3.6.1"
|
||||
extism = { version = "1.30.0", optional = true }
|
||||
toml = { version = "1.1.2", optional = true }
|
||||
file-rotate = { version = "0.7.6", optional = true }
|
||||
ort = { version = "2.0.0-rc.12", default-features = false, features = ["load-dynamic", "ndarray", "tracing", "api-24"], optional = true }
|
||||
ndarray = { version = "0.17.2", optional = true }
|
||||
|
||||
[features]
|
||||
default = []
|
||||
@@ -95,6 +97,12 @@ plugins = ["dep:extism", "dep:toml", "dep:file-rotate"]
|
||||
# this lets one `cargo build` produce both `oxicloud` and `load-seed`
|
||||
# without recompiling oxicloud with mockall in scope.
|
||||
load_seed_bin = []
|
||||
# Real ONNX-backed face analyzer (detector + embedder) for the People feature.
|
||||
# Opt-in: pulls `ort` (ONNX Runtime, load-dynamic — dlopen's libonnxruntime at
|
||||
# runtime) + `ndarray`, a heavy stack most deployments won't use. Activation also
|
||||
# requires OXICLOUD_ENABLE_FACES=true *and* operator-provided ONNX models; without
|
||||
# this feature the People pipeline falls back to the inert NoopFaceAnalyzer.
|
||||
faces-onnx = ["dep:ort", "dep:ndarray"]
|
||||
|
||||
[lints.rust]
|
||||
unexpected_cfgs = { level = "warn", check-cfg = ['cfg(integration_tests)'] }
|
||||
|
||||
+37
@@ -224,6 +224,43 @@ DATABASE_URL=postgres://postgres:postgres@localhost:5432/oxicloud
|
||||
# Set to false to prevent users from browsing the user directory.
|
||||
#OXICLOUD_EXPOSE_SYSTEM_USERS=true
|
||||
|
||||
# ── People (face recognition) ────────────────────────────────────────────
|
||||
# Biometric data (GDPR Art. 9) — OFF by default, opt-in per deployment.
|
||||
# Detects faces and clusters them into people in the photo library.
|
||||
#
|
||||
# Requires ALL of:
|
||||
# 1. a binary built with the `faces-onnx` cargo feature
|
||||
# (`cargo build --release --features faces-onnx`),
|
||||
# 2. OXICLOUD_ENABLE_FACES=true,
|
||||
# 3. the ONNX Runtime shared library + two operator-provided ONNX models
|
||||
# (a SCRFD/RetinaFace detector with 5-point landmarks, and an ArcFace
|
||||
# 512-d embedder — e.g. InsightFace `buffalo_l`). Models are NOT shipped.
|
||||
# Without all three, the People pipeline stays inert (no-op analyzer) and the
|
||||
# server still boots; the People tab stays hidden in the UI.
|
||||
#OXICLOUD_ENABLE_FACES=false
|
||||
|
||||
# Path to libonnxruntime.{so,dylib,dll}. Falls back to ORT_DYLIB_PATH.
|
||||
# Use the ONNX Runtime build matching this app's `ort` crate (>= 1.24).
|
||||
#OXICLOUD_FACES_ORT_DYLIB=/opt/onnxruntime/lib/libonnxruntime.so
|
||||
|
||||
# Face detector model (SCRFD/RetinaFace, 5-point landmarks).
|
||||
#OXICLOUD_FACES_DETECTOR_MODEL=/var/lib/oxicloud/models/scrfd_10g_bnkps.onnx
|
||||
|
||||
# Face embedder model (ArcFace, 112x112 input -> 512-d output).
|
||||
#OXICLOUD_FACES_EMBEDDER_MODEL=/var/lib/oxicloud/models/w600k_r50.onnx
|
||||
|
||||
# Detector square input size in px (default: 640)
|
||||
#OXICLOUD_FACES_DET_SIZE=640
|
||||
|
||||
# Minimum detector confidence to keep a face, 0..1 (default: 0.5)
|
||||
#OXICLOUD_FACES_DET_THRESHOLD=0.5
|
||||
|
||||
# IoU threshold for non-maximum suppression, 0..1 (default: 0.4)
|
||||
#OXICLOUD_FACES_NMS_THRESHOLD=0.4
|
||||
|
||||
# ONNX Runtime intra-op threads; 0 = let ONNX Runtime decide (default: 0)
|
||||
#OXICLOUD_FACES_INTRA_THREADS=0
|
||||
|
||||
# WASM plugin runtime (Extism). Requires a binary built with the `plugins`
|
||||
# cargo feature (`cargo run --features plugins`); without that feature these
|
||||
# vars are inert. Untrusted plugins run sandboxed: no filesystem, no network,
|
||||
|
||||
@@ -905,6 +905,53 @@ impl Default for FeaturesConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// Face-recognition (People) model configuration.
|
||||
///
|
||||
/// Only consulted when the `faces-onnx` cargo feature is compiled in *and*
|
||||
/// [`FeaturesConfig::enable_faces`] is true; otherwise the inert
|
||||
/// `NoopFaceAnalyzer` is used regardless of these values. The ONNX Runtime
|
||||
/// dylib and both model files are operator-provided at runtime (never
|
||||
/// committed) — when any is unset or fails to load, the People pipeline
|
||||
/// silently falls back to the no-op analyzer and the server still boots.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct FacesConfig {
|
||||
/// `libonnxruntime.{so,dylib,dll}`. Falls back to the `ORT_DYLIB_PATH`
|
||||
/// environment variable when unset. Env: `OXICLOUD_FACES_ORT_DYLIB`.
|
||||
pub ort_dylib: Option<PathBuf>,
|
||||
/// SCRFD/RetinaFace detector model with 5-point landmarks.
|
||||
/// Env: `OXICLOUD_FACES_DETECTOR_MODEL`.
|
||||
pub detector_model: Option<PathBuf>,
|
||||
/// ArcFace embedder model (112×112 → 512-d).
|
||||
/// Env: `OXICLOUD_FACES_EMBEDDER_MODEL`.
|
||||
pub embedder_model: Option<PathBuf>,
|
||||
/// Detector square input size in pixels (default 640).
|
||||
/// Env: `OXICLOUD_FACES_DET_SIZE`.
|
||||
pub det_size: u32,
|
||||
/// Minimum detector confidence to keep a face (default 0.5).
|
||||
/// Env: `OXICLOUD_FACES_DET_THRESHOLD`.
|
||||
pub det_threshold: f32,
|
||||
/// IoU threshold for non-max suppression (default 0.4).
|
||||
/// Env: `OXICLOUD_FACES_NMS_THRESHOLD`.
|
||||
pub nms_threshold: f32,
|
||||
/// ONNX Runtime intra-op threads (0 = let ORT decide).
|
||||
/// Env: `OXICLOUD_FACES_INTRA_THREADS`.
|
||||
pub intra_threads: usize,
|
||||
}
|
||||
|
||||
impl Default for FacesConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
ort_dylib: None,
|
||||
detector_model: None,
|
||||
embedder_model: None,
|
||||
det_size: 640,
|
||||
det_threshold: 0.5,
|
||||
nms_threshold: 0.4,
|
||||
intra_threads: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Content-search configuration (embedded Tantivy index over file names and
|
||||
/// extracted file content).
|
||||
///
|
||||
@@ -1070,6 +1117,8 @@ pub struct AppConfig {
|
||||
pub content_search: ContentSearchConfig,
|
||||
/// WASM plugin runtime configuration
|
||||
pub plugins: PluginConfig,
|
||||
/// Face-recognition (People) model configuration
|
||||
pub faces: FacesConfig,
|
||||
}
|
||||
|
||||
/// Server-side i18n knobs.
|
||||
@@ -1123,6 +1172,7 @@ impl Default for AppConfig {
|
||||
i18n: I18nConfig::default(),
|
||||
content_search: ContentSearchConfig::default(),
|
||||
plugins: PluginConfig::default(),
|
||||
faces: FacesConfig::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1397,6 +1447,43 @@ impl AppConfig {
|
||||
config.features.enable_faces = val;
|
||||
}
|
||||
|
||||
// Faces (People) ONNX runtime + models — operator-provided at runtime.
|
||||
if let Ok(v) = env::var("OXICLOUD_FACES_ORT_DYLIB").or_else(|_| env::var("ORT_DYLIB_PATH"))
|
||||
&& !v.is_empty()
|
||||
{
|
||||
config.faces.ort_dylib = Some(PathBuf::from(v));
|
||||
}
|
||||
if let Ok(v) = env::var("OXICLOUD_FACES_DETECTOR_MODEL")
|
||||
&& !v.is_empty()
|
||||
{
|
||||
config.faces.detector_model = Some(PathBuf::from(v));
|
||||
}
|
||||
if let Ok(v) = env::var("OXICLOUD_FACES_EMBEDDER_MODEL")
|
||||
&& !v.is_empty()
|
||||
{
|
||||
config.faces.embedder_model = Some(PathBuf::from(v));
|
||||
}
|
||||
if let Ok(v) = env::var("OXICLOUD_FACES_DET_SIZE").map(|v| v.parse::<u32>())
|
||||
&& let Ok(val) = v
|
||||
{
|
||||
config.faces.det_size = val;
|
||||
}
|
||||
if let Ok(v) = env::var("OXICLOUD_FACES_DET_THRESHOLD").map(|v| v.parse::<f32>())
|
||||
&& let Ok(val) = v
|
||||
{
|
||||
config.faces.det_threshold = val;
|
||||
}
|
||||
if let Ok(v) = env::var("OXICLOUD_FACES_NMS_THRESHOLD").map(|v| v.parse::<f32>())
|
||||
&& let Ok(val) = v
|
||||
{
|
||||
config.faces.nms_threshold = val;
|
||||
}
|
||||
if let Ok(v) = env::var("OXICLOUD_FACES_INTRA_THREADS").map(|v| v.parse::<usize>())
|
||||
&& let Ok(val) = v
|
||||
{
|
||||
config.faces.intra_threads = val;
|
||||
}
|
||||
|
||||
// Content search (embedded Tantivy index)
|
||||
if let Ok(v) = env::var("OXICLOUD_ENABLE_CONTENT_SEARCH").map(|v| v.parse::<bool>())
|
||||
&& let Ok(val) = v
|
||||
|
||||
+54
-4
@@ -814,15 +814,16 @@ impl AppServiceFactory {
|
||||
service
|
||||
}
|
||||
|
||||
/// Creates the face-indexing lifecycle hook (People feature). Uses the
|
||||
/// default no-op analyzer until the operator wires a real ONNX model.
|
||||
/// Creates the face-indexing lifecycle hook (People feature). Picks the
|
||||
/// real ONNX analyzer when the `faces-onnx` feature is compiled in and the
|
||||
/// operator has configured the runtime + models; otherwise the inert no-op
|
||||
/// analyzer (see [`Self::build_face_analyzer`]).
|
||||
pub fn create_face_indexing_service(
|
||||
&self,
|
||||
db_pool: &Arc<PgPool>,
|
||||
) -> Arc<crate::infrastructure::services::face_indexing_service::FaceIndexingService> {
|
||||
let blob_root = self.storage_path.join(".blobs");
|
||||
let analyzer: Arc<dyn crate::application::ports::face_ports::FaceAnalyzerPort> =
|
||||
Arc::new(crate::infrastructure::services::noop_face_analyzer::NoopFaceAnalyzer);
|
||||
let analyzer = self.build_face_analyzer();
|
||||
Arc::new(
|
||||
crate::infrastructure::services::face_indexing_service::FaceIndexingService::new(
|
||||
db_pool.clone(),
|
||||
@@ -832,6 +833,55 @@ impl AppServiceFactory {
|
||||
)
|
||||
}
|
||||
|
||||
/// Selects the face analyzer. With the `faces-onnx` feature and a fully
|
||||
/// configured runtime + models, loads the real ONNX analyzer; any missing
|
||||
/// piece or load failure degrades gracefully to the no-op analyzer (logged)
|
||||
/// so startup never fails on biometric configuration.
|
||||
fn build_face_analyzer(
|
||||
&self,
|
||||
) -> Arc<dyn crate::application::ports::face_ports::FaceAnalyzerPort> {
|
||||
#[cfg(feature = "faces-onnx")]
|
||||
{
|
||||
let f = &self.config.faces;
|
||||
if let (Some(dylib), Some(detector), Some(embedder)) = (
|
||||
f.ort_dylib.as_ref(),
|
||||
f.detector_model.as_ref(),
|
||||
f.embedder_model.as_ref(),
|
||||
) {
|
||||
use crate::infrastructure::services::onnx_face_analyzer::{
|
||||
OnnxFaceAnalyzer, OnnxLoadConfig,
|
||||
};
|
||||
let cfg = OnnxLoadConfig {
|
||||
dylib,
|
||||
detector,
|
||||
embedder,
|
||||
det_size: f.det_size,
|
||||
det_threshold: f.det_threshold,
|
||||
nms_threshold: f.nms_threshold,
|
||||
intra_threads: f.intra_threads,
|
||||
};
|
||||
match OnnxFaceAnalyzer::load(&cfg) {
|
||||
Ok(analyzer) => {
|
||||
tracing::info!("Face analyzer: ONNX models loaded");
|
||||
return Arc::new(analyzer);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Face analyzer: failed to load ONNX models ({e}); \
|
||||
falling back to no-op analyzer"
|
||||
);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
tracing::info!(
|
||||
"Face analyzer: faces-onnx compiled but runtime/models not fully \
|
||||
configured; using no-op analyzer"
|
||||
);
|
||||
}
|
||||
}
|
||||
Arc::new(crate::infrastructure::services::noop_face_analyzer::NoopFaceAnalyzer)
|
||||
}
|
||||
|
||||
/// Creates the People (faces) read/clustering service.
|
||||
pub fn create_people_service(&self, db_pool: &Arc<PgPool>) -> Arc<PeopleService> {
|
||||
let repo = Arc::new(
|
||||
|
||||
@@ -0,0 +1,473 @@
|
||||
//! Pure geometry + post-processing for the ONNX face pipeline.
|
||||
//!
|
||||
//! Everything here is plain Rust (no `ort`, no `ndarray`) so it compiles in the
|
||||
//! default build and is exercised by `cargo test` — the error-prone numerical
|
||||
//! parts (SCRFD anchor decode, NMS, 5-point similarity alignment, the affine
|
||||
//! warp, normalization) are unit-tested in isolation, while the untestable ONNX
|
||||
//! session calls live behind the `faces-onnx` feature in `onnx_face_analyzer`.
|
||||
//!
|
||||
//! The pipeline mirrors InsightFace's reference implementation:
|
||||
//! SCRFD detector (distance-to-box anchors over strides 8/16/32) → 5-point
|
||||
//! similarity transform onto the canonical 112×112 ArcFace template → ArcFace
|
||||
//! embedder → L2-normalized 512-d vector.
|
||||
|
||||
use image::RgbImage;
|
||||
|
||||
/// One detected face in **detector-input pixel** coordinates (before scaling
|
||||
/// back to the original image): an axis-aligned box `[x1, y1, x2, y2]`, the
|
||||
/// five facial landmarks, and the detector confidence.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct Detection {
|
||||
pub bbox: [f32; 4],
|
||||
pub kps: [[f32; 2]; 5],
|
||||
pub score: f32,
|
||||
}
|
||||
|
||||
/// A 2×3 affine transform mapping an output/template coordinate to a source
|
||||
/// coordinate: `src = (a·ox + b·oy + tx, c·ox + d·oy + ty)`. Used to sample the
|
||||
/// source image when warping an aligned face crop.
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub struct Affine {
|
||||
pub a: f32,
|
||||
pub b: f32,
|
||||
pub c: f32,
|
||||
pub d: f32,
|
||||
pub tx: f32,
|
||||
pub ty: f32,
|
||||
}
|
||||
|
||||
/// Canonical ArcFace 5-point template for a 112×112 crop
|
||||
/// (left eye, right eye, nose, left mouth, right mouth).
|
||||
pub const ARCFACE_TEMPLATE: [[f32; 2]; 5] = [
|
||||
[38.2946, 51.6963],
|
||||
[73.5318, 51.5014],
|
||||
[56.0252, 71.7366],
|
||||
[41.5493, 92.3655],
|
||||
[70.7299, 92.2041],
|
||||
];
|
||||
|
||||
/// Aligned-crop side length expected by the ArcFace embedder.
|
||||
pub const ALIGN_SIZE: u32 = 112;
|
||||
|
||||
/// Letterbox geometry for the detector: the largest scale that fits a
|
||||
/// `w0 × h0` image into a `det × det` square without distortion, plus the
|
||||
/// resulting (possibly smaller) dimensions placed at the top-left.
|
||||
///
|
||||
/// Returns `(new_w, new_h, scale)` where `scale = min(det/w0, det/h0)` and
|
||||
/// detector-space coordinates map back to the original by dividing by `scale`.
|
||||
pub fn letterbox(w0: u32, h0: u32, det: u32) -> (u32, u32, f32) {
|
||||
if w0 == 0 || h0 == 0 {
|
||||
return (0, 0, 1.0);
|
||||
}
|
||||
let scale = (det as f32 / w0 as f32).min(det as f32 / h0 as f32);
|
||||
let new_w = ((w0 as f32 * scale).round() as u32).clamp(1, det);
|
||||
let new_h = ((h0 as f32 * scale).round() as u32).clamp(1, det);
|
||||
(new_w, new_h, scale)
|
||||
}
|
||||
|
||||
/// `NCHW`, RGB, float input tensor for an ONNX model: `(px − mean) · scale`,
|
||||
/// channel-major (all R, then all G, then all B). Length is `3 · w · h`.
|
||||
pub fn chw_normalized(img: &RgbImage, mean: f32, scale: f32) -> Vec<f32> {
|
||||
let (w, h) = (img.width() as usize, img.height() as usize);
|
||||
let mut out = vec![0.0f32; 3 * w * h];
|
||||
let plane = w * h;
|
||||
for (i, px) in img.pixels().enumerate() {
|
||||
out[i] = (px[0] as f32 - mean) * scale;
|
||||
out[plane + i] = (px[1] as f32 - mean) * scale;
|
||||
out[2 * plane + i] = (px[2] as f32 - mean) * scale;
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Decode one SCRFD feature-map stride into detections, appending those above
|
||||
/// `threshold` to `out`. All coordinates are in detector-input pixels.
|
||||
///
|
||||
/// `scores` is `[n]`, `bbox` is `[n·4]` (left, top, right, bottom *distances*,
|
||||
/// already multiplied by `stride`), `kps` (when present) is `[n·10]`
|
||||
/// (5 × (dx, dy) distances, already multiplied by `stride`), where
|
||||
/// `n = feat_h · feat_w · num_anchors`. Anchor centers follow InsightFace's
|
||||
/// row-major `mgrid` order with `num_anchors` consecutive duplicates.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn decode_stride(
|
||||
scores: &[f32],
|
||||
bbox: &[f32],
|
||||
kps: Option<&[f32]>,
|
||||
stride: u32,
|
||||
feat_h: u32,
|
||||
feat_w: u32,
|
||||
num_anchors: u32,
|
||||
threshold: f32,
|
||||
out: &mut Vec<Detection>,
|
||||
) {
|
||||
let stride_f = stride as f32;
|
||||
let mut idx = 0usize;
|
||||
for y in 0..feat_h {
|
||||
for x in 0..feat_w {
|
||||
let cx = x as f32 * stride_f;
|
||||
let cy = y as f32 * stride_f;
|
||||
for _ in 0..num_anchors {
|
||||
if idx >= scores.len() {
|
||||
return;
|
||||
}
|
||||
let score = scores[idx];
|
||||
if score >= threshold {
|
||||
let b = idx * 4;
|
||||
if b + 3 < bbox.len() {
|
||||
let det_bbox = [
|
||||
cx - bbox[b],
|
||||
cy - bbox[b + 1],
|
||||
cx + bbox[b + 2],
|
||||
cy + bbox[b + 3],
|
||||
];
|
||||
let mut det_kps = [[0.0f32; 2]; 5];
|
||||
if let Some(kps) = kps {
|
||||
let k = idx * 10;
|
||||
if k + 9 < kps.len() {
|
||||
for (p, slot) in det_kps.iter_mut().enumerate() {
|
||||
*slot = [cx + kps[k + p * 2], cy + kps[k + p * 2 + 1]];
|
||||
}
|
||||
}
|
||||
}
|
||||
out.push(Detection {
|
||||
bbox: det_bbox,
|
||||
kps: det_kps,
|
||||
score,
|
||||
});
|
||||
}
|
||||
}
|
||||
idx += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Intersection-over-union of two `[x1, y1, x2, y2]` boxes.
|
||||
pub fn iou(a: &[f32; 4], b: &[f32; 4]) -> f32 {
|
||||
let x1 = a[0].max(b[0]);
|
||||
let y1 = a[1].max(b[1]);
|
||||
let x2 = a[2].min(b[2]);
|
||||
let y2 = a[3].min(b[3]);
|
||||
let iw = (x2 - x1).max(0.0);
|
||||
let ih = (y2 - y1).max(0.0);
|
||||
let inter = iw * ih;
|
||||
let area_a = (a[2] - a[0]).max(0.0) * (a[3] - a[1]).max(0.0);
|
||||
let area_b = (b[2] - b[0]).max(0.0) * (b[3] - b[1]).max(0.0);
|
||||
let union = area_a + area_b - inter;
|
||||
if union <= 0.0 { 0.0 } else { inter / union }
|
||||
}
|
||||
|
||||
/// Greedy non-maximum suppression: keep highest-scoring boxes, drop any whose
|
||||
/// IoU with an already-kept box exceeds `iou_thresh`. Returns the kept
|
||||
/// detections, highest score first.
|
||||
pub fn nms(mut dets: Vec<Detection>, iou_thresh: f32) -> Vec<Detection> {
|
||||
dets.sort_by(|a, b| b.score.total_cmp(&a.score));
|
||||
let mut keep: Vec<Detection> = Vec::with_capacity(dets.len());
|
||||
for d in dets {
|
||||
if keep.iter().all(|k| iou(&k.bbox, &d.bbox) <= iou_thresh) {
|
||||
keep.push(d);
|
||||
}
|
||||
}
|
||||
keep
|
||||
}
|
||||
|
||||
/// Least-squares similarity transform (scale + rotation + translation, no
|
||||
/// shear, no reflection) mapping `src` landmarks onto `dst`, returned as its
|
||||
/// **inverse** affine (output/template coordinate → source coordinate) ready
|
||||
/// for backward-warp sampling.
|
||||
///
|
||||
/// Solved in closed form via the complex-number formulation: with points as
|
||||
/// complex numbers, `w = Σ (b'ᵢ · conj(a'ᵢ)) / Σ |a'ᵢ|²` and `t = mean_b −
|
||||
/// w·mean_a`, which is equivalent to the Umeyama solution InsightFace obtains
|
||||
/// from `skimage.SimilarityTransform`.
|
||||
pub fn similarity_transform_inverse(src: &[[f32; 2]; 5], dst: &[[f32; 2]; 5]) -> Affine {
|
||||
let n = 5.0f32;
|
||||
let (mut max, mut may, mut mbx, mut mby) = (0.0f32, 0.0f32, 0.0f32, 0.0f32);
|
||||
for i in 0..5 {
|
||||
max += src[i][0];
|
||||
may += src[i][1];
|
||||
mbx += dst[i][0];
|
||||
mby += dst[i][1];
|
||||
}
|
||||
max /= n;
|
||||
may /= n;
|
||||
mbx /= n;
|
||||
mby /= n;
|
||||
|
||||
// num = Σ b'·conj(a') (complex), den = Σ |a'|² (real)
|
||||
let (mut num_re, mut num_im, mut den) = (0.0f32, 0.0f32, 0.0f32);
|
||||
for i in 0..5 {
|
||||
let ax = src[i][0] - max;
|
||||
let ay = src[i][1] - may;
|
||||
let bx = dst[i][0] - mbx;
|
||||
let by = dst[i][1] - mby;
|
||||
// b' · conj(a') = (bx + i·by)(ax − i·ay)
|
||||
num_re += bx * ax + by * ay;
|
||||
num_im += by * ax - bx * ay;
|
||||
den += ax * ax + ay * ay;
|
||||
}
|
||||
let den = if den.abs() < 1e-12 { 1e-12 } else { den };
|
||||
// w = num/den (forward scale·rotation)
|
||||
let wr = num_re / den;
|
||||
let wi = num_im / den;
|
||||
// t = mean_b − w·mean_a
|
||||
let tr = mbx - (wr * max - wi * may);
|
||||
let ti = mby - (wi * max + wr * may);
|
||||
|
||||
// Inverse of the similarity: src = Ainv·(out − t), Ainv = [[wr,wi],[−wi,wr]]/|w|²
|
||||
let det = wr * wr + wi * wi;
|
||||
let g = if det.abs() < 1e-12 { 0.0 } else { 1.0 / det };
|
||||
Affine {
|
||||
a: g * wr,
|
||||
b: g * wi,
|
||||
c: -g * wi,
|
||||
d: g * wr,
|
||||
tx: -g * (wr * tr + wi * ti),
|
||||
ty: g * (wi * tr - wr * ti),
|
||||
}
|
||||
}
|
||||
|
||||
/// Warp `img` into an `ALIGN_SIZE × ALIGN_SIZE` aligned face crop using the
|
||||
/// inverse affine from [`similarity_transform_inverse`], sampling bilinearly
|
||||
/// and clamping to the image edge.
|
||||
pub fn warp_to_aligned(img: &RgbImage, inv: &Affine) -> RgbImage {
|
||||
let (w, h) = (img.width(), img.height());
|
||||
let mut out = RgbImage::new(ALIGN_SIZE, ALIGN_SIZE);
|
||||
for oy in 0..ALIGN_SIZE {
|
||||
for ox in 0..ALIGN_SIZE {
|
||||
let sx = inv.a * ox as f32 + inv.b * oy as f32 + inv.tx;
|
||||
let sy = inv.c * ox as f32 + inv.d * oy as f32 + inv.ty;
|
||||
let px = bilinear_sample(img, sx, sy, w, h);
|
||||
out.put_pixel(ox, oy, px);
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Bilinear RGB sample at floating `(x, y)`, clamping out-of-bounds reads to
|
||||
/// the nearest edge.
|
||||
fn bilinear_sample(img: &RgbImage, x: f32, y: f32, w: u32, h: u32) -> image::Rgb<u8> {
|
||||
let x = x.clamp(0.0, (w - 1) as f32);
|
||||
let y = y.clamp(0.0, (h - 1) as f32);
|
||||
let x0 = x.floor() as u32;
|
||||
let y0 = y.floor() as u32;
|
||||
let x1 = (x0 + 1).min(w - 1);
|
||||
let y1 = (y0 + 1).min(h - 1);
|
||||
let dx = x - x0 as f32;
|
||||
let dy = y - y0 as f32;
|
||||
let p00 = img.get_pixel(x0, y0);
|
||||
let p10 = img.get_pixel(x1, y0);
|
||||
let p01 = img.get_pixel(x0, y1);
|
||||
let p11 = img.get_pixel(x1, y1);
|
||||
let mut out = [0u8; 3];
|
||||
for (ch, slot) in out.iter_mut().enumerate() {
|
||||
let top = p00[ch] as f32 * (1.0 - dx) + p10[ch] as f32 * dx;
|
||||
let bot = p01[ch] as f32 * (1.0 - dx) + p11[ch] as f32 * dx;
|
||||
*slot = (top * (1.0 - dy) + bot * dy).round().clamp(0.0, 255.0) as u8;
|
||||
}
|
||||
image::Rgb(out)
|
||||
}
|
||||
|
||||
/// In-place L2 normalization. A zero vector is left unchanged.
|
||||
pub fn l2_normalize(v: &mut [f32]) {
|
||||
let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
if norm > 1e-12 {
|
||||
for x in v.iter_mut() {
|
||||
*x /= norm;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Variance of the discrete Laplacian over the luminance of an RGB crop — a
|
||||
/// cheap focus/sharpness proxy (higher = sharper). Used as a face quality
|
||||
/// score for cover selection and gating.
|
||||
pub fn laplacian_variance(img: &RgbImage) -> f32 {
|
||||
let (w, h) = (img.width() as i64, img.height() as i64);
|
||||
if w < 3 || h < 3 {
|
||||
return 0.0;
|
||||
}
|
||||
let lum = |x: i64, y: i64| -> f32 {
|
||||
let p = img.get_pixel(x as u32, y as u32);
|
||||
0.299 * p[0] as f32 + 0.587 * p[1] as f32 + 0.114 * p[2] as f32
|
||||
};
|
||||
let mut vals = Vec::with_capacity(((w - 2) * (h - 2)) as usize);
|
||||
for y in 1..h - 1 {
|
||||
for x in 1..w - 1 {
|
||||
let l = 4.0 * lum(x, y) - lum(x - 1, y) - lum(x + 1, y) - lum(x, y - 1) - lum(x, y + 1);
|
||||
vals.push(l);
|
||||
}
|
||||
}
|
||||
let n = vals.len() as f32;
|
||||
if n == 0.0 {
|
||||
return 0.0;
|
||||
}
|
||||
let mean = vals.iter().sum::<f32>() / n;
|
||||
vals.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / n
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn letterbox_fits_and_preserves_aspect() {
|
||||
// Landscape 1000×500 into 640 → width-bound, scale 0.64.
|
||||
let (nw, nh, s) = letterbox(1000, 500, 640);
|
||||
assert_eq!(nw, 640);
|
||||
assert_eq!(nh, 320);
|
||||
assert!((s - 0.64).abs() < 1e-6);
|
||||
// Square fills exactly.
|
||||
let (nw, nh, s) = letterbox(800, 800, 640);
|
||||
assert_eq!((nw, nh), (640, 640));
|
||||
assert!((s - 0.8).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn letterbox_degenerate_is_safe() {
|
||||
assert_eq!(letterbox(0, 10, 640), (0, 0, 1.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chw_layout_and_normalization() {
|
||||
let mut img = RgbImage::new(2, 1);
|
||||
img.put_pixel(0, 0, image::Rgb([127, 0, 255]));
|
||||
img.put_pixel(1, 0, image::Rgb([128, 255, 0]));
|
||||
let t = chw_normalized(&img, 127.5, 1.0 / 128.0);
|
||||
// Length = 3 channels × 2 px.
|
||||
assert_eq!(t.len(), 6);
|
||||
// R plane first, then G, then B (NCHW).
|
||||
assert!((t[0] - (127.0 - 127.5) / 128.0).abs() < 1e-6);
|
||||
assert!((t[1] - (128.0 - 127.5) / 128.0).abs() < 1e-6);
|
||||
assert!((t[2] - (0.0 - 127.5) / 128.0).abs() < 1e-6); // G of px0
|
||||
assert!((t[4] - (255.0 - 127.5) / 128.0).abs() < 1e-6); // B of px0
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn distance_decode_recovers_box_and_kps() {
|
||||
// 1×2 grid, stride 8, 1 anchor → cell centers (0,0) then (8,0).
|
||||
let scores = [0.9f32, 0.9];
|
||||
// distances left/top/right/bottom (already × stride), identical per cell.
|
||||
let bbox = [2.0, 1.0, 3.0, 4.0, 2.0, 1.0, 3.0, 4.0];
|
||||
let kps: Vec<f32> = vec![
|
||||
1.0, 1.0, 2.0, 2.0, 0.0, 0.0, -1.0, 1.0, 1.0, -1.0, // cell 0
|
||||
1.0, 1.0, 2.0, 2.0, 0.0, 0.0, -1.0, 1.0, 1.0, -1.0, // cell 1
|
||||
];
|
||||
let mut out = Vec::new();
|
||||
decode_stride(&scores, &bbox, Some(&kps), 8, 1, 2, 1, 0.5, &mut out);
|
||||
assert_eq!(out.len(), 2);
|
||||
// Cell 0, center (0,0): box = center ± distances, kps = center + offset.
|
||||
assert_eq!(out[0].bbox, [-2.0, -1.0, 3.0, 4.0]);
|
||||
assert_eq!(out[0].kps[0], [1.0, 1.0]);
|
||||
assert_eq!(out[0].kps[1], [2.0, 2.0]);
|
||||
// Cell 1, center (8,0): anchor center advanced by one stride in x.
|
||||
assert_eq!(out[1].bbox, [8.0 - 2.0, -1.0, 8.0 + 3.0, 4.0]);
|
||||
assert_eq!(out[1].kps[0], [9.0, 1.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_thresholds_out_low_scores() {
|
||||
let scores = [0.2f32, 0.8];
|
||||
let bbox = [0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0];
|
||||
let mut out = Vec::new();
|
||||
// 1×2 grid, 1 anchor → two cells.
|
||||
decode_stride(&scores, &bbox, None, 8, 1, 2, 1, 0.5, &mut out);
|
||||
assert_eq!(out.len(), 1);
|
||||
assert!((out[0].score - 0.8).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn iou_and_nms() {
|
||||
let a = [0.0, 0.0, 10.0, 10.0];
|
||||
let b = [0.0, 0.0, 10.0, 10.0];
|
||||
assert!((iou(&a, &b) - 1.0).abs() < 1e-6);
|
||||
let c = [100.0, 100.0, 110.0, 110.0];
|
||||
assert_eq!(iou(&a, &c), 0.0);
|
||||
|
||||
let dets = vec![
|
||||
Detection {
|
||||
bbox: a,
|
||||
kps: [[0.0; 2]; 5],
|
||||
score: 0.9,
|
||||
},
|
||||
Detection {
|
||||
bbox: b,
|
||||
kps: [[0.0; 2]; 5],
|
||||
score: 0.8,
|
||||
}, // dup of a
|
||||
Detection {
|
||||
bbox: c,
|
||||
kps: [[0.0; 2]; 5],
|
||||
score: 0.7,
|
||||
}, // separate
|
||||
];
|
||||
let kept = nms(dets, 0.4);
|
||||
assert_eq!(kept.len(), 2);
|
||||
assert!((kept[0].score - 0.9).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn similarity_identity() {
|
||||
let inv = similarity_transform_inverse(&ARCFACE_TEMPLATE, &ARCFACE_TEMPLATE);
|
||||
assert!((inv.a - 1.0).abs() < 1e-4);
|
||||
assert!(inv.b.abs() < 1e-4);
|
||||
assert!(inv.c.abs() < 1e-4);
|
||||
assert!((inv.d - 1.0).abs() < 1e-4);
|
||||
assert!(inv.tx.abs() < 1e-3);
|
||||
assert!(inv.ty.abs() < 1e-3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn similarity_pure_translation() {
|
||||
// src = dst shifted by (+10, +5); inverse must map out→src by the same shift.
|
||||
let mut src = ARCFACE_TEMPLATE;
|
||||
for p in &mut src {
|
||||
p[0] += 10.0;
|
||||
p[1] += 5.0;
|
||||
}
|
||||
let inv = similarity_transform_inverse(&src, &ARCFACE_TEMPLATE);
|
||||
assert!((inv.a - 1.0).abs() < 1e-4);
|
||||
assert!(inv.b.abs() < 1e-4);
|
||||
assert!((inv.tx - 10.0).abs() < 1e-3);
|
||||
assert!((inv.ty - 5.0).abs() < 1e-3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn warp_identity_preserves_template_region() {
|
||||
// A 112×112 gradient warped by identity returns (close to) itself.
|
||||
let mut img = RgbImage::new(ALIGN_SIZE, ALIGN_SIZE);
|
||||
for y in 0..ALIGN_SIZE {
|
||||
for x in 0..ALIGN_SIZE {
|
||||
img.put_pixel(x, y, image::Rgb([x as u8, y as u8, 128]));
|
||||
}
|
||||
}
|
||||
let inv = similarity_transform_inverse(&ARCFACE_TEMPLATE, &ARCFACE_TEMPLATE);
|
||||
let out = warp_to_aligned(&img, &inv);
|
||||
let a = out.get_pixel(40, 60);
|
||||
assert!((a[0] as i32 - 40).abs() <= 1);
|
||||
assert!((a[1] as i32 - 60).abs() <= 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn l2_normalize_unit_length() {
|
||||
let mut v = vec![3.0f32, 4.0];
|
||||
l2_normalize(&mut v);
|
||||
assert!((v[0] - 0.6).abs() < 1e-6);
|
||||
assert!((v[1] - 0.8).abs() < 1e-6);
|
||||
let mut z = vec![0.0f32, 0.0];
|
||||
l2_normalize(&mut z); // unchanged, no NaN
|
||||
assert_eq!(z, vec![0.0, 0.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn laplacian_variance_sharp_vs_flat() {
|
||||
let flat = RgbImage::from_pixel(8, 8, image::Rgb([100, 100, 100]));
|
||||
assert!(laplacian_variance(&flat) < 1e-3);
|
||||
let mut checker = RgbImage::new(8, 8);
|
||||
for y in 0..8 {
|
||||
for x in 0..8 {
|
||||
let v = if (x + y) % 2 == 0 { 0 } else { 255 };
|
||||
checker.put_pixel(x, y, image::Rgb([v, v, v]));
|
||||
}
|
||||
}
|
||||
assert!(laplacian_variance(&checker) > 1000.0);
|
||||
}
|
||||
}
|
||||
@@ -6,6 +6,7 @@ pub mod compression_service;
|
||||
pub mod dedup_service;
|
||||
pub mod encrypted_blob_backend;
|
||||
pub mod exif_service;
|
||||
pub mod face_geometry;
|
||||
pub mod face_indexing_service;
|
||||
pub mod file_content_cache;
|
||||
pub mod file_system_i18n_service;
|
||||
@@ -20,6 +21,8 @@ pub mod mock_email_sender;
|
||||
pub mod nextcloud_chunked_upload_service;
|
||||
pub mod noop_face_analyzer;
|
||||
pub mod oidc_service;
|
||||
#[cfg(feature = "faces-onnx")]
|
||||
pub mod onnx_face_analyzer;
|
||||
pub mod password_hasher;
|
||||
pub mod path_resolver_service;
|
||||
pub mod path_service;
|
||||
|
||||
@@ -0,0 +1,340 @@
|
||||
//! ONNX-backed face analyzer (SCRFD detector + ArcFace embedder).
|
||||
//!
|
||||
//! Compiled only with the `faces-onnx` cargo feature. Mirrors the
|
||||
//! immich/InsightFace pipeline: detect faces + 5-point landmarks (SCRFD),
|
||||
//! similarity-align each face to the canonical 112×112 template, then embed
|
||||
//! (ArcFace) into an L2-normalized 512-d vector. All inference runs on a
|
||||
//! blocking thread (`spawn_blocking`) so it never stalls a Tokio worker, and
|
||||
//! each ONNX session is serialized behind a `Mutex` (ORT's `run` needs `&mut`).
|
||||
//!
|
||||
//! The heavy numerical post-processing lives in [`super::face_geometry`] (plain
|
||||
//! Rust, unit-tested); this module only wires it to ONNX Runtime.
|
||||
//!
|
||||
//! **Models are operator-provided at runtime, never committed.** `load` returns
|
||||
//! an error (→ caller falls back to the no-op analyzer) if the ONNX Runtime
|
||||
//! dylib or either model file is missing or incompatible — the server still
|
||||
//! boots. The dylib is loaded via [`ort::init_from`] (a fallible path) rather
|
||||
//! than ORT's lazy loader, which would `panic` on a missing library (fatal
|
||||
//! under `panic = "abort"`).
|
||||
|
||||
use std::path::Path;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use image::RgbImage;
|
||||
use ort::session::Session;
|
||||
use ort::value::Tensor;
|
||||
|
||||
use super::face_geometry as geom;
|
||||
use crate::application::ports::face_ports::FaceAnalyzerPort;
|
||||
use crate::common::errors::DomainError;
|
||||
use crate::domain::entities::face::{BoundingBox, DetectedFace, EMBEDDING_DIM};
|
||||
|
||||
/// SCRFD pyramid strides for the 3- and 5-level model variants.
|
||||
const STRIDES_3: [u32; 3] = [8, 16, 32];
|
||||
const STRIDES_5: [u32; 5] = [8, 16, 32, 64, 128];
|
||||
|
||||
/// Discard faces smaller than this (original-image pixels) — embeddings of tiny
|
||||
/// faces are unreliable.
|
||||
const MIN_FACE_PX: f32 = 24.0;
|
||||
/// Hard cap on faces processed per image (bounds work on crowd shots).
|
||||
const MAX_FACES: usize = 64;
|
||||
|
||||
/// Output layout of an InsightFace SCRFD model, inferred from its output count.
|
||||
#[derive(Clone, Copy)]
|
||||
struct ScrfdLayout {
|
||||
/// Feature-map count per output kind (3 for strides 8/16/32, 5 with 64/128).
|
||||
fmc: usize,
|
||||
num_anchors: u32,
|
||||
use_kps: bool,
|
||||
}
|
||||
|
||||
impl ScrfdLayout {
|
||||
fn from_num_outputs(n: usize) -> Option<Self> {
|
||||
match n {
|
||||
6 => Some(Self {
|
||||
fmc: 3,
|
||||
num_anchors: 2,
|
||||
use_kps: false,
|
||||
}),
|
||||
9 => Some(Self {
|
||||
fmc: 3,
|
||||
num_anchors: 2,
|
||||
use_kps: true,
|
||||
}),
|
||||
10 => Some(Self {
|
||||
fmc: 5,
|
||||
num_anchors: 1,
|
||||
use_kps: false,
|
||||
}),
|
||||
15 => Some(Self {
|
||||
fmc: 5,
|
||||
num_anchors: 1,
|
||||
use_kps: true,
|
||||
}),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn strides(&self) -> &'static [u32] {
|
||||
if self.fmc == 3 {
|
||||
&STRIDES_3
|
||||
} else {
|
||||
&STRIDES_5
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Where to find the runtime + models, plus detector knobs. Borrowed paths;
|
||||
/// nothing is retained after [`OnnxFaceAnalyzer::load`].
|
||||
pub struct OnnxLoadConfig<'a> {
|
||||
/// Path to `libonnxruntime.{so,dylib,dll}`.
|
||||
pub dylib: &'a Path,
|
||||
/// SCRFD detector `.onnx`.
|
||||
pub detector: &'a Path,
|
||||
/// ArcFace embedder `.onnx`.
|
||||
pub embedder: &'a Path,
|
||||
pub det_size: u32,
|
||||
pub det_threshold: f32,
|
||||
pub nms_threshold: f32,
|
||||
/// ORT intra-op threads (0 = let ONNX Runtime decide).
|
||||
pub intra_threads: usize,
|
||||
}
|
||||
|
||||
struct Inner {
|
||||
detector: Mutex<Session>,
|
||||
embedder: Mutex<Session>,
|
||||
layout: ScrfdLayout,
|
||||
det_size: u32,
|
||||
det_threshold: f32,
|
||||
nms_threshold: f32,
|
||||
}
|
||||
|
||||
/// Real face analyzer. Cheap to clone (`Arc` inside).
|
||||
#[derive(Clone)]
|
||||
pub struct OnnxFaceAnalyzer {
|
||||
inner: Arc<Inner>,
|
||||
}
|
||||
|
||||
fn dom(e: impl std::fmt::Display) -> DomainError {
|
||||
DomainError::internal_error("Faces", e.to_string())
|
||||
}
|
||||
|
||||
fn build_session(path: &Path, intra_threads: usize) -> Result<Session, DomainError> {
|
||||
let mut builder = Session::builder().map_err(dom)?;
|
||||
if intra_threads > 0 {
|
||||
builder = builder.with_intra_threads(intra_threads).map_err(dom)?;
|
||||
}
|
||||
builder.commit_from_file(path).map_err(dom)
|
||||
}
|
||||
|
||||
impl OnnxFaceAnalyzer {
|
||||
/// Load the ONNX Runtime dylib and both models. Returns an error (caller
|
||||
/// falls back to the no-op analyzer) on any missing/incompatible artifact.
|
||||
pub fn load(cfg: &OnnxLoadConfig<'_>) -> Result<Self, DomainError> {
|
||||
// Fallible dylib load — populates ORT's global handle so later calls
|
||||
// never hit the panicking lazy loader.
|
||||
ort::init_from(cfg.dylib)
|
||||
.map_err(|e| dom(format!("ONNX Runtime dylib: {e}")))?
|
||||
.commit();
|
||||
|
||||
let detector = build_session(cfg.detector, cfg.intra_threads)?;
|
||||
let embedder = build_session(cfg.embedder, cfg.intra_threads)?;
|
||||
|
||||
let n_out = detector.outputs().len();
|
||||
let layout = ScrfdLayout::from_num_outputs(n_out).ok_or_else(|| {
|
||||
dom(format!(
|
||||
"detector has {n_out} outputs; expected an SCRFD model (6/9/10/15)"
|
||||
))
|
||||
})?;
|
||||
if !layout.use_kps {
|
||||
tracing::warn!(
|
||||
target: "oxicloud::faces",
|
||||
"SCRFD model has no landmark outputs; face alignment will be approximate"
|
||||
);
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
target: "oxicloud::faces",
|
||||
"ONNX face analyzer ready (detector {} outputs, embedder loaded, det_size={})",
|
||||
n_out, cfg.det_size
|
||||
);
|
||||
|
||||
Ok(Self {
|
||||
inner: Arc::new(Inner {
|
||||
detector: Mutex::new(detector),
|
||||
embedder: Mutex::new(embedder),
|
||||
layout,
|
||||
det_size: cfg.det_size,
|
||||
det_threshold: cfg.det_threshold,
|
||||
nms_threshold: cfg.nms_threshold,
|
||||
}),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Inner {
|
||||
/// Full synchronous pipeline for one encoded image.
|
||||
fn analyze_blocking(&self, image_bytes: &[u8]) -> Result<Vec<DetectedFace>, DomainError> {
|
||||
let orig = image::load_from_memory(image_bytes)
|
||||
.map_err(|e| dom(format!("decode image: {e}")))?
|
||||
.to_rgb8();
|
||||
let (w0, h0) = (orig.width(), orig.height());
|
||||
if w0 == 0 || h0 == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let dets = self.detect(&orig)?;
|
||||
|
||||
let mut faces = Vec::new();
|
||||
for det in dets.into_iter().take(MAX_FACES) {
|
||||
let fw = det.bbox[2] - det.bbox[0];
|
||||
let fh = det.bbox[3] - det.bbox[1];
|
||||
if fw < MIN_FACE_PX || fh < MIN_FACE_PX {
|
||||
continue;
|
||||
}
|
||||
let Some(embedding) = self.embed(&orig, &det)? else {
|
||||
continue;
|
||||
};
|
||||
let aligned_quality = {
|
||||
let inv = geom::similarity_transform_inverse(&det.kps, &geom::ARCFACE_TEMPLATE);
|
||||
let aligned = geom::warp_to_aligned(&orig, &inv);
|
||||
geom::laplacian_variance(&aligned)
|
||||
};
|
||||
let x = (det.bbox[0] / w0 as f32).clamp(0.0, 1.0);
|
||||
let y = (det.bbox[1] / h0 as f32).clamp(0.0, 1.0);
|
||||
let bw = (fw / w0 as f32).clamp(0.0, 1.0);
|
||||
let bh = (fh / h0 as f32).clamp(0.0, 1.0);
|
||||
faces.push(DetectedFace {
|
||||
bbox: BoundingBox { x, y, w: bw, h: bh },
|
||||
det_score: det.score,
|
||||
quality: Some(aligned_quality),
|
||||
embedding,
|
||||
});
|
||||
}
|
||||
Ok(faces)
|
||||
}
|
||||
|
||||
/// Run SCRFD and return detections in **original-image pixels**.
|
||||
fn detect(&self, orig: &RgbImage) -> Result<Vec<geom::Detection>, DomainError> {
|
||||
let det = self.det_size;
|
||||
let (nw, nh, scale) = geom::letterbox(orig.width(), orig.height(), det);
|
||||
let resized = image::imageops::resize(orig, nw, nh, image::imageops::FilterType::Triangle);
|
||||
let mut canvas = RgbImage::new(det, det);
|
||||
image::imageops::overlay(&mut canvas, &resized, 0, 0);
|
||||
let input = geom::chw_normalized(&canvas, 127.5, 1.0 / 128.0);
|
||||
let tensor =
|
||||
Tensor::from_array(([1_i64, 3, det as i64, det as i64], input)).map_err(dom)?;
|
||||
|
||||
let layout = self.layout;
|
||||
let total = layout.fmc * if layout.use_kps { 3 } else { 2 };
|
||||
let raw: Vec<Vec<f32>> = {
|
||||
let mut sess = self
|
||||
.detector
|
||||
.lock()
|
||||
.map_err(|_| dom("detector mutex poisoned"))?;
|
||||
let outputs = sess.run(ort::inputs![tensor]).map_err(dom)?;
|
||||
(0..total)
|
||||
.map(|i| {
|
||||
outputs[i]
|
||||
.try_extract_tensor::<f32>()
|
||||
.map(|(_, data)| data.to_vec())
|
||||
.map_err(dom)
|
||||
})
|
||||
.collect::<Result<_, _>>()?
|
||||
};
|
||||
|
||||
let mut dets = Vec::new();
|
||||
for (si, &stride) in layout.strides().iter().enumerate() {
|
||||
let scores = &raw[si];
|
||||
let bbox: Vec<f32> = raw[layout.fmc + si]
|
||||
.iter()
|
||||
.map(|v| v * stride as f32)
|
||||
.collect();
|
||||
let kps: Option<Vec<f32>> = if layout.use_kps {
|
||||
Some(
|
||||
raw[2 * layout.fmc + si]
|
||||
.iter()
|
||||
.map(|v| v * stride as f32)
|
||||
.collect(),
|
||||
)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let feat = det / stride;
|
||||
geom::decode_stride(
|
||||
scores,
|
||||
&bbox,
|
||||
kps.as_deref(),
|
||||
stride,
|
||||
feat,
|
||||
feat,
|
||||
layout.num_anchors,
|
||||
self.det_threshold,
|
||||
&mut dets,
|
||||
);
|
||||
}
|
||||
|
||||
// Scale detector-space coordinates back to the original image.
|
||||
let inv_scale = if scale.abs() < 1e-9 { 1.0 } else { 1.0 / scale };
|
||||
for d in &mut dets {
|
||||
for v in &mut d.bbox {
|
||||
*v *= inv_scale;
|
||||
}
|
||||
for k in &mut d.kps {
|
||||
k[0] *= inv_scale;
|
||||
k[1] *= inv_scale;
|
||||
}
|
||||
}
|
||||
Ok(geom::nms(dets, self.nms_threshold))
|
||||
}
|
||||
|
||||
/// Align one detection and run the ArcFace embedder. Returns `None` if the
|
||||
/// embedder produces an unexpected output length.
|
||||
fn embed(
|
||||
&self,
|
||||
orig: &RgbImage,
|
||||
det: &geom::Detection,
|
||||
) -> Result<Option<Vec<f32>>, DomainError> {
|
||||
let inv = geom::similarity_transform_inverse(&det.kps, &geom::ARCFACE_TEMPLATE);
|
||||
let aligned = geom::warp_to_aligned(orig, &inv);
|
||||
let input = geom::chw_normalized(&aligned, 127.5, 1.0 / 127.5);
|
||||
let size = geom::ALIGN_SIZE as i64;
|
||||
let tensor = Tensor::from_array(([1_i64, 3, size, size], input)).map_err(dom)?;
|
||||
|
||||
let mut embedding: Vec<f32> = {
|
||||
let mut sess = self
|
||||
.embedder
|
||||
.lock()
|
||||
.map_err(|_| dom("embedder mutex poisoned"))?;
|
||||
let outputs = sess.run(ort::inputs![tensor]).map_err(dom)?;
|
||||
let (_, data) = outputs[0].try_extract_tensor::<f32>().map_err(dom)?;
|
||||
data.to_vec()
|
||||
};
|
||||
if embedding.len() != EMBEDDING_DIM {
|
||||
tracing::warn!(
|
||||
target: "oxicloud::faces",
|
||||
"embedder returned {} dims, expected {EMBEDDING_DIM}; skipping face",
|
||||
embedding.len()
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
geom::l2_normalize(&mut embedding);
|
||||
Ok(Some(embedding))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl FaceAnalyzerPort for OnnxFaceAnalyzer {
|
||||
fn is_ready(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn analyze(&self, image_bytes: &[u8]) -> Result<Vec<DetectedFace>, DomainError> {
|
||||
let inner = self.inner.clone();
|
||||
let bytes = image_bytes.to_vec();
|
||||
tokio::task::spawn_blocking(move || inner.analyze_blocking(&bytes))
|
||||
.await
|
||||
.map_err(|e| dom(format!("inference task join: {e}")))?
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user