diff --git a/Cargo.lock b/Cargo.lock index 23cda13ca..cca32c373 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1205,6 +1205,7 @@ dependencies = [ "effectcraft-keyframe", "effectcraft-project", "effectcraft-raster", + "effectcraft-segment", "effectcraft-track", "log", "rayon", @@ -1229,6 +1230,7 @@ dependencies = [ "effectcraft-psd", "effectcraft-raster", "effectcraft-render", + "effectcraft-segment", "effectcraft-svg", "effectcraft-text", "effectcraft-time", @@ -1541,6 +1543,15 @@ dependencies = [ "serde_json", ] +[[package]] +name = "effectcraft-segment" +version = "0.2.0" +dependencies = [ + "miniz_oxide 0.8.9", + "ndarray", + "rayon", +] + [[package]] name = "effectcraft-svg" version = "0.2.0" @@ -1578,6 +1589,7 @@ name = "effectcraft-track" version = "0.2.0" dependencies = [ "effectcraft-raster", + "effectcraft-segment", "rayon", "serde", "serde_json", @@ -3050,6 +3062,16 @@ dependencies = [ "libc", ] +[[package]] +name = "matrixmultiply" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f607c237553f086e7043417a51df26b2eb899d3caff94e6a67592ff992fedc7" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "memchr" version = "2.8.3" @@ -3167,6 +3189,21 @@ dependencies = [ "thiserror 2.0.21", ] +[[package]] +name = "ndarray" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + [[package]] name = "ndk" version = "0.9.0" @@ -4169,6 +4206,12 @@ dependencies = [ "objc2-quartz-core 0.3.2", ] +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + [[package]] name = "rayon" version = "1.12.0" diff --git a/Cargo.toml b/Cargo.toml index e3f926ffe..13156f35a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,6 +22,7 @@ effectcraft-project = { path = "crates/project" } effectcraft-text = { path = "crates/text" } effectcraft-effects = { path = "crates/effects" } effectcraft-track = { path = "crates/track" } +effectcraft-segment = { path = "crates/segment" } effectcraft-model = { path = "crates/model" } effectcraft-render = { path = "crates/render" } effectcraft-gpu = { path = "crates/gpu" } diff --git a/README.md b/README.md index fa1d6fc69..a59b29d3a 100644 --- a/README.md +++ b/README.md @@ -196,8 +196,9 @@ tests, layering, asset attribution and the WebAssembly build). See [CONTRIBUTING ## Where it stands EffectCraft is young. The core of After Effects is in place and gets deeper every week: what is -done and what comes next is in the [ROADMAP](ROADMAP.md). Still to come: learned models for -Roto Brush and face tracking (both use classical methods today), GPU versions of the remaining +done and what comes next is in the [ROADMAP](ROADMAP.md). Roto Brush can use a trained model +(MobileSAM, Apache-2.0, a 40 MB download you opt into in Settings ▸ Roto Brush). Still to come: a +learned face tracker (face tracking uses classical methods today), GPU versions of the remaining CPU-only effects, the Advanced 3D Extended Viewer, and shared-memory threads in the web build. The web build ([docs/web.md](docs/web.md)) runs the full app in the browser, with renders, analyses and viewer frames in Web Workers. Projects are saved as `.ecproj`, readable versioned JSON; After diff --git a/ROADMAP.md b/ROADMAP.md index 5936aa9a9..e17091605 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -19,7 +19,7 @@ once. | M10 | Export: render queue, H.264, ProRes, image sequences, GIF, audio | Done | | M11 | Animation tools: audio playback, meters and waveforms, Lottie import and export, presets, Motion Sketch, Wiggler, Smoother | Done | | M12 | Performance: layer cache, parallel and GPU compositing, disk cache, motion tracking (done); GPU versions of the remaining CPU-only effects | Mostly done | -| M13 | Puppet tools, paint, Roto Brush, motion tracking, a plug-in API, Timeline depth | Done (Roto Brush and face tracking use classical models; learned models are still to come) | +| M13 | Puppet tools, paint, Roto Brush, motion tracking, a plug-in API, Timeline depth | Done (Roto Brush 2.0 / 3.0 can use a trained model: MobileSAM, open source, downloaded on demand, in a swappable model module; face tracking uses a classical model, a learned one is still to come) | | M15 | The web app (WebAssembly, WebGPU) | Done: browser storage, Web Audio, renders and analyses in Web Workers, viewer frames and GPU effects in frame workers with their own WebGPU devices, the disk cache in the Origin Private File System, a storage manager, offline install ([docs/web.md](docs/web.md)) | | M14 | Built for agents: MCP server, command-line tool, control channel; Settings, keyboard shortcut editor, auto-save and crash recovery ([docs/preferences.md](docs/preferences.md)) | Done | @@ -28,11 +28,11 @@ once. Measured feature by feature in [docs/parity.md](docs/parity.md) (4 October 2026): - **≈ 99% of After Effects' features, weighted by importance**, counting partial features as half - done (≈ 99.8% with per-feature fractions): 89 of 92 features done, 3 partial, none missing; every + done (≈ 99.9% with per-feature fractions): 90 of 92 features done, 2 partial, none missing; every essential (P0) feature done; all 306 effects implemented in full, 280 of them on the GPU. - **≈ 8.5 agent-hours of work remain**: about **2–3 hours of wall-clock time** with five Claude - Opus 5.5 agents in parallel, ≈ 1 hour without learned models for Roto Brush and face tracking. -- What is left: GPU kernels for the remaining visual effects, and learned-model quality for Roto - Brush and face tracking ([parity.md](docs/parity.md#highest-value-gaps-in-order)). + Opus 5.5 agents in parallel (estimated before Roto Brush's trained model landed in M13.35). +- What is left: GPU kernels for the remaining visual effects, and a learned face tracker + ([parity.md](docs/parity.md#highest-value-gaps-in-order)). Come tell us what matters most to you on [Discord](https://discord.gg/artcraft). diff --git a/apps/effectcraft/src/main.rs b/apps/effectcraft/src/main.rs index a1923f826..7a56944e7 100644 --- a/apps/effectcraft/src/main.rs +++ b/apps/effectcraft/src/main.rs @@ -72,7 +72,7 @@ fn main() -> eframe::Result { session.autosave.background = true; // Settings, shortcut presets and the crash-recovery sentinel live in the platform // config directory. Agent-driven runs (`--control`) skip crash recovery. - if let Some(dir) = config_dir() { + if let Some(dir) = effectcraft_host::config_dir() { // WebAssembly effect plug-ins in /Plug-ins load before the menus are built. let plugins = dir.join("Plug-ins"); if plugins.is_dir() { @@ -163,24 +163,6 @@ impl eframe::App for Desktop { } } -/// The platform config directory for EffectCraft (`EFFECTCRAFT_CONFIG_DIR` overrides): -/// `~/Library/Application Support/EffectCraft` (macOS), `%APPDATA%\EffectCraft` (Windows), -/// `$XDG_CONFIG_HOME/effectcraft` or `~/.config/effectcraft` (Linux and others). -fn config_dir() -> Option { - use std::path::PathBuf; - if let Some(d) = std::env::var_os("EFFECTCRAFT_CONFIG_DIR") { - return Some(PathBuf::from(d)); - } - let home = || std::env::var_os("HOME").map(PathBuf::from); - if cfg!(target_os = "macos") { - return home().map(|h| h.join("Library/Application Support/EffectCraft")); - } - if cfg!(target_os = "windows") { - return std::env::var_os("APPDATA").map(|a| PathBuf::from(a).join("EffectCraft")); - } - std::env::var_os("XDG_CONFIG_HOME").map(PathBuf::from).or_else(|| home().map(|h| h.join(".config"))).map(|c| c.join("effectcraft")) -} - /// Driven by an agent (`--control`): don't activate the app on launch, so the user's keyboard /// focus stays where it is (macOS; winit activates ignoring other apps by default). fn agent_event_loop(agent: bool) -> Option { diff --git a/crates/automation/src/tests.rs b/crates/automation/src/tests.rs index f78447b9d..61ec7f8fe 100644 --- a/crates/automation/src/tests.rs +++ b/crates/automation/src/tests.rs @@ -393,3 +393,17 @@ fn high_level_tools_add_effect_and_report_state() { let (c, err) = call(&mut s, "add_effect", json!({"layer": l, "effect": "No Such Effect"})); assert!(err && c[0]["text"].as_str().unwrap().contains("unknown effect"), "{c:?}"); } + +/// Roto Brush's trained models over MCP: listed with their licence, chosen, refused when the +/// file isn't the published one. +#[test] +fn roto_models_over_mcp() { + let mut s = server(); + let list = call_json(&mut s, "execute_command", json!({"command": "roto.models", "params": {}})); + let models = list["models"].as_array().unwrap(); + assert!(models.iter().any(|m| m["id"] == "mobilesam" && m["licence"] == "Apache-2.0" && m["installed"] == false), "{list}"); + let r = call_json(&mut s, "execute_command", json!({"command": "roto.model.select", "params": {"id": "mobilesam"}})); + assert_eq!(r["model"], "mobilesam"); + let (c, err) = call(&mut s, "execute_command", json!({"command": "roto.model.select", "params": {"id": "nope"}})); + assert!(err && c[0]["text"].as_str().unwrap().contains("unknown model"), "{c:?}"); +} diff --git a/crates/effects/Cargo.toml b/crates/effects/Cargo.toml index 9b0653609..c37d09949 100644 --- a/crates/effects/Cargo.toml +++ b/crates/effects/Cargo.toml @@ -15,6 +15,7 @@ effectcraft-color = { workspace = true } effectcraft-keyframe = { workspace = true } effectcraft-project = { workspace = true } effectcraft-raster = { workspace = true } +effectcraft-segment = { workspace = true } effectcraft-track = { workspace = true } rayon = { workspace = true } serde = { workspace = true } diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index 63d261b19..d514e80a1 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -16,15 +16,19 @@ //! decontamination of the edge colours. Outside the segmentation span the layer is transparent; //! without strokes the effect passes its input through. //! -//! *Version* is stored but all versions use the same (classic, graph-cut) engine; *Quality ▸ -//! Best* segments at a higher working resolution with more iterations. +//! *Version*: 1.0 is the classic graph-cut engine; 2.0 and 3.0 also use the trained model chosen +//! in Settings ▸ Roto Brush ([`set_model`], an `effectcraft_segment::MaskModel` such as +//! MobileSAM), falling back to the classic engine when none is installed. The model's id is part +//! of the chain seed, so choosing another recomputes the mattes. *Quality ▸ Best* segments at a +//! higher working resolution with more iterations. use std::collections::{BTreeMap, HashMap, VecDeque}; -use std::sync::{Arc, Mutex, OnceLock}; +use std::sync::{Arc, Mutex, OnceLock, RwLock}; use effectcraft_keyframe::Value; use effectcraft_project::ParamUi; use effectcraft_raster::Image; +use effectcraft_segment::MaskModel; use effectcraft_track::roto::refine::{self, MatteParams}; use effectcraft_track::roto::{self as rb, FrameSeg, RotoData, SegOpts, rle}; use rayon::prelude::*; @@ -173,8 +177,29 @@ pub fn matte_params(params: &Params) -> MatteParams { } } +fn model_slot() -> &'static RwLock>> { + static M: OnceLock>>> = OnceLock::new(); + M.get_or_init(|| RwLock::new(None)) +} + +/// The trained model Roto Brush 2.0 / 3.0 use (Settings ▸ Roto Brush; `None`: the classic +/// engine for every version). +pub fn set_model(m: Option>) { + if let Ok(mut slot) = model_slot().write() { + *slot = m; + } +} + +/// The model an instance segments with: its version's choice of the installed one. +pub fn model_for(params: &Params) -> Option> { + if e(params, "version") == 0 { + return None; + } + model_slot().read().ok().and_then(|m| m.clone()) +} + /// Seed of the chain keys: what the frames are (the engine-maintained Input Key), the layer -/// size, the frame rate and the segmentation settings. +/// size, the frame rate, the segmentation settings and the model. pub fn seed(params: &Params, layer_size: [f64; 2], fps: f64) -> u64 { let mut h = 0xcbf2_9ce4_8422_2325; fnv(&mut h, s(params, INPUT_KEY).as_bytes()); @@ -189,6 +214,7 @@ pub fn seed(params: &Params, layer_size: [f64; 2], fps: f64) -> u64 { fnv(&mut h, &v.to_bits().to_le_bytes()); } fnv(&mut h, &[e(params, "quality") as u8, e(params, "version") as u8]); + fnv(&mut h, model_for(params).map_or(effectcraft_segment::CLASSICAL, |m| m.info().id).as_bytes()); h } @@ -275,13 +301,15 @@ pub struct Chain { pub keys: BTreeMap, pub opts: SegOpts, pub refine_radius: f64, + /// The trained model (Version 2.0 / 3.0 with one installed). + pub model: Option>, } impl Chain { pub fn new(params: &Params, layer_size: [f64; 2], fps: f64, scale: f64) -> Chain { let data = data(params); let keys = data.chain_keys(seed(params, layer_size, fps)); - Chain { refine_radius: data.refine_radius(), data, keys, opts: seg_opts(params, scale) } + Chain { refine_radius: data.refine_radius(), data, keys, opts: seg_opts(params, scale), model: model_for(params) } } /// Frame `f`'s segmentation at this chain's scale: cached, or computed from the nearest @@ -331,7 +359,7 @@ impl Chain { let img = frame(g)?; let strokes = self.data.strokes_at(g); let seg = match (&cur, g == base) { - (_, true) | (None, _) => rb::segment(&img, &strokes, None, &self.opts, self.refine_radius), + (_, true) | (None, _) => rb::segment_with(&img, &strokes, None, &self.opts, self.refine_radius, self.model.as_deref()), (Some(p), false) => { let src = self.data.source_of(g)?; let pimg = match prev_img.take() { @@ -341,7 +369,7 @@ impl Chain { if pimg.width != img.width || pimg.height != img.height { return None; } - rb::propagate(&pimg, p, &img, &strokes, &self.opts, self.refine_radius) + rb::propagate_with(&pimg, p, &img, &strokes, &self.opts, self.refine_radius, self.model.as_deref()) } }; let seg = Arc::new(seg); diff --git a/crates/engine/Cargo.toml b/crates/engine/Cargo.toml index 629c9c21b..68e761ab0 100644 --- a/crates/engine/Cargo.toml +++ b/crates/engine/Cargo.toml @@ -20,6 +20,7 @@ effectcraft-path = { workspace = true } effectcraft-text = { workspace = true } effectcraft-effects = { workspace = true } effectcraft-render = { workspace = true } +effectcraft-segment = { workspace = true } effectcraft-track = { workspace = true } effectcraft-model = { workspace = true } effectcraft-lottie = { workspace = true } diff --git a/crates/engine/src/commands/mod.rs b/crates/engine/src/commands/mod.rs index 54236f41a..edceb2534 100644 --- a/crates/engine/src/commands/mod.rs +++ b/crates/engine/src/commands/mod.rs @@ -133,6 +133,7 @@ pub fn command_specs() -> &'static [CommandSpec] { SPECS.get_or_init(|| { let mut v = Vec::new(); v.extend(file::specs()); + v.extend(crate::roto_models::specs()); v.extend(edit::specs()); v.extend(comp::specs()); v.extend(layer::specs()); diff --git a/crates/engine/src/lib.rs b/crates/engine/src/lib.rs index d0fdfc444..3a5295f96 100644 --- a/crates/engine/src/lib.rs +++ b/crates/engine/src/lib.rs @@ -34,6 +34,7 @@ pub mod psd_import; pub mod remote; pub mod render_queue; pub mod roto; +pub mod roto_models; pub mod scriptui; mod session_settings; pub mod shortcuts; @@ -63,6 +64,7 @@ pub use effectcraft_keyframe as keyframe; pub use effectcraft_project as project; pub use effectcraft_raster as raster; pub use effectcraft_render as render; +pub use effectcraft_segment as segment; pub use effectcraft_text as text; pub use effectcraft_time as time; pub use effectcraft_track as track; @@ -347,6 +349,10 @@ pub struct Session { pub camera_pending: Vec<(ItemId, LayerId, Uid)>, /// The running (or finished, not yet polled) Roto Brush propagation / Freeze. pub roto_job: Option, + /// Roto Brush's trained models (Settings ▸ Roto Brush). + pub roto_models: roto_models::Models, + /// Where installed models live (hosts set it; else `models` next to the settings). + pub models_dir: Option, /// Roto Brush instances edited since their last propagation: (comp, layer, effect uid). pub roto_pending: Vec<(ItemId, LayerId, Uid)>, pub events: Vec, @@ -463,6 +469,8 @@ impl Default for Session { camera_job: None, camera_pending: vec![], roto_job: None, + roto_models: Default::default(), + models_dir: None, roto_pending: vec![], events: vec![], journal: vec![], @@ -1084,3 +1092,5 @@ mod tests_paint; mod tests_panels; #[cfg(test)] mod tests_roto; +#[cfg(test)] +mod tests_roto_models; diff --git a/crates/engine/src/prefs.rs b/crates/engine/src/prefs.rs index a56ce6d61..f8e585c2c 100644 --- a/crates/engine/src/prefs.rs +++ b/crates/engine/src/prefs.rs @@ -209,6 +209,15 @@ page!(ThreeD { realtime_shadows: bool = true, }); +page!( + /// Settings ▸ Roto Brush. + RotoPrefs { + /// The segmentation model Roto Brush 2.0 / 3.0 use: `classical` (built in) or a model id + /// from `effectcraft_segment::MODELS` (once installed). + model: String = "classical".into(), + } +); + page!(Scripting { allow_scripts_write_files: bool = false, warn_executing_files: bool = true, @@ -272,6 +281,7 @@ pub struct Prefs { #[serde(rename = "threeD")] pub three_d: ThreeD, pub scripting: Scripting, + pub roto: RotoPrefs, /// File ▸ Open Recent, newest first. pub recent_projects: Vec, /// File ▸ Import Recent Footage, newest first. @@ -312,6 +322,7 @@ impl Default for Prefs { video: Video::default(), three_d: ThreeD::default(), scripting: Scripting::default(), + roto: RotoPrefs::default(), recent_projects: vec![], recent_footage: vec![], recent_presets: vec![], @@ -709,6 +720,7 @@ pub fn section_key(page: &str) -> Option<&'static str> { "video" => "video", "3d" | "threeD" => "threeD", "scripting" => "scripting", + "roto" => "roto", _ => return None, }) } @@ -735,6 +747,7 @@ pub fn page_id(name: &str) -> Option<&'static str> { "video" | "videopreview" => "video", "3d" | "threed" => "3d", "scripting" | "scriptingexpressions" => "scripting", + "roto" | "rotobrush" | "models" => "roto", _ => return None, }; Some(id) @@ -822,6 +835,8 @@ pub enum Item { /// The browser's storage manager (web app only, `storage.*`): usage and quota, persistent /// storage, Clear buttons. Not shown where there is no browser storage. BrowserStorage, + /// Roto Brush's segmentation models (`roto.models`): licence, size, install, choose. + RotoModels, } /// One page of the Settings dialog. @@ -1096,6 +1111,17 @@ pub fn pages() -> Vec { Note("Expressions use the JavaScript engine. Script errors report their file and line in the Script Console; there is no step debugger."), ], }, + Page { + id: "roto", + title: "Roto Brush", + items: vec![ + Section("Segmentation Model"), + Note( + "Roto Brush 2.0 and 3.0 use the model chosen here (Version 1.0 always uses the classic engine). Models are open source, downloaded only when you ask, and checked against their published SHA-256.", + ), + RotoModels, + ], + }, ] } @@ -1148,6 +1174,7 @@ pub fn pages_json() -> Value { Item::AudioDevices { key } => json!({"type": "device", "key": key, "label": "Default Output"}), Item::Note(t) => json!({"note": t}), Item::BrowserStorage => json!({"browserStorage": "storage.info / storage.persist / storage.clear (web app)"}), + Item::RotoModels => json!({"rotoModels": "roto.models / roto.model.select / roto.model.download / roto.model.install / roto.model.remove", "key": "roto.model"}), }) .collect(); json!({"id": p.id, "title": p.title, "items": items}) diff --git a/crates/engine/src/roto.rs b/crates/engine/src/roto.rs index c5beda06d..476254da5 100644 --- a/crates/engine/src/roto.rs +++ b/crates/engine/src/roto.rs @@ -636,6 +636,7 @@ impl Session { /// Frontends call this every frame: finishes jobs and (with `auto`) propagates edited /// instances in the background. pub fn poll_roto(&mut self, auto: bool) -> bool { + self.poll_roto_models(); let changed = self.poll_roto_job(); if auto && self.roto_job.is_none() && self.offloaded_roto().is_none() && self.state.roto.auto_propagate { while let Some((c, l, e)) = self.roto_pending.first().copied() { diff --git a/crates/engine/src/roto_models.rs b/crates/engine/src/roto_models.rs new file mode 100644 index 000000000..ef4608406 --- /dev/null +++ b/crates/engine/src/roto_models.rs @@ -0,0 +1,283 @@ +//! Roto Brush's trained models (Settings ▸ Roto Brush, `roto.model*` commands): which models the +//! registry offers (`effectcraft_segment::MODELS`, all open source), which are installed (in the +//! `models` folder next to the settings), installing one from a file or downloading the official +//! weights (verified against the registry's SHA-256 either way), choosing one, and loading it in +//! the background into the Roto Brush effect (`effects::roto::set_model`). +//! +//! Weights are never bundled: the install stays small and nothing is fetched unless asked for. +//! Downloads use the system's `curl` (shipped with Windows 10+, macOS and most Linux systems), +//! so no networking code is linked into EffectCraft; without it, the file can be downloaded in a +//! browser and installed with Install from File. The browser build segments with the built-in +//! engine only, for now. + +use std::path::{Path, PathBuf}; +use std::sync::{Arc, Mutex}; + +use effectcraft_segment::{self as seg, ModelInfo}; +use serde_json::{Value, json}; + +use crate::commands::{CommandSpec, always, bad, str_p}; +use crate::{EngineError, Result, Session, cmd, query}; + +/// What the background worker reports. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct Status { + /// What is happening now ("Downloading MobileSAM…"). + pub busy: Option, + /// The model loaded into Roto Brush. + pub active: Option, + /// The last failure, for the Settings page. + pub error: Option, + /// A finished job's message, for a toast (taken by [`Session::poll_roto_models`]). + pub done: Option, +} + +/// The session's model state. +#[derive(Default)] +pub struct Models { + pub status: Arc>, +} + +impl Models { + pub fn status(&self) -> Status { + self.status.lock().map(|s| s.clone()).unwrap_or_default() + } +} + +fn set(status: &Arc>, f: impl FnOnce(&mut Status)) { + if let Ok(mut s) = status.lock() { + f(&mut s); + } +} + +impl Session { + /// The folder installed models live in: the host's, else `models` next to the settings. + pub fn models_dir(&self) -> Option { + self.models_dir.clone().or_else(|| Some(self.config.as_ref()?.dir()?.join("models"))) + } + + /// Where `m` is installed, if it is (a file of the published size). + pub fn model_path(&self, m: &ModelInfo) -> Option { + let p = self.models_dir()?.join(m.file_name); + std::fs::metadata(&p).ok().filter(|md| md.len() == m.size).map(|_| p) + } + + /// Load (or drop) the model Settings ▸ Roto Brush chooses; called whenever settings change. + pub fn apply_roto_model(&mut self) { + let want = self.prefs.roto.model.clone(); + let st = self.roto_models.status(); + if want == seg::CLASSICAL || seg::info(&want).is_none() { + if st.active.is_some() { + crate::effects::roto::set_model(None); + set(&self.roto_models.status, |s| s.active = None); + } + return; + } + if st.active.as_deref() == Some(want.as_str()) || st.busy.is_some() { + return; + } + let Some(info) = seg::info(&want) else { return }; + let Some(path) = self.model_path(info) else { + crate::effects::roto::set_model(None); + set(&self.roto_models.status, |s| s.active = None); + return; + }; + spawn(&self.roto_models.status, format!("Loading {}…", info.name), move |status| { + let m = load_file(&path, info.id)?; + crate::effects::roto::set_model(Some(m)); + set(status, |s| s.active = Some(info.id.to_string())); + Ok(format!("Roto Brush: {} loaded", info.name)) + }); + } + + /// Finish what the model worker did: a toast, and loading a model that just arrived. + pub fn poll_roto_models(&mut self) { + let done = self.roto_models.status.lock().ok().and_then(|mut s| s.done.take()); + if let Some(msg) = done { + let error = msg.starts_with("Roto Brush model:"); + self.events.push(crate::Event::Toast { message: msg, error }); + self.apply_roto_model(); + } + } +} + +fn load_file(path: &Path, id: &str) -> std::result::Result, String> { + let bytes = std::fs::read(path).map_err(|e| format!("cannot read {}: {e}", path.display()))?; + seg::load(id, &bytes) +} + +/// Run `job` on a worker thread, reporting through `status` (desktop; the browser build has no +/// trained models yet). +fn spawn(status: &Arc>, what: String, job: impl FnOnce(&Arc>) -> std::result::Result + Send + 'static) { + set(status, |s| { + s.busy = Some(what); + s.error = None; + }); + let st = status.clone(); + let run = move || { + let r = job(&st); + set(&st, |s| { + s.busy = None; + match r { + Ok(msg) => s.done = Some(msg), + Err(e) => { + s.done = Some(format!("Roto Brush model: {e}")); + s.error = Some(e); + } + } + }); + }; + if cfg!(target_arch = "wasm32") { + set(status, |s| { + s.busy = None; + s.error = Some("trained models are available in the desktop app".into()); + }); + return; + } + if let Err(e) = std::thread::Builder::new().name("ec-roto-model".into()).spawn(run) { + set(status, |s| { + s.busy = None; + s.error = Some(format!("cannot start: {e}")); + }); + } +} + +/// Copy verified weights into the models folder (atomically). +fn install_bytes(dir: &Path, m: &ModelInfo, bytes: &[u8]) -> std::result::Result { + if bytes.len() as u64 != m.size || seg::sha256::hex(bytes) != m.sha256 { + return Err(format!("this is not the published {} file (size or SHA-256 differs)", m.name)); + } + let path = dir.join(m.file_name); + crate::config::FileOps::write(&crate::config::StdFiles, &path, bytes).map_err(|e| format!("cannot write {}: {e}", path.display()))?; + Ok(path) +} + +fn model_p(p: &Value, cmd: &str) -> Result<&'static ModelInfo> { + let id = str_p(p, "id").ok_or_else(|| bad(cmd, "missing `id`"))?; + seg::info(id).ok_or_else(|| bad(cmd, format!("unknown model `{id}` (see roto.models)"))) +} + +fn list(s: &mut Session, _: &Value) -> Result { + let st = s.roto_models.status(); + let chosen = s.prefs.roto.model.clone(); + let mut v = vec![json!({ + "id": seg::CLASSICAL, + "name": "Classic (graph cut)", + "description": "Built in: colour models and graph cuts, propagated by optical flow. Always available.", + "licence": "MIT OR Apache-2.0", + "installed": true, + "selected": chosen == seg::CLASSICAL, + "active": st.active.is_none(), + })]; + for m in seg::MODELS { + v.push(json!({ + "id": m.id, + "name": m.name, + "description": m.description, + "licence": m.licence, + "licenceUrl": m.licence_url, + "homepage": m.homepage, + "url": m.url, + "size": m.size, + "sha256": m.sha256, + "installed": s.model_path(m).is_some(), + "selected": chosen == m.id, + "active": st.active.as_deref() == Some(m.id), + })); + } + Ok(json!({"models": v, "busy": st.busy, "error": st.error, "folder": s.models_dir().map(|d| d.display().to_string())})) +} + +fn select(s: &mut Session, p: &Value) -> Result { + let id = str_p(p, "id").ok_or_else(|| bad("roto.model.select", "missing `id`"))?; + if id != seg::CLASSICAL && seg::info(id).is_none() { + return Err(bad("roto.model.select", format!("unknown model `{id}` (see roto.models)"))); + } + s.prefs.roto.model = id.to_string(); + s.save_prefs(); + s.prefs_changed(); + Ok(json!({"model": id})) +} + +fn install(s: &mut Session, p: &Value) -> Result { + let cmd = "roto.model.install"; + let path = str_p(p, "path").ok_or_else(|| bad(cmd, "missing `path` (the downloaded weights file)"))?; + let dir = s.models_dir().ok_or_else(|| bad(cmd, "no settings folder to install into"))?; + let bytes = std::fs::read(path).map_err(|e| EngineError::Other(format!("cannot read {path}: {e}")))?; + // Which model it is: the one asked for, or the one whose checksum it has. + let m = match str_p(p, "id") { + Some(_) => model_p(p, cmd)?, + None => { + let h = seg::sha256::hex(&bytes); + seg::MODELS.iter().find(|m| m.sha256 == h).ok_or_else(|| bad(cmd, "this file is not a published model in the registry (SHA-256 differs)"))? + } + }; + let dest = install_bytes(&dir, m, &bytes).map_err(|e| bad(cmd, e))?; + s.apply_roto_model(); + Ok(json!({"id": m.id, "path": dest.display().to_string()})) +} + +fn download(s: &mut Session, p: &Value) -> Result { + let cmd = "roto.model.download"; + let m = model_p(p, cmd)?; + let dir = s.models_dir().ok_or_else(|| bad(cmd, "no settings folder to install into"))?; + if s.roto_models.status().busy.is_some() { + return Err(bad(cmd, "a model is already downloading or loading")); + } + spawn(&s.roto_models.status, format!("Downloading {} ({} MB)…", m.name, m.size / 1_000_000), move |_| { + std::fs::create_dir_all(&dir).map_err(|e| format!("cannot create {}: {e}", dir.display()))?; + let tmp = dir.join(format!("{}.part", m.file_name)); + let out = std::process::Command::new("curl") + .args(["-fsSL", "--retry", "2", "--connect-timeout", "20", "-o"]) + .arg(&tmp) + .arg(m.url) + .output() + .map_err(|e| format!("downloading needs curl ({e}); or download {} in a browser and use Install from File", m.url))?; + if !out.status.success() { + let _ = std::fs::remove_file(&tmp); + return Err(format!("download failed: {}", String::from_utf8_lossy(&out.stderr).trim())); + } + let bytes = std::fs::read(&tmp).map_err(|e| format!("cannot read the download: {e}"))?; + let _ = std::fs::remove_file(&tmp); + install_bytes(&dir, m, &bytes)?; + Ok(format!("Roto Brush: {} downloaded and verified", m.name)) + }); + Ok(json!({"id": m.id, "downloading": true})) +} + +fn remove(s: &mut Session, p: &Value) -> Result { + let m = model_p(p, "roto.model.remove")?; + let Some(path) = s.model_path(m) else { return Ok(json!({"removed": false})) }; + if s.roto_models.status().active.as_deref() == Some(m.id) { + crate::effects::roto::set_model(None); + set(&s.roto_models.status, |st| st.active = None); + } + std::fs::remove_file(&path).map_err(|e| EngineError::Other(format!("cannot remove {}: {e}", path.display())))?; + Ok(json!({"removed": true})) +} + +pub fn specs() -> Vec { + vec![ + query!("roto.models", "Roto Brush Models", "{}", list), + cmd!("roto.model.select", "Use Roto Brush Model", [], None, "{id: classical|mobilesam} (Roto Brush 2.0 / 3.0 use it)", always, select), + cmd!( + "roto.model.install", + "Install Roto Brush Model", + [], + None, + "{path: weights file, id?} — verified against the registry's SHA-256", + always, + install + ), + cmd!( + "roto.model.download", + "Download Roto Brush Model", + [], + None, + "{id} — the official weights, verified (desktop; uses the system curl)", + always, + download + ), + cmd!("roto.model.remove", "Remove Roto Brush Model", [], None, "{id}", always, remove), + ] +} diff --git a/crates/engine/src/session_settings.rs b/crates/engine/src/session_settings.rs index 4afd9691d..6f1eef7ed 100644 --- a/crates/engine/src/session_settings.rs +++ b/crates/engine/src/session_settings.rs @@ -38,6 +38,8 @@ impl Session { let levels = self.prefs.general.undo_levels.max(1) as usize; let current = self.project.clone(); self.history.trim(levels, ¤t); + // Settings ▸ Roto Brush: load (or drop) the chosen trained model. + self.apply_roto_model(); self.prefs_revision += 1; } diff --git a/crates/engine/src/tests_roto_models.rs b/crates/engine/src/tests_roto_models.rs new file mode 100644 index 000000000..8977e7242 --- /dev/null +++ b/crates/engine/src/tests_roto_models.rs @@ -0,0 +1,57 @@ +//! Settings ▸ Roto Brush: the model registry, choosing, installing (verified) and loading a +//! trained model into Roto Brush. + +use std::sync::Arc; + +use serde_json::json; + +use crate::Session; +use crate::config::{ConfigStore, DirConfig}; + +fn session() -> (Session, std::path::PathBuf) { + let d = std::env::temp_dir().join(format!("ec-roto-models-{}-{}", std::process::id(), line!())); + let store: Arc = Arc::new(DirConfig::new(d.join("config"))); + (Session { config: Some(store), ..Default::default() }, d) +} + +#[test] +fn models_are_listed_chosen_and_verified() { + let (mut s, d) = session(); + let list = s.execute_checked("roto.models", json!({})).unwrap(); + let ids: Vec<&str> = list["models"].as_array().unwrap().iter().map(|m| m["id"].as_str().unwrap()).collect(); + assert_eq!(ids, ["classical", "mobilesam"]); + let sam = &list["models"][1]; + assert_eq!((sam["licence"].as_str(), sam["installed"].as_bool()), (Some("Apache-2.0"), Some(false))); + assert!(sam["url"].as_str().unwrap().starts_with("https://")); + // Unknown ids and files that aren't the published weights are refused. + assert!(s.execute_checked("roto.model.select", json!({"id": "nope"})).is_err()); + let bogus = d.join("bogus.pt"); + std::fs::create_dir_all(&d).unwrap(); + std::fs::write(&bogus, b"not the weights").unwrap(); + let e = s.execute_checked("roto.model.install", json!({"path": bogus.to_string_lossy()})).unwrap_err().to_string(); + assert!(e.contains("SHA-256"), "{e}"); + assert!(s.execute_checked("roto.model.install", json!({"path": bogus.to_string_lossy(), "id": "mobilesam"})).is_err()); + // Choosing a model that isn't installed keeps the classic engine. + s.execute_checked("roto.model.select", json!({"id": "mobilesam"})).unwrap(); + assert_eq!(s.prefs.roto.model, "mobilesam"); + assert!(s.roto_models.status().active.is_none()); + assert!(crate::prefs::pages().iter().any(|p| p.id == "roto")); + // With the real weights (EFFECTCRAFT_MOBILESAM=path/to/mobile_sam.pt): install, load, use. + if let Ok(path) = std::env::var("EFFECTCRAFT_MOBILESAM") { + s.execute_checked("roto.model.install", json!({"path": path})).unwrap(); + for _ in 0..200 { + s.poll_roto_models(); + if s.roto_models.status().active.is_some() { + break; + } + std::thread::sleep(std::time::Duration::from_millis(20)); + } + assert_eq!(s.roto_models.status().active.as_deref(), Some("mobilesam")); + // Roto Brush 3.0 (the default version) now segments with it. + let params = crate::effects::Params { values: [("version".to_string(), effectcraft_keyframe::Value::Enum(2))].into_iter().collect() }; + assert_eq!(crate::effects::roto::model_for(¶ms).map(|m| m.info().id), Some("mobilesam")); + s.execute_checked("roto.model.select", json!({"id": "classical"})).unwrap(); + assert!(crate::effects::roto::model_for(¶ms).is_none()); + } + let _ = std::fs::remove_dir_all(d); +} diff --git a/crates/host/src/lib.rs b/crates/host/src/lib.rs index be26b169a..54e2f602c 100644 --- a/crates/host/src/lib.rs +++ b/crates/host/src/lib.rs @@ -81,6 +81,8 @@ impl Exporter for FileExporter { /// A new session with media, import, expressions, scripting and export enabled. pub fn session() -> Session { Session { + // The desktop app, the CLI and the MCP server share installed Roto Brush models. + models_dir: config_dir().map(|d| d.join("models")), exporter: Some(Arc::new(FileExporter::default())), footage: Arc::new(effectcraft_media::MediaPool::new()), importer: Some(Arc::new(MediaImporter)), @@ -93,6 +95,24 @@ pub fn session() -> Session { } } +/// The platform config directory for EffectCraft (`EFFECTCRAFT_CONFIG_DIR` overrides): +/// `~/Library/Application Support/EffectCraft` (macOS), `%APPDATA%\EffectCraft` (Windows), +/// `$XDG_CONFIG_HOME/effectcraft` or `~/.config/effectcraft` (Linux and others). +pub fn config_dir() -> Option { + use std::path::PathBuf; + if let Some(d) = std::env::var_os("EFFECTCRAFT_CONFIG_DIR") { + return Some(PathBuf::from(d)); + } + let home = || std::env::var_os("HOME").map(PathBuf::from); + if cfg!(target_os = "macos") { + return home().map(|h| h.join("Library/Application Support/EffectCraft")); + } + if cfg!(target_os = "windows") { + return std::env::var_os("APPDATA").map(|a| PathBuf::from(a).join("EffectCraft")); + } + std::env::var_os("XDG_CONFIG_HOME").map(PathBuf::from).or_else(|| home().map(|h| h.join(".config"))).map(|c| c.join("effectcraft")) +} + #[cfg(test)] mod tests { use serde_json::json; diff --git a/crates/segment/Cargo.toml b/crates/segment/Cargo.toml new file mode 100644 index 000000000..b83e749f9 --- /dev/null +++ b/crates/segment/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "effectcraft-segment" +description = "Swappable trained segmentation models for Roto Brush: the MaskModel interface, a registry of open-licensed models, a PyTorch checkpoint reader and a pure-Rust MobileSAM (TinyViT + SAM prompt encoder and mask decoder)" +version.workspace = true +edition.workspace = true +license.workspace = true +rust-version.workspace = true + +[lints] +workspace = true + +[dependencies] +miniz_oxide = "0.8" +ndarray = { version = "0.17", default-features = false, features = ["std"] } +rayon = { workspace = true } diff --git a/crates/segment/README.md b/crates/segment/README.md new file mode 100644 index 000000000..7a533392f --- /dev/null +++ b/crates/segment/README.md @@ -0,0 +1,38 @@ +# effectcraft-segment + +Trained segmentation models for Roto Brush, behind one swappable interface. + +- **`MaskModel`**: what Roto Brush asks of a model: given a frame and prompts (foreground / + background points, a box, a prior mask), the foreground probability of every pixel. Roto Brush + (`effectcraft-track`) only talks to this trait, so models can be added or swapped without + touching it. The classical graph-cut segmenter stays built in as the fallback. +- **`MODELS`**: the registry. Every entry must be open source under a licence compatible with + EffectCraft's MIT OR Apache-2.0 (`ALLOWED_LICENCES`), with its homepage, licence URL, the + official weights URL, size and SHA-256. `load` refuses anything else. +- **Weights are never bundled.** They are downloaded on demand (or installed from a file) by the + engine (`roto.model.*`, Settings ▸ Roto Brush) and verified against the registry's SHA-256 + (`sha256`, dependency-free). +- **`pt`**: reads PyTorch checkpoints (`torch.save` zip archives and the pickle subset state dicts + use), so official weights load exactly as published. +- **`mobilesam`**: MobileSAM in plain Rust: the TinyViT-5M image encoder and Segment Anything's + prompt encoder and two-way-transformer mask decoder, on `nn` kernels (rayon, and ndarray's safe + `general_mat_mul`, which picks AVX/FMA kernels at run time). No ML framework, no unsafe code. + +## Models + +| id | Model | Licence | Weights | Size | +|---|---|---|---|---| +| `mobilesam` | [MobileSAM](https://github.com/ChaoningZhang/MobileSAM) (C. Zhang et al., 2023): Segment Anything with a 5M-parameter TinyViT encoder | Apache-2.0 | `mobile_sam.pt` from the official repository | 40.7 MB | + +The implementation follows the published architectures (MobileSAM, TinyViT and Segment Anything +papers and their Apache-2.0 reference code); its output was checked against an independent +implementation (Hugging Face candle's, MIT/Apache-2.0) on the same weights: embeddings within 0.3% +RMS (exact vs. tanh GELU) and 99.9% of mask pixels equal. + +## Adding a model + +1. Pick a model whose code **and** weights are open source under an allowed licence. +2. Add a `ModelInfo` to `MODELS` (pin the official file's SHA-256 and size). +3. Implement `MaskModel` for it and add its loader to `load`. +4. Test it against a reference implementation; Roto Brush, Settings and the `roto.model.*` + commands pick it up from the registry. diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs new file mode 100644 index 000000000..9db6069a8 --- /dev/null +++ b/crates/segment/src/lib.rs @@ -0,0 +1,211 @@ +//! Trained segmentation models for Roto Brush, behind one swappable interface. +//! +//! - [`MaskModel`]: what Roto Brush asks of a model, the foreground probability of each pixel of a +//! frame given prompts (foreground / background points from the strokes, a box, a prior mask). +//! Anything implementing it can be plugged in; the classical graph-cut segmenter in +//! `effectcraft-track` stays the built-in fallback when no model is chosen. +//! - [`MODELS`]: the registry. Every entry is open source under a licence compatible with +//! EffectCraft's (MIT OR Apache-2.0), with its source, size and SHA-256. Weights are never +//! bundled: they are fetched on demand (or installed from a file), verified ([`sha256`]) and +//! loaded by [`load`]. +//! - [`pt`]: a PyTorch checkpoint reader, so official weights load as published. +//! - [`mobilesam`]: MobileSAM (Apache-2.0) in plain Rust ([`nn`] kernels on rayon). +#![forbid(unsafe_code)] +#![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::todo, clippy::unreachable)] + +pub mod mobilesam; +pub mod nn; +pub mod pt; + +use std::sync::Arc; + +pub type Result = std::result::Result; + +/// What a model is asked: coordinates in frame pixels. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct Prompt { + /// Points and whether each is foreground (`true`) or background. + pub points: Vec<([f32; 2], bool)>, + /// A box around the object: `[x0, y0, x1, y1]`. + pub bbox: Option<[f32; 4]>, + /// A prior: the foreground probability of every pixel (`w×h`), e.g. the previous frame's + /// matte carried forward. + pub mask: Option>, +} + +/// A trained segmentation model (swappable: Roto Brush only uses this interface). +pub trait MaskModel: Send + Sync { + fn info(&self) -> &'static ModelInfo; + /// Foreground probability (0–1) of each pixel of a frame (straight RGB 0–1, row-major `w×h`). + fn segment(&self, rgb: &[[f32; 3]], w: usize, h: usize, prompt: &Prompt) -> Result>; +} + +/// A model in the registry: what it is, its licence and where its weights come from. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct ModelInfo { + pub id: &'static str, + pub name: &'static str, + pub description: &'static str, + /// SPDX licence of the weights and the architecture. + pub licence: &'static str, + pub licence_url: &'static str, + /// The project's page. + pub homepage: &'static str, + /// The official weights file. + pub url: &'static str, + pub file_name: &'static str, + pub size: u64, + /// SHA-256 of the weights file (lower-case hex). + pub sha256: &'static str, +} + +/// Licences a model may have to be listed (open source, compatible with MIT OR Apache-2.0). +pub const ALLOWED_LICENCES: &[&str] = &["MIT", "Apache-2.0", "BSD-2-Clause", "BSD-3-Clause", "ISC", "Zlib", "CC0-1.0"]; + +pub const MOBILE_SAM: ModelInfo = ModelInfo { + id: "mobilesam", + name: "MobileSAM", + description: "Segment Anything distilled to a 5M-parameter TinyViT encoder (C. Zhang et al., 2023). Turns your strokes into a clean matte and tidies every propagated frame.", + licence: "Apache-2.0", + licence_url: "https://github.com/ChaoningZhang/MobileSAM/blob/master/LICENSE", + homepage: "https://github.com/ChaoningZhang/MobileSAM", + url: "https://raw.githubusercontent.com/ChaoningZhang/MobileSAM/master/weights/mobile_sam.pt", + file_name: "mobile_sam.pt", + size: 40_728_226, + sha256: "6dbb90523a35330fedd7f1d3dfc66f995213d81b29a5ca8108dbcdd4e37d6c2f", +}; + +/// The models Roto Brush can use besides the built-in classical segmenter. +pub const MODELS: &[ModelInfo] = &[MOBILE_SAM]; + +/// The id of the built-in classical (graph cut) segmenter: no weights, always available. +pub const CLASSICAL: &str = "classical"; + +pub fn info(id: &str) -> Option<&'static ModelInfo> { + MODELS.iter().find(|m| m.id == id) +} + +/// Check a weights file against the registry and build the model. +pub fn load(id: &str, bytes: &[u8]) -> Result> { + let info = info(id).ok_or_else(|| format!("unknown model `{id}`"))?; + if !ALLOWED_LICENCES.contains(&info.licence) { + return Err(format!("{}: licence {} is not allowed", info.name, info.licence)); + } + let got = sha256::hex(bytes); + if got != info.sha256 { + return Err(format!("{}: the file does not match (SHA-256 {got}, expected {})", info.name, info.sha256)); + } + match id { + "mobilesam" => { + let w = pt::read(bytes)?; + Ok(Arc::new(Sam { net: mobilesam::MobileSam::from_weights(&w)? })) + } + _ => Err(format!("no loader for `{id}`")), + } +} + +struct Sam { + net: mobilesam::MobileSam, +} + +impl MaskModel for Sam { + fn info(&self) -> &'static ModelInfo { + &MOBILE_SAM + } + fn segment(&self, rgb: &[[f32; 3]], w: usize, h: usize, prompt: &Prompt) -> Result> { + let e = self.net.embed(rgb, w, h)?; + Ok(self.net.predict(&e, prompt).0) + } +} + +/// SHA-256 (FIPS 180-4), to verify downloaded weights without another dependency. +pub mod sha256 { + const K: [u32; 64] = [ + 0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5, 0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, + 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174, 0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da, + 0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967, 0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, + 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85, 0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070, + 0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3, 0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, + 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2, + ]; + + pub fn digest(data: &[u8]) -> [u8; 32] { + let mut h: [u32; 8] = [0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19]; + let bits = (data.len() as u64).wrapping_mul(8); + let mut tail = data[data.len() - data.len() % 64..].to_vec(); + tail.push(0x80); + while tail.len() % 64 != 56 { + tail.push(0); + } + tail.extend_from_slice(&bits.to_be_bytes()); + for block in data[..data.len() - data.len() % 64].chunks_exact(64).chain(tail.chunks_exact(64)) { + let mut w = [0u32; 64]; + for (i, c) in block.chunks_exact(4).enumerate() { + w[i] = u32::from_be_bytes([c[0], c[1], c[2], c[3]]); + } + for i in 16..64 { + let s0 = w[i - 15].rotate_right(7) ^ w[i - 15].rotate_right(18) ^ (w[i - 15] >> 3); + let s1 = w[i - 2].rotate_right(17) ^ w[i - 2].rotate_right(19) ^ (w[i - 2] >> 10); + w[i] = w[i - 16].wrapping_add(s0).wrapping_add(w[i - 7]).wrapping_add(s1); + } + let [mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut hh] = h; + for i in 0..64 { + let s1 = e.rotate_right(6) ^ e.rotate_right(11) ^ e.rotate_right(25); + let ch = (e & f) ^ (!e & g); + let t1 = hh.wrapping_add(s1).wrapping_add(ch).wrapping_add(K[i]).wrapping_add(w[i]); + let s0 = a.rotate_right(2) ^ a.rotate_right(13) ^ a.rotate_right(22); + let maj = (a & b) ^ (a & c) ^ (b & c); + let t2 = s0.wrapping_add(maj); + hh = g; + g = f; + f = e; + e = d.wrapping_add(t1); + d = c; + c = b; + b = a; + a = t1.wrapping_add(t2); + } + for (x, y) in h.iter_mut().zip([a, b, c, d, e, f, g, hh]) { + *x = x.wrapping_add(y); + } + } + let mut out = [0u8; 32]; + for (i, v) in h.iter().enumerate() { + out[i * 4..i * 4 + 4].copy_from_slice(&v.to_be_bytes()); + } + out + } + + pub fn hex(data: &[u8]) -> String { + digest(data).iter().map(|b| format!("{b:02x}")).collect() + } + + #[cfg(test)] + mod tests { + #[test] + fn known_vectors() { + assert_eq!(super::hex(b""), "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"); + assert_eq!(super::hex(b"abc"), "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"); + let long = b"abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq"; + assert_eq!(super::hex(long), "248d6a61d20638b8e5c026930c3e6039a33ce45964ff2167f6ecedd419db06c1"); + assert_eq!(super::hex(&vec![b'a'; 1_000_000]), "cdc76e5c9914fb9281a1c7e284d73e67f1809a48a497200e046d39ccc7112cd0"); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn registry_is_open_source_and_verifies_files() { + for m in MODELS { + assert!(ALLOWED_LICENCES.contains(&m.licence), "{}", m.id); + assert_eq!(m.sha256.len(), 64); + assert!(m.url.starts_with("https://")); + } + // A file that isn't the published one is refused before anything is parsed. + assert!(load("mobilesam", b"not the weights").err().is_some_and(|e| e.contains("does not match"))); + assert!(load("nope", b"").is_err()); + } +} diff --git a/crates/segment/src/mobilesam.rs b/crates/segment/src/mobilesam.rs new file mode 100644 index 000000000..22d98ba7c --- /dev/null +++ b/crates/segment/src/mobilesam.rs @@ -0,0 +1,677 @@ +//! MobileSAM (C. Zhang et al., "Faster Segment Anything: Towards Lightweight SAM for Mobile +//! Applications", 2023; Apache-2.0) in plain Rust: the TinyViT-5M image encoder (Wu et al., +//! "TinyViT", ECCV 2022) feeding Segment Anything's prompt encoder and two-way-transformer mask +//! decoder (Kirillov et al., "Segment Anything", ICCV 2023), loaded straight from the official +//! `mobile_sam.pt` checkpoint. The image goes in at 1024 pixels on its long side; the encoder +//! gives a 64×64×256 embedding, cached per image so that new prompts on the same frame only run +//! the (fast) decoder. + +use std::f32::consts::PI; +use std::sync::Mutex; + +use rayon::prelude::*; + +use crate::nn::{self, Conv, Depthwise, Linear, Norm, Upconv}; +use crate::pt::Weights; +use crate::{Prompt, Result}; + +const IMG: usize = 1024; +const EMB: usize = 64; +const DIM: usize = 256; +const LOW: usize = 256; +const MEAN: [f32; 3] = [123.675, 116.28, 103.53]; +const STD: [f32; 3] = [58.395, 57.12, 57.375]; + +/// Weight lookup by name with the expected element count. +struct W<'a>(&'a Weights); + +impl W<'_> { + fn get(&self, name: &str, n: usize) -> Result<&[f32]> { + let t = self.0.get(name).ok_or_else(|| format!("checkpoint is missing {name}"))?; + if t.data.len() != n { + return Err(format!("{name}: expected {n} values, found {}", t.data.len())); + } + Ok(&t.data) + } + fn linear(&self, p: &str, out: usize, inp: usize, bias: bool) -> Result { + let w = self.get(&format!("{p}.weight"), out * inp)?; + let b = if bias { Some(self.get(&format!("{p}.bias"), out)?) } else { None }; + Linear::new(w, b, out, inp).ok_or_else(|| format!("{p}: bad shape")) + } + fn norm(&self, p: &str, c: usize, eps: f32) -> Result { + Ok(Norm { w: self.get(&format!("{p}.weight"), c)?.to_vec(), b: self.get(&format!("{p}.bias"), c)?.to_vec(), eps }) + } + /// A conv + batch norm (eval) folded into one weight and bias: `[out × inp/groups × k × k]`. + fn conv_bn(&self, p: &str, out: usize, per: usize) -> Result<(Vec, Vec)> { + let w = self.get(&format!("{p}.c.weight"), out * per)?; + let g = self.get(&format!("{p}.bn.weight"), out)?; + let b = self.get(&format!("{p}.bn.bias"), out)?; + let m = self.get(&format!("{p}.bn.running_mean"), out)?; + let v = self.get(&format!("{p}.bn.running_var"), out)?; + let mut wf = w.to_vec(); + let mut bf = vec![0.0; out]; + for o in 0..out { + let s = g[o] / (v[o] + 1e-5).sqrt(); + wf[o * per..(o + 1) * per].iter_mut().for_each(|x| *x *= s); + bf[o] = b[o] - m[o] * s; + } + Ok((wf, bf)) + } + fn pointwise_bn(&self, p: &str, out: usize, inp: usize) -> Result { + let (w, b) = self.conv_bn(p, out, inp)?; + Linear::new(&w, Some(&b), out, inp).ok_or_else(|| format!("{p}: bad shape")) + } + fn depthwise_bn(&self, p: &str, c: usize, stride: usize) -> Result { + let (w, b) = self.conv_bn(p, c, 9)?; + Depthwise::new(&w, Some(&b), c, stride).ok_or_else(|| format!("{p}: bad shape")) + } + fn conv3_bn(&self, p: &str, out: usize, inp: usize, stride: usize) -> Result { + let (w, b) = self.conv_bn(p, out, inp * 9)?; + Conv::new(&w, Some(&b), out, inp, 3, stride, 1).ok_or_else(|| format!("{p}: bad shape")) + } +} + +/// MBConv: 1×1 expand, depthwise 3×3, 1×1 project, residual. +struct MbConv { + c1: Linear, + c2: Depthwise, + c3: Linear, +} + +impl MbConv { + fn forward(&self, x: &[f32], h: usize, w: usize) -> Vec { + let mut t = self.c1.forward(x, h * w); + nn::gelu(&mut t); + let (mut t, _, _) = self.c2.forward(&t, h, w); + nn::gelu(&mut t); + let mut t = self.c3.forward(&t, h * w); + nn::add(&mut t, x); + nn::gelu(&mut t); + t + } +} + +/// Patch merging between stages: 1×1, depthwise 3×3 (stride 2, or 1 into the last stage), 1×1. +struct Merge { + c1: Linear, + c2: Depthwise, + c3: Linear, +} + +impl Merge { + fn forward(&self, x: &[f32], h: usize, w: usize) -> (Vec, usize, usize) { + let mut t = self.c1.forward(x, h * w); + nn::gelu(&mut t); + let (mut t, oh, ow) = self.c2.forward(&t, h, w); + nn::gelu(&mut t); + (self.c3.forward(&t, oh * ow), oh, ow) + } +} + +/// A TinyViT block: windowed attention (with learned relative-position biases), a depthwise +/// "local conv", then an MLP. +struct Block { + norm: Norm, + qkv: Linear, + proj: Linear, + /// `[heads × n × n]` for the `ws × ws` window. + bias: Vec, + heads: usize, + ws: usize, + local: Depthwise, + mlp_norm: Norm, + fc1: Linear, + fc2: Linear, +} + +impl Block { + fn new(w: &W, p: &str, dim: usize, heads: usize, ws: usize) -> Result { + let n = ws * ws; + // Offsets between window positions, numbered in first-seen order (as TinyViT does). + let points: Vec<(i64, i64)> = (0..ws as i64).flat_map(|r| (0..ws as i64).map(move |c| (r, c))).collect(); + let mut offsets: Vec<(i64, i64)> = vec![]; + let mut idx = Vec::with_capacity(n * n); + for a in &points { + for b in &points { + let o = ((a.0 - b.0).abs(), (a.1 - b.1).abs()); + let i = match offsets.iter().position(|x| *x == o) { + Some(i) => i, + None => { + offsets.push(o); + offsets.len() - 1 + } + }; + idx.push(i); + } + } + let table = w.get(&format!("{p}.attn.attention_biases"), heads * offsets.len())?; + let no = offsets.len(); + let bias = (0..heads).flat_map(|h| idx.iter().map(move |i| table[h * no + i])).collect::>(); + Ok(Block { + norm: w.norm(&format!("{p}.attn.norm"), dim, 1e-5)?, + qkv: w.linear(&format!("{p}.attn.qkv"), 3 * dim, dim, true)?, + proj: w.linear(&format!("{p}.attn.proj"), dim, dim, true)?, + bias, + heads, + ws, + local: w.depthwise_bn(&format!("{p}.local_conv"), dim, 1)?, + mlp_norm: w.norm(&format!("{p}.mlp.norm"), dim, 1e-5)?, + fc1: w.linear(&format!("{p}.mlp.fc1"), 4 * dim, dim, true)?, + fc2: w.linear(&format!("{p}.mlp.fc2"), dim, 4 * dim, true)?, + }) + } + + fn forward(&self, x: &mut Vec, h: usize, w: usize) { + let dim = self.proj.out; + let (ws, heads) = (self.ws, self.heads); + let hd = dim / heads; + // Pad the grid to whole windows (zeros, as TinyViT does), then norm + qkv once. + let (ph, pw) = (h.div_ceil(ws) * ws, w.div_ceil(ws) * ws); + let mut pad = vec![0.0f32; ph * pw * dim]; + for y in 0..h { + pad[y * pw * dim..(y * pw + w) * dim].copy_from_slice(&x[y * w * dim..(y + 1) * w * dim]); + } + self.norm.forward(&mut pad); + let qkv = self.qkv.forward(&pad, ph * pw); + // Per window: gather q, k, v (each head's qkv is [q | k | v]), attend. + let (nwy, nwx) = (ph / ws, pw / ws); + let n = ws * ws; + let outs: Vec> = (0..nwy * nwx) + .into_par_iter() + .map(|win| { + let (wy, wx) = (win / nwx, win % nwx); + let (mut q, mut k, mut v) = (vec![0.0f32; n * dim], vec![0.0f32; n * dim], vec![0.0f32; n * dim]); + for t in 0..n { + let (ty, tx) = (wy * ws + t / ws, wx * ws + t % ws); + let src = &qkv[(ty * pw + tx) * 3 * dim..][..3 * dim]; + for hh in 0..heads { + let s = &src[hh * 3 * hd..][..3 * hd]; + q[t * dim + hh * hd..][..hd].copy_from_slice(&s[..hd]); + k[t * dim + hh * hd..][..hd].copy_from_slice(&s[hd..2 * hd]); + v[t * dim + hh * hd..][..hd].copy_from_slice(&s[2 * hd..]); + } + } + nn::attention(&q, &k, &v, n, n, heads, hd, Some(&self.bias)) + }) + .collect(); + // Un-window into the real grid, project, residual. + let mut a = vec![0.0f32; h * w * dim]; + for (win, o) in outs.iter().enumerate() { + let (wy, wx) = (win / nwx, win % nwx); + for t in 0..n { + let (ty, tx) = (wy * ws + t / ws, wx * ws + t % ws); + if ty < h && tx < w { + a[(ty * w + tx) * dim..][..dim].copy_from_slice(&o[t * dim..][..dim]); + } + } + } + let a = self.proj.forward(&a, h * w); + nn::add(x, &a); + // Local depthwise conv (no residual), then the MLP with its residual. + let (lc, _, _) = self.local.forward(x, h, w); + *x = lc; + let mut m = x.clone(); + self.mlp_norm.forward(&mut m); + let mut m = self.fc1.forward(&m, h * w); + nn::gelu(&mut m); + let m = self.fc2.forward(&m, h * w); + nn::add(x, &m); + } +} + +/// The TinyViT-5M encoder with SAM's neck. +struct Encoder { + pe1: Conv, + pe2: Conv, + stage0: Vec, + merges: Vec, + stages: Vec>, + neck1: Linear, + neck_ln1: Norm, + neck2: Conv, + neck_ln2: Norm, +} + +const DIMS: [usize; 4] = [64, 128, 160, 320]; +const DEPTHS: [usize; 4] = [2, 2, 6, 2]; +const HEADS: [usize; 4] = [2, 4, 5, 10]; +const WINDOWS: [usize; 4] = [7, 7, 14, 7]; + +impl Encoder { + fn new(w: &W) -> Result { + let p = "image_encoder"; + let stage0 = (0..DEPTHS[0]) + .map(|i| { + let b = format!("{p}.layers.0.blocks.{i}"); + Ok(MbConv { + c1: w.pointwise_bn(&format!("{b}.conv1"), DIMS[0] * 4, DIMS[0])?, + c2: w.depthwise_bn(&format!("{b}.conv2"), DIMS[0] * 4, 1)?, + c3: w.pointwise_bn(&format!("{b}.conv3"), DIMS[0], DIMS[0] * 4)?, + }) + }) + .collect::>>()?; + let mut merges = vec![]; + let mut stages = vec![]; + for l in 0..3 { + let (inp, out) = (DIMS[l], DIMS[l + 1]); + let d = format!("{p}.layers.{l}.downsample"); + merges.push(Merge { + c1: w.pointwise_bn(&format!("{d}.conv1"), out, inp)?, + c2: w.depthwise_bn(&format!("{d}.conv2"), out, if out == 320 { 1 } else { 2 })?, + c3: w.pointwise_bn(&format!("{d}.conv3"), out, out)?, + }); + let s = l + 1; + stages + .push((0..DEPTHS[s]).map(|i| Block::new(w, &format!("{p}.layers.{s}.blocks.{i}"), DIMS[s], HEADS[s], WINDOWS[s])).collect::>>()?); + } + let n2 = w.get(&format!("{p}.neck.2.weight"), DIM * DIM * 9)?; + Ok(Encoder { + pe1: w.conv3_bn(&format!("{p}.patch_embed.seq.0"), DIMS[0] / 2, 3, 2)?, + pe2: w.conv3_bn(&format!("{p}.patch_embed.seq.2"), DIMS[0], DIMS[0] / 2, 2)?, + stage0, + merges, + stages, + neck1: w.linear(&format!("{p}.neck.0"), DIM, DIMS[3], false)?, + neck_ln1: w.norm(&format!("{p}.neck.1"), DIM, 1e-6)?, + neck2: Conv::new(n2, None, DIM, DIM, 3, 1, 1).ok_or("neck: bad shape")?, + neck_ln2: w.norm(&format!("{p}.neck.3"), DIM, 1e-6)?, + }) + } + + /// `IMG×IMG×3` normalised pixels → `EMB×EMB×DIM`. + fn forward(&self, x: &[f32]) -> Vec { + let (mut t, h, w) = self.pe1.forward(x, IMG, IMG); + nn::gelu(&mut t); + let (mut t, mut h, mut w) = { + let (t2, h2, w2) = self.pe2.forward(&t, h, w); + (t2, h2, w2) + }; + for b in &self.stage0 { + t = b.forward(&t, h, w); + } + for (m, blocks) in self.merges.iter().zip(&self.stages) { + let (t2, h2, w2) = m.forward(&t, h, w); + (t, h, w) = (t2, h2, w2); + for b in blocks { + b.forward(&mut t, h, w); + } + } + let mut t = self.neck1.forward(&t, h * w); + self.neck_ln1.forward(&mut t); + let (mut t, _, _) = self.neck2.forward(&t, h, w); + self.neck_ln2.forward(&mut t); + t + } +} + +/// SAM's attention with optional down-projection (`internal = DIM / rate`). +struct Attn { + q: Linear, + k: Linear, + v: Linear, + out: Linear, +} + +impl Attn { + fn new(w: &W, p: &str, rate: usize) -> Result { + let i = DIM / rate; + Ok(Attn { + q: w.linear(&format!("{p}.q_proj"), i, DIM, true)?, + k: w.linear(&format!("{p}.k_proj"), i, DIM, true)?, + v: w.linear(&format!("{p}.v_proj"), i, DIM, true)?, + out: w.linear(&format!("{p}.out_proj"), DIM, i, true)?, + }) + } + fn forward(&self, q: &[f32], nq: usize, k: &[f32], v: &[f32], nk: usize) -> Vec { + let (q, k, v) = (self.q.forward(q, nq), self.k.forward(k, nk), self.v.forward(v, nk)); + let heads = 8; + let o = nn::attention(&q, &k, &v, nq, nk, heads, self.q.out / heads, None); + self.out.forward(&o, nq) + } +} + +fn plus(a: &[f32], b: &[f32]) -> Vec { + a.iter().zip(b).map(|(x, y)| x + y).collect() +} + +struct TwoWayBlock { + self_attn: Attn, + norm1: Norm, + t2i: Attn, + norm2: Norm, + lin1: Linear, + lin2: Linear, + norm3: Norm, + i2t: Attn, + norm4: Norm, + skip_pe: bool, +} + +/// The decoder: prompt encoder, two-way transformer, mask head. +struct Decoder { + gauss: Vec, + /// `EMB×EMB×DIM` positional encoding of the embedding grid. + dense_pe: Vec, + point_embed: [Vec; 4], + not_a_point: Vec, + no_mask: Vec, + md1: Conv, + md_ln1: Norm, + md2: Conv, + md_ln2: Norm, + md3: Linear, + blocks: Vec, + final_attn: Attn, + norm_final: Norm, + iou_token: Vec, + mask_tokens: Vec, + up1: Upconv, + up_ln: Norm, + up2: Upconv, + hyper: Vec<[Linear; 3]>, + iou_head: [Linear; 3], +} + +impl Decoder { + fn new(w: &W) -> Result { + let pe = "prompt_encoder"; + let md = "mask_decoder"; + let gauss = w.get(&format!("{pe}.pe_layer.positional_encoding_gaussian_matrix"), 2 * DIM / 2)?.to_vec(); + let pt = |i: usize| w.get(&format!("{pe}.point_embeddings.{i}.weight"), DIM).map(<[f32]>::to_vec); + let conv = |p: &str, out: usize, inp: usize, ks: usize| -> Result { + let wt = w.get(&format!("{p}.weight"), out * inp * ks * ks)?; + let b = w.get(&format!("{p}.bias"), out)?; + Conv::new(wt, Some(b), out, inp, ks, 2, 0).ok_or_else(|| format!("{p}: bad shape")) + }; + let blocks = (0..2) + .map(|i| { + let p = format!("{md}.transformer.layers.{i}"); + Ok(TwoWayBlock { + self_attn: Attn::new(w, &format!("{p}.self_attn"), 1)?, + norm1: w.norm(&format!("{p}.norm1"), DIM, 1e-5)?, + t2i: Attn::new(w, &format!("{p}.cross_attn_token_to_image"), 2)?, + norm2: w.norm(&format!("{p}.norm2"), DIM, 1e-5)?, + lin1: w.linear(&format!("{p}.mlp.lin1"), 2048, DIM, true)?, + lin2: w.linear(&format!("{p}.mlp.lin2"), DIM, 2048, true)?, + norm3: w.norm(&format!("{p}.norm3"), DIM, 1e-5)?, + i2t: Attn::new(w, &format!("{p}.cross_attn_image_to_token"), 2)?, + norm4: w.norm(&format!("{p}.norm4"), DIM, 1e-5)?, + skip_pe: i == 0, + }) + }) + .collect::>>()?; + let upconv = |p: &str, inp: usize, out: usize| -> Result { + Upconv::new(w.get(&format!("{p}.weight"), inp * out * 4)?, w.get(&format!("{p}.bias"), out)?, inp, out).ok_or_else(|| format!("{p}: bad shape")) + }; + let mlp3 = |p: &str, out: usize| -> Result<[Linear; 3]> { + Ok([ + w.linear(&format!("{p}.layers.0"), DIM, DIM, true)?, + w.linear(&format!("{p}.layers.1"), DIM, DIM, true)?, + w.linear(&format!("{p}.layers.2"), out, DIM, true)?, + ]) + }; + let mut d = Decoder { + gauss, + dense_pe: vec![], + point_embed: [pt(0)?, pt(1)?, pt(2)?, pt(3)?], + not_a_point: w.get(&format!("{pe}.not_a_point_embed.weight"), DIM)?.to_vec(), + no_mask: w.get(&format!("{pe}.no_mask_embed.weight"), DIM)?.to_vec(), + md1: conv(&format!("{pe}.mask_downscaling.0"), 4, 1, 2)?, + md_ln1: w.norm(&format!("{pe}.mask_downscaling.1"), 4, 1e-6)?, + md2: conv(&format!("{pe}.mask_downscaling.3"), 16, 4, 2)?, + md_ln2: w.norm(&format!("{pe}.mask_downscaling.4"), 16, 1e-6)?, + md3: w.linear(&format!("{pe}.mask_downscaling.6"), DIM, 16, true)?, + blocks, + final_attn: Attn::new(w, &format!("{md}.transformer.final_attn_token_to_image"), 2)?, + norm_final: w.norm(&format!("{md}.transformer.norm_final_attn"), DIM, 1e-5)?, + iou_token: w.get(&format!("{md}.iou_token.weight"), DIM)?.to_vec(), + mask_tokens: w.get(&format!("{md}.mask_tokens.weight"), 4 * DIM)?.to_vec(), + up1: upconv(&format!("{md}.output_upscaling.0"), DIM, DIM / 4)?, + up_ln: w.norm(&format!("{md}.output_upscaling.1"), DIM / 4, 1e-6)?, + up2: upconv(&format!("{md}.output_upscaling.3"), DIM / 4, DIM / 8)?, + hyper: (0..4).map(|i| mlp3(&format!("{md}.output_hypernetworks_mlps.{i}"), DIM / 8)).collect::>>()?, + iou_head: mlp3(&format!("{md}.iou_prediction_head"), 4)?, + }; + d.dense_pe = (0..EMB * EMB).flat_map(|i| d.pe(((i % EMB) as f32 + 0.5) / EMB as f32, ((i / EMB) as f32 + 0.5) / EMB as f32)).collect(); + Ok(d) + } + + /// Random Fourier positional encoding of a point in [0, 1]². + fn pe(&self, x: f32, y: f32) -> Vec { + let half = DIM / 2; + let (cx, cy) = (2.0 * x - 1.0, 2.0 * y - 1.0); + let a: Vec = (0..half).map(|j| 2.0 * PI * (cx * self.gauss[j] + cy * self.gauss[half + j])).collect(); + a.iter().map(|v| v.sin()).chain(a.iter().map(|v| v.cos())).collect() + } + + fn mlp3(layers: &[Linear; 3], x: &[f32]) -> Vec { + let mut t = layers[0].forward(x, 1); + nn::relu(&mut t); + let mut t = layers[1].forward(&t, 1); + nn::relu(&mut t); + layers[2].forward(&t, 1) + } + + /// Mask logits (`LOW×LOW`, in the padded `IMG` square) and their predicted IoU for the + /// prompt (coordinates already in `IMG` space). + fn forward(&self, emb: &[f32], points: &[([f32; 2], bool)], bbox: Option<[f32; 4]>, mask: Option<&[f32]>, multimask: bool) -> (Vec, f32) { + // Sparse prompt tokens. + let mut sparse: Vec = vec![]; + for (p, fg) in points { + let mut e = self.pe((p[0] + 0.5) / IMG as f32, (p[1] + 0.5) / IMG as f32); + let add = &self.point_embed[usize::from(*fg)]; + e.iter_mut().zip(add).for_each(|(a, b)| *a += b); + sparse.extend(e); + } + match bbox { + Some(b) => { + for (i, (x, y)) in [(b[0], b[1]), (b[2], b[3])].into_iter().enumerate() { + let mut e = self.pe((x + 0.5) / IMG as f32, (y + 0.5) / IMG as f32); + e.iter_mut().zip(&self.point_embed[2 + i]).for_each(|(a, b)| *a += b); + sparse.extend(e); + } + } + // Points without a box get a padding "not a point". + None if !points.is_empty() => sparse.extend(&self.not_a_point), + None => {} + } + // Dense prompt: the mask input, or "no mask". + let dense: Vec = match mask { + Some(m) if m.len() == LOW * LOW => { + let (mut t, h, w) = self.md1.forward(m, LOW, LOW); + self.md_ln1.forward(&mut t); + nn::gelu(&mut t); + let (mut t, h, w) = self.md2.forward(&t, h, w); + self.md_ln2.forward(&mut t); + nn::gelu(&mut t); + self.md3.forward(&t, h * w) + } + _ => (0..EMB * EMB).flat_map(|_| self.no_mask.iter().copied()).collect(), + }; + // Tokens: IoU, four mask tokens, prompts. + let mut tokens: Vec = self.iou_token.clone(); + tokens.extend(&self.mask_tokens); + tokens.extend(&sparse); + let nt = tokens.len() / DIM; + let nk = EMB * EMB; + let mut keys = plus(emb, &dense); + let pe = &self.dense_pe; + let qpe = tokens.clone(); + let mut queries = tokens; + for b in &self.blocks { + // Self attention. + queries = if b.skip_pe { + b.self_attn.forward(&queries, nt, &queries, &queries, nt) + } else { + let q = plus(&queries, &qpe); + plus(&queries, &b.self_attn.forward(&q, nt, &q, &queries, nt)) + }; + b.norm1.forward(&mut queries); + // Tokens to image. + let q = plus(&queries, &qpe); + let k = plus(&keys, pe); + queries = plus(&queries, &b.t2i.forward(&q, nt, &k, &keys, nk)); + b.norm2.forward(&mut queries); + // MLP. + let mut m = b.lin1.forward(&queries, nt); + nn::relu(&mut m); + queries = plus(&queries, &b.lin2.forward(&m, nt)); + b.norm3.forward(&mut queries); + // Image to tokens. + let q = plus(&queries, &qpe); + let k = plus(&keys, pe); + keys = plus(&keys, &b.i2t.forward(&k, nk, &q, &queries, nt)); + b.norm4.forward(&mut keys); + } + let q = plus(&queries, &qpe); + let k = plus(&keys, pe); + queries = plus(&queries, &self.final_attn.forward(&q, nt, &k, &keys, nk)); + self.norm_final.forward(&mut queries); + // Upscale the image tokens 4× and dot them with each mask token's hypernetwork output. + let mut up = self.up1.forward(&keys, EMB, EMB); + self.up_ln.forward(&mut up); + nn::gelu(&mut up); + let mut up = self.up2.forward(&up, 2 * EMB, 2 * EMB); + nn::gelu(&mut up); + let c = DIM / 8; + let iou = Self::mlp3(&self.iou_head, &queries[..DIM]); + let pick = if multimask { (1..4).max_by(|a, b| iou[*a].total_cmp(&iou[*b])).unwrap_or(0) } else { 0 }; + let hyper = Self::mlp3(&self.hyper[pick], &queries[(1 + pick) * DIM..][..DIM]); + let logits: Vec = up.par_chunks(c).map(|px| px.iter().zip(&hyper).map(|(a, b)| a * b).sum()).collect(); + (logits, iou.get(pick).copied().unwrap_or(0.0)) + } +} + +/// An image's embedding and how it was fitted into the model's square. +#[derive(Clone)] +pub struct Embedding { + pub data: Vec, + pub w: usize, + pub h: usize, + /// Image pixels → model pixels. + pub scale: f32, +} + +pub struct MobileSam { + enc: Encoder, + dec: Decoder, + /// The last image embedded (by content hash): new prompts on the same frame reuse it. + cache: Mutex>, +} + +fn hash_rgb(rgb: &[[f32; 3]], w: usize, h: usize) -> u64 { + let mut s: u64 = 0xcbf2_9ce4_8422_2325 ^ ((w as u64) << 32 | h as u64); + for p in rgb.iter().step_by(7) { + for v in p { + s = (s ^ v.to_bits() as u64).wrapping_mul(0x0100_0000_01b3); + } + } + s +} + +impl MobileSam { + pub fn from_weights(w: &Weights) -> Result { + let w = W(w); + Ok(MobileSam { enc: Encoder::new(&w)?, dec: Decoder::new(&w)?, cache: Mutex::new(None) }) + } + + /// Embed an image (straight RGB 0–1, row-major `w×h`), or reuse the last embedding. + pub fn embed(&self, rgb: &[[f32; 3]], w: usize, h: usize) -> Result { + if w == 0 || h == 0 || rgb.len() != w * h { + return Err("empty image".into()); + } + let key = hash_rgb(rgb, w, h); + if let Ok(c) = self.cache.lock() + && let Some((k, e)) = c.as_ref() + && *k == key + { + return Ok(e.clone()); + } + let scale = IMG as f32 / w.max(h) as f32; + let (rw, rh) = (((w as f32 * scale).round() as usize).clamp(1, IMG), ((h as f32 * scale).round() as usize).clamp(1, IMG)); + // Resize (box filter when shrinking, bilinear when growing), normalise, pad. + let mut x = vec![0.0f32; IMG * IMG * 3]; + x.par_chunks_mut(IMG * 3).enumerate().take(rh).for_each(|(y, row)| { + for xo in 0..rw { + let px = sample(rgb, w, h, (xo as f32 + 0.5) / scale, (y as f32 + 0.5) / scale, 1.0 / scale); + for c in 0..3 { + row[xo * 3 + c] = (px[c] * 255.0 - MEAN[c]) / STD[c]; + } + } + }); + let e = Embedding { data: self.enc.forward(&x), w, h, scale }; + if let Ok(mut c) = self.cache.lock() { + *c = Some((key, e.clone())); + } + Ok(e) + } + + /// Foreground probability (`w×h`) for a prompt on an embedded image, and the predicted IoU. + pub fn predict(&self, e: &Embedding, prompt: &Prompt) -> (Vec, f32) { + let s = e.scale; + let points: Vec<([f32; 2], bool)> = prompt.points.iter().map(|(p, fg)| ([p[0] * s, p[1] * s], *fg)).collect(); + let bbox = prompt.bbox.map(|b| b.map(|v| v * s)); + // A prior mask (probabilities at image size) → logits on the low-res grid. + let mask: Option> = prompt.mask.as_ref().filter(|m| m.len() == e.w * e.h).map(|m| { + (0..LOW * LOW) + .map(|i| { + let (gx, gy) = (((i % LOW) as f32 + 0.5) * 4.0 / s, ((i / LOW) as f32 + 0.5) * 4.0 / s); + if gx >= e.w as f32 || gy >= e.h as f32 { + return -20.0; + } + let p = m[(gy as usize).min(e.h - 1) * e.w + (gx as usize).min(e.w - 1)]; + ((p - 0.5) * 40.0).clamp(-20.0, 20.0) + }) + .collect() + }); + let multimask = points.len() == 1 && bbox.is_none() && mask.is_none(); + let (logits, iou) = self.dec.forward(&e.data, &points, bbox, mask.as_deref(), multimask); + // Low-res logits (the padded square) → image pixels, bilinear, then sigmoid. + let prob: Vec = (0..e.w * e.h) + .into_par_iter() + .map(|i| { + let (x, y) = ((i % e.w) as f32 + 0.5, (i / e.w) as f32 + 0.5); + let (u, v) = (x * s / 4.0 - 0.5, y * s / 4.0 - 0.5); + let l = bilinear(&logits, LOW, LOW, u, v); + 1.0 / (1.0 + (-l).exp()) + }) + .collect(); + (prob, iou) + } +} + +fn bilinear(g: &[f32], w: usize, h: usize, u: f32, v: f32) -> f32 { + let (u, v) = (u.clamp(0.0, (w - 1) as f32), v.clamp(0.0, (h - 1) as f32)); + let (x0, y0) = (u.floor() as usize, v.floor() as usize); + let (x1, y1) = ((x0 + 1).min(w - 1), (y0 + 1).min(h - 1)); + let (fx, fy) = (u - x0 as f32, v - y0 as f32); + let at = |x: usize, y: usize| g.get(y * w + x).copied().unwrap_or(0.0); + (at(x0, y0) * (1.0 - fx) + at(x1, y0) * fx) * (1.0 - fy) + (at(x0, y1) * (1.0 - fx) + at(x1, y1) * fx) * fy +} + +/// The image's colour around `(x, y)`: a `foot`-pixel box average when shrinking, else bilinear. +fn sample(rgb: &[[f32; 3]], w: usize, h: usize, x: f32, y: f32, foot: f32) -> [f32; 3] { + if foot <= 1.0 { + let (u, v) = ((x - 0.5).clamp(0.0, (w - 1) as f32), (y - 0.5).clamp(0.0, (h - 1) as f32)); + let (x0, y0) = (u.floor() as usize, v.floor() as usize); + let (x1, y1) = ((x0 + 1).min(w - 1), (y0 + 1).min(h - 1)); + let (fx, fy) = (u - x0 as f32, v - y0 as f32); + let p = |x: usize, y: usize| rgb[y * w + x]; + return [0, 1, 2].map(|c| (p(x0, y0)[c] * (1.0 - fx) + p(x1, y0)[c] * fx) * (1.0 - fy) + (p(x0, y1)[c] * (1.0 - fx) + p(x1, y1)[c] * fx) * fy); + } + let r = foot / 2.0; + let (xa, xb) = (((x - r).floor().max(0.0)) as usize, ((x + r).ceil() as usize).min(w)); + let (ya, yb) = (((y - r).floor().max(0.0)) as usize, ((y + r).ceil() as usize).min(h)); + let mut s = [0.0f32; 3]; + let mut n = 0.0; + for yy in ya..yb.max(ya + 1).min(h) { + for xx in xa..xb.max(xa + 1).min(w) { + let p = rgb[yy * w + xx]; + for c in 0..3 { + s[c] += p[c]; + } + n += 1.0; + } + } + s.map(|v| v / f32::max(n, 1.0)) +} diff --git a/crates/segment/src/nn.rs b/crates/segment/src/nn.rs new file mode 100644 index 000000000..1628a0ae9 --- /dev/null +++ b/crates/segment/src/nn.rs @@ -0,0 +1,432 @@ +//! The few neural-network kernels MobileSAM needs, on token-major feature maps (`[pixels, channels]` +//! row-major, i.e. NHWC with a batch of one): dense layers on a parallel GEMM (ndarray's safe +//! `general_mat_mul`, which picks AVX/FMA kernels at run time), convolutions as patch matrices +//! times weights, depthwise and transposed convolutions, layer norm, exact GELU and softmax. +//! Weights are re-laid out once at load so every inner loop runs over contiguous channels. + +use ndarray::linalg::general_mat_mul; +use ndarray::{ArrayView2, ArrayViewMut2}; +use rayon::prelude::*; + +/// Rows of the left matrix per parallel task. +const ROW_BLOCK: usize = 64; + +/// `out[m×n] += a[m×k] · b[k×n]` (row-major), rows split across threads. +pub fn gemm(a: &[f32], m: usize, k: usize, b: &[f32], n: usize, out: &mut [f32]) { + if m == 0 || n == 0 || a.len() < m * k || b.len() < k * n || out.len() < m * n { + return; + } + let Ok(bv) = ArrayView2::from_shape((k, n), &b[..k * n]) else { return }; + let block = |r0: usize, oc: &mut [f32]| { + let rows = oc.len() / n; + if let (Ok(av), Ok(mut cv)) = (ArrayView2::from_shape((rows, k), &a[r0 * k..(r0 + rows) * k]), ArrayViewMut2::from_shape((rows, n), oc)) { + general_mat_mul(1.0, &av, &bv, 1.0, &mut cv); + } + }; + if m <= ROW_BLOCK { + block(0, &mut out[..m * n]); + } else { + out[..m * n].par_chunks_mut(ROW_BLOCK * n).enumerate().for_each(|(i, oc)| block(i * ROW_BLOCK, oc)); + } +} + +/// A dense layer (also a 1×1 convolution): `y = x·Wᵀ + b`, with `W` stored transposed. +#[derive(Clone, Debug, Default)] +pub struct Linear { + /// `[inp × out]`. + pub wt: Vec, + pub b: Vec, + pub inp: usize, + pub out: usize, +} + +impl Linear { + /// From PyTorch's `[out × inp]` weight (a 1×1 conv's `[out × inp × 1 × 1]` too). + pub fn new(w: &[f32], b: Option<&[f32]>, out: usize, inp: usize) -> Option { + if w.len() != out * inp || b.is_some_and(|b| b.len() != out) { + return None; + } + let mut wt = vec![0.0; inp * out]; + for o in 0..out { + for i in 0..inp { + wt[i * out + o] = w[o * inp + i]; + } + } + Some(Linear { wt, b: b.map_or_else(|| vec![0.0; out], <[f32]>::to_vec), inp, out }) + } + + /// `rows` tokens of `inp` channels → `rows × out`. + pub fn forward(&self, x: &[f32], rows: usize) -> Vec { + let mut y: Vec = (0..rows).flat_map(|_| self.b.iter().copied()).collect(); + gemm(x, rows, self.inp, &self.wt, self.out, &mut y); + y + } +} + +/// A `ks×ks` convolution (any stride and padding, one group) on `h×w×inp` → `oh×ow×out`, as a +/// patch matrix times the weights, built a block of output rows at a time. +#[derive(Clone, Debug, Default)] +pub struct Conv { + /// `[(ky·ks + kx)·inp + ci] × out`. + pub wt: Vec, + pub b: Vec, + pub inp: usize, + pub out: usize, + pub ks: usize, + pub stride: usize, + pub pad: usize, +} + +impl Conv { + /// From PyTorch's `[out × inp × ks × ks]`. + pub fn new(w: &[f32], b: Option<&[f32]>, out: usize, inp: usize, ks: usize, stride: usize, pad: usize) -> Option { + if w.len() != out * inp * ks * ks || b.is_some_and(|b| b.len() != out) || stride == 0 { + return None; + } + let kk = ks * ks * inp; + let mut wt = vec![0.0; kk * out]; + for o in 0..out { + for ci in 0..inp { + for t in 0..ks * ks { + wt[(t * inp + ci) * out + o] = w[(o * inp + ci) * ks * ks + t]; + } + } + } + Some(Conv { wt, b: b.map_or_else(|| vec![0.0; out], <[f32]>::to_vec), inp, out, ks, stride, pad }) + } + + pub fn out_size(&self, n: usize) -> usize { + (n + 2 * self.pad).saturating_sub(self.ks) / self.stride + 1 + } + + pub fn forward(&self, x: &[f32], h: usize, w: usize) -> (Vec, usize, usize) { + let (oh, ow) = (self.out_size(h), self.out_size(w)); + let kk = self.ks * self.ks * self.inp; + let mut y: Vec = (0..oh * ow).flat_map(|_| self.b.iter().copied()).collect(); + if x.len() < h * w * self.inp { + return (y, oh, ow); + } + y.par_chunks_mut(ow * self.out).enumerate().for_each(|(oy, row)| { + // This output row's patches. + let mut patch = vec![0.0f32; ow * kk]; + for ox in 0..ow { + for ky in 0..self.ks { + let iy = (oy * self.stride + ky) as isize - self.pad as isize; + if iy < 0 || iy >= h as isize { + continue; + } + for kx in 0..self.ks { + let ix = (ox * self.stride + kx) as isize - self.pad as isize; + if ix < 0 || ix >= w as isize { + continue; + } + let src = (iy as usize * w + ix as usize) * self.inp; + let dst = ox * kk + (ky * self.ks + kx) * self.inp; + patch[dst..dst + self.inp].copy_from_slice(&x[src..src + self.inp]); + } + } + } + if let (Ok(av), Ok(bv), Ok(mut cv)) = ( + ArrayView2::from_shape((ow, kk), &patch[..]), + ArrayView2::from_shape((kk, self.out), &self.wt[..]), + ArrayViewMut2::from_shape((ow, self.out), row), + ) { + general_mat_mul(1.0, &av, &bv, 1.0, &mut cv); + } + }); + (y, oh, ow) + } +} + +/// A depthwise 3×3 convolution (padding 1, stride 1 or 2) on `h×w×c`. +#[derive(Clone, Debug, Default)] +pub struct Depthwise { + /// `[tap × c]`. + pub w: Vec, + pub b: Vec, + pub c: usize, + pub stride: usize, +} + +impl Depthwise { + /// From PyTorch's `[c × 1 × 3 × 3]`. + pub fn new(w: &[f32], b: Option<&[f32]>, c: usize, stride: usize) -> Option { + if w.len() != c * 9 || b.is_some_and(|b| b.len() != c) || stride == 0 { + return None; + } + let mut wt = vec![0.0; 9 * c]; + for ch in 0..c { + for t in 0..9 { + wt[t * c + ch] = w[ch * 9 + t]; + } + } + Some(Depthwise { w: wt, b: b.map_or_else(|| vec![0.0; c], <[f32]>::to_vec), c, stride }) + } + + pub fn forward(&self, x: &[f32], h: usize, w: usize) -> (Vec, usize, usize) { + let c = self.c; + let (oh, ow) = ((h + 2 - 3) / self.stride + 1, (w + 2 - 3) / self.stride + 1); + let mut y = vec![0.0f32; oh * ow * c]; + if x.len() < h * w * c { + return (y, oh, ow); + } + y.par_chunks_mut(ow * c).enumerate().for_each(|(oy, row)| { + for ox in 0..ow { + let out = &mut row[ox * c..(ox + 1) * c]; + out.copy_from_slice(&self.b); + for ky in 0..3 { + let iy = (oy * self.stride + ky) as isize - 1; + if iy < 0 || iy >= h as isize { + continue; + } + for kx in 0..3 { + let ix = (ox * self.stride + kx) as isize - 1; + if ix < 0 || ix >= w as isize { + continue; + } + let src = &x[(iy as usize * w + ix as usize) * c..][..c]; + let tap = &self.w[(ky * 3 + kx) * c..][..c]; + for ((o, s), k) in out.iter_mut().zip(src).zip(tap) { + *o += s * k; + } + } + } + } + }); + (y, oh, ow) + } +} + +/// A 2×2, stride-2 transposed convolution (each input pixel paints a 2×2 block): +/// `h×w×inp` → `2h×2w×out`. +#[derive(Clone, Debug, Default)] +pub struct Upconv { + /// One dense layer per output tap (dy, dx). + pub taps: Vec, + pub out: usize, +} + +impl Upconv { + /// From PyTorch's ConvTranspose2d `[inp × out × 2 × 2]`. + pub fn new(w: &[f32], b: &[f32], inp: usize, out: usize) -> Option { + if w.len() != inp * out * 4 || b.len() != out { + return None; + } + let taps = (0..4) + .map(|t| { + let mut wt = vec![0.0; inp * out]; + for i in 0..inp { + for o in 0..out { + wt[i * out + o] = w[(i * out + o) * 4 + t]; + } + } + Linear { wt, b: b.to_vec(), inp, out } + }) + .collect(); + Some(Upconv { taps, out }) + } + + pub fn forward(&self, x: &[f32], h: usize, w: usize) -> Vec { + let (ow, c) = (2 * w, self.out); + let mut y = vec![0.0f32; 4 * h * w * c]; + for (t, lin) in self.taps.iter().enumerate() { + let (dy, dx) = (t / 2, t % 2); + let part = lin.forward(x, h * w); + for p in 0..h * w { + let (py, px) = (p / w, p % w); + let dst = ((2 * py + dy) * ow + 2 * px + dx) * c; + if let (Some(d), Some(s)) = (y.get_mut(dst..dst + c), part.get(p * c..(p + 1) * c)) { + d.copy_from_slice(s); + } + } + } + y + } +} + +/// Layer norm over the channels of each token. +#[derive(Clone, Debug, Default)] +pub struct Norm { + pub w: Vec, + pub b: Vec, + pub eps: f32, +} + +impl Norm { + pub fn forward(&self, x: &mut [f32]) { + let c = self.w.len(); + if c == 0 { + return; + } + x.par_chunks_mut(c).for_each(|t| { + let mean = t.iter().sum::() / c as f32; + let var = t.iter().map(|v| (v - mean) * (v - mean)).sum::() / c as f32; + let inv = 1.0 / (var + self.eps).sqrt(); + for ((v, w), b) in t.iter_mut().zip(&self.w).zip(&self.b) { + *v = (*v - mean) * inv * w + b; + } + }); + } +} + +/// The error function (Abramowitz & Stegun 7.1.26 refined: |error| < 1.5e-7). +fn erf(x: f32) -> f32 { + let s = x.signum(); + let x = x.abs() as f64; + let t = 1.0 / (1.0 + 0.327_591_1 * x); + let y = 1.0 - (((((1.061_405_429 * t - 1.453_152_027) * t) + 1.421_413_741) * t - 0.284_496_736) * t + 0.254_829_592) * t * (-x * x).exp(); + s * y as f32 +} + +/// GELU (the exact, erf form, as PyTorch's `nn.GELU`), in place. +pub fn gelu(x: &mut [f32]) { + x.par_chunks_mut(4096).for_each(|c| { + for v in c { + *v = 0.5 * *v * (1.0 + erf(*v * std::f32::consts::FRAC_1_SQRT_2)); + } + }); +} + +pub fn relu(x: &mut [f32]) { + for v in x { + *v = v.max(0.0); + } +} + +/// `a += b`. +pub fn add(a: &mut [f32], b: &[f32]) { + a.par_chunks_mut(4096).zip(b.par_chunks(4096)).for_each(|(a, b)| { + for (x, y) in a.iter_mut().zip(b) { + *x += y; + } + }); +} + +/// Softmax of each row of `n` values, in place. +pub fn softmax_rows(x: &mut [f32], n: usize) { + if n == 0 { + return; + } + for row in x.chunks_mut(n) { + let m = row.iter().copied().fold(f32::NEG_INFINITY, f32::max); + let mut s = 0.0; + for v in row.iter_mut() { + *v = (*v - m).exp(); + s += *v; + } + let inv = 1.0 / s.max(1e-30); + for v in row.iter_mut() { + *v *= inv; + } + } +} + +/// Multi-head scaled dot-product attention of `nq` queries over `nk` keys/values, each `heads × +/// hd` wide (token-major, heads side by side), plus an optional per-head bias `[heads × nq × nk]`. +pub fn attention(q: &[f32], k: &[f32], v: &[f32], nq: usize, nk: usize, heads: usize, hd: usize, bias: Option<&[f32]>) -> Vec { + let dim = heads * hd; + let mut out = vec![0.0f32; nq * dim]; + if q.len() < nq * dim || k.len() < nk * dim || v.len() < nk * dim { + return out; + } + let scale = 1.0 / (hd as f32).sqrt(); + for h in 0..heads { + // This head's q, kᵀ and v as contiguous matrices. + let qh: Vec = (0..nq).flat_map(|i| q[i * dim + h * hd..][..hd].iter().map(|x| x * scale)).collect(); + let mut kt = vec![0.0f32; hd * nk]; + for j in 0..nk { + for d in 0..hd { + kt[d * nk + j] = k[j * dim + h * hd + d]; + } + } + let vh: Vec = (0..nk).flat_map(|j| v[j * dim + h * hd..][..hd].iter().copied()).collect(); + let mut s = match bias.and_then(|b| b.get(h * nq * nk..(h + 1) * nq * nk)) { + Some(b) => b.to_vec(), + None => vec![0.0; nq * nk], + }; + gemm(&qh, nq, hd, &kt, nk, &mut s); + softmax_rows(&mut s, nk); + let mut oh = vec![0.0f32; nq * hd]; + gemm(&s, nq, nk, &vh, hd, &mut oh); + for i in 0..nq { + out[i * dim + h * hd..][..hd].copy_from_slice(&oh[i * hd..][..hd]); + } + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn gemm_matches_naive() { + let (m, k, n) = (130, 17, 9); + let a: Vec = (0..m * k).map(|i| ((i * 7) % 13) as f32 - 6.0).collect(); + let b: Vec = (0..k * n).map(|i| ((i * 5) % 11) as f32 - 5.0).collect(); + let mut c = vec![1.0f32; m * n]; + gemm(&a, m, k, &b, n, &mut c); + for i in 0..m { + for j in 0..n { + let want: f32 = 1.0 + (0..k).map(|p| a[i * k + p] * b[p * n + j]).sum::(); + assert_eq!(c[i * n + j], want); + } + } + } + + #[test] + fn convolutions_match_direct_sums() { + // 3×3 stride-2 conv on a 5×4×2 input vs. the definition. + let (h, w, inp, out) = (5, 4, 2, 3); + let x: Vec = (0..h * w * inp).map(|i| (i % 7) as f32 * 0.5).collect(); + let wt: Vec = (0..out * inp * 9).map(|i| ((i % 5) as f32 - 2.0) * 0.25).collect(); + let b = [0.1, -0.2, 0.3]; + let conv = Conv::new(&wt, Some(&b), out, inp, 3, 2, 1).unwrap(); + let (y, oh, ow) = conv.forward(&x, h, w); + assert_eq!((oh, ow), (3, 2)); + for oy in 0..oh { + for ox in 0..ow { + for o in 0..out { + let mut s = b[o]; + for ci in 0..inp { + for ky in 0..3 { + for kx in 0..3 { + let (iy, ix) = ((oy * 2 + ky) as isize - 1, (ox * 2 + kx) as isize - 1); + if iy >= 0 && ix >= 0 && iy < h as isize && ix < w as isize { + s += x[(iy as usize * w + ix as usize) * inp + ci] * wt[((o * inp + ci) * 3 + ky) * 3 + kx]; + } + } + } + } + assert!((y[(oy * ow + ox) * out + o] - s).abs() < 1e-5); + } + } + } + // Depthwise == a grouped conv: compare with Conv on each channel alone. + let dw_w: Vec = (0..inp * 9).map(|i| (i as f32 - 8.0) * 0.1).collect(); + let dw = Depthwise::new(&dw_w, None, inp, 1).unwrap(); + let (yd, _, _) = dw.forward(&x, h, w); + for ci in 0..inp { + let xc: Vec = (0..h * w).map(|p| x[p * inp + ci]).collect(); + let single = Conv::new(&dw_w[ci * 9..ci * 9 + 9], None, 1, 1, 3, 1, 1).unwrap(); + let (yc, _, _) = single.forward(&xc, h, w); + for p in 0..h * w { + assert!((yd[p * inp + ci] - yc[p]).abs() < 1e-5); + } + } + } + + #[test] + fn gelu_and_softmax() { + let mut x = [-3.0f32, -1.0, 0.0, 1.0, 3.0]; + gelu(&mut x); + // PyTorch's exact GELU values. + let want = [-0.004_049_7, -0.158_655_3, 0.0, 0.841_344_7, 2.995_950_3]; + for (a, b) in x.iter().zip(want) { + assert!((a - b).abs() < 1e-5, "{a} {b}"); + } + let mut s = [1.0f32, 2.0, 3.0]; + softmax_rows(&mut s, 3); + assert!((s.iter().sum::() - 1.0).abs() < 1e-6 && s[2] > s[1]); + } +} diff --git a/crates/segment/src/pt.rs b/crates/segment/src/pt.rs new file mode 100644 index 000000000..16a17fbe9 --- /dev/null +++ b/crates/segment/src/pt.rs @@ -0,0 +1,499 @@ +//! Reading PyTorch checkpoints (`torch.save` of a state dict) without PyTorch: the file is a zip +//! archive (stored or deflated entries, zip64 sizes) holding `data.pkl`, a pickle (protocol 2–4) +//! whose tensors are `torch._utils._rebuild_tensor_v2(storage, offset, size, stride, …)` calls +//! with storages as persistent ids pointing at `data/` entries. Only what state dicts use is +//! understood; anything else is an error, never a panic. + +use std::collections::HashMap; + +/// A tensor read from a checkpoint: its shape and its values as f32 (row-major, contiguous). +#[derive(Clone, Debug, PartialEq)] +pub struct Tensor { + pub shape: Vec, + pub data: Vec, +} + +/// A checkpoint's tensors by name (non-float tensors such as batch-norm counters are skipped). +pub type Weights = HashMap; + +type Result = std::result::Result; + +fn u16_at(b: &[u8], o: usize) -> Result { + b.get(o..o + 2).map(|s| u16::from_le_bytes([s[0], s[1]]) as usize).ok_or_else(|| "truncated zip".to_string()) +} +fn u32_at(b: &[u8], o: usize) -> Result { + b.get(o..o + 4).map(|s| u32::from_le_bytes([s[0], s[1], s[2], s[3]]) as u64).ok_or_else(|| "truncated zip".to_string()) +} +fn u64_at(b: &[u8], o: usize) -> Result { + b.get(o..o + 8).and_then(|s| s.try_into().ok()).map(u64::from_le_bytes).ok_or_else(|| "truncated zip".to_string()) +} + +/// One archive entry: where its data starts, its sizes and compression. +struct Entry { + data: usize, + size: usize, + stored: usize, + method: usize, +} + +/// The entries of a zip archive by name. +fn zip_entries(b: &[u8]) -> Result> { + // End of central directory: the last `PK\x05\x06` (a comment may follow it). + let eocd = + (0..b.len().saturating_sub(21)).rev().find(|&i| b.get(i..i + 4) == Some(b"PK\x05\x06")).ok_or("not a zip archive (no end of central directory)")?; + let mut count = u16_at(b, eocd + 10)? as u64; + let mut cd = u32_at(b, eocd + 16)?; + // Zip64: the locator precedes the record. + if (count == 0xffff || cd == 0xffff_ffff) && eocd >= 20 && b.get(eocd - 20..eocd - 16) == Some(b"PK\x06\x07") { + let rec = u64_at(b, eocd - 12)? as usize; + if b.get(rec..rec + 4) == Some(b"PK\x06\x06") { + count = u64_at(b, rec + 32)?; + cd = u64_at(b, rec + 48)?; + } + } + let mut out = HashMap::new(); + let mut p = cd as usize; + for _ in 0..count.min(1 << 20) { + if b.get(p..p + 4) != Some(b"PK\x01\x02") { + return Err("damaged zip central directory".into()); + } + let method = u16_at(b, p + 10)?; + let mut stored = u32_at(b, p + 20)?; + let mut size = u32_at(b, p + 24)?; + let (nl, xl, cl) = (u16_at(b, p + 28)?, u16_at(b, p + 30)?, u16_at(b, p + 32)?); + let mut local = u32_at(b, p + 42)?; + let name = String::from_utf8_lossy(b.get(p + 46..p + 46 + nl).ok_or("truncated zip")?).to_string(); + // Zip64 extended information: the 0xffffffff fields, in order. + let mut x = p + 46 + nl; + let xend = x + xl; + while x + 4 <= xend { + let (id, len) = (u16_at(b, x)?, u16_at(b, x + 2)?); + if id == 1 { + let mut f = x + 4; + for v in [&mut size, &mut stored, &mut local] { + if *v == 0xffff_ffff { + *v = u64_at(b, f)?; + f += 8; + } + } + } + x += 4 + len; + } + let l = local as usize; + if b.get(l..l + 4) != Some(b"PK\x03\x04") { + return Err(format!("damaged zip entry {name}")); + } + let data = l + 30 + u16_at(b, l + 26)? + u16_at(b, l + 28)?; + out.insert(name, Entry { data, size: size as usize, stored: stored as usize, method }); + p += 46 + nl + xl + cl; + } + Ok(out) +} + +/// An entry's bytes (stored, or raw-deflated). +fn read_entry<'a>(b: &'a [u8], e: &Entry) -> Result> { + let raw = b.get(e.data..e.data + e.stored).ok_or("truncated zip entry")?; + match e.method { + 0 => Ok(std::borrow::Cow::Borrowed(raw)), + 8 => { + miniz_oxide::inflate::decompress_to_vec_with_limit(raw, e.size.max(1)).map(std::borrow::Cow::Owned).map_err(|e| format!("bad deflate data: {e:?}")) + } + m => Err(format!("unsupported zip compression {m}")), + } +} + +/// Pickle values (what state dicts use). +#[derive(Clone, Debug)] +enum Obj { + None, + Bool, + Int(i64), + Float, + Str(String), + Tuple(Vec), + List(Vec), + Dict(Vec<(Obj, Obj)>), + Global(String, String), + /// A storage: (dtype class name, entry key). + Storage(String, String), + /// `_rebuild_tensor_v2`: storage, element offset, size, stride. + Tensor { + dtype: String, + key: String, + offset: usize, + size: Vec, + stride: Vec, + }, + /// A call we don't interpret. + Other, + Mark, +} + +fn int_list(o: &Obj) -> Option> { + match o { + Obj::Tuple(v) | Obj::List(v) => v.iter().map(|x| if let Obj::Int(i) = x { usize::try_from(*i).ok() } else { None }).collect(), + _ => None, + } +} + +/// Run the pickle program and return its result. +fn unpickle(p: &[u8]) -> Result { + let mut stack: Vec = Vec::new(); + let mut memo: HashMap = HashMap::new(); + let mut i = 0usize; + let take = |i: &mut usize, n: usize| -> Result<&[u8]> { + let s = p.get(*i..*i + n).ok_or("truncated pickle")?; + *i += n; + Ok(s) + }; + let pop = |stack: &mut Vec| stack.pop().ok_or_else(|| "pickle stack underflow".to_string()); + let pop_mark = |stack: &mut Vec| -> Result> { + let m = stack.iter().rposition(|o| matches!(o, Obj::Mark)).ok_or("pickle: no mark")?; + let items = stack.split_off(m + 1); + stack.pop(); + Ok(items) + }; + let line = |i: &mut usize| -> Result { + let rest = p.get(*i..).ok_or("truncated pickle")?; + let n = rest.iter().position(|c| *c == b'\n').ok_or("truncated pickle")?; + let s = String::from_utf8_lossy(&rest[..n]).to_string(); + *i += n + 1; + Ok(s) + }; + loop { + let op = *p.get(i).ok_or("truncated pickle")?; + i += 1; + match op { + 0x80 => { + take(&mut i, 1)?; + } + 0x95 => { + take(&mut i, 8)?; + } + b'}' => stack.push(Obj::Dict(vec![])), + b']' => stack.push(Obj::List(vec![])), + b')' => stack.push(Obj::Tuple(vec![])), + b'(' => stack.push(Obj::Mark), + b'N' => stack.push(Obj::None), + 0x88 | 0x89 => stack.push(Obj::Bool), + b'K' => stack.push(Obj::Int(take(&mut i, 1)?[0] as i64)), + b'M' => { + let s = take(&mut i, 2)?; + stack.push(Obj::Int(u16::from_le_bytes([s[0], s[1]]) as i64)); + } + b'J' => { + let s = take(&mut i, 4)?; + stack.push(Obj::Int(i32::from_le_bytes([s[0], s[1], s[2], s[3]]) as i64)); + } + 0x8a => { + let n = take(&mut i, 1)?[0] as usize; + let s = take(&mut i, n)?; + let mut v: i64 = 0; + for (k, byte) in s.iter().take(8).enumerate() { + v |= (*byte as i64) << (8 * k); + } + if n > 0 && n < 8 && s[n - 1] & 0x80 != 0 { + v -= 1i64 << (8 * n); + } + stack.push(Obj::Int(v)); + } + b'G' => { + take(&mut i, 8)?; + stack.push(Obj::Float); + } + b'X' => { + let s = take(&mut i, 4)?; + let n = u32::from_le_bytes([s[0], s[1], s[2], s[3]]) as usize; + stack.push(Obj::Str(String::from_utf8_lossy(take(&mut i, n)?).to_string())); + } + 0x8c => { + let n = take(&mut i, 1)?[0] as usize; + stack.push(Obj::Str(String::from_utf8_lossy(take(&mut i, n)?).to_string())); + } + b'C' => { + let n = take(&mut i, 1)?[0] as usize; + take(&mut i, n)?; + stack.push(Obj::Other); + } + b'B' => { + let s = take(&mut i, 4)?; + let n = u32::from_le_bytes([s[0], s[1], s[2], s[3]]) as usize; + take(&mut i, n)?; + stack.push(Obj::Other); + } + b'c' => { + let (m, n) = (line(&mut i)?, line(&mut i)?); + stack.push(Obj::Global(m, n)); + } + 0x93 => { + let n = pop(&mut stack)?; + let m = pop(&mut stack)?; + match (m, n) { + (Obj::Str(m), Obj::Str(n)) => stack.push(Obj::Global(m, n)), + _ => return Err("pickle: bad STACK_GLOBAL".into()), + } + } + b'q' => { + let k = take(&mut i, 1)?[0] as u64; + memo.insert(k, stack.last().cloned().ok_or("pickle stack underflow")?); + } + b'r' => { + let s = take(&mut i, 4)?; + memo.insert(u32::from_le_bytes([s[0], s[1], s[2], s[3]]) as u64, stack.last().cloned().ok_or("pickle stack underflow")?); + } + 0x94 => { + let k = memo.len() as u64; + memo.insert(k, stack.last().cloned().ok_or("pickle stack underflow")?); + } + b'h' => { + let k = take(&mut i, 1)?[0] as u64; + stack.push(memo.get(&k).cloned().ok_or("pickle: bad memo")?); + } + b'j' => { + let s = take(&mut i, 4)?; + stack.push(memo.get(&(u32::from_le_bytes([s[0], s[1], s[2], s[3]]) as u64)).cloned().ok_or("pickle: bad memo")?); + } + b't' => { + let items = pop_mark(&mut stack)?; + stack.push(Obj::Tuple(items)); + } + 0x85..=0x87 => { + let n = (op - 0x84) as usize; + if stack.len() < n { + return Err("pickle stack underflow".into()); + } + let items = stack.split_off(stack.len() - n); + stack.push(Obj::Tuple(items)); + } + b'Q' => { + // Persistent id: ('storage', dtype class, key, location, numel). + let pid = pop(&mut stack)?; + let Obj::Tuple(t) = pid else { return Err("pickle: bad persistent id".into()) }; + match (t.get(1), t.get(2)) { + (Some(Obj::Global(_, dtype)), Some(Obj::Str(key))) => stack.push(Obj::Storage(dtype.clone(), key.clone())), + _ => stack.push(Obj::Other), + } + } + b'R' | 0x81 => { + let args = pop(&mut stack)?; + let f = pop(&mut stack)?; + let args = match args { + Obj::Tuple(a) => a, + _ => vec![], + }; + stack.push(match f { + Obj::Global(m, n) if m == "torch._utils" && n == "_rebuild_tensor_v2" => { + match (args.first(), args.get(1), args.get(2).and_then(int_list), args.get(3).and_then(int_list)) { + (Some(Obj::Storage(dtype, key)), Some(Obj::Int(off)), Some(size), Some(stride)) => { + Obj::Tensor { dtype: dtype.clone(), key: key.clone(), offset: usize::try_from(*off).unwrap_or(0), size, stride } + } + _ => Obj::Other, + } + } + Obj::Global(m, n) if m == "collections" && n == "OrderedDict" => Obj::Dict(vec![]), + _ => Obj::Other, + }); + } + b'b' => { + pop(&mut stack)?; + } + b's' => { + let v = pop(&mut stack)?; + let k = pop(&mut stack)?; + if let Some(Obj::Dict(d)) = stack.last_mut() { + d.push((k, v)); + } + } + b'u' => { + let items = pop_mark(&mut stack)?; + if let Some(Obj::Dict(d)) = stack.last_mut() { + let mut it = items.into_iter(); + while let (Some(k), Some(v)) = (it.next(), it.next()) { + d.push((k, v)); + } + } + } + b'a' => { + let v = pop(&mut stack)?; + if let Some(Obj::List(l)) = stack.last_mut() { + l.push(v); + } + } + b'e' => { + let items = pop_mark(&mut stack)?; + if let Some(Obj::List(l)) = stack.last_mut() { + l.extend(items); + } + } + b'.' => return pop(&mut stack), + other => return Err(format!("pickle: unsupported opcode 0x{other:02x}")), + } + } +} + +/// Element bytes and conversion of a storage dtype. +fn to_f32(dtype: &str, bytes: &[u8]) -> Option> { + match dtype { + "FloatStorage" => Some(bytes.chunks_exact(4).map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect()), + "HalfStorage" => Some(bytes.chunks_exact(2).map(|c| half_to_f32(u16::from_le_bytes([c[0], c[1]]))).collect()), + "BFloat16Storage" => Some(bytes.chunks_exact(2).map(|c| f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16)).collect()), + _ => None, + } +} + +fn half_to_f32(h: u16) -> f32 { + let (s, e, m) = ((h >> 15) as u32, ((h >> 10) & 0x1f) as u32, (h & 0x3ff) as u32); + let bits = match (e, m) { + (0, 0) => s << 31, + (0, m) => { + // Subnormal: normalise. + let mut e = 127 - 15 + 1; + let mut m = m; + while m & 0x400 == 0 { + m <<= 1; + e -= 1; + } + (s << 31) | (e << 23) | ((m & 0x3ff) << 13) + } + (0x1f, m) => (s << 31) | (0xff << 23) | (m << 13), + (e, m) => (s << 31) | ((e + 127 - 15) << 23) | (m << 13), + }; + f32::from_bits(bits) +} + +/// Read a PyTorch checkpoint's float tensors. A dict holding the state dict under `model` or +/// `state_dict` is unwrapped. +pub fn read(bytes: &[u8]) -> Result { + let entries = zip_entries(bytes)?; + let (pkl_name, pkl) = entries.iter().find(|(n, _)| n.ends_with("data.pkl")).ok_or("not a PyTorch checkpoint (no data.pkl)")?; + let prefix = pkl_name.trim_end_matches("data.pkl").to_string(); + let root = unpickle(&read_entry(bytes, pkl)?)?; + let dict = match root { + Obj::Dict(d) => d, + _ => return Err("the checkpoint is not a state dict".into()), + }; + let dict = match dict.iter().find(|(k, _)| matches!(k, Obj::Str(s) if s == "model" || s == "state_dict")) { + Some((_, Obj::Dict(inner))) => inner.clone(), + _ => dict, + }; + let mut storages: HashMap> = HashMap::new(); + let mut out = Weights::new(); + for (k, v) in dict { + let (Obj::Str(name), Obj::Tensor { dtype, key, offset, size, stride }) = (k, v) else { continue }; + if !storages.contains_key(&key) { + let e = entries.get(&format!("{prefix}data/{key}")).ok_or_else(|| format!("missing storage {key}"))?; + storages.insert(key.clone(), read_entry(bytes, e)?); + } + let Some(values) = storages.get(&key).and_then(|b| to_f32(&dtype, b)) else { continue }; + let n: usize = size.iter().product(); + // Gather (any strides) into a contiguous row-major tensor. + let mut data = Vec::with_capacity(n); + let mut idx = vec![0usize; size.len()]; + for _ in 0..n { + let at = offset + idx.iter().zip(&stride).map(|(i, s)| i * s).sum::(); + data.push(*values.get(at).ok_or_else(|| format!("{name}: data out of range"))?); + for d in (0..idx.len()).rev() { + idx[d] += 1; + if idx[d] < size[d] { + break; + } + idx[d] = 0; + } + } + out.insert(name, Tensor { shape: size, data }); + } + Ok(out) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A zip archive with stored entries (what `torch.save` writes). + fn zip(files: &[(&str, &[u8])]) -> Vec { + let mut out = vec![]; + let mut central = vec![]; + for (name, data) in files { + let off = out.len() as u32; + out.extend_from_slice(b"PK\x03\x04"); + out.extend_from_slice(&[20, 0, 0, 0, 0, 0, 0, 0, 0, 0]); + out.extend_from_slice(&[0; 4]); // crc (unchecked) + out.extend_from_slice(&(data.len() as u32).to_le_bytes()); + out.extend_from_slice(&(data.len() as u32).to_le_bytes()); + out.extend_from_slice(&(name.len() as u16).to_le_bytes()); + out.extend_from_slice(&0u16.to_le_bytes()); + out.extend_from_slice(name.as_bytes()); + out.extend_from_slice(data); + central.extend_from_slice(b"PK\x01\x02"); + central.extend_from_slice(&[20, 0, 20, 0, 0, 0, 0, 0, 0, 0, 0, 0]); + central.extend_from_slice(&[0; 4]); + central.extend_from_slice(&(data.len() as u32).to_le_bytes()); + central.extend_from_slice(&(data.len() as u32).to_le_bytes()); + central.extend_from_slice(&(name.len() as u16).to_le_bytes()); + central.extend_from_slice(&[0; 12]); + central.extend_from_slice(&off.to_le_bytes()); + central.extend_from_slice(name.as_bytes()); + } + let cd = out.len() as u32; + out.extend_from_slice(¢ral); + out.extend_from_slice(b"PK\x05\x06"); + out.extend_from_slice(&[0; 4]); + out.extend_from_slice(&(files.len() as u16).to_le_bytes()); + out.extend_from_slice(&(files.len() as u16).to_le_bytes()); + out.extend_from_slice(&(central.len() as u32).to_le_bytes()); + out.extend_from_slice(&cd.to_le_bytes()); + out.extend_from_slice(&[0, 0]); + out + } + + /// The pickle `torch.save({'w': t.t()})` writes for a 2×3 tensor stored transposed (strides + /// (1, 2)), plus an int64 counter that is skipped. + fn pickle() -> Vec { + let mut p = vec![0x80, 2, b'}', b'q', 0, b'(']; + let s = |p: &mut Vec, t: &str| { + p.push(b'X'); + p.extend_from_slice(&(t.len() as u32).to_le_bytes()); + p.extend_from_slice(t.as_bytes()); + }; + for (name, dtype, key, size, stride) in [("w", "FloatStorage", "0", [2u8, 3], [1u8, 2]), ("n", "LongStorage", "1", [1, 1], [1, 1])] { + s(&mut p, name); + // The call's argument mark, then the persistent id tuple's. + p.extend_from_slice(b"ctorch._utils\n_rebuild_tensor_v2\n(("); + s(&mut p, "storage"); + p.extend_from_slice(format!("ctorch\n{dtype}\n").as_bytes()); + s(&mut p, key); + s(&mut p, "cpu"); + p.extend_from_slice(&[b'K', 6, b't', b'Q', b'K', 0]); + p.extend_from_slice(&[b'K', size[0], b'K', size[1], 0x86, b'K', stride[0], b'K', stride[1], 0x86]); + p.extend_from_slice(&[0x89, b'c']); + p.extend_from_slice(b"collections\nOrderedDict\n"); + p.extend_from_slice(b")RtR"); + } + p.extend_from_slice(b"u."); + p + } + + #[test] + fn reads_a_state_dict_with_strides() { + let storage: Vec = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0].iter().flat_map(|v| v.to_le_bytes()).collect(); + let pkl = pickle(); + let file = zip(&[("ckpt/data.pkl", &pkl), ("ckpt/data/0", &storage), ("ckpt/data/1", &7i64.to_le_bytes())]); + let w = read(&file).unwrap(); + // Strides (1, 2) over [1..6]: rows are [1, 3, 5] and [2, 4, 6]. + assert_eq!(w.get("w"), Some(&Tensor { shape: vec![2, 3], data: vec![1.0, 3.0, 5.0, 2.0, 4.0, 6.0] })); + assert!(!w.contains_key("n"), "integer tensors are skipped"); + } + + #[test] + fn damaged_files_are_errors() { + assert!(read(b"not a zip").is_err()); + let pkl = pickle(); + let file = zip(&[("ckpt/data.pkl", &pkl[..pkl.len() - 5])]); + assert!(read(&file).is_err()); + // Truncate the storage: the tensor reads out of range. + let file = zip(&[("ckpt/data.pkl", &pkl), ("ckpt/data/0", &[0u8; 8]), ("ckpt/data/1", &[0u8; 8])]); + assert!(read(&file).is_err()); + assert_eq!(half_to_f32(0x3c00), 1.0); + assert_eq!(half_to_f32(0xc000), -2.0); + } +} diff --git a/crates/track/Cargo.toml b/crates/track/Cargo.toml index 33ffacae1..cd9daae36 100644 --- a/crates/track/Cargo.toml +++ b/crates/track/Cargo.toml @@ -11,6 +11,7 @@ workspace = true [dependencies] effectcraft-raster = { workspace = true } +effectcraft-segment = { workspace = true } rayon = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } diff --git a/crates/track/src/roto/mod.rs b/crates/track/src/roto/mod.rs index 7103624be..7138a93a3 100644 --- a/crates/track/src/roto/mod.rs +++ b/crates/track/src/roto/mod.rs @@ -19,6 +19,13 @@ //! colour models and the warped matte as a shape prior. Correction strokes on a frame re-cut //! the whole frame with the warped matte as a soft prior, and propagation restarts from there. //! +//! - **Trained models** (optional, swappable: `effectcraft_segment::MaskModel`, e.g. MobileSAM): +//! [`segment_with`] / [`propagate_with`] prompt the model with points along the strokes, or with +//! the warped matte (its box, points deep inside it and the matte itself), and its foreground +//! probability becomes a strong prior for the same graph cut, so strokes stay hard constraints +//! and edges still snap to colour. A model that fails, or whose mask disagrees with the flow, is +//! ignored for that frame. +//! //! Mattes are in layer pixels at a given scale (1 = full resolution). Strokes are stored in //! layer pixels ([`RotoData`], serde) and keyed per frame so that each frame's result is a //! deterministic function of its *chain*: the base frame's strokes, then every frame's strokes @@ -33,6 +40,7 @@ pub mod rle; use std::collections::BTreeMap; use effectcraft_raster::Image; +use effectcraft_segment::{MaskModel, Prompt}; use rayon::prelude::*; use serde::{Deserialize, Serialize}; @@ -328,6 +336,8 @@ const HARD_BG: u8 = 2; /// Graph-cut energy weights. const GAMMA: f64 = 50.0; const PRIOR: f64 = 1.0; +/// A trained model's matte as a prior: it outweighs the colour models. +const MODEL_PRIOR: f64 = 4.0; /// Fixed-point scale of the energies. const K: f64 = 16.0; const INF: i64 = 1 << 40; @@ -356,8 +366,9 @@ fn beta(img: &[[f32; 3]], free: &[bool], w: usize, h: usize) -> f64 { if s <= 1e-12 { 1.0 } else { n / (2.0 * s) } } -/// One min-cut over the `free` pixels; others keep `labels`. `prior`: P(foreground) per pixel. -fn cut(img: &[[f32; 3]], w: usize, h: usize, labels: &mut [u8], free: &[bool], hard: &[u8], fg: &Gmm, bg: &Gmm, prior: Option<&[f32]>) { +/// One min-cut over the `free` pixels; others keep `labels`. `prior`: P(foreground) per pixel +/// and how much it weighs against the colour models. +fn cut(img: &[[f32; 3]], w: usize, h: usize, labels: &mut [u8], free: &[bool], hard: &[u8], fg: &Gmm, bg: &Gmm, prior: Option<(&[f32], f64)>) { let mut idx = vec![u32::MAX; w * h]; let mut nodes = vec![]; for i in 0..w * h { @@ -375,10 +386,10 @@ fn cut(img: &[[f32; 3]], w: usize, h: usize, labels: &mut [u8], free: &[bool], h .map(|&i| { let c = img[i]; let (mut cf, mut cb) = (fg.cost(c), bg.cost(c)); - if let Some(p) = prior { + if let Some((p, wt)) = prior { let pf = (p[i] as f64).clamp(0.02, 0.98); - cf += PRIOR * -pf.ln(); - cb += PRIOR * -(1.0 - pf).ln(); + cf += wt * -pf.ln(); + cb += wt * -(1.0 - pf).ln(); } (cf, cb) }) @@ -568,6 +579,8 @@ pub struct Prior { pub bg: Gmm, /// Warped Refine Edge band. pub refine: Vec, + /// How much the prior weighs against the colour models. + pub weight: f64, } /// Segment a frame from its strokes (and a propagated prior): GrabCut iterations on a reduced @@ -593,7 +606,10 @@ pub fn segment(img: &Image, strokes: &[&Stroke], prior: Option<&Prior>, opts: &S } }) .collect(); - let soft_prior: Option> = prior.map(|p| matting::gauss(&p.soft, w, h, (opts.search_radius * opts.scale / 4.0).max(1.0))); + // A warped matte is uncertain by about the search radius: blur it. A model's matte is not. + let soft_prior: Option> = + prior.map(|p| if p.weight > PRIOR { p.soft.clone() } else { matting::gauss(&p.soft, w, h, (opts.search_radius * opts.scale / 4.0).max(1.0)) }); + let pw = prior.map_or(PRIOR, |p| p.weight); // Initial colour models: the strokes (and the propagated matte). let (fg0, bg0) = match prior { Some(p) => { @@ -666,7 +682,7 @@ pub fn segment(img: &Image, strokes: &[&Stroke], prior: Option<&Prior>, opts: &S }; let free = vec![true; cw * ch]; for it in 0..iters { - cut(&crgb, cw, ch, &mut cl, &free, &chard, &fgm, &bgm, cprior.as_deref()); + cut(&crgb, cw, ch, &mut cl, &free, &chard, &fgm, &bgm, cprior.as_deref().map(|c| (c, pw))); if it + 1 < iters { let (a, b) = fit_models(&crgb, &cl); if !a.is_empty() && !b.is_empty() { @@ -688,11 +704,11 @@ pub fn segment(img: &Image, strokes: &[&Stroke], prior: Option<&Prior>, opts: &S let sd = matting::signed_distance(&labels, w, h); let band = f as f32 + 1.5; let free: Vec = sd.iter().map(|d| d.abs() <= band).collect(); - cut(&rgb, w, h, &mut labels, &free, &hard, &fgm, &bgm, soft_prior.as_deref()); + cut(&rgb, w, h, &mut labels, &free, &hard, &fgm, &bgm, soft_prior.as_deref().map(|c| (c, pw))); } else { let free = vec![true; w * h]; for it in 0..iters { - cut(&rgb, w, h, &mut labels, &free, &hard, &fgm, &bgm, soft_prior.as_deref()); + cut(&rgb, w, h, &mut labels, &free, &hard, &fgm, &bgm, soft_prior.as_deref().map(|c| (c, pw))); if it + 1 < iters { let (a, b) = fit_models(&rgb, &labels); if !a.is_empty() && !b.is_empty() { @@ -768,7 +784,7 @@ pub fn warp_prior(prev_img: &Image, prev: &FrameSeg, next_img: &Image, opts: &Se (soft, r) }) .collect(); - Prior { soft: res.iter().map(|r| r.0).collect(), fg: prev.fg.clone(), bg: prev.bg.clone(), refine: res.iter().map(|r| r.1).collect() } + Prior { soft: res.iter().map(|r| r.0).collect(), fg: prev.fg.clone(), bg: prev.bg.clone(), refine: res.iter().map(|r| r.1).collect(), weight: PRIOR } } /// Propagate the previous frame's segmentation to the next frame: warp it, then re-cut in a @@ -788,12 +804,133 @@ pub fn propagate(prev_img: &Image, prev: &FrameSeg, next_img: &Image, strokes: & let hard = vec![0u8; w * h]; let soft = matting::gauss(&prior.soft, w, h, (r as f64 / 4.0).max(1.0)); let (fg, bg) = if prev.fg.is_empty() || prev.bg.is_empty() { fit_models(&rgb, &labels) } else { (prev.fg.clone(), prev.bg.clone()) }; - cut(&rgb, w, h, &mut labels, &free, &hard, &fg, &bg, Some(&soft)); + cut(&rgb, w, h, &mut labels, &free, &hard, &fg, &bg, Some((&soft, PRIOR))); let seeds = core_of(&prior.soft, w, h, r as f64); clean(&mut labels, &seeds, &vec![false; w * h], w, h); finish(next_img, &rgb, labels, strokes, Some(&prior.refine), opts, refine_radius) } +/// Points spread along a stroke (matte pixels), about one per stroke width, at most `max`. +fn points_along(s: &Stroke, scale: f64, max: usize) -> Vec<[f32; 2]> { + let pts: Vec<[f64; 2]> = s.points.iter().map(|p| [p[0] * scale, p[1] * scale]).collect(); + let len: f64 = pts.windows(2).map(|w| (w[1][0] - w[0][0]).hypot(w[1][1] - w[0][1])).sum(); + let step = (2.0 * s.radius * scale).max(len / max.max(1) as f64).max(1.0); + let mut out = pts.first().map(|p| vec![[p[0] as f32, p[1] as f32]]).unwrap_or_default(); + let mut carry = 0.0; + for w in pts.windows(2) { + let d = (w[1][0] - w[0][0]).hypot(w[1][1] - w[0][1]); + let mut t = step - carry; + while t <= d && out.len() < max { + let f = t / d.max(1e-9); + out.push([(w[0][0] + (w[1][0] - w[0][0]) * f) as f32, (w[0][1] + (w[1][1] - w[0][1]) * f) as f32]); + t += step; + } + carry = (carry + d) % step; + } + out +} + +/// A model's prompt from the strokes: points along the foreground and background strokes. +fn stroke_prompt(strokes: &[&Stroke], scale: f64) -> Prompt { + let mut points = vec![]; + for (kind, fg) in [(StrokeKind::Fg, true), (StrokeKind::Bg, false)] { + let ss: Vec<&&Stroke> = strokes.iter().filter(|s| s.kind == kind).collect(); + let per = (16 / ss.len().max(1)).max(1); + for s in ss { + points.extend(points_along(s, scale, per).into_iter().map(|p| (p, fg))); + } + } + Prompt { points, ..Default::default() } +} + +/// Add a prior matte to a prompt: its box (a little larger), up to four points spread through +/// its core, and the matte itself. +fn add_prior_prompt(prompt: &mut Prompt, soft: &[f32], w: usize, h: usize) -> bool { + let bin: Vec = soft.iter().map(|v| (*v >= 0.5) as u8).collect(); + let (mut x0, mut y0, mut x1, mut y1) = (usize::MAX, usize::MAX, 0usize, 0usize); + for (i, v) in bin.iter().enumerate() { + if *v != 0 { + let (x, y) = (i % w, i / w); + (x0, y0, x1, y1) = (x0.min(x), y0.min(y), x1.max(x), y1.max(y)); + } + } + if x0 > x1 { + return false; + } + let (mx, my) = (((x1 - x0) as f32 * 0.05).max(4.0), ((y1 - y0) as f32 * 0.05).max(4.0)); + prompt.bbox = Some([(x0 as f32 - mx).max(0.0), (y0 as f32 - my).max(0.0), (x1 as f32 + 1.0 + mx).min(w as f32), (y1 as f32 + 1.0 + my).min(h as f32)]); + // Farthest-point picks among the core (deep inside), starting from the deepest pixel. + let sd = matting::signed_distance(&bin, w, h); + let deep = sd.iter().copied().fold(0.0f32, f32::max); + let core: Vec = (0..w * h).filter(|i| sd[*i] >= (deep * 0.5).max(1.0)).step_by(((w * h) / 20_000).max(1)).collect(); + let mut picks: Vec = core.iter().copied().max_by(|a, b| sd[*a].total_cmp(&sd[*b])).into_iter().collect(); + while picks.len() < 4 { + let far = core.iter().copied().max_by_key(|c| picks.iter().map(|p| (c % w).abs_diff(p % w).pow(2) + (c / w).abs_diff(p / w).pow(2)).min().unwrap_or(0)); + match far { + Some(f) if !picks.contains(&f) => picks.push(f), + _ => break, + } + } + prompt.points.extend(picks.into_iter().map(|i| ([(i % w) as f32 + 0.5, (i / w) as f32 + 0.5], true))); + prompt.mask = Some(soft.to_vec()); + true +} + +/// [`segment`] with a trained model (if any): its matte for the strokes (and the propagated +/// prior) becomes a strong prior for the cut. +pub fn segment_with(img: &Image, strokes: &[&Stroke], prior: Option<&Prior>, opts: &SegOpts, refine_radius: f64, model: Option<&dyn MaskModel>) -> FrameSeg { + let (w, h) = (img.width as usize, img.height as usize); + let Some(m) = model.filter(|_| w > 0 && h > 0) else { return segment(img, strokes, prior, opts, refine_radius) }; + let mut prompt = stroke_prompt(strokes, opts.scale); + if let Some(p) = prior { + add_prior_prompt(&mut prompt, &p.soft, w, h); + } + if !prompt.points.iter().any(|(_, fg)| *fg) { + return segment(img, strokes, prior, opts, refine_radius); + } + match m.segment(&rgb_of(img), w, h, &prompt) { + Ok(soft) if soft.len() == w * h => { + let (fg, bg) = prior.map_or_else(Default::default, |p| (p.fg.clone(), p.bg.clone())); + let refine = prior.map_or_else(|| vec![0; w * h], |p| p.refine.clone()); + segment(img, strokes, Some(&Prior { soft, fg, bg, refine, weight: MODEL_PRIOR }), opts, refine_radius) + } + _ => segment(img, strokes, prior, opts, refine_radius), + } +} + +/// [`propagate`] with a trained model (if any): the model, prompted with the flow-warped matte, +/// re-segments the frame; when it disagrees with the flow (IoU under 0.5) the classical result +/// stands. +pub fn propagate_with( + prev_img: &Image, + prev: &FrameSeg, + next_img: &Image, + strokes: &[&Stroke], + opts: &SegOpts, + refine_radius: f64, + model: Option<&dyn MaskModel>, +) -> FrameSeg { + let (w, h) = (next_img.width as usize, next_img.height as usize); + let Some(m) = model.filter(|_| w > 0 && h > 0) else { return propagate(prev_img, prev, next_img, strokes, opts, refine_radius) }; + let prior = warp_prior(prev_img, prev, next_img, opts); + let mut prompt = stroke_prompt(strokes, opts.scale); + if !add_prior_prompt(&mut prompt, &prior.soft, w, h) { + return propagate(prev_img, prev, next_img, strokes, opts, refine_radius); + } + match m.segment(&rgb_of(next_img), w, h, &prompt) { + Ok(soft) if soft.len() == w * h => { + let a: Vec = soft.iter().map(|v| (*v >= 0.5) as u8).collect(); + let b: Vec = prior.soft.iter().map(|v| (*v >= 0.5) as u8).collect(); + if iou(&a, &b) < 0.5 { + return propagate(prev_img, prev, next_img, strokes, opts, refine_radius); + } + let p = Prior { soft, weight: MODEL_PRIOR, ..prior }; + segment(next_img, strokes, Some(&p), opts, refine_radius) + } + _ => propagate(prev_img, prev, next_img, strokes, opts, refine_radius), + } +} + /// Intersection over union of two binary mattes. pub fn iou(a: &[u8], b: &[u8]) -> f64 { let (mut i, mut u) = (0usize, 0usize); diff --git a/crates/track/src/roto/tests.rs b/crates/track/src/roto/tests.rs index 15bd1f68b..4b335142d 100644 --- a/crates/track/src/roto/tests.rs +++ b/crates/track/src/roto/tests.rs @@ -335,3 +335,94 @@ fn roto_perf_1080p() { let _ = matte_alpha(&rgb, &seg2, None, None, &MatteParams::default(), 20.0, 1.0); eprintln!("refine 1080p: {:?}", t.elapsed()); } + +/// A stand-in trained model: the foreground is a fixed rectangle (or nothing, or an error). +struct FakeModel(Option<[usize; 4]>, bool); + +impl effectcraft_segment::MaskModel for FakeModel { + fn info(&self) -> &'static effectcraft_segment::ModelInfo { + &effectcraft_segment::MOBILE_SAM + } + fn segment(&self, _: &[[f32; 3]], w: usize, h: usize, prompt: &Prompt) -> Result, String> { + if self.1 { + return Err("broken".into()); + } + assert!(prompt.points.iter().any(|(_, fg)| *fg), "prompted with foreground points"); + Ok((0..w * h) + .map(|i| match self.0 { + Some([x0, y0, x1, y1]) => ((i % w) >= x0 && (i % w) < x1 && (i / w) >= y0 && (i / w) < y1) as u8 as f32 * 0.98 + 0.01, + None => 0.01, + }) + .collect()) + } +} + +/// A grey bar on a two-tone background whose halves match the bar's halves: colour alone can't +/// find its edge, a model can. +#[test] +fn a_model_prior_shapes_the_cut_and_bad_models_are_ignored() { + let (w, h) = (160usize, 100usize); + let rect = [40, 30, 120, 70]; + let mut img = Image::new(w as u32, h as u32); + let mut gt = vec![0u8; w * h]; + for y in 0..h { + for x in 0..w { + let inside = x >= rect[0] && x < rect[2] && y >= rect[1] && y < rect[3]; + // Left half dark, right half light, inside or out; a faint stripe pattern inside. + let base = if x < w / 2 { 0.2 } else { 0.8 }; + let v = if inside { base + 0.02 * ((x / 3) % 2) as f32 } else { base }; + img.set(x as u32, y as u32, [v, v, v, 1.0]); + gt[y * w + x] = inside as u8; + } + } + let st = [line(StrokeKind::Fg, 0, [50.0, 50.0], [110.0, 50.0], 4.0)]; + let refs: Vec<&Stroke> = st.iter().collect(); + let opts = SegOpts::default(); + let classical = segment(&img, &refs, None, &opts, 0.0); + let good = FakeModel(Some(rect), false); + let with = segment_with(&img, &refs, None, &opts, 0.0, Some(&good)); + let (jc, jm) = (iou(&classical.matte, >), iou(&with.matte, >)); + assert!(jm > 0.95 && jm > jc + 0.2, "model {jm} vs classical {jc}"); + // A failing model: the classical result, unchanged. + let broken = FakeModel(None, true); + assert_eq!(segment_with(&img, &refs, None, &opts, 0.0, Some(&broken)), classical); + // Propagation: a model that loses the object (disagrees with the flow) is ignored. + let next = img.clone(); + let lost = FakeModel(None, false); + let a = propagate_with(&img, &with, &next, &[], &opts, 0.0, Some(&lost)); + let b = propagate(&img, &with, &next, &[], &opts, 0.0); + assert_eq!(a, b); + // A good model keeps the bar through propagation. + let kept = propagate_with(&img, &with, &next, &[], &opts, 0.0, Some(&good)); + assert!(iou(&kept.matte, >) > 0.95, "{}", iou(&kept.matte, >)); +} + +/// With the real MobileSAM weights (`EFFECTCRAFT_MOBILESAM=path/to/mobile_sam.pt`, else +/// skipped): the base frame and 20 propagated frames of the moving textured disk stay at least +/// as accurate as the classic engine. +#[test] +fn mobilesam_segments_and_propagates_the_disk() { + let Ok(path) = std::env::var("EFFECTCRAFT_MOBILESAM") else { return }; + let model = effectcraft_segment::load("mobilesam", &std::fs::read(path).unwrap()).unwrap(); + let m = Some(model.as_ref()); + let (img0, gt0) = frame_at(center(0)); + let st = base_strokes(0); + let refs: Vec<&Stroke> = st.iter().collect(); + let opts = SegOpts::default(); + let base = segment_with(&img0, &refs, None, &opts, 0.0, m); + let jb = iou(&base.matte, >0); + let (mut seg, mut cseg) = (base, segment(&img0, &refs, None, &opts, 0.0)); + let mut prev = img0; + let (mut worst, mut cworst): (f64, f64) = (1.0, 1.0); + let t = std::time::Instant::now(); + for f in 1..=20 { + let (img, gt) = frame_at(center(f)); + seg = propagate_with(&prev, &seg, &img, &[], &opts, 0.0, m); + cseg = propagate(&prev, &cseg, &img, &[], &opts, 0.0); + worst = worst.min(iou(&seg.matte, >)); + cworst = cworst.min(iou(&cseg.matte, >)); + prev = img; + } + eprintln!("MobileSAM: base IoU {jb:.4}, worst propagated {worst:.4} (classic {cworst:.4}), {:.0} ms/frame", t.elapsed().as_secs_f64() * 1000.0 / 20.0); + assert!(jb > 0.95 && worst > 0.93 && worst >= cworst - 0.01, "base {jb}, worst {worst} vs classic {cworst}"); +} diff --git a/crates/ui-egui/src/panels/settings.rs b/crates/ui-egui/src/panels/settings.rs index 1abceb569..6fcc70ca6 100644 --- a/crates/ui-egui/src/panels/settings.rs +++ b/crates/ui-egui/src/panels/settings.rs @@ -28,6 +28,8 @@ enum Act { Run(String, Value), PickFolder(String), PickProject(String), + /// Install from File (a downloaded weights file). + PickModel, } fn hex(c: [u8; 3]) -> String { @@ -61,6 +63,7 @@ fn page_ui(app: &mut EffectcraftApp, ui: &mut egui::Ui, page: &Page, cur: &Value } Item::Labels => labels_ui(app, ui, cur, acts), Item::BrowserStorage => browser_storage_ui(app, ui, t, acts), + Item::RotoModels => roto_models_ui(app, ui, t, acts), Item::AudioDevices { key } => { let devices: Vec = app.hooks.audio_devices.as_ref().map(|f| f()).unwrap_or_default(); let sel = get(key).as_str().unwrap_or("").to_string(); @@ -209,6 +212,97 @@ pub fn human_bytes(b: u64) -> String { /// Settings ▸ Disk ▸ Browser Storage (the web app's storage manager, `storage.*`): where the /// data lives, the origin's usage and quota, persistent storage, and Clear buttons. +/// Settings ▸ Roto Brush: the segmentation models (`roto.models`): choose one; download, install +/// from a file or remove the trained ones; their licence, size and source. +fn roto_models_ui(app: &mut EffectcraftApp, ui: &mut egui::Ui, t: &Tokens, acts: &mut Vec) { + let Ok(info) = app.session.execute("roto.models", json!({})) else { return }; + let web = cfg!(target_arch = "wasm32"); + let busy = info["busy"].as_str().map(str::to_string); + for m in info["models"].as_array().into_iter().flatten() { + let id = m["id"].as_str().unwrap_or_default().to_string(); + let name = m["name"].as_str().unwrap_or_default(); + let (installed, selected, active) = + (m["installed"].as_bool().unwrap_or(false), m["selected"].as_bool().unwrap_or(false), m["active"].as_bool().unwrap_or(false)); + ui.add_space(6.0); + egui::Frame::new().fill(t.field_bg).corner_radius(6.0).inner_margin(egui::Margin::same(10)).show(ui, |ui| { + ui.set_width(ui.available_width()); + ui.horizontal(|ui| { + let r = ui.radio(selected, RichText::new(name).font(Tokens::semibold(12.5))); + reg(app, &format!("settings.roto.model.{id}"), &r, name); + if r.clicked() && !selected { + acts.push(Act::Run("roto.model.select".into(), json!({"id": id}))); + } + let classic = id == effectcraft_engine::segment::CLASSICAL; + let state = match (active, installed, selected) { + (true, _, _) => "In use", + _ if classic => "Built in", + (false, true, true) if busy.is_some() => "Loading…", + (false, true, _) => "Installed", + (false, false, true) => "Not installed: Roto Brush uses the classic engine", + _ => "Not installed", + }; + ui.label(RichText::new(state).color(if active { t.accent } else { t.text_dim })); + }); + ui.label(RichText::new(m["description"].as_str().unwrap_or_default()).color(t.text_dim)); + ui.horizontal_wrapped(|ui| { + ui.label(RichText::new(format!("Licence: {}", m["licence"].as_str().unwrap_or_default())).color(t.text_dim)); + if let Some(size) = m["size"].as_u64() { + ui.label(RichText::new(format!("· {}", human_bytes(size))).color(t.text_dim)); + } + for (label, key) in [("Project page", "homepage"), ("Licence text", "licenceUrl")] { + if let Some(url) = m[key].as_str() { + let r = ui.link(label); + if r.clicked() { + ui.ctx().open_url(egui::OpenUrl::new_tab(url)); + } + } + } + }); + if id == effectcraft_engine::segment::CLASSICAL { + return; + } + ui.horizontal(|ui| { + if !installed && !web { + let r = ui.add_enabled(busy.is_none(), egui::Button::new("Download")); + reg(app, &format!("settings.roto.download.{id}"), &r, "Download"); + if r.clicked() { + acts.push(Act::Run("roto.model.download".into(), json!({"id": id}))); + } + } + if !web { + let r = ui.add_enabled(busy.is_none(), egui::Button::new("Install from File…")); + reg(app, &format!("settings.roto.install.{id}"), &r, "Install from File"); + if r.clicked() { + acts.push(Act::PickModel); + } + } + if installed { + let r = ui.add_enabled(busy.is_none(), egui::Button::new("Remove")); + reg(app, &format!("settings.roto.remove.{id}"), &r, "Remove"); + if r.clicked() { + acts.push(Act::Run("roto.model.remove".into(), json!({"id": id}))); + } + } + }); + }); + } + if let Some(b) = &busy { + ui.horizontal(|ui| { + ui.spinner(); + ui.label(b); + }); + ui.ctx().request_repaint_after(std::time::Duration::from_millis(250)); + } + if let Some(e) = info["error"].as_str() { + ui.label(RichText::new(e).color(t.danger)); + } + if web { + ui.label(RichText::new("Trained models are available in the desktop app; the browser uses the classic engine.").color(t.text_dim)); + } else if let Some(f) = info["folder"].as_str() { + ui.label(RichText::new(format!("Models folder: {f}")).color(t.text_faint)); + } +} + fn browser_storage_ui(app: &mut EffectcraftApp, ui: &mut egui::Ui, t: &Tokens, acts: &mut Vec) { let Some(host) = app.session.storage.clone() else { return }; let info = host.info(); @@ -419,6 +513,14 @@ pub fn show(app: &mut EffectcraftApp, ctx: &egui::Context, t: &Tokens) { app.session.prefs_changed(); } } + Act::PickModel => { + let picked = app.hooks.pick_files.as_ref().map(|f| f(&["pt"])).unwrap_or_default(); + if let Some(path) = picked.into_iter().next() + && let Err(e) = crate::menus::invoke(app, ctx, "roto.model.install", json!({"path": path})) + { + app.ui.status = e; + } + } } } if cancel { diff --git a/docs/agents.md b/docs/agents.md index 37c39c992..3d7250e07 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -295,6 +295,21 @@ state, job progress and, for a frame, the matte's area, centroid, RLE matte and reference. Matte settings are ordinary properties (`effects/#1/rotoBrushMatte/searchRadius`, `effects/#1/refineEdgeMatte/decontaminateEdgeColors`, …). +Trained model (Roto Brush 2.0 / 3.0): +- `roto.models` lists the built-in engine and every registered model, with its 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}` chooses which one Roto Brush uses. +- `roto.model.remove {id}` deletes the installed weights. +- The `version` property (`effects/#1/version`: 0 = 1.0, 1 = 2.0, 2 = 3.0) picks between the + classic engine (1.0) and the chosen model (2.0, 3.0). + +```sh +effectcraft-cli run clip.ecproj roto.model.install '{"path":"mobile_sam.pt"}' roto.model.select '{"id":"mobilesam"}' \ + roto.stroke '{"layer":"#1","points":[[300,200],[360,230]],"radius":10}' roto.propagate '{"layer":"#1","wait":true}' --save +``` + ```sh effectcraft-cli run clip.ecproj roto.stroke '{"layer":"#1","points":[[300,200],[360,230]],"radius":10}' \ roto.stroke '{"layer":"#1","kind":"bg","points":[[40,40],[600,40]],"radius":12}' \ diff --git a/docs/architecture.md b/docs/architecture.md index a03ff5876..3ca8343c5 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -22,6 +22,7 @@ cpal or muda. Everything in L0 to L4, the egui UI and the web app also build for | L0 | `av1enc` | AV1 encoder from the AV1 bitstream specification: Main profile 8/10-bit 4:2:0, key + inter frames (quarter-pel motion), deblocking, for MP4 (`av01`) and WebM export | | L1 | `raster` | Premultiplied float images, sampling, affine and projective warps, blurs, compositing (parallel with rayon) | | L1 | `keyframe` | Animated values, keyframes with temporal ease and spatial Bezier, roving, hold, velocity | +| L1 | `segment` | Swappable trained segmentation models for Roto Brush: the `MaskModel` interface, a registry of open-licensed models (licence, URL, size, SHA-256), a PyTorch checkpoint reader, SHA-256, and MobileSAM (TinyViT + Segment Anything decoder) on rayon and ndarray's safe GEMM | | L1 | `psd` | Photoshop PSD/PSB reader (layers, groups, masks, blend modes, text, layer effects, adjustment layers; 8/16/32-bit; RGB/CMYK/Gray/Lab) and a minimal writer, from Adobe's published format specification | | L1 | `path` | Bezier paths, path operators (trim, offset, round corners, zig zag, twist, merge…), stroking, coverage masks | | L2 | `project` | The document: items, compositions, layers, the property tree, render queue model, `.ecproj` serde | @@ -392,7 +393,10 @@ parameter with the base frame and segmentation span. Each frame is segmented by (Boykov–Jolly hard constraints, GrabCut colour mixtures, our own Boykov–Kolmogorov max-flow, coarse to fine) and propagated to the next frame by warping the matte with block optical flow and re-cutting in a Search Radius band; frames with correction strokes are re-cut with the warped -matte as a soft prior. Refine Edge bands get guided-filter + closed-form matting and +matte as a soft prior. With Version 2.0 / 3.0 and a trained model chosen (`effectcraft_segment`, +Settings ▸ Roto Brush), the model's foreground probability, prompted by the strokes or by the +warped matte, becomes a strong prior for the same cut (falling back to the classic result when it +disagrees with the flow); the model's id is part of the chain seed. Refine Edge bands get guided-filter + closed-form matting and foreground-colour decontamination. Segmentations are derived data: each frame's chain key (Input Key, settings, strokes from the base frame out) indexes a process-wide cache that renders fill on demand (`self_at`) and `roto.propagate` fills in the background; `Session::edit` keeps the Input diff --git a/docs/parity.md b/docs/parity.md index 797111275..5b5fef5be 100644 --- a/docs/parity.md +++ b/docs/parity.md @@ -2,17 +2,17 @@ How close EffectCraft is to After Effects 2026, feature by feature, and how much work is left. -## Current status (audit at commit `d39c0e8`, 4 October 2026; updated 5 October for M3.9–M4.12) +## Current status (audit at commit `d39c0e8`, 4 October 2026; updated 5 October for M3.9–M4.12 and M13.35) | Measure | Value | |---|---| -| **Feature parity, weighted by tier** (P0 ×3, P1 ×2, P2 ×1) | **≈ 99%** counting every partial feature as half done; ≈ 99.8% using per-feature fractions | -| Unweighted | ≈ 98.4% (half-credit) / 99.6% (fractions) | -| P0 / P1 / P2 | 100% / 98% / 90% (half-credit); 100% / 99.7% / 97.5% (fractions) | -| Features done / partial / missing | 89 / 3 / 0 of 92 | +| **Feature parity, weighted by tier** (P0 ×3, P1 ×2, P2 ×1) | **≈ 99.4%** counting every partial feature as half done; ≈ 99.9% using per-feature fractions | +| Unweighted | ≈ 98.9% (half-credit) / 99.9% (fractions) | +| P0 / P1 / P2 | 100% / 98% / 95% (half-credit); 100% / 99.8% / 99% (fractions) | +| Features done / partial / missing | 90 / 2 / 0 of 92 (MSK-4 Roto Brush done in M13.35) | | **Effects** | **306** effects, every one implemented in full ([effects.md](effects.md)); 280 run on the GPU | | Disabled menu entries left | 1, on purpose: Import ▸ Vanishing Point (.vpe), whose format has no public specification | -| Remaining work | ≈ 8.5 agent-hours at the pace measured so far (≈ 26 on the conservative audit scale); ≈ 3.5 without the learned-model items | +| Remaining work | ≈ 8.5 agent-hours at the audit (≈ 26 on the conservative audit scale), before M13.35 landed Roto Brush's trained model; the learned face tracker is the one learned-model item left | | **Wall-clock estimate** | **≈ 2–3 hours** with five agents in parallel; ≈ 1 hour without the learned-model items | Partial features: @@ -20,7 +20,6 @@ Partial features: | id | Tier | Done | What is missing | |---|---|---|---| | EFF-5 GPU effects | P1 | 0.98 | Fractal (its escape iteration needs f64: WGSL has none and Metal's fast math defeats double-f32 emulation) and the analysis / tool effects (Camera-Shake Deblur, Detail-preserving Upscale, Puppet, Rolling Shutter Repair, Warp Stabilizer, Mocha Shape, Roto Brush, Match Grain, Paint, Camera Tracker, Face Tracker) render on the CPU. M13.28 moved the remaining pixel effects (CC blur / glass family, Color Link, Cineon, HDR, Inner/Outer Key, Basic 3D…) and the geometry and audio generators (Lightning, Advanced Lightning, Beam, Lens Flare, Radio Waves, Vegas, Stroke, Scribble, Write-on, Paint Bucket, Eyedropper Fill, CC Glue Gun, CC Threads, Audio Spectrum / Waveform, Basic / Path Text) onto the GPU; M13.29 the particle effects (CC Particle World / Systems II, Particle Playground, CC Ball Action, CC Pixel Polly, CC Scatterize) and Curl Noise, and removed the Shatter wireframe, Foam extras, Timewarp matte, strong Warp Fisheye / Twist and custom `.ocio` fallbacks. The 13 audio filters stay on the CPU by design | -| MSK-4 Roto Brush | P2 | 0.85 | Segmentation is classical (graph cut + optical flow), not a learned model | | TRK-3 Face tracking | P2 | 0.9 | Classical fitter: weak on profile and occluded faces; Rolling Shutter Ripple is approximated | Since this audit: the project operations (M3.9–M3.14) and preview (M4.5–M4.12) passes fixed @@ -165,7 +164,7 @@ the disk cache in the formats wave; the 3D Camera Tracker in M12.6.) | Effects | ≈ 85% | 3.0 | GPU versions of the remaining effects (280 run on the GPU since M13.29, EFF-5: M13.28 added the remaining pixel effects and the geometry and audio generators (Lightning, Advanced Lightning, Beam, Lens Flare, Radio Waves, Vegas, Stroke, Scribble, Write-on, Paint Bucket, Eyedropper Fill, CC Glue Gun, CC Threads, Audio Spectrum / Waveform, Basic / Path Text); M13.29 added the particle effects (CC Particle World / Systems II, Particle Playground, CC Ball Action, CC Pixel Polly, CC Scatterize, Curl Noise) and removed the fallbacks of Shatter's wireframe views, Foam's extras, Timewarp's Matte Layer, strong Warp Fisheye / Twist bends and custom `.ocio` configs; M13.23 added the 3D Channel and Immersive Video families, Apply Color LUT, the OCIO effects, Color Profile Converter and the simulations' render passes (Card Wipe included); M13.22 added the CC light family, the CC transitions, Block Dissolve, Radial Shadow, CC Bender / Blobbylize / Cylinder / Sphere / Spotlight / Environment, 3D Glasses, Numbers, Timecode and the time effects joined the 166 of M13.13; Fractal and the analysis / tool effects still render on the CPU) and the missing controls listed as partial in [effects.md](effects.md) (every After Effects effect exists since M9.11, M12.5 and M12.6; parameter names, order, twirl-downs, popups, units and defaults were aligned in M9.12) | | Interface | 75% | 5.0 | more Learn tutorials and pixel-level fidelity of dialogs (View ▸ New Viewer landed in M13.32; an automation id on every control of Composition and Solid Settings, Group / Ungroup Shapes in the Layer menu, Window ▸ Learn, the flowchart shortcuts and Reveal in Explorer landed in M3.10 and M3.13; the Home ▸ Templates gallery (eight original built-in templates, user templates from File ▸ Save as Template…) and View ▸ Simulate Output ▸ My Custom RGB… landed in M13.25; Timeline layer reordering by drag, a non-snapping viewer pan, a working rename field and twirl arrows, and the full set of property reveal shortcuts — double presses, Alt+Shift keyframes, Ctrl+` — landed in M13.17–M13.20; the Home ▸ Learn tab with interactive tutorials and a UI fidelity pass landed in M13.5 UI completion; visual editors for Lumetri RGB / hue-saturation curves, Colorama's output cycle wheel, Glow's colour map and Reshape's correspondence points (viewer handles), a scrolling Preview panel and AE-style Composition / Timeline tabs (close, label swatch, viewer lock) landed in M13.10; Timeline outline and Project panel columns scroll horizontally, the Layer Style dialog, ROI resize handles, Pan Behind snapping and 3D Reference Axes landed in M13.5; native macOS menu bar, Timeline columns/search/reveal-add, Home screen with recent projects and all AE workspaces landed; viewer rulers/snapping/channels/snapshots landed in M0.13) | | Project | ≈ 68% | 5.5 | OCIO displays beyond the built-in tone map; relative footage paths and relinking when a project moves (unsaved-changes prompts, a modified mark that follows undo and project files that name their version and always reopen landed in M3.9 and M3.14; auto-save and crash recovery in M3.7; Color Engine with OCIO/ACES working spaces, HDR compand/tone mapping, Rec. 2100 PQ/HLG output, Feet + Frames, display colour management, Simulate Output and the locked viewer landed in M7.7; proxies and Interpret Footage fields / pixel aspect / alpha guess landed: PRJ-8, PRJ-3) | -| Masks & roto | 74% | 5.0 | Roto Brush's learned (3.0) segmentation model (variable-width mask feather points with the Mask Feather tool landed in M13.5; mask tracking and Mask Interpolation landed in M6.6; Roto Brush & Refine Edge with graph-cut segmentation, flow propagation, edge matting, decontamination and Freeze in M6.7) | +| Masks & roto | 74% | 5.0 | (audit estimates; Roto Brush 2.0 / 3.0's trained model, MobileSAM in the swappable `segment` module, has since landed in M13.35; variable-width mask feather points with the Mask Feather tool landed in M13.5; mask tracking and Mask Interpolation landed in M6.6; Roto Brush & Refine Edge with graph-cut segmentation, flow propagation, edge matting, decontamination and Freeze in M6.7) | | Preview | 77% | 2.8 | GPU effects for the audio-driven effects and the stepped particle systems; content-keyed RAM preview frames (an edit drops every comp's), audio scrubbing, more snapshot slots (RAM preview keys with the render options, least-recently-shown eviction, Auto resolution while playing, Cache Before Playback that fits, an achieved-fps readout, Cache Frames When Idle and purges that clear what they name landed in M4.5–M4.8; 3D channel, VR, OCIO, Card Wipe and the simulations' render passes run on the GPU since M13.23; Advanced 3D layers with blend modes / track mattes / Preserve Transparency and environment backgrounds composite on the GPU since M13.22; pooled GPU textures and fused quantisation since M13.13: Lower Third GPU warm 21 → 4 ms/frame, and Auto now picks the CPU or the GPU per comp from measured frame times; Advanced 3D runs — raster, motion blur, iris depth of field, compositing — and wireframes run on the GPU since M13.11: CPU warm 3394 → GPU warm 290 ms/frame at 1920×1080 (11.7×; Half 968 → 85 ms), GPU ≠ CPU on 0.005 % of pixels, on an M4 Pro under load; Classic 3D bokeh depth of field — iris shapes, highlights, fringe, progressive blur on tilted planes — runs in WGSL since M13.6: 3D Showcase GPU warm 205 → 54 ms/frame at full size on an M4 Pro; the Preview panel's five shortcuts with their own Include / Loop / Cache Before Playback / Range / Play From / Frame Rate / Skip / Resolution / Full Screen / stop options and `playback.settings.get/set` landed in M13.8; Classic 3D runs and adjustment layers composite on the GPU since M12.7; persistent disk cache with the blue cache bar landed; region of interest, snapshots, exposure and Fast Previews landed in M0.13) | | Tracking | ≈ 96% | 0.3 | face tracking is a classical (skin model + feature components + shape model) fitter, not a learned detector: profile views and occluded faces are weak; Rolling Shutter Ripple is approximated by Subspace Warp's mesh density (face tracking (Outline Only / Detailed Features with Face Track Points and Extract & Copy Face Measurements), Subspace Warp's content-preserving mesh warp on subspace-smoothed trajectories, and radial lens distortion (k1, k2) in the camera bundle adjustment with Undistort Footage landed in M13.3; Rolling Shutter Repair in M9.11; point tracker, mask tracking, Warp Stabilizer and the 3D Camera Tracker in M6.x / M12.5 / M12.6) | | Paint | 0% | 4.0 | Brush, Clone Stamp, Eraser | @@ -434,6 +433,29 @@ stretch) and shown within its In–Out range. Hovering names them ("beat (marker double-click opens the nested comp at the marker. Agents read the same list with `markers.nested {layer}`. +Since (M13.35): Roto Brush can use a trained segmentation model; MSK-4 is done. +- The model sits behind a swappable interface: `effectcraft_segment::MaskModel` plus a registry + where every entry must have an open-source licence compatible with ours, a source URL, a size + and a SHA-256. +- The first model is **MobileSAM** (Apache-2.0: a TinyViT-5M encoder with Segment Anything's + decoder), running in plain Rust (rayon and ndarray's safe GEMM, no ML framework). It reads the + official `mobile_sam.pt` checkpoint directly and checks it against a pinned SHA-256. +- Roto Brush **Version 2.0 / 3.0** uses the model chosen in **Settings ▸ Roto Brush**; Version + 1.0 and an uninstalled model use the classic graph cut. +- On the base frame, points along the strokes prompt it. When propagating, the flow-warped matte + prompts it (its box, points deep inside it and the matte itself). The model's probability map + becomes a strong prior for the same graph cut, so strokes stay hard constraints and edges still + snap to colour. A model whose mask disagrees with the flow (IoU under 0.5) is ignored for that + frame. +- Weights are never bundled. Download (verified; uses the system `curl`) or Install from File + puts them in the settings' `models` folder, and they load in the background. +- On the moving-disc test sequence it scores IoU 0.989 on the base frame and at worst 0.980 over + 20 propagated frames (classic: 0.973). The encoder takes about 0.7 s a frame on a 32-thread + CPU, 1.3 s on 4 threads; its embedding is cached, so new strokes on the same frame only rerun + the decoder (about 0.1 s). +- Agents use `roto.models`, `roto.model.select`, `roto.model.download`, `roto.model.install` and + `roto.model.remove`. The browser build uses the classic engine for now. + Since (M13.32): View ▸ New Viewer (Alt+Shift+N, also in the Composition panel menu) opens another Composition viewer as its own panel, as in After Effects. The viewer in use is locked, so it keeps its comp. One viewer is active: it is the interactive one and shows the active comp. diff --git a/docs/preferences.md b/docs/preferences.md index 60b380843..2dc8e57c0 100644 --- a/docs/preferences.md +++ b/docs/preferences.md @@ -139,6 +139,10 @@ Commands (CLI, MCP, control channel): the Timeline's expression editor - `scripting.errorBanner`: a banner along the bottom of the Composition panel names the first failing expression (click: reveal it) +- `roto.model` (Settings ▸ Roto Brush): the segmentation model Roto Brush 2.0 / 3.0 use, + `classical` (built in) or `mobilesam` (once installed). The page lists every registered model + with its licence, size, source and status, with Download, Install from File… and Remove + (`roto.models` / `roto.model.*`); weights live in the `models` folder next to the settings ### Display-only settings diff --git a/xtask/src/main.rs b/xtask/src/main.rs index e7c6a38c6..051f6ecb7 100644 --- a/xtask/src/main.rs +++ b/xtask/src/main.rs @@ -32,6 +32,7 @@ const LAYERS: &[(&str, u8)] = &[ ("keyframe", 1), ("path", 1), ("psd", 1), + ("segment", 1), ("project", 2), ("text", 2), ("effects", 2),