From 8a6d48cdc8910a7da1351757efd1cb9740f67040 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Sat, 17 Oct 2026 20:28:45 +0000 Subject: [PATCH] S2.74: bound pending multi-object tracking masks --- crates/effects/src/roto.rs | 113 ++++++++++++++++++++--- crates/effects/src/roto/pending.rs | 138 +++++++++++++++++++++++++++++ docs/roto-brush-sam2.md | 5 ++ 3 files changed, 244 insertions(+), 12 deletions(-) create mode 100644 crates/effects/src/roto/pending.rs diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index e3f6da171..e8cae42c3 100644 --- a/crates/effects/src/roto.rs +++ b/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>, cancel: &dyn Fn() -> bool, progress: &mut dyn FnMut(usize), + ) -> Result, 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>, + cancel: &dyn Fn() -> bool, + progress: &mut dyn FnMut(usize), + max_pending_bytes: usize, ) -> Result, 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>> = BTreeMap::new(); + let mut pending = pending::PendingMasks::new(objects.len(), max_pending_bytes); let finish = |g: i64, parts: Vec>, 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> { + 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>, + 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(¶ms, [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(); diff --git a/crates/effects/src/roto/pending.rs b/crates/effects/src/roto/pending.rs new file mode 100644 index 000000000..256789ee9 --- /dev/null +++ b/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>>, + 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::>()).saturating_add(256) + } + + fn mask_bytes(mask: &TrackedFrame) -> usize { + mask.prob.capacity().saturating_mul(size_of::()) + } + + /// 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>>, 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>> { + 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()); + } +} diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index 72c48e4e1..a80052bab 100644 --- a/docs/roto-brush-sam2.md +++ b/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