S2.65: prompt SAM 2 with a box from the classic matte on stroke frames

On real footage the stroke frame came out partial (clicks make SAM 2 pick a
part of the object). Measured on a multi-part figure: clicks only IoU 0.41,
with a box around the figure IoU 1.00. Chain::track now adds the box of the
classic engine's matte of the strokes to each prompt; a box stroke wins.

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

crates/effects/src/roto.rs

@@ -553,17 +553,19 @@ impl Chain {
     ) -> 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();
+        // A prompt is a few clicks from the strokes and a box around what the classic engine makes of
+        // them on that frame (clicks alone pick a part of the object, not all of it).
+        let mut prompts: Vec<(i64, effectcraft_segment::Prompt)> = vec![];
+        for f in self.data.stroke_frames().into_iter().filter(|f| self.data.in_span(*f)) {
+            let strokes = self.data.strokes_at(f);
+            let p = match frame(f) {
+                Some(img) => rb::tracker_prompt_boxed(&strokes, self.opts.scale, &rb::segment(&img, &strokes, None, &self.opts, self.refine_radius)),
+                None => rb::tracker_prompt(&strokes, self.opts.scale),
+            };
+            if p.points.iter().any(|(_, fg)| *fg) {
+                prompts.push((f, p));
+            }
+        }
         if prompts.is_empty() {
             return Err("no foreground strokes to follow".into());
         }

crates/sam2/tests/figure.rs

@@ -0,0 +1,77 @@
+use effectcraft_sam2::model::ImageModel;
+use effectcraft_sam2::prompt::Prompt;
+use effectcraft_sam2::weights::Weights;
+
+fn figure(w: usize, h: usize) -> (Vec<[f32; 3]>, Vec<bool>) {
+    let mut rgb = vec![[0f32; 3]; w * h];
+    let mut gt = vec![false; w * h];
+    let part = |x: f32, y: f32| -> Option<[f32; 3]> {
+        let (cx, _cy) = (128.0, 0.0);
+        let r = |x0: f32, y0: f32, x1: f32, y1: f32| x >= x0 && x < x1 && y >= y0 && y < y1;
+        if (x - cx).powi(2) + (y - 40.0).powi(2) < 18.0f32.powi(2) {
+            return Some([0.8, 0.6, 0.5]);
+        } // head
+        if r(104.0, 60.0, 152.0, 130.0) {
+            return Some([0.2, 0.3, 0.7 + 0.1 * (x / 7.0).sin()]);
+        } // torso
+        if r(76.0, 62.0, 104.0, 74.0) || r(152.0, 62.0, 180.0, 74.0) {
+            return Some([0.8, 0.6, 0.5]);
+        } // arms
+        if r(108.0, 130.0, 124.0, 210.0) || r(132.0, 130.0, 148.0, 210.0) {
+            return Some([0.15, 0.15, 0.2]);
+        } // legs
+        if r(102.0, 210.0, 126.0, 224.0) || r(130.0, 210.0, 154.0, 224.0) {
+            return Some([0.95, 0.95, 0.95]);
+        } // white shoes
+        None
+    };
+    for y in 0..h {
+        for x in 0..w {
+            let bg = 0.45 + 0.1 * (x as f32 / 11.0).sin() * (y as f32 / 9.0).cos();
+            let p = part(x as f32, y as f32);
+            rgb[y * w + x] = p.unwrap_or([bg, bg * 0.9, bg * 0.8]);
+            gt[y * w + x] = p.is_some();
+        }
+    }
+    (rgb, gt)
+}
+
+fn iou(p: &[f32], gt: &[bool]) -> f64 {
+    let (mut i, mut u) = (0usize, 0usize);
+    for (a, b) in p.iter().zip(gt) {
+        let a = *a > 0.0;
+        i += usize::from(a && *b);
+        u += usize::from(a || *b);
+    }
+    i as f64 / u.max(1) as f64
+}
+
+/// A multi-part figure (torso, head, legs, arms, white shoes in different colours): clicks alone make
+/// SAM 2 pick the torso and head, a box around the figure makes it cover all of it. This is why Roto
+/// Brush prompts the tracker with a box from the classic engine's matte (found on real footage, where
+/// the stroke frame came out partial).
+#[test]
+fn a_box_makes_sam2_cover_every_part_of_the_object() {
+    let wf = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../sam2-tiny/model.safetensors");
+    let wf = std::env::var_os("EFFECTCRAFT_SAM2_DIR").map_or(wf, |d| std::path::PathBuf::from(d).join("model.safetensors"));
+    if !wf.exists() {
+        eprintln!("skipped: no SAM 2 weights");
+        return;
+    }
+    let model = ImageModel::load(&Weights::open(&wf).unwrap()).unwrap();
+    let (w, h) = (256usize, 256usize);
+    let (rgb, gt) = figure(w, h);
+    let enc = model.encode(&effectcraft_sam2::mask::preprocess(&rgb, w, h).unwrap()).unwrap();
+    let s = 4.0f32;
+    let clicks = vec![([128.0 * s, 75.0 * s], true), ([128.0 * s, 100.0 * s], true), ([128.0 * s, 125.0 * s], true), ([20.0 * s, 20.0 * s], false)];
+    let mask = |p: &Prompt| -> (f64, f32) {
+        let out = model.segment(&enc, p, false).unwrap();
+        let m = out.masks.narrow(1, 0, 1).unwrap().upsample_bilinear2d(h, w, false).unwrap().flatten_all().unwrap().to_vec1::<f32>().unwrap();
+        (iou(&m, &gt), out.iou.flatten_all().unwrap().to_vec1::<f32>().unwrap()[0])
+    };
+    let (points_only, _) = mask(&Prompt { points: clicks.clone(), bbox: None });
+    let (boxed, predicted) = mask(&Prompt { points: clicks, bbox: Some([70.0 * s, 18.0 * s, 186.0 * s, 228.0 * s]) });
+    eprintln!("clicks only: IoU {points_only:.2}; with the box: IoU {boxed:.2} (predicted quality {predicted:.2})");
+    assert!(points_only < 0.6, "the clicks-only failure this test documents is gone ({points_only}): revisit the box prompt");
+    assert!(boxed > 0.95 && predicted > 0.7, "boxed IoU {boxed}, predicted {predicted}");
+}

crates/track/src/roto.rs

@@ -923,6 +923,29 @@ pub fn tracker_prompt(strokes: &[&Stroke], scale: f64) -> Prompt {
     Prompt { points, bbox, ..Default::default() }
 }
 
+/// [`tracker_prompt`] plus a box around the classic engine's matte of the same strokes. Clicks alone
+/// make SAM 2 pick the part they land on (a torso, not the legs and shoes: measured IoU 0.41 on a
+/// multi-part figure, against 1.00 with the box); the box says how far the object reaches. A box
+/// stroke the user drew wins over the matte's.
+pub fn tracker_prompt_boxed(strokes: &[&Stroke], scale: f64, classic: &FrameSeg) -> Prompt {
+    let mut p = tracker_prompt(strokes, scale);
+    if p.bbox.is_none() {
+        let (w, h) = (classic.w, classic.h);
+        let (mut x0, mut y0, mut x1, mut y1) = (usize::MAX, usize::MAX, 0usize, 0usize);
+        for (i, v) in classic.matte.iter().enumerate() {
+            if *v != 0 && w > 0 {
+                let (x, y) = (i % w, i / w);
+                (x0, y0, x1, y1) = (x0.min(x), y0.min(y), x1.max(x), y1.max(y));
+            }
+        }
+        if x0 <= x1 {
+            let (mx, my) = (((x1 - x0) as f32 * 0.05).max(4.0), ((y1 - y0) as f32 * 0.05).max(4.0));
+            p.bbox = Some([(x0 as f32 - mx).max(0.0), (y0 as f32 - my).max(0.0), (x1 as f32 + 1.0 + mx).min(w as f32), (y1 as f32 + 1.0 + my).min(h as f32)]);
+        }
+    }
+    p
+}
+
 /// Most foreground / background points a tracker prompt carries.
 pub const TRACKER_FG_POINTS: usize = 3;
 pub const TRACKER_BG_POINTS: usize = 3;

crates/track/src/roto/tests.rs

@@ -612,3 +612,26 @@ fn tracked_segmentation_handles_absent_objects_and_bad_maps() {
     let wrong = segment_tracked(&img, &refs, vec![1.0; 3], false, &opts, 0.0);
     assert!(iou(&wrong.matte, &gt) > 0.95);
 }
+
+#[test]
+fn tracker_prompts_carry_the_box_of_the_classic_matte() {
+    let (img, gt) = frame_at(center(0));
+    let st = base_strokes(0);
+    let refs: Vec<&Stroke> = st.iter().collect();
+    let classic = segment(&img, &refs, None, &SegOpts::default(), 0.0);
+    let p = tracker_prompt_boxed(&refs, 1.0, &classic);
+    let b = p.bbox.expect("a box");
+    let c = center(0);
+    // the box reaches the disk's edges, with a margin
+    assert!(b[0] <= (c[0] - R) as f32 && b[0] > (c[0] - R) as f32 - 12.0, "{b:?}");
+    assert!(b[2] >= (c[0] + R) as f32 && b[2] < (c[0] + R) as f32 + 12.0, "{b:?}");
+    assert!(b[1] <= (c[1] - R) as f32 && b[3] >= (c[1] + R) as f32, "{b:?}");
+    assert!(iou(&classic.matte, &gt) > 0.95);
+    // no matte, no box; a box the user drew wins
+    let none = tracker_prompt_boxed(&refs, 1.0, &FrameSeg::empty(W as usize, H as usize));
+    assert_eq!(none.bbox, None);
+    let bx = Stroke { kind: StrokeKind::Box, frame: 0, radius: 0.0, points: vec![[10.0, 10.0], [50.0, 60.0]] };
+    let mut with_box: Vec<&Stroke> = refs.clone();
+    with_box.push(&bx);
+    assert_eq!(tracker_prompt_boxed(&with_box, 1.0, &classic).bbox, Some([10.0, 10.0, 50.0, 60.0]));
+}