Compare commits

...

15 Commits

Author SHA1 Message Date
bflatastic e43bf66775 M13.36: face tracking trained model — MediaPipe Face Landmarker on a pure-Rust TFLite interpreter
- effectcraft-segment: a FaceModel interface (find a face in a region, follow it; a Topology
  maps the model's points to the tracker's landmarks and outline), a TensorFlow Lite reader and
  interpreter (float ops, float16 weights folded; matches ai-edge-litert to ~1e-6), and
  MediaPipe Face Landmarker (Google, Apache-2.0: BlazeFace + Face Mesh V2, 3.8 MB, pinned
  SHA-256). The registry gains a task and authors per model; installed weights get a NOTICE.
- effectcraft-track: FaceTracker::new_with uses the model, falls back to the classical fitter
  when it finds no face; chin and jaw come from the outline either way so measurements agree.
- Engine: roto_models becomes models (both tasks); face.models / face.model.* commands,
  Settings ▸ Face Tracking, `wait` on select/install/download for scripts; track.mask reports
  faceModel.
- Credits: ATTRIBUTION.md lists the optional models and the Apache-2.0 specs the code follows.
- Docs: README, ROADMAP (G7), gaps, parity (TRK-3 score unchanged pending G1), agents,
  preferences, architecture, segment README.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-05 20:39:45 +07:00
BFlat 4624c00308
Merge pull request #58 from storytold/m13-roto-ml
M13.35: Roto Brush trained model — swappable segment module, MobileSAM in pure Rust
2026-10-05 17:57:36 +07:00
bflatastic 174ab2b6af Merge main (gap assessment) into M13.35: keep parity numbers, record G7's decision and evidence
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-05 17:57:09 +07:00
bflatastic 7579de21c6 M13.35: Roto Brush trained model — swappable segment module, MobileSAM in pure Rust
- New L1 crate effectcraft-segment: the MaskModel interface (prompts → foreground probability),
  a registry of open-licensed models (licence must be MIT/Apache/BSD/…; URL, size, SHA-256),
  a PyTorch checkpoint reader (zip + pickle subset), dependency-free SHA-256, and MobileSAM
  (TinyViT-5M + SAM prompt encoder / two-way transformer / mask decoder) on rayon and ndarray's
  safe GEMM. Loads the official mobile_sam.pt as published; checked against candle
  (embedding 0.3% RMS, 99.9% equal mask pixels); ~0.7 s/frame encoder at 32 threads, cached
  per frame, ~0.1 s decoder.
- track::roto: segment_with / propagate_with prompt a model with the strokes or the flow-warped
  matte (box, core points, mask); its probability is a strong, unblurred prior for the same
  graph cut (strokes stay hard constraints); a failing model or one disagreeing with the flow
  (IoU < 0.5) is ignored. Prior gets a weight.
- effects::roto: Version 2.0 / 3.0 use the model chosen in Settings (set_model / model_for);
  the model id is part of the chain seed.
- engine: roto_models (roto.models, roto.model.select / install / download / remove), weights in
  the models folder (host::config_dir()/models, shared by the app, CLI and MCP), background
  loading, downloads through the system curl, verified; Settings ▸ Roto Brush page (prefs
  roto.model). host::config_dir moved from the desktop app.
- Tests: segment (kernels vs direct sums, checkpoint reader, SHA-256 vectors, registry), track
  (fake model: prior shapes the cut, failing / disagreeing models ignored; real weights with
  EFFECTCRAFT_MOBILESAM: base IoU 0.989, propagated ≥ 0.980 vs classic 0.973), engine (registry,
  verification, selection; real install + load), MCP.
- Docs: README, ROADMAP, parity (MSK-4 done: 90/2/0), preferences, agents, architecture,
  crates/segment/README.md.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-05 17:55:02 +07:00
Brandon Thomas aa4656b543 Docs: honest gap assessment (docs/gaps.md), roadmap workstreams G1–G9, README status
parity.md's ≈99% measures breadth only. docs/gaps.md assesses real-world
readiness by dimension (fidelity, .aep interop, stability, platform
reliability, performance, media, AI tools, maturity) with evidence, and
sets prioritised workstreams with done criteria. ROADMAP, README, parity.md,
AGENTS.md and CLAUDE.md point agents at it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-05 06:32:19 -04:00
BFlat 11d4885764
Merge pull request #57 from storytold/m13-audio-scrub
M13.34: audio scrubbing (Ctrl/Cmd-drag the current time)
2026-10-05 15:19:12 +07:00
bflatastic 5afc25d0d8 M13.34: audio scrubbing (Ctrl/Cmd-drag the current time)
- audio::AudioScrub: an open output that plays a faded snippet (one frame, 30–100 ms) of the
  comp's mix at each new frame, replacing what is still queued; holding still repeats nothing.
  The mixing is shared with the preview feeder (mix_block); AudioFeed::clear.
- The Timeline ruler: Ctrl/Cmd-drag scrubs audio (EffectcraftApp::scrub_audio opens the output
  through the audio device hook on first use). The Audio panel meters follow the snippets; an
  idle second closes the output.
- playback.scrubAudio {time?} for agents.
- Tests: ui_audio_scrub (snippet length, fades, same-frame no-op, replace not pile up, silent
  comps, Cmd-drag in the ruler, idle close) with a fake device. Docs: parity.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-05 15:18:54 +07:00
BFlat 9607be471b
Merge pull request #56 from storytold/m13-nested-markers
M13.33: nested comp markers on the precomp layer bar
2026-10-05 15:12:22 +07:00
bflatastic f013bdd8d7 M13.33: nested comp markers on the precomp layer bar
- commands::markers::nested_markers: a precomp layer's comp markers at their times in the
  containing comp (start and stretch in closed form; with time remapping or a stretch over
  protected regions, the first frame whose source time reaches them), within In–Out.
- markers.nested {layer}: the same list for agents.
- Timeline: drawn on the precomp bar as outlined read-only markers (positions cached per
  revision), hover names them, double-click opens the nested comp at the marker; automation id
  timeline.layer.<id>.nestedMarker.<i>.
- Tests: engine (start, stretch, trim, time remap, non-precomp error), UI (draw, label, open).
  Docs: parity, control protocol.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-05 15:12:06 +07:00
BFlat 1ac55c33a3
Merge pull request #55 from storytold/m13-new-viewer
M13.32: View ▸ New Viewer — several Composition viewers
2026-10-05 15:00:30 +07:00
bflatastic 97390f15fb M13.32: View ▸ New Viewer — several Composition viewers
- PanelKind::Viewer(n): more Composition viewers (panels::viewers). One is active (the full
  interactive viewer, showing the active comp); the others show their own comp's frame at that
  comp's current time (CPU or GPU frames, their own textures) and become active on a click or
  their tab, their comp becoming the active comp.
- view.newViewer (View menu, the Composition panel menu, Alt+Shift+N): a new viewer of the
  active comp next to the one in use, which is locked. Opening a comp while the active viewer
  is locked shows it in an unlocked viewer, or a new one when every viewer is locked. Closing
  the active viewer hands over to another with its comp. Tabs are named after their comp, with
  their own lock and swatch; Preview full screen maximizes the active viewer.
- UI state: viewers (comp per viewer) and active_viewer (serde).
- Tests: ui_viewers (lock on New Viewer, routing, click to activate, all locked → new viewer,
  close hand-over, tab titles; an ignored snapshot). The bounded-drag frames test asserts what
  holds when freed addresses are reused. Docs: parity, control protocol.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-05 15:00:11 +07:00
BFlat 6a5abad753
Merge pull request #54 from storytold/m4-content-keyed-frames
M4.14: content-keyed RAM preview frames (M4 done); proxies in the disk cache key
2026-10-05 14:39:33 +07:00
bflatastic 26c714db7f M4.14: content-keyed RAM preview frames; proxies in the disk cache's content key
- FrameKey's identity is the comp's content (frames::comp_content), not the project revision:
  an edit keeps the cached frames of comps it doesn't touch and undo finds the frames of the
  state it returns to. The identity hashes the addresses of the comps a comp draws (~0.1 ms on
  the 1500-layer test project; the full content hash took ~117 ms), by value the footage, solids,
  proxies and project settings; with expressions, every item. Frames::content_of keeps a
  project snapshot per identity with frames (or in flight), so Arc::make_mut always copies an
  edited comp and its address can't be reused meanwhile; past 64 identities the least recently
  used go with their frames (long drags).
- drop_stale is gone (frames of other contents are evicted least-recently-shown first).
- disk_cache::comp_content_key covers items' proxies and Use Proxy (it served stale frames).
- Tests: frames (content keys, edits elsewhere and undo, bounded drags, in-place edit safety),
  ui_ram_preview (end to end), the disk cache proxy key. Docs: roadmap (M4 done), parity.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-05 14:39:17 +07:00
bflatastic 4f7723a7b6 M13.31: Home size column test uses an unambiguous size (12.5 KB rounds to even)
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-05 14:33:44 +07:00
BFlat 0f1cad8bbd
Merge pull request #53 from storytold/ci-packaging-lint
CI: fix Packaging lint (SC2015 in release.yml)
2026-10-05 13:52:21 +07:00
70 changed files with 5986 additions and 172 deletions

View File

@ -107,6 +107,7 @@ a command, everything is agent-drivable, and the quality gates (`cargo xtask ci`
## See also
- [docs/gaps.md](docs/gaps.md): where EffectCraft falls short of After Effects and the prioritised workstreams (G1–G9); read it before choosing work
- [CONTRIBUTING.md](CONTRIBUTING.md): setup, gates, commits, how to add things
- [docs/architecture.md](docs/architecture.md): layers, data model, commands, pipeline
- [docs/testing.md](docs/testing.md): oracle tests, criteria, benchmarks

View File

@ -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

View File

@ -6,7 +6,8 @@ EffectCraft is a clean-room, open-source, pure-Rust motion graphics and visual e
`plan/` is maintainer-local (gitignored). Public equivalents: [`ROADMAP.md`](ROADMAP.md), [`docs/`](docs/).
1. Read `plan/STATUS.md` (current milestone, next task, running agents, blockers).
2. Read the task in `plan/execution-plan.md` §3, the relevant section of `plan/architecture.md`, and the crate README/docs you touch. Feature ids: `plan/aftereffects/feature-catalog.md`; effects: `plan/aftereffects/effects-list.md`; UI reference: `plan/aftereffects/README.md`.
3. Follow the autonomous operation protocol (`plan/execution-plan.md` §7). Don't stop to ask.
3. Read [`docs/gaps.md`](docs/gaps.md): the honest assessment of where we fall short and the prioritised workstreams. Fidelity, reliability and stability (G1–G3) come before new features.
4. Follow the autonomous operation protocol (`plan/execution-plan.md` §7). Don't stop to ask.
## Non-negotiables
**Read [`AGENTS.md`](AGENTS.md) first; its rules override everything here.** No Adobe assets, every asset licensed + attributed (`cargo xtask assets`).

43
Cargo.lock generated
View File

@ -1205,6 +1205,7 @@ dependencies = [
"effectcraft-keyframe",
"effectcraft-project",
"effectcraft-raster",
"effectcraft-segment",
"effectcraft-track",
"log",
"rayon",
@ -1229,6 +1230,7 @@ dependencies = [
"effectcraft-psd",
"effectcraft-raster",
"effectcraft-render",
"effectcraft-segment",
"effectcraft-svg",
"effectcraft-text",
"effectcraft-time",
@ -1541,6 +1543,15 @@ dependencies = [
"serde_json",
]
[[package]]
name = "effectcraft-segment"
version = "0.2.0"
dependencies = [
"miniz_oxide 0.8.9",
"ndarray",
"rayon",
]
[[package]]
name = "effectcraft-svg"
version = "0.2.0"
@ -1578,6 +1589,7 @@ name = "effectcraft-track"
version = "0.2.0"
dependencies = [
"effectcraft-raster",
"effectcraft-segment",
"rayon",
"serde",
"serde_json",
@ -3050,6 +3062,16 @@ dependencies = [
"libc",
]
[[package]]
name = "matrixmultiply"
version = "0.3.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3f607c237553f086e7043417a51df26b2eb899d3caff94e6a67592ff992fedc7"
dependencies = [
"autocfg",
"rawpointer",
]
[[package]]
name = "memchr"
version = "2.8.3"
@ -3167,6 +3189,21 @@ dependencies = [
"thiserror 2.0.21",
]
[[package]]
name = "ndarray"
version = "0.17.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d"
dependencies = [
"matrixmultiply",
"num-complex",
"num-integer",
"num-traits",
"portable-atomic",
"portable-atomic-util",
"rawpointer",
]
[[package]]
name = "ndk"
version = "0.9.0"
@ -4169,6 +4206,12 @@ dependencies = [
"objc2-quartz-core 0.3.2",
]
[[package]]
name = "rawpointer"
version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
[[package]]
name = "rayon"
version = "1.12.0"

View File

@ -22,6 +22,7 @@ effectcraft-project = { path = "crates/project" }
effectcraft-text = { path = "crates/text" }
effectcraft-effects = { path = "crates/effects" }
effectcraft-track = { path = "crates/track" }
effectcraft-segment = { path = "crates/segment" }
effectcraft-model = { path = "crates/model" }
effectcraft-render = { path = "crates/render" }
effectcraft-gpu = { path = "crates/gpu" }

View File

@ -195,13 +195,25 @@ tests, layering, asset attribution and the WebAssembly build). See [CONTRIBUTING
## Where it stands
EffectCraft is young. The core of After Effects is in place and gets deeper every week: what is
done and what comes next is in the [ROADMAP](ROADMAP.md). Still to come: learned models for
Roto Brush and face tracking (both use classical methods today), GPU versions of the remaining
CPU-only effects, the Advanced 3D Extended Viewer, and shared-memory threads in the web build. The
web build ([docs/web.md](docs/web.md)) runs the full app in the browser, with renders, analyses
and viewer frames in Web Workers. Projects are saved as `.ecproj`, readable versioned JSON; After
Effects `.aep` files cannot be opened.
EffectCraft is young: its first commit was on 1 October 2026. Nearly every After Effects feature
exists, including all of its effects, but existing is not the same as behaving exactly like After
Effects, and we haven't measured that yet. Today it's a good place to try things and to tell us
what breaks. It isn't yet a replacement for After Effects on client work. In particular:
- **After Effects projects can't be opened.** Projects are saved as `.ecproj`, readable versioned
JSON, but EffectCraft can't open `.aep` / `.aepx` files, and After Effects plug-ins don't run.
- **Behaviour still differs from After Effects in places.** Nothing yet compares our renders with
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.
- **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.
[Bug reports](https://github.com/storytold/effectcraft/issues) are the most useful thing you can
send us right now.
## How it's made

View File

@ -10,7 +10,7 @@ once.
| M1 | Keyframes: temporal and spatial interpolation, Easy Ease, roving, velocity | Done |
| M2 | Compositing: 38 blend modes, track mattes, parenting, adjustment layers | Done |
| M3 | Project operations: settings dialogs, layer commands, `.ecproj`, undo, After Effects menu bar | Done (unsaved-changes prompts, Solid / Layer Settings, lock-aware layer commands, safe project files; View ▸ New Viewer still opens only the one locked viewer) |
| M4 | Preview: precomps, motion blur, cached playback | Mostly done (Preserve frame rate / resolution when nested, one motion blur gate, RAM preview fixes, Cache Frames When Idle, motion blur of collapsed precomps and of animated shape / text content; left: content-keyed RAM preview frames) |
| M4 | Preview: precomps, motion blur, cached playback | Done (Preserve frame rate / resolution when nested, one motion blur gate, RAM preview fixes, Cache Frames When Idle, motion blur of collapsed precomps and of animated shape / text content, content-keyed RAM preview frames) |
| M5 | Timeline depth: graph editor, keyframe clipboard and dialogs, time remapping, pick-whips, expression editor | Done |
| M6 | Shapes, masks and footage: shape operators, masks and the pen tool, video and image import | Done |
| M7 | 3D: 3D layers, cameras, lights, shadows, depth of field, 3D views and camera tools | In review |
@ -19,20 +19,45 @@ once.
| M10 | Export: render queue, H.264, ProRes, image sequences, GIF, audio | Done |
| M11 | Animation tools: audio playback, meters and waveforms, Lottie import and export, presets, Motion Sketch, Wiggler, Smoother | Done |
| M12 | Performance: layer cache, parallel and GPU compositing, disk cache, motion tracking (done); GPU versions of the remaining CPU-only effects | Mostly done |
| M13 | Puppet tools, paint, Roto Brush, motion tracking, a plug-in API, Timeline depth | Done (Roto Brush and face tracking use classical models; learned models are still to come) |
| M13 | Puppet tools, paint, Roto Brush, motion tracking, a plug-in API, Timeline depth | Done (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 |
## How far from full parity
Measured feature by feature in [docs/parity.md](docs/parity.md) (4 October 2026):
There are two honest answers, and they are far apart. The full assessment, with the evidence, is
in [docs/gaps.md](docs/gaps.md).
- **≈ 99% of After Effects' features, weighted by importance**, counting partial features as half
done (≈ 99.8% with per-feature fractions): 89 of 92 features done, 3 partial, none missing; every
essential (P0) feature done; all 306 effects implemented in full, 280 of them on the GPU.
- **≈ 8.5 agent-hours of work remain**: about **2–3 hours of wall-clock time** with five Claude
Opus 5.5 agents in parallel, ≈ 1 hour without learned models for Roto Brush and face tracking.
- What is left: GPU kernels for the remaining visual effects, and learned-model quality for Roto
Brush and face tracking; in previews, content-keyed RAM preview frames ([parity.md](docs/parity.md#highest-value-gaps-in-order)).
- **Breadth: ≈ 99%.** Almost every After Effects feature exists: 89 of 92 catalogued features
done and 3 partial, all 306 effects implemented, 280 of them on the GPU
([docs/parity.md](docs/parity.md)). This is our own checklist, graded by the agents that built
the features.
- **Real use: ≈ 30–50% (estimated, not yet measured).** Whether someone who uses After Effects for
a living can do client work in EffectCraft. What holds it back:
- **Fidelity is unmeasured.** No test compares our output with After Effects itself, and users
are still finding behaviour bugs in features marked done.
- **After Effects projects can't be opened.** EffectCraft can't read `.aep` / `.aepx` files,
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 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
In priority order (details and "done" criteria in [docs/gaps.md](docs/gaps.md)):
| | Workstream | Why |
|---|---|---|
| G1 | Measure fidelity against After Effects: a corpus of test projects rendered and sampled in both apps, scored per feature | Turns every estimate into a measurement; a feature that behaves differently is not done |
| G2 | Real-user reliability on macOS, Windows and Linux; fix every open user issue | Basic workflows must work everywhere before more features matter |
| G3 | Stability: a green gate on the current toolchain, fuzzing of every input | A crash loses people's work |
| 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 (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 |
Come tell us what matters most to you on [Discord](https://discord.gg/artcraft).

View File

@ -72,7 +72,7 @@ fn main() -> eframe::Result {
session.autosave.background = true;
// Settings, shortcut presets and the crash-recovery sentinel live in the platform
// config directory. Agent-driven runs (`--control`) skip crash recovery.
if let Some(dir) = config_dir() {
if let Some(dir) = effectcraft_host::config_dir() {
// WebAssembly effect plug-ins in <config>/Plug-ins load before the menus are built.
let plugins = dir.join("Plug-ins");
if plugins.is_dir() {
@ -163,24 +163,6 @@ impl eframe::App for Desktop {
}
}
/// The platform config directory for EffectCraft (`EFFECTCRAFT_CONFIG_DIR` overrides):
/// `~/Library/Application Support/EffectCraft` (macOS), `%APPDATA%\EffectCraft` (Windows),
/// `$XDG_CONFIG_HOME/effectcraft` or `~/.config/effectcraft` (Linux and others).
fn config_dir() -> Option<std::path::PathBuf> {
use std::path::PathBuf;
if let Some(d) = std::env::var_os("EFFECTCRAFT_CONFIG_DIR") {
return Some(PathBuf::from(d));
}
let home = || std::env::var_os("HOME").map(PathBuf::from);
if cfg!(target_os = "macos") {
return home().map(|h| h.join("Library/Application Support/EffectCraft"));
}
if cfg!(target_os = "windows") {
return std::env::var_os("APPDATA").map(|a| PathBuf::from(a).join("EffectCraft"));
}
std::env::var_os("XDG_CONFIG_HOME").map(PathBuf::from).or_else(|| home().map(|h| h.join(".config"))).map(|c| c.join("effectcraft"))
}
/// Driven by an agent (`--control`): don't activate the app on launch, so the user's keyboard
/// focus stays where it is (macOS; winit activates ignoring other apps by default).
fn agent_event_loop(agent: bool) -> Option<eframe::EventLoopBuilderHook> {

View File

@ -393,3 +393,19 @@ fn high_level_tools_add_effect_and_report_state() {
let (c, err) = call(&mut s, "add_effect", json!({"layer": l, "effect": "No Such Effect"}));
assert!(err && c[0]["text"].as_str().unwrap().contains("unknown effect"), "{c:?}");
}
/// Trained models over MCP (Roto Brush and face tracking): listed with their authors and
/// licence, chosen, refused when unknown.
#[test]
fn trained_models_over_mcp() {
let mut s = server();
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:?}");
}
}

View File

@ -15,6 +15,7 @@ effectcraft-color = { workspace = true }
effectcraft-keyframe = { workspace = true }
effectcraft-project = { workspace = true }
effectcraft-raster = { workspace = true }
effectcraft-segment = { workspace = true }
effectcraft-track = { workspace = true }
rayon = { workspace = true }
serde = { workspace = true }

View File

@ -16,15 +16,19 @@
//! decontamination of the edge colours. Outside the segmentation span the layer is transparent;
//! without strokes the effect passes its input through.
//!
//! *Version* is stored but all versions use the same (classic, graph-cut) engine; *Quality ▸
//! Best* segments at a higher working resolution with more iterations.
//! *Version*: 1.0 is the classic graph-cut engine; 2.0 and 3.0 also use the trained model chosen
//! in Settings ▸ Roto Brush ([`set_model`], an `effectcraft_segment::MaskModel` such as
//! MobileSAM), falling back to the classic engine when none is installed. The model's id is part
//! of the chain seed, so choosing another recomputes the mattes. *Quality ▸ Best* segments at a
//! higher working resolution with more iterations.
use std::collections::{BTreeMap, HashMap, VecDeque};
use std::sync::{Arc, Mutex, OnceLock};
use std::sync::{Arc, Mutex, OnceLock, RwLock};
use effectcraft_keyframe::Value;
use effectcraft_project::ParamUi;
use effectcraft_raster::Image;
use effectcraft_segment::MaskModel;
use effectcraft_track::roto::refine::{self, MatteParams};
use effectcraft_track::roto::{self as rb, FrameSeg, RotoData, SegOpts, rle};
use rayon::prelude::*;
@ -173,8 +177,29 @@ pub fn matte_params(params: &Params) -> MatteParams {
}
}
fn model_slot() -> &'static RwLock<Option<Arc<dyn MaskModel>>> {
static M: OnceLock<RwLock<Option<Arc<dyn MaskModel>>>> = OnceLock::new();
M.get_or_init(|| RwLock::new(None))
}
/// The trained model Roto Brush 2.0 / 3.0 use (Settings ▸ Roto Brush; `None`: the classic
/// engine for every version).
pub fn set_model(m: Option<Arc<dyn MaskModel>>) {
if let Ok(mut slot) = model_slot().write() {
*slot = m;
}
}
/// The model an instance segments with: its version's choice of the installed one.
pub fn model_for(params: &Params) -> Option<Arc<dyn MaskModel>> {
if e(params, "version") == 0 {
return None;
}
model_slot().read().ok().and_then(|m| m.clone())
}
/// Seed of the chain keys: what the frames are (the engine-maintained Input Key), the layer
/// size, the frame rate and the segmentation settings.
/// size, the frame rate, the segmentation settings and the model.
pub fn seed(params: &Params, layer_size: [f64; 2], fps: f64) -> u64 {
let mut h = 0xcbf2_9ce4_8422_2325;
fnv(&mut h, s(params, INPUT_KEY).as_bytes());
@ -189,6 +214,7 @@ pub fn seed(params: &Params, layer_size: [f64; 2], fps: f64) -> u64 {
fnv(&mut h, &v.to_bits().to_le_bytes());
}
fnv(&mut h, &[e(params, "quality") as u8, e(params, "version") as u8]);
fnv(&mut h, model_for(params).map_or(effectcraft_segment::CLASSICAL, |m| m.info().id).as_bytes());
h
}
@ -275,13 +301,15 @@ pub struct Chain {
pub keys: BTreeMap<i64, u64>,
pub opts: SegOpts,
pub refine_radius: f64,
/// The trained model (Version 2.0 / 3.0 with one installed).
pub model: Option<Arc<dyn MaskModel>>,
}
impl Chain {
pub fn new(params: &Params, layer_size: [f64; 2], fps: f64, scale: f64) -> Chain {
let data = data(params);
let keys = data.chain_keys(seed(params, layer_size, fps));
Chain { refine_radius: data.refine_radius(), data, keys, opts: seg_opts(params, scale) }
Chain { refine_radius: data.refine_radius(), data, keys, opts: seg_opts(params, scale), model: model_for(params) }
}
/// Frame `f`'s segmentation at this chain's scale: cached, or computed from the nearest
@ -331,7 +359,7 @@ impl Chain {
let img = frame(g)?;
let strokes = self.data.strokes_at(g);
let seg = match (&cur, g == base) {
(_, true) | (None, _) => rb::segment(&img, &strokes, None, &self.opts, self.refine_radius),
(_, true) | (None, _) => rb::segment_with(&img, &strokes, None, &self.opts, self.refine_radius, self.model.as_deref()),
(Some(p), false) => {
let src = self.data.source_of(g)?;
let pimg = match prev_img.take() {
@ -341,7 +369,7 @@ impl Chain {
if pimg.width != img.width || pimg.height != img.height {
return None;
}
rb::propagate(&pimg, p, &img, &strokes, &self.opts, self.refine_radius)
rb::propagate_with(&pimg, p, &img, &strokes, &self.opts, self.refine_radius, self.model.as_deref())
}
};
let seg = Arc::new(seg);

View File

@ -20,6 +20,7 @@ effectcraft-path = { workspace = true }
effectcraft-text = { workspace = true }
effectcraft-effects = { workspace = true }
effectcraft-render = { workspace = true }
effectcraft-segment = { workspace = true }
effectcraft-track = { workspace = true }
effectcraft-model = { workspace = true }
effectcraft-lottie = { workspace = true }

View File

@ -46,6 +46,14 @@ pub fn specs() -> Vec<CommandSpec> {
),
fe!("playback.cacheWhenIdle", "Cache Frames When Idle", ["Composition", "Preview"], None, "{value?}", always),
fe!("playback.audio", "Audio", ["Composition", "Preview"], None, "{value?: include audio in previews}", always),
fe!(
"playback.scrubAudio",
"Scrub Audio",
[],
None,
"{time? (s, default the current time)} — plays one frame of the comp's audio there (Ctrl/Cmd-drag the current time)",
has_comp
),
// Viewer.
fe!("view.zoomIn", "Zoom In", ["View"], Some("."), "{}", has_comp),
fe!("view.zoomOut", "Zoom Out", ["View"], Some(","), "{}", has_comp),
@ -64,6 +72,14 @@ pub fn specs() -> Vec<CommandSpec> {
fe!("view.options", "View Options...", ["View"], Some("Cmd+Alt+U"), "{}", has_comp),
fe!("view.layerControls", "Show Layer Controls", ["View"], Some("Cmd+Shift+H"), "{value?}", always),
fe!("view.fullScreen", "Enter Full Screen", ["View"], Some("Ctrl+Cmd+F"), "{}", always),
fe!(
"view.newViewer",
"New Viewer",
["View"],
Some("Alt+Shift+N"),
"{} — another Composition viewer of the active comp; the one in use is locked",
has_comp
),
// Panels and workspaces.
fe!(
"window.panel",

View File

@ -5,7 +5,7 @@
//! comp seconds; layer markers are stored in layer time. All commands are undoable.
use effectcraft_color::Label;
use effectcraft_project::{Comp, CuePoint, ItemId, LayerId, LayerSource, Marker};
use effectcraft_project::{Comp, CuePoint, ItemId, Layer, LayerId, LayerSource, Marker, Project};
use effectcraft_time::Tick;
use serde_json::{Value, json};
@ -85,6 +85,56 @@ fn list(s: &mut Session, p: &Value) -> Result<Value> {
Ok(Value::Array(ms.iter().enumerate().map(|(i, m)| marker_json(c, o, i, m)).collect()))
}
/// The nested comp's markers on a precomp layer, at their times in the containing comp `comp`
/// and within the layer's In–Out range, as After Effects shows them on the precomp layer's bar:
/// `(comp time, index in the nested comp, marker)`. The layer's start time and stretch map them;
/// with time remapping, or a stretch over protected regions (Responsive Design), a marker sits at
/// the first frame the layer shows its time.
pub fn nested_markers<'a>(project: &'a Project, comp_id: ItemId, comp: &Comp, layer: &Layer) -> Vec<(Tick, usize, &'a Marker)> {
let LayerSource::Comp { item } = layer.source else { return vec![] };
let Some(nc) = project.comp(item) else { return vec![] };
if nc.markers.is_empty() || item == comp_id {
return vec![];
}
let shown = |t: Tick| t >= layer.in_point && t < layer.out_point;
let stretched_regions = (layer.stretch - 100.0).abs() > 1e-9 && nc.markers.iter().any(|m| m.protected);
if layer.props.get("timeRemap").is_none() && !stretched_regions {
return nc.markers.iter().enumerate().map(|(i, m)| (layer.comp_time(m.time), i, m)).filter(|(t, _, _)| shown(*t)).collect();
}
// The layer's source time at every frame of its range; each marker at the first frame that
// reaches it (within rounding: a thousandth of a nested frame).
let fr = comp.frame_rate;
let (f0, f1) = (fr.frame_at(layer.in_point), fr.frame_at(layer.out_point));
let frames: Vec<(Tick, Tick)> = (f0..f1.min(f0.saturating_add(100_000)))
.map(|f| {
let t = fr.tick_of(f);
(t, effectcraft_render::EvalCtx::new(project, comp_id, comp, t).source_time(layer))
})
.collect();
let eps = nc.frame_duration().0 / 1000;
nc.markers
.iter()
.enumerate()
.filter_map(|(i, m)| frames.iter().find(|(_, src)| src.0.saturating_add(eps) >= m.time.0).map(|(t, _)| (*t, i, m)))
.filter(|(t, _, _)| shown(*t))
.collect()
}
/// `markers.nested`: the nested comp's markers on a precomp layer (see [`nested_markers`]).
fn nested(s: &mut Session, p: &Value) -> Result<Value> {
let cmd = "markers.nested";
let cid = comp_id(s, p)?;
let c = s.project.comp(cid).ok_or(EngineError::NoComp)?;
let lid = p.get("layer").and_then(|v| resolve_layer(c, v)).ok_or_else(|| bad(cmd, "missing `layer` (a precomp layer)"))?;
let l = c.layer(lid).ok_or_else(|| bad(cmd, "no such layer"))?;
let LayerSource::Comp { item } = l.source else { return Err(bad(cmd, "not a precomp layer")) };
let v: Vec<Value> = nested_markers(&s.project, cid, c, l)
.into_iter()
.map(|(t, i, m)| json!({"time": t.seconds(), "comp": item.0, "index": i, "nestedTime": m.time.seconds(), "duration": m.duration.seconds(), "comment": m.comment, "label": m.label.name()}))
.collect();
Ok(json!(v))
}
fn locked(c: &Comp, o: Owner) -> bool {
matches!(o, Owner::Layer(l) if c.layer(l).is_some_and(|l| l.markers_locked))
}
@ -299,6 +349,12 @@ fn has_precomp_layer(s: &Session) -> std::result::Result<(), String> {
pub fn specs() -> Vec<CommandSpec> {
vec![
query!("markers.list", "List Markers", "{layer? (omit for composition markers)}", list),
query!(
"markers.nested",
"Nested Comp Markers",
"{layer: a precomp layer, comp?} — its comp's markers at their times in this comp (as on the layer bar)",
nested
),
cmd!(
"markers.set",
"Marker Settings",

View File

@ -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()),

View File

@ -34,7 +34,7 @@ mod layer_time;
pub(crate) mod link;
mod liquify;
mod lottie;
mod markers;
pub mod markers;
mod mask;
pub mod mask_interp;
pub(crate) mod model3d;
@ -133,6 +133,7 @@ pub fn command_specs() -> &'static [CommandSpec] {
SPECS.get_or_init(|| {
let mut v = Vec::new();
v.extend(file::specs());
v.extend(crate::models::specs());
v.extend(edit::specs());
v.extend(comp::specs());
v.extend(layer::specs());

View File

@ -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;
@ -63,6 +64,7 @@ pub use effectcraft_keyframe as keyframe;
pub use effectcraft_project as project;
pub use effectcraft_raster as raster;
pub use effectcraft_render as render;
pub use effectcraft_segment as segment;
pub use effectcraft_text as text;
pub use effectcraft_time as time;
pub use effectcraft_track as track;
@ -347,6 +349,10 @@ pub struct Session {
pub camera_pending: Vec<(ItemId, LayerId, Uid)>,
/// The running (or finished, not yet polled) Roto Brush propagation / Freeze.
pub roto_job: Option<roto::RotoJob>,
/// 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).
pub roto_pending: Vec<(ItemId, LayerId, Uid)>,
pub events: Vec<Event>,
@ -463,6 +469,8 @@ impl Default for Session {
camera_job: None,
camera_pending: vec![],
roto_job: None,
models: Default::default(),
models_dir: None,
roto_pending: vec![],
events: vec![],
journal: vec![],
@ -1079,6 +1087,8 @@ 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;

View File

@ -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;

View File

@ -1046,6 +1046,7 @@ Animation
Reveal Properties with Animation | anim.reveal {"kind":"animation"}
Reveal All Modified Properties | anim.reveal {"kind":"modified"}
View
New Viewer | view.newViewer
Split with New Locked Viewer | view.splitLockedViewer
---
Zoom In | view.zoomIn

463
crates/engine/src/models.rs Normal file
View File

@ -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(&notice_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)),
]
}

View File

@ -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)
})

View File

@ -209,6 +209,24 @@ page!(ThreeD {
realtime_shadows: bool = true,
});
page!(
/// Settings ▸ Roto Brush.
RotoPrefs {
/// The segmentation model Roto Brush 2.0 / 3.0 use: `classical` (built in) or a model id
/// from `effectcraft_segment::MODELS` (once installed).
model: String = "classical".into(),
}
);
page!(
/// 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,
@ -272,6 +290,8 @@ pub struct Prefs {
#[serde(rename = "threeD")]
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.
@ -312,6 +332,8 @@ impl Default for Prefs {
video: Video::default(),
three_d: ThreeD::default(),
scripting: Scripting::default(),
roto: RotoPrefs::default(),
face: FacePrefs::default(),
recent_projects: vec![],
recent_footage: vec![],
recent_presets: vec![],
@ -709,6 +731,8 @@ pub fn section_key(page: &str) -> Option<&'static str> {
"video" => "video",
"3d" | "threeD" => "threeD",
"scripting" => "scripting",
"roto" => "roto",
"face" => "face",
_ => return None,
})
}
@ -735,6 +759,8 @@ pub fn page_id(name: &str) -> Option<&'static str> {
"video" | "videopreview" => "video",
"3d" | "threed" => "3d",
"scripting" | "scriptingexpressions" => "scripting",
"roto" | "rotobrush" | "models" => "roto",
"face" | "facetracking" => "face",
_ => return None,
};
Some(id)
@ -822,6 +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,
/// A task's trained models (`roto.models`, `face.models`): authors, licence, size, install,
/// choose.
Models(effectcraft_segment::Task),
}
/// One page of the Settings dialog.
@ -1096,6 +1125,28 @@ pub fn pages() -> Vec<Page> {
Note("Expressions use the JavaScript engine. Script errors report their file and line in the Script Console; there is no step debugger."),
],
},
Page {
id: "roto",
title: "Roto Brush",
items: vec![
Section("Segmentation Model"),
Note(
"Roto Brush 2.0 and 3.0 use the model chosen here (Version 1.0 always uses the classic engine). Models are open source, downloaded only when you ask, and checked against their published SHA-256.",
),
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),
],
},
]
}
@ -1148,6 +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::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})

View File

@ -636,6 +636,7 @@ impl Session {
/// Frontends call this every frame: finishes jobs and (with `auto`) propagates edited
/// instances in the background.
pub fn poll_roto(&mut self, auto: bool) -> bool {
self.poll_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() {

View File

@ -38,6 +38,8 @@ impl Session {
let levels = self.prefs.general.undo_levels.max(1) as usize;
let current = self.project.clone();
self.history.trim(levels, &current);
// Settings ▸ Roto Brush / Face Tracking: load (or drop) the chosen trained models.
self.apply_models();
self.prefs_revision += 1;
}

View File

@ -119,4 +119,13 @@ fn footage_changes_change_the_salt() {
// Unrelated comps do not change it.
s.execute("comp.new", json!({"name": "Other", "width": 10, "height": 10, "frameRate": 30, "duration": 1})).unwrap();
assert_eq!(effectcraft_render::disk_cache::comp_content_key(&s.project, cid), k2);
// A proxy, and its Use Proxy switch, change the pixels.
let set_proxy = |s: &mut Session, enabled: bool| {
let it = std::sync::Arc::make_mut(&mut s.project).items.get_mut(&cid).unwrap();
it.proxy = Some(Box::new(effectcraft_project::Proxy { footage: effectcraft_project::Footage::default(), enabled }));
effectcraft_render::disk_cache::comp_content_key(&s.project, cid)
};
let k3 = set_proxy(&mut s, true);
assert_ne!(k3, k2);
assert_ne!(set_proxy(&mut s, false), k3);
}

View File

@ -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);
}

View File

@ -147,3 +147,37 @@ fn precompose_leave_and_move_attributes() {
assert_eq!((nl.in_point.seconds(), nl.out_point.seconds()), (1.0, 3.0));
assert_eq!(s.active_comp_id(), Some(inner));
}
/// A precomp layer shows its comp's markers on its bar, mapped through its timing.
#[test]
fn nested_comp_markers_follow_the_precomp_layers_timing() {
let mut s = crate::Session::default();
let pre = s.execute("comp.new", json!({"name": "Pre", "width": 32, "height": 32, "frameRate": 10, "duration": 4})).unwrap()["comp"].as_u64().unwrap();
for (t, c) in [(1.0, "beat"), (3.0, "drop")] {
s.execute("markers.set", json!({"new": true, "time": t, "comment": c})).unwrap();
}
s.execute("comp.new", json!({"name": "Main", "width": 32, "height": 32, "frameRate": 10, "duration": 20})).unwrap();
let l = s.execute("layer.addItem", json!({"item": pre})).unwrap()["layer"].as_u64().unwrap();
s.execute("layer.timing", json!({"layers": [l], "start": 2.0})).unwrap();
let times = |s: &mut crate::Session| -> Vec<(f64, String)> {
let v = s.execute_checked("markers.nested", json!({"layer": l})).unwrap();
v.as_array().unwrap().iter().map(|m| (m["time"].as_f64().unwrap(), m["comment"].as_str().unwrap().to_string())).collect()
};
assert_eq!(times(&mut s), vec![(3.0, "beat".to_string()), (5.0, "drop".to_string())]);
// Stretched to 200 %: twice as far from the start.
s.execute("layer.timeStretch", json!({"layers": [l], "percent": 200})).unwrap();
assert_eq!(times(&mut s), vec![(4.0, "beat".to_string()), (8.0, "drop".to_string())]);
s.execute("layer.timeStretch", json!({"layers": [l], "percent": 100})).unwrap();
// Trimmed: markers outside the In–Out range are not on the bar.
s.execute("layer.timing", json!({"layers": [l], "out": 4.0})).unwrap();
assert_eq!(times(&mut s), vec![(3.0, "beat".to_string())]);
s.execute("layer.timing", json!({"layers": [l], "out": 6.0})).unwrap();
// Time remapping at half speed: nested 1 s is reached 2 s into the layer.
s.execute("layer.enableTimeRemap", json!({"layers": [l]})).unwrap();
s.execute("prop.addKey", json!({"layer": l, "path": "timeRemap", "time": 4.0, "value": 2.0})).unwrap();
let t = times(&mut s);
assert_eq!(t.first(), Some(&(4.0, "beat".to_string())), "{t:?}");
// Not a precomp layer.
let solid = s.execute("layer.newSolid", json!({"color": "#ffffff"})).unwrap()["layer"].as_u64().unwrap();
assert!(s.execute("markers.nested", json!({"layer": solid})).is_err());
}

View File

@ -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(&params).map(|m| m.info().id), Some("mobilesam"));
s.execute_checked("roto.model.select", json!({"id": "classical"})).unwrap();
assert!(crate::effects::roto::model_for(&params).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);
}

View File

@ -81,6 +81,8 @@ impl Exporter for FileExporter {
/// A new session with media, import, expressions, scripting and export enabled.
pub fn session() -> Session {
Session {
// The desktop app, the CLI and the MCP server share installed Roto Brush models.
models_dir: config_dir().map(|d| d.join("models")),
exporter: Some(Arc::new(FileExporter::default())),
footage: Arc::new(effectcraft_media::MediaPool::new()),
importer: Some(Arc::new(MediaImporter)),
@ -93,6 +95,24 @@ pub fn session() -> Session {
}
}
/// The platform config directory for EffectCraft (`EFFECTCRAFT_CONFIG_DIR` overrides):
/// `~/Library/Application Support/EffectCraft` (macOS), `%APPDATA%\EffectCraft` (Windows),
/// `$XDG_CONFIG_HOME/effectcraft` or `~/.config/effectcraft` (Linux and others).
pub fn config_dir() -> Option<std::path::PathBuf> {
use std::path::PathBuf;
if let Some(d) = std::env::var_os("EFFECTCRAFT_CONFIG_DIR") {
return Some(PathBuf::from(d));
}
let home = || std::env::var_os("HOME").map(PathBuf::from);
if cfg!(target_os = "macos") {
return home().map(|h| h.join("Library/Application Support/EffectCraft"));
}
if cfg!(target_os = "windows") {
return std::env::var_os("APPDATA").map(|a| PathBuf::from(a).join("EffectCraft"));
}
std::env::var_os("XDG_CONFIG_HOME").map(PathBuf::from).or_else(|| home().map(|h| h.join(".config"))).map(|c| c.join("effectcraft"))
}
#[cfg(test)]
mod tests {
use serde_json::json;

View File

@ -778,8 +778,9 @@ pub fn footage_salt(project: &effectcraft_project::Project) -> u64 {
}
/// Content key of a composition's pixels: project render settings, the comp and everything it
/// uses (nested comps, footage, solids, footage files' size and time). With expressions in play
/// (which can read any comp) the whole project counts.
/// uses (nested comps, footage, solids, their proxies and Use Proxy switches, footage and proxy
/// files' size and time). With expressions in play (which can read any comp) the whole project
/// counts.
pub fn comp_content_key(project: &effectcraft_project::Project, comp: effectcraft_project::ItemId) -> u128 {
use effectcraft_project::{ItemKind, LayerSource};
let mut h = Hash128::default();
@ -813,6 +814,10 @@ pub fn comp_content_key(project: &effectcraft_project::Project, comp: effectcraf
if let ItemKind::Footage(f) = &it.kind {
footage_stamp(&mut h, f);
}
hash_debug(&mut h, &it.proxy);
if let Some(px) = &it.proxy {
footage_stamp(&mut h, &px.footage);
}
}
h.write_u64(comp.0);
h.finish()

15
crates/segment/Cargo.toml Normal file
View File

@ -0,0 +1,15 @@
[package]
name = "effectcraft-segment"
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
rust-version.workspace = true
[lints]
workspace = true
[dependencies]
miniz_oxide = "0.8"
ndarray = { version = "0.17", default-features = false, features = ["std"] }
rayon = { workspace = true }

64
crates/segment/README.md Normal file
View File

@ -0,0 +1,64 @@
# effectcraft-segment
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 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.*` / `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.
- **`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 | 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 |
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. 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.

115
crates/segment/src/face.rs Normal file
View File

@ -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));
}
}

272
crates/segment/src/lib.rs Normal file
View File

@ -0,0 +1,272 @@
//! 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 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;
pub type Result<T> = std::result::Result<T, String>;
/// What a model is asked: coordinates in frame pixels.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct Prompt {
/// Points and whether each is foreground (`true`) or background.
pub points: Vec<([f32; 2], bool)>,
/// A box around the object: `[x0, y0, x1, y1]`.
pub bbox: Option<[f32; 4]>,
/// A prior: the foreground probability of every pixel (`w×h`), e.g. the previous frame's
/// matte carried forward.
pub mask: Option<Vec<f32>>,
}
/// A trained segmentation model (swappable: Roto Brush only uses this interface).
pub trait MaskModel: Send + Sync {
fn info(&self) -> &'static ModelInfo;
/// Foreground probability (0–1) of each pixel of a frame (straight RGB 0–1, row-major `w×h`).
fn segment(&self, rgb: &[[f32; 3]], w: usize, h: usize, prompt: &Prompt) -> Result<Vec<f32>>;
}
/// 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,
/// The project's page.
pub homepage: &'static str,
/// The official weights file.
pub url: &'static str,
pub file_name: &'static str,
pub size: u64,
/// SHA-256 of the weights file (lower-case hex).
pub sha256: &'static str,
}
/// Licences a model may have to be listed (open source, compatible with MIT OR Apache-2.0).
pub const ALLOWED_LICENCES: &[&str] = &["MIT", "Apache-2.0", "BSD-2-Clause", "BSD-3-Clause", "ISC", "Zlib", "CC0-1.0"];
pub const MOBILE_SAM: ModelInfo = ModelInfo {
id: "mobilesam",
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",
url: "https://raw.githubusercontent.com/ChaoningZhang/MobileSAM/master/weights/mobile_sam.pt",
file_name: "mobile_sam.pt",
size: 40_728_226,
sha256: "6dbb90523a35330fedd7f1d3dfc66f995213d81b29a5ca8108dbcdd4e37d6c2f",
};
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 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<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));
}
let got = sha256::hex(bytes);
if got != info.sha256 {
return Err(format!("{}: the file does not match (SHA-256 {got}, expected {})", info.name, info.sha256));
}
match id {
"mobilesam" => {
let w = pt::read(bytes)?;
Ok(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,
}
impl MaskModel for Sam {
fn info(&self) -> &'static ModelInfo {
&MOBILE_SAM
}
fn segment(&self, rgb: &[[f32; 3]], w: usize, h: usize, prompt: &Prompt) -> Result<Vec<f32>> {
let e = self.net.embed(rgb, w, h)?;
Ok(self.net.predict(&e, prompt).0)
}
}
/// SHA-256 (FIPS 180-4), to verify downloaded weights without another dependency.
pub mod sha256 {
const K: [u32; 64] = [
0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5, 0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3,
0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174, 0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967, 0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13,
0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85, 0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3, 0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208,
0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2,
];
pub fn digest(data: &[u8]) -> [u8; 32] {
let mut h: [u32; 8] = [0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19];
let bits = (data.len() as u64).wrapping_mul(8);
let mut tail = data[data.len() - data.len() % 64..].to_vec();
tail.push(0x80);
while tail.len() % 64 != 56 {
tail.push(0);
}
tail.extend_from_slice(&bits.to_be_bytes());
for block in data[..data.len() - data.len() % 64].chunks_exact(64).chain(tail.chunks_exact(64)) {
let mut w = [0u32; 64];
for (i, c) in block.chunks_exact(4).enumerate() {
w[i] = u32::from_be_bytes([c[0], c[1], c[2], c[3]]);
}
for i in 16..64 {
let s0 = w[i - 15].rotate_right(7) ^ w[i - 15].rotate_right(18) ^ (w[i - 15] >> 3);
let s1 = w[i - 2].rotate_right(17) ^ w[i - 2].rotate_right(19) ^ (w[i - 2] >> 10);
w[i] = w[i - 16].wrapping_add(s0).wrapping_add(w[i - 7]).wrapping_add(s1);
}
let [mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut hh] = h;
for i in 0..64 {
let s1 = e.rotate_right(6) ^ e.rotate_right(11) ^ e.rotate_right(25);
let ch = (e & f) ^ (!e & g);
let t1 = hh.wrapping_add(s1).wrapping_add(ch).wrapping_add(K[i]).wrapping_add(w[i]);
let s0 = a.rotate_right(2) ^ a.rotate_right(13) ^ a.rotate_right(22);
let maj = (a & b) ^ (a & c) ^ (b & c);
let t2 = s0.wrapping_add(maj);
hh = g;
g = f;
f = e;
e = d.wrapping_add(t1);
d = c;
c = b;
b = a;
a = t1.wrapping_add(t2);
}
for (x, y) in h.iter_mut().zip([a, b, c, d, e, f, g, hh]) {
*x = x.wrapping_add(y);
}
}
let mut out = [0u8; 32];
for (i, v) in h.iter().enumerate() {
out[i * 4..i * 4 + 4].copy_from_slice(&v.to_be_bytes());
}
out
}
pub fn hex(data: &[u8]) -> String {
digest(data).iter().map(|b| format!("{b:02x}")).collect()
}
#[cfg(test)]
mod tests {
#[test]
fn known_vectors() {
assert_eq!(super::hex(b""), "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855");
assert_eq!(super::hex(b"abc"), "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad");
let long = b"abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq";
assert_eq!(super::hex(long), "248d6a61d20638b8e5c026930c3e6039a33ce45964ff2167f6ecedd419db06c1");
assert_eq!(super::hex(&vec![b'a'; 1_000_000]), "cdc76e5c9914fb9281a1c7e284d73e67f1809a48a497200e046d39ccc7112cd0");
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn registry_is_open_source_and_verifies_files() {
for m in MODELS {
assert!(ALLOWED_LICENCES.contains(&m.licence), "{}", m.id);
assert_eq!(m.sha256.len(), 64);
assert!(m.url.starts_with("https://") && 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());
}
}

View File

@ -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);
}
}

View File

@ -0,0 +1,677 @@
//! MobileSAM (C. Zhang et al., "Faster Segment Anything: Towards Lightweight SAM for Mobile
//! Applications", 2023; Apache-2.0) in plain Rust: the TinyViT-5M image encoder (Wu et al.,
//! "TinyViT", ECCV 2022) feeding Segment Anything's prompt encoder and two-way-transformer mask
//! decoder (Kirillov et al., "Segment Anything", ICCV 2023), loaded straight from the official
//! `mobile_sam.pt` checkpoint. The image goes in at 1024 pixels on its long side; the encoder
//! gives a 64×64×256 embedding, cached per image so that new prompts on the same frame only run
//! the (fast) decoder.
use std::f32::consts::PI;
use std::sync::Mutex;
use rayon::prelude::*;
use crate::nn::{self, Conv, Depthwise, Linear, Norm, Upconv};
use crate::pt::Weights;
use crate::{Prompt, Result};
const IMG: usize = 1024;
const EMB: usize = 64;
const DIM: usize = 256;
const LOW: usize = 256;
const MEAN: [f32; 3] = [123.675, 116.28, 103.53];
const STD: [f32; 3] = [58.395, 57.12, 57.375];
/// Weight lookup by name with the expected element count.
struct W<'a>(&'a Weights);
impl W<'_> {
fn get(&self, name: &str, n: usize) -> Result<&[f32]> {
let t = self.0.get(name).ok_or_else(|| format!("checkpoint is missing {name}"))?;
if t.data.len() != n {
return Err(format!("{name}: expected {n} values, found {}", t.data.len()));
}
Ok(&t.data)
}
fn linear(&self, p: &str, out: usize, inp: usize, bias: bool) -> Result<Linear> {
let w = self.get(&format!("{p}.weight"), out * inp)?;
let b = if bias { Some(self.get(&format!("{p}.bias"), out)?) } else { None };
Linear::new(w, b, out, inp).ok_or_else(|| format!("{p}: bad shape"))
}
fn norm(&self, p: &str, c: usize, eps: f32) -> Result<Norm> {
Ok(Norm { w: self.get(&format!("{p}.weight"), c)?.to_vec(), b: self.get(&format!("{p}.bias"), c)?.to_vec(), eps })
}
/// A conv + batch norm (eval) folded into one weight and bias: `[out × inp/groups × k × k]`.
fn conv_bn(&self, p: &str, out: usize, per: usize) -> Result<(Vec<f32>, Vec<f32>)> {
let w = self.get(&format!("{p}.c.weight"), out * per)?;
let g = self.get(&format!("{p}.bn.weight"), out)?;
let b = self.get(&format!("{p}.bn.bias"), out)?;
let m = self.get(&format!("{p}.bn.running_mean"), out)?;
let v = self.get(&format!("{p}.bn.running_var"), out)?;
let mut wf = w.to_vec();
let mut bf = vec![0.0; out];
for o in 0..out {
let s = g[o] / (v[o] + 1e-5).sqrt();
wf[o * per..(o + 1) * per].iter_mut().for_each(|x| *x *= s);
bf[o] = b[o] - m[o] * s;
}
Ok((wf, bf))
}
fn pointwise_bn(&self, p: &str, out: usize, inp: usize) -> Result<Linear> {
let (w, b) = self.conv_bn(p, out, inp)?;
Linear::new(&w, Some(&b), out, inp).ok_or_else(|| format!("{p}: bad shape"))
}
fn depthwise_bn(&self, p: &str, c: usize, stride: usize) -> Result<Depthwise> {
let (w, b) = self.conv_bn(p, c, 9)?;
Depthwise::new(&w, Some(&b), c, stride).ok_or_else(|| format!("{p}: bad shape"))
}
fn conv3_bn(&self, p: &str, out: usize, inp: usize, stride: usize) -> Result<Conv> {
let (w, b) = self.conv_bn(p, out, inp * 9)?;
Conv::new(&w, Some(&b), out, inp, 3, stride, 1).ok_or_else(|| format!("{p}: bad shape"))
}
}
/// MBConv: 1×1 expand, depthwise 3×3, 1×1 project, residual.
struct MbConv {
c1: Linear,
c2: Depthwise,
c3: Linear,
}
impl MbConv {
fn forward(&self, x: &[f32], h: usize, w: usize) -> Vec<f32> {
let mut t = self.c1.forward(x, h * w);
nn::gelu(&mut t);
let (mut t, _, _) = self.c2.forward(&t, h, w);
nn::gelu(&mut t);
let mut t = self.c3.forward(&t, h * w);
nn::add(&mut t, x);
nn::gelu(&mut t);
t
}
}
/// Patch merging between stages: 1×1, depthwise 3×3 (stride 2, or 1 into the last stage), 1×1.
struct Merge {
c1: Linear,
c2: Depthwise,
c3: Linear,
}
impl Merge {
fn forward(&self, x: &[f32], h: usize, w: usize) -> (Vec<f32>, usize, usize) {
let mut t = self.c1.forward(x, h * w);
nn::gelu(&mut t);
let (mut t, oh, ow) = self.c2.forward(&t, h, w);
nn::gelu(&mut t);
(self.c3.forward(&t, oh * ow), oh, ow)
}
}
/// A TinyViT block: windowed attention (with learned relative-position biases), a depthwise
/// "local conv", then an MLP.
struct Block {
norm: Norm,
qkv: Linear,
proj: Linear,
/// `[heads × n × n]` for the `ws × ws` window.
bias: Vec<f32>,
heads: usize,
ws: usize,
local: Depthwise,
mlp_norm: Norm,
fc1: Linear,
fc2: Linear,
}
impl Block {
fn new(w: &W, p: &str, dim: usize, heads: usize, ws: usize) -> Result<Block> {
let n = ws * ws;
// Offsets between window positions, numbered in first-seen order (as TinyViT does).
let points: Vec<(i64, i64)> = (0..ws as i64).flat_map(|r| (0..ws as i64).map(move |c| (r, c))).collect();
let mut offsets: Vec<(i64, i64)> = vec![];
let mut idx = Vec::with_capacity(n * n);
for a in &points {
for b in &points {
let o = ((a.0 - b.0).abs(), (a.1 - b.1).abs());
let i = match offsets.iter().position(|x| *x == o) {
Some(i) => i,
None => {
offsets.push(o);
offsets.len() - 1
}
};
idx.push(i);
}
}
let table = w.get(&format!("{p}.attn.attention_biases"), heads * offsets.len())?;
let no = offsets.len();
let bias = (0..heads).flat_map(|h| idx.iter().map(move |i| table[h * no + i])).collect::<Vec<f32>>();
Ok(Block {
norm: w.norm(&format!("{p}.attn.norm"), dim, 1e-5)?,
qkv: w.linear(&format!("{p}.attn.qkv"), 3 * dim, dim, true)?,
proj: w.linear(&format!("{p}.attn.proj"), dim, dim, true)?,
bias,
heads,
ws,
local: w.depthwise_bn(&format!("{p}.local_conv"), dim, 1)?,
mlp_norm: w.norm(&format!("{p}.mlp.norm"), dim, 1e-5)?,
fc1: w.linear(&format!("{p}.mlp.fc1"), 4 * dim, dim, true)?,
fc2: w.linear(&format!("{p}.mlp.fc2"), dim, 4 * dim, true)?,
})
}
fn forward(&self, x: &mut Vec<f32>, h: usize, w: usize) {
let dim = self.proj.out;
let (ws, heads) = (self.ws, self.heads);
let hd = dim / heads;
// Pad the grid to whole windows (zeros, as TinyViT does), then norm + qkv once.
let (ph, pw) = (h.div_ceil(ws) * ws, w.div_ceil(ws) * ws);
let mut pad = vec![0.0f32; ph * pw * dim];
for y in 0..h {
pad[y * pw * dim..(y * pw + w) * dim].copy_from_slice(&x[y * w * dim..(y + 1) * w * dim]);
}
self.norm.forward(&mut pad);
let qkv = self.qkv.forward(&pad, ph * pw);
// Per window: gather q, k, v (each head's qkv is [q | k | v]), attend.
let (nwy, nwx) = (ph / ws, pw / ws);
let n = ws * ws;
let outs: Vec<Vec<f32>> = (0..nwy * nwx)
.into_par_iter()
.map(|win| {
let (wy, wx) = (win / nwx, win % nwx);
let (mut q, mut k, mut v) = (vec![0.0f32; n * dim], vec![0.0f32; n * dim], vec![0.0f32; n * dim]);
for t in 0..n {
let (ty, tx) = (wy * ws + t / ws, wx * ws + t % ws);
let src = &qkv[(ty * pw + tx) * 3 * dim..][..3 * dim];
for hh in 0..heads {
let s = &src[hh * 3 * hd..][..3 * hd];
q[t * dim + hh * hd..][..hd].copy_from_slice(&s[..hd]);
k[t * dim + hh * hd..][..hd].copy_from_slice(&s[hd..2 * hd]);
v[t * dim + hh * hd..][..hd].copy_from_slice(&s[2 * hd..]);
}
}
nn::attention(&q, &k, &v, n, n, heads, hd, Some(&self.bias))
})
.collect();
// Un-window into the real grid, project, residual.
let mut a = vec![0.0f32; h * w * dim];
for (win, o) in outs.iter().enumerate() {
let (wy, wx) = (win / nwx, win % nwx);
for t in 0..n {
let (ty, tx) = (wy * ws + t / ws, wx * ws + t % ws);
if ty < h && tx < w {
a[(ty * w + tx) * dim..][..dim].copy_from_slice(&o[t * dim..][..dim]);
}
}
}
let a = self.proj.forward(&a, h * w);
nn::add(x, &a);
// Local depthwise conv (no residual), then the MLP with its residual.
let (lc, _, _) = self.local.forward(x, h, w);
*x = lc;
let mut m = x.clone();
self.mlp_norm.forward(&mut m);
let mut m = self.fc1.forward(&m, h * w);
nn::gelu(&mut m);
let m = self.fc2.forward(&m, h * w);
nn::add(x, &m);
}
}
/// The TinyViT-5M encoder with SAM's neck.
struct Encoder {
pe1: Conv,
pe2: Conv,
stage0: Vec<MbConv>,
merges: Vec<Merge>,
stages: Vec<Vec<Block>>,
neck1: Linear,
neck_ln1: Norm,
neck2: Conv,
neck_ln2: Norm,
}
const DIMS: [usize; 4] = [64, 128, 160, 320];
const DEPTHS: [usize; 4] = [2, 2, 6, 2];
const HEADS: [usize; 4] = [2, 4, 5, 10];
const WINDOWS: [usize; 4] = [7, 7, 14, 7];
impl Encoder {
fn new(w: &W) -> Result<Encoder> {
let p = "image_encoder";
let stage0 = (0..DEPTHS[0])
.map(|i| {
let b = format!("{p}.layers.0.blocks.{i}");
Ok(MbConv {
c1: w.pointwise_bn(&format!("{b}.conv1"), DIMS[0] * 4, DIMS[0])?,
c2: w.depthwise_bn(&format!("{b}.conv2"), DIMS[0] * 4, 1)?,
c3: w.pointwise_bn(&format!("{b}.conv3"), DIMS[0], DIMS[0] * 4)?,
})
})
.collect::<Result<Vec<_>>>()?;
let mut merges = vec![];
let mut stages = vec![];
for l in 0..3 {
let (inp, out) = (DIMS[l], DIMS[l + 1]);
let d = format!("{p}.layers.{l}.downsample");
merges.push(Merge {
c1: w.pointwise_bn(&format!("{d}.conv1"), out, inp)?,
c2: w.depthwise_bn(&format!("{d}.conv2"), out, if out == 320 { 1 } else { 2 })?,
c3: w.pointwise_bn(&format!("{d}.conv3"), out, out)?,
});
let s = l + 1;
stages
.push((0..DEPTHS[s]).map(|i| Block::new(w, &format!("{p}.layers.{s}.blocks.{i}"), DIMS[s], HEADS[s], WINDOWS[s])).collect::<Result<Vec<_>>>()?);
}
let n2 = w.get(&format!("{p}.neck.2.weight"), DIM * DIM * 9)?;
Ok(Encoder {
pe1: w.conv3_bn(&format!("{p}.patch_embed.seq.0"), DIMS[0] / 2, 3, 2)?,
pe2: w.conv3_bn(&format!("{p}.patch_embed.seq.2"), DIMS[0], DIMS[0] / 2, 2)?,
stage0,
merges,
stages,
neck1: w.linear(&format!("{p}.neck.0"), DIM, DIMS[3], false)?,
neck_ln1: w.norm(&format!("{p}.neck.1"), DIM, 1e-6)?,
neck2: Conv::new(n2, None, DIM, DIM, 3, 1, 1).ok_or("neck: bad shape")?,
neck_ln2: w.norm(&format!("{p}.neck.3"), DIM, 1e-6)?,
})
}
/// `IMG×IMG×3` normalised pixels → `EMB×EMB×DIM`.
fn forward(&self, x: &[f32]) -> Vec<f32> {
let (mut t, h, w) = self.pe1.forward(x, IMG, IMG);
nn::gelu(&mut t);
let (mut t, mut h, mut w) = {
let (t2, h2, w2) = self.pe2.forward(&t, h, w);
(t2, h2, w2)
};
for b in &self.stage0 {
t = b.forward(&t, h, w);
}
for (m, blocks) in self.merges.iter().zip(&self.stages) {
let (t2, h2, w2) = m.forward(&t, h, w);
(t, h, w) = (t2, h2, w2);
for b in blocks {
b.forward(&mut t, h, w);
}
}
let mut t = self.neck1.forward(&t, h * w);
self.neck_ln1.forward(&mut t);
let (mut t, _, _) = self.neck2.forward(&t, h, w);
self.neck_ln2.forward(&mut t);
t
}
}
/// SAM's attention with optional down-projection (`internal = DIM / rate`).
struct Attn {
q: Linear,
k: Linear,
v: Linear,
out: Linear,
}
impl Attn {
fn new(w: &W, p: &str, rate: usize) -> Result<Attn> {
let i = DIM / rate;
Ok(Attn {
q: w.linear(&format!("{p}.q_proj"), i, DIM, true)?,
k: w.linear(&format!("{p}.k_proj"), i, DIM, true)?,
v: w.linear(&format!("{p}.v_proj"), i, DIM, true)?,
out: w.linear(&format!("{p}.out_proj"), DIM, i, true)?,
})
}
fn forward(&self, q: &[f32], nq: usize, k: &[f32], v: &[f32], nk: usize) -> Vec<f32> {
let (q, k, v) = (self.q.forward(q, nq), self.k.forward(k, nk), self.v.forward(v, nk));
let heads = 8;
let o = nn::attention(&q, &k, &v, nq, nk, heads, self.q.out / heads, None);
self.out.forward(&o, nq)
}
}
fn plus(a: &[f32], b: &[f32]) -> Vec<f32> {
a.iter().zip(b).map(|(x, y)| x + y).collect()
}
struct TwoWayBlock {
self_attn: Attn,
norm1: Norm,
t2i: Attn,
norm2: Norm,
lin1: Linear,
lin2: Linear,
norm3: Norm,
i2t: Attn,
norm4: Norm,
skip_pe: bool,
}
/// The decoder: prompt encoder, two-way transformer, mask head.
struct Decoder {
gauss: Vec<f32>,
/// `EMB×EMB×DIM` positional encoding of the embedding grid.
dense_pe: Vec<f32>,
point_embed: [Vec<f32>; 4],
not_a_point: Vec<f32>,
no_mask: Vec<f32>,
md1: Conv,
md_ln1: Norm,
md2: Conv,
md_ln2: Norm,
md3: Linear,
blocks: Vec<TwoWayBlock>,
final_attn: Attn,
norm_final: Norm,
iou_token: Vec<f32>,
mask_tokens: Vec<f32>,
up1: Upconv,
up_ln: Norm,
up2: Upconv,
hyper: Vec<[Linear; 3]>,
iou_head: [Linear; 3],
}
impl Decoder {
fn new(w: &W) -> Result<Decoder> {
let pe = "prompt_encoder";
let md = "mask_decoder";
let gauss = w.get(&format!("{pe}.pe_layer.positional_encoding_gaussian_matrix"), 2 * DIM / 2)?.to_vec();
let pt = |i: usize| w.get(&format!("{pe}.point_embeddings.{i}.weight"), DIM).map(<[f32]>::to_vec);
let conv = |p: &str, out: usize, inp: usize, ks: usize| -> Result<Conv> {
let wt = w.get(&format!("{p}.weight"), out * inp * ks * ks)?;
let b = w.get(&format!("{p}.bias"), out)?;
Conv::new(wt, Some(b), out, inp, ks, 2, 0).ok_or_else(|| format!("{p}: bad shape"))
};
let blocks = (0..2)
.map(|i| {
let p = format!("{md}.transformer.layers.{i}");
Ok(TwoWayBlock {
self_attn: Attn::new(w, &format!("{p}.self_attn"), 1)?,
norm1: w.norm(&format!("{p}.norm1"), DIM, 1e-5)?,
t2i: Attn::new(w, &format!("{p}.cross_attn_token_to_image"), 2)?,
norm2: w.norm(&format!("{p}.norm2"), DIM, 1e-5)?,
lin1: w.linear(&format!("{p}.mlp.lin1"), 2048, DIM, true)?,
lin2: w.linear(&format!("{p}.mlp.lin2"), DIM, 2048, true)?,
norm3: w.norm(&format!("{p}.norm3"), DIM, 1e-5)?,
i2t: Attn::new(w, &format!("{p}.cross_attn_image_to_token"), 2)?,
norm4: w.norm(&format!("{p}.norm4"), DIM, 1e-5)?,
skip_pe: i == 0,
})
})
.collect::<Result<Vec<_>>>()?;
let upconv = |p: &str, inp: usize, out: usize| -> Result<Upconv> {
Upconv::new(w.get(&format!("{p}.weight"), inp * out * 4)?, w.get(&format!("{p}.bias"), out)?, inp, out).ok_or_else(|| format!("{p}: bad shape"))
};
let mlp3 = |p: &str, out: usize| -> Result<[Linear; 3]> {
Ok([
w.linear(&format!("{p}.layers.0"), DIM, DIM, true)?,
w.linear(&format!("{p}.layers.1"), DIM, DIM, true)?,
w.linear(&format!("{p}.layers.2"), out, DIM, true)?,
])
};
let mut d = Decoder {
gauss,
dense_pe: vec![],
point_embed: [pt(0)?, pt(1)?, pt(2)?, pt(3)?],
not_a_point: w.get(&format!("{pe}.not_a_point_embed.weight"), DIM)?.to_vec(),
no_mask: w.get(&format!("{pe}.no_mask_embed.weight"), DIM)?.to_vec(),
md1: conv(&format!("{pe}.mask_downscaling.0"), 4, 1, 2)?,
md_ln1: w.norm(&format!("{pe}.mask_downscaling.1"), 4, 1e-6)?,
md2: conv(&format!("{pe}.mask_downscaling.3"), 16, 4, 2)?,
md_ln2: w.norm(&format!("{pe}.mask_downscaling.4"), 16, 1e-6)?,
md3: w.linear(&format!("{pe}.mask_downscaling.6"), DIM, 16, true)?,
blocks,
final_attn: Attn::new(w, &format!("{md}.transformer.final_attn_token_to_image"), 2)?,
norm_final: w.norm(&format!("{md}.transformer.norm_final_attn"), DIM, 1e-5)?,
iou_token: w.get(&format!("{md}.iou_token.weight"), DIM)?.to_vec(),
mask_tokens: w.get(&format!("{md}.mask_tokens.weight"), 4 * DIM)?.to_vec(),
up1: upconv(&format!("{md}.output_upscaling.0"), DIM, DIM / 4)?,
up_ln: w.norm(&format!("{md}.output_upscaling.1"), DIM / 4, 1e-6)?,
up2: upconv(&format!("{md}.output_upscaling.3"), DIM / 4, DIM / 8)?,
hyper: (0..4).map(|i| mlp3(&format!("{md}.output_hypernetworks_mlps.{i}"), DIM / 8)).collect::<Result<Vec<_>>>()?,
iou_head: mlp3(&format!("{md}.iou_prediction_head"), 4)?,
};
d.dense_pe = (0..EMB * EMB).flat_map(|i| d.pe(((i % EMB) as f32 + 0.5) / EMB as f32, ((i / EMB) as f32 + 0.5) / EMB as f32)).collect();
Ok(d)
}
/// Random Fourier positional encoding of a point in [0, 1]².
fn pe(&self, x: f32, y: f32) -> Vec<f32> {
let half = DIM / 2;
let (cx, cy) = (2.0 * x - 1.0, 2.0 * y - 1.0);
let a: Vec<f32> = (0..half).map(|j| 2.0 * PI * (cx * self.gauss[j] + cy * self.gauss[half + j])).collect();
a.iter().map(|v| v.sin()).chain(a.iter().map(|v| v.cos())).collect()
}
fn mlp3(layers: &[Linear; 3], x: &[f32]) -> Vec<f32> {
let mut t = layers[0].forward(x, 1);
nn::relu(&mut t);
let mut t = layers[1].forward(&t, 1);
nn::relu(&mut t);
layers[2].forward(&t, 1)
}
/// Mask logits (`LOW×LOW`, in the padded `IMG` square) and their predicted IoU for the
/// prompt (coordinates already in `IMG` space).
fn forward(&self, emb: &[f32], points: &[([f32; 2], bool)], bbox: Option<[f32; 4]>, mask: Option<&[f32]>, multimask: bool) -> (Vec<f32>, f32) {
// Sparse prompt tokens.
let mut sparse: Vec<f32> = vec![];
for (p, fg) in points {
let mut e = self.pe((p[0] + 0.5) / IMG as f32, (p[1] + 0.5) / IMG as f32);
let add = &self.point_embed[usize::from(*fg)];
e.iter_mut().zip(add).for_each(|(a, b)| *a += b);
sparse.extend(e);
}
match bbox {
Some(b) => {
for (i, (x, y)) in [(b[0], b[1]), (b[2], b[3])].into_iter().enumerate() {
let mut e = self.pe((x + 0.5) / IMG as f32, (y + 0.5) / IMG as f32);
e.iter_mut().zip(&self.point_embed[2 + i]).for_each(|(a, b)| *a += b);
sparse.extend(e);
}
}
// Points without a box get a padding "not a point".
None if !points.is_empty() => sparse.extend(&self.not_a_point),
None => {}
}
// Dense prompt: the mask input, or "no mask".
let dense: Vec<f32> = match mask {
Some(m) if m.len() == LOW * LOW => {
let (mut t, h, w) = self.md1.forward(m, LOW, LOW);
self.md_ln1.forward(&mut t);
nn::gelu(&mut t);
let (mut t, h, w) = self.md2.forward(&t, h, w);
self.md_ln2.forward(&mut t);
nn::gelu(&mut t);
self.md3.forward(&t, h * w)
}
_ => (0..EMB * EMB).flat_map(|_| self.no_mask.iter().copied()).collect(),
};
// Tokens: IoU, four mask tokens, prompts.
let mut tokens: Vec<f32> = self.iou_token.clone();
tokens.extend(&self.mask_tokens);
tokens.extend(&sparse);
let nt = tokens.len() / DIM;
let nk = EMB * EMB;
let mut keys = plus(emb, &dense);
let pe = &self.dense_pe;
let qpe = tokens.clone();
let mut queries = tokens;
for b in &self.blocks {
// Self attention.
queries = if b.skip_pe {
b.self_attn.forward(&queries, nt, &queries, &queries, nt)
} else {
let q = plus(&queries, &qpe);
plus(&queries, &b.self_attn.forward(&q, nt, &q, &queries, nt))
};
b.norm1.forward(&mut queries);
// Tokens to image.
let q = plus(&queries, &qpe);
let k = plus(&keys, pe);
queries = plus(&queries, &b.t2i.forward(&q, nt, &k, &keys, nk));
b.norm2.forward(&mut queries);
// MLP.
let mut m = b.lin1.forward(&queries, nt);
nn::relu(&mut m);
queries = plus(&queries, &b.lin2.forward(&m, nt));
b.norm3.forward(&mut queries);
// Image to tokens.
let q = plus(&queries, &qpe);
let k = plus(&keys, pe);
keys = plus(&keys, &b.i2t.forward(&k, nk, &q, &queries, nt));
b.norm4.forward(&mut keys);
}
let q = plus(&queries, &qpe);
let k = plus(&keys, pe);
queries = plus(&queries, &self.final_attn.forward(&q, nt, &k, &keys, nk));
self.norm_final.forward(&mut queries);
// Upscale the image tokens 4× and dot them with each mask token's hypernetwork output.
let mut up = self.up1.forward(&keys, EMB, EMB);
self.up_ln.forward(&mut up);
nn::gelu(&mut up);
let mut up = self.up2.forward(&up, 2 * EMB, 2 * EMB);
nn::gelu(&mut up);
let c = DIM / 8;
let iou = Self::mlp3(&self.iou_head, &queries[..DIM]);
let pick = if multimask { (1..4).max_by(|a, b| iou[*a].total_cmp(&iou[*b])).unwrap_or(0) } else { 0 };
let hyper = Self::mlp3(&self.hyper[pick], &queries[(1 + pick) * DIM..][..DIM]);
let logits: Vec<f32> = up.par_chunks(c).map(|px| px.iter().zip(&hyper).map(|(a, b)| a * b).sum()).collect();
(logits, iou.get(pick).copied().unwrap_or(0.0))
}
}
/// An image's embedding and how it was fitted into the model's square.
#[derive(Clone)]
pub struct Embedding {
pub data: Vec<f32>,
pub w: usize,
pub h: usize,
/// Image pixels → model pixels.
pub scale: f32,
}
pub struct MobileSam {
enc: Encoder,
dec: Decoder,
/// The last image embedded (by content hash): new prompts on the same frame reuse it.
cache: Mutex<Option<(u64, Embedding)>>,
}
fn hash_rgb(rgb: &[[f32; 3]], w: usize, h: usize) -> u64 {
let mut s: u64 = 0xcbf2_9ce4_8422_2325 ^ ((w as u64) << 32 | h as u64);
for p in rgb.iter().step_by(7) {
for v in p {
s = (s ^ v.to_bits() as u64).wrapping_mul(0x0100_0000_01b3);
}
}
s
}
impl MobileSam {
pub fn from_weights(w: &Weights) -> Result<MobileSam> {
let w = W(w);
Ok(MobileSam { enc: Encoder::new(&w)?, dec: Decoder::new(&w)?, cache: Mutex::new(None) })
}
/// Embed an image (straight RGB 0–1, row-major `w×h`), or reuse the last embedding.
pub fn embed(&self, rgb: &[[f32; 3]], w: usize, h: usize) -> Result<Embedding> {
if w == 0 || h == 0 || rgb.len() != w * h {
return Err("empty image".into());
}
let key = hash_rgb(rgb, w, h);
if let Ok(c) = self.cache.lock()
&& let Some((k, e)) = c.as_ref()
&& *k == key
{
return Ok(e.clone());
}
let scale = IMG as f32 / w.max(h) as f32;
let (rw, rh) = (((w as f32 * scale).round() as usize).clamp(1, IMG), ((h as f32 * scale).round() as usize).clamp(1, IMG));
// Resize (box filter when shrinking, bilinear when growing), normalise, pad.
let mut x = vec![0.0f32; IMG * IMG * 3];
x.par_chunks_mut(IMG * 3).enumerate().take(rh).for_each(|(y, row)| {
for xo in 0..rw {
let px = sample(rgb, w, h, (xo as f32 + 0.5) / scale, (y as f32 + 0.5) / scale, 1.0 / scale);
for c in 0..3 {
row[xo * 3 + c] = (px[c] * 255.0 - MEAN[c]) / STD[c];
}
}
});
let e = Embedding { data: self.enc.forward(&x), w, h, scale };
if let Ok(mut c) = self.cache.lock() {
*c = Some((key, e.clone()));
}
Ok(e)
}
/// Foreground probability (`w×h`) for a prompt on an embedded image, and the predicted IoU.
pub fn predict(&self, e: &Embedding, prompt: &Prompt) -> (Vec<f32>, f32) {
let s = e.scale;
let points: Vec<([f32; 2], bool)> = prompt.points.iter().map(|(p, fg)| ([p[0] * s, p[1] * s], *fg)).collect();
let bbox = prompt.bbox.map(|b| b.map(|v| v * s));
// A prior mask (probabilities at image size) → logits on the low-res grid.
let mask: Option<Vec<f32>> = prompt.mask.as_ref().filter(|m| m.len() == e.w * e.h).map(|m| {
(0..LOW * LOW)
.map(|i| {
let (gx, gy) = (((i % LOW) as f32 + 0.5) * 4.0 / s, ((i / LOW) as f32 + 0.5) * 4.0 / s);
if gx >= e.w as f32 || gy >= e.h as f32 {
return -20.0;
}
let p = m[(gy as usize).min(e.h - 1) * e.w + (gx as usize).min(e.w - 1)];
((p - 0.5) * 40.0).clamp(-20.0, 20.0)
})
.collect()
});
let multimask = points.len() == 1 && bbox.is_none() && mask.is_none();
let (logits, iou) = self.dec.forward(&e.data, &points, bbox, mask.as_deref(), multimask);
// Low-res logits (the padded square) → image pixels, bilinear, then sigmoid.
let prob: Vec<f32> = (0..e.w * e.h)
.into_par_iter()
.map(|i| {
let (x, y) = ((i % e.w) as f32 + 0.5, (i / e.w) as f32 + 0.5);
let (u, v) = (x * s / 4.0 - 0.5, y * s / 4.0 - 0.5);
let l = bilinear(&logits, LOW, LOW, u, v);
1.0 / (1.0 + (-l).exp())
})
.collect();
(prob, iou)
}
}
fn bilinear(g: &[f32], w: usize, h: usize, u: f32, v: f32) -> f32 {
let (u, v) = (u.clamp(0.0, (w - 1) as f32), v.clamp(0.0, (h - 1) as f32));
let (x0, y0) = (u.floor() as usize, v.floor() as usize);
let (x1, y1) = ((x0 + 1).min(w - 1), (y0 + 1).min(h - 1));
let (fx, fy) = (u - x0 as f32, v - y0 as f32);
let at = |x: usize, y: usize| g.get(y * w + x).copied().unwrap_or(0.0);
(at(x0, y0) * (1.0 - fx) + at(x1, y0) * fx) * (1.0 - fy) + (at(x0, y1) * (1.0 - fx) + at(x1, y1) * fx) * fy
}
/// The image's colour around `(x, y)`: a `foot`-pixel box average when shrinking, else bilinear.
fn sample(rgb: &[[f32; 3]], w: usize, h: usize, x: f32, y: f32, foot: f32) -> [f32; 3] {
if foot <= 1.0 {
let (u, v) = ((x - 0.5).clamp(0.0, (w - 1) as f32), (y - 0.5).clamp(0.0, (h - 1) as f32));
let (x0, y0) = (u.floor() as usize, v.floor() as usize);
let (x1, y1) = ((x0 + 1).min(w - 1), (y0 + 1).min(h - 1));
let (fx, fy) = (u - x0 as f32, v - y0 as f32);
let p = |x: usize, y: usize| rgb[y * w + x];
return [0, 1, 2].map(|c| (p(x0, y0)[c] * (1.0 - fx) + p(x1, y0)[c] * fx) * (1.0 - fy) + (p(x0, y1)[c] * (1.0 - fx) + p(x1, y1)[c] * fx) * fy);
}
let r = foot / 2.0;
let (xa, xb) = (((x - r).floor().max(0.0)) as usize, ((x + r).ceil() as usize).min(w));
let (ya, yb) = (((y - r).floor().max(0.0)) as usize, ((y + r).ceil() as usize).min(h));
let mut s = [0.0f32; 3];
let mut n = 0.0;
for yy in ya..yb.max(ya + 1).min(h) {
for xx in xa..xb.max(xa + 1).min(w) {
let p = rgb[yy * w + xx];
for c in 0..3 {
s[c] += p[c];
}
n += 1.0;
}
}
s.map(|v| v / f32::max(n, 1.0))
}

445
crates/segment/src/nn.rs Normal file
View File

@ -0,0 +1,445 @@
//! The few neural-network kernels MobileSAM needs, on token-major feature maps (`[pixels, channels]`
//! row-major, i.e. NHWC with a batch of one): dense layers on a parallel GEMM (ndarray's safe
//! `general_mat_mul`, which picks AVX/FMA kernels at run time), convolutions as patch matrices
//! times weights, depthwise and transposed convolutions, layer norm, exact GELU and softmax.
//! Weights are re-laid out once at load so every inner loop runs over contiguous channels.
use ndarray::linalg::general_mat_mul;
use ndarray::{ArrayView2, ArrayViewMut2};
use rayon::prelude::*;
/// Rows of the left matrix per parallel task.
const ROW_BLOCK: usize = 64;
/// `out[m×n] += a[m×k] · b[k×n]` (row-major), rows split across threads.
pub fn gemm(a: &[f32], m: usize, k: usize, b: &[f32], n: usize, out: &mut [f32]) {
if m == 0 || n == 0 || a.len() < m * k || b.len() < k * n || out.len() < m * n {
return;
}
let Ok(bv) = ArrayView2::from_shape((k, n), &b[..k * n]) else { return };
let block = |r0: usize, oc: &mut [f32]| {
let rows = oc.len() / n;
if let (Ok(av), Ok(mut cv)) = (ArrayView2::from_shape((rows, k), &a[r0 * k..(r0 + rows) * k]), ArrayViewMut2::from_shape((rows, n), oc)) {
general_mat_mul(1.0, &av, &bv, 1.0, &mut cv);
}
};
if m <= ROW_BLOCK {
block(0, &mut out[..m * n]);
} else {
out[..m * n].par_chunks_mut(ROW_BLOCK * n).enumerate().for_each(|(i, oc)| block(i * ROW_BLOCK, oc));
}
}
/// A dense layer (also a 1×1 convolution): `y = x·Wᵀ + b`, with `W` stored transposed.
#[derive(Clone, Debug, Default)]
pub struct Linear {
/// `[inp × out]`.
pub wt: Vec<f32>,
pub b: Vec<f32>,
pub inp: usize,
pub out: usize,
}
impl Linear {
/// From PyTorch's `[out × inp]` weight (a 1×1 conv's `[out × inp × 1 × 1]` too).
pub fn new(w: &[f32], b: Option<&[f32]>, out: usize, inp: usize) -> Option<Linear> {
if w.len() != out * inp || b.is_some_and(|b| b.len() != out) {
return None;
}
let mut wt = vec![0.0; inp * out];
for o in 0..out {
for i in 0..inp {
wt[i * out + o] = w[o * inp + i];
}
}
Some(Linear { wt, b: b.map_or_else(|| vec![0.0; out], <[f32]>::to_vec), inp, out })
}
/// `rows` tokens of `inp` channels → `rows × out`.
pub fn forward(&self, x: &[f32], rows: usize) -> Vec<f32> {
let mut y: Vec<f32> = (0..rows).flat_map(|_| self.b.iter().copied()).collect();
gemm(x, rows, self.inp, &self.wt, self.out, &mut y);
y
}
}
/// A `ks×ks` convolution (any stride and padding, one group) on `h×w×inp` → `oh×ow×out`, as a
/// patch matrix times the weights, built a block of output rows at a time.
#[derive(Clone, Debug, Default)]
pub struct Conv {
/// `[(ky·ks + kx)·inp + ci] × out`.
pub wt: Vec<f32>,
pub b: Vec<f32>,
pub inp: usize,
pub out: usize,
pub ks: usize,
pub stride: usize,
pub pad: usize,
}
impl Conv {
/// From PyTorch's `[out × inp × ks × ks]`.
pub fn new(w: &[f32], b: Option<&[f32]>, out: usize, inp: usize, ks: usize, stride: usize, pad: usize) -> Option<Conv> {
if w.len() != out * inp * ks * ks || b.is_some_and(|b| b.len() != out) || stride == 0 {
return None;
}
let kk = ks * ks * inp;
let mut wt = vec![0.0; kk * out];
for o in 0..out {
for ci in 0..inp {
for t in 0..ks * ks {
wt[(t * inp + ci) * out + o] = w[(o * inp + ci) * ks * ks + t];
}
}
}
Some(Conv { wt, b: b.map_or_else(|| vec![0.0; out], <[f32]>::to_vec), inp, out, ks, stride, pad })
}
pub fn out_size(&self, n: usize) -> usize {
(n + 2 * self.pad).saturating_sub(self.ks) / self.stride + 1
}
pub fn forward(&self, x: &[f32], h: usize, w: usize) -> (Vec<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 || 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 - 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 - pad[1] as isize;
if ix < 0 || ix >= w as isize {
continue;
}
let src = (iy as usize * w + ix as usize) * self.inp;
let dst = ox * kk + (ky * self.ks + kx) * self.inp;
patch[dst..dst + self.inp].copy_from_slice(&x[src..src + self.inp]);
}
}
}
if let (Ok(av), Ok(bv), Ok(mut cv)) = (
ArrayView2::from_shape((ow, kk), &patch[..]),
ArrayView2::from_shape((kk, self.out), &self.wt[..]),
ArrayViewMut2::from_shape((ow, self.out), row),
) {
general_mat_mul(1.0, &av, &bv, 1.0, &mut cv);
}
});
y
}
}
/// 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,
}
impl Depthwise {
/// From PyTorch's `[c × 1 × 3 × 3]`.
pub fn new(w: &[f32], b: Option<&[f32]>, c: usize, stride: usize) -> Option<Depthwise> {
if w.len() != c * 9 || b.is_some_and(|b| b.len() != c) || stride == 0 {
return None;
}
let mut wt = vec![0.0; 9 * c];
for ch in 0..c {
for t in 0..9 {
wt[t * c + ch] = w[ch * 9 + t];
}
}
Some(Depthwise { w: wt, b: b.map_or_else(|| vec![0.0; c], <[f32]>::to_vec), c, ks: 3, stride })
}
/// 3×3 with padding 1.
pub fn forward(&self, x: &[f32], h: usize, w: usize) -> (Vec<f32>, usize, usize) {
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 || 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..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..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 * ks + kx) * c..][..c];
for ((o, s), k) in out.iter_mut().zip(src).zip(tap) {
*o += s * k;
}
}
}
}
});
y
}
}
/// A 2×2, stride-2 transposed convolution (each input pixel paints a 2×2 block):
/// `h×w×inp` → `2h×2w×out`.
#[derive(Clone, Debug, Default)]
pub struct Upconv {
/// One dense layer per output tap (dy, dx).
pub taps: Vec<Linear>,
pub out: usize,
}
impl Upconv {
/// From PyTorch's ConvTranspose2d `[inp × out × 2 × 2]`.
pub fn new(w: &[f32], b: &[f32], inp: usize, out: usize) -> Option<Upconv> {
if w.len() != inp * out * 4 || b.len() != out {
return None;
}
let taps = (0..4)
.map(|t| {
let mut wt = vec![0.0; inp * out];
for i in 0..inp {
for o in 0..out {
wt[i * out + o] = w[(i * out + o) * 4 + t];
}
}
Linear { wt, b: b.to_vec(), inp, out }
})
.collect();
Some(Upconv { taps, out })
}
pub fn forward(&self, x: &[f32], h: usize, w: usize) -> Vec<f32> {
let (ow, c) = (2 * w, self.out);
let mut y = vec![0.0f32; 4 * h * w * c];
for (t, lin) in self.taps.iter().enumerate() {
let (dy, dx) = (t / 2, t % 2);
let part = lin.forward(x, h * w);
for p in 0..h * w {
let (py, px) = (p / w, p % w);
let dst = ((2 * py + dy) * ow + 2 * px + dx) * c;
if let (Some(d), Some(s)) = (y.get_mut(dst..dst + c), part.get(p * c..(p + 1) * c)) {
d.copy_from_slice(s);
}
}
}
y
}
}
/// Layer norm over the channels of each token.
#[derive(Clone, Debug, Default)]
pub struct Norm {
pub w: Vec<f32>,
pub b: Vec<f32>,
pub eps: f32,
}
impl Norm {
pub fn forward(&self, x: &mut [f32]) {
let c = self.w.len();
if c == 0 {
return;
}
x.par_chunks_mut(c).for_each(|t| {
let mean = t.iter().sum::<f32>() / c as f32;
let var = t.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / c as f32;
let inv = 1.0 / (var + self.eps).sqrt();
for ((v, w), b) in t.iter_mut().zip(&self.w).zip(&self.b) {
*v = (*v - mean) * inv * w + b;
}
});
}
}
/// The error function (Abramowitz & Stegun 7.1.26 refined: |error| < 1.5e-7).
fn erf(x: f32) -> f32 {
let s = x.signum();
let x = x.abs() as f64;
let t = 1.0 / (1.0 + 0.327_591_1 * x);
let y = 1.0 - (((((1.061_405_429 * t - 1.453_152_027) * t) + 1.421_413_741) * t - 0.284_496_736) * t + 0.254_829_592) * t * (-x * x).exp();
s * y as f32
}
/// GELU (the exact, erf form, as PyTorch's `nn.GELU`), in place.
pub fn gelu(x: &mut [f32]) {
x.par_chunks_mut(4096).for_each(|c| {
for v in c {
*v = 0.5 * *v * (1.0 + erf(*v * std::f32::consts::FRAC_1_SQRT_2));
}
});
}
pub fn relu(x: &mut [f32]) {
for v in x {
*v = v.max(0.0);
}
}
/// `a += b`.
pub fn add(a: &mut [f32], b: &[f32]) {
a.par_chunks_mut(4096).zip(b.par_chunks(4096)).for_each(|(a, b)| {
for (x, y) in a.iter_mut().zip(b) {
*x += y;
}
});
}
/// Softmax of each row of `n` values, in place.
pub fn softmax_rows(x: &mut [f32], n: usize) {
if n == 0 {
return;
}
for row in x.chunks_mut(n) {
let m = row.iter().copied().fold(f32::NEG_INFINITY, f32::max);
let mut s = 0.0;
for v in row.iter_mut() {
*v = (*v - m).exp();
s += *v;
}
let inv = 1.0 / s.max(1e-30);
for v in row.iter_mut() {
*v *= inv;
}
}
}
/// Multi-head scaled dot-product attention of `nq` queries over `nk` keys/values, each `heads ×
/// hd` wide (token-major, heads side by side), plus an optional per-head bias `[heads × nq × nk]`.
pub fn attention(q: &[f32], k: &[f32], v: &[f32], nq: usize, nk: usize, heads: usize, hd: usize, bias: Option<&[f32]>) -> Vec<f32> {
let dim = heads * hd;
let mut out = vec![0.0f32; nq * dim];
if q.len() < nq * dim || k.len() < nk * dim || v.len() < nk * dim {
return out;
}
let scale = 1.0 / (hd as f32).sqrt();
for h in 0..heads {
// This head's q, kᵀ and v as contiguous matrices.
let qh: Vec<f32> = (0..nq).flat_map(|i| q[i * dim + h * hd..][..hd].iter().map(|x| x * scale)).collect();
let mut kt = vec![0.0f32; hd * nk];
for j in 0..nk {
for d in 0..hd {
kt[d * nk + j] = k[j * dim + h * hd + d];
}
}
let vh: Vec<f32> = (0..nk).flat_map(|j| v[j * dim + h * hd..][..hd].iter().copied()).collect();
let mut s = match bias.and_then(|b| b.get(h * nq * nk..(h + 1) * nq * nk)) {
Some(b) => b.to_vec(),
None => vec![0.0; nq * nk],
};
gemm(&qh, nq, hd, &kt, nk, &mut s);
softmax_rows(&mut s, nk);
let mut oh = vec![0.0f32; nq * hd];
gemm(&s, nq, nk, &vh, hd, &mut oh);
for i in 0..nq {
out[i * dim + h * hd..][..hd].copy_from_slice(&oh[i * hd..][..hd]);
}
}
out
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn gemm_matches_naive() {
let (m, k, n) = (130, 17, 9);
let a: Vec<f32> = (0..m * k).map(|i| ((i * 7) % 13) as f32 - 6.0).collect();
let b: Vec<f32> = (0..k * n).map(|i| ((i * 5) % 11) as f32 - 5.0).collect();
let mut c = vec![1.0f32; m * n];
gemm(&a, m, k, &b, n, &mut c);
for i in 0..m {
for j in 0..n {
let want: f32 = 1.0 + (0..k).map(|p| a[i * k + p] * b[p * n + j]).sum::<f32>();
assert_eq!(c[i * n + j], want);
}
}
}
#[test]
fn convolutions_match_direct_sums() {
// 3×3 stride-2 conv on a 5×4×2 input vs. the definition.
let (h, w, inp, out) = (5, 4, 2, 3);
let x: Vec<f32> = (0..h * w * inp).map(|i| (i % 7) as f32 * 0.5).collect();
let wt: Vec<f32> = (0..out * inp * 9).map(|i| ((i % 5) as f32 - 2.0) * 0.25).collect();
let b = [0.1, -0.2, 0.3];
let conv = Conv::new(&wt, Some(&b), out, inp, 3, 2, 1).unwrap();
let (y, oh, ow) = conv.forward(&x, h, w);
assert_eq!((oh, ow), (3, 2));
for oy in 0..oh {
for ox in 0..ow {
for o in 0..out {
let mut s = b[o];
for ci in 0..inp {
for ky in 0..3 {
for kx in 0..3 {
let (iy, ix) = ((oy * 2 + ky) as isize - 1, (ox * 2 + kx) as isize - 1);
if iy >= 0 && ix >= 0 && iy < h as isize && ix < w as isize {
s += x[(iy as usize * w + ix as usize) * inp + ci] * wt[((o * inp + ci) * 3 + ky) * 3 + kx];
}
}
}
}
assert!((y[(oy * ow + ox) * out + o] - s).abs() < 1e-5);
}
}
}
// Depthwise == a grouped conv: compare with Conv on each channel alone.
let dw_w: Vec<f32> = (0..inp * 9).map(|i| (i as f32 - 8.0) * 0.1).collect();
let dw = Depthwise::new(&dw_w, None, inp, 1).unwrap();
let (yd, _, _) = dw.forward(&x, h, w);
for ci in 0..inp {
let xc: Vec<f32> = (0..h * w).map(|p| x[p * inp + ci]).collect();
let single = Conv::new(&dw_w[ci * 9..ci * 9 + 9], None, 1, 1, 3, 1, 1).unwrap();
let (yc, _, _) = single.forward(&xc, h, w);
for p in 0..h * w {
assert!((yd[p * inp + ci] - yc[p]).abs() < 1e-5);
}
}
}
#[test]
fn gelu_and_softmax() {
let mut x = [-3.0f32, -1.0, 0.0, 1.0, 3.0];
gelu(&mut x);
// PyTorch's exact GELU values.
let want = [-0.004_049_7, -0.158_655_3, 0.0, 0.841_344_7, 2.995_950_3];
for (a, b) in x.iter().zip(want) {
assert!((a - b).abs() < 1e-5, "{a} {b}");
}
let mut s = [1.0f32, 2.0, 3.0];
softmax_rows(&mut s, 3);
assert!((s.iter().sum::<f32>() - 1.0).abs() < 1e-6 && s[2] > s[1]);
}
}

505
crates/segment/src/pt.rs Normal file
View File

@ -0,0 +1,505 @@
//! Reading PyTorch checkpoints (`torch.save` of a state dict) without PyTorch: the file is a zip
//! archive (stored or deflated entries, zip64 sizes) holding `data.pkl`, a pickle (protocol 2–4)
//! whose tensors are `torch._utils._rebuild_tensor_v2(storage, offset, size, stride, …)` calls
//! with storages as persistent ids pointing at `data/<key>` entries. Only what state dicts use is
//! understood; anything else is an error, never a panic.
use std::collections::HashMap;
/// A tensor read from a checkpoint: its shape and its values as f32 (row-major, contiguous).
#[derive(Clone, Debug, PartialEq)]
pub struct Tensor {
pub shape: Vec<usize>,
pub data: Vec<f32>,
}
/// A checkpoint's tensors by name (non-float tensors such as batch-norm counters are skipped).
pub type Weights = HashMap<String, Tensor>;
type Result<T> = std::result::Result<T, String>;
fn u16_at(b: &[u8], o: usize) -> Result<usize> {
b.get(o..o + 2).map(|s| u16::from_le_bytes([s[0], s[1]]) as usize).ok_or_else(|| "truncated zip".to_string())
}
fn u32_at(b: &[u8], o: usize) -> Result<u64> {
b.get(o..o + 4).map(|s| u32::from_le_bytes([s[0], s[1], s[2], s[3]]) as u64).ok_or_else(|| "truncated zip".to_string())
}
fn u64_at(b: &[u8], o: usize) -> Result<u64> {
b.get(o..o + 8).and_then(|s| s.try_into().ok()).map(u64::from_le_bytes).ok_or_else(|| "truncated zip".to_string())
}
/// One archive entry: where its data starts, its sizes and compression.
struct Entry {
data: usize,
size: usize,
stored: usize,
method: usize,
}
/// The entries of a zip archive by name.
fn zip_entries(b: &[u8]) -> Result<HashMap<String, Entry>> {
// End of central directory: the last `PK\x05\x06` (a comment may follow it).
let eocd =
(0..b.len().saturating_sub(21)).rev().find(|&i| b.get(i..i + 4) == Some(b"PK\x05\x06")).ok_or("not a zip archive (no end of central directory)")?;
let mut count = u16_at(b, eocd + 10)? as u64;
let mut cd = u32_at(b, eocd + 16)?;
// Zip64: the locator precedes the record.
if (count == 0xffff || cd == 0xffff_ffff) && eocd >= 20 && b.get(eocd - 20..eocd - 16) == Some(b"PK\x06\x07") {
let rec = u64_at(b, eocd - 12)? as usize;
if b.get(rec..rec + 4) == Some(b"PK\x06\x06") {
count = u64_at(b, rec + 32)?;
cd = u64_at(b, rec + 48)?;
}
}
let mut out = HashMap::new();
let mut p = cd as usize;
for _ in 0..count.min(1 << 20) {
if b.get(p..p + 4) != Some(b"PK\x01\x02") {
return Err("damaged zip central directory".into());
}
let method = u16_at(b, p + 10)?;
let mut stored = u32_at(b, p + 20)?;
let mut size = u32_at(b, p + 24)?;
let (nl, xl, cl) = (u16_at(b, p + 28)?, u16_at(b, p + 30)?, u16_at(b, p + 32)?);
let mut local = u32_at(b, p + 42)?;
let name = String::from_utf8_lossy(b.get(p + 46..p + 46 + nl).ok_or("truncated zip")?).to_string();
// Zip64 extended information: the 0xffffffff fields, in order.
let mut x = p + 46 + nl;
let xend = x + xl;
while x + 4 <= xend {
let (id, len) = (u16_at(b, x)?, u16_at(b, x + 2)?);
if id == 1 {
let mut f = x + 4;
for v in [&mut size, &mut stored, &mut local] {
if *v == 0xffff_ffff {
*v = u64_at(b, f)?;
f += 8;
}
}
}
x += 4 + len;
}
let l = local as usize;
if b.get(l..l + 4) != Some(b"PK\x03\x04") {
return Err(format!("damaged zip entry {name}"));
}
let data = l + 30 + u16_at(b, l + 26)? + u16_at(b, l + 28)?;
out.insert(name, Entry { data, size: size as usize, stored: stored as usize, method });
p += 46 + nl + xl + cl;
}
Ok(out)
}
/// An entry's bytes (stored, or raw-deflated).
fn read_entry<'a>(b: &'a [u8], e: &Entry) -> Result<std::borrow::Cow<'a, [u8]>> {
let raw = b.get(e.data..e.data + e.stored).ok_or("truncated zip entry")?;
match e.method {
0 => Ok(std::borrow::Cow::Borrowed(raw)),
8 => {
miniz_oxide::inflate::decompress_to_vec_with_limit(raw, e.size.max(1)).map(std::borrow::Cow::Owned).map_err(|e| format!("bad deflate data: {e:?}"))
}
m => Err(format!("unsupported zip compression {m}")),
}
}
/// 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 {
None,
Bool,
Int(i64),
Float,
Str(String),
Tuple(Vec<Obj>),
List(Vec<Obj>),
Dict(Vec<(Obj, Obj)>),
Global(String, String),
/// A storage: (dtype class name, entry key).
Storage(String, String),
/// `_rebuild_tensor_v2`: storage, element offset, size, stride.
Tensor {
dtype: String,
key: String,
offset: usize,
size: Vec<usize>,
stride: Vec<usize>,
},
/// A call we don't interpret.
Other,
Mark,
}
fn int_list(o: &Obj) -> Option<Vec<usize>> {
match o {
Obj::Tuple(v) | Obj::List(v) => v.iter().map(|x| if let Obj::Int(i) = x { usize::try_from(*i).ok() } else { None }).collect(),
_ => None,
}
}
/// Run the pickle program and return its result.
fn unpickle(p: &[u8]) -> Result<Obj> {
let mut stack: Vec<Obj> = Vec::new();
let mut memo: HashMap<u64, Obj> = HashMap::new();
let mut i = 0usize;
let take = |i: &mut usize, n: usize| -> Result<&[u8]> {
let s = p.get(*i..*i + n).ok_or("truncated pickle")?;
*i += n;
Ok(s)
};
let pop = |stack: &mut Vec<Obj>| stack.pop().ok_or_else(|| "pickle stack underflow".to_string());
let pop_mark = |stack: &mut Vec<Obj>| -> Result<Vec<Obj>> {
let m = stack.iter().rposition(|o| matches!(o, Obj::Mark)).ok_or("pickle: no mark")?;
let items = stack.split_off(m + 1);
stack.pop();
Ok(items)
};
let line = |i: &mut usize| -> Result<String> {
let rest = p.get(*i..).ok_or("truncated pickle")?;
let n = rest.iter().position(|c| *c == b'\n').ok_or("truncated pickle")?;
let s = String::from_utf8_lossy(&rest[..n]).to_string();
*i += n + 1;
Ok(s)
};
loop {
let op = *p.get(i).ok_or("truncated pickle")?;
i += 1;
match op {
0x80 => {
take(&mut i, 1)?;
}
0x95 => {
take(&mut i, 8)?;
}
b'}' => stack.push(Obj::Dict(vec![])),
b']' => stack.push(Obj::List(vec![])),
b')' => stack.push(Obj::Tuple(vec![])),
b'(' => stack.push(Obj::Mark),
b'N' => stack.push(Obj::None),
0x88 | 0x89 => stack.push(Obj::Bool),
b'K' => stack.push(Obj::Int(take(&mut i, 1)?[0] as i64)),
b'M' => {
let s = take(&mut i, 2)?;
stack.push(Obj::Int(u16::from_le_bytes([s[0], s[1]]) as i64));
}
b'J' => {
let s = take(&mut i, 4)?;
stack.push(Obj::Int(i32::from_le_bytes([s[0], s[1], s[2], s[3]]) as i64));
}
0x8a => {
let n = take(&mut i, 1)?[0] as usize;
let s = take(&mut i, n)?;
let mut v: i64 = 0;
for (k, byte) in s.iter().take(8).enumerate() {
v |= (*byte as i64) << (8 * k);
}
if n > 0 && n < 8 && s[n - 1] & 0x80 != 0 {
v -= 1i64 << (8 * n);
}
stack.push(Obj::Int(v));
}
b'G' => {
take(&mut i, 8)?;
stack.push(Obj::Float);
}
b'X' => {
let s = take(&mut i, 4)?;
let n = u32::from_le_bytes([s[0], s[1], s[2], s[3]]) as usize;
stack.push(Obj::Str(String::from_utf8_lossy(take(&mut i, n)?).to_string()));
}
0x8c => {
let n = take(&mut i, 1)?[0] as usize;
stack.push(Obj::Str(String::from_utf8_lossy(take(&mut i, n)?).to_string()));
}
b'C' => {
let n = take(&mut i, 1)?[0] as usize;
take(&mut i, n)?;
stack.push(Obj::Other);
}
b'B' => {
let s = take(&mut i, 4)?;
let n = u32::from_le_bytes([s[0], s[1], s[2], s[3]]) as usize;
take(&mut i, n)?;
stack.push(Obj::Other);
}
b'c' => {
let (m, n) = (line(&mut i)?, line(&mut i)?);
stack.push(Obj::Global(m, n));
}
0x93 => {
let n = pop(&mut stack)?;
let m = pop(&mut stack)?;
match (m, n) {
(Obj::Str(m), Obj::Str(n)) => stack.push(Obj::Global(m, n)),
_ => return Err("pickle: bad STACK_GLOBAL".into()),
}
}
b'q' => {
let k = take(&mut i, 1)?[0] as u64;
memo.insert(k, stack.last().cloned().ok_or("pickle stack underflow")?);
}
b'r' => {
let s = take(&mut i, 4)?;
memo.insert(u32::from_le_bytes([s[0], s[1], s[2], s[3]]) as u64, stack.last().cloned().ok_or("pickle stack underflow")?);
}
0x94 => {
let k = memo.len() as u64;
memo.insert(k, stack.last().cloned().ok_or("pickle stack underflow")?);
}
b'h' => {
let k = take(&mut i, 1)?[0] as u64;
stack.push(memo.get(&k).cloned().ok_or("pickle: bad memo")?);
}
b'j' => {
let s = take(&mut i, 4)?;
stack.push(memo.get(&(u32::from_le_bytes([s[0], s[1], s[2], s[3]]) as u64)).cloned().ok_or("pickle: bad memo")?);
}
b't' => {
let items = pop_mark(&mut stack)?;
stack.push(Obj::Tuple(items));
}
0x85..=0x87 => {
let n = (op - 0x84) as usize;
if stack.len() < n {
return Err("pickle stack underflow".into());
}
let items = stack.split_off(stack.len() - n);
stack.push(Obj::Tuple(items));
}
b'Q' => {
// Persistent id: ('storage', dtype class, key, location, numel).
let pid = pop(&mut stack)?;
let Obj::Tuple(t) = pid else { return Err("pickle: bad persistent id".into()) };
match (t.get(1), t.get(2)) {
(Some(Obj::Global(_, dtype)), Some(Obj::Str(key))) => stack.push(Obj::Storage(dtype.clone(), key.clone())),
_ => stack.push(Obj::Other),
}
}
b'R' | 0x81 => {
let args = pop(&mut stack)?;
let f = pop(&mut stack)?;
let args = match args {
Obj::Tuple(a) => a,
_ => vec![],
};
stack.push(match f {
Obj::Global(m, n) if m == "torch._utils" && n == "_rebuild_tensor_v2" => {
match (args.first(), args.get(1), args.get(2).and_then(int_list), args.get(3).and_then(int_list)) {
(Some(Obj::Storage(dtype, key)), Some(Obj::Int(off)), Some(size), Some(stride)) => {
Obj::Tensor { dtype: dtype.clone(), key: key.clone(), offset: usize::try_from(*off).unwrap_or(0), size, stride }
}
_ => Obj::Other,
}
}
Obj::Global(m, n) if m == "collections" && n == "OrderedDict" => Obj::Dict(vec![]),
_ => Obj::Other,
});
}
b'b' => {
pop(&mut stack)?;
}
b's' => {
let v = pop(&mut stack)?;
let k = pop(&mut stack)?;
if let Some(Obj::Dict(d)) = stack.last_mut() {
d.push((k, v));
}
}
b'u' => {
let items = pop_mark(&mut stack)?;
if let Some(Obj::Dict(d)) = stack.last_mut() {
let mut it = items.into_iter();
while let (Some(k), Some(v)) = (it.next(), it.next()) {
d.push((k, v));
}
}
}
b'a' => {
let v = pop(&mut stack)?;
if let Some(Obj::List(l)) = stack.last_mut() {
l.push(v);
}
}
b'e' => {
let items = pop_mark(&mut stack)?;
if let Some(Obj::List(l)) = stack.last_mut() {
l.extend(items);
}
}
b'.' => return pop(&mut stack),
other => return Err(format!("pickle: unsupported opcode 0x{other:02x}")),
}
}
}
/// Element bytes and conversion of a storage dtype.
fn to_f32(dtype: &str, bytes: &[u8]) -> Option<Vec<f32>> {
match dtype {
"FloatStorage" => Some(bytes.chunks_exact(4).map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect()),
"HalfStorage" => Some(bytes.chunks_exact(2).map(|c| half_to_f32(u16::from_le_bytes([c[0], c[1]]))).collect()),
"BFloat16Storage" => Some(bytes.chunks_exact(2).map(|c| f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16)).collect()),
_ => None,
}
}
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, m) => {
// Subnormal: normalise.
let mut e = 127 - 15 + 1;
let mut m = m;
while m & 0x400 == 0 {
m <<= 1;
e -= 1;
}
(s << 31) | (e << 23) | ((m & 0x3ff) << 13)
}
(0x1f, m) => (s << 31) | (0xff << 23) | (m << 13),
(e, m) => (s << 31) | ((e + 127 - 15) << 23) | (m << 13),
};
f32::from_bits(bits)
}
/// Read a PyTorch checkpoint's float tensors. A dict holding the state dict under `model` or
/// `state_dict` is unwrapped.
pub fn read(bytes: &[u8]) -> Result<Weights> {
let entries = zip_entries(bytes)?;
let (pkl_name, pkl) = entries.iter().find(|(n, _)| n.ends_with("data.pkl")).ok_or("not a PyTorch checkpoint (no data.pkl)")?;
let prefix = pkl_name.trim_end_matches("data.pkl").to_string();
let root = unpickle(&read_entry(bytes, pkl)?)?;
let dict = match root {
Obj::Dict(d) => d,
_ => return Err("the checkpoint is not a state dict".into()),
};
let dict = match dict.iter().find(|(k, _)| matches!(k, Obj::Str(s) if s == "model" || s == "state_dict")) {
Some((_, Obj::Dict(inner))) => inner.clone(),
_ => dict,
};
let mut storages: HashMap<String, std::borrow::Cow<[u8]>> = HashMap::new();
let mut out = Weights::new();
for (k, v) in dict {
let (Obj::Str(name), Obj::Tensor { dtype, key, offset, size, stride }) = (k, v) else { continue };
if !storages.contains_key(&key) {
let e = entries.get(&format!("{prefix}data/{key}")).ok_or_else(|| format!("missing storage {key}"))?;
storages.insert(key.clone(), read_entry(bytes, e)?);
}
let Some(values) = storages.get(&key).and_then(|b| to_f32(&dtype, b)) else { continue };
let n: usize = size.iter().product();
// Gather (any strides) into a contiguous row-major tensor.
let mut data = Vec::with_capacity(n);
let mut idx = vec![0usize; size.len()];
for _ in 0..n {
let at = offset + idx.iter().zip(&stride).map(|(i, s)| i * s).sum::<usize>();
data.push(*values.get(at).ok_or_else(|| format!("{name}: data out of range"))?);
for d in (0..idx.len()).rev() {
idx[d] += 1;
if idx[d] < size[d] {
break;
}
idx[d] = 0;
}
}
out.insert(name, Tensor { shape: size, data });
}
Ok(out)
}
#[cfg(test)]
mod tests {
use super::*;
/// A zip archive with stored entries (what `torch.save` writes).
fn zip(files: &[(&str, &[u8])]) -> Vec<u8> {
let mut out = vec![];
let mut central = vec![];
for (name, data) in files {
let off = out.len() as u32;
out.extend_from_slice(b"PK\x03\x04");
out.extend_from_slice(&[20, 0, 0, 0, 0, 0, 0, 0, 0, 0]);
out.extend_from_slice(&[0; 4]); // crc (unchecked)
out.extend_from_slice(&(data.len() as u32).to_le_bytes());
out.extend_from_slice(&(data.len() as u32).to_le_bytes());
out.extend_from_slice(&(name.len() as u16).to_le_bytes());
out.extend_from_slice(&0u16.to_le_bytes());
out.extend_from_slice(name.as_bytes());
out.extend_from_slice(data);
central.extend_from_slice(b"PK\x01\x02");
central.extend_from_slice(&[20, 0, 20, 0, 0, 0, 0, 0, 0, 0, 0, 0]);
central.extend_from_slice(&[0; 4]);
central.extend_from_slice(&(data.len() as u32).to_le_bytes());
central.extend_from_slice(&(data.len() as u32).to_le_bytes());
central.extend_from_slice(&(name.len() as u16).to_le_bytes());
central.extend_from_slice(&[0; 12]);
central.extend_from_slice(&off.to_le_bytes());
central.extend_from_slice(name.as_bytes());
}
let cd = out.len() as u32;
out.extend_from_slice(&central);
out.extend_from_slice(b"PK\x05\x06");
out.extend_from_slice(&[0; 4]);
out.extend_from_slice(&(files.len() as u16).to_le_bytes());
out.extend_from_slice(&(files.len() as u16).to_le_bytes());
out.extend_from_slice(&(central.len() as u32).to_le_bytes());
out.extend_from_slice(&cd.to_le_bytes());
out.extend_from_slice(&[0, 0]);
out
}
/// The pickle `torch.save({'w': t.t()})` writes for a 2×3 tensor stored transposed (strides
/// (1, 2)), plus an int64 counter that is skipped.
fn pickle() -> Vec<u8> {
let mut p = vec![0x80, 2, b'}', b'q', 0, b'('];
let s = |p: &mut Vec<u8>, t: &str| {
p.push(b'X');
p.extend_from_slice(&(t.len() as u32).to_le_bytes());
p.extend_from_slice(t.as_bytes());
};
for (name, dtype, key, size, stride) in [("w", "FloatStorage", "0", [2u8, 3], [1u8, 2]), ("n", "LongStorage", "1", [1, 1], [1, 1])] {
s(&mut p, name);
// The call's argument mark, then the persistent id tuple's.
p.extend_from_slice(b"ctorch._utils\n_rebuild_tensor_v2\n((");
s(&mut p, "storage");
p.extend_from_slice(format!("ctorch\n{dtype}\n").as_bytes());
s(&mut p, key);
s(&mut p, "cpu");
p.extend_from_slice(&[b'K', 6, b't', b'Q', b'K', 0]);
p.extend_from_slice(&[b'K', size[0], b'K', size[1], 0x86, b'K', stride[0], b'K', stride[1], 0x86]);
p.extend_from_slice(&[0x89, b'c']);
p.extend_from_slice(b"collections\nOrderedDict\n");
p.extend_from_slice(b")RtR");
}
p.extend_from_slice(b"u.");
p
}
#[test]
fn reads_a_state_dict_with_strides() {
let storage: Vec<u8> = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0].iter().flat_map(|v| v.to_le_bytes()).collect();
let pkl = pickle();
let file = zip(&[("ckpt/data.pkl", &pkl), ("ckpt/data/0", &storage), ("ckpt/data/1", &7i64.to_le_bytes())]);
let w = read(&file).unwrap();
// Strides (1, 2) over [1..6]: rows are [1, 3, 5] and [2, 4, 6].
assert_eq!(w.get("w"), Some(&Tensor { shape: vec![2, 3], data: vec![1.0, 3.0, 5.0, 2.0, 4.0, 6.0] }));
assert!(!w.contains_key("n"), "integer tensors are skipped");
}
#[test]
fn damaged_files_are_errors() {
assert!(read(b"not a zip").is_err());
let pkl = pickle();
let file = zip(&[("ckpt/data.pkl", &pkl[..pkl.len() - 5])]);
assert!(read(&file).is_err());
// Truncate the storage: the tensor reads out of range.
let file = zip(&[("ckpt/data.pkl", &pkl), ("ckpt/data/0", &[0u8; 8]), ("ckpt/data/1", &[0u8; 8])]);
assert!(read(&file).is_err());
assert_eq!(half_to_f32(0x3c00), 1.0);
assert_eq!(half_to_f32(0xc000), -2.0);
}
}

View File

@ -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);
}
}
}

View File

@ -11,6 +11,7 @@ workspace = true
[dependencies]
effectcraft-raster = { workspace = true }
effectcraft-segment = { workspace = true }
rayon = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }

View File

@ -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}");
}
}

View File

@ -19,6 +19,13 @@
//! colour models and the warped matte as a shape prior. Correction strokes on a frame re-cut
//! the whole frame with the warped matte as a soft prior, and propagation restarts from there.
//!
//! - **Trained models** (optional, swappable: `effectcraft_segment::MaskModel`, e.g. MobileSAM):
//! [`segment_with`] / [`propagate_with`] prompt the model with points along the strokes, or with
//! the warped matte (its box, points deep inside it and the matte itself), and its foreground
//! probability becomes a strong prior for the same graph cut, so strokes stay hard constraints
//! and edges still snap to colour. A model that fails, or whose mask disagrees with the flow, is
//! ignored for that frame.
//!
//! Mattes are in layer pixels at a given scale (1 = full resolution). Strokes are stored in
//! layer pixels ([`RotoData`], serde) and keyed per frame so that each frame's result is a
//! deterministic function of its *chain*: the base frame's strokes, then every frame's strokes
@ -33,6 +40,7 @@ pub mod rle;
use std::collections::BTreeMap;
use effectcraft_raster::Image;
use effectcraft_segment::{MaskModel, Prompt};
use rayon::prelude::*;
use serde::{Deserialize, Serialize};
@ -328,6 +336,8 @@ const HARD_BG: u8 = 2;
/// Graph-cut energy weights.
const GAMMA: f64 = 50.0;
const PRIOR: f64 = 1.0;
/// A trained model's matte as a prior: it outweighs the colour models.
const MODEL_PRIOR: f64 = 4.0;
/// Fixed-point scale of the energies.
const K: f64 = 16.0;
const INF: i64 = 1 << 40;
@ -356,8 +366,9 @@ fn beta(img: &[[f32; 3]], free: &[bool], w: usize, h: usize) -> f64 {
if s <= 1e-12 { 1.0 } else { n / (2.0 * s) }
}
/// One min-cut over the `free` pixels; others keep `labels`. `prior`: P(foreground) per pixel.
fn cut(img: &[[f32; 3]], w: usize, h: usize, labels: &mut [u8], free: &[bool], hard: &[u8], fg: &Gmm, bg: &Gmm, prior: Option<&[f32]>) {
/// One min-cut over the `free` pixels; others keep `labels`. `prior`: P(foreground) per pixel
/// and how much it weighs against the colour models.
fn cut(img: &[[f32; 3]], w: usize, h: usize, labels: &mut [u8], free: &[bool], hard: &[u8], fg: &Gmm, bg: &Gmm, prior: Option<(&[f32], f64)>) {
let mut idx = vec![u32::MAX; w * h];
let mut nodes = vec![];
for i in 0..w * h {
@ -375,10 +386,10 @@ fn cut(img: &[[f32; 3]], w: usize, h: usize, labels: &mut [u8], free: &[bool], h
.map(|&i| {
let c = img[i];
let (mut cf, mut cb) = (fg.cost(c), bg.cost(c));
if let Some(p) = prior {
if let Some((p, wt)) = prior {
let pf = (p[i] as f64).clamp(0.02, 0.98);
cf += PRIOR * -pf.ln();
cb += PRIOR * -(1.0 - pf).ln();
cf += wt * -pf.ln();
cb += wt * -(1.0 - pf).ln();
}
(cf, cb)
})
@ -568,6 +579,8 @@ pub struct Prior {
pub bg: Gmm,
/// Warped Refine Edge band.
pub refine: Vec<u8>,
/// How much the prior weighs against the colour models.
pub weight: f64,
}
/// Segment a frame from its strokes (and a propagated prior): GrabCut iterations on a reduced
@ -593,7 +606,10 @@ pub fn segment(img: &Image, strokes: &[&Stroke], prior: Option<&Prior>, opts: &S
}
})
.collect();
let soft_prior: Option<Vec<f32>> = prior.map(|p| matting::gauss(&p.soft, w, h, (opts.search_radius * opts.scale / 4.0).max(1.0)));
// A warped matte is uncertain by about the search radius: blur it. A model's matte is not.
let soft_prior: Option<Vec<f32>> =
prior.map(|p| if p.weight > PRIOR { p.soft.clone() } else { matting::gauss(&p.soft, w, h, (opts.search_radius * opts.scale / 4.0).max(1.0)) });
let pw = prior.map_or(PRIOR, |p| p.weight);
// Initial colour models: the strokes (and the propagated matte).
let (fg0, bg0) = match prior {
Some(p) => {
@ -666,7 +682,7 @@ pub fn segment(img: &Image, strokes: &[&Stroke], prior: Option<&Prior>, opts: &S
};
let free = vec![true; cw * ch];
for it in 0..iters {
cut(&crgb, cw, ch, &mut cl, &free, &chard, &fgm, &bgm, cprior.as_deref());
cut(&crgb, cw, ch, &mut cl, &free, &chard, &fgm, &bgm, cprior.as_deref().map(|c| (c, pw)));
if it + 1 < iters {
let (a, b) = fit_models(&crgb, &cl);
if !a.is_empty() && !b.is_empty() {
@ -688,11 +704,11 @@ pub fn segment(img: &Image, strokes: &[&Stroke], prior: Option<&Prior>, opts: &S
let sd = matting::signed_distance(&labels, w, h);
let band = f as f32 + 1.5;
let free: Vec<bool> = sd.iter().map(|d| d.abs() <= band).collect();
cut(&rgb, w, h, &mut labels, &free, &hard, &fgm, &bgm, soft_prior.as_deref());
cut(&rgb, w, h, &mut labels, &free, &hard, &fgm, &bgm, soft_prior.as_deref().map(|c| (c, pw)));
} else {
let free = vec![true; w * h];
for it in 0..iters {
cut(&rgb, w, h, &mut labels, &free, &hard, &fgm, &bgm, soft_prior.as_deref());
cut(&rgb, w, h, &mut labels, &free, &hard, &fgm, &bgm, soft_prior.as_deref().map(|c| (c, pw)));
if it + 1 < iters {
let (a, b) = fit_models(&rgb, &labels);
if !a.is_empty() && !b.is_empty() {
@ -768,7 +784,7 @@ pub fn warp_prior(prev_img: &Image, prev: &FrameSeg, next_img: &Image, opts: &Se
(soft, r)
})
.collect();
Prior { soft: res.iter().map(|r| r.0).collect(), fg: prev.fg.clone(), bg: prev.bg.clone(), refine: res.iter().map(|r| r.1).collect() }
Prior { soft: res.iter().map(|r| r.0).collect(), fg: prev.fg.clone(), bg: prev.bg.clone(), refine: res.iter().map(|r| r.1).collect(), weight: PRIOR }
}
/// Propagate the previous frame's segmentation to the next frame: warp it, then re-cut in a
@ -788,12 +804,133 @@ pub fn propagate(prev_img: &Image, prev: &FrameSeg, next_img: &Image, strokes: &
let hard = vec![0u8; w * h];
let soft = matting::gauss(&prior.soft, w, h, (r as f64 / 4.0).max(1.0));
let (fg, bg) = if prev.fg.is_empty() || prev.bg.is_empty() { fit_models(&rgb, &labels) } else { (prev.fg.clone(), prev.bg.clone()) };
cut(&rgb, w, h, &mut labels, &free, &hard, &fg, &bg, Some(&soft));
cut(&rgb, w, h, &mut labels, &free, &hard, &fg, &bg, Some((&soft, PRIOR)));
let seeds = core_of(&prior.soft, w, h, r as f64);
clean(&mut labels, &seeds, &vec![false; w * h], w, h);
finish(next_img, &rgb, labels, strokes, Some(&prior.refine), opts, refine_radius)
}
/// Points spread along a stroke (matte pixels), about one per stroke width, at most `max`.
fn points_along(s: &Stroke, scale: f64, max: usize) -> Vec<[f32; 2]> {
let pts: Vec<[f64; 2]> = s.points.iter().map(|p| [p[0] * scale, p[1] * scale]).collect();
let len: f64 = pts.windows(2).map(|w| (w[1][0] - w[0][0]).hypot(w[1][1] - w[0][1])).sum();
let step = (2.0 * s.radius * scale).max(len / max.max(1) as f64).max(1.0);
let mut out = pts.first().map(|p| vec![[p[0] as f32, p[1] as f32]]).unwrap_or_default();
let mut carry = 0.0;
for w in pts.windows(2) {
let d = (w[1][0] - w[0][0]).hypot(w[1][1] - w[0][1]);
let mut t = step - carry;
while t <= d && out.len() < max {
let f = t / d.max(1e-9);
out.push([(w[0][0] + (w[1][0] - w[0][0]) * f) as f32, (w[0][1] + (w[1][1] - w[0][1]) * f) as f32]);
t += step;
}
carry = (carry + d) % step;
}
out
}
/// A model's prompt from the strokes: points along the foreground and background strokes.
fn stroke_prompt(strokes: &[&Stroke], scale: f64) -> Prompt {
let mut points = vec![];
for (kind, fg) in [(StrokeKind::Fg, true), (StrokeKind::Bg, false)] {
let ss: Vec<&&Stroke> = strokes.iter().filter(|s| s.kind == kind).collect();
let per = (16 / ss.len().max(1)).max(1);
for s in ss {
points.extend(points_along(s, scale, per).into_iter().map(|p| (p, fg)));
}
}
Prompt { points, ..Default::default() }
}
/// Add a prior matte to a prompt: its box (a little larger), up to four points spread through
/// its core, and the matte itself.
fn add_prior_prompt(prompt: &mut Prompt, soft: &[f32], w: usize, h: usize) -> bool {
let bin: Vec<u8> = soft.iter().map(|v| (*v >= 0.5) as u8).collect();
let (mut x0, mut y0, mut x1, mut y1) = (usize::MAX, usize::MAX, 0usize, 0usize);
for (i, v) in bin.iter().enumerate() {
if *v != 0 {
let (x, y) = (i % w, i / w);
(x0, y0, x1, y1) = (x0.min(x), y0.min(y), x1.max(x), y1.max(y));
}
}
if x0 > x1 {
return false;
}
let (mx, my) = (((x1 - x0) as f32 * 0.05).max(4.0), ((y1 - y0) as f32 * 0.05).max(4.0));
prompt.bbox = Some([(x0 as f32 - mx).max(0.0), (y0 as f32 - my).max(0.0), (x1 as f32 + 1.0 + mx).min(w as f32), (y1 as f32 + 1.0 + my).min(h as f32)]);
// Farthest-point picks among the core (deep inside), starting from the deepest pixel.
let sd = matting::signed_distance(&bin, w, h);
let deep = sd.iter().copied().fold(0.0f32, f32::max);
let core: Vec<usize> = (0..w * h).filter(|i| sd[*i] >= (deep * 0.5).max(1.0)).step_by(((w * h) / 20_000).max(1)).collect();
let mut picks: Vec<usize> = core.iter().copied().max_by(|a, b| sd[*a].total_cmp(&sd[*b])).into_iter().collect();
while picks.len() < 4 {
let far = core.iter().copied().max_by_key(|c| picks.iter().map(|p| (c % w).abs_diff(p % w).pow(2) + (c / w).abs_diff(p / w).pow(2)).min().unwrap_or(0));
match far {
Some(f) if !picks.contains(&f) => picks.push(f),
_ => break,
}
}
prompt.points.extend(picks.into_iter().map(|i| ([(i % w) as f32 + 0.5, (i / w) as f32 + 0.5], true)));
prompt.mask = Some(soft.to_vec());
true
}
/// [`segment`] with a trained model (if any): its matte for the strokes (and the propagated
/// prior) becomes a strong prior for the cut.
pub fn segment_with(img: &Image, strokes: &[&Stroke], prior: Option<&Prior>, opts: &SegOpts, refine_radius: f64, model: Option<&dyn MaskModel>) -> FrameSeg {
let (w, h) = (img.width as usize, img.height as usize);
let Some(m) = model.filter(|_| w > 0 && h > 0) else { return segment(img, strokes, prior, opts, refine_radius) };
let mut prompt = stroke_prompt(strokes, opts.scale);
if let Some(p) = prior {
add_prior_prompt(&mut prompt, &p.soft, w, h);
}
if !prompt.points.iter().any(|(_, fg)| *fg) {
return segment(img, strokes, prior, opts, refine_radius);
}
match m.segment(&rgb_of(img), w, h, &prompt) {
Ok(soft) if soft.len() == w * h => {
let (fg, bg) = prior.map_or_else(Default::default, |p| (p.fg.clone(), p.bg.clone()));
let refine = prior.map_or_else(|| vec![0; w * h], |p| p.refine.clone());
segment(img, strokes, Some(&Prior { soft, fg, bg, refine, weight: MODEL_PRIOR }), opts, refine_radius)
}
_ => segment(img, strokes, prior, opts, refine_radius),
}
}
/// [`propagate`] with a trained model (if any): the model, prompted with the flow-warped matte,
/// re-segments the frame; when it disagrees with the flow (IoU under 0.5) the classical result
/// stands.
pub fn propagate_with(
prev_img: &Image,
prev: &FrameSeg,
next_img: &Image,
strokes: &[&Stroke],
opts: &SegOpts,
refine_radius: f64,
model: Option<&dyn MaskModel>,
) -> FrameSeg {
let (w, h) = (next_img.width as usize, next_img.height as usize);
let Some(m) = model.filter(|_| w > 0 && h > 0) else { return propagate(prev_img, prev, next_img, strokes, opts, refine_radius) };
let prior = warp_prior(prev_img, prev, next_img, opts);
let mut prompt = stroke_prompt(strokes, opts.scale);
if !add_prior_prompt(&mut prompt, &prior.soft, w, h) {
return propagate(prev_img, prev, next_img, strokes, opts, refine_radius);
}
match m.segment(&rgb_of(next_img), w, h, &prompt) {
Ok(soft) if soft.len() == w * h => {
let a: Vec<u8> = soft.iter().map(|v| (*v >= 0.5) as u8).collect();
let b: Vec<u8> = prior.soft.iter().map(|v| (*v >= 0.5) as u8).collect();
if iou(&a, &b) < 0.5 {
return propagate(prev_img, prev, next_img, strokes, opts, refine_radius);
}
let p = Prior { soft, weight: MODEL_PRIOR, ..prior };
segment(next_img, strokes, Some(&p), opts, refine_radius)
}
_ => propagate(prev_img, prev, next_img, strokes, opts, refine_radius),
}
}
/// Intersection over union of two binary mattes.
pub fn iou(a: &[u8], b: &[u8]) -> f64 {
let (mut i, mut u) = (0usize, 0usize);

View File

@ -335,3 +335,96 @@ fn roto_perf_1080p() {
let _ = matte_alpha(&rgb, &seg2, None, None, &MatteParams::default(), 20.0, 1.0);
eprintln!("refine 1080p: {:?}", t.elapsed());
}
/// A stand-in trained model: the foreground is a fixed rectangle (or nothing, or an error).
struct FakeModel(Option<[usize; 4]>, bool);
impl effectcraft_segment::MaskModel for FakeModel {
fn info(&self) -> &'static effectcraft_segment::ModelInfo {
&effectcraft_segment::MOBILE_SAM
}
fn segment(&self, _: &[[f32; 3]], w: usize, h: usize, prompt: &Prompt) -> Result<Vec<f32>, String> {
if self.1 {
return Err("broken".into());
}
assert!(prompt.points.iter().any(|(_, fg)| *fg), "prompted with foreground points");
Ok((0..w * h)
.map(|i| match self.0 {
Some([x0, y0, x1, y1]) => ((i % w) >= x0 && (i % w) < x1 && (i / w) >= y0 && (i / w) < y1) as u8 as f32 * 0.98 + 0.01,
None => 0.01,
})
.collect())
}
}
/// A grey bar on a two-tone background whose halves match the bar's halves: colour alone can't
/// find its edge, a model can.
#[test]
fn a_model_prior_shapes_the_cut_and_bad_models_are_ignored() {
let (w, h) = (160usize, 100usize);
let rect = [40, 30, 120, 70];
let mut img = Image::new(w as u32, h as u32);
let mut gt = vec![0u8; w * h];
for y in 0..h {
for x in 0..w {
let inside = x >= rect[0] && x < rect[2] && y >= rect[1] && y < rect[3];
// Left half dark, right half light, inside or out; a faint stripe pattern inside.
let base = if x < w / 2 { 0.2 } else { 0.8 };
let v = if inside { base + 0.02 * ((x / 3) % 2) as f32 } else { base };
img.set(x as u32, y as u32, [v, v, v, 1.0]);
gt[y * w + x] = inside as u8;
}
}
let st = [line(StrokeKind::Fg, 0, [50.0, 50.0], [110.0, 50.0], 4.0)];
let refs: Vec<&Stroke> = st.iter().collect();
let opts = SegOpts::default();
let classical = segment(&img, &refs, None, &opts, 0.0);
let good = FakeModel(Some(rect), false);
let with = segment_with(&img, &refs, None, &opts, 0.0, Some(&good));
let (jc, jm) = (iou(&classical.matte, &gt), iou(&with.matte, &gt));
assert!(jm > 0.95 && jm > jc + 0.2, "model {jm} vs classical {jc}");
// A failing model: the classical result, unchanged.
let broken = FakeModel(None, true);
assert_eq!(segment_with(&img, &refs, None, &opts, 0.0, Some(&broken)), classical);
// Propagation: a model that loses the object (disagrees with the flow) is ignored.
let next = img.clone();
let lost = FakeModel(None, false);
let a = propagate_with(&img, &with, &next, &[], &opts, 0.0, Some(&lost));
let b = propagate(&img, &with, &next, &[], &opts, 0.0);
assert_eq!(a, b);
// A good model keeps the bar through propagation.
let kept = propagate_with(&img, &with, &next, &[], &opts, 0.0, Some(&good));
assert!(iou(&kept.matte, &gt) > 0.95, "{}", iou(&kept.matte, &gt));
}
/// With the real MobileSAM weights (`EFFECTCRAFT_MOBILESAM=path/to/mobile_sam.pt`, else
/// skipped): the base frame and 20 propagated frames of the moving textured disk stay at least
/// as accurate as the classic engine.
#[test]
fn mobilesam_segments_and_propagates_the_disk() {
let Ok(path) = std::env::var("EFFECTCRAFT_MOBILESAM") else { return };
let 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);
let refs: Vec<&Stroke> = st.iter().collect();
let opts = SegOpts::default();
let base = segment_with(&img0, &refs, None, &opts, 0.0, m);
let jb = iou(&base.matte, &gt0);
let (mut seg, mut cseg) = (base, segment(&img0, &refs, None, &opts, 0.0));
let mut prev = img0;
let (mut worst, mut cworst): (f64, f64) = (1.0, 1.0);
let t = std::time::Instant::now();
for f in 1..=20 {
let (img, gt) = frame_at(center(f));
seg = propagate_with(&prev, &seg, &img, &[], &opts, 0.0, m);
cseg = propagate(&prev, &cseg, &img, &[], &opts, 0.0);
worst = worst.min(iou(&seg.matte, &gt));
cworst = cworst.min(iou(&cseg.matte, &gt));
prev = img;
}
eprintln!("MobileSAM: base IoU {jb:.4}, worst propagated {worst:.4} (classic {cworst:.4}), {:.0} ms/frame", t.elapsed().as_secs_f64() * 1000.0 / 20.0);
assert!(jb > 0.95 && worst > 0.93 && worst >= cworst - 0.01, "base {jb}, worst {worst} vs classic {cworst}");
}

View File

@ -14,6 +14,9 @@
//! than drifting.
//!
//! The feed also tracks per-channel peaks for the Audio panel's VU meters ([`Meter`]).
//!
//! **Scrubbing:** Ctrl/Cmd-dragging the current-time indicator plays a short snippet of the mix
//! at each new frame ([`AudioScrub`], `playback.scrubAudio`), as in After Effects.
use std::collections::VecDeque;
use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
@ -93,6 +96,13 @@ impl AudioFeed {
}
}
/// Drop what is still queued (scrubbing replaces the snippet still playing).
pub fn clear(&self) {
if let Ok(mut q) = self.queue.lock() {
q.clear();
}
}
pub fn push(&self, samples: &[f32]) {
if let Ok(mut q) = self.queue.lock() {
q.extend(samples.iter().copied());
@ -278,18 +288,82 @@ fn feed_block(src: &RenderSource, comp: ItemId, feed: &AudioFeed, span: &PlaySpa
cursor = span.wa;
}
let n = ((span.wb - cursor) as usize).min(BLOCK);
let t = sample_time(cursor, rate);
let buf = if mix == rate {
effectcraft_engine::render::audio::mix_comp(&src.project, src.footage.as_ref(), src.expr.as_deref(), comp, t, n, rate)
} else {
let m = (n as u64 * mix as u64).div_ceil(rate as u64) as usize + 1;
let b = effectcraft_engine::render::audio::mix_comp(&src.project, src.footage.as_ref(), src.expr.as_deref(), comp, t, m, mix);
resample_stereo(&b, mix, rate, n)
};
feed.push(&buf);
feed.push(&mix_block(src, comp, sample_time(cursor, rate), n, rate, mix));
cursor + n as i64
}
/// `n` stereo frames of `comp`'s mix from comp time `t` at the device `rate`, mixed at `mix` Hz
/// (resampled when they differ).
fn mix_block(src: &RenderSource, comp: ItemId, t: Tick, n: usize, rate: u32, mix: u32) -> Vec<f32> {
let render =
|frames: usize, at: u32| effectcraft_engine::render::audio::mix_comp(&src.project, src.footage.as_ref(), src.expr.as_deref(), comp, t, frames, at);
if mix == rate {
return render(n, rate);
}
let m = (n as u64 * mix as u64).div_ceil(rate as u64) as usize + 1;
resample_stereo(&render(m, mix), mix, rate, n)
}
/// Audio scrubbing: an open output that plays a short, faded snippet of the comp's mix each time
/// the current time moves to another frame (each snippet replaces what is still queued, so the
/// sound follows the pointer without lagging behind).
pub struct AudioScrub {
pub feed: Arc<AudioFeed>,
device: Box<dyn AudioDevice>,
pub rate: u32,
/// The comp and frame last played (holding still repeats nothing).
last: Option<(ItemId, i64)>,
/// When it last played (UI seconds): an idle scrub closes the device.
pub last_used: f64,
}
impl AudioScrub {
pub fn open(mut device: Box<dyn AudioDevice>, now: f64) -> Result<AudioScrub, String> {
let rate = device.sample_rate().max(8000);
let feed = Arc::new(AudioFeed::default());
device.start(feed.clone())?;
Ok(AudioScrub { feed, device, rate, last: None, last_used: now })
}
/// Play the frame of `comp` at `t` (one frame long, 30–100 ms) unless it was the last one
/// played. Returns whether a snippet was queued.
pub fn play(&mut self, src: &RenderSource, comp: ItemId, t: Tick, mix_rate: u32, now: f64) -> bool {
self.last_used = now;
let Some(c) = src.project.comp(comp) else { return false };
let frame = c.frame_rate.frame_at(t);
if self.last == Some((comp, frame)) {
return false;
}
self.last = Some((comp, frame));
let rate = self.rate;
let mix = if mix_rate == 0 { rate } else { mix_rate.clamp(8000, 192_000) };
let n = (c.frame_duration().seconds().clamp(0.03, 0.1) * rate as f64) as usize;
let mut buf = mix_block(src, comp, c.frame_rate.tick_of(frame), n, rate, mix);
// 3 ms fades: no clicks where snippets meet.
let fade = (rate as usize * 3 / 1000).clamp(1, n / 2 + 1);
let frames = buf.len() / 2;
for i in 0..frames {
let g = (i.min(frames - 1 - i) as f32 / fade as f32).min(1.0);
buf[2 * i] *= g;
buf[2 * i + 1] *= g;
}
self.feed.clear();
self.feed.push(&buf);
true
}
/// Every UI frame while open (Web Audio moves queued samples to the output here).
pub fn pump(&mut self) {
self.device.pump();
}
}
impl Drop for AudioScrub {
fn drop(&mut self) {
self.device.stop();
}
}
/// Linear resampling of interleaved stereo from `from` Hz to `n` frames at `to` Hz.
pub fn resample_stereo(src: &[f32], from: u32, to: u32, n: usize) -> Vec<f32> {
let frames = src.len() / 2;

View File

@ -38,6 +38,8 @@ pub enum PanelKind {
EssentialGraphics,
/// A dockable ScriptUI panel (a script from the ScriptUI Panels folder), by script window id.
ScriptPanel(u32),
/// Another Composition viewer (View ▸ New Viewer), by viewer id (`panels::viewers`).
Viewer(u32),
LumetriScopes,
Footage,
MediaBrowser,
@ -116,6 +118,7 @@ impl PanelKind {
PanelKind::ScriptConsole => "Script Console",
PanelKind::EssentialGraphics => "Essential Graphics",
PanelKind::ScriptPanel(_) => "ScriptUI Panel",
PanelKind::Viewer(_) => "Composition",
PanelKind::LumetriScopes => "Lumetri Scopes",
PanelKind::Footage => "Footage",
PanelKind::MediaBrowser => "Media Browser",

View File

@ -25,9 +25,13 @@ impl EffectcraftApp {
out.push((PanelKind::ScriptPanel(w.id), t));
}
}
// Each Composition viewer is named after the comp it shows.
for id in std::iter::once(0).chain(self.ui.viewers.keys().copied().filter(|v| *v != 0)) {
if let Some(t) = panels::viewers::title(self, id) {
out.push((panels::viewers::panel(id), t));
}
}
let Some(comp) = self.session.active_comp() else { return out };
let cname = self.session.active_comp_id().and_then(|id| self.session.project.item(id)).map(|i| i.name.clone()).unwrap_or_default();
out.push((PanelKind::Composition, format!("Composition {cname}")));
if let Some(l) = self.session.state.selected_layers.first().and_then(|id| comp.layer(*id)) {
out.push((PanelKind::EffectControls, format!("Effect Controls {}", l.name)));
out.push((PanelKind::Properties, format!("Properties: {}", l.name)));
@ -64,10 +68,14 @@ impl EffectcraftApp {
fn tab_decos(&self) -> Vec<(PanelKind, dock::TabDeco)> {
let Some(cid) = self.session.active_comp_id() else { return vec![] };
let swatch = self.comp_swatch(cid);
let mut v: Vec<(PanelKind, dock::TabDeco)> = [PanelKind::Composition, PanelKind::Timeline]
.into_iter()
.map(|p| (p, dock::TabDeco { swatch, locked: self.ui.locked_tabs.contains(&p.id()), viewer: true }))
.collect();
let mut v: Vec<(PanelKind, dock::TabDeco)> =
vec![(PanelKind::Timeline, dock::TabDeco { swatch, locked: self.ui.locked_tabs.contains(&PanelKind::Timeline.id()), viewer: true })];
// Every Composition viewer: its comp's swatch and its own lock.
for id in std::iter::once(0).chain(self.ui.viewers.keys().copied().filter(|v| *v != 0)) {
let p = panels::viewers::panel(id);
let swatch = panels::viewers::comp_of(self, id).map_or(self.tokens.text_faint, |c| self.comp_swatch(c));
v.push((p, dock::TabDeco { swatch, locked: self.ui.locked_tabs.contains(&p.id()), viewer: true }));
}
// Effect Controls carries the selected layer's label colour.
if let Some(l) = self.session.active_comp().and_then(|c| self.session.state.selected_layers.first().and_then(|id| c.layer(*id)))
&& l.label != effectcraft_engine::color::Label::None
@ -123,6 +131,7 @@ impl EffectcraftApp {
}
self.ui.dock.close(p);
self.edit_layout(|l| l.unfloat(p));
panels::viewers::on_close(self, p);
if self.ui.maximized == Some(p) {
self.ui.maximized = None;
}
@ -141,7 +150,8 @@ impl EffectcraftApp {
fn panel_body(&mut self, ui: &mut egui::Ui, p: PanelKind, rect: Rect) {
self.auto.add(&format!("panel.{}", p.id()), rect, p.title());
let mut content = rect;
if p == PanelKind::Composition && !self.ui.start_screen && panels::precomp::has_flow(self) {
let viewer = p == panels::viewers::active_panel(self);
if viewer && !self.ui.start_screen && panels::precomp::has_flow(self) {
// Composition Navigator: the flow of nested comps above the viewer.
let nav = Rect::from_min_size(content.min, vec2(content.width(), panels::precomp::NAV_H));
let mut child = ui.new_child(egui::UiBuilder::new().max_rect(nav).id_salt("comp-navigator"));
@ -152,7 +162,7 @@ impl EffectcraftApp {
let mut child = ui.new_child(egui::UiBuilder::new().max_rect(content).id_salt(("panel", p.id())));
child.set_clip_rect(content.intersect(ui.clip_rect()));
panels::show(self, &mut child, p, content);
if p == PanelKind::Composition && !self.ui.start_screen {
if viewer && !self.ui.start_screen {
panels::anim_tools::sketch_overlay(self, &mut child);
}
}
@ -220,7 +230,13 @@ impl EffectcraftApp {
DockAction::ToggleStacked(p) => {
self.ui.dock.toggle_stacked(p);
}
DockAction::Focus(p) => self.ui.focused = p,
DockAction::Focus(p) => {
self.ui.focused = p;
// A click in (or on the tab of) another Composition viewer makes it active.
if let Some(id) = panels::viewers::id_of(p) {
panels::viewers::activate(self, &ctx, id);
}
}
DockAction::Close(p) => self.close_panel(p),
DockAction::PanelMenu(p, pos) => {
ctx.data_mut(|d| d.insert_temp(egui::Id::new("panel-menu"), (p, pos)));

View File

@ -1,11 +1,11 @@
//! Background frame rendering and the RAM preview cache.
//!
//! Frames render on a dedicated thread pool from immutable `Arc<Project>` snapshots, so the UI
//! never waits for the compositor. Results land in a memory-budgeted cache keyed by (project
//! revision, comp, frame, scale, view, render options); the timeline draws the cached range as
//! the green cache bar. When the budget is full the least recently shown frames go first, and
//! frames of an older revision are dropped as soon as the project changes (they can never be
//! shown again).
//! never waits for the compositor. Results land in a memory-budgeted cache keyed by the comp's
//! *content* ([`comp_content`]: the comp and everything it uses), frame, scale, view and render
//! options; the timeline draws the cached range as the green cache bar. An edit keeps the frames
//! of every comp it doesn't touch, and undo finds the frames of the state it returns to, as in
//! After Effects. When the budget is full the least recently shown frames go first.
//!
//! wasm32 has no threads: jobs wait in the same priority queue and [`Frames::pump`] hands them
//! to a [`RemoteFrames`] (the browser's frame worker, a second engine instance fed with project
@ -22,16 +22,19 @@
//! registers with egui-wgpu and draws directly; pixels are read back only when something needs
//! them (Info panel, eyedroppers, histograms).
use std::collections::{HashMap, HashSet, VecDeque};
use std::collections::{BTreeSet, HashMap, HashSet, VecDeque};
use std::sync::{Arc, Mutex};
use effectcraft_engine::project::{ItemId, Project};
use effectcraft_engine::project::{ItemId, ItemKind, LayerSource, Node, Project, PropGroup};
use effectcraft_engine::render::disk_cache::{self, DiskCache};
use effectcraft_engine::render::{ExprHost, FootageSource, LayerCache, RenderOpts, Renderer};
use effectcraft_engine::time::Tick;
use effectcraft_gpu::{DisplayFrame, Gpu};
use rayon::prelude::*;
/// Comp identities (and the project snapshots behind them) the RAM preview keeps frames of.
const MAX_KEEPERS: usize = 64;
/// GPU frames kept for RAM preview (bytes of video memory).
const GPU_BUDGET: usize = 1 << 30;
@ -57,9 +60,14 @@ impl FrameImage {
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
/// A viewer frame: its comp's content, the frame, scale, view and render options. `revision` is
/// the project revision it was asked for at (job bookkeeping); it is not part of the identity, so
/// a later revision with the same content finds the frame.
#[derive(Clone, Copy, Debug)]
pub struct FrameKey {
pub revision: u64,
/// The comp's content key ([`Frames::content_of`]).
pub content: u64,
pub comp: u64,
pub frame: i64,
/// Render scale × 1000.
@ -71,12 +79,31 @@ pub struct FrameKey {
}
impl FrameKey {
/// Whether `self` is a frame of the same comp, revision, scale, view and options as `base`.
/// Everything but the revision.
fn id(&self) -> (u64, u64, i64, u32, u64, u64) {
(self.content, self.comp, self.frame, self.scale, self.view, self.opts)
}
/// Whether `self` is a frame of the same comp content, scale, view and options as `base`.
pub fn same_series(&self, base: &FrameKey) -> bool {
FrameKey { frame: base.frame, ..*self } == *base
}
}
impl PartialEq for FrameKey {
fn eq(&self, o: &FrameKey) -> bool {
self.id() == o.id()
}
}
impl Eq for FrameKey {}
impl std::hash::Hash for FrameKey {
fn hash<H: std::hash::Hasher>(&self, h: &mut H) {
self.id().hash(h);
}
}
struct Cache {
map: HashMap<FrameKey, FrameImage>,
order: VecDeque<FrameKey>,
@ -129,6 +156,14 @@ impl Cache {
self.order.push_back(k);
}
}
/// Drop every frame of the identities in `gone`.
fn drop_contents(&mut self, gone: &HashSet<u64>) {
let keys: Vec<FrameKey> = self.map.keys().filter(|k| gone.contains(&k.content)).copied().collect();
for k in &keys {
self.remove(k);
}
self.order.retain(|k| !gone.contains(&k.content));
}
fn evict_to_budget(&mut self) {
while self.bytes > self.budget {
let Some(old) = self.order.pop_front() else { break };
@ -173,6 +208,54 @@ fn content_key(keys: &ContentKeys, project: &Project, revision: u64, comp: u64)
k
}
/// The identity of a comp's pixels in `project`, cheap enough to take at every revision: the
/// addresses of the comps it draws (an edit replaces a comp's `Arc`: `Arc::make_mut` copies it
/// while another snapshot holds it, and [`Frames::content_of`] keeps one for every identity it
/// hands out), and by value the project settings and the footage, solids and proxies it uses.
/// With expressions in its comps (which can read anything) every item counts.
pub fn comp_content(project: &Project, comp: ItemId) -> u64 {
use std::hash::{Hash, Hasher};
fn has_expression(g: &PropGroup) -> bool {
g.children.iter().any(|c| match c {
Node::Prop(p) => p.expr.is_some(),
Node::Group(sub) => has_expression(sub),
})
}
let mut h = std::collections::hash_map::DefaultHasher::new();
format!("{:?}", project.settings).hash(&mut h);
let mut seen = BTreeSet::new();
let mut stack = vec![comp];
let mut expressions = false;
while let Some(id) = stack.pop() {
if !seen.insert(id) {
continue;
}
let Some(ItemKind::Comp(c)) = project.item(id).map(|i| &i.kind) else { continue };
for l in &c.layers {
if let LayerSource::Comp { item } | LayerSource::Footage { item } | LayerSource::Solid { item } = l.source {
stack.push(item);
}
expressions = expressions || has_expression(&l.props);
}
}
let ids: Vec<ItemId> = if expressions { project.items.keys().copied().collect() } else { seen.into_iter().collect() };
for id in ids {
let Some(it) = project.item(id) else { continue };
id.0.hash(&mut h);
if expressions {
// (`comp("Name")`, `footage("Name")`)
it.name.hash(&mut h);
}
match &it.kind {
ItemKind::Comp(c) => (Arc::as_ptr(c) as usize).hash(&mut h),
k => format!("{k:?}").hash(&mut h),
}
format!("{:?}", it.proxy).hash(&mut h);
}
comp.0.hash(&mut h);
h.finish()
}
/// Hash of render options (part of [`FrameKey`] and of the disk key).
pub fn opts_hash(o: &RenderOpts) -> u64 {
let mut h = disk_cache::Hash128::default();
@ -259,6 +342,11 @@ pub struct Frames {
/// Render time of the last viewer (urgent) frame in ms, for the Info panel / perf readout.
pub last_ms: Arc<Mutex<f64>>,
content_keys: ContentKeys,
/// [`comp_content`] identities by (revision, comp).
identities: Mutex<HashMap<(u64, u64), u64>>,
/// A project snapshot for each identity handed out, with when it was last asked for (see
/// [`Frames::content_of`]); the counter.
keepers: Mutex<(u64, HashMap<u64, (u64, Arc<Project>)>)>,
/// Where frames render instead of the UI thread (wasm32), and how many it is rendering.
remote: Option<Arc<dyn RemoteFrames>>,
remote_busy: Arc<Mutex<usize>>,
@ -284,6 +372,8 @@ impl Default for Frames {
ctx: None,
last_ms: Arc::new(Mutex::new(0.0)),
content_keys: Arc::default(),
identities: Mutex::default(),
keepers: Mutex::default(),
remote: None,
remote_busy: Arc::default(),
}
@ -325,16 +415,44 @@ impl Frames {
Some(img)
}
/// Drop frames of project revisions before `revision` (and their queued jobs): an edit
/// makes them unreachable, so they only take memory from frames that can still be shown.
pub fn drop_stale(&self, revision: u64) {
if let Ok(mut c) = self.cache.lock() {
let stale: Vec<FrameKey> = c.map.keys().filter(|k| k.revision < revision).copied().collect();
for k in &stale {
c.remove(k);
/// The [`comp_content`] identity of `comp` in `project` at `revision` (taken once per
/// revision). The project is kept while frames of that identity are cached or rendering, so
/// the comp addresses it hashes can't be reused by other comps meanwhile.
pub fn content_of(&self, project: &Arc<Project>, revision: u64, comp: ItemId) -> u64 {
let known = self.identities.lock().ok().and_then(|m| m.get(&(revision, comp.0)).copied());
let id = known.unwrap_or_else(|| {
let id = comp_content(project, comp);
if let Ok(mut m) = self.identities.lock() {
if m.len() > 256 {
m.clear();
}
m.insert((revision, comp.0), id);
}
id
});
if let Ok(mut guard) = self.keepers.lock() {
let (clock, k) = &mut *guard;
*clock += 1;
let now = *clock;
k.entry(id).or_insert_with(|| (now, project.clone())).0 = now;
if k.len() > MAX_KEEPERS {
let inflight: HashSet<u64> = self.inflight.lock().map(|s| s.iter().map(|f| f.content).collect()).unwrap_or_default();
if let Ok(mut c) = self.cache.lock() {
// Snapshots of identities without frames go; past the limit the least
// recently asked for go too, with their frames (a long drag leaves many).
let cached: HashSet<u64> = c.map.keys().map(|f| f.content).collect();
k.retain(|i, _| *i == id || cached.contains(i) || inflight.contains(i));
if k.len() > MAX_KEEPERS {
let mut ages: Vec<(u64, u64)> = k.iter().filter(|(i, _)| **i != id && !inflight.contains(i)).map(|(i, (t, _))| (*t, *i)).collect();
ages.sort_unstable();
let gone: HashSet<u64> = ages.into_iter().take(k.len().saturating_sub(MAX_KEEPERS / 2)).map(|(_, i)| i).collect();
k.retain(|i, _| !gone.contains(i));
c.drop_contents(&gone);
}
}
}
c.order.retain(|k| k.revision >= revision);
}
id
}
pub fn is_cached(&self, k: &FrameKey) -> bool {
@ -345,7 +463,7 @@ impl Frames {
self.inflight.lock().map(|s| s.len()).unwrap_or(0)
}
/// Cached frame numbers of the series of `base` (its comp, revision, scale, view and render
/// Cached frame numbers of the series of `base` (its comp content, scale, view and render
/// options) — for the cache bar.
pub fn cached_frames(&self, base: &FrameKey) -> Vec<i64> {
let Ok(c) = self.cache.lock() else { return vec![] };
@ -744,8 +862,8 @@ impl Queue {
mod tests {
use super::*;
fn key(revision: u64, frame: i64, opts: u64) -> FrameKey {
FrameKey { revision, comp: 1, frame, scale: 1000, view: 0, opts }
fn key(content: u64, frame: i64, opts: u64) -> FrameKey {
FrameKey { revision: content, content, comp: 1, frame, scale: 1000, view: 0, opts }
}
fn img() -> FrameImage {
@ -771,17 +889,73 @@ mod tests {
}
#[test]
fn stale_revisions_are_dropped_and_series_are_kept_apart() {
fn frames_are_keyed_by_content_and_series_are_kept_apart() {
let f = Frames::default();
insert(&f, key(1, 0, 0));
insert(&f, key(2, 0, 0));
insert(&f, key(2, 1, 7));
// Other render options (Draft, shadows…) are another series.
// Other content or render options (Draft, shadows…) are another series.
assert_eq!(f.cached_frames(&key(2, 0, 0)), vec![0]);
assert_eq!(f.cached_frames(&key(2, 0, 7)), vec![1]);
f.drop_stale(2);
assert!(!f.is_cached(&key(1, 0, 0)) && f.is_cached(&key(2, 0, 0)));
let c = f.cache.lock().unwrap();
assert_eq!((c.order.len(), c.bytes), (2, 2 * 16), "the memory is given back");
// A later revision with the same content finds the frame (an edit elsewhere, an undo).
assert!(f.is_cached(&FrameKey { revision: 99, ..key(1, 0, 0) }));
assert!(!f.is_cached(&key(3, 0, 0)));
}
#[test]
fn an_edit_keeps_the_frames_of_comps_it_does_not_touch() {
use effectcraft_engine::Session;
use serde_json::json;
let mut s = Session::default();
let a = ItemId(s.execute("comp.new", json!({"name": "A", "width": 8, "height": 8, "duration": 1})).unwrap()["comp"].as_u64().unwrap());
let b = ItemId(s.execute("comp.new", json!({"name": "B", "width": 8, "height": 8, "duration": 1})).unwrap()["comp"].as_u64().unwrap());
let f = Frames::default();
let at = |s: &Session, c: ItemId| f.content_of(&s.project, s.revision, c);
let (a0, b0) = (at(&s, a), at(&s, b));
// Edit B: A's content (and frames) stay, B's change.
s.execute("layer.newSolid", json!({"comp": b.0, "color": "#ff0000"})).unwrap();
assert_eq!(at(&s, a), a0);
assert_ne!(at(&s, b), b0);
// Undo: B is back to the content its old frames were rendered for.
s.execute("edit.undo", json!({})).unwrap();
assert_eq!(at(&s, b), b0);
}
#[test]
fn a_long_drag_keeps_a_bounded_number_of_states() {
use effectcraft_engine::Session;
use serde_json::json;
let mut s = Session::default();
let a = ItemId(s.execute("comp.new", json!({"name": "A", "width": 8, "height": 8, "duration": 1})).unwrap()["comp"].as_u64().unwrap());
let f = Frames::default();
let mut last = None;
for i in 0..200 {
s.execute("comp.settings", json!({"width": 8 + i, "merge": "drag"})).unwrap();
let k = FrameKey { content: f.content_of(&s.project, s.revision, a), ..key(0, 0, 0) };
insert(&f, k);
last = Some(k);
}
// The snapshots and the frames kept stay bounded (a dropped identity's frames go with it,
// so a later state that reuses its addresses can't be shown an old frame).
assert!(f.keepers.lock().unwrap().1.len() <= MAX_KEEPERS);
assert!(f.cache.lock().unwrap().map.len() <= MAX_KEEPERS);
assert!(f.is_cached(&last.unwrap()), "the current state's frame stays");
}
#[test]
fn an_identity_keeps_its_comps_so_an_edit_in_place_still_changes_it() {
use effectcraft_engine::Session;
use serde_json::json;
let mut s = Session::default();
let a = ItemId(s.execute("comp.new", json!({"name": "A", "width": 8, "height": 8, "duration": 1})).unwrap()["comp"].as_u64().unwrap());
// A project nothing else holds (no undo history): make_mut would edit the comp in place,
// at the same address, if the identity hadn't kept a snapshot holding it.
let mut p = Arc::new((*s.project).clone());
drop(s);
let f = Frames::default();
let before = f.content_of(&p, 1, a);
let comp = Arc::make_mut(&mut p).comp_mut(a).unwrap();
comp.width = 16;
assert_ne!(f.content_of(&p, 2, a), before);
}
}

View File

@ -187,6 +187,8 @@ pub struct EffectcraftApp {
pub(crate) toast: Option<(String, f64)>,
/// Texture of the last CPU frame shown in the viewer and the key it came from.
pub(crate) viewer_tex: Option<(egui::TextureHandle, FrameKey)>,
/// The frames the other (passive) Composition viewers show, by viewer id.
pub(crate) passive_tex: std::collections::HashMap<u32, (FrameKey, panels::viewers::PassiveTexture)>,
/// The last GPU frame shown: its egui texture id (registered with egui-wgpu), key and texture.
pub(crate) viewer_native: Option<(egui::TextureId, FrameKey, Arc<effectcraft_gpu::DisplayFrame>)>,
/// The texture the viewer draws and its frame key (CPU or GPU frame).
@ -209,6 +211,8 @@ pub struct EffectcraftApp {
pub integrated_titlebar: bool,
/// Audio preview while playing (audio clock drives playback).
pub audio: Option<audio::AudioPlayback>,
/// Audio scrubbing's open output (Ctrl/Cmd-drag the current time, `playback.scrubAudio`).
pub scrub: Option<audio::AudioScrub>,
/// Audio panel VU meters.
pub meter: audio::Meter,
/// Waveform peak summaries per footage item (see `panels::waveform`).
@ -261,6 +265,7 @@ impl EffectcraftApp {
last_activity: (0.0, 0),
toast: None,
viewer_tex: None,
passive_tex: std::collections::HashMap::new(),
viewer_native: None,
viewer_shown: None,
viewer_image: None,
@ -273,6 +278,7 @@ impl EffectcraftApp {
dialog_state: Default::default(),
integrated_titlebar: false,
audio: None,
scrub: None,
meter: Default::default(),
waveforms: Default::default(),
last_reveal: None,
@ -377,7 +383,7 @@ impl EffectcraftApp {
}
if !self.ui.dock.contains(p) {
let near = match p {
PanelKind::Layer | PanelKind::Flowchart => PanelKind::Composition,
PanelKind::Layer | PanelKind::Flowchart | PanelKind::Viewer(_) => PanelKind::Composition,
PanelKind::RenderQueue => PanelKind::Timeline,
PanelKind::EffectControls | PanelKind::History => PanelKind::Project,
_ => PanelKind::EffectsPresets,
@ -511,11 +517,12 @@ impl EffectcraftApp {
FrameKey { frame, ..self.frame_series(comp, scale) }
}
/// The key of frame 0 of `comp`'s viewer frames at `scale` (the revision, view and render
/// The key of frame 0 of `comp`'s viewer frames at `scale` (the content, view and render
/// options are the same for every frame: loops compute this once).
pub fn frame_series(&self, comp: ItemId, scale: f64) -> FrameKey {
FrameKey {
revision: self.session.revision,
content: self.frames.content_of(&self.session.project, self.session.revision, comp),
comp: comp.0,
frame: 0,
scale: (scale * 1000.0).round() as u32,
@ -524,14 +531,11 @@ impl EffectcraftApp {
}
}
/// The frame series the viewer of `comp` shows now (its scale, view and options, at the
/// current revision): what the cache bars count.
/// The frame series the viewer of `comp` shows now (its scale, view and options, for the
/// comp's current content): what the cache bars count.
pub fn shown_series(&self, comp: ItemId) -> FrameKey {
match self.viewer_shown.as_ref().map(|(_, k)| *k) {
Some(k) if k.comp == comp.0 && k.revision == self.session.revision => FrameKey { frame: 0, ..k },
Some(k) => self.frame_series(comp, k.scale as f64 / 1000.0),
None => self.frame_series(comp, 1.0),
}
let scale = self.viewer_shown.as_ref().filter(|(_, k)| k.comp == comp.0).map_or(1000, |(_, k)| k.scale);
self.frame_series(comp, scale as f64 / 1000.0)
}
/// Hash of the comp viewer's 3D view camera (0 for the active camera view).
@ -621,8 +625,9 @@ impl EffectcraftApp {
self.session.set_time(fr.tick_of(plan.first));
}
let restore_max = preset.full_screen.then_some(self.ui.maximized);
if preset.full_screen && self.ui.dock.contains(PanelKind::Composition) {
self.ui.maximized = Some(PanelKind::Composition);
let viewer = panels::viewers::active_panel(self);
if preset.full_screen && self.ui.dock.contains(viewer) {
self.ui.maximized = Some(viewer);
}
self.playback = Playback {
playing: true,
@ -675,6 +680,30 @@ impl EffectcraftApp {
}
}
/// Audio scrubbing: play the frame of `comp` at `t` (see [`audio::AudioScrub`]), opening the
/// output on first use. Silent comps and missing devices do nothing.
pub fn scrub_audio(&mut self, comp: ItemId, t: Tick, now: f64) {
if self.playback.playing || !effectcraft_engine::render::audio::comp_has_audio(&self.session.project, comp) {
return;
}
if self.scrub.is_none() {
let out = audio::AudioOutput::from_prefs(&self.session.prefs);
let Some(dev) = self.hooks.audio_device.as_ref().and_then(|f| f(&out)) else { return };
match audio::AudioScrub::open(dev, now) {
Ok(s) => self.scrub = Some(s),
Err(e) => {
log::warn!("audio scrubbing: {e}");
return;
}
}
}
let src = self.render_source();
let mix = self.session.prefs.audio.preview_sample_rate;
if let Some(s) = &mut self.scrub {
s.play(&src, comp, t, mix, now);
}
}
pub fn stop(&mut self) {
let was = std::mem::replace(&mut self.playback.playing, false);
self.audio = None;
@ -843,6 +872,14 @@ impl EffectcraftApp {
if let Some(a) = &mut self.audio {
a.pump();
self.meter.update(a.feed.take_peaks(), now);
} else if let Some(s) = &mut self.scrub {
// Scrubbing: the meters follow the snippets; an idle second closes the output.
s.pump();
self.meter.update(s.feed.take_peaks(), now);
if now - s.last_used > 1.0 {
self.scrub = None;
}
ctx.request_repaint_after(std::time::Duration::from_millis(50));
} else if self.meter.active() {
self.meter.update([0.0; 2], now);
ctx.request_repaint();
@ -1120,19 +1157,14 @@ impl EffectcraftApp {
fn handle_events(&mut self, ctx: &egui::Context) {
for ev in self.session.drain_events() {
match ev {
effectcraft_engine::Event::OpenComp(_) => {
// The comp's viewer and its Timeline tab come to the front.
for p in [PanelKind::Composition, PanelKind::Timeline] {
if !self.ui.locked_tabs.contains(&p.id()) {
self.ui.dock.activate(p);
}
}
effectcraft_engine::Event::OpenComp(c) => {
panels::viewers::on_open_comp(self, c);
self.ui.timeline.pps = None;
}
effectcraft_engine::Event::Toast { message, .. } => self.toast = Some((message, ctx.input(|i| i.time))),
effectcraft_engine::Event::OpenUrl(url) => ctx.open_url(egui::OpenUrl::new_tab(url)),
// Frames of earlier revisions can't be shown again.
effectcraft_engine::Event::ProjectChanged { revision } => self.frames.drop_stale(revision),
// Frames are keyed by content: an edit keeps those of comps it doesn't touch.
effectcraft_engine::Event::ProjectChanged { .. } => {}
effectcraft_engine::Event::Frontend { command, params } => {
if let Err(e) = crate::menus::frontend(self, ctx, &command, params) {
self.ui.status = e;

View File

@ -562,6 +562,15 @@ pub fn frontend(app: &mut EffectcraftApp, ctx: &egui::Context, id: &str, p: Valu
Value::Null
}
// Window ▸ <ScriptUI panel>: dock (or bring forward) the panel the script built.
// Audio scrubbing: one frame of audio at `time` (default the current time).
"playback.scrubAudio" => {
let cid = app.session.active_comp_id().ok_or("no composition")?;
let t = p.get("time").and_then(Value::as_f64).map_or(app.session.time(), effectcraft_engine::time::Tick::from_seconds_f64);
app.scrub_audio(cid, t, now);
json!({"playing": app.scrub.is_some()})
}
// View ▸ New Viewer.
"view.newViewer" => json!({"viewer": crate::panels::viewers::new_viewer(app)}),
"window.scriptPanel" => {
let id = p.get("window").and_then(Value::as_u64).ok_or("no ScriptUI panel window")? as u32;
app.show_panel(PanelKind::ScriptPanel(id));

View File

@ -530,7 +530,7 @@ mod tests {
assert_eq!(format_ago(now, now - 3 * 86_400), "3 days ago");
assert_eq!(format_ago(now, now - 30 * 86_400), "2026-08-22");
assert_eq!(format_ago(now, now + 100), "Just now");
assert_eq!((format_size(900), format_size(12_800), format_size(3_500_000)), ("900 B".into(), "13 KB".into(), "3.3 MB".into()));
assert_eq!((format_size(900), format_size(13_000), format_size(3_500_000)), ("900 B".into(), "13 KB".into(), "3.3 MB".into()));
assert_eq!((kind_of("/a/B.ecproj"), kind_of("c.AEP"), kind_of("x")), ("EffectCraft project", "After Effects project", "Project"));
}

View File

@ -367,14 +367,61 @@ pub(crate) fn comp_markers(app: &mut EffectcraftApp, ui: &mut egui::Ui, comp: &C
}
}
/// Layer markers on a layer's row `r`.
/// Layer markers on a layer's row `r`; on a precomp layer, its comp's markers too (read-only,
/// outlined: hover names them, a double-click opens the nested comp at the marker).
pub(crate) fn layer_markers(app: &mut EffectcraftApp, ui: &mut egui::Ui, clip: Rect, comp: &Comp, layer: &Layer, tm: TMap, r: Rect) {
let p = ui.painter().with_clip_rect(clip);
nested_markers(app, ui, &p, comp, layer, tm, r);
for (i, m) in layer.markers.iter().enumerate() {
marker(app, ui, &p, comp, tm, m, layer.comp_time(m.time).seconds(), MarkerRef { layer: Some(layer.id.0), index: i }, r.min.y + 3.0, r.max.y - 3.0);
}
}
/// The nested comp's markers on a precomp layer's bar (`markers.nested`).
fn nested_markers(app: &mut EffectcraftApp, ui: &mut egui::Ui, p: &egui::Painter, comp: &Comp, layer: &Layer, tm: TMap, r: Rect) {
let effectcraft_engine::project::LayerSource::Comp { item } = layer.source else { return };
let Some(cid) = app.session.active_comp_id() else { return };
let project = app.session.project.clone();
let Some(nc) = project.comp(item) else { return };
if nc.markers.is_empty() {
return;
}
// Where they land, once per revision (with time remapping that samples the layer's frames).
let (key, rev) = (egui::Id::new(("nested-markers", cid.0, layer.id.0)), app.session.revision);
let cached = ui.ctx().data(|d| d.get_temp::<(u64, Vec<(Tick, usize)>)>(key)).filter(|(r, _)| *r == rev);
let at = cached.map(|(_, v)| v).unwrap_or_else(|| {
let v: Vec<(Tick, usize)> =
effectcraft_engine::commands::markers::nested_markers(&project, cid, comp, layer).into_iter().map(|(t, i, _)| (t, i)).collect();
ui.ctx().data_mut(|d| d.insert_temp(key, (rev, v.clone())));
v
});
let list: Vec<(Tick, usize, &Marker)> = at.into_iter().filter_map(|(t, i)| Some((t, i, nc.markers.get(i)?))).collect();
let t = app.tokens;
let name = project.item(item).map(|i| i.name.clone()).unwrap_or_default();
let (y0, h) = (r.min.y + 3.0, (r.height() - 6.0).min(10.0));
for (ct, i, m) in list {
let x = tm.x(ct.seconds());
let col = if m.label == effectcraft_engine::color::Label::None { Color32::from_rgb(0xd8, 0xd8, 0x60) } else { t.label(m.label) }.gamma_multiply(0.75);
let ym = y0 + h * 0.6;
let pts = vec![pos2(x - 4.0, y0), pos2(x + 4.0, y0), pos2(x + 4.0, ym), pos2(x, y0 + h), pos2(x - 4.0, ym)];
p.add(egui::Shape::closed_line(pts, Stroke::new(1.2, col)));
if !m.comment.is_empty() {
p.text(pos2(x + 7.0, y0 + h / 2.0), Align2::LEFT_CENTER, m.comment.lines().next().unwrap_or_default(), Tokens::ui(10.0), t.text_dim);
}
let hit = Rect::from_min_max(pos2(x - 5.0, y0 - 1.0), pos2(x + 5.0, r.max.y - 3.0));
let tip = if m.comment.is_empty() { format!("Marker in {name}") } else { format!("{} (marker in {name})", m.comment) };
let resp = ui.interact(hit, egui::Id::new(("nested-marker", layer.id.0, i)), Sense::click()).on_hover_text(&tip);
app.auto.add(&format!("timeline.layer.{}.nestedMarker.{i}", layer.id.0), hit, &tip);
if resp.double_clicked() {
let opened =
app.session.execute("comp.open", json!({"comp": item.0})).and_then(|_| app.session.execute("time.set", json!({"time": m.time.seconds()})));
if let Err(e) = opened {
app.ui.status = e.to_string();
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;

View File

@ -51,6 +51,7 @@ pub mod viewer;
pub mod viewer_overlays;
pub mod viewer_text;
pub mod viewer_tools;
pub mod viewers;
pub mod waveform;
use effectcraft_engine::Session;
@ -82,7 +83,8 @@ pub fn timecode(session: &Session, comp: &Comp, t: Tick) -> String {
pub fn show(app: &mut EffectcraftApp, ui: &mut egui::Ui, p: PanelKind, rect: Rect) {
match p {
PanelKind::Composition => viewer::show(app, ui, rect),
PanelKind::Composition => viewers::show(app, ui, 0, rect),
PanelKind::Viewer(n) => viewers::show(app, ui, n, rect),
PanelKind::Timeline => timeline::show(app, ui, rect),
PanelKind::Project => project::show(app, ui, rect),
PanelKind::EffectControls => effect_controls::show(app, ui, rect),
@ -161,8 +163,9 @@ pub fn panel_menu_popup(app: &mut EffectcraftApp, ui: &mut egui::Ui) {
close = true;
}
}
PanelKind::Composition => {
PanelKind::Composition | PanelKind::Viewer(_) => {
for (label, cmd) in [
("New Viewer", "view.newViewer"),
("Composition Settings…", "app.compSettings"),
("View Options…", "view.layerControls"),
("Show Grid", "view.grid"),

View File

@ -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,6 +29,8 @@ enum Act {
Run(String, Value),
PickFolder(String),
PickProject(String),
/// Install from File (a downloaded weights file) for a task's models.
PickModel(Task),
}
fn hex(c: [u8; 3]) -> String {
@ -61,6 +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::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();
@ -207,6 +211,105 @@ pub fn human_bytes(b: u64) -> String {
if u == 0 { format!("{b} bytes") } else { format!("{v:.1} {}", units[u]) }
}
/// 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);
egui::Frame::new().fill(t.field_bg).corner_radius(6.0).inner_margin(egui::Margin::same(10)).show(ui, |ui| {
ui.set_width(ui.available_width());
ui.horizontal(|ui| {
let r = ui.radio(selected, RichText::new(name).font(Tokens::semibold(12.5)));
reg(app, &format!("settings.{prefix}.model.{id}"), &r, name);
if r.clicked() && !selected {
acts.push(Act::Run(format!("{prefix}.model.select"), json!({"id": id})));
}
let state = match (active, installed, selected) {
(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", "licenceUrl")] {
if let Some(url) = m[key].as_str() {
let r = ui.link(label);
if r.clicked() {
ui.ctx().open_url(egui::OpenUrl::new_tab(url));
}
}
}
});
if classic {
return;
}
ui.horizontal(|ui| {
if !installed && !web {
let r = ui.add_enabled(busy.is_none(), egui::Button::new("Download"));
reg(app, &format!("settings.{prefix}.download.{id}"), &r, "Download");
if r.clicked() {
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.{prefix}.install.{id}"), &r, "Install from File");
if r.clicked() {
acts.push(Act::PickModel(task));
}
}
if installed {
let r = ui.add_enabled(busy.is_none(), egui::Button::new("Remove"));
reg(app, &format!("settings.{prefix}.remove.{id}"), &r, "Remove");
if r.clicked() {
acts.push(Act::Run(format!("{prefix}.model.remove"), json!({"id": id})));
}
}
});
});
}
if let Some(b) = &busy {
ui.horizontal(|ui| {
ui.spinner();
ui.label(b);
});
ui.ctx().request_repaint_after(std::time::Duration::from_millis(250));
}
if let Some(e) = info["error"].as_str() {
ui.label(RichText::new(e).color(t.danger));
}
if web {
ui.label(RichText::new("Trained models are available in the desktop app; the browser uses the classic engine.").color(t.text_dim));
} else if let Some(f) = info["folder"].as_str() {
ui.label(RichText::new(format!("Models folder: {f}")).color(t.text_faint));
}
}
/// 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>) {
@ -419,6 +522,18 @@ pub fn show(app: &mut EffectcraftApp, ctx: &egui::Context, t: &Tokens) {
app.session.prefs_changed();
}
}
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, cmd, json!({"path": path}))
{
app.ui.status = e;
}
}
}
}
if cancel {

View File

@ -1048,6 +1048,15 @@ pub fn show(app: &mut EffectcraftApp, ui: &mut egui::Ui, rect: Rect) {
let tt = Tick::from_seconds_f64(secs);
app.session.set_time(tt);
app.stop();
// Ctrl/Cmd-drag scrubs the audio too (After Effects).
if rresp.dragged()
&& ui.input(|i| i.modifiers.command)
&& let Some(cid) = app.session.active_comp_id()
{
let now = ui.input(|i| i.time);
let t = app.session.time();
app.scrub_audio(cid, t, now);
}
}
// ---- column headers (right-click: show/hide columns; click the name header: Source/Layer Name).

View File

@ -520,6 +520,9 @@ pub fn show(app: &mut EffectcraftApp, ui: &mut egui::Ui, rect: Rect) {
let stale = app.viewer_shown.as_ref().is_none_or(|(_, k)| *k != key);
if stale {
show_frame(app, &ctx, key, img);
} else if let Some((_, k)) = app.viewer_shown.as_mut() {
// The same frame at a later revision (an edit elsewhere, an undo): it shows that one.
k.revision = key.revision;
}
}
if app.ui.viewer.transparency_grid {

View File

@ -0,0 +1,253 @@
//! Several Composition viewers (View ▸ New Viewer), as in After Effects. Viewer 0 is the
//! Composition panel, the others are `PanelKind::Viewer(n)` panels. One viewer is *active*: it
//! is the full interactive viewer and shows the active comp. The others show their own comp's
//! frame at that comp's current time; a click (or their tab) makes them active and their comp
//! the active one. New Viewer locks the viewer in use, so it keeps its comp. Opening a comp while
//! the active viewer is locked shows it in an unlocked viewer, or in a new one if every viewer is
//! locked.
use std::sync::Arc;
use effectcraft_engine::project::ItemId;
use egui::{Align2, Color32, Rect, Sense, pos2, vec2};
use serde_json::json;
use crate::EffectcraftApp;
use crate::dock::PanelKind;
use crate::frames::{FrameImage, FrameKey};
use crate::theme::Tokens;
/// A passive viewer's texture: the frame it shows.
pub enum PassiveTexture {
Cpu(egui::TextureHandle),
Gpu(egui::TextureId, Arc<effectcraft_gpu::DisplayFrame>),
}
impl PassiveTexture {
fn id(&self) -> egui::TextureId {
match self {
PassiveTexture::Cpu(h) => h.id(),
PassiveTexture::Gpu(id, _) => *id,
}
}
}
/// The dock panel of viewer `id` (0 = the Composition panel).
pub fn panel(id: u32) -> PanelKind {
if id == 0 { PanelKind::Composition } else { PanelKind::Viewer(id) }
}
/// The viewer a panel is, if it is one.
pub fn id_of(p: PanelKind) -> Option<u32> {
match p {
PanelKind::Composition => Some(0),
PanelKind::Viewer(n) => Some(n),
_ => None,
}
}
/// The active viewer's panel.
pub fn active_panel(app: &EffectcraftApp) -> PanelKind {
panel(app.ui.active_viewer)
}
pub fn locked(app: &EffectcraftApp, id: u32) -> bool {
app.ui.locked_tabs.contains(&panel(id).id())
}
/// Viewer `id`'s panel is docked or floating.
fn open(app: &EffectcraftApp, id: u32) -> bool {
let p = panel(id);
app.ui.dock.contains(p) || app.ui.floating.iter().any(|f| f.panels.contains(&p))
}
/// The comp viewer `id` shows: the active comp in the active viewer, else its own.
pub fn comp_of(app: &EffectcraftApp, id: u32) -> Option<ItemId> {
if id == app.ui.active_viewer {
return app.session.active_comp_id();
}
app.ui.viewers.get(&id).copied().flatten().map(ItemId).filter(|c| app.session.project.comp(*c).is_some())
}
/// Draw viewer `id`: the active one is the full viewer; another shows its comp's frame and
/// becomes the active viewer when clicked.
pub fn show(app: &mut EffectcraftApp, ui: &mut egui::Ui, id: u32, rect: Rect) {
if id == app.ui.active_viewer {
let c = app.session.active_comp_id().map(|c| c.0);
app.ui.viewers.insert(id, c);
super::viewer::show(app, ui, rect);
return;
}
let t = app.tokens;
let ctx = ui.ctx().clone();
let p = ui.painter().with_clip_rect(rect);
p.rect_filled(rect, 0.0, t.pasteboard);
let resp = ui.interact(rect, egui::Id::new(("passive-viewer", id)), Sense::click()).on_hover_text("Click to make this viewer active");
app.auto.add(&format!("viewers.{id}"), rect, "Composition viewer (click to make it active)");
match comp_of(app, id).and_then(|cid| Some((cid, app.session.project.comp(cid)?.clone()))) {
Some((cid, c)) => {
let (cw, ch) = ((c.width.max(1) as f64 * c.pixel_aspect.max(0.01)) as f32, c.height.max(1) as f32);
let fit = ((rect.width() - 24.0) / cw).min((rect.height() - 24.0) / ch).max(0.01);
let r = Rect::from_center_size(rect.center(), vec2(cw * fit, ch * fit));
// Render at the size it is drawn (eighths, so resizing reuses frames), at most 100%.
let scale = ((fit * ctx.pixels_per_point()) as f64 * 8.0).ceil().clamp(1.0, 8.0) / 8.0;
let frame = c.frame_rate.frame_at(app.session.time_of(cid));
app.request_frame(cid, frame, scale);
let key = app.frame_key(cid, frame, scale);
if let Some(img) = app.frames.get(&key) {
set_texture(app, &ctx, id, key, img);
}
match app.passive_tex.get(&id).filter(|(k, _)| k.comp == cid.0) {
Some((_, tex)) => {
p.image(tex.id(), r, Rect::from_min_max(pos2(0.0, 0.0), pos2(1.0, 1.0)), Color32::WHITE);
}
None => {
let bg = c.background;
p.rect_filled(r, 0.0, Color32::from_rgb((bg[0] * 255.0) as u8, (bg[1] * 255.0) as u8, (bg[2] * 255.0) as u8));
ctx.request_repaint_after(std::time::Duration::from_millis(50));
}
}
p.rect_stroke(r, 0.0, egui::Stroke::new(1.0, t.separator), egui::StrokeKind::Outside);
}
None => {
p.text(rect.center(), Align2::CENTER_CENTER, "Click to show the active composition here", Tokens::ui(12.0), t.text_faint);
}
}
if resp.clicked() {
activate(app, &ctx, id);
}
}
/// Keep the frame viewer `id` shows in its texture.
fn set_texture(app: &mut EffectcraftApp, ctx: &egui::Context, id: u32, key: FrameKey, img: FrameImage) {
if app.passive_tex.get(&id).is_some_and(|(k, _)| *k == key) {
return;
}
let opts = super::viewer::zoom_texture_options(app.session.prefs.viewer_zoom_smooth());
let tex = match (img, app.passive_tex.remove(&id)) {
(FrameImage::Cpu(img), Some((_, PassiveTexture::Cpu(mut h)))) if h.size() == img.size => {
h.set((*img).clone(), opts);
PassiveTexture::Cpu(h)
}
(FrameImage::Cpu(img), old) => {
free(app, old.map(|(_, t)| t));
PassiveTexture::Cpu(ctx.load_texture(format!("viewer-{id}"), (*img).clone(), opts))
}
(FrameImage::Gpu(f), old) => {
let Some(rs) = &app.wgpu else { return };
let view = f.texture.create_view(&Default::default());
let filter = if app.session.prefs.viewer_zoom_smooth() { eframe::wgpu::FilterMode::Linear } else { eframe::wgpu::FilterMode::Nearest };
let mut ren = rs.renderer.write();
let tid = match old {
Some((_, PassiveTexture::Gpu(tid, _))) => {
ren.update_egui_texture_from_wgpu_texture(&rs.device, &view, filter, tid);
tid
}
_ => ren.register_native_texture(&rs.device, &view, filter),
};
PassiveTexture::Gpu(tid, f)
}
};
app.passive_tex.insert(id, (key, tex));
}
/// Give a passive viewer's texture back (GPU textures are registered with the renderer).
fn free(app: &EffectcraftApp, tex: Option<PassiveTexture>) {
if let (Some(PassiveTexture::Gpu(id, _)), Some(rs)) = (tex, &app.wgpu) {
rs.renderer.write().free_texture(&id);
}
}
/// Make viewer `id` the active one: its comp becomes the active comp.
pub fn activate(app: &mut EffectcraftApp, ctx: &egui::Context, id: u32) {
if id == app.ui.active_viewer || !open(app, id) {
return;
}
let active = app.session.active_comp_id();
let target = comp_of(app, id);
app.ui.viewers.insert(app.ui.active_viewer, active.map(|c| c.0));
app.ui.active_viewer = id;
match target {
Some(c) if Some(c) != active => {
if let Err(e) = crate::menus::invoke(app, ctx, "comp.open", json!({"comp": c.0})) {
app.ui.status = e;
}
}
_ => {
app.ui.viewers.insert(id, active.map(|c| c.0));
}
}
}
/// A comp was opened (comp.open, a Timeline tab, a precomp double-click): it shows in the active
/// viewer, unless that one is locked to another comp: then in an unlocked viewer, or a new one.
pub fn on_open_comp(app: &mut EffectcraftApp, comp: ItemId) {
let av = app.ui.active_viewer;
let shown = app.ui.viewers.get(&av).copied().flatten().filter(|s| app.session.project.comp(ItemId(*s)).is_some());
if locked(app, av) && shown.is_some_and(|s| s != comp.0) {
let free = app.ui.viewers.keys().copied().find(|v| *v != av && !locked(app, *v) && open(app, *v));
app.ui.active_viewer = match free {
Some(v) => v,
None => add_viewer(app, av),
};
}
app.ui.viewers.insert(app.ui.active_viewer, Some(comp.0));
// The comp's viewer and its Timeline tab come to the front (unless locked).
for p in [active_panel(app), PanelKind::Timeline] {
if !app.ui.locked_tabs.contains(&p.id()) {
app.ui.dock.activate(p);
}
}
}
/// View ▸ New Viewer: a new viewer of the active comp next to the one in use, which is locked
/// (so it keeps its comp). Returns the new viewer's id.
pub fn new_viewer(app: &mut EffectcraftApp) -> u32 {
let prev = app.ui.active_viewer;
let active = app.session.active_comp_id().map(|c| c.0);
app.ui.viewers.insert(prev, active);
app.ui.locked_tabs.insert(panel(prev).id());
let id = add_viewer(app, prev);
app.ui.viewers.insert(id, active);
app.ui.active_viewer = id;
app.show_panel(panel(id));
id
}
/// Dock a new viewer panel as a tab next to viewer `near`.
fn add_viewer(app: &mut EffectcraftApp, near: u32) -> u32 {
let id = app.ui.viewers.keys().copied().max().unwrap_or(0).saturating_add(1).max(1);
app.ui.viewers.insert(id, None);
let p = panel(id);
let anchor = if app.ui.dock.contains(panel(near)) { panel(near) } else { PanelKind::Composition };
app.ui.dock.open_near(p, anchor);
id
}
/// A viewer panel was closed: forget it (the Composition panel keeps its slot), and if it was the
/// active viewer another open one takes over.
pub fn on_close(app: &mut EffectcraftApp, p: PanelKind) {
let Some(id) = id_of(p) else { return };
let tex = app.passive_tex.remove(&id).map(|(_, t)| t);
free(app, tex);
if id != 0 {
app.ui.viewers.remove(&id);
app.ui.locked_tabs.remove(&p.id());
}
if app.ui.active_viewer == id {
let next = std::iter::once(0).chain(app.ui.viewers.keys().copied()).find(|v| *v != id && open(app, *v)).unwrap_or(0);
app.ui.active_viewer = next;
// The viewer taking over keeps its comp (which becomes the active comp).
if let Some(c) = app.ui.viewers.get(&next).copied().flatten().map(ItemId).filter(|c| app.session.project.comp(*c).is_some())
&& app.session.active_comp_id() != Some(c)
{
app.session.open_comp(c);
}
}
}
/// The tab label of viewer `id` ("Composition Intro").
pub fn title(app: &EffectcraftApp, id: u32) -> Option<String> {
let c = comp_of(app, id)?;
Some(format!("Composition {}", app.session.project.item(c)?.name))
}

View File

@ -542,6 +542,13 @@ pub struct UiState {
/// bring a locked panel forward (the tab's lock icon, as After Effects' viewer lock).
#[serde(default)]
pub locked_tabs: BTreeSet<String>,
/// The comp each Composition viewer shows, by viewer id (0 = the Composition panel; see
/// `panels::viewers`).
#[serde(default)]
pub viewers: std::collections::BTreeMap<u32, Option<u64>>,
/// The active viewer (the interactive one, showing the active comp).
#[serde(default)]
pub active_viewer: u32,
pub focused: PanelKind,
pub viewer: ViewerState,
pub timeline: TimelineState,
@ -650,6 +657,8 @@ impl Default for UiState {
saved_floating: Default::default(),
maximized: None,
locked_tabs: BTreeSet::new(),
viewers: Default::default(),
active_viewer: 0,
focused: PanelKind::Composition,
viewer: ViewerState::default(),
timeline: TimelineState::default(),

View File

@ -0,0 +1,109 @@
//! Audio scrubbing: Ctrl/Cmd-dragging the current time plays one frame of the comp's mix at each
//! new frame (each snippet replacing what is still queued); holding still repeats nothing; silent
//! comps open no output; an idle scrub closes it.
use std::sync::{Arc, Mutex};
use effectcraft_engine::Session;
use effectcraft_engine::project::ItemId;
use effectcraft_engine::time::Tick;
use effectcraft_ui_egui::EffectcraftApp;
use effectcraft_ui_egui::audio::{AudioDevice, AudioFeed};
use egui::{Event, Modifiers, pos2};
use egui_kittest::Harness;
use serde_json::json;
/// An output that records the feed it is given and whether it was stopped.
#[derive(Clone, Default)]
struct Fake(Arc<Mutex<(Option<Arc<AudioFeed>>, bool)>>);
impl AudioDevice for Fake {
fn sample_rate(&self) -> u32 {
48000
}
fn start(&mut self, feed: Arc<AudioFeed>) -> Result<(), String> {
self.0.lock().unwrap().0 = Some(feed);
Ok(())
}
fn stop(&mut self) {
self.0.lock().unwrap().1 = true;
}
}
/// A 30 fps comp whose solid sounds a tone (Audio switch on).
fn setup(audible: bool, fake: &Fake) -> (EffectcraftApp, ItemId) {
let mut s = Session::default();
s.execute("comp.new", json!({"name": "A", "width": 64, "height": 36, "frameRate": 30, "duration": 2})).unwrap();
let cid = s.active_comp_id().unwrap();
let l = s.execute("layer.newSolid", json!({"color": "#808080"})).unwrap()["layer"].as_u64().unwrap();
if audible {
s.execute("effect.apply", json!({"layers": [l], "effect": "ec.audio.tone"})).unwrap();
s.execute("layer.setSwitch", json!({"layers": [l], "switch": "audio", "value": true})).unwrap();
}
let mut app = EffectcraftApp::new(s);
let f = fake.clone();
app.hooks.audio_device = Some(Box::new(move |_| Some(Box::new(f.clone()) as Box<dyn AudioDevice>)));
(app, cid)
}
#[test]
fn scrubbing_plays_one_frame_per_new_frame() {
let fake = Fake::default();
let (mut app, cid) = setup(true, &fake);
let t = |s: f64| Tick::from_seconds_f64(s);
app.scrub_audio(cid, t(0.5), 0.0);
let feed = fake.0.lock().unwrap().0.clone().expect("the output opened");
// One 30 fps frame at 48 kHz, and not silence.
assert_eq!(feed.queued_frames(), 1600);
let mut out = vec![0.0f32; 3200];
feed.pull(&mut out);
assert!(out.iter().any(|v| v.abs() > 0.01), "the tone is there");
// Faded in: the first sample is silent.
assert_eq!(out[0], 0.0);
// The same frame again plays nothing; the next frame replaces what is left.
app.scrub_audio(cid, t(0.51), 0.1);
assert_eq!(feed.queued_frames(), 0);
feed.push(&[0.0; 1000]);
app.scrub_audio(cid, t(0.6), 0.2);
assert_eq!(feed.queued_frames(), 1600, "replaced, not piled up");
// A silent comp opens nothing.
let silent = Fake::default();
let (mut quiet, qid) = setup(false, &silent);
quiet.scrub_audio(qid, t(0.5), 0.0);
assert!(silent.0.lock().unwrap().0.is_none() && quiet.scrub.is_none());
}
#[test]
fn cmd_dragging_the_current_time_scrubs_and_an_idle_scrub_closes() {
let fake = Fake::default();
let (app, _) = setup(true, &fake);
let mut h = Harness::builder().with_size(egui::vec2(1400.0, 900.0)).build_eframe(|_| app);
h.run_steps(3);
let cti = h.state().auto.find("timeline.cti").expect("current time indicator").clone();
let from = pos2(cti.rect[0] + cti.rect[2] / 2.0, cti.rect[1] + cti.rect[3] / 2.0);
let cmd = Modifiers::COMMAND;
h.event(Event::ModifiersChanged(cmd));
h.event(Event::PointerMoved(from));
h.step();
h.event(Event::PointerButton { pos: from, button: egui::PointerButton::Primary, pressed: true, modifiers: cmd });
h.step();
for k in 1..=4 {
h.event(Event::PointerMoved(from + egui::vec2(40.0 * k as f32, 0.0)));
h.step();
}
assert!(h.state().scrub.is_some(), "Cmd-drag opened the scrub output");
assert!(fake.0.lock().unwrap().0.as_ref().is_some_and(|f| f.queued_frames() > 0));
h.event(Event::PointerButton { pos: from + egui::vec2(160.0, 0.0), button: egui::PointerButton::Primary, pressed: false, modifiers: cmd });
h.step();
h.event(Event::ModifiersChanged(Modifiers::NONE));
h.step();
// An idle second closes the output.
for _ in 0..90 {
h.step();
std::thread::sleep(std::time::Duration::from_millis(15));
if h.state().scrub.is_none() {
break;
}
}
assert!(h.state().scrub.is_none() && fake.0.lock().unwrap().1, "closed and stopped");
}

View File

@ -0,0 +1,53 @@
//! RAM preview frames are keyed by content: editing one comp keeps the cached frames of another,
//! and undo finds the frames of the state it returns to.
use effectcraft_engine::Session;
use effectcraft_engine::project::ItemId;
use effectcraft_ui_egui::EffectcraftApp;
use egui_kittest::Harness;
use serde_json::json;
fn settle(h: &mut Harness<'_, EffectcraftApp>) {
for _ in 0..600 {
h.step();
if h.state().frames.inflight() == 0 && h.state().frames.last_ms.lock().map(|v| *v > 0.0).unwrap_or(false) {
break;
}
std::thread::sleep(std::time::Duration::from_millis(10));
}
h.run_steps(2);
}
fn shown_cached(h: &Harness<'_, EffectcraftApp>, comp: u64) -> bool {
let app = h.state();
let key = effectcraft_ui_egui::frames::FrameKey { frame: 0, ..app.shown_series(ItemId(comp)) };
app.frames.is_cached(&key)
}
#[test]
fn an_edit_in_another_comp_keeps_the_cached_frames_and_undo_finds_them_again() {
let mut s = Session::default();
let b = s.execute("comp.new", json!({"name": "B", "width": 64, "height": 36, "duration": 1})).unwrap()["comp"].as_u64().unwrap();
let a = s.execute("comp.new", json!({"name": "A", "width": 64, "height": 36, "duration": 1})).unwrap()["comp"].as_u64().unwrap();
s.execute("layer.newSolid", json!({"color": "#3080ff"})).unwrap();
let mut h = Harness::builder().with_size(egui::vec2(1400.0, 900.0)).build_eframe(|_| EffectcraftApp::new(s));
settle(&mut h);
assert!(shown_cached(&h, a), "A's frame rendered");
// Edit B (open it, add a layer), then come back to A: its frame is still in RAM.
h.state_mut().session.execute("comp.open", json!({"comp": b})).unwrap();
h.state_mut().session.execute("layer.newSolid", json!({"color": "#ff8030"})).unwrap();
settle(&mut h);
let b_frame = h.state().shown_series(ItemId(b));
h.state_mut().session.execute("comp.open", json!({"comp": a})).unwrap();
assert!(shown_cached(&h, a), "A's frame survived the edit in B");
// Edit A, then undo: the frame of the state undo returns to is found again.
let before = h.state().shown_series(ItemId(a));
h.state_mut().session.execute("layer.newSolid", json!({"color": "#20c040"})).unwrap();
settle(&mut h);
assert_ne!(h.state().shown_series(ItemId(a)), before, "A's content changed");
h.state_mut().session.execute("edit.undo", json!({})).unwrap();
assert_eq!(h.state().shown_series(ItemId(a)), before);
assert!(shown_cached(&h, a), "undo finds the earlier frame");
// B's frame is still there too.
assert!(h.state().frames.is_cached(&effectcraft_ui_egui::frames::FrameKey { frame: 0, ..b_frame }));
}

View File

@ -144,3 +144,25 @@ fn tab_drag_snapshot() {
std::fs::create_dir_all(dir).unwrap();
h.render().expect("render").save(format!("{dir}/tab-drag.png")).unwrap();
}
/// A precomp layer's bar shows its comp's markers; a double-click opens that comp at the marker.
#[test]
fn nested_comp_markers_show_on_the_precomp_bar() {
let (mut h, pre, main) = harness();
{
let s = &mut h.state_mut().session;
s.execute("comp.open", json!({"comp": pre})).unwrap();
s.execute("markers.set", json!({"new": true, "time": 1.0, "comment": "beat"})).unwrap();
s.execute("comp.open", json!({"comp": main})).unwrap();
}
h.run_steps(4);
let l = h.state().session.project.comp(ItemId(main)).unwrap().layers[0].id.0;
let id = format!("timeline.layer.{l}.nestedMarker.0");
let m = rect(&h, &id);
let e = h.state().auto.find(&id).unwrap().clone();
assert_eq!(e.label, "beat (marker in Pre)");
click_n(&mut h, m.center(), 2);
assert_eq!(h.state().session.active_comp_id(), Some(ItemId(pre)));
// (29.97 fps: the marker's frame is at 1.001 s)
assert!((h.state().session.time().seconds() - 1.0).abs() < 0.034, "{}", h.state().session.time().seconds());
}

View File

@ -0,0 +1,115 @@
//! View ▸ New Viewer: several Composition viewers. New Viewer locks the viewer in use; opening a
//! comp goes to an unlocked viewer (a new one when every viewer is locked); a click on another
//! viewer makes it the active one; closing a viewer forgets it.
use effectcraft_engine::Session;
use effectcraft_engine::project::ItemId;
use effectcraft_ui_egui::EffectcraftApp;
use effectcraft_ui_egui::dock::{PanelKind, Zone};
use effectcraft_ui_egui::panels::viewers;
use egui::{Event, pos2};
use egui_kittest::Harness;
use serde_json::{Value, json};
fn invoke(h: &mut Harness<'_, EffectcraftApp>, id: &str, p: Value) {
let ctx = h.ctx.clone();
effectcraft_ui_egui::menus::invoke(h.state_mut(), &ctx, id, p).unwrap();
h.run_steps(3);
}
fn click(h: &mut Harness<'_, EffectcraftApp>, id: &str) {
let e = h.state().auto.find(id).unwrap_or_else(|| panic!("no {id}")).clone();
let p = pos2(e.rect[0] + e.rect[2] / 2.0, e.rect[1] + e.rect[3] / 2.0);
h.input_mut().events.push(Event::PointerMoved(p));
h.step();
h.input_mut().events.push(Event::PointerButton { pos: p, button: egui::PointerButton::Primary, pressed: true, modifiers: Default::default() });
h.input_mut().events.push(Event::PointerButton { pos: p, button: egui::PointerButton::Primary, pressed: false, modifiers: Default::default() });
h.run_steps(3);
}
fn shows(h: &Harness<'_, EffectcraftApp>, viewer: u32) -> Option<u64> {
viewers::comp_of(h.state(), viewer).map(|c| c.0)
}
#[test]
fn new_viewers_lock_route_comps_and_activate_on_click() {
let mut s = Session::default();
let new =
|s: &mut Session, name: &str| s.execute("comp.new", json!({"name": name, "width": 64, "height": 36, "duration": 1})).unwrap()["comp"].as_u64().unwrap();
let c = new(&mut s, "C");
let b = new(&mut s, "B");
let a = new(&mut s, "A");
let mut h = Harness::builder().with_size(egui::vec2(1600.0, 1000.0)).build_eframe(|_| EffectcraftApp::new(s));
h.run_steps(3);
// View ▸ New Viewer: viewer 1 shows A too, viewer 0 (the Composition panel) is locked.
invoke(&mut h, "view.newViewer", json!({}));
assert!(h.state().ui.dock.contains(PanelKind::Viewer(1)));
assert!(viewers::locked(h.state(), 0) && !viewers::locked(h.state(), 1));
assert_eq!((h.state().ui.active_viewer, shows(&h, 0), shows(&h, 1)), (1, Some(a), Some(a)));
// Side by side, so both draw.
assert!(h.state_mut().edit_layout(|l| l.dock(PanelKind::Viewer(1), PanelKind::Composition, Zone::Right)));
h.state_mut().ui.dock.activate(PanelKind::Composition);
h.run_steps(3);
// Opening B shows it in the active (unlocked) viewer; the locked one keeps A.
invoke(&mut h, "comp.open", json!({"comp": b}));
assert_eq!((h.state().ui.active_viewer, shows(&h, 0), shows(&h, 1)), (1, Some(a), Some(b)));
assert!(h.state().auto.find("viewers.0").is_some(), "the other viewer draws passively");
// A click on viewer 0 makes it active, and A the active comp.
click(&mut h, "viewers.0");
assert_eq!(h.state().ui.active_viewer, 0);
assert_eq!(h.state().session.active_comp_id(), Some(ItemId(a)));
assert_eq!(shows(&h, 1), Some(b));
// Viewer 0 is locked to A: opening C goes to the unlocked viewer 1.
invoke(&mut h, "comp.open", json!({"comp": c}));
assert_eq!((h.state().ui.active_viewer, shows(&h, 0), shows(&h, 1)), (1, Some(a), Some(c)));
// Every viewer locked: opening B makes a new viewer.
h.state_mut().ui.locked_tabs.insert(PanelKind::Viewer(1).id());
invoke(&mut h, "comp.open", json!({"comp": b}));
assert_eq!(h.state().ui.active_viewer, 2);
assert!(h.state().ui.dock.contains(PanelKind::Viewer(2)));
assert_eq!((shows(&h, 0), shows(&h, 1), shows(&h, 2)), (Some(a), Some(c), Some(b)));
// Closing the active viewer forgets it; another open viewer takes over.
h.state_mut().close_panel(PanelKind::Viewer(2));
h.run_steps(2);
assert!(!h.state().ui.viewers.contains_key(&2));
// Viewer 0 takes over with its own comp (it is locked to A), which becomes the active comp.
assert_eq!(h.state().ui.active_viewer, 0);
assert_eq!((shows(&h, 0), h.state().session.active_comp_id()), (Some(a), Some(ItemId(a))));
// The tab names the comp each viewer shows.
let titles: Vec<String> = h
.state()
.auto
.elements
.iter()
.filter(|e| e.id.starts_with("panel.tab.Composition") || e.id.starts_with("panel.tab.Viewer"))
.map(|e| e.label.clone())
.collect();
assert!(titles.iter().any(|t| t == "Composition A") && titles.iter().any(|t| t == "Composition C"), "{titles:?}");
}
/// `cargo test -p effectcraft-ui-egui --test ui_viewers -- --ignored`: two viewers side by side
/// (`target/test-out/viewers.png`).
#[test]
#[ignore]
fn viewers_snapshot() {
let mut s = Session::default();
s.execute("file.openDemoProject", json!({})).unwrap();
let other = s.project.items.values().filter(|i| i.as_comp().is_some()).map(|i| i.id).find(|id| Some(*id) != s.active_comp_id()).unwrap();
let mut h = Harness::builder().with_size(egui::vec2(1600.0, 1000.0)).build_eframe(|_| EffectcraftApp::new(s));
h.run_steps(3);
invoke(&mut h, "view.newViewer", json!({}));
assert!(h.state_mut().edit_layout(|l| l.dock(PanelKind::Viewer(1), PanelKind::Composition, Zone::Right)));
h.state_mut().ui.dock.activate(PanelKind::Composition);
invoke(&mut h, "comp.open", json!({"comp": other.0}));
for _ in 0..300 {
h.step();
if h.state().frames.inflight() == 0 {
break;
}
std::thread::sleep(std::time::Duration::from_millis(20));
}
h.run_steps(4);
let dir = concat!(env!("CARGO_MANIFEST_DIR"), "/../../target/test-out");
std::fs::create_dir_all(dir).unwrap();
h.render().unwrap().save(format!("{dir}/viewers.png")).unwrap();
}

View File

@ -137,7 +137,15 @@ fn gpu_worker_renders_effect_frames_and_the_disk_cache_survives_a_reload() {
s.execute("prop.set", json!({"layer": "#1", "path": "effects/#1/blurriness", "value": 12})).unwrap();
let opts = RenderOpts { backend: effectcraft_engine::render::Backend::Gpu, ..Default::default() };
let t = effectcraft_engine::time::Tick::ZERO;
let key = |comp: effectcraft_engine::project::ItemId| FrameKey { revision: s.revision, comp: comp.0, frame: 0, scale: 1000, view: 0, opts: 0 };
let key = |comp: effectcraft_engine::project::ItemId| FrameKey {
revision: s.revision,
content: effectcraft_ui_egui::frames::comp_content(&s.project, comp),
comp: comp.0,
frame: 0,
scale: 1000,
view: 0,
opts: 0,
};
let disk = Arc::new(Mutex::new(HashMap::new()));
let worker = Arc::new(GpuWorker::new(&s, worker_gpu.clone(), disk.clone()));
let mut frames = Frames::default();

View File

@ -105,16 +105,24 @@ fn viewer_frames_render_remotely_from_project_diffs() {
let comp = s.active_comp_id().unwrap();
let c = s.project.comp(comp).unwrap().clone();
let opts = RenderOpts { scale: 0.25, guides: true, ..Default::default() };
let key = |rev: u64, f: i64| FrameKey { revision: rev, comp: comp.0, frame: f, scale: 250, view: 0, opts: 0 };
let key = |s: &Session, f: i64| FrameKey {
revision: s.revision,
content: effectcraft_ui_egui::frames::comp_content(&s.project, comp),
comp: comp.0,
frame: f,
scale: 250,
view: 0,
opts: 0,
};
let worker = Arc::new(FakeWorker::new(&s, 2));
let mut frames = Frames::default();
frames.set_remote(Some(worker.clone()));
assert!(frames.remote_active());
// Three prefetch frames, then the viewer's frame.
for f in 1..=3 {
frames.request(&source(&s), key(s.revision, f), comp, c.frame_rate.tick_of(f), opts);
frames.request(&source(&s), key(&s, f), comp, c.frame_rate.tick_of(f), opts);
}
frames.request_urgent(&source(&s), key(s.revision, 10), comp, c.frame_rate.tick_of(10), opts);
frames.request_urgent(&source(&s), key(&s, 10), comp, c.frame_rate.tick_of(10), opts);
assert_eq!(frames.inflight(), 4);
// Nothing renders until dispatched (no frame threads with a remote renderer); the viewer's
// frame goes first, then prefetch in request order.
@ -126,15 +134,15 @@ fn viewer_frames_render_remotely_from_project_diffs() {
// Frames arrive: the cache holds exactly what the page would have rendered.
worker.deliver();
assert_eq!(frames.remote_busy(), 0);
let img = match frames.get(&key(s.revision, 10)).expect("viewer frame cached") {
let img = match frames.get(&key(&s, 10)).expect("viewer frame cached") {
FrameImage::Cpu(c) => c,
FrameImage::Gpu(_) => unreachable!(),
};
let local = effectcraft_ui_egui::frames::to_color_image(&s.render(comp, c.frame_rate.tick_of(10), opts));
assert_eq!(img.size, local.size);
assert_eq!(img.pixels, local.pixels);
assert!(frames.is_cached(&key(s.revision, 1)));
assert!(!frames.is_cached(&key(s.revision, 2)));
assert!(frames.is_cached(&key(&s, 1)));
assert!(!frames.is_cached(&key(&s, 2)));
// The rest; an edit later travels as a patch.
frames.dispatch_remote();
worker.deliver();
@ -142,19 +150,19 @@ fn viewer_frames_render_remotely_from_project_diffs() {
worker.deliver();
assert_eq!(frames.inflight(), 0);
s.execute("layer.newSolid", json!({"color": "#ff8800"})).unwrap();
frames.request_urgent(&source(&s), key(s.revision, 10), comp, c.frame_rate.tick_of(10), opts);
frames.request_urgent(&source(&s), key(&s, 10), comp, c.frame_rate.tick_of(10), opts);
frames.dispatch_remote();
worker.deliver();
assert_eq!(worker.syncs.lock().unwrap().last(), Some(&"patch"));
let FrameImage::Cpu(img) = frames.get(&key(s.revision, 10)).unwrap() else { unreachable!() };
let FrameImage::Cpu(img) = frames.get(&key(&s, 10)).unwrap() else { unreachable!() };
let local = effectcraft_ui_egui::frames::to_color_image(&s.render(comp, c.frame_rate.tick_of(10), opts));
assert_eq!(img.pixels, local.pixels);
// A lost frame is released (requested again later), not cached.
frames.request(&source(&s), key(s.revision, 20), comp, c.frame_rate.tick_of(20), opts);
frames.request(&source(&s), key(&s, 20), comp, c.frame_rate.tick_of(20), opts);
frames.dispatch_remote();
worker.lose();
assert_eq!(frames.inflight(), 0);
assert!(!frames.is_cached(&key(s.revision, 20)));
assert!(!frames.is_cached(&key(&s, 20)));
// Without slots the remote renderer is unusable (frames render here again).
frames.set_remote(Some(Arc::new(FakeWorker::new(&s, 0))));
assert!(!frames.remote_active());

View File

@ -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
@ -295,6 +309,22 @@ state, job progress and, for a frame, the matte's area, centroid, RLE matte and
reference. Matte settings are ordinary properties (`effects/#1/rotoBrushMatte/searchRadius`,
`effects/#1/refineEdgeMatte/decontaminateEdgeColors`, …).
Trained model (Roto Brush 2.0 / 3.0):
- `roto.models` lists the built-in engine and every registered model, with its 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. 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","wait":true}' \
roto.stroke '{"layer":"#1","points":[[300,200],[360,230]],"radius":10}' roto.propagate '{"layer":"#1","wait":true}' --save
```
```sh
effectcraft-cli run clip.ecproj roto.stroke '{"layer":"#1","points":[[300,200],[360,230]],"radius":10}' \
roto.stroke '{"layer":"#1","kind":"bg","points":[[40,40],[600,40]],"radius":12}' \

View File

@ -22,6 +22,7 @@ cpal or muda. Everything in L0 to L4, the egui UI and the web app also build for
| L0 | `av1enc` | AV1 encoder from the AV1 bitstream specification: Main profile 8/10-bit 4:2:0, key + inter frames (quarter-pel motion), deblocking, for MP4 (`av01`) and WebM export |
| L1 | `raster` | Premultiplied float images, sampling, affine and projective warps, blurs, compositing (parallel with rayon) |
| L1 | `keyframe` | Animated values, keyframes with temporal ease and spatial Bezier, roving, hold, velocity |
| L1 | `segment` | Swappable trained 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 |
@ -334,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
@ -392,7 +397,10 @@ parameter with the base frame and segmentation span. Each frame is segmented by
(Boykov–Jolly hard constraints, GrabCut colour mixtures, our own Boykov–Kolmogorov max-flow,
coarse to fine) and propagated to the next frame by warping the matte with block optical flow and
re-cutting in a Search Radius band; frames with correction strokes are re-cut with the warped
matte as a soft prior. Refine Edge bands get guided-filter + closed-form matting and
matte as a soft prior. With Version 2.0 / 3.0 and a trained model chosen (`effectcraft_segment`,
Settings ▸ Roto Brush), the model's foreground probability, prompted by the strokes or by the
warped matte, becomes a strong prior for the same cut (falling back to the classic result when it
disagrees with the flow); the model's id is part of the chain seed. Refine Edge bands get guided-filter + closed-form matting and
foreground-colour decontamination. Segmentations are derived data: each frame's chain key (Input
Key, settings, strokes from the base frame out) indexes a process-wide cache that renders fill on
demand (`self_at`) and `roto.propagate` fills in the background; `Session::edit` keeps the Input

View File

@ -86,9 +86,10 @@ Points are logical window coordinates. A target is `{id}` (the element's centre;
Element ids are stable, for example `tools.Selection`, `panel.Timeline`, `panel.tab.EffectControls`,
`panel.tab.Timeline.<comp>` (one Timeline tab per open comp; `panel.tab.Timeline` is the shown
one), `header.workspace.Animation`, `project.item.<id>` (a double-click opens the item, Enter
one), `panel.tab.Viewer(<n>)` and `viewers.<n>` (View ▸ New Viewer's other Composition viewers;
`viewers.0` is the Composition panel while another viewer is active), `header.workspace.Animation`, `project.item.<id>` (a double-click opens the item, Enter
renames it), `viewer.comp`, `viewer.handle.<layer>.<i>`,
`timeline.layer.<id>.bar`, `timeline.layer.<id>.twirl`, `timeline.layer.<id>.row` (click selects,
`timeline.layer.<id>.bar`, `timeline.layer.<id>.nestedMarker.<i>` (a precomp's comp markers), `timeline.layer.<id>.twirl`, `timeline.layer.<id>.row` (click selects,
Enter or a double-click renames, drag reorders), `timeline.prop.<uid>.stopwatch`,
`timeline.key.<uid>.<frame>`, `timeline.cti`, `effectControls.prop.<uid>.value` and
`effects.item.<name>`. The Composition viewer adds `viewer.magnification`, `viewer.resolution`,

138
docs/gaps.md Normal file
View File

@ -0,0 +1,138 @@
# Where EffectCraft falls short
An honest assessment of how far EffectCraft is from being a real replacement for After Effects,
and the work that closes the gap. This document is meant for contributors and agents choosing
what to work on. The [ROADMAP](../ROADMAP.md) summarises it; [parity.md](parity.md) is the
feature-by-feature checklist it builds on.
*Assessed 5 October 2026. Estimates marked "≈" are judgements from the evidence listed, not
measurements. Update this file when the evidence changes.*
## Two different questions
[parity.md](parity.md) reports **≈ 99%**. That number answers one question: *does each After
Effects feature exist?* It scores a 92-item catalogue we wrote ourselves, graded by the same agents
that built the features, with partial features counting half. As a measure of breadth it is fair,
and the breadth is real: every one of After Effects' effects exists by name, along with the
menus, panels, 3D, tracking, expressions, scripting and export.
It does not answer the question users care about: *can someone who uses After Effects for a living
do real work in EffectCraft?* Nobody has measured that yet. Our best estimate is **≈ 30–50%**,
limited mainly by four things: projects we can't open, behaviour nobody has checked against After
Effects, reliability on platforms other than macOS, and the third-party plug-ins professional
projects depend on.
The checklist is also internally inconsistent: parity.md's per-area table still shows Interface
75%, Project 68%, Animation 70% and Paint 0% next to the 99% headline. Some rows are stale; none
of them have been reconciled.
## By dimension
| Dimension | Estimate | Evidence |
|---|---|---|
| Breadth of features | ≈ 95%+ | All 306 effects, the After Effects menus and panels, Classic and Advanced 3D, tracking, expressions, scripting, render queue (parity.md) |
| Behaves like After Effects | unmeasured, ≈ 60–80% | One outside contributor found four bugs in features marked done within a day (PRs #6–#10): Hold keyframes eased the motion into them, `keyInSpatialTangent` / `keyOutSpatialTangent` had the wrong names, a zero frame rate crashed, Find and Enter Full Screen shared Ctrl+F off macOS. There are about 10 After Effects reference captures in total. No test compares our renders or property values with After Effects itself; the tests check what the agents believed After Effects does |
| Opening existing After Effects work | ≈ 0% | `.aep` / `.aepx` projects cannot be opened. Third-party After Effects plug-ins cannot run. Expressions and the scripting object model are strong, so scripts and expressions carry over |
| Stability | improving | 23 never-crash PRs landed on 4 October; [AGENTS.md](../AGENTS.md) "Never crash" now binds every crate. On 5 October `cargo xtask ci` failed on main under Rust 1.99's clippy (a fix is in progress) |
| 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% | 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
Each workstream lists what "done" means, so progress can be measured rather than self-graded.
### G1. Measure fidelity against After Effects
The most important missing piece: it turns every other estimate here into a measurement.
- Build a corpus of original test projects (our own content, no Adobe assets), one or more per
feature id in `plan/aftereffects/feature-catalog.md`: keyframe interpolation of every kind,
blend modes, track mattes, each effect at default and at non-default settings, text animators,
expressions, 3D, motion blur, time remapping.
- Drive After Effects through ExtendScript (see CLAUDE.md) to record property values at sampled
times and render reference frames; render the same projects in EffectCraft headless.
- Compare values exactly and frames with a perceptual metric; publish a per-feature fidelity
score and a fidelity column in parity.md.
- Clean room: After Effects output stays local in `plan/aftereffects/ref/` (gitignored). Commit
only our projects, the harness and the scores, never After Effects frames.
- Done when: every P0 and P1 feature has at least one corpus project, scores are reproducible from
one command, and parity.md reports breadth and measured fidelity side by side.
### G2. Real-user reliability on every platform
- Fix the open user issues (#41–#47), each with a regression test.
- Run the headless snapshot and control-channel checks on Linux (X11 and Wayland) and Windows, not
only macOS. Interactions that only fail with real input (docking drags, viewer pan, drag-and-drop
import, inline rename) need scripted input tests.
- Done when: no open bug blocks a basic workflow (import, arrange, animate, preview, render) on
any of the three desktop platforms, and every bug users report gets triaged within a day.
### G3. Stability and a green main
- Keep `cargo xtask ci` green on the current stable toolchain. A toolchain release that breaks the
gate is a P0.
- Fuzz the inputs we don't control: project files, imported media and documents, expressions,
scripts, control-channel and MCP requests. Every crash found becomes a regression test
(AGENTS.md "Never crash").
- Done when: the gate is green on stable and the fuzzers run regularly without new crashes.
### G4. Open After Effects projects
- `.aep` / `.aepx` import is the largest single barrier to switching. It needs an **owner
decision** on clean-room scope (whether and how the format may be studied), like the open
`.prproj` question in `plan/STATUS.md`.
- Also in this area: relative footage paths and relinking when a project moves (parity.md, Project).
- Done when: the decision is recorded and, if approved, a corpus of real-world-shaped projects
opens with its comps, layers, keyframes, effects and expressions intact.
### G5. Performance at real-world scale
- Benchmark projects at 1080p and 4K with many layers, heavy effects, long footage and nested
comps; record preview frame times, RAM preview fill and render times on reference machines.
- Compare with After Effects on the same machine where possible.
- Done when: benchmark numbers are tracked over time and regressions fail a check.
### G6. Media depth
- Encoder efficiency: B-frames and multi-reference for HEVC / AV1, SAO / CDEF / restoration, Opus
FEC / DTX.
- Verify decode of the camera and phone formats people actually bring, including variable frame
rate; document what is not supported.
### G7. Learned models for Roto Brush and face tracking
- Decided on 5 October 2026 for Roto Brush: licensed weights with pure-Rust inference, if the
impact on the system is low, in a composable, swappable module. Done in M13.35: the
`effectcraft-segment` crate (a `MaskModel` interface, a registry that only accepts
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: 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
- After Effects SDK plug-ins cannot run in EffectCraft. Our own WebAssembly plug-in API exists
([plugins.md](plugins.md)). Grow it: documentation, examples, and original effects that cover
what the most common third-party plug-ins are used for (particles, glows, 3D objects, sabers).
### G9. Reach
- Localisation of the interface, accessibility (keyboard navigation, screen readers, contrast),
user documentation and tutorials.
## For agents choosing work
1. Read `plan/STATUS.md` for owner blockers and running work, then this file.
2. Prefer G1–G3 over new features. A feature that exists but behaves differently from After
Effects is not done.
3. When you fix a behaviour bug in a feature parity.md marks as done, add it to the evidence
above. That is how this assessment stays honest.
4. Don't raise parity.md's numbers without evidence: a G1 score, a user-facing check on more than
one platform, or both.

View File

@ -2,6 +2,14 @@
How close EffectCraft is to After Effects 2026, feature by feature, and how much work is left.
> **Read this as a measure of breadth, not of real-world readiness.** It checks whether each
> feature exists, using our own catalogue, graded by the agents that built it. It does not measure
> whether the feature behaves like After Effects (nothing compares the two yet), whether After
> Effects projects open (they don't), or how the app holds up for real users on every platform.
> Those gaps, and the work that closes them, are in [gaps.md](gaps.md). Don't raise the numbers
> here without the evidence gaps.md asks for. The per-area table further down has stale rows that
> have not been reconciled with the headline.
## Current status (audit at commit `d39c0e8`, 4 October 2026; updated 5 October for M3.9–M4.12)
| Measure | Value |
@ -20,8 +28,8 @@ Partial features:
| id | Tier | Done | What is missing |
|---|---|---|---|
| EFF-5 GPU effects | P1 | 0.98 | Fractal (its escape iteration needs f64: WGSL has none and Metal's fast math defeats double-f32 emulation) and the analysis / tool effects (Camera-Shake Deblur, Detail-preserving Upscale, Puppet, Rolling Shutter Repair, Warp Stabilizer, Mocha Shape, Roto Brush, Match Grain, Paint, Camera Tracker, Face Tracker) render on the CPU. M13.28 moved the remaining pixel effects (CC blur / glass family, Color Link, Cineon, HDR, Inner/Outer Key, Basic 3D…) and the geometry and audio generators (Lightning, Advanced Lightning, Beam, Lens Flare, Radio Waves, Vegas, Stroke, Scribble, Write-on, Paint Bucket, Eyedropper Fill, CC Glue Gun, CC Threads, Audio Spectrum / Waveform, Basic / Path Text) onto the GPU; M13.29 the particle effects (CC Particle World / Systems II, Particle Playground, CC Ball Action, CC Pixel Polly, CC Scatterize) and Curl Noise, and removed the Shatter wireframe, Foam extras, Timewarp matte, strong Warp Fisheye / Twist and custom `.ocio` fallbacks. The 13 audio filters stay on the CPU by design |
| MSK-4 Roto Brush | P2 | 0.85 | Segmentation is classical (graph cut + optical flow), not a learned model |
| TRK-3 Face tracking | P2 | 0.9 | Classical fitter: weak on profile and occluded faces; Rolling Shutter Ripple is approximated |
| 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 | 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
@ -163,11 +171,11 @@ the disk cache in the formats wave; the 3D Camera Tracker in M12.6.)
| Web | 98% | 0.1 | No shared-memory threads inside one engine instance (the decided design is one engine instance per worker; a threaded build needs nightly `build-std`) (M13.30: Render Queue, analyses and Content-Aware Fill render on each job worker's own WebGPU device in passes, particles and Advanced 3D read back under keys, layer buffers in the browser's disk cache; browser storage, Web Audio, Web Worker renders/analyses, WebGPU viewer and offline install landed in M15.2; viewer frames in frame workers fed by project diffs, Roto Brush propagation in a worker, non-blocking `wait: true` jobs and a browser Media Browser (File System Access folders, browser storage) in M13.10; WEB-1 in M13.24: GPU effects and the GPU compositor in the frame workers on their own WebGPU devices with deferred readbacks (frames render in passes; Backend Auto per comp), the disk cache in the Origin Private File System (written by the workers with sync access handles, LRU under the settings' limit, served after a reload), the storage manager (Settings ▸ Disk ▸ Browser Storage, `storage.info` / `storage.persist` / `storage.clear`), and a WGSL constant Chrome rejected (it disabled every GPU kernel in Chrome) fixed) |
| 3D | 88% | 5.5 | multi-view layouts, the Extended Viewer for Advanced 3D comps (Classic 3D Extended Viewer landed in M13.5 UI completion); collapsed precomps of another size seen through the parent's camera render (fixed in M13.2); stereo rigs, orbit nulls, lights controlled by the camera, cameras/lights from glTF models, environment backgrounds, Advanced 3D motion blur, blend modes and track mattes landed in M7.7; Classic 3D iris-shaped bokeh with highlights, progressive depth of field on tilted layers and the focus-link commands landed in M13.5; Advanced 3D depth of field with the iris and highlight options, collapsed precomps as real Advanced 3D geometry and extruded text/shape strokes landed in M13.8; Advanced 3D (glTF/OBJ models, primitives, extruded text and shapes, PBR, image-based light, shadow maps, GPU rasteriser) in M7.4–M7.6; Advanced 3D end to end on the GPU (motion blur, iris depth of field, compositing) in M13.11 |
| Effects | ≈ 85% | 3.0 | GPU versions of the remaining effects (280 run on the GPU since M13.29, EFF-5: M13.28 added the remaining pixel effects and the geometry and audio generators (Lightning, Advanced Lightning, Beam, Lens Flare, Radio Waves, Vegas, Stroke, Scribble, Write-on, Paint Bucket, Eyedropper Fill, CC Glue Gun, CC Threads, Audio Spectrum / Waveform, Basic / Path Text); M13.29 added the particle effects (CC Particle World / Systems II, Particle Playground, CC Ball Action, CC Pixel Polly, CC Scatterize, Curl Noise) and removed the fallbacks of Shatter's wireframe views, Foam's extras, Timewarp's Matte Layer, strong Warp Fisheye / Twist bends and custom `.ocio` configs; M13.23 added the 3D Channel and Immersive Video families, Apply Color LUT, the OCIO effects, Color Profile Converter and the simulations' render passes (Card Wipe included); M13.22 added the CC light family, the CC transitions, Block Dissolve, Radial Shadow, CC Bender / Blobbylize / Cylinder / Sphere / Spotlight / Environment, 3D Glasses, Numbers, Timecode and the time effects joined the 166 of M13.13; Fractal and the analysis / tool effects still render on the CPU) and the missing controls listed as partial in [effects.md](effects.md) (every After Effects effect exists since M9.11, M12.5 and M12.6; parameter names, order, twirl-downs, popups, units and defaults were aligned in M9.12) |
| Interface | 75% | 5.0 | more Learn tutorials and pixel-level fidelity of dialogs; View ▸ New Viewer (several unlocked viewers) (an automation id on every control of Composition and Solid Settings, Group / Ungroup Shapes in the Layer menu, Window ▸ Learn, the flowchart shortcuts and Reveal in Explorer landed in M3.10 and M3.13; the Home ▸ Templates gallery (eight original built-in templates, user templates from File ▸ Save as Template…) and View ▸ Simulate Output ▸ My Custom RGB… landed in M13.25; Timeline layer reordering by drag, a non-snapping viewer pan, a working rename field and twirl arrows, and the full set of property reveal shortcuts — double presses, Alt+Shift keyframes, Ctrl+` — landed in M13.17–M13.20; the Home ▸ Learn tab with interactive tutorials and a UI fidelity pass landed in M13.5 UI completion; visual editors for Lumetri RGB / hue-saturation curves, Colorama's output cycle wheel, Glow's colour map and Reshape's correspondence points (viewer handles), a scrolling Preview panel and AE-style Composition / Timeline tabs (close, label swatch, viewer lock) landed in M13.10; Timeline outline and Project panel columns scroll horizontally, the Layer Style dialog, ROI resize handles, Pan Behind snapping and 3D Reference Axes landed in M13.5; native macOS menu bar, Timeline columns/search/reveal-add, Home screen with recent projects and all AE workspaces landed; viewer rulers/snapping/channels/snapshots landed in M0.13) |
| Interface | 75% | 5.0 | more Learn tutorials and pixel-level fidelity of dialogs (View ▸ New Viewer landed in M13.32; an automation id on every control of Composition and Solid Settings, Group / Ungroup Shapes in the Layer menu, Window ▸ Learn, the flowchart shortcuts and Reveal in Explorer landed in M3.10 and M3.13; the Home ▸ Templates gallery (eight original built-in templates, user templates from File ▸ Save as Template…) and View ▸ Simulate Output ▸ My Custom RGB… landed in M13.25; Timeline layer reordering by drag, a non-snapping viewer pan, a working rename field and twirl arrows, and the full set of property reveal shortcuts — double presses, Alt+Shift keyframes, Ctrl+` — landed in M13.17–M13.20; the Home ▸ Learn tab with interactive tutorials and a UI fidelity pass landed in M13.5 UI completion; visual editors for Lumetri RGB / hue-saturation curves, Colorama's output cycle wheel, Glow's colour map and Reshape's correspondence points (viewer handles), a scrolling Preview panel and AE-style Composition / Timeline tabs (close, label swatch, viewer lock) landed in M13.10; Timeline outline and Project panel columns scroll horizontally, the Layer Style dialog, ROI resize handles, Pan Behind snapping and 3D Reference Axes landed in M13.5; native macOS menu bar, Timeline columns/search/reveal-add, Home screen with recent projects and all AE workspaces landed; viewer rulers/snapping/channels/snapshots landed in M0.13) |
| Project | ≈ 68% | 5.5 | OCIO displays beyond the built-in tone map; relative footage paths and relinking when a project moves (unsaved-changes prompts, a modified mark that follows undo and project files that name their version and always reopen landed in M3.9 and M3.14; auto-save and crash recovery in M3.7; Color Engine with OCIO/ACES working spaces, HDR compand/tone mapping, Rec. 2100 PQ/HLG output, Feet + Frames, display colour management, Simulate Output and the locked viewer landed in M7.7; proxies and Interpret Footage fields / pixel aspect / alpha guess landed: PRJ-8, PRJ-3) |
| Masks & roto | 74% | 5.0 | Roto Brush's learned (3.0) segmentation model (variable-width mask feather points with the Mask Feather tool landed in M13.5; mask tracking and Mask Interpolation landed in M6.6; Roto Brush & Refine Edge with graph-cut segmentation, flow propagation, edge matting, decontamination and Freeze in M6.7) |
| Masks & roto | 74% | 5.0 | (audit estimates; Roto Brush 2.0 / 3.0's trained model, MobileSAM in the swappable `segment` module, has since landed in M13.35; variable-width mask feather points with the Mask Feather tool landed in M13.5; mask tracking and Mask Interpolation landed in M6.6; Roto Brush & Refine Edge with graph-cut segmentation, flow propagation, edge matting, decontamination and Freeze in M6.7) |
| Preview | 77% | 2.8 | GPU effects for the audio-driven effects and the stepped particle systems; content-keyed RAM preview frames (an edit drops every comp's), audio scrubbing, more snapshot slots (RAM preview keys with the render options, least-recently-shown eviction, Auto resolution while playing, Cache Before Playback that fits, an achieved-fps readout, Cache Frames When Idle and purges that clear what they name landed in M4.5–M4.8; 3D channel, VR, OCIO, Card Wipe and the simulations' render passes run on the GPU since M13.23; Advanced 3D layers with blend modes / track mattes / Preserve Transparency and environment backgrounds composite on the GPU since M13.22; pooled GPU textures and fused quantisation since M13.13: Lower Third GPU warm 21 → 4 ms/frame, and Auto now picks the CPU or the GPU per comp from measured frame times; Advanced 3D runs — raster, motion blur, iris depth of field, compositing — and wireframes run on the GPU since M13.11: CPU warm 3394 → GPU warm 290 ms/frame at 1920×1080 (11.7×; Half 968 → 85 ms), GPU ≠ CPU on 0.005 % of pixels, on an M4 Pro under load; Classic 3D bokeh depth of field — iris shapes, highlights, fringe, progressive blur on tilted planes — runs in WGSL since M13.6: 3D Showcase GPU warm 205 → 54 ms/frame at full size on an M4 Pro; the Preview panel's five shortcuts with their own Include / Loop / Cache Before Playback / Range / Play From / Frame Rate / Skip / Resolution / Full Screen / stop options and `playback.settings.get/set` landed in M13.8; Classic 3D runs and adjustment layers composite on the GPU since M12.7; persistent disk cache with the blue cache bar landed; region of interest, snapshots, exposure and Fast Previews landed in M0.13) |
| Tracking | ≈ 96% | 0.3 | face tracking is a classical (skin model + feature components + shape model) fitter, not a learned detector: profile views and occluded faces are weak; Rolling Shutter Ripple is approximated by Subspace Warp's mesh density (face tracking (Outline Only / Detailed Features with Face Track Points and Extract & Copy Face Measurements), Subspace Warp's content-preserving mesh warp on subspace-smoothed trajectories, and radial lens distortion (k1, k2) in the camera bundle adjustment with Undistort Footage landed in M13.3; Rolling Shutter Repair in M9.11; point tracker, mask tracking, Warp Stabilizer and the 3D Camera Tracker in M6.x / M12.5 / M12.6) |
| 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
@ -422,6 +430,78 @@ layers); the Graph Editor has Auto-Select Graph Type (the default: the speed gra
spatial properties are shown, else the value graph; choosing Value or Speed turns it off). Graph
Editor keys are small squares whatever their interpolation, as in After Effects.
Since (M13.34): audio scrubbing. Ctrl-dragging the current-time indicator (Cmd on macOS) plays
one frame (30–100 ms) of the comp's mix at each new frame, as in After Effects. Snippets fade in
and out and replace what is still queued, so the sound follows the pointer without lagging.
Holding still repeats nothing, the Audio panel meters follow, and the output closes after an idle
second. Agents call `playback.scrubAudio {time}`.
Since (M13.33): a precomp layer's bar shows its comp's markers, as outlined, read-only markers.
They are mapped through the layer's timing (start, stretch, time remapping, Responsive Design
stretch) and shown within its In–Out range. Hovering names them ("beat (marker in Pre)"), and a
double-click opens the nested comp at the marker. Agents read the same list with
`markers.nested {layer}`.
Since (M13.35): Roto Brush can use a trained segmentation model. MSK-4 keeps its score until G1
measures it against After Effects ([gaps.md](gaps.md)).
- The model sits behind a swappable interface: `effectcraft_segment::MaskModel` plus a registry
where every entry must have an open-source licence compatible with ours, a source URL, a size
and a SHA-256.
- The first model is **MobileSAM** (Apache-2.0: a TinyViT-5M encoder with Segment Anything's
decoder), running in plain Rust (rayon and ndarray's safe GEMM, no ML framework). It reads the
official `mobile_sam.pt` checkpoint directly and checks it against a pinned SHA-256.
- Roto Brush **Version 2.0 / 3.0** uses the model chosen in **Settings ▸ Roto Brush**; Version
1.0 and an uninstalled model use the classic graph cut.
- On the base frame, points along the strokes prompt it. When propagating, the flow-warped matte
prompts it (its box, points deep inside it and the matte itself). The model's probability map
becomes a strong prior for the same graph cut, so strokes stay hard constraints and edges still
snap to colour. A model whose mask disagrees with the flow (IoU under 0.5) is ignored for that
frame.
- Weights are never bundled. Download (verified; uses the system `curl`) or Install from File
puts them in the settings' `models` folder, and they load in the background.
- On the moving-disc test sequence it scores IoU 0.989 on the base frame and at worst 0.980 over
20 propagated frames (classic: 0.973). The encoder takes about 0.7 s a frame on a 32-thread
CPU, 1.3 s on 4 threads; its embedding is cached, so new strokes on the same frame only rerun
the decoder (about 0.1 s).
- Agents use `roto.models`, `roto.model.select`, `roto.model.download`, `roto.model.install` and
`roto.model.remove`. The browser build uses the classic engine for now.
Since (M13.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.
The others show their own comp's frame at that comp's current time, and a click (or their tab)
makes them active. Opening a comp while the active viewer is locked shows it in an unlocked
viewer, or in a new one if every viewer is locked. Closing a viewer hands over to another one,
with that viewer's comp. Each viewer's tab is named after its comp, with its own lock.
Since (M4.14): RAM preview frames are keyed by their comp's content, not the project revision. An
edit keeps the cached frames (the green bar) of every comp it doesn't touch, and undo finds the
frames of the state it returns to, as in After Effects. The identity costs about 0.1 ms on the
1500-layer test project: it hashes the addresses of the comps a comp draws, and the RAM cache
keeps a project snapshot for every identity with frames, so those addresses can't be reused
meanwhile. Footage, solids, proxies and the project settings are hashed by value. At most 64
states are kept, so a long drag doesn't hold on to every intermediate state. The disk cache's
content key now also covers proxies and Use Proxy.
Since (M13.30–M13.31): the Tools bar and About dialog carry the ArtCraft mark. The
Home screen is laid out like After Effects' and covers the whole workspace. A left rail holds New
Project / Open Project, the Home, Templates and Learn pages and, at its foot, the community
@ -732,9 +812,10 @@ are pinned bit for bit by golden hashes (`crates/effects/tests/particle_golden.r
15. ~~Preferences and a shortcut editor that can rebind; real Wiggler, Smoother and Motion Sketch; the marker dialog.~~ (second and third waves)
16. ~~Motion blur of collapsed precomps and of animated content inside a layer (shape paths, text
animators).~~ (M4.13)
17. Content-keyed RAM preview frames, so an edit keeps the frames of comps it doesn't touch.
18. View ▸ New Viewer (several unlocked Composition viewers), nested comp markers on the
precomp layer bar, audio scrubbing.
17. ~~Content-keyed RAM preview frames, so an edit keeps the frames of comps it doesn't touch.~~
(M4.14)
18. ~~View ▸ New Viewer (several Composition viewers)~~ (M13.32); ~~nested comp markers on the
precomp layer bar~~ (M13.33), ~~audio scrubbing~~ (M13.34).
At the original audit the engine underneath (keyframes, expressions, shape operators, text
animators, Classic 3D, all 38 blend modes, the render queue) was already deep, and most of what was

View File

@ -139,6 +139,13 @@ Commands (CLI, MCP, control channel):
the Timeline's expression editor
- `scripting.errorBanner`: a banner along the bottom of the Composition panel names the first
failing expression (click: reveal it)
- `roto.model` (Settings ▸ Roto Brush): the segmentation model Roto Brush 2.0 / 3.0 use,
`classical` (built in) or `mobilesam` (once installed). The page lists every registered model
with its 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

View File

@ -32,6 +32,7 @@ const LAYERS: &[(&str, u8)] = &[
("keyframe", 1),
("path", 1),
("psd", 1),
("segment", 1),
("project", 2),
("text", 2),
("effects", 2),