diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index f4f5e6bd0..869e2c132 100644 --- a/crates/effects/src/roto.rs +++ b/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); diff --git a/crates/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index b18cd386d..4ba557781 100644 --- a/crates/engine/src/tests/roto.rs +++ b/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>, _: &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::>(); - 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>, _: &dyn Fn() -> bool, - _: &mut dyn FnMut(i64, TrackedFrame), + _: &mut dyn FnMut(usize, i64, TrackedFrame), ) -> effectcraft_segment::Result<()> { Err("the tracker broke".into()) } diff --git a/crates/sam2/src/session.rs b/crates/sam2/src/session.rs index 3440ff379..4b20b5c64 100644 --- a/crates/sam2/src/session.rs +++ b/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, 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, + 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 = self.objects.iter().filter(|(_, o)| !o.cond.is_empty()).map(|(id, _)| *id).collect(); + let ids: Vec = 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())); } diff --git a/crates/sam2/src/track.rs b/crates/sam2/src/track.rs index 9320fbbe4..a46f7f6c6 100644 --- a/crates/sam2/src/track.rs +++ b/crates/sam2/src/track.rs @@ -50,32 +50,41 @@ fn to_frame(m: &FrameMask, size: (usize, usize)) -> Result { 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> = Default::default(); + let mut first_of: HashMap = 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 = 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>, 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()) } diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index ce1fd14d8..da7dd4706 100644 --- a/crates/sam2/tests/reference.rs +++ b/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> = (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::>()); + 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()); } diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index 37b3a4ca0..0a9a0494a 100644 --- a/crates/segment/src/lib.rs +++ b/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>, +} + +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>, cancel: &dyn Fn() -> bool, - out: &mut dyn FnMut(i64, TrackedFrame), + out: &mut dyn FnMut(usize, i64, TrackedFrame), ) -> Result<()>; } diff --git a/crates/track/src/roto/bench.rs b/crates/track/src/roto/bench.rs index 810bf34e5..2ce0663df 100644 --- a/crates/track/src/roto/bench.rs +++ b/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> = 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![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![] };