Files
DarkRoom/core/dr-face/src/embedding.rs
T
dtourolle 84fade99ec Put the developer docs under docs/dev and index the folder for users first
docs/ had 26 developer documents flat beside the manual, and the two
audiences are very differently sized: most readers want the manual and
the gesture reference, a few want the register, the designs and the
measurements. The manual and gestures.md stay at the top; everything for
someone changing the code moves to docs/dev/, and the two documents that
name their own successors — the v0.1 milestone and the UI-refinement plan
— go to docs/dev/archive/ rather than being deleted, since both are still
cited. docs/README.md is the index, users first.

Every reference follows: code comments, Cargo manifests, the workflows,
the pre-commit hook, the bench and traceability tools (which locate the
repo root by docs/dev/requirements.md now), packaging, the Docker READMEs,
CLAUDE.md, CONTRIBUTING.md and the README. The matrix links one level
deeper and is regenerated. Links out of the moved documents into the tree
gain a level; a link checker over every Markdown file finds none broken.
2026-09-20 21:16:03 +02:00

342 lines
13 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! What an embedder produces, and how it is stored (docs/dev/faces.md §6).
//!
//! Deliberately **model-free**: the vector, its identity, its comparison and
//! its storage encoding are arithmetic, and `calibrate` and `cluster` are built
//! on them. Keeping them out of the `inference` feature is what lets the part
//! of this subsystem most likely to be subtly wrong be tested on a machine with
//! no weights on it.
//!
//! [`crate::embed::Embedder`] is the thing that needs a model, and it lives
//! behind the feature.
/// Embedding dimensionality. Fixed by the model family, not a parameter.
pub const EMBEDDING_DIM: usize = 512;
/// The shortest raw embedding a face may be *compared against*.
///
/// # What the length of the vector says
///
/// ArcFace is trained on the direction of its output and nothing else, and
/// the length it leaves behind turns out to be a free quality signal: the
/// magnitude grows with how recognisable the crop was to the model, and a
/// blurred, occluded, badly lit or hard-profile face comes out short. MagFace
/// (Meng et al., CVPR 2021) made that the training objective; the plain
/// ArcFace heads this crate runs already show it, weaker but usable, which is
/// why it is worth keeping the number the normalisation discards.
///
/// # Why it gates the gallery and not the face
///
/// A short vector is a bad *reference*: it sits nearer the centre of the
/// sphere than a real identity does and matches a little of everyone, which
/// is exactly the face that welds two people together in a clustering pass.
/// It is not a bad *probe* — the face is still real, still somebody, and
/// comparing it against good references is the only way it will ever be named.
/// So a face below this floor is compared against the gallery and never
/// becomes part of it: see `cluster::Candidate::in_gallery`.
///
/// 14 is the operating point for `w600k_mbf`, whose norms on the reference
/// library run from about 8 on a blur to the high 20s on a clean portrait. A
/// face whose quality was never recorded — indexed before the number was kept
/// — is not gated, because a rule that cannot be checked should admit, not
/// exclude.
pub const MIN_GALLERY_QUALITY: f32 = 14.0;
/// Whether an embedding of this quality may serve as a reference.
///
/// `None` is "not measured", and is admitted: the rule is about a number that
/// was read and found short, not about a number that is missing.
pub fn in_gallery(quality: Option<f32>) -> bool {
quality.is_none_or(|q| q >= MIN_GALLERY_QUALITY)
}
/// Which model produced an embedding.
///
/// Embeddings from different models are not comparable, and this is the one
/// mistake that produces plausible-looking garbage rather than an error — so
/// the id travels *with* the vector rather than beside it, and
/// [`Embedding::cosine`] refuses a cross-model comparison.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ModelId(pub std::sync::Arc<str>);
impl ModelId {
pub fn new(s: impl Into<std::sync::Arc<str>>) -> Self {
Self(s.into())
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl std::fmt::Display for ModelId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
/// A 512-d L2-normalised face embedding.
#[derive(Debug, Clone, PartialEq)]
pub struct Embedding {
pub model: ModelId,
pub v: Box<[f32; EMBEDDING_DIM]>,
}
impl Embedding {
/// Cosine similarity, which for unit vectors is the plain dot product.
///
/// `None` when the two came from different models. That is a real
/// possibility in a library indexed across a model upgrade, and the
/// alternative — returning a number — is the failure mode
/// `faces.model_id` exists to prevent.
pub fn cosine(&self, other: &Embedding) -> Option<f32> {
if self.model != other.model {
return None;
}
Some(dot(&self.v, &other.v))
}
/// Storage form: `512 × f16`, 1 KB per face (catalog.md §10.1).
///
/// This writes the unit vector. What the catalog stores is the raw one —
/// `embed::Embedded::to_f16_bytes` — because the length is the quality
/// and a unit vector has none left to read.
pub fn to_f16_bytes(&self) -> Vec<u8> {
self.to_f16_bytes_scaled(1.0)
}
/// The unit vector scaled by `length`, as `512 × f16`.
pub(crate) fn to_f16_bytes_scaled(&self, length: f32) -> Vec<u8> {
let mut out = Vec::with_capacity(EMBEDDING_DIM * 2);
for &x in self.v.iter() {
out.extend_from_slice(&f32_to_f16_bits(x * length).to_le_bytes());
}
out
}
/// Read back from storage, re-normalising.
///
/// The f16 round-trip perturbs a unit vector by ~1e-3 in cosine — three
/// orders below the separation between a match and a non-match — but the
/// drift is free to remove and invisible if left, so it is removed here
/// rather than remembered at every call site. The same pass is what turns
/// a stored raw vector back into the unit one every comparison expects.
pub fn from_f16_bytes(model: ModelId, bytes: &[u8]) -> Option<Self> {
read_f16_bytes(model, bytes).map(|(e, _)| e)
}
}
/// Read a stored vector back, with the length it was stored at.
///
/// The length is the quality where the blob is a raw one, and ~1 where it is
/// a unit vector from before raw vectors were stored — which is why the
/// catalog keeps the quality beside the blob rather than deriving it from
/// this: a unit vector reads as a quality of 1, not as "unmeasured".
pub fn read_f16_bytes(model: ModelId, bytes: &[u8]) -> Option<(Embedding, f32)> {
if bytes.len() != EMBEDDING_DIM * 2 {
return None;
}
let mut v = Box::new([0.0_f32; EMBEDDING_DIM]);
for (i, chunk) in bytes.chunks_exact(2).enumerate() {
v[i] = f16_bits_to_f32(u16::from_le_bytes([chunk[0], chunk[1]]));
}
let length = normalise(&mut v);
Some((Embedding { model, v }, length))
}
fn dot(a: &[f32; EMBEDDING_DIM], b: &[f32; EMBEDDING_DIM]) -> f32 {
a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
}
/// Scale `v` to unit length, and return the length it had.
///
/// The length is the one thing about the raw output that survives being
/// thrown away by everything downstream, and it is a quality signal
/// ([`MIN_GALLERY_QUALITY`]) — so it comes back out rather than being lost
/// here.
pub(crate) fn normalise(v: &mut [f32; EMBEDDING_DIM]) -> f32 {
// Clamped rather than checked: a zero-norm embedding is a broken model,
// not a runtime condition worth an error path, and dividing by 1e-6 keeps
// the NaN out of the catalog.
let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt().max(1e-6);
for x in v.iter_mut() {
*x /= norm;
}
norm
}
// ── f16 ───────────────────────────────────────────────────────────────────
//
// Hand-rolled rather than pulling in `half`: two functions over a format that
// has not changed since 2008, used at exactly one boundary. The dependency
// policy (D13, D1) makes the bar for a new crate high, and this is well under
// it.
fn f32_to_f16_bits(x: f32) -> u16 {
let bits = x.to_bits();
let sign = ((bits >> 16) & 0x8000) as u16;
let exp = ((bits >> 23) & 0xff) as i32 - 127 + 15;
let mant = bits & 0x007f_ffff;
if exp >= 0x1f {
// Overflow, inf, or NaN. No component of an embedding exceeds its
// length, and the lengths this model produces are in the tens, so
// this is the broken-model path; infinity is the honest answer, not a
// clamp that hides it.
return sign
| 0x7c00
| if mant != 0 && exp == 0x1f + 112 {
0x200
} else {
0
};
}
if exp <= 0 {
// Subnormal or underflow. A component of a unit 512-vector is ~0.04
// and a stored one is that times the length, nowhere near here, so
// this branch exists for correctness rather than for traffic.
if exp < -10 {
return sign;
}
let mant = mant | 0x0080_0000;
let shift = (14 - exp) as u32;
let half = (mant >> shift) as u16;
// Round to nearest, ties to even.
let rem = mant & ((1 << shift) - 1);
let tie = 1 << (shift - 1);
let round = u16::from(rem > tie || (rem == tie && (half & 1) == 1));
return sign | (half + round);
}
let half = ((exp as u16) << 10) | (mant >> 13) as u16;
let rem = mant & 0x1fff;
let round = u16::from(rem > 0x1000 || (rem == 0x1000 && (half & 1) == 1));
sign | (half + round)
}
fn f16_bits_to_f32(h: u16) -> f32 {
let sign = ((h & 0x8000) as u32) << 16;
let exp = ((h >> 10) & 0x1f) as u32;
let mant = (h & 0x03ff) as u32;
if exp == 0 {
if mant == 0 {
return f32::from_bits(sign);
}
// Subnormal: renormalise into f32's range.
let mut e = -1_i32;
let mut m = mant;
while m & 0x0400 == 0 {
m <<= 1;
e -= 1;
}
let m = m & 0x03ff;
return f32::from_bits(sign | (((127 - 15 + 1 + e) as u32) << 23) | (m << 13));
}
if exp == 0x1f {
return f32::from_bits(sign | 0x7f80_0000 | (mant << 13));
}
f32::from_bits(sign | ((exp + 127 - 15) << 23) | (mant << 13))
}
#[cfg(test)]
mod tests {
use super::*;
fn unit(seed: u32) -> Embedding {
let mut v = Box::new([0.0_f32; EMBEDDING_DIM]);
let mut s = seed.wrapping_mul(2_654_435_761).wrapping_add(1);
for x in v.iter_mut() {
s = s.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
*x = (s >> 8) as f32 / (1u32 << 23) as f32 - 0.5;
}
normalise(&mut v);
Embedding {
model: ModelId::new("test"),
v,
}
}
#[test]
fn a_normalised_embedding_has_cosine_one_with_itself() {
let e = unit(7);
assert!((e.cosine(&e).unwrap() - 1.0).abs() < 1e-5);
}
#[test]
fn embeddings_from_different_models_do_not_compare() {
let a = unit(1);
let mut b = unit(1);
b.model = ModelId::new("other");
assert_eq!(
a.cosine(&b),
None,
"a cross-model cosine must not be a number"
);
}
/// The claim docs/dev/faces.md §6 makes about the storage format: the f16
/// round-trip costs ~1e-3 of cosine, three orders below the separation
/// between a match and a non-match.
#[test]
fn f16_round_trip_preserves_the_embedding() {
for seed in 0..16 {
let e = unit(seed);
let back = Embedding::from_f16_bytes(e.model.clone(), &e.to_f16_bytes()).unwrap();
let cos = e.cosine(&back).unwrap();
assert!(cos > 0.9999, "seed {seed}: round-trip cosine {cos}");
}
}
#[test]
fn normalising_reports_the_length_it_removed() {
let mut v = Box::new([0.0_f32; EMBEDDING_DIM]);
v[0] = 3.0;
v[1] = 4.0;
let norm = normalise(&mut v);
assert!((norm - 5.0).abs() < 1e-6, "norm {norm}");
assert!((v[0] - 0.6).abs() < 1e-6 && (v[1] - 0.8).abs() < 1e-6);
}
/// The gate admits what it cannot measure: a face from before the number
/// was kept is not a face that was found wanting.
#[test]
fn an_unmeasured_quality_is_admitted_to_the_gallery() {
assert!(in_gallery(None));
assert!(in_gallery(Some(MIN_GALLERY_QUALITY)));
assert!(in_gallery(Some(27.5)));
assert!(!in_gallery(Some(MIN_GALLERY_QUALITY - 0.01)));
assert!(!in_gallery(Some(8.0)));
}
/// The storage form carries the length, and the length comes back out —
/// without touching the direction every comparison is made on.
#[test]
fn a_raw_vector_round_trips_with_its_length() {
let e = unit(3);
let raw = e.to_f16_bytes_scaled(21.5);
let (back, length) = read_f16_bytes(e.model.clone(), &raw).unwrap();
assert!((length - 21.5).abs() < 0.05, "length {length}");
assert!(e.cosine(&back).unwrap() > 0.9999);
// A unit vector from an older store reads as length 1, not as an
// error — see `read_f16_bytes` on why that is not "unmeasured".
let (_, one) = read_f16_bytes(e.model.clone(), &e.to_f16_bytes()).unwrap();
assert!((one - 1.0).abs() < 1e-2, "length {one}");
}
#[test]
fn f16_round_trip_rejects_a_wrong_length_blob() {
assert!(Embedding::from_f16_bytes(ModelId::new("m"), &[0u8; 100]).is_none());
}
#[test]
fn f16_handles_the_values_an_embedding_actually_contains() {
// Components of a unit 512-vector cluster around ±1/sqrt(512) ≈ 0.044.
for &x in &[0.0_f32, 1.0, -1.0, 0.044_194_17, -0.044_194_17, 1e-3, -7e-4] {
let back = f16_bits_to_f32(f32_to_f16_bits(x));
assert!(
(back - x).abs() <= 1e-3 * x.abs().max(1e-3),
"{x} round-tripped to {back}"
);
}
}
}