From 0e0db685bf81e8558e7f1f539e038adc650b9cf3 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 20:24:08 +0200 Subject: [PATCH] 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 --- crates/effects/src/roto.rs | 30 ++++++++++++---- crates/engine/src/tests/roto.rs | 61 +++++++++++++++++++++++++++++++++ docs/roto-brush-sam2.md | 9 +++++ 3 files changed, 94 insertions(+), 6 deletions(-) diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index 869e2c132..c31c4836b 100644 --- a/crates/effects/src/roto.rs +++ b/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, 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 = (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( diff --git a/crates/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index 4ba557781..7d18e5c98 100644 --- a/crates/engine/src/tests/roto.rs +++ b/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)>>, +} + +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> { + 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>, + 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}"); +} diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index d5216e6d4..a4e58a36d 100644 --- a/docs/roto-brush-sam2.md +++ b/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