From b5460f653b8dd0b49d51dc1c568502e7a4c1d318 Mon Sep 17 00:00:00 2001 From: tylerjenningsw Date: Sun, 18 Oct 2026 12:15:11 -0400 Subject: [PATCH] Roto Brush: a correction keeps the model prompted with what is left of a stroke (#1646) Dropping the points a later stroke of the other kind painted over could leave a stroke with none, even though most of it was still painted: a small Alt tap in the middle of a large foreground dot removed its only point, and the selected model was skipped for the classical cut. Points are now kept by the resolved ownership under them, and a stroke that loses every sampled point but still owns pixels is prompted from the one deepest inside them. A stroke painted over completely still prompts nothing. --- crates/track/src/roto.rs | 81 ++++++++++++++++++++++++---------- crates/track/src/roto/tests.rs | 64 +++++++++++++++++++++++---- 2 files changed, 113 insertions(+), 32 deletions(-) diff --git a/crates/track/src/roto.rs b/crates/track/src/roto.rs index 0072ad223..998f5100d 100644 --- a/crates/track/src/roto.rs +++ b/crates/track/src/roto.rs @@ -375,12 +375,6 @@ fn sweep(s: &Stroke, scale: f64, w: usize, h: usize, mut paint: impl FnMut(usize } } -/// Whether `s` covers the point `p` (matte pixels at `scale`). -fn covers(s: &Stroke, scale: f64, p: [f64; 2]) -> bool { - let (r, segs) = segments(s, scale); - segs.iter().any(|(a, b)| segment_distance(p, *a, *b) <= r) -} - /// Pixels covered by `strokes` (of the given kinds) at `scale`. pub fn rasterize(strokes: &[&Stroke], kinds: &[StrokeKind], scale: f64, w: usize, h: usize) -> Vec { let mut m = vec![false; w * h]; @@ -900,27 +894,66 @@ 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 (but -/// not where a later stroke of the other kind painted over them), and the box of a -/// [`StrokeKind::Box`] stroke (matte pixels). -pub fn stroke_prompt(strokes: &[&Stroke], scale: f64) -> Prompt { - let seg: Vec<&Stroke> = strokes.iter().copied().filter(|s| s.kind.is_segmentation()).collect(); +/// 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 +/// sampled along it, but not everywhere, is prompted from the middle of what is left of it. +pub fn stroke_prompt(strokes: &[&Stroke], hard: &[u8], scale: f64, w: usize, h: usize) -> Prompt { let mut points = vec![]; - for (kind, fg) in [(StrokeKind::Fg, true), (StrokeKind::Bg, false)] { - let ss: Vec<(usize, &Stroke)> = seg.iter().copied().enumerate().filter(|(_, s)| s.kind == kind).collect(); + for (kind, own, other) in [(StrokeKind::Fg, HARD_FG, HARD_BG), (StrokeKind::Bg, HARD_BG, HARD_FG)] { + let ss: Vec<&&Stroke> = strokes.iter().filter(|s| s.kind == kind).collect(); let per = (16 / ss.len().max(1)).max(1); - for (k, s) in ss { - let kept = |p: &[f32; 2]| { - let p = [p[0] as f64, p[1] as f64]; - seg.iter().skip(k + 1).rev().find(|t| covers(t, scale, p)).is_none_or(|t| t.kind == kind) - }; - points.extend(points_along(s, scale, per).into_iter().filter(kept).map(|p| (p, fg))); + for s in ss { + let along = points_along(s, scale, per); + let mut kept: Vec<[f32; 2]> = along.iter().copied().filter(|p| owner(hard, w, h, *p) != Some(other)).collect(); + if kept.is_empty() && !along.is_empty() { + kept.extend(deepest_owned(s, hard, own, scale, w, h)); + } + points.extend(kept.into_iter().map(|p| (p, kind == StrokeKind::Fg))); } } let bbox = strokes.iter().filter(|s| s.kind == StrokeKind::Box).find_map(|s| box_of(s, scale)); Prompt { points, bbox, ..Default::default() } } +/// The value of `hard` under the point `p` (`None` outside the frame). +fn owner(hard: &[u8], w: usize, h: usize, p: [f32; 2]) -> Option { + if !(p[0] >= 0.0 && p[1] >= 0.0) { + return None; + } + let (x, y) = (p[0] as usize, p[1] as usize); + if x >= w || y >= h { + return None; + } + hard.get(y * w + x).copied() +} + +/// The centre of the pixel deepest inside what `s` still owns (where `hard` is `own` under it), +/// if anything. +fn deepest_owned(s: &Stroke, hard: &[u8], own: u8, scale: f64, w: usize, h: usize) -> Option<[f32; 2]> { + let mut owned = vec![]; + sweep(s, scale, w, h, |i| { + if hard.get(i) == Some(&own) { + owned.push((i % w, i / w)); + } + }); + let (x0, x1) = (owned.iter().map(|p| p.0).min()?, owned.iter().map(|p| p.0).max()?); + let (y0, y1) = (owned.iter().map(|p| p.1).min()?, owned.iter().map(|p| p.1).max()?); + // One pixel of margin all round, so depth is measured to the nearest pixel it doesn't own. + let (cw, ch) = (x1 - x0 + 3, y1 - y0 + 3); + let at = |(x, y): (usize, usize)| (y - y0 + 1) * cw + (x - x0 + 1); + let mut free = vec![true; cw * ch]; + for p in &owned { + if let Some(v) = free.get_mut(at(*p)) { + *v = false; + } + } + let depth = matting::distance(&free, cw, ch); + let d = |p: &(usize, usize)| depth.get(at(*p)).copied().unwrap_or(0.0); + let (x, y) = owned.iter().copied().max_by(|a, b| d(a).total_cmp(&d(b)))?; + Some([x as f32 + 0.5, y as f32 + 0.5]) +} + /// 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]> { @@ -980,9 +1013,10 @@ fn add_prior_prompt(prompt: &mut Prompt, soft: &[f32], hard: &[u8], w: usize, h: pub fn segment_with(img: &Image, strokes: &[&Stroke], prior: Option<&Prior>, opts: &SegOpts, refine_radius: f64, model: Option<&dyn MaskModel>) -> FrameSeg { let (w, h) = (img.width as usize, img.height as usize); let Some(m) = model.filter(|_| w > 0 && h > 0) else { return segment(img, strokes, prior, opts, refine_radius) }; - let mut prompt = stroke_prompt(strokes, opts.scale); + let hard = resolve(strokes, StrokeKind::Fg, StrokeKind::Bg, opts.scale, w, h); + let mut prompt = stroke_prompt(strokes, &hard, opts.scale, w, h); if let Some(p) = prior { - add_prior_prompt(&mut prompt, &p.soft, &resolve(strokes, StrokeKind::Fg, StrokeKind::Bg, opts.scale, w, h), w, h); + add_prior_prompt(&mut prompt, &p.soft, &hard, w, h); } if !prompt.points.iter().any(|(_, fg)| *fg) { return segment(img, strokes, prior, opts, refine_radius); @@ -1012,8 +1046,9 @@ pub fn propagate_with( let (w, h) = (next_img.width as usize, next_img.height as usize); let Some(m) = model.filter(|_| w > 0 && h > 0) else { return propagate(prev_img, prev, next_img, strokes, opts, refine_radius) }; let prior = warp_prior(prev_img, prev, next_img, opts); - let mut prompt = stroke_prompt(strokes, opts.scale); - if !add_prior_prompt(&mut prompt, &prior.soft, &resolve(strokes, StrokeKind::Fg, StrokeKind::Bg, opts.scale, w, h), w, h) { + let hard = resolve(strokes, StrokeKind::Fg, StrokeKind::Bg, opts.scale, w, h); + let mut prompt = stroke_prompt(strokes, &hard, opts.scale, w, h); + if !add_prior_prompt(&mut prompt, &prior.soft, &hard, w, h) { return propagate(prev_img, prev, next_img, strokes, opts, refine_radius); } match m.segment(&rgb_of(next_img), w, h, &prompt) { diff --git a/crates/track/src/roto/tests.rs b/crates/track/src/roto/tests.rs index f5170e754..f0059c651 100644 --- a/crates/track/src/roto/tests.rs +++ b/crates/track/src/roto/tests.rs @@ -420,9 +420,8 @@ impl effectcraft_segment::MaskModel for FakeModel { } /// A grey bar on a two-tone background whose halves match the bar's halves: colour alone can't -/// find its edge, a model can. -#[test] -fn a_model_prior_shapes_the_cut_and_bad_models_are_ignored() { +/// find its edge, a model can. The frame, its truth and the bar's rectangle. +fn bar_frame() -> (Image, Vec, [usize; 4]) { let (w, h) = (160usize, 100usize); let rect = [40, 30, 120, 70]; let mut img = Image::new(w as u32, h as u32); @@ -437,6 +436,12 @@ fn a_model_prior_shapes_the_cut_and_bad_models_are_ignored() { gt[y * w + x] = inside as u8; } } + (img, gt, rect) +} + +#[test] +fn a_model_prior_shapes_the_cut_and_bad_models_are_ignored() { + let (img, gt, rect) = bar_frame(); let st = [line(StrokeKind::Fg, 0, [50.0, 50.0], [110.0, 50.0], 4.0)]; let refs: Vec<&Stroke> = st.iter().collect(); let opts = SegOpts::default(); @@ -497,17 +502,17 @@ fn mobilesam_segments_and_propagates_the_disk() { 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); + let p = stroke_prompt(&[&fg, &bx], &[], 0.5, 0, 0); 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_eq!(stroke_prompt(&[&fg, &b], &[], 1.0, 0, 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 mut q = stroke_prompt(&[&fg, &bx], &[], 1.0, 0, 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])); @@ -520,16 +525,57 @@ fn box_stroke_is_the_model_box_prompt() { fn model_prompts_follow_the_latest_stroke() { let green = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 4.0, points: vec![[50.0, 40.0]] }; let red = Stroke { kind: StrokeKind::Bg, ..green.clone() }; - assert_eq!(stroke_prompt(&[&green, &red], 1.0).points, vec![([50.0, 40.0], false)]); - assert_eq!(stroke_prompt(&[&green, &red, &green], 1.0).points, vec![([50.0, 40.0], true); 2]); let (w, h) = (100, 80); + let prompt = |st: &[&Stroke]| stroke_prompt(st, &resolve(st, StrokeKind::Fg, StrokeKind::Bg, 1.0, w, h), 1.0, w, h).points; + assert_eq!(prompt(&[&green, &red]), vec![([50.0, 40.0], false)]); + assert_eq!(prompt(&[&green, &red, &green]), vec![([50.0, 40.0], true); 2]); let soft: Vec = (0..w * h).map(|i| ((i % w) < 50) as u8 as f32).collect(); let erase = Stroke { kind: StrokeKind::Bg, frame: 0, radius: 15.0, points: vec![[25.0, 40.0]] }; let hard = resolve(&[&erase], StrokeKind::Fg, StrokeKind::Bg, 1.0, w, h); let mut q = Prompt::default(); assert!(add_prior_prompt(&mut q, &soft, &hard, w, h)); assert!(!q.points.is_empty()); - assert!(q.points.iter().all(|(p, fg)| *fg && !covers(&erase, 1.0, [p[0] as f64, p[1] as f64])), "{:?}", q.points); + assert!(q.points.iter().all(|(p, fg)| *fg && owner(&hard, w, h, *p) != Some(HARD_BG)), "{:?}", q.points); let mask = q.mask.unwrap(); assert!(hard.iter().zip(&mask).all(|(hd, m)| *hd != HARD_BG || *m == 0.0)); } + +/// A stand-in model that counts its calls; its foreground is a fixed rectangle. +struct CountingModel([usize; 4], std::sync::atomic::AtomicUsize); + +impl effectcraft_segment::MaskModel for CountingModel { + fn info(&self) -> &'static effectcraft_segment::ModelInfo { + &effectcraft_segment::MOBILE_SAM + } + fn segment(&self, _: &[[f32; 3]], w: usize, h: usize, _: &Prompt) -> Result, String> { + self.1.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let [x0, y0, x1, y1] = self.0; + Ok((0..w * h).map(|i| ((i % w) >= x0 && (i % w) < x1 && (i / w) >= y0 && (i / w) < y1) as u8 as f32 * 0.98 + 0.01).collect()) + } +} + +/// A small background correction in the middle of a large foreground dot removes the only point +/// sampled along the dot, but most of the dot is still foreground: the model is still prompted +/// with it, from a point inside what is left. A dot painted over completely prompts nothing. +#[test] +fn a_correction_keeps_the_model_prompted_with_what_is_left() { + let (img, gt, rect) = bar_frame(); + let (w, h) = (img.width as usize, img.height as usize); + let dot = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 15.0, points: vec![[80.0, 50.0]] }; + let tap = Stroke { kind: StrokeKind::Bg, radius: 2.0, ..dot.clone() }; + let refs = [&dot, &tap]; + let model = CountingModel(rect, Default::default()); + let seg = segment_with(&img, &refs, None, &SegOpts::default(), 0.0, Some(&model)); + assert_eq!(model.1.load(std::sync::atomic::Ordering::Relaxed), 1, "the model was prompted"); + assert!(iou(&seg.matte, >) > 0.95, "IoU {}", iou(&seg.matte, >)); + let hard = resolve(&refs, StrokeKind::Fg, StrokeKind::Bg, 1.0, w, h); + let fg: Vec<[f32; 2]> = stroke_prompt(&refs, &hard, 1.0, w, h).points.into_iter().filter(|p| p.1).map(|p| p.0).collect(); + let [q] = fg.as_slice() else { panic!("one foreground point, got {fg:?}") }; + let r = (q[0] - 80.0).hypot(q[1] - 50.0); + assert!(r > 5.0 && r < 13.0, "{q:?} sits inside what is left of the dot"); + assert_eq!(hard[q[1] as usize * w + q[0] as usize], HARD_FG); + let cover = Stroke { radius: 16.0, ..tap }; + let refs = [&dot, &cover]; + 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)); +}