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

@@ -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<bool> {
     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<f32> = 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<u8> = 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) {

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<f32> = (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<f32> = (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));
+}