diff --git a/apps/effectcraft-cli/src/main.rs b/apps/effectcraft-cli/src/main.rs index be066ecd1..e6a3352dc 100644 --- a/apps/effectcraft-cli/src/main.rs +++ b/apps/effectcraft-cli/src/main.rs @@ -27,7 +27,7 @@ //! (--gpu: CPU vs GPU ms/frame for every comp at Full and Half) //! effectcraft-cli bench --ops [--small] [--layers N] [--comps N] [--footage N] everyday-operation timings //! on a large generated project (open, save, auto-save, undo/redo, timeline, Project panel) -//! effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors] [--mobilesam mobile_sam.pt] [--fixes] [--json] Roto Brush trackers on synthetic clips with known mattes +//! effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors] [--mobilesam mobile_sam.pt] [--fixes] [--clip FRAMES --truth MASKS] [--json] Roto Brush trackers on synthetic clips with known mattes //! (classical engine, MobileSAM and SAM 2.1 Tiny when installed; IoU, flicker, re-appearance, seconds per frame) //! effectcraft-cli script FILE.jsx [F.ecproj] | --eval CODE run an After Effects-style script //! effectcraft-cli mcp [--autosave | --bridge PORT] MCP server on stdio @@ -126,7 +126,7 @@ const USAGE: &str = "usage: effectcraft-cli Result<(), Failure> { let fixes = args.flag("--fixes"); let measure = |sc: &bench::Scenario, run: &dyn Fn(&bench::Scenario) -> Option| if fixes { bench::with_fixes(sc, 0.9, 6, run) } else { run(sc) }; let mut runs = vec![]; - for sc in bench::scenarios(frames) { + let mut scenarios = bench::scenarios(frames); + // `--clip FRAMES_DIR --truth MASKS_DIR`: the user's own clip (image sequences, same file order) with its true mattes. + if let Some(dir) = args.opt("--clip") { + let truth = args.opt("--truth").ok_or_else(|| Failure::Error("--clip needs --truth MASKS_DIR".into()))?; + let name = std::path::Path::new(dir).file_name().map_or_else(|| "clip".to_string(), |n| n.to_string_lossy().chars().take(10).collect()); + scenarios = vec![load_bench_clip(&name, dir, truth, frames).map_err(Failure::Error)?]; + } + for sc in scenarios { runs.extend(measure(&sc, &|s| Some(bench::run_flow(s, "classical", None)))); if let Some(m) = &mobilesam { runs.extend(measure(&sc, &|s| Some(bench::run_flow(s, "mobilesam", Some(m.as_ref()))))); @@ -1055,6 +1064,39 @@ fn bench_roto(args: &Args) -> Result<(), Failure> { Ok(()) } +/// The image files of `dir` (png, jpg, jpeg), in file-name order. +fn image_files(dir: &str) -> Result, String> { + let mut files: Vec = std::fs::read_dir(dir) + .map_err(|e| format!("cannot read {dir}: {e}"))? + .filter_map(|e| e.ok().map(|e| e.path())) + .filter(|p| p.extension().and_then(|e| e.to_str()).is_some_and(|e| ["png", "jpg", "jpeg"].contains(&e.to_ascii_lowercase().as_str()))) + .collect(); + files.sort(); + Ok(files) +} + +/// A benchmark scenario from the user's frames and truth masks (at most `limit` frames, frames at most 480 px +/// wide; a mask pixel above 127 is the object). Nothing is copied or kept. +fn load_bench_clip(name: &str, frames_dir: &str, truth_dir: &str, limit: usize) -> Result { + let (frames, masks) = (image_files(frames_dir)?, image_files(truth_dir)?); + if frames.is_empty() || frames.len() != masks.len() { + return Err(format!("{} frames in {frames_dir} but {} masks in {truth_dir}: they must match one to one", frames.len(), masks.len())); + } + let n = frames.len().min(limit.clamp(1, 240)); + let (mut imgs, mut truth) = (vec![], vec![]); + for (f, m) in frames.iter().zip(&masks).take(n) { + let rgba = image::open(f).map_err(|e| format!("{}: {e}", f.display()))?.to_rgba8(); + let scale = (480.0 / f64::from(rgba.width().max(1))).min(1.0); + let (w, h) = (((f64::from(rgba.width()) * scale).round() as u32).max(1), ((f64::from(rgba.height()) * scale).round() as u32).max(1)); + let rgba = if (w, h) == rgba.dimensions() { rgba } else { image::imageops::resize(&rgba, w, h, image::imageops::FilterType::Triangle) }; + let mask = image::open(m).map_err(|e| format!("{}: {e}", m.display()))?.to_luma8(); + let mask = image::imageops::resize(&mask, w, h, image::imageops::FilterType::Nearest); + imgs.push(effectcraft_render::Image::from_rgba8(w, h, rgba.as_raw())); + truth.push(mask.as_raw().iter().map(|v| u8::from(*v > 127)).collect()); + } + effectcraft_track::roto::bench::scenario_from(name, imgs, truth) +} + /// `bench --ops`: everyday operations on a large generated project (200 comps, 5,000 layers in /// the main comp, 300 footage items, 20-deep precomp nesting, expressions): startup to first /// frame, open, save, auto-save, small and large edits with undo/redo, first frame, timeline diff --git a/apps/effectcraft-cli/tests/it/cli.rs b/apps/effectcraft-cli/tests/it/cli.rs index 7a770fa87..5348f2a3a 100644 --- a/apps/effectcraft-cli/tests/it/cli.rs +++ b/apps/effectcraft-cli/tests/it/cli.rs @@ -805,3 +805,44 @@ fn render_without_progress_writes_no_progress_lines() { let stderr = render_demo_stderr("no-progress.gif", &[]); assert!(!stderr.lines().any(|l| l.starts_with("progress")), "{stderr}"); } + +/// `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()); +} diff --git a/crates/sam2/src/mask.rs b/crates/sam2/src/mask.rs index 94cf8b40a..4578af0f9 100644 --- a/crates/sam2/src/mask.rs +++ b/crates/sam2/src/mask.rs @@ -55,11 +55,12 @@ pub struct Sam2Mask { model: Arc, last: Mutex)>>, used: Mutex, + pub(crate) peak: std::sync::atomic::AtomicU64, } impl Sam2Mask { pub fn new(model: Arc) -> Sam2Mask { - Sam2Mask { model, last: Mutex::new(None), used: Mutex::new(std::time::Instant::now()) } + Sam2Mask { model, last: Mutex::new(None), used: Mutex::new(std::time::Instant::now()), peak: std::sync::atomic::AtomicU64::new(0) } } /// Note that the model is being used (for [`MaskModel::idle_secs`]). @@ -121,6 +122,10 @@ impl MaskModel for Sam2Mask { Some(self.model.device_name()) } + fn peak_state_bytes(&self) -> Option { + Some(self.peak.load(std::sync::atomic::Ordering::Relaxed)) + } + fn idle_secs(&self) -> Option { Some(self.used.lock().unwrap_or_else(PoisonError::into_inner).elapsed().as_secs_f64()) } diff --git a/crates/sam2/src/session.rs b/crates/sam2/src/session.rs index 9de8fedc2..336518bd5 100644 --- a/crates/sam2/src/session.rs +++ b/crates/sam2/src/session.rs @@ -290,6 +290,7 @@ pub struct VideoTracker<'a> { objects: BTreeMap, cache: VecDeque<(usize, Arc)>, budget: Budget, + peak: usize, } fn sine_embedding_1d(positions: &[f32], dim: usize) -> Result { @@ -313,7 +314,7 @@ impl<'a> VideoTracker<'a> { if n == 0 || n > MAX_FRAMES { return Err(Error::Input(format!("a clip needs between 1 and {MAX_FRAMES} frames, got {n}"))); } - Ok(VideoTracker { model, frames, objects: BTreeMap::new(), cache: VecDeque::new(), budget: Budget::default() }) + Ok(VideoTracker { model, frames, objects: BTreeMap::new(), cache: VecDeque::new(), budget: Budget::default(), peak: 0 }) } pub fn num_frames(&self) -> usize { @@ -331,8 +332,14 @@ impl<'a> VideoTracker<'a> { frames + self.cache.len() * FEATURE_BYTES } - fn check_budget(&self) -> Result<()> { + /// The most [`VideoTracker::state_bytes`] reached so far. + pub fn peak_state_bytes(&self) -> usize { + self.peak + } + + fn check_budget(&mut self) -> Result<()> { let used = self.state_bytes(); + self.peak = self.peak.max(used); if used > self.budget.max_state_bytes { return Err(Error::Budget(format!("the tracker holds {} MB, over its {} MB limit", used >> 20, self.budget.max_state_bytes >> 20))); } diff --git a/crates/sam2/src/track.rs b/crates/sam2/src/track.rs index 2476ddb74..1aff06d9d 100644 --- a/crates/sam2/src/track.rs +++ b/crates/sam2/src/track.rs @@ -139,6 +139,7 @@ impl Sam2Mask { })?; } } + self.peak.store(tracker.peak_state_bytes() as u64, std::sync::atomic::Ordering::Relaxed); failure.map_or(Ok(()), Err) } } diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index 16f5cfd24..d121df840 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -752,3 +752,25 @@ fn gpu_memory_encoder_matches_the_cpu_one() { // wrong shapes are errors assert!(enc.encode(&features.narrow(1, 0, 8).unwrap(), mask).is_err()); } + +/// The tracker reports the most memory it held (memory banks, kept results, cached frames), and the model +/// reports it for the last tracking run. +#[test] +fn the_tracker_reports_its_peak_memory() { + use effectcraft_segment::{ClipFrame, MaskModel, Prompt, TrackRequest}; + let Some((w, _)) = load() else { return }; + let model = effectcraft_sam2::mask::Sam2Mask::new(std::sync::Arc::new(effectcraft_sam2::session::VideoModel::load(&w).unwrap())); + assert_eq!(model.peak_state_bytes(), Some(0)); + let (fw, fh) = (64usize, 48usize); + let frame = std::sync::Arc::new(ClipFrame { + rgb: (0..fw * fh).map(|i| if (i % fw).abs_diff(32) < 12 && (i / fw).abs_diff(24) < 12 { [0.9, 0.2, 0.1] } else { [0.2, 0.3, 0.5] }).collect(), + w: fw, + h: fh, + }); + let mut get = |_f: i64| Some(frame.clone()); + let req = TrackRequest::single(0, 3, vec![(0, Prompt { points: vec![([32.0, 24.0], true)], bbox: None, mask: None })]); + model.tracker().unwrap().track(&req, &mut get, &|| false, &mut |_, _, _| {}).unwrap(); + let peak = model.peak_state_bytes().unwrap(); + // three stored frames (about 2 MB each) and the cached frame features + assert!(peak > 6 << 20 && peak < 400 << 20, "peak {peak}"); +} diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index 3fff1eec4..3ebf964f9 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -49,6 +49,10 @@ pub trait MaskModel: Send + Sync { fn tracker(&self) -> Option<&dyn ObjectTracker> { None } + /// The most memory (bytes) the last tracking run held for its memory banks and cached frames (`None`: not tracked). + fn peak_state_bytes(&self) -> Option { + None + } /// What runs the model: the GPU's name or "CPU" (`None`: it does not say). fn device_name(&self) -> Option { None diff --git a/crates/track/src/roto/bench.rs b/crates/track/src/roto/bench.rs index 58aea2385..2dafb9f3f 100644 --- a/crates/track/src/roto/bench.rs +++ b/crates/track/src/roto/bench.rs @@ -20,10 +20,13 @@ const R: f64 = 32.0; /// A clip with the true matte of the visible object on every frame. #[derive(Clone)] pub struct Scenario { - pub name: &'static str, - pub what: &'static str, + pub name: String, + pub what: String, + /// Frame size in pixels (the synthetic clips are `W × H`; a real clip has its own). + pub w: usize, + pub h: usize, pub frames: Vec, - /// Visible-object matte per frame (`W×H`, 0/1); empty when the object is out of view. + /// Visible-object matte per frame (`w × h`, 0/1); empty when the object is out of view. pub truth: Vec>, /// Frame the strokes are painted on. pub base: usize, @@ -33,7 +36,7 @@ pub struct Scenario { /// One tracker's result on one scenario. #[derive(Clone, Debug, Serialize)] pub struct Run { - pub scenario: &'static str, + pub scenario: String, pub tracker: String, pub mean_iou: f64, pub min_iou: f64, @@ -46,6 +49,8 @@ pub struct Run { /// Correction strokes needed before every frame reaches the target IoU (`>N`: not within N); set by /// [`with_fixes`], `None` when not measured. pub fixes: Option, + /// The most memory (MB) the tracker held for its memory banks and cached frames (trackers that have none: `None`). + pub peak_mb: Option, /// The matte of every frame (not serialised). #[serde(skip)] pub mattes: Vec>, @@ -136,7 +141,46 @@ fn scenario(name: &'static str, what: &'static str, n: usize, look: impl Fn(usiz frames.push(img); truth.push(gt); } - Scenario { name, what, frames, truth, base: 0, strokes: strokes_for(look(0).c) } + Scenario { name: name.to_string(), what: what.to_string(), w: W as usize, h: H as usize, frames, truth, base: 0, strokes: strokes_for(look(0).c) } +} + +/// A scenario from the user's own clip: frames and their true mattes (0/1, same size as the frames). The strokes +/// are made from the first frame with an object: a short foreground stroke through its centre and background +/// strokes along the top and bottom edges where the truth is empty. Nothing is bundled: the clip stays with the +/// user (a licensed dataset, say), and so does its licence. +pub fn scenario_from(name: &str, frames: Vec, truth: Vec>) -> std::result::Result { + let first = frames.first().ok_or("a clip needs at least one frame")?; + let (w, h) = (first.width as usize, first.height as usize); + if w == 0 + || h == 0 + || frames.len() != truth.len() + || frames.iter().any(|f| (f.width as usize, f.height as usize) != (w, h)) + || truth.iter().any(|t| t.len() != w * h) + { + return Err("frames and truth masks must be the same size and count".into()); + } + let base = truth.iter().position(|t| t.iter().any(|v| *v != 0)).ok_or("no frame has the object in its truth mask")?; + let t = truth.get(base).map(Vec::as_slice).unwrap_or(&[]); + let (mut x0, mut x1, mut sx, mut sy, mut n) = (usize::MAX, 0usize, 0.0f64, 0.0f64, 0.0f64); + for i in t.iter().enumerate().filter(|(_, v)| **v != 0).map(|(i, _)| i) { + let (x, y) = (i % w, i / w); + (x0, x1) = (x0.min(x), x1.max(x)); + sx += x as f64 + 0.5; + sy += y as f64 + 0.5; + n += 1.0; + } + let (cx, cy) = (sx / n, sy / n); + let reach = (((x1 - x0) as f64) * 0.15).max(2.0); + let radius = ((w.min(h) as f64) * 0.03).max(3.0); + let mut strokes = vec![line(StrokeKind::Fg, base as i64, [cx - reach, cy], [cx + reach, cy], radius)]; + for y in [h as f64 * 0.04, h as f64 * 0.96] { + let row = y as usize; + let free = t.get(row * w..(row + 1) * w).is_some_and(|r| r.iter().all(|v| *v == 0)); + if free { + strokes.push(line(StrokeKind::Bg, base as i64, [w as f64 * 0.05, y], [w as f64 * 0.95, y], radius)); + } + } + Ok(Scenario { name: name.to_string(), what: "a clip supplied by the user".to_string(), w, h, frames, truth, base, strokes }) } /// The benchmark clips, `n` frames each (at least 12, so every event happens). @@ -181,7 +225,7 @@ fn metrics(sc: &Scenario, tracker: &str, mattes: &[Vec], secs: f64) -> Run { (!tail.is_empty()).then(|| tail.iter().sum::() / tail.len() as f64) }); Run { - scenario: sc.name, + scenario: sc.name.clone(), tracker: tracker.to_string(), mean_iou: per_frame.iter().sum::() / n, min_iou: per_frame.iter().copied().fold(1.0, f64::min), @@ -190,6 +234,7 @@ fn metrics(sc: &Scenario, tracker: &str, mattes: &[Vec], secs: f64) -> Run { after_return, per_frame, fixes: None, + peak_mb: None, mattes: mattes.to_vec(), } } @@ -217,7 +262,7 @@ pub fn run_flow(sc: &Scenario, tracker: &str, model: Option<&dyn MaskModel>) -> } } let secs = start.elapsed().as_secs_f64(); - let empty = vec![0u8; (W * H) as usize]; + let empty = vec![0u8; sc.w * sc.h]; let mattes: Vec> = segs.into_iter().map(|s| s.map_or_else(|| empty.clone(), |s| s.matte)).collect(); metrics(sc, tracker, &mattes, secs) } @@ -239,10 +284,10 @@ pub fn run_tracker(sc: &Scenario, tracker: &str, model: &dyn MaskModel) -> Optio } let req = TrackRequest::single(0, sc.frames.len(), prompts); let start = Instant::now(); - let mut mattes: Vec> = vec![vec![0u8; (W * H) as usize]; sc.frames.len()]; + let mut mattes: Vec> = vec![vec![0u8; sc.w * sc.h]; sc.frames.len()]; let mut get = |f: i64| usize::try_from(f).ok().and_then(|i| clip.get(i).cloned()); let r = t.track(&req, &mut get, &|| false, &mut |_, f, tf| { - let (Ok(i), true) = (usize::try_from(f), tf.prob.len() == (W * H) as usize) else { return }; + let (Ok(i), true) = (usize::try_from(f), tf.prob.len() == sc.w * sc.h) else { return }; let Some(img) = sc.frames.get(i) else { return }; let strokes = at(i); if let Some(slot) = mattes.get_mut(i) { @@ -250,7 +295,9 @@ pub fn run_tracker(sc: &Scenario, tracker: &str, model: &dyn MaskModel) -> Optio } }); r.ok()?; - Some(metrics(sc, tracker, &mattes, start.elapsed().as_secs_f64())) + let mut run = metrics(sc, tracker, &mattes, start.elapsed().as_secs_f64()); + run.peak_mb = model.peak_state_bytes().map(|b| b as f64 / (1 << 20) as f64); + Some(run) } /// The correction a user would paint on frame `f` where `pred` is wrong: a short foreground stroke through @@ -260,8 +307,8 @@ fn correction(sc: &Scenario, f: usize, pred: &[u8]) -> Vec { let centre = |keep: &dyn Fn(usize) -> bool| -> Option<[f64; 2]> { let (mut sx, mut sy, mut n) = (0.0, 0.0, 0.0); for i in (0..truth.len().min(pred.len())).filter(|i| keep(*i)) { - sx += (i % W as usize) as f64 + 0.5; - sy += (i / W as usize) as f64 + 0.5; + sx += (i % sc.w.max(1)) as f64 + 0.5; + sy += (i / sc.w.max(1)) as f64 + 0.5; n += 1.0; } (n > 0.0).then(|| [sx / n, sy / n]) @@ -307,12 +354,14 @@ pub fn with_fixes(sc: &Scenario, target: f64, max: usize, run: &dyn Fn(&Scenario /// A fixed-width table of the runs, grouped by scenario. pub fn table(runs: &[Run]) -> String { - let mut out = - format!("{:<10} {:<12} {:>8} {:>8} {:>8} {:>10} {:>9} {:>6}\n", "clip", "tracker", "meanIoU", "minIoU", "flicker", "after ret.", "s/frame", "fixes"); + let mut out = format!( + "{:<10} {:<12} {:>8} {:>8} {:>8} {:>10} {:>9} {:>6} {:>7}\n", + "clip", "tracker", "meanIoU", "minIoU", "flicker", "after ret.", "s/frame", "fixes", "peakMB" + ); for r in runs { let back = r.after_return.map_or_else(|| "-".to_string(), |v| format!("{v:.3}")); out.push_str(&format!( - "{:<10} {:<12} {:>8.3} {:>8.3} {:>8.3} {:>10} {:>9.3} {:>6}\n", + "{:<10} {:<12} {:>8.3} {:>8.3} {:>8.3} {:>10} {:>9.3} {:>6} {:>7}\n", r.scenario, r.tracker, r.mean_iou, @@ -320,7 +369,8 @@ pub fn table(runs: &[Run]) -> String { r.flicker, back, r.secs_per_frame, - r.fixes.as_deref().unwrap_or("-") + r.fixes.as_deref().unwrap_or("-"), + r.peak_mb.map_or_else(|| "-".to_string(), |v| format!("{v:.0}")) )); } out @@ -373,6 +423,23 @@ mod tests { assert!(table(&[easy, hard]).contains("fixes")); } + #[test] + fn a_users_clip_becomes_a_scenario_and_bad_ones_are_refused() { + let all = scenarios(24); + let src = all.iter().find(|s| s.name == "move").unwrap(); + let sc = scenario_from("mine", src.frames.clone(), src.truth.clone()).unwrap(); + assert_eq!((sc.w, sc.h, sc.base), (W as usize, H as usize, 0)); + assert!(sc.strokes.iter().any(|s| s.kind == StrokeKind::Fg) && sc.strokes.iter().any(|s| s.kind == StrokeKind::Bg)); + let run = run_flow(&sc, "classical", None); + assert!(run.mean_iou > 0.9, "{}", run.mean_iou); + assert!(scenario_from("x", vec![], vec![]).is_err()); + assert!(scenario_from("x", src.frames.clone(), src.truth[..3].to_vec()).is_err()); + assert!(scenario_from("x", src.frames.clone(), vec![vec![0; sc.w * sc.h]; src.frames.len()]).is_err(), "no object anywhere"); + let mut short = src.truth.clone(); + short[0].pop(); + assert!(scenario_from("x", src.frames.clone(), short).is_err()); + } + #[test] fn tiny_clip_requests_are_clamped() { assert!(scenarios(0).iter().all(|s| s.frames.len() == 12));