S2.43: track several objects in one SAM 2 pass

TrackRequest takes a list of objects (each with its prompts on its frames) and
the tracker reports (object, frame, result); frame encodings are shared, objects
may start on different frames. VideoTracker::propagate_objects runs a subset.
Limits: 32 objects. Test: two discs, one prompted on frame 0 and one on frame 2.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

crates/effects/src/roto.rs

@@ -569,7 +569,7 @@ impl Chain {
         if prompts.is_empty() {
             return Err("no foreground strokes to follow".into());
         }
-        let req = TrackRequest { first: self.data.span[0], count: (self.data.span[1] - self.data.span[0] + 1) as usize, prompts };
+        let req = TrackRequest::single(self.data.span[0], (self.data.span[1] - self.data.span[0] + 1) as usize, prompts);
         let frame = std::cell::RefCell::new(frame);
         let mut summary = TrackSummary::default();
         tracker.track(
@@ -579,7 +579,7 @@ impl Chain {
                 Some(Arc::new(ClipFrame { rgb: rb::rgb_of(&img), w: img.width as usize, h: img.height as usize }))
             },
             cancel,
-            &mut |g, t| {
+            &mut |_, g, t| {
                 let (Some(key), Some(img)) = (self.keys.get(&g).copied(), (frame.borrow_mut())(g)) else { return };
                 let strokes = self.data.strokes_at(g);
                 let seg = rb::segment_tracked(&img, &strokes, t.prob, t.occluded, &self.opts, self.refine_radius);

crates/engine/src/tests/roto.rs

@@ -778,7 +778,7 @@ impl effectcraft_segment::ObjectTracker for Exact {
         req: &effectcraft_segment::TrackRequest,
         frame: &mut dyn FnMut(i64) -> Option<std::sync::Arc<effectcraft_segment::ClipFrame>>,
         _: &dyn Fn() -> bool,
-        out: &mut dyn FnMut(i64, effectcraft_segment::TrackedFrame),
+        out: &mut dyn FnMut(usize, i64, effectcraft_segment::TrackedFrame),
     ) -> effectcraft_segment::Result<()> {
         for g in req.first..req.first + req.count as i64 {
             let Some(img) = frame(g) else { return Err(format!("no frame {g}")) };
@@ -786,7 +786,7 @@ impl effectcraft_segment::ObjectTracker for Exact {
             let prob = (0..img.w * img.h)
                 .map(|i| ((((i % img.w) as f64 + 0.5 - c[0]).hypot((i / img.w) as f64 + 0.5 - c[1])) <= R) as u8 as f32)
                 .collect::<Vec<f32>>();
-            out(g, effectcraft_segment::TrackedFrame { prob, w: img.w, h: img.h, confidence: 0.99, occluded: false });
+            out(0, g, effectcraft_segment::TrackedFrame { prob, w: img.w, h: img.h, confidence: 0.99, occluded: false });
         }
         Ok(())
     }
@@ -889,7 +889,7 @@ fn a_failing_tracker_falls_back_to_the_classic_engine() {
             _: &TrackRequest,
             _: &mut dyn FnMut(i64) -> Option<Arc<ClipFrame>>,
             _: &dyn Fn() -> bool,
-            _: &mut dyn FnMut(i64, TrackedFrame),
+            _: &mut dyn FnMut(usize, i64, TrackedFrame),
         ) -> effectcraft_segment::Result<()> {
             Err("the tracker broke".into())
         }

crates/sam2/src/session.rs

@@ -334,11 +334,24 @@ impl<'a> VideoTracker<'a> {
     /// 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<usize>, cancel: &Cancel, out: &mut dyn FnMut(FrameMask)) -> Result<()> {
+        self.propagate_objects(None, start, direction, max_frames, cancel, out)
+    }
+
+    /// [`VideoTracker::propagate`] for the listed objects only (`None`: every prompted object).
+    pub fn propagate_objects(
+        &mut self,
+        only: Option<&[ObjectId]>,
+        start: usize,
+        direction: Direction,
+        max_frames: Option<usize>,
+        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<ObjectId> = self.objects.iter().filter(|(_, o)| !o.cond.is_empty()).map(|(id, _)| *id).collect();
+        let ids: Vec<ObjectId> = self.objects.iter().filter(|(id, o)| !o.cond.is_empty() && only.is_none_or(|l| l.contains(id))).map(|(id, _)| *id).collect();
         if ids.is_empty() {
             return Err(Error::Input("no object has been prompted".into()));
         }

crates/sam2/src/track.rs

@@ -50,32 +50,41 @@ fn to_frame(m: &FrameMask, size: (usize, usize)) -> Result<TrackedFrame> {
     Ok(TrackedFrame { prob, w, h, confidence: (presence * m.iou.clamp(0.0, 1.0)).clamp(0.0, 1.0), occluded: m.occluded })
 }
 
+/// Most objects one request may track.
+const MAX_OBJECTS: usize = 32;
+
 impl Sam2Mask {
-    fn run_track(&self, req: &TrackRequest, frames: Frames<'_>, cancel: &dyn Fn() -> bool, out: &mut dyn FnMut(i64, TrackedFrame)) -> Result<()> {
+    fn run_track(&self, req: &TrackRequest, frames: Frames<'_>, cancel: &dyn Fn() -> bool, out: &mut dyn FnMut(usize, 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()));
+        if req.objects.is_empty() || req.objects.len() > MAX_OBJECTS || req.objects.iter().any(|o| o.is_empty()) {
+            return Err(Error::Input(format!("tracking needs 1 to {MAX_OBJECTS} objects, each with a prompt on at least one frame")));
         }
         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));
+        for (oi, prompts) in req.objects.iter().enumerate() {
+            for (frame, p) in 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")))?;
+                if !sizes.borrow().contains_key(&idx) {
+                    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((oi, 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 {
+        let object_id = |oi: usize| ObjectId(oi as u32 + 1);
+        // each object's first prompted frame: its pass starts there, forwards and backwards
+        let mut starts: std::collections::BTreeMap<usize, Vec<ObjectId>> = Default::default();
+        let mut first_of: HashMap<usize, usize> = HashMap::new();
+        for (oi, idx, p) in prompt_frames {
             if cancel() {
                 return Err(Error::Cancelled);
             }
@@ -85,34 +94,38 @@ impl Sam2Mask {
                 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);
+            tracker.prompt(idx, object_id(oi), &prompt)?;
+            let e = first_of.entry(oi).or_insert(idx);
+            *e = (*e).min(idx);
+        }
+        for (oi, idx) in first_of {
+            starts.entry(idx).or_default().push(object_id(oi));
         }
-        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();
             }
+            let oi = (m.object.0 as usize).saturating_sub(1);
             match to_frame(&m, sizes.borrow().get(&m.frame).copied().unwrap_or((0, 0))) {
-                Ok(t) => out(req.first + m.frame as i64, t),
+                Ok(t) => out(oi, 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);
-            })?;
+        for (start, ids) in starts {
+            tracker.propagate_objects(Some(&ids), start, Direction::Forward, None, &flag, &mut emit)?;
+            if start > 0 {
+                // the start frame was reported by the forward pass; the backward pass passes it again
+                tracker.propagate_objects(Some(&ids), start, Direction::Backward, None, &flag, &mut |m| {
+                    if m.frame != start {
+                        emit(m);
+                    }
+                })?;
+            }
         }
         failure.map_or(Ok(()), Err)
     }
@@ -124,7 +137,7 @@ impl ObjectTracker for Sam2Mask {
         req: &TrackRequest,
         frames: &mut dyn FnMut(i64) -> Option<Arc<ClipFrame>>,
         cancel: &dyn Fn() -> bool,
-        out: &mut dyn FnMut(i64, TrackedFrame),
+        out: &mut dyn FnMut(usize, i64, TrackedFrame),
     ) -> effectcraft_segment::Result<()> {
         self.run_track(req, frames, cancel, out).map_err(|e| e.to_string())
     }

crates/sam2/tests/reference.rs

@@ -364,13 +364,14 @@ fn sam2_tracks_a_clip_through_the_object_tracker_interface() {
     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 req = TrackRequest::single(100, n, 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| {
+        .track(&req, &mut get, &|| false, &mut |o, f, t| {
+            assert_eq!(o, 0);
             assert!(seen.insert(f, t).is_none(), "frame {f} reported twice");
         })
         .unwrap();
@@ -391,9 +392,67 @@ fn sam2_tracks_a_clip_through_the_object_tracker_interface() {
         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());
+    assert!(model.tracker().unwrap().track(&req, &mut get, &|| true, &mut |_, _, _| {}).is_err());
+    let none = TrackRequest::single(100, n, vec![]);
+    assert!(model.tracker().unwrap().track(&none, &mut get, &|| false, &mut |_, _, _| {}).is_err());
+    let outside = TrackRequest::single(100, n, 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());
+}
+
+/// Two objects in one pass: each follows its own disc, every frame is reported once per object, and
+/// objects may start on different frames.
+#[test]
+fn sam2_tracks_two_objects_in_one_pass() {
+    use effectcraft_segment::{ClipFrame, MaskModel, Prompt, TrackRequest};
+    let Some((w, _)) = load() else { return };
+    let (fw, fh, n) = (160usize, 120usize, 4usize);
+    // A red disc moving right and a blue one moving left, never touching.
+    let a = |t: usize| (36.0 + 6.0 * t as f32, 40.0f32);
+    let b = |t: usize| (124.0 - 6.0 * t as f32, 84.0f32);
+    let hit = |c: (f32, f32), x: usize, y: usize| ((x as f32 - c.0).powi(2) + (y as f32 - c.1).powi(2)).sqrt() < 20.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.45 + 0.1 * (x as f32 / 9.0).sin() * (y as f32 / 7.0).cos();
+                    rgb[y * fw + x] = if hit(a(t), x, y) {
+                        [0.9, 0.15, 0.1]
+                    } else if hit(b(t), x, y) {
+                        [0.1, 0.25, 0.9]
+                    } else {
+                        [bg, bg * 0.9, bg * 0.8]
+                    };
+                }
+            }
+            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()));
+    let click = |c: (f32, f32)| Prompt { points: vec![([c.0, c.1], true)], bbox: None, mask: None };
+    // object 0 is prompted on frame 0, object 1 on frame 2
+    let req = TrackRequest { first: 10, count: n, objects: vec![vec![(10, click(a(0)))], vec![(12, click(b(2)))]] };
+    let mut get = |f: i64| frames.get((f - 10) as usize).cloned();
+    let mut seen: std::collections::BTreeMap<(usize, i64), effectcraft_segment::TrackedFrame> = Default::default();
+    model.tracker().unwrap().track(&req, &mut get, &|| false, &mut |o, f, t| assert!(seen.insert((o, f), t).is_none(), "object {o} frame {f} twice")).unwrap();
+    assert_eq!(seen.len(), 2 * n, "{:?}", seen.keys().collect::<Vec<_>>());
+    for ((o, f), t) in &seen {
+        let k = (f - 10) as usize;
+        let c = if *o == 0 { a(k) } else { b(k) };
+        let (mut i, mut u) = (0usize, 0usize);
+        for y in 0..fh {
+            for x in 0..fw {
+                let (p, q) = (t.prob[y * fw + x] > 0.5, hit(c, x, y));
+                i += usize::from(p && q);
+                u += usize::from(p || q);
+            }
+        }
+        let j = i as f64 / u.max(1) as f64;
+        eprintln!("object {o} frame {f}: IoU {j:.3}, confidence {:.2}", t.confidence);
+        assert!(j > 0.85, "object {o} frame {f}: IoU {j}");
+    }
+    let too_many = TrackRequest { first: 10, count: n, objects: vec![vec![(10, click(a(0)))]; 40] };
+    assert!(model.tracker().unwrap().track(&too_many, &mut get, &|| false, &mut |_, _, _| {}).is_err());
+    let empty = TrackRequest { first: 10, count: n, objects: vec![] };
+    assert!(model.tracker().unwrap().track(&empty, &mut get, &|| false, &mut |_, _, _| {}).is_err());
 }

crates/segment/src/lib.rs

@@ -59,13 +59,22 @@ pub struct ClipFrame {
     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.
+/// What to track: the clip's frames `first..first + count` (the host's frame numbers) and, for
+/// each object, its prompts on their frames. Objects are tracked in one pass, so the cost of
+/// encoding each frame is shared.
 #[derive(Clone, Debug)]
 pub struct TrackRequest {
     pub first: i64,
     pub count: usize,
-    pub prompts: Vec<(i64, Prompt)>,
+    /// One entry per object: its prompts, each on its frame.
+    pub objects: Vec<Vec<(i64, Prompt)>>,
+}
+
+impl TrackRequest {
+    /// A request for a single object.
+    pub fn single(first: i64, count: usize, prompts: Vec<(i64, Prompt)>) -> TrackRequest {
+        TrackRequest { first, count, objects: vec![prompts] }
+    }
 }
 
 /// A tracker's answer for one frame.
@@ -84,14 +93,15 @@ pub struct TrackedFrame {
 /// 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.
+    /// Track the objects. `frames(i)` supplies frame `i` (`None`: unavailable); `cancel()` is polled
+    /// between frames; `out(object, i, result)` is called as soon as each object's frame is finished
+    /// (`object` indexes `req.objects`).
     fn track(
         &self,
         req: &TrackRequest,
         frames: &mut dyn FnMut(i64) -> Option<Arc<ClipFrame>>,
         cancel: &dyn Fn() -> bool,
-        out: &mut dyn FnMut(i64, TrackedFrame),
+        out: &mut dyn FnMut(usize, i64, TrackedFrame),
     ) -> Result<()>;
 }
 

crates/track/src/roto/bench.rs

@@ -220,11 +220,11 @@ pub fn run_tracker(sc: &Scenario, tracker: &str, model: &dyn MaskModel) -> Optio
     let opts = SegOpts::default();
     let refs: Vec<&Stroke> = sc.strokes.iter().collect();
     let clip: Vec<Arc<ClipFrame>> = sc.frames.iter().map(|i| Arc::new(ClipFrame { rgb: rgb_of(i), w: i.width as usize, h: i.height as usize })).collect();
-    let req = TrackRequest { first: 0, count: sc.frames.len(), prompts: vec![(sc.base as i64, tracker_prompt(&refs, 1.0))] };
+    let req = TrackRequest::single(0, sc.frames.len(), vec![(sc.base as i64, tracker_prompt(&refs, 1.0))]);
     let start = Instant::now();
     let mut mattes: Vec<Vec<u8>> = vec![vec![0u8; (W * H) as usize]; sc.frames.len()];
     let mut get = |f: i64| usize::try_from(f).ok().and_then(|i| clip.get(i).cloned());
-    let r = t.track(&req, &mut get, &|| false, &mut |f, tf| {
+    let r = t.track(&req, &mut get, &|| false, &mut |_, f, tf| {
         let (Ok(i), true) = (usize::try_from(f), tf.prob.len() == (W * H) as usize) else { return };
         let Some(img) = sc.frames.get(i) else { return };
         let strokes: Vec<&Stroke> = if i == sc.base { refs.clone() } else { vec![] };