Roto Brush: propagate on the band bounding box (#2335)

* Roto Brush: propagate on the band bounding box

propagate() ran the signed distance, gauss, graph cut, model fit and
clean over the whole frame. Work on the crop of the band bounding box
(|sd| <= r plus the core, padded) and paste the labels back; fall back
to the full frame when the crop is not smaller. warp_prior stays
full-frame.

Fixes #2323

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

* Roto Brush: band crop padded by the prior blur, models fitted on the full frame

Review follow-up for the band bounding-box crop in propagate:

- The crop pad is now the reach of the prior blur (three box passes of
  matting::gauss, exposed as matting::gauss_reach) instead of 1 px, so every
  band pixel sees the same blur window as on the full frame.
- When the previous frame has no colour models they are fitted on the full
  frame before cropping, as the full-frame path does.
- The free pixels and the clean seeds are cut out of the full-frame signed
  distance instead of being recomputed on the crop; the seeds equal
  core_of(prior.soft), which saves one signed distance per frame.
- The full-frame path and the crop share recut_band/cut_band; full frame is
  just the whole rectangle. A frame with no band and no core (no
  foreground) skips the cut, which was a no-op.
- The w - 1 / h - 1 underflow on zero-sized frames is gone (band_rect
  returns None); regression test for 4x0, 0x4 and 0x0 (panicked before in
  debug builds).

Tests: crop vs full frame labels on a disk, a disk at the frame corner, a
donut, r = 2 and r = 3 with prior edges wider than the band, and no
foreground, each with fitted, given and equal colour models; the padded crop
blurs the prior like the full frame at every band pixel while a 1 px pad
does not.

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

---------

Co-authored-by: Claude Sonnet 5.5 <noreply@anthropic.com>
Co-authored-by: FredApps <153173854+FredApps@users.noreply.github.com>

crates/track/src/roto.rs

@@ -901,6 +901,87 @@ pub fn warp_prior(prev_img: &Image, prev: &FrameSeg, next_img: &Image, opts: &Se
     Prior { soft: res.iter().map(|r| r.0).collect(), fg: prev.fg.clone(), bg: prev.bg.clone(), refine: res.iter().map(|r| r.1).collect(), weight: PRIOR }
 }
 
+/// The inputs of a band re-cut on one (sub)image of `size`.
+struct Band<'a> {
+    size: (usize, usize),
+    rgb: &'a [[f32; 3]],
+    /// The warped soft matte (the prior, before blurring).
+    soft: &'a [f32],
+    free: &'a [bool],
+    seeds: &'a [bool],
+}
+
+/// Graph cut of the free pixels with the blurred prior, then [`clean`].
+fn cut_band(b: &Band, labels: &mut [u8], sigma: f64, models: (&Gmm, &Gmm)) {
+    let (w, h) = b.size;
+    let soft = matting::gauss(b.soft, w, h, sigma);
+    cut(b.rgb, b.size, labels, b.free, &vec![0u8; w * h], models, Some((&soft, PRIOR)));
+    clean(labels, b.seeds, &vec![false; w * h], w, h);
+}
+
+/// The half-open rectangle `[x0, y0, x1, y1)` around the pixels with `sd >= -r` (the band
+/// `|sd| <= r` and the core `sd > r`), grown by `pad` and clipped to the frame; `None` when
+/// there are none.
+fn band_rect(sd: &[f32], size: (usize, usize), r: f32, pad: usize) -> Option<[usize; 4]> {
+    let (w, h) = size;
+    let (mut x0, mut y0, mut x1, mut y1) = (usize::MAX, usize::MAX, 0, 0);
+    for (i, d) in sd.iter().enumerate() {
+        if *d >= -r {
+            let (x, y) = (i % w.max(1), i / w.max(1));
+            (x0, y0, x1, y1) = (x0.min(x), y0.min(y), x1.max(x), y1.max(y));
+        }
+    }
+    (x0 <= x1).then(|| [x0.saturating_sub(pad), y0.saturating_sub(pad), x1.saturating_add(pad + 1).min(w), y1.saturating_add(pad + 1).min(h)])
+}
+
+/// The pixels of `rect` (row-major) of a `w`-wide plane.
+fn crop_of<T: Copy>(src: &[T], w: usize, rect: [usize; 4]) -> Vec<T> {
+    let [x0, y0, x1, y1] = rect;
+    (y0..y1).flat_map(|y| src.get(y * w + x0..y * w + x1).unwrap_or(&[]).iter().copied()).collect()
+}
+
+/// Re-cut the band `|sd| <= r` around the edge of `labels` (the warped matte) with `soft` as the
+/// prior; outside it the labels stay. `models`: the colour models, fitted on the whole frame when
+/// `None`. With `crop` the cut runs on the bounding box of the band and the core, padded by the
+/// reach of the prior blur: outside that box every pixel is background more than `r` from the
+/// edge, the free pixels, their neighbours and the seeds are the same, and every free pixel sees
+/// the same blur window as on the full frame.
+fn recut_band(rgb: &[[f32; 3]], size: (usize, usize), labels: &mut [u8], soft: &[f32], r: f32, models: Option<(&Gmm, &Gmm)>, crop: bool) {
+    let (w, h) = size;
+    let fitted;
+    let models = match models {
+        Some(m) => m,
+        None => {
+            fitted = fit_models(rgb, labels);
+            (&fitted.0, &fitted.1)
+        }
+    };
+    let sd = matting::signed_distance(labels, w, h);
+    let free: Vec<bool> = sd.iter().map(|d| d.abs() <= r).collect();
+    // The seeds of `core_of(soft, ..)`: `labels` is `soft >= 0.5`, so this is its core.
+    let core: Vec<bool> = sd.iter().map(|d| *d > r).collect();
+    let seeds: Vec<bool> = if core.iter().any(|c| *c) { core } else { labels.iter().map(|v| *v != 0).collect() };
+    let sigma = (r as f64 / 4.0).max(1.0);
+    let full = [0, 0, w, h];
+    let rect = if crop { band_rect(&sd, size, r, matting::gauss_reach(sigma).max(1)) } else { Some(full) };
+    // No band and no core: there is no foreground, so nothing is free and `clean` keeps all.
+    let Some(rect) = rect else { return };
+    if rect == full {
+        cut_band(&Band { size, rgb, soft, free: &free, seeds: &seeds }, labels, sigma, models);
+        return;
+    }
+    let [x0, y0, x1, y1] = rect;
+    let cw = x1 - x0;
+    let (rgb_c, soft_c, free_c, seeds_c) = (crop_of(rgb, w, rect), crop_of(soft, w, rect), crop_of(&free, w, rect), crop_of(&seeds, w, rect));
+    let mut cl = crop_of(labels, w, rect);
+    cut_band(&Band { size: (cw, y1 - y0), rgb: &rgb_c, soft: &soft_c, free: &free_c, seeds: &seeds_c }, &mut cl, sigma, models);
+    for (y, row) in (y0..y1).zip(cl.chunks(cw.max(1))) {
+        if let Some(dst) = labels.get_mut(y * w + x0..y * w + x1) {
+            dst.copy_from_slice(row);
+        }
+    }
+}
+
 /// Propagate the previous frame's segmentation to the next frame: warp it, then re-cut in a
 /// band of Search Radius pixels around the warped boundary. Strokes on the next frame re-cut
 /// the whole frame with the warped matte as a soft prior instead.
@@ -912,15 +993,9 @@ pub fn propagate(prev_img: &Image, prev: &FrameSeg, next_img: &Image, strokes: &
     }
     let rgb = rgb_of(next_img);
     let mut labels: Vec<u8> = prior.soft.iter().map(|v| (*v >= 0.5) as u8).collect();
-    let sd = matting::signed_distance(&labels, w, h);
     let r = (opts.search_radius * opts.scale).max(2.0) as f32;
-    let free: Vec<bool> = sd.iter().map(|d| d.abs() <= r).collect();
-    let hard = vec![0u8; w * h];
-    let soft = matting::gauss(&prior.soft, w, h, (r as f64 / 4.0).max(1.0));
-    let (fg, bg) = if prev.fg.is_empty() || prev.bg.is_empty() { fit_models(&rgb, &labels) } else { (prev.fg.clone(), prev.bg.clone()) };
-    cut(&rgb, (w, h), &mut labels, &free, &hard, (&fg, &bg), Some((&soft, PRIOR)));
-    let seeds = core_of(&prior.soft, w, h, r as f64);
-    clean(&mut labels, &seeds, &vec![false; w * h], w, h);
+    let models = (!prev.fg.is_empty() && !prev.bg.is_empty()).then_some((&prev.fg, &prev.bg));
+    recut_band(&rgb, (w, h), &mut labels, &prior.soft, r, models, true);
     finish(next_img, &rgb, labels, strokes, Some(&prior.refine), opts, refine_radius)
 }
 

crates/track/src/roto/matting.rs

@@ -145,13 +145,21 @@ pub fn box_mean(src: &[f32], w: usize, h: usize, r: usize) -> Vec<f32> {
     out
 }
 
+/// Box radius of [`gauss`]: three passes of it match the variance, 3·((2r+1)² − 1)/12 = σ².
+/// `None` when the blur is a copy.
+fn gauss_box(sigma: f64) -> Option<usize> {
+    (sigma >= 0.3).then(|| (((4.0 * sigma * sigma + 1.0).sqrt() - 1.0) / 2.0).round().max(1.0) as usize)
+}
+
+/// How far (pixels, per axis) [`gauss`] reaches: an output pixel depends only on the input
+/// within this distance.
+pub fn gauss_reach(sigma: f64) -> usize {
+    gauss_box(sigma).map_or(0, |r| 3 * r)
+}
+
 /// Approximate Gaussian blur (three box passes) with standard deviation `sigma` pixels.
 pub fn gauss(src: &[f32], w: usize, h: usize, sigma: f64) -> Vec<f32> {
-    if sigma < 0.3 {
-        return src.to_vec();
-    }
-    // Box radius whose three passes match the variance: 3·((2r+1)² − 1)/12 = σ².
-    let r = (((4.0 * sigma * sigma + 1.0).sqrt() - 1.0) / 2.0).round().max(1.0) as usize;
+    let Some(r) = gauss_box(sigma) else { return src.to_vec() };
     let a = box_mean(src, w, h, r);
     let b = box_mean(&a, w, h, r);
     box_mean(&b, w, h, r)

crates/track/src/roto/tests.rs

@@ -722,3 +722,111 @@ fn the_union_of_objects_is_their_largest_extent() {
     let odd = FrameSeg::empty(2, 2);
     assert_eq!(FrameSeg::union(vec![seg(&[1]), odd]).matte, vec![0, 255, 0, 0]);
 }
+
+/// A scene for the band re-cut: `rgb` of a `w`×`h` frame where the foreground is `inside`, with
+/// a background that turns red left of `x = 30` (far from most shapes), and a soft prior that is
+/// `inside` shifted by (3, 2) px with an edge `soft_edge` px wide.
+fn band_scene(w: usize, h: usize, inside: &dyn Fn(f64, f64) -> f64, soft_edge: f64) -> (Vec<[f32; 3]>, Vec<f32>) {
+    let mut rgb = vec![[0.0f32; 3]; w * h];
+    let mut soft = vec![0.0f32; w * h];
+    for y in 0..h {
+        for x in 0..w {
+            let (px, py) = (x as f64 + 0.5, y as f64 + 0.5);
+            let a = inside(px, py).clamp(0.0, 1.0) as f32;
+            let f = fg_color(px, py);
+            let b = if px < 30.0 { [0.9, 0.1, 0.1] } else { bg_color(px, py) };
+            rgb[y * w + x] = [f[0] * a + b[0] * (1.0 - a), f[1] * a + b[1] * (1.0 - a), f[2] * a + b[2] * (1.0 - a)];
+            soft[y * w + x] = (inside(px - 3.0, py - 2.0) / soft_edge + 0.5).clamp(0.0, 1.0) as f32;
+        }
+    }
+    (rgb, soft)
+}
+
+#[test]
+fn band_crop_matches_the_full_frame() {
+    let (w, h) = (200usize, 150usize);
+    // Signed coverage: > 0.5 inside, falling off by 1 per pixel across the edge.
+    let disk = |cx: f64, cy: f64, r: f64| move |x: f64, y: f64| r + 0.5 - (x - cx).hypot(y - cy);
+    let mid = disk(120.0, 80.0, 30.0);
+    let corner = disk(185.0, 10.0, 28.0);
+    let donut = move |x: f64, y: f64| (40.0 + 0.5 - (x - 110.0).hypot(y - 75.0)).min((x - 110.0).hypot(y - 75.0) - 14.0 + 0.5);
+    let nothing = |_: f64, _: f64| -10.0;
+    let shapes: [(&str, &dyn Fn(f64, f64) -> f64, f32, f64); 8] = [
+        ("disk", &mid, 8.0, 1.0),
+        ("frame corner", &corner, 8.0, 1.0),
+        ("donut", &donut, 6.0, 1.0),
+        ("r = 2, soft edge", &mid, 2.0, 3.0),
+        // A prior edge wider than the band: the blur sees prior values outside it.
+        ("r = 2, wide soft edge", &mid, 2.0, 10.0),
+        ("r = 3, wide soft edge at the corner", &corner, 3.0, 14.0),
+        ("r = 12, soft edge", &mid, 12.0, 4.0),
+        ("no foreground", &nothing, 8.0, 1.0),
+    ];
+    let mut changed = 0;
+    for (name, inside, r, edge) in shapes {
+        let (rgb, soft) = band_scene(w, h, inside, edge);
+        let start: Vec<u8> = soft.iter().map(|v| (*v >= 0.5) as u8).collect();
+        let fitted = fit_models(&rgb, &start);
+        // Equal models leave the cut to the prior and the smoothness term, which makes it
+        // sensitive to the blurred prior.
+        for models in [None, Some((&fitted.0, &fitted.1)), Some((&fitted.1, &fitted.1))] {
+            let (mut full, mut crop) = (start.clone(), start.clone());
+            recut_band(&rgb, (w, h), &mut full, &soft, r, models, false);
+            recut_band(&rgb, (w, h), &mut crop, &soft, r, models, true);
+            let diff = full.iter().zip(&crop).filter(|(a, b)| a != b).count();
+            assert_eq!(diff, 0, "{name}, models {}: {diff} pixels differ", if models.is_some() { "given" } else { "fitted" });
+            changed += full.iter().zip(&start).filter(|(a, b)| a != b).count();
+        }
+    }
+    // The cut did move edges (the prior is offset from the shapes).
+    assert!(changed > 100, "{changed}");
+}
+
+/// Padded by the blur's reach, the crop blurs the prior exactly as the full frame does at every
+/// band pixel (up to the summation rounding of the box sums); a 1 px pad does not when the prior
+/// edge is wider than the band.
+#[test]
+fn band_rect_pad_covers_the_prior_blur() {
+    let (w, h) = (200usize, 150usize);
+    let mid = |x: f64, y: f64| 30.5 - (x - 120.0).hypot(y - 80.0);
+    let corner = |x: f64, y: f64| 28.5 - (x - 185.0).hypot(y - 10.0);
+    let cases: [(&dyn Fn(f64, f64) -> f64, f32, f64); 4] = [(&mid, 8.0, 1.0), (&mid, 2.0, 10.0), (&corner, 3.0, 14.0), (&mid, 12.0, 30.0)];
+    for (inside, r, edge) in cases {
+        let (_, soft) = band_scene(w, h, inside, edge);
+        let labels: Vec<u8> = soft.iter().map(|v| (*v >= 0.5) as u8).collect();
+        let sd = matting::signed_distance(&labels, w, h);
+        let sigma = (r as f64 / 4.0).max(1.0);
+        let full = matting::gauss(&soft, w, h, sigma);
+        let worst = |pad: usize| -> f32 {
+            let rect = band_rect(&sd, (w, h), r, pad).unwrap();
+            let [x0, y0, x1, y1] = rect;
+            let blurred = matting::gauss(&crop_of(&soft, w, rect), x1 - x0, y1 - y0, sigma);
+            let mut worst = 0.0f32;
+            for y in y0..y1 {
+                for x in x0..x1 {
+                    if sd[y * w + x].abs() <= r {
+                        worst = worst.max((blurred[(y - y0) * (x1 - x0) + x - x0] - full[y * w + x]).abs());
+                    }
+                }
+            }
+            worst
+        };
+        let reach = matting::gauss_reach(sigma);
+        assert!(worst(reach) < 1e-6, "r {r}, edge {edge}: {}", worst(reach));
+        if edge > 2.0 * r as f64 {
+            assert!(worst(1) > 1e-4, "r {r}, edge {edge}: a 1 px pad should differ, {}", worst(1));
+        }
+    }
+}
+
+#[test]
+fn propagate_survives_empty_frames() {
+    let opts = SegOpts::default();
+    for (w, h) in [(4u32, 0u32), (0, 4), (0, 0)] {
+        let img = Image::new(w, h);
+        let prev = FrameSeg::empty(w as usize, h as usize);
+        let next = propagate(&img, &prev, &img, &[], &opts, 0.0);
+        assert_eq!((next.w, next.h), (w as usize, h as usize));
+        assert!(next.matte.is_empty());
+    }
+}