From 818ce2ab427be38010308488ee838f5eb4ae05a8 Mon Sep 17 00:00:00 2001 From: tylerjenningsw Date: Sun, 18 Oct 2026 11:28:13 -0400 Subject: [PATCH] Roto Brush: model prompts follow the latest stroke (#1646) A trained model is no longer prompted with both labels at one spot: points along a stroke that a later stroke of the other kind painted over are dropped, and the propagated matte's points, box and mask take this frame's strokes over it. --- crates/track/src/roto.rs | 42 ++++++++++++++++++++++++++-------- crates/track/src/roto/tests.rs | 23 ++++++++++++++++++- 2 files changed, 54 insertions(+), 11 deletions(-) diff --git a/crates/track/src/roto.rs b/crates/track/src/roto.rs index 9e5d8d6e1..0072ad223 100644 --- a/crates/track/src/roto.rs +++ b/crates/track/src/roto.rs @@ -375,6 +375,12 @@ 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]; @@ -894,15 +900,21 @@ 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, and -/// the box of a [`StrokeKind::Box`] stroke (matte pixels). +/// 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(); 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(); + let ss: Vec<(usize, &Stroke)> = seg.iter().copied().enumerate().filter(|(_, s)| s.kind == kind).collect(); let per = (16 / ss.len().max(1)).max(1); - for s in ss { - points.extend(points_along(s, scale, per).into_iter().map(|p| (p, fg))); + 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))); } } let bbox = strokes.iter().filter(|s| s.kind == StrokeKind::Box).find_map(|s| box_of(s, scale)); @@ -918,8 +930,18 @@ fn box_of(s: &Stroke, scale: f64) -> Option<[f32; 4]> { } /// 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 { +/// four points spread through its core, and the matte itself, all with this frame's strokes +/// (`hard`, from [`resolve`]) painted over it. +fn add_prior_prompt(prompt: &mut Prompt, soft: &[f32], hard: &[u8], w: usize, h: usize) -> bool { + let soft: Vec = soft + .iter() + .enumerate() + .map(|(i, v)| match hard.get(i) { + Some(&HARD_FG) => 1.0, + Some(&HARD_BG) => 0.0, + _ => *v, + }) + .collect(); 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); for (i, v) in bin.iter().enumerate() { @@ -949,7 +971,7 @@ fn add_prior_prompt(prompt: &mut Prompt, soft: &[f32], w: usize, h: usize) -> bo } } prompt.points.extend(picks.into_iter().map(|i| ([(i % w) as f32 + 0.5, (i / w) as f32 + 0.5], true))); - prompt.mask = Some(soft.to_vec()); + prompt.mask = Some(soft); true } @@ -960,7 +982,7 @@ pub fn segment_with(img: &Image, strokes: &[&Stroke], prior: Option<&Prior>, opt 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); if let Some(p) = prior { - add_prior_prompt(&mut prompt, &p.soft, w, h); + add_prior_prompt(&mut prompt, &p.soft, &resolve(strokes, StrokeKind::Fg, StrokeKind::Bg, opts.scale, w, h), w, h); } if !prompt.points.iter().any(|(_, fg)| *fg) { return segment(img, strokes, prior, opts, refine_radius); @@ -991,7 +1013,7 @@ pub fn propagate_with( 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, w, h) { + if !add_prior_prompt(&mut prompt, &prior.soft, &resolve(strokes, StrokeKind::Fg, StrokeKind::Bg, opts.scale, w, h), 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 1922dd0a2..f5170e754 100644 --- a/crates/track/src/roto/tests.rs +++ b/crates/track/src/roto/tests.rs @@ -509,6 +509,27 @@ fn box_stroke_is_the_model_box_prompt() { 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!(add_prior_prompt(&mut q, &soft, &[], 100, 80)); assert_eq!(q.bbox, Some([20.0, 10.0, 80.0, 60.0])); } + +/// A model is never prompted against the latest correction: stroke points that a later stroke +/// of the other kind painted over are dropped, and the propagated matte's points and mask leave +/// out what this frame's background strokes removed. +#[test] +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 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); + let mask = q.mask.unwrap(); + assert!(hard.iter().zip(&mask).all(|(hd, m)| *hd != HARD_BG || *m == 0.0)); +}