mirror of
https://github.com/storytold/effectcraft.git
synced 2026-10-14 19:27:05 -04:00
Merge pull request #61 from storytold/m13-face-ml
M13.36: face tracking trained model (MediaPipe Face Landmarker, pure Rust)
This commit is contained in:
33 files changed
+2298
-485
No files matched your search
@@ -70,6 +70,34 @@ Screenshots of EffectCraft itself, rendered headlessly, showing only procedurall
|
||||
(sidecars `.attribution`).
|
||||
- Example effect plug-in: `examples/plugins/posterize-bands` (plug-in API v1, WebAssembly).
|
||||
|
||||
## Trained models (optional downloads, not bundled)
|
||||
|
||||
The repository and the installers contain no model weights. These models are downloaded only when
|
||||
a user asks (Settings ▸ Roto Brush, Settings ▸ Face Tracking, or the `roto.model.*` /
|
||||
`face.model.*` commands), unmodified from their authors' official location, checked against a
|
||||
pinned SHA-256 and stored in the user's `models` folder with a `<file>.NOTICE.txt` naming their
|
||||
authors, licence and source. EffectCraft runs them with its own Rust implementation and is not
|
||||
affiliated with or endorsed by their authors.
|
||||
|
||||
| Model | Authors | Source | Licence |
|
||||
|---|---|---|---|
|
||||
| MobileSAM (`mobile_sam.pt`) | Chaoning Zhang et al., building on Meta AI's Segment Anything and Microsoft's TinyViT | https://github.com/ChaoningZhang/MobileSAM | Apache-2.0 |
|
||||
| MediaPipe Face Landmarker (`face_landmarker.task`: BlazeFace short-range detector, Face Mesh V2, blendshapes) | Google LLC | https://developers.google.com/edge/mediapipe/solutions/vision/face_landmarker (model cards: https://storage.googleapis.com/mediapipe-assets/MediaPipe%20BlazeFace%20Model%20Card%20(Short%20Range).pdf, https://storage.googleapis.com/mediapipe-assets/Model%20Card%20MediaPipe%20Face%20Mesh%20V2.pdf) | Apache-2.0 |
|
||||
|
||||
## Formats and parameters from other open-source projects
|
||||
|
||||
Written in Rust for EffectCraft; no code is copied. They follow published specifications and
|
||||
configuration from these Apache-2.0 projects:
|
||||
|
||||
- `crates/segment/src/tflite.rs` reads TensorFlow Lite files using the field layout of TensorFlow
|
||||
Lite's `schema.fbs` (The TensorFlow Authors, https://github.com/tensorflow/tensorflow,
|
||||
Apache-2.0).
|
||||
- `crates/segment/src/mediapipe.rs` follows MediaPipe's published face landmarker graph: the SSD
|
||||
anchor layout, detection decoding, crop sizes and the face-mesh point indices it uses (Google
|
||||
LLC, https://github.com/google-ai-edge/mediapipe, Apache-2.0).
|
||||
- `crates/segment/src/mobilesam.rs` follows the MobileSAM, TinyViT and Segment Anything
|
||||
architectures and their Apache-2.0 reference code.
|
||||
|
||||
## First-party brand marks
|
||||
|
||||
The ArtCraft name, wordmark and logos are trademarks of the ArtCraft Team. They are not open source and are not
|
||||
|
||||
@@ -206,9 +206,9 @@ what breaks. It isn't yet a replacement for After Effects on client work. In par
|
||||
After Effects automatically, so please report differences.
|
||||
- **macOS is the most tested platform.** Linux and Windows users have hit basic interaction
|
||||
bugs, which we're fixing first.
|
||||
- **Face tracking uses classical methods.** Roto Brush can use a trained model (MobileSAM, an
|
||||
optional 40 MB download in Settings ▸ Roto Brush); it is new and not yet compared with After
|
||||
Effects.
|
||||
- **The trained models are new.** Roto Brush can use MobileSAM (an optional 40 MB download in
|
||||
Settings ▸ Roto Brush) and face tracking can use Google's MediaPipe Face Landmarker (an optional
|
||||
3.8 MB download in Settings ▸ Face Tracking). Neither is compared with After Effects yet.
|
||||
|
||||
The [ROADMAP](ROADMAP.md) has the plan and [docs/gaps.md](docs/gaps.md) the full, honest
|
||||
assessment. The web build ([docs/web.md](docs/web.md)) runs the full app in the browser.
|
||||
|
||||
+5
-4
@@ -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 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) |
|
||||
| M13 | Puppet tools, paint, Roto Brush, motion tracking, a plug-in API, Timeline depth | Done (trained models, open source, downloaded on demand, in a swappable model module: MobileSAM for Roto Brush 2.0 / 3.0, MediaPipe Face Landmarker for face tracking) |
|
||||
| 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 |
|
||||
|
||||
@@ -40,8 +40,9 @@ in [docs/gaps.md](docs/gaps.md).
|
||||
and third-party After Effects plug-ins can't run.
|
||||
- **Reliability is uneven across platforms.** Early Linux users hit basic problems: viewer
|
||||
panning, panel docking, drag-and-drop import.
|
||||
- **AI tools are behind.** Face tracking uses classical methods. Roto Brush's trained model
|
||||
(MobileSAM, optional, M13.35) is new and not yet compared with After Effects.
|
||||
- **AI tools are new.** Roto Brush's trained model (MobileSAM, M13.35) and face tracking's
|
||||
(MediaPipe Face Landmarker, M13.36) are optional downloads, not yet compared with After
|
||||
Effects.
|
||||
|
||||
## Where we're going
|
||||
|
||||
@@ -55,7 +56,7 @@ In priority order (details and "done" criteria in [docs/gaps.md](docs/gaps.md)):
|
||||
| G4 | Open After Effects projects (`.aep` / `.aepx`, pending an owner decision on clean-room scope); relinking moved footage | The biggest barrier to switching |
|
||||
| G5 | Performance benchmarks at 1080p and 4K on real-world-sized projects | Unknown today |
|
||||
| G6 | Media depth: encoder efficiency, camera and phone formats | Files are larger than from mature encoders; camera formats are unverified |
|
||||
| G7 | Learned models: Roto Brush started (M13.35: MobileSAM, open source, optional download, swappable model module); face tracking still classical | Quality is well below After Effects' and unmeasured |
|
||||
| G7 | Learned models: Roto Brush (M13.35: MobileSAM) and face tracking (M13.36: MediaPipe Face Landmarker), open source, optional downloads, swappable model module; next, measure them | Quality against After Effects is unmeasured |
|
||||
| G8 | Plug-in ecosystem on our WebAssembly plug-in API | After Effects plug-ins can't run here |
|
||||
| G9 | Localisation, accessibility, user documentation | Reach beyond English-speaking power users |
|
||||
|
||||
|
||||
@@ -394,16 +394,18 @@ fn high_level_tools_add_effect_and_report_state() {
|
||||
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.
|
||||
/// Trained models over MCP (Roto Brush and face tracking): listed with their authors and
|
||||
/// licence, chosen, refused when unknown.
|
||||
#[test]
|
||||
fn roto_models_over_mcp() {
|
||||
fn trained_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:?}");
|
||||
for (prefix, id) in [("roto", "mobilesam"), ("face", "mediapipe-face")] {
|
||||
let list = call_json(&mut s, "execute_command", json!({"command": format!("{prefix}.models"), "params": {}}));
|
||||
let models = list["models"].as_array().unwrap();
|
||||
assert!(models.iter().any(|m| m["id"] == id && m["licence"] == "Apache-2.0" && m["authors"].is_string() && m["installed"] == false), "{list}");
|
||||
let r = call_json(&mut s, "execute_command", json!({"command": format!("{prefix}.model.select"), "params": {"id": id}}));
|
||||
assert_eq!(r["model"], id);
|
||||
let (c, err) = call(&mut s, "execute_command", json!({"command": format!("{prefix}.model.select"), "params": {"id": "nope"}}));
|
||||
assert!(err && c[0]["text"].as_str().unwrap().contains("unknown model"), "{c:?}");
|
||||
}
|
||||
}
|
||||
@@ -187,6 +187,7 @@ fn track_mask(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
path,
|
||||
method,
|
||||
times: times.clone(),
|
||||
face_model: if method.is_face() { s.models.face() } else { None },
|
||||
};
|
||||
s.state.mask_track_method = method;
|
||||
let wait = b_p(p, "wait").unwrap_or(false);
|
||||
@@ -195,6 +196,8 @@ fn track_mask(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
Ok(json!({
|
||||
"mask": uid,
|
||||
"method": method.id(),
|
||||
// The engine face tracking tries first (the classic one when the model finds no face).
|
||||
"faceModel": method.is_face().then(|| s.models.face().map_or(effectcraft_segment::CLASSICAL, |m| m.info().id)),
|
||||
"frames": times.len() - 1,
|
||||
"start": times[0].seconds(),
|
||||
"end": times.last().map(|t| t.seconds()),
|
||||
|
||||
@@ -133,7 +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(crate::models::specs());
|
||||
v.extend(edit::specs());
|
||||
v.extend(comp::specs());
|
||||
v.extend(layer::specs());
|
||||
|
||||
@@ -26,6 +26,7 @@ pub mod mask_track;
|
||||
pub mod media_browser;
|
||||
pub mod media_cache;
|
||||
pub mod menus;
|
||||
pub mod models;
|
||||
pub mod offload;
|
||||
pub mod perf;
|
||||
pub mod prefs;
|
||||
@@ -34,7 +35,6 @@ 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;
|
||||
@@ -349,8 +349,8 @@ 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::RotoJob>,
|
||||
/// Roto Brush's trained models (Settings ▸ Roto Brush).
|
||||
pub roto_models: roto_models::Models,
|
||||
/// Trained models for Roto Brush and face tracking (Settings ▸ Roto Brush / Face Tracking).
|
||||
pub models: models::Models,
|
||||
/// Where installed models live (hosts set it; else `models` next to the settings).
|
||||
pub models_dir: Option<std::path::PathBuf>,
|
||||
/// Roto Brush instances edited since their last propagation: (comp, layer, effect uid).
|
||||
@@ -469,7 +469,7 @@ impl Default for Session {
|
||||
camera_job: None,
|
||||
camera_pending: vec![],
|
||||
roto_job: None,
|
||||
roto_models: Default::default(),
|
||||
models: Default::default(),
|
||||
models_dir: None,
|
||||
roto_pending: vec![],
|
||||
events: vec![],
|
||||
@@ -1087,10 +1087,10 @@ pub fn font_menu(prefs: &prefs::Prefs) -> Vec<FontRow> {
|
||||
#[cfg(test)]
|
||||
mod tests_history;
|
||||
#[cfg(test)]
|
||||
mod tests_models;
|
||||
#[cfg(test)]
|
||||
mod tests_paint;
|
||||
#[cfg(test)]
|
||||
mod tests_panels;
|
||||
#[cfg(test)]
|
||||
mod tests_roto;
|
||||
#[cfg(test)]
|
||||
mod tests_roto_models;
|
||||
@@ -198,6 +198,8 @@ pub(crate) struct MaskWork {
|
||||
pub method: MaskMethod,
|
||||
/// Comp times; the first is the start frame.
|
||||
pub times: Vec<Tick>,
|
||||
/// Face tracking's trained model (Settings ▸ Face Tracking), if one is in use.
|
||||
pub face_model: Option<Arc<dyn effectcraft_segment::face::FaceModel>>,
|
||||
}
|
||||
|
||||
async fn run_work(w: MaskWork, shared: &MaskTrackShared) {
|
||||
@@ -275,7 +277,7 @@ async fn run_face<F: std::future::Future<Output = Option<(Arc<effectcraft_raster
|
||||
) {
|
||||
let detailed = w.method == MaskMethod::FaceDetailed;
|
||||
let b = trk::mask::bounds(&flatten(&w.path, 8));
|
||||
let Some((mut tracker, fit)) = trk::face::FaceTracker::new(&trk::Frame { img: &first.0, offset: first.1 }, b) else {
|
||||
let Some((mut tracker, fit)) = trk::face::FaceTracker::new_with(&trk::Frame { img: &first.0, offset: first.1 }, b, w.face_model.clone()) else {
|
||||
let mut s = lock(&shared.state);
|
||||
s.error = Some("no face found inside the mask (draw the mask around a face)".into());
|
||||
s.finished = true;
|
||||
|
||||
@@ -0,0 +1,463 @@
|
||||
//! Trained models (Settings ▸ Roto Brush and Settings ▸ Face Tracking; `roto.model*` and
|
||||
//! `face.model*` commands): which models the registry offers (`effectcraft_segment::MODELS`, all
|
||||
//! open source), which are installed (in the `models` folder), installing one from a file or
|
||||
//! downloading the official weights (verified against the registry's SHA-256 either way, and kept
|
||||
//! with a notice naming their authors and licence), choosing one per task, and loading it in the
|
||||
//! background: into the Roto Brush effect (`effects::roto::set_model`) or the face tracker
|
||||
//! ([`Models::face`]).
|
||||
//!
|
||||
//! 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 uses the built-in engines
|
||||
//! only, for now.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use effectcraft_segment::face::FaceModel;
|
||||
use effectcraft_segment::{self as seg, Loaded, ModelInfo, Task};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::commands::{CommandSpec, always, b_p, bad, str_p};
|
||||
use crate::{EngineError, Result, Session, cmd, query};
|
||||
|
||||
/// What a task's background worker reports.
|
||||
#[derive(Clone, Debug, Default, PartialEq)]
|
||||
pub struct Status {
|
||||
/// What is happening now ("Downloading MobileSAM…").
|
||||
pub busy: Option<String>,
|
||||
/// The model in use.
|
||||
pub active: Option<String>,
|
||||
/// The last failure, for the Settings page.
|
||||
pub error: Option<String>,
|
||||
/// A finished job's message and whether it failed, for a toast (taken by
|
||||
/// [`Session::poll_models`]).
|
||||
pub done: Option<(String, bool)>,
|
||||
}
|
||||
|
||||
/// The session's trained models: per task, the worker's status; and the face model in use.
|
||||
#[derive(Default)]
|
||||
pub struct Models {
|
||||
roto: Arc<Mutex<Status>>,
|
||||
face: Arc<Mutex<Status>>,
|
||||
face_model: FaceSlot,
|
||||
}
|
||||
|
||||
impl Models {
|
||||
fn lane(&self, task: Task) -> &Arc<Mutex<Status>> {
|
||||
match task {
|
||||
Task::Mask => &self.roto,
|
||||
Task::Face => &self.face,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn status(&self, task: Task) -> Status {
|
||||
self.lane(task).lock().map(|s| s.clone()).unwrap_or_default()
|
||||
}
|
||||
|
||||
/// The face model face tracking uses (`None`: the classic tracker).
|
||||
pub fn face(&self) -> Option<Arc<dyn FaceModel>> {
|
||||
self.face_model.lock().ok().and_then(|m| m.clone())
|
||||
}
|
||||
|
||||
fn set_active(&self, task: Task, model: Option<Loaded>) {
|
||||
activate(self.lane(task), &self.face_model, task, model);
|
||||
}
|
||||
}
|
||||
|
||||
type FaceSlot = Arc<Mutex<Option<Arc<dyn FaceModel>>>>;
|
||||
|
||||
/// Put a loaded model to use (or drop `task`'s): Roto Brush's lives in the effect, the face
|
||||
/// tracker's in `face`.
|
||||
fn activate(lane: &Arc<Mutex<Status>>, face: &FaceSlot, task: Task, model: Option<Loaded>) {
|
||||
let id = model.as_ref().map(|m| match m {
|
||||
Loaded::Mask(m) => m.info().id.to_string(),
|
||||
Loaded::Face(m) => m.info().id.to_string(),
|
||||
});
|
||||
match (task, model) {
|
||||
(Task::Mask, Some(Loaded::Mask(m))) => crate::effects::roto::set_model(Some(m)),
|
||||
(Task::Mask, _) => crate::effects::roto::set_model(None),
|
||||
(Task::Face, m) => {
|
||||
if let Ok(mut slot) = face.lock() {
|
||||
*slot = match m {
|
||||
Some(Loaded::Face(f)) => Some(f),
|
||||
_ => None,
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
set(lane, |s| s.active = id);
|
||||
}
|
||||
|
||||
fn set(status: &Arc<Mutex<Status>>, f: impl FnOnce(&mut Status)) {
|
||||
if let Ok(mut s) = status.lock() {
|
||||
f(&mut s);
|
||||
}
|
||||
}
|
||||
|
||||
/// How a task shows in Settings, toasts and the command list.
|
||||
struct Kind {
|
||||
/// Command prefix and settings key section.
|
||||
prefix: &'static str,
|
||||
label: &'static str,
|
||||
classic_name: &'static str,
|
||||
classic_description: &'static str,
|
||||
}
|
||||
|
||||
const ROTO: Kind = Kind {
|
||||
prefix: "roto",
|
||||
label: "Roto Brush",
|
||||
classic_name: "Classic (graph cut)",
|
||||
classic_description: "Built in: colour models and graph cuts, propagated by optical flow. Always available.",
|
||||
};
|
||||
|
||||
const FACE: Kind = Kind {
|
||||
prefix: "face",
|
||||
label: "Face Tracking",
|
||||
classic_name: "Classic (shape model)",
|
||||
classic_description: "Built in: a skin colour model, facial feature shapes and a fitted face shape model. Best on front-facing, evenly lit faces. Always available.",
|
||||
};
|
||||
|
||||
fn kind(task: Task) -> &'static Kind {
|
||||
match task {
|
||||
Task::Mask => &ROTO,
|
||||
Task::Face => &FACE,
|
||||
}
|
||||
}
|
||||
|
||||
impl Session {
|
||||
/// The folder installed models live in: the host's, else `models` next to the settings.
|
||||
pub fn models_dir(&self) -> Option<PathBuf> {
|
||||
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<PathBuf> {
|
||||
let p = self.models_dir()?.join(m.file_name);
|
||||
std::fs::metadata(&p).ok().filter(|md| md.len() == m.size).map(|_| p)
|
||||
}
|
||||
|
||||
/// The model chosen for `task` (Settings).
|
||||
pub fn chosen_model(&self, task: Task) -> &str {
|
||||
match task {
|
||||
Task::Mask => &self.prefs.roto.model,
|
||||
Task::Face => &self.prefs.face.model,
|
||||
}
|
||||
}
|
||||
|
||||
fn choose_model(&mut self, task: Task, id: &str) {
|
||||
let slot = match task {
|
||||
Task::Mask => &mut self.prefs.roto.model,
|
||||
Task::Face => &mut self.prefs.face.model,
|
||||
};
|
||||
*slot = id.to_string();
|
||||
}
|
||||
|
||||
/// Load (or drop) the models Settings chooses; called whenever settings change.
|
||||
pub fn apply_models(&mut self) {
|
||||
for task in [Task::Mask, Task::Face] {
|
||||
self.apply_model(task, false);
|
||||
}
|
||||
}
|
||||
|
||||
/// Load (or drop) `task`'s chosen model; `wait`: before returning (scripts and agents).
|
||||
fn apply_model(&mut self, task: Task, wait: bool) {
|
||||
let want = self.chosen_model(task).to_string();
|
||||
let lane = self.models.lane(task).clone();
|
||||
let st = self.models.status(task);
|
||||
let info = seg::info(&want).filter(|m| m.task == task);
|
||||
let Some(info) = info else {
|
||||
if st.active.is_some() {
|
||||
self.models.set_active(task, None);
|
||||
}
|
||||
return;
|
||||
};
|
||||
if st.active.as_deref() == Some(info.id) || st.busy.is_some() {
|
||||
return;
|
||||
}
|
||||
let Some(path) = self.model_path(info) else {
|
||||
self.models.set_active(task, None);
|
||||
return;
|
||||
};
|
||||
let (face_slot, label) = (self.models.face_model.clone(), kind(task).label);
|
||||
spawn(&lane, label, format!("Loading {}…", info.name), wait, move |status| {
|
||||
let m = load_file(&path, info.id)?;
|
||||
activate(status, &face_slot, task, Some(m));
|
||||
Ok(format!("{label}: {} loaded", info.name))
|
||||
});
|
||||
}
|
||||
|
||||
/// Finish what the model workers did: a toast, and loading a model that just arrived.
|
||||
pub fn poll_models(&mut self) {
|
||||
for task in [Task::Mask, Task::Face] {
|
||||
let done = self.models.lane(task).lock().ok().and_then(|mut s| s.done.take());
|
||||
if let Some((message, error)) = done {
|
||||
self.events.push(crate::Event::Toast { message, error });
|
||||
self.apply_model(task, false);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn load_file(path: &Path, id: &str) -> std::result::Result<Loaded, 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 (`wait`: right here), reporting through `status` (desktop; the
|
||||
/// browser build has no trained models yet).
|
||||
fn spawn(
|
||||
status: &Arc<Mutex<Status>>,
|
||||
label: &'static str,
|
||||
what: String,
|
||||
wait: bool,
|
||||
job: impl FnOnce(&Arc<Mutex<Status>>) -> std::result::Result<String, String> + Send + 'static,
|
||||
) {
|
||||
if cfg!(target_arch = "wasm32") {
|
||||
set(status, |s| s.error = Some("trained models are available in the desktop app".into()));
|
||||
return;
|
||||
}
|
||||
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, false)),
|
||||
Err(e) => {
|
||||
s.done = Some((format!("{label} model: {e}"), true));
|
||||
s.error = Some(e);
|
||||
}
|
||||
}
|
||||
});
|
||||
};
|
||||
if wait {
|
||||
run();
|
||||
} else if let Err(e) = std::thread::Builder::new().name("ec-model".into()).spawn(run) {
|
||||
set(status, |s| {
|
||||
s.busy = None;
|
||||
s.error = Some(format!("cannot start: {e}"));
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// The notice kept next to `m`'s weights.
|
||||
fn notice_path(dir: &Path, m: &ModelInfo) -> PathBuf {
|
||||
dir.join(format!("{}.NOTICE.txt", m.file_name))
|
||||
}
|
||||
|
||||
/// Copy verified weights into the models folder (atomically), with their notice.
|
||||
fn install_bytes(dir: &Path, m: &ModelInfo, bytes: &[u8]) -> std::result::Result<PathBuf, String> {
|
||||
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));
|
||||
}
|
||||
use crate::config::FileOps;
|
||||
let path = dir.join(m.file_name);
|
||||
let files = crate::config::StdFiles;
|
||||
files.write(&path, bytes).map_err(|e| format!("cannot write {}: {e}", path.display()))?;
|
||||
files.write(¬ice_path(dir, m), seg::notice(m).as_bytes()).map_err(|e| format!("cannot write the notice: {e}"))?;
|
||||
Ok(path)
|
||||
}
|
||||
|
||||
fn model_p(p: &Value, cmd: &str, task: Task) -> Result<&'static ModelInfo> {
|
||||
let id = str_p(p, "id").ok_or_else(|| bad(cmd, "missing `id`"))?;
|
||||
seg::info(id).filter(|m| m.task == task).ok_or_else(|| bad(cmd, format!("unknown model `{id}` (see {}.models)", kind(task).prefix)))
|
||||
}
|
||||
|
||||
fn list(s: &mut Session, task: Task) -> Result<Value> {
|
||||
let k = kind(task);
|
||||
let st = s.models.status(task);
|
||||
let chosen = s.chosen_model(task).to_string();
|
||||
let mut v = vec![json!({
|
||||
"id": seg::CLASSICAL,
|
||||
"name": k.classic_name,
|
||||
"description": k.classic_description,
|
||||
"licence": "MIT OR Apache-2.0",
|
||||
"installed": true,
|
||||
"selected": chosen == seg::CLASSICAL,
|
||||
"active": st.active.is_none(),
|
||||
})];
|
||||
for m in seg::models(task) {
|
||||
v.push(json!({
|
||||
"id": m.id,
|
||||
"name": m.name,
|
||||
"description": m.description,
|
||||
"authors": m.authors,
|
||||
"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())}))
|
||||
}
|
||||
|
||||
/// After a `wait`: the job's failure, as the command's.
|
||||
fn waited(s: &Session, task: Task, cmd: &str) -> Result<()> {
|
||||
match s.models.status(task).error {
|
||||
Some(e) => Err(bad(cmd, e)),
|
||||
None => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
fn select(s: &mut Session, p: &Value, task: Task) -> Result<Value> {
|
||||
let cmd = format!("{}.model.select", kind(task).prefix);
|
||||
let id = str_p(p, "id").ok_or_else(|| bad(&cmd, "missing `id`"))?;
|
||||
if id != seg::CLASSICAL {
|
||||
model_p(p, &cmd, task)?;
|
||||
}
|
||||
let wait = b_p(p, "wait").unwrap_or(false);
|
||||
s.choose_model(task, id);
|
||||
if wait {
|
||||
s.apply_model(task, true);
|
||||
waited(s, task, &cmd)?;
|
||||
}
|
||||
s.save_prefs();
|
||||
s.prefs_changed();
|
||||
Ok(json!({"model": id, "active": s.models.status(task).active}))
|
||||
}
|
||||
|
||||
fn install(s: &mut Session, p: &Value, task: Task) -> Result<Value> {
|
||||
let cmd = format!("{}.model.install", kind(task).prefix);
|
||||
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, task)?,
|
||||
None => {
|
||||
let h = seg::sha256::hex(&bytes);
|
||||
seg::models(task).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))?;
|
||||
let wait = b_p(p, "wait").unwrap_or(false);
|
||||
s.apply_model(task, wait);
|
||||
if wait {
|
||||
waited(s, task, &cmd)?;
|
||||
}
|
||||
Ok(json!({"id": m.id, "path": dest.display().to_string()}))
|
||||
}
|
||||
|
||||
fn download(s: &mut Session, p: &Value, task: Task) -> Result<Value> {
|
||||
let k = kind(task);
|
||||
let cmd = format!("{}.model.download", k.prefix);
|
||||
let m = model_p(p, &cmd, task)?;
|
||||
let dir = s.models_dir().ok_or_else(|| bad(&cmd, "no settings folder to install into"))?;
|
||||
if s.models.status(task).busy.is_some() {
|
||||
return Err(bad(&cmd, "a model is already downloading or loading"));
|
||||
}
|
||||
let size = if m.size >= 10_000_000 { format!("{} MB", m.size / 1_000_000) } else { format!("{:.1} MB", m.size as f64 / 1e6) };
|
||||
let wait = b_p(p, "wait").unwrap_or(false);
|
||||
spawn(s.models.lane(task), k.label, format!("Downloading {} ({size})…", m.name), wait, 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!("{}: {} downloaded and verified", k.label, m.name))
|
||||
});
|
||||
if wait {
|
||||
waited(s, task, &cmd)?;
|
||||
s.apply_model(task, true);
|
||||
waited(s, task, &cmd)?;
|
||||
return Ok(json!({"id": m.id, "installed": true}));
|
||||
}
|
||||
Ok(json!({"id": m.id, "downloading": true}))
|
||||
}
|
||||
|
||||
fn remove(s: &mut Session, p: &Value, task: Task) -> Result<Value> {
|
||||
let m = model_p(p, &format!("{}.model.remove", kind(task).prefix), task)?;
|
||||
let Some(path) = s.model_path(m) else { return Ok(json!({"removed": false})) };
|
||||
if s.models.status(task).active.as_deref() == Some(m.id) {
|
||||
s.models.set_active(task, None);
|
||||
}
|
||||
std::fs::remove_file(&path).map_err(|e| EngineError::Other(format!("cannot remove {}: {e}", path.display())))?;
|
||||
if let Some(dir) = path.parent() {
|
||||
let _ = std::fs::remove_file(notice_path(dir, m));
|
||||
}
|
||||
Ok(json!({"removed": true}))
|
||||
}
|
||||
|
||||
pub fn specs() -> Vec<CommandSpec> {
|
||||
vec![
|
||||
query!("roto.models", "Roto Brush Models", "{}", |s, _| list(s, Task::Mask)),
|
||||
cmd!(
|
||||
"roto.model.select",
|
||||
"Use Roto Brush Model",
|
||||
[],
|
||||
None,
|
||||
"{id: classical|mobilesam, wait?} (Roto Brush 2.0 / 3.0 use it; wait: loaded before returning)",
|
||||
always,
|
||||
|s, p| select(s, p, Task::Mask)
|
||||
),
|
||||
cmd!(
|
||||
"roto.model.install",
|
||||
"Install Roto Brush Model",
|
||||
[],
|
||||
None,
|
||||
"{path: weights file, id?, wait?} — verified against the registry's SHA-256",
|
||||
always,
|
||||
|s, p| install(s, p, Task::Mask)
|
||||
),
|
||||
cmd!(
|
||||
"roto.model.download",
|
||||
"Download Roto Brush Model",
|
||||
[],
|
||||
None,
|
||||
"{id, wait?} — the official weights, verified (desktop; uses the system curl)",
|
||||
always,
|
||||
|s, p| download(s, p, Task::Mask)
|
||||
),
|
||||
cmd!("roto.model.remove", "Remove Roto Brush Model", [], None, "{id}", always, |s, p| remove(s, p, Task::Mask)),
|
||||
query!("face.models", "Face Tracking Models", "{}", |s, _| list(s, Task::Face)),
|
||||
cmd!(
|
||||
"face.model.select",
|
||||
"Use Face Tracking Model",
|
||||
[],
|
||||
None,
|
||||
"{id: classical|mediapipe-face, wait?} (face tracking uses it; wait: loaded before returning)",
|
||||
always,
|
||||
|s, p| select(s, p, Task::Face)
|
||||
),
|
||||
cmd!(
|
||||
"face.model.install",
|
||||
"Install Face Tracking Model",
|
||||
[],
|
||||
None,
|
||||
"{path: model file, id?, wait?} — verified against the registry's SHA-256",
|
||||
always,
|
||||
|s, p| install(s, p, Task::Face)
|
||||
),
|
||||
cmd!(
|
||||
"face.model.download",
|
||||
"Download Face Tracking Model",
|
||||
[],
|
||||
None,
|
||||
"{id, wait?} — the official model, verified (desktop; uses the system curl)",
|
||||
always,
|
||||
|s, p| download(s, p, Task::Face)
|
||||
),
|
||||
cmd!("face.model.remove", "Remove Face Tracking Model", [], None, "{id}", always, |s, p| remove(s, p, Task::Face)),
|
||||
]
|
||||
}
|
||||
@@ -519,6 +519,7 @@ async fn run_job(s: &mut Session, job: WorkerJob, post: Post) -> Result<(), Stri
|
||||
path,
|
||||
method,
|
||||
times,
|
||||
face_model: if method.is_face() { s.models.face() } else { None },
|
||||
};
|
||||
s.start_mask_track(work, mask, direction, true)
|
||||
})
|
||||
|
||||
@@ -218,6 +218,15 @@ page!(
|
||||
}
|
||||
);
|
||||
|
||||
page!(
|
||||
/// Settings ▸ Face Tracking.
|
||||
FacePrefs {
|
||||
/// The model Track Mask ▸ Face Tracking uses: `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,
|
||||
@@ -282,6 +291,7 @@ pub struct Prefs {
|
||||
pub three_d: ThreeD,
|
||||
pub scripting: Scripting,
|
||||
pub roto: RotoPrefs,
|
||||
pub face: FacePrefs,
|
||||
/// File ▸ Open Recent, newest first.
|
||||
pub recent_projects: Vec<String>,
|
||||
/// File ▸ Import Recent Footage, newest first.
|
||||
@@ -323,6 +333,7 @@ impl Default for Prefs {
|
||||
three_d: ThreeD::default(),
|
||||
scripting: Scripting::default(),
|
||||
roto: RotoPrefs::default(),
|
||||
face: FacePrefs::default(),
|
||||
recent_projects: vec![],
|
||||
recent_footage: vec![],
|
||||
recent_presets: vec![],
|
||||
@@ -721,6 +732,7 @@ pub fn section_key(page: &str) -> Option<&'static str> {
|
||||
"3d" | "threeD" => "threeD",
|
||||
"scripting" => "scripting",
|
||||
"roto" => "roto",
|
||||
"face" => "face",
|
||||
_ => return None,
|
||||
})
|
||||
}
|
||||
@@ -748,6 +760,7 @@ pub fn page_id(name: &str) -> Option<&'static str> {
|
||||
"3d" | "threed" => "3d",
|
||||
"scripting" | "scriptingexpressions" => "scripting",
|
||||
"roto" | "rotobrush" | "models" => "roto",
|
||||
"face" | "facetracking" => "face",
|
||||
_ => return None,
|
||||
};
|
||||
Some(id)
|
||||
@@ -835,8 +848,9 @@ 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,
|
||||
/// A task's trained models (`roto.models`, `face.models`): authors, licence, size, install,
|
||||
/// choose.
|
||||
Models(effectcraft_segment::Task),
|
||||
}
|
||||
|
||||
/// One page of the Settings dialog.
|
||||
@@ -1119,7 +1133,18 @@ pub fn pages() -> Vec<Page> {
|
||||
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,
|
||||
Models(effectcraft_segment::Task::Mask),
|
||||
],
|
||||
},
|
||||
Page {
|
||||
id: "face",
|
||||
title: "Face Tracking",
|
||||
items: vec![
|
||||
Section("Face Model"),
|
||||
Note(
|
||||
"Face tracking in the Tracker panel (Outline Only and Detailed Features) uses the model chosen here. Models are open source, downloaded only when you ask, and checked against their published SHA-256.",
|
||||
),
|
||||
Models(effectcraft_segment::Task::Face),
|
||||
],
|
||||
},
|
||||
]
|
||||
@@ -1174,7 +1199,10 @@ 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"}),
|
||||
Item::Models(task) => {
|
||||
let p = if *task == effectcraft_segment::Task::Face { "face" } else { "roto" };
|
||||
json!({"models": format!("{p}.models / {p}.model.select / {p}.model.download / {p}.model.install / {p}.model.remove"), "key": format!("{p}.model")})
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
json!({"id": p.id, "title": p.title, "items": items})
|
||||
|
||||
@@ -636,7 +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();
|
||||
self.poll_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() {
|
||||
|
||||
@@ -1,283 +0,0 @@
|
||||
//! 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<String>,
|
||||
/// The model loaded into Roto Brush.
|
||||
pub active: Option<String>,
|
||||
/// The last failure, for the Settings page.
|
||||
pub error: Option<String>,
|
||||
/// A finished job's message, for a toast (taken by [`Session::poll_roto_models`]).
|
||||
pub done: Option<String>,
|
||||
}
|
||||
|
||||
/// The session's model state.
|
||||
#[derive(Default)]
|
||||
pub struct Models {
|
||||
pub status: Arc<Mutex<Status>>,
|
||||
}
|
||||
|
||||
impl Models {
|
||||
pub fn status(&self) -> Status {
|
||||
self.status.lock().map(|s| s.clone()).unwrap_or_default()
|
||||
}
|
||||
}
|
||||
|
||||
fn set(status: &Arc<Mutex<Status>>, 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<PathBuf> {
|
||||
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<PathBuf> {
|
||||
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<Arc<dyn seg::MaskModel>, 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<Mutex<Status>>, what: String, job: impl FnOnce(&Arc<Mutex<Status>>) -> std::result::Result<String, String> + 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<PathBuf, String> {
|
||||
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<Value> {
|
||||
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<Value> {
|
||||
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<Value> {
|
||||
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<Value> {
|
||||
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<Value> {
|
||||
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<CommandSpec> {
|
||||
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),
|
||||
]
|
||||
}
|
||||
@@ -38,8 +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();
|
||||
// Settings ▸ Roto Brush / Face Tracking: load (or drop) the chosen trained models.
|
||||
self.apply_models();
|
||||
self.prefs_revision += 1;
|
||||
}
|
||||
|
||||
|
||||
@@ -104,3 +104,37 @@ fn face_tracking_keys_outline_points_and_measurements() {
|
||||
let layer = s.active_comp().unwrap().layer(clip).unwrap();
|
||||
assert!(layer.props.find_group(m3).unwrap().get("path").unwrap().keys.is_empty());
|
||||
}
|
||||
|
||||
/// With MediaPipe Face Landmarker installed (`EFFECTCRAFT_FACE_LANDMARKER` = path to
|
||||
/// `face_landmarker.task`, else skipped): Track Mask ▸ Face Tracking uses it, and the pupils and
|
||||
/// eye corners it keys sit on the synthetic face's.
|
||||
#[test]
|
||||
fn face_tracking_uses_the_chosen_model() {
|
||||
let Ok(model) = std::env::var("EFFECTCRAFT_FACE_LANDMARKER") else { return };
|
||||
let (mut s, clip, _) = setup(|f| render_synthetic(W, H, &face(f), 5));
|
||||
let dir = std::env::temp_dir().join(format!("ec-face-model-{}", std::process::id()));
|
||||
s.models_dir = Some(dir.clone());
|
||||
s.execute_checked("face.model.install", json!({"path": model})).unwrap();
|
||||
// `wait`: loaded before the track starts (as scripts need).
|
||||
let r = s.execute_checked("face.model.select", json!({"id": "mediapipe-face", "wait": true})).unwrap();
|
||||
assert_eq!(r["active"], "mediapipe-face");
|
||||
let c = face(0).center;
|
||||
let r = s.execute("mask.new", json!({"layer": clip.0, "vertices": ShapePath::rect(c, 120.0, 160.0).vertices, "closed": true})).unwrap();
|
||||
let mask = r["mask"].as_u64().unwrap();
|
||||
s.execute("time.set", json!({"time": 0})).unwrap();
|
||||
let r = s.execute_checked("track.mask", json!({"layer": clip.0, "mask": mask, "method": "faceDetailed", "direction": "forward", "wait": true})).unwrap();
|
||||
assert_eq!((r["frames"].as_u64(), r["faceModel"].as_str()), (Some(24), Some("mediapipe-face")), "{r}");
|
||||
let layer = s.active_comp().unwrap().layer(clip).unwrap().clone();
|
||||
let pts = layer.effects().unwrap().groups().find(|g| g.match_id == effectcraft_effects::face_track::POINTS_ID).expect("Face Track Points");
|
||||
let mut worst: f64 = 0.0;
|
||||
for f in [0, 8, 16, 24] {
|
||||
let truth = synth_landmarks(&face(f));
|
||||
for id in ["leftPupil", "rightPupil", "leftEyeOuter", "rightEyeOuter"] {
|
||||
let i = effectcraft_track::face::landmark(id).unwrap();
|
||||
let KV::Vec2(v) = pts.get(id).unwrap().value_at(frame_time(f)) else { panic!() };
|
||||
worst = worst.max((v[0] - truth[i][0]).hypot(v[1] - truth[i][1]) / face(f).height);
|
||||
}
|
||||
}
|
||||
assert!(worst < 0.05, "worst eye error {worst} face heights");
|
||||
let _ = std::fs::remove_dir_all(dir);
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
//! Settings ▸ Roto Brush and Settings ▸ Face Tracking: the model registry, choosing, installing
|
||||
//! (verified, with its notice) and loading trained models.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use effectcraft_segment::Task;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::Session;
|
||||
use crate::config::{ConfigStore, DirConfig};
|
||||
|
||||
fn session(tag: &str) -> (Session, std::path::PathBuf) {
|
||||
let d = std::env::temp_dir().join(format!("ec-models-{tag}-{}", std::process::id()));
|
||||
let store: Arc<dyn ConfigStore> = Arc::new(DirConfig::new(d.join("config")));
|
||||
(Session { config: Some(store), ..Default::default() }, d)
|
||||
}
|
||||
|
||||
/// Poll until `task`'s model is in use (or a few seconds pass).
|
||||
fn wait_loaded(s: &mut Session, task: Task) {
|
||||
for _ in 0..400 {
|
||||
s.poll_models();
|
||||
if s.models.status(task).active.is_some() {
|
||||
return;
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_millis(20));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn models_are_listed_chosen_and_verified() {
|
||||
let (mut s, d) = session("list");
|
||||
std::fs::create_dir_all(&d).unwrap();
|
||||
let bogus = d.join("bogus.bin");
|
||||
std::fs::write(&bogus, b"not the weights").unwrap();
|
||||
for (prefix, model, task) in [("roto", "mobilesam", Task::Mask), ("face", "mediapipe-face", Task::Face)] {
|
||||
let list = s.execute_checked(&format!("{prefix}.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", model]);
|
||||
let m = &list["models"][1];
|
||||
assert_eq!((m["licence"].as_str(), m["installed"].as_bool()), (Some("Apache-2.0"), Some(false)));
|
||||
assert!(m["url"].as_str().unwrap().starts_with("https://") && !m["authors"].as_str().unwrap().is_empty());
|
||||
// Unknown ids, the other task's models and files that aren't the published weights are
|
||||
// refused.
|
||||
let other = if task == Task::Mask { "mediapipe-face" } else { "mobilesam" };
|
||||
for id in ["nope", other] {
|
||||
assert!(s.execute_checked(&format!("{prefix}.model.select"), json!({"id": id})).is_err(), "{prefix} {id}");
|
||||
}
|
||||
let path = bogus.to_string_lossy();
|
||||
let e = s.execute_checked(&format!("{prefix}.model.install"), json!({"path": path})).unwrap_err().to_string();
|
||||
assert!(e.contains("SHA-256"), "{e}");
|
||||
assert!(s.execute_checked(&format!("{prefix}.model.install"), json!({"path": path, "id": model})).is_err());
|
||||
// Choosing a model that isn't installed keeps the classic engine.
|
||||
s.execute_checked(&format!("{prefix}.model.select"), json!({"id": model})).unwrap();
|
||||
assert_eq!(s.chosen_model(task), model);
|
||||
assert!(s.models.status(task).active.is_none());
|
||||
assert!(crate::prefs::pages().iter().any(|p| p.id == prefix));
|
||||
}
|
||||
assert!(s.models.face().is_none());
|
||||
assert_eq!(crate::prefs::page_id("Face Tracking"), Some("face"));
|
||||
// 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();
|
||||
wait_loaded(&mut s, Task::Mask);
|
||||
assert_eq!(s.models.status(Task::Mask).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());
|
||||
}
|
||||
// EFFECTCRAFT_FACE_LANDMARKER=path/to/face_landmarker.task: installed with its notice,
|
||||
// loaded for face tracking, dropped when the classic tracker is chosen, removed.
|
||||
if let Ok(path) = std::env::var("EFFECTCRAFT_FACE_LANDMARKER") {
|
||||
let r = s.execute_checked("face.model.install", json!({"path": path})).unwrap();
|
||||
assert_eq!(r["id"], "mediapipe-face");
|
||||
let dir = s.models_dir().unwrap();
|
||||
let notice = std::fs::read_to_string(dir.join("face_landmarker.task.NOTICE.txt")).unwrap();
|
||||
assert!(notice.contains("Google LLC") && notice.contains("Apache-2.0"), "{notice}");
|
||||
wait_loaded(&mut s, Task::Face);
|
||||
assert_eq!(s.models.face().map(|m| m.info().id), Some("mediapipe-face"));
|
||||
s.execute_checked("face.model.select", json!({"id": "classical"})).unwrap();
|
||||
assert!(s.models.face().is_none());
|
||||
assert_eq!(s.execute_checked("face.model.remove", json!({"id": "mediapipe-face"})).unwrap()["removed"], true);
|
||||
assert!(!dir.join("face_landmarker.task").exists() && !dir.join("face_landmarker.task.NOTICE.txt").exists());
|
||||
}
|
||||
let _ = std::fs::remove_dir_all(d);
|
||||
}
|
||||
@@ -1,57 +0,0 @@
|
||||
//! 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<dyn ConfigStore> = 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);
|
||||
}
|
||||
@@ -1,6 +1,6 @@
|
||||
[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)"
|
||||
description = "Swappable trained models for Roto Brush and face tracking: the MaskModel and FaceModel interfaces, a registry of open-licensed models, PyTorch checkpoint and TensorFlow Lite readers, a pure-Rust MobileSAM and MediaPipe Face Landmarker"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
|
||||
+45
-19
@@ -1,38 +1,64 @@
|
||||
# effectcraft-segment
|
||||
|
||||
Trained segmentation models for Roto Brush, behind one swappable interface.
|
||||
Trained models for Roto Brush and face tracking, each kind 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.
|
||||
- **`face::FaceModel`**: what face tracking asks of a model: find a face in a region of a frame,
|
||||
then follow it from frame to frame. A model reports its own points plus a `Topology` naming
|
||||
which are the tracker's landmarks and which trace the face outline, so the tracker never
|
||||
depends on one model's layout. The classical face tracker 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.
|
||||
EffectCraft's MIT OR Apache-2.0 (`ALLOWED_LICENCES`), with its task, authors, 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).
|
||||
engine (`roto.model.*` / `face.model.*`, Settings ▸ Roto Brush / Face Tracking), verified
|
||||
against the registry's SHA-256 (`sha256`, dependency-free) and kept with a notice naming their
|
||||
authors, licence and source (`notice`).
|
||||
- **`pt`**: reads PyTorch checkpoints (`torch.save` zip archives and the pickle subset state dicts
|
||||
use), so official weights load exactly as published.
|
||||
- **`tflite`**: reads TensorFlow Lite models (the FlatBuffer laid out by TensorFlow Lite's
|
||||
`schema.fbs`) and runs the float operators small vision models use (convolutions, depthwise
|
||||
convolutions, pooling, PReLU, padding, concatenation; float16 weights dequantised at load).
|
||||
- **`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.
|
||||
prompt encoder and two-way-transformer mask decoder.
|
||||
- **`mediapipe`**: MediaPipe Face Landmarker: the BlazeFace short-range detector (SSD anchors,
|
||||
weighted non-maximum suppression) and Face Mesh V2 (478 points), with the rotated crops
|
||||
MediaPipe's face landmarker graph uses.
|
||||
|
||||
All of it runs 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 |
|
||||
| id | Task | Model | Authors | Licence | File | Size |
|
||||
|---|---|---|---|---|---|---|
|
||||
| `mobilesam` | Roto Brush | [MobileSAM](https://github.com/ChaoningZhang/MobileSAM): Segment Anything with a 5M-parameter TinyViT encoder | Chaoning Zhang et al., building on Meta AI's Segment Anything and Microsoft's TinyViT | Apache-2.0 | `mobile_sam.pt` from the official repository | 40.7 MB |
|
||||
| `mediapipe-face` | Face tracking | [MediaPipe Face Landmarker](https://developers.google.com/edge/mediapipe/solutions/vision/face_landmarker): BlazeFace detector and Face Mesh V2 | Google LLC | Apache-2.0 ([model card](https://storage.googleapis.com/mediapipe-assets/Model%20Card%20MediaPipe%20Face%20Mesh%20V2.pdf)) | `face_landmarker.task` (float16, version 1) from Google's model storage | 3.8 MB |
|
||||
|
||||
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.
|
||||
How they were checked:
|
||||
|
||||
- **MobileSAM**: 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.
|
||||
- **MediaPipe Face Landmarker**: the interpreter's outputs match TensorFlow Lite's own
|
||||
(`ai-edge-litert` 2.2.0) on both models to about one part in a million (a test pins them). The
|
||||
whole pipeline's 478 points on a test portrait are 0.85 px from MediaPipe's own Python pipeline
|
||||
on average (2.2 px at worst). Finding a face takes about 45 ms; following it, about 20 ms a
|
||||
frame (32-thread CPU).
|
||||
|
||||
Tests that need the weights run when `EFFECTCRAFT_MOBILESAM` / `EFFECTCRAFT_FACE_LANDMARKER`
|
||||
point at the files, and are skipped otherwise.
|
||||
|
||||
## 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.
|
||||
1. Pick a model whose code **and** weights are open source under an allowed licence. Check the
|
||||
weights' own terms: many models with MIT code ship weights trained on research-only data.
|
||||
2. Add a `ModelInfo` to `MODELS` (its task and authors; pin the official file's SHA-256 and size),
|
||||
and credit it in the root `ATTRIBUTION.md`.
|
||||
3. Implement `MaskModel` or `FaceModel` for it and add its loader to `load`.
|
||||
4. Test it against a reference implementation; Settings, the `roto.model.*` / `face.model.*`
|
||||
commands and the tools pick it up from the registry.
|
||||
@@ -0,0 +1,115 @@
|
||||
//! Trained face models for face tracking, behind one swappable interface.
|
||||
//!
|
||||
//! - [`FaceModel`]: what the face tracker (`effectcraft-track`) asks of a model: find a face in a
|
||||
//! region of a frame, then follow it from frame to frame. A model reports its own points (a
|
||||
//! mesh, for MediaPipe) plus a [`Topology`] saying which of them are the tracker's named
|
||||
//! landmarks and which trace the face outline, so the tracker never depends on one model's
|
||||
//! layout. The classical tracker stays built in as the fallback when no model is chosen.
|
||||
//! - [`Roi`] and [`crop`]: the rotated square crops face models look at.
|
||||
|
||||
use crate::{ModelInfo, Result};
|
||||
|
||||
/// Which of a model's points are the tracker's landmarks and its face outline.
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
pub struct Topology {
|
||||
/// The points tracing the face outline, in order around the face.
|
||||
pub outline: &'static [usize],
|
||||
/// (landmark id, as the tracker names them, e.g. `leftEyeInner`; point index). "Left" is
|
||||
/// image-left. Chin and jaw come from the outline.
|
||||
pub landmarks: &'static [(&'static str, usize)],
|
||||
}
|
||||
|
||||
/// A face in a frame: the model's points (frame pixels) and its confidence (0–1).
|
||||
#[derive(Clone, Debug, Default, PartialEq)]
|
||||
pub struct Face {
|
||||
pub points: Vec<[f32; 2]>,
|
||||
pub score: f32,
|
||||
}
|
||||
|
||||
/// A trained face model (swappable: the face tracker only uses this interface). Frames are
|
||||
/// straight RGB 0–1, row-major `w×h`, coordinates frame pixels.
|
||||
pub trait FaceModel: Send + Sync {
|
||||
fn info(&self) -> &'static ModelInfo;
|
||||
fn topology(&self) -> &'static Topology;
|
||||
/// The most confident face whose centre lies in `region` (`[x0, y0, x1, y1]`).
|
||||
fn find(&self, rgb: &[[f32; 3]], w: usize, h: usize, region: [f32; 4]) -> Result<Option<Face>>;
|
||||
/// The face `prev` was in the previous frame, in this one; `None` when it is lost.
|
||||
fn follow(&self, rgb: &[[f32; 3]], w: usize, h: usize, prev: &Face) -> Result<Option<Face>>;
|
||||
}
|
||||
|
||||
/// A rotated square region of a frame: centre, side and angle (radians, the direction of the
|
||||
/// square's x axis in the frame, y down).
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
pub struct Roi {
|
||||
pub center: [f32; 2],
|
||||
pub size: f32,
|
||||
pub angle: f32,
|
||||
}
|
||||
|
||||
impl Roi {
|
||||
/// Frame position of the normalised square position `(u, v)` (0–1).
|
||||
pub fn to_frame(&self, u: f32, v: f32) -> [f32; 2] {
|
||||
let (s, c) = self.angle.sin_cos();
|
||||
let (x, y) = ((u - 0.5) * self.size, (v - 0.5) * self.size);
|
||||
[self.center[0] + c * x - s * y, self.center[1] + s * x + c * y]
|
||||
}
|
||||
}
|
||||
|
||||
/// Sample `roi` into an `n×n` RGB tensor (bilinear; zero outside the frame), each channel mapped
|
||||
/// from 0–1 to `lo`–`hi`.
|
||||
pub fn crop(rgb: &[[f32; 3]], w: usize, h: usize, roi: &Roi, n: usize, [lo, hi]: [f32; 2]) -> Vec<f32> {
|
||||
let mut out = vec![0.0f32; n * n * 3];
|
||||
if rgb.len() < w * h || n == 0 {
|
||||
return out;
|
||||
}
|
||||
let px = |x: i64, y: i64| -> [f32; 3] {
|
||||
if x < 0 || y < 0 || x >= w as i64 || y >= h as i64 { [0.0; 3] } else { rgb.get(y as usize * w + x as usize).copied().unwrap_or([0.0; 3]) }
|
||||
};
|
||||
for j in 0..n {
|
||||
for i in 0..n {
|
||||
let p = roi.to_frame((i as f32 + 0.5) / n as f32, (j as f32 + 0.5) / n as f32);
|
||||
let (x, y) = (p[0] - 0.5, p[1] - 0.5);
|
||||
let (x0, y0) = (x.floor(), y.floor());
|
||||
let (fx, fy) = (x - x0, y - y0);
|
||||
let (x0, y0) = (x0 as i64, y0 as i64);
|
||||
let (a, b, c, d) = (px(x0, y0), px(x0 + 1, y0), px(x0, y0 + 1), px(x0 + 1, y0 + 1));
|
||||
for k in 0..3 {
|
||||
let v = (a[k] * (1.0 - fx) + b[k] * fx) * (1.0 - fy) + (c[k] * (1.0 - fx) + d[k] * fx) * fy;
|
||||
out[(j * n + i) * 3 + k] = lo + v * (hi - lo);
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
pub(crate) fn sigmoid(x: f32) -> f32 {
|
||||
1.0 / (1.0 + (-x).exp())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn roi_maps_corners_and_rotates() {
|
||||
let r = Roi { center: [10.0, 20.0], size: 4.0, angle: 0.0 };
|
||||
assert_eq!(r.to_frame(0.0, 0.0), [8.0, 18.0]);
|
||||
assert_eq!(r.to_frame(1.0, 1.0), [12.0, 22.0]);
|
||||
// A quarter turn: the square's x axis points down the frame.
|
||||
let q = Roi { angle: std::f32::consts::FRAC_PI_2, ..r };
|
||||
let p = q.to_frame(1.0, 0.5);
|
||||
assert!((p[0] - 10.0).abs() < 1e-5 && (p[1] - 22.0).abs() < 1e-5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn crop_samples_pixel_centres() {
|
||||
// A 4×4 ramp cropped 1:1 reproduces it; outside the frame is zero.
|
||||
let rgb: Vec<[f32; 3]> = (0..16).map(|i| [i as f32 / 16.0, 0.0, 1.0]).collect();
|
||||
let c = crop(&rgb, 4, 4, &Roi { center: [2.0, 2.0], size: 4.0, angle: 0.0 }, 4, [0.0, 1.0]);
|
||||
for i in 0..16 {
|
||||
assert!((c[i * 3] - i as f32 / 16.0).abs() < 1e-6);
|
||||
}
|
||||
let c = crop(&rgb, 4, 4, &Roi { center: [-10.0, -10.0], size: 4.0, angle: 0.0 }, 2, [-1.0, 1.0]);
|
||||
assert!(c.iter().all(|v| *v == -1.0));
|
||||
}
|
||||
}
|
||||
+74
-13
@@ -1,21 +1,28 @@
|
||||
//! Trained segmentation models for Roto Brush, behind one swappable interface.
|
||||
//! Trained models for Roto Brush and face tracking, each kind 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.
|
||||
//! - [`face::FaceModel`]: what face tracking asks of a model: find a face in a region, follow it
|
||||
//! from frame to frame. The classical face tracker stays the built-in fallback.
|
||||
//! - [`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).
|
||||
//! EffectCraft's (MIT OR Apache-2.0), with its authors, 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; [`tflite`]: a TensorFlow Lite reader and interpreter;
|
||||
//! so official weights load as published.
|
||||
//! - [`mobilesam`]: MobileSAM (Apache-2.0) and [`mediapipe`]: MediaPipe Face Landmarker
|
||||
//! (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 face;
|
||||
pub mod mediapipe;
|
||||
pub mod mobilesam;
|
||||
pub mod nn;
|
||||
pub mod pt;
|
||||
pub mod tflite;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
@@ -40,12 +47,24 @@ pub trait MaskModel: Send + Sync {
|
||||
fn segment(&self, rgb: &[[f32; 3]], w: usize, h: usize, prompt: &Prompt) -> Result<Vec<f32>>;
|
||||
}
|
||||
|
||||
/// A model in the registry: what it is, its licence and where its weights come from.
|
||||
/// What a model is for.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum Task {
|
||||
/// Roto Brush ([`MaskModel`]).
|
||||
Mask,
|
||||
/// Face tracking ([`face::FaceModel`]).
|
||||
Face,
|
||||
}
|
||||
|
||||
/// A model in the registry: what it is, who made it, its licence and where its weights come from.
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
pub struct ModelInfo {
|
||||
pub id: &'static str,
|
||||
pub task: Task,
|
||||
pub name: &'static str,
|
||||
pub description: &'static str,
|
||||
/// Who made it (shown with the model, and in ATTRIBUTION.md).
|
||||
pub authors: &'static str,
|
||||
/// SPDX licence of the weights and the architecture.
|
||||
pub licence: &'static str,
|
||||
pub licence_url: &'static str,
|
||||
@@ -64,8 +83,10 @@ pub const ALLOWED_LICENCES: &[&str] = &["MIT", "Apache-2.0", "BSD-2-Clause", "BS
|
||||
|
||||
pub const MOBILE_SAM: ModelInfo = ModelInfo {
|
||||
id: "mobilesam",
|
||||
task: Task::Mask,
|
||||
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.",
|
||||
authors: "Chaoning Zhang et al. (MobileSAM), building on Meta AI's Segment Anything and Microsoft's TinyViT",
|
||||
licence: "Apache-2.0",
|
||||
licence_url: "https://github.com/ChaoningZhang/MobileSAM/blob/master/LICENSE",
|
||||
homepage: "https://github.com/ChaoningZhang/MobileSAM",
|
||||
@@ -75,18 +96,46 @@ pub const MOBILE_SAM: ModelInfo = ModelInfo {
|
||||
sha256: "6dbb90523a35330fedd7f1d3dfc66f995213d81b29a5ca8108dbcdd4e37d6c2f",
|
||||
};
|
||||
|
||||
/// The models Roto Brush can use besides the built-in classical segmenter.
|
||||
pub const MODELS: &[ModelInfo] = &[MOBILE_SAM];
|
||||
pub const FACE_LANDMARKER: ModelInfo = ModelInfo {
|
||||
id: "mediapipe-face",
|
||||
task: Task::Face,
|
||||
name: "MediaPipe Face Landmarker",
|
||||
description: "Google's BlazeFace detector and 478-point Face Mesh V2. Follows faces turned to the side, partly covered or in difficult light, where the classic tracker loses them.",
|
||||
authors: "Google LLC (MediaPipe)",
|
||||
licence: "Apache-2.0",
|
||||
licence_url: "https://storage.googleapis.com/mediapipe-assets/Model%20Card%20MediaPipe%20Face%20Mesh%20V2.pdf",
|
||||
homepage: "https://developers.google.com/edge/mediapipe/solutions/vision/face_landmarker",
|
||||
url: "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/1/face_landmarker.task",
|
||||
file_name: "face_landmarker.task",
|
||||
size: 3_758_596,
|
||||
sha256: "64184e229b263107bc2b804c6625db1341ff2bb731874b0bcc2fe6544e0bc9ff",
|
||||
};
|
||||
|
||||
/// The id of the built-in classical (graph cut) segmenter: no weights, always available.
|
||||
/// The trained models, besides the built-in classical engines.
|
||||
pub const MODELS: &[ModelInfo] = &[MOBILE_SAM, FACE_LANDMARKER];
|
||||
|
||||
/// The id of the built-in classical engines (graph cut segmenter, shape-model face tracker): no
|
||||
/// weights, always available.
|
||||
pub const CLASSICAL: &str = "classical";
|
||||
|
||||
pub fn info(id: &str) -> Option<&'static ModelInfo> {
|
||||
MODELS.iter().find(|m| m.id == id)
|
||||
}
|
||||
|
||||
/// The models for `task`.
|
||||
pub fn models(task: Task) -> impl Iterator<Item = &'static ModelInfo> {
|
||||
MODELS.iter().filter(move |m| m.task == task)
|
||||
}
|
||||
|
||||
/// A loaded model.
|
||||
#[derive(Clone)]
|
||||
pub enum Loaded {
|
||||
Mask(Arc<dyn MaskModel>),
|
||||
Face(Arc<dyn face::FaceModel>),
|
||||
}
|
||||
|
||||
/// Check a weights file against the registry and build the model.
|
||||
pub fn load(id: &str, bytes: &[u8]) -> Result<Arc<dyn MaskModel>> {
|
||||
pub fn load(id: &str, bytes: &[u8]) -> Result<Loaded> {
|
||||
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));
|
||||
@@ -98,12 +147,21 @@ pub fn load(id: &str, bytes: &[u8]) -> Result<Arc<dyn MaskModel>> {
|
||||
match id {
|
||||
"mobilesam" => {
|
||||
let w = pt::read(bytes)?;
|
||||
Ok(Arc::new(Sam { net: mobilesam::MobileSam::from_weights(&w)? }))
|
||||
Ok(Loaded::Mask(Arc::new(Sam { net: mobilesam::MobileSam::from_weights(&w)? })))
|
||||
}
|
||||
"mediapipe-face" => Ok(Loaded::Face(Arc::new(mediapipe::FaceLandmarker::from_task(bytes)?))),
|
||||
_ => Err(format!("no loader for `{id}`")),
|
||||
}
|
||||
}
|
||||
|
||||
/// The notice kept next to installed weights: what they are, who made them, their licence.
|
||||
pub fn notice(m: &ModelInfo) -> String {
|
||||
format!(
|
||||
"{}\n\nAuthors: {}\nLicence: {} (full text: https://spdx.org/licenses/{}.html; as stated by the authors: {})\nSource: {}\nProject: {}\n\nDownloaded unmodified from the source above. EffectCraft runs it with its own implementation\nand is not affiliated with or endorsed by its authors.\n",
|
||||
m.name, m.authors, m.licence, m.licence, m.licence_url, m.url, m.homepage
|
||||
)
|
||||
}
|
||||
|
||||
struct Sam {
|
||||
net: mobilesam::MobileSam,
|
||||
}
|
||||
@@ -202,8 +260,11 @@ mod tests {
|
||||
for m in MODELS {
|
||||
assert!(ALLOWED_LICENCES.contains(&m.licence), "{}", m.id);
|
||||
assert_eq!(m.sha256.len(), 64);
|
||||
assert!(m.url.starts_with("https://"));
|
||||
assert!(m.url.starts_with("https://") && m.licence_url.starts_with("https://") && m.homepage.starts_with("https://"));
|
||||
assert!(!m.authors.is_empty() && notice(m).contains(m.authors));
|
||||
}
|
||||
assert_eq!(models(Task::Face).map(|m| m.id).collect::<Vec<_>>(), ["mediapipe-face"]);
|
||||
assert!(load("mediapipe-face", b"nope").err().is_some_and(|e| e.contains("does not match")));
|
||||
// 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());
|
||||
|
||||
@@ -0,0 +1,312 @@
|
||||
//! MediaPipe Face Landmarker (Google LLC, Apache-2.0) in plain Rust: the official
|
||||
//! `face_landmarker.task` bundle's BlazeFace short-range detector and Face Mesh V2 (478 points)
|
||||
//! run on the [`tflite`](crate::tflite) interpreter.
|
||||
//!
|
||||
//! The pipeline follows MediaPipe's published face landmarker graph (Apache-2.0): the detector
|
||||
//! sees a 128×128 crop scaled to −1…1 and decodes 896 SSD anchors (strides 8, 16, 16, 16; two
|
||||
//! anchors per cell per layer) with weighted non-maximum suppression; the mesh sees a 256×256
|
||||
//! crop (0…1) of a square 1.5× the face, rotated so the eyes are level (detector keypoints 0 → 1,
|
||||
//! then mesh points 33 → 263), and reports a presence score. The landmark indices in
|
||||
//! [`TOPOLOGY`] are points of MediaPipe's canonical face mesh.
|
||||
|
||||
use crate::face::{Face, FaceModel, Roi, Topology, crop, sigmoid};
|
||||
use crate::tflite::Model;
|
||||
use crate::{FACE_LANDMARKER, ModelInfo, Result};
|
||||
|
||||
pub const DETECTOR: &str = "face_detector.tflite";
|
||||
pub const MESH: &str = "face_landmarks_detector.tflite";
|
||||
|
||||
const DET_SIZE: usize = 128;
|
||||
const MESH_SIZE: usize = 256;
|
||||
const MESH_POINTS: usize = 478;
|
||||
/// Detection and presence thresholds (MediaPipe's defaults).
|
||||
const MIN_SCORE: f32 = 0.5;
|
||||
const NMS_IOU: f32 = 0.3;
|
||||
/// Crops are this much larger than the face.
|
||||
const ROI_SCALE: f32 = 1.5;
|
||||
/// Mesh points whose direction levels the crop: the outer eye corners.
|
||||
const EYE_LINE: [usize; 2] = [33, 263];
|
||||
|
||||
/// MediaPipe's face mesh: the face oval (`FACEMESH_FACE_OVAL`, in order) and the points the
|
||||
/// tracker's landmarks are (inner lips for the mouth, iris centres for the pupils).
|
||||
pub const TOPOLOGY: Topology = Topology {
|
||||
outline: &[
|
||||
10, 338, 297, 332, 284, 251, 389, 356, 454, 323, 361, 288, 397, 365, 379, 378, 400, 377, 152, 148, 176, 149, 150, 136, 172, 58, 132, 93, 234, 127, 162,
|
||||
21, 54, 103, 67, 109,
|
||||
],
|
||||
landmarks: &[
|
||||
("leftEyebrowInner", 107),
|
||||
("leftEyebrowMiddle", 105),
|
||||
("leftEyebrowOuter", 70),
|
||||
("rightEyebrowInner", 336),
|
||||
("rightEyebrowMiddle", 334),
|
||||
("rightEyebrowOuter", 300),
|
||||
("leftEyeInner", 133),
|
||||
("leftEyeOuter", 33),
|
||||
("leftEyeTop", 159),
|
||||
("leftEyeBottom", 145),
|
||||
("leftPupil", 468),
|
||||
("rightEyeInner", 362),
|
||||
("rightEyeOuter", 263),
|
||||
("rightEyeTop", 386),
|
||||
("rightEyeBottom", 374),
|
||||
("rightPupil", 473),
|
||||
("noseTip", 1),
|
||||
("leftNostril", 79),
|
||||
("rightNostril", 309),
|
||||
("mouthLeft", 61),
|
||||
("mouthRight", 291),
|
||||
("mouthTop", 13),
|
||||
("mouthBottom", 14),
|
||||
],
|
||||
};
|
||||
|
||||
/// A detected face: box `[x0, y0, x1, y1]`, six keypoints (eyes, nose tip, mouth, ears; image
|
||||
/// left first) and score, in frame pixels.
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
pub struct Detection {
|
||||
pub bbox: [f32; 4],
|
||||
pub keypoints: [[f32; 2]; 6],
|
||||
pub score: f32,
|
||||
}
|
||||
|
||||
pub struct FaceLandmarker {
|
||||
detector: Model,
|
||||
mesh: Model,
|
||||
/// Anchor centres (normalised).
|
||||
anchors: Vec<[f32; 2]>,
|
||||
}
|
||||
|
||||
/// SSD anchor centres: per stride group, `fm×fm` cells with two anchors per layer each.
|
||||
fn anchors() -> Vec<[f32; 2]> {
|
||||
let mut out = vec![];
|
||||
for (stride, layers) in [(8usize, 1usize), (16, 3)] {
|
||||
let fm = DET_SIZE.div_ceil(stride);
|
||||
for y in 0..fm {
|
||||
for x in 0..fm {
|
||||
for _ in 0..2 * layers {
|
||||
out.push([(x as f32 + 0.5) / fm as f32, (y as f32 + 0.5) / fm as f32]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn iou(a: &[f32; 4], b: &[f32; 4]) -> f32 {
|
||||
let iw = (a[2].min(b[2]) - a[0].max(b[0])).max(0.0);
|
||||
let ih = (a[3].min(b[3]) - a[1].max(b[1])).max(0.0);
|
||||
let inter = iw * ih;
|
||||
let union = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter;
|
||||
if union > 0.0 { inter / union } else { 0.0 }
|
||||
}
|
||||
|
||||
/// Weighted non-maximum suppression: each cluster of overlapping detections becomes their
|
||||
/// score-weighted average, with the best score.
|
||||
fn weighted_nms(mut dets: Vec<Detection>) -> Vec<Detection> {
|
||||
dets.sort_by(|a, b| b.score.total_cmp(&a.score));
|
||||
let mut out = vec![];
|
||||
while let Some(top) = dets.first().copied() {
|
||||
let (cluster, rest): (Vec<Detection>, Vec<Detection>) = dets.into_iter().partition(|d| iou(&d.bbox, &top.bbox) > NMS_IOU);
|
||||
dets = rest;
|
||||
let total: f32 = cluster.iter().map(|d| d.score).sum();
|
||||
if total <= 0.0 {
|
||||
out.push(top);
|
||||
continue;
|
||||
}
|
||||
let mut avg = Detection { bbox: [0.0; 4], keypoints: [[0.0; 2]; 6], score: top.score };
|
||||
for d in &cluster {
|
||||
let k = d.score / total;
|
||||
for i in 0..4 {
|
||||
avg.bbox[i] += d.bbox[i] * k;
|
||||
}
|
||||
for (a, p) in avg.keypoints.iter_mut().zip(&d.keypoints) {
|
||||
a[0] += p[0] * k;
|
||||
a[1] += p[1] * k;
|
||||
}
|
||||
}
|
||||
out.push(avg);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// The crop a face's mesh is computed in: `bbox` grown to a square 1.5× its long side, turned so
|
||||
/// `from → to` is level.
|
||||
fn roi_of(bbox: [f32; 4], from: [f32; 2], to: [f32; 2]) -> Roi {
|
||||
Roi {
|
||||
center: [(bbox[0] + bbox[2]) / 2.0, (bbox[1] + bbox[3]) / 2.0],
|
||||
size: (bbox[2] - bbox[0]).max(bbox[3] - bbox[1]) * ROI_SCALE,
|
||||
angle: (to[1] - from[1]).atan2(to[0] - from[0]),
|
||||
}
|
||||
}
|
||||
|
||||
fn bounds(points: &[[f32; 2]]) -> [f32; 4] {
|
||||
points
|
||||
.iter()
|
||||
.fold([f32::INFINITY, f32::INFINITY, f32::NEG_INFINITY, f32::NEG_INFINITY], |b, p| [b[0].min(p[0]), b[1].min(p[1]), b[2].max(p[0]), b[3].max(p[1])])
|
||||
}
|
||||
|
||||
impl FaceLandmarker {
|
||||
/// From the official `face_landmarker.task` bundle (a zip archive of `.tflite` models).
|
||||
pub fn from_task(bytes: &[u8]) -> Result<FaceLandmarker> {
|
||||
let detector = Model::read(&crate::pt::zip_file(bytes, DETECTOR)?)?;
|
||||
let mesh = Model::read(&crate::pt::zip_file(bytes, MESH)?)?;
|
||||
let anchors = anchors();
|
||||
if detector.input_shape(0) != Some(&[1, DET_SIZE, DET_SIZE, 3]) || detector.output(0).map(|o| o.0) != Some(&[1, anchors.len(), 16]) {
|
||||
return Err("unexpected face detector".into());
|
||||
}
|
||||
if mesh.input_shape(0) != Some(&[1, MESH_SIZE, MESH_SIZE, 3]) || mesh.output(0).map(|o| o.0.iter().product::<usize>()) != Some(MESH_POINTS * 3) {
|
||||
return Err("unexpected face mesh".into());
|
||||
}
|
||||
Ok(FaceLandmarker { detector, mesh, anchors })
|
||||
}
|
||||
|
||||
/// Faces in the square `roi` of a frame (frame pixels), best first.
|
||||
pub fn detect(&self, rgb: &[[f32; 3]], w: usize, h: usize, roi: &Roi) -> Result<Vec<Detection>> {
|
||||
let x = crop(rgb, w, h, roi, DET_SIZE, [-1.0, 1.0]);
|
||||
let out = self.detector.run(&[&x])?;
|
||||
let (Some(reg), Some(cls)) = (out.first(), out.get(1)) else { return Err("face detector: missing outputs".into()) };
|
||||
let s = DET_SIZE as f32;
|
||||
let mut dets = vec![];
|
||||
for (i, a) in self.anchors.iter().enumerate() {
|
||||
let (Some(r), Some(&c)) = (reg.get(i * 16..i * 16 + 16), cls.get(i)) else { break };
|
||||
let score = sigmoid(c.clamp(-100.0, 100.0));
|
||||
if score < MIN_SCORE {
|
||||
continue;
|
||||
}
|
||||
let (cx, cy, bw, bh) = (r[0] / s + a[0], r[1] / s + a[1], r[2] / s, r[3] / s);
|
||||
let to = |u: f32, v: f32| roi.to_frame(u, v);
|
||||
let corners =
|
||||
[to(cx - bw / 2.0, cy - bh / 2.0), to(cx + bw / 2.0, cy - bh / 2.0), to(cx - bw / 2.0, cy + bh / 2.0), to(cx + bw / 2.0, cy + bh / 2.0)];
|
||||
let mut keypoints = [[0.0f32; 2]; 6];
|
||||
for (k, p) in keypoints.iter_mut().enumerate() {
|
||||
*p = to(r[4 + 2 * k] / s + a[0], r[5 + 2 * k] / s + a[1]);
|
||||
}
|
||||
dets.push(Detection { bbox: bounds(&corners), keypoints, score });
|
||||
}
|
||||
Ok(weighted_nms(dets))
|
||||
}
|
||||
|
||||
/// The face mesh in `roi`: 478 points (frame pixels) and the presence score.
|
||||
pub fn mesh(&self, rgb: &[[f32; 3]], w: usize, h: usize, roi: &Roi) -> Result<Face> {
|
||||
let x = crop(rgb, w, h, roi, MESH_SIZE, [0.0, 1.0]);
|
||||
let out = self.mesh.run(&[&x])?;
|
||||
let (Some(pts), Some(flag)) = (out.first(), out.get(1).and_then(|f| f.first())) else { return Err("face mesh: missing outputs".into()) };
|
||||
let s = MESH_SIZE as f32;
|
||||
let points = pts.chunks_exact(3).take(MESH_POINTS).map(|p| roi.to_frame(p[0] / s, p[1] / s)).collect();
|
||||
Ok(Face { points, score: sigmoid(*flag) })
|
||||
}
|
||||
|
||||
/// The crop that frames a face found earlier.
|
||||
fn roi_of_face(face: &Face) -> Option<Roi> {
|
||||
let (a, b) = (face.points.get(EYE_LINE[0])?, face.points.get(EYE_LINE[1])?);
|
||||
Some(roi_of(bounds(&face.points), *a, *b))
|
||||
}
|
||||
|
||||
/// The mesh from `roi`, refined once from its own points (a detector box is coarse).
|
||||
fn mesh_refined(&self, rgb: &[[f32; 3]], w: usize, h: usize, roi: &Roi) -> Result<Option<Face>> {
|
||||
let first = self.mesh(rgb, w, h, roi)?;
|
||||
if first.score < MIN_SCORE {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(r) = Self::roi_of_face(&first) else { return Ok(None) };
|
||||
let refined = self.mesh(rgb, w, h, &r)?;
|
||||
Ok(Some(if refined.score >= MIN_SCORE { refined } else { first }))
|
||||
}
|
||||
}
|
||||
|
||||
impl FaceModel for FaceLandmarker {
|
||||
fn info(&self) -> &'static ModelInfo {
|
||||
&FACE_LANDMARKER
|
||||
}
|
||||
|
||||
fn topology(&self) -> &'static Topology {
|
||||
&TOPOLOGY
|
||||
}
|
||||
|
||||
fn find(&self, rgb: &[[f32; 3]], w: usize, h: usize, region: [f32; 4]) -> Result<Option<Face>> {
|
||||
let inside = |d: &Detection| {
|
||||
let c = [(d.bbox[0] + d.bbox[2]) / 2.0, (d.bbox[1] + d.bbox[3]) / 2.0];
|
||||
c[0] >= region[0] && c[0] <= region[2] && c[1] >= region[1] && c[1] <= region[3]
|
||||
};
|
||||
// Around the region first (a face filling the mask), then the whole frame (a small face).
|
||||
let side = (region[2] - region[0]).max(region[3] - region[1]);
|
||||
let around = Roi { center: [(region[0] + region[2]) / 2.0, (region[1] + region[3]) / 2.0], size: side * ROI_SCALE, angle: 0.0 };
|
||||
let whole = Roi { center: [w as f32 / 2.0, h as f32 / 2.0], size: w.max(h) as f32, angle: 0.0 };
|
||||
for roi in [around, whole] {
|
||||
if !roi.size.is_finite() || roi.size <= 4.0 {
|
||||
continue;
|
||||
}
|
||||
if let Some(d) = self.detect(rgb, w, h, &roi)?.into_iter().find(inside) {
|
||||
return self.mesh_refined(rgb, w, h, &roi_of(d.bbox, d.keypoints[0], d.keypoints[1]));
|
||||
}
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn follow(&self, rgb: &[[f32; 3]], w: usize, h: usize, prev: &Face) -> Result<Option<Face>> {
|
||||
let Some(roi) = Self::roi_of_face(prev) else { return Ok(None) };
|
||||
let face = self.mesh(rgb, w, h, &roi)?;
|
||||
if face.score >= MIN_SCORE {
|
||||
return Ok(Some(face));
|
||||
}
|
||||
// Lost by the mesh (a fast move): look for it again near where it was.
|
||||
let b = bounds(&prev.points);
|
||||
let (gw, gh) = ((b[2] - b[0]) * 0.5, (b[3] - b[1]) * 0.5);
|
||||
self.find(rgb, w, h, [b[0] - gw, b[1] - gh, b[2] + gw, b[3] + gh])
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn anchors_match_the_detector_layout() {
|
||||
let a = anchors();
|
||||
assert_eq!(a.len(), 896);
|
||||
// Stride 8: 16×16 cells, two anchors each; then stride 16: 8×8 cells, six each.
|
||||
assert_eq!(a[0], [0.5 / 16.0, 0.5 / 16.0]);
|
||||
assert_eq!(a[1], a[0]);
|
||||
assert_eq!(a[2], [1.5 / 16.0, 0.5 / 16.0]);
|
||||
assert_eq!(a[512], [0.5 / 8.0, 0.5 / 8.0]);
|
||||
assert_eq!(a[517], a[512]);
|
||||
assert_eq!(a[518], [1.5 / 8.0, 0.5 / 8.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn nms_averages_overlapping_detections() {
|
||||
let d = |x: f32, s: f32| Detection { bbox: [x, 0.0, x + 10.0, 10.0], keypoints: [[x, 0.0]; 6], score: s };
|
||||
let out = weighted_nms(vec![d(0.0, 0.9), d(1.0, 0.9), d(50.0, 0.6)]);
|
||||
assert_eq!(out.len(), 2);
|
||||
assert!((out[0].bbox[0] - 0.5).abs() < 1e-5 && out[0].score == 0.9);
|
||||
assert_eq!(out[1].bbox[0], 50.0);
|
||||
}
|
||||
|
||||
/// The official bundle (`EFFECTCRAFT_FACE_LANDMARKER` = path to `face_landmarker.task`) on a
|
||||
/// fixed input matches TensorFlow Lite's own interpreter (`ai-edge-litert` 2.2.0, run once
|
||||
/// to record these sums).
|
||||
#[test]
|
||||
fn matches_tensorflow_lite_on_the_official_weights() {
|
||||
let Some(path) = std::env::var_os("EFFECTCRAFT_FACE_LANDMARKER") else { return };
|
||||
let bytes = std::fs::read(path).unwrap();
|
||||
let m = FaceLandmarker::from_task(&bytes).unwrap();
|
||||
let input = |n: usize| -> Vec<f32> { (0..n).map(|i| ((i as f32) * 0.013).sin() * 0.5).collect() };
|
||||
let close = |got: &[f32], want: f64| {
|
||||
let s: f64 = got.iter().map(|v| *v as f64).sum();
|
||||
assert!((s - want).abs() <= 1e-5 * want.abs().max(1.0), "{s} vs {want}");
|
||||
};
|
||||
let d = m.detector.run(&[&input(DET_SIZE * DET_SIZE * 3)]).unwrap();
|
||||
close(&d[0], 54_202.040_273);
|
||||
close(&d[1], -3_061.866_384);
|
||||
let y = m.mesh.run(&[&input(MESH_SIZE * MESH_SIZE * 3)]).unwrap();
|
||||
close(&y[0], 119_864.749_257);
|
||||
close(&y[1], -19.777_233);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn topology_indices_are_mesh_points() {
|
||||
assert!(TOPOLOGY.outline.iter().chain(TOPOLOGY.landmarks.iter().map(|l| &l.1)).all(|&i| i < MESH_POINTS));
|
||||
assert_eq!(TOPOLOGY.outline.len(), 36);
|
||||
}
|
||||
}
|
||||
+29
-16
@@ -101,22 +101,28 @@ impl Conv {
|
||||
|
||||
pub fn forward(&self, x: &[f32], h: usize, w: usize) -> (Vec<f32>, usize, usize) {
|
||||
let (oh, ow) = (self.out_size(h), self.out_size(w));
|
||||
(self.forward_sized(x, h, w, oh, ow, [self.pad; 2]), oh, ow)
|
||||
}
|
||||
|
||||
/// An `oh×ow` output with `pad` (top, left) zero rows and columns before the input; what
|
||||
/// falls past the input's far edges is zero too (so asymmetric "same" padding works).
|
||||
pub fn forward_sized(&self, x: &[f32], h: usize, w: usize, oh: usize, ow: usize, pad: [usize; 2]) -> Vec<f32> {
|
||||
let kk = self.ks * self.ks * self.inp;
|
||||
let mut y: Vec<f32> = (0..oh * ow).flat_map(|_| self.b.iter().copied()).collect();
|
||||
if x.len() < h * w * self.inp {
|
||||
return (y, oh, ow);
|
||||
if x.len() < h * w * self.inp || self.wt.len() < kk * self.out {
|
||||
return y;
|
||||
}
|
||||
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;
|
||||
let iy = (oy * self.stride + ky) as isize - pad[0] 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;
|
||||
let ix = (ox * self.stride + kx) as isize - pad[1] as isize;
|
||||
if ix < 0 || ix >= w as isize {
|
||||
continue;
|
||||
}
|
||||
@@ -134,17 +140,18 @@ impl Conv {
|
||||
general_mat_mul(1.0, &av, &bv, 1.0, &mut cv);
|
||||
}
|
||||
});
|
||||
(y, oh, ow)
|
||||
y
|
||||
}
|
||||
}
|
||||
|
||||
/// A depthwise 3×3 convolution (padding 1, stride 1 or 2) on `h×w×c`.
|
||||
/// A depthwise `ks×ks` convolution (one filter per channel) on `h×w×c`.
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct Depthwise {
|
||||
/// `[tap × c]`.
|
||||
pub w: Vec<f32>,
|
||||
pub b: Vec<f32>,
|
||||
pub c: usize,
|
||||
pub ks: usize,
|
||||
pub stride: usize,
|
||||
}
|
||||
|
||||
@@ -160,32 +167,38 @@ impl Depthwise {
|
||||
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 })
|
||||
Some(Depthwise { w: wt, b: b.map_or_else(|| vec![0.0; c], <[f32]>::to_vec), c, ks: 3, stride })
|
||||
}
|
||||
|
||||
/// 3×3 with padding 1.
|
||||
pub fn forward(&self, x: &[f32], h: usize, w: usize) -> (Vec<f32>, usize, usize) {
|
||||
let c = self.c;
|
||||
let (oh, ow) = ((h + 2 - 3) / self.stride + 1, (w + 2 - 3) / self.stride + 1);
|
||||
(self.forward_sized(x, h, w, oh, ow, [1, 1]), oh, ow)
|
||||
}
|
||||
|
||||
/// An `oh×ow` output with `pad` (top, left) before the input, as [`Conv::forward_sized`].
|
||||
pub fn forward_sized(&self, x: &[f32], h: usize, w: usize, oh: usize, ow: usize, pad: [usize; 2]) -> Vec<f32> {
|
||||
let (c, ks) = (self.c, self.ks);
|
||||
let mut y = vec![0.0f32; oh * ow * c];
|
||||
if x.len() < h * w * c {
|
||||
return (y, oh, ow);
|
||||
if x.len() < h * w * c || self.w.len() < ks * ks * c || self.b.len() < c {
|
||||
return y;
|
||||
}
|
||||
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;
|
||||
for ky in 0..ks {
|
||||
let iy = (oy * self.stride + ky) as isize - pad[0] as isize;
|
||||
if iy < 0 || iy >= h as isize {
|
||||
continue;
|
||||
}
|
||||
for kx in 0..3 {
|
||||
let ix = (ox * self.stride + kx) as isize - 1;
|
||||
for kx in 0..ks {
|
||||
let ix = (ox * self.stride + kx) as isize - pad[1] as isize;
|
||||
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];
|
||||
let tap = &self.w[(ky * ks + kx) * c..][..c];
|
||||
for ((o, s), k) in out.iter_mut().zip(src).zip(tap) {
|
||||
*o += s * k;
|
||||
}
|
||||
@@ -193,7 +206,7 @@ impl Depthwise {
|
||||
}
|
||||
}
|
||||
});
|
||||
(y, oh, ow)
|
||||
y
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -102,6 +102,12 @@ fn read_entry<'a>(b: &'a [u8], e: &Entry) -> Result<std::borrow::Cow<'a, [u8]>>
|
||||
}
|
||||
}
|
||||
|
||||
/// The bytes of the archive entry `name` (model bundles are zip archives too).
|
||||
pub(crate) fn zip_file<'a>(b: &'a [u8], name: &str) -> Result<std::borrow::Cow<'a, [u8]>> {
|
||||
let entries = zip_entries(b)?;
|
||||
read_entry(b, entries.get(name).ok_or_else(|| format!("{name} is missing from the archive"))?)
|
||||
}
|
||||
|
||||
/// Pickle values (what state dicts use).
|
||||
#[derive(Clone, Debug)]
|
||||
enum Obj {
|
||||
@@ -341,7 +347,7 @@ fn to_f32(dtype: &str, bytes: &[u8]) -> Option<Vec<f32>> {
|
||||
}
|
||||
}
|
||||
|
||||
fn half_to_f32(h: u16) -> f32 {
|
||||
pub(crate) 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,0 +1,660 @@
|
||||
//! Running TensorFlow Lite models without TensorFlow: a reader for the `.tflite` FlatBuffer
|
||||
//! (fields as numbered in TensorFlow Lite's published `schema.fbs`, Apache-2.0) and an
|
||||
//! interpreter for the float operators small vision models use: `CONV_2D`, `DEPTHWISE_CONV_2D`,
|
||||
//! `MAX_POOL_2D`, `ADD`, `PAD`, `CONCATENATION`, `RESHAPE`, `RELU`, `RELU6`, `PRELU`, `LOGISTIC`
|
||||
//! and `DEQUANTIZE` of float16 weights (folded at load). Weights are re-laid out once at load for
|
||||
//! the [`nn`](crate::nn) kernels. Anything else (other operators, quantised tensors, dilation) is a
|
||||
//! load error, never a panic: models are pinned by SHA-256, so what loads once always loads.
|
||||
|
||||
use crate::nn::{Conv, Depthwise, Linear};
|
||||
use crate::pt::half_to_f32;
|
||||
|
||||
type Result<T> = std::result::Result<T, String>;
|
||||
|
||||
// ---------------------------------------------------------------- FlatBuffers
|
||||
|
||||
/// A FlatBuffer: bounds-checked little-endian reads.
|
||||
#[derive(Clone, Copy)]
|
||||
struct Fb<'a>(&'a [u8]);
|
||||
|
||||
/// A table: its position in the buffer.
|
||||
#[derive(Clone, Copy)]
|
||||
struct Table(usize);
|
||||
|
||||
impl<'a> Fb<'a> {
|
||||
fn bytes<const N: usize>(&self, at: usize) -> Result<[u8; N]> {
|
||||
self.0.get(at..at.checked_add(N).ok_or("bad offset")?).and_then(|s| s.try_into().ok()).ok_or_else(|| "truncated model".to_string())
|
||||
}
|
||||
fn u8(&self, at: usize) -> Result<u8> {
|
||||
Ok(self.bytes::<1>(at)?[0])
|
||||
}
|
||||
fn u16(&self, at: usize) -> Result<u16> {
|
||||
Ok(u16::from_le_bytes(self.bytes(at)?))
|
||||
}
|
||||
fn u32(&self, at: usize) -> Result<u32> {
|
||||
Ok(u32::from_le_bytes(self.bytes(at)?))
|
||||
}
|
||||
fn i32(&self, at: usize) -> Result<i32> {
|
||||
Ok(i32::from_le_bytes(self.bytes(at)?))
|
||||
}
|
||||
fn u64(&self, at: usize) -> Result<u64> {
|
||||
Ok(u64::from_le_bytes(self.bytes(at)?))
|
||||
}
|
||||
/// Follow the offset stored at `at`.
|
||||
fn deref(&self, at: usize) -> Result<usize> {
|
||||
at.checked_add(self.u32(at)? as usize).ok_or_else(|| "bad offset".to_string())
|
||||
}
|
||||
fn root(&self) -> Result<Table> {
|
||||
Ok(Table(self.deref(0)?))
|
||||
}
|
||||
/// Where field `id` of `t` is stored (absent fields have their default).
|
||||
fn field(&self, t: Table, id: usize) -> Result<Option<usize>> {
|
||||
let vt = (t.0 as i64 - self.i32(t.0)? as i64).try_into().map_err(|_| "bad vtable")?;
|
||||
let vt_len = self.u16(vt)? as usize;
|
||||
let slot = 4 + 2 * id;
|
||||
if slot + 2 > vt_len {
|
||||
return Ok(None);
|
||||
}
|
||||
let off = self.u16(vt + slot)? as usize;
|
||||
Ok((off != 0).then_some(t.0 + off))
|
||||
}
|
||||
fn int(&self, t: Table, id: usize, default: i32) -> Result<i32> {
|
||||
self.field(t, id)?.map_or(Ok(default), |p| self.i32(p))
|
||||
}
|
||||
fn byte(&self, t: Table, id: usize, default: u8) -> Result<u8> {
|
||||
self.field(t, id)?.map_or(Ok(default), |p| self.u8(p))
|
||||
}
|
||||
fn table(&self, t: Table, id: usize) -> Result<Option<Table>> {
|
||||
self.field(t, id)?.map(|p| self.deref(p).map(Table)).transpose()
|
||||
}
|
||||
/// A vector field: (start of its elements, length).
|
||||
fn vector(&self, t: Table, id: usize) -> Result<(usize, usize)> {
|
||||
let Some(p) = self.field(t, id)? else { return Ok((0, 0)) };
|
||||
let v = self.deref(p)?;
|
||||
let n = self.u32(v)? as usize;
|
||||
if v.saturating_add(4).saturating_add(n) > self.0.len() {
|
||||
return Err("truncated model".into());
|
||||
}
|
||||
Ok((v + 4, n))
|
||||
}
|
||||
fn tables(&self, t: Table, id: usize) -> Result<Vec<Table>> {
|
||||
let (at, n) = self.vector(t, id)?;
|
||||
(0..n).map(|i| self.deref(at + 4 * i).map(Table)).collect()
|
||||
}
|
||||
fn ints(&self, t: Table, id: usize) -> Result<Vec<i32>> {
|
||||
let (at, n) = self.vector(t, id)?;
|
||||
(0..n).map(|i| self.i32(at + 4 * i)).collect()
|
||||
}
|
||||
fn string(&self, t: Table, id: usize) -> Result<String> {
|
||||
let (at, n) = self.vector(t, id)?;
|
||||
Ok(String::from_utf8_lossy(self.0.get(at..at + n).ok_or("truncated model")?).into_owned())
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------- the model
|
||||
|
||||
/// Operator codes (`BuiltinOperator`).
|
||||
mod op {
|
||||
pub const ADD: i32 = 0;
|
||||
pub const CONCATENATION: i32 = 2;
|
||||
pub const CONV_2D: i32 = 3;
|
||||
pub const DEPTHWISE_CONV_2D: i32 = 4;
|
||||
pub const DEQUANTIZE: i32 = 6;
|
||||
pub const LOGISTIC: i32 = 14;
|
||||
pub const MAX_POOL_2D: i32 = 17;
|
||||
pub const RELU: i32 = 19;
|
||||
pub const RELU6: i32 = 21;
|
||||
pub const RESHAPE: i32 = 22;
|
||||
pub const PAD: i32 = 34;
|
||||
pub const PRELU: i32 = 54;
|
||||
}
|
||||
|
||||
/// The most values a tensor may hold.
|
||||
const MAX_TENSOR: usize = 1 << 28;
|
||||
|
||||
/// Tensor types (`TensorType`).
|
||||
const FLOAT32: u8 = 0;
|
||||
const FLOAT16: u8 = 1;
|
||||
const INT32: u8 = 2;
|
||||
|
||||
/// A fused activation (`ActivationFunctionType`).
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
enum Act {
|
||||
None,
|
||||
Relu,
|
||||
ReluN1To1,
|
||||
Relu6,
|
||||
Tanh,
|
||||
}
|
||||
|
||||
impl Act {
|
||||
fn from(code: u8) -> Result<Act> {
|
||||
Ok(match code {
|
||||
0 => Act::None,
|
||||
1 => Act::Relu,
|
||||
2 => Act::ReluN1To1,
|
||||
3 => Act::Relu6,
|
||||
4 => Act::Tanh,
|
||||
c => return Err(format!("unsupported fused activation {c}")),
|
||||
})
|
||||
}
|
||||
fn apply(self, x: &mut [f32]) {
|
||||
let f: fn(f32) -> f32 = match self {
|
||||
Act::None => return,
|
||||
Act::Relu => |v| v.max(0.0),
|
||||
Act::ReluN1To1 => |v| v.clamp(-1.0, 1.0),
|
||||
Act::Relu6 => |v| v.clamp(0.0, 6.0),
|
||||
Act::Tanh => f32::tanh,
|
||||
};
|
||||
for v in x {
|
||||
*v = f(*v);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// One operator, its weights laid out for the kernels.
|
||||
#[derive(Clone, Debug)]
|
||||
enum Op {
|
||||
/// A 1×1, stride-1 convolution: a dense layer over the pixels.
|
||||
Pointwise(Linear, Act),
|
||||
Conv(Conv, [usize; 2], Act),
|
||||
Depthwise(Depthwise, [usize; 2], Act),
|
||||
/// Kernel, stride, padding before (top, left).
|
||||
MaxPool {
|
||||
k: [usize; 2],
|
||||
stride: [usize; 2],
|
||||
pad: [usize; 2],
|
||||
act: Act,
|
||||
},
|
||||
Add(Act),
|
||||
/// Constant padding (zeros), per dimension (before, after).
|
||||
Pad(Vec<[usize; 2]>),
|
||||
Concat {
|
||||
axis: usize,
|
||||
act: Act,
|
||||
},
|
||||
Reshape,
|
||||
Act(Act),
|
||||
Prelu(Vec<f32>),
|
||||
Logistic,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct Step {
|
||||
op: Op,
|
||||
/// Runtime inputs (tensor ids).
|
||||
inputs: Vec<usize>,
|
||||
output: usize,
|
||||
}
|
||||
|
||||
/// A loaded model, ready to run (batch of one, static shapes).
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Model {
|
||||
shapes: Vec<Vec<usize>>,
|
||||
names: Vec<String>,
|
||||
/// Constant float tensors (weights folded into the steps are not kept).
|
||||
consts: Vec<Option<Vec<f32>>>,
|
||||
steps: Vec<Step>,
|
||||
inputs: Vec<usize>,
|
||||
outputs: Vec<usize>,
|
||||
/// After which step each tensor is last needed (to free memory as the model runs).
|
||||
last_use: Vec<usize>,
|
||||
}
|
||||
|
||||
fn index(i: i32, n: usize) -> Result<usize> {
|
||||
usize::try_from(i).ok().filter(|&i| i < n).ok_or_else(|| format!("bad tensor index {i}"))
|
||||
}
|
||||
|
||||
/// "Same" padding: output size and padding before.
|
||||
fn same(input: usize, k: usize, stride: usize) -> (usize, usize) {
|
||||
let out = input.div_ceil(stride);
|
||||
let total = ((out.saturating_sub(1)) * stride + k).saturating_sub(input);
|
||||
(out, total / 2)
|
||||
}
|
||||
|
||||
/// Output size and padding before, for `padding` (0 same, 1 valid).
|
||||
fn window(padding: u8, input: usize, k: usize, stride: usize) -> Result<(usize, usize)> {
|
||||
match padding {
|
||||
0 => Ok(same(input, k, stride)),
|
||||
1 if input >= k => Ok(((input - k) / stride + 1, 0)),
|
||||
_ => Err("bad window".into()),
|
||||
}
|
||||
}
|
||||
|
||||
impl Model {
|
||||
/// Read a `.tflite` file.
|
||||
pub fn read(bytes: &[u8]) -> Result<Model> {
|
||||
let fb = Fb(bytes);
|
||||
let root = fb.root()?;
|
||||
let codes: Vec<i32> = fb.tables(root, 1)?.into_iter().map(|c| Ok(fb.int(c, 3, 0)?.max(fb.byte(c, 0, 0)? as i8 as i32))).collect::<Result<_>>()?;
|
||||
let buffers = fb.tables(root, 4)?;
|
||||
let graphs = fb.tables(root, 2)?;
|
||||
let g = *graphs.first().ok_or("the model has no graph")?;
|
||||
let tensors = fb.tables(g, 0)?;
|
||||
let n = tensors.len();
|
||||
let mut shapes = Vec::with_capacity(n);
|
||||
let mut names = Vec::with_capacity(n);
|
||||
// Constant data: float tensors as f32, int tensors (pads, shapes) as i32.
|
||||
let mut floats: Vec<Option<Vec<f32>>> = vec![None; n];
|
||||
let mut ints: Vec<Option<Vec<i32>>> = vec![None; n];
|
||||
for (i, t) in tensors.iter().enumerate() {
|
||||
let shape: Vec<usize> = fb
|
||||
.ints(*t, 0)?
|
||||
.into_iter()
|
||||
.map(|d| usize::try_from(d).map_err(|_| "dynamic shapes are not supported"))
|
||||
.collect::<std::result::Result<_, _>>()?;
|
||||
let ty = fb.byte(*t, 1, FLOAT32)?;
|
||||
let buf = fb.int(*t, 2, 0)? as u32 as usize;
|
||||
names.push(fb.string(*t, 3)?);
|
||||
// Every size computed from shapes later stays in range.
|
||||
let count = shape.iter().try_fold(1usize, |a, &d| a.checked_mul(d)).filter(|&c| c <= MAX_TENSOR).ok_or("a tensor is too large")?;
|
||||
if fb.field(*t, 4)?.is_some() && fb.table(*t, 4)?.is_some_and(|q| fb.vector(q, 2).is_ok_and(|v| v.1 > 0)) {
|
||||
return Err(format!("{}: quantised tensors are not supported", names[i]));
|
||||
}
|
||||
if buf > 0 {
|
||||
let b = *buffers.get(buf).ok_or("bad buffer index")?;
|
||||
let (mut at, mut len) = fb.vector(b, 0)?;
|
||||
if len == 0 {
|
||||
// Large models keep data after the FlatBuffer: offset and size from its start.
|
||||
let (o, s) = (fb.field(b, 1)?.map(|p| fb.u64(p)).transpose()?.unwrap_or(0), fb.field(b, 2)?.map(|p| fb.u64(p)).transpose()?.unwrap_or(0));
|
||||
(at, len) = (o as usize, s as usize);
|
||||
}
|
||||
if len > 0 {
|
||||
let data = bytes.get(at..at.checked_add(len).ok_or("bad buffer")?).ok_or("truncated model")?;
|
||||
match ty {
|
||||
FLOAT32 if data.len() == 4 * count => {
|
||||
floats[i] = Some(data.chunks_exact(4).map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect())
|
||||
}
|
||||
FLOAT16 if data.len() == 2 * count => {
|
||||
floats[i] = Some(data.chunks_exact(2).map(|c| half_to_f32(u16::from_le_bytes([c[0], c[1]]))).collect())
|
||||
}
|
||||
INT32 if data.len() == 4 * count => {
|
||||
ints[i] = Some(data.chunks_exact(4).map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect())
|
||||
}
|
||||
_ => return Err(format!("{}: unsupported tensor data", names[i])),
|
||||
}
|
||||
}
|
||||
}
|
||||
shapes.push(shape);
|
||||
}
|
||||
let inputs = fb.ints(g, 1)?.into_iter().map(|i| index(i, n)).collect::<Result<Vec<_>>>()?;
|
||||
let outputs = fb.ints(g, 2)?.into_iter().map(|i| index(i, n)).collect::<Result<Vec<_>>>()?;
|
||||
let mut steps = vec![];
|
||||
for o in fb.tables(g, 3)? {
|
||||
let code = *codes.get(fb.int(o, 0, 0)? as u32 as usize).ok_or("bad operator code")?;
|
||||
// Optional inputs are -1.
|
||||
let ins: Vec<Option<usize>> = fb.ints(o, 1)?.into_iter().map(|i| if i < 0 { Ok(None) } else { index(i, n).map(Some) }).collect::<Result<_>>()?;
|
||||
let outs = fb.ints(o, 2)?;
|
||||
let output = index(*outs.first().ok_or("an operator has no output")?, n)?;
|
||||
// Float16 weights: fold their DEQUANTIZE.
|
||||
if code == op::DEQUANTIZE
|
||||
&& let Some(f) = ins.first().copied().flatten().and_then(|i| floats[i].clone())
|
||||
{
|
||||
floats[output] = Some(f);
|
||||
continue;
|
||||
}
|
||||
let opts = fb.table(o, 4)?;
|
||||
let opt_int = |id: usize, default: i32| opts.map_or(Ok(default), |t| fb.int(t, id, default));
|
||||
let opt_byte = |id: usize, default: u8| opts.map_or(Ok(default), |t| fb.byte(t, id, default));
|
||||
let input = |k: usize| ins.get(k).copied().flatten().ok_or_else(|| format!("{}: missing input {k}", names[output]));
|
||||
let weights = |k: usize| -> Result<Vec<f32>> {
|
||||
let id = input(k)?;
|
||||
floats[id].clone().ok_or_else(|| format!("{}: input {k} is not constant", names[output]))
|
||||
};
|
||||
let bias = |count: usize| -> Result<Vec<f32>> {
|
||||
match ins.get(2).copied().flatten() {
|
||||
Some(id) => floats[id].clone().filter(|b| b.len() == count).ok_or_else(|| format!("{}: bad bias", names[output])),
|
||||
None => Ok(vec![0.0; count]),
|
||||
}
|
||||
};
|
||||
let in_shape = |k: usize| -> Result<Vec<usize>> { Ok(shapes[input(k)?].clone()) };
|
||||
let out_shape = shapes[output].clone();
|
||||
let stride = |w: usize, h: usize| -> Result<[usize; 2]> {
|
||||
let s = [opt_int(h, 1)?, opt_int(w, 1)?];
|
||||
if s.iter().any(|v| *v < 1) {
|
||||
return Err("bad stride".into());
|
||||
}
|
||||
Ok(s.map(|v| v as usize))
|
||||
};
|
||||
let mut runtime = vec![input(0)?];
|
||||
let op = match code {
|
||||
op::DEQUANTIZE => Op::Reshape,
|
||||
op::CONV_2D | op::DEPTHWISE_CONV_2D => {
|
||||
let depthwise = code == op::DEPTHWISE_CONV_2D;
|
||||
let (padding, s) = (opt_byte(0, 0)?, stride(1, 2)?);
|
||||
let act_id = if depthwise { 4 } else { 3 };
|
||||
let dil = if depthwise { [opt_int(6, 1)?, opt_int(5, 1)?] } else { [opt_int(5, 1)?, opt_int(4, 1)?] };
|
||||
if dil != [1, 1] {
|
||||
return Err(format!("{}: dilated convolutions are not supported", names[output]));
|
||||
}
|
||||
let act = Act::from(opt_byte(act_id, 0)?)?;
|
||||
let (x, f) = (in_shape(0)?, shapes[input(1)?].clone());
|
||||
let (&[1, h, w, cin], &[fo, kh, kw, fi], &[1, oh, ow, cout]) = (&x[..], &f[..], &out_shape[..]) else {
|
||||
return Err(format!("{}: unexpected convolution shapes", names[output]));
|
||||
};
|
||||
if kh != kw || s[0] != s[1] {
|
||||
return Err(format!("{}: only square kernels and strides are supported", names[output]));
|
||||
}
|
||||
let (ey, py) = window(padding, h, kh, s[0])?;
|
||||
let (ex, px) = window(padding, w, kw, s[1])?;
|
||||
if (ey, ex) != (oh, ow) {
|
||||
return Err(format!("{}: output size disagrees with the padding", names[output]));
|
||||
}
|
||||
let wt = weights(1)?;
|
||||
let b = bias(cout)?;
|
||||
if depthwise {
|
||||
if fo != 1 || fi != cout || cin != cout {
|
||||
return Err(format!("{}: depth multipliers are not supported", names[output]));
|
||||
}
|
||||
Op::Depthwise(Depthwise { w: wt, b, c: cout, ks: kh, stride: s[0] }, [py, px], act)
|
||||
} else {
|
||||
if fo != cout || fi != cin {
|
||||
return Err(format!("{}: unexpected filter shape", names[output]));
|
||||
}
|
||||
// OHWI → [(ky·ks + kx)·inp + ci] × out.
|
||||
let kk = kh * kw * cin;
|
||||
let mut t = vec![0.0f32; kk * cout];
|
||||
for o in 0..cout {
|
||||
for p in 0..kk {
|
||||
t[p * cout + o] = *wt.get(o * kk + p).ok_or("bad filter")?;
|
||||
}
|
||||
}
|
||||
if kh == 1 && s[0] == 1 {
|
||||
Op::Pointwise(Linear { wt: t, b, inp: cin, out: cout }, act)
|
||||
} else {
|
||||
Op::Conv(Conv { wt: t, b, inp: cin, out: cout, ks: kh, stride: s[0], pad: 0 }, [py, px], act)
|
||||
}
|
||||
}
|
||||
}
|
||||
op::MAX_POOL_2D => {
|
||||
let (padding, s) = (opt_byte(0, 0)?, stride(1, 2)?);
|
||||
let k = [opt_int(4, 1)?, opt_int(3, 1)?];
|
||||
if k.iter().any(|v| *v < 1) {
|
||||
return Err("bad pool size".into());
|
||||
}
|
||||
let k = k.map(|v| v as usize);
|
||||
let x = in_shape(0)?;
|
||||
let (&[1, h, w, _], &[1, oh, ow, _]) = (&x[..], &out_shape[..]) else { return Err("unexpected pool shapes".into()) };
|
||||
let (ey, py) = window(padding, h, k[0], s[0])?;
|
||||
let (ex, px) = window(padding, w, k[1], s[1])?;
|
||||
if (ey, ex) != (oh, ow) {
|
||||
return Err(format!("{}: output size disagrees with the padding", names[output]));
|
||||
}
|
||||
Op::MaxPool { k, stride: s, pad: [py, px], act: Act::from(opt_byte(5, 0)?)? }
|
||||
}
|
||||
op::ADD => {
|
||||
runtime.push(input(1)?);
|
||||
Op::Add(Act::from(opt_byte(0, 0)?)?)
|
||||
}
|
||||
op::PAD => {
|
||||
let p = ints[input(1)?].clone().ok_or("PAD: paddings must be constant")?;
|
||||
let pads = p
|
||||
.chunks_exact(2)
|
||||
.map(|c| Ok([usize::try_from(c[0]).map_err(|_| "bad pad")?, usize::try_from(c[1]).map_err(|_| "bad pad")?]))
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
Op::Pad(pads)
|
||||
}
|
||||
op::CONCATENATION => {
|
||||
runtime = (0..ins.len()).map(input).collect::<Result<_>>()?;
|
||||
let rank = out_shape.len() as i32;
|
||||
let a = opt_int(0, 0)?;
|
||||
let axis = if a < 0 { a + rank } else { a };
|
||||
Op::Concat { axis: index(axis, out_shape.len())?, act: Act::from(opt_byte(1, 0)?)? }
|
||||
}
|
||||
op::RESHAPE => Op::Reshape,
|
||||
op::RELU => Op::Act(Act::Relu),
|
||||
op::RELU6 => Op::Act(Act::Relu6),
|
||||
op::LOGISTIC => Op::Logistic,
|
||||
op::PRELU => {
|
||||
let alpha = weights(1)?;
|
||||
let c = *out_shape.last().ok_or("PRELU: bad shape")?;
|
||||
if alpha.len() != c {
|
||||
return Err(format!("{}: only per-channel PReLU is supported", names[output]));
|
||||
}
|
||||
Op::Prelu(alpha)
|
||||
}
|
||||
c => return Err(format!("operator {c} is not supported")),
|
||||
};
|
||||
steps.push(Step { op, inputs: runtime, output });
|
||||
}
|
||||
let mut last_use = vec![usize::MAX; n];
|
||||
for (k, s) in steps.iter().enumerate() {
|
||||
for &i in &s.inputs {
|
||||
last_use[i] = k;
|
||||
}
|
||||
}
|
||||
for &o in &outputs {
|
||||
last_use[o] = usize::MAX;
|
||||
}
|
||||
Ok(Model { shapes, names, consts: floats, steps, inputs, outputs, last_use })
|
||||
}
|
||||
|
||||
/// The shape of input `k`.
|
||||
pub fn input_shape(&self, k: usize) -> Option<&[usize]> {
|
||||
self.inputs.get(k).and_then(|&i| self.shapes.get(i)).map(Vec::as_slice)
|
||||
}
|
||||
|
||||
/// The shape and name of output `k`.
|
||||
pub fn output(&self, k: usize) -> Option<(&[usize], &str)> {
|
||||
let &i = self.outputs.get(k)?;
|
||||
Some((self.shapes.get(i)?.as_slice(), self.names.get(i)?.as_str()))
|
||||
}
|
||||
|
||||
/// Run the model on its inputs (row-major, NHWC); returns its outputs.
|
||||
pub fn run(&self, inputs: &[&[f32]]) -> Result<Vec<Vec<f32>>> {
|
||||
let mut vals: Vec<Option<Vec<f32>>> = vec![None; self.shapes.len()];
|
||||
for (k, &id) in self.inputs.iter().enumerate() {
|
||||
let x = inputs.get(k).ok_or("missing model input")?;
|
||||
if x.len() != self.shapes[id].iter().product::<usize>() {
|
||||
return Err(format!("input {k}: expected {:?}", self.shapes[id]));
|
||||
}
|
||||
vals[id] = Some(x.to_vec());
|
||||
}
|
||||
for (k, s) in self.steps.iter().enumerate() {
|
||||
let y = {
|
||||
let get = |i: usize| -> Result<&[f32]> {
|
||||
vals.get(i)
|
||||
.and_then(Option::as_deref)
|
||||
.or_else(|| self.consts.get(i).and_then(Option::as_deref))
|
||||
.ok_or_else(|| format!("{}: input not computed", self.names[s.output]))
|
||||
};
|
||||
let x = get(s.inputs[0])?;
|
||||
let xs = &self.shapes[s.inputs[0]];
|
||||
let ys = &self.shapes[s.output];
|
||||
let (h, w) = (xs.get(1).copied().unwrap_or(1), xs.get(2).copied().unwrap_or(1));
|
||||
let (oh, ow) = (ys.get(1).copied().unwrap_or(1), ys.get(2).copied().unwrap_or(1));
|
||||
let y = match &s.op {
|
||||
Op::Pointwise(l, a) => {
|
||||
let mut y = l.forward(x, h * w);
|
||||
a.apply(&mut y);
|
||||
y
|
||||
}
|
||||
Op::Conv(c, pad, a) => {
|
||||
let mut y = c.forward_sized(x, h, w, oh, ow, *pad);
|
||||
a.apply(&mut y);
|
||||
y
|
||||
}
|
||||
Op::Depthwise(d, pad, a) => {
|
||||
let mut y = d.forward_sized(x, h, w, oh, ow, *pad);
|
||||
a.apply(&mut y);
|
||||
y
|
||||
}
|
||||
Op::MaxPool { k, stride, pad, act } => {
|
||||
let mut y = max_pool(x, [h, w], *xs.last().unwrap_or(&1), [oh, ow], *k, *stride, *pad);
|
||||
act.apply(&mut y);
|
||||
y
|
||||
}
|
||||
Op::Add(a) => {
|
||||
let b = get(s.inputs[1])?;
|
||||
if b.len() != x.len() {
|
||||
return Err(format!("{}: broadcasting ADD is not supported", self.names[s.output]));
|
||||
}
|
||||
let mut y: Vec<f32> = x.iter().zip(b).map(|(p, q)| p + q).collect();
|
||||
a.apply(&mut y);
|
||||
y
|
||||
}
|
||||
Op::Pad(p) => pad(x, xs, p)?,
|
||||
Op::Concat { axis, act } => {
|
||||
let parts = s.inputs.iter().map(|&i| Ok((get(i)?, &self.shapes[i]))).collect::<Result<Vec<_>>>()?;
|
||||
let mut y = concat(&parts, *axis)?;
|
||||
act.apply(&mut y);
|
||||
y
|
||||
}
|
||||
Op::Reshape => x.to_vec(),
|
||||
Op::Act(a) => {
|
||||
let mut y = x.to_vec();
|
||||
a.apply(&mut y);
|
||||
y
|
||||
}
|
||||
Op::Prelu(alpha) => {
|
||||
let c = alpha.len().max(1);
|
||||
x.iter().enumerate().map(|(i, &v)| if v >= 0.0 { v } else { v * alpha[i % c] }).collect()
|
||||
}
|
||||
Op::Logistic => x.iter().map(|v| 1.0 / (1.0 + (-v).exp())).collect(),
|
||||
};
|
||||
if y.len() != ys.iter().product::<usize>() {
|
||||
return Err(format!("{}: computed {} values for shape {ys:?}", self.names[s.output], y.len()));
|
||||
}
|
||||
y
|
||||
};
|
||||
vals[s.output] = Some(y);
|
||||
for &i in &s.inputs {
|
||||
if self.last_use[i] == k {
|
||||
vals[i] = None;
|
||||
}
|
||||
}
|
||||
}
|
||||
self.outputs.iter().map(|&o| vals[o].take().or_else(|| self.consts[o].clone()).ok_or_else(|| "an output was not computed".to_string())).collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// Max pooling over `k` windows; window cells outside the input are ignored.
|
||||
fn max_pool(x: &[f32], [h, w]: [usize; 2], c: usize, [oh, ow]: [usize; 2], k: [usize; 2], s: [usize; 2], pad: [usize; 2]) -> Vec<f32> {
|
||||
let mut y = vec![f32::NEG_INFINITY; oh * ow * c];
|
||||
for oy in 0..oh {
|
||||
for ox in 0..ow {
|
||||
let out = &mut y[(oy * ow + ox) * c..][..c];
|
||||
for ky in 0..k[0] {
|
||||
let Some(iy) = (oy * s[0] + ky).checked_sub(pad[0]).filter(|&v| v < h) else { continue };
|
||||
for kx in 0..k[1] {
|
||||
let Some(ix) = (ox * s[1] + kx).checked_sub(pad[1]).filter(|&v| v < w) else { continue };
|
||||
let Some(src) = x.get((iy * w + ix) * c..(iy * w + ix + 1) * c) else { continue };
|
||||
for (o, v) in out.iter_mut().zip(src) {
|
||||
*o = o.max(*v);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
y
|
||||
}
|
||||
|
||||
/// Zero padding of a row-major tensor of shape `shape`.
|
||||
fn pad(x: &[f32], shape: &[usize], pads: &[[usize; 2]]) -> Result<Vec<f32>> {
|
||||
if pads.len() != shape.len() {
|
||||
return Err("PAD: rank mismatch".into());
|
||||
}
|
||||
let out: Vec<usize> = shape.iter().zip(pads).map(|(d, p)| d + p[0] + p[1]).collect();
|
||||
let mut y = vec![0.0f32; out.iter().product()];
|
||||
let rank = shape.len();
|
||||
let inner = *shape.last().unwrap_or(&1);
|
||||
let rows = x.len() / inner.max(1);
|
||||
// Copy each innermost row to its padded place.
|
||||
let mut idx = vec![0usize; rank.saturating_sub(1)];
|
||||
for r in 0..rows {
|
||||
let mut at = 0usize;
|
||||
for d in 0..rank {
|
||||
let i = if d + 1 == rank { pads[d][0] } else { idx[d] + pads[d][0] };
|
||||
at = at * out[d] + i;
|
||||
}
|
||||
if let (Some(dst), Some(src)) = (y.get_mut(at..at + inner), x.get(r * inner..(r + 1) * inner)) {
|
||||
dst.copy_from_slice(src);
|
||||
}
|
||||
for d in (0..idx.len()).rev() {
|
||||
idx[d] += 1;
|
||||
if idx[d] < shape[d] {
|
||||
break;
|
||||
}
|
||||
idx[d] = 0;
|
||||
}
|
||||
}
|
||||
Ok(y)
|
||||
}
|
||||
|
||||
/// Concatenate along `axis`.
|
||||
fn concat(parts: &[(&[f32], &Vec<usize>)], axis: usize) -> Result<Vec<f32>> {
|
||||
let Some((_, s0)) = parts.first() else { return Ok(vec![]) };
|
||||
let outer: usize = s0.get(..axis).ok_or("CONCATENATION: bad axis")?.iter().product();
|
||||
let mut y = Vec::with_capacity(parts.iter().map(|p| p.0.len()).sum());
|
||||
for o in 0..outer {
|
||||
for (x, s) in parts {
|
||||
let chunk: usize = s.get(axis..).ok_or("CONCATENATION: bad axis")?.iter().product();
|
||||
y.extend_from_slice(x.get(o * chunk..(o + 1) * chunk).ok_or("CONCATENATION: bad input")?);
|
||||
}
|
||||
}
|
||||
Ok(y)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn same_padding_matches_tensorflow() {
|
||||
// TensorFlow pads the extra row after: 128 → 64 with a 3×3 stride-2 window pads 0 before.
|
||||
assert_eq!(same(128, 3, 2), (64, 0));
|
||||
assert_eq!(same(128, 5, 2), (64, 1));
|
||||
assert_eq!(same(16, 3, 1), (16, 1));
|
||||
assert_eq!(same(7, 2, 2), (4, 0));
|
||||
assert_eq!(window(1, 4, 2, 2), Ok((2, 0)));
|
||||
assert!(window(1, 1, 2, 2).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn asymmetric_conv_padding() {
|
||||
// A 3×3 stride-2 "same" conv on 4×4: pads nothing before, one after.
|
||||
let x: Vec<f32> = (0..16).map(|v| v as f32).collect();
|
||||
let ones = vec![1.0f32; 9];
|
||||
let c = Conv { wt: ones.clone(), b: vec![0.0], inp: 1, out: 1, ks: 3, stride: 2, pad: 0 };
|
||||
let y = c.forward_sized(&x, 4, 4, 2, 2, [0, 0]);
|
||||
// Top-left window rows 0–2, cols 0–2; the bottom-right one is cut by the edge.
|
||||
assert_eq!(y, vec![45.0, 39.0, 66.0, 50.0]);
|
||||
let d = Depthwise { w: ones, b: vec![0.0], c: 1, ks: 3, stride: 2 };
|
||||
assert_eq!(d.forward_sized(&x, 4, 4, 2, 2, [0, 0]), y);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pool_pad_and_concat() {
|
||||
let x: Vec<f32> = (0..16).map(|v| v as f32).collect();
|
||||
assert_eq!(max_pool(&x, [4, 4], 1, [2, 2], [2, 2], [2, 2], [0, 0]), vec![5.0, 7.0, 13.0, 15.0]);
|
||||
// Same padding on 3×3 → 2×2: the last window is cut by the edge.
|
||||
let x3: Vec<f32> = (0..9).map(|v| v as f32).collect();
|
||||
assert_eq!(max_pool(&x3, [3, 3], 1, [2, 2], [2, 2], [2, 2], [0, 0]), vec![4.0, 5.0, 7.0, 8.0]);
|
||||
// Channel padding (residual connections widen the channels with zeros).
|
||||
let p = pad(&[1.0, 2.0, 3.0, 4.0], &[1, 1, 2, 2], &[[0, 0], [0, 0], [0, 0], [0, 1]]).unwrap();
|
||||
assert_eq!(p, vec![1.0, 2.0, 0.0, 3.0, 4.0, 0.0]);
|
||||
let p = pad(&[1.0, 2.0], &[1, 2], &[[0, 0], [1, 0]]).unwrap();
|
||||
assert_eq!(p, vec![0.0, 1.0, 2.0]);
|
||||
let a = [1.0, 2.0, 3.0, 4.0];
|
||||
let b = [5.0, 6.0];
|
||||
let (sa, sb) = (vec![1, 2, 2], vec![1, 1, 2]);
|
||||
assert_eq!(concat(&[(&a, &sa), (&b, &sb)], 1).unwrap(), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
|
||||
let (sa, sb) = (vec![2, 2], vec![2, 1]);
|
||||
assert_eq!(concat(&[(&a, &sa), (&b, &sb)], 1).unwrap(), vec![1.0, 2.0, 5.0, 3.0, 4.0, 6.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn damaged_files_are_errors_not_panics() {
|
||||
assert!(Model::read(b"").is_err());
|
||||
assert!(Model::read(&[0xff; 64]).is_err());
|
||||
// Every prefix and a few corruptions of a hand-made buffer.
|
||||
let mut fb = vec![12, 0, 0, 0, 0, 0, 6, 0, 8, 0, 4, 0, 6, 0, 0, 0, 0, 0, 0, 0];
|
||||
for n in 0..fb.len() {
|
||||
let _ = Model::read(&fb[..n]);
|
||||
}
|
||||
for i in 0..fb.len() {
|
||||
fb[i] ^= 0x5a;
|
||||
let _ = Model::read(&fb);
|
||||
}
|
||||
}
|
||||
}
|
||||
+273
-22
@@ -23,12 +23,19 @@
|
||||
//! ([`synth_shape`]: eye and mouth openness, mouth width, brow raise, head yaw and pitch) by
|
||||
//! principal component analysis — no external data or weights.
|
||||
//!
|
||||
//! With a trained model (Settings ▸ Face Tracking: an `effectcraft_segment::face::FaceModel` such
|
||||
//! as MediaPipe Face Landmarker), the model finds the face in the mask and follows it; its
|
||||
//! outline and named points replace steps 1–4, and chin and jaw come from its outline the same
|
||||
//! way, so measurements mean the same with either engine ([`FaceTracker::new_with`]). When the
|
||||
//! model finds no face in the mask, the classical pipeline runs.
|
||||
//!
|
||||
//! Coordinates are layer pixels; the face frame is centred on the face ellipse, with the
|
||||
//! half-height as unit, `u` right and `v` down (image-left is "Left").
|
||||
|
||||
use std::sync::OnceLock;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
|
||||
use effectcraft_raster::Image;
|
||||
use effectcraft_segment::face::{Face, FaceModel, Topology};
|
||||
|
||||
use crate::Frame;
|
||||
use crate::camtrack::linalg::{cholesky_solve, sym_eigen};
|
||||
@@ -435,7 +442,7 @@ pub fn fit_shape(detected: &[[f64; 2]; N], w: &[f64; N], init: [f64; 4]) -> ([[f
|
||||
// ---------------------------------------------------------------- tracking
|
||||
|
||||
/// The skin colour model (luma, r and g chromaticity: median and spread).
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq)]
|
||||
struct Skin {
|
||||
med: [f32; 3],
|
||||
tol: [f32; 3],
|
||||
@@ -475,6 +482,69 @@ pub struct FaceTracker {
|
||||
prev: FaceFit,
|
||||
/// Previous landmarks in the face frame (the expected feature positions).
|
||||
expect: [[f64; 2]; N],
|
||||
/// The trained model following the face, if one does.
|
||||
learned: Option<Learned>,
|
||||
}
|
||||
|
||||
/// A trained model and where it last saw the face.
|
||||
#[derive(Clone)]
|
||||
struct Learned {
|
||||
model: Arc<dyn FaceModel>,
|
||||
face: Face,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for Learned {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("Learned").field("model", &self.model.info().id).finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
/// A frame as trained models see it: straight RGB, size.
|
||||
fn model_input(frame: &Frame) -> (Vec<[f32; 3]>, usize, usize) {
|
||||
(crate::roto::rgb_of(frame.img), frame.img.width as usize, frame.img.height as usize)
|
||||
}
|
||||
|
||||
/// Where the outline meets the direction of canonical face-frame point `q` from `center`.
|
||||
fn outline_point(outline: &[[f64; 2]], center: [f64; 2], roll: f64, q: [f64; 2]) -> [f64; 2] {
|
||||
let dir = to_layer([0.0; 2], 1.0, roll, q);
|
||||
let rad = ray_radius(outline, center, dir[1].atan2(dir[0]));
|
||||
let len = dir[0].hypot(dir[1]).max(1e-12);
|
||||
[center[0] + dir[0] / len * rad, center[1] + dir[1] / len * rad]
|
||||
}
|
||||
|
||||
/// A trained model's face as a fit: its outline and named points (all measured); chin and jaw
|
||||
/// where the outline meets their canonical directions, as in the classical fit.
|
||||
fn fit_from_face(face: &Face, topo: &Topology, offset: [f64; 2]) -> Option<FaceFit> {
|
||||
let pt = |i: usize| face.points.get(i).map(|p| [p[0] as f64 - offset[0], p[1] as f64 - offset[1]]);
|
||||
let outline: Vec<[f64; 2]> = topo.outline.iter().map(|&i| pt(i)).collect::<Option<_>>()?;
|
||||
let (center, b, _, ellipse_roll) = polygon_ellipse(&outline)?;
|
||||
if !b.is_finite() || b <= 1.0 {
|
||||
return None;
|
||||
}
|
||||
let mut detected = [false; N];
|
||||
let mut named = [None; N];
|
||||
for (id, i) in topo.landmarks {
|
||||
if let (Some(k), Some(p)) = (landmark(id), pt(*i)) {
|
||||
named[k] = Some(p);
|
||||
}
|
||||
}
|
||||
// Roll: the line between the outer eye corners; else the outline's axis.
|
||||
let roll = match (named[EYE[0][1]], named[EYE[1][1]]) {
|
||||
(Some(l), Some(r)) => (r[1] - l[1]).atan2(r[0] - l[0]).to_degrees(),
|
||||
_ => ellipse_roll,
|
||||
};
|
||||
let mut landmarks = LANDMARKS.map(|l| to_layer(center, b, roll, l.2));
|
||||
for k in 0..N {
|
||||
if let Some(p) = named[k] {
|
||||
landmarks[k] = p;
|
||||
detected[k] = true;
|
||||
}
|
||||
}
|
||||
for k in [CHIN, JAW[0], JAW[1]] {
|
||||
landmarks[k] = outline_point(&outline, center, roll, LANDMARKS[k].2);
|
||||
detected[k] = true;
|
||||
}
|
||||
Some(FaceFit { landmarks, detected, outline, center, height: b, roll })
|
||||
}
|
||||
|
||||
/// A pixel window of a frame in layer coordinates.
|
||||
@@ -622,7 +692,7 @@ impl FaceTracker {
|
||||
height: hh.max(hw / ASPECT),
|
||||
roll: 0.0,
|
||||
};
|
||||
let mut t = FaceTracker { skin, prev: guess, expect: canon };
|
||||
let mut t = FaceTracker { skin, prev: guess, expect: canon, learned: None };
|
||||
let fit = t.fit(frame)?;
|
||||
// A face shows at least two of: left eye, right eye, mouth.
|
||||
let found = [EYE[0][4], EYE[1][4], MOUTH[0]].iter().filter(|i| fit.detected[**i]).count();
|
||||
@@ -633,6 +703,29 @@ impl FaceTracker {
|
||||
Some((t, fit))
|
||||
}
|
||||
|
||||
/// [`FaceTracker::new`] with a trained model, if any: the model finds the face in `region`
|
||||
/// and follows it. Without a model, or when it finds no face there, the classical tracker.
|
||||
pub fn new_with(frame: &Frame, region: [f64; 4], model: Option<Arc<dyn FaceModel>>) -> Option<(FaceTracker, FaceFit)> {
|
||||
if let Some(model) = model {
|
||||
let (rgb, w, h) = model_input(frame);
|
||||
let (ox, oy) = (frame.offset[0], frame.offset[1]);
|
||||
let r = [region[0] + ox, region[1] + oy, region[2] + ox, region[3] + oy].map(|v| v as f32);
|
||||
if let Ok(Some(face)) = model.find(&rgb, w, h, r)
|
||||
&& let Some(fit) = fit_from_face(&face, model.topology(), frame.offset)
|
||||
{
|
||||
let mut t = FaceTracker { skin: Skin::default(), prev: fit.clone(), expect: LANDMARKS.map(|l| l.2), learned: Some(Learned { model, face }) };
|
||||
t.accept(&fit);
|
||||
return Some((t, fit));
|
||||
}
|
||||
}
|
||||
Self::new(frame, region)
|
||||
}
|
||||
|
||||
/// The trained model following the face (`None`: the classical tracker).
|
||||
pub fn model(&self) -> Option<&'static effectcraft_segment::ModelInfo> {
|
||||
self.learned.as_ref().map(|l| l.model.info())
|
||||
}
|
||||
|
||||
/// The fit of the last tracked frame.
|
||||
pub fn last(&self) -> &FaceFit {
|
||||
&self.prev
|
||||
@@ -645,7 +738,16 @@ impl FaceTracker {
|
||||
|
||||
/// Track into the next frame (either direction). `None` when the face is lost.
|
||||
pub fn step(&mut self, frame: &Frame) -> Option<FaceFit> {
|
||||
let fit = self.fit(frame)?;
|
||||
let fit = match &mut self.learned {
|
||||
Some(l) => {
|
||||
let (rgb, w, h) = model_input(frame);
|
||||
let face = l.model.follow(&rgb, w, h, &l.face).ok().flatten()?;
|
||||
let fit = fit_from_face(&face, l.model.topology(), frame.offset)?;
|
||||
l.face = face;
|
||||
fit
|
||||
}
|
||||
None => self.fit(frame)?,
|
||||
};
|
||||
self.accept(&fit);
|
||||
Some(fit)
|
||||
}
|
||||
@@ -894,13 +996,7 @@ impl FaceTracker {
|
||||
}
|
||||
// Chin and jaw from the outline, along their canonical directions.
|
||||
for i in [CHIN, JAW[0], JAW[1]] {
|
||||
let q = LANDMARKS[i].2;
|
||||
let dir = to_layer([0.0; 2], 1.0, roll, q);
|
||||
let ang = dir[1].atan2(dir[0]);
|
||||
let rad = ray_radius(&outline, center, ang);
|
||||
let len = dir[0].hypot(dir[1]);
|
||||
let p = [center[0] + dir[0] / len * rad, center[1] + dir[1] / len * rad];
|
||||
det[i] = p;
|
||||
det[i] = outline_point(&outline, center, roll, LANDMARKS[i].2);
|
||||
wgt[i] = 1.0;
|
||||
}
|
||||
// Shape model: fill the missing landmarks and replace inconsistent detections.
|
||||
@@ -1033,6 +1129,19 @@ pub fn measure(l: &[[f64; 2]; N], reference_height: f64) -> [f64; 14] {
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Pearson correlation.
|
||||
fn corr(a: &[f64], b: &[f64]) -> f64 {
|
||||
let n = a.len() as f64;
|
||||
let (ma, mb) = (a.iter().sum::<f64>() / n, b.iter().sum::<f64>() / n);
|
||||
let (mut sab, mut saa, mut sbb) = (0.0, 0.0, 0.0);
|
||||
for (x, y) in a.iter().zip(b) {
|
||||
sab += (x - ma) * (y - mb);
|
||||
saa += (x - ma).powi(2);
|
||||
sbb += (y - mb).powi(2);
|
||||
}
|
||||
sab / (saa * sbb).sqrt()
|
||||
}
|
||||
|
||||
fn frame_params(t: usize) -> FaceParams {
|
||||
let tf = t as f64;
|
||||
FaceParams {
|
||||
@@ -1110,17 +1219,6 @@ mod tests {
|
||||
assert!(worst < 0.08, "worst landmark error {worst:.3}");
|
||||
assert!(outline_err < 0.04, "outline error {outline_err:.3}");
|
||||
// Measurements follow the generator: mouth openness and yaw correlate strongly.
|
||||
let corr = |a: &[f64], b: &[f64]| {
|
||||
let n = a.len() as f64;
|
||||
let (ma, mb) = (a.iter().sum::<f64>() / n, b.iter().sum::<f64>() / n);
|
||||
let (mut sab, mut saa, mut sbb) = (0.0, 0.0, 0.0);
|
||||
for (x, y) in a.iter().zip(b) {
|
||||
sab += (x - ma) * (y - mb);
|
||||
saa += (x - ma).powi(2);
|
||||
sbb += (y - mb).powi(2);
|
||||
}
|
||||
sab / (saa * sbb).sqrt()
|
||||
};
|
||||
let mo: Vec<f64> = meas.iter().map(|m| m.1[10]).collect();
|
||||
let mt: Vec<f64> = meas.iter().map(|m| m.0.mouth_open).collect();
|
||||
assert!(corr(&mo, &mt) > 0.95, "mouth openness r = {}", corr(&mo, &mt));
|
||||
@@ -1144,4 +1242,157 @@ mod tests {
|
||||
assert!((m[5] - 20.0).abs() < 1e-6);
|
||||
assert!((m[4] - 10.0).abs() < 0.5 && (m[3] + 5.0).abs() < 0.5, "{m:?}");
|
||||
}
|
||||
|
||||
/// A stand-in trained model: it "sees" the synthetic clip's true landmarks (points
|
||||
/// `0..23` in `LANDMARKS` order) and its true face ellipse (points `N..N + 36`), frame after
|
||||
/// frame, until `lose_at`; `blind`, it finds no face.
|
||||
struct Fake {
|
||||
frame: std::sync::atomic::AtomicUsize,
|
||||
lose_at: usize,
|
||||
blind: bool,
|
||||
}
|
||||
|
||||
fn fake(lose_at: usize, blind: bool) -> Arc<Fake> {
|
||||
Arc::new(Fake { frame: 0.into(), lose_at, blind })
|
||||
}
|
||||
|
||||
const OUTLINE_POINTS: usize = 36;
|
||||
static FAKE_OUTLINE: [usize; OUTLINE_POINTS] = {
|
||||
let mut a = [0; OUTLINE_POINTS];
|
||||
let mut i = 0;
|
||||
while i < OUTLINE_POINTS {
|
||||
a[i] = N + i;
|
||||
i += 1;
|
||||
}
|
||||
a
|
||||
};
|
||||
static FAKE_LANDMARKS: [(&str, usize); CHIN] = {
|
||||
let mut a = [("", 0); CHIN];
|
||||
let mut i = 0;
|
||||
while i < CHIN {
|
||||
a[i] = (LANDMARKS[i].0, i);
|
||||
i += 1;
|
||||
}
|
||||
a
|
||||
};
|
||||
static FAKE_TOPOLOGY: Topology = Topology { outline: &FAKE_OUTLINE, landmarks: &FAKE_LANDMARKS };
|
||||
static FAKE_INFO: effectcraft_segment::ModelInfo = effectcraft_segment::ModelInfo { id: "fake", name: "Fake", ..effectcraft_segment::FACE_LANDMARKER };
|
||||
|
||||
impl Fake {
|
||||
fn face(t: usize) -> Face {
|
||||
let p = frame_params(t);
|
||||
let mut points: Vec<[f32; 2]> = synth_landmarks(&p).iter().map(|q| [q[0] as f32, q[1] as f32]).collect();
|
||||
for k in 0..OUTLINE_POINTS {
|
||||
let a = k as f64 / OUTLINE_POINTS as f64 * std::f64::consts::TAU;
|
||||
let q = to_layer(p.center, p.height, p.roll, [ASPECT * a.sin(), -a.cos()]);
|
||||
points.push([q[0] as f32, q[1] as f32]);
|
||||
}
|
||||
Face { points, score: 1.0 }
|
||||
}
|
||||
}
|
||||
|
||||
impl FaceModel for Fake {
|
||||
fn info(&self) -> &'static effectcraft_segment::ModelInfo {
|
||||
&FAKE_INFO
|
||||
}
|
||||
fn topology(&self) -> &'static Topology {
|
||||
&FAKE_TOPOLOGY
|
||||
}
|
||||
fn find(&self, _: &[[f32; 3]], _: usize, _: usize, r: [f32; 4]) -> effectcraft_segment::Result<Option<Face>> {
|
||||
let c = frame_params(0).center;
|
||||
let inside = (r[0] as f64) < c[0] && c[0] < r[2] as f64 && (r[1] as f64) < c[1] && c[1] < r[3] as f64;
|
||||
Ok((inside && !self.blind).then(|| Fake::face(0)))
|
||||
}
|
||||
fn follow(&self, _: &[[f32; 3]], _: usize, _: usize, _: &Face) -> effectcraft_segment::Result<Option<Face>> {
|
||||
let t = self.frame.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
|
||||
Ok((t < self.lose_at).then(|| Fake::face(t)))
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_trained_model_drives_the_fit() {
|
||||
let (w, h) = (320, 240);
|
||||
let p0 = frame_params(0);
|
||||
let img0 = render_synthetic(w, h, &p0, 3);
|
||||
let frame0 = Frame { img: &img0, offset: [0.0; 2] };
|
||||
let region = [p0.center[0] - 60.0, p0.center[1] - 75.0, p0.center[0] + 60.0, p0.center[1] + 75.0];
|
||||
let (mut tr, fit0) = FaceTracker::new_with(&frame0, region, Some(fake(30, false))).expect("face found");
|
||||
assert_eq!(tr.model().map(|m| m.id), Some("fake"));
|
||||
let mut fits = vec![fit0];
|
||||
for t in 1..30 {
|
||||
let img = render_synthetic(w, h, &frame_params(t), 3);
|
||||
fits.push(tr.step(&Frame { img: &img, offset: [0.0; 2] }).unwrap_or_else(|| panic!("lost at {t}")));
|
||||
}
|
||||
let mut meas = vec![];
|
||||
for (t, f) in fits.iter().enumerate() {
|
||||
let p = frame_params(t);
|
||||
let truth = synth_landmarks(&p);
|
||||
// The model's points are the landmarks; chin and jaw come from its outline.
|
||||
for i in 0..N {
|
||||
let e = (f.landmarks[i][0] - truth[i][0]).hypot(f.landmarks[i][1] - truth[i][1]) / p.height;
|
||||
assert!(e < if i < CHIN { 1e-4 } else { 0.06 }, "frame {t} {}: {e}", LANDMARKS[i].0);
|
||||
}
|
||||
assert!(f.detected.iter().all(|d| *d) && f.outline.len() == OUTLINE_POINTS);
|
||||
assert!((f.height - p.height).abs() < 0.03 * p.height, "frame {t}: height {} vs {}", f.height, p.height);
|
||||
assert!((f.roll - p.roll).abs() < 3.0, "frame {t}: roll {} vs {}", f.roll, p.roll);
|
||||
meas.push((p, measure(&f.landmarks, fits[0].height)));
|
||||
}
|
||||
let pick = |k: usize| meas.iter().map(|m| m.1[k]).collect::<Vec<_>>();
|
||||
let mt: Vec<f64> = meas.iter().map(|m| m.0.mouth_open).collect();
|
||||
assert!(corr(&pick(10), &mt) > 0.95);
|
||||
let rt: Vec<f64> = meas.iter().map(|m| m.0.roll).collect();
|
||||
assert!(corr(&pick(5), &rt) > 0.98);
|
||||
// The model loses the face: so does the tracker.
|
||||
let img = render_synthetic(w, h, &frame_params(30), 3);
|
||||
assert!(tr.step(&Frame { img: &img, offset: [0.0; 2] }).is_none());
|
||||
// A model that finds no face in the mask leaves the face to the classical tracker.
|
||||
let (classic, _) = FaceTracker::new_with(&frame0, region, Some(fake(30, true))).expect("the classic tracker finds it");
|
||||
assert!(classic.model().is_none());
|
||||
// Layer offsets: frame pixels = layer + offset; fits stay in layer pixels.
|
||||
let shifted = Frame { img: &img0, offset: [10.0, -4.0] };
|
||||
let moved = [region[0] - 10.0, region[1] + 4.0, region[2] - 10.0, region[3] + 4.0];
|
||||
let (_, f) = FaceTracker::new_with(&shifted, moved, Some(fake(9, false))).expect("found");
|
||||
let truth = synth_landmarks(&p0);
|
||||
assert!((f.landmarks[EYE[0][4]][0] - (truth[EYE[0][4]][0] - 10.0)).abs() < 1e-3);
|
||||
}
|
||||
|
||||
/// With the official MediaPipe Face Landmarker (`EFFECTCRAFT_FACE_LANDMARKER` = path to
|
||||
/// `face_landmarker.task`, else skipped): it recognises the synthetic face and follows it
|
||||
/// through the moving, rolling clip. Brows and lips are drawn differently from where a real
|
||||
/// face's mesh points sit, so only the points both agree on are held to ground truth: eyes and
|
||||
/// chin tightly, the mouth corners (the cartoon mouth's ends, as it turns) more loosely.
|
||||
#[test]
|
||||
fn mediapipe_follows_the_synthetic_face() {
|
||||
let Ok(path) = std::env::var("EFFECTCRAFT_FACE_LANDMARKER") else { return };
|
||||
let Ok(effectcraft_segment::Loaded::Face(model)) = effectcraft_segment::load("mediapipe-face", &std::fs::read(path).unwrap()) else {
|
||||
panic!("not a face model")
|
||||
};
|
||||
let (w, h) = (320, 240);
|
||||
let p0 = frame_params(0);
|
||||
let img0 = render_synthetic(w, h, &p0, 3);
|
||||
let region = [p0.center[0] - 60.0, p0.center[1] - 75.0, p0.center[0] + 60.0, p0.center[1] + 75.0];
|
||||
let (mut tr, fit0) = FaceTracker::new_with(&Frame { img: &img0, offset: [0.0; 2] }, region, Some(model)).expect("face found");
|
||||
assert_eq!(tr.model().map(|m| m.id), Some("mediapipe-face"));
|
||||
let mut fits = vec![fit0];
|
||||
let n = 40;
|
||||
for t in 1..n {
|
||||
let img = render_synthetic(w, h, &frame_params(t), 3);
|
||||
fits.push(tr.step(&Frame { img: &img, offset: [0.0; 2] }).unwrap_or_else(|| panic!("lost at {t}")));
|
||||
}
|
||||
let held = [(EYE[0][1], 0.04), (EYE[0][4], 0.04), (EYE[1][1], 0.04), (EYE[1][4], 0.04), (CHIN, 0.04), (MOUTH[0], 0.08), (MOUTH[1], 0.08)];
|
||||
for (t, f) in fits.iter().enumerate() {
|
||||
let p = frame_params(t);
|
||||
let truth = synth_landmarks(&p);
|
||||
for &(i, bound) in &held {
|
||||
let e = (f.landmarks[i][0] - truth[i][0]).hypot(f.landmarks[i][1] - truth[i][1]) / p.height;
|
||||
assert!(e < bound, "frame {t} {}: {e:.3} face heights", LANDMARKS[i].0);
|
||||
}
|
||||
assert!((f.roll - p.roll).abs() < 3.0, "frame {t}: roll {} vs {}", f.roll, p.roll);
|
||||
}
|
||||
// Head Scale follows the face's size (relative to the first frame).
|
||||
let h0 = frame_params(0).height;
|
||||
let base = measure(&fits[0].landmarks, 1.0)[2];
|
||||
let worst_scale = fits.iter().enumerate().map(|(t, f)| (measure(&f.landmarks, 1.0)[2] / base - frame_params(t).height / h0).abs()).fold(0.0, f64::max);
|
||||
assert!(worst_scale < 0.05, "Head Scale off by {worst_scale:.3}");
|
||||
}
|
||||
}
|
||||
@@ -403,7 +403,9 @@ fn a_model_prior_shapes_the_cut_and_bad_models_are_ignored() {
|
||||
#[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 effectcraft_segment::Loaded::Mask(model) = effectcraft_segment::load("mobilesam", &std::fs::read(path).unwrap()).unwrap() else {
|
||||
panic!("not a mask model")
|
||||
};
|
||||
let m = Some(model.as_ref());
|
||||
let (img0, gt0) = frame_at(center(0));
|
||||
let st = base_strokes(0);
|
||||
|
||||
@@ -8,6 +8,7 @@
|
||||
//! `settings.cancel`, `settings.previous`, `settings.next`.
|
||||
|
||||
use effectcraft_engine::prefs::{Item, Kind, Page, pages};
|
||||
use effectcraft_engine::segment::Task;
|
||||
use egui::{Color32, RichText, vec2};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
@@ -28,8 +29,8 @@ enum Act {
|
||||
Run(String, Value),
|
||||
PickFolder(String),
|
||||
PickProject(String),
|
||||
/// Install from File (a downloaded weights file).
|
||||
PickModel,
|
||||
/// Install from File (a downloaded weights file) for a task's models.
|
||||
PickModel(Task),
|
||||
}
|
||||
|
||||
fn hex(c: [u8; 3]) -> String {
|
||||
@@ -63,7 +64,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::Models(task) => models_ui(app, ui, t, acts, task),
|
||||
Item::AudioDevices { key } => {
|
||||
let devices: Vec<String> = app.hooks.audio_devices.as_ref().map(|f| f()).unwrap_or_default();
|
||||
let sel = get(key).as_str().unwrap_or("").to_string();
|
||||
@@ -210,17 +211,21 @@ pub fn human_bytes(b: u64) -> String {
|
||||
if u == 0 { format!("{b} bytes") } else { format!("{v:.1} {}", units[u]) }
|
||||
}
|
||||
|
||||
/// 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<Act>) {
|
||||
let Ok(info) = app.session.execute("roto.models", json!({})) else { return };
|
||||
/// Settings ▸ Roto Brush and Settings ▸ Face Tracking: a task's models (`roto.models`,
|
||||
/// `face.models`): choose one; download, install from a file or remove the trained ones; who made
|
||||
/// them, their licence, size and source.
|
||||
fn models_ui(app: &mut EffectcraftApp, ui: &mut egui::Ui, t: &Tokens, acts: &mut Vec<Act>, task: Task) {
|
||||
let (prefix, user) = match task {
|
||||
Task::Mask => ("roto", "Roto Brush"),
|
||||
Task::Face => ("face", "Face tracking"),
|
||||
};
|
||||
let Ok(info) = app.session.execute(&format!("{prefix}.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 classic = id == effectcraft_engine::segment::CLASSICAL;
|
||||
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);
|
||||
@@ -228,28 +233,30 @@ fn roto_models_ui(app: &mut EffectcraftApp, ui: &mut egui::Ui, t: &Tokens, acts:
|
||||
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);
|
||||
reg(app, &format!("settings.{prefix}.model.{id}"), &r, name);
|
||||
if r.clicked() && !selected {
|
||||
acts.push(Act::Run("roto.model.select".into(), json!({"id": id})));
|
||||
acts.push(Act::Run(format!("{prefix}.model.select"), 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",
|
||||
(true, _, _) => "In use".to_string(),
|
||||
_ if classic => "Built in".to_string(),
|
||||
(false, true, true) if busy.is_some() => "Loading…".to_string(),
|
||||
(false, true, _) => "Installed".to_string(),
|
||||
(false, false, true) => format!("Not installed: {user} uses the classic engine"),
|
||||
_ => "Not installed".to_string(),
|
||||
};
|
||||
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));
|
||||
if let Some(by) = m["authors"].as_str() {
|
||||
ui.label(RichText::new(format!("By {by}")).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")] {
|
||||
for (label, key) in [("Project page", "homepage"), ("Licence", "licenceUrl")] {
|
||||
if let Some(url) = m[key].as_str() {
|
||||
let r = ui.link(label);
|
||||
if r.clicked() {
|
||||
@@ -258,29 +265,29 @@ fn roto_models_ui(app: &mut EffectcraftApp, ui: &mut egui::Ui, t: &Tokens, acts:
|
||||
}
|
||||
}
|
||||
});
|
||||
if id == effectcraft_engine::segment::CLASSICAL {
|
||||
if classic {
|
||||
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");
|
||||
reg(app, &format!("settings.{prefix}.download.{id}"), &r, "Download");
|
||||
if r.clicked() {
|
||||
acts.push(Act::Run("roto.model.download".into(), json!({"id": id})));
|
||||
acts.push(Act::Run(format!("{prefix}.model.download"), 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");
|
||||
reg(app, &format!("settings.{prefix}.install.{id}"), &r, "Install from File");
|
||||
if r.clicked() {
|
||||
acts.push(Act::PickModel);
|
||||
acts.push(Act::PickModel(task));
|
||||
}
|
||||
}
|
||||
if installed {
|
||||
let r = ui.add_enabled(busy.is_none(), egui::Button::new("Remove"));
|
||||
reg(app, &format!("settings.roto.remove.{id}"), &r, "Remove");
|
||||
reg(app, &format!("settings.{prefix}.remove.{id}"), &r, "Remove");
|
||||
if r.clicked() {
|
||||
acts.push(Act::Run("roto.model.remove".into(), json!({"id": id})));
|
||||
acts.push(Act::Run(format!("{prefix}.model.remove"), json!({"id": id})));
|
||||
}
|
||||
}
|
||||
});
|
||||
@@ -303,6 +310,8 @@ fn roto_models_ui(app: &mut EffectcraftApp, ui: &mut egui::Ui, t: &Tokens, acts:
|
||||
}
|
||||
}
|
||||
|
||||
/// 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.
|
||||
fn browser_storage_ui(app: &mut EffectcraftApp, ui: &mut egui::Ui, t: &Tokens, acts: &mut Vec<Act>) {
|
||||
let Some(host) = app.session.storage.clone() else { return };
|
||||
let info = host.info();
|
||||
@@ -513,10 +522,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();
|
||||
Act::PickModel(task) => {
|
||||
// The registry's file types for the task (`.pt`, `.task`).
|
||||
let mut exts: Vec<&str> = effectcraft_engine::segment::models(task).filter_map(|m| m.file_name.rsplit_once('.').map(|x| x.1)).collect();
|
||||
exts.dedup();
|
||||
let picked = app.hooks.pick_files.as_ref().map(|f| f(&exts)).unwrap_or_default();
|
||||
let cmd = if task == Task::Face { "face.model.install" } else { "roto.model.install" };
|
||||
if let Some(path) = picked.into_iter().next()
|
||||
&& let Err(e) = crate::menus::invoke(app, ctx, "roto.model.install", json!({"path": path}))
|
||||
&& let Err(e) = crate::menus::invoke(app, ctx, cmd, json!({"path": path}))
|
||||
{
|
||||
app.ui.status = e;
|
||||
}
|
||||
|
||||
+19
-4
@@ -207,6 +207,20 @@ orientation X/Y/Z, eye openness, eyebrow distance from eye, mouth openness, widt
|
||||
those keys on the keyframe clipboard and returns them per frame plus a tab-separated `clipboard`
|
||||
text.
|
||||
|
||||
Face tracking's trained model works like Roto Brush's (below): `face.models` lists the built-in
|
||||
tracker and MediaPipe Face Landmarker with its authors and licence; `face.model.download {id}` /
|
||||
`face.model.install {path, id?}` install it (verified); `face.model.select {id:
|
||||
classical|mediapipe-face}` chooses it; `face.model.remove {id}` deletes it. Downloads and loading
|
||||
run in the background; `"wait": true` on download, install or select returns once the model is in
|
||||
use (or with its error), which scripts need before tracking. `track.mask` with a
|
||||
face method returns `faceModel`, the engine it tries first (the classical one when the model finds
|
||||
no face in the mask).
|
||||
|
||||
```sh
|
||||
effectcraft-cli run clip.ecproj face.model.download '{"id":"mediapipe-face","wait":true}' face.model.select '{"id":"mediapipe-face","wait":true}' \
|
||||
track.mask '{"layer":"#1","mask":1,"method":"faceDetailed","direction":"forward","wait":true}' --save
|
||||
```
|
||||
|
||||
`mask.interpolate` (Window ▸ Mask Interpolation ▸ Apply) adds in-between Mask Path keys between
|
||||
each pair of selected Mask Path keys (or the existing keys at `times`, in seconds), giving both
|
||||
ends the same vertex count with a matched correspondence. Options: `keyframeRate` (number or
|
||||
@@ -296,17 +310,18 @@ reference. Matte settings are ordinary properties (`effects/#1/rotoBrushMatte/se
|
||||
`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.models` lists the built-in engine and every registered model, with its authors,
|
||||
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.select {id: classical|mobilesam}` chooses which one Roto Brush uses. Loading runs
|
||||
in the background; `"wait": true` (also on download and install) returns once it is in use.
|
||||
- `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"}' \
|
||||
effectcraft-cli run clip.ecproj roto.model.install '{"path":"mobile_sam.pt"}' roto.model.select '{"id":"mobilesam","wait":true}' \
|
||||
roto.stroke '{"layer":"#1","points":[[300,200],[360,230]],"radius":10}' roto.propagate '{"layer":"#1","wait":true}' --save
|
||||
```
|
||||
|
||||
|
||||
@@ -22,7 +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 | `segment` | Swappable trained models for Roto Brush and face tracking: the `MaskModel` and `FaceModel` interfaces, a registry of open-licensed models (authors, licence, URL, size, SHA-256), PyTorch checkpoint and TensorFlow Lite readers (with an interpreter), SHA-256, MobileSAM (TinyViT + Segment Anything decoder) and MediaPipe Face Landmarker (BlazeFace + Face Mesh V2) 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 |
|
||||
@@ -335,6 +335,10 @@ components inside it, and an active-shape point distribution model trained on sy
|
||||
shapes (our own generator, no external weights) that fills and checks the landmarks. The outline
|
||||
keys the Mask Path; Detailed Features keys a Face Track Points effect, and
|
||||
`track.extractFaceMeasurements` derives a keyed Face Measurements effect (and copies its keys).
|
||||
With a trained face model chosen (`effectcraft_segment::face::FaceModel`, Settings ▸ Face
|
||||
Tracking; `Session::models` hands it to the mask track), `FaceTracker::new_with` lets the model
|
||||
find the face in the mask and follow it; the model's outline and named points replace the
|
||||
classical steps, and chin and jaw still come from the outline so the measurements agree.
|
||||
|
||||
**Warp Stabilizer** (`effects::warp_stab`, `engine::warp`): `warp.analyze` renders the layer's
|
||||
input to the effect (`Renderer::layer_input`: source, masks and the effects above it) for every
|
||||
|
||||
+6
-2
@@ -37,7 +37,7 @@ of them have been reconciled.
|
||||
| Real-user experience | ≈ 70%, uneven by platform | Issues #41–#47 (all from the Linux AppImage, 5 October): panning the viewer snaps back, panels can't be resized or rearranged, drag-and-drop and double-click import don't work, layer rename gets stuck, the Layer Settings arrow does nothing, the Project panel clips, the Wayland window icon is generic. parity.md counts docking as done. Most checking happens on macOS. No localisation, no accessibility work |
|
||||
| Performance | unknown against After Effects | Internal numbers only (e.g. Advanced 3D 290 ms/frame at 1080p on the GPU, an M4 Pro under load). Nothing benchmarked against After Effects; no large real projects (4K footage, hundreds of layers) tested |
|
||||
| Media formats | ≈ 80% | H.264, ProRes, HEVC, AV1, image sequences and audio exist. The new HEVC / AV1 encoders have no B-frames, multi-reference or SAO / CDEF, so files are larger than from mature encoders. Camera formats (BRAW, R3D, ProRes RAW, variable-frame-rate phone video) are unverified |
|
||||
| AI-assisted tools | ≈ 50% | Face tracking uses a classical shape-model fitter. Roto Brush 2.0 / 3.0 can use MobileSAM since M13.35 (pure-Rust inference, an optional download): on our synthetic moving-disc test it scores IoU 0.989 on the base frame and ≥ 0.980 over 20 propagated frames (classic: 0.973), but nothing compares it with After Effects yet |
|
||||
| AI-assisted tools | ≈ 50% | Both tools can use trained models (pure-Rust inference, optional downloads), but nothing compares them with After Effects yet. Roto Brush 2.0 / 3.0 with MobileSAM (M13.35) scores IoU 0.989 on the base frame of our synthetic moving-disc test and ≥ 0.980 over 20 propagated frames (classic: 0.973). Face tracking with MediaPipe Face Landmarker (M13.36) matches Google's own pipeline to 0.85 px on average on a test portrait; on our synthetic clip the eyes and chin stay within 4% of the face height |
|
||||
| Maturity | early | First commit 1 October 2026. ≈ 285,000 lines of almost entirely agent-written Rust, ≈ 2,050 tests, 8 external issue reports so far. After Effects has around 30 years of edge cases behind it |
|
||||
|
||||
## Where we're going: workstreams in priority order
|
||||
@@ -110,7 +110,11 @@ The most important missing piece: it turns every other estimate here into a meas
|
||||
open-source licences, MobileSAM under Apache-2.0), weights downloaded on demand and verified,
|
||||
Settings ▸ Roto Brush to choose. Next: measure it against After Effects (G1); faster encoders
|
||||
(GPU) for long shots.
|
||||
- Face tracking: still an owner decision (a learned detector, or the classical fitter as final).
|
||||
- Face tracking: decided the same way on 5 October 2026, done in M13.36: MediaPipe Face
|
||||
Landmarker (Google, Apache-2.0: BlazeFace detector and the 478-point Face Mesh V2) on a pure-Rust
|
||||
TensorFlow Lite interpreter, behind a `FaceModel` interface in the same crate, a 3.8 MB optional
|
||||
download chosen in Settings ▸ Face Tracking. The classical fitter stays as the fallback. Next:
|
||||
measure both against After Effects (G1).
|
||||
|
||||
### G8. Plug-in ecosystem
|
||||
|
||||
|
||||
+21
-2
@@ -29,7 +29,7 @@ Partial features:
|
||||
|---|---|---|---|
|
||||
| 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 | Roto Brush 2.0 / 3.0 can use a trained model since M13.35 (MobileSAM, optional download); its quality has not been compared with After Effects (G1), so the score stands |
|
||||
| TRK-3 Face tracking | P2 | 0.9 | Classical fitter: weak on profile and occluded faces; Rolling Shutter Ripple is approximated |
|
||||
| TRK-3 Face tracking | P2 | 0.9 | Can use a trained model since M13.36 (MediaPipe Face Landmarker, optional download); the classical fitter is weak on profile and occluded faces; neither is compared with After Effects (G1), so the score stands; Rolling Shutter Ripple is approximated |
|
||||
|
||||
Since this audit: the project operations (M3.9–M3.14) and preview (M4.5–M4.12) passes fixed
|
||||
behaviour inside features already counted as done, so the counts above stand: unsaved-changes
|
||||
@@ -175,7 +175,7 @@ the disk cache in the formats wave; the 3D Camera Tracker in M12.6.)
|
||||
| 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 | (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) |
|
||||
| Tracking | ≈ 96% | 0.3 | face tracking's trained model (MediaPipe Face Landmarker, M13.36) is an optional download; without it the classical (skin model + feature components + shape model) fitter is weak on profile views and occluded faces; 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 |
|
||||
|
||||
## Effects still missing
|
||||
@@ -466,6 +466,25 @@ measures it against After Effects ([gaps.md](gaps.md)).
|
||||
- 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.36): face tracking can use a trained model too. TRK-3 keeps its score until G1 measures
|
||||
it against After Effects ([gaps.md](gaps.md)).
|
||||
- The model sits behind `effectcraft_segment::face::FaceModel` (find a face in a region, follow
|
||||
it frame to frame), in the same registry as Roto Brush's. A model reports its own points and
|
||||
says which are the tracker's landmarks and which trace the face outline, so models can be
|
||||
swapped without touching the tracker.
|
||||
- The first model is **MediaPipe Face Landmarker** (Google, Apache-2.0): the BlazeFace
|
||||
short-range detector and Face Mesh V2 (478 points). Its official `face_landmarker.task` bundle
|
||||
(3.8 MB, pinned SHA-256) runs on our own TensorFlow Lite interpreter in plain Rust, which matches
|
||||
TensorFlow Lite's to about one part in a million.
|
||||
- **Settings ▸ Face Tracking** chooses it. Outline Only keys the mesh's face oval into the mask;
|
||||
Detailed Features keys the mesh's eye, pupil, brow, nose and lip points; chin and jaw come from
|
||||
the outline as with the classical fitter, so Extract & Copy Face Measurements means the same with
|
||||
either. When the model finds no face in the mask, the classical fitter runs.
|
||||
- On a test portrait its points are 0.85 px from Google's own pipeline on average (2.2 px at
|
||||
worst). Finding a face takes about 45 ms; following it, about 20 ms a frame (32-thread CPU).
|
||||
- Agents use `face.models`, `face.model.select`, `face.model.download`, `face.model.install` and
|
||||
`face.model.remove`; `track.mask` reports the engine it tried first as `faceModel`.
|
||||
|
||||
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.
|
||||
|
||||
+5
-2
@@ -141,8 +141,11 @@ Commands (CLI, MCP, control channel):
|
||||
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
|
||||
with its authors, 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, each with a `.NOTICE.txt` naming its authors and licence
|
||||
- `face.model` (Settings ▸ Face Tracking): the model face tracking uses, `classical` (built in) or
|
||||
`mediapipe-face` (once installed); the same page of models (`face.models` / `face.model.*`)
|
||||
|
||||
### Display-only settings
|
||||
|
||||
|
||||
Reference in new issue
Block a user