S2.74: bound pending multi-object tracking masks

crates/effects/src/roto.rs

@@ -40,6 +40,8 @@ use crate::{Buf, EffectCtx, EffectSpec, Params, num, p, popup, slider};
 
 pub use effectcraft_track::roto::{FrameSeg, Stroke, StrokeKind};
 
+mod pending;
+
 /// The effect's id.
 pub const ID: &str = "ec.matte.rotobrush";
 /// Hidden parameters.
@@ -607,6 +609,16 @@ impl Chain {
         frame: &mut dyn FnMut(i64) -> Option<Arc<Image>>,
         cancel: &dyn Fn() -> bool,
         progress: &mut dyn FnMut(usize),
+    ) -> Result<Option<TrackSummary>, String> {
+        self.track_with_pending_limit(frame, cancel, progress, pending::MAX_PENDING_BYTES)
+    }
+
+    fn track_with_pending_limit(
+        &self,
+        frame: &mut dyn FnMut(i64) -> Option<Arc<Image>>,
+        cancel: &dyn Fn() -> bool,
+        progress: &mut dyn FnMut(usize),
+        max_pending_bytes: usize,
     ) -> Result<Option<TrackSummary>, String> {
         let Some(tracker) = self.model.as_ref().and_then(|m| m.tracker()) else { return Ok(None) };
         let (Some(_), true) = (self.data.base, self.data.span[1] >= self.data.span[0]) else { return Err("no base frame".into()) };
@@ -657,7 +669,7 @@ impl Chain {
         let frame = std::cell::RefCell::new(frame);
         let mut summary = TrackSummary::default();
         // A frame is finished when every object has reported it (the objects' mattes are united).
-        let mut pending: BTreeMap<i64, Vec<Option<effectcraft_segment::TrackedFrame>>> = BTreeMap::new();
+        let mut pending = pending::PendingMasks::new(objects.len(), max_pending_bytes);
         let finish = |g: i64, parts: Vec<Option<effectcraft_segment::TrackedFrame>>, summary: &mut TrackSummary, progress: &mut dyn FnMut(usize)| {
             let (Some(key), Some(img)) = (self.keys.get(&g).copied(), (frame.borrow_mut())(g)) else { return };
             let mut segs: Vec<(u32, effectcraft_track::roto::FrameSeg)> = vec![];
@@ -694,27 +706,37 @@ impl Chain {
             summary.min_confidence = if summary.frames == 1 { confidence } else { summary.min_confidence.min(confidence) };
             progress(summary.frames);
         };
-        let n_objects = objects.len();
-        tracker.track(
+        let pending_failed = std::cell::Cell::new(false);
+        let mut pending_error = None;
+        let result = tracker.track(
             &req,
             &mut |g| {
                 let img = (frame.borrow_mut())(g)?;
                 Some(Arc::new(ClipFrame { rgb: rb::rgb_of(&img), w: img.width as usize, h: img.height as usize }))
             },
-            cancel,
+            &|| cancel() || pending_failed.get(),
             &mut |o, g, t| {
-                let parts = pending.entry(g).or_insert_with(|| vec![None; n_objects]);
-                if let Some(slot) = parts.get_mut(o) {
-                    *slot = Some(t);
+                // A tracker that reports object-by-object can otherwise retain a full clip of
+                // probability maps. Stop buffering immediately, even if it ignores cancellation.
+                if pending_failed.get() {
+                    return;
                 }
-                if parts.iter().all(Option::is_some) {
-                    let parts = pending.remove(&g).unwrap_or_default();
-                    finish(g, parts, &mut summary, progress);
+                match pending.insert(o, g, t) {
+                    Ok(Some(parts)) => finish(g, parts, &mut summary, progress),
+                    Ok(None) => {}
+                    Err(e) => {
+                        pending_error = Some(e);
+                        pending_failed.set(true);
+                    }
                 }
             },
-        )?;
+        );
+        if let Some(e) = pending_error {
+            return Err(e);
+        }
+        result?;
         // frames some object never reported (it was not followed there): united from the rest
-        for (g, parts) in std::mem::take(&mut pending) {
+        for (g, parts) in pending.into_incomplete() {
             finish(g, parts, &mut summary, progress);
         }
         Ok(Some(summary))
@@ -1285,6 +1307,73 @@ mod tests {
         }
     }
 
+    /// The ordering SAM 2 uses when objects are first prompted on different frames: each object
+    /// reports its entire span before the next. No weights or large allocations are needed.
+    #[derive(Default)]
+    struct SeparatePasses {
+        emitted: std::sync::atomic::AtomicUsize,
+        cancelled: std::sync::atomic::AtomicBool,
+    }
+
+    impl effectcraft_segment::MaskModel for SeparatePasses {
+        fn info(&self) -> &'static effectcraft_segment::ModelInfo {
+            &effectcraft_segment::SAM2_TINY
+        }
+        fn segment(&self, _: &[[f32; 3]], _: usize, _: usize, _: &effectcraft_segment::Prompt) -> effectcraft_segment::Result<Vec<f32>> {
+            Err("no weights".into())
+        }
+        fn tracker(&self) -> Option<&dyn effectcraft_segment::ObjectTracker> {
+            Some(self)
+        }
+    }
+
+    impl effectcraft_segment::ObjectTracker for SeparatePasses {
+        fn track(
+            &self,
+            req: &effectcraft_segment::TrackRequest,
+            frames: &mut dyn FnMut(i64) -> Option<Arc<ClipFrame>>,
+            cancel: &dyn Fn() -> bool,
+            out: &mut dyn FnMut(usize, i64, effectcraft_segment::TrackedFrame),
+        ) -> effectcraft_segment::Result<()> {
+            use std::sync::atomic::Ordering::Relaxed;
+            assert_ne!(req.objects[0][0].0, req.objects[1][0].0);
+            for o in 0..req.objects.len() {
+                for g in req.first..req.first + req.count as i64 {
+                    if cancel() {
+                        self.cancelled.store(true, Relaxed);
+                        return Err("tracking cancelled".into());
+                    }
+                    let img = frames(g).ok_or("no frame")?;
+                    self.emitted.fetch_add(1, Relaxed);
+                    out(o, g, effectcraft_segment::TrackedFrame { prob: vec![0.5; img.w * img.h], w: img.w, h: img.h, confidence: 0.9, occluded: false });
+                }
+            }
+            Ok(())
+        }
+    }
+
+    #[test]
+    fn separate_object_passes_cancel_at_the_pending_mask_budget() {
+        use std::sync::atomic::Ordering::Relaxed;
+        let mut d = RotoData::default();
+        for object in 0..2 {
+            d.add(Stroke { kind: StrokeKind::Fg, frame: i64::from(object), radius: 3.0, points: vec![[35.0, 40.0]], object }, 8, [0, 24]);
+        }
+        let mut params = params_with(&d);
+        params.values.insert(INPUT_KEY.into(), Value::Str("pending-budget-regression".into()));
+        let mut chain = Chain::new(&params, [100.0, 80.0], 25.0, 1.0);
+        let fake = Arc::new(SeparatePasses::default());
+        chain.model = Some(fake.clone());
+        let mut get = |f: i64| Some(Arc::new(frame(f as f64 / 25.0)));
+        let mut completed = 0;
+        // Each map takes 32,000 bytes; two incomplete frames exceed this small test budget.
+        let error = chain.track_with_pending_limit(&mut get, &|| false, &mut |n| completed = n, 64_000).unwrap_err();
+        assert!(error.contains("pending object masks") && error.contains("memory budget"), "{error}");
+        assert!(fake.cancelled.load(Relaxed), "the model must see cancellation before allocating the rest of the clip");
+        assert_eq!(fake.emitted.load(Relaxed), 2);
+        assert_eq!(completed, 0, "partial objects are not complete tracked frames");
+    }
+
     #[test]
     fn objects_are_followed_in_one_pass_and_their_mattes_united() {
         let mut d = RotoData::default();

crates/effects/src/roto/pending.rs

@@ -0,0 +1,138 @@
+//! Bounded assembly of the objects a tracker reports separately for each frame.
+
+use std::collections::BTreeMap;
+use std::mem::size_of;
+
+use effectcraft_segment::TrackedFrame;
+
+/// Separate from the model's memory banks: probability maps are full-resolution and can otherwise
+/// grow with the entire clip when objects start on different frames.
+pub(super) const MAX_PENDING_BYTES: usize = 256 << 20;
+
+pub(super) struct PendingMasks {
+    frames: BTreeMap<i64, Vec<Option<TrackedFrame>>>,
+    objects: usize,
+    bytes: usize,
+    limit: usize,
+}
+
+impl PendingMasks {
+    pub(super) fn new(objects: usize, limit: usize) -> Self {
+        Self { frames: BTreeMap::new(), objects, bytes: 0, limit }
+    }
+
+    /// Probability allocation plus a conservative allowance for the map node and object slots.
+    fn frame_bytes(&self) -> usize {
+        self.objects.saturating_mul(size_of::<Option<TrackedFrame>>()).saturating_add(256)
+    }
+
+    fn mask_bytes(mask: &TrackedFrame) -> usize {
+        mask.prob.capacity().saturating_mul(size_of::<f32>())
+    }
+
+    /// Complete frames leave immediately; incomplete frames may never exceed the budget. A rejected
+    /// insertion leaves the buffer unchanged, including when it replaces an object's previous result.
+    pub(super) fn insert(&mut self, object: usize, frame: i64, mask: TrackedFrame) -> Result<Option<Vec<Option<TrackedFrame>>>, String> {
+        if object >= self.objects {
+            return Err("tracker reported an unknown object".into());
+        }
+        if self.objects == 1 {
+            return Ok(Some(vec![Some(mask)]));
+        }
+        let existing = self.frames.get(&frame);
+        let old_bytes = existing.and_then(|p| p.get(object)).and_then(Option::as_ref).map_or(0, Self::mask_bytes);
+        let added = Self::mask_bytes(&mask).saturating_add(if existing.is_none() { self.frame_bytes() } else { 0 });
+        let bytes = self.bytes.saturating_sub(old_bytes).checked_add(added);
+        let Some(bytes) = bytes.filter(|n| *n <= self.limit) else {
+            return Err(format!("pending object masks exceed their {} MiB memory budget", self.limit >> 20));
+        };
+        let parts = self.frames.entry(frame).or_insert_with(|| vec![None; self.objects]);
+        if let Some(slot) = parts.get_mut(object) {
+            *slot = Some(mask);
+        }
+        self.bytes = bytes;
+        if parts.iter().all(Option::is_some) {
+            let parts = self.frames.remove(&frame).unwrap_or_default();
+            let used = parts.iter().flatten().fold(self.frame_bytes(), |n, m| n.saturating_add(Self::mask_bytes(m)));
+            self.bytes = self.bytes.saturating_sub(used);
+            return Ok(Some(parts));
+        }
+        Ok(None)
+    }
+
+    pub(super) fn into_incomplete(self) -> BTreeMap<i64, Vec<Option<TrackedFrame>>> {
+        self.frames
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn mask(value: f32) -> TrackedFrame {
+        TrackedFrame { prob: vec![value; 16], w: 4, h: 4, confidence: value, occluded: false }
+    }
+
+    #[test]
+    fn different_start_passes_stop_before_buffering_the_whole_clip() {
+        let mut pending = PendingMasks::new(2, 2048);
+        let mut accepted = 0;
+        // SAM 2 can report all of object 0 before starting object 1. A long clip must stop
+        // at the memory limit even though no frame has received all its objects yet.
+        for frame in 0..10_000 {
+            match pending.insert(0, frame, mask(0.9)) {
+                Ok(None) => accepted += 1,
+                Ok(Some(_)) => panic!("no frame is complete yet"),
+                Err(e) => {
+                    assert!(e.contains("memory budget"), "{e}");
+                    break;
+                }
+            }
+        }
+        assert!(accepted > 0 && accepted < 10_000);
+        assert_eq!(pending.frames.len(), accepted);
+        assert!(pending.bytes <= pending.limit);
+        let before = pending.bytes;
+        assert!(pending.insert(0, 10_001, mask(0.8)).is_err());
+        assert_eq!(pending.bytes, before, "rejection must not retain the new map");
+    }
+
+    #[test]
+    fn interleaved_objects_release_the_budget_and_keep_their_masks() {
+        let mut pending = PendingMasks::new(2, 2048);
+        for frame in 0..10_000 {
+            let first = (frame % 2) as usize;
+            assert!(pending.insert(first, frame, mask(first as f32)).unwrap().is_none());
+            let parts = pending.insert(1 - first, frame, mask((1 - first) as f32)).unwrap().unwrap();
+            assert_eq!(parts[0].as_ref().unwrap().prob, vec![0.0; 16]);
+            assert_eq!(parts[1].as_ref().unwrap().prob, vec![1.0; 16]);
+            assert!(pending.frames.is_empty());
+            assert_eq!(pending.bytes, 0);
+        }
+    }
+
+    #[test]
+    fn replacements_charge_capacity_and_partial_frames_are_preserved() {
+        let mut pending = PendingMasks::new(2, 2048);
+        pending.insert(0, 7, mask(0.1)).unwrap();
+        let before = pending.bytes;
+        pending.insert(0, 7, mask(0.2)).unwrap();
+        assert_eq!(pending.bytes, before, "a duplicate replaces rather than accumulates");
+        let mut oversized = mask(0.3);
+        oversized.prob.reserve(2048);
+        assert!(pending.insert(0, 7, oversized).is_err(), "capacity, not length, owns the memory");
+        assert_eq!(pending.bytes, before);
+        let frames = pending.into_incomplete();
+        assert_eq!(frames[&7][0].as_ref().unwrap().prob, vec![0.2; 16]);
+        assert!(frames[&7][1].is_none());
+    }
+
+    #[test]
+    fn single_objects_need_no_pending_budget_and_unknown_objects_are_rejected() {
+        let mut pending = PendingMasks::new(1, 0);
+        assert!(pending.insert(0, 0, mask(0.9)).unwrap().is_some());
+        assert_eq!(pending.bytes, 0);
+        assert!(pending.insert(1, 0, mask(0.9)).is_err());
+        assert!(pending.frames.is_empty());
+    }
+}

docs/roto-brush-sam2.md

@@ -120,6 +120,11 @@ stored frame, pruned to the frames later ones can still look at) and refuses to
 one is set, with a clear error instead of running out of memory. Two runs of the same prompts give bit-identical masks on
 the CPU.
 
+Pending per-object probability maps have a separate 256 MiB budget, including their allocated capacity and frame
+slots. This bounds the full-resolution results waiting to be united when objects first prompted on different frames
+report in separate passes. Reaching the budget cancels tracking and reports the reason through the existing
+`trackerNote`; propagation falls back to the classic engine. A single object's results need no pending buffer.
+
 ## How it works
 `crates/sam2` is a Rust port (on [candle](https://github.com/huggingface/candle), native only) of the Hugging Face
 `transformers` SAM 2 implementation, Apache-2.0 (see its `NOTICE`). Every stage is checked against PyTorch by