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(_)))); +}