diff --git a/crates/effects/src/object_matte.rs b/crates/effects/src/object_matte.rs index 3b7b7c28..9ee33a90 100644 --- a/crates/effects/src/object_matte.rs +++ b/crates/effects/src/object_matte.rs @@ -5,8 +5,9 @@ //! The prompt is the **Object Point** (layer pixels) and, with **Use Box**, a box around the //! object. On the **Reference Frame** it is segmented by the same engine as Roto Brush: the //! trained model chosen in Settings ▸ Roto Brush (MobileSAM, `effectcraft_segment`) when one is -//! installed, the classical graph cut otherwise (the point is a foreground stroke, the box's -//! edges, or without a box the layer's border, background strokes). With **Track Object** the +//! installed, the classical graph cut otherwise. The model gets the point as a foreground point +//! and the box as its box prompt; the graph cut gets the point as a foreground stroke and the +//! box's edges (without a box, the layer's border) as background strokes. With **Track Object** the //! matte is then carried frame by frame up to **Frames Before** / **Frames After** away (optical //! flow, re-segmented by the model prompted with the carried matte); without it every frame is //! segmented on its own from the prompt (useful when the point is animated). @@ -119,9 +120,10 @@ pub fn tracks(params: &Params) -> bool { } /// The prompt as Roto Brush strokes on layer frame `frame`: the Object Point as a foreground -/// stroke; the box's edges as background strokes; without a box, the layer's border when there -/// is no model (the graph cut needs background samples; a model does not, and a border stroke -/// would cut objects touching the frame edge). `None` when the point is unusable. +/// stroke; the box as a box stroke (the model's box prompt) with a model, its edges as +/// background strokes for the graph cut without one; without a box, the layer's border when +/// there is no model (the graph cut needs background samples; a model does not, and a border +/// stroke would cut objects touching the frame edge). `None` when the point is unusable. pub fn prompt_strokes(params: &Params, layer_size: [f64; 2], frame: i64, with_model: bool) -> Option> { let [w, h] = layer_size; if !(w.is_finite() && h.is_finite() && w >= 1.0 && h >= 1.0) { @@ -144,7 +146,12 @@ pub fn prompt_strokes(params: &Params, layer_size: [f64; 2], frame: i64, with_mo if pt[0] - m <= x0 || pt[0] + m >= x1 || pt[1] - m <= y0 || pt[1] + m >= y1 { return None; } - out.push(Stroke { kind: StrokeKind::Bg, frame, radius: EDGE_RADIUS, points: rect(x0, y0, x1, y1) }); + out.push(if with_model { + // The model takes the box itself as its box prompt. + Stroke { kind: StrokeKind::Box, frame, radius: 0.0, points: vec![[x0, y0], [x1, y1]] } + } else { + Stroke { kind: StrokeKind::Bg, frame, radius: EDGE_RADIUS, points: rect(x0, y0, x1, y1) } + }); } None if !with_model => { let i = EDGE_RADIUS; @@ -438,6 +445,8 @@ mod tests { #[derive(Default)] struct Stub { calls: std::sync::atomic::AtomicUsize, + /// The box prompt of the last call. + bbox: std::sync::Mutex>, } impl MaskModel for Stub { fn info(&self) -> &'static effectcraft_segment::ModelInfo { @@ -445,6 +454,7 @@ mod tests { } fn segment(&self, _: &[[f32; 3]], w: usize, h: usize, prompt: &effectcraft_segment::Prompt) -> effectcraft_segment::Result> { self.calls.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + *self.bbox.lock().unwrap() = prompt.bbox; let Some(([cx, cy], _)) = prompt.points.iter().find(|(_, fg)| *fg) else { return Err("no point".into()) }; Ok((0..w * h).map(|i| (((i % w) as f32 + 0.5 - cx).hypot((i / w) as f32 + 0.5 - cy) < 15.0) as u8 as f32).collect()) } @@ -464,5 +474,39 @@ mod tests { let FrameMatte::Alpha(a) = frame_matte(&p, [100.0, 80.0], 25.0, 0, 1.0, &img, Some(stub.clone()), &mut frame, &|| false) else { panic!("no matte") }; assert!(a[40 * 100 + 40] > 0.5 && a[5 * 100 + 5] < 0.5); assert_eq!(stub.calls.load(std::sync::atomic::Ordering::Relaxed), 1); + assert_eq!(*stub.bbox.lock().unwrap(), None); + } + + #[test] + fn the_box_reaches_the_model_as_a_box_prompt() { + let p = params(&[ + ("objectPoint", Value::Vec2([90.0, 70.0])), + ("useBox", Value::Bool(true)), + ("boxTopLeft", Value::Vec2([100.0, 80.0])), + ("boxBottomRight", Value::Vec2([74.0, 54.0])), + ]); + // With a model: the point and a box stroke (corners ordered); without: the box's edges + // as a background stroke, as before. + let with = prompt_strokes(&p, [100.0, 80.0], 0, true).unwrap(); + assert_eq!(with.iter().map(|s| s.kind).collect::>(), [StrokeKind::Fg, StrokeKind::Box]); + assert_eq!(with[1].points, [[74.0, 54.0], [100.0, 80.0]]); + let without = prompt_strokes(&p, [100.0, 80.0], 0, false).unwrap(); + assert_eq!(without.iter().map(|s| s.kind).collect::>(), [StrokeKind::Fg, StrokeKind::Bg]); + // The model is called with the box itself (matte pixels at scale 1) and only one + // foreground point: no background points along the edges. + let stub = Arc::new(Stub::default()); + let img = frame_at(0); + let mut frame = |g: i64| Some(Arc::new(frame_at(g))); + let FrameMatte::Alpha(a) = frame_matte(&p, [100.0, 80.0], 25.0, 0, 1.0, &img, Some(stub.clone()), &mut frame, &|| false) else { panic!("no matte") }; + assert_eq!(*stub.bbox.lock().unwrap(), Some([74.0, 54.0, 100.0, 80.0])); + assert!(a[70 * 100 + 90] > 0.5 && a[40 * 100 + 40] < 0.5); + // At half resolution the box is scaled with the frame (another Input Key, so the + // full-resolution result above is not reused). + let mut p = p; + p.values.insert(INPUT_KEY.into(), Value::Str("half".into())); + let half = effectcraft_raster::resample(&img, 50, 40); + let mut frame = |g: i64| Some(Arc::new(effectcraft_raster::resample(&frame_at(g), 50, 40))); + let _ = frame_matte(&p, [100.0, 80.0], 25.0, 0, 0.5, &half, Some(stub.clone()), &mut frame, &|| false); + assert_eq!(*stub.bbox.lock().unwrap(), Some([37.0, 27.0, 50.0, 40.0])); } } diff --git a/crates/track/src/roto/mod.rs b/crates/track/src/roto/mod.rs index 7138a93a..4efb4335 100644 --- a/crates/track/src/roto/mod.rs +++ b/crates/track/src/roto/mod.rs @@ -59,6 +59,9 @@ pub enum StrokeKind { Refine, /// Refine Edge tool with Alt: remove from the band. RefineErase, + /// A box around the object (Object Matte): two opposite corners in `points`. Trained models + /// get it as their box prompt; the graph cut ignores it. + Box, } impl StrokeKind { @@ -830,8 +833,9 @@ fn points_along(s: &Stroke, scale: f64, max: usize) -> Vec<[f32; 2]> { out } -/// A model's prompt from the strokes: points along the foreground and background strokes. -fn stroke_prompt(strokes: &[&Stroke], scale: f64) -> Prompt { +/// A model's prompt from the strokes: points along the foreground and background strokes, and +/// the box of a [`StrokeKind::Box`] stroke (matte pixels). +pub fn stroke_prompt(strokes: &[&Stroke], scale: f64) -> Prompt { let mut points = vec![]; for (kind, fg) in [(StrokeKind::Fg, true), (StrokeKind::Bg, false)] { let ss: Vec<&&Stroke> = strokes.iter().filter(|s| s.kind == kind).collect(); @@ -840,11 +844,20 @@ fn stroke_prompt(strokes: &[&Stroke], scale: f64) -> Prompt { points.extend(points_along(s, scale, per).into_iter().map(|p| (p, fg))); } } - Prompt { points, ..Default::default() } + let bbox = strokes.iter().filter(|s| s.kind == StrokeKind::Box).find_map(|s| box_of(s, scale)); + Prompt { points, bbox, ..Default::default() } } -/// Add a prior matte to a prompt: its box (a little larger), up to four points spread through -/// its core, and the matte itself. +/// The box `[x0, y0, x1, y1]` (matte pixels) of a box stroke: its first two points as opposite +/// corners; `None` when they are missing, not finite or the box is empty. +fn box_of(s: &Stroke, scale: f64) -> Option<[f32; 4]> { + let (a, c) = (s.points.first()?, s.points.get(1)?); + let v = [a[0].min(c[0]), a[1].min(c[1]), a[0].max(c[0]), a[1].max(c[1])].map(|v| (v * scale) as f32); + (v.iter().all(|x| x.is_finite()) && v[2] > v[0] && v[3] > v[1]).then_some(v) +} + +/// Add a prior matte to a prompt: its box (a little larger; unless the prompt has one), up to +/// four points spread through its core, and the matte itself. fn add_prior_prompt(prompt: &mut Prompt, soft: &[f32], w: usize, h: usize) -> bool { let bin: Vec = soft.iter().map(|v| (*v >= 0.5) as u8).collect(); let (mut x0, mut y0, mut x1, mut y1) = (usize::MAX, usize::MAX, 0usize, 0usize); @@ -858,7 +871,10 @@ fn add_prior_prompt(prompt: &mut Prompt, soft: &[f32], w: usize, h: usize) -> bo return false; } let (mx, my) = (((x1 - x0) as f32 * 0.05).max(4.0), ((y1 - y0) as f32 * 0.05).max(4.0)); - prompt.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)]); + // A box the user drew on this frame wins over the prior's. + if prompt.bbox.is_none() { + prompt.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)]); + } // Farthest-point picks among the core (deep inside), starting from the deepest pixel. let sd = matting::signed_distance(&bin, w, h); let deep = sd.iter().copied().fold(0.0f32, f32::max); diff --git a/crates/track/src/roto/tests.rs b/crates/track/src/roto/tests.rs index f92e67d5..65d1743f 100644 --- a/crates/track/src/roto/tests.rs +++ b/crates/track/src/roto/tests.rs @@ -428,3 +428,25 @@ fn mobilesam_segments_and_propagates_the_disk() { eprintln!("MobileSAM: base IoU {jb:.4}, worst propagated {worst:.4} (classic {cworst:.4}), {:.0} ms/frame", t.elapsed().as_secs_f64() * 1000.0 / 20.0); assert!(jb > 0.95 && worst > 0.93 && worst >= cworst - 0.01, "base {jb}, worst {worst} vs classic {cworst}"); } + +/// A box stroke becomes the model's box prompt (scaled, corners in any order); the graph cut +/// ignores it; a prior's box does not replace it. +#[test] +fn box_stroke_is_the_model_box_prompt() { + let fg = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 2.0, points: vec![[50.0, 40.0]] }; + let bx = Stroke { kind: StrokeKind::Box, frame: 0, radius: 0.0, points: vec![[80.0, 60.0], [20.0, 10.0]] }; + let p = stroke_prompt(&[&fg, &bx], 0.5); + assert_eq!(p.bbox, Some([10.0, 5.0, 40.0, 30.0])); + assert_eq!(p.points, vec![([25.0, 20.0], true)]); + // Degenerate or non-finite boxes are no prompt. + for pts in [vec![[20.0, 10.0]], vec![[20.0, 10.0], [20.0, 50.0]], vec![[f64::NAN, 10.0], [80.0, 60.0]]] { + let b = Stroke { points: pts, ..bx.clone() }; + assert_eq!(stroke_prompt(&[&fg, &b], 1.0).bbox, None); + } + assert!(!StrokeKind::Box.is_segmentation()); + assert!(rasterize(&[&bx], &[StrokeKind::Fg, StrokeKind::Bg], 1.0, 100, 80).iter().all(|v| !v)); + let mut q = stroke_prompt(&[&fg, &bx], 1.0); + let soft: Vec = (0..100 * 80).map(|i| ((i % 100) < 10) as u8 as f32).collect(); + assert!(add_prior_prompt(&mut q, &soft, 100, 80)); + assert_eq!(q.bbox, Some([20.0, 10.0, 80.0, 60.0])); +}