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

SAM 2 roto tracker: review fixes (direction-limited tracking, non-blocking model reload, per-object start passes, loader dispatch, CLI errors, hints)

Browse Source 
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RDvxCS8xcR28VArjzn1x7U
This commit is contained in:
Claude
2026-10-30 08:43:32 +01:00
committed by Kristian Myram
parent 6d10a6e932
commit 66d8e5e948
19 changed files with 410 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/effectcraft-cli/src/main.rs
+30 -10
View File
@@ -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<T>(
flag: &str,
named: Option<&str>,
installed: Option<std::path::PathBuf>,
build: impl Fn(&[u8]) -> Result<T, String>,
) -> Result<Option<T>, 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<std::sync::Arc<dyn MaskModel>> =
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<std::sync::Arc<dyn MaskModel>> = 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<std::sync::Arc<dyn MaskModel>> =
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<dyn MaskModel>);
bench_weights("--sam2", args.opt("--sam2"), installed_sam2, |bytes| effectcraft_sam2::mask::load(bytes).map(|m| m as std::sync::Arc<dyn MaskModel>))?;
if sam2.is_none() {
note!("bench --roto: no SAM 2.1 size is installed; no SAM 2 column");
}
apps/effectcraft-cli/tests/it/cli.rs
+64
View File
@@ -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() {
crates/effects/src/catalog.rs
+2 -1
View File
@@ -128,7 +128,8 @@ pub fn effects_markdown() -> String {
s.push_str("# Effects\n\n");
s.push_str("<!-- Generated from the effect registry (crates/effects/src/catalog.rs); do not edit by hand.\n");
s.push_str(" Regenerate: UPDATE_DOCS=1 cargo test -p effectcraft-effects --lib effects_doc_is_current -->\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();
crates/effects/src/roto.rs
+39 -3
View File
@@ -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<Option<TrackSummary>, 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<Arc<Image>>,
cancel: &dyn Fn() -> bool,
progress: &mut dyn FnMut(usize),
) -> Result<Option<TrackSummary>, 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<Arc<Image>>,
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<i64> = (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<Option<effectcraft_segment::TrackedFrame>>, 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(&params, [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<i64> = (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();
crates/engine/src/analysis/models.rs
+7 -3
View File
@@ -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);
}
}
crates/engine/src/analysis/roto.rs
+1 -1
View File
@@ -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;
crates/engine/src/tests/models.rs
+28
View File
@@ -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.
crates/engine/src/tests/roto.rs
+3
View File
@@ -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.
crates/sam2/src/gpu/mod.rs
+10 -1
View File
@@ -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<std::result::Result<Arc<Gpu>, 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<Arc<Gpu>> {
static DEVICE: std::sync::OnceLock<std::result::Result<Arc<Gpu>, 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<bool> {
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<Gpu> {
let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle());
crates/sam2/src/mask.rs
+32 -7
View File
@@ -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<String> {
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<bool>) -> Option<String> {
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();
crates/sam2/src/session.rs
+109 -6
View File
@@ -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<FrameMask> {
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<Item = (ObjectId, usize)> + '_ {
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::<usize, usize>::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");
}
}
crates/sam2/src/track.rs
+16 -15
View File
@@ -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<usize, Vec<ObjectId>> = Default::default();
let mut first_of: HashMap<usize, usize> = 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<Error> = 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,
}
}
}
crates/segment/README.md
+3
View File
@@ -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:
crates/segment/src/lib.rs
+36 -5
View File
@@ -281,14 +281,34 @@ pub fn load(id: &str, bytes: &[u8]) -> Result<Loaded> {
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<Loader> {
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);
}
}
}
crates/ui-egui/tests/it/main.rs
+1
View File
@@ -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;
crates/ui-egui/tests/it/roto_model_guard.rs
+26
View File
@@ -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<dyn MaskModel>) -> 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);
}
}
crates/ui-egui/tests/ui_roto_objects.rs
+1 -2
View File
@@ -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);
}
crates/ui-egui/tests/ui_roto_span_confidence.rs
+1 -2
View File
@@ -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);
}
docs/agents.md
+1 -1
View File
@@ -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`