diff --git a/Cargo.lock b/Cargo.lock index f837cab2..90b2d8e6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -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" diff --git a/Cargo.toml b/Cargo.toml index 10b9e41a..18b4e207 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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)'] } diff --git a/example.env b/example.env index e1bbc0ae..a4913e46 100644 --- a/example.env +++ b/example.env @@ -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, diff --git a/src/common/config.rs b/src/common/config.rs index 3c28d46f..84bad869 100644 --- a/src/common/config.rs +++ b/src/common/config.rs @@ -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, + /// SCRFD/RetinaFace detector model with 5-point landmarks. + /// Env: `OXICLOUD_FACES_DETECTOR_MODEL`. + pub detector_model: Option, + /// ArcFace embedder model (112×112 → 512-d). + /// Env: `OXICLOUD_FACES_EMBEDDER_MODEL`. + pub embedder_model: Option, + /// 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::()) + && let Ok(val) = v + { + config.faces.det_size = val; + } + if let Ok(v) = env::var("OXICLOUD_FACES_DET_THRESHOLD").map(|v| v.parse::()) + && let Ok(val) = v + { + config.faces.det_threshold = val; + } + if let Ok(v) = env::var("OXICLOUD_FACES_NMS_THRESHOLD").map(|v| v.parse::()) + && let Ok(val) = v + { + config.faces.nms_threshold = val; + } + if let Ok(v) = env::var("OXICLOUD_FACES_INTRA_THREADS").map(|v| v.parse::()) + && 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::()) && let Ok(val) = v diff --git a/src/common/di.rs b/src/common/di.rs index d75cbea7..0ca12004 100644 --- a/src/common/di.rs +++ b/src/common/di.rs @@ -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, ) -> Arc { let blob_root = self.storage_path.join(".blobs"); - let analyzer: Arc = - 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 { + #[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) -> Arc { let repo = Arc::new( diff --git a/src/infrastructure/services/face_geometry.rs b/src/infrastructure/services/face_geometry.rs new file mode 100644 index 00000000..f607cbad --- /dev/null +++ b/src/infrastructure/services/face_geometry.rs @@ -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 { + 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, +) { + 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, iou_thresh: f32) -> Vec { + dets.sort_by(|a, b| b.score.total_cmp(&a.score)); + let mut keep: Vec = 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 { + 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::().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::() / n; + vals.iter().map(|v| (v - mean) * (v - mean)).sum::() / 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 = 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); + } +} diff --git a/src/infrastructure/services/mod.rs b/src/infrastructure/services/mod.rs index b41b6edc..67c52e70 100644 --- a/src/infrastructure/services/mod.rs +++ b/src/infrastructure/services/mod.rs @@ -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; diff --git a/src/infrastructure/services/onnx_face_analyzer.rs b/src/infrastructure/services/onnx_face_analyzer.rs new file mode 100644 index 00000000..f58d73c5 --- /dev/null +++ b/src/infrastructure/services/onnx_face_analyzer.rs @@ -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 { + 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, + embedder: Mutex, + 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, +} + +fn dom(e: impl std::fmt::Display) -> DomainError { + DomainError::internal_error("Faces", e.to_string()) +} + +fn build_session(path: &Path, intra_threads: usize) -> Result { + 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 { + // 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, 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, 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> = { + 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::() + .map(|(_, data)| data.to_vec()) + .map_err(dom) + }) + .collect::>()? + }; + + let mut dets = Vec::new(); + for (si, &stride) in layout.strides().iter().enumerate() { + let scores = &raw[si]; + let bbox: Vec = raw[layout.fmc + si] + .iter() + .map(|v| v * stride as f32) + .collect(); + let kps: Option> = 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>, 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 = { + 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::().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, 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}")))? + } +}