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>

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.

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<Vec<f32>> {
         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.

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<Tensor>;
 }
@@ -129,9 +129,9 @@ struct FrameFeatures {
     high: (Tensor, Tensor),
 }
 
-pub struct VideoTracker {
+pub struct VideoTracker<'a> {
     model: Arc<VideoModel>,
-    frames: Box<dyn FrameSource>,
+    frames: Box<dyn FrameSource + 'a>,
     objects: BTreeMap<ObjectId, ObjectState>,
     cache: VecDeque<(usize, Arc<FrameFeatures>)>,
 }
@@ -151,8 +151,8 @@ fn scalar(t: &Tensor) -> Result<f32> {
     Ok(t.flatten_all()?.to_vec1::<f32>()?.first().copied().unwrap_or(0.0))
 }
 
-impl VideoTracker {
-    pub fn new(model: Arc<VideoModel>, frames: Box<dyn FrameSource>) -> Result<VideoTracker> {
+impl<'a> VideoTracker<'a> {
+    pub fn new(model: Arc<VideoModel>, frames: Box<dyn FrameSource + 'a>) -> Result<VideoTracker<'a>> {
         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}")));

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<Arc<ClipFrame>>;
+
+/// 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<Frames<'a>>,
+    sizes: Rc<RefCell<HashMap<usize, (usize, usize)>>>,
+}
+
+impl FrameSource for Source<'_> {
+    fn num_frames(&self) -> usize {
+        self.count
+    }
+
+    fn frame(&self, index: usize) -> Result<Tensor> {
+        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<TrackedFrame> {
+    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::<f32>()?;
+    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<Error> = 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<Arc<ClipFrame>>,
+        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())
+    }
+}

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

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<Vec<f32>>;
+    /// 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<f32>,
+    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<Arc<ClipFrame>>,
+        cancel: &dyn Fn() -> bool,
+        out: &mut dyn FnMut(i64, TrackedFrame),
+    ) -> Result<()>;
 }
 
 /// What a model is for.