From 5b0cc57f87670035eef83439e08146ed959a7b37 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:35:01 +0200 Subject: [PATCH] S2.23: video tracker session Per-object memory banks, prompts on any frame, forward and backward propagation with streaming results, cancellation, object pointers with temporal encoding, pruning of memories no later frame can use, and checked input handling. Hostile-input and cancel tests. Co-Authored-By: Claude Sonnet 5.5 --- crates/sam2/src/lib.rs | 1 + crates/sam2/src/model.rs | 10 + crates/sam2/src/prompt.rs | 8 + crates/sam2/src/session.rs | 363 +++++++++++++++++++++++++++++++++ crates/sam2/tests/reference.rs | 53 +++++ 5 files changed, 435 insertions(+) create mode 100644 crates/sam2/src/session.rs diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index 53156c7ab..4cb6f6f3e 100644 --- a/crates/sam2/src/lib.rs +++ b/crates/sam2/src/lib.rs @@ -19,6 +19,7 @@ pub mod model; pub mod neck; pub mod nn; pub mod prompt; +pub mod session; pub mod weights; /// Everything that can go wrong in this crate. Never a panic. diff --git a/crates/sam2/src/model.rs b/crates/sam2/src/model.rs index 13481fdb9..c30cbb5b9 100644 --- a/crates/sam2/src/model.rs +++ b/crates/sam2/src/model.rs @@ -71,6 +71,16 @@ impl ImageModel { self.decoder.forward(&image, &self.image_pos, &sparse, &dense, &enc.high, multimask) } + /// Decode against an already memory-conditioned embedding `[1, 256, g, g]` (video tracking). + pub fn decode_embedded(&self, image: &Tensor, high: &(Tensor, Tensor), sparse: &Tensor, dense: &Tensor, multimask: bool) -> Result { + self.decoder.forward(image, &self.image_pos, sparse, dense, high, multimask) + } + + /// The learned embedding added to a frame's features when there is no memory yet. + pub fn no_memory(&self) -> &Tensor { + &self.no_memory + } + pub fn device(&self) -> Device { Device::Cpu } diff --git a/crates/sam2/src/prompt.rs b/crates/sam2/src/prompt.rs index 55dd901a2..51aade1c7 100644 --- a/crates/sam2/src/prompt.rs +++ b/crates/sam2/src/prompt.rs @@ -146,6 +146,14 @@ impl PromptEncoder { Ok(Tensor::from_vec(tokens, (1, n, DIM), &Device::Cpu)?) } + /// Tokens for "no prompt" on a tracked frame: the reference feeds one point labelled -1 and + /// then pads another, so two `not_a_point` tokens, `[1, 2, 256]`. + pub fn sparse_none(&self) -> Result { + let mut v = self.not_a_point.clone(); + v.extend_from_slice(&self.not_a_point); + Ok(Tensor::from_vec(v, (1, 2, DIM), &Device::Cpu)?) + } + /// The dense embedding for "no mask prompt", `[1, 256, GRID, GRID]`. pub fn dense_empty(&self) -> Result { Ok(self.no_mask.reshape((1, DIM, 1, 1))?.broadcast_as((1, DIM, GRID, GRID))?.contiguous()?) diff --git a/crates/sam2/src/session.rs b/crates/sam2/src/session.rs new file mode 100644 index 000000000..45d1ffd0a --- /dev/null +++ b/crates/sam2/src/session.rs @@ -0,0 +1,363 @@ +// Modified work: ported from Hugging Face `transformers` `models/sam2_video/modeling_sam2_video.py` +// (Sam2VideoModel and Sam2VideoInferenceSession: memory-conditioned single-frame inference, +// memory bank selection, object pointers, propagation; Apache-2.0, Copyright The HuggingFace +// Inc. team) to Rust on candle. See NOTICE. +//! The video tracker: per-object memory banks, prompts on any frame, propagation forwards and +//! backwards with streaming results and cancellation. + +use std::collections::{BTreeMap, VecDeque}; +use std::path::Path; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; + +use candle_core::{DType, Tensor}; +use candle_nn::{Linear, ops::sigmoid}; + +use crate::decoder::FeedForward; +use crate::memattn::MemoryAttention; +use crate::memory::{Memory, MemoryEncoder}; +use crate::model::ImageModel; +use crate::nn::{apply, linear}; +use crate::prompt::{INPUT_SIZE, Prompt}; +use crate::weights::Weights; +use crate::{Error, Result}; + +/// Mask logit used where the object is absent (the reference's `NO_OBJ_SCORE`). +const NO_OBJECT: f32 = -1024.0; +/// Frames of memory the attention looks back over. +const MEMORY_FRAMES: usize = 7; +/// Object pointers kept per query. +const MAX_POINTERS: usize = 16; +/// Largest clip the session accepts. +const MAX_FRAMES: usize = 1_000_000; +const GRID: usize = 64; + +/// Frames of the clip, already resized to the model input and normalised: `[1, 3, 1024, 1024]`. +pub trait FrameSource: Send { + fn num_frames(&self) -> usize; + fn frame(&self, index: usize) -> Result; +} + +/// Cooperative cancellation shared with the caller. +#[derive(Clone, Default)] +pub struct Cancel(Arc); + +impl Cancel { + pub fn new() -> Cancel { + Cancel::default() + } + pub fn cancel(&self) { + self.0.store(true, Ordering::Relaxed); + } + pub fn is_cancelled(&self) -> bool { + self.0.load(Ordering::Relaxed) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ObjectId(pub u32); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Direction { + Forward, + Backward, +} + +/// One object's mask on one frame. +#[derive(Clone, Debug, PartialEq)] +pub struct FrameMask { + pub frame: usize, + pub object: ObjectId, + /// Mask logits at 256 x 256 (positive is inside the object); `-1024` everywhere when the + /// object is absent. + pub logits: Vec, + /// The model's own estimate of the mask quality (0..1). + pub iou: f32, + /// Object-presence logit; at or below zero the object is not in the frame. + pub object_score: f32, + pub occluded: bool, +} + +/// The weights of everything the video tracker needs. +pub struct VideoModel { + pub image: ImageModel, + memory_encoder: MemoryEncoder, + memory_attention: MemoryAttention, + temporal_pos: Vec, + no_object_pointer: Tensor, + pointer_proj: FeedForward, + pointer_time_proj: Linear, +} + +impl VideoModel { + pub fn load(w: &Weights) -> Result { + Ok(VideoModel { + image: ImageModel::load(w)?, + memory_encoder: MemoryEncoder::load(w)?, + memory_attention: MemoryAttention::load(w)?, + temporal_pos: w.get("memory_temporal_positional_encoding", &[MEMORY_FRAMES, 1, 1, 64])?.flatten_all()?.to_vec1::()?, + no_object_pointer: w.get("no_object_pointer", &[1, 256])?, + pointer_proj: FeedForward::load(w, "object_pointer_proj", 256, 256, 256, 3, false)?, + pointer_time_proj: linear(w, "temporal_positional_encoding_projection_layer", 256, 64)?, + }) + } + + pub fn open(path: &Path) -> Result { + VideoModel::load(&Weights::open(path)?) + } +} + +/// What the tracker keeps for a frame it has processed. +struct FrameOut { + memory: Memory, + /// `[1, 256]` + pointer: Tensor, + /// The mask of a conditioning frame (returned again, unchanged, when propagation passes it). + result: Option, +} + +#[derive(Default)] +struct ObjectState { + cond: BTreeMap, + non_cond: BTreeMap, +} + +struct FrameFeatures { + /// Lowest-resolution level `[1, 256, 64, 64]` and its position encoding. + low: Tensor, + pos: Tensor, + high: (Tensor, Tensor), +} + +pub struct VideoTracker { + model: Arc, + frames: Box, + objects: BTreeMap, + cache: VecDeque<(usize, Arc)>, +} + +fn sine_embedding_1d(positions: &[f32], dim: usize) -> Result { + let half = dim / 2; + let dim_t: Vec = (0..half).map(|i| 10000f32.powf(2.0 * (i / 2) as f32 / half as f32)).collect(); + let mut out = Vec::with_capacity(positions.len() * dim); + for p in positions { + out.extend(dim_t.iter().map(|d| (p / d).sin())); + out.extend(dim_t.iter().map(|d| (p / d).cos())); + } + Ok(Tensor::from_vec(out, (positions.len(), dim), &candle_core::Device::Cpu)?) +} + +fn scalar(t: &Tensor) -> Result { + Ok(t.flatten_all()?.to_vec1::()?.first().copied().unwrap_or(0.0)) +} + +impl VideoTracker { + pub fn new(model: Arc, frames: Box) -> Result { + let n = frames.num_frames(); + if n == 0 || n > MAX_FRAMES { + return Err(Error::Input(format!("a clip needs between 1 and {MAX_FRAMES} frames, got {n}"))); + } + Ok(VideoTracker { model, frames, objects: BTreeMap::new(), cache: VecDeque::new() }) + } + + pub fn num_frames(&self) -> usize { + self.frames.num_frames() + } + + /// Forget every object and cached frame (the clip stays). + pub fn reset(&mut self) { + self.objects.clear(); + self.cache.clear(); + } + + fn features(&mut self, frame: usize) -> Result> { + if let Some((_, f)) = self.cache.iter().find(|(i, _)| *i == frame) { + return Ok(f.clone()); + } + let pixels = self.frames.frame(frame)?; + let enc = self.model.image.encode(&pixels)?; + let (low, pos) = match (enc.features.features.get(2), enc.features.positions.get(2)) { + (Some(l), Some(p)) => (l.clone(), p.clone()), + _ => return Err(Error::Compute("vision encoder returned too few levels".into())), + }; + let f = Arc::new(FrameFeatures { low, pos, high: enc.high }); + self.cache.push_back((frame, f.clone())); + while self.cache.len() > 2 { + self.cache.pop_front(); + } + Ok(f) + } + + /// The memory bank for `frame`: memory frames and object pointers with their position + /// encodings, plus how many trailing tokens are pointers. + fn memory_bank(&self, obj: &ObjectState, frame: usize, reverse: bool) -> Result<(Tensor, Tensor, usize)> { + let (mut feats, mut poss) = (vec![], vec![]); + let temporal = |off: usize| -> Result { + // The reference indexes with `offset - 1`, so the conditioning frames (offset 0) use the last row. + let row = if off == 0 { MEMORY_FRAMES - 1 } else { off - 1 }; + let v = self.model.temporal_pos.get(row * 64..(row + 1) * 64).ok_or_else(|| Error::Weights("temporal position table".into()))?; + Ok(Tensor::from_slice(v, (1, 64), &candle_core::Device::Cpu)?) + }; + for out in obj.cond.values() { + feats.push(out.memory.features.clone()); + poss.push(out.memory.pos.broadcast_add(&temporal(0)?)?); + } + for off in (1..MEMORY_FRAMES).rev() { + let prev = if reverse { frame.checked_add(off) } else { frame.checked_sub(off) }; + if let Some(out) = prev.and_then(|p| obj.non_cond.get(&p)) { + feats.push(out.memory.features.clone()); + poss.push(out.memory.pos.broadcast_add(&temporal(off)?)?); + } + } + // Object pointers: conditioning frames on the already-tracked side, then recent frames. + let n_frames = self.frames.num_frames(); + let max_ptrs = n_frames.min(MAX_POINTERS); + let (mut offsets, mut pointers): (Vec, Vec) = (vec![], vec![]); + let sign = if reverse { -1.0f32 } else { 1.0 }; + for (t, out) in &obj.cond { + if if reverse { *t >= frame } else { *t <= frame } { + offsets.push((frame as f32 - *t as f32) * sign); + pointers.push(out.pointer.clone()); + } + } + for off in 1..max_ptrs { + let r = if reverse { frame.checked_add(off) } else { frame.checked_sub(off) }; + let Some(r) = r.filter(|r| *r < n_frames) else { break }; + if let Some(out) = obj.non_cond.get(&r) { + offsets.push(off as f32); + pointers.push(out.pointer.clone()); + } + } + let mut n_pointer_tokens = 0; + if !pointers.is_empty() { + let max_diff = (max_ptrs.saturating_sub(1)).max(1) as f32; + let norm: Vec = offsets.iter().map(|o| o / max_diff).collect(); + let pe = apply(&self.model.pointer_time_proj, &sine_embedding_1d(&norm, 256)?)?; + let ptrs = Tensor::cat(&pointers, 0)?; // [P, 256] + let p = ptrs.dim(0)?; + let tokens = ptrs.reshape((p * 4, 64))?; + let pe = pe.unsqueeze(1)?.broadcast_as((p, 4, 64))?.reshape((p * 4, 64))?; + n_pointer_tokens = p * 4; + feats.push(tokens); + poss.push(pe.contiguous()?); + } + if feats.is_empty() { + return Err(Error::Input("the object has no conditioning frame; add a prompt first".into())); + } + Ok((Tensor::cat(&feats, 0)?, Tensor::cat(&poss, 0)?, n_pointer_tokens)) + } + + /// Run the model on one frame for one object and store the result. + fn step(&mut self, id: ObjectId, frame: usize, prompt: Option<&Prompt>, reverse: bool, conditioning: bool) -> Result { + let feats = self.features(frame)?; + let model = self.model.clone(); + let tokens = |t: &Tensor| -> Result { Ok(t.flatten_from(2)?.squeeze(0)?.t()?.contiguous()?) }; + let conditioned = if conditioning { + feats.low.broadcast_add(model.image.no_memory())? + } else { + let obj = self.objects.get(&id).ok_or_else(|| Error::Input("unknown object".into()))?; + let (mem, mem_pos, n_ptr) = self.memory_bank(obj, frame, reverse)?; + let out = model.memory_attention.forward(&tokens(&feats.low)?, &tokens(&feats.pos)?, &mem, &mem_pos, n_ptr)?; + out.t()?.contiguous()?.reshape((1, 256, GRID, GRID))? + }; + let num_points = prompt.map_or(0, |p| p.points.len() + 2 * usize::from(p.bbox.is_some())); + let multimask = num_points <= 1; + let sparse = match prompt { + Some(p) => model.image.prompt.sparse(p)?, + None => model.image.prompt.sparse_none()?, + }; + let dense = model.image.prompt.dense_empty()?; + let out = model.image.decode_embedded(&conditioned, &feats.high, &sparse, &dense, multimask)?; + + let score = scalar(&out.object_score)?; + let present = score > 0.0; + let ious = out.iou.flatten_all()?.to_vec1::()?; + let best = if multimask { ious.iter().enumerate().fold((0usize, f32::NEG_INFINITY), |b, (i, v)| if *v > b.1 { (i, *v) } else { b }).0 } else { 0 }; + let low = out.masks.narrow(1, best, 1)?; + let low = if present { low } else { Tensor::full(NO_OBJECT, low.dims(), low.device())? }; + let token = if multimask { out.sam_tokens.narrow(1, best, 1)? } else { out.sam_tokens.narrow(1, 0, 1)? }.squeeze(1)?; + let pointer = if present { model.pointer_proj.forward(&token)? } else { model.no_object_pointer.clone() }; + + let high = low.upsample_bilinear2d(INPUT_SIZE, INPUT_SIZE, false)?; + let for_memory = if prompt.is_some() { high.gt(0f32)?.to_dtype(DType::F32)? } else { sigmoid(&high)? }; + let for_memory = ((for_memory * 20.0)? - 10.0)?; + let memory = model.memory_encoder.memory(&feats.low, &for_memory, present)?; + + let iou = ious.get(best).copied().unwrap_or(0.0); + let result = FrameMask { frame, object: id, logits: low.flatten_all()?.to_vec1::()?, iou, object_score: score, occluded: !present }; + let state = self.objects.entry(id).or_default(); + let entry = FrameOut { memory, pointer, result: conditioning.then(|| result.clone()) }; + if conditioning { + state.non_cond.remove(&frame); + state.cond.insert(frame, entry); + } else { + state.non_cond.insert(frame, entry); + Self::prune(state, frame); + } + Ok(result) + } + + /// Drop memories that no later frame in either direction can still look at. + fn prune(state: &mut ObjectState, frame: usize) { + let keep = |f: usize, cond: &BTreeMap| f.abs_diff(frame) <= MAX_POINTERS || cond.keys().any(|c| f.abs_diff(*c) <= MEMORY_FRAMES); + let doomed: Vec = state.non_cond.keys().copied().filter(|f| !keep(*f, &state.cond)).collect(); + for f in doomed { + state.non_cond.remove(&f); + } + } + + fn check_frame(&self, frame: usize) -> Result<()> { + if frame >= self.frames.num_frames() { + return Err(Error::Input(format!("frame {frame} is outside the clip of {} frames", self.frames.num_frames()))); + } + Ok(()) + } + + /// Prompt an object on a frame (clicks and/or a box). The frame becomes one of the object's + /// conditioning frames when it has not been tracked yet; the mask comes back at once. + pub fn prompt(&mut self, frame: usize, object: ObjectId, prompt: &Prompt) -> Result { + self.check_frame(frame)?; + prompt.validate()?; + if prompt.points.is_empty() && prompt.bbox.is_none() { + return Err(Error::Input("a prompt needs at least one point or a box".into())); + } + let tracked = self.objects.get(&object).is_some_and(|o| o.non_cond.contains_key(&frame)); + if tracked { + // A correction on a tracked frame is conditioned on the memory like any tracked frame. + let reverse = false; + return self.step(object, frame, Some(prompt), reverse, false); + } + self.objects.entry(object).or_default(); + self.step(object, frame, Some(prompt), false, true) + } + + /// Track every object from `start` through `max_frames` frames (the rest of the clip when + /// `None`) in `direction`, calling `out` for each object on each frame as it is finished. + pub fn propagate(&mut self, start: usize, direction: Direction, max_frames: Option, cancel: &Cancel, out: &mut dyn FnMut(FrameMask)) -> Result<()> { + self.check_frame(start)?; + let n = self.frames.num_frames(); + let reverse = direction == Direction::Backward; + let span = max_frames.unwrap_or(n).min(n); + let ids: Vec = self.objects.iter().filter(|(_, o)| !o.cond.is_empty()).map(|(id, _)| *id).collect(); + if ids.is_empty() { + return Err(Error::Input("no object has been prompted".into())); + } + for k in 0..=span { + let frame = if reverse { start.checked_sub(k) } else { Some(start + k).filter(|f| *f < n) }; + let Some(frame) = frame else { break }; + for id in &ids { + if cancel.is_cancelled() { + return Err(Error::Cancelled); + } + let cached = self.objects.get(id).is_some_and(|o| o.cond.contains_key(&frame)); + let stored = self.objects.get(id).and_then(|o| o.cond.get(&frame)).and_then(|o| o.result.clone()); + let mask = match stored { + Some(m) if cached => m, + _ => self.step(*id, frame, None, reverse, false)?, + }; + out(mask); + } + } + Ok(()) + } +} diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index d6d457871..7e480af10 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -194,3 +194,56 @@ fn memory_attention_matches_pytorch() { 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 { + Ok(self.0.narrow(0, i, 1)?.contiguous()?) + } +} + +#[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 { + 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(_)))); +}