From a62d59da4bf7b5489d4e86a0652fffc0b7619e7d Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:47:59 +0200 Subject: [PATCH] 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 --- crates/sam2/src/lib.rs | 1 + crates/sam2/src/mask.rs | 4 + crates/sam2/src/session.rs | 10 +-- crates/sam2/src/track.rs | 131 +++++++++++++++++++++++++++++++++ crates/sam2/tests/reference.rs | 60 +++++++++++++++ crates/segment/src/lib.rs | 48 ++++++++++++ 6 files changed, 249 insertions(+), 5 deletions(-) create mode 100644 crates/sam2/src/track.rs diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index 0d7700a0e..6a0eb7fbf 100644 --- a/crates/sam2/src/lib.rs +++ b/crates/sam2/src/lib.rs @@ -21,6 +21,7 @@ pub mod neck; pub mod nn; pub mod prompt; pub mod session; +pub mod track; pub mod weights; /// Everything that can go wrong in this crate. Never a panic. diff --git a/crates/sam2/src/mask.rs b/crates/sam2/src/mask.rs index 8bb309da0..a0eac9cd6 100644 --- a/crates/sam2/src/mask.rs +++ b/crates/sam2/src/mask.rs @@ -105,6 +105,10 @@ impl MaskModel for Sam2Mask { fn segment(&self, rgb: &[[f32; 3]], w: usize, h: usize, prompt: &effectcraft_segment::Prompt) -> effectcraft_segment::Result> { self.probability(rgb, w, h, prompt).map_err(|e| e.to_string()) } + + fn tracker(&self) -> Option<&dyn effectcraft_segment::ObjectTracker> { + Some(self) + } } /// Verify the installed file against the registry and build the model. diff --git a/crates/sam2/src/session.rs b/crates/sam2/src/session.rs index 45d1ffd0a..3440ff379 100644 --- a/crates/sam2/src/session.rs +++ b/crates/sam2/src/session.rs @@ -33,7 +33,7 @@ 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 { +pub trait FrameSource { fn num_frames(&self) -> usize; fn frame(&self, index: usize) -> Result; } @@ -129,9 +129,9 @@ struct FrameFeatures { high: (Tensor, Tensor), } -pub struct VideoTracker { +pub struct VideoTracker<'a> { model: Arc, - frames: Box, + frames: Box, objects: BTreeMap, cache: VecDeque<(usize, Arc)>, } @@ -151,8 +151,8 @@ 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 { +impl<'a> VideoTracker<'a> { + 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}"))); diff --git a/crates/sam2/src/track.rs b/crates/sam2/src/track.rs new file mode 100644 index 000000000..9320fbbe4 --- /dev/null +++ b/crates/sam2/src/track.rs @@ -0,0 +1,131 @@ +// Modified work: adapts the ported SAM 2 video tracker to EffectCraft's `ObjectTracker` +// interface (Apache-2.0, see NOTICE). +//! SAM 2 as the Roto Brush's object tracker: prompts on some frames, the whole clip tracked +//! forwards and backwards, each frame reported as it is finished. + +use std::cell::RefCell; +use std::collections::HashMap; +use std::rc::Rc; +use std::sync::Arc; + +use candle_core::Tensor; +use effectcraft_segment::{ClipFrame, ObjectTracker, TrackRequest, TrackedFrame}; + +use crate::mask::{Sam2Mask, preprocess}; +use crate::prompt::{INPUT_SIZE, Prompt}; +use crate::session::{Cancel, Direction, FrameMask, FrameSource, ObjectId, VideoTracker}; +use crate::{Error, Result}; + +/// Most frames one tracking request may cover. +const MAX_CLIP: usize = 100_000; + +type Frames<'a> = &'a mut dyn FnMut(i64) -> Option>; + +/// Feeds the host's frames to the tracker, remembering each frame's size for the output. +struct Source<'a> { + first: i64, + count: usize, + fetch: RefCell>, + sizes: Rc>>, +} + +impl FrameSource for Source<'_> { + fn num_frames(&self) -> usize { + self.count + } + + fn frame(&self, index: usize) -> Result { + let f = (self.fetch.borrow_mut())(self.first + index as i64) + .ok_or_else(|| Error::Input(format!("frame {} is not available", self.first + index as i64)))?; + self.sizes.borrow_mut().insert(index, (f.w, f.h)); + preprocess(&f.rgb, f.w, f.h) + } +} + +fn to_frame(m: &FrameMask, size: (usize, usize)) -> Result { + let (w, h) = size; + let logits = Tensor::from_slice(&m.logits, (1, 1, 256, 256), &candle_core::Device::Cpu)?; + let prob = candle_nn::ops::sigmoid(&logits.upsample_bilinear2d(h, w, false)?)?.flatten_all()?.to_vec1::()?; + let presence = 1.0 / (1.0 + (-m.object_score).exp()); + Ok(TrackedFrame { prob, w, h, confidence: (presence * m.iou.clamp(0.0, 1.0)).clamp(0.0, 1.0), occluded: m.occluded }) +} + +impl Sam2Mask { + fn run_track(&self, req: &TrackRequest, frames: Frames<'_>, cancel: &dyn Fn() -> bool, out: &mut dyn FnMut(i64, TrackedFrame)) -> Result<()> { + if req.count == 0 || req.count > MAX_CLIP { + return Err(Error::Input(format!("a tracking request needs 1 to {MAX_CLIP} frames, got {}", req.count))); + } + if req.prompts.is_empty() { + return Err(Error::Input("tracking needs a prompt on at least one frame".into())); + } + let sizes = Rc::new(RefCell::new(HashMap::new())); + let source = Source { first: req.first, count: req.count, fetch: RefCell::new(frames), sizes: sizes.clone() }; + // Prompt coordinates are in the frame's pixels and the model works at 1024 x 1024, so the + // prompted frames' sizes are needed first. + let mut prompt_frames = vec![]; + for (frame, p) in &req.prompts { + let idx = usize::try_from(frame - req.first) + .ok() + .filter(|i| *i < req.count) + .ok_or_else(|| Error::Input(format!("prompt frame {frame} is outside the clip")))?; + let f = (source.fetch.borrow_mut())(*frame).ok_or_else(|| Error::Input(format!("frame {frame} is not available")))?; + sizes.borrow_mut().insert(idx, (f.w, f.h)); + prompt_frames.push((idx, p)); + } + let mut tracker = VideoTracker::new(self.video_model().clone(), Box::new(source))?; + let object = ObjectId(1); + let mut prompted = vec![]; + for (idx, p) in prompt_frames { + if cancel() { + return Err(Error::Cancelled); + } + let (w, h) = sizes.borrow().get(&idx).copied().ok_or_else(|| Error::Compute("frame size unknown".into()))?; + let (sx, sy) = (INPUT_SIZE as f32 / w as f32, INPUT_SIZE as f32 / h as f32); + let prompt = Prompt { + points: p.points.iter().map(|(c, fg)| ([c[0] * sx, c[1] * sy], *fg)).collect(), + bbox: p.bbox.map(|b| [b[0] * sx, b[1] * sy, b[2] * sx, b[3] * sy]), + }; + tracker.prompt(idx, object, &prompt)?; + prompted.push(idx); + } + let first = prompted.iter().copied().min().unwrap_or(0); + let flag = Cancel::new(); + let mut failure: Option = None; + let mut emit = |m: FrameMask| { + if cancel() { + flag.cancel(); + } + match to_frame(&m, sizes.borrow().get(&m.frame).copied().unwrap_or((0, 0))) { + Ok(t) => out(req.first + m.frame as i64, t), + Err(e) => { + failure.get_or_insert(e); + flag.cancel(); + } + } + }; + tracker.propagate(first, Direction::Forward, None, &flag, &mut emit)?; + if first > 0 { + // frame `first` was reported by the forward pass; the backward pass starts there again + let mut skip_first = true; + tracker.propagate(first, Direction::Backward, None, &flag, &mut |m| { + if std::mem::take(&mut skip_first) { + return; + } + emit(m); + })?; + } + failure.map_or(Ok(()), Err) + } +} + +impl ObjectTracker for Sam2Mask { + fn track( + &self, + req: &TrackRequest, + frames: &mut dyn FnMut(i64) -> Option>, + cancel: &dyn Fn() -> bool, + out: &mut dyn FnMut(i64, TrackedFrame), + ) -> effectcraft_segment::Result<()> { + self.run_track(req, frames, cancel, out).map_err(|e| e.to_string()) + } +} diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index 63c1a06ed..ce1fd14d8 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -337,3 +337,63 @@ fn sam2_works_as_a_mask_model() { 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> = (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::>(), (100..100 + n as i64).collect::>()); + 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()); +} diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index dbfe07c1f..37b3a4ca0 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -45,6 +45,54 @@ pub trait MaskModel: Send + Sync { fn info(&self) -> &'static ModelInfo; /// Foreground probability (0–1) of each pixel of a frame (straight RGB 0–1, row-major `w×h`). fn segment(&self, rgb: &[[f32; 3]], w: usize, h: usize, prompt: &Prompt) -> Result>; + /// A video tracker built on the same weights, if the model has one. + fn tracker(&self) -> Option<&dyn ObjectTracker> { + None + } +} + +/// One decoded frame handed to a tracker (straight RGB 0–1, row-major `w×h`). +#[derive(Clone, Debug, Default)] +pub struct ClipFrame { + pub rgb: Vec<[f32; 3]>, + pub w: usize, + pub h: usize, +} + +/// What to track: the clip's frames `first..first + count` (the host's frame numbers) and the +/// prompts of the one object, each on its frame. +#[derive(Clone, Debug)] +pub struct TrackRequest { + pub first: i64, + pub count: usize, + pub prompts: Vec<(i64, Prompt)>, +} + +/// A tracker's answer for one frame. +#[derive(Clone, Debug, Default)] +pub struct TrackedFrame { + /// Foreground probability of every pixel (`w×h`); all zero when the object is absent. + pub prob: Vec, + pub w: usize, + pub h: usize, + /// How sure the tracker is (0–1), for the timeline. + pub confidence: f32, + /// The object is not in this frame (occluded or out of view). + pub occluded: bool, +} + +/// A video tracker with memory (SAM 2): given prompts on some frames, it follows the object +/// through the clip in both directions. Optional: a [`MaskModel`] may offer one. +pub trait ObjectTracker: Send + Sync { + /// Track the object. `frames(i)` supplies frame `i` (`None`: unavailable); `cancel()` is polled + /// between frames; `out(i, result)` is called as soon as each frame is finished. + fn track( + &self, + req: &TrackRequest, + frames: &mut dyn FnMut(i64) -> Option>, + cancel: &dyn Fn() -> bool, + out: &mut dyn FnMut(i64, TrackedFrame), + ) -> Result<()>; } /// What a model is for.