98 lines
4.7 KiB
Rust
98 lines
4.7 KiB
Rust
//! The port against the reference implementation (Hugging Face `transformers`, fp32).
|
||
//!
|
||
//! Needs the checkpoint and a reference file made by `tools/sam3_reference.py`:
|
||
//! `LIGHTCRAFT_SAM3_DIR=<model dir> LIGHTCRAFT_SAM3_REF=<ref.safetensors> cargo test -p lightcraft-segment --release -- --nocapture`.
|
||
//! Skipped (passes) when either is missing.
|
||
|
||
use std::collections::HashMap;
|
||
use std::path::PathBuf;
|
||
|
||
use candle_core::{Device, IndexOp, Tensor};
|
||
use lightcraft_segment::{Click, Sam3, tokenizer::Tokenizer};
|
||
|
||
fn setup() -> Option<(PathBuf, HashMap<String, Tensor>)> {
|
||
let dir = PathBuf::from(std::env::var_os("LIGHTCRAFT_SAM3_DIR")?);
|
||
let reference = PathBuf::from(std::env::var_os("LIGHTCRAFT_SAM3_REF")?);
|
||
let r = candle_core::safetensors::load(&reference, &Device::Cpu).ok()?;
|
||
Some((dir, r))
|
||
}
|
||
|
||
fn vec(t: &Tensor) -> Vec<f32> {
|
||
t.to_device(&Device::Cpu).unwrap().flatten_all().unwrap().to_vec1::<f32>().unwrap()
|
||
}
|
||
|
||
/// Max |a − b| relative to the reference's max |b|.
|
||
fn rel_err(name: &str, a: &[f32], b: &[f32]) -> f32 {
|
||
assert_eq!(a.len(), b.len(), "{name}: length");
|
||
let scale = b.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
|
||
let err = a.iter().zip(b).fold(0f32, |m, (x, y)| m.max((x - y).abs()));
|
||
eprintln!("{name:>18}: max abs err {err:.2e}, rel {:.2e} (scale {scale:.2e})", err / scale);
|
||
err / scale
|
||
}
|
||
|
||
/// Fraction of pixels where the two logit maps disagree on the mask (> 0).
|
||
fn mask_disagreement(a: &[f32], b: &[f32]) -> f32 {
|
||
a.iter().zip(b).filter(|(x, y)| (**x > 0.0) != (**y > 0.0)).count() as f32 / a.len() as f32
|
||
}
|
||
|
||
#[test]
|
||
fn matches_reference() {
|
||
let Some((dir, r)) = setup() else {
|
||
eprintln!("skipped: set LIGHTCRAFT_SAM3_DIR and LIGHTCRAFT_SAM3_REF");
|
||
return;
|
||
};
|
||
let mut model = Sam3::load(&dir).unwrap();
|
||
eprintln!("device: {:?}", model.device());
|
||
let t = std::time::Instant::now();
|
||
let mut enc = model.encode_pixels(&r["pixel_values"]).unwrap();
|
||
let backbone = enc.backbone().permute((0, 2, 3, 1)).unwrap();
|
||
assert!(rel_err("backbone", &vec(&backbone), &vec(&r["backbone"])) < 1e-3);
|
||
eprintln!("backbone: {:?}", t.elapsed());
|
||
|
||
// points: two clicks (single-mask output), one click (best of three)
|
||
let clicks = [Click { x: 0.5, y: 0.55, positive: true }, Click { x: 0.2, y: 0.2, positive: false }];
|
||
let t = std::time::Instant::now();
|
||
let p = model.segment_clicks(&mut enc, &clicks).unwrap();
|
||
eprintln!("clicks (incl. tracker load + features): {:?}", t.elapsed());
|
||
let want = vec(&r["trk_masks"]);
|
||
assert!(rel_err("2-click logits", &p.logits, &want) < 1e-2);
|
||
assert!(mask_disagreement(&p.logits, &want) < 1e-3);
|
||
let t = std::time::Instant::now();
|
||
let p = model.segment_clicks(&mut enc, &clicks[..1]).unwrap();
|
||
eprintln!("one more click: {:?}", t.elapsed());
|
||
let iou = vec(&r["trk_multi_iou"]);
|
||
let best = (0..3).max_by(|a, b| iou[*a].partial_cmp(&iou[*b]).unwrap()).unwrap();
|
||
let want = vec(&r["trk_multi_masks"].i((0, 0, best)).unwrap());
|
||
assert!(rel_err("1-click logits", &p.logits, &want) < 1e-2);
|
||
assert!(mask_disagreement(&p.logits, &want) < 1e-3);
|
||
|
||
// tokenizer
|
||
let tok = Tokenizer::load(&dir).unwrap();
|
||
let cases = ["person", "Red Car", "a dog's toy", "sky, clouds!", "café 2 people", "tree-line"];
|
||
for (i, c) in cases.iter().enumerate() {
|
||
let (ids, _) = tok.encode(c);
|
||
let want: Vec<i64> = r[&format!("tok{i}")].flatten_all().unwrap().to_vec1().unwrap();
|
||
assert_eq!(ids.iter().map(|v| *v as i64).collect::<Vec<_>>(), want, "tokens of {c:?}");
|
||
}
|
||
|
||
// text prompts
|
||
for (i, text) in ["moon", "sky", "bird"].iter().enumerate() {
|
||
let t = std::time::Instant::now();
|
||
let out = model.detect(&mut enc, text).unwrap();
|
||
eprintln!("detect {text:?}: {:?}", t.elapsed());
|
||
assert!(rel_err("text features", &vec(&out.text), &vec(&r[&format!("det{i}_text")])) < 1e-3);
|
||
assert!(rel_err("query logits", &out.logits, &vec(&r[&format!("det{i}_logits")])) < 1e-2);
|
||
let presence = vec(&r[&format!("det{i}_presence")])[0];
|
||
assert!((out.presence - presence).abs() < 1e-2, "presence {} vs {presence}", out.presence);
|
||
let boxes: Vec<f32> = out.boxes.iter().flatten().copied().collect();
|
||
assert!(rel_err("boxes", &boxes, &vec(&r[&format!("det{i}_boxes")])) < 1e-2);
|
||
// the best instance's mask
|
||
let logits = vec(&r[&format!("det{i}_logits")]);
|
||
let q = (0..logits.len()).max_by(|a, b| logits[*a].partial_cmp(&logits[*b]).unwrap()).unwrap();
|
||
let ours = vec(&out.masks.i(q).unwrap());
|
||
let want = vec(&r[&format!("det{i}_masks")].i((0, q)).unwrap());
|
||
assert!(rel_err("best mask", &ours, &want) < 2e-2);
|
||
assert!(mask_disagreement(&ours, &want) < 2e-3);
|
||
}
|
||
}
|