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

Merge pull request #1804 from tylerjenningsw/fix/roto-stroke-order

Browse Source 
Roto Brush: apply correction strokes in drawing order (#1646)
This commit is contained in:
BFlat
2026-10-30 01:21:32 +07:00
committed by GitHub
parent 8c7c8bee8c b5460f653b
commit 7175a20f63
4 changed files with 286 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/engine/src/tests/roto.rs
+27
View File
@@ -164,6 +164,33 @@ fn alpha_at(s: &Session, f: u32) -> Vec<f32> {
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");
crates/render/src/disk_cache.rs
+1 -1
View File
@@ -721,7 +721,7 @@ pub fn read_layer_entry(data: &[u8]) -> Option<Buf> {
/// 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()
crates/track/src/roto.rs
+137 -55
View File
@@ -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<bool> {
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<bool> {
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<u8> {
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;
@@ -662,21 +700,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<u8> = 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<bool> = hard.iter().map(|v| *v == HARD_FG).collect();
let bgs: Vec<bool> = 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<Vec<f32>> =
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)) });
@@ -788,8 +814,7 @@ fn core_of(soft: &[f32], w: usize, h: usize, r: f64) -> Vec<bool> {
fn finish(_img: &Image, rgb: &[[f32; 3]], labels: Vec<u8>, 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<Vec<f32>> = 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<u8> = (0..w * h)
@@ -798,7 +823,7 @@ fn finish(_img: &Image, rgb: &[[f32; 3]], labels: Vec<u8>, 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 }
@@ -877,21 +902,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, and
/// the box of a [`StrokeKind::Box`] stroke (matte pixels).
pub fn stroke_prompt(strokes: &[&Stroke], scale: f64) -> Prompt {
/// 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)] {
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 s in ss {
points.extend(points_along(s, scale, per).into_iter().map(|p| (p, fg)));
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]> {
@@ -901,8 +971,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<f32> = 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<u8> = 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() {
@@ -932,7 +1012,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
}
@@ -941,9 +1021,10 @@ fn add_prior_prompt(prompt: &mut Prompt, soft: &[f32], w: usize, h: usize) -> bo
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, 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);
@@ -973,8 +1054,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, 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
+121 -7
View File
@@ -182,6 +182,28 @@ fn a_correction_stroke_fixes_a_bad_frame() {
assert!(iou(&next.matte, &gt2) > 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::<Vec<_>>(), 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, &gt) > 0.95, "IoU {}", iou(&restored.matte, &gt));
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<f32>) {
let mut img = Image::new(W, H);
@@ -265,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::<Vec<_>>(), 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);
@@ -373,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);
@@ -390,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();
@@ -450,18 +502,80 @@ 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!(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() };
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 && 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));
}