dtourolle
facb44cb55
The 106 points the eye boxes were cut from, stored beside the reading as 16-bit fixed point over the frame: 424 bytes a face, a seventh of a pixel on a 6000-pixel frame, where f16 at the same size would have been six. Derived data like the embedding, kept for the same reason — it cost a fetch and a model run, and the next per-face pass should run from the catalog. Shards carry it; a peer's shard from before it is still read.
262 lines
9.9 KiB
Rust
262 lines
9.9 KiB
Rust
//! TRACES: FR-CULL-8a
|
||
//! Dense facial landmarks — InsightFace's `2d106det` (docs/faces.md §17.2).
|
||
//!
|
||
//! SCRFD's five points place a face; they do not place an eye. Its eye
|
||
//! point is loose enough that a window centred on it left the eye in a
|
||
//! corner on turned and smiling heads, and two model-free ways of
|
||
//! re-centring it made things worse. So a second model draws the eye's lid
|
||
//! contour, and the eye box is cut from that.
|
||
//!
|
||
//! **Why this one.** Three were measured on the same faces — MediaPipe Face
|
||
//! Mesh V2, PIPNet and this — and tied on what the eye classifier made of
|
||
//! their boxes (22 of 25 open eyes read open, against 19 from the SCRFD
|
||
//! point). This is the cheapest of the three by a wide margin (5 MB, 106
|
||
//! points, ~24 ms in tract), and it is under the grant the detector and
|
||
//! embedder already carry rather than a new one to read.
|
||
//!
|
||
//! # Pre-processing
|
||
//!
|
||
//! Ported from InsightFace's `landmark.py`: a square crop centred on the
|
||
//! detector box, 1.5× its longer edge, resized to 192; **RGB in 0..255**
|
||
//! (the graph carries its own `bn_data` normalisation, so `input_mean` is
|
||
//! 0 and `input_std` 1); 106 `(x, y)` in −1..1 mapped back through
|
||
//! `(p + 1) · 96`. The graph's batch dimension is the literal `None` and
|
||
//! is pinned to 1 by `tools/fix-face-model-shapes.sh`, like the embedder's.
|
||
//!
|
||
//! # The layout
|
||
//!
|
||
//! Checked by drawing the points on the reference faces rather than taken
|
||
//! from a diagram: the subject's right eye (image-left) is points 33–42,
|
||
//! the left 87–96, ten each round the lids.
|
||
|
||
use ndarray::Array4;
|
||
|
||
use crate::align::crop_box;
|
||
use crate::{install_backend, FaceError, Pixels};
|
||
|
||
/// The graph's input edge, in pixels.
|
||
pub const INPUT_EDGE: usize = 192;
|
||
|
||
/// How many points the model returns.
|
||
pub const POINTS: usize = 106;
|
||
|
||
/// The crop's edge as a multiple of the detector box's longer edge.
|
||
const CROP_SCALE: f32 = 1.5;
|
||
|
||
/// The span of the frame, in long-edge units, the packed form covers: a
|
||
/// quarter of the frame outside each edge.
|
||
pub const PACKED_RANGE: (f32, f32) = (-0.25, 1.25);
|
||
|
||
/// Bytes the packed form of one face's landmarks takes.
|
||
pub const PACKED_BYTES: usize = POINTS * 4;
|
||
|
||
/// Point indices of the subject's right eye's lid contour (image-left).
|
||
pub const RIGHT_EYE: [usize; 10] = [33, 34, 35, 36, 37, 38, 39, 40, 41, 42];
|
||
/// Point indices of the subject's left eye's lid contour (image-right).
|
||
pub const LEFT_EYE: [usize; 10] = [87, 88, 89, 90, 91, 92, 93, 94, 95, 96];
|
||
|
||
/// The 106 points of one face, in **source pixels**.
|
||
#[derive(Debug, Clone, PartialEq)]
|
||
pub struct Landmarks {
|
||
pub points: [(f32, f32); POINTS],
|
||
}
|
||
|
||
impl Landmarks {
|
||
/// Storage form: `106 × (x, y)` as little-endian **`u16` fixed point**
|
||
/// over the frame, 424 bytes.
|
||
///
|
||
/// Each coordinate is normalised by `long_edge` like the five points the
|
||
/// catalog already keeps, then mapped over [`PACKED_RANGE`] — a quarter
|
||
/// of the frame either side of it, because a landmark on a face at the
|
||
/// edge does land outside the image — onto 0..65535. That is 0.14 source
|
||
/// pixels on a 6000-pixel frame. `f16` would be the same size and worse:
|
||
/// its three significant figures near 1.0 are six pixels at that scale,
|
||
/// and the eye contour this is kept for is drawn to the pixel.
|
||
pub fn to_packed_bytes(&self, long_edge: f32) -> Vec<u8> {
|
||
let (lo, hi) = PACKED_RANGE;
|
||
let pack = |v: f32| -> [u8; 2] {
|
||
let t = ((v / long_edge - lo) / (hi - lo)).clamp(0.0, 1.0);
|
||
((t * 65535.0).round() as u16).to_le_bytes()
|
||
};
|
||
let mut out = Vec::with_capacity(POINTS * 4);
|
||
for &(x, y) in &self.points {
|
||
out.extend_from_slice(&pack(x));
|
||
out.extend_from_slice(&pack(y));
|
||
}
|
||
out
|
||
}
|
||
|
||
/// [`Self::to_packed_bytes`] read back, into source pixels of a frame
|
||
/// with this `long_edge`. `None` for a blob of the wrong length.
|
||
pub fn from_packed_bytes(bytes: &[u8], long_edge: f32) -> Option<Self> {
|
||
if bytes.len() != POINTS * 4 {
|
||
return None;
|
||
}
|
||
let (lo, hi) = PACKED_RANGE;
|
||
let unpack = |b: &[u8]| -> f32 {
|
||
let t = u16::from_le_bytes([b[0], b[1]]) as f32 / 65535.0;
|
||
(t * (hi - lo) + lo) * long_edge
|
||
};
|
||
let mut points = [(0.0_f32, 0.0_f32); POINTS];
|
||
for (i, p) in points.iter_mut().enumerate() {
|
||
let at = i * 4;
|
||
*p = (unpack(&bytes[at..at + 2]), unpack(&bytes[at + 2..at + 4]));
|
||
}
|
||
Some(Self { points })
|
||
}
|
||
|
||
/// The lid contour of the subject's right eye.
|
||
pub fn right_eye(&self) -> [(f32, f32); 10] {
|
||
RIGHT_EYE.map(|i| self.points[i])
|
||
}
|
||
|
||
/// The lid contour of the subject's left eye.
|
||
pub fn left_eye(&self) -> [(f32, f32); 10] {
|
||
LEFT_EYE.map(|i| self.points[i])
|
||
}
|
||
}
|
||
|
||
/// A loaded `2d106det` graph.
|
||
pub struct Landmarker {
|
||
session: ort::session::Session,
|
||
}
|
||
|
||
impl Landmarker {
|
||
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 input = session.inputs().first().ok_or(FaceError::WrongModel {
|
||
expected: "2d106det",
|
||
detail: "model has no inputs".into(),
|
||
})?;
|
||
let shape: Option<Vec<i64>> = input.dtype().tensor_shape().map(|s| s.to_vec());
|
||
let want = [1, 3, INPUT_EDGE as i64, INPUT_EDGE as i64];
|
||
if shape.as_deref() != Some(&want[..]) {
|
||
return Err(FaceError::WrongModel {
|
||
expected: "2d106det",
|
||
detail: format!(
|
||
"input '{}' is {:?}, expected {:?} (batch pinned to 1)",
|
||
input.name(),
|
||
shape,
|
||
want
|
||
),
|
||
});
|
||
}
|
||
let out = session.outputs().first().ok_or(FaceError::WrongModel {
|
||
expected: "2d106det",
|
||
detail: "model has no outputs".into(),
|
||
})?;
|
||
let last: Option<i64> = out.dtype().tensor_shape().and_then(|d| d.last().copied());
|
||
if last != Some((POINTS * 2) as i64) {
|
||
return Err(FaceError::WrongModel {
|
||
expected: "2d106det",
|
||
detail: format!(
|
||
"output '{}' is {:?}-wide, expected {}",
|
||
out.name(),
|
||
last,
|
||
POINTS * 2
|
||
),
|
||
});
|
||
}
|
||
Ok(Self { session })
|
||
}
|
||
|
||
/// The landmarks of the face in `bbox` — `(x0, y0, x1, y1)` in source
|
||
/// pixels, the detector's box — read from the source.
|
||
///
|
||
/// `None` for a box with no area or a buffer that is not the size it
|
||
/// claims, as every crop here.
|
||
pub fn landmarks(
|
||
&mut self,
|
||
px: Pixels<'_>,
|
||
width: usize,
|
||
height: usize,
|
||
bbox: (f32, f32, f32, f32),
|
||
) -> Result<Option<Landmarks>, FaceError> {
|
||
let (w, h) = (bbox.2 - bbox.0, bbox.3 - bbox.1);
|
||
let side = w.max(h) * CROP_SCALE;
|
||
let (cx, cy) = ((bbox.0 + bbox.2) / 2.0, (bbox.1 + bbox.3) / 2.0);
|
||
let (x0, y0) = (cx - side / 2.0, cy - side / 2.0);
|
||
let Some(crop) = crop_box(
|
||
px,
|
||
width,
|
||
height,
|
||
(x0, y0, side, side),
|
||
INPUT_EDGE,
|
||
INPUT_EDGE,
|
||
) else {
|
||
return Ok(None);
|
||
};
|
||
|
||
let e = INPUT_EDGE;
|
||
let mut input = Array4::<f32>::zeros((1, 3, e, e));
|
||
for y in 0..e {
|
||
for x in 0..e {
|
||
for c in 0..3 {
|
||
input[[0, c, y, x]] = crop[(y * e + x) * 3 + c] * 255.0;
|
||
}
|
||
}
|
||
}
|
||
let outputs = self
|
||
.session
|
||
.run(ort::inputs![
|
||
ort::value::Tensor::from_array(input).map_err(FaceError::Inference)?
|
||
])
|
||
.map_err(FaceError::Inference)?;
|
||
let (_, data) = outputs[0]
|
||
.try_extract_tensor::<f32>()
|
||
.map_err(FaceError::Inference)?;
|
||
if data.len() < POINTS * 2 {
|
||
return Err(FaceError::WrongModel {
|
||
expected: "2d106det",
|
||
detail: format!("got {} values, expected {}", data.len(), POINTS * 2),
|
||
});
|
||
}
|
||
|
||
// −1..1 in the crop → crop pixels → source pixels.
|
||
let scale = side / e as f32;
|
||
let half = e as f32 / 2.0;
|
||
let mut points = [(0.0_f32, 0.0_f32); POINTS];
|
||
for (i, p) in points.iter_mut().enumerate() {
|
||
let (u, v) = ((data[2 * i] + 1.0) * half, (data[2 * i + 1] + 1.0) * half);
|
||
*p = (x0 + u * scale, y0 + v * scale);
|
||
}
|
||
Ok(Some(Landmarks { points }))
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
/// Packed and unpacked, every point comes back within a fifth of a
|
||
/// source pixel on a 6000-pixel frame — including one outside the
|
||
/// image, which a face at the edge does produce.
|
||
#[test]
|
||
fn dense_landmarks_round_trip_through_their_packed_bytes() {
|
||
let mut points = [(0.0_f32, 0.0_f32); POINTS];
|
||
for (i, p) in points.iter_mut().enumerate() {
|
||
*p = (i as f32 * 37.3 - 200.0, 5900.0 - i as f32 * 11.1);
|
||
}
|
||
let lm = Landmarks { points };
|
||
let bytes = lm.to_packed_bytes(6000.0);
|
||
assert_eq!(bytes.len(), PACKED_BYTES);
|
||
assert_eq!(PACKED_BYTES, 424);
|
||
let back = Landmarks::from_packed_bytes(&bytes, 6000.0).unwrap();
|
||
for (a, b) in lm.points.iter().zip(back.points.iter()) {
|
||
assert!((a.0 - b.0).abs() < 0.2, "{} vs {}", a.0, b.0);
|
||
assert!((a.1 - b.1).abs() < 0.2, "{} vs {}", a.1, b.1);
|
||
}
|
||
assert!(Landmarks::from_packed_bytes(&bytes[..100], 6000.0).is_none());
|
||
}
|
||
}
|