Files
effectcraft/crates/sam2/tests/reference.rs
T
Kristian MyramandClaude Sonnet 5.5 a62d59da4b S2.40: SAM 2 as the Roto Brush object tracker
ObjectTracker interface in the segment crate (prompts on some frames, the
clip tracked both ways, each frame streamed with confidence and an occluded
flag), implemented by SAM 2. MaskModel::tracker() is None for other models.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
2026-10-30 08:43:30 +01:00

400 lines
19 KiB
Rust

//! Stage-by-stage comparison with the PyTorch / transformers reference (`tools/sam2_reference.py`).
//!
//! Needs the `facebook/sam2.1-hiera-tiny` weights (default `<repo>/sam2-tiny`, or the folder in
//! `EFFECTCRAFT_SAM2_DIR`) and the fixture tensors (`<repo>/plan/sam2-ref/image.safetensors`, or
//! `EFFECTCRAFT_SAM2_REF`). Without them every test here skips, so CI stays green.
#![cfg(not(target_arch = "wasm32"))]
use std::collections::HashMap;
use std::path::PathBuf;
use candle_core::{Device, Tensor};
use effectcraft_sam2::hiera::{Hiera, HieraConfig};
use effectcraft_sam2::weights::Weights;
fn repo() -> PathBuf {
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../..")
}
fn weights_file() -> PathBuf {
std::env::var_os("EFFECTCRAFT_SAM2_DIR").map_or_else(|| repo().join("sam2-tiny"), PathBuf::from).join("model.safetensors")
}
fn fixture_file() -> PathBuf {
std::env::var_os("EFFECTCRAFT_SAM2_REF").map_or_else(|| repo().join("plan/sam2-ref/image.safetensors"), PathBuf::from)
}
/// The weights and the reference tensors, or `None` (the test then skips).
fn load() -> Option<(Weights, HashMap<String, Tensor>)> {
let (wf, rf) = (weights_file(), fixture_file());
if !wf.exists() || !rf.exists() {
eprintln!("skipped: SAM 2 weights or reference tensors not found ({} / {})", wf.display(), rf.display());
return None;
}
let w = Weights::open(&wf).unwrap();
let r = candle_core::safetensors::load(&rf, &Device::Cpu).unwrap();
Some((w, r))
}
fn max_abs(a: &Tensor, b: &Tensor) -> f32 {
assert_eq!(a.dims(), b.dims(), "shape mismatch");
(a - b).unwrap().abs().unwrap().flatten_all().unwrap().max(0).unwrap().to_scalar::<f32>().unwrap()
}
/// f32 activations differ from PyTorch by accumulation order only; the Hiera blocks stay well
/// inside this bound (measured max 2.6e-5, 2026-10-15).
const HIERA_TOL: f32 = 1e-4;
#[test]
fn hiera_matches_pytorch_block_by_block() {
let Some((w, r)) = load() else { return };
let hiera = Hiera::load(&w, &HieraConfig::tiny()).unwrap();
let pixels = r.get("input.pixel_values").unwrap();
let (emb, blocks) = hiera.forward_all(pixels).unwrap();
let e = max_abs(&emb, r.get("vision_encoder.backbone.patch_embed.0").unwrap());
eprintln!("patch_embed max abs error {e:e}");
assert!(e < 1e-4, "patch embed error {e}");
assert_eq!(blocks.len(), 12);
for (i, b) in blocks.iter().enumerate() {
let want = r.get(&format!("vision_encoder.backbone.blocks.{i}.0")).unwrap();
let e = max_abs(b, want);
eprintln!("block {i}: max abs error {e:e}");
assert!(e < HIERA_TOL, "block {i} error {e} (tolerance {HIERA_TOL})");
}
}
#[test]
fn hiera_is_deterministic_and_rejects_bad_sizes() {
let Some((w, r)) = load() else { return };
let hiera = Hiera::load(&w, &HieraConfig::tiny()).unwrap();
let pixels = r.get("input.pixel_values").unwrap();
let a = hiera.forward(pixels).unwrap();
let b = hiera.forward(pixels).unwrap();
assert_eq!(a.len(), 4);
for (x, y) in a.iter().zip(&b) {
assert_eq!(max_abs(x, y), 0.0);
}
let odd = pixels.narrow(2, 0, 1022).unwrap();
assert!(hiera.forward(&odd).is_err());
}
/// The neck adds 1x1 convolutions and nearest-neighbour sums to the Hiera output; the sine
/// position encodings are exact functions of the grid size.
const NECK_TOL: f32 = 1e-4;
#[test]
fn neck_and_vision_encoder_match_pytorch() {
let Some((w, r)) = load() else { return };
let enc = effectcraft_sam2::neck::VisionEncoder::load_tiny(&w).unwrap();
let pixels = r.get("input.pixel_values").unwrap();
let stages = enc.hiera.forward(pixels).unwrap();
let out = enc.neck.forward(&stages).unwrap();
assert_eq!((out.features.len(), out.positions.len()), (4, 4));
for (i, f) in out.features.iter().enumerate() {
let e = max_abs(f, r.get(&format!("vision_encoder.neck.{i}")).unwrap());
eprintln!("neck feature {i}: max abs error {e:e}");
assert!(e < NECK_TOL, "neck feature {i} error {e}");
}
for (i, p) in out.positions.iter().enumerate() {
let e = max_abs(p, r.get(&format!("vision_encoder.neck.{}", i + 4)).unwrap());
eprintln!("neck position {i}: max abs error {e:e}");
assert!(e < NECK_TOL, "neck position {i} error {e}");
}
// The vision encoder keeps the three highest-resolution levels, highest first.
let feats = enc.forward(pixels).unwrap();
let dims: Vec<usize> = feats.features.iter().map(|f| f.dim(2).unwrap()).collect();
assert_eq!(dims, vec![256, 128, 64]);
let e = max_abs(feats.features.first().unwrap(), r.get("vision_encoder.neck.3").unwrap());
assert!(e < NECK_TOL);
}
/// Click-prompted single image: prompt embeddings and decoder outputs against `Sam2Model`.
const DECODER_TOL: f32 = 1e-3;
#[test]
fn prompt_encoder_and_decoder_match_pytorch() {
use effectcraft_sam2::model::ImageModel;
use effectcraft_sam2::prompt::Prompt;
let Some((w, r)) = load() else { return };
let model = ImageModel::load(&w).unwrap();
let pixels = r.get("input.pixel_values").unwrap();
let prompt = Prompt { points: vec![([460.0, 512.0], true)], bbox: None };
let (sparse, dense) = model.embed_prompt(&prompt).unwrap();
let want_sparse = r.get("prompt_encoder.0").unwrap().squeeze(0).unwrap();
let e = max_abs(&sparse, &want_sparse);
eprintln!("sparse prompt: max abs error {e:e}");
assert!(e < 1e-5, "sparse prompt error {e}");
let e = max_abs(&dense, r.get("prompt_encoder.1").unwrap());
assert!(e < 1e-6, "dense prompt error {e}");
let enc = model.encode(pixels).unwrap();
let out = model.segment(&enc, &prompt, true).unwrap();
let want_masks = r.get("mask_decoder.0").unwrap().squeeze(0).unwrap();
let e = max_abs(&out.masks, &want_masks);
eprintln!("masks: max abs error {e:e} (logit range {:?})", want_masks.flatten_all().unwrap().max(0).unwrap().to_scalar::<f32>());
assert!(e < DECODER_TOL, "mask logits error {e}");
let e = max_abs(&out.iou, &r.get("mask_decoder.1").unwrap().squeeze(0).unwrap());
eprintln!("iou: max abs error {e:e}");
assert!(e < 1e-4, "iou error {e}");
let e = max_abs(&out.object_score, &r.get("mask_decoder.3").unwrap().squeeze(0).unwrap());
eprintln!("object score: max abs error {e:e}");
assert!(e < 1e-4, "object score error {e}");
let e = max_abs(&out.sam_tokens, &r.get("mask_decoder.2").unwrap().squeeze(0).unwrap());
assert!(e < 1e-4, "sam token error {e}");
}
fn video_fixture_file() -> PathBuf {
std::env::var_os("EFFECTCRAFT_SAM2_VIDEO_REF").map_or_else(|| repo().join("plan/sam2-ref/video.safetensors"), PathBuf::from)
}
/// The weights and the video reference tensors (`tools/sam2_video_reference.py`), or `None`.
fn load_video() -> Option<(Weights, HashMap<String, Tensor>)> {
let (wf, rf) = (weights_file(), video_fixture_file());
if !wf.exists() || !rf.exists() {
eprintln!("skipped: SAM 2 weights or video reference tensors not found ({} / {})", wf.display(), rf.display());
return None;
}
Some((Weights::open(&wf).unwrap(), candle_core::safetensors::load(&rf, &Device::Cpu).unwrap()))
}
/// Measured max 3.5e-6 (memory features), 2.6e-5 (memory attention), 2026-10-15.
const MEMORY_TOL: f32 = 1e-4;
#[test]
fn memory_encoder_matches_pytorch() {
let Some((w, r)) = load_video() else { return };
let enc = effectcraft_sam2::memory::MemoryEncoder::load(&w).unwrap();
let (feat, pos) = enc.encode(r.get("mem.enc.vision_features").unwrap(), r.get("mem.enc.masks").unwrap()).unwrap();
let e = max_abs(&feat, r.get("mem.enc.out_features").unwrap());
eprintln!("memory features: max abs error {e:e}");
assert!(e < MEMORY_TOL, "memory feature error {e}");
let e = max_abs(&pos, r.get("mem.enc.out_pos").unwrap());
eprintln!("memory position encoding: max abs error {e:e}");
assert!(e < 1e-5, "memory position error {e}");
}
#[test]
fn memory_attention_matches_pytorch() {
let Some((w, r)) = load_video() else { return };
let att = effectcraft_sam2::memattn::MemoryAttention::load(&w).unwrap();
let get = |k: &str| r.get(k).unwrap().squeeze(1).unwrap();
let pointers = r.get("mem.attn.num_pointers").unwrap().to_vec1::<f32>().unwrap()[0] as usize;
let out = att
.forward(
&get("mem.attn.current_vision_features"),
&get("mem.attn.current_vision_position_embeddings"),
&get("mem.attn.memory"),
&get("mem.attn.memory_posision_embeddings"),
pointers,
)
.unwrap();
let want = r.get("mem.attn.out").unwrap().squeeze(0).unwrap().squeeze(0).unwrap();
let e = max_abs(&out, &want);
eprintln!("memory attention: max abs error {e:e} (pointers {pointers})");
assert!(e < MEMORY_TOL, "memory attention error {e}");
}
/// Frames of the video fixture, served one at a time.
struct Clip(Tensor);
impl effectcraft_sam2::session::FrameSource for Clip {
fn num_frames(&self) -> usize {
self.0.dim(0).unwrap()
}
fn frame(&self, i: usize) -> effectcraft_sam2::Result<Tensor> {
Ok(self.0.narrow(0, i, 1)?.contiguous()?)
}
}
fn iou(a: &[f32], b: &[f32]) -> f64 {
let (mut i, mut u) = (0usize, 0usize);
for (x, y) in a.iter().zip(b) {
let (x, y) = (*x > 0.0, *y > 0.0);
i += usize::from(x && y);
u += usize::from(x || y);
}
if u == 0 { 1.0 } else { i as f64 / u as f64 }
}
/// Through the clip the memory bank feeds back on itself (bfloat16 memories, rotary attention), so
/// the bound on logits is looser than for a single stage; the masks themselves must agree.
/// Measured max 0.0038 over 9 frames forwards and backwards, 2026-10-15.
const VIDEO_LOGIT_TOL: f32 = 0.02;
const VIDEO_MIN_IOU: f64 = 0.995;
#[test]
fn video_tracking_matches_sam2_video_model() {
use effectcraft_sam2::prompt::Prompt;
use effectcraft_sam2::session::{Cancel, Direction, ObjectId, VideoModel, VideoTracker};
let Some((w, r)) = load_video() else { return };
let model = std::sync::Arc::new(VideoModel::load(&w).unwrap());
let clip = r.get("clip.pixel_values").unwrap().clone();
let mut tracker = VideoTracker::new(model, Box::new(Clip(clip))).unwrap();
let obj = ObjectId(1);
let first = tracker.prompt(2, obj, &Prompt { points: vec![([416.0, 512.0], true)], bbox: None }).unwrap();
let check = |tag: &str, m: &effectcraft_sam2::session::FrameMask| {
let want = r.get(&format!("video.{tag}.{}.pred_masks", m.frame)).unwrap().flatten_all().unwrap().to_vec1::<f32>().unwrap();
let want_score = r.get(&format!("video.{tag}.{}.score", m.frame)).unwrap().flatten_all().unwrap().to_vec1::<f32>().unwrap()[0];
let err = m.logits.iter().zip(&want).map(|(a, b)| (a - b).abs()).fold(0f32, f32::max);
let j = iou(&m.logits, &want);
eprintln!("{tag} frame {}: logit error {err:.4}, IoU {j:.4}, score {:.3} (want {want_score:.3})", m.frame, m.object_score);
assert!(j >= VIDEO_MIN_IOU, "{tag} frame {} IoU {j}", m.frame);
assert!(err < VIDEO_LOGIT_TOL, "{tag} frame {} logit error {err}", m.frame);
assert!((m.object_score - want_score).abs() < 0.05, "{tag} frame {} score {} vs {want_score}", m.frame, m.object_score);
assert_eq!(m.occluded, want_score <= 0.0);
};
check("prompt", &first);
let cancel = Cancel::new();
let mut fwd = vec![];
tracker.propagate(2, Direction::Forward, None, &cancel, &mut |m| fwd.push(m)).unwrap();
assert_eq!(fwd.iter().map(|m| m.frame).collect::<Vec<_>>(), vec![2, 3, 4, 5, 6]);
for m in &fwd {
check("fwd", m);
}
let mut rev = vec![];
tracker.propagate(2, Direction::Backward, None, &cancel, &mut |m| rev.push(m)).unwrap();
assert_eq!(rev.iter().map(|m| m.frame).collect::<Vec<_>>(), vec![2, 1, 0]);
for m in &rev {
check("rev", m);
}
}
#[test]
fn tracker_rejects_hostile_input_and_honours_cancel() {
use effectcraft_sam2::Error;
use effectcraft_sam2::prompt::Prompt;
use effectcraft_sam2::session::{Cancel, Direction, ObjectId, VideoModel, VideoTracker};
let Some((w, r)) = load_video() else { return };
let model = std::sync::Arc::new(VideoModel::load(&w).unwrap());
let clip = r.get("clip.pixel_values").unwrap().narrow(0, 0, 3).unwrap();
let mut t = VideoTracker::new(model.clone(), Box::new(Clip(clip.clone()))).unwrap();
let click = Prompt { points: vec![([416.0, 512.0], true)], bbox: None };
let obj = ObjectId(7);
// outside the clip, empty and non-finite prompts, nothing prompted yet
assert!(matches!(t.prompt(3, obj, &click), Err(Error::Input(_))));
assert!(matches!(t.prompt(0, obj, &Prompt::default()), Err(Error::Input(_))));
assert!(matches!(t.prompt(0, obj, &Prompt { points: vec![([f32::NAN, 1.0], true)], bbox: None }), Err(Error::Input(_))));
let mut sink = |_m| {};
assert!(matches!(t.propagate(0, Direction::Forward, None, &Cancel::new(), &mut sink), Err(Error::Input(_))));
assert!(matches!(t.propagate(9, Direction::Forward, None, &Cancel::new(), &mut sink), Err(Error::Input(_))));
// an empty clip is refused; a cancelled run stops with Cancelled and keeps the tracker usable
assert!(VideoTracker::new(model.clone(), Box::new(Clip(clip.narrow(0, 0, 0).unwrap()))).is_err());
t.prompt(0, obj, &click).unwrap();
let cancel = Cancel::new();
cancel.cancel();
assert!(matches!(t.propagate(0, Direction::Forward, None, &cancel, &mut sink), Err(Error::Cancelled)));
let mut seen = vec![];
t.propagate(0, Direction::Forward, Some(1), &Cancel::new(), &mut |m| seen.push(m.frame)).unwrap();
assert_eq!(seen, vec![0, 1]);
// a frame source that returns the wrong size is an error, not a crash
struct Bad;
impl effectcraft_sam2::session::FrameSource for Bad {
fn num_frames(&self) -> usize {
2
}
fn frame(&self, _i: usize) -> effectcraft_sam2::Result<Tensor> {
Ok(Tensor::zeros((1, 3, 64, 64), candle_core::DType::F32, &Device::Cpu)?)
}
}
let mut bad = VideoTracker::new(model, Box::new(Bad)).unwrap();
assert!(matches!(bad.prompt(0, obj, &click), Err(Error::Input(_))));
}
#[test]
fn sam2_works_as_a_mask_model() {
use effectcraft_segment::{MaskModel, Prompt};
let Some((w, _)) = load() else { return };
let (fw, fh) = (256usize, 192usize);
let mut rgb = vec![[0f32; 3]; fw * fh];
let inside = |x: usize, y: usize| ((x as f32 - 120.0).powi(2) + (y as f32 - 96.0).powi(2)).sqrt() < 50.0;
for y in 0..fh {
for x in 0..fw {
let bg = 0.3 + 0.2 * (x as f32 / 9.0).sin() * (y as f32 / 7.0).cos();
rgb[y * fw + x] = if inside(x, y) { [0.85 + 0.1 * (x as f32 / 5.0).sin(), 0.4, 0.15] } else { [bg, bg * 0.8 + 0.1, bg * 0.6 + 0.3] };
}
}
let model = effectcraft_sam2::mask::Sam2Mask::new(std::sync::Arc::new(effectcraft_sam2::session::VideoModel::load(&w).unwrap()));
let prompt = Prompt { points: vec![([120.0, 96.0], true)], bbox: None, mask: None };
let p = model.segment(&rgb, fw, fh, &prompt).unwrap();
assert_eq!(p.len(), fw * fh);
assert!(p.iter().all(|v| (0.0..=1.0).contains(v)));
let (mut i, mut u) = (0usize, 0usize);
for y in 0..fh {
for x in 0..fw {
let (a, b) = (p[y * fw + x] > 0.5, inside(x, y));
i += usize::from(a && b);
u += usize::from(a || b);
}
}
let j = i as f64 / u as f64;
eprintln!("SAM 2 as a MaskModel: disc IoU {j:.3}");
assert!(j > 0.9, "IoU {j}");
// The same frame again is served from the cached encoding, and bad prompts are errors.
let again = model.segment(&rgb, fw, fh, &Prompt { points: vec![([20.0, 20.0], false), ([120.0, 96.0], true)], bbox: None, mask: None }).unwrap();
assert_eq!(again.len(), fw * fh);
assert!(model.segment(&rgb, fw, fh, &Prompt::default()).is_err());
assert!(model.segment(&rgb, fw + 1, fh, &prompt).is_err());
}
#[test]
fn sam2_tracks_a_clip_through_the_object_tracker_interface() {
use effectcraft_segment::{ClipFrame, MaskModel, Prompt, TrackRequest};
let Some((w, _)) = load() else { return };
let (fw, fh, n) = (128usize, 96usize, 6usize);
let centre = |t: usize| (40.0 + 8.0 * t as f32, 48.0f32);
let inside = |t: usize, x: usize, y: usize| {
let c = centre(t);
((x as f32 - c.0).powi(2) + (y as f32 - c.1).powi(2)).sqrt() < 26.0
};
let frames: Vec<std::sync::Arc<ClipFrame>> = (0..n)
.map(|t| {
let mut rgb = vec![[0f32; 3]; fw * fh];
for y in 0..fh {
for x in 0..fw {
let bg = 0.3 + 0.2 * (x as f32 / 9.0).sin() * (y as f32 / 7.0).cos();
rgb[y * fw + x] =
if inside(t, x, y) { [0.85 + 0.1 * ((x as f32 - centre(t).0) / 5.0).sin(), 0.4, 0.15] } else { [bg, bg * 0.8 + 0.1, bg * 0.6 + 0.3] };
}
}
std::sync::Arc::new(ClipFrame { rgb, w: fw, h: fh })
})
.collect();
let model = effectcraft_sam2::mask::Sam2Mask::new(std::sync::Arc::new(effectcraft_sam2::session::VideoModel::load(&w).unwrap()));
assert!(model.tracker().is_some());
let c = centre(1);
let req = TrackRequest { first: 100, count: n, prompts: vec![(101, Prompt { points: vec![([c.0, c.1], true)], bbox: None, mask: None })] };
let mut seen = std::collections::BTreeMap::new();
let mut get = |f: i64| frames.get((f - 100) as usize).cloned();
model
.tracker()
.unwrap()
.track(&req, &mut get, &|| false, &mut |f, t| {
assert!(seen.insert(f, t).is_none(), "frame {f} reported twice");
})
.unwrap();
assert_eq!(seen.keys().copied().collect::<Vec<_>>(), (100..100 + n as i64).collect::<Vec<_>>());
for (f, t) in &seen {
assert_eq!((t.w, t.h, t.prob.len()), (fw, fh, fw * fh));
let k = (f - 100) as usize;
let (mut i, mut u) = (0usize, 0usize);
for y in 0..fh {
for x in 0..fw {
let (a, b) = (t.prob[y * fw + x] > 0.5, inside(k, x, y));
i += usize::from(a && b);
u += usize::from(a || b);
}
}
let j = i as f64 / u.max(1) as f64;
eprintln!("frame {f}: IoU {j:.3}, confidence {:.2}, occluded {}", t.confidence, t.occluded);
assert!(j > 0.85 && !t.occluded && t.confidence > 0.3, "frame {f}: IoU {j}");
}
// Cancelling before the work starts, and bad requests, are errors.
assert!(model.tracker().unwrap().track(&req, &mut get, &|| true, &mut |_, _| {}).is_err());
let none = TrackRequest { first: 100, count: n, prompts: vec![] };
assert!(model.tracker().unwrap().track(&none, &mut get, &|| false, &mut |_, _| {}).is_err());
let outside = TrackRequest { first: 100, count: n, prompts: vec![(500, Prompt { points: vec![([1.0, 1.0], true)], bbox: None, mask: None })] };
assert!(model.tracker().unwrap().track(&outside, &mut get, &|| false, &mut |_, _| {}).is_err());
}