diff --git a/crates/segment/src/face.rs b/crates/segment/src/face.rs index df0474b8..f49f38cc 100644 --- a/crates/segment/src/face.rs +++ b/crates/segment/src/face.rs @@ -55,26 +55,50 @@ impl Roi { } } -/// Sample `roi` into an `n×n` RGB tensor (bilinear; zero outside the frame), each channel mapped -/// from 0–1 to `lo`–`hi`. -pub fn crop(rgb: &[[f32; 3]], w: usize, h: usize, roi: &Roi, n: usize, [lo, hi]: [f32; 2]) -> Vec { - let mut out = vec![0.0f32; n * n * 3]; - if rgb.len() < w * h || n == 0 { +/// What a [`crop`] reads outside the frame. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Border { + /// Black (MediaPipe's face detector input). + Zero, + /// The nearest edge pixel (MediaPipe's face mesh input). + Replicate, +} + +/// Sample `roi` into an `n×n` RGB tensor (bilinear), each channel mapped from 0–1 to `lo`–`hi`. +/// +/// The sampling follows MediaPipe's CPU `ImageToTensor`, so the landmarks match MediaPipe's to a +/// fraction of a pixel: the square's corners map to the output corners `(0,0)` to `(n,n)` (output +/// pixel `i` is at `u = i/n`, not at its centre), the source position is read as a pixel-centre +/// index with no half-pixel shift, and the result is rounded to a whole 0–255 value before the range +/// map. Outside the frame it reads `border`. A frame whose buffer is shorter than `w·h`, or whose +/// size overflows, yields `lo` everywhere. +pub fn crop(rgb: &[[f32; 3]], w: usize, h: usize, roi: &Roi, n: usize, [lo, hi]: [f32; 2], border: Border) -> Vec { + let mut out = vec![lo; n.saturating_mul(n).saturating_mul(3)]; + if w == 0 || h == 0 || n == 0 || w.checked_mul(h).is_none_or(|len| rgb.len() < len) { return out; } + let (wi, hi_) = (w as i64, h as i64); + // In range after the border rule, so the index is below `w·h` (which fits in `usize`). let px = |x: i64, y: i64| -> [f32; 3] { - if x < 0 || y < 0 || x >= w as i64 || y >= h as i64 { [0.0; 3] } else { rgb.get(y as usize * w + x as usize).copied().unwrap_or([0.0; 3]) } + let (x, y) = match border { + Border::Zero if x < 0 || y < 0 || x >= wi || y >= hi_ => return [0.0; 3], + Border::Zero => (x, y), + Border::Replicate => (x.clamp(0, wi - 1), y.clamp(0, hi_ - 1)), + }; + rgb.get(y as usize * w + x as usize).copied().unwrap_or([0.0; 3]) }; for j in 0..n { for i in 0..n { - let p = roi.to_frame((i as f32 + 0.5) / n as f32, (j as f32 + 0.5) / n as f32); - let (x, y) = (p[0] - 0.5, p[1] - 0.5); + let [x, y] = roi.to_frame(i as f32 / n as f32, j as f32 / n as f32); let (x0, y0) = (x.floor(), y.floor()); let (fx, fy) = (x - x0, y - y0); + // Saturating casts: NaN becomes 0 and an absurd position reads the border. let (x0, y0) = (x0 as i64, y0 as i64); - let (a, b, c, d) = (px(x0, y0), px(x0 + 1, y0), px(x0, y0 + 1), px(x0 + 1, y0 + 1)); + let (x1, y1) = (x0.saturating_add(1), y0.saturating_add(1)); + let (a, b, c, d) = (px(x0, y0), px(x1, y0), px(x0, y1), px(x1, y1)); for k in 0..3 { let v = (a[k] * (1.0 - fx) + b[k] * fx) * (1.0 - fy) + (c[k] * (1.0 - fx) + d[k] * fx) * fy; + let v = if v.is_finite() { (v * 255.0).round().clamp(0.0, 255.0) / 255.0 } else { 0.0 }; out[(j * n + i) * 3 + k] = lo + v * (hi - lo); } } @@ -102,14 +126,48 @@ mod tests { } #[test] - fn crop_samples_pixel_centres() { - // A 4×4 ramp cropped 1:1 reproduces it; outside the frame is zero. - let rgb: Vec<[f32; 3]> = (0..16).map(|i| [i as f32 / 16.0, 0.0, 1.0]).collect(); - let c = crop(&rgb, 4, 4, &Roi { center: [2.0, 2.0], size: 4.0, angle: 0.0 }, 4, [0.0, 1.0]); + fn crop_follows_mediapipe_corner_aligned_sampling() { + // A 4×4 ramp cropped 1:1 reproduces it: output pixel i is at u = i/n, a source pixel centre. + let rgb: Vec<[f32; 3]> = (0..16).map(|i| [(i * 16) as f32 / 255.0, 0.0, 1.0]).collect(); + let at = |cx: f32| crop(&rgb, 4, 4, &Roi { center: [cx, 2.0], size: 4.0, angle: 0.0 }, 4, [0.0, 255.0], Border::Zero)[0]; + let c = crop(&rgb, 4, 4, &Roi { center: [2.0, 2.0], size: 4.0, angle: 0.0 }, 4, [0.0, 1.0], Border::Zero); for i in 0..16 { - assert!((c[i * 3] - i as f32 / 16.0).abs() < 1e-6); + assert!((c[i * 3] - (i * 16) as f32 / 255.0).abs() < 1e-6); } - let c = crop(&rgb, 4, 4, &Roi { center: [-10.0, -10.0], size: 4.0, angle: 0.0 }, 2, [-1.0, 1.0]); + // Half a pixel to the right is halfway between ramp values 0 and 16. + assert!((at(2.5) - 8.0).abs() < 1e-3, "{}", at(2.5)); + // The result is rounded to a whole 0-255 value: a third of the way from 0 to 16 is 5.33. + assert!((at(2.0 + 1.0 / 3.0) - 5.0).abs() < 1e-3, "{}", at(2.0 + 1.0 / 3.0)); + } + + #[test] + fn crop_borders_are_zero_or_replicated() { + let rgb: Vec<[f32; 3]> = (0..16).map(|i| [(i * 16) as f32 / 255.0, 0.0, 1.0]).collect(); + let far = Roi { center: [-10.0, -10.0], size: 4.0, angle: 0.0 }; + let c = crop(&rgb, 4, 4, &far, 2, [-1.0, 1.0], Border::Zero); assert!(c.iter().all(|v| *v == -1.0)); + // Replicated, the top-left pixel (0, 0, 255) fills the crop. + let c = crop(&rgb, 4, 4, &far, 2, [0.0, 255.0], Border::Replicate); + for p in c.chunks(3) { + assert!(p[0].abs() < 1e-3 && p[1].abs() < 1e-3 && (p[2] - 255.0).abs() < 1e-3, "{p:?}"); + } + // Half outside on the left: zero blends toward black, replicate keeps the edge value. + let edge = Roi { center: [0.0, 2.0], size: 4.0, angle: 0.0 }; + let z = crop(&rgb, 4, 4, &edge, 4, [0.0, 255.0], Border::Zero); + let r = crop(&rgb, 4, 4, &edge, 4, [0.0, 255.0], Border::Replicate); + assert!(z[2] == 0.0 && (r[2] - 255.0).abs() < 1e-3, "{} {}", z[2], r[2]); + } + + #[test] + fn crop_survives_absurd_frame_sizes() { + let rgb = [[0.5f32; 3]; 4]; + let roi = Roi { center: [1e30, -1e30], size: 1e30, angle: 1.0 }; + for (w, h) in [(usize::MAX, 2), (2, usize::MAX), (usize::MAX, usize::MAX), (1 << 40, 1 << 40), (0, 4), (4, 0)] { + for border in [Border::Zero, Border::Replicate] { + let c = crop(&rgb, w, h, &roi, 4, [-1.0, 1.0], border); + assert_eq!(c.len(), 48); + assert!(c.iter().all(|v| *v == -1.0), "{w}x{h}: an unreadable frame gives `lo`"); + } + } } } diff --git a/crates/segment/src/mediapipe.rs b/crates/segment/src/mediapipe.rs index 5798bc05..3c9e9ebf 100644 --- a/crates/segment/src/mediapipe.rs +++ b/crates/segment/src/mediapipe.rs @@ -9,7 +9,7 @@ //! then mesh points 33 → 263), and reports a presence score. The landmark indices in //! [`TOPOLOGY`] are points of MediaPipe's canonical face mesh. -use crate::face::{Face, FaceModel, Roi, Topology, crop, sigmoid}; +use crate::face::{Border, Face, FaceModel, Roi, Topology, crop, sigmoid}; use crate::tflite::Model; use crate::{FACE_LANDMARKER, ModelInfo, Result}; @@ -166,7 +166,7 @@ impl FaceLandmarker { /// Faces in the square `roi` of a frame (frame pixels), best first. pub fn detect(&self, rgb: &[[f32; 3]], w: usize, h: usize, roi: &Roi) -> Result> { - let x = crop(rgb, w, h, roi, DET_SIZE, [-1.0, 1.0]); + let x = crop(rgb, w, h, roi, DET_SIZE, [-1.0, 1.0], Border::Zero); let out = self.detector.run(&[&x])?; let (Some(reg), Some(cls)) = (out.first(), out.get(1)) else { return Err("face detector: missing outputs".into()) }; let s = DET_SIZE as f32; @@ -192,7 +192,7 @@ impl FaceLandmarker { /// The face mesh in `roi`: 478 points (frame pixels) and the presence score. pub fn mesh(&self, rgb: &[[f32; 3]], w: usize, h: usize, roi: &Roi) -> Result { - let x = crop(rgb, w, h, roi, MESH_SIZE, [0.0, 1.0]); + let x = crop(rgb, w, h, roi, MESH_SIZE, [0.0, 1.0], Border::Replicate); let out = self.mesh.run(&[&x])?; let (Some(pts), Some(flag)) = (out.first(), out.get(1).and_then(|f| f.first())) else { return Err("face mesh: missing outputs".into()) }; let s = MESH_SIZE as f32;