Files
DarkRoom/core/dr-face/src/detect.rs
T
dtourolleandClaude Opus 5 1d4c3348af Let a 32-pixel face count, and move the blur floor with it
64 source pixels was too strict: it threw away 70% of everything the detector
finds, and plenty of what it took were faces a person could name.

Lowering it is not a one-line change, because the two floors are coupled. A face
under 112 pixels is *upsampled* to reach the embedder and upsampling invents no
edges, so a small face scores low on sharpness however crisp the original was.
Re-measured over the reference library with `face_index --quality`:

    min crop   min sharp   size cut   blur cut       kept
          32       0.000        44%         0%        56%
          32       0.002        44%         3%        53%
          32       0.005        44%         8%        48%
          32       0.010        44%        16%        40%
          32       0.020        44%        27%        30%
          64       0.020        70%         7%        23%

Holding the blur floor at 0.020 while dropping the size floor to 32 would have
rejected a further 27% — for being small rather than for being blurred — and
kept only 30%, barely more than the 23% the strict pair kept. Most of the point
of lowering the size floor would have gone straight back out through the other
gate.

0.005 removes 8% of what the size floor leaves, which is the same job 0.020 was
doing at 64 (7%): the large-but-soft face this gate exists for. Together they
now keep 48% of what the detector finds, against 23% before.

The box pre-filter follows down to 24, staying below what the real floor accepts
so it cannot reject a face that would have cleared 32.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-28 08:30:09 +02:00

442 lines
17 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! 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,
/// Cheap pre-filter: smallest box to keep, in source pixels on the shorter
/// edge.
///
/// **Not the real size floor** — [`DetectOptions::min_source_px`] is, and
/// it is measured on the aligned crop rather than the box. This one exists
/// only to throw away the obviously hopeless before paying for a warp, so
/// it is deliberately set *below* what the real floor will accept: the
/// aligned crop spans roughly 1.3x the box's shorter edge, so 24 here
/// cannot reject a face that would have cleared 32 there.
pub min_face_px: f32,
/// Smallest face the embedder may be given, in **source pixels across the
/// aligned crop** — `crop_px` in the catalog.
///
/// The honest statement of "a face must be at least 32x32", because this is
/// the number of real pixels behind the 112x112 the model actually sees.
/// The box's own size is not that: the ArcFace template reaches past the
/// box for forehead and chin, so a 64-pixel box and a 64-pixel crop are
/// different faces.
///
/// Below this the crop was upsampled to reach the embedder, and upsampling
/// invents no detail — the embedding is of a soft, stretched face and is
/// correspondingly untrustworthy.
///
/// Applied after alignment, so it lives with the sharpness floor rather
/// than with the detector. See [`DetectOptions::min_sharpness`].
pub min_source_px: f32,
/// Least acceptable [`crate::align::Aligned112::sharpness`].
///
/// Applied after alignment rather than here, because it is a property of
/// the warped crop the embedder receives and not of the box. The pipeline
/// that enforces it is `dr_ui::faces::index_proxy`; it lives on this struct
/// so that every quality decision about a face is configured in one place
/// and a caller cannot enable one gate while forgetting the other.
///
/// Zero disables it, which is what a measurement run wants.
///
/// # It has to move with the size floor
///
/// The two are coupled, because an upsampled face scores low here whatever
/// its original sharpness. Measured over the reference library, with the
/// size floor at 32 source pixels:
///
/// | min sharpness | of what the size floor left, this removes |
/// |---|---|
/// | 0.002 | 3% |
/// | 0.005 | 8% |
/// | 0.010 | 16% |
/// | 0.020 | 27% |
///
/// At a 64-pixel floor, 0.020 removed 7% — the same *kind* of face, the
/// large-but-soft one this gate exists for. Holding 0.020 while dropping
/// the size floor to 32 would have thrown away a quarter of the newly
/// admitted faces for being small rather than for being blurred, undoing
/// most of the point of lowering it. 0.005 removes 8% at 32, which is the
/// same job.
pub min_sharpness: f32,
}
impl Default for DetectOptions {
fn default() -> Self {
Self {
confidence: 0.5,
nms_iou: 0.4,
min_face_px: 24.0,
min_source_px: 32.0,
min_sharpness: 0.005,
}
}
}
/// 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);
}
}