From aae936e9bdda9973bdf51e05c901cc0f7de7a79d Mon Sep 17 00:00:00 2001 From: tylerjenningsw Date: Sun, 18 Oct 2026 11:24:57 -0400 Subject: [PATCH 1/6] Roto Brush: foreground and background strokes apply in drawing order (#1646) A green stroke over an earlier Alt (red) stroke now restores the area: where strokes overlap, the one painted last sets the hard constraint, for the cut, the colour samples and the cleanup seeds alike. --- crates/track/src/roto.rs | 106 ++++++++++++++++++++------------- crates/track/src/roto/tests.rs | 22 +++++++ 2 files changed, 88 insertions(+), 40 deletions(-) diff --git a/crates/track/src/roto.rs b/crates/track/src/roto.rs index 305065d2c..acd1b0234 100644 --- a/crates/track/src/roto.rs +++ b/crates/track/src/roto.rs @@ -339,38 +339,76 @@ pub fn rgb_of(img: &Image) -> Vec<[f32; 3]> { img.data.par_iter().map(|p| if p[3] > 1e-6 { [p[0] / p[3], p[1] / p[3], p[2] / p[3]] } else { [0.0; 3] }).collect() } -/// Pixels covered by `strokes` (of the given kinds) at `scale`. -pub fn rasterize(strokes: &[&Stroke], kinds: &[StrokeKind], scale: f64, w: usize, h: usize) -> Vec { - let mut m = vec![false; w * h]; - for s in strokes.iter().filter(|s| kinds.contains(&s.kind)) { - let r = (s.radius * scale).max(0.5); - let pts: Vec<[f64; 2]> = s.points.iter().map(|p| [p[0] * scale, p[1] * scale]).collect(); - let segs: Vec<([f64; 2], [f64; 2])> = if pts.len() == 1 { vec![(pts[0], pts[0])] } else { pts.windows(2).map(|w| (w[0], w[1])).collect() }; - for (a, b) in segs { - let x0 = (a[0].min(b[0]) - r).floor().max(0.0) as usize; - let x1 = ((a[0].max(b[0]) + r).ceil().max(0.0) as usize).min(w); - let y0 = (a[1].min(b[1]) - r).floor().max(0.0) as usize; - let y1 = ((a[1].max(b[1]) + r).ceil().max(0.0) as usize).min(h); - let d = [b[0] - a[0], b[1] - a[1]]; - let l2 = d[0] * d[0] + d[1] * d[1]; - for y in y0..y1 { - for x in x0..x1 { - let p = [x as f64 + 0.5, y as f64 + 0.5]; - let t = if l2 > 0.0 { (((p[0] - a[0]) * d[0] + (p[1] - a[1]) * d[1]) / l2).clamp(0.0, 1.0) } else { 0.0 }; - let q = [a[0] + d[0] * t, a[1] + d[1] * t]; - if (p[0] - q[0]).hypot(p[1] - q[1]) <= r { - m[y * w + x] = true; - } +/// A stroke's radius and segments at `scale`. +fn segments(s: &Stroke, scale: f64) -> (f64, Vec<([f64; 2], [f64; 2])>) { + let pts: Vec<[f64; 2]> = s.points.iter().map(|p| [p[0] * scale, p[1] * scale]).collect(); + let segs = match pts.as_slice() { + [p] => vec![(*p, *p)], + _ => pts.windows(2).map(|w| (w[0], w[1])).collect(), + }; + ((s.radius * scale).max(0.5), segs) +} + +/// Distance from `p` to the segment `a`–`b`. +fn segment_distance(p: [f64; 2], a: [f64; 2], b: [f64; 2]) -> f64 { + let d = [b[0] - a[0], b[1] - a[1]]; + let l2 = d[0] * d[0] + d[1] * d[1]; + let t = if l2 > 0.0 { (((p[0] - a[0]) * d[0] + (p[1] - a[1]) * d[1]) / l2).clamp(0.0, 1.0) } else { 0.0 }; + (p[0] - a[0] - d[0] * t).hypot(p[1] - a[1] - d[1] * t) +} + +/// Call `paint` with every pixel `s` covers at `scale`. +fn sweep(s: &Stroke, scale: f64, w: usize, h: usize, mut paint: impl FnMut(usize)) { + let (r, segs) = segments(s, scale); + for (a, b) in segs { + let x0 = (a[0].min(b[0]) - r).floor().max(0.0) as usize; + let x1 = ((a[0].max(b[0]) + r).ceil().max(0.0) as usize).min(w); + let y0 = (a[1].min(b[1]) - r).floor().max(0.0) as usize; + let y1 = ((a[1].max(b[1]) + r).ceil().max(0.0) as usize).min(h); + for y in y0..y1 { + for x in x0..x1 { + if segment_distance([x as f64 + 0.5, y as f64 + 0.5], a, b) <= r { + paint(y * w + x); } } } } +} + +/// Pixels covered by `strokes` (of the given kinds) at `scale`. +pub fn rasterize(strokes: &[&Stroke], kinds: &[StrokeKind], scale: f64, w: usize, h: usize) -> Vec { + let mut m = vec![false; w * h]; + for s in strokes.iter().filter(|s| kinds.contains(&s.kind)) { + sweep(s, scale, w, h, |i| m[i] = true); + } m } -/// Hard constraints: 1 = foreground, 2 = background, 0 = free. -const HARD_FG: u8 = 1; -const HARD_BG: u8 = 2; +/// What a pair of opposing strokes (add / remove) leaves on each pixel. +const ON: u8 = 1; +const OFF: u8 = 2; + +/// [`ON`] where an `on` stroke covers a pixel, [`OFF`] where an `off` stroke does, 0 elsewhere. +/// Strokes apply in drawing order, so where they overlap the last one painted wins: a stroke +/// corrects the ones before it. +pub fn resolve(strokes: &[&Stroke], on: StrokeKind, off: StrokeKind, scale: f64, w: usize, h: usize) -> Vec { + let mut m = vec![0; w * h]; + for s in strokes { + let v = if s.kind == on { + ON + } else if s.kind == off { + OFF + } else { + continue; + }; + sweep(s, scale, w, h, |i| m[i] = v); + } + m +} + +/// Hard constraints: foreground, background, else free. +const HARD_FG: u8 = ON; +const HARD_BG: u8 = OFF; /// Graph-cut energy weights. const GAMMA: f64 = 50.0; @@ -630,21 +668,9 @@ pub fn segment(img: &Image, strokes: &[&Stroke], prior: Option<&Prior>, opts: &S return FrameSeg::empty(w, h); } let rgb = rgb_of(img); - let fgs = rasterize(strokes, &[StrokeKind::Fg], opts.scale, w, h); - let bgs = rasterize(strokes, &[StrokeKind::Bg], opts.scale, w, h); - let hard: Vec = fgs - .iter() - .zip(&bgs) - .map(|(f, b)| { - if *b { - HARD_BG - } else if *f { - HARD_FG - } else { - 0 - } - }) - .collect(); + let hard = resolve(strokes, StrokeKind::Fg, StrokeKind::Bg, opts.scale, w, h); + let fgs: Vec = hard.iter().map(|v| *v == HARD_FG).collect(); + let bgs: Vec = hard.iter().map(|v| *v == HARD_BG).collect(); // A warped matte is uncertain by about the search radius: blur it. A model's matte is not. let soft_prior: Option> = prior.map(|p| if p.weight > PRIOR { p.soft.clone() } else { matting::gauss(&p.soft, w, h, (opts.search_radius * opts.scale / 4.0).max(1.0)) }); diff --git a/crates/track/src/roto/tests.rs b/crates/track/src/roto/tests.rs index 3c4645cb2..2b4add27c 100644 --- a/crates/track/src/roto/tests.rs +++ b/crates/track/src/roto/tests.rs @@ -182,6 +182,28 @@ fn a_correction_stroke_fixes_a_bad_frame() { assert!(iou(&next.matte, >2) > 0.93); } +/// A stroke corrects the ones painted before it: foreground over an earlier background stroke +/// restores the area, and background over that removes it again. +#[test] +fn later_strokes_override_earlier_ones() { + let c = center(0); + let (img, gt) = frame_at(c); + let spot = [c[0], c[1] + 20.0]; + let i = spot[1] as usize * W as usize + spot[0] as usize; + let red = line(StrokeKind::Bg, 0, [spot[0] - 6.0, spot[1]], [spot[0] + 6.0, spot[1]], 6.0); + let green = Stroke { kind: StrokeKind::Fg, ..red.clone() }; + let seg = |st: &[Stroke]| segment(&img, &st.iter().collect::>(), None, &SegOpts::default(), 0.0); + let mut st = base_strokes(0); + st.push(red.clone()); + assert_eq!(seg(&st).matte[i], 0, "background over the disk removes it"); + st.push(green); + let restored = seg(&st); + assert_eq!(restored.matte[i], 1, "foreground over the background stroke restores it"); + assert!(iou(&restored.matte, >) > 0.95, "IoU {}", iou(&restored.matte, >)); + st.push(red); + assert_eq!(seg(&st).matte[i], 0, "background again removes it again"); +} + /// A soft-edged disk (alpha ramp over `width` pixels) over a background. fn soft_frame(c: [f64; 2], width: f64, fg: [f32; 3], bg: [f32; 3]) -> (Image, Vec) { let mut img = Image::new(W, H); From af995ea26aea6ec41522e5903933248cf0474c27 Mon Sep 17 00:00:00 2001 From: tylerjenningsw Date: Sun, 18 Oct 2026 11:25:54 -0400 Subject: [PATCH 2/6] Roto Brush: Refine Edge strokes apply in drawing order (#1646) A Refine Edge stroke over an earlier Alt (erase) stroke now restores the band there; the band carried from the previous frame stays under this frame's strokes. --- crates/track/src/roto.rs | 5 ++--- crates/track/src/roto/tests.rs | 25 +++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/crates/track/src/roto.rs b/crates/track/src/roto.rs index acd1b0234..9e5d8d6e1 100644 --- a/crates/track/src/roto.rs +++ b/crates/track/src/roto.rs @@ -806,8 +806,7 @@ fn core_of(soft: &[f32], w: usize, h: usize, r: f64) -> Vec { fn finish(_img: &Image, rgb: &[[f32; 3]], labels: Vec, strokes: &[&Stroke], warped_refine: Option<&[u8]>, opts: &SegOpts, refine_radius: f64) -> FrameSeg { let (w, h) = (_img.width as usize, _img.height as usize); let (fg, bg) = fit_models(rgb, &labels); - let add = rasterize(strokes, &[StrokeKind::Refine], opts.scale, w, h); - let erase = rasterize(strokes, &[StrokeKind::RefineErase], opts.scale, w, h); + let painted = resolve(strokes, StrokeKind::Refine, StrokeKind::RefineErase, opts.scale, w, h); let near: Option> = warped_refine.filter(|r| r.iter().any(|v| *v != 0)).map(|_| matting::signed_distance(&labels, w, h)); let rr = (refine_radius * opts.scale + 2.0) as f32; let refine: Vec = (0..w * h) @@ -816,7 +815,7 @@ fn finish(_img: &Image, rgb: &[[f32; 3]], labels: Vec, strokes: &[&Stroke], (Some(r), Some(sd)) => r[i] != 0 && sd[i].abs() <= rr, _ => false, }; - ((carried || add[i]) && !erase[i]) as u8 + ((carried && painted[i] != OFF) || painted[i] == ON) as u8 }) .collect(); FrameSeg { w, h, matte: labels, refine, fg, bg } diff --git a/crates/track/src/roto/tests.rs b/crates/track/src/roto/tests.rs index 2b4add27c..1922dd0a2 100644 --- a/crates/track/src/roto/tests.rs +++ b/crates/track/src/roto/tests.rs @@ -287,6 +287,31 @@ fn refine_edge_recovers_a_soft_edge_and_decontaminates() { assert!(a2.iter().zip(&seg.matte).all(|(x, y)| (*x - *y as f32).abs() < 1e-6)); } +/// Refine Edge strokes correct each other the same way: refine over an earlier erase restores +/// the band there. +#[test] +fn later_refine_strokes_override_earlier_erases() { + let c = [100.0, 75.0]; + let (img, _) = soft_frame(c, 10.0, [0.9, 0.25, 0.1], [0.1, 0.3, 0.9]); + let spot = [c[0] + R, c[1]]; + let i = spot[1] as usize * W as usize + spot[0] as usize; + let erase = Stroke { kind: StrokeKind::RefineErase, frame: 0, radius: 6.0, points: vec![spot] }; + let refine = Stroke { kind: StrokeKind::Refine, ..erase.clone() }; + let band = |st: &[Stroke]| segment(&img, &st.iter().collect::>(), None, &SegOpts::default(), 9.0).refine[i]; + let mut st = vec![ + line(StrokeKind::Fg, 0, [c[0] - 15.0, c[1]], [c[0] + 15.0, c[1]], 5.0), + line(StrokeKind::Bg, 0, [c[0] - 60.0, c[1] - 55.0], [c[0] + 60.0, c[1] - 55.0], 5.0), + ring(0, c, R, 9.0), + ]; + assert_eq!(band(&st), 1); + st.push(erase.clone()); + assert_eq!(band(&st), 0, "erase removes the band"); + st.push(refine); + assert_eq!(band(&st), 1, "refine over the erase restores it"); + st.push(erase); + assert_eq!(band(&st), 0, "erase again removes it again"); +} + #[test] fn reduce_chatter_and_motion_blur_use_the_neighbours() { let (w, h) = (40usize, 20usize); From 818ce2ab427be38010308488ee838f5eb4ae05a8 Mon Sep 17 00:00:00 2001 From: tylerjenningsw Date: Sun, 18 Oct 2026 11:28:13 -0400 Subject: [PATCH 3/6] Roto Brush: model prompts follow the latest stroke (#1646) A trained model is no longer prompted with both labels at one spot: points along a stroke that a later stroke of the other kind painted over are dropped, and the propagated matte's points, box and mask take this frame's strokes over it. --- crates/track/src/roto.rs | 42 ++++++++++++++++++++++++++-------- crates/track/src/roto/tests.rs | 23 ++++++++++++++++++- 2 files changed, 54 insertions(+), 11 deletions(-) diff --git a/crates/track/src/roto.rs b/crates/track/src/roto.rs index 9e5d8d6e1..0072ad223 100644 --- a/crates/track/src/roto.rs +++ b/crates/track/src/roto.rs @@ -375,6 +375,12 @@ fn sweep(s: &Stroke, scale: f64, w: usize, h: usize, mut paint: impl FnMut(usize } } +/// Whether `s` covers the point `p` (matte pixels at `scale`). +fn covers(s: &Stroke, scale: f64, p: [f64; 2]) -> bool { + let (r, segs) = segments(s, scale); + segs.iter().any(|(a, b)| segment_distance(p, *a, *b) <= r) +} + /// Pixels covered by `strokes` (of the given kinds) at `scale`. pub fn rasterize(strokes: &[&Stroke], kinds: &[StrokeKind], scale: f64, w: usize, h: usize) -> Vec { let mut m = vec![false; w * h]; @@ -894,15 +900,21 @@ fn points_along(s: &Stroke, scale: f64, max: usize) -> Vec<[f32; 2]> { out } -/// A model's prompt from the strokes: points along the foreground and background strokes, and -/// the box of a [`StrokeKind::Box`] stroke (matte pixels). +/// A model's prompt from the strokes: points along the foreground and background strokes (but +/// not where a later stroke of the other kind painted over them), and the box of a +/// [`StrokeKind::Box`] stroke (matte pixels). pub fn stroke_prompt(strokes: &[&Stroke], scale: f64) -> Prompt { + let seg: Vec<&Stroke> = strokes.iter().copied().filter(|s| s.kind.is_segmentation()).collect(); let mut points = vec![]; for (kind, fg) in [(StrokeKind::Fg, true), (StrokeKind::Bg, false)] { - let ss: Vec<&&Stroke> = strokes.iter().filter(|s| s.kind == kind).collect(); + let ss: Vec<(usize, &Stroke)> = seg.iter().copied().enumerate().filter(|(_, s)| s.kind == kind).collect(); let per = (16 / ss.len().max(1)).max(1); - for s in ss { - points.extend(points_along(s, scale, per).into_iter().map(|p| (p, fg))); + for (k, s) in ss { + let kept = |p: &[f32; 2]| { + let p = [p[0] as f64, p[1] as f64]; + seg.iter().skip(k + 1).rev().find(|t| covers(t, scale, p)).is_none_or(|t| t.kind == kind) + }; + points.extend(points_along(s, scale, per).into_iter().filter(kept).map(|p| (p, fg))); } } let bbox = strokes.iter().filter(|s| s.kind == StrokeKind::Box).find_map(|s| box_of(s, scale)); @@ -918,8 +930,18 @@ fn box_of(s: &Stroke, scale: f64) -> Option<[f32; 4]> { } /// Add a prior matte to a prompt: its box (a little larger; unless the prompt has one), up to -/// four points spread through its core, and the matte itself. -fn add_prior_prompt(prompt: &mut Prompt, soft: &[f32], w: usize, h: usize) -> bool { +/// four points spread through its core, and the matte itself, all with this frame's strokes +/// (`hard`, from [`resolve`]) painted over it. +fn add_prior_prompt(prompt: &mut Prompt, soft: &[f32], hard: &[u8], w: usize, h: usize) -> bool { + let soft: Vec = soft + .iter() + .enumerate() + .map(|(i, v)| match hard.get(i) { + Some(&HARD_FG) => 1.0, + Some(&HARD_BG) => 0.0, + _ => *v, + }) + .collect(); let bin: Vec = soft.iter().map(|v| (*v >= 0.5) as u8).collect(); let (mut x0, mut y0, mut x1, mut y1) = (usize::MAX, usize::MAX, 0usize, 0usize); for (i, v) in bin.iter().enumerate() { @@ -949,7 +971,7 @@ fn add_prior_prompt(prompt: &mut Prompt, soft: &[f32], w: usize, h: usize) -> bo } } prompt.points.extend(picks.into_iter().map(|i| ([(i % w) as f32 + 0.5, (i / w) as f32 + 0.5], true))); - prompt.mask = Some(soft.to_vec()); + prompt.mask = Some(soft); true } @@ -960,7 +982,7 @@ pub fn segment_with(img: &Image, strokes: &[&Stroke], prior: Option<&Prior>, opt let Some(m) = model.filter(|_| w > 0 && h > 0) else { return segment(img, strokes, prior, opts, refine_radius) }; let mut prompt = stroke_prompt(strokes, opts.scale); if let Some(p) = prior { - add_prior_prompt(&mut prompt, &p.soft, w, h); + add_prior_prompt(&mut prompt, &p.soft, &resolve(strokes, StrokeKind::Fg, StrokeKind::Bg, opts.scale, w, h), w, h); } if !prompt.points.iter().any(|(_, fg)| *fg) { return segment(img, strokes, prior, opts, refine_radius); @@ -991,7 +1013,7 @@ pub fn propagate_with( let Some(m) = model.filter(|_| w > 0 && h > 0) else { return propagate(prev_img, prev, next_img, strokes, opts, refine_radius) }; let prior = warp_prior(prev_img, prev, next_img, opts); let mut prompt = stroke_prompt(strokes, opts.scale); - if !add_prior_prompt(&mut prompt, &prior.soft, w, h) { + if !add_prior_prompt(&mut prompt, &prior.soft, &resolve(strokes, StrokeKind::Fg, StrokeKind::Bg, opts.scale, w, h), w, h) { return propagate(prev_img, prev, next_img, strokes, opts, refine_radius); } match m.segment(&rgb_of(next_img), w, h, &prompt) { diff --git a/crates/track/src/roto/tests.rs b/crates/track/src/roto/tests.rs index 1922dd0a2..f5170e754 100644 --- a/crates/track/src/roto/tests.rs +++ b/crates/track/src/roto/tests.rs @@ -509,6 +509,27 @@ fn box_stroke_is_the_model_box_prompt() { assert!(rasterize(&[&bx], &[StrokeKind::Fg, StrokeKind::Bg], 1.0, 100, 80).iter().all(|v| !v)); let mut q = stroke_prompt(&[&fg, &bx], 1.0); let soft: Vec = (0..100 * 80).map(|i| ((i % 100) < 10) as u8 as f32).collect(); - assert!(add_prior_prompt(&mut q, &soft, 100, 80)); + assert!(add_prior_prompt(&mut q, &soft, &[], 100, 80)); assert_eq!(q.bbox, Some([20.0, 10.0, 80.0, 60.0])); } + +/// A model is never prompted against the latest correction: stroke points that a later stroke +/// of the other kind painted over are dropped, and the propagated matte's points and mask leave +/// out what this frame's background strokes removed. +#[test] +fn model_prompts_follow_the_latest_stroke() { + let green = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 4.0, points: vec![[50.0, 40.0]] }; + let red = Stroke { kind: StrokeKind::Bg, ..green.clone() }; + assert_eq!(stroke_prompt(&[&green, &red], 1.0).points, vec![([50.0, 40.0], false)]); + assert_eq!(stroke_prompt(&[&green, &red, &green], 1.0).points, vec![([50.0, 40.0], true); 2]); + let (w, h) = (100, 80); + let soft: Vec = (0..w * h).map(|i| ((i % w) < 50) as u8 as f32).collect(); + let erase = Stroke { kind: StrokeKind::Bg, frame: 0, radius: 15.0, points: vec![[25.0, 40.0]] }; + let hard = resolve(&[&erase], StrokeKind::Fg, StrokeKind::Bg, 1.0, w, h); + let mut q = Prompt::default(); + assert!(add_prior_prompt(&mut q, &soft, &hard, w, h)); + assert!(!q.points.is_empty()); + assert!(q.points.iter().all(|(p, fg)| *fg && !covers(&erase, 1.0, [p[0] as f64, p[1] as f64])), "{:?}", q.points); + let mask = q.mask.unwrap(); + assert!(hard.iter().zip(&mask).all(|(hd, m)| *hd != HARD_BG || *m == 0.0)); +} From 0f8a65f6c2db145bd3099ad827f9b468cc257d4e Mon Sep 17 00:00:00 2001 From: tylerjenningsw Date: Sun, 18 Oct 2026 11:31:48 -0400 Subject: [PATCH 4/6] Disk cache: drop segmentations and frames made before strokes applied in order (#1646) Segmentations and rendered frames are keyed by their strokes, which now segment differently where strokes overlap; bumping both versions keeps old entries from showing the old result. Frozen mattes are stored in the project and are untouched. --- crates/render/src/disk_cache.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/render/src/disk_cache.rs b/crates/render/src/disk_cache.rs index 4fa2e1cb3..61026abbe 100644 --- a/crates/render/src/disk_cache.rs +++ b/crates/render/src/disk_cache.rs @@ -652,7 +652,7 @@ pub fn read_layer_entry(data: &[u8]) -> Option { /// The disk key of a segmentation: its chain key and scale. fn seg_key(key: u64, scale: f64) -> u128 { let mut h = Hash128::default(); - h.write(b"segmentation-v1"); + h.write(b"segmentation-v2"); h.write_u64(key); h.write_u64(scale.to_bits()); h.finish() @@ -772,7 +772,7 @@ fn decode_layer(p: &[u8]) -> Option { // ---------------------------------------------------------------- content keys /// Bump when rendering changes in a way that makes old disk entries wrong. -const RENDER_VERSION: &str = "effectcraft-render-2"; +const RENDER_VERSION: &str = "effectcraft-render-3"; fn hash_debug(h: &mut Hash128, v: &impl std::fmt::Debug) { h.write(format!("{v:?}").as_bytes()); From b2cace1320a15248743f3fde253792b143a3a35c Mon Sep 17 00:00:00 2001 From: tylerjenningsw Date: Sun, 18 Oct 2026 11:42:01 -0400 Subject: [PATCH 5/6] Roto Brush: test that a stroke corrects earlier ones through the commands (#1646) Green, Alt (red) over part of the object, then green there again, as the Layer panel sends them through roto.stroke: the render follows the last stroke, and Undo / Redo step between the two results. --- crates/engine/src/tests/roto.rs | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/crates/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index 423af1d2f..64ee30d5e 100644 --- a/crates/engine/src/tests/roto.rs +++ b/crates/engine/src/tests/roto.rs @@ -164,6 +164,33 @@ fn alpha_at(s: &Session, f: u32) -> Vec { img.data.iter().map(|p| p[3]).collect() } +/// Green, then Alt (red) over part of the object, then green there again: the last stroke wins, +/// in the render and through Undo / Redo (#1646). +#[test] +fn roto_brush_strokes_correct_earlier_ones() { + let (mut s, clip, _) = painted(center, "disk-1646.mov"); + let c = center(0); + let spot = [c[0], c[1] + 25.0]; + let i = spot[1] as usize * W as usize + spot[0] as usize; + let paint = |s: &mut Session, kind: &str| { + s.execute_checked( + "roto.stroke", + json!({"layer": clip.0, "kind": kind, "points": line([spot[0] - 6.0, spot[1]], [spot[0] + 6.0, spot[1]]), "radius": 7}), + ) + .unwrap(); + }; + assert!(alpha_at(&s, 0)[i] > 0.5, "the spot starts inside the matte"); + paint(&mut s, "bg"); + assert!(alpha_at(&s, 0)[i] < 0.5, "Alt removes it"); + paint(&mut s, "fg"); + assert!(alpha_at(&s, 0)[i] > 0.5, "green over the Alt stroke restores it"); + assert!(iou_at(&mut s, clip, 0, c) > 0.95); + assert!(s.undo()); + assert!(alpha_at(&s, 0)[i] < 0.5, "undo brings the removal back"); + assert!(s.redo()); + assert!(alpha_at(&s, 0)[i] > 0.5, "redo restores it again"); +} + #[test] fn roto_brush_segments_and_propagates() { let (mut s, clip, uid) = painted(center, "disk-a.mov"); From b5460f653b8dd0b49d51dc1c568502e7a4c1d318 Mon Sep 17 00:00:00 2001 From: tylerjenningsw Date: Sun, 18 Oct 2026 12:15:11 -0400 Subject: [PATCH 6/6] Roto Brush: a correction keeps the model prompted with what is left of a stroke (#1646) Dropping the points a later stroke of the other kind painted over could leave a stroke with none, even though most of it was still painted: a small Alt tap in the middle of a large foreground dot removed its only point, and the selected model was skipped for the classical cut. Points are now kept by the resolved ownership under them, and a stroke that loses every sampled point but still owns pixels is prompted from the one deepest inside them. A stroke painted over completely still prompts nothing. --- crates/track/src/roto.rs | 81 ++++++++++++++++++++++++---------- crates/track/src/roto/tests.rs | 64 +++++++++++++++++++++++---- 2 files changed, 113 insertions(+), 32 deletions(-) diff --git a/crates/track/src/roto.rs b/crates/track/src/roto.rs index 0072ad223..998f5100d 100644 --- a/crates/track/src/roto.rs +++ b/crates/track/src/roto.rs @@ -375,12 +375,6 @@ fn sweep(s: &Stroke, scale: f64, w: usize, h: usize, mut paint: impl FnMut(usize } } -/// Whether `s` covers the point `p` (matte pixels at `scale`). -fn covers(s: &Stroke, scale: f64, p: [f64; 2]) -> bool { - let (r, segs) = segments(s, scale); - segs.iter().any(|(a, b)| segment_distance(p, *a, *b) <= r) -} - /// Pixels covered by `strokes` (of the given kinds) at `scale`. pub fn rasterize(strokes: &[&Stroke], kinds: &[StrokeKind], scale: f64, w: usize, h: usize) -> Vec { let mut m = vec![false; w * h]; @@ -900,27 +894,66 @@ fn points_along(s: &Stroke, scale: f64, max: usize) -> Vec<[f32; 2]> { out } -/// A model's prompt from the strokes: points along the foreground and background strokes (but -/// not where a later stroke of the other kind painted over them), and the box of a -/// [`StrokeKind::Box`] stroke (matte pixels). -pub fn stroke_prompt(strokes: &[&Stroke], scale: f64) -> Prompt { - let seg: Vec<&Stroke> = strokes.iter().copied().filter(|s| s.kind.is_segmentation()).collect(); +/// A model's prompt from the strokes: points along the foreground and background strokes, but +/// not where a later stroke of the other kind painted over them (`hard`, from [`resolve`]), and +/// the box of a [`StrokeKind::Box`] stroke (matte pixels). A stroke painted over at every point +/// sampled along it, but not everywhere, is prompted from the middle of what is left of it. +pub fn stroke_prompt(strokes: &[&Stroke], hard: &[u8], scale: f64, w: usize, h: usize) -> Prompt { let mut points = vec![]; - for (kind, fg) in [(StrokeKind::Fg, true), (StrokeKind::Bg, false)] { - let ss: Vec<(usize, &Stroke)> = seg.iter().copied().enumerate().filter(|(_, s)| s.kind == kind).collect(); + for (kind, own, other) in [(StrokeKind::Fg, HARD_FG, HARD_BG), (StrokeKind::Bg, HARD_BG, HARD_FG)] { + let ss: Vec<&&Stroke> = strokes.iter().filter(|s| s.kind == kind).collect(); let per = (16 / ss.len().max(1)).max(1); - for (k, s) in ss { - let kept = |p: &[f32; 2]| { - let p = [p[0] as f64, p[1] as f64]; - seg.iter().skip(k + 1).rev().find(|t| covers(t, scale, p)).is_none_or(|t| t.kind == kind) - }; - points.extend(points_along(s, scale, per).into_iter().filter(kept).map(|p| (p, fg))); + for s in ss { + let along = points_along(s, scale, per); + let mut kept: Vec<[f32; 2]> = along.iter().copied().filter(|p| owner(hard, w, h, *p) != Some(other)).collect(); + if kept.is_empty() && !along.is_empty() { + kept.extend(deepest_owned(s, hard, own, scale, w, h)); + } + points.extend(kept.into_iter().map(|p| (p, kind == StrokeKind::Fg))); } } let bbox = strokes.iter().filter(|s| s.kind == StrokeKind::Box).find_map(|s| box_of(s, scale)); Prompt { points, bbox, ..Default::default() } } +/// The value of `hard` under the point `p` (`None` outside the frame). +fn owner(hard: &[u8], w: usize, h: usize, p: [f32; 2]) -> Option { + if !(p[0] >= 0.0 && p[1] >= 0.0) { + return None; + } + let (x, y) = (p[0] as usize, p[1] as usize); + if x >= w || y >= h { + return None; + } + hard.get(y * w + x).copied() +} + +/// The centre of the pixel deepest inside what `s` still owns (where `hard` is `own` under it), +/// if anything. +fn deepest_owned(s: &Stroke, hard: &[u8], own: u8, scale: f64, w: usize, h: usize) -> Option<[f32; 2]> { + let mut owned = vec![]; + sweep(s, scale, w, h, |i| { + if hard.get(i) == Some(&own) { + owned.push((i % w, i / w)); + } + }); + let (x0, x1) = (owned.iter().map(|p| p.0).min()?, owned.iter().map(|p| p.0).max()?); + let (y0, y1) = (owned.iter().map(|p| p.1).min()?, owned.iter().map(|p| p.1).max()?); + // One pixel of margin all round, so depth is measured to the nearest pixel it doesn't own. + let (cw, ch) = (x1 - x0 + 3, y1 - y0 + 3); + let at = |(x, y): (usize, usize)| (y - y0 + 1) * cw + (x - x0 + 1); + let mut free = vec![true; cw * ch]; + for p in &owned { + if let Some(v) = free.get_mut(at(*p)) { + *v = false; + } + } + let depth = matting::distance(&free, cw, ch); + let d = |p: &(usize, usize)| depth.get(at(*p)).copied().unwrap_or(0.0); + let (x, y) = owned.iter().copied().max_by(|a, b| d(a).total_cmp(&d(b)))?; + Some([x as f32 + 0.5, y as f32 + 0.5]) +} + /// The box `[x0, y0, x1, y1]` (matte pixels) of a box stroke: its first two points as opposite /// corners; `None` when they are missing, not finite or the box is empty. fn box_of(s: &Stroke, scale: f64) -> Option<[f32; 4]> { @@ -980,9 +1013,10 @@ fn add_prior_prompt(prompt: &mut Prompt, soft: &[f32], hard: &[u8], w: usize, h: pub fn segment_with(img: &Image, strokes: &[&Stroke], prior: Option<&Prior>, opts: &SegOpts, refine_radius: f64, model: Option<&dyn MaskModel>) -> FrameSeg { let (w, h) = (img.width as usize, img.height as usize); let Some(m) = model.filter(|_| w > 0 && h > 0) else { return segment(img, strokes, prior, opts, refine_radius) }; - let mut prompt = stroke_prompt(strokes, opts.scale); + let hard = resolve(strokes, StrokeKind::Fg, StrokeKind::Bg, opts.scale, w, h); + let mut prompt = stroke_prompt(strokes, &hard, opts.scale, w, h); if let Some(p) = prior { - add_prior_prompt(&mut prompt, &p.soft, &resolve(strokes, StrokeKind::Fg, StrokeKind::Bg, opts.scale, w, h), w, h); + add_prior_prompt(&mut prompt, &p.soft, &hard, w, h); } if !prompt.points.iter().any(|(_, fg)| *fg) { return segment(img, strokes, prior, opts, refine_radius); @@ -1012,8 +1046,9 @@ pub fn propagate_with( let (w, h) = (next_img.width as usize, next_img.height as usize); let Some(m) = model.filter(|_| w > 0 && h > 0) else { return propagate(prev_img, prev, next_img, strokes, opts, refine_radius) }; let prior = warp_prior(prev_img, prev, next_img, opts); - let mut prompt = stroke_prompt(strokes, opts.scale); - if !add_prior_prompt(&mut prompt, &prior.soft, &resolve(strokes, StrokeKind::Fg, StrokeKind::Bg, opts.scale, w, h), w, h) { + let hard = resolve(strokes, StrokeKind::Fg, StrokeKind::Bg, opts.scale, w, h); + let mut prompt = stroke_prompt(strokes, &hard, opts.scale, w, h); + if !add_prior_prompt(&mut prompt, &prior.soft, &hard, w, h) { return propagate(prev_img, prev, next_img, strokes, opts, refine_radius); } match m.segment(&rgb_of(next_img), w, h, &prompt) { diff --git a/crates/track/src/roto/tests.rs b/crates/track/src/roto/tests.rs index f5170e754..f0059c651 100644 --- a/crates/track/src/roto/tests.rs +++ b/crates/track/src/roto/tests.rs @@ -420,9 +420,8 @@ impl effectcraft_segment::MaskModel for FakeModel { } /// A grey bar on a two-tone background whose halves match the bar's halves: colour alone can't -/// find its edge, a model can. -#[test] -fn a_model_prior_shapes_the_cut_and_bad_models_are_ignored() { +/// find its edge, a model can. The frame, its truth and the bar's rectangle. +fn bar_frame() -> (Image, Vec, [usize; 4]) { let (w, h) = (160usize, 100usize); let rect = [40, 30, 120, 70]; let mut img = Image::new(w as u32, h as u32); @@ -437,6 +436,12 @@ fn a_model_prior_shapes_the_cut_and_bad_models_are_ignored() { gt[y * w + x] = inside as u8; } } + (img, gt, rect) +} + +#[test] +fn a_model_prior_shapes_the_cut_and_bad_models_are_ignored() { + let (img, gt, rect) = bar_frame(); let st = [line(StrokeKind::Fg, 0, [50.0, 50.0], [110.0, 50.0], 4.0)]; let refs: Vec<&Stroke> = st.iter().collect(); let opts = SegOpts::default(); @@ -497,17 +502,17 @@ fn mobilesam_segments_and_propagates_the_disk() { fn box_stroke_is_the_model_box_prompt() { let fg = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 2.0, points: vec![[50.0, 40.0]] }; let bx = Stroke { kind: StrokeKind::Box, frame: 0, radius: 0.0, points: vec![[80.0, 60.0], [20.0, 10.0]] }; - let p = stroke_prompt(&[&fg, &bx], 0.5); + let p = stroke_prompt(&[&fg, &bx], &[], 0.5, 0, 0); assert_eq!(p.bbox, Some([10.0, 5.0, 40.0, 30.0])); assert_eq!(p.points, vec![([25.0, 20.0], true)]); // Degenerate or non-finite boxes are no prompt. for pts in [vec![[20.0, 10.0]], vec![[20.0, 10.0], [20.0, 50.0]], vec![[f64::NAN, 10.0], [80.0, 60.0]]] { let b = Stroke { points: pts, ..bx.clone() }; - assert_eq!(stroke_prompt(&[&fg, &b], 1.0).bbox, None); + assert_eq!(stroke_prompt(&[&fg, &b], &[], 1.0, 0, 0).bbox, None); } assert!(!StrokeKind::Box.is_segmentation()); assert!(rasterize(&[&bx], &[StrokeKind::Fg, StrokeKind::Bg], 1.0, 100, 80).iter().all(|v| !v)); - let mut q = stroke_prompt(&[&fg, &bx], 1.0); + let mut q = stroke_prompt(&[&fg, &bx], &[], 1.0, 0, 0); let soft: Vec = (0..100 * 80).map(|i| ((i % 100) < 10) as u8 as f32).collect(); assert!(add_prior_prompt(&mut q, &soft, &[], 100, 80)); assert_eq!(q.bbox, Some([20.0, 10.0, 80.0, 60.0])); @@ -520,16 +525,57 @@ fn box_stroke_is_the_model_box_prompt() { fn model_prompts_follow_the_latest_stroke() { let green = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 4.0, points: vec![[50.0, 40.0]] }; let red = Stroke { kind: StrokeKind::Bg, ..green.clone() }; - assert_eq!(stroke_prompt(&[&green, &red], 1.0).points, vec![([50.0, 40.0], false)]); - assert_eq!(stroke_prompt(&[&green, &red, &green], 1.0).points, vec![([50.0, 40.0], true); 2]); let (w, h) = (100, 80); + let prompt = |st: &[&Stroke]| stroke_prompt(st, &resolve(st, StrokeKind::Fg, StrokeKind::Bg, 1.0, w, h), 1.0, w, h).points; + assert_eq!(prompt(&[&green, &red]), vec![([50.0, 40.0], false)]); + assert_eq!(prompt(&[&green, &red, &green]), vec![([50.0, 40.0], true); 2]); let soft: Vec = (0..w * h).map(|i| ((i % w) < 50) as u8 as f32).collect(); let erase = Stroke { kind: StrokeKind::Bg, frame: 0, radius: 15.0, points: vec![[25.0, 40.0]] }; let hard = resolve(&[&erase], StrokeKind::Fg, StrokeKind::Bg, 1.0, w, h); let mut q = Prompt::default(); assert!(add_prior_prompt(&mut q, &soft, &hard, w, h)); assert!(!q.points.is_empty()); - assert!(q.points.iter().all(|(p, fg)| *fg && !covers(&erase, 1.0, [p[0] as f64, p[1] as f64])), "{:?}", q.points); + assert!(q.points.iter().all(|(p, fg)| *fg && owner(&hard, w, h, *p) != Some(HARD_BG)), "{:?}", q.points); let mask = q.mask.unwrap(); assert!(hard.iter().zip(&mask).all(|(hd, m)| *hd != HARD_BG || *m == 0.0)); } + +/// A stand-in model that counts its calls; its foreground is a fixed rectangle. +struct CountingModel([usize; 4], std::sync::atomic::AtomicUsize); + +impl effectcraft_segment::MaskModel for CountingModel { + fn info(&self) -> &'static effectcraft_segment::ModelInfo { + &effectcraft_segment::MOBILE_SAM + } + fn segment(&self, _: &[[f32; 3]], w: usize, h: usize, _: &Prompt) -> Result, String> { + self.1.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let [x0, y0, x1, y1] = self.0; + Ok((0..w * h).map(|i| ((i % w) >= x0 && (i % w) < x1 && (i / w) >= y0 && (i / w) < y1) as u8 as f32 * 0.98 + 0.01).collect()) + } +} + +/// A small background correction in the middle of a large foreground dot removes the only point +/// sampled along the dot, but most of the dot is still foreground: the model is still prompted +/// with it, from a point inside what is left. A dot painted over completely prompts nothing. +#[test] +fn a_correction_keeps_the_model_prompted_with_what_is_left() { + let (img, gt, rect) = bar_frame(); + let (w, h) = (img.width as usize, img.height as usize); + let dot = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 15.0, points: vec![[80.0, 50.0]] }; + let tap = Stroke { kind: StrokeKind::Bg, radius: 2.0, ..dot.clone() }; + let refs = [&dot, &tap]; + let model = CountingModel(rect, Default::default()); + let seg = segment_with(&img, &refs, None, &SegOpts::default(), 0.0, Some(&model)); + assert_eq!(model.1.load(std::sync::atomic::Ordering::Relaxed), 1, "the model was prompted"); + assert!(iou(&seg.matte, >) > 0.95, "IoU {}", iou(&seg.matte, >)); + let hard = resolve(&refs, StrokeKind::Fg, StrokeKind::Bg, 1.0, w, h); + let fg: Vec<[f32; 2]> = stroke_prompt(&refs, &hard, 1.0, w, h).points.into_iter().filter(|p| p.1).map(|p| p.0).collect(); + let [q] = fg.as_slice() else { panic!("one foreground point, got {fg:?}") }; + let r = (q[0] - 80.0).hypot(q[1] - 50.0); + assert!(r > 5.0 && r < 13.0, "{q:?} sits inside what is left of the dot"); + assert_eq!(hard[q[1] as usize * w + q[0] as usize], HARD_FG); + let cover = Stroke { radius: 16.0, ..tap }; + let refs = [&dot, &cover]; + let hard = resolve(&refs, StrokeKind::Fg, StrokeKind::Bg, 1.0, w, h); + assert!(stroke_prompt(&refs, &hard, 1.0, w, h).points.iter().all(|p| !p.1)); +}