dtourolleandClaude Opus 5
2ac069a6b3
Wires dr-face to dr-catalog: a background sweep that reads the proxy the grid already built, detects, aligns, embeds and stores, then a clustering pass that turns those embeddings into suggested people. Detection runs on the Large thumbnail tier and nowhere else. That is what makes the feature affordable -- a browsed library has already paid for its proxies, so face indexing adds no RAW decode that was not already happening -- and it is why an image whose proxy is missing is skipped rather than fetched: requesting one here would put face indexing on the network path FR-CULL-8 keeps it off. The sweep keeps no cursor. It asks the catalog what is missing, so it resumes after process death with no repeated work beyond the in-flight image, and cancelling is dropping the receiver. recluster writes only the suggested half. Confirmed faces go in as anchors and come back untouched, and a cluster of one stays nameless -- naming every stray face would fill the People view with noise the user then has to dismiss. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
390 lines
14 KiB
Rust
390 lines
14 KiB
Rust
//! SCRFD face detection (docs/faces.md §4).
|
||
//!
|
||
//! One forward pass produces a box, a confidence and **five landmarks** per
|
||
//! face — the landmarks being the reason for this detector rather than a
|
||
//! general one, since [`crate::align`] cannot work without them.
|
||
//!
|
||
//! # The graph must have fixed input dimensions
|
||
//!
|
||
//! InsightFace ships `det_500m.onnx` with a dynamic H/W input, and **tract
|
||
//! cannot parse it in that form** — it fails at node #0. The same file run
|
||
//! through `tools/fix-face-model-shapes.sh` loads cleanly. Its outputs were
|
||
//! already static at 640, so 640 is not a choice made here: it is the shape
|
||
//! the export was always going to run at.
|
||
|
||
use ndarray::Array4;
|
||
|
||
use crate::{install_backend, FaceError};
|
||
|
||
/// The graph's input edge, in pixels. See the module note: not configurable.
|
||
pub const INPUT_EDGE: usize = 640;
|
||
|
||
/// Strides, in the order SCRFD emits them.
|
||
const ALL_STRIDES: [usize; 4] = [8, 16, 32, 64];
|
||
|
||
/// Anchors per feature-map location.
|
||
const ANCHORS: usize = 2;
|
||
|
||
/// How detection is tuned.
|
||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||
pub struct DetectOptions {
|
||
/// Minimum detector confidence.
|
||
///
|
||
/// Deliberately *not* the low threshold `dr-segment` chose. There a false
|
||
/// positive costs one spurious row in a list the user is picking from;
|
||
/// here it costs a face in the People view to reject and — worse — a
|
||
/// garbage embedding that can bridge two real clusters into one. A false
|
||
/// negative is recoverable by re-indexing with a better model; a polluted
|
||
/// cluster graph, once the user has confirmed faces inside it, is not.
|
||
pub confidence: f32,
|
||
/// Box IoU above which two detections are judged to be the same face.
|
||
pub nms_iou: f32,
|
||
/// Smallest face to keep, in source pixels on the shorter box edge.
|
||
///
|
||
/// Below this there is not enough face left to align reliably — see
|
||
/// docs/faces.md §7 for what the embedder actually receives at each size.
|
||
pub min_face_px: f32,
|
||
}
|
||
|
||
impl Default for DetectOptions {
|
||
fn default() -> Self {
|
||
Self {
|
||
confidence: 0.5,
|
||
nms_iou: 0.4,
|
||
min_face_px: 40.0,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// One detected face, in **source image pixels**.
|
||
///
|
||
/// Pixels rather than the normalised form the catalog stores, because the
|
||
/// caller still has to crop from this image. Normalisation happens at the
|
||
/// storage boundary, where the long edge is known to be the right divisor.
|
||
#[derive(Debug, Clone, PartialEq)]
|
||
pub struct Detection {
|
||
/// `(x0, y0, x1, y1)`.
|
||
pub bbox: (f32, f32, f32, f32),
|
||
/// Five points in the detector's own order — see [`crate::align`], which
|
||
/// consumes them without reordering.
|
||
pub landmarks: [(f32, f32); 5],
|
||
pub confidence: f32,
|
||
}
|
||
|
||
impl Detection {
|
||
pub fn width(&self) -> f32 {
|
||
self.bbox.2 - self.bbox.0
|
||
}
|
||
pub fn height(&self) -> f32 {
|
||
self.bbox.3 - self.bbox.1
|
||
}
|
||
}
|
||
|
||
/// A loaded SCRFD graph.
|
||
pub struct Detector {
|
||
session: ort::session::Session,
|
||
/// Feature-map count: 3 for strides {8,16,32}, 4 for {8,16,32,64}.
|
||
///
|
||
/// Discovered from the output count rather than assumed, because both
|
||
/// exports exist and hardcoding 3 silently ignores the largest faces a
|
||
/// four-stride model finds.
|
||
fmc: usize,
|
||
}
|
||
|
||
impl Detector {
|
||
pub fn from_path(path: impl AsRef<std::path::Path>) -> Result<Self, FaceError> {
|
||
let bytes = std::fs::read(path).map_err(FaceError::ModelRead)?;
|
||
Self::from_bytes(&bytes)
|
||
}
|
||
|
||
pub fn from_bytes(bytes: &[u8]) -> Result<Self, FaceError> {
|
||
install_backend();
|
||
|
||
let session = ort::session::Session::builder()
|
||
.map_err(FaceError::Inference)?
|
||
.commit_from_memory(bytes)
|
||
.map_err(FaceError::Inference)?;
|
||
|
||
let n_out = session.outputs().len();
|
||
if n_out % 3 != 0 || !(9..=12).contains(&n_out) {
|
||
return Err(FaceError::WrongModel {
|
||
expected: "InsightFace SCRFD",
|
||
detail: format!("expected 9 or 12 outputs, got {n_out}"),
|
||
});
|
||
}
|
||
let fmc = n_out / 3;
|
||
|
||
// The check that actually distinguishes the models. YuNet also has
|
||
// twelve outputs in three strides, so the count proves nothing — its
|
||
// groups are cls/obj/bbox/kps where SCRFD's are score/bbox/kps, and
|
||
// decoding one as the other yields a page of plausible numbers rather
|
||
// than an error. The last dimension is what separates them.
|
||
for (group, expected_last) in [1_i64, 4, 10].into_iter().enumerate() {
|
||
for s in 0..fmc {
|
||
let idx = group * fmc + s;
|
||
let out = &session.outputs()[idx];
|
||
let last: Option<i64> =
|
||
out.dtype().tensor_shape().and_then(|d| d.last().copied());
|
||
if last != Some(expected_last) {
|
||
return Err(FaceError::WrongModel {
|
||
expected: "InsightFace SCRFD",
|
||
detail: format!(
|
||
"output '{}' last dim is {:?}, expected {expected_last} \
|
||
(a YuNet export fails exactly here)",
|
||
out.name(),
|
||
last
|
||
),
|
||
});
|
||
}
|
||
}
|
||
}
|
||
|
||
Ok(Self { session, fmc })
|
||
}
|
||
|
||
/// Stride levels this graph emits.
|
||
pub fn strides(&self) -> &'static [usize] {
|
||
&ALL_STRIDES[..self.fmc]
|
||
}
|
||
|
||
/// Find the faces in an image.
|
||
///
|
||
/// `rgb` is tightly packed `f32` RGB in `0.0..=1.0`, row-major — the same
|
||
/// convention `dr-segment` and [`crate::align`] use.
|
||
pub fn detect(
|
||
&mut self,
|
||
rgb: &[f32],
|
||
width: usize,
|
||
height: usize,
|
||
options: &DetectOptions,
|
||
) -> Result<Vec<Detection>, FaceError> {
|
||
if width == 0 || height == 0 {
|
||
return Ok(Vec::new());
|
||
}
|
||
if rgb.len() != width * height * 3 {
|
||
return Err(FaceError::ImageShape {
|
||
expected: width * height * 3,
|
||
got: rgb.len(),
|
||
});
|
||
}
|
||
|
||
let lb = Letterbox::fit(width as f32, height as f32);
|
||
let input = lb.sample(rgb, width, height);
|
||
|
||
let outputs = self
|
||
.session
|
||
.run(ort::inputs![
|
||
ort::value::Tensor::from_array(input).map_err(FaceError::Inference)?
|
||
])
|
||
.map_err(FaceError::Inference)?;
|
||
|
||
let mut raw: Vec<Detection> = Vec::new();
|
||
|
||
for (si, &stride) in ALL_STRIDES[..self.fmc].iter().enumerate() {
|
||
let (_, scores) = outputs[si]
|
||
.try_extract_tensor::<f32>()
|
||
.map_err(FaceError::Inference)?;
|
||
let (_, boxes) = outputs[self.fmc + si]
|
||
.try_extract_tensor::<f32>()
|
||
.map_err(FaceError::Inference)?;
|
||
let (_, kps) = outputs[self.fmc * 2 + si]
|
||
.try_extract_tensor::<f32>()
|
||
.map_err(FaceError::Inference)?;
|
||
|
||
let fw = INPUT_EDGE / stride;
|
||
let fh = INPUT_EDGE / stride;
|
||
let s = stride as f32;
|
||
|
||
for r in 0..fh {
|
||
for c in 0..fw {
|
||
for a in 0..ANCHORS {
|
||
let idx = (r * fw + c) * ANCHORS + a;
|
||
let score = scores[idx];
|
||
if score < options.confidence {
|
||
continue;
|
||
}
|
||
|
||
// Anchor centre in input space, then distance-to-box
|
||
// decoding: the four regressed values are distances
|
||
// left/top/right/bottom in units of the stride.
|
||
let (cx, cy) = ((c * stride) as f32, (r * stride) as f32);
|
||
let b = &boxes[idx * 4..idx * 4 + 4];
|
||
let (x0, y0) = lb.into_source(cx - b[0] * s, cy - b[1] * s);
|
||
let (x1, y1) = lb.into_source(cx + b[2] * s, cy + b[3] * s);
|
||
|
||
let k = &kps[idx * 10..idx * 10 + 10];
|
||
let mut landmarks = [(0.0_f32, 0.0_f32); 5];
|
||
for (p, lm) in landmarks.iter_mut().enumerate() {
|
||
*lm = lb.into_source(cx + k[p * 2] * s, cy + k[p * 2 + 1] * s);
|
||
}
|
||
|
||
raw.push(Detection {
|
||
bbox: (x0, y0, x1, y1),
|
||
landmarks,
|
||
confidence: score,
|
||
});
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
let mut kept = non_max_suppress(raw, options.nms_iou);
|
||
|
||
// Size floor last, on the *merged* boxes: a face that only clears the
|
||
// floor once NMS has picked the best of its overlapping detections
|
||
// should be kept.
|
||
kept.retain(|d| d.width().min(d.height()) >= options.min_face_px);
|
||
|
||
// No cap on the count. The reference implementation keeps the ten
|
||
// largest, which is right for a film frame where background extras are
|
||
// noise; it is wrong for a photo library, where a group shot with
|
||
// thirty faces is precisely the picture worth indexing.
|
||
Ok(kept)
|
||
}
|
||
}
|
||
|
||
/// Greedy NMS across all strides together.
|
||
fn non_max_suppress(mut dets: Vec<Detection>, iou_threshold: f32) -> Vec<Detection> {
|
||
dets.sort_by(|a, b| b.confidence.total_cmp(&a.confidence));
|
||
let mut kept: Vec<Detection> = Vec::new();
|
||
for d in dets {
|
||
if kept.iter().all(|k| iou(&k.bbox, &d.bbox) <= iou_threshold) {
|
||
kept.push(d);
|
||
}
|
||
}
|
||
kept
|
||
}
|
||
|
||
fn iou(a: &(f32, f32, f32, f32), b: &(f32, f32, f32, f32)) -> f32 {
|
||
let ix = (a.2.min(b.2) - a.0.max(b.0)).max(0.0);
|
||
let iy = (a.3.min(b.3) - a.1.max(b.1)).max(0.0);
|
||
let inter = ix * iy;
|
||
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
|
||
}
|
||
}
|
||
|
||
/// How the image is fitted into the graph's fixed square input.
|
||
///
|
||
/// The forward and inverse mappings live in one struct on purpose:
|
||
/// docs/faces.md §4.1 notes that what matters is not *where* the padding goes
|
||
/// but that the two agree. A mismatch offsets every box and landmark by the
|
||
/// padding, producing detections that look plausible and embeddings that
|
||
/// quietly cluster badly three stages later.
|
||
#[derive(Debug, Clone, Copy)]
|
||
struct Letterbox {
|
||
/// Input pixels per source pixel.
|
||
scale: f32,
|
||
pad_x: f32,
|
||
pad_y: f32,
|
||
}
|
||
|
||
impl Letterbox {
|
||
fn fit(w: f32, h: f32) -> Self {
|
||
let scale = (INPUT_EDGE as f32 / w).min(INPUT_EDGE as f32 / h);
|
||
Self {
|
||
scale,
|
||
pad_x: (INPUT_EDGE as f32 - w * scale) * 0.5,
|
||
pad_y: (INPUT_EDGE as f32 - h * scale) * 0.5,
|
||
}
|
||
}
|
||
|
||
/// Resample into `[1, 3, 640, 640]`, normalised as the weights expect.
|
||
///
|
||
/// `(x·255 − 127.5) / 128` — note `/128`, not `/127.5`. The reference
|
||
/// implementation this is ported from uses `/128` for both models, and
|
||
/// every measured number in docs/faces.md §1 came from it.
|
||
///
|
||
/// Padding is grey, matching the reference's `114`: the value the network
|
||
/// reads least as an edge, where black would draw a hard border across the
|
||
/// frame and invite a detection along it.
|
||
fn sample(&self, rgb: &[f32], width: usize, height: usize) -> Array4<f32> {
|
||
const PAD: f32 = 114.0;
|
||
let norm = |v: f32| (v * 255.0 - 127.5) / 128.0;
|
||
|
||
let mut input =
|
||
Array4::<f32>::from_elem((1, 3, INPUT_EDGE, INPUT_EDGE), (PAD - 127.5) / 128.0);
|
||
|
||
for iy in 0..INPUT_EDGE {
|
||
let sy = (iy as f32 + 0.5 - self.pad_y) / self.scale - 0.5;
|
||
if sy < -0.5 || sy > height as f32 - 0.5 {
|
||
continue;
|
||
}
|
||
for ix in 0..INPUT_EDGE {
|
||
let sx = (ix as f32 + 0.5 - self.pad_x) / self.scale - 0.5;
|
||
if sx < -0.5 || sx > width as f32 - 0.5 {
|
||
continue;
|
||
}
|
||
let (x0f, y0f) = (sx.floor(), sy.floor());
|
||
let (fx, fy) = (sx - x0f, sy - y0f);
|
||
let x0 = (x0f as isize).clamp(0, width as isize - 1) as usize;
|
||
let y0 = (y0f as isize).clamp(0, height as isize - 1) as usize;
|
||
let x1 = (x0 + 1).min(width - 1);
|
||
let y1 = (y0 + 1).min(height - 1);
|
||
|
||
for c in 0..3 {
|
||
let at = |x: usize, y: usize| rgb[(y * width + x) * 3 + c];
|
||
let top = at(x0, y0) * (1.0 - fx) + at(x1, y0) * fx;
|
||
let bot = at(x0, y1) * (1.0 - fx) + at(x1, y1) * fx;
|
||
input[[0, c, iy, ix]] = norm(top * (1.0 - fy) + bot * fy);
|
||
}
|
||
}
|
||
}
|
||
|
||
input
|
||
}
|
||
|
||
/// Input-space point back to source pixels.
|
||
fn into_source(self, x: f32, y: f32) -> (f32, f32) {
|
||
((x - self.pad_x) / self.scale, (y - self.pad_y) / self.scale)
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn letterbox_round_trips_a_point() {
|
||
let lb = Letterbox::fit(1024.0, 683.0);
|
||
for &(x, y) in &[(0.0_f32, 0.0_f32), (512.0, 341.0), (1023.0, 682.0)] {
|
||
let (bx, by) = lb.into_source(x * lb.scale + lb.pad_x, y * lb.scale + lb.pad_y);
|
||
assert!((bx - x).abs() < 1e-2, "{bx} vs {x}");
|
||
assert!((by - y).abs() < 1e-2, "{by} vs {y}");
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn letterbox_centres_the_short_axis() {
|
||
let lb = Letterbox::fit(640.0, 320.0);
|
||
assert!((lb.scale - 1.0).abs() < 1e-6);
|
||
assert!(lb.pad_x.abs() < 1e-6);
|
||
assert!((lb.pad_y - 160.0).abs() < 1e-6);
|
||
}
|
||
|
||
#[test]
|
||
fn nms_keeps_the_confident_box_and_drops_its_duplicate() {
|
||
let d = |x: f32, conf: f32| Detection {
|
||
bbox: (x, 0.0, x + 100.0, 100.0),
|
||
landmarks: [(0.0, 0.0); 5],
|
||
confidence: conf,
|
||
};
|
||
let kept = non_max_suppress(vec![d(0.0, 0.8), d(5.0, 0.9), d(500.0, 0.7)], 0.4);
|
||
assert_eq!(kept.len(), 2);
|
||
assert!((kept[0].confidence - 0.9).abs() < 1e-6);
|
||
assert!((kept[1].bbox.0 - 500.0).abs() < 1e-6);
|
||
}
|
||
|
||
#[test]
|
||
fn iou_of_a_box_with_itself_is_one_and_with_a_disjoint_box_is_zero() {
|
||
let a = (0.0, 0.0, 10.0, 10.0);
|
||
assert!((iou(&a, &a) - 1.0).abs() < 1e-6);
|
||
assert!(iou(&a, &(100.0, 100.0, 110.0, 110.0)) < 1e-6);
|
||
}
|
||
}
|