mirror of
https://github.com/storytold/lightcraft.git
synced 2026-10-24 18:18:12 -06:00
Daniel Voshart
24f87a91a0
Keep current upstream UI behavior and the later detector contract checks. Validate candidate models before replacing working weights and records, with a synthetic regression that failed before the fix. Remove the unrelated GPU equivalence-test relaxation. All seven quality gates pass with the pinned hosted fonts. Generated with Codex.
353 lines
15 KiB
Rust
353 lines
15 KiB
Rust
//! Installed recognition models on LightCraft's own checked, single-thread float32 CPU engine.
|
|
//! The detector and recognisers share graph validation and convolution kernels.
|
|
|
|
use std::panic::{AssertUnwindSafe, catch_unwind};
|
|
use std::path::Path;
|
|
use std::sync::Arc;
|
|
use web_time::Instant;
|
|
|
|
use crate::{
|
|
graph,
|
|
net::{Net, Tensor},
|
|
};
|
|
use serde::Serialize;
|
|
|
|
use crate::manifest::{Colour, ModelManifest, OutputSpec};
|
|
|
|
#[derive(Debug, thiserror::Error)]
|
|
pub enum RuntimeError {
|
|
#[error("the model could not be loaded: {0}")]
|
|
Load(String),
|
|
#[error("the model failed while running: {0}")]
|
|
Run(String),
|
|
#[error("{0}")]
|
|
Input(String),
|
|
}
|
|
|
|
/// What `self_test` found.
|
|
#[derive(Clone, Debug, PartialEq, Serialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub struct SelfTest {
|
|
pub ok: bool,
|
|
/// Milliseconds to load and optimise the model.
|
|
pub load_ms: f64,
|
|
/// Median milliseconds to embed one face.
|
|
pub embed_ms: f64,
|
|
pub dimension: usize,
|
|
/// Each check: what was checked and whether it held.
|
|
pub checks: Vec<(String, bool)>,
|
|
}
|
|
|
|
type Plan = Net;
|
|
|
|
/// A loaded recognition model.
|
|
#[derive(Clone)]
|
|
pub struct Embedder {
|
|
plan: Arc<Plan>,
|
|
manifest: ModelManifest,
|
|
width: usize,
|
|
height: usize,
|
|
dim: usize,
|
|
load_ms: f64,
|
|
}
|
|
|
|
fn guarded<T>(what: &str, f: impl FnOnce() -> Result<T, RuntimeError>) -> Result<T, RuntimeError> {
|
|
match catch_unwind(AssertUnwindSafe(f)) {
|
|
Ok(r) => r,
|
|
Err(_) => Err(RuntimeError::Run(format!("the runtime gave up on this model while {what}"))),
|
|
}
|
|
}
|
|
|
|
impl Embedder {
|
|
/// Load the `.onnx` at `path` as described by `manifest` (an embedder).
|
|
pub fn load(path: &Path, manifest: &ModelManifest) -> Result<Embedder, RuntimeError> {
|
|
let OutputSpec::Embedding { dim } = manifest.output else {
|
|
return Err(RuntimeError::Load("this is not a recognition model".into()));
|
|
};
|
|
let (w, h) = (manifest.input.width as usize, manifest.input.height as usize);
|
|
let started = Instant::now();
|
|
crate::manifest::validate(manifest).map_err(|e| RuntimeError::Load(e.to_string()))?;
|
|
let mut file = std::fs::File::open(path).map_err(|e| RuntimeError::Load(e.to_string()))?;
|
|
let length = file.metadata().map_err(|e| RuntimeError::Load(e.to_string()))?.len();
|
|
if length == 0 || length > 512 * 1024 * 1024 {
|
|
return Err(RuntimeError::Load("the model file is empty or exceeds 512 MiB".into()));
|
|
}
|
|
let mut bytes = Vec::new();
|
|
bytes
|
|
.try_reserve_exact(usize::try_from(length).map_err(|_| RuntimeError::Load("model size overflow".into()))?)
|
|
.map_err(|_| RuntimeError::Load("not enough memory for the model".into()))?;
|
|
use std::io::Read;
|
|
(&mut file).take(512 * 1024 * 1024 + 1).read_to_end(&mut bytes).map_err(|e| RuntimeError::Load(e.to_string()))?;
|
|
let graph = graph::load(&bytes).map_err(|e| RuntimeError::Load(e.to_string()))?;
|
|
let input = graph.inputs.first().ok_or_else(|| RuntimeError::Load("no input".into()))?;
|
|
let want = [1, 3, h, w];
|
|
if input.1.len() != 4 || input.1.iter().zip(want).any(|(&d, w)| d != 0 && d != w) {
|
|
return Err(RuntimeError::Load("the graph input differs from the manifest's [1,3,H,W]".into()));
|
|
}
|
|
let plan = Net::new(graph).map_err(|e| RuntimeError::Load(e.to_string()))?;
|
|
let shapes = plan.validate_input(&want).map_err(|e| RuntimeError::Load(e.to_string()))?;
|
|
if shapes.len() != 1 || shapes.first().is_none_or(|s| s.first() != Some(&1) || s.iter().product::<usize>() != dim as usize) {
|
|
return Err(RuntimeError::Load("the graph output differs from the manifest's embedding dimension".into()));
|
|
}
|
|
Ok(Embedder {
|
|
plan: Arc::new(plan),
|
|
manifest: manifest.clone(),
|
|
width: w,
|
|
height: h,
|
|
dim: dim as usize,
|
|
load_ms: started.elapsed().as_secs_f64() * 1000.0,
|
|
})
|
|
}
|
|
|
|
/// The size of the aligned face picture the model wants (width, height).
|
|
pub fn input_size(&self) -> (usize, usize) {
|
|
(self.width, self.height)
|
|
}
|
|
|
|
pub fn manifest(&self) -> &ModelManifest {
|
|
&self.manifest
|
|
}
|
|
|
|
/// The face as a unit-length vector. `rgb` is an aligned face at exactly [`Self::input_size`], 8-bit RGB.
|
|
pub fn embed(&self, rgb: &[u8]) -> Result<Vec<f32>, RuntimeError> {
|
|
if rgb.len() != self.width * self.height * 3 {
|
|
return Err(RuntimeError::Input("the face picture is not the size the model wants".into()));
|
|
}
|
|
let (mean, std, bgr) = (self.manifest.input.mean, self.manifest.input.std, self.manifest.input.colour == Colour::Bgr);
|
|
let plane = self.width * self.height;
|
|
let mut data = vec![0.0f32; 3 * plane];
|
|
for (i, px) in rgb.as_chunks::<3>().0.iter().enumerate() {
|
|
for c in 0..3 {
|
|
let src = if bgr { 2 - c } else { c };
|
|
if let (Some(v), Some(m), Some(s), Some(o)) = (px.get(src), mean.get(c), std.get(c), data.get_mut(c * plane + i)) {
|
|
*o = (f32::from(*v) - m) / s;
|
|
}
|
|
}
|
|
}
|
|
let out = guarded("running it", || {
|
|
let input = Tensor::new(vec![1, 3, self.height, self.width], data).map_err(|e| RuntimeError::Input(e.to_string()))?;
|
|
let mut result = self.plan.run(input).map_err(|e| RuntimeError::Run(e.to_string()))?;
|
|
if result.len() != 1 {
|
|
return Err(RuntimeError::Run("the model must give one embedding".into()));
|
|
}
|
|
Ok(result.pop().ok_or_else(|| RuntimeError::Run("no output".into()))?.1.data)
|
|
})?;
|
|
if out.len() != self.dim {
|
|
return Err(RuntimeError::Run(format!("the model gave {} numbers per face, not the {} its description says", out.len(), self.dim)));
|
|
}
|
|
if out.iter().any(|v| !v.is_finite()) {
|
|
return Err(RuntimeError::Run("the model gave numbers that are not finite".into()));
|
|
}
|
|
let norm = out.iter().map(|&v| f64::from(v).powi(2)).sum::<f64>().sqrt();
|
|
if !norm.is_finite() || norm < 1e-9 {
|
|
return Err(RuntimeError::Run("the model gave an all-zero face".into()));
|
|
}
|
|
Ok(out.into_iter().map(|v| (f64::from(v) / norm) as f32).collect())
|
|
}
|
|
|
|
/// Check the model does something sensible before it is ever used on a photo: the right number of finite
|
|
/// numbers, the same picture twice gives the same face, and two different pictures give different faces.
|
|
pub fn self_test(&self) -> SelfTest {
|
|
let mut checks: Vec<(String, bool)> = Vec::new();
|
|
let n = self.width * self.height * 3;
|
|
// two deterministic, different test pictures: noise, and a diagonal gradient with a colour cast
|
|
let mut seed = 0x2545_f491u32;
|
|
let noise: Vec<u8> = (0..n)
|
|
.map(|_| {
|
|
seed ^= seed << 13;
|
|
seed ^= seed >> 17;
|
|
seed ^= seed << 5;
|
|
(seed >> 8) as u8
|
|
})
|
|
.collect();
|
|
let gradient: Vec<u8> = (0..n)
|
|
.map(|i| (((i / 3) % self.width + (i / 3) / self.width) * 255 / (self.width + self.height) + (i % 3) * 40).min(255) as u8)
|
|
.collect();
|
|
let mut times = Vec::new();
|
|
let mut run = |img: &[u8]| {
|
|
let t = Instant::now();
|
|
let r = self.embed(img);
|
|
times.push(t.elapsed().as_secs_f64() * 1000.0);
|
|
r
|
|
};
|
|
let (a, a2, b) = (run(&noise), run(&noise), run(&gradient));
|
|
let dot = |x: &[f32], y: &[f32]| x.iter().zip(y).map(|(p, q)| p * q).sum::<f32>();
|
|
match (&a, &a2, &b) {
|
|
(Ok(a), Ok(a2), Ok(b)) => {
|
|
checks.push((format!("gives {} finite numbers per face", a.len()), true));
|
|
checks.push(("gives the same face for the same picture".into(), dot(a, a2) > 0.9999));
|
|
checks.push(("gives different faces for different pictures".into(), dot(a, b) < 0.9995));
|
|
}
|
|
(Err(e), _, _) | (_, Err(e), _) | (_, _, Err(e)) => checks.push((format!("runs: {e}"), false)),
|
|
}
|
|
times.sort_by(|x, y| x.total_cmp(y));
|
|
SelfTest {
|
|
ok: checks.iter().all(|(_, ok)| *ok),
|
|
load_ms: self.load_ms,
|
|
embed_ms: times.get(times.len() / 2).copied().unwrap_or(0.0),
|
|
dimension: self.dim,
|
|
checks,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::manifest::{InputSpec, Licence, OutputSpec, Role, Thresholds};
|
|
|
|
fn manifest(dim: u32) -> ModelManifest {
|
|
ModelManifest {
|
|
id: "tiny".into(),
|
|
name: "Tiny".into(),
|
|
version: "1".into(),
|
|
role: Role::Embedder,
|
|
licence: Licence::default(),
|
|
source: None,
|
|
sha256: None,
|
|
size_bytes: None,
|
|
provenance: String::new(),
|
|
input: InputSpec::default(),
|
|
output: OutputSpec::Embedding { dim },
|
|
thresholds: Thresholds::default(),
|
|
speed: None,
|
|
}
|
|
}
|
|
|
|
fn temp_model(name: &str, bytes: &[u8]) -> std::path::PathBuf {
|
|
let p = std::env::temp_dir().join(format!("lc-faces-runtime-{name}-{}.onnx", std::process::id()));
|
|
std::fs::write(&p, bytes).unwrap();
|
|
p
|
|
}
|
|
|
|
#[test]
|
|
fn a_working_model_loads_embeds_and_passes_the_self_test() {
|
|
let p = temp_model("ok", &crate::synthetic::tiny_embedder_model(16));
|
|
let e = Embedder::load(&p, &manifest(16)).unwrap();
|
|
assert_eq!(e.input_size(), (112, 112));
|
|
let v = e.embed(&vec![100u8; 112 * 112 * 3]).unwrap();
|
|
assert_eq!(v.len(), 16);
|
|
assert!((v.iter().map(|x| x * x).sum::<f32>() - 1.0).abs() < 1e-4, "unit length");
|
|
let t = e.self_test();
|
|
assert!(t.ok, "{t:?}");
|
|
assert_eq!(t.dimension, 16);
|
|
let _ = std::fs::remove_file(p);
|
|
}
|
|
|
|
#[test]
|
|
fn wrong_sizes_wrong_dimensions_and_broken_files_are_errors() {
|
|
let p = temp_model("dim", &crate::synthetic::tiny_embedder_model(16));
|
|
let e = Embedder::load(&p, &manifest(16)).unwrap();
|
|
assert!(e.embed(&[0u8; 10]).is_err(), "wrong picture size");
|
|
// the manifest promises 32 numbers, the model gives 16
|
|
assert!(Embedder::load(&p, &manifest(32)).is_err());
|
|
// not a model, an empty file, a graph with no layers, a detector-shaped manifest
|
|
for (name, bytes) in [("junk", b"nonsense".to_vec()), ("empty", Vec::new()), ("nolayers", crate::synthetic::embedder_model(16))] {
|
|
let q = temp_model(name, &bytes);
|
|
assert!(Embedder::load(&q, &manifest(16)).is_err(), "{name}");
|
|
let _ = std::fs::remove_file(q);
|
|
}
|
|
let mut detector = manifest(16);
|
|
detector.output = OutputSpec::Detector { decoder: "yunet-v2".into() };
|
|
assert!(Embedder::load(&p, &detector).is_err());
|
|
assert!(Embedder::load(&std::env::temp_dir().join("does-not-exist.onnx"), &manifest(16)).is_err());
|
|
let _ = std::fs::remove_file(p);
|
|
}
|
|
|
|
#[test]
|
|
fn truncated_and_corrupted_models_never_panic() {
|
|
let good = crate::synthetic::tiny_embedder_model(16);
|
|
for n in (0..good.len()).step_by(7) {
|
|
let q = temp_model("trunc", &good[..n]);
|
|
let _ = Embedder::load(&q, &manifest(16));
|
|
let _ = std::fs::remove_file(q);
|
|
}
|
|
for at in (0..good.len()).step_by(11) {
|
|
let mut bad = good.clone();
|
|
if let Some(b) = bad.get_mut(at) {
|
|
*b = b.wrapping_add(0x55);
|
|
}
|
|
let q = temp_model("flip", &bad);
|
|
if let Ok(e) = Embedder::load(&q, &manifest(16)) {
|
|
let _ = e.embed(&vec![7u8; 112 * 112 * 3]);
|
|
}
|
|
let _ = std::fs::remove_file(q);
|
|
}
|
|
}
|
|
|
|
/// Opt-in: `LC_FACE_MODELS=<folder> cargo test -p lightcraft-faces -- --ignored --nocapture`
|
|
/// self-tests every known recognition model found there.
|
|
#[test]
|
|
#[ignore = "needs real model files: set LC_FACE_MODELS"]
|
|
fn real_models_pass_the_self_test() {
|
|
let Some(dir) = std::env::var_os("LC_FACE_MODELS") else { return };
|
|
for entry in std::fs::read_dir(dir).unwrap().flatten() {
|
|
let path = entry.path();
|
|
let Ok(sha) = crate::hash::sha256_file(&path) else { continue };
|
|
let Some(m) = crate::known::lookup(&sha).filter(|m| m.role == Role::Embedder) else { continue };
|
|
let e = Embedder::load(&path, &m).unwrap();
|
|
let t = e.self_test();
|
|
println!("{}: ok {} | load {:.0} ms | {:.0} ms per face | {} numbers | {:?}", m.id, t.ok, t.load_ms, t.embed_ms, t.dimension, t.checks);
|
|
assert!(t.ok);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn bgr_changes_the_input() {
|
|
let p = temp_model("bgr", &crate::synthetic::tiny_embedder_model(16));
|
|
let rgb = Embedder::load(&p, &manifest(16)).unwrap();
|
|
let mut m = manifest(16);
|
|
m.input.colour = crate::manifest::Colour::Bgr;
|
|
let bgr = Embedder::load(&p, &m).unwrap();
|
|
// a picture with different channel levels: swapping the order must change the embedding
|
|
let img: Vec<u8> = (0..112 * 112).flat_map(|_| [200u8, 100, 20]).collect();
|
|
let (a, b) = (rgb.embed(&img).unwrap(), bgr.embed(&img).unwrap());
|
|
let cos: f32 = a.iter().zip(&b).map(|(x, y)| x * y).sum();
|
|
assert!(cos < 0.9999, "{cos}");
|
|
let _ = std::fs::remove_file(p);
|
|
}
|
|
#[test]
|
|
fn large_finite_embeddings_have_a_finite_unit_norm() {
|
|
let p = temp_model("large-norm", &crate::synthetic::norm_overflow_model());
|
|
let e = Embedder::load(&p, &manifest(3)).unwrap();
|
|
let out = e.embed(&vec![255; 112 * 112 * 3]).unwrap();
|
|
assert!(out.iter().all(|v| v.is_finite()));
|
|
assert!((out.iter().map(|v| v * v).sum::<f32>() - 1.0).abs() < 1e-6);
|
|
let _ = std::fs::remove_file(p);
|
|
}
|
|
#[test]
|
|
fn invalid_manifest_numbers_are_rejected_before_preprocessing() {
|
|
let p = temp_model("invalid-manifest", &crate::synthetic::tiny_embedder_model(16));
|
|
for case in 0..5 {
|
|
let mut m = manifest(16);
|
|
match case {
|
|
0 => m.input.std = [0.0; 3],
|
|
1 => m.input.mean = [f32::NAN; 3],
|
|
2 => m.input.width = 0,
|
|
3 => m.input.width = u32::MAX,
|
|
_ => m.output = OutputSpec::Embedding { dim: 0 },
|
|
};
|
|
assert!(Embedder::load(&p, &m).is_err());
|
|
}
|
|
let _ = std::fs::remove_file(p);
|
|
}
|
|
#[test]
|
|
fn concurrent_faces_share_weights_and_give_deterministic_embeddings() {
|
|
let p = temp_model("concurrent", &crate::synthetic::tiny_embedder_model(16));
|
|
let e = Embedder::load(&p, &manifest(16)).unwrap();
|
|
let input = vec![123; 112 * 112 * 3];
|
|
let want = e.embed(&input).unwrap();
|
|
let threads: Vec<_> = (0..8)
|
|
.map(|_| {
|
|
let e = e.clone();
|
|
let input = input.clone();
|
|
std::thread::spawn(move || e.embed(&input))
|
|
})
|
|
.collect();
|
|
for thread in threads {
|
|
assert_eq!(thread.join().unwrap().unwrap(), want);
|
|
}
|
|
let _ = std::fs::remove_file(p);
|
|
}
|
|
}
|