aleleba/effectcraft
1
0
Fork 0
mirror of https://github.com/storytold/effectcraft.git synced 2026-10-29 20:47:18 -06:00
Code Issues Packages Projects Releases Wiki Activity

Roto Brush: a correction keeps the model prompted with what is left of a stroke (#1646)

Browse Source 
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.
This commit is contained in:
tylerjenningsw
2026-10-27 14:13:40 -04:00
parent b2cace1320
commit b5460f653b
2 changed files with 113 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/track/src/roto.rs
+58 -23
View File
@@ -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<bool> {
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<u8> {
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) {
crates/track/src/roto/tests.rs
+55 -9
View File
@@ -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<u8>, [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<f32> = (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<f32> = (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<Vec<f32>, 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, &gt) > 0.95, "IoU {}", iou(&seg.matte, &gt));
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));
}