S2.41: Roto Brush follows the object with SAM 2

When the chosen model has a tracker, roto.propagate runs one tracking pass over
the span: prompts from the strokes (a handful of points, which SAM 2 trusts),
the clip tracked both ways, each frame's probability map refined by the same
graph cut with the strokes as hard constraints. Results are cached under their
own keys with a per-frame confidence and an occluded flag; roto.status reports
them (source: tracker|classical). Without a tracker nothing changes.
Real weights: IoU 0.98 on all 21 frames, confidence 0.99.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

crates/effects/src/roto.rs

@@ -29,7 +29,7 @@ use std::sync::{Arc, Mutex, OnceLock, RwLock, Weak};
 use effectcraft_keyframe::Value;
 use effectcraft_project::ParamUi;
 use effectcraft_raster::Image;
-use effectcraft_segment::MaskModel;
+use effectcraft_segment::{ClipFrame, MaskModel, TrackRequest};
 use effectcraft_track::roto::refine::{self, MatteParams};
 use effectcraft_track::roto::{self as rb, RotoData, SegOpts, rle};
 use rayon::prelude::*;
@@ -298,9 +298,39 @@ fn from_store(key: u64, scale: f64) -> Option<Arc<FrameSeg>> {
     Some(seg)
 }
 
+/// Tracker results live under the chain key mixed with this, so a frame the tracker followed and
+/// one the classic engine computed never share a cache slot.
+const TRACK_SALT: u64 = 0x7472_6163_6b65_7231;
+
+fn tracked_key(key: u64) -> u64 {
+    key.rotate_left(17) ^ TRACK_SALT
+}
+
+/// How far to trust a frame the tracker followed.
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub struct FrameTrust {
+    /// 0–1, from the tracker's own object score and mask quality.
+    pub confidence: f32,
+    /// The object is not in this frame (an empty matte).
+    pub occluded: bool,
+}
+
+fn trust_cache() -> &'static Mutex<HashMap<u64, FrameTrust>> {
+    static T: OnceLock<Mutex<HashMap<u64, FrameTrust>>> = OnceLock::new();
+    T.get_or_init(|| Mutex::new(HashMap::new()))
+}
+
+/// The tracker's confidence in the frame with chain key `key`, if the tracker followed it.
+pub fn trust(key: u64) -> Option<FrameTrust> {
+    trust_cache().lock().unwrap_or_else(|e| e.into_inner()).get(&tracked_key(key)).copied()
+}
+
 /// Drop the cached segmentations of these chain keys (every scale).
 pub fn forget(keys: impl IntoIterator<Item = u64>) {
-    let keys: std::collections::HashSet<u64> = keys.into_iter().collect();
+    let mut trusts = trust_cache().lock().unwrap_or_else(|e| e.into_inner());
+    let keys: std::collections::HashSet<u64> = keys.into_iter().flat_map(|k| [k, tracked_key(k)]).collect();
+    trusts.retain(|k, _| !keys.contains(k));
+    drop(trusts);
     let mut c = cache().lock().unwrap_or_else(|e| e.into_inner());
     let gone: Vec<(u64, u64)> = c.map.keys().filter(|k| keys.contains(&k.0)).copied().collect();
     for k in gone {
@@ -313,6 +343,7 @@ pub fn forget(keys: impl IntoIterator<Item = u64>) {
 
 /// Drop every cached segmentation (Edit ▸ Purge).
 pub fn purge() {
+    trust_cache().lock().unwrap_or_else(|e| e.into_inner()).clear();
     let mut c = cache().lock().unwrap_or_else(|e| e.into_inner());
     c.map.clear();
     c.order.clear();
@@ -333,6 +364,16 @@ pub fn layer_grid(buf: &Buf, layer_size: [f64; 2]) -> Image {
     img
 }
 
+/// What a tracking pass found.
+#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct TrackSummary {
+    pub frames: usize,
+    /// Frames where the object was absent.
+    pub occluded: usize,
+    pub min_confidence: f32,
+}
+
 /// Everything needed to compute an instance's segmentations.
 pub struct Chain {
     pub data: Arc<RotoData>,
@@ -357,11 +398,16 @@ impl Chain {
         let sc = self.opts.scale;
         let lookup = |g: i64| -> Option<Arc<FrameSeg>> {
             let k = *self.keys.get(&g)?;
+            // A frame the tracker followed wins over the classic engine's (memory only: its
+            // confidence is not persisted).
+            if let Some(s) = cached(tracked_key(k), sc) {
+                return Some(s);
+            }
             if let Some(s) = cached(k, sc).or_else(|| from_store(k, sc)) {
                 return Some(s);
             }
             if sc != 1.0 {
-                let full = cached(k, 1.0).or_else(|| from_store(k, 1.0))?;
+                let full = cached(tracked_key(k), 1.0).or_else(|| cached(k, 1.0)).or_else(|| from_store(k, 1.0))?;
                 let s = Arc::new(downscale(&full, sc));
                 store(k, sc, s.clone());
                 return Some(s);
@@ -407,7 +453,8 @@ impl Chain {
                     if pimg.width != img.width || pimg.height != img.height {
                         return None;
                     }
-                    rb::propagate_with(&pimg, p, &img, &strokes, &self.opts, self.refine_radius, self.model.as_deref())
+                    // A tracker is too slow to prompt on every frame (its pass is `track`); the classic flow carries the rest.
+                    rb::propagate_with(&pimg, p, &img, &strokes, &self.opts, self.refine_radius, if self.has_tracker() { None } else { self.model.as_deref() })
                 }
             };
             let seg = Arc::new(seg);
@@ -424,7 +471,64 @@ impl Chain {
 
     /// Whether frame `f` is computed (at full resolution or at `scale`).
     pub fn is_cached(&self, f: i64, scale: f64) -> bool {
-        self.keys.get(&f).is_some_and(|k| cached(*k, 1.0).is_some() || cached(*k, scale).is_some())
+        self.keys.get(&f).is_some_and(|k| [*k, tracked_key(*k)].iter().any(|k| cached(*k, 1.0).is_some() || cached(*k, scale).is_some()))
+    }
+
+    /// Whether the chosen model can follow an object through a clip (SAM 2).
+    pub fn has_tracker(&self) -> bool {
+        self.model.as_ref().is_some_and(|m| m.tracker().is_some())
+    }
+
+    /// Follow the object through the whole span with the model's tracker, from the strokes on
+    /// their frames, and cache every frame's segmentation (under the tracker's keys, with its
+    /// confidence). The strokes stay hard constraints and the edges are refined at full
+    /// resolution; frames the tracker can't give are left to [`Chain::ensure`]'s classic engine.
+    /// `progress` gets the number of frames finished. `Ok(None)`: no tracker, nothing done.
+    pub fn track(
+        &self,
+        frame: &mut dyn FnMut(i64) -> Option<Arc<Image>>,
+        cancel: &dyn Fn() -> bool,
+        progress: &mut dyn FnMut(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()) };
+        let prompts: Vec<(i64, effectcraft_segment::Prompt)> = self
+            .data
+            .stroke_frames()
+            .into_iter()
+            .filter(|f| self.data.in_span(*f))
+            .filter_map(|f| {
+                let strokes = self.data.strokes_at(f);
+                let p = rb::tracker_prompt(&strokes, self.opts.scale);
+                p.points.iter().any(|(_, fg)| *fg).then_some((f, p))
+            })
+            .collect();
+        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 frame = std::cell::RefCell::new(frame);
+        let mut summary = TrackSummary::default();
+        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,
+            &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);
+                store(tracked_key(key), self.opts.scale, Arc::new(seg));
+                trust_cache().lock().unwrap_or_else(|e| e.into_inner()).insert(tracked_key(key), FrameTrust { confidence: t.confidence, occluded: t.occluded });
+                summary.frames += 1;
+                summary.occluded += usize::from(t.occluded);
+                summary.min_confidence = if summary.frames == 1 { t.confidence } else { summary.min_confidence.min(t.confidence) };
+                progress(summary.frames);
+            },
+        )?;
+        Ok(Some(summary))
     }
 }
 

crates/engine/src/analysis/models.rs

@@ -188,6 +188,11 @@ impl Session {
         });
     }
 
+    /// The Roto Brush model in use (`None`: the classic engine).
+    pub fn models_roto_active(&self) -> Option<String> {
+        self.models.status(Task::Mask).active
+    }
+
     /// Finish what the model workers did: a toast, and loading a model that just arrived.
     pub fn poll_models(&mut self) {
         for task in [Task::Mask, Task::Face] {

crates/engine/src/analysis/roto.rs

@@ -112,6 +112,10 @@ pub struct RotoProgress {
     pub error: Option<String>,
     /// The job works on an Object Matte (else a Roto Brush & Refine Edge).
     pub object_matte: bool,
+    /// What the model's tracker found, when one followed the object (SAM 2).
+    pub tracked: Option<fx::TrackSummary>,
+    /// Why the tracker was not used, when a tracking model was chosen but the classic engine ran.
+    pub tracker_note: Option<String>,
 }
 
 impl RotoProgress {
@@ -436,6 +440,31 @@ async fn run_work(w: Work, shared: &RotoShared) {
     let model = fx::installed_model();
     let key = if w.object_matte { om::frozen_key(&w.params, w.size, w.fps, model.as_deref()) } else { fx::frozen_key(&w.params, w.size, w.fps) };
     let mut frozen = FrozenCache { key, ..Default::default() };
+    // A tracking model (SAM 2) follows the object through the whole span in one pass; the loop
+    // below then finds every frame cached. Without one (or if it fails) nothing changes: the
+    // classic engine runs frame by frame.
+    if w.task == RotoTask::Propagate && !w.object_matte && chain.has_tracker() {
+        let mut get = |g: i64| src.get(g);
+        let mut report = |n: usize| {
+            let mut s = lock(&shared.state);
+            s.done = n as u64;
+            s.elapsed = t0.elapsed().as_secs_f64();
+            s.fps = if s.elapsed > 0.0 { s.done as f64 / s.elapsed } else { 0.0 };
+            crate::offload::report(s.done, s.total);
+        };
+        let result = chain.track(&mut get, &cancel, &mut report);
+        if cancel() {
+            let mut s = lock(&shared.state);
+            s.cancelled = true;
+            s.finished = true;
+            return;
+        }
+        let mut s = lock(&shared.state);
+        match result {
+            Ok(summary) => s.tracked = summary,
+            Err(e) => s.tracker_note = Some(format!("tracker unavailable, classic engine used: {e}")),
+        }
+    }
     for (k, f) in w.frames.iter().enumerate() {
         if cancel() {
             let mut s = lock(&shared.state);

crates/engine/src/cmd/roto_cmds.rs

@@ -463,6 +463,19 @@ fn status(s: &mut Session, p: &Value) -> Result<Value> {
             }
         }
     }
+    // Where each frame's matte came from, and how far to trust it: frames a tracker followed
+    // report its confidence and whether the object was absent; the rest came from the classic engine.
+    {
+        let size = rt::layer_size(&s.project, &comp, &layer);
+        let chain = fx::Chain::new(&params_static(&g), size, comp.frame_rate.as_f64(), 1.0);
+        let tracked: serde_json::Map<String, Value> = chain
+            .keys
+            .iter()
+            .filter_map(|(f, k)| fx::trust(*k).map(|t| (f.to_string(), json!({"confidence": t.confidence, "occluded": t.occluded}))))
+            .collect();
+        out["tracker"] = json!({"model": s.models_roto_active(), "frames": tracked});
+        out["source"] = json!(if tracked.contains_key(&frame.to_string()) { "tracker" } else { "classical" });
+    }
     let computed = s.roto_computed(cid, lid, uid);
     out["computed"] = json!(computed.iter().filter(|(_, c)| **c).map(|(f, _)| *f).collect::<Vec<_>>());
     Ok(out)

crates/engine/src/tests/roto.rs

@@ -668,3 +668,41 @@ fn object_matte_mattes_persist_in_the_disk_cache() {
     assert!(om::ensure(&moved, size, 25.0, 0, None, &mut |_| None, &|| false).is_none());
     let _ = std::fs::remove_dir_all(&dir);
 }
+
+/// SAM 2 tracks the clip end to end (needs the real weights: `EFFECTCRAFT_SAM2=path/to/model.safetensors`).
+/// Every frame comes from the tracker with a confidence, the mattes are right, and nothing changes
+/// once the model is deselected.
+#[cfg(not(target_arch = "wasm32"))]
+#[test]
+fn sam2_tracks_the_clip_through_roto_propagate() {
+    let Ok(weights) = std::env::var("EFFECTCRAFT_SAM2") else { return };
+    let dir = std::env::temp_dir().join(format!("ec-roto-sam2-{}", std::process::id()));
+    let (mut s, clip, _) = painted(center, "disk-sam2.mov");
+    s.models_dir = Some(dir.join("models"));
+    std::fs::create_dir_all(dir.join("models")).unwrap();
+    s.execute_checked("roto.model.install", json!({"path": weights})).unwrap();
+    s.execute_checked("roto.model.select", json!({"id": "sam2.1-hiera-tiny", "wait": true})).unwrap();
+    assert_eq!(s.models.status(effectcraft_segment::Task::Mask).active.as_deref(), Some("sam2.1-hiera-tiny"));
+    let r = s.execute_checked("roto.propagate", json!({"layer": clip.0, "direction": "both", "wait": true})).unwrap();
+    assert_eq!(r["frames"], 21, "{r}");
+    let mut worst: f64 = 1.0;
+    for f in 0..=20 {
+        let j = iou_at(&mut s, clip, f, center(f));
+        let st = s.execute("roto.status", json!({"layer": clip.0, "frame": f})).unwrap();
+        eprintln!("frame {f}: IoU {j:.3}, {:?}", st["tracker"]["frames"][f.to_string()]);
+        worst = worst.min(j);
+    }
+    assert!(worst > 0.9, "worst IoU over the span with SAM 2: {worst}");
+    let st = s.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap();
+    assert_eq!(st["source"], "tracker", "{st}");
+    let frames = st["tracker"]["frames"].as_object().unwrap();
+    assert_eq!(frames.len(), 21);
+    assert!(frames.values().all(|t| t["occluded"] == false && t["confidence"].as_f64().unwrap() > 0.3), "{frames:?}");
+    assert_eq!(st["tracker"]["model"], "sam2.1-hiera-tiny");
+    // Back to the classic engine: same strokes, its own segmentation, no tracker data.
+    s.execute_checked("roto.model.select", json!({"id": "classical", "wait": true})).unwrap();
+    let st = s.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap();
+    assert_eq!(st["source"], "classical", "{st}");
+    assert_eq!(st["tracker"]["frames"].as_object().unwrap().len(), 0);
+    let _ = std::fs::remove_dir_all(dir);
+}

crates/track/src/roto.rs

@@ -903,6 +903,47 @@ fn points_along(s: &Stroke, scale: f64, max: usize) -> Vec<[f32; 2]> {
     out
 }
 
+/// A tracker's prompt from a frame's strokes (points along the foreground and background strokes,
+/// in frame pixels at `scale`).
+///
+/// SAM 2 was trained on a handful of clicks: with sixteen or more points it doubts its own mask
+/// (measured: confidence 0.25 with 16 points, 0.00 with 32, 0.99 with 1 to 7), so this picks at
+/// most [`TRACKER_FG_POINTS`] foreground and [`TRACKER_BG_POINTS`] background points, spread
+/// along the strokes.
+pub fn tracker_prompt(strokes: &[&Stroke], scale: f64) -> Prompt {
+    let mut points = vec![];
+    for (kind, fg, max) in [(StrokeKind::Fg, true, TRACKER_FG_POINTS), (StrokeKind::Bg, false, TRACKER_BG_POINTS)] {
+        let all: Vec<[f32; 2]> = strokes.iter().filter(|s| s.kind == kind).flat_map(|s| points_along(s, scale, 3)).collect();
+        let n = all.len().min(max);
+        // evenly spaced picks, keeping the first and the last
+        points.extend((0..n).filter_map(|i| all.get(if n > 1 { i * (all.len() - 1) / (n - 1) } else { all.len() / 2 }).map(|p| (*p, fg))));
+    }
+    // a box the user drew on the frame prompts the tracker too
+    let bbox = strokes.iter().filter(|s| s.kind == StrokeKind::Box).find_map(|s| box_of(s, scale));
+    Prompt { points, bbox, ..Default::default() }
+}
+
+/// Most foreground / background points a tracker prompt carries.
+pub const TRACKER_FG_POINTS: usize = 3;
+pub const TRACKER_BG_POINTS: usize = 3;
+
+/// A frame's segmentation from a tracker's probability map: the map is the shape prior of the
+/// same graph cut as everywhere else, so the strokes stay hard constraints and the edges snap
+/// to colour. An absent object gives an empty matte (strokes on the frame still win); a map of
+/// the wrong size is ignored.
+pub fn segment_tracked(img: &Image, strokes: &[&Stroke], prob: Vec<f32>, occluded: bool, opts: &SegOpts, refine_radius: f64) -> FrameSeg {
+    let (w, h) = (img.width as usize, img.height as usize);
+    if prob.len() != w * h || w == 0 || h == 0 {
+        return segment(img, strokes, None, opts, refine_radius);
+    }
+    let has_fg = strokes.iter().any(|s| s.kind == StrokeKind::Fg);
+    if (occluded || !prob.iter().any(|p| *p >= 0.5)) && !has_fg {
+        return FrameSeg::empty(w, h);
+    }
+    let prior = Prior { soft: prob, fg: Gmm::default(), bg: Gmm::default(), refine: vec![0; w * h], weight: MODEL_PRIOR };
+    segment(img, strokes, Some(&prior), opts, refine_radius)
+}
+
 /// A model's prompt from the strokes: points along the foreground and background strokes, but
 /// not where a later stroke of the other kind painted over them (`hard`, from [`resolve`]), and
 /// the box of a [`StrokeKind::Box`] stroke (matte pixels). A stroke painted over at every point

crates/track/src/roto/tests.rs

@@ -579,3 +579,36 @@ fn a_correction_keeps_the_model_prompted_with_what_is_left() {
     let hard = resolve(&refs, StrokeKind::Fg, StrokeKind::Bg, 1.0, w, h);
     assert!(stroke_prompt(&refs, &hard, 1.0, w, h).points.iter().all(|p| !p.1));
 }
+
+#[test]
+fn tracker_prompts_carry_a_handful_of_points() {
+    let st = base_strokes(0);
+    let refs: Vec<&Stroke> = st.iter().collect();
+    let p = tracker_prompt(&refs, 1.0);
+    let (fg, bg) = (p.points.iter().filter(|(_, f)| *f).count(), p.points.iter().filter(|(_, f)| !*f).count());
+    assert!((1..=TRACKER_FG_POINTS).contains(&fg) && (1..=TRACKER_BG_POINTS).contains(&bg), "{fg} fg, {bg} bg");
+    // a long, finely sampled stroke still gives at most the limit, and an empty one gives none
+    let long = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 2.0, points: (0..400).map(|i| [i as f64 * 0.4, 50.0]).collect() };
+    assert!(tracker_prompt(&[&long], 1.0).points.len() <= TRACKER_FG_POINTS);
+    assert!(tracker_prompt(&[], 1.0).points.is_empty());
+}
+
+#[test]
+fn tracked_segmentation_handles_absent_objects_and_bad_maps() {
+    let (img, gt) = frame_at(center(0));
+    let st = base_strokes(0);
+    let refs: Vec<&Stroke> = st.iter().collect();
+    let opts = SegOpts::default();
+    let n = (W * H) as usize;
+    // a good map (the true disk) reproduces it, with the strokes as constraints
+    let soft: Vec<f32> = gt.iter().map(|v| f32::from(*v)).collect();
+    let seg = segment_tracked(&img, &refs, soft, false, &opts, 0.0);
+    assert!(iou(&seg.matte, &gt) > 0.95);
+    // an absent object (no strokes on the frame) is an empty matte, whatever the map says
+    let none = segment_tracked(&img, &[], vec![1.0; n], true, &opts, 0.0);
+    assert_eq!(none.area(), 0);
+    assert_eq!(segment_tracked(&img, &[], vec![0.0; n], false, &opts, 0.0).area(), 0);
+    // a map of the wrong size is ignored: the strokes alone segment the frame
+    let wrong = segment_tracked(&img, &refs, vec![1.0; 3], false, &opts, 0.0);
+    assert!(iou(&wrong.matte, &gt) > 0.95);
+}