diff --git a/crates/sam2/src/session.rs b/crates/sam2/src/session.rs index c629a7b20..2e45cc407 100644 --- a/crates/sam2/src/session.rs +++ b/crates/sam2/src/session.rs @@ -133,11 +133,34 @@ struct FrameFeatures { high: (Tensor, Tensor), } +/// What a tracker may hold and take. A call that would exceed it stops with [`Error::Budget`]. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Budget { + /// Memory banks, results and cached frame features (bytes). + pub max_state_bytes: usize, + /// Wall-clock seconds one `propagate` call may take (`None`: no limit). + pub max_seconds: Option, +} + +impl Default for Budget { + fn default() -> Budget { + Budget { max_state_bytes: 2 << 30, max_seconds: None } + } +} + +/// Bytes of one stored frame: memory features and position encoding `[4096, 64]`, the pointer. +const FRAME_BYTES: usize = 2 * GRID * GRID * 64 * 4 + 256 * 4; +/// Bytes of a conditioning frame's kept result (`256 x 256` logits). +const RESULT_BYTES: usize = 4 * GRID * 4 * GRID * 4; +/// Bytes of one frame's cached neck features (lowest level and position, two skip levels). +const FEATURE_BYTES: usize = (2 * 256 * GRID * GRID + 32 * 16 * GRID * GRID + 64 * 4 * GRID * GRID) * 4; + pub struct VideoTracker<'a> { model: Arc, frames: Box, objects: BTreeMap, cache: VecDeque<(usize, Arc)>, + budget: Budget, } fn sine_embedding_1d(positions: &[f32], dim: usize) -> Result { @@ -161,13 +184,32 @@ impl<'a> VideoTracker<'a> { if n == 0 || n > MAX_FRAMES { return Err(Error::Input(format!("a clip needs between 1 and {MAX_FRAMES} frames, got {n}"))); } - Ok(VideoTracker { model, frames, objects: BTreeMap::new(), cache: VecDeque::new() }) + Ok(VideoTracker { model, frames, objects: BTreeMap::new(), cache: VecDeque::new(), budget: Budget::default() }) } pub fn num_frames(&self) -> usize { self.frames.num_frames() } + /// Replace the limits (see [`Budget`]). + pub fn set_budget(&mut self, budget: Budget) { + self.budget = budget; + } + + /// Bytes held now: stored frames, kept results and cached features. + pub fn state_bytes(&self) -> usize { + let frames: usize = self.objects.values().map(|o| (o.cond.len() + o.non_cond.len()) * FRAME_BYTES + o.cond.len() * RESULT_BYTES).sum(); + frames + self.cache.len() * FEATURE_BYTES + } + + fn check_budget(&self) -> Result<()> { + let used = self.state_bytes(); + if used > self.budget.max_state_bytes { + return Err(Error::Budget(format!("the tracker holds {} MB, over its {} MB limit", used >> 20, self.budget.max_state_bytes >> 20))); + } + Ok(()) + } + /// Forget every object and cached frame (the clip stays). pub fn reset(&mut self) { self.objects.clear(); @@ -298,6 +340,7 @@ impl<'a> VideoTracker<'a> { state.non_cond.insert(frame, entry); Self::prune(state, frame); } + self.check_budget()?; Ok(result) } @@ -383,6 +426,7 @@ impl<'a> VideoTracker<'a> { state.non_cond.insert(frame, entry); Self::prune(state, frame); } + self.check_budget()?; Ok(result) } @@ -410,6 +454,7 @@ impl<'a> VideoTracker<'a> { if ids.is_empty() { return Err(Error::Input("no object has been prompted".into())); } + let started = std::time::Instant::now(); for k in 0..=span { let frame = if reverse { start.checked_sub(k) } else { Some(start + k).filter(|f| *f < n) }; let Some(frame) = frame else { break }; @@ -417,6 +462,9 @@ impl<'a> VideoTracker<'a> { if cancel.is_cancelled() { return Err(Error::Cancelled); } + if let Some(limit) = self.budget.max_seconds.filter(|l| started.elapsed().as_secs_f64() > *l) { + return Err(Error::Budget(format!("tracking took more than {limit} s"))); + } let cached = self.objects.get(id).is_some_and(|o| o.cond.contains_key(&frame)); let stored = self.objects.get(id).and_then(|o| o.cond.get(&frame)).and_then(|o| o.result.clone()); let mask = match stored { diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index 4d4377374..230de7389 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -544,3 +544,51 @@ fn a_mask_prompt_through_the_object_tracker_interface() { let bad = TrackRequest::single(0, n, vec![(1, Prompt { points: vec![], bbox: None, mask: Some(vec![1.0; 5]) })]); assert!(model.tracker().unwrap().track(&bad, &mut get, &|| false, &mut |_, _, _| {}).is_err()); } + +/// The same prompts on the same frames give the same masks, bit for bit, run after run (CPU). +#[test] +fn tracking_is_deterministic() { + use effectcraft_sam2::prompt::Prompt; + use effectcraft_sam2::session::{Cancel, Direction, ObjectId, VideoModel, VideoTracker}; + let Some((w, r)) = load_video() else { return }; + let model = std::sync::Arc::new(VideoModel::load(&w).unwrap()); + let clip = r.get("clip.pixel_values").unwrap().narrow(0, 1, 3).unwrap(); + let run = || { + let mut t = VideoTracker::new(model.clone(), Box::new(Clip(clip.clone()))).unwrap(); + t.prompt(0, ObjectId(1), &Prompt { points: vec![([416.0, 512.0], true)], bbox: None }).unwrap(); + let mut out = vec![]; + t.propagate(0, Direction::Forward, None, &Cancel::new(), &mut |m| out.push(m)).unwrap(); + out + }; + let (a, b) = (run(), run()); + assert_eq!(a.len(), 3); + for (x, y) in a.iter().zip(&b) { + assert_eq!(x.logits, y.logits, "frame {} differs between runs", x.frame); + assert_eq!((x.iou.to_bits(), x.object_score.to_bits()), (y.iou.to_bits(), y.object_score.to_bits())); + } +} + +/// Hard limits: a tracker over its memory budget or past its time budget stops with Error::Budget. +#[test] +fn the_tracker_stops_at_its_budget() { + use effectcraft_sam2::Error; + use effectcraft_sam2::prompt::Prompt; + use effectcraft_sam2::session::{Budget, Cancel, Direction, ObjectId, VideoModel, VideoTracker}; + let Some((w, r)) = load_video() else { return }; + let model = std::sync::Arc::new(VideoModel::load(&w).unwrap()); + let clip = r.get("clip.pixel_values").unwrap().narrow(0, 1, 3).unwrap(); + let click = Prompt { points: vec![([416.0, 512.0], true)], bbox: None }; + // memory: one frame's worth is more than a 1 MB limit + let mut t = VideoTracker::new(model.clone(), Box::new(Clip(clip.clone()))).unwrap(); + t.set_budget(Budget { max_state_bytes: 1 << 20, max_seconds: None }); + assert!(matches!(t.prompt(0, ObjectId(1), &click), Err(Error::Budget(_)))); + // default limit: fine, and the bytes held are reported + let mut t = VideoTracker::new(model.clone(), Box::new(Clip(clip.clone()))).unwrap(); + t.prompt(0, ObjectId(1), &click).unwrap(); + let held = t.state_bytes(); + assert!(held > 2 << 20 && held < 200 << 20, "{held}"); + // time: a limit of zero stops the first propagation + t.set_budget(Budget { max_state_bytes: 2 << 30, max_seconds: Some(0.0) }); + let r = t.propagate(0, Direction::Forward, None, &Cancel::new(), &mut |_| {}); + assert!(matches!(r, Err(Error::Budget(_))), "{r:?}"); +}