From 1fb51600d5c93a74c8742b18ff741731f51fbcc3 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 20:13:38 +0200 Subject: [PATCH] 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 --- crates/effects/src/roto.rs | 24 ++++++----- crates/sam2/tests/figure.rs | 77 ++++++++++++++++++++++++++++++++++ crates/track/src/roto.rs | 23 ++++++++++ crates/track/src/roto/tests.rs | 23 ++++++++++ 4 files changed, 136 insertions(+), 11 deletions(-) create mode 100644 crates/sam2/tests/figure.rs diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index 9596b6d50..f4f5e6bd0 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -553,17 +553,19 @@ impl Chain { ) -> Result, 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()); } diff --git a/crates/sam2/tests/figure.rs b/crates/sam2/tests/figure.rs new file mode 100644 index 000000000..2ca4d0b13 --- /dev/null +++ b/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) { + 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::().unwrap(); + (iou(&m, >), out.iou.flatten_all().unwrap().to_vec1::().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}"); +} diff --git a/crates/track/src/roto.rs b/crates/track/src/roto.rs index ce52a79a4..000f7b55f 100644 --- a/crates/track/src/roto.rs +++ b/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; diff --git a/crates/track/src/roto/tests.rs b/crates/track/src/roto/tests.rs index c08306060..b7d8c7809 100644 --- a/crates/track/src/roto/tests.rs +++ b/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, >) > 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, >) > 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])); +}