shifathosn
/
effectcraft
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #388 from KazW/fix/mediapipe-crop-sampling 

segment: sample crops like MediaPipe's ImageToTensor; no overflow in crop
This commit is contained in:
Brandon Thomas 2026-10-09 09:25:23 -04:00
parent d3684216e4 5304437629
commit a3cfc3e014
2 changed files with 76 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

88
crates/segment/src/face.rs
View File

@ -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<f32> {
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<f32> {
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`");
}
}
}
}

6
crates/segment/src/mediapipe.rs
View File

@ -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<Vec<Detection>> {
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<Face> {
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;