#1539: preserve stroke corrections and current main APIs

crates/track/src/roto.rs

@@ -950,18 +950,18 @@ fn points_along(s: &Stroke, scale: f64, max: usize) -> Vec<[f32; 2]> {
 /// SAM 2 was trained on a handful of clicks: with sixteen or more points it doubts its own mask
 /// (measured: confidence 0.25 with 16 points, 0.00 with 32, 0.99 with 1 to 7), so this picks at
 /// most [`TRACKER_FG_POINTS`] foreground and [`TRACKER_BG_POINTS`] background points, spread
-/// along the strokes.
-pub fn tracker_prompt(strokes: &[&Stroke], scale: f64) -> Prompt {
+/// along the strokes. Like [`stroke_prompt`], later corrections win over earlier paint.
+pub fn tracker_prompt(strokes: &[&Stroke], scale: f64, w: usize, h: usize) -> Prompt {
+    let hard = resolve(strokes, StrokeKind::Fg, StrokeKind::Bg, scale, w, h);
+    let prompt = stroke_prompt(strokes, &hard, scale, w, h);
     let mut points = vec![];
-    for (kind, fg, max) in [(StrokeKind::Fg, true, TRACKER_FG_POINTS), (StrokeKind::Bg, false, TRACKER_BG_POINTS)] {
-        let all: Vec<[f32; 2]> = strokes.iter().filter(|s| s.kind == kind).flat_map(|s| points_along(s, scale, 3)).collect();
+    for (fg, max) in [(true, TRACKER_FG_POINTS), (false, TRACKER_BG_POINTS)] {
+        let all: Vec<[f32; 2]> = prompt.points.iter().filter_map(|(p, label)| (*label == fg).then_some(*p)).collect();
         let n = all.len().min(max);
         // evenly spaced picks, keeping the first and the last
         points.extend((0..n).filter_map(|i| all.get(if n > 1 { i * (all.len() - 1) / (n - 1) } else { all.len() / 2 }).map(|p| (*p, fg))));
     }
-    // a box the user drew on the frame prompts the tracker too
-    let bbox = strokes.iter().filter(|s| s.kind == StrokeKind::Box).find_map(|s| box_of(s, scale));
-    Prompt { points, bbox, ..Default::default() }
+    Prompt { points, ..prompt }
 }
 
 /// [`tracker_prompt`] plus a box around the classic engine's matte of the same strokes. Clicks alone
@@ -969,7 +969,7 @@ pub fn tracker_prompt(strokes: &[&Stroke], scale: f64) -> Prompt {
 /// 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);
+    let mut p = tracker_prompt(strokes, scale, classic.w, classic.h);
     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);

crates/track/src/roto/tests.rs

@@ -591,13 +591,46 @@ fn a_correction_keeps_the_model_prompted_with_what_is_left() {
 fn tracker_prompts_carry_a_handful_of_points() {
     let st = base_strokes(0);
     let refs: Vec<&Stroke> = st.iter().collect();
-    let p = tracker_prompt(&refs, 1.0);
+    let p = tracker_prompt(&refs, 1.0, W as usize, H as usize);
     let (fg, bg) = (p.points.iter().filter(|(_, f)| *f).count(), p.points.iter().filter(|(_, f)| !*f).count());
     assert!((1..=TRACKER_FG_POINTS).contains(&fg) && (1..=TRACKER_BG_POINTS).contains(&bg), "{fg} fg, {bg} bg");
     // a long, finely sampled stroke still gives at most the limit, and an empty one gives none
     let long = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 2.0, points: (0..400).map(|i| [i as f64 * 0.4, 50.0]).collect(), object: 0 };
-    assert!(tracker_prompt(&[&long], 1.0).points.len() <= TRACKER_FG_POINTS);
-    assert!(tracker_prompt(&[], 1.0).points.is_empty());
+    assert!(tracker_prompt(&[&long], 1.0, W as usize, H as usize).points.len() <= TRACKER_FG_POINTS);
+    assert!(tracker_prompt(&[], 1.0, W as usize, H as usize).points.is_empty());
+}
+
+/// Tracker clicks follow the same correction order as the classic matte and model prompts.
+#[test]
+fn tracker_prompts_follow_latest_corrections_and_keep_remaining_foreground() {
+    let (w, h) = (100, 80);
+    let green = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 15.0, points: vec![[50.0, 40.0]], object: 0 };
+    let red = Stroke { kind: StrokeKind::Bg, radius: 16.0, ..green.clone() };
+    for strokes in [&[&green, &red][..], &[&red, &green][..]] {
+        let classic = FrameSeg::empty(w, h);
+        for prompt in [tracker_prompt(strokes, 1.0, w, h), tracker_prompt_boxed(strokes, 1.0, &classic)] {
+            let expected = std::ptr::eq(strokes[1], &green);
+            assert_eq!(prompt.points.iter().any(|(_, fg)| *fg), expected, "{:?}", prompt.points);
+            let hard = resolve(strokes, StrokeKind::Fg, StrokeKind::Bg, 1.0, w, h);
+            for (point, fg) in &prompt.points {
+                assert_eq!(owner(&hard, w, h, *point), Some(if *fg { HARD_FG } else { HARD_BG }));
+            }
+            assert!(!prompt.points.is_empty());
+        }
+    }
+    // A small correction removes the sampled centre, while most of the foreground dot remains.
+    let tap = Stroke { radius: 2.0, ..red };
+    let strokes = [&green, &tap];
+    let hard = resolve(&strokes, StrokeKind::Fg, StrokeKind::Bg, 1.0, w, h);
+    for prompt in [tracker_prompt(&strokes, 1.0, w, h), tracker_prompt_boxed(&strokes, 1.0, &FrameSeg::empty(w, h))] {
+        assert!(prompt.points.iter().any(|(_, fg)| *fg));
+        assert!(prompt.points.iter().any(|(_, fg)| !*fg));
+        for (point, fg) in &prompt.points {
+            assert_eq!(owner(&hard, w, h, *point), Some(if *fg { HARD_FG } else { HARD_BG }));
+        }
+        assert!(prompt.points.iter().filter(|(_, fg)| *fg).count() <= TRACKER_FG_POINTS);
+        assert!(prompt.points.iter().filter(|(_, fg)| !*fg).count() <= TRACKER_BG_POINTS);
+    }
 }
 
 #[test]