diff --git a/crates/effects/src/migrate.rs b/crates/effects/src/migrate.rs index 35d0960e0..ecea0fe19 100644 --- a/crates/effects/src/migrate.rs +++ b/crates/effects/src/migrate.rs @@ -765,7 +765,13 @@ not a stroke", let data = effectcraft_track::roto::RotoData { base: Some(0), span: [0, 10], - strokes: vec![effectcraft_track::roto::Stroke { kind: effectcraft_track::roto::StrokeKind::Fg, frame: 0, radius: 5.0, points: vec![[2.0, 2.0]] }], + strokes: vec![effectcraft_track::roto::Stroke { + kind: effectcraft_track::roto::StrokeKind::Fg, + frame: 0, + radius: 5.0, + points: vec![[2.0, 2.0]], + object: 0, + }], }; roto.get_mut(crate::roto::STROKES).unwrap().value = Value::Str(data.to_json()); to_effect_space(spec, &mut roto, SIZE, D, &own); diff --git a/crates/effects/src/object_matte.rs b/crates/effects/src/object_matte.rs index 4b5e1fd42..3a75bec1b 100644 --- a/crates/effects/src/object_matte.rs +++ b/crates/effects/src/object_matte.rs @@ -144,7 +144,7 @@ pub fn prompt_strokes(params: &Params, layer_size: [f64; 2], frame: i64, with_mo if !inside(pt) { return None; } - let mut out = vec![Stroke { kind: StrokeKind::Fg, frame, radius: POINT_RADIUS, points: vec![pt] }]; + let mut out = vec![Stroke { kind: StrokeKind::Fg, frame, radius: POINT_RADIUS, points: vec![pt], object: 0 }]; let rect = |x0: f64, y0: f64, x1: f64, y1: f64| vec![[x0, y0], [x1, y0], [x1, y1], [x0, y1], [x0, y0]]; let boxed = b(params, "useBox").then(|| Some((v2(params, "boxTopLeft")?, v2(params, "boxBottomRight")?))).flatten(); match boxed { @@ -158,9 +158,9 @@ pub fn prompt_strokes(params: &Params, layer_size: [f64; 2], frame: i64, with_mo } out.push(if with_model { // The model takes the box itself as its box prompt. - Stroke { kind: StrokeKind::Box, frame, radius: 0.0, points: vec![[x0, y0], [x1, y1]] } + Stroke { kind: StrokeKind::Box, frame, radius: 0.0, points: vec![[x0, y0], [x1, y1]], object: 0 } } else { - Stroke { kind: StrokeKind::Bg, frame, radius: EDGE_RADIUS, points: rect(x0, y0, x1, y1) } + Stroke { kind: StrokeKind::Bg, frame, radius: EDGE_RADIUS, points: rect(x0, y0, x1, y1), object: 0 } }); } None if !with_model => { @@ -168,7 +168,7 @@ pub fn prompt_strokes(params: &Params, layer_size: [f64; 2], frame: i64, with_mo if w <= 4.0 * i || h <= 4.0 * i { return None; } - out.push(Stroke { kind: StrokeKind::Bg, frame, radius: EDGE_RADIUS, points: rect(i, i, w - i, h - i) }); + out.push(Stroke { kind: StrokeKind::Bg, frame, radius: EDGE_RADIUS, points: rect(i, i, w - i, h - i), object: 0 }); } None => {} } diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index c31c4836b..ffa41f01e 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -564,32 +564,72 @@ impl Chain { let (Some(&need_lo), Some(&need_hi)) = (needed.first(), needed.last()) else { return Ok(Some(TrackSummary::default())); }; - // A prompt is a few clicks from the strokes and a box around what the classic engine makes of - // them on that frame (clicks alone pick a part of the object, not all of it). - let mut all_prompts: Vec<(i64, effectcraft_segment::Prompt)> = vec![]; - for f in self.data.stroke_frames().into_iter().filter(|f| self.data.in_span(*f)) { - let strokes = self.data.strokes_at(f); - let p = match frame(f) { - Some(img) => rb::tracker_prompt_boxed(&strokes, self.opts.scale, &rb::segment(&img, &strokes, None, &self.opts, self.refine_radius)), - None => rb::tracker_prompt(&strokes, self.opts.scale), - }; - if p.points.iter().any(|(_, fg)| *fg) { - all_prompts.push((f, p)); + // Each object has its own prompts: a few clicks from its strokes and a box around what the classic engine + // makes of them on that frame (clicks alone pick a part of the object, not all of it). + let objects = self.data.objects(); + let mut by_object: Vec> = vec![]; + for &obj in &objects { + let mut prompts = vec![]; + for f in self.data.stroke_frames().into_iter().filter(|f| self.data.in_span(*f)) { + let strokes = self.data.strokes_of(f, obj); + let p = match frame(f) { + Some(img) => rb::tracker_prompt_boxed(&strokes, self.opts.scale, &rb::segment(&img, &strokes, None, &self.opts, self.refine_radius)), + None => rb::tracker_prompt(&strokes, self.opts.scale), + }; + if p.points.iter().any(|(_, fg)| *fg) { + prompts.push((f, p)); + } } + by_object.push(prompts); } - if all_prompts.is_empty() { + // objects without a foreground prompt cannot be followed + let followed: Vec<(u32, Vec<(i64, effectcraft_segment::Prompt)>)> = objects.into_iter().zip(by_object).filter(|(_, p)| !p.is_empty()).collect(); + if followed.is_empty() { return Err("no foreground strokes to follow".into()); } - // The window to track: from the nearest prompted frame at or before the first frame needed - // (a correction restarts the object there) to the last frame needed, reaching the first - // prompted frame when everything needed lies before it. - let first_prompt = all_prompts.first().map_or(need_lo, |p| p.0); - let lo = all_prompts.iter().map(|p| p.0).filter(|f| *f <= need_lo).max().unwrap_or(need_lo); - let hi = if need_hi < first_prompt { first_prompt } else { need_hi }; - let prompts: Vec<(i64, effectcraft_segment::Prompt)> = all_prompts.into_iter().filter(|p| p.0 >= lo && p.0 <= hi).collect(); - let req = TrackRequest::single(lo, (hi - lo + 1) as usize, prompts); + // The window to track: every object restarts at the nearest of its prompted frames at or before the first + // frame needed (a correction restarts the object there), so the window starts at the earliest restart; it + // ends at the last frame needed, reaching the farthest first prompt when an object's needed frames lie + // before its prompt. + let lo = followed.iter().map(|(_, p)| p.iter().map(|q| q.0).filter(|f| *f <= need_lo).max().unwrap_or(need_lo)).min().unwrap_or(need_lo); + let first_prompt = followed.iter().filter_map(|(_, p)| p.first().map(|q| q.0)).max().unwrap_or(need_hi); + let hi = need_hi.max(first_prompt); + let objects: Vec = followed.iter().map(|(o, _)| *o).collect(); + let prompts: Vec> = + followed.into_iter().map(|(_, p)| p.into_iter().filter(|q| q.0 >= lo && q.0 <= hi).collect()).collect(); + let req = TrackRequest { first: lo, count: (hi - lo + 1) as usize, objects: prompts }; let frame = std::cell::RefCell::new(frame); let mut summary = TrackSummary::default(); + // A frame is finished when every object has reported it (the objects' mattes are united). + let mut pending: BTreeMap>> = BTreeMap::new(); + let finish = |g: i64, parts: Vec>, summary: &mut TrackSummary, progress: &mut dyn FnMut(usize)| { + let (Some(key), Some(img)) = (self.keys.get(&g).copied(), (frame.borrow_mut())(g)) else { return }; + let mut segs = vec![]; + let (mut confidence, mut present) = (f32::INFINITY, false); + let mut worst = f32::INFINITY; + for (obj, t) in objects.iter().zip(parts) { + let Some(t) = t else { continue }; + let strokes = self.data.strokes_of(g, *obj); + segs.push(rb::segment_tracked(&img, &strokes, t.prob, t.occluded, &self.opts, self.refine_radius)); + worst = worst.min(t.confidence); + if !t.occluded { + present = true; + confidence = confidence.min(t.confidence); + } + } + if segs.is_empty() { + return; + } + // a frame is occluded when every object is absent; otherwise the doubtful object sets the confidence + let (occluded, confidence) = (!present, if present { confidence } else { worst }); + store(tracked_key(key), self.opts.scale, Arc::new(effectcraft_track::roto::FrameSeg::union(segs))); + trust_cache().lock().unwrap_or_else(|e| e.into_inner()).insert(tracked_key(key), FrameTrust { confidence, occluded }); + summary.frames += 1; + summary.occluded += usize::from(occluded); + summary.min_confidence = if summary.frames == 1 { confidence } else { summary.min_confidence.min(confidence) }; + progress(summary.frames); + }; + let n_objects = objects.len(); tracker.track( &req, &mut |g| { @@ -597,18 +637,21 @@ impl Chain { Some(Arc::new(ClipFrame { rgb: rb::rgb_of(&img), w: img.width as usize, h: img.height as usize })) }, cancel, - &mut |_, g, t| { - let (Some(key), Some(img)) = (self.keys.get(&g).copied(), (frame.borrow_mut())(g)) else { return }; - let strokes = self.data.strokes_at(g); - let seg = rb::segment_tracked(&img, &strokes, t.prob, t.occluded, &self.opts, self.refine_radius); - store(tracked_key(key), self.opts.scale, Arc::new(seg)); - trust_cache().lock().unwrap_or_else(|e| e.into_inner()).insert(tracked_key(key), FrameTrust { confidence: t.confidence, occluded: t.occluded }); - summary.frames += 1; - summary.occluded += usize::from(t.occluded); - summary.min_confidence = if summary.frames == 1 { t.confidence } else { summary.min_confidence.min(t.confidence) }; - progress(summary.frames); + &mut |o, g, t| { + let parts = pending.entry(g).or_insert_with(|| vec![None; n_objects]); + if let Some(slot) = parts.get_mut(o) { + *slot = Some(t); + } + if parts.iter().all(Option::is_some) { + let parts = pending.remove(&g).unwrap_or_default(); + finish(g, parts, &mut summary, progress); + } }, )?; + // frames some object never reported (it was not followed there): united from the rest + for (g, parts) in std::mem::take(&mut pending) { + finish(g, parts, &mut summary, progress); + } Ok(Some(summary)) } } @@ -1038,8 +1081,8 @@ mod tests { #[test] fn renders_the_segmented_disk_on_demand() { let mut d = RotoData::default(); - d.add(Stroke { kind: StrokeKind::Fg, frame: 0, radius: 3.0, points: vec![[35.0, 40.0], [45.0, 40.0]] }, 20, [0, 24]); - d.add(Stroke { kind: StrokeKind::Bg, frame: 0, radius: 3.0, points: vec![[5.0, 5.0], [95.0, 5.0], [95.0, 75.0]] }, 20, [0, 24]); + d.add(Stroke { kind: StrokeKind::Fg, frame: 0, radius: 3.0, points: vec![[35.0, 40.0], [45.0, 40.0]], object: 0 }, 20, [0, 24]); + d.add(Stroke { kind: StrokeKind::Bg, frame: 0, radius: 3.0, points: vec![[5.0, 5.0], [95.0, 5.0], [95.0, 75.0]], object: 0 }, 20, [0, 24]); let params = params_with(&d); let host = Seq; let run = |t: f64| { @@ -1083,6 +1126,91 @@ mod tests { assert!(o.data.iter().all(|p| p[3] == 0.0)); } + /// A stand-in tracker for two objects: object 0 is the moving disk, object 1 a fixed square. It records how many + /// objects (and prompts per object) it was asked for. + struct TwoObjects(std::sync::Mutex>>); + impl effectcraft_segment::MaskModel for TwoObjects { + fn info(&self) -> &'static effectcraft_segment::ModelInfo { + &effectcraft_segment::SAM2_TINY + } + fn segment(&self, _: &[[f32; 3]], _: usize, _: usize, _: &effectcraft_segment::Prompt) -> effectcraft_segment::Result> { + Err("no weights".into()) + } + fn tracker(&self) -> Option<&dyn effectcraft_segment::ObjectTracker> { + Some(self) + } + } + impl effectcraft_segment::ObjectTracker for TwoObjects { + fn track( + &self, + req: &effectcraft_segment::TrackRequest, + frame: &mut dyn FnMut(i64) -> Option>, + _: &dyn Fn() -> bool, + out: &mut dyn FnMut(usize, i64, effectcraft_segment::TrackedFrame), + ) -> effectcraft_segment::Result<()> { + self.0.lock().unwrap().push(req.objects.iter().map(Vec::len).collect()); + for g in req.first..req.first + req.count as i64 { + let img = frame(g).ok_or("no frame")?; + let c = [40.0 + 2.0 * g as f64, 40.0]; + // the objects report in reverse order on odd frames: the union must not depend on it + let order: Vec = if g % 2 == 1 { (0..req.objects.len()).rev().collect() } else { (0..req.objects.len()).collect() }; + for o in order { + let prob = (0..img.w * img.h) + .map(|i| { + let (x, y) = ((i % img.w) as f64 + 0.5, (i / img.w) as f64 + 0.5); + let inside = if o == 0 { (x - c[0]).hypot(y - c[1]) <= 15.0 } else { (8.0..22.0).contains(&x) && (8.0..22.0).contains(&y) }; + f32::from(inside) + }) + .collect(); + out(o, g, effectcraft_segment::TrackedFrame { prob, w: img.w, h: img.h, confidence: if o == 0 { 0.9 } else { 0.4 }, occluded: false }); + } + } + Ok(()) + } + } + + #[test] + fn objects_are_followed_in_one_pass_and_their_mattes_united() { + let mut d = RotoData::default(); + let stroke = |kind, points: Vec<[f64; 2]>, object| Stroke { kind, frame: 0, radius: 3.0, points, object }; + d.add(stroke(StrokeKind::Fg, vec![[35.0, 40.0], [45.0, 40.0]], 0), 8, [0, 24]); + d.add(stroke(StrokeKind::Fg, vec![[12.0, 15.0], [18.0, 15.0]], 1), 8, [0, 24]); + d.add(stroke(StrokeKind::Bg, vec![[85.0, 70.0], [95.0, 70.0]], 0), 8, [0, 24]); + assert_eq!(d.objects(), vec![0, 1]); + // the Refine Edge strokes belong to every object; the others only to their own + assert_eq!(d.strokes_of(0, 1).len(), 1); + assert_eq!(d.strokes_of(0, 0).len(), 2); + let params = params_with(&d); + let mut chain = Chain::new(¶ms, [100.0, 80.0], 25.0, 1.0); + let fake = Arc::new(TwoObjects(Default::default())); + chain.model = Some(fake.clone()); + // the clip: the moving disk and a fixed yellow square + let mut get = |f: i64| { + let mut img = frame(f as f64 / 25.0 + 0.001); + for y in 8..22 { + for x in 8..22 { + img.set(x, y, [0.95, 0.9, 0.1 + ((x * 3 + y * 5) % 7) as f32 / 70.0, 1.0]); + } + } + Some(Arc::new(img)) + }; + let mut seen = 0; + let summary = chain.track(&mut get, &|| false, &mut |n| seen = n).unwrap().unwrap(); + // one request for both objects; object 0 has its foreground and background prompt on the frame, object 1 its own + assert_eq!(fake.0.lock().unwrap().len(), 1); + assert_eq!(fake.0.lock().unwrap()[0], vec![1, 1]); + assert_eq!(summary.frames, 9); + assert_eq!(seen, 9); + // the least sure object sets the frame's confidence + assert!((summary.min_confidence - 0.4).abs() < 1e-6, "{}", summary.min_confidence); + assert_eq!(summary.occluded, 0); + let seg = cached(tracked_key(chain.keys[&3]), 1.0).expect("frame 3 is tracked"); + let at = |x: usize, y: usize| seg.matte[y * 100 + x]; + assert!(at(46, 40) != 0, "the moving disk (frame 3 centre x = 46)"); + assert!(at(15, 15) != 0, "the fixed square"); + assert_eq!(at(60, 70), 0, "nothing elsewhere"); + } + /// Instances saved before the propagation controls got their own twirl-down (and motion /// blur / decontamination their nested ones) load with their values in the new places. #[test] diff --git a/crates/engine/src/analysis/roto.rs b/crates/engine/src/analysis/roto.rs index 9cc950e07..626de3030 100644 --- a/crates/engine/src/analysis/roto.rs +++ b/crates/engine/src/analysis/roto.rs @@ -51,6 +51,9 @@ pub struct RotoOptions { pub boundary_color: [f64; 3], /// Start propagating in the background after each stroke. pub auto_propagate: bool, + /// The object new strokes belong to (0 is the first). With a tracker every object is followed on its own + /// and their mattes are united; the classic engine treats all foreground strokes as one object. + pub object: u32, } impl Default for RotoOptions { @@ -63,10 +66,14 @@ impl Default for RotoOptions { overlay_opacity: 50.0, boundary_color: [1.0, 0.25, 0.75], auto_propagate: true, + object: 0, } } } +/// The highest object id (SAM 2 follows up to 32 objects in one pass). +pub const MAX_OBJECT_ID: u32 = 31; + pub const VIEWS: [&str; 4] = ["alphaBoundary", "alpha", "alphaOverlay", "none"]; /// What a job does. diff --git a/crates/engine/src/cmd/roto_cmds.rs b/crates/engine/src/cmd/roto_cmds.rs index 5849fa2c8..5124d413d 100644 --- a/crates/engine/src/cmd/roto_cmds.rs +++ b/crates/engine/src/cmd/roto_cmds.rs @@ -27,7 +27,7 @@ pub fn specs() -> Vec { "Roto Brush Stroke", [], None, - "{layer?, kind?: fg|bg|refine|refineErase, points: [[x, y], …] (layer pixels), frame? (layer frame; default current), radius? (layer px; default the tool's diameter / 2), effect?}", + "{layer?, kind?: fg|bg|refine|refineErase, points: [[x, y], …] (layer pixels), frame? (layer frame; default current), radius? (layer px; default the tool's diameter / 2), effect?, object? (0–31; default the current object: with a tracker each object is followed on its own and the mattes are united)}", has_comp, stroke ), @@ -43,14 +43,22 @@ pub fn specs() -> Vec { cmd!("roto.span", "Set Segmentation Span", [], None, "{layer?, effect?, start?, end? (layer frames)}", has_comp, span), cmd!("roto.freeze", "Freeze", [], None, "{layer?, effect?, wait?}", has_comp, freeze), cmd!("roto.unfreeze", "Unfreeze", [], None, "{layer?, effect?}", has_comp, unfreeze), - cmd!("roto.clearStrokes", "Remove Roto Brush Strokes", [], None, "{layer?, effect?, frame? (only this layer frame), kind?}", has_comp, clear_strokes), + cmd!( + "roto.clearStrokes", + "Remove Roto Brush Strokes", + [], + None, + "{layer?, effect?, frame? (only this layer frame), kind?, object? (only this object's strokes)}", + has_comp, + clear_strokes + ), cmd!("roto.cancel", "Stop Roto Brush", [], None, "{}", is_running, |s, _| Ok(json!({"stopped": s.stop_roto()}))), cmd!( "roto.options", "Roto Brush Options", [], None, - "{diameter?, refineDiameter?, view?: alphaBoundary|alpha|alphaOverlay|none, overlayColor?: [r,g,b], overlayOpacity?, boundaryColor?: [r,g,b], autoPropagate?}", + "{diameter?, refineDiameter?, view?: alphaBoundary|alpha|alphaOverlay|none, overlayColor?: [r,g,b], overlayOpacity?, boundaryColor?: [r,g,b], autoPropagate?, object? (the object new strokes belong to, 0–31)}", always, options ), @@ -232,7 +240,8 @@ fn stroke(s: &mut Session, p: &Value) -> Result { if existing.is_none() && !kind.is_segmentation() { return Err(bad(cmd, "draw a foreground stroke with the Roto Brush tool first")); } - let st = Stroke { kind, frame, radius, points }; + let object = f_p(p, "object").map_or(s.state.roto.object, |v| v.clamp(0.0, rt::MAX_OBJECT_ID as f64).round() as u32); + let st = Stroke { kind, frame, radius, points, object }; let label = if kind.is_segmentation() { "Roto Brush Stroke" } else { "Refine Edge Stroke" }; let (uid, d) = s.edit(label, None, |proj, _| { let mut next = proj.next_id; @@ -263,7 +272,7 @@ fn stroke(s: &mut Session, p: &Value) -> Result { s.state.selected_layers = vec![lid]; } Ok(json!({ - "layer": lid.0, "effect": uid, "frame": frame, "kind": kind, "radius": radius, + "layer": lid.0, "effect": uid, "frame": frame, "kind": kind, "radius": radius, "object": object, "base": d.base, "span": d.span, "strokes": d.strokes.len(), })) } @@ -369,8 +378,9 @@ fn clear_strokes(s: &mut Session, p: &Value) -> Result { Some(k) => Some(StrokeKind::from_name(k).ok_or_else(|| bad(cmd, format!("unknown kind `{k}`")))?), None => None, }; + let object = f_p(p, "object").map(|v| v.round() as u32); let n0 = d.strokes.len(); - d.strokes.retain(|st| !(frame.is_none_or(|f| st.frame == f) && kind.is_none_or(|k| st.kind == k))); + d.strokes.retain(|st| !(frame.is_none_or(|f| st.frame == f) && kind.is_none_or(|k| st.kind == k) && object.is_none_or(|o| st.object == o))); if d.strokes.iter().all(|s| !s.kind.is_segmentation()) { d.base = None; } else if d.base.is_some_and(|b| !d.strokes.iter().any(|s| s.frame == b && s.kind.is_segmentation())) { @@ -412,6 +422,9 @@ fn options(s: &mut Session, p: &Value) -> Result { if let Some(b) = b_p(p, "autoPropagate") { o.auto_propagate = b; } + if let Some(v) = f_p(p, "object") { + o.object = v.clamp(0.0, rt::MAX_OBJECT_ID as f64).round() as u32; + } s.state.roto = o.clone(); Ok(serde_json::to_value(o).unwrap_or(Value::Null)) } @@ -430,6 +443,8 @@ fn status(s: &mut Session, p: &Value) -> Result { out["base"] = json!(d.base); out["span"] = json!(d.span); out["strokes"] = json!(d.strokes.len()); + out["objects"] = json!(d.objects().iter().map(|o| json!({"id": o, "strokes": d.strokes.iter().filter(|s| s.object == *o).count()})).collect::>()); + out["currentObject"] = json!(s.state.roto.object); out["strokeFrames"] = json!(d.stroke_frames()); out["frozen"] = json!(frozen_of(&g)); out["pending"] = json!(s.roto_pending.contains(&(cid, lid, uid))); diff --git a/crates/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index 66baee148..e4db67647 100644 --- a/crates/engine/src/tests/roto.rs +++ b/crates/engine/src/tests/roto.rs @@ -1007,3 +1007,28 @@ fn a_correction_re_tracks_only_the_frames_it_changes() { assert_eq!(trust.len(), 21); assert!(trust.values().all(|t| t.confidence > 0.9 && !t.occluded)); } + +#[test] +fn strokes_belong_to_the_current_object() { + let mut s = Session::default(); + s.execute("comp.new", json!({"width": 64, "height": 64, "frameRate": 25, "duration": 1})).unwrap(); + let layer = s.execute("layer.newSolid", json!({"width": 64, "height": 64})).unwrap()["layer"].as_u64().unwrap(); + s.execute("roto.stroke", json!({"layer": layer, "points": [[10, 10]], "radius": 4})).unwrap(); + // the current object is an option shared by the tool and the commands + s.execute("roto.options", json!({"object": 1})).unwrap(); + s.execute("roto.stroke", json!({"layer": layer, "points": [[50, 50]], "radius": 4})).unwrap(); + s.execute("roto.stroke", json!({"layer": layer, "kind": "bg", "points": [[30, 30]], "radius": 4})).unwrap(); + // an explicit object wins over the option, and is clamped to the supported range + s.execute("roto.stroke", json!({"layer": layer, "points": [[20, 50]], "radius": 4, "object": 99})).unwrap(); + let st = s.execute("roto.status", json!({"layer": layer})).unwrap(); + assert_eq!(st["currentObject"], 1); + assert_eq!(st["objects"], json!([{"id": 0, "strokes": 1}, {"id": 1, "strokes": 2}, {"id": 31, "strokes": 1}])); + // removing one object's strokes leaves the others + let r = s.execute("roto.clearStrokes", json!({"layer": layer, "object": 1})).unwrap(); + assert_eq!(r["removed"], 2); + let st = s.execute("roto.status", json!({"layer": layer})).unwrap(); + assert_eq!(st["objects"], json!([{"id": 0, "strokes": 1}, {"id": 31, "strokes": 1}])); + // the object survives a save and a load + let json = s.project.to_json(); + assert!(json.contains("\\\"object\\\":31"), "the project file carries the object"); +} diff --git a/crates/track/src/roto.rs b/crates/track/src/roto.rs index 000f7b55f..18c6eaa77 100644 --- a/crates/track/src/roto.rs +++ b/crates/track/src/roto.rs @@ -88,6 +88,15 @@ pub struct Stroke { pub frame: i64, pub radius: f64, pub points: Vec<[f64; 2]>, + /// Which object of the instance the stroke belongs to (0 for the first, and for every older project). + /// With a tracker each object is followed on its own and the mattes are combined; the classic engine + /// treats all foreground strokes as one object. + #[serde(default, skip_serializing_if = "is_zero")] + pub object: u32, +} + +fn is_zero(v: &u32) -> bool { + *v == 0 } /// A Roto Brush instance's user data: strokes, the base frame and the segmentation span @@ -130,6 +139,18 @@ impl RotoData { v.dedup(); v } + /// The objects that have a foreground stroke, in order. + pub fn objects(&self) -> Vec { + let mut v: Vec = self.strokes.iter().filter(|s| s.kind == StrokeKind::Fg).map(|s| s.object).collect(); + v.sort_unstable(); + v.dedup(); + v + } + /// The strokes of `object` on `frame`: its own foreground, background and box strokes, and the Refine Edge + /// strokes (which belong to the whole matte). + pub fn strokes_of(&self, frame: i64, object: u32) -> Vec<&Stroke> { + self.strokes.iter().filter(|s| s.frame == frame && (s.object == object || matches!(s.kind, StrokeKind::Refine | StrokeKind::RefineErase))).collect() + } pub fn in_span(&self, f: i64) -> bool { self.base.is_some() && f >= self.span[0] && f <= self.span[1] } @@ -224,6 +245,26 @@ pub struct FrameSeg { } impl FrameSeg { + /// The union of several objects' segmentations of one frame (the largest keeps its colour models). The + /// parts must share one size; any that does not is ignored. + pub fn union(parts: Vec) -> FrameSeg { + let Some(first) = parts.first() else { return FrameSeg::default() }; + let (w, h) = (first.w, first.h); + let mut parts: Vec = parts.into_iter().filter(|p| p.w == w && p.h == h && p.matte.len() == w * h).collect(); + parts.sort_by_key(|p| std::cmp::Reverse(p.area())); + let mut it = parts.into_iter(); + let Some(mut out) = it.next() else { return FrameSeg::default() }; + for p in it { + for (a, b) in out.matte.iter_mut().zip(&p.matte) { + *a = (*a).max(*b); + } + for (a, b) in out.refine.iter_mut().zip(&p.refine) { + *a = (*a).max(*b); + } + } + out + } + pub fn empty(w: usize, h: usize) -> FrameSeg { let Some(n) = w.checked_mul(h).filter(|n| *n <= rle::MAX_PLANE_BYTES) else { return FrameSeg::default(); diff --git a/crates/track/src/roto/bench.rs b/crates/track/src/roto/bench.rs index 2dafb9f3f..c086bfd57 100644 --- a/crates/track/src/roto/bench.rs +++ b/crates/track/src/roto/bench.rs @@ -122,7 +122,7 @@ fn line(kind: StrokeKind, frame: i64, a: [f64; 2], b: [f64; 2], radius: f64) -> let t = f64::from(k) / 10.0; [a[0] + (b[0] - a[0]) * t, a[1] + (b[1] - a[1]) * t] }; - Stroke { kind, frame, radius, points: (0..=10).map(pt).collect() } + Stroke { kind, frame, radius, points: (0..=10).map(pt).collect(), object: 0 } } /// The strokes an artist would paint on the base frame: one through the disk, two around it. diff --git a/crates/track/src/roto/tests.rs b/crates/track/src/roto/tests.rs index b7d8c7809..50ccaeaa4 100644 --- a/crates/track/src/roto/tests.rs +++ b/crates/track/src/roto/tests.rs @@ -65,7 +65,13 @@ pub(crate) fn frame_at(c: [f64; 2]) -> (Image, Vec) { } fn line(kind: StrokeKind, frame: i64, a: [f64; 2], b: [f64; 2], radius: f64) -> Stroke { - Stroke { kind, frame, radius, points: (0..=10).map(|k| [a[0] + (b[0] - a[0]) * k as f64 / 10.0, a[1] + (b[1] - a[1]) * k as f64 / 10.0]).collect() } + Stroke { + kind, + frame, + radius, + points: (0..=10).map(|k| [a[0] + (b[0] - a[0]) * k as f64 / 10.0, a[1] + (b[1] - a[1]) * k as f64 / 10.0]).collect(), + object: 0, + } } /// The strokes an artist would paint on frame `f`: one through the disk, one around it. @@ -231,6 +237,7 @@ fn ring(frame: i64, c: [f64; 2], r: f64, radius: f64) -> Stroke { frame, radius, points: (0..=48).map(|k| k as f64 / 48.0 * std::f64::consts::TAU).map(|a| [c[0] + r * a.cos(), c[1] + r * a.sin()]).collect(), + object: 0, } } @@ -295,7 +302,7 @@ fn later_refine_strokes_override_earlier_erases() { 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 erase = Stroke { kind: StrokeKind::RefineErase, frame: 0, radius: 6.0, points: vec![spot], object: 0 }; 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![ @@ -389,7 +396,7 @@ fn roto_perf_1080p() { let mut seg2 = next.clone(); let c = center(1); let ring: Vec<[f64; 2]> = (0..=64).map(|k| k as f64 / 64.0 * std::f64::consts::TAU).map(|t| [(c[0] + R * t.cos()) * s, (c[1] + R * t.sin()) * s]).collect(); - let st = [Stroke { kind: StrokeKind::Refine, frame: 1, radius: 20.0, points: ring }]; + let st = [Stroke { kind: StrokeKind::Refine, frame: 1, radius: 20.0, points: ring, object: 0 }]; let r = rasterize(&st.iter().collect::>(), &[StrokeKind::Refine], 1.0, 1920, 1080); seg2.refine = r.iter().map(|v| *v as u8).collect(); let rgb = rgb_of(&b); @@ -500,8 +507,8 @@ fn mobilesam_segments_and_propagates_the_disk() { /// ignores it; a prior's box does not replace it. #[test] 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 fg = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 2.0, points: vec![[50.0, 40.0]], object: 0 }; + let bx = Stroke { kind: StrokeKind::Box, frame: 0, radius: 0.0, points: vec![[80.0, 60.0], [20.0, 10.0]], object: 0 }; 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)]); @@ -523,14 +530,14 @@ fn box_stroke_is_the_model_box_prompt() { /// 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 green = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 4.0, points: vec![[50.0, 40.0]], object: 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 = (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 erase = Stroke { kind: StrokeKind::Bg, frame: 0, radius: 15.0, points: vec![[25.0, 40.0]], object: 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)); @@ -561,7 +568,7 @@ impl effectcraft_segment::MaskModel for CountingModel { 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 dot = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 15.0, points: vec![[80.0, 50.0]], object: 0 }; let tap = Stroke { kind: StrokeKind::Bg, radius: 2.0, ..dot.clone() }; let refs = [&dot, &tap]; let model = CountingModel(rect, Default::default()); @@ -588,7 +595,7 @@ fn tracker_prompts_carry_a_handful_of_points() { let (fg, bg) = (p.points.iter().filter(|(_, f)| *f).count(), p.points.iter().filter(|(_, f)| !*f).count()); assert!((1..=TRACKER_FG_POINTS).contains(&fg) && (1..=TRACKER_BG_POINTS).contains(&bg), "{fg} fg, {bg} bg"); // a long, finely sampled stroke still gives at most the limit, and an empty one gives none - let long = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 2.0, points: (0..400).map(|i| [i as f64 * 0.4, 50.0]).collect() }; + let long = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 2.0, points: (0..400).map(|i| [i as f64 * 0.4, 50.0]).collect(), object: 0 }; assert!(tracker_prompt(&[&long], 1.0).points.len() <= TRACKER_FG_POINTS); assert!(tracker_prompt(&[], 1.0).points.is_empty()); } @@ -630,8 +637,55 @@ fn tracker_prompts_carry_the_box_of_the_classic_matte() { // no matte, no box; a box the user drew wins let none = tracker_prompt_boxed(&refs, 1.0, &FrameSeg::empty(W as usize, H as usize)); assert_eq!(none.bbox, None); - let bx = Stroke { kind: StrokeKind::Box, frame: 0, radius: 0.0, points: vec![[10.0, 10.0], [50.0, 60.0]] }; + let bx = Stroke { kind: StrokeKind::Box, frame: 0, radius: 0.0, points: vec![[10.0, 10.0], [50.0, 60.0]], object: 0 }; let mut with_box: Vec<&Stroke> = refs.clone(); with_box.push(&bx); assert_eq!(tracker_prompt_boxed(&with_box, 1.0, &classic).bbox, Some([10.0, 10.0, 50.0, 60.0])); } + +#[test] +fn strokes_carry_an_object_and_old_data_is_object_zero() { + // saved before objects existed: no field, and the field is not written for object 0 + let old = r#"{"base":0,"span":[0,4],"strokes":[{"kind":"fg","frame":0,"radius":3.0,"points":[[1.0,2.0]]}]}"#; + let d = RotoData::from_json(old); + assert_eq!(d.strokes.first().map(|s| s.object), Some(0)); + assert!(!d.to_json().contains("object"), "{}", d.to_json()); + // the keys of existing projects do not change + let same = RotoData::from_json(&d.to_json()); + assert_eq!(d.chain_keys(1), same.chain_keys(1)); + + let mut d = d; + d.add(Stroke { kind: StrokeKind::Fg, frame: 0, radius: 3.0, points: vec![[9.0, 9.0]], object: 2 }, 4, [0, 10]); + d.add(Stroke { kind: StrokeKind::Bg, frame: 0, radius: 3.0, points: vec![[0.0, 0.0]], object: 2 }, 4, [0, 10]); + d.add(Stroke { kind: StrokeKind::Refine, frame: 0, radius: 3.0, points: vec![[5.0, 5.0]], object: 0 }, 4, [0, 10]); + assert_eq!(d.objects(), vec![0, 2]); + assert!(d.to_json().contains("\"object\":2")); + assert_eq!(RotoData::from_json(&d.to_json()), d); + // an object's own strokes, plus the Refine Edge strokes that belong to the whole matte + assert_eq!(d.strokes_of(0, 2).len(), 3); + assert_eq!(d.strokes_of(0, 0).len(), 2); + assert!(d.strokes_of(0, 1).len() == 1); + // moving a stroke to another object changes the frame keys (the tracked frames are stale) + let before = d.chain_keys(1); + if let Some(s) = d.strokes.first_mut() { + s.object = 1; + } + assert_ne!(before, d.chain_keys(1)); +} + +#[test] +fn the_union_of_objects_is_their_largest_extent() { + let seg = |cells: &[usize]| { + let mut s = FrameSeg::empty(4, 1); + for c in cells { + s.matte[*c] = 255; + } + s + }; + let u = FrameSeg::union(vec![seg(&[0]), seg(&[2, 3]), seg(&[3])]); + assert_eq!(u.matte, vec![255, 0, 255, 255]); + assert_eq!(FrameSeg::union(vec![]).w, 0); + // a part of another size is ignored + let odd = FrameSeg::empty(2, 2); + assert_eq!(FrameSeg::union(vec![seg(&[1]), odd]).matte, vec![0, 255, 0, 0]); +} diff --git a/crates/ui-egui/src/panels/roto_tool.rs b/crates/ui-egui/src/panels/roto_tool.rs index f8cd383c6..bf47b558d 100644 --- a/crates/ui-egui/src/panels/roto_tool.rs +++ b/crates/ui-egui/src/panels/roto_tool.rs @@ -8,7 +8,7 @@ use effectcraft_engine::render::{EvalCtx, RenderOpts, Renderer}; use egui::{Align2, Color32, Pos2, Rect, Sense, Stroke, StrokeKind, pos2, vec2}; use serde_json::json; -use crate::i18n::tr; +use crate::i18n::{tr, tr_args}; use crate::theme::Tokens; use crate::{EffectcraftApp, widgets}; @@ -373,6 +373,31 @@ pub fn bar_buttons(app: &mut EffectcraftApp, ui: &mut egui::Ui, p: &egui::Painte } app.auto.add("layerPanel.roto.freeze", r, label); x += 80.0; + // Objects: strokes belong to the current object; with a tracker each is followed on its own and the mattes are united. + let objects = fx::data(¶ms).objects(); + let current = app.session.state.roto.object; + let known = objects.contains(¤t); + let r = Rect::from_min_size(pos2(x, bar.center().y - 10.0), vec2(86.0, 20.0)); + let label = tr_args("Object {}", &[&(current + 1)]); + if widgets::text_button(ui, r, &label, false, &t, egui::Id::new(("roto-object", uid))).clicked() { + // cycle through the objects that have strokes (and the empty one being drawn) + let mut all = objects.clone(); + if !known { + all.push(current); + all.sort_unstable(); + } + let next = all.iter().copied().find(|o| *o > current).or_else(|| all.first().copied()).unwrap_or(0); + let _ = app.session.execute("roto.options", json!({"object": next})); + } + app.auto.add("layerPanel.roto.object", r, &label); + x += 90.0; + let r = Rect::from_min_size(pos2(x, bar.center().y - 10.0), vec2(26.0, 20.0)); + if widgets::text_button(ui, r, "+", false, &t, egui::Id::new(("roto-object-new", uid))).clicked() { + let next = objects.iter().copied().chain([current]).max().map_or(0, |m| m + 1); + let _ = app.session.execute("roto.options", json!({"object": next})); + } + app.auto.add("layerPanel.roto.newObject", r, "New object"); + x += 34.0; if running { let r = Rect::from_min_size(pos2(x, bar.center().y - 10.0), vec2(60.0, 20.0)); if widgets::text_button(ui, r, tr("Stop"), false, &t, egui::Id::unique("roto-stop")).clicked() { diff --git a/crates/ui-egui/tests/ui_roto_span_confidence.rs b/crates/ui-egui/tests/ui_roto_span_confidence.rs index de561014a..f3a8bcd9b 100644 --- a/crates/ui-egui/tests/ui_roto_span_confidence.rs +++ b/crates/ui-egui/tests/ui_roto_span_confidence.rs @@ -82,6 +82,15 @@ fn the_span_bar_shows_the_trackers_confidence() { h.run_steps(4); assert_eq!(h.state().ui.focused, PanelKind::Layer); + // the object picker is in the bar; "+" starts a new object, which the next stroke would belong to + assert!(h.state().auto.find("layerPanel.roto.object").is_some(), "the object picker is on screen"); + assert!(h.state().auto.find("layerPanel.roto.newObject").is_some()); + assert_eq!(h.state().session.state.roto.object, 0); + h.state_mut().session.execute("roto.options", json!({"object": 1})).unwrap(); + h.run_steps(2); + assert_eq!(h.state().session.state.roto.object, 1); + h.state_mut().session.execute("roto.options", json!({"object": 0})).unwrap(); + h.run_steps(2); let track = h.state().auto.find("layerPanel.rotoSpan").expect("the span bar is on screen").rect; let (comp, layer) = { let comp = h.state().session.active_comp().unwrap().clone(); diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index 8b1931494..656574df6 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -29,6 +29,23 @@ strokes and span they came from are unchanged. Changing a stroke or the span mak Choosing another model, or the classic engine, leaves them in place and `roto.status` says so; a classic Propagate replaces them with classic mattes. Object Matte saves nothing. +## Several objects in one instance +A Roto Brush instance can hold several objects. The **Object N** button in the Layer panel's bar picks the object new +strokes belong to (it cycles through the objects that have strokes), and **+** starts a new one. Agents: `roto.options +{object}` sets the current object, `roto.stroke` takes `object` (0–31; default the current one), `roto.clearStrokes` +takes `object` to remove one object's strokes, and `roto.status` lists `objects` (`{id, strokes}`) and `currentObject`. + +With SAM 2, Propagate follows every object in **one pass** (the frame is encoded once, whatever the number of objects), each +from its own foreground and background strokes, and the mattes are united into the instance's one matte. Each object keeps +its own occlusion: the frame counts as absent only when every object is, and otherwise the least sure object sets the +frame's confidence on the span bar. A background stroke belongs to its object, so it separates that object from what +lies under it without cutting a hole in the others. Refine Edge strokes belong to the whole matte. The classic engine has +no objects: all foreground strokes make one object, and background strokes apply to it as before. Saved projects are +unchanged (strokes without an object are object 0, and nothing is written for object 0). + +This is not several Roto Brush instances sharing a pass: each instance reads the layer after the effects above it, so +instances can see different pixels. Objects inside one instance always see the same ones. + ## Sizes Four checkpoints are in the registry (Settings ▸ Roto Brush); the size is detected from the file, and each is pinned to a Hugging Face commit with its SHA-256. All run the same tracker; a bigger backbone costs time and memory.