diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index 0d7700a0e..6a0eb7fbf 100644 --- a/crates/sam2/src/lib.rs +++ b/crates/sam2/src/lib.rs @@ -21,6 +21,7 @@ pub mod neck; pub mod nn; pub mod prompt; pub mod session; +pub mod track; pub mod weights; /// Everything that can go wrong in this crate. Never a panic. diff --git a/crates/sam2/src/mask.rs b/crates/sam2/src/mask.rs index 8bb309da0..a0eac9cd6 100644 --- a/crates/sam2/src/mask.rs +++ b/crates/sam2/src/mask.rs @@ -105,6 +105,10 @@ impl MaskModel for Sam2Mask { fn segment(&self, rgb: &[[f32; 3]], w: usize, h: usize, prompt: &effectcraft_segment::Prompt) -> effectcraft_segment::Result> { self.probability(rgb, w, h, prompt).map_err(|e| e.to_string()) } + + fn tracker(&self) -> Option<&dyn effectcraft_segment::ObjectTracker> { + Some(self) + } } /// Verify the installed file against the registry and build the model. diff --git a/crates/sam2/src/session.rs b/crates/sam2/src/session.rs index 45d1ffd0a..3440ff379 100644 --- a/crates/sam2/src/session.rs +++ b/crates/sam2/src/session.rs @@ -33,7 +33,7 @@ const MAX_FRAMES: usize = 1_000_000; const GRID: usize = 64; /// Frames of the clip, already resized to the model input and normalised: `[1, 3, 1024, 1024]`. -pub trait FrameSource: Send { +pub trait FrameSource { fn num_frames(&self) -> usize; fn frame(&self, index: usize) -> Result; } @@ -129,9 +129,9 @@ struct FrameFeatures { high: (Tensor, Tensor), } -pub struct VideoTracker { +pub struct VideoTracker<'a> { model: Arc, - frames: Box, + frames: Box, objects: BTreeMap, cache: VecDeque<(usize, Arc)>, } @@ -151,8 +151,8 @@ fn scalar(t: &Tensor) -> Result { Ok(t.flatten_all()?.to_vec1::()?.first().copied().unwrap_or(0.0)) } -impl VideoTracker { - pub fn new(model: Arc, frames: Box) -> Result { +impl<'a> VideoTracker<'a> { + pub fn new(model: Arc, frames: Box) -> Result> { let n = frames.num_frames(); if n == 0 || n > MAX_FRAMES { return Err(Error::Input(format!("a clip needs between 1 and {MAX_FRAMES} frames, got {n}"))); diff --git a/crates/sam2/src/track.rs b/crates/sam2/src/track.rs new file mode 100644 index 000000000..9320fbbe4 --- /dev/null +++ b/crates/sam2/src/track.rs @@ -0,0 +1,131 @@ +// Modified work: adapts the ported SAM 2 video tracker to EffectCraft's `ObjectTracker` +// interface (Apache-2.0, see NOTICE). +//! SAM 2 as the Roto Brush's object tracker: prompts on some frames, the whole clip tracked +//! forwards and backwards, each frame reported as it is finished. + +use std::cell::RefCell; +use std::collections::HashMap; +use std::rc::Rc; +use std::sync::Arc; + +use candle_core::Tensor; +use effectcraft_segment::{ClipFrame, ObjectTracker, TrackRequest, TrackedFrame}; + +use crate::mask::{Sam2Mask, preprocess}; +use crate::prompt::{INPUT_SIZE, Prompt}; +use crate::session::{Cancel, Direction, FrameMask, FrameSource, ObjectId, VideoTracker}; +use crate::{Error, Result}; + +/// Most frames one tracking request may cover. +const MAX_CLIP: usize = 100_000; + +type Frames<'a> = &'a mut dyn FnMut(i64) -> Option>; + +/// Feeds the host's frames to the tracker, remembering each frame's size for the output. +struct Source<'a> { + first: i64, + count: usize, + fetch: RefCell>, + sizes: Rc>>, +} + +impl FrameSource for Source<'_> { + fn num_frames(&self) -> usize { + self.count + } + + fn frame(&self, index: usize) -> Result { + let f = (self.fetch.borrow_mut())(self.first + index as i64) + .ok_or_else(|| Error::Input(format!("frame {} is not available", self.first + index as i64)))?; + self.sizes.borrow_mut().insert(index, (f.w, f.h)); + preprocess(&f.rgb, f.w, f.h) + } +} + +fn to_frame(m: &FrameMask, size: (usize, usize)) -> Result { + let (w, h) = size; + let logits = Tensor::from_slice(&m.logits, (1, 1, 256, 256), &candle_core::Device::Cpu)?; + let prob = candle_nn::ops::sigmoid(&logits.upsample_bilinear2d(h, w, false)?)?.flatten_all()?.to_vec1::()?; + let presence = 1.0 / (1.0 + (-m.object_score).exp()); + Ok(TrackedFrame { prob, w, h, confidence: (presence * m.iou.clamp(0.0, 1.0)).clamp(0.0, 1.0), occluded: m.occluded }) +} + +impl Sam2Mask { + fn run_track(&self, req: &TrackRequest, frames: Frames<'_>, cancel: &dyn Fn() -> bool, out: &mut dyn FnMut(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())); + } + 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)); + } + 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 { + if cancel() { + return Err(Error::Cancelled); + } + let (w, h) = sizes.borrow().get(&idx).copied().ok_or_else(|| Error::Compute("frame size unknown".into()))?; + let (sx, sy) = (INPUT_SIZE as f32 / w as f32, INPUT_SIZE as f32 / h as f32); + let prompt = Prompt { + 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); + } + 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(); + } + match to_frame(&m, sizes.borrow().get(&m.frame).copied().unwrap_or((0, 0))) { + Ok(t) => out(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); + })?; + } + failure.map_or(Ok(()), Err) + } +} + +impl ObjectTracker for Sam2Mask { + fn track( + &self, + req: &TrackRequest, + frames: &mut dyn FnMut(i64) -> Option>, + cancel: &dyn Fn() -> bool, + out: &mut dyn FnMut(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 63c1a06ed..ce1fd14d8 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -337,3 +337,63 @@ fn sam2_works_as_a_mask_model() { assert!(model.segment(&rgb, fw, fh, &Prompt::default()).is_err()); assert!(model.segment(&rgb, fw + 1, fh, &prompt).is_err()); } + +#[test] +fn sam2_tracks_a_clip_through_the_object_tracker_interface() { + use effectcraft_segment::{ClipFrame, MaskModel, Prompt, TrackRequest}; + let Some((w, _)) = load() else { return }; + let (fw, fh, n) = (128usize, 96usize, 6usize); + let centre = |t: usize| (40.0 + 8.0 * t as f32, 48.0f32); + let inside = |t: usize, x: usize, y: usize| { + let c = centre(t); + ((x as f32 - c.0).powi(2) + (y as f32 - c.1).powi(2)).sqrt() < 26.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.3 + 0.2 * (x as f32 / 9.0).sin() * (y as f32 / 7.0).cos(); + rgb[y * fw + x] = + if inside(t, x, y) { [0.85 + 0.1 * ((x as f32 - centre(t).0) / 5.0).sin(), 0.4, 0.15] } else { [bg, bg * 0.8 + 0.1, bg * 0.6 + 0.3] }; + } + } + 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())); + 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 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| { + assert!(seen.insert(f, t).is_none(), "frame {f} reported twice"); + }) + .unwrap(); + assert_eq!(seen.keys().copied().collect::>(), (100..100 + n as i64).collect::>()); + for (f, t) in &seen { + assert_eq!((t.w, t.h, t.prob.len()), (fw, fh, fw * fh)); + let k = (f - 100) as usize; + let (mut i, mut u) = (0usize, 0usize); + for y in 0..fh { + for x in 0..fw { + let (a, b) = (t.prob[y * fw + x] > 0.5, inside(k, x, y)); + i += usize::from(a && b); + u += usize::from(a || b); + } + } + let j = i as f64 / u.max(1) as f64; + eprintln!("frame {f}: IoU {j:.3}, confidence {:.2}, occluded {}", t.confidence, t.occluded); + 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()); +} diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index dbfe07c1f..37b3a4ca0 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -45,6 +45,54 @@ pub trait MaskModel: Send + Sync { fn info(&self) -> &'static ModelInfo; /// Foreground probability (0–1) of each pixel of a frame (straight RGB 0–1, row-major `w×h`). fn segment(&self, rgb: &[[f32; 3]], w: usize, h: usize, prompt: &Prompt) -> Result>; + /// A video tracker built on the same weights, if the model has one. + fn tracker(&self) -> Option<&dyn ObjectTracker> { + None + } +} + +/// One decoded frame handed to a tracker (straight RGB 0–1, row-major `w×h`). +#[derive(Clone, Debug, Default)] +pub struct ClipFrame { + pub rgb: Vec<[f32; 3]>, + pub w: usize, + 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. +#[derive(Clone, Debug)] +pub struct TrackRequest { + pub first: i64, + pub count: usize, + pub prompts: Vec<(i64, Prompt)>, +} + +/// A tracker's answer for one frame. +#[derive(Clone, Debug, Default)] +pub struct TrackedFrame { + /// Foreground probability of every pixel (`w×h`); all zero when the object is absent. + pub prob: Vec, + pub w: usize, + pub h: usize, + /// How sure the tracker is (0–1), for the timeline. + pub confidence: f32, + /// The object is not in this frame (occluded or out of view). + pub occluded: bool, +} + +/// 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. + fn track( + &self, + req: &TrackRequest, + frames: &mut dyn FnMut(i64) -> Option>, + cancel: &dyn Fn() -> bool, + out: &mut dyn FnMut(i64, TrackedFrame), + ) -> Result<()>; } /// What a model is for.