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.
265 lines
9.7 KiB
Rust
265 lines
9.7 KiB
Rust
//! TRACES: FR-CULL-8a
|
||
//! The two small classifiers behind a face's eye state (docs/faces.md §17).
|
||
//!
|
||
//! **OCEC** — *open closed eyes classification*, Hyodo 2025 — reads one
|
||
//! 40×24 eye and answers P(open). **SGC** — *sunglasses classification*,
|
||
//! Hyodo 2026 — reads a 48×48 head and answers P(sunglasses); it is shown
|
||
//! two framings of each face and the higher answer stands, for the reason
|
||
//! [`crate::align::SUNGLASSES_WINDOWS`] gives. Both are
|
||
//! depthwise-separable CNNs of a few hundred kilobytes, both MIT with their
|
||
//! weights, and both were exported with BatchNorm already folded, which is
|
||
//! about the friendliest graph tract can be handed.
|
||
//!
|
||
//! Neither takes a plain buffer. [`EyeClassifier::classify`] takes an
|
||
//! [`EyePatch`] and [`SunglassesClassifier::classify`] a [`HeadViews`], each
|
||
//! constructible only by the crop in [`crate::align`] that puts the right
|
||
//! pixels in it — the same defence [`crate::embed::Embedder`] makes with
|
||
//! [`crate::align::Aligned112`], for the same reason: a classifier handed the
|
||
//! wrong region returns a confident probability of nothing. Where the eye
|
||
//! box comes from is [`crate::landmarks`]; [`EyeModels::read`] is the whole
|
||
//! chain.
|
||
//!
|
||
//! # The graphs must have a fixed batch
|
||
//!
|
||
//! Both ship with a dynamic batch dimension, which tract will not analyse.
|
||
//! `tools/fix-face-model-shapes.sh` pins it to 1, exactly as it does for the
|
||
//! embedder; the shipped files are the pinned ones.
|
||
//!
|
||
//! # Pre-processing
|
||
//!
|
||
//! Read off the reference demos rather than assumed: RGB, `x / 255`, NCHW,
|
||
//! the crop resized to the input with bilinear interpolation and **without**
|
||
//! preserving its aspect. [`crate::align`]'s crops arrive already at the
|
||
//! input size in `0..=1`, so there is nothing left to do but lay them out.
|
||
|
||
use ndarray::Array4;
|
||
|
||
use crate::align::{
|
||
eye_box, eye_patch, head_views, EyePatch, HeadViews, EYE_PATCH_HEIGHT, EYE_PATCH_WIDTH,
|
||
SUNGLASSES_EDGE,
|
||
};
|
||
use crate::eyes::{Eye, EyeReading};
|
||
use crate::landmarks::{Landmarker, Landmarks};
|
||
use crate::{install_backend, FaceError, Pixels};
|
||
|
||
/// A loaded OCEC graph.
|
||
pub struct EyeClassifier {
|
||
session: ort::session::Session,
|
||
}
|
||
|
||
/// A loaded SGC graph.
|
||
pub struct SunglassesClassifier {
|
||
session: ort::session::Session,
|
||
}
|
||
|
||
/// Open a single-input, single-output classifier and check it is the shape
|
||
/// the crop feeding it will be.
|
||
///
|
||
/// The check is against the *input*, because that is where these two graphs
|
||
/// differ from each other and from everything else in this crate: an SGC file
|
||
/// given to the eye classifier would otherwise be resized into by an eye
|
||
/// patch, and answer. `expected` names the model in the error.
|
||
fn open_classifier(
|
||
bytes: &[u8],
|
||
expected: &'static str,
|
||
(h, w): (usize, usize),
|
||
) -> Result<ort::session::Session, 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,
|
||
detail: "model has no inputs".into(),
|
||
})?;
|
||
let shape: Option<Vec<i64>> = input.dtype().tensor_shape().map(|s| s.to_vec());
|
||
let want = [1, 3, h as i64, w as i64];
|
||
if shape.as_deref() != Some(&want[..]) {
|
||
return Err(FaceError::WrongModel {
|
||
expected,
|
||
detail: format!(
|
||
"input '{}' is {:?}, expected {:?} (batch pinned to 1)",
|
||
input.name(),
|
||
shape,
|
||
want
|
||
),
|
||
});
|
||
}
|
||
if session.outputs().len() != 1 {
|
||
return Err(FaceError::WrongModel {
|
||
expected,
|
||
detail: format!("{} outputs, expected one", session.outputs().len()),
|
||
});
|
||
}
|
||
Ok(session)
|
||
}
|
||
|
||
/// Lay a `h × w` RGB crop out as the `[1, 3, h, w]` tensor both graphs take.
|
||
fn to_nchw(pixels: &[f32], h: usize, w: usize) -> Array4<f32> {
|
||
let mut input = Array4::<f32>::zeros((1, 3, h, w));
|
||
for y in 0..h {
|
||
for x in 0..w {
|
||
for c in 0..3 {
|
||
input[[0, c, y, x]] = pixels[(y * w + x) * 3 + c];
|
||
}
|
||
}
|
||
}
|
||
input
|
||
}
|
||
|
||
/// Run a one-number classifier and read its sigmoid back, clamped.
|
||
fn run_scalar(
|
||
session: &mut ort::session::Session,
|
||
input: Array4<f32>,
|
||
expected: &'static str,
|
||
) -> Result<f32, FaceError> {
|
||
let outputs = 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)?;
|
||
let Some(&p) = data.first() else {
|
||
return Err(FaceError::WrongModel {
|
||
expected,
|
||
detail: "empty output".into(),
|
||
});
|
||
};
|
||
// The graph ends in a sigmoid, so this is a clamp against rounding and
|
||
// nothing more — the reference demo does the same.
|
||
Ok(p.clamp(0.0, 1.0))
|
||
}
|
||
|
||
impl EyeClassifier {
|
||
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> {
|
||
Ok(Self {
|
||
session: open_classifier(bytes, "OCEC", (EYE_PATCH_HEIGHT, EYE_PATCH_WIDTH))?,
|
||
})
|
||
}
|
||
|
||
/// P(open) for one eye.
|
||
pub fn classify(&mut self, eye: &EyePatch) -> Result<f32, FaceError> {
|
||
let input = to_nchw(eye.pixels(), EYE_PATCH_HEIGHT, EYE_PATCH_WIDTH);
|
||
run_scalar(&mut self.session, input, "OCEC")
|
||
}
|
||
}
|
||
|
||
impl SunglassesClassifier {
|
||
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> {
|
||
Ok(Self {
|
||
session: open_classifier(bytes, "SGC", (SUNGLASSES_EDGE, SUNGLASSES_EDGE))?,
|
||
})
|
||
}
|
||
|
||
/// P(sunglasses) for one head: the highest answer over its framings.
|
||
pub fn classify(&mut self, head: &HeadViews) -> Result<f32, FaceError> {
|
||
let mut best = 0.0_f32;
|
||
for view in head.views() {
|
||
let input = to_nchw(view, SUNGLASSES_EDGE, SUNGLASSES_EDGE);
|
||
best = best.max(run_scalar(&mut self.session, input, "SGC")?);
|
||
}
|
||
Ok(best)
|
||
}
|
||
}
|
||
|
||
/// The three models behind a reading, which is how every caller holds them.
|
||
///
|
||
/// One struct rather than three optional parameters, because a partial
|
||
/// reading is not a reading: an eye state with no sunglasses number behind
|
||
/// it is exactly the beach-photograph failure [`crate::eyes`] describes, and
|
||
/// an eye box without the landmarks is the loose one this module replaced.
|
||
/// The models load together or not at all.
|
||
pub struct EyeModels {
|
||
pub landmarks: Landmarker,
|
||
pub eyes: EyeClassifier,
|
||
pub sunglasses: SunglassesClassifier,
|
||
}
|
||
|
||
impl EyeModels {
|
||
pub fn from_paths(
|
||
landmarks: impl AsRef<std::path::Path>,
|
||
eyes: impl AsRef<std::path::Path>,
|
||
sunglasses: impl AsRef<std::path::Path>,
|
||
) -> Result<Self, FaceError> {
|
||
Ok(Self {
|
||
landmarks: Landmarker::from_path(landmarks)?,
|
||
eyes: EyeClassifier::from_path(eyes)?,
|
||
sunglasses: SunglassesClassifier::from_path(sunglasses)?,
|
||
})
|
||
}
|
||
|
||
/// Read one face's eyes, and hand back the dense landmarks it read them
|
||
/// from.
|
||
///
|
||
/// `bbox` is the detector's `(x0, y0, x1, y1)` and `landmarks5` its five
|
||
/// points, both in source pixels; the buffer is the one the aligned
|
||
/// crop was taken from, so an eye is read from the same pixels the
|
||
/// embedder saw the face in. `None` where nothing could be cut — a
|
||
/// degenerate box or landmarks — which the caller stores as "not read".
|
||
///
|
||
/// The landmarks come back because they cost a model run the caller will
|
||
/// not want to pay twice: stored beside the reading, a later pass over
|
||
/// faces — head pose, expression — has them without the original.
|
||
pub fn read(
|
||
&mut self,
|
||
px: Pixels<'_>,
|
||
width: usize,
|
||
height: usize,
|
||
bbox: (f32, f32, f32, f32),
|
||
landmarks5: &[(f32, f32); 5],
|
||
) -> Result<Option<(EyeReading, Landmarks)>, FaceError> {
|
||
let Some(lm) = self.landmarks.landmarks(px, width, height, bbox)? else {
|
||
return Ok(None);
|
||
};
|
||
let Some(head) = head_views(px, width, height, landmarks5) else {
|
||
return Ok(None);
|
||
};
|
||
let mut eye = |contour: &[(f32, f32)]| -> Result<Eye, FaceError> {
|
||
// A hidden eye's contour can collapse to no width. Its numbers
|
||
// are then zero — no pixels, no sharpness — which is what the
|
||
// rule in `crate::eyes` reads as "not readable".
|
||
let Some(b) = eye_box(contour) else {
|
||
return Ok(Eye {
|
||
open: 0.0,
|
||
px: 0.0,
|
||
sharpness: 0.0,
|
||
});
|
||
};
|
||
let Some(patch) = eye_patch(px, width, height, b) else {
|
||
return Ok(Eye {
|
||
open: 0.0,
|
||
px: 0.0,
|
||
sharpness: 0.0,
|
||
});
|
||
};
|
||
Ok(Eye {
|
||
open: self.eyes.classify(&patch)?,
|
||
px: patch.source_px(),
|
||
sharpness: patch.sharpness(),
|
||
})
|
||
};
|
||
let right = eye(&lm.right_eye())?;
|
||
let left = eye(&lm.left_eye())?;
|
||
let reading = EyeReading {
|
||
right,
|
||
left,
|
||
sunglasses: self.sunglasses.classify(&head)?,
|
||
};
|
||
Ok(Some((reading, lm)))
|
||
}
|
||
}
|