From 66d8e5e94871719e1c3bc238c0747a332dbfd7f9 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Oct 2026 13:15:09 +0000 Subject: [PATCH] SAM 2 roto tracker: review fixes (direction-limited tracking, non-blocking model reload, per-object start passes, loader dispatch, CLI errors, hints) Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01RDvxCS8xcR28VArjzn1x7U --- apps/effectcraft-cli/src/main.rs | 40 ++++-- apps/effectcraft-cli/tests/it/cli.rs | 64 ++++++++++ crates/effects/src/catalog.rs | 3 +- crates/effects/src/roto.rs | 42 ++++++- crates/engine/src/analysis/models.rs | 10 +- crates/engine/src/analysis/roto.rs | 2 +- crates/engine/src/tests/models.rs | 28 +++++ crates/engine/src/tests/roto.rs | 3 + crates/sam2/src/gpu/mod.rs | 11 +- crates/sam2/src/mask.rs | 39 ++++-- crates/sam2/src/session.rs | 115 +++++++++++++++++- crates/sam2/src/track.rs | 31 ++--- crates/segment/README.md | 3 + crates/segment/src/lib.rs | 41 ++++++- crates/ui-egui/tests/it/main.rs | 1 + crates/ui-egui/tests/it/roto_model_guard.rs | 26 ++++ crates/ui-egui/tests/ui_roto_objects.rs | 3 +- .../ui-egui/tests/ui_roto_span_confidence.rs | 3 +- docs/agents.md | 2 +- 19 files changed, 410 insertions(+), 57 deletions(-) create mode 100644 crates/ui-egui/tests/it/roto_model_guard.rs diff --git a/apps/effectcraft-cli/src/main.rs b/apps/effectcraft-cli/src/main.rs index 1c528c8fb..dc47671e2 100644 --- a/apps/effectcraft-cli/src/main.rs +++ b/apps/effectcraft-cli/src/main.rs @@ -1013,26 +1013,46 @@ fn bench_cmd(args: &Args) -> Result<(), Failure> { Ok(()) } +/// The model a bench option asks for: the file named on the command line, where failing to read or +/// build it is an error saying why (not "not installed"); else the installed file, where a file that +/// cannot be used just leaves the model out. +fn bench_weights( + flag: &str, + named: Option<&str>, + installed: Option, + build: impl Fn(&[u8]) -> Result, +) -> Result, Failure> { + match named { + Some(path) => { + let bytes = std::fs::read(path).map_err(|e| Failure::Error(format!("{flag} {path}: {e}")))?; + build(&bytes).map(Some).map_err(|e| Failure::Error(format!("{flag} {path}: {e}"))) + } + None => Ok(installed.and_then(|p| std::fs::read(p).ok()).and_then(|bytes| build(&bytes).ok())), + } +} + /// `bench --roto`: the Roto Brush's trackers on synthetic clips whose true mattes are known. fn bench_roto(args: &Args) -> Result<(), Failure> { use effectcraft_segment::{Loaded, MOBILE_SAM, MaskModel}; use effectcraft_track::roto::bench; let frames = args.num("--frames")?.unwrap_or(24.0).clamp(12.0, 240.0) as usize; let session = effectcraft_host::session(); - // `--mobilesam FILE` uses the weights at FILE; otherwise the installed ones. - let mobilesam_file = args.opt("--mobilesam").map(std::path::PathBuf::from).or_else(|| session.model_path(&MOBILE_SAM)); - let mobilesam: Option> = - mobilesam_file.and_then(|p| std::fs::read(p).ok()).and_then(|bytes| effectcraft_segment::load(MOBILE_SAM.id, &bytes).ok()).and_then(|m| match m { - Loaded::Mask(m) => Some(m), - _ => None, - }); + // `--mobilesam FILE` uses the weights at FILE (a file that cannot be used is an error); + // otherwise the installed ones. + let mobilesam: Option> = bench_weights("--mobilesam", args.opt("--mobilesam"), session.model_path(&MOBILE_SAM), |bytes| { + match effectcraft_segment::load(MOBILE_SAM.id, bytes)? { + Loaded::Mask(m) => Ok(m), + _ => Err("not a mask model".to_string()), + } + })?; if mobilesam.is_none() { note!("bench --roto: MobileSAM is not installed; comparing the classical engine only"); } - // `--sam2 FILE` uses the weights at FILE; otherwise the installed ones. - let sam2_file = args.opt("--sam2").map(std::path::PathBuf::from).or_else(|| effectcraft_sam2::mask::SIZES.iter().find_map(|m| session.model_path(m))); + // `--sam2 FILE` uses the weights at FILE (a file that cannot be used is an error); otherwise + // the installed ones. + let installed_sam2 = effectcraft_sam2::mask::SIZES.iter().find_map(|m| session.model_path(m)); let sam2: Option> = - sam2_file.and_then(|p| std::fs::read(p).ok()).and_then(|bytes| effectcraft_sam2::mask::load(&bytes).ok()).map(|m| m as std::sync::Arc); + bench_weights("--sam2", args.opt("--sam2"), installed_sam2, |bytes| effectcraft_sam2::mask::load(bytes).map(|m| m as std::sync::Arc))?; if sam2.is_none() { note!("bench --roto: no SAM 2.1 size is installed; no SAM 2 column"); } diff --git a/apps/effectcraft-cli/tests/it/cli.rs b/apps/effectcraft-cli/tests/it/cli.rs index 5348f2a3a..ed6d98086 100644 --- a/apps/effectcraft-cli/tests/it/cli.rs +++ b/apps/effectcraft-cli/tests/it/cli.rs @@ -678,6 +678,70 @@ fn render_bitrate_applies_to_vp9_webm() { let _ = std::fs::remove_dir_all(&root); } +/// `bench --roto --clip FRAMES --truth MASKS` runs the trackers on the user's own image sequences: here a bright disc +/// drifting over a dark textured frame, with its true masks. Mismatched sequences are refused with a message. +#[test] +fn bench_roto_runs_on_a_users_clip() { + let (frames, masks) = (tmp("clip-frames"), tmp("clip-masks")); + for d in [&frames, &masks] { + let _ = std::fs::remove_dir_all(d); + std::fs::create_dir_all(d).unwrap(); + } + let (w, h) = (96u32, 72u32); + for f in 0..6u32 { + let (cx, cy) = (30.0 + 4.0 * f as f32, 36.0f32); + let (mut img, mut mask) = (image::RgbImage::new(w, h), image::GrayImage::new(w, h)); + for y in 0..h { + for x in 0..w { + let inside = ((x as f32 - cx).powi(2) + (y as f32 - cy).powi(2)).sqrt() < 14.0; + let tex = 40 + ((x * 7 + y * 13) % 30) as u8; + img.put_pixel(x, y, if inside { image::Rgb([230, 120, 40 + (x % 20) as u8]) } else { image::Rgb([tex / 2, tex, tex + 20]) }); + mask.put_pixel(x, y, image::Luma([if inside { 255 } else { 0 }])); + } + } + img.save(frames.join(format!("f{f:03}.png"))).unwrap(); + mask.save(masks.join(format!("f{f:03}.png"))).unwrap(); + } + let (f, m) = (frames.to_str().unwrap(), masks.to_str().unwrap()); + let out = bin().args(["bench", "--roto", "--clip", f, "--truth", m, "--frames", "6"]).output().unwrap(); + let text = String::from_utf8_lossy(&out.stdout).to_string(); + assert!(out.status.success(), "{text}\n{}", String::from_utf8_lossy(&out.stderr)); + let row = text.lines().find(|l| l.contains("classical")).unwrap_or_else(|| panic!("no classical row: {text}")); + let mean: f64 = row.split_whitespace().nth(2).unwrap().parse().unwrap(); + assert!(mean > 0.8, "classical mean IoU on the user's clip: {mean}\n{text}"); + // one mask short: refused + std::fs::remove_file(masks.join("f005.png")).unwrap(); + let out = bin().args(["bench", "--roto", "--clip", f, "--truth", m]).output().unwrap(); + assert!(!out.status.success()); + assert!(String::from_utf8_lossy(&out.stderr).contains("must match"), "{}", String::from_utf8_lossy(&out.stderr)); + // --clip without --truth + let out = bin().args(["bench", "--roto", "--clip", f]).output().unwrap(); + assert!(!out.status.success()); +} + +/// `bench --roto --sam2 FILE` / `--mobilesam FILE` say why a file named on the command line cannot be +/// used (a missing file, one that is not the published weights), instead of "not installed". +#[test] +fn bench_roto_reports_why_a_named_model_file_failed() { + let dir = tmp("bench-weights"); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + let bogus = dir.join("bogus.bin"); + std::fs::write(&bogus, b"not the weights").unwrap(); + for flag in ["--sam2", "--mobilesam"] { + let missing = dir.join("missing.bin"); + let out = bin().args(["bench", "--roto", flag, missing.to_str().unwrap()]).output().unwrap(); + let err = String::from_utf8_lossy(&out.stderr).to_string(); + assert!(!out.status.success(), "{flag}: {err}"); + assert!(err.contains(flag) && err.contains("missing.bin"), "{flag} names the file that could not be read: {err}"); + let out = bin().args(["bench", "--roto", flag, bogus.to_str().unwrap()]).output().unwrap(); + let err = String::from_utf8_lossy(&out.stderr).to_string(); + assert!(!out.status.success(), "{flag}: {err}"); + assert!(err.contains("does not match"), "{flag} says the file is not the published weights: {err}"); + } + let _ = std::fs::remove_dir_all(&dir); +} + /// Render uses the same empty/project/demo selection as the other CLI commands. #[test] fn render_empty_does_not_load_the_demo() { diff --git a/crates/effects/src/catalog.rs b/crates/effects/src/catalog.rs index 332968baf..7fcb024bf 100644 --- a/crates/effects/src/catalog.rs +++ b/crates/effects/src/catalog.rs @@ -128,7 +128,8 @@ pub fn effects_markdown() -> String { s.push_str("# Effects\n\n"); s.push_str("\n\n"); - let params: usize = r.iter().map(|e| e.params.len()).sum(); + // The parameters the table's rows list: the shown ones, not the hidden bookkeeping ones. + let params: usize = r.iter().map(|e| e.params.iter().filter(|p| !matches!(p.ui, effectcraft_project::ParamUi::Hidden)).count()).sum(); let gpu = r.iter().filter(|e| e.gpu).count(); let float = r.iter().filter(|e| e.float).count(); let partial = r.iter().filter(|e| PARTIAL.iter().any(|(id, _)| *id == e.id)).count(); diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index 5bcb7386d..235c26afa 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -588,9 +588,12 @@ impl Chain { cur } - /// Whether frame `f` is computed (at full resolution or at `scale`). + /// Whether frame `f` is computed (at full resolution or at `scale`): in this session's caches, + /// or in the tracking saved with the project, which needs no work to be used even before it has + /// been decoded (and whichever model Settings choose now). pub fn is_cached(&self, f: i64, scale: f64) -> bool { self.keys.get(&f).is_some_and(|k| [*k, tracked_key(*k)].iter().any(|k| cached(*k, 1.0).is_some() || cached(*k, scale).is_some())) + || self.saved.as_ref().is_some_and(|t| t.cache.frames.contains_key(&f) && t.keys.contains_key(&f)) } /// Whether the chosen model can follow an object through a clip (SAM 2). @@ -618,11 +621,24 @@ impl Chain { cancel: &dyn Fn() -> bool, progress: &mut dyn FnMut(usize), ) -> Result, String> { - self.track_with_pending_limit(frame, cancel, progress, pending::MAX_PENDING_BYTES) + self.track_frames(None, frame, cancel, progress) + } + + /// [`Chain::track`] limited to `wanted` frames (a one-direction Propagate): the window, the work and the + /// progress count cover only those frames of the span. + pub fn track_frames( + &self, + wanted: Option<&[i64]>, + frame: &mut dyn FnMut(i64) -> Option>, + cancel: &dyn Fn() -> bool, + progress: &mut dyn FnMut(usize), + ) -> Result, String> { + self.track_with_pending_limit(wanted, frame, cancel, progress, pending::MAX_PENDING_BYTES) } fn track_with_pending_limit( &self, + wanted: Option<&[i64]>, frame: &mut dyn FnMut(i64) -> Option>, cancel: &dyn Fn() -> bool, progress: &mut dyn FnMut(usize), @@ -635,6 +651,7 @@ impl Chain { // frames from it on (both directions), so only those come back here. let (lo_span, hi_span) = (self.data.span[0], self.data.span[1]); let needed: Vec = (lo_span..=hi_span) + .filter(|f| wanted.is_none_or(|w| w.contains(f))) .filter(|f| self.keys.get(f).is_some_and(|k| self.saved_seg(*f, self.opts.scale).is_none() && cached(tracked_key(*k), self.opts.scale).is_none())) .collect(); let (Some(&need_lo), Some(&need_hi)) = (needed.first(), needed.last()) else { @@ -679,6 +696,9 @@ impl Chain { // A frame is finished when every object has reported it (the objects' mattes are united). let mut pending = pending::PendingMasks::new(objects.len(), max_pending_bytes); let finish = |g: i64, parts: Vec>, summary: &mut TrackSummary, progress: &mut dyn FnMut(usize)| { + if !wanted.is_none_or(|w| w.contains(&g)) { + return; + } let (Some(key), Some(img)) = (self.keys.get(&g).copied(), (frame.borrow_mut())(g)) else { return }; let mut segs: Vec<(u32, effectcraft_track::roto::FrameSeg)> = vec![]; let (mut confidence, mut present) = (f32::INFINITY, false); @@ -1375,13 +1395,29 @@ mod tests { let mut get = |f: i64| Some(Arc::new(frame(f as f64 / 25.0))); let mut completed = 0; // Each map takes 32,000 bytes; two incomplete frames exceed this small test budget. - let error = chain.track_with_pending_limit(&mut get, &|| false, &mut |n| completed = n, 64_000).unwrap_err(); + let error = chain.track_with_pending_limit(None, &mut get, &|| false, &mut |n| completed = n, 64_000).unwrap_err(); assert!(error.contains("pending object masks") && error.contains("memory budget"), "{error}"); assert!(fake.cancelled.load(Relaxed), "the model must see cancellation before allocating the rest of the clip"); assert_eq!(fake.emitted.load(Relaxed), 2); assert_eq!(completed, 0, "partial objects are not complete tracked frames"); } + #[test] + fn a_one_direction_pass_tracks_and_reports_only_its_frames() { + let mut d = RotoData::default(); + d.add(Stroke { kind: StrokeKind::Fg, frame: 12, radius: 3.0, points: vec![[35.0, 40.0]], object: 0 }, 12, [0, 24]); + let params = params_with(&d); + let mut chain = Chain::new(¶ms, [100.0, 80.0], 25.0, 1.0); + chain.model = Some(Arc::new(TwoObjects(Default::default()))); + let mut get = |f: i64| Some(Arc::new(frame(f as f64 / 25.0))); + let forward: Vec = (12..=24).collect(); + let mut last = 0; + let summary = chain.track_frames(Some(&forward), &mut get, &|| false, &mut |n| last = n).unwrap().unwrap(); + assert_eq!((summary.frames, last), (forward.len(), forward.len()), "progress never exceeds the frames asked for"); + let _ = chain.ensure(5, &mut get, &|| false); + assert!(tracks_of(&chain).is_none_or(|t| t.frames.keys().all(|f| *f >= 12)), "earlier frames were not tracked"); + } + #[test] fn propagation_replaces_saved_tracking_from_a_different_model() { let mut d = RotoData::default(); diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index d166702ac..ae15fe1d3 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -311,11 +311,15 @@ impl Session { } } - /// Load the chosen Roto Brush model again if it was unloaded for being idle (before a stroke uses it). + /// Start loading the chosen Roto Brush model again if it was unloaded for being idle (before a stroke uses it). + /// The load runs on the model worker, so the stroke is never blocked on it (the classic engine segments it + /// meanwhile and `poll_models` installs the model when it arrives), and a model that already failed to load is + /// not retried on every stroke: choosing it again clears the error. pub fn reload_roto_model(&mut self) { let chosen = self.chosen_model(Task::Mask).to_string(); - if chosen != seg::CLASSICAL && self.models.status(Task::Mask).active.is_none() && self.models.status(Task::Mask).busy.is_none() { - self.apply_model(Task::Mask, true); + let st = self.models.status(Task::Mask); + if chosen != seg::CLASSICAL && st.active.is_none() && st.busy.is_none() && st.error.is_none() { + self.apply_model(Task::Mask, false); } } diff --git a/crates/engine/src/analysis/roto.rs b/crates/engine/src/analysis/roto.rs index 7a240f9df..a07b0010e 100644 --- a/crates/engine/src/analysis/roto.rs +++ b/crates/engine/src/analysis/roto.rs @@ -465,7 +465,7 @@ async fn run_work(w: Work, shared: &RotoShared) { s.fps = if s.elapsed > 0.0 { s.done as f64 / s.elapsed } else { 0.0 }; crate::offload::report(s.done, s.total); }; - let result = chain.track(&mut get, &cancel, &mut report); + let result = chain.track_frames(Some(&w.frames), &mut get, &cancel, &mut report); if cancel() { let mut s = lock(&shared.state); s.cancelled = true; diff --git a/crates/engine/src/tests/models.rs b/crates/engine/src/tests/models.rs index ed0eecf9b..84bbd2381 100644 --- a/crates/engine/src/tests/models.rs +++ b/crates/engine/src/tests/models.rs @@ -161,6 +161,34 @@ fn an_idle_roto_model_is_unloaded() { let _ = std::fs::remove_dir_all(d); } +/// A stroke reloading an unloaded model never waits for the load and never retries one that already failed. +#[cfg(not(target_arch = "wasm32"))] +#[test] +fn reloading_for_a_stroke_does_not_block_or_retry_a_failed_load() { + super::roto::hold_roto_cache_test_lock(); + let (mut s, d) = session("reload-failed"); + let info = effectcraft_segment::info("mobilesam").unwrap(); + let dir = s.models_dir().unwrap(); + std::fs::create_dir_all(&dir).unwrap(); + // the right size, the wrong content: found as installed, then fails to load + std::fs::File::create(dir.join(info.file_name)).unwrap().set_len(info.size).unwrap(); + s.execute_checked("roto.model.select", json!({"id": "mobilesam"})).unwrap(); + let t0 = std::time::Instant::now(); + while s.models.status(Task::Mask).busy.is_some() && t0.elapsed().as_secs() < 30 { + std::thread::sleep(std::time::Duration::from_millis(5)); + } + let failed = s.models.status(Task::Mask).error; + assert!(failed.is_some(), "the unreadable file fails to load"); + s.reload_roto_model(); + let st = s.models.status(Task::Mask); + assert!(st.busy.is_none(), "a failed model is not loaded again by every stroke"); + assert_eq!(st.error, failed); + // an unloaded, not failed model loads on the worker: the call returns with the load still pending or done + s.execute_checked("roto.model.select", json!({"id": "classical"})).unwrap(); + crate::effects::roto::set_model(None); + let _ = std::fs::remove_dir_all(d); +} + /// With `EFFECTCRAFT_DOWNLOAD_SIZE=sam2.1-hiera-small` (or another registry id): the in-app download (Settings ▸ Roto /// Brush ▸ Download) fetches the pinned file over the network, verifies it, installs it with its notice and loads it. /// Skipped without the variable: it needs the network and hundreds of megabytes. diff --git a/crates/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index b28490a57..6c57d5c29 100644 --- a/crates/engine/src/tests/roto.rs +++ b/crates/engine/src/tests/roto.rs @@ -705,6 +705,7 @@ fn sam2_tracks_the_clip_through_roto_propagate() { let st = s.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap(); assert_eq!(st["source"], "tracker", "{st}"); assert!(st["notice"].as_str().unwrap().contains("Propagate again"), "{st}"); + assert_eq!(st["computed"].as_array().unwrap().len(), 21, "still computed under another model: {st}"); // A classic Propagate: its own segmentation, no tracker data. s.execute_checked("roto.propagate", json!({"layer": clip.0, "direction": "both", "wait": true})).unwrap(); let st = s.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap(); @@ -808,6 +809,8 @@ fn saved_tracking_is_kept_in_the_project_and_applies_when_it_reopens() { assert_eq!(st["source"], "tracker", "{st}"); assert_eq!(st["tracker"]["saved"], sam2, "{st}"); assert_eq!(st["tracker"]["frames"].as_object().unwrap().len(), 21, "{st}"); + // The span bar's "computed" frames include the saved ones before anything has been decoded. + assert_eq!(st["computed"].as_array().unwrap().len(), 21, "{st}"); let j = iou_at(&mut s2, clip, 12, center(12)); assert!(j > 0.9, "IoU of a saved frame: {j}"); // The render path (Chain::ensure) reads the same frame. diff --git a/crates/sam2/src/gpu/mod.rs b/crates/sam2/src/gpu/mod.rs index 152c67c50..fab8573e1 100644 --- a/crates/sam2/src/gpu/mod.rs +++ b/crates/sam2/src/gpu/mod.rs @@ -96,15 +96,24 @@ pub fn half_weights() -> bool { std::env::var("EFFECTCRAFT_SAM2_F16").is_ok_and(|v| matches!(v.trim(), "1" | "on" | "true")) } +/// The process's one GPU device, once something has asked for it. +static DEVICE: std::sync::OnceLock, String>> = std::sync::OnceLock::new(); + impl Gpu { /// The process's one GPU device, created on first use. Several devices at once upset some Vulkan loaders /// (a crash when two threads create or drop them together), and one is all SAM 2 needs: models come and go, /// the device stays. An error (no GPU) is remembered too. pub fn shared() -> Result> { - static DEVICE: std::sync::OnceLock, String>> = std::sync::OnceLock::new(); DEVICE.get_or_init(|| Gpu::new().map(Arc::new).map_err(|e| e.to_string())).clone().map_err(Error::Compute) } + /// Whether the shared device has been looked for yet, and found: `None` before anything has asked + /// for it (this never creates it, so a caller that only wants to describe the GPU does not stall + /// on the adapter and shader compilation), then `Some(true)` or `Some(false)` (no GPU). + pub fn probe() -> Option { + DEVICE.get().map(|d| d.is_ok()) + } + /// A device on the best real GPU (a software adapter is refused: the CPU path is faster). fn new() -> Result { let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle()); diff --git a/crates/sam2/src/mask.rs b/crates/sam2/src/mask.rs index 8470bd4f3..911ff2e00 100644 --- a/crates/sam2/src/mask.rs +++ b/crates/sam2/src/mask.rs @@ -156,17 +156,29 @@ pub fn recommended() -> &'static ModelInfo { /// A one-line note for the Settings list: what was measured for this size, and what the bigger ones cost without a GPU. pub fn hint(id: &str) -> Option { - let gpu = crate::gpu::Gpu::shared().is_ok(); + hint_for(id, crate::gpu::Gpu::probe()) +} + +/// [`hint`] given whether a GPU is known (`None`: not looked for yet). The list is built on the +/// Settings path, so it never creates the GPU device itself: that happens when a model is loaded, +/// and after that the notes say whether one was found. +fn hint_for(id: &str, gpu: Option) -> Option { Some( match id { "sam2.1-hiera-tiny" => "Recommended: the smallest download, and no larger size tracked better on the benchmark clips.", "sam2.1-hiera-small" => "As fast as Tiny on a GPU and no better on the benchmark clips; worth a try on difficult footage.", - "sam2.1-hiera-base-plus" if !gpu => "Slow without a GPU (several seconds a frame); none was found.", - "sam2.1-hiera-base-plus" => "About 15% slower than Tiny on a GPU and no better on the benchmark clips.", - "sam2.1-hiera-large" if !gpu => "Slow without a GPU (many seconds a frame); none was found.", - "sam2.1-hiera-large" => { - "About 1.7 times Tiny's time on a GPU; it lost the object after it re-entered in one benchmark clip. Try it on difficult footage." - } + "sam2.1-hiera-base-plus" => match gpu { + Some(false) => "Slow without a GPU (several seconds a frame); none was found.", + Some(true) => "About 15% slower than Tiny on a GPU and no better on the benchmark clips.", + None => "About 15% slower than Tiny on a GPU and no better on the benchmark clips; several seconds a frame without one.", + }, + "sam2.1-hiera-large" => match gpu { + Some(false) => "Slow without a GPU (many seconds a frame); none was found.", + Some(true) => { + "About 1.7 times Tiny's time on a GPU; it lost the object after it re-entered in one benchmark clip. Try it on difficult footage." + } + None => "About 1.7 times Tiny's time on a GPU, many seconds a frame without one; it lost the object after it re-entered in one benchmark clip.", + }, _ => return None, } .to_string(), @@ -191,6 +203,19 @@ mod tests { assert!(v.iter().all(|x| x.is_finite())); } + /// The Settings list asks for hints without creating the GPU device: before anything has loaded a + /// model there is no answer yet, and the notes say so rather than guess. + #[test] + fn hints_do_not_need_the_gpu_to_have_been_created() { + for id in ["sam2.1-hiera-base-plus", "sam2.1-hiera-large"] { + let (unknown, none, found) = (hint_for(id, None), hint_for(id, Some(false)), hint_for(id, Some(true))); + assert!(unknown.is_some() && none.is_some() && found.is_some()); + assert!(none.as_deref().is_some_and(|h| h.contains("none was found")), "{id}: {none:?}"); + assert!(unknown.as_deref().is_some_and(|h| h.contains("without one") && !h.contains("none was found")), "{id}: {unknown:?}"); + } + assert!(hint_for("mobilesam", None).is_none()); + } + #[test] fn a_wrong_file_is_refused_before_parsing() { let e = load(b"not the weights").err().unwrap_or_default(); diff --git a/crates/sam2/src/session.rs b/crates/sam2/src/session.rs index 336518bd5..54a39d212 100644 --- a/crates/sam2/src/session.rs +++ b/crates/sam2/src/session.rs @@ -615,15 +615,118 @@ impl<'a> VideoTracker<'a> { if let Some(limit) = self.budget.max_seconds.filter(|l| started.elapsed().as_secs_f64() > *l) { return Err(Error::Budget(format!("tracking took more than {limit} s"))); } - let cached = self.objects.get(id).is_some_and(|o| o.cond.contains_key(&frame)); - let stored = self.objects.get(id).and_then(|o| o.cond.get(&frame)).and_then(|o| o.result.clone()); - let mask = match stored { - Some(m) if cached => m, - _ => self.step(*id, frame, None, reverse, false)?, - }; + let mask = self.mask_at(*id, frame, reverse)?; out(mask); } } Ok(()) } + + /// Object `id` on `frame`: the result stored with its prompt there, else the tracker's step. + fn mask_at(&mut self, id: ObjectId, frame: usize, reverse: bool) -> Result { + let stored = self.objects.get(&id).and_then(|o| o.cond.get(&frame)).and_then(|o| o.result.clone()); + match stored { + Some(m) => Ok(m), + None => self.step(id, frame, None, reverse, false), + } + } + + /// Track every listed object from its own prompted frame (`starts`: object, frame) in one pass + /// in `direction`, frame by frame with all the objects that have started: a frame is finished + /// for all of them before the next begins, so a caller that assembles the objects of a frame + /// does not hold frames waiting for a pass that has not got there. A backward pass leaves out + /// each object's own start frame (the forward pass reported it). + pub fn propagate_each(&mut self, starts: &[(ObjectId, usize)], direction: Direction, cancel: &Cancel, out: &mut dyn FnMut(FrameMask)) -> Result<()> { + let n = self.frames.num_frames(); + for (id, start) in starts { + self.check_frame(*start)?; + if self.objects.get(id).is_none_or(|o| o.cond.is_empty()) { + return Err(Error::Input("no object has been prompted".into())); + } + } + if starts.is_empty() { + return Err(Error::Input("no object has been prompted".into())); + } + let reverse = direction == Direction::Backward; + let started = std::time::Instant::now(); + for (id, frame) in pass_order(starts, n, reverse) { + if cancel.is_cancelled() { + return Err(Error::Cancelled); + } + if let Some(limit) = self.budget.max_seconds.filter(|l| started.elapsed().as_secs_f64() > *l) { + return Err(Error::Budget(format!("tracking took more than {limit} s"))); + } + let mask = self.mask_at(id, frame, reverse)?; + out(mask); + } + Ok(()) + } +} + +/// The (object, frame) steps of [`VideoTracker::propagate_each`], in order: frame by frame from the +/// first start (the last, backwards), each frame for the objects whose pass has reached it. +fn pass_order(starts: &[(ObjectId, usize)], frames: usize, reverse: bool) -> impl Iterator + '_ { + let first = if reverse { starts.iter().map(|s| s.1).max() } else { starts.iter().map(|s| s.1).min() }; + let span = match first { + Some(f) if reverse => f + 1, + Some(f) => frames.saturating_sub(f), + None => 0, + }; + (0..span).flat_map(move |k| { + let frame = if reverse { first.unwrap_or(0).saturating_sub(k) } else { first.unwrap_or(0) + k }; + starts.iter().filter_map(move |(id, start)| { + let reached = if reverse { *start > frame } else { *start <= frame }; + reached.then_some((*id, frame)) + }) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn order(starts: &[(u32, usize)], frames: usize, reverse: bool) -> Vec<(u32, usize)> { + let starts: Vec<(ObjectId, usize)> = starts.iter().map(|(o, s)| (ObjectId(*o), *s)).collect(); + pass_order(&starts, frames, reverse).map(|(o, f)| (o.0, f)).collect() + } + + /// Forward, every frame is finished for all the objects that have started before the next frame + /// begins; an object joins the pass at its own start frame. + #[test] + fn the_forward_pass_goes_frame_by_frame_over_the_objects_that_have_started() { + assert_eq!(order(&[(1, 0), (2, 3)], 6, false), [(1, 0), (1, 1), (1, 2), (1, 3), (2, 3), (1, 4), (2, 4), (1, 5), (2, 5)]); + // Prompts on the same frame go together; the pass starts at the earliest start. + assert_eq!(order(&[(1, 2), (2, 2)], 4, false), [(1, 2), (2, 2), (1, 3), (2, 3)]); + assert_eq!(order(&[(1, 2)], 4, false), [(1, 2), (1, 3)]); + } + + /// Backward, from the last start down to frame 0, leaving out each object's own start frame. + #[test] + fn the_backward_pass_goes_frame_by_frame_and_skips_each_objects_own_start() { + assert_eq!(order(&[(1, 0), (2, 3)], 6, true), [(2, 2), (2, 1), (2, 0)]); + assert_eq!(order(&[(1, 2), (2, 3)], 6, true), [(2, 2), (1, 1), (2, 1), (1, 0), (2, 0)]); + assert!(order(&[(1, 0)], 6, true).is_empty(), "nothing before frame 0"); + } + + /// A frame is finished for every object before the next begins, however far apart the prompts + /// are: the number of frames waiting for an object never grows with the clip. + #[test] + fn no_frame_waits_for_an_object_that_has_not_reached_it() { + let steps = order(&[(1, 0), (2, 1000)], 2000, false); + let mut waiting = 0usize; + let mut worst = 0usize; + let mut seen = std::collections::BTreeMap::::new(); + for (_, f) in steps { + let n = seen.entry(f).or_default(); + *n += 1; + // From frame 1000 on, both objects report each frame back to back. + if f >= 1000 && *n == 2 { + waiting -= 1; + } else if f >= 1000 { + waiting += 1; + } + worst = worst.max(waiting); + } + assert!(worst <= 1, "{worst} frames waiting"); + } } diff --git a/crates/sam2/src/track.rs b/crates/sam2/src/track.rs index 1aff06d9d..ad920efe7 100644 --- a/crates/sam2/src/track.rs +++ b/crates/sam2/src/track.rs @@ -83,7 +83,6 @@ impl Sam2Mask { let mut tracker = VideoTracker::new(self.video_model().clone(), Box::new(source))?; let object_id = |oi: usize| ObjectId(oi as u32 + 1); // each object's first prompted frame: its pass starts there, forwards and backwards - let mut starts: std::collections::BTreeMap> = Default::default(); let mut first_of: HashMap = HashMap::new(); for (oi, idx, p) in prompt_frames { if cancel() { @@ -110,9 +109,8 @@ impl Sam2Mask { let e = first_of.entry(oi).or_insert(idx); *e = (*e).min(idx); } - for (oi, idx) in first_of { - starts.entry(idx).or_default().push(object_id(oi)); - } + let mut starts: Vec<(ObjectId, usize)> = first_of.into_iter().map(|(oi, idx)| (object_id(oi), idx)).collect(); + starts.sort_by_key(|(id, _)| id.0); let flag = Cancel::new(); let mut failure: Option = None; let mut emit = |m: FrameMask| { @@ -128,19 +126,22 @@ impl Sam2Mask { } } }; - for (start, ids) in starts { - tracker.propagate_objects(Some(&ids), start, Direction::Forward, None, &flag, &mut emit)?; - if start > 0 { - // the start frame was reported by the forward pass; the backward pass passes it again - tracker.propagate_objects(Some(&ids), start, Direction::Backward, None, &flag, &mut |m| { - if m.frame != start { - emit(m); - } - })?; + // One pass each way over all the objects, frame by frame, so the host can finish a frame (and + // let go of its masks) as soon as every object has it, however far apart the prompts are. + let passes = (|| { + tracker.propagate_each(&starts, Direction::Forward, &flag, &mut emit)?; + if starts.iter().any(|(_, start)| *start > 0) { + tracker.propagate_each(&starts, Direction::Backward, &flag, &mut emit)?; } - } + Ok(()) + })(); self.peak.store(tracker.peak_state_bytes() as u64, std::sync::atomic::Ordering::Relaxed); - failure.map_or(Ok(()), Err) + // A frame that could not be converted cancelled the passes; that is the error to report, not + // the cancellation it caused. + match failure { + Some(e) => Err(e), + None => passes, + } } } diff --git a/crates/segment/README.md b/crates/segment/README.md index 4962f8ef8..e6e4c6c98 100644 --- a/crates/segment/README.md +++ b/crates/segment/README.md @@ -37,6 +37,9 @@ kernels at run time). No ML framework, no unsafe code. |---|---|---|---|---|---|---| | `mobilesam` | Roto Brush | [MobileSAM](https://github.com/ChaoningZhang/MobileSAM): Segment Anything with a 5M-parameter TinyViT encoder | Chaoning Zhang et al., building on Meta AI's Segment Anything and Microsoft's TinyViT | Apache-2.0 | `mobile_sam.pt` from the official repository | 40.7 MB | | `sam2.1-hiera-tiny` | Roto Brush | [SAM 2.1 Tiny](https://github.com/facebookresearch/sam2): Hiera-Tiny video object tracker with streaming memory | Meta AI (Nikhila Ravi et al.) | Apache-2.0 | `model.safetensors` of `facebook/sam2.1-hiera-tiny` (Hugging Face, pinned to a commit); loaded by `effectcraft-sam2` | 155.9 MB | +| `sam2.1-hiera-small` | Roto Brush | [SAM 2.1 Small](https://github.com/facebookresearch/sam2): Hiera-Small video object tracker with streaming memory | Meta AI (Nikhila Ravi et al.) | Apache-2.0 | `model.safetensors` of `facebook/sam2.1-hiera-small` (Hugging Face, pinned to a commit); loaded by `effectcraft-sam2` | 184.3 MB | +| `sam2.1-hiera-base-plus` | Roto Brush | [SAM 2.1 Base+](https://github.com/facebookresearch/sam2): Hiera-Base+ video object tracker with streaming memory | Meta AI (Nikhila Ravi et al.) | Apache-2.0 | `model.safetensors` of `facebook/sam2.1-hiera-base-plus` (Hugging Face, pinned to a commit); loaded by `effectcraft-sam2` | 323.5 MB | +| `sam2.1-hiera-large` | Roto Brush | [SAM 2.1 Large](https://github.com/facebookresearch/sam2): Hiera-Large video object tracker with streaming memory | Meta AI (Nikhila Ravi et al.) | Apache-2.0 | `model.safetensors` of `facebook/sam2.1-hiera-large` (Hugging Face, pinned to a commit); loaded by `effectcraft-sam2` | 897.9 MB | | `mediapipe-face` | Face tracking | [MediaPipe Face Landmarker](https://developers.google.com/edge/mediapipe/solutions/vision/face_landmarker): BlazeFace detector and Face Mesh V2 | Google LLC | Apache-2.0 ([model card](https://storage.googleapis.com/mediapipe-assets/Model%20Card%20MediaPipe%20Face%20Mesh%20V2.pdf)) | `face_landmarker.task` (float16, version 1) from Google's model storage | 3.8 MB | How they were checked: diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index fb16f01ce..04894cb2b 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -281,14 +281,34 @@ pub fn load(id: &str, bytes: &[u8]) -> Result { if got != info.sha256 { return Err(format!("{}: the file does not match (SHA-256 {got}, expected {})", info.name, info.sha256)); } - match id { - "mobilesam" => { + match loader(id) { + Some(Loader::MobileSam) => { let w = pt::read(bytes)?; Ok(Loaded::Mask(Arc::new(Sam { net: mobilesam::MobileSam::from_weights(&w)? }))) } - "mediapipe-face" => Ok(Loaded::Face(Arc::new(mediapipe::FaceLandmarker::from_task(bytes)?))), - "sam2.1-hiera-tiny" => Err(format!("{} is loaded by the desktop application (effectcraft-sam2)", info.name)), - _ => Err(format!("no loader for `{id}`")), + Some(Loader::MediaPipe) => Ok(Loaded::Face(Arc::new(mediapipe::FaceLandmarker::from_task(bytes)?))), + Some(Loader::Desktop) => Err(format!("{} is loaded by the desktop application (effectcraft-sam2)", info.name)), + None => Err(format!("no loader for `{id}`")), + } +} + +/// What builds a registered model. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Loader { + MobileSam, + MediaPipe, + /// Built outside this crate, by the desktop application (every SAM 2.1 size). + Desktop, +} + +/// The loader for model `id`: one answer for every size of a family, so a new size cannot fall +/// through to "no loader" while its sibling gets the helpful error. +fn loader(id: &str) -> Option { + match id { + "mobilesam" => Some(Loader::MobileSam), + "mediapipe-face" => Some(Loader::MediaPipe), + id if id.starts_with("sam2.1-hiera-") => Some(Loader::Desktop), + _ => None, } } @@ -412,4 +432,15 @@ mod tests { assert!(load("mobilesam", b"not the weights").err().is_some_and(|e| e.contains("does not match"))); assert!(load("nope", b"").is_err()); } + + /// Every registered model has a loader, so none falls through to "no loader" after its checksum + /// matches; every SAM 2.1 size is loaded by the desktop application. + #[test] + fn every_registered_model_has_a_loader() { + for m in MODELS { + let l = loader(m.id); + assert!(l.is_some(), "{} has no loader", m.id); + assert_eq!(l == Some(Loader::Desktop), m.id.starts_with("sam2.1-"), "{}", m.id); + } + } } diff --git a/crates/ui-egui/tests/it/main.rs b/crates/ui-egui/tests/it/main.rs index 4fc34d9ff..9ef4c4b54 100644 --- a/crates/ui-egui/tests/it/main.rs +++ b/crates/ui-egui/tests/it/main.rs @@ -3,6 +3,7 @@ mod gpu_test; mod native_menu_text_keys; mod render_queue_snapshot; +mod roto_model_guard; mod shared_docking_ui; mod ui_3d; mod ui_accessibility; diff --git a/crates/ui-egui/tests/it/roto_model_guard.rs b/crates/ui-egui/tests/it/roto_model_guard.rs new file mode 100644 index 000000000..a65805b55 --- /dev/null +++ b/crates/ui-egui/tests/it/roto_model_guard.rs @@ -0,0 +1,26 @@ +//! The Roto Brush model is process-global (`effects::roto::set_model`), and every test in this binary shares the +//! process: a test that installs a scripted model holds this guard, so another cannot swap or clear the model +//! mid-run, and the model is cleared again even when the test panics. + +use std::sync::{Arc, Mutex, MutexGuard, PoisonError}; + +use effectcraft_segment::MaskModel; + +static LOCK: Mutex<()> = Mutex::new(()); + +pub struct InstalledModel { + _exclusive: MutexGuard<'static, ()>, +} + +/// Install `model` for the roto effect until the returned guard drops. +pub fn install(model: Arc) -> InstalledModel { + let guard = LOCK.lock().unwrap_or_else(PoisonError::into_inner); + effectcraft_engine::effects::roto::set_model(Some(model)); + InstalledModel { _exclusive: guard } +} + +impl Drop for InstalledModel { + fn drop(&mut self) { + effectcraft_engine::effects::roto::set_model(None); + } +} diff --git a/crates/ui-egui/tests/ui_roto_objects.rs b/crates/ui-egui/tests/ui_roto_objects.rs index 5d16b0ada..252ecfb3d 100644 --- a/crates/ui-egui/tests/ui_roto_objects.rs +++ b/crates/ui-egui/tests/ui_roto_objects.rs @@ -66,7 +66,7 @@ fn the_boundary_view_outlines_each_object_and_show_object_picks_one() { for (object, x) in [(0, 22.0), (1, 58.0)] { s.execute("roto.stroke", json!({"layer": id, "object": object, "points": [[x, 38.0], [x, 42.0]], "radius": 4})).unwrap(); } - effectcraft_engine::effects::roto::set_model(Some(Arc::new(TwoRegions))); + let _model = crate::roto_model_guard::install(Arc::new(TwoRegions)); s.execute("roto.propagate", json!({"layer": id, "direction": "both", "wait": true})).unwrap(); let st = s.execute("roto.status", json!({"layer": id})).unwrap(); assert_eq!(st["objects"], json!([{"id": 0, "strokes": 1}, {"id": 1, "strokes": 1}]), "{st}"); @@ -96,5 +96,4 @@ fn the_boundary_view_outlines_each_object_and_show_object_picks_one() { let one = h.render().expect("the panel renders"); let (a, b) = (count(&one, pink), count(&one, cyan)); assert!(a < 10 && b > 50, "only object 1 is shown: {a} pink, {b} cyan pixels"); - effectcraft_engine::effects::roto::set_model(None); } diff --git a/crates/ui-egui/tests/ui_roto_span_confidence.rs b/crates/ui-egui/tests/ui_roto_span_confidence.rs index f3a8bcd9b..04379f1ae 100644 --- a/crates/ui-egui/tests/ui_roto_span_confidence.rs +++ b/crates/ui-egui/tests/ui_roto_span_confidence.rs @@ -70,7 +70,7 @@ fn the_span_bar_shows_the_trackers_confidence() { let id = s.active_comp().unwrap().layers.iter().find(|l| l.name == "Box").unwrap().id.0; s.execute("roto.stroke", json!({"layer": id, "points": [[20.0, 40.0], [60.0, 40.0]], "radius": 6})).unwrap(); s.execute("roto.stroke", json!({"layer": id, "kind": "bg", "points": [[5.0, 5.0], [75.0, 5.0]], "radius": 4})).unwrap(); - effectcraft_engine::effects::roto::set_model(Some(Arc::new(Scripted))); + let _model = crate::roto_model_guard::install(Arc::new(Scripted)); // the model stays loaded while the panel draws (its id is part of the cache keys) s.execute("roto.propagate", json!({"layer": id, "direction": "both", "wait": true})).unwrap(); @@ -115,5 +115,4 @@ fn the_span_bar_shows_the_trackers_confidence() { let got = sample(f); assert!(close(got, want), "frame {f} ({what}): drew {got:?}, wanted about {want:?}"); } - effectcraft_engine::effects::roto::set_model(None); } diff --git a/docs/agents.md b/docs/agents.md index ddd730557..ef88acccb 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -506,7 +506,7 @@ Trained model (Roto Brush 2.0 / 3.0): licence, size, URL, SHA-256 and whether it is installed, selected and active. - `roto.model.download {id}` fetches and verifies the official weights in the background; `roto.model.install {path, id?}` installs a file you already have. -- `roto.model.select {id: classical|mobilesam|sam2.1-hiera-tiny}` chooses which one Roto Brush uses (SAM 2 also tracks the object through the clip, see [roto-brush-sam2.md](roto-brush-sam2.md); `roto.status` then reports `source` and per-frame `tracker.frames` with `confidence` and `occluded`). Loading runs +- `roto.model.select {id: classical|mobilesam|sam2.1-hiera-tiny|sam2.1-hiera-small|sam2.1-hiera-base-plus|sam2.1-hiera-large}` chooses which one Roto Brush uses (SAM 2 also tracks the object through the clip, see [roto-brush-sam2.md](roto-brush-sam2.md); `roto.status` then reports `source` and per-frame `tracker.frames` with `confidence` and `occluded`). Loading runs in the background; `"wait": true` (also on download and install) returns once it is in use. - `roto.propagate` never falls back to the classic engine silently: a chosen model that is not installed, or failed to load, is an error, and without `"wait": true` so is a model still loading. `roto.status` reports `model`