S2.45: a correction stroke re-tracks only the frames it changes

Chain::track tracks the frames of the span without a tracked result, in a window
that restarts at the nearest prompted frame; unchanged spans are not tracked
again. Test with a recording tracker: first pass 0..20, no second pass, then a
correction on frame 8 tracks 8..20 with that prompt only.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

crates/effects/src/roto.rs

@@ -540,8 +540,9 @@ impl Chain {
         self.model.as_ref().is_some_and(|m| m.tracker().is_some())
     }
 
-    /// Follow the object through the whole span with the model's tracker, from the strokes on
-    /// their frames, and cache every frame's segmentation (under the tracker's keys, with its
+    /// Follow the object through the span with the model's tracker, from the strokes on their
+    /// frames (only the frames without a tracked result yet: after a correction stroke, the frames it
+    /// changes, tracked on from the nearest prompted frame), and cache every frame's segmentation (under the tracker's keys, with its
     /// confidence). The strokes stay hard constraints and the edges are refined at full
     /// resolution; frames the tracker can't give are left to [`Chain::ensure`]'s classic engine.
     /// `progress` gets the number of frames finished. `Ok(None)`: no tracker, nothing done.
@@ -553,9 +554,19 @@ impl Chain {
     ) -> Result<Option<TrackSummary>, String> {
         let Some(tracker) = self.model.as_ref().and_then(|m| m.tracker()) else { return Ok(None) };
         let (Some(_), true) = (self.data.base, self.data.span[1] >= self.data.span[0]) else { return Err("no base frame".into()) };
+        // Frames of the span that already have a tracked result (this session's, or the saved
+        // tracking) are kept: only the rest is tracked. A correction stroke changes the keys of the
+        // frames from it on (both directions), so only those come back here.
+        let (lo_span, hi_span) = (self.data.span[0], self.data.span[1]);
+        let needed: Vec<i64> = (lo_span..=hi_span)
+            .filter(|f| self.keys.get(f).is_some_and(|k| self.saved_seg(*f, self.opts.scale).is_none() && cached(tracked_key(*k), self.opts.scale).is_none()))
+            .collect();
+        let (Some(&need_lo), Some(&need_hi)) = (needed.first(), needed.last()) else {
+            return Ok(Some(TrackSummary::default()));
+        };
         // A prompt is a few clicks from the strokes and a box around what the classic engine makes of
         // them on that frame (clicks alone pick a part of the object, not all of it).
-        let mut prompts: Vec<(i64, effectcraft_segment::Prompt)> = vec![];
+        let mut all_prompts: Vec<(i64, effectcraft_segment::Prompt)> = vec![];
         for f in self.data.stroke_frames().into_iter().filter(|f| self.data.in_span(*f)) {
             let strokes = self.data.strokes_at(f);
             let p = match frame(f) {
@@ -563,13 +574,20 @@ impl Chain {
                 None => rb::tracker_prompt(&strokes, self.opts.scale),
             };
             if p.points.iter().any(|(_, fg)| *fg) {
-                prompts.push((f, p));
+                all_prompts.push((f, p));
             }
         }
-        if prompts.is_empty() {
+        if all_prompts.is_empty() {
             return Err("no foreground strokes to follow".into());
         }
-        let req = TrackRequest::single(self.data.span[0], (self.data.span[1] - self.data.span[0] + 1) as usize, prompts);
+        // The window to track: from the nearest prompted frame at or before the first frame needed
+        // (a correction restarts the object there) to the last frame needed, reaching the first
+        // prompted frame when everything needed lies before it.
+        let first_prompt = all_prompts.first().map_or(need_lo, |p| p.0);
+        let lo = all_prompts.iter().map(|p| p.0).filter(|f| *f <= need_lo).max().unwrap_or(need_lo);
+        let hi = if need_hi < first_prompt { first_prompt } else { need_hi };
+        let prompts: Vec<(i64, effectcraft_segment::Prompt)> = all_prompts.into_iter().filter(|p| p.0 >= lo && p.0 <= hi).collect();
+        let req = TrackRequest::single(lo, (hi - lo + 1) as usize, prompts);
         let frame = std::cell::RefCell::new(frame);
         let mut summary = TrackSummary::default();
         tracker.track(

crates/engine/src/tests/roto.rs

@@ -940,3 +940,64 @@ fn without_a_tracker_the_chain_is_exactly_the_classic_engine() {
     }
     forget(&s, clip, uid);
 }
+
+/// Records the windows it is asked to track, and follows the disc exactly like [`Exact`].
+#[derive(Default)]
+struct Recorder {
+    requests: std::sync::Mutex<Vec<(i64, usize, Vec<i64>)>>,
+}
+
+impl effectcraft_segment::MaskModel for Recorder {
+    fn info(&self) -> &'static effectcraft_segment::ModelInfo {
+        &effectcraft_segment::SAM2_TINY
+    }
+    fn segment(&self, _: &[[f32; 3]], _: usize, _: usize, _: &effectcraft_segment::Prompt) -> effectcraft_segment::Result<Vec<f32>> {
+        Err("no weights".into())
+    }
+    fn tracker(&self) -> Option<&dyn effectcraft_segment::ObjectTracker> {
+        Some(self)
+    }
+}
+
+impl effectcraft_segment::ObjectTracker for Recorder {
+    fn track(
+        &self,
+        req: &effectcraft_segment::TrackRequest,
+        frame: &mut dyn FnMut(i64) -> Option<std::sync::Arc<effectcraft_segment::ClipFrame>>,
+        cancel: &dyn Fn() -> bool,
+        out: &mut dyn FnMut(usize, i64, effectcraft_segment::TrackedFrame),
+    ) -> effectcraft_segment::Result<()> {
+        let prompted = req.objects.iter().flatten().map(|(f, _)| *f).collect();
+        self.requests.lock().unwrap().push((req.first, req.count, prompted));
+        effectcraft_segment::ObjectTracker::track(&Exact, req, frame, cancel, out)
+    }
+}
+
+/// A correction stroke re-tracks only what it changes: the frames from it on, restarted at the
+/// corrected frame; the frames before it keep their tracked results.
+#[test]
+fn a_correction_re_tracks_only_the_frames_it_changes() {
+    let (mut s, clip, _) = painted(center, "disk-retrack.mov");
+    let rec = std::sync::Arc::new(Recorder::default());
+    effectcraft_effects::roto::set_model(Some(rec.clone()));
+    let run = |s: &mut Session| s.execute_checked("roto.propagate", json!({"layer": clip.0, "direction": "forward", "wait": true}));
+    let first = run(&mut s);
+    assert_eq!(rec.requests.lock().unwrap().as_slice(), [(0, 21, vec![0])], "the first pass tracks the whole span");
+    first.unwrap();
+    // Nothing changed: nothing to track.
+    run(&mut s).unwrap();
+    assert_eq!(rec.requests.lock().unwrap().len(), 1, "an unchanged span is not tracked again");
+    // A correction on frame 8: the frames from 8 on are tracked again, restarting at frame 8.
+    let c = center(8);
+    s.execute_checked("roto.stroke", json!({"layer": clip.0, "frame": 8, "points": line([c[0] - 20.0, c[1]], [c[0] + 20.0, c[1]]), "radius": 6})).unwrap();
+    let again = run(&mut s);
+    let reqs = rec.requests.lock().unwrap().clone();
+    effectcraft_effects::roto::set_model(None);
+    again.unwrap();
+    assert_eq!(reqs.len(), 2);
+    assert_eq!(reqs[1], (8, 13, vec![8]), "the window runs from the corrected frame to the end of the span");
+    // The frames before the correction are still the first pass's, the rest the second's.
+    let st = s.execute("roto.status", json!({"layer": clip.0, "frame": 3})).unwrap();
+    assert_eq!(st["source"], "tracker", "{st}");
+    assert_eq!(st["tracker"]["frames"].as_object().unwrap().len(), 21, "{st}");
+}

docs/roto-brush-sam2.md

@@ -29,6 +29,15 @@ strokes and span they came from are unchanged. Changing a stroke or the span mak
 Choosing another model, or the classic engine, leaves them in place and `roto.status` says so; a classic Propagate
 replaces them with classic mattes. Object Matte saves nothing.
 
+## Corrections re-track only what they change
+A correction stroke changes the keys of the frames from it on (in both directions from the base), so Propagate tracks
+only those again: the window starts at the nearest prompted frame at or before the first frame it needs (the corrected
+frame, which restarts the object with its own prompt) and runs to the last frame it needs. Frames before the correction
+keep their tracked results, and a Propagate with nothing changed does no tracking. The older non-conditioning memory
+frames before the corrected frame are not rebuilt, so frames just after it can differ slightly from a from-scratch pass.
+Tracked frames of a session are in memory; a project reopened in a new process re-tracks the whole span after a stroke
+change (its saved tracking is stale as a whole).
+
 ## How it works
 `crates/sam2` is a Rust port (on [candle](https://github.com/huggingface/candle), native only) of the Hugging Face
 `transformers` SAM 2 implementation, Apache-2.0 (see its `NOTICE`). Every stage is checked against PyTorch by