diff --git a/crates/track/src/roto.rs b/crates/track/src/roto.rs index c742654e9..3bab65dff 100644 --- a/crates/track/src/roto.rs +++ b/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(src: &[T], w: usize, rect: [usize; 4]) -> Vec { + 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 = 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 = sd.iter().map(|d| *d > r).collect(); + let seeds: Vec = 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 = 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 = 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) } diff --git a/crates/track/src/roto/matting.rs b/crates/track/src/roto/matting.rs index 3f5327085..ff6e33162 100644 --- a/crates/track/src/roto/matting.rs +++ b/crates/track/src/roto/matting.rs @@ -145,13 +145,21 @@ pub fn box_mean(src: &[f32], w: usize, h: usize, r: usize) -> Vec { 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 { + (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 { - 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) diff --git a/crates/track/src/roto/tests.rs b/crates/track/src/roto/tests.rs index bfc00df74..e50dbbce8 100644 --- a/crates/track/src/roto/tests.rs +++ b/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) { + 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", ¬hing, 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 = 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 = 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()); + } +}