diff --git a/Cargo.lock b/Cargo.lock index 86d4bca6..3fc4e721 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2296,6 +2296,15 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "heic-rs" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adf9dde634bb6d5b718bacd01bef80192f2334306b2f785519a0136cca7b9618" +dependencies = [ + "rayon", +] + [[package]] name = "hermit-abi" version = "0.5.3" @@ -2837,6 +2846,7 @@ dependencies = [ "jxl-oxide", "lightcraft-color", "lightcraft-geom", + "lightcraft-heif", "lightcraft-raster", "moxcms 0.9.1", "png", @@ -2882,6 +2892,8 @@ dependencies = [ "lightcraft-codecs", "lightcraft-color", "lightcraft-develop", + "lightcraft-faces", + "lightcraft-fetch", "lightcraft-geom", "lightcraft-gpu", "lightcraft-merge", @@ -2903,6 +2915,20 @@ dependencies = [ "web-time", ] +[[package]] +name = "lightcraft-faces" +version = "0.4.0" +dependencies = [ + "matrixmultiply", + "ndarray", + "proptest", + "serde", + "serde_json", + "sha2", + "thiserror 2.0.21", + "web-time", +] + [[package]] name = "lightcraft-fetch" version = "0.4.0" @@ -2938,6 +2964,13 @@ dependencies = [ "wgpu", ] +[[package]] +name = "lightcraft-heif" +version = "0.4.0" +dependencies = [ + "heic-rs", +] + [[package]] name = "lightcraft-mcp" version = "0.4.0" @@ -2986,6 +3019,7 @@ dependencies = [ "lightcraft-develop", "lightcraft-geom", "lightcraft-raster", + "lightcraft-raw", "lightcraft-scenes", "proptest", "rayon", @@ -3022,6 +3056,7 @@ name = "lightcraft-raw" version = "0.4.0" dependencies = [ "jpeg-encoder", + "jxl-oxide", "lightcraft-color", "lightcraft-geom", "lightcraft-meta", @@ -3035,6 +3070,7 @@ dependencies = [ "thiserror 2.0.21", "zune-core", "zune-jpeg", + "zune-jpegxl", ] [[package]] @@ -3084,6 +3120,7 @@ dependencies = [ "lightcraft-color", "lightcraft-develop", "lightcraft-engine", + "lightcraft-faces", "lightcraft-geom", "lightcraft-meta", "lightcraft-pipeline", @@ -3174,6 +3211,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 = "maybe-rayon" version = "0.1.1" @@ -3327,6 +3374,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" @@ -4545,6 +4607,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" @@ -6752,6 +6820,15 @@ dependencies = [ "zune-core", ] +[[package]] +name = "zune-jpegxl" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cba11ecd07cc23351500e840ae03841b7e4997923ec8e4832ee3ae199ea0b6b4" +dependencies = [ + "zune-core", +] + [[package]] name = "zvariant" version = "5.15.0" diff --git a/Cargo.toml b/Cargo.toml index 77b20666..058da8f9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,6 +12,7 @@ repository = "https://github.com/storytold/lightcraft" [workspace.dependencies] # internal lightcraft-geom = { path = "crates/geom" } +lightcraft-faces = { path = "crates/faces" } lightcraft-sysmem = { path = "crates/sysmem" } lightcraft-color = { path = "crates/color" } lightcraft-raster = { path = "crates/raster" } @@ -30,10 +31,12 @@ lightcraft-engine = { path = "crates/engine" } lightcraft-ui-egui = { path = "crates/ui-egui" } lightcraft-mcp = { path = "crates/mcp" } lightcraft-fetch = { path = "crates/fetch" } +lightcraft-heif = { path = "crates/heif" } lightcraft-segment = { path = "crates/segment" } # foundation thiserror = "2" +sha2 = "0.10" serde = { version = "1", features = ["derive", "rc"] } serde_json = { version = "1", features = ["float_roundtrip"] } log = "0.4" @@ -61,6 +64,7 @@ tiff = "0.11" image-webp = "0.2" moxcms = { version = "0.9", default-features = false, features = ["lut"] } jxl-oxide = { version = "0.12", default-features = false } +zune-jpegxl = { version = "0.5", default-features = false, features = ["std"] } # MIT/Apache-2.0/Zlib lossless JXL encoder (test fixtures only) ravif = { version = "0.13", default-features = false, features = ["threading"] } # no "asm": it needs nasm (non-Rust) # gpu (crates/gpu; native backends only, no GL: everything stays pure Rust) diff --git a/README.md b/README.md index 5aacbee0..e93683c6 100644 --- a/README.md +++ b/README.md @@ -260,7 +260,8 @@ lightcraft --control 7980 ~/Pictures/trip - **Pure Rust, no C.** Our own RAW decoders (DNG, Canon CR2/CR3, Sony ARW, Nikon NEF, Fujifilm RAF incl. X-Trans, Panasonic RW2 / Leica RWL, Pentax PEF, Olympus ORF), our own colour science, our own pipeline. JPEG, PNG, TIFF, WebP, - PSD composites and JPEG XL open today. + PSD composites and JPEG XL open today; HEIC/HEIF photos too with `--features heif` (HEVC is a build-time choice, + as in PhotoCraft). - **Scene-referred & wide-gamut.** Linear Rec.2020 float internally, Bradford-adapted white balance, gamut mapping instead of clipping, a filmic shoulder for raw and pixel-exact pass-through for JPEGs you haven't touched. - **Resolution-independent edits.** Radii and brush sizes are relative to the image, so a 400 px preview, your diff --git a/ROADMAP.md b/ROADMAP.md index 5e793e8b..f816f7d1 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -25,7 +25,7 @@ duplicate Local entries #22, black GPU exports on an Intel iGPU #78). | **Feature checklist** | 79.7% | P0 core and P1 nearly complete: import (Add / Copy / Move, templates, devices), library, grid/loupe/compare/survey, every Edit slider, curves, colour grading, masking tools, crop/Upright, heal/clone, presets/profiles, versions/history, sync, export, menus, shortcuts | P0: full-resolution zoom detail; P1: lens-profile database, content-aware fill (patch synthesis), video playback/trim | | **RAW coverage** (formats people shoot) | ~55% | DNG (all kinds), CR2, CR3 lossless Bayer and version 0x100/0x200 C-RAW (M50/R100/R8 pixel-exact corpus), ARW, NEF (uncompressed + Huffman lossless/lossy), RAF (uncompressed, lossless/lossy compressed Bayer and X-Trans), uncompressed ORF, RW2 / RWL / Panasonic RAW (every raw format, checked on 178 files from 118 Panasonic and Leica bodies), PEF; every container's embedded preview (incl. CR3) | **CR3** (every Canon since ~2018), compressed ORF, Nikon lossy-after-split, Canon sRAW, HEIC/AVIF. NEFs labelled compressed but stored uncompressed (Z 6 packed 14-bit, D850 12-bit uncompressed). Per-model verification is thin (~55 corpus files vs >1,000 models) | | **Colour & image quality** | ~55–65% | Pipeline is complete and fast; GPU path CPU-exact within 1/255 | **No measured camera calibration database**: ARW, NEF, RW2, RAF, CR3, CR2 and PEF have guarded camera-JPEG colour estimates, with bundled ILCE-7M4, X-H2S and X-T4 profiles (docs/camera-preview-colour.md); other non-DNG raws and rejected estimates use a neutral matrix. Colour fidelity remains incomplete. No lens-profile database. No measured fidelity against Lightroom (tone, highlights, texture/clarity, NR, sharpening are tuned by eye) | -| **AI & computational** | ~15–20% | Assisted culling (focus, bursts), auto tone, HDR/panorama merge; subject/sky/background masks as classical heuristics | Real segmentation masks (subject, sky, people, objects, landscape, depth), AI denoise, super resolution, lens blur, generative remove, faces/people, natural-language search. **Blocked on a model strategy** (licensable weights or our own training; pure-Rust inference is feasible) | +| **AI & computational** | ~15–20% | Assisted culling (focus, bursts), auto tone, HDR/panorama merge; subject/sky/background masks as classical heuristics, or the photo's own segmentation mattes when it ships them (DNG semantic masks, e.g. iPhone ProRAW's sky matte); **faces (MVP, [docs/faces.md](docs/faces.md))**: a pure-Rust detector (a 232 KB opt-in download), opt-in recognition models (one-click download or add a file; they run on the shared pure-Rust CPU interpreter; pinned SFace/AuraFace embeddings agree with an independent reference runtime on 40 public crops; one-thread medians 34–36 / 233–252 ms (AuraFace remains expensive)), a background scan that finds and embeds faces, name suggestions, a person's page with look-alike faces to confirm, an Unnamed faces section (look-alikes together, select a group and name it at once) | Real segmentation masks (subject, sky, people, objects, landscape, depth), AI denoise, super resolution, lens blur, generative remove, natural-language search; for faces: clusters with separators between groups, small faces in large photos (the detector looks at 640 px), thresholds calibrated on live faces, writing face regions back to XMP. **Blocked on a model strategy** (licensable weights or our own training; pure-Rust inference is feasible) | | **Workflow & library** | ~85% (single machine) | Robust catalog (journal + snapshots, background compaction, crash-tested), 85k-photo libraries stay responsive, Local browsing with automatic cleanup, XMP interop, keywords, smart albums, Move import | Opening an 85k library takes 1.7–4.7 s; no cloud sync (out of scope), no tablet companion (#74, roadmap), shared albums, publish services, tethering | | **Classic modules** | ~30% | Geotagging from GPX track logs, soft proofing (partial), slideshow (basic) | **Map view, Book, Print, Slideshow module, Web, publish services**: ~40 tracker rows ⬜ | | **HDR & video** | 0% | | HDR edit/display/export; video play/trim/edit/export | @@ -59,7 +59,7 @@ Priorities, in order. Each points at tracker rows in [`docs/parity.md`](docs/par (references stay in the local `plan/`), then tune against the numbers. 5. **Lens profiles of our own** (LR-EDIT-OPTICS-PROFILE, P1). 6. **AI model strategy** (maintainer decision): which permissively licensed models (or our own training) for - segmentation masks and denoise; then pure-Rust inference. Unblocks M12 and Enhance. Object / Describe masks now + segmentation masks and denoise; then pure-Rust inference. Unblocks M12 and Enhance. Face detection and recognition run on our own checked Rust engine. Object / Describe masks now run Meta's SAM 3 in pure Rust (optional, user-consented download under the non-OSI SAM License; CDN mirrors still to be set up, docs/ai-masks.md); Subject / Sky / People and Enhance are still open. 7. **Then:** HDR (Q), the Classic output modules (Print first, then Map view, Book, Slideshow), video (R), localisation @@ -124,7 +124,7 @@ trip, slideshow, auto import…) in ≈ 12 h on 2026-10-02; four to six parallel | Work package | Tracker rows | Agent-hours | Risk | |---|---|---:|---| | Remaining P0/P1 UI and library features (folder rename/move, keyword painter, people view…) | ≈ 8 | 5–10 | low | -| Raw codecs: CR3 (CRX), compressed ORF, NEF lossy-after-split, HEIC/AVIF decode, JPEG XL DNG | LR-IMP-FORMATS | 40–80 | **high** — clean-room black-box analysis, no permissive specs | +| Raw codecs: CR3 (CRX), compressed ORF, NEF lossy-after-split, HEIC/AVIF decode | LR-IMP-FORMATS | 40–80 | **high** — clean-room black-box analysis, no permissive specs | | Lens-profile database of our own (calibration targets, fitting, data) | LR-EDIT-OPTICS-PROFILE | 15–30 | data collection | | Video: playback, trim, edits, export (pure-Rust decode, ideally shared with FilmCraft) | R. Video | 20–40 | medium | | AI: subject / sky / background / people / object masks, object-aware remove, AI denoise, super resolution, lens blur, people & faces, natural-language search, culling | ≈ 30 | 80–150 | **high** — permissively licensed weights, pure-Rust inference, maybe training | @@ -174,7 +174,8 @@ private 3:2 & 4:3 examples verified through native rendering & full-resolution 1 CR3 lossless Bayer and version 0x100/0x200 C-RAW sensor decoding is verified pixel-exactly on Canon M50, R100 and R8; any other CR3, or one that fails to decode, opens from its embedded JPEG as before. See [`docs/cr3.md`](docs/cr3.md) for coding limits, reference provenance and colour limitations. Decoded (CC0 corpus from raw.pixls.us, `cargo xtask corpus --download`, `crates/raw/tests/corpus.rs`): DNG (uncompressed, -LJ92, lossy JPEG / Smart Previews, Deflate, float, linear), CR2, ARW (uncompressed, ARW2, LJ92; as-shot white balance and black level of pre-2017 bodies from the enciphered +LJ92, lossy JPEG / Smart Previews, Deflate, float, linear, DNG 1.7 JPEG XL tiles: lossless ones bit-exact on synthetic files, +lossy ones decode but no real file has verified them), CR2, ARW (uncompressed, ARW2, LJ92; as-shot white balance and black level of pre-2017 bodies from the enciphered maker-note `Tag2010` and the encrypted `SR2SubIFD`, both recovered by black-box analysis, `crates/raw/src/vendor/arw.rs`), NEF/NRW uncompressed and Huffman-compressed (lossless, lossy type 1/2, 12/14-bit), RAF uncompressed and lossless/lossy compressed (Bayer and X-Trans, 14/16-bit real-file verification), RW2 / Leica RWL / Panasonic RAW in every raw format (compressed formats 4 and 6, the prefix-coded strips of format 8, packed formats 2, 5 and 7, the 16-bit words of the 2005–2007 bodies; all recovered by black-box analysis of 178 CC0 files from @@ -185,7 +186,7 @@ Not decoded yet — preview only (no permissively licensed description; black-bo files): rows above the split decode with the regular lossy table; from the split row on a different code is used that our black-box analysis has not recovered yet (none of the four regular tables fits, also not with byte alignment or reset predictors at the split row). The other Nikon Huffman variants are decoded (`crates/raw/src/vendor/nefc.rs` documents the analysis). -- **Olympus compressed ORF**, **Canon CR3 unverified CRX variants** (M11.1), **Canon sRAW/mRAW**, lossy DNG. +- **Olympus compressed ORF**, **Canon CR3 unverified CRX variants** (M11.1), **Canon sRAW/mRAW**. **Camera colour matrices:** ARW, NEF, RW2, RAF, CR3, CR2 and PEF files can use guarded, separate chromaticity and tone estimates from their own embedded JPEG (see `docs/camera-preview-colour.md`); this is a per-file camera-look estimate with relative WB, not measured calibration or absolute-Kelvin WB. Other non-DNG raws and rejected fits use the documented neutral fallback (camera RGB ≈ linear sRGB, flagged `matrix_is_fallback`) with the file's as-shot white-balance multipliers. Clean sources to evaluate next: manufacturer diff --git a/apps/lightcraft-cli/Cargo.toml b/apps/lightcraft-cli/Cargo.toml index 3f60fe32..4612b9a5 100644 --- a/apps/lightcraft-cli/Cargo.toml +++ b/apps/lightcraft-cli/Cargo.toml @@ -25,7 +25,12 @@ serde_json = { workspace = true } dhat = { workspace = true, optional = true } [features] +default = [] # Heap profiling: `cargo run --release -p lightcraft-cli --features dhat-heap -- snapshot …` counts # every allocation (live/peak bytes in `library.memory` / `ui.inspect`) and writes dhat-heap.json # (open it in DHAT's viewer: allocation sites by bytes at the peak). Slows allocation-heavy work. dhat-heap = ["dep:dhat"] + +# Open HEIC/HEIF photos (the optional lightcraft-heif crate). Not default so distributors choose +# it (HEVC patents); build with `--features heif` to include it. +heif = ["lightcraft-codecs/heif"] diff --git a/apps/lightcraft-cli/src/main.rs b/apps/lightcraft-cli/src/main.rs index 6c5e6e44..15edac61 100644 --- a/apps/lightcraft-cli/src/main.rs +++ b/apps/lightcraft-cli/src/main.rs @@ -728,19 +728,19 @@ fn snapshot(args: &[String]) -> Result<(), String> { if script.is_none() && output.is_none() { return Err("snapshot: give -o OUT.png and/or --script FILE".into()); } - if !(scale > 0.0 && size[0] >= 1.0 && size[1] >= 1.0) { - return Err("snapshot: bad --size/--scale".into()); - } + lightcraft_ui_egui::headless::viewport_pixels(size, scale).map_err(|e| format!("snapshot: bad --size/--scale: {e}"))?; let t0 = Instant::now(); let mut session = match &library { Some(dir) => { - let mut s = Session::new().with_fs(); + let mut s = Session::new().with_fs().with_default_face_models(); s.open_library(dir, false).map_err(|e| library_error(dir, e))?; s } - None if files.is_empty() => Session::with_demo().with_fs(), - None => Session::new().with_fs(), + None if files.is_empty() => Session::with_demo().with_fs().with_default_face_models(), + None => Session::new().with_fs().with_default_face_models(), }; + // Apply compute policy before imports or the first scripted query can discover an adapter. + session.execute("app.gpu", &json!({"enabled": false})).map_err(|e| e.to_string())?; if !files.is_empty() { let ti = Instant::now(); let r = session.execute("library.import", &json!({"paths": expand_paths(&files)})).map_err(|e| e.to_string())?; @@ -755,7 +755,11 @@ fn snapshot(args: &[String]) -> Result<(), String> { })), ..Default::default() }; - let app = lightcraft_ui_egui::LightcraftApp::new(session, services); + let mut app = lightcraft_ui_egui::LightcraftApp::new(session, services); + // Snapshot sessions promise no GPU: disable photo compute as well as the compositor + // before the first frame can start background adapter discovery. + // This UI state is session-local: the CLI neither loads nor saves desktop ui.json. + app.ui.settings.gpu = false; let mut h = Headless::new(app, size, scale); let timeout = Duration::from_secs(60); let mut shots = 0usize; diff --git a/apps/lightcraft-cli/tests/cli.rs b/apps/lightcraft-cli/tests/cli.rs index d1b178e4..2009e818 100644 --- a/apps/lightcraft-cli/tests/cli.rs +++ b/apps/lightcraft-cli/tests/cli.rs @@ -233,6 +233,58 @@ fn snapshot_script_failure_exits_non_zero() { assert_eq!((d.width, d.height), (320, 200)); } +#[test] +fn snapshot_ui_zoom_keeps_requested_pixel_dimensions() { + let script = tmp("snap-ui-zoom.jsonl"); + std::fs::write( + &script, + format!("{}\n{}\n", json!({"method": "ui.key", "params": {"key": "Plus", "cmd": true}}), json!({"method": "ui.inspect"})), + ) + .unwrap(); + for (size, scale, expected) in [("400x240", "1", (400, 240)), ("345x200", "0.9", (311, 180))] { + let out = tmp(&format!("snap-ui-zoom-{scale}.png")); + let o = Command::new(BIN) + .args(["snapshot", "--script", script.to_str().unwrap(), "-o", out.to_str().unwrap(), "--size", size, "--scale", scale]) + .output() + .unwrap(); + assert!(o.status.success(), "{}", String::from_utf8_lossy(&o.stderr)); + let replies: Vec = String::from_utf8_lossy(&o.stdout).lines().map(|line| serde_json::from_str(line).unwrap()).collect(); + assert!(replies.iter().all(|reply| reply["ok"] == true), "{replies:?}"); + assert!(replies[1]["result"]["pixelsPerPoint"].as_f64().unwrap() > scale.parse::().unwrap()); + let decoded = lightcraft_codecs::decode(&std::fs::read(out).unwrap(), Default::default()).unwrap(); + assert_eq!((decoded.width, decoded.height), expected); + } +} + +/// Headless snapshots default to CPU photo rendering as well as CPU UI rasterization. +#[test] +fn snapshot_defaults_to_cpu_without_gpu_environment_overrides() { + let script = tmp("default-cpu.jsonl"); + std::fs::write( + &script, + format!("{}\n{}\n", json!({"method": "engine.execute", "params": {"command": "app.gpu"}}), json!({"method": "ui.screenshot"})), + ) + .unwrap(); + let image = tmp("default-cpu.png"); + let output = Command::new(BIN) + .env_remove("LIGHTCRAFT_GPU") + .env_remove("LIGHTCRAFT_GPU_BACKEND") + .env_remove("WGPU_BACKEND") + .args(["snapshot", "--demo", "--script", script.to_str().unwrap(), "-o", image.to_str().unwrap(), "--size", "480x320"]) + .output() + .unwrap(); + assert!(output.status.success(), "status {}: {}", output.status, String::from_utf8_lossy(&output.stderr)); + let replies: Vec = String::from_utf8_lossy(&output.stdout).lines().map(|line| serde_json::from_str(line).unwrap()).collect(); + assert!(replies.iter().all(|reply| reply["ok"] == true), "{replies:?}"); + let gpu = &replies[0]["result"]; + assert_eq!(gpu["enabled"], false, "{gpu}"); + assert_eq!(gpu["available"], false, "{gpu}"); + assert_eq!(gpu["adapter"], Value::Null, "no device is created: {gpu}"); + assert!(gpu["reason"].as_str().is_some_and(|reason| reason.contains("preference")), "{gpu}"); + let decoded = lightcraft_codecs::decode(&std::fs::read(image).unwrap(), Default::default()).unwrap(); + assert_eq!((decoded.width, decoded.height), (480, 320)); +} + /// Issue #136: with the GPU switched off from the environment the UI starts and renders on the CPU, /// without creating a GPU device (no driver is loaded), and says why. #[test] @@ -412,3 +464,79 @@ fn warnings_are_logged_on_stderr() { let off = run("off"); assert!(!off.contains("names no known backend"), "{off}"); } + +#[test] +fn snapshot_rejects_invalid_dimensions_before_creating_an_image() { + for (index, (size, scale)) in + [("NaNx100", "1"), ("100x100", "inf"), ("9000x9000", "1"), ("9000x100", "2"), ("100x100", "0.001")].into_iter().enumerate() + { + let image = tmp(&format!("invalid-viewport-{index}.png")); + let output = Command::new(BIN).args(["snapshot", "--size", size, "--scale", scale, "-o"]).arg(&image).output().unwrap(); + assert!(!output.status.success()); + assert!(String::from_utf8_lossy(&output.stderr).contains("snapshot: bad --size/--scale"), "{}", String::from_utf8_lossy(&output.stderr)); + assert!(!image.exists()); + } +} + +#[test] +fn snapshot_preserves_fractional_scale_rounding() { + let image = tmp("fractional-scale-rounding.png"); + let output = + Command::new(BIN).env("LIGHTCRAFT_GPU", "0").args(["snapshot", "--size", "345x200", "--scale", "0.9", "-o"]).arg(&image).output().unwrap(); + assert!(output.status.success(), "{}", String::from_utf8_lossy(&output.stderr)); + let decoded = lightcraft_codecs::decode(&std::fs::read(image).unwrap(), Default::default()).unwrap(); + assert_eq!((decoded.width, decoded.height), (311, 180)); +} + +/// Snapshot's CPU setting must never replace the desktop preference, including with --library. +#[cfg(target_os = "linux")] +#[test] +fn snapshot_leaves_desktop_gpu_preferences_untouched() { + let dir = tmp("snapshot-desktop-preferences"); + let library = dir.join("library"); + let mut session = lightcraft_engine::Session::new().with_fs(); + session.open_library(&library, true).unwrap(); + drop(session); + let config = dir.join("config"); + let app_config = config.join("lightcraft"); + std::fs::create_dir_all(&app_config).unwrap(); + let ui_path = app_config.join("ui.json"); + let mut ui = lightcraft_ui_egui::UiState::default(); + ui.settings.gpu = true; + let saved = serde_json::to_vec_pretty(&ui).unwrap(); + std::fs::write(&ui_path, &saved).unwrap(); + let script = dir.join("preferences.jsonl"); + std::fs::write( + &script, + format!("{}\n{}\n", json!({"method": "engine.execute", "params": {"command": "app.gpu"}}), json!({"method": "ui.inspect"})), + ) + .unwrap(); + for gpu_environment in [None, Some("1")] { + let image = dir.join(format!("{}.png", gpu_environment.unwrap_or("default"))); + let mut cmd = Command::new(BIN); + cmd.env("XDG_CONFIG_HOME", &config).env_remove("LIGHTCRAFT_GPU").env_remove("LIGHTCRAFT_GPU_BACKEND").env_remove("WGPU_BACKEND"); + if let Some(value) = gpu_environment { + cmd.env("LIGHTCRAFT_GPU", value); + } + let output = cmd + .args(["snapshot", "--library"]) + .arg(&library) + .args(["--script"]) + .arg(&script) + .args(["-o"]) + .arg(&image) + .args(["--size", "480x320"]) + .output() + .unwrap(); + assert!(output.status.success(), "{}", String::from_utf8_lossy(&output.stderr)); + let replies: Vec = String::from_utf8_lossy(&output.stdout).lines().map(|line| serde_json::from_str(line).unwrap()).collect(); + assert!(replies.iter().all(|reply| reply["ok"] == true), "{replies:?}"); + assert_eq!(replies[0]["result"]["enabled"], false); + assert_eq!(replies[0]["result"]["available"], false); + assert_eq!(replies[0]["result"]["adapter"], Value::Null); + assert_eq!(replies[1]["result"]["ui"]["settings"]["gpu"], false); + assert_eq!(std::fs::read(&ui_path).unwrap(), saved, "snapshot changed the desktop preferences"); + assert!(!app_config.join("ui.json.tmp").exists()); + assert!(image.exists()); + } +} diff --git a/apps/lightcraft-web/Cargo.toml b/apps/lightcraft-web/Cargo.toml index a3331a60..aaa14a64 100644 --- a/apps/lightcraft-web/Cargo.toml +++ b/apps/lightcraft-web/Cargo.toml @@ -71,3 +71,9 @@ web-sys = { workspace = true, features = [ "WorkerType", "WritableStream", ] } + +[features] +# Open HEIC/HEIF photos (the optional lightcraft-heif crate). Not default so distributors choose +# it (HEVC patents); build with `--features heif` to include it. On wasm32 panics abort, so +# heic-rs's `catch_unwind` guard does not apply there. +heif = ["lightcraft-codecs/heif"] diff --git a/apps/lightcraft-web/src/web.rs b/apps/lightcraft-web/src/web.rs index 9ffbf170..a1c9d76d 100644 --- a/apps/lightcraft-web/src/web.rs +++ b/apps/lightcraft-web/src/web.rs @@ -209,6 +209,8 @@ fn services(originals: Originals, backend: Option, files: Files, frozen pick_preset_files: None, pick_tracklog: None, pick_lightroom_catalog: None, + // Face models need a folder to live in; the web has none. + pick_model_file: None, save_preset_file: None, pick_curve_preset_files: None, save_curve_preset_file: None, diff --git a/apps/lightcraft-web/src/wire.rs b/apps/lightcraft-web/src/wire.rs index 81a76e6c..b08ae150 100644 --- a/apps/lightcraft-web/src/wire.rs +++ b/apps/lightcraft-web/src/wire.rs @@ -368,7 +368,11 @@ mod tests { let a = core.render(&wire, None).unwrap(); let b = job.run().rendered.unwrap(); assert_eq!((a.image.width, a.image.height), (160, 120)); - assert_eq!(a.image.data, b.image.data); + assert_eq!((b.image.width, b.image.height), (160, 120)); + // the worker renders on the CPU; the engine job uses a GPU when the machine has one, which + // stays within docs/gpu-pipeline.md's bound (≤ 3 LSB per channel) rather than bit-exact + let max = a.image.data.iter().zip(&b.image.data).flat_map(|(p, q)| (0..3).map(move |c| p[c].abs_diff(q[c]))).max(); + assert!(max.is_some_and(|m| m <= 3), "max channel difference {max:?}"); } // a job from before windows existed has none diff --git a/apps/lightcraft/Cargo.toml b/apps/lightcraft/Cargo.toml index a29299da..f9173c2a 100644 --- a/apps/lightcraft/Cargo.toml +++ b/apps/lightcraft/Cargo.toml @@ -9,6 +9,12 @@ default-run = "lightcraft" [lints] workspace = true +[features] +default = [] +# Open HEIC/HEIF photos (the optional lightcraft-heif crate). Not default so distributors choose +# it (HEVC patents); build with `--features heif` to include it. +heif = ["lightcraft-codecs/heif"] + [dependencies] lightcraft-engine = { workspace = true, features = ["sam"] } lightcraft-sysmem = { workspace = true } diff --git a/apps/lightcraft/src/main.rs b/apps/lightcraft/src/main.rs index 4c7bb424..be44cc33 100644 --- a/apps/lightcraft/src/main.rs +++ b/apps/lightcraft/src/main.rs @@ -400,6 +400,14 @@ fn services(ctx: egui::Context, log_file: Option<&std::path::Path>) -> Services .map(|p| p.to_string_lossy().to_string()) .collect() })), + pick_model_file: Some(Box::new(|| { + rfd::FileDialog::new() + .set_title(lightcraft_ui_egui::i18n::tr("Add a Face Recognition Model")) + .add_filter_nocase(lightcraft_ui_egui::i18n::tr("Face models (ONNX)"), &["onnx"]) + .pick_file() + .map(|p| vec![p.to_string_lossy().to_string()]) + .unwrap_or_default() + })), pick_tracklog: Some(Box::new(|| { rfd::FileDialog::new() .set_title(lightcraft_ui_egui::i18n::tr("Auto-Tag from Tracklog")) @@ -446,12 +454,19 @@ fn services(ctx: egui::Context, log_file: Option<&std::path::Path>) -> Services } } +/// The library session with the folder face models are kept in (`/models`). +fn open_session(in_memory: bool, dir: Option, seed_demo: bool) -> (Session, Option) { + let (mut s, problem) = open_library_session(in_memory, dir, seed_demo); + s.face_models_dir = lightcraft_engine::config::default_face_models_dir(); + (s, problem) +} + /// The persistent library session, or with `--memory` an in-memory one (demo photos). /// /// If the library can't be opened the session is empty and in memory — never seeded with the /// demo photos, never written anywhere — and the problem is returned: the window then says so /// and offers Try Again / Choose Another Library… / Continue Without Saving / Quit (issue #100). -fn open_session(in_memory: bool, dir: Option, seed_demo: bool) -> (Session, Option) { +fn open_library_session(in_memory: bool, dir: Option, seed_demo: bool) -> (Session, Option) { if in_memory { return (if seed_demo { Session::with_demo() } else { Session::new() }.with_fs().with_system_clock(), None); } diff --git a/contributors/people.toml b/contributors/people.toml index 89c6db20..fbda1f24 100644 --- a/contributors/people.toml +++ b/contributors/people.toml @@ -21,3 +21,6 @@ real_name = "Brandon Thomas" display_name = "Brandon Thomas" sync_github_name = false + +[people.sanyamsingh001] +display_name = "sanyamsingh001" diff --git a/crates/codecs/Cargo.toml b/crates/codecs/Cargo.toml index 68d5e8cd..18d00ac3 100644 --- a/crates/codecs/Cargo.toml +++ b/crates/codecs/Cargo.toml @@ -12,6 +12,7 @@ workspace = true lightcraft-geom = { workspace = true } lightcraft-color = { workspace = true } lightcraft-raster = { workspace = true } +lightcraft-heif = { workspace = true, optional = true } serde = { workspace = true } bytemuck = { workspace = true } thiserror = { workspace = true } @@ -34,6 +35,10 @@ ravif = { workspace = true, optional = true } default = ["parallel", "jxl", "avif"] parallel = ["dep:rayon", "lightcraft-raster/parallel", "jpeg-decoder/rayon"] jxl = ["dep:jxl-oxide"] +# HEIC/HEIF decode via the optional lightcraft-heif crate (heic-rs, pure Rust). Off by default so +# distributors choose it at build time (HEVC patents are a distributor's call, same policy as +# PhotoCraft); `cargo xtask ci` clippies and tests it (`--features heif`), releases don't ship it yet. +heif = ["dep:lightcraft-heif"] ## AVIF encoding via ravif/rav1e (pure Rust, BSD). Native targets only; a no-op on wasm32. avif = ["dep:ravif"] diff --git a/crates/codecs/src/heif.rs b/crates/codecs/src/heif.rs new file mode 100644 index 00000000..fe91c0af --- /dev/null +++ b/crates/codecs/src/heif.rs @@ -0,0 +1,125 @@ +//! HEIF / HEIC (the iPhone and Mac photo format), read-only. +//! +//! Decoding lives in the optional `lightcraft-heif` crate (heic-rs, a pure-Rust HEVC +//! still-picture decoder), enabled by this crate's `heif` feature — off by default because HEVC +//! is patent-encumbered and whether a build carries an HEVC decoder is the distributor's call +//! (the same policy as the sibling PhotoCraft). Without the feature, HEIF files are still +//! detected and opening one is an [`Error::Unsupported`] error, never a panic. +//! +//! HEIF records orientation in the container (`irot`/`imir`, plus a `clap` crop), not in EXIF, +//! so the decoder applies them and `orientation` is reported as 1: the EXIF Orientation tag only +//! mirrors what the container says, and applying it as well would turn the photo twice. +//! The samples come back in the container's own primaries with its ICC profile (if any); like +//! every decoder here, the colour interpretation happens in [`crate::convert::finish`]. + +#[cfg(feature = "heif")] +use crate::convert::{Buf, Meta, Model, Raw, finish}; +use crate::{DecodeOptions, Decoded, Error, Format, Result}; + +const F: Format = Format::Heif; + +/// The reason given when this build has no HEIF decoder. +#[cfg(not(feature = "heif"))] +pub(crate) const NOT_IN_BUILD: &str = "HEIC/HEIF support isn't included in this build of LightCraft"; + +/// The reason a HEIF file that no build can open is refused (this crate decodes HEVC stills +/// only; the wrapper names the exact construct in its own errors). +#[cfg(feature = "heif")] +const NOT_DECODED: &str = "the image is not a decodable HEVC still (image sequence, overlay, alpha-less grid or AVIF-coded item)"; + +#[cfg(not(feature = "heif"))] +pub(crate) fn decode(_bytes: &[u8], _opts: &DecodeOptions) -> Result { + Err(Error::Unsupported(F, NOT_IN_BUILD)) +} + +#[cfg(feature = "heif")] +fn err(e: lightcraft_heif::Error) -> Error { + match e { + // This crate's Unsupported takes a fixed reason; the dynamic wording lives in the + // lightcraft-heif error itself. + lightcraft_heif::Error::Unsupported(_) => Error::Unsupported(F, NOT_DECODED), + lightcraft_heif::Error::Limit(m) | lightcraft_heif::Error::Malformed(m) => Error::Malformed(F, m), + } +} + +#[cfg(feature = "heif")] +pub(crate) fn decode(bytes: &[u8], opts: &DecodeOptions) -> Result { + // The container alone: the declared size is checked before any pixel is decoded. + let info = lightcraft_heif::probe(bytes).map_err(err)?; + if (info.width as u64).saturating_mul(info.height as u64) > opts.max_pixels { + return Err(Error::TooLarge(info.width as u64, info.height as u64)); + } + let options = lightcraft_heif::Options { max_pixels: opts.max_pixels, apply_transforms: true }; + let decoded = lightcraft_heif::decode(bytes, &options).map_err(err)?; + // The decoder re-reads the container; trust its output's shape, not the probe's. + let buf = if decoded.sixteen_bit { + let mut v = Vec::with_capacity(decoded.data.len() / 2); + for pair in decoded.data.as_chunks::<2>().0 { + v.push(u16::from_ne_bytes(*pair)); + } + Buf::U16(v) + } else { + Buf::U8(decoded.data) + }; + let raw = Raw { + width: decoded.width as usize, + height: decoded.height as usize, + model: Model::Rgb, + alpha: decoded.has_alpha, + premultiplied: false, + buf, + bit_depth: if decoded.sixteen_bit { 16 } else { 8 }, + }; + let meta = Meta { icc: decoded.icc, exif: decoded.exif, xmp: decoded.xmp, orientation: Some(1), ..Default::default() }; + finish(F, raw, meta, (decoded.width, decoded.height), opts) +} + +#[cfg(test)] +mod tests { + use crate::{Format, sniff}; + + /// The container header both builds recognise, whatever the feature says. + const HEIC_HEADER: &[u8] = b"\0\0\0\x18ftypheic\0\0\0\0mif1heic"; + + #[test] + fn heif_is_always_recognised() { + assert_eq!(sniff(HEIC_HEADER), Some(Format::Heif)); + } + + #[test] + fn can_decode_follows_the_feature() { + assert_eq!(Format::Heif.can_decode(), cfg!(feature = "heif")); + } + + #[cfg(not(feature = "heif"))] + #[test] + fn without_the_feature_heif_is_a_clear_unsupported_error() { + let r = crate::decode(HEIC_HEADER, crate::DecodeOptions::default()); + assert!(matches!(&r, Err(crate::Error::Unsupported(Format::Heif, why)) if why.contains("isn't included in this build")), "{r:?}"); + } + + /// Found by PhotoCraft's `decode_heif` fuzz target: a malformed box makes heic-rs 0.1.1 slice + /// out of range (`boxes.rs:130`). It must be an error, not a crash. + #[cfg(feature = "heif")] + const HEIC_RS_BOX_PANIC: [u8; 72] = [ + 0x00, 0x00, 0x00, 0x24, 0x66, 0x74, 0x79, 0x70, 0x68, 0x65, 0x69, 0x63, 0x00, 0x00, 0x00, 0x00, 0x6d, 0x69, 0x66, 0x31, 0x4d, 0x69, 0x50, + 0x72, 0x6d, 0x69, 0x61, 0x66, 0x4d, 0x69, 0x48, 0x42, 0x72, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x6d, 0x65, 0x74, 0x61, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x75, 0x75, 0x69, 0x64, 0x00, 0x00, 0x00, 0x00, 0x03, 0x08, 0x08, 0x08, 0x00, 0x03, 0x01, 0x03, 0x70, + 0x00, 0xa8, 0x00, + ]; + + #[cfg(feature = "heif")] + #[test] + fn a_heic_rs_panic_is_a_malformed_error() { + let r = crate::decode(&HEIC_RS_BOX_PANIC, crate::DecodeOptions::default()); + assert!(matches!(&r, Err(crate::Error::Malformed(Format::Heif, _))), "{r:?}"); + } + + #[cfg(feature = "heif")] + #[test] + fn an_image_sequence_is_unsupported() { + // An ftyp with no meta box: a video track, not a still. + let r = crate::decode(b"\0\0\0\x18ftypmsf1\0\0\0\0msf1hevc", crate::DecodeOptions::default()); + assert!(matches!(&r, Err(crate::Error::Unsupported(Format::Heif, _))), "{r:?}"); + } +} diff --git a/crates/codecs/src/lib.rs b/crates/codecs/src/lib.rs index 15990329..169f40c1 100644 --- a/crates/codecs/src/lib.rs +++ b/crates/codecs/src/lib.rs @@ -19,6 +19,7 @@ mod convert; pub mod encode; pub mod exif; +mod heif; pub mod icc; mod jpeg; pub mod jpeg_par; @@ -227,8 +228,8 @@ pub fn decode_unguarded(bytes: &[u8], format: Format, opts: &DecodeOptions) -> R Format::Jxl => jxl::decode(bytes, &opts), #[cfg(not(feature = "jxl"))] Format::Jxl => Err(Error::Unsupported(format, "built without the `jxl` feature")), + Format::Heif => heif::decode(bytes, &opts), Format::Avif => Err(Error::Unsupported(format, "no pure-Rust, permissively licensed AV1 decoder yet")), - Format::Heif => Err(Error::Unsupported(format, "no pure-Rust, permissively licensed HEVC decoder yet")), Format::RawTiffLike | Format::RawOther => Err(Error::Unsupported(format, "camera raw: decode with lightcraft-raw")), } } diff --git a/crates/codecs/src/sniff.rs b/crates/codecs/src/sniff.rs index a02d87d2..50fa3ccf 100644 --- a/crates/codecs/src/sniff.rs +++ b/crates/codecs/src/sniff.rs @@ -66,7 +66,8 @@ impl Format { pub fn can_decode(self) -> bool { match self { Format::Jxl => cfg!(feature = "jxl"), - Format::Avif | Format::Heif | Format::RawTiffLike | Format::RawOther => false, + Format::Heif => cfg!(feature = "heif"), + Format::Avif | Format::RawTiffLike | Format::RawOther => false, _ => true, } } diff --git a/crates/engine/Cargo.toml b/crates/engine/Cargo.toml index a916d180..7a6ab439 100644 --- a/crates/engine/Cargo.toml +++ b/crates/engine/Cargo.toml @@ -22,6 +22,7 @@ lightcraft-preview = { workspace = true } lightcraft-codecs = { workspace = true } lightcraft-raw = { workspace = true } lightcraft-meta = { workspace = true } +lightcraft-faces = { workspace = true } miniz_oxide = { workspace = true } lightcraft-merge = { workspace = true } rayon = { workspace = true } @@ -35,6 +36,10 @@ harfrust = { version = "0.12", default-features = false, features = ["std"] } unicode-segmentation = "1.13" lightcraft-segment = { workspace = true, optional = true } +# model downloads run natively only (the web build fetches no models) +[target.'cfg(not(target_arch = "wasm32"))'.dependencies] +lightcraft-fetch = { workspace = true } + [features] # AI masks with SAM 3 (Object / Describe). Off for the web and CLI builds. sam = ["dep:lightcraft-segment"] diff --git a/crates/engine/src/cmd/face_detect.rs b/crates/engine/src/cmd/face_detect.rs new file mode 100644 index 00000000..ffdedbef --- /dev/null +++ b/crates/engine/src/cmd/face_detect.rs @@ -0,0 +1,184 @@ +//! `faces.detect`: find faces with the YuNet detector the user has downloaded (Settings > Faces). +//! +//! The photo is rendered upright and uncropped with default settings (so edits, crops and rotations you have +//! made do not matter), the detector runs on it, and the faces come back as fractions of the photo's width +//! and height, the same frame face regions from XMP use. Unless `apply` is false, they are added to the photo as +//! unnamed face regions in one undoable step; earlier detections are replaced, and regions that came from XMP (or +//! that you drew or named) are never touched. Like every region edit it is catalog-only: no sidecar is written. + +use std::path::PathBuf; +use std::sync::{Arc, Mutex, PoisonError}; +use std::time::Instant; + +use lightcraft_catalog::Op; +use lightcraft_develop::DevelopSettings; +use lightcraft_faces::known::YUNET_ID; +use lightcraft_faces::yunet::{Detector, Face, Options}; +use lightcraft_meta::{Region, RegionKind}; +use serde_json::{Value, json}; + +use super::{CommandSpec, bad, cmd, has_selection}; +use crate::{Result, Session}; + +const C: &str = "faces.detect"; +/// Regions made by this command say so in their description; only those are replaced on a new run. +pub(crate) const MARK: &str = "Detected by "; +pub(crate) const LABEL: &str = "YuNet 2023mar"; +/// Long edge of the picture the detector looks at (it shrinks it to its own 640 anyway). +const LOOK_EDGE: usize = 1280; +/// YuNet is 232 KB; a file this much larger in its folder is not it. +const MAX_MODEL_BYTES: u64 = 8 << 20; +/// How much a detection must overlap a region the photo already has to count as the same face. +const EXISTING_IOU: f64 = 0.3; + +/// What to tell someone who has not downloaded the detector yet. +const NOT_INSTALLED: &str = "finding faces needs the YuNet face detector (232 KB, MIT): download it in Settings > Faces"; + +/// The installed YuNet, loaded once and kept (a new one is loaded when the model folder changes). +/// (also used by the background scan, which aligns faces by the detector's landmarks and finds faces in photos that have none) +pub(crate) fn detector(s: &Session) -> std::result::Result, String> { + static LOADED: Mutex)>> = Mutex::new(None); + let home = s.face_models_dir.as_deref().ok_or("this build has nowhere to keep face models (the desktop app does)")?.join(YUNET_ID); + let path = home.join("model.onnx"); + // installing writes the manifest last, so a folder without one is not (yet) a model + if !path.is_file() || !home.join("face-model.json").is_file() { + return Err(NOT_INSTALLED.into()); + } + let mut loaded = LOADED.lock().unwrap_or_else(PoisonError::into_inner); + if let Some((p, d)) = loaded.as_ref() + && *p == path + { + return Ok(d.clone()); + } + let len = std::fs::metadata(&path).map_err(|e| format!("could not read the detector: {e}"))?.len(); + if len > MAX_MODEL_BYTES { + return Err("the detector's file is far bigger than YuNet: remove it in Settings > Faces and download it again".into()); + } + let bytes = std::fs::read(&path).map_err(|e| format!("could not read the detector: {e}"))?; + let d = Arc::new(Detector::new(&bytes).map_err(|e| format!("{e}: remove the detector in Settings > Faces and download it again"))?); + *loaded = Some((path, d.clone())); + Ok(d) +} + +fn face_json(f: &Face) -> Value { + json!({ + "rect": {"x0": f.x0, "y0": f.y0, "x1": f.x1, "y1": f.y1}, + "score": f.score, + "landmarks": f.landmarks.iter().map(|(x, y)| json!([x, y])).collect::>(), + }) +} + +/// A detected face that lies on a region the photo already has (from XMP, drawn or named) is that region, not a new face. +fn overlaps_existing(f: &Face, regions: &[Region]) -> bool { + regions.iter().any(|r| { + let (w, h) = ( + (r.rect.x1.min(f64::from(f.x1)) - r.rect.x0.max(f64::from(f.x0))).max(0.0), + (r.rect.y1.min(f64::from(f.y1)) - r.rect.y0.max(f64::from(f.y0))).max(0.0), + ); + let inter = w * h; + let union = (r.rect.x1 - r.rect.x0) * (r.rect.y1 - r.rect.y0) + f64::from((f.x1 - f.x0) * (f.y1 - f.y0)) - inter; + union > 0.0 && inter / union >= EXISTING_IOU + }) +} + +pub(crate) fn is_detected(r: &Region) -> bool { + r.description.as_deref().is_some_and(|d| d.starts_with(MARK)) +} + +/// The photo as the faces see it: upright, uncropped, default settings, as 8-bit RGB with its size, its long edge at +/// most `edge`. +pub(crate) fn render_rgb(s: &mut Session, id: lightcraft_catalog::PhotoId, edge: usize) -> std::result::Result<(Vec, usize, usize), String> { + let settings = DevelopSettings::default(); + let job = s.preview_job(id, edge, edge, false, &settings).ok_or("no such photo")?; + let image = job.run().rendered?.image; + let rgb: Vec = image.data.iter().flat_map(|px| [px[0], px[1], px[2]]).collect(); + Ok((rgb, image.width, image.height)) +} + +fn detect(s: &mut Session, p: &Value) -> Result { + let apply = p.get("apply").and_then(Value::as_bool).unwrap_or(true); + let defaults = Options::default(); + let opts = Options { + score: p.get("score").and_then(Value::as_f64).map_or(defaults.score, |v| v as f32), + nms_iou: p.get("nmsIou").and_then(Value::as_f64).map_or(defaults.nms_iou, |v| v as f32), + max_faces: p.get("maxFaces").and_then(Value::as_u64).map_or(defaults.max_faces, |v| usize::try_from(v).unwrap_or(defaults.max_faces)), + }; + let detector = detector(s).map_err(|e| bad(C, e))?; + let started = Instant::now(); + let (mut results, mut ops) = (Vec::new(), Vec::new()); + for id in s.targets(p) { + let (rgb, width, height) = match render_rgb(s, id, LOOK_EDGE) { + Ok(r) => r, + Err(e) => { + results.push(json!({"id": id.0, "error": e})); + continue; + } + }; + let faces = match detector.detect(&rgb, width, height, &opts) { + Ok(f) => f, + Err(e) => { + results.push(json!({"id": id.0, "error": e.to_string()})); + continue; + } + }; + if apply && let Some(photo) = s.catalog.photo(id) { + let mut meta = photo.meta.clone(); + meta.regions.retain(|r| !is_detected(r)); + let fresh: Vec<&Face> = faces.iter().filter(|f| !overlaps_existing(f, &meta.regions)).collect(); + for f in fresh { + meta.regions.push(Region { + rect: lightcraft_geom::Rect { x0: f64::from(f.x0), y0: f64::from(f.y0), x1: f64::from(f.x1), y1: f64::from(f.y1) }, + kind: RegionKind::Face, + name: None, + description: Some(format!("{MARK}{LABEL}")), + }); + } + ops.push(Op::SetMeta { id, meta: Box::new(meta) }); + } + results.push(json!({"id": id.0, "faces": faces.iter().map(face_json).collect::>()})); + } + let applied = !ops.is_empty(); + if applied { + s.commit("Detect Faces", Op::Batch { ops })?; + // regions are not written to XMP, so a sidecar must not be rewritten for this + s.skip_auto_write = true; + } + Ok(json!({"detector": LABEL, "applied": applied, "photos": results, "ms": started.elapsed().as_secs_f64() * 1000.0})) +} + +pub fn specs() -> Vec { + vec![cmd!( + "faces.detect", + "Detect Faces", + ["Photo"], + None, + "{ids?, apply?: bool (default true), score?: 0.6, nmsIou?: 0.3, maxFaces?} → {photos: [{id, faces: [{rect, score, landmarks}]}]}. Finds faces with the YuNet detector (download it first: `faces.models.download {id: \"yunet-2023mar\"}`) and, unless `apply` is false, adds them as unnamed face regions (one undo step), replacing earlier detections but never regions from XMP or ones you named; the sidecar is not touched", + has_selection, + detect + )] +} + +#[cfg(test)] +mod tests { + use super::*; + + fn face(x0: f32, y0: f32, x1: f32, y1: f32) -> Face { + Face { x0, y0, x1, y1, score: 0.9, landmarks: [(0.0, 0.0); 5] } + } + fn region(x0: f64, y0: f64, x1: f64, y1: f64) -> Region { + Region { rect: lightcraft_geom::Rect { x0, y0, x1, y1 }, kind: RegionKind::Face, name: Some("Ann".into()), description: None } + } + + #[test] + fn a_detection_on_an_existing_region_is_not_a_new_face() { + let existing = [region(0.40, 0.20, 0.60, 0.50), region(0.10, 0.10, 0.20, 0.20)]; + // the same face, a bit larger or shifted (a Lightroom box is not YuNet's box) + assert!(overlaps_existing(&face(0.38, 0.18, 0.64, 0.54), &existing)); + assert!(overlaps_existing(&face(0.42, 0.22, 0.58, 0.48), &existing)); + // another face, a sliver of overlap, and no regions at all + assert!(!overlaps_existing(&face(0.70, 0.20, 0.90, 0.50), &existing)); + assert!(!overlaps_existing(&face(0.59, 0.49, 0.80, 0.80), &existing)); + assert!(!overlaps_existing(&face(0.0, 0.0, 0.0, 0.0), &existing), "an empty box overlaps nothing"); + assert!(!overlaps_existing(&face(0.4, 0.2, 0.6, 0.5), &[])); + } +} diff --git a/crates/engine/src/cmd/face_models.rs b/crates/engine/src/cmd/face_models.rs new file mode 100644 index 00000000..44b359e4 --- /dev/null +++ b/crates/engine/src/cmd/face_models.rs @@ -0,0 +1,511 @@ +//! Face model commands (`faces.*`): which models exist, inspecting a file the user dropped in, installing +//! it once its licence is accepted, removing it, and the faces on/off setting. +//! +//! Models live in the folder the host provides (`Session::face_models_dir`): one subfolder per model with +//! `model.onnx`, `face-model.json` (its manifest) and `installed.json` (what the user accepted), plus +//! `settings.json` beside them. Model files and manifests are hostile input: sizes are capped, ids are +//! validated before they become folder names, and nothing outside a model's own folder is ever touched. + +use std::io::Read; +use std::path::{Path, PathBuf}; + +use lightcraft_faces::catalog::{self, Catalog}; +use lightcraft_faces::hash::sha256_file; +use lightcraft_faces::manifest::{self, MAX_MANIFEST_BYTES, MAX_MODEL_BYTES, ModelManifest, Role}; +use lightcraft_faces::suggest::{Suggestion, suggest_with}; +use lightcraft_faces::{known, onnx}; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; + +use super::{CommandSpec, always, bad, cmd, str_param}; +use crate::face_download::{self, State}; +use crate::{EngineError, Result, Session}; + +#[derive(Default, Serialize, Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub(super) struct FaceSettings { + pub(super) enabled: bool, + /// The installed recogniser in use. + pub(super) embedder: Option, +} + +fn fail(what: &str, e: impl std::fmt::Display) -> EngineError { + EngineError::Other(format!("{what}: {e}")) +} + +fn models_dir(s: &Session, cmd: &str) -> Result { + s.face_models_dir.clone().ok_or_else(|| bad(cmd, "this build has nowhere to keep face models (the desktop app does)")) +} + +/// Read at most `max` bytes of a small file. +fn read_capped(path: &Path, max: usize) -> Option> { + let mut buf = Vec::new(); + std::fs::File::open(path).ok()?.take(max as u64 + 1).read_to_end(&mut buf).ok()?; + (buf.len() <= max).then_some(buf) +} + +pub(super) fn read_settings(dir: &Path) -> FaceSettings { + read_capped(&dir.join("settings.json"), 4096).and_then(|b| serde_json::from_slice(&b).ok()).unwrap_or_default() +} + +fn write_atomic(path: &Path, bytes: &[u8]) -> Result<()> { + let part = path.with_extension("part"); + std::fs::write(&part, bytes).map_err(|e| fail(&format!("could not write {}", path.display()), e))?; + std::fs::rename(&part, path).map_err(|e| { + let _ = std::fs::remove_file(&part); + fail(&format!("could not write {}", path.display()), e) + }) +} + +fn write_settings(dir: &Path, st: &FaceSettings) -> Result<()> { + std::fs::create_dir_all(dir).map_err(|e| fail("could not create the face models folder", e))?; + write_atomic(&dir.join("settings.json"), &serde_json::to_vec_pretty(st).map_err(|e| fail("settings", e))?) +} + +pub(super) struct Installed { + pub(super) manifest: ModelManifest, + pub(super) accepted: Value, +} + +/// The models in the folder. A folder that is not a valid model (wrong name, no model file, a manifest that +/// does not validate or does not match its folder) is ignored, never trusted. +pub(super) fn installed_models(dir: &Path) -> Vec { + let Ok(entries) = std::fs::read_dir(dir) else { return Vec::new() }; + let mut out = Vec::new(); + for e in entries.flatten() { + let p = e.path(); + let Some(name) = p.file_name().and_then(|n| n.to_str()).map(str::to_string) else { continue }; + if !manifest::valid_id(&name) || !p.is_dir() || !p.join("model.onnx").is_file() { + continue; + } + let Some(m) = read_capped(&p.join("face-model.json"), MAX_MANIFEST_BYTES).and_then(|b| manifest::parse(&b).ok()) else { continue }; + if m.id != name { + continue; + } + let accepted = read_capped(&p.join("installed.json"), 4096).and_then(|b| serde_json::from_slice(&b).ok()).unwrap_or(Value::Null); + out.push(Installed { manifest: m, accepted }); + } + out.sort_by(|a, b| a.manifest.id.cmp(&b.manifest.id)); + out +} + +/// Read the user's catalog (`catalog.json` in the models folder) anew: a few kilobytes, so it is always current. +fn refresh_catalog(s: &mut Session) { + let Some(dir) = s.face_models_dir.clone() else { return }; + let file = dir.join("catalog.json"); + s.face_catalog = match read_capped(&file, catalog::MAX_CATALOG_BYTES) { + Some(bytes) => catalog::parse(&bytes), + None if file.is_file() => catalog::parse(&vec![b' '; catalog::MAX_CATALOG_BYTES + 1]), + None => Catalog::default(), + }; +} + +/// The manifests of the catalog's models (what a file's hash is matched against, besides the built-in ones). +fn catalog_manifests(s: &Session) -> Vec { + s.face_catalog.entries.iter().map(|e| e.manifest.clone()).collect() +} + +/// How to fetch a model: the built-in address, or the one in the user's catalog. +fn download_spec(s: &Session, id: &str) -> Option { + known::download(id).or_else(|| s.face_catalog.entries.iter().find(|e| e.manifest.id == id).and_then(catalog::Entry::download)) +} + +fn row(s: &Session, m: &ModelManifest, installed: bool, selected: bool, accepted: &Value) -> Value { + let from_catalog = s.face_catalog.entries.iter().any(|e| e.manifest.id == m.id); + json!({ + "id": m.id, + "name": m.name, + "role": m.role, + "licence": m.licence, + "provenance": m.provenance, + "source": m.source, + "sizeBytes": m.size_bytes, + "sha256": m.sha256, + // how fast it is as a multiple of a ResNet-100 model (computers differ, ratios between models hardly do), and + // that in words; null when nobody has measured it + "speed": m.speed, + "speedText": m.speed.map(manifest::describe_speed).filter(|t| !t.is_empty()), + "known": from_catalog || known::all().iter().any(|k| k.id == m.id), + // from the user's own catalog file + "fromCatalog": from_catalog, + // a pinned address LightCraft can fetch it from (the user presses Download), and which site that is + "downloadHost": download_spec(s, &m.id).map(|d| known::host(&d.url).to_string()), + "installed": installed, + "selected": selected, + "accepted": accepted, + }) +} + +struct Inspection { + sha256: String, + size: u64, + file_name: String, + suggestion: Suggestion, +} + +/// Look at a model file. `known_sha` is its SHA-256 when something already checked it (a finished download), so a +/// file of hundreds of megabytes is not read again just to be hashed. +fn inspect_file(path: &str, cmd: &str, known_sha: Option<&str>, extra: &[ModelManifest]) -> Result { + let p = Path::new(path); + let meta = std::fs::metadata(p).map_err(|e| bad(cmd, format!("cannot read `{path}`: {e}")))?; + if !meta.is_file() { + return Err(bad(cmd, format!("`{path}` is not a file"))); + } + if meta.len() == 0 || meta.len() > MAX_MODEL_BYTES { + return Err(bad(cmd, "a model file must be between 1 byte and 512 MiB")); + } + let sha256 = match known_sha { + Some(h) => h.to_string(), + None => sha256_file(p).map_err(|e| fail(&format!("could not read {path}"), e))?, + }; + let file_name = p.file_name().map(|n| n.to_string_lossy().into_owned()).unwrap_or_default(); + let suggestion = match onnx::probe_path(p) { + Ok(info) => suggest_with(&info, &sha256, meta.len(), &file_name, extra), + Err(e) => Suggestion::Unsupported(e.to_string()), + }; + Ok(Inspection { sha256, size: meta.len(), file_name, suggestion }) +} + +fn list(s: &mut Session, _: &Value) -> Result { + refresh_catalog(s); + let dir = s.face_models_dir.clone(); + let (settings, installed) = match &dir { + Some(d) => (read_settings(d), installed_models(d)), + None => (FaceSettings::default(), Vec::new()), + }; + let on_disk = |id: &str| installed.iter().find(|i| i.manifest.id == id); + let mut models: Vec = Vec::new(); + for k in known::all() { + let found = on_disk(&k.id); + models.push(row(s, &k, found.is_some(), settings.embedder.as_deref() == Some(k.id.as_str()), found.map_or(&Value::Null, |f| &f.accepted))); + } + // the user's own catalog + for e in &s.face_catalog.entries { + let found = on_disk(&e.manifest.id); + let chosen = settings.embedder.as_deref() == Some(e.manifest.id.as_str()); + models.push(row(s, &e.manifest, found.is_some(), chosen, found.map_or(&Value::Null, |f| &f.accepted))); + } + for i in &installed { + let listed = known::all().iter().any(|k| k.id == i.manifest.id) || s.face_catalog.entries.iter().any(|e| e.manifest.id == i.manifest.id); + if !listed { + models.push(row(s, &i.manifest, true, settings.embedder.as_deref() == Some(i.manifest.id.as_str()), &i.accepted)); + } + } + Ok(json!({ + "dir": dir.as_ref().map(|d| d.display().to_string()), + "enabled": settings.enabled, + "embedder": settings.embedder, + // Native builds always have the pure-Rust runner; model files/background scanning remain native-only. + "runtime": cfg!(not(target_arch = "wasm32")), + "models": models, + // the user's own list of models to download: where it goes, how many it holds, and what was wrong with the rest + "catalog": { + "path": dir.as_ref().map(|d| d.join("catalog.json").display().to_string()), + "models": s.face_catalog.entries.len(), + "errors": s.face_catalog.errors, + }, + })) +} + +fn inspect(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.models.inspect"; + let path = str_param(p, "path").ok_or_else(|| bad(C, "missing `path`"))?; + refresh_catalog(s); + let ins = inspect_file(path, C, None, &catalog_manifests(s))?; + let installed = s.face_models_dir.as_deref().map(installed_models).unwrap_or_default(); + let (kind, model, assumptions, reason) = match &ins.suggestion { + Suggestion::Known(m) => ("known", Some(m), Vec::new(), None), + Suggestion::Draft { manifest, assumptions } => ("draft", Some(manifest), assumptions.clone(), None), + Suggestion::Unsupported(why) => ("unsupported", None, Vec::new(), Some(why.clone())), + }; + let already = model.is_some_and(|m| installed.iter().any(|i| i.manifest.sha256 == m.sha256 && i.manifest.id == m.id)); + Ok(json!({ + "path": path, + "fileName": ins.file_name, + "sizeBytes": ins.size, + "sha256": ins.sha256, + "kind": kind, + "model": model.map(|m| row(s, m, already, false, &Value::Null)), + "assumptions": assumptions, + "reason": reason, + "alreadyInstalled": already, + })) +} + +fn install(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.models.install"; + let path = str_param(p, "path").ok_or_else(|| bad(C, "missing `path`"))?; + let acknowledged = p.get("acknowledged").and_then(Value::as_bool).unwrap_or(false); + let activate = p.get("activate").and_then(Value::as_bool).unwrap_or(true); + refresh_catalog(s); + let ins = inspect_file(path, C, None, &catalog_manifests(s))?; + install_file(s, C, Path::new(path), ins, acknowledged, activate, false) +} + +/// Install the inspected model file `src`. A recognition model that passes its self-test becomes the one in use and +/// recognition is switched on (`activate`, the default): installing a model is how someone says they want it. +/// `verified` is for a download that was already checked against its recorded hash: it is moved into place instead +/// of being copied and read again. +fn install_file(s: &mut Session, c: &str, src: &Path, ins: Inspection, acknowledged: bool, activate: bool, verified: bool) -> Result { + let dir = models_dir(s, c)?; + let m = match ins.suggestion { + Suggestion::Known(m) | Suggestion::Draft { manifest: m, .. } => m, + Suggestion::Unsupported(why) => return Err(bad(c, format!("this model cannot be used yet: {why}"))), + }; + if !matches!(m.role, Role::Embedder | Role::Detector) { + return Err(bad(c, "only face detection and recognition models can be installed")); + } + manifest::validate(&m).map_err(|e| bad(c, e.to_string()))?; + if !acknowledged { + return Err(bad(c, "the licence has not been accepted: show the user its terms and pass `acknowledged: true` once they agree")); + } + let home = dir.join(&m.id); + std::fs::create_dir_all(&home).map_err(|e| fail("could not create the model's folder", e))?; + let (part, final_path) = (home.join("model.onnx.part"), home.join("model.onnx")); + let moved = verified && std::fs::rename(src, &part).is_ok(); + let placed = if moved { + Ok(()) + } else { + std::fs::copy(src, &part).map_err(|e| fail("could not copy the model (is the disk full?)", e)).and_then(|_| match sha256_file(&part) { + Ok(h) if h == ins.sha256 => Ok(()), + Ok(_) => Err(EngineError::Other("the copy does not match the original (the file changed while copying?)".into())), + Err(e) => Err(fail("could not check the copy", e)), + }) + }; + if let Err(e) = placed { + let _ = std::fs::remove_file(&part); + return Err(e); + } + // Validate the candidate before replacing a working installation. + let test = match self_test(&part, &m) { + Ok(t) => t, + Err(e) => { + let _ = std::fs::remove_file(&part); + let _ = std::fs::remove_dir(&home); + return Err(bad(c, e)); + } + }; + if let Err(e) = std::fs::rename(&part, &final_path) { + let _ = std::fs::remove_file(&part); + return Err(fail("could not finish the copy", e)); + } + write_atomic(&home.join("face-model.json"), &serde_json::to_vec_pretty(&m).map_err(|e| fail("manifest", e))?)?; + let accepted = json!({"acceptedAt": (s.clock)(), "licence": m.licence.name, "commercial": m.licence.commercial, "fileName": ins.file_name, "selfTest": test}); + write_atomic(&home.join("installed.json"), &serde_json::to_vec_pretty(&accepted).map_err(|e| fail("record", e))?)?; + // a file that was downloaded and waited in the staging folder has served: the installed copy is the model now + if src.parent() == Some(dir.join(face_download::STAGING).as_path()) { + let _ = std::fs::remove_file(src); + s.face_downloads.staged_file_used(src); + } + // the newest model is the one in use (the earlier one stays installed, and `faces.models.select` switches back) + if activate && m.role == Role::Embedder { + let mut st = read_settings(&dir); + st.embedder = Some(m.id.clone()); + st.enabled = true; + write_settings(&dir, &st)?; + } + let now = read_settings(&dir); + let selected = now.embedder.as_deref() == Some(m.id.as_str()); + Ok(json!({"installed": row(s, &m, true, selected, &accepted), "active": {"embedder": now.embedder, "enabled": now.enabled}})) +} + +/// What an error says, without the command it came from (for a line shown to the user). +fn plain(e: &EngineError) -> String { + match e { + EngineError::BadParams { msg, .. } => msg.clone(), + other => other.to_string(), + } +} + +/// Install every download that has arrived: each was accepted when it was started, so it is installed, chosen and +/// switched on without another question. Called by `faces.models.downloads` and by every frame's `faces.pump`. +pub(super) fn finish_downloads(s: &mut Session) { + const C: &str = "faces.models.download"; + let extra = catalog_manifests(s); + for (id, path, sha) in s.face_downloads.finished() { + let outcome = inspect_file(&path.to_string_lossy(), C, Some(&sha), &extra).and_then(|ins| install_file(s, C, &path, ins, true, true, true)); + match outcome { + Ok(_) => s.face_downloads.set_outcome(&id, State::Installed), + Err(e) => { + // a file that cannot be installed is of no use to anyone: throw it away, say why + let _ = std::fs::remove_file(&path); + s.face_downloads.set_outcome(&id, State::Failed(plain(&e))); + } + } + } +} + +/// `faces.models.download {id, acknowledged: true}`: fetch a model LightCraft knows an address for, in the background; +/// when it has arrived and checked out it is installed. +fn download(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.models.download"; + let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?; + let dir = models_dir(s, C)?; + refresh_catalog(s); + let spec = download_spec(s, id).ok_or_else(|| bad(C, "LightCraft has no download for that model: get the file from its page, then add it"))?; + if cfg!(target_arch = "wasm32") { + return Err(bad(C, "the web build downloads no models")); + } + if p.get("acknowledged").and_then(Value::as_bool) != Some(true) { + return Err(bad( + C, + "the licence has not been accepted: show the user the model's terms (see `faces.models.list`) and pass `acknowledged: true` once they agree", + )); + } + if installed_models(&dir).iter().any(|i| i.manifest.id == spec.id) { + return Err(bad(C, "that model is already installed")); + } + let host = spec.host().to_string(); + s.face_downloads.start(spec, &dir).map_err(|e| bad(C, e))?; + Ok(json!({"started": id, "from": host})) +} + +fn download_row(s: &Session, id: &str, state: &State) -> Value { + let from = download_spec(s, id).map(|d| d.host().to_string()); + match state { + State::Running { bytes, total } => json!({"id": id, "state": "running", "bytes": bytes, "total": total, "from": from}), + State::Done { path, .. } => json!({"id": id, "state": "done", "path": path.display().to_string(), "from": from}), + State::Installed => json!({"id": id, "state": "installed", "from": from}), + State::Failed(why) => json!({"id": id, "state": "failed", "error": why, "from": from}), + State::Cancelled => json!({"id": id, "state": "cancelled", "from": from}), + } +} + +/// `faces.models.downloads`: where each download stands (a finished one is installed first). +fn downloads(s: &mut Session, _: &Value) -> Result { + finish_downloads(s); + let all: Vec = s.face_downloads.snapshot().iter().map(|(id, st)| download_row(s, id, st)).collect(); + Ok(json!({"running": s.face_downloads.running(), "downloads": all})) +} + +/// `faces.models.downloadCancel {id}`: stop a download, or clear a finished one (a staged file is deleted). +fn download_cancel(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.models.downloadCancel"; + let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?; + Ok(json!({"discarded": s.face_downloads.discard(id)})) +} + +/// Load the model and run its self-test: `Ok(null)` in a build without the recognition runtime (nothing can +/// be checked), `Ok(result)` when it passes, `Err(reason)` when it does not load or fails a check. +#[cfg(not(target_arch = "wasm32"))] +fn self_test(path: &Path, m: &ModelManifest) -> std::result::Result { + if m.role == Role::Detector { + let bytes = read_capped(path, 8 << 20).ok_or("the detector cannot be read within its 8 MiB limit")?; + let detector = lightcraft_faces::yunet::Detector::new(&bytes).map_err(|e| e.to_string())?; + if Some(detector.input_side()) != usize::try_from(m.input.width).ok() + || m.input.width != m.input.height + || !matches!(&m.output, manifest::OutputSpec::Detector { decoder } if decoder == "yunet-v2") + { + return Err("unsupported detector input or output decoder".into()); + } + detector.detect(&[0u8; 12], 2, 2, &lightcraft_faces::yunet::Options::default()).map_err(|e| e.to_string())?; + return Ok(json!({"ok": true})); + } + let embedder = lightcraft_faces::runtime::Embedder::load(path, m).map_err(|e| e.to_string())?; + let t = embedder.self_test(); + if !t.ok { + let failed: Vec<&str> = t.checks.iter().filter(|(_, ok)| !*ok).map(|(c, _)| c.as_str()).collect(); + return Err(format!("the model failed its self-test: it does not {}", failed.join(", and does not "))); + } + serde_json::to_value(&t).map_err(|e| e.to_string()) +} + +#[cfg(target_arch = "wasm32")] +fn self_test(_: &Path, _: &ModelManifest) -> std::result::Result { + Ok(Value::Null) +} + +/// `faces.models.test {id}`: run an installed model's self-test again. +fn test(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.models.test"; + let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?; + let dir = models_dir(s, C)?; + if !cfg!(not(target_arch = "wasm32")) { + return Err(bad(C, "this build cannot run face recognition models")); + } + let found = installed_models(&dir).into_iter().find(|i| i.manifest.id == id).ok_or_else(|| bad(C, "that model is not installed"))?; + if found.manifest.role != Role::Embedder { + return Err(bad(C, "only recognition models have a self-test")); + } + let result = self_test(&dir.join(id).join("model.onnx"), &found.manifest); + let (ok, detail) = match &result { + Ok(v) => (true, v.clone()), + Err(e) => (false, json!(e)), + }; + // keep the latest outcome with the model's record + let mut record = found.accepted; + if let Some(o) = record.as_object_mut() { + o.insert("selfTest".into(), if ok { detail.clone() } else { json!({"ok": false, "error": detail}) }); + write_atomic(&dir.join(id).join("installed.json"), &serde_json::to_vec_pretty(&record).map_err(|e| fail("record", e))?)?; + } + Ok(json!({"id": id, "ok": ok, "result": detail})) +} + +fn remove(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.models.remove"; + let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?; + let dir = models_dir(s, C)?; + if !manifest::valid_id(id) { + return Err(bad(C, "not a model id")); + } + let home = dir.join(id); + // only a folder that is a model (it has a manifest) is ever deleted + if !home.join("face-model.json").is_file() { + return Err(bad(C, "that model is not installed")); + } + std::fs::remove_dir_all(&home).map_err(|e| fail("could not remove the model", e))?; + let mut st = read_settings(&dir); + if st.embedder.as_deref() == Some(id) { + // another installed recogniser takes over, so recognition keeps working; with none left nothing is chosen + st.embedder = installed_models(&dir).into_iter().find(|i| i.manifest.id != id && i.manifest.role == Role::Embedder).map(|i| i.manifest.id); + write_settings(&dir, &st)?; + } + Ok(json!({"removed": id, "embedder": st.embedder})) +} + +fn select(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.models.select"; + let dir = models_dir(s, C)?; + let id = match p.get("id") { + None | Some(Value::Null) => None, + Some(v) => Some(v.as_str().ok_or_else(|| bad(C, "`id` must be a model id or null"))?.to_string()), + }; + if let Some(id) = &id + && !installed_models(&dir).iter().any(|i| &i.manifest.id == id && i.manifest.role == Role::Embedder) + { + return Err(bad(C, "that recognition model is not installed")); + } + let mut st = read_settings(&dir); + st.embedder = id; + write_settings(&dir, &st)?; + Ok(json!({"embedder": st.embedder})) +} + +fn enable(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.enable"; + let dir = models_dir(s, C)?; + let mut st = read_settings(&dir); + st.enabled = p.get("enabled").and_then(Value::as_bool).unwrap_or(!st.enabled); + write_settings(&dir, &st)?; + // the background scan notices at once, not within the second it trusts its last look + #[cfg(not(target_arch = "wasm32"))] + { + s.faces.enabled_seen = None; + } + Ok(json!({"enabled": st.enabled})) +} + +pub fn specs() -> Vec { + vec![ + cmd!(query "faces.models.list", "Face Models", [], None, "{} → {dir, enabled, embedder, runtime, models: [{id, name, role, licence{name, commercial, url, notice}, provenance, source, sizeBytes, known, installed, selected}]}", always, list), + cmd!(query "faces.models.inspect", "Inspect Face Model File", [], None, "{path} → what a .onnx file is: {kind: known | draft | unsupported, model, assumptions, reason, alreadyInstalled}; installs nothing", always, inspect), + cmd!(query "faces.models.install", "Install Face Model", [], None, "{path, acknowledged: true, activate?: true} → {installed, active: {embedder, enabled}} — copy a .onnx face recognition model into the models folder. `acknowledged` must be true: the user has been shown its licence (see inspect) and accepted it. Unless `activate` is false the model becomes the one in use and recognition is switched on; an earlier model stays installed", always, install), + cmd!(query "faces.models.download", "Download Face Model", [], None, "{id, acknowledged: true} → {started, from} — fetch a recognition model LightCraft has a pinned address for (see `downloadHost` in the list), in the background over pure-Rust https (lightcraft-fetch). `acknowledged` must be true: the user has been shown the model's terms and accepted them. It is checked against its size and SHA-256, then installed, chosen and switched on by itself; `faces.models.downloads` shows how far it is", always, download), + cmd!(query "faces.models.downloads", "Face Model Downloads", [], None, "{} → {running, downloads: [{id, state: running | done | installed | failed | cancelled, bytes, total, error, from}]} — also installs any download that has arrived; `installed` stays listed until `faces.models.downloadCancel` clears it", always, downloads), + cmd!(query "faces.models.downloadCancel", "Cancel Face Model Download", [], None, "{id} → {discarded} — stop a download, or delete a finished one that was not installed", always, download_cancel), + cmd!(query "faces.models.test", "Test Face Model", [], None, "{id} → {ok, result: {loadMs, embedMs, dimension, checks}} — load an installed recognition model and check it gives sensible faces; needs the recognition runtime", always, test), + cmd!(query "faces.models.remove", "Remove Face Model", [], None, "{id} — delete an installed model", always, remove), + cmd!(query "faces.models.select", "Choose Face Recognition Model", [], None, "{id: installed recogniser | null}", always, select), + cmd!(query "faces.enable", "Face Recognition On/Off", [], None, "{enabled?: bool} (toggles when omitted)", always, enable), + ] +} diff --git a/crates/engine/src/cmd/face_recognize.rs b/crates/engine/src/cmd/face_recognize.rs new file mode 100644 index 00000000..9efdb146 --- /dev/null +++ b/crates/engine/src/cmd/face_recognize.rs @@ -0,0 +1,1030 @@ +//! Recognition commands: `faces.index` (embed the faces in the library), `faces.suggest` (who might an unnamed +//! face be?), `faces.person` (one person's faces, and the unnamed faces that look like them), `faces.unnamed` (every +//! unnamed face, look-alikes together) and `faces.setName` / `faces.nameFaces` (name one face, or many at once, which is +//! what teaches the suggestions). +//! +//! Everything here is a suggestion for the user to confirm; nothing is ever named automatically. All of it is +//! catalog-only (no sidecar is written) and `faces.setName` is undoable. + +use lightcraft_catalog::Op; +use serde_json::{Value, json}; + +use super::{CommandSpec, bad, cmd}; +use crate::{Result, Session}; + +/// The `name` parameter of the naming commands: a name, or none (`null`, empty) to clear it. +fn name_param(p: &Value, c: &str) -> Result> { + let name = match p.get("name") { + None | Some(Value::Null) => None, + Some(v) => Some(v.as_str().ok_or_else(|| bad(c, "`name` must be text or null"))?.trim().to_string()).filter(|n| !n.is_empty()), + }; + if let Some(n) = &name + && (n.chars().count() > 200 || n.chars().any(char::is_control)) + { + return Err(bad(c, "a name is at most 200 characters, without control characters")); + } + Ok(name) +} + +/// Give a face region its name. Naming makes the face yours: a later "Detect Faces" run no longer replaces it. +fn apply_name(region: &mut lightcraft_meta::Region, name: &Option) { + region.name = name.clone(); + if name.is_some() && region.description.as_deref().is_some_and(|d| d.starts_with(super::face_detect::MARK)) { + region.description = None; + } +} + +/// `faces.nameFaces {faces: [{photo, index}], name}`: name (or clear) many face regions at once, as one undo step: what +/// naming a group of look-alikes does. +fn name_faces(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.nameFaces"; + let name = name_param(p, C)?; + let list = p.get("faces").and_then(Value::as_array).ok_or_else(|| bad(C, "missing `faces`"))?; + if list.len() > 5000 { + return Err(bad(C, "at most 5000 faces at a time")); + } + let mut by_photo: std::collections::BTreeMap> = std::collections::BTreeMap::new(); + for f in list { + let (Some(photo), Some(index)) = + (f.get("photo").and_then(Value::as_u64), f.get("index").and_then(Value::as_u64).and_then(|i| usize::try_from(i).ok())) + else { + return Err(bad(C, "each face needs a `photo` and an `index`")); + }; + by_photo.entry(photo).or_default().push(index); + } + let (mut ops, mut named) = (Vec::new(), 0usize); + for (photo, indexes) in by_photo { + let id = lightcraft_catalog::PhotoId(photo); + let Some(photo) = s.catalog.photo(id) else { continue }; + let mut meta = photo.meta.clone(); + let before = named; + for i in indexes { + if let Some(region) = meta.regions.get_mut(i).filter(|r| r.kind == lightcraft_meta::RegionKind::Face) { + apply_name(region, &name); + named += 1; + } + } + if named > before { + ops.push(Op::SetMeta { id, meta: Box::new(meta) }); + } + } + if ops.is_empty() { + return Err(bad(C, "none of those faces exist")); + } + let photos = ops.len(); + s.commit(if name.is_some() { "Name Faces" } else { "Clear Face Names" }, Op::Batch { ops })?; + s.skip_auto_write = true; + Ok(json!({"named": named, "photos": photos, "name": name})) +} + +/// `faces.setName {id?, index, name}`: name (or, with `null` or an empty name, un-name) one face region. Naming +/// makes the face yours: a later "Detect Faces" run no longer replaces it. +fn set_name(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.setName"; + let id = p.get("id").and_then(Value::as_u64).map(lightcraft_catalog::PhotoId).or_else(|| s.active()).ok_or_else(|| bad(C, "no photo"))?; + let index = p.get("index").and_then(Value::as_u64).and_then(|i| usize::try_from(i).ok()).ok_or_else(|| bad(C, "missing or invalid `index`"))?; + let name = name_param(p, C)?; + let mut meta = s.catalog.photo(id).ok_or_else(|| bad(C, "no such photo"))?.meta.clone(); + let region = meta.regions.get_mut(index).ok_or_else(|| bad(C, "no such region"))?; + apply_name(region, &name); + s.commit(if name.is_some() { "Name Face" } else { "Clear Face Name" }, Op::SetMeta { id, meta: Box::new(meta) })?; + s.skip_auto_write = true; + Ok(json!({"id": id.0, "index": index, "name": name})) +} + +/// Most faces `faces.unnamed` lists. +const MAX_UNNAMED: usize = 3000; + +/// How the unnamed faces were laid out. +struct Arranged { + /// Faces that look alike are next to each other (only the ones already embedded can be). + ordered: bool, + /// Recognition is on and a model is loaded. + ready: bool, + /// A name for some faces, when the faces already named say so with confidence: `(photo, index)` to `(name, score)`. + suggestions: std::collections::HashMap<(u64, usize), (String, f32)>, +} + +/// `faces.unnamed {limit?: 600}`: the faces nobody has named, to be named a group at a time. With recognition running, +/// faces that look alike are next to each other, and a face the named faces recognise carries the name they suggest. +fn unnamed(s: &mut Session, p: &Value) -> Result { + let limit = p.get("limit").and_then(Value::as_u64).map_or(600, |n| usize::try_from(n).unwrap_or(600)).clamp(1, MAX_UNNAMED); + let mut all: Vec<(u64, usize, lightcraft_geom::Rect)> = Vec::new(); + for ph in s.catalog.photos().filter(|ph| ph.in_library()) { + for (index, r) in ph.meta.regions.iter().enumerate() { + if r.kind == lightcraft_meta::RegionKind::Face && r.name.is_none() { + all.push((ph.id.0, index, r.rect)); + } + } + } + all.sort_by_key(|(photo, index, _)| (*photo, *index)); + let total = all.len(); + let arranged = imp::arrange_unnamed(s, &mut all, limit); + all.truncate(limit); + let faces: Vec = all + .iter() + .map(|(photo, index, rect)| { + let view = s.face_view(lightcraft_catalog::PhotoId(*photo), *rect); + let suggestion = arranged.suggestions.get(&(*photo, *index)).map(|(name, score)| json!({"name": name, "score": score})); + json!({"photo": photo, "index": index, "rect": rect_json(rect), "view": rect_json(&view), "suggestion": suggestion}) + }) + .collect(); + Ok(json!({"total": total, "faces": faces, "ordered": arranged.ordered, "ready": arranged.ready})) +} + +/// What looking for faces that resemble a person found. +struct Similar { + /// `{photo, index, rect, score}` of unnamed faces, the most alike first. + items: Vec, + /// Recognition is on and a model is loaded, so an empty list means "none yet" rather than "not looking". + ready: bool, + /// Photos the background scan has still to do. + pending: usize, +} + +fn rect_json(r: &lightcraft_geom::Rect) -> Value { + json!({"x0": r.x0, "y0": r.y0, "x1": r.x1, "y1": r.y1}) +} + +/// Most confirmed faces `faces.person` lists. +const MAX_CONFIRMED: usize = 5000; + +/// `faces.person {name, more?}`: the faces named `name` (case-insensitively) across the library, and, when recognition is +/// running, the unnamed faces that look most like them, to be confirmed one by one. +fn person(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.person"; + let name = p.get("name").and_then(Value::as_str).map(str::trim).filter(|n| !n.is_empty()).ok_or_else(|| bad(C, "missing `name`"))?; + let key = name.to_lowercase(); + let limit = p.get("more").and_then(Value::as_u64).map_or(120, |n| usize::try_from(n).unwrap_or(120)).min(1000); + let (mut shown, mut total, mut confirmed) = (name.to_string(), 0usize, Vec::new()); + for ph in s.catalog.photos().filter(|ph| ph.in_library()) { + for (index, r) in ph.meta.regions.iter().enumerate() { + let Some(n) = r.name.as_deref().filter(|n| r.kind == lightcraft_meta::RegionKind::Face && n.trim().to_lowercase() == key) else { + continue; + }; + if total == 0 { + shown = n.trim().to_string(); + } + total += 1; + if confirmed.len() < MAX_CONFIRMED { + confirmed.push(json!({"photo": ph.id.0, "index": index, "rect": rect_json(&r.rect), "view": rect_json(&s.face_view(ph.id, r.rect))})); + } + } + } + let similar = if total > 0 { imp::similar(s, &key, limit) } else { Similar { items: Vec::new(), ready: false, pending: 0 } }; + Ok( + json!({"name": shown, "total": total, "confirmed": confirmed, "more": similar.items, "ready": similar.ready, "pendingPhotos": similar.pending}), + ) +} + +#[cfg(not(target_arch = "wasm32"))] +mod imp { + use std::sync::Arc; + use std::time::{Duration, Instant}; + + use lightcraft_catalog::{Op, Photo, PhotoId}; + use lightcraft_develop::DevelopSettings; + use lightcraft_faces::matching; + use lightcraft_faces::runtime::Embedder; + use lightcraft_geom::Rect; + use lightcraft_meta::{Region, RegionKind}; + use serde_json::{Value, json}; + + use super::super::bad; + use super::super::face_detect::{LABEL, MARK}; + use super::super::face_models::{finish_downloads, installed_models, read_settings}; + use super::Similar; + use crate::faces_worker::{Done, EDGE, PREFETCH, Pace, Prepared, Worker, max_workers, process, target_workers}; + use crate::{Result, Session}; + + /// Defaults for suggestions, until a model's own threshold is known. + const DEFAULT_THRESHOLD: f32 = 0.45; + const DEFAULT_MARGIN: f32 = 0.06; + /// How long a look at the models folder is trusted by the per-frame `faces.pump`. + const RECHECK: Duration = Duration::from_secs(1); + /// How often the background pump writes new embeddings to the cache file (appending is cheap; this keeps it to a few + /// writes a minute, and a session that ends writes the rest). + const SAVE_EVERY: Duration = Duration::from_secs(3); + + /// The chosen recognition model, loaded, with the index reset to it when it changed. `fresh` looks at the models + /// folder again; otherwise a look made in the last second is trusted (the UI asks every frame). + fn current_with(s: &mut Session, cmd: &str, fresh: bool) -> Result> { + if !fresh + && let Some((_, e)) = &s.faces.embedder + && s.faces.checked.is_some_and(|t| t.elapsed() < RECHECK) + { + return Ok(e.clone()); + } + let dir = s.face_models_dir.clone().ok_or_else(|| bad(cmd, "this build has nowhere to keep face models"))?; + let id = read_settings(&dir).embedder.ok_or_else(|| bad(cmd, "no recognition model is chosen: add one in Settings > Faces and press Use"))?; + let installed = installed_models(&dir) + .into_iter() + .find(|i| i.manifest.id == id) + .ok_or_else(|| bad(cmd, "the chosen recognition model is not installed"))?; + let tag = format!("{id}@{}", installed.manifest.sha256.as_deref().unwrap_or("unknown")); + s.faces.checked = Some(Instant::now()); + // which photos the detector has searched is kept beside the library (or in memory when it is not on disk) + let scan_path = s.library.as_ref().filter(|l| l.on_disk).map(|l| l.dir.join("face-scanned.bin")); + if s.faces.scanned.path() != scan_path.as_deref() { + s.faces.scanned.reset(scan_path); + } + if let Some((have, e)) = &s.faces.embedder + && *have == tag + { + return Ok(e.clone()); + } + let embedder = Arc::new(Embedder::load(&dir.join(&id).join("model.onnx"), &installed.manifest).map_err(|e| bad(cmd, e.to_string()))?); + let dim = match installed.manifest.output { + lightcraft_faces::OutputSpec::Embedding { dim } => dim as usize, + _ => return Err(bad(cmd, "the chosen model is not a recognition model")), + }; + let cache = s.library.as_ref().filter(|l| l.on_disk).map(|l| l.dir.join(format!("face-embeddings-{id}.bin"))); + s.faces.index.reset(&tag, dim, cache); + s.faces.queue.clear(); + s.faces.queue_stamp = None; + s.faces.in_flight.clear(); + s.faces.embedder = Some((tag, embedder.clone())); + Ok(embedder) + } + + fn current(s: &mut Session, cmd: &str) -> Result> { + current_with(s, cmd, true) + } + + /// Whether the background scan should look for faces in this photo: it is in the library, has no face regions at all + /// (nothing from XMP, nothing drawn or named, no earlier detection) and has not been searched yet. + fn needs_search(s: &Session, photo: &Photo) -> bool { + photo.in_library() + && !photo.meta.regions.iter().any(|r| r.kind == RegionKind::Face) + && !s.faces.scanned.contains(photo.id.0) + && !s.faces.scan_failed.contains(&photo.id.0) + } + + /// The photo's faces that still need embedding (and, with `search`, the faces to find in it first), with the render job + /// for its picture. `None` when there is nothing to do. + fn prepare(s: &mut Session, id: PhotoId, search: bool) -> Option { + let tag = s.faces.index.tag.clone(); + let photo = s.catalog.photo(id)?.clone(); + let todo: Vec = + photo.meta.regions.iter().filter(|r| r.kind == RegionKind::Face && s.faces.index.needs(id.0, &r.rect)).map(|r| r.rect).collect(); + // finding faces needs the detector the user may not have downloaded: such a photo is left unsearched, not marked as searched + let detector = crate::cmd::face_detect::detector(s).ok(); + let detect = search && detector.is_some() && needs_search(s, &photo); + if todo.is_empty() && !detect { + return None; + } + let job = s.preview_job(id, EDGE, EDGE, false, &DevelopSettings::default())?; + let preview = s.scan_preview_of(&photo); + Some(Prepared { tag, epoch: s.faces.epoch, id, job, preview, todo, detect, detector }) + } + + /// The faces the detector found in a photo that had none become its face regions: unnamed, marked as detected (so a + /// manual Detect Faces replaces them), and not an undo step, since the user did not do this. + fn add_found(s: &mut Session, id: PhotoId, found: &[Rect]) { + if found.is_empty() { + return; + } + let Some(photo) = s.catalog.photo(id) else { return }; + // the user (or an import) got there first: leave their faces alone + if photo.meta.regions.iter().any(|r| r.kind == RegionKind::Face) { + return; + } + let mut meta = photo.meta.clone(); + for rect in found { + meta.regions.push(Region { rect: *rect, kind: RegionKind::Face, name: None, description: Some(format!("{MARK}{LABEL}")) }); + } + if s.apply_system(Op::SetMeta { id, meta: Box::new(meta) }).is_ok() { + // regions are not written to XMP, so no sidecar may be rewritten for this + s.skip_auto_write = true; + } + } + + /// Take a finished photo into the index (ignored when it was made for another model or another library). Returns + /// (embedded, failed). + fn ingest(s: &mut Session, d: Done) -> (usize, usize) { + s.faces.in_flight.remove(&d.id); + if d.epoch != s.faces.epoch || d.tag != s.faces.index.tag { + return (0, 0); + } + if d.detect_wanted { + match &d.found { + Some(found) => { + s.faces.scanned.insert(d.id.0); + add_found(s, d.id, found); + } + None => { + s.faces.scan_failed.insert(d.id.0); + } + } + } + let (mut embedded, mut failed) = (0, 0); + for (rect, v, view) in d.results { + match v { + Some(v) if v.len() == s.faces.index.dim => { + s.faces.index.insert_with_view(d.id.0, &rect, v, &view); + embedded += 1; + } + _ => { + s.faces.index.skip(d.id.0, &rect); + failed += 1; + } + } + } + (embedded, failed) + } + + /// Take in whatever the worker has finished. + fn collect(s: &mut Session) -> (usize, usize) { + let finished = s.faces.worker.as_ref().map(Worker::finished).unwrap_or_default(); + finished.into_iter().fold((0, 0), |(e, f), d| { + let (de, df) = ingest(s, d); + (e + de, f + df) + }) + } + + /// Embed every face region of `id` the index does not have yet, here and now. Returns (embedded, failed). + fn embed_photo(s: &mut Session, id: PhotoId, e: &Embedder) -> Result<(usize, usize)> { + let Some(job) = prepare(s, id, false) else { return Ok((0, 0)) }; + Ok(ingest(s, process(job, e))) + } + + /// Photos with a face that needs embedding (and, with `search`, photos to look for faces in): those the caller asked + /// for first, then those with named faces (they are the gallery everything else is compared with), then the rest, and + /// last the photos that only need searching. + fn pending(s: &Session, first: &[PhotoId], search: bool) -> Vec { + let mut out: Vec<(u8, PhotoId)> = Vec::new(); + for p in s.catalog.photos().filter(|p| p.in_library()) { + let faces: Vec<_> = p.meta.regions.iter().filter(|r| r.kind == RegionKind::Face).collect(); + let embed = faces.iter().any(|r| s.faces.index.needs(p.id.0, &r.rect)); + if embed || (search && needs_search(s, p)) { + let rank = if first.contains(&p.id) { + 0 + } else if faces.iter().any(|r| r.name.is_some()) { + 1 + } else if embed { + 2 + } else { + 3 + }; + out.push((rank, p.id)); + } + } + out.sort_by_key(|(rank, id)| (*rank, id.0)); + out.into_iter().map(|(_, id)| id).collect() + } + + pub fn index(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.index"; + let embedder = current(s, C)?; + collect(s); + let budget = Duration::from_millis(p.get("budgetMs").and_then(Value::as_u64).unwrap_or(200).min(600_000)); + let first: Vec = + p.get("ids").and_then(Value::as_array).map(|a| a.iter().filter_map(Value::as_u64).map(PhotoId).collect()).unwrap_or_default(); + let todo = pending(s, &first, false); + let started = Instant::now(); + let (mut embedded, mut failed, mut done) = (0, 0, 0); + for id in &todo { + // at least one photo per call, so a slow model still makes progress; none when only the status is asked for + if budget.is_zero() || (done > 0 && started.elapsed() >= budget) { + break; + } + let (e, f) = embed_photo(s, *id, &embedder)?; + (embedded, failed, done) = (embedded + e, failed + f, done + 1); + } + if let Err(e) = s.faces.index.save() { + log::warn!("could not save the face embeddings: {e}"); + } + Ok(json!({ + "model": s.faces.index.tag, + "embedded": embedded, + "failed": failed, + "photosDone": done, + "pendingPhotos": todo.len().saturating_sub(done), + "indexedFaces": s.faces.index.len(), + "ms": started.elapsed().as_secs_f64() * 1000.0, + })) + } + + fn str_scope(p: &Value) -> Option<&str> { + p.get("scope").and_then(Value::as_str) + } + + /// Seconds since 1970 of an ISO capture time (`2024-09-20T10:48:45…`, any zone suffix ignored): only for telling + /// shots taken close together, so the zone does not matter. + fn capture_secs(iso: &str) -> Option { + let n = |a: usize, b: usize| iso.get(a..b)?.parse::().ok(); + let (y, m, d, hh, mm, ss) = (n(0, 4)?, n(5, 7)?, n(8, 10)?, n(11, 13)?, n(14, 16)?, n(17, 19)?); + if !(1..=12).contains(&m) || !(1..=31).contains(&d) { + return None; + } + // days from civil (proleptic Gregorian) + let y = if m <= 2 { y - 1 } else { y }; + let era = y.div_euclid(400); + let yoe = y - era * 400; + let doy = (153 * (if m > 2 { m - 3 } else { m + 9 }) + 2) / 5 + d - 1; + let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy; + Some((era * 146_097 + doe - 719_468) * 86_400 + hh * 3600 + mm * 60 + ss) + } + + /// `faces.evaluate`: leave-one-out over the library's named faces, to see how well recognition works on *your* + /// photos and which threshold to trust. Each named face in turn is treated as unknown and ranked against the + /// other named faces, ignoring those from shots taken within `burstSecs` of it (near-duplicates would make + /// everything look easy). Reports counts only. + pub fn evaluate(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.evaluate"; + let embedder = current(s, C)?; + collect(s); + let burst = p.get("burstSecs").and_then(Value::as_i64).unwrap_or(5).clamp(0, 86_400 * 7); + let margin = p.get("margin").and_then(Value::as_f64).filter(|v| v.is_finite()).map_or(DEFAULT_MARGIN, |v| v as f32); + let thresholds: Vec = p + .get("thresholds") + .and_then(Value::as_array) + .map(|a| a.iter().filter_map(Value::as_f64).filter(|v| v.is_finite()).map(|v| v as f32).take(64).collect()) + .filter(|v: &Vec| !v.is_empty()) + .unwrap_or_else(|| (4..=16).map(|i| i as f32 * 0.05).collect()); + let budget = Duration::from_millis(p.get("budgetMs").and_then(Value::as_u64).unwrap_or(0).min(3_600_000)); + let started = Instant::now(); + // `scope: "visible"` looks only at the photos the current filter shows (say one year, or one album) + let scope: Option> = (str_scope(p) == Some("visible")).then(|| s.visible_cloned().into_iter().collect()); + let in_scope = |id: PhotoId| scope.as_ref().is_none_or(|v| v.contains(&id)); + let todo: Vec = pending(s, &[], false).into_iter().filter(|id| in_scope(*id)).collect(); + let mut done = 0; + for id in &todo { + if started.elapsed() >= budget { + break; + } + embed_photo(s, *id, &embedder)?; + done += 1; + } + // every named face with its embedding, the photo it is in, and when that was taken + struct Named { + name: String, + emb: Arc<[f32]>, + photo: PhotoId, + when: Option, + } + let named: Vec = s + .catalog + .photos() + .filter(|ph| ph.in_library() && in_scope(ph.id)) + .flat_map(|ph| { + let when = ph.captured.as_deref().and_then(capture_secs); + ph.meta + .regions + .iter() + .filter(|r| r.kind == RegionKind::Face) + .filter_map(|r| Some(Named { name: r.name.clone()?, emb: s.faces.index.get(ph.id.0, &r.rect)?.clone(), photo: ph.id, when })) + .collect::>() + }) + .collect(); + let near = |a: &Named, b: &Named| a.photo == b.photo || matches!((a.when, b.when), (Some(x), Some(y)) if (x - y).abs() <= burst); + let mut sweep: Vec<(f32, usize, usize)> = thresholds.iter().map(|t| (*t, 0, 0)).collect(); + let (mut queries, mut top1, mut unmatchable) = (0usize, 0usize, 0usize); + for q in &named { + let others: Vec<&Named> = named.iter().filter(|g| !near(q, g)).collect(); + if !others.iter().any(|g| g.name.eq_ignore_ascii_case(&q.name)) { + unmatchable += 1; + continue; + } + let refs: Vec<(&str, &[f32])> = others.iter().map(|g| (g.name.as_str(), &g.emb[..])).collect(); + let ranked = matching::rank(&q.emb, &refs); + queries += 1; + if ranked.first().is_some_and(|r| r.name.eq_ignore_ascii_case(&q.name)) { + top1 += 1; + } + for (t, suggested, correct) in sweep.iter_mut() { + if let Some(sg) = matching::suggest(&ranked, *t, margin) { + *suggested += 1; + if sg.name.eq_ignore_ascii_case(&q.name) { + *correct += 1; + } + } + } + } + let pct = |n: usize, d: usize| if d == 0 { Value::Null } else { json!((n as f64 / d as f64 * 1000.0).round() / 1000.0) }; + Ok(json!({ + "model": s.faces.index.tag, + "namedFaces": named.len(), + "queries": queries, + "unmatchable": unmatchable, + "pendingPhotos": todo.len().saturating_sub(done), + "top1": pct(top1, queries), + "margin": margin, + "burstSecs": burst, + "sweep": sweep.iter().map(|(t, sg, ok)| json!({"threshold": t, "suggested": sg, "correct": ok, "precision": pct(*ok, *sg), "recall": pct(*ok, queries)})).collect::>(), + })) + } + + /// `faces.pump`: cheap to call every frame. Takes in what the background worker has finished and, when it is idle, + /// hands it the next photo with faces to embed. Does nothing (and says so) until recognition is switched on in + /// Settings with a model chosen. + pub fn pump(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.pump"; + // how hard the scan may work now, from what the user is doing (the app says); nothing new is started while paused, + // and what is already running finishes + let pace = Pace::parse(p.get("pace").and_then(Value::as_str)); + let target = target_workers(pace); + // a model that has finished downloading is installed and switched on here, whether or not recognition was on + finish_downloads(s); + let enabled = s.face_models_dir.clone().is_some_and(|d| enabled_cached(s, &d)); + if !enabled { + return Ok(json!({"active": false})); + } + if s.faces.retry_at.is_some_and(|t| Instant::now() < t) { + return Ok(json!({"active": false})); + } + let Ok(embedder) = current_with(s, C, false) else { + // the chosen model is gone or broken: look again in a few seconds, not every frame + s.faces.retry_at = Some(Instant::now() + Duration::from_secs(3)); + return Ok(json!({"active": false})); + }; + s.faces.retry_at = None; + let (embedded, _) = collect(s); + if (s.faces.index.has_unsaved() || s.faces.scanned.has_unsaved()) && s.faces.saved_at.is_none_or(|t| t.elapsed() >= SAVE_EVERY) { + if let Err(e) = s.faces.index.save() { + log::warn!("could not save the face embeddings: {e}"); + } + if let Err(e) = s.faces.scanned.save() { + log::warn!("could not save the list of searched photos: {e}"); + } + s.faces.saved_at = Some(Instant::now()); + } + if s.faces.in_flight.is_empty() && s.faces.queue.is_empty() && s.faces.queue_stamp != Some(s.catalog.revision) { + // the queue is taken from the end, so it is made last first + let mut queue = pending(s, &[], true); + queue.reverse(); + s.faces.queue = queue; + s.faces.queue_stamp = Some(s.catalog.revision); + } + if target > 0 && s.faces.worker.is_none() && !s.faces.queue.is_empty() { + s.faces.worker = Worker::start(max_workers()); + } + // the workers are allowed `target` photos at once; a couple more wait behind them, so a worker that finishes has its + // next photo at once instead of waiting for the next frame to hand it one + if let Some(w) = &s.faces.worker { + w.allow(target, pace); + } + while target > 0 && s.faces.in_flight.len() < target + PREFETCH { + let Some(id) = s.faces.queue.pop() else { break }; + let Some(job) = prepare(s, id, true) else { continue }; + if s.faces.worker.as_ref().is_some_and(|w| w.submit(embedder.clone(), job)) { + s.faces.in_flight.insert(id); + } else { + break; + } + } + Ok(json!({ + "active": true, + "inFlight": s.faces.in_flight.len(), + "pendingPhotos": s.faces.queue.len() + s.faces.in_flight.len(), + "indexedFaces": s.faces.index.len(), + "searchedPhotos": s.faces.scanned.len(), + "workers": target, + "embedded": embedded, + })) + } + + /// The unnamed faces that look most like the person whose lower-cased name is `key`, best first, at most `limit`: each + /// is scored like a suggestion (the mean of its two closest faces of theirs) and kept above four fifths of the + /// model's suggestion bar, since a person looks at these before anything is named. Only faces already embedded can be + /// found; `ready` says whether recognition is running at all. + pub fn similar(s: &mut Session, key: &str, limit: usize) -> Similar { + let none = Similar { items: Vec::new(), ready: false, pending: 0 }; + let Some(dir) = s.face_models_dir.clone() else { return none }; + if !enabled_cached(s, &dir) { + return none; + } + let Ok(embedder) = current_with(s, "faces.person", false) else { return none }; + let pending = s.faces.queue.len() + s.faces.in_flight.len(); + let (mut gallery, mut candidates): (Vec>, Vec<(u64, usize, Rect, Arc<[f32]>)>) = (Vec::new(), Vec::new()); + for ph in s.catalog.photos().filter(|ph| ph.in_library()) { + for (index, r) in ph.meta.regions.iter().enumerate().filter(|(_, r)| r.kind == RegionKind::Face) { + let Some(e) = s.faces.index.get(ph.id.0, &r.rect) else { continue }; + match r.name.as_deref() { + Some(n) if n.trim().to_lowercase() == key => gallery.push(e.clone()), + Some(_) => {} + None => candidates.push((ph.id.0, index, r.rect, e.clone())), + } + } + } + let ready = Similar { items: Vec::new(), ready: true, pending }; + if gallery.is_empty() { + return ready; + } + let bar = embedder.manifest().thresholds.match_cosine.unwrap_or(DEFAULT_THRESHOLD) * 0.8; + let refs: Vec<(&str, &[f32])> = gallery.iter().map(|e| ("person", &e[..])).collect(); + let mut scored: Vec<(f32, u64, usize, Rect)> = candidates + .into_iter() + .filter_map(|(photo, index, rect, e)| { + let score = matching::rank(&e, &refs).first()?.score; + (score >= bar).then_some((score, photo, index, rect)) + }) + .collect(); + scored.sort_by(|a, b| b.0.total_cmp(&a.0).then(a.1.cmp(&b.1)).then(a.2.cmp(&b.2))); + scored.truncate(limit); + let items = scored + .into_iter() + .map(|(score, photo, index, rect)| { + let view = s.face_view(PhotoId(photo), rect); + json!({"photo": photo, "index": index, "rect": super::rect_json(&rect), "view": super::rect_json(&view), "score": score}) + }) + .collect(); + Similar { items, ..ready } + } + + /// Most faces put in order of likeness (the order is a pass over every pair of them). + const MAX_ARRANGED: usize = 1500; + + /// Faces in an order that puts look-alikes together: start anywhere, then always go to the closest face not yet placed. + fn nearest_chain(embeddings: &[&[f32]]) -> Vec { + let n = embeddings.len(); + let (mut placed, mut order) = (vec![false; n], Vec::with_capacity(n)); + let mut current = 0; + if n == 0 { + return order; + } + placed[0] = true; + order.push(0); + for _ in 1..n { + let mut best: Option<(usize, f32)> = None; + for (j, e) in embeddings.iter().enumerate() { + if placed[j] { + continue; + } + let c = matching::cosine(embeddings[current], e).unwrap_or(-1.0); + if best.is_none_or(|(_, b)| c > b) { + best = Some((j, c)); + } + } + let Some((j, _)) = best else { break }; + placed[j] = true; + order.push(j); + current = j; + } + order + } + + /// Put the unnamed faces `all` in order of likeness (the ones already embedded; the rest follow in photo order) and + /// find a name for those the named faces recognise, for the first `limit` of them. + pub fn arrange_unnamed(s: &mut Session, all: &mut Vec<(u64, usize, Rect)>, limit: usize) -> super::Arranged { + let none = super::Arranged { ordered: false, ready: false, suggestions: Default::default() }; + let Some(dir) = s.face_models_dir.clone() else { return none }; + if !enabled_cached(s, &dir) { + return none; + } + let Ok(embedder) = current_with(s, "faces.unnamed", false) else { return none }; + let (mut embedded, mut plain) = (Vec::new(), Vec::new()); + for f in all.drain(..) { + match s.faces.index.get(f.0, &f.2) { + Some(e) => embedded.push((f, e.clone())), + None => plain.push(f), + } + } + let beyond = embedded.split_off(embedded.len().min(MAX_ARRANGED)); + let order = nearest_chain(&embedded.iter().map(|(_, e)| &e[..]).collect::>()); + let ordered: Vec<((u64, usize, Rect), Arc<[f32]>)> = order.into_iter().filter_map(|i| embedded.get(i).cloned()).collect(); + // names the faces already named suggest, as `faces.suggest` would + let gallery: Vec<(String, Arc<[f32]>)> = s + .catalog + .photos() + .filter(|ph| ph.in_library()) + .flat_map(|ph| { + ph.meta + .regions + .iter() + .filter(|r| r.kind == RegionKind::Face) + .filter_map(|r| Some((r.name.clone()?, s.faces.index.get(ph.id.0, &r.rect)?.clone()))) + .collect::>() + }) + .collect(); + let refs: Vec<(&str, &[f32])> = gallery.iter().map(|(n, e)| (n.as_str(), &e[..])).collect(); + let threshold = embedder.manifest().thresholds.match_cosine.unwrap_or(DEFAULT_THRESHOLD); + let mut suggestions = std::collections::HashMap::new(); + if !refs.is_empty() { + for ((photo, index, _), e) in ordered.iter().take(limit) { + if let Some(x) = matching::suggest(&matching::rank(e, &refs), threshold, DEFAULT_MARGIN) { + suggestions.insert((*photo, *index), (x.name, x.score)); + } + } + } + let grouped = !ordered.is_empty(); + all.extend(ordered.into_iter().map(|(f, _)| f)); + all.extend(beyond.into_iter().map(|(f, _)| f)); + all.extend(plain); + super::Arranged { ordered: grouped, ready: true, suggestions } + } + + /// Whether face recognition is switched on, looked up at most once a second. + fn enabled_cached(s: &mut Session, dir: &std::path::Path) -> bool { + if let Some((t, on)) = s.faces.enabled_seen + && t.elapsed() < RECHECK + { + return on; + } + let on = read_settings(dir).enabled; + s.faces.enabled_seen = Some((Instant::now(), on)); + on + } + + pub fn suggest(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.suggest"; + let embedder = current(s, C)?; + collect(s); + let targets = s.targets(p); + let model_threshold = embedder.manifest().thresholds.match_cosine; + let num = |k: &str| p.get(k).and_then(Value::as_f64).filter(|v| v.is_finite()).map(|v| v as f32); + let threshold = num("threshold").or(model_threshold).unwrap_or(DEFAULT_THRESHOLD); + let margin = num("margin").unwrap_or(DEFAULT_MARGIN); + // the photos asked about are embedded first, then the named faces they are compared with, within the budget + let budget = Duration::from_millis(p.get("budgetMs").and_then(Value::as_u64).unwrap_or(1500).min(600_000)); + let started = Instant::now(); + let todo = pending(s, &targets, false); + let mut done = 0; + for id in &todo { + // a budget of zero only looks at what is already embedded + if budget.is_zero() || (done > 0 && started.elapsed() >= budget) { + break; + } + embed_photo(s, *id, &embedder)?; + done += 1; + } + let pending_photos: Vec = todo.iter().skip(done).map(|id| id.0).collect(); + let gallery: Vec<(String, Arc<[f32]>)> = s + .catalog + .photos() + .filter(|ph| ph.in_library()) + .flat_map(|ph| { + ph.meta + .regions + .iter() + .filter(|r| r.kind == RegionKind::Face) + .filter_map(|r| Some((r.name.clone()?, s.faces.index.get(ph.id.0, &r.rect)?.clone()))) + .collect::>() + }) + .collect(); + let refs: Vec<(&str, &[f32])> = gallery.iter().map(|(n, e)| (n.as_str(), &e[..])).collect(); + let mut photos = Vec::new(); + for id in &targets { + let Some(photo) = s.catalog.photo(*id) else { continue }; + let mut faces = Vec::new(); + for (index, r) in photo.meta.regions.iter().enumerate().filter(|(_, r)| r.kind == RegionKind::Face && r.name.is_none()) { + let Some(q) = s.faces.index.get(id.0, &r.rect) else { + faces.push(json!({"index": index, "pending": true})); + continue; + }; + let ranked = matching::rank(q, &refs); + let suggestion = matching::suggest(&ranked, threshold, margin); + faces.push(json!({ + "index": index, + "suggestion": suggestion.as_ref().map(|x| json!({"name": x.name, "score": x.score, "runnerUp": x.runner_up.as_ref().map(|(n, sc)| json!({"name": n, "score": sc}))})), + "candidates": ranked.iter().take(3).map(|c| json!({"name": c.name, "score": c.score, "faces": c.faces})).collect::>(), + })); + } + photos.push(json!({"id": id.0, "faces": faces})); + } + Ok( + json!({"model": s.faces.index.tag, "threshold": threshold, "margin": margin, "galleryFaces": gallery.len(), "pendingPhotos": pending_photos, "photos": photos}), + ) + } + #[cfg(test)] + mod tests { + use super::*; + use crate::tests_face_recognize::{install_yunet, region, set_regions, setup, temp, two_photos}; + + fn unit(parts: &[(usize, f32)]) -> Vec { + let mut v = vec![0.0; 64]; + for (i, x) in parts { + v[*i] = *x; + } + let n = v.iter().map(|x| x * x).sum::().sqrt(); + v.iter().map(|x| x / n).collect() + } + + /// A session with a tiny model loaded (so the index has its model's tag). + fn loaded(name: &str) -> (Session, std::path::PathBuf) { + let d = temp(name); + let (mut s, _) = setup(&d, 64); + s.execute("faces.index", &json!({"budgetMs": 0})).unwrap(); + (s, d) + } + + fn finished(s: &Session, id: PhotoId, found: Option>, results: Vec<(Rect, Option>, Rect)>) -> Done { + Done { tag: s.faces.index.tag.clone(), epoch: s.faces.epoch, id, results, detect_wanted: true, found } + } + + #[test] + fn faces_found_by_the_scan_become_unnamed_regions_without_an_undo_step() { + let (mut s, d) = loaded("found"); + let (a, b) = two_photos(&s); + set_regions(&mut s, b, vec![region(0.1, Some("Ann"))]); + let (undo, redo) = (s.undo.len(), s.redo.len()); + let found = region(0.3, None).rect; + let done = finished(&s, a, Some(vec![found]), vec![(found, Some(unit(&[(0, 1.0)])), found)]); + assert_eq!(ingest(&mut s, done), (1, 0)); + let regions = &s.catalog.photo(a).unwrap().meta.regions; + assert_eq!(regions.len(), 1); + assert!(regions[0].kind == RegionKind::Face && regions[0].name.is_none() && regions[0].rect == found); + assert!(regions[0].description.as_deref().unwrap().starts_with(MARK), "marked as detected"); + // not something the user did, and its embedding is already there + assert_eq!((s.undo.len(), s.redo.len()), (undo, redo)); + assert!(s.skip_auto_write, "no sidecar is rewritten for it"); + assert!(s.faces.scanned.contains(a.0) && s.faces.index.contains(a.0, &found)); + // a photo that has faces by now (the user, an import) keeps exactly those + let other = region(0.6, None).rect; + let done = finished(&s, b, Some(vec![other]), vec![(other, None, other)]); + ingest(&mut s, done); + assert_eq!(s.catalog.photo(b).unwrap().meta.regions.len(), 1); + assert!(s.faces.scanned.contains(b.0)); + let _ = std::fs::remove_dir_all(&d); + } + + #[test] + fn a_photo_that_could_not_be_searched_is_left_alone_for_the_session() { + let (mut s, d) = loaded("scanfail"); + let (a, _) = two_photos(&s); + assert!(needs_search(&s, s.catalog.photo(a).unwrap())); + let done = finished(&s, a, None, vec![]); + ingest(&mut s, done); + assert!(!s.faces.scanned.contains(a.0), "not remembered as searched: the file may come back"); + assert!(!needs_search(&s, s.catalog.photo(a).unwrap()), "but not retried on every pass"); + assert!(!pending(&s, &[], true).contains(&a)); + let _ = std::fs::remove_dir_all(&d); + } + + #[test] + fn only_the_background_scan_looks_for_faces_and_only_in_photos_without_any() { + let (mut s, d) = loaded("search"); + // finding faces needs the detector, which is not bundled; without it the scan only embeds + let detector = install_yunet(&mut s); + let (a, b) = two_photos(&s); + set_regions(&mut s, b, vec![region(0.1, Some("Ann")), region(0.5, None)]); + // `faces.index` embeds what is boxed; only the scan (search = true) also searches the rest, last + assert_eq!(pending(&s, &[], false), vec![b]); + let scan = pending(&s, &[], true); + assert_eq!(scan.first(), Some(&b), "named faces first"); + assert!(scan.contains(&a) && scan.len() > 2, "the unsearched photos follow"); + assert!(prepare(&mut s, a, false).is_none(), "nothing to embed in a photo without boxes"); + if detector { + assert!(prepare(&mut s, a, true).is_some_and(|p| p.detect && p.todo.is_empty())); + } else { + assert!(prepare(&mut s, a, true).is_none(), "no detector: nothing to do, and not marked as searched"); + } + assert!(prepare(&mut s, b, true).is_some_and(|p| !p.detect && p.todo.len() == 2), "a photo with boxes is not searched"); + // once searched, a photo is not searched again + s.faces.scanned.insert(a.0); + assert!(!pending(&s, &[], true).contains(&a)); + let _ = std::fs::remove_dir_all(&d); + } + + #[test] + fn opening_another_library_forgets_the_old_ones_faces() { + let (mut s, d) = loaded("newlib"); + let (a, _) = two_photos(&s); + let rect = region(0.2, None).rect; + s.faces.index.insert(a.0, &rect, unit(&[(0, 1.0)])); + s.faces.scanned.insert(a.0); + let epoch = s.faces.epoch; + let old = finished(&s, a, Some(vec![rect]), vec![(rect, Some(unit(&[(1, 1.0)])), rect)]); + s.open_library(d.join("second"), true).unwrap(); + assert_eq!((s.faces.index.len(), s.faces.scanned.len(), s.faces.queue.len(), s.faces.in_flight.len()), (0, 0, 0, 0)); + assert!(s.faces.epoch != epoch && s.faces.embedder.is_none()); + // a worker that was still busy with the old library cannot put its faces into the new one + assert_eq!(ingest(&mut s, old), (0, 0)); + assert!(s.faces.index.len() == 0 && !s.faces.scanned.contains(a.0)); + let _ = std::fs::remove_dir_all(&d); + } + + #[test] + fn the_unnamed_faces_that_look_like_a_person_come_best_first() { + let (mut s, d) = loaded("similar"); + let (a, b) = two_photos(&s); + let c = s.catalog.photos().map(|p| p.id).find(|id| *id != a && *id != b).unwrap(); + set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, Some("Bob"))]); + set_regions(&mut s, b, vec![region(0.1, None), region(0.5, None)]); + set_regions(&mut s, c, vec![region(0.3, None)]); + let put = |s: &mut Session, id: PhotoId, x: f64, v: Vec| s.faces.index.insert(id.0, ®ion(x, None).rect, v); + put(&mut s, a, 0.1, unit(&[(0, 1.0)])); // Ann + put(&mut s, a, 0.5, unit(&[(1, 1.0)])); // Bob + put(&mut s, b, 0.1, unit(&[(0, 0.95), (1, 0.05)])); // looks just like Ann + put(&mut s, b, 0.5, unit(&[(1, 1.0)])); // looks like Bob, not Ann + put(&mut s, c, 0.3, unit(&[(0, 0.8), (1, 0.6)])); // somewhat like Ann + let r = similar(&mut s, "ann", 10); + assert!(r.ready); + let found: Vec<(u64, u64)> = r.items.iter().map(|i| (i["photo"].as_u64().unwrap(), i["index"].as_u64().unwrap())).collect(); + assert_eq!(found, vec![(b.0, 0), (c.0, 0)], "best first; the face that looks like Bob is left out"); + assert!(r.items[0]["score"].as_f64().unwrap() > r.items[1]["score"].as_f64().unwrap()); + assert_eq!(similar(&mut s, "ann", 1).items.len(), 1, "capped"); + assert!(similar(&mut s, "nobody", 10).items.is_empty()); + let _ = std::fs::remove_dir_all(&d); + } + } +} + +#[cfg(target_arch = "wasm32")] +mod imp { + use serde_json::Value; + + use super::super::bad; + use crate::{Result, Session}; + + pub fn index(_: &mut Session, _: &Value) -> Result { + Err(bad("faces.index", "this build cannot run face recognition")) + } + + pub fn suggest(_: &mut Session, _: &Value) -> Result { + Err(bad("faces.suggest", "this build cannot run face recognition")) + } + + pub fn evaluate(_: &mut Session, _: &Value) -> Result { + Err(bad("faces.evaluate", "this build cannot run face recognition")) + } + + pub fn pump(_: &mut Session, _: &Value) -> Result { + Ok(serde_json::json!({"active": false})) + } + + pub fn similar(_: &mut Session, _: &str, _: usize) -> super::Similar { + super::Similar { items: Vec::new(), ready: false, pending: 0 } + } + + pub fn arrange_unnamed(_: &mut Session, _: &mut Vec<(u64, usize, lightcraft_geom::Rect)>, _: usize) -> super::Arranged { + super::Arranged { ordered: false, ready: false, suggestions: Default::default() } + } +} + +pub fn specs() -> Vec { + vec![ + cmd!( + "faces.index", + "Index Faces", + [], + None, + "{budgetMs?: 200, ids?: [photo ids to do first]} → {embedded, photosDone, pendingPhotos, indexedFaces, ms} — embed the faces in the library with the chosen recognition model, as many photos as fit in the time (at least one); call again until `pendingPhotos` is 0. `budgetMs: 0` only reports", + super::always, + imp::index + ), + cmd!( + "faces.suggest", + "Suggest Names for Faces", + [], + None, + "{ids?, threshold?, margin?, budgetMs?} → {photos: [{id, faces: [{index, suggestion: {name, score, runnerUp} | null, candidates: [{name, score}]}]}]} — who the unnamed faces in these photos might be, from the faces already named in the library. Only suggestions: nothing is named", + super::always, + imp::suggest + ), + cmd!( + query "faces.pump", + "Index Faces in the Background", + [], + None, + "{pace?: pause | light | normal | full} → {active, inFlight, pendingPhotos, indexedFaces, searchedPhotos, workers} — cheap to call every frame: takes in what the background workers have finished and gives them the next photos, as many at once as the pace allows (none when paused, one when light, half the machine's threads when normal, four fifths when full); inactive until face recognition is on in Settings with a model chosen", + super::always, + imp::pump + ), + cmd!( + "faces.evaluate", + "Evaluate Face Recognition", + [], + None, + "{budgetMs?: 0, burstSecs?: 5, margin?, thresholds?: [..], scope?: \"visible\"} → {queries, top1, sweep: [{threshold, suggested, correct, precision, recall}]} — leave-one-out over your named faces: each is ranked against the others (ignoring shots within burstSecs of it) to show how well the chosen model recognises people in *your* photos and which threshold to trust. Faces not yet embedded are done within budgetMs (`faces.index` does them all); counts only", + super::always, + imp::evaluate + ), + cmd!( + query "faces.unnamed", + "Unnamed Faces", + [], + None, + "{limit?: 600} → {total, faces: [{photo, index, rect, view, suggestion: {name, score} | null}], ordered, ready} — the faces nobody has named; with recognition running, look-alikes are next to each other and a face the named ones recognise carries the name they suggest. `view` is the box to show the face by. Name them with `faces.nameFaces`", + super::always, + unnamed + ), + cmd!( + "faces.nameFaces", + "Name Faces", + [], + None, + "{faces: [{photo, index}], name} → {named, photos} — name (or, with a null or empty name, clear) many face regions as one undoable step; the sidecar is not touched", + super::always, + name_faces + ), + cmd!( + query "faces.person", + "Faces of a Person", + [], + None, + "{name, more?: 120} → {name, total, confirmed: [{photo, index, rect}], more: [{photo, index, rect, score}], ready, pendingPhotos} — the faces named `name` across the library, and the unnamed faces that look most like them (best first, already embedded ones only, once recognition is on); confirm one with `faces.setName`", + super::always, + person + ), + cmd!( + "faces.setName", + "Name Face", + [], + None, + "{id?, index, name} — name a face region (null or empty clears it); undoable; the sidecar is not touched", + super::always, + set_name + ), + ] +} diff --git a/crates/engine/src/cmd/mod.rs b/crates/engine/src/cmd/mod.rs index 9fc83116..303a92d4 100644 --- a/crates/engine/src/cmd/mod.rs +++ b/crates/engine/src/cmd/mod.rs @@ -13,6 +13,9 @@ pub mod curves; mod develop; mod edit; mod export; +pub(crate) mod face_detect; +mod face_models; +mod face_recognize; pub mod filters; pub mod keywords; pub mod library; @@ -140,6 +143,9 @@ pub fn command_specs() -> &'static [CommandSpec] { v.extend(missing::specs()); v.extend(metadata::specs()); v.extend(filters::specs()); + v.extend(face_models::specs()); + v.extend(face_detect::specs()); + v.extend(face_recognize::specs()); v }) } diff --git a/crates/engine/src/config.rs b/crates/engine/src/config.rs new file mode 100644 index 00000000..933dc22d --- /dev/null +++ b/crates/engine/src/config.rs @@ -0,0 +1,37 @@ +//! Where LightCraft keeps per-user data on this machine: the config folder (app settings, face models). +//! +//! Hosts that have a file system (the desktop app, the CLI, the MCP server) share these defaults so a model +//! installed from one is there for the others. Nothing here is applied automatically: a [`Session`] has no +//! face-models folder until a host asks for one ([`Session::with_default_face_models`]), so tests stay +//! hermetic. + +use std::path::PathBuf; + +use crate::Session; + +/// LightCraft's per-user config folder (`…/LightCraft`), if the platform tells us where. +pub fn config_dir() -> Option { + if cfg!(target_os = "macos") { + std::env::var_os("HOME").map(|h| PathBuf::from(h).join("Library/Application Support/LightCraft")) + } else if cfg!(windows) { + std::env::var_os("APPDATA").map(|a| PathBuf::from(a).join("LightCraft")) + } else { + std::env::var_os("XDG_CONFIG_HOME") + .map(PathBuf::from) + .or_else(|| std::env::var_os("HOME").map(|h| PathBuf::from(h).join(".config"))) + .map(|c| c.join("lightcraft")) + } +} + +/// Where face models are kept: `$LIGHTCRAFT_FACE_MODELS` if set, else `/models`. +pub fn default_face_models_dir() -> Option { + std::env::var_os("LIGHTCRAFT_FACE_MODELS").filter(|v| !v.is_empty()).map(PathBuf::from).or_else(|| config_dir().map(|d| d.join("models"))) +} + +impl Session { + /// Keep face models in the shared default folder ([`default_face_models_dir`]). + pub fn with_default_face_models(mut self) -> Self { + self.face_models_dir = default_face_models_dir(); + self + } +} diff --git a/crates/engine/src/face_download.rs b/crates/engine/src/face_download.rs new file mode 100644 index 00000000..0c6ed4ee --- /dev/null +++ b/crates/engine/src/face_download.rs @@ -0,0 +1,372 @@ +//! Downloading a face model the user asked for (Settings ▸ Faces ▸ Download). +//! +//! The transfer is [`lightcraft_fetch`] (pure-Rust HTTPS with resume, size and SHA-256 checks), on a background thread, +//! so the UI never waits for the network. Only a model listed in [`lightcraft_faces::known::download`] can be fetched: +//! the address is a pinned constant, never something typed or passed in, and no token is ever sent with it. What arrives +//! matches the model's recorded size and SHA-256 or it is thrown away; it then waits in a staging folder until the model +//! is installed (`faces.models.install`, or by itself once the user has accepted the model's terms in the download). +//! +//! Failure is a message, never a panic: no network, a stalled or truncated transfer, a full disk, a file that is not the +//! one expected, and a panic inside the thread itself all end as [`State::Failed`]. + +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex, PoisonError}; +use std::time::Duration; + +use lightcraft_faces::known::Download; + +/// The folder inside the models folder where a downloaded file waits to be installed. Its name starts with a dot, so it +/// can never be a model's own folder (ids start with a letter or digit). +pub const STAGING: &str = ".downloads"; + +/// Where a download stands. +#[derive(Clone, Debug, PartialEq)] +pub enum State { + /// Fetching (`bytes` of `total` so far); the file is checked as the last step. + Running { + bytes: u64, + total: u64, + }, + /// Fetched and verified (its SHA-256 is `sha256`), waiting in the staging folder at `path` to be installed. + Done { + path: PathBuf, + sha256: String, + }, + /// Installed and in use: nothing more to do but tell the user. + Installed, + Failed(String), + Cancelled, +} + +struct Shared { + state: Mutex, + cancel: AtomicBool, +} + +impl Shared { + fn new(total: u64) -> Self { + Shared { state: Mutex::new(State::Running { bytes: 0, total }), cancel: AtomicBool::new(false) } + } + + fn set(&self, s: State) { + *self.state.lock().unwrap_or_else(PoisonError::into_inner) = s; + } + + fn get(&self) -> State { + self.state.lock().unwrap_or_else(PoisonError::into_inner).clone() + } +} + +/// The downloads of one session, by model id. +#[derive(Default)] +pub struct Downloads { + jobs: HashMap>, +} + +impl Downloads { + /// Start fetching `spec` into `/.downloads/`. Different models can download side by side; the same + /// one is started once. + pub fn start(&mut self, spec: Download, models_dir: &Path) -> Result<(), String> { + if self.jobs.get(spec.id.as_str()).is_some_and(|j| matches!(j.get(), State::Running { .. })) { + return Err("that model is already downloading".into()); + } + let staging = models_dir.join(STAGING); + std::fs::create_dir_all(&staging).map_err(|e| format!("could not create the download folder: {e}"))?; + let shared = Arc::new(Shared::new(spec.size_bytes)); + let worker = shared.clone(); + let id = spec.id.clone(); + std::thread::Builder::new() + .name("face-model-download".into()) + .spawn(move || { + let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| run(&spec, spec.base(), &staging, &worker))); + if outcome.is_err() { + worker.set(State::Failed("the download stopped unexpectedly".into())); + } + }) + .map_err(|e| format!("could not start the download: {e}"))?; + self.jobs.insert(id, shared); + Ok(()) + } + + /// A download that has arrived and been checked, as if the transfer had just ended (tests of installing). + #[cfg(test)] + pub(crate) fn arrived(&mut self, id: &str, path: PathBuf, sha256: &str) { + let j = Shared::new(0); + j.set(State::Done { path, sha256: sha256.to_string() }); + self.jobs.insert(id.to_string(), Arc::new(j)); + } + + /// Every download this session has started and not discarded, by model id. + pub fn snapshot(&self) -> Vec<(String, State)> { + let mut all: Vec<_> = self.jobs.iter().map(|(id, j)| (id.clone(), j.get())).filter(|(_, s)| *s != State::Cancelled).collect(); + all.sort_by(|a, b| a.0.cmp(&b.0)); + all + } + + pub fn running(&self) -> bool { + self.jobs.values().any(|j| matches!(j.get(), State::Running { .. })) + } + + /// Stop a running download, or discard a finished one (its staged file is deleted). `true` if there was one. + pub fn discard(&mut self, id: &str) -> bool { + let Some(job) = self.jobs.get(id) else { return false }; + job.cancel.store(true, Ordering::Relaxed); + match job.get() { + // the thread notices, deletes the partial file and marks itself cancelled + State::Running { .. } => {} + State::Done { path, .. } => { + let _ = std::fs::remove_file(path); + self.jobs.remove(id); + } + State::Installed | State::Failed(_) | State::Cancelled => { + self.jobs.remove(id); + } + } + true + } + + /// The staged file at `path` has been installed (its copy is the model now) or is no longer wanted. + pub fn staged_file_used(&mut self, path: &Path) { + self.jobs.retain(|_, j| !matches!(j.get(), State::Done { path: p, .. } if p == path)); + } + + /// Downloads that have arrived and been checked, ready to install: (model id, staged file, its SHA-256). + pub fn finished(&self) -> Vec<(String, PathBuf, String)> { + let mut out: Vec<_> = self + .jobs + .iter() + .filter_map(|(id, j)| match j.get() { + State::Done { path, sha256 } => Some((id.clone(), path, sha256)), + _ => None, + }) + .collect(); + out.sort_by(|a, b| a.0.cmp(&b.0)); + out + } + + /// Record how installing a finished download went (`Installed` or `Failed`). + pub fn set_outcome(&mut self, id: &str, state: State) { + match self.jobs.get(id) { + Some(j) => j.set(state), + None => { + let j = Shared::new(0); + j.set(state); + self.jobs.insert(id.to_string(), Arc::new(j)); + } + } + } +} + +impl Drop for Downloads { + /// Quitting while a download runs stops it instead of leaving it behind. + fn drop(&mut self) { + for job in self.jobs.values() { + job.cancel.store(true, Ordering::Relaxed); + } + let started = std::time::Instant::now(); + while self.running() && started.elapsed() < Duration::from_secs(1) { + std::thread::sleep(Duration::from_millis(10)); + } + } +} + +/// What went wrong, in words the user can act on. +#[cfg(not(target_arch = "wasm32"))] +fn explain(e: &lightcraft_fetch::DownloadError) -> String { + use lightcraft_fetch::DownloadError; + match e { + DownloadError::AllMirrorsFailed { errors, .. } => { + let detail: String = errors.join("; ").chars().take(300).collect(); + format!("could not download the model ({detail}): check the internet connection, or get the file from its page and add it yourself") + } + other => other.to_string(), + } +} + +/// Fetch `spec` from `mirror` (the address without the file name) into `staging`, check it and leave it there. +#[cfg(not(target_arch = "wasm32"))] +fn fetch(spec: &Download, mirror: &str, staging: &Path, shared: &Shared) -> Result<(), lightcraft_fetch::DownloadError> { + let file = lightcraft_fetch::FileSpec { name: &spec.file_name, size: Some(spec.size_bytes), sha256: Some(&spec.sha256), max: spec.size_bytes }; + lightcraft_fetch::download(&[file], &[mirror.to_string()], staging, &lightcraft_fetch::Options::default(), &shared.cancel, &mut |p| { + shared.set(State::Running { bytes: p.done.min(spec.size_bytes), total: spec.size_bytes }); + }) +} + +#[cfg(target_arch = "wasm32")] +fn fetch(_: &Download, _: &str, _: &Path, _: &Shared) -> Result<(), String> { + Err("the web build downloads no models".into()) +} + +/// The whole job: reuse a verified earlier copy, else fetch, check and stage. Always ends in a final [`State`]. +fn run(spec: &Download, mirror: &str, staging: &Path, shared: &Shared) { + let target = staging.join(&spec.file_name); + let part = staging.join(format!("{}.part", spec.file_name)); + match fetch(spec, mirror, staging, shared) { + Ok(()) => shared.set(State::Done { path: target, sha256: spec.sha256.to_string() }), + #[cfg(not(target_arch = "wasm32"))] + Err(lightcraft_fetch::DownloadError::Cancelled) => { + let _ = std::fs::remove_file(&part); + shared.set(State::Cancelled); + } + #[cfg(not(target_arch = "wasm32"))] + Err(e) => { + let _ = std::fs::remove_file(&part); + shared.set(State::Failed(explain(&e))); + } + #[cfg(target_arch = "wasm32")] + Err(e) => { + let _ = std::fs::remove_file(&part); + shared.set(State::Failed(e)); + } + } +} + +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::*; + use lightcraft_faces::hash::sha256_hex; + use std::io::{BufRead, BufReader, Write}; + use std::net::TcpListener; + + fn scratch(name: &str) -> PathBuf { + let d = std::env::temp_dir().join(format!("lc-face-dl-{name}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&d); + std::fs::create_dir_all(&d).unwrap(); + d + } + + fn spec_for(bytes: &[u8]) -> Download { + Download { + id: "test-model".into(), + url: "https://example.invalid/m/test-model.onnx".into(), + file_name: "test-model.onnx".into(), + size_bytes: bytes.len() as u64, + sha256: sha256_hex(bytes), + } + } + + /// A local server that answers every request with `status` and `body` (the whole file, no ranges). + fn serve(status: &'static str, body: Vec) -> String { + let l = TcpListener::bind("127.0.0.1:0").unwrap(); + let base = format!("http://{}/m", l.local_addr().unwrap()); + std::thread::spawn(move || { + for mut s in l.incoming().flatten() { + let mut r = BufReader::new(s.try_clone().unwrap()); + loop { + let mut line = String::new(); + if r.read_line(&mut line).unwrap_or(0) == 0 || line.trim().is_empty() { + break; + } + } + let _ = s.write_all(format!("HTTP/1.1 {status}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len()).as_bytes()); + let _ = s.write_all(&body); + } + }); + base + } + + #[test] + fn fetches_checks_and_stages_a_file() { + let dir = scratch("ok"); + let bytes: Vec = (0..200_000u32).map(|i| (i % 251) as u8).collect(); + let (staging, spec) = (dir.join(STAGING), spec_for(&bytes)); + std::fs::create_dir_all(&staging).unwrap(); + let mirror = serve("200 OK", bytes.clone()); + let shared = Shared::new(spec.size_bytes); + run(&spec, &mirror, &staging, &shared); + let State::Done { path, sha256 } = shared.get() else { panic!("{:?}", shared.get()) }; + assert_eq!((std::fs::read(&path).unwrap(), sha256), (bytes, spec.sha256.to_string())); + assert!(!staging.join("test-model.onnx.part").exists()); + // asked again with the copy still staged: no fetch needed (nothing listens there) + let again = Shared::new(spec.size_bytes); + run(&spec, "http://127.0.0.1:1/m", &staging, &again); + assert!(matches!(again.get(), State::Done { .. }), "{:?}", again.get()); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn a_file_that_does_not_match_is_thrown_away() { + let dir = scratch("bad"); + let staging = dir.join(STAGING); + std::fs::create_dir_all(&staging).unwrap(); + // the right size and the wrong content + let spec = spec_for(&[7u8; 5000]); + let shared = Shared::new(spec.size_bytes); + run(&spec, &serve("200 OK", vec![8u8; 5000]), &staging, &shared); + assert!(matches!(shared.get(), State::Failed(_)), "{:?}", shared.get()); + // nothing is left behind: not the file, not a partial one + assert_eq!(std::fs::read_dir(&staging).unwrap().count(), 0); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn an_unreachable_or_refusing_server_is_a_message() { + let dir = scratch("down"); + let staging = dir.join(STAGING); + std::fs::create_dir_all(&staging).unwrap(); + let spec = spec_for(b"abc"); + for mirror in ["http://127.0.0.1:1/m".to_string(), serve("404 Not Found", b"not found".to_vec())] { + let shared = Shared::new(3); + run(&spec, &mirror, &staging, &shared); + let State::Failed(why) = shared.get() else { panic!("{:?}", shared.get()) }; + assert!(why.contains("internet connection"), "{why}"); + } + assert_eq!(std::fs::read_dir(&staging).unwrap().count(), 0); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn a_cancelled_download_leaves_nothing() { + let dir = scratch("cancel"); + let staging = dir.join(STAGING); + std::fs::create_dir_all(&staging).unwrap(); + let spec = spec_for(b"abc"); + let shared = Shared::new(3); + shared.cancel.store(true, Ordering::Relaxed); + run(&spec, &serve("200 OK", b"abc".to_vec()), &staging, &shared); + assert_eq!(shared.get(), State::Cancelled); + assert_eq!(std::fs::read_dir(&staging).unwrap().count(), 0); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn dropping_the_session_stops_a_running_download() { + let shared = Arc::new(Shared::new(10)); + // a stand-in for the download thread: notices the cancel flag the way the downloader does + let worker = shared.clone(); + let handle = std::thread::spawn(move || { + while !worker.cancel.load(Ordering::Relaxed) { + std::thread::sleep(Duration::from_millis(5)); + } + worker.set(State::Cancelled); + }); + let mut d = Downloads::default(); + d.jobs.insert("m".into(), shared.clone()); + assert!(d.running()); + let started = std::time::Instant::now(); + drop(d); + assert!(started.elapsed() < Duration::from_millis(900), "quitting must not wait for a transfer to finish"); + assert_eq!(shared.get(), State::Cancelled); + handle.join().unwrap(); + } + + #[test] + fn discarding_a_finished_download_deletes_its_file() { + let dir = scratch("discard"); + let file = dir.join("staged.onnx"); + std::fs::write(&file, b"x").unwrap(); + let mut d = Downloads::default(); + let shared = Arc::new(Shared::new(1)); + shared.set(State::Done { path: file.clone(), sha256: String::new() }); + d.jobs.insert("m".into(), shared); + assert_eq!(d.snapshot().len(), 1); + assert!(!d.running()); + assert!(d.discard("m")); + assert!(!file.exists()); + assert!(d.snapshot().is_empty()); + assert!(!d.discard("m"), "nothing left to discard"); + let _ = std::fs::remove_dir_all(&dir); + } +} diff --git a/crates/engine/src/faces_index.rs b/crates/engine/src/faces_index.rs new file mode 100644 index 00000000..c250f868 --- /dev/null +++ b/crates/engine/src/faces_index.rs @@ -0,0 +1,455 @@ +//! The face index: one embedding per face region, kept in memory and cached beside the library. +//! +//! Embedding a face costs a model run (tens to hundreds of milliseconds), so each region is embedded once and +//! remembered, keyed by its photo and its box. The cache file (`face-embeddings-.bin` in the library folder, +//! one per model, so switching back to an earlier model does not start over) belongs to one model: its header names the +//! model and its file hash, and a file made by another model, an older format, or a damaged one is simply ignored and +//! rebuilt. Embeddings are kept out of the catalog on purpose: they are large, rebuildable, and not comparable between +//! models. +//! +//! Each face also keeps the box to cut its picture from (its "view"): the detector's box when the detector found the same +//! face, so every face is shown equally close however loosely or tightly its own region was drawn (a Lightroom box, an +//! older tool's box), else the region's own box. +//! +//! Layout (little-endian): `"LCFE"`, version `2`, tag length `u16`, the tag (`@`), dimension `u32`, +//! then records of `photo u64`, four `i32` (the region's box in ten-thousandths), four `i32` (the view box likewise), and +//! `dim` `f32`. + +use std::collections::HashMap; +use std::io::Write; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use lightcraft_geom::Rect; + +const MAGIC: &[u8; 4] = b"LCFE"; +const VERSION: u8 = 2; +/// Largest cache file we read. +const MAX_FILE: u64 = 1 << 30; + +pub(crate) type Key = (u64, [i32; 4]); + +/// One face's record: its embedding and its view box (quantized like a key). +type Entry = (Arc<[f32]>, [i32; 4]); + +/// A quantized box as a rectangle again. +fn rect_of(k: [i32; 4]) -> Rect { + let f = |v: i32| f64::from(v) / 10_000.0; + Rect { x0: f(k[0]), y0: f(k[1]), x1: f(k[2]), y1: f(k[3]) } +} + +/// A region's box as the integer key it is cached under. +pub(crate) fn region_key(r: &Rect) -> [i32; 4] { + let q = |v: f64| if v.is_finite() { (v.clamp(-1e3, 1e3) * 10_000.0).round() as i32 } else { 0 }; + [q(r.x0), q(r.y0), q(r.x1), q(r.y1)] +} + +#[derive(Default)] +pub(crate) struct Index { + /// `@` of the model the embeddings belong to. + pub tag: String, + pub dim: usize, + entries: HashMap, + unsaved: Vec<(Key, Entry)>, + /// Faces that could not be embedded this session (so they are not retried on every call). + skipped: std::collections::HashSet, + /// The cache file on disk belongs to this model, so new records can be appended to it. + file_matches: bool, + pub path: Option, +} + +impl Index { + pub fn len(&self) -> usize { + self.entries.len() + } + + pub fn get(&self, photo: u64, rect: &Rect) -> Option<&Arc<[f32]>> { + self.entries.get(&(photo, region_key(rect))).map(|e| &e.0) + } + + /// The box to cut this face's picture from, once the scan has looked at it. + pub fn view(&self, photo: u64, rect: &Rect) -> Option { + self.entries.get(&(photo, region_key(rect))).map(|e| rect_of(e.1)) + } + + #[cfg(test)] + pub fn contains(&self, photo: u64, rect: &Rect) -> bool { + self.entries.contains_key(&(photo, region_key(rect))) + } + + /// Whether this face still has to be embedded. + pub fn needs(&self, photo: u64, rect: &Rect) -> bool { + let key = (photo, region_key(rect)); + !self.entries.contains_key(&key) && !self.skipped.contains(&key) + } + + pub fn skip(&mut self, photo: u64, rect: &Rect) { + self.skipped.insert((photo, region_key(rect))); + } + + #[cfg(test)] + pub fn insert(&mut self, photo: u64, rect: &Rect, embedding: Vec) { + self.insert_with_view(photo, rect, embedding, rect); + } + + /// Remember a face's embedding and the box to show it by. + pub fn insert_with_view(&mut self, photo: u64, rect: &Rect, embedding: Vec, view: &Rect) { + let entry: Entry = (embedding.into(), region_key(view)); + let key = (photo, region_key(rect)); + self.entries.insert(key, entry.clone()); + self.unsaved.push((key, entry)); + } + + /// Whether there are embeddings not yet written to the cache file. + pub fn has_unsaved(&self) -> bool { + !self.unsaved.is_empty() + } + + /// Start over for `tag` (`dim` numbers per face), loading what the cache file at `path` holds for it. What the + /// previous model's index had not written yet is written first, so switching models loses nothing. + pub fn reset(&mut self, tag: &str, dim: usize, path: Option) { + let _ = self.save(); + self.tag = tag.to_string(); + self.dim = dim; + self.path = path; + self.entries.clear(); + self.unsaved.clear(); + self.skipped.clear(); + self.file_matches = false; + if let Some(p) = self.path.clone() { + self.load(&p); + } + } + + fn load(&mut self, path: &Path) { + let Ok(meta) = std::fs::metadata(path) else { return }; + if meta.len() > MAX_FILE { + return; + } + let Ok(bytes) = std::fs::read(path) else { return }; + let Some((tag_len, rest)) = bytes + .strip_prefix(MAGIC) + .and_then(|r| r.split_first()) + .filter(|(v, _)| **v == VERSION) + .and_then(|(_, r)| r.split_at_checked(2)) + .map(|(l, r)| (usize::from(u16::from_le_bytes([l[0], l[1]])), r)) + else { + return; + }; + let Some((tag, rest)) = rest.split_at_checked(tag_len) else { return }; + let Some((dim, mut rest)) = rest.split_at_checked(4).map(|(d, r)| (u32::from_le_bytes([d[0], d[1], d[2], d[3]]) as usize, r)) else { return }; + if tag != self.tag.as_bytes() || dim != self.dim || dim == 0 { + return; + } + self.file_matches = true; + let record = 8 + 16 + 16 + dim * 4; + while let Some((rec, tail)) = rest.split_at_checked(record) { + rest = tail; + let photo = u64::from_le_bytes([rec[0], rec[1], rec[2], rec[3], rec[4], rec[5], rec[6], rec[7]]); + let int = |at: usize| i32::from_le_bytes([rec[at], rec[at + 1], rec[at + 2], rec[at + 3]]); + let key = [int(8), int(12), int(16), int(20)]; + let view = [int(24), int(28), int(32), int(36)]; + let emb: Vec = rec[40..].as_chunks::<4>().0.iter().map(|c| f32::from_le_bytes(*c)).collect(); + if emb.iter().all(|v| v.is_finite()) { + self.entries.insert((photo, key), (emb.into(), view)); + } + } + } + + fn header(&self) -> Vec { + let tag = self.tag.as_bytes(); + let mut h = Vec::with_capacity(11 + tag.len()); + h.extend_from_slice(MAGIC); + h.push(VERSION); + h.extend_from_slice(&u16::try_from(tag.len()).unwrap_or(0).to_le_bytes()); + h.extend_from_slice(tag); + h.extend_from_slice(&(self.dim as u32).to_le_bytes()); + h + } + + fn record(&self, key: &Key, entry: &Entry) -> Vec { + let mut r = Vec::with_capacity(40 + entry.0.len() * 4); + r.extend_from_slice(&key.0.to_le_bytes()); + key.1.iter().chain(entry.1.iter()).for_each(|k| r.extend_from_slice(&k.to_le_bytes())); + entry.0.iter().for_each(|v| r.extend_from_slice(&v.to_le_bytes())); + r + } + + /// Write what was added since the last save: appended when the file already belongs to this model, else the + /// whole index is written fresh (atomically). A failure is returned, the records stay pending. + pub fn save(&mut self) -> std::io::Result<()> { + let Some(path) = self.path.clone() else { + self.unsaved.clear(); + return Ok(()); + }; + if self.unsaved.is_empty() { + return Ok(()); + } + if self.file_matches { + let mut f = std::fs::OpenOptions::new().append(true).open(&path)?; + let mut buf = Vec::new(); + for (k, e) in &self.unsaved { + buf.extend(self.record(k, e)); + } + f.write_all(&buf)?; + } else { + let mut buf = self.header(); + for (k, e) in &self.entries { + buf.extend(self.record(k, e)); + } + let part = path.with_extension("part"); + std::fs::write(&part, &buf)?; + std::fs::rename(&part, &path)?; + self.file_matches = true; + } + self.unsaved.clear(); + Ok(()) + } +} + +impl Drop for Index { + fn drop(&mut self) { + let _ = self.save(); + } +} + +const SCAN_MAGIC: &[u8; 4] = b"LCFS"; +const SCAN_VERSION: u8 = 1; +/// The detector a search was made with: a record made by another one is ignored (a better detector may find more). +pub(crate) const SCAN_TAG: &str = "yunet-2023mar"; +/// Largest record file we read. +const SCAN_MAX_FILE: u64 = 64 << 20; + +/// The photos the detector has already searched for faces, so each is searched once, however many times LightCraft +/// starts. Kept beside the library (`face-scanned.bin`: a header naming the detector, then photo ids, appended to). +#[derive(Default)] +pub(crate) struct Scanned { + set: std::collections::HashSet, + unsaved: Vec, + path: Option, + file_matches: bool, +} + +impl Scanned { + pub fn path(&self) -> Option<&Path> { + self.path.as_deref() + } + + pub fn contains(&self, photo: u64) -> bool { + self.set.contains(&photo) + } + + pub fn len(&self) -> usize { + self.set.len() + } + + pub fn insert(&mut self, photo: u64) { + if self.set.insert(photo) { + self.unsaved.push(photo); + } + } + + pub fn has_unsaved(&self) -> bool { + !self.unsaved.is_empty() + } + + /// Start over with the record at `path` (written first: what was not yet saved is not lost). + pub fn reset(&mut self, path: Option) { + let _ = self.save(); + self.set.clear(); + self.unsaved.clear(); + self.file_matches = false; + self.path = path; + if let Some(p) = self.path.clone() { + self.load(&p); + } + } + + fn header() -> Vec { + let mut h = Vec::new(); + h.extend_from_slice(SCAN_MAGIC); + h.push(SCAN_VERSION); + h.extend_from_slice(&(SCAN_TAG.len() as u16).to_le_bytes()); + h.extend_from_slice(SCAN_TAG.as_bytes()); + h + } + + fn load(&mut self, path: &Path) { + let Ok(meta) = std::fs::metadata(path) else { return }; + if meta.len() > SCAN_MAX_FILE { + return; + } + let Ok(bytes) = std::fs::read(path) else { return }; + let head = Self::header(); + let Some(rest) = bytes.strip_prefix(head.as_slice()) else { return }; + self.file_matches = true; + // a half-written last id (a crash while appending) is ignored + self.set.extend(rest.as_chunks::<8>().0.iter().map(|c| u64::from_le_bytes(*c))); + } + + /// Write what was added since the last save: appended when the file is this detector's, else the whole record + /// is written fresh (atomically). A failure is returned and the ids stay pending. + pub fn save(&mut self) -> std::io::Result<()> { + let Some(path) = self.path.clone() else { + self.unsaved.clear(); + return Ok(()); + }; + if self.unsaved.is_empty() { + return Ok(()); + } + if self.file_matches { + let mut f = std::fs::OpenOptions::new().append(true).open(&path)?; + let buf: Vec = self.unsaved.iter().flat_map(|id| id.to_le_bytes()).collect(); + f.write_all(&buf)?; + } else { + let mut buf = Self::header(); + buf.extend(self.set.iter().flat_map(|id| id.to_le_bytes())); + let part = path.with_extension("part"); + std::fs::write(&part, &buf)?; + std::fs::rename(&part, &path)?; + self.file_matches = true; + } + self.unsaved.clear(); + Ok(()) + } +} + +impl Drop for Scanned { + fn drop(&mut self) { + let _ = self.save(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn rect(x: f64) -> Rect { + Rect { x0: x, y0: 0.1, x1: x + 0.2, y1: 0.4 } + } + + fn temp(name: &str) -> PathBuf { + std::env::temp_dir().join(format!("lc-faceindex-{name}-{}.bin", std::process::id())) + } + + #[test] + fn embeddings_survive_a_save_and_load_and_append() { + let path = temp("roundtrip"); + let _ = std::fs::remove_file(&path); + let mut a = Index::default(); + a.reset("m@abc", 3, Some(path.clone())); + a.insert(7, &rect(0.1), vec![1.0, 0.0, 0.5]); + a.insert(8, &rect(0.3), vec![0.0, 1.0, 0.25]); + a.save().unwrap(); + a.insert(9, &rect(0.5), vec![0.5, 0.5, 0.5]); + a.save().unwrap(); // appended + let mut b = Index::default(); + b.reset("m@abc", 3, Some(path.clone())); + assert_eq!(b.len(), 3); + assert_eq!(b.get(8, &rect(0.3)).map(|e| e.to_vec()), Some(vec![0.0, 1.0, 0.25])); + assert!(b.contains(9, &rect(0.5)) && !b.contains(9, &rect(0.6)) && !b.contains(1, &rect(0.5))); + let _ = std::fs::remove_file(path); + } + + #[test] + fn another_model_a_damaged_file_or_a_torn_tail_are_handled() { + let path = temp("other"); + let _ = std::fs::remove_file(&path); + let mut a = Index::default(); + a.reset("m@abc", 3, Some(path.clone())); + a.insert(1, &rect(0.1), vec![1.0, 0.0, 0.0]); + a.insert(2, &rect(0.2), vec![0.0, 1.0, 0.0]); + a.save().unwrap(); + // another model: nothing is loaded, and saving replaces the file for the new model + let mut other = Index::default(); + other.reset("n@def", 4, Some(path.clone())); + assert_eq!(other.len(), 0); + other.insert(5, &rect(0.1), vec![1.0, 0.0, 0.0, 0.0]); + other.save().unwrap(); + let mut again = Index::default(); + again.reset("n@def", 4, Some(path.clone())); + assert_eq!(again.len(), 1); + let mut old = Index::default(); + old.reset("m@abc", 3, Some(path.clone())); + assert_eq!(old.len(), 0, "the old model's records are gone"); + // a torn last record loses only that record + let bytes = std::fs::read(&path).unwrap(); + std::fs::write(&path, &bytes[..bytes.len() - 5]).unwrap(); + let mut torn = Index::default(); + torn.reset("n@def", 4, Some(path.clone())); + assert_eq!(torn.len(), 0); + // garbage of every kind is ignored + for junk in [Vec::new(), b"LCFE".to_vec(), b"LCFE\x01\xff\xff".to_vec(), vec![0xff; 200], b"LCFE\x02\x00\x00".to_vec()] { + std::fs::write(&path, &junk).unwrap(); + let mut j = Index::default(); + j.reset("n@def", 4, Some(path.clone())); + assert_eq!(j.len(), 0); + } + // and a missing file is fine + let _ = std::fs::remove_file(&path); + let mut m = Index::default(); + m.reset("n@def", 4, Some(path.clone())); + assert_eq!(m.len(), 0); + } + + #[test] + fn non_finite_boxes_and_vectors_do_not_break_the_keys() { + assert_eq!(region_key(&Rect { x0: f64::NAN, y0: f64::INFINITY, x1: 1e300, y1: -1e300 }), [0, 0, 10_000_000, -10_000_000]); + let path = temp("nan"); + let _ = std::fs::remove_file(&path); + let mut a = Index::default(); + a.reset("m@abc", 2, Some(path.clone())); + a.insert(1, &rect(0.1), vec![f32::NAN, 1.0]); + a.save().unwrap(); + let mut b = Index::default(); + b.reset("m@abc", 2, Some(path.clone())); + assert_eq!(b.len(), 0, "a non-finite vector is not loaded back"); + let _ = std::fs::remove_file(path); + } +} + +/// What the recognition commands keep between calls: the loaded model and the embeddings made with it. +#[derive(Default)] +pub(crate) struct FacesState { + /// `@` and the loaded model. + pub embedder: Option<(String, Arc)>, + pub index: Index, + /// When the models folder was last looked at (the per-frame pump trusts a look for a second). + pub checked: Option, + /// Whether recognition was switched on when last looked up, and when. + pub enabled_seen: Option<(std::time::Instant, bool)>, + pub worker: Option, + /// After a failure to load the chosen model, the background pump leaves it alone until then. + pub retry_at: Option, + /// Photos the worker is busy with. + pub in_flight: std::collections::HashSet, + /// Photos still to hand to the worker, and the catalog revision this list was made at. + pub queue: Vec, + pub queue_stamp: Option, + /// When the index was last written to its cache file by the background pump. + pub saved_at: Option, + /// Bumped when another library is opened: what a worker was making for the old one is dropped. + pub epoch: u64, + /// Photos already searched for faces. + pub scanned: Scanned, + /// Photos that could not be searched this session (unreadable files): left alone until the next launch. + pub scan_failed: std::collections::HashSet, +} + +impl FacesState { + /// Another library is being opened: what was learned about the old one is saved, and nothing of it is kept (photo + /// ids are only unique within a library). The loaded model stays; the index is reloaded for the new library. + pub fn library_changed(&mut self) { + self.index.reset("", 0, None); + self.scanned.reset(None); + self.scan_failed.clear(); + self.embedder = None; + self.checked = None; + self.queue.clear(); + self.queue_stamp = None; + self.in_flight.clear(); + self.retry_at = None; + self.saved_at = None; + self.epoch = self.epoch.wrapping_add(1); + } +} diff --git a/crates/engine/src/faces_worker.rs b/crates/engine/src/faces_worker.rs new file mode 100644 index 00000000..81ff6ada --- /dev/null +++ b/crates/engine/src/faces_worker.rs @@ -0,0 +1,483 @@ +//! The heavy part of scanning faces, run off the session thread, on a few threads at once. +//! +//! For one photo: render it (upright, uncropped, default settings), find its faces with the YuNet detector (when installed) +//! so each region can be aligned by landmarks, cut each region out and run the recognition model. A photo that has +//! no face regions at all and has not been searched yet also gets its detections back, with their embeddings, so the +//! scan decodes each photo once. The session thread only prepares the job (which needs the catalog) and later ingests +//! the result, so the window never waits for a model. + +use std::panic::{AssertUnwindSafe, catch_unwind}; +use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering}; +use std::sync::mpsc::{Receiver, Sender, channel}; +use std::sync::{Arc, Condvar, Mutex, PoisonError}; +use std::time::Duration; + +use lightcraft_catalog::PhotoId; +use lightcraft_faces::align::{Rgb, align_to_template, crop_box}; +use lightcraft_faces::runtime::Embedder; +use lightcraft_faces::yunet::{Detector, Face, Options}; +use lightcraft_geom::Rect; + +use crate::media::{PreviewLoader, RenderJob}; + +/// Long edge of the picture faces are cut from. +pub(crate) const EDGE: usize = 2048; +/// A box from an XMP sidecar is cut with this much room around it (there are no landmarks to align by). +pub(crate) const BOX_PAD: f32 = 1.5; +/// How much a detected face must overlap a region to lend it its landmarks. +pub(crate) const MATCH_IOU: f32 = 0.3; +/// A detection must be this sure to become a face region (the bar `faces.detect` uses). +pub(crate) const REGION_SCORE: f32 = 0.6; +/// A raw's embedded preview is used for the scan only when its long edge is at least this (a thumbnail is too small to +/// find faces in; the raw itself is decoded then). +const MIN_PREVIEW_EDGE: usize = 1000; + +/// Most photos ever worked on at once. +const MAX_WORKERS: usize = 24; +/// What one photo in progress holds: the file, its decoded picture and the copies the detector and recogniser read. +/// Measured: peak memory grew by about 180 MB for each extra worker. +const WORKER_BYTES: usize = 180 << 20; +/// Photos queued behind the ones running, so a worker that finishes never waits for the next frame to be fed. +pub(crate) const PREFETCH: usize = 2; + +/// How hard the background scan may work right now, as the app judges from what the user is doing. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum Pace { + /// Start nothing new: the user is dragging, typing or scrolling. + Pause = 0, + /// One photo at a time: the user is around, or the window is minimized. + Light = 1, + /// Half the machine: the user is idle. + Normal = 2, + /// Most of the machine: the user is watching the progress. + Full = 3, +} + +impl Pace { + fn from_u8(n: u8) -> Pace { + match n { + 0 => Pace::Pause, + 1 => Pace::Light, + 3 => Pace::Full, + _ => Pace::Normal, + } + } + + /// `pause`, `light`, `normal` or `full`; anything else (or nothing) is `normal`. + pub fn parse(name: Option<&str>) -> Pace { + match name { + Some("pause") => Pace::Pause, + Some("light") => Pace::Light, + Some("full") => Pace::Full, + _ => Pace::Normal, + } + } +} + +/// How many photos are worked on at once at `pace` on a machine with `cores` threads (as the system reports them, which +/// respects container limits and affinity). `cap` is a limit the user set (`LIGHTCRAFT_FACE_THREADS`). Memory is not +/// decided here: decodes wait at the process-wide memory gate ([`crate::memory::work_gate`]), sized from the RAM. +pub(crate) fn workers_for(cores: usize, pace: Pace, cap: Option) -> usize { + let n = match pace { + Pace::Pause => return 0, + Pace::Light => 1, + Pace::Normal => cores / 2, + Pace::Full => cores * 4 / 5, + } + .clamp(1, MAX_WORKERS); + cap.map_or(n, |c| n.min(c.clamp(1, MAX_WORKERS))) +} + +fn cores() -> usize { + std::thread::available_parallelism().map_or(1, |n| n.get()) +} + +fn cap_from_env() -> Option { + std::env::var("LIGHTCRAFT_FACE_THREADS").ok().and_then(|v| v.trim().parse::().ok()) +} + +/// The most photos the memory budget can hold in progress at once: half of it (the rest is the window's, the caches', +/// the renders'), a worker at a time. Raising the budget (`LIGHTCRAFT_MEMORY_MB`) raises this. +fn memory_workers() -> usize { + (crate::memory::budget() / 2 / WORKER_BYTES).max(1) +} + +/// Photos to work on at once at `pace` on this machine: by the processor's threads and the pace, and never more than +/// the memory allows. A limit the user sets (`LIGHTCRAFT_FACE_THREADS`) replaces the memory one. +pub(crate) fn target_workers(pace: Pace) -> usize { + workers_for(cores(), pace, Some(cap_from_env().unwrap_or_else(memory_workers))) +} + +/// Threads the worker pool starts with: enough for the fullest pace (the idle ones cost nothing). +pub(crate) fn max_workers() -> usize { + target_workers(Pace::Full) +} + +/// Threads for the scan's own parallel work (decoding a picture, developing it) at `pace` on a machine with `cores` threads: +/// two when light, half the machine when normal, four fifths when watched. A photo's parallel parts are most of its cost, +/// so this, more than the number of photos at once, is how much of the processor the scan takes. Measured on a 32-thread +/// desktop (4 photos at once): 8 threads scanned 2.2 photos a second, 16 threads 5.8, 25 threads 6.4. +fn pool_threads_for(cores: usize, pace: Pace, cap: Option) -> usize { + let n = match pace { + Pace::Pause | Pace::Light => 2, + Pace::Normal => cores / 2, + Pace::Full => cores * 4 / 5, + } + .clamp(2, 32); + cap.map_or(n, |c| n.min(c.max(1))) +} + +/// One photo's faces to embed (and, if `detect`, to find first). +pub(crate) struct Prepared { + /// `@` the embeddings are for. + pub tag: String, + /// The session's face epoch when this was prepared: another library has been opened since if it differs. + pub epoch: u64, + pub id: PhotoId, + pub job: RenderJob, + /// A raw's embedded preview, tried before the render job (much faster than decoding the raw). + pub preview: Option<(String, PreviewLoader)>, + pub todo: Vec, + /// The photo has no face regions and has not been searched: look for faces too. + pub detect: bool, + /// The YuNet detector, when the user has installed it: faces are aligned by its landmarks, and it finds the faces of + /// photos that have none. Without it a scan only embeds the faces the photos already have, cut by their boxes. + pub detector: Option>, +} + +/// The embeddings made for a [`Prepared`]; `None` for a face that could not be embedded. +pub(crate) struct Done { + pub tag: String, + pub epoch: u64, + pub id: PhotoId, + /// Each face's region, its embedding (`None` when it could not be made) and the box to show it by. + pub results: Vec<(Rect, Option>, Rect)>, + /// Faces were to be looked for in this photo. + pub detect_wanted: bool, + /// What the detector found (their embeddings are in `results`); `None` when the photo could not be searched. + pub found: Option>, +} + +fn overlap(a: &Rect, f: &Face) -> f32 { + let (w, h) = ((a.x1.min(f64::from(f.x1)) - a.x0.max(f64::from(f.x0))).max(0.0), (a.y1.min(f64::from(f.y1)) - a.y0.max(f64::from(f.y0))).max(0.0)); + let inter = w * h; + let union = (a.x1 - a.x0) * (a.y1 - a.y0) + f64::from((f.x1 - f.x0) * (f.y1 - f.y0)) - inter; + if union > 0.0 { (inter / union) as f32 } else { 0.0 } +} + +/// A detection as a region box: inside the photo and not a sliver; `None` for anything else (it is only ever data). +fn face_rect(f: &Face) -> Option { + if ![f.x0, f.y0, f.x1, f.y1].iter().all(|v| v.is_finite()) { + return None; + } + let c = |v: f32| f64::from(v.clamp(0.0, 1.0)); + let r = Rect { x0: c(f.x0), y0: c(f.y0), x1: c(f.x1), y1: c(f.y1) }; + (r.x1 - r.x0 > 0.005 && r.y1 - r.y0 > 0.005).then_some(r) +} + +/// Do the work for one photo. Never panics: whatever goes wrong, its faces come back as `None`. +pub(crate) fn process(p: Prepared, embedder: &Embedder) -> Done { + let Prepared { tag, epoch, id, job, preview, todo, detect, detector } = p; + let fallback = |todo: Vec| todo.into_iter().map(|r| (r, None, r)).collect::>(); + let rects = todo.clone(); + let run = catch_unwind(AssertUnwindSafe(|| { + let from_preview = preview.and_then(|(path, load)| load(&path, EDGE)).filter(|img| img.width.max(img.height) >= MIN_PREVIEW_EDGE); + let image = match from_preview { + Some(img) => img, + None => match job.run().rendered { + Ok(rendered) => rendered.image, + Err(_) => return (fallback(todo), None), + }, + }; + let (w, h) = (image.width, image.height); + let rgb: Vec = image.data.iter().flat_map(|px| [px[0], px[1], px[2]]).collect(); + // the detector's landmarks, where it finds the same face: faces are aligned by them, not cut by a box + let detected = detector.as_deref().and_then(|d| d.detect(&rgb, w, h, &Options { score: 0.5, ..Options::default() }).ok()); + let found = detected.clone().unwrap_or_default(); + let img = Rgb { data: &rgb, width: w, height: h }; + let (iw, ih) = embedder.input_size(); + let scaled = |f: &Face| f.landmarks.map(|(x, y)| (x * w as f32, y * h as f32)); + let mut results: Vec<(Rect, Option>, Rect)> = todo + .into_iter() + .map(|rect| { + let best = found.iter().map(|f| (overlap(&rect, f), f)).filter(|(o, _)| *o >= MATCH_IOU).max_by(|a, b| a.0.total_cmp(&b.0)); + let crop = match best { + Some((_, f)) => align_to_template(&img, &scaled(f), iw, ih), + None => crop_box( + &img, + rect.x0 as f32 * w as f32, + rect.y0 as f32 * h as f32, + rect.x1 as f32 * w as f32, + rect.y1 as f32 * h as f32, + BOX_PAD, + iw, + ih, + ), + }; + let v = crop.and_then(|c| embedder.embed(&c).ok()); + // shown by the detector's box where it found this face: every face then has the same tightness + let view = best.and_then(|(_, f)| face_rect(f)).unwrap_or(rect); + (rect, v, view) + }) + .collect(); + // a photo with no face regions: what the detector is sure of becomes its faces, embedded in the same pass + let new_faces = if detect { + detected.map(|faces| { + let mut rects = Vec::new(); + for f in faces.iter().filter(|f| f.score >= REGION_SCORE) { + let Some(rect) = face_rect(f) else { continue }; + let v = align_to_template(&img, &scaled(f), iw, ih).and_then(|c| embedder.embed(&c).ok()); + results.push((rect, v, rect)); + rects.push(rect); + } + rects + }) + } else { + None + }; + (results, new_faces) + })); + let (results, found) = run.unwrap_or_else(|_| (fallback(rects), None)); + Done { tag, epoch, id, results, detect_wanted: detect, found } +} + +type Job = (Arc, Prepared); + +/// How many photos may be worked on at once, changed at any moment by the pace: a thread that has a job waits here until +/// fewer than that are running. A pace of zero holds every job where it is, with nothing decoded. +struct Limit { + allowed: AtomicUsize, + /// The pace the limit was last set for (a [`Pace`] as a number): the pool a photo's parallel work uses. + pace: AtomicU8, + running: Mutex, + wake: Condvar, +} + +impl Limit { + fn new() -> Limit { + Limit { allowed: AtomicUsize::new(0), pace: AtomicU8::new(Pace::Normal as u8), running: Mutex::new(0), wake: Condvar::new() } + } + + fn set(&self, n: usize, pace: Pace) { + self.pace.store(pace as u8, Ordering::Relaxed); + self.allowed.store(n, Ordering::Relaxed); + self.wake.notify_all(); + } + + fn pace(&self) -> Pace { + Pace::from_u8(self.pace.load(Ordering::Relaxed)) + } + + /// Wait for a place; the limit is looked at again whenever it changes and at least every 100 ms. + fn enter(&self) { + let mut running = self.running.lock().unwrap_or_else(PoisonError::into_inner); + while *running >= self.allowed.load(Ordering::Relaxed) { + running = self.wake.wait_timeout(running, Duration::from_millis(100)).unwrap_or_else(PoisonError::into_inner).0; + } + *running += 1; + } + + fn leave(&self) { + let mut running = self.running.lock().unwrap_or_else(PoisonError::into_inner); + *running = running.saturating_sub(1); + drop(running); + self.wake.notify_one(); + } +} + +/// The threads a photo's parallel work (decoding, developing a picture) runs on at `pace`: one pool for each of three sizes +/// ([`pool_threads_for`]), made when first needed, whose idle threads cost nothing. rayon's global pool serves the +/// interactive work (the loupe, exports) and has no priorities, so scan work queued there would make a slider drag wait +/// behind it; here a light scan can use only two threads however large the machine is. +fn scan_pool(pace: Pace) -> Option<&'static rayon::ThreadPool> { + static LIGHT: std::sync::OnceLock> = std::sync::OnceLock::new(); + static NORMAL: std::sync::OnceLock> = std::sync::OnceLock::new(); + static FULL: std::sync::OnceLock> = std::sync::OnceLock::new(); + let (cell, pace) = match pace { + Pace::Pause | Pace::Light => (&LIGHT, Pace::Light), + Pace::Normal => (&NORMAL, Pace::Normal), + Pace::Full => (&FULL, Pace::Full), + }; + cell.get_or_init(|| { + rayon::ThreadPoolBuilder::new() + .num_threads(pool_threads_for(cores(), pace, cap_from_env())) + .thread_name(move |i| format!("lightcraft-faces-{}-{i}", pace as u8)) + .build() + .ok() + }) + .as_ref() +} + +/// Run one photo on the pool for the pace, so its parallel parts stay there. +/// +/// It must NOT wait at the memory gate ([`crate::memory::work_gate`], which background decodes do): a job running on a +/// pool thread that waits inside a nested rayon call can pick up another job to run on top of itself, and if that one waits +/// for memory the first holds, neither can finish (seen as a few photos that never come back). The memory a scan may use is +/// instead limited by how many photos run at once ([`memory_workers`]); its decodes count as interactive ones at the gate, +/// which never wait. +fn process_in_background(prepared: Prepared, embedder: &Embedder, pace: Pace) -> Done { + let work = || process(prepared, embedder); + match scan_pool(pace) { + Some(pool) => pool.install(work), + None => work(), + } +} + +/// Threads that run [`process`] for the jobs they are given, each on one photo at a time. +pub(crate) struct Worker { + jobs: Sender, + done: Receiver, + limit: Arc, +} + +impl Worker { + pub fn start(threads: usize) -> Option { + let (jobs, job_rx) = channel::(); + let (done_tx, done) = channel::(); + let job_rx = Arc::new(Mutex::new(job_rx)); + let limit = Arc::new(Limit::new()); + let mut started = 0; + for i in 0..threads.clamp(1, MAX_WORKERS) { + let (rx, tx, limit) = (job_rx.clone(), done_tx.clone(), limit.clone()); + let spawned = std::thread::Builder::new().name(format!("lightcraft-faces-{i}")).spawn(move || { + loop { + // the lock is held only while waiting for a job, never while working on one; the threads end when + // the session (the only sender) is dropped + let next = rx.lock().unwrap_or_else(PoisonError::into_inner).recv(); + let Ok((embedder, prepared)) = next else { break }; + limit.enter(); + let done = process_in_background(prepared, &embedder, limit.pace()); + limit.leave(); + if tx.send(done).is_err() { + break; + } + } + }); + if spawned.is_ok() { + started += 1; + } + } + (started > 0).then_some(Worker { jobs, done, limit }) + } + + /// How many photos may be worked on at once from now on (what is running is not interrupted). + pub fn allow(&self, n: usize, pace: Pace) { + self.limit.set(n, pace); + } + + pub fn submit(&self, embedder: Arc, prepared: Prepared) -> bool { + self.jobs.send((embedder, prepared)).is_ok() + } + + /// Everything finished since the last call. + pub fn finished(&self) -> Vec { + self.done.try_iter().collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn the_pace_sets_how_much_of_the_machine_the_scan_may_use() { + use Pace::*; + // a 32-thread desktop: nothing while the user works, one photo around them, half when idle, most when watched + assert_eq!([Pause, Light, Normal, Full].map(|p| workers_for(32, p, None)), [0, 1, 16, 24]); + assert_eq!([Pause, Light, Normal, Full].map(|p| workers_for(8, p, None)), [0, 1, 4, 6]); + assert_eq!([Pause, Light, Normal, Full].map(|p| workers_for(4, p, None)), [0, 1, 2, 3]); + // small machines still get one worker (except while paused) + assert_eq!([Pause, Light, Normal, Full].map(|p| workers_for(1, p, None)), [0, 1, 1, 1]); + assert_eq!(workers_for(0, Full, None), 1); + assert_eq!(workers_for(10_000, Full, None), MAX_WORKERS); + // a limit set by the user caps every pace, and is itself kept sane + assert_eq!([Light, Normal, Full].map(|p| workers_for(32, p, Some(3))), [1, 3, 3]); + assert_eq!((workers_for(32, Full, Some(0)), workers_for(32, Pause, Some(5))), (1, 0)); + assert_eq!(workers_for(32, Full, Some(usize::MAX)), 24); + } + + #[test] + fn the_limit_holds_threads_to_the_pace_and_can_change_at_any_moment() { + let limit = Arc::new(Limit::new()); + let (now, peak) = (Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0))); + let started = Arc::new(AtomicUsize::new(0)); + let threads: Vec<_> = (0..6) + .map(|_| { + let (limit, now, peak, started) = (limit.clone(), now.clone(), peak.clone(), started.clone()); + std::thread::spawn(move || { + limit.enter(); + started.fetch_add(1, Ordering::SeqCst); + let n = now.fetch_add(1, Ordering::SeqCst) + 1; + peak.fetch_max(n, Ordering::SeqCst); + std::thread::sleep(Duration::from_millis(40)); + now.fetch_sub(1, Ordering::SeqCst); + limit.leave(); + }) + }) + .collect(); + // paused: nothing starts + std::thread::sleep(Duration::from_millis(150)); + assert_eq!(started.load(Ordering::SeqCst), 0); + // two at a time, never more + limit.set(2, Pace::Normal); + std::thread::sleep(Duration::from_millis(100)); + assert!(started.load(Ordering::SeqCst) >= 2); + // raised, the rest go; every thread finishes + limit.set(6, Pace::Full); + for t in threads { + t.join().unwrap(); + } + assert_eq!(started.load(Ordering::SeqCst), 6); + assert!(peak.load(Ordering::SeqCst) >= 2, "it did run in parallel"); + assert_eq!(*limit.running.lock().unwrap(), 0, "every place was given back"); + } + + #[test] + fn memory_limits_the_workers_and_a_users_limit_replaces_it() { + // the default budget (a quarter of the RAM, at most 1.5 GiB) holds a few photos in progress, not dozens + assert!((1..=8).contains(&memory_workers()), "{}", memory_workers()); + assert!(target_workers(Pace::Full) >= 1); + assert_eq!(target_workers(Pace::Pause), 0); + } + + #[test] + fn pace_names_are_forgiving() { + assert_eq!(["pause", "light", "normal", "full"].map(|n| Pace::parse(Some(n))), [Pace::Pause, Pace::Light, Pace::Normal, Pace::Full]); + assert_eq!( + (Pace::parse(None), Pace::parse(Some("")), Pace::parse(Some("FULL")), Pace::parse(Some("turbo"))), + (Pace::Normal, Pace::Normal, Pace::Normal, Pace::Normal) + ); + } + + #[test] + fn the_scans_parallel_work_gets_a_pool_the_size_of_the_pace() { + use Pace::*; + assert_eq!([Pause, Light, Normal, Full].map(|p| pool_threads_for(32, p, None)), [2, 2, 16, 25]); + assert_eq!([Light, Normal, Full].map(|p| pool_threads_for(8, p, None)), [2, 4, 6]); + assert_eq!([Light, Normal, Full].map(|p| pool_threads_for(2, p, None)), [2, 2, 2]); + assert_eq!(pool_threads_for(10_000, Full, None), 32); + assert_eq!([Light, Normal, Full].map(|p| pool_threads_for(32, p, Some(3))), [2, 3, 3]); + // the pools exist and are the sizes asked for + for pace in [Light, Normal, Full] { + let pool = scan_pool(pace).expect("a pool can be built"); + assert_eq!(pool.install(rayon::current_num_threads), pool_threads_for(cores(), pace, cap_from_env())); + } + // a paused scan has the light pool; the pace is remembered as a number and read back + assert_eq!(scan_pool(Pause).map(|p| p as *const _), scan_pool(Light).map(|p| p as *const _)); + assert_eq!([0, 1, 2, 3, 9].map(Pace::from_u8), [Pause, Light, Normal, Full, Normal]); + } + + #[test] + fn detections_become_boxes_only_when_they_make_sense() { + let f = |x0, y0, x1, y1| Face { x0, y0, x1, y1, score: 0.9, landmarks: [(0.0, 0.0); 5] }; + assert!(face_rect(&f(0.2, 0.2, 0.4, 0.5)).is_some()); + // a box that spills over the edge is brought back inside + let r = face_rect(&f(-0.1, 0.1, 0.3, 1.2)).unwrap(); + assert_eq!((r.x0, r.y1), (0.0, 1.0)); + // slivers, inverted boxes and non-numbers are dropped, never turned into regions + for bad in [f(0.2, 0.2, 0.2, 0.5), f(0.4, 0.2, 0.2, 0.5), f(f32::NAN, 0.2, 0.4, 0.5), f(0.2, 0.2, f32::INFINITY, 0.5), f(1.5, 0.2, 2.0, 0.5)] + { + assert!(face_rect(&bad).is_none(), "{bad:?}"); + } + } +} diff --git a/crates/engine/src/files.rs b/crates/engine/src/files.rs index c620be92..01de8a7a 100644 --- a/crates/engine/src/files.rs +++ b/crates/engine/src/files.rs @@ -245,7 +245,8 @@ fn load_bytes_now(bytes: std::borrow::Cow<'_, [u8]>, max_edge: usize) -> Result< raw.opcodes.list3.retain(|op| !op.is_lens_correction()); let xy = lightcraft_raw::color::as_shot_white_xy(&raw); let t = lightcraft_raw::color::camera_transform(&raw, xy); - let camera_look = crate::camera_preview::fit_preview(&raw, &bytes, &t); + // the source's segmentation mattes (DNG semantic masks), read while the preview is fitted + let (camera_look, mattes) = rayon::join(|| crate::camera_preview::fit_preview(&raw, &bytes, &t), || dng_mattes(&bytes, &raw)); drop(bytes); // Previews and thumbnails bin the mosaic straight to (about) the size they need; only // larger levels (exports, 1:1) demosaic the whole sensor. @@ -306,9 +307,18 @@ fn load_bytes_now(bytes: std::borrow::Cow<'_, [u8]>, max_edge: usize) -> Result< } let (temp, tint) = xy_to_temp_tint(xy); let relative = crate::camera_preview::file_local_look(raw.format) && t.matrix_is_fallback; + // White balance re-evaluates the file's own colour model (when it has one and no + // file-local look matrix sits on top of it) + let camera_color = (!t.matrix_is_fallback && camera_look.is_none()).then(|| { + let tags = lightcraft_raw::ColorData { profile: Default::default(), ..raw.color.clone() }; + Arc::new(lightcraft_pipeline::CameraColor { tags, developed_for: xy }) + }); let camera_tone = camera_look.as_ref().map(|p| p.tone).or_else(|| raw.color.profile.tone_curve.as_ref().and_then(dng_tone_curve)); let (temp, tint) = if relative { (6500.0, 0.0) } else { (temp.round(), tint.round()) }; - return Ok((img, SourceInfo { raw: true, as_shot_temp: temp, as_shot_tint: tint, lens, relative_wb: relative, camera_tone })); + return Ok(( + img, + SourceInfo { raw: true, as_shot_temp: temp, as_shot_tint: tint, lens, relative_wb: relative, camera_color, camera_tone, mattes }, + )); } let d = lightcraft_codecs::decode(&bytes, fit_box(max_edge)).map_err(|e| e.to_string())?; drop(bytes); @@ -317,6 +327,42 @@ fn load_bytes_now(bytes: std::borrow::Cow<'_, [u8]>, max_edge: usize) -> Result< Ok((img.into_oriented(Orientation::from_exif(d.orientation)), SourceInfo::default())) } +/// The semantic masks of a DNG that AI masks understand, over the developed image (default crop, +/// oriented like it). Person mattes without a confidently selected pixel are left out (the +/// Subject mask then falls back to its heuristic); a sky matte counts even when empty (no sky). +fn dng_mattes(bytes: &[u8], raw: &lightcraft_raw::RawImage) -> Option> { + use lightcraft_pipeline::masks::{MatteKind, Mattes}; + if raw.format != lightcraft_raw::RawFormat::Dng { + return None; + } + let mut mattes = Mattes::default(); + for m in lightcraft_raw::semantic_masks(bytes) { + let Some(kind) = matte_kind(&m.name) else { continue }; + let Some(img) = m.developed(raw.active_area, raw.crop) else { continue }; + if kind != MatteKind::Sky && !img.data.iter().any(|&v| v >= 128) { + continue; + } + mattes.push(kind, img.into_oriented(raw.orientation)); + } + (!mattes.is_empty()).then(|| Arc::new(mattes)) +} + +/// What a semantic mask selects, by its `SemanticName`: Apple's (iPhone ProRAW) are named after +/// their auxiliary image types, `urn:com:apple:photo::aux:`. Other names: `None`. +fn matte_kind(name: &str) -> Option { + use lightcraft_pipeline::masks::MatteKind::*; + let kind = name.strip_prefix("urn:com:apple:photo:")?.rsplit(':').next()?; + Some(match kind { + "semanticskymatte" => Sky, + "semanticskinmatte" => Skin, + "semantichairmatte" => Hair, + "semanticteethmatte" => Teeth, + "semanticglassesmatte" => Glasses, + "portraiteffectsmatte" => Person, + _ => return None, + }) +} + /// A DNG `ProfileToneCurve` (linear in, linear out, 1.0 = white after exposure compensation) as the /// finish stage's camera tone curve: 32 knots, log-spaced over 12 stops below white; above white /// the camera tone's shoulder continues it. @@ -404,8 +450,16 @@ fn fit_box(max_edge: usize) -> lightcraft_codecs::DecodeOptions { } /// Filesystem-backed embedded-preview hook (native). +/// +/// Reads the whole file to find the preview, so it holds the memory gate for twice the file's size (the file and the +/// decoded preview): background work (the face scan) waits for room, interactive work is counted and never waits. pub fn fs_preview_loader() -> PreviewLoader { - Arc::new(|path: &str, max_edge: usize| embedded_preview_srgb(&std::fs::read(path).ok()?, max_edge)) + Arc::new(|path: &str, max_edge: usize| { + let len = std::fs::metadata(path).map(|m| usize::try_from(m.len()).unwrap_or(usize::MAX)).unwrap_or(0); + let gate = crate::memory::work_gate(); + let _permit = if crate::memory::is_background() { gate.acquire(len.saturating_mul(2)) } else { gate.acquire_urgent(len.saturating_mul(2)) }; + embedded_preview_srgb(&std::fs::read(path).ok()?, max_edge) + }) } /// Filesystem-backed hooks (native). On the web the host installs bytes-based hooks instead. @@ -417,7 +471,7 @@ pub fn fs_preview_loader() -> PreviewLoader { /// yields to them. pub fn fs_hooks() -> (FileLoader, FileProbe) { let loader: FileLoader = Arc::new(|path: &str, max_edge: usize| { - let len = std::fs::metadata(path).map(|m| m.len() as usize).unwrap_or(0); + let len = std::fs::metadata(path).map(|m| usize::try_from(m.len()).unwrap_or(usize::MAX)).unwrap_or(0); let weight = len * if max_edge <= crate::media::SourceLevel::Thumb.max_edge() { 3 } else { 6 }; let gate = crate::memory::work_gate(); let _permit = if crate::memory::is_background() { gate.acquire(weight) } else { gate.acquire_urgent(weight) }; @@ -428,7 +482,7 @@ pub fn fs_hooks() -> (FileLoader, FileProbe) { r }); let probe: FileProbe = Arc::new(|path: &str| { - let len = std::fs::metadata(path).map(|m| m.len() as usize).unwrap_or(0); + let len = std::fs::metadata(path).map(|m| usize::try_from(m.len()).unwrap_or(usize::MAX)).unwrap_or(0); let _permit = crate::memory::work_gate().acquire(len); let bytes = std::fs::read(path).map_err(|e| format!("{path}: {e}"))?; probe_bytes(path, &bytes) @@ -676,8 +730,12 @@ mod tests { assert!(info.camera_tone.is_none()); // a map that removes all saturation, and a tone curve let grey = HsvTable { hue_divisions: 4, sat_divisions: 2, val_divisions: 1, data: vec![[0.0, 0.0, 1.0]; 8], srgb_value: false }; - raw.color.profile = - ProfileLook { hue_sat_map: [Some(grey), None], look_table: None, tone_curve: ToneCurve::from_tag(&[0.0, 0.0, 0.18, 0.3, 1.0, 1.0]) }; + raw.color.profile = ProfileLook { + hue_sat_map: [Some(grey), None], + look_table: None, + tone_curve: ToneCurve::from_tag(&[0.0, 0.0, 0.18, 0.3, 1.0, 1.0]), + gain_table_map: None, + }; let with = lightcraft_raw::write_dng(&raw, &Default::default()).unwrap(); let (after, info) = load_bytes(&with, 64).unwrap(); assert!(sat(&after) < 1e-3, "saturation {} → {}", sat(&before), sat(&after)); @@ -695,6 +753,119 @@ mod tests { assert!((mean(&dim) / mean(&before) - 0.5).abs() < 0.02, "{} vs {}", mean(&dim), mean(&before)); } + /// iPhone ProRAW-style semantic masks drive the Sky mask: Apple's sky matte (here the right + /// half of the sensor, i.e. the bottom of the photo once turned 90° clockwise — where the sky + /// heuristic would never look) is used; a matte with an unknown name is ignored. + #[test] + fn dng_sky_matte_drives_the_sky_mask() { + use lightcraft_develop::{DevelopSettings, LocalAdjustments, Mask, MaskComponent, MaskOp, MaskShape}; + use lightcraft_pipeline::{RenderRequest, render}; + use lightcraft_tiff::tags::{self as t, photometric}; + use lightcraft_tiff::{ByteOrder, IfdBuilder, ImageData, TiffWriter, Value}; + let (w, h) = (16u32, 8u32); + let image = |ifd: &mut IfdBuilder, w: u32, h: u32, cpp: u16, bits: u16, data: Vec| { + ifd.set(t::IMAGE_WIDTH, Value::Long(vec![w])); + ifd.set(t::IMAGE_LENGTH, Value::Long(vec![h])); + ifd.set(t::BITS_PER_SAMPLE, Value::Short(vec![bits; cpp as usize])); + ifd.set(t::SAMPLES_PER_PIXEL, Value::Short(vec![cpp])); + ifd.set(t::COMPRESSION, Value::Short(vec![1])); + ifd.set_image(ImageData::Strips { rows_per_strip: h, strips: vec![data] }); + }; + let mut raw = IfdBuilder::new(); + raw.set(t::NEW_SUBFILE_TYPE, Value::Long(vec![0])); + raw.set(t::PHOTOMETRIC, Value::Short(vec![photometric::LINEAR_RAW])); + image(&mut raw, w, h, 3, 16, 12000u16.to_le_bytes().repeat((w * h * 3) as usize)); + let matte = |name: &str, right: bool| { + let mut m = IfdBuilder::new(); + m.set(t::NEW_SUBFILE_TYPE, Value::Long(vec![t::SUBFILE_SEMANTIC_MASK])); + m.set(t::PHOTOMETRIC, Value::Short(vec![photometric::MASK])); + m.set(t::SEMANTIC_NAME, Value::Ascii(name.into())); + image(&mut m, w / 2, h / 2, 1, 8, (0..w / 2 * h / 2).map(|i| if (i % (w / 2) >= w / 4) == right { 255 } else { 0 }).collect()); + m + }; + let mut ifd0 = IfdBuilder::new(); + ifd0.set(t::NEW_SUBFILE_TYPE, Value::Long(vec![1])); + ifd0.set(t::DNG_VERSION, Value::Byte(vec![1, 6, 0, 0])); + ifd0.set(t::MAKE, Value::Ascii("Apple".into())); + ifd0.set(t::ORIENTATION, Value::Short(vec![6])); + ifd0.set(t::COLOR_MATRIX_1, Value::SRational(vec![(1, 1), (0, 1), (0, 1), (0, 1), (1, 1), (0, 1), (0, 1), (0, 1), (1, 1)])); + ifd0.set(t::CALIBRATION_ILLUMINANT_1, Value::Short(vec![21])); + ifd0.set(t::PHOTOMETRIC, Value::Short(vec![photometric::RGB])); + image(&mut ifd0, 4, 2, 3, 8, vec![128; 24]); + ifd0.add_sub_ifd(raw); + ifd0.add_sub_ifd(matte("urn:com:apple:photo:2020:aux:semanticskymatte", true)); + ifd0.add_sub_ifd(matte("urn:com:apple:photo:2020:aux:semanticsomethingelse", false)); + let bytes = TiffWriter::new(ByteOrder::Little, false).write(&[ifd0]).unwrap(); + + let (img, info) = load_bytes(&bytes, 64).unwrap(); + assert_eq!((img.width, img.height), (8, 16)); + assert!(info.mattes.is_some()); + let sky = Mask { + components: vec![MaskComponent { name: None, op: MaskOp::Add, invert: false, shape: MaskShape::Sky }], + adjust: LocalAdjustments { exposure: -3.0, ..Default::default() }, + ..Default::default() + }; + let s = DevelopSettings { masks: vec![sky], ..Default::default() }; + let req = RenderRequest::fit(8, 16); + let green = |s: &DevelopSettings, info: &SourceInfo, y: usize| render(&img, info, s, &req).image.get(4, y)[1] as i32; + let plain = DevelopSettings::default(); + assert_eq!(green(&s, &info, 2), green(&plain, &info, 2), "no sky at the top"); + assert!(green(&s, &info, 13) + 40 < green(&plain, &info, 13), "the sky at the bottom is darkened"); + // without the matte the heuristic looks at the top of the frame instead + let heuristic = SourceInfo { mattes: None, ..info.clone() }; + assert!(green(&s, &heuristic, 2) < green(&plain, &heuristic, 2)); + assert_eq!(green(&s, &heuristic, 13), green(&plain, &heuristic, 13)); + } + + /// A DNG's white balance is re-evaluated through its own colour model: the decoder hands the + /// pipeline the file's colour tags (without the profile look, applied at load) and the white + /// the pixels were developed for, and a custom white balance renders a neutral differently + /// from the Bradford adaptation it replaces. + #[test] + fn dng_white_balance_uses_the_camera_colour_model() { + use lightcraft_pipeline::{RenderRequest, render}; + let dng = crate::tests_xmp::synthetic_dng_with(None, Default::default()); + let raw = lightcraft_raw::decode(&dng).unwrap(); + let (img, info) = load_bytes(&dng, 64).unwrap(); + let cc = info.camera_color.as_ref().expect("a DNG with a colour matrix"); + assert_eq!(cc.developed_for, lightcraft_raw::color::as_shot_white_xy(&raw)); + assert!(cc.tags.profile.is_empty()); + assert_eq!(cc.tags.color_matrix, raw.color.color_matrix); + let mut s = lightcraft_develop::DevelopSettings::default(); + (s.wb.mode, s.wb.temp, s.wb.tint) = (lightcraft_develop::WbMode::Custom, 3000.0, 20.0); + let req = RenderRequest::fit(32, 32); + let camera = render(&img, &info, &s, &req).image; + let adapted = render(&img, &SourceInfo { camera_color: None, ..info.clone() }, &s, &req).image; + assert_ne!(camera.data, adapted.data); + } + + /// Apple ProRAW: like Lightroom Classic, the default render ignores the file's + /// `ProfileGainTableMap` (Apple's local tone mapping); a map that would double every pixel + /// changes nothing. + #[test] + fn dng_gain_table_map_is_not_rendered() { + use lightcraft_raw::gaintable::GainTableMap; + let plain = crate::tests_xmp::synthetic_dng_with(None, Default::default()); + let mut raw = lightcraft_raw::decode(&plain).unwrap(); + let (before, _) = load_bytes(&lightcraft_raw::write_dng(&raw, &Default::default()).unwrap(), 64).unwrap(); + raw.color.profile.gain_table_map = Some(GainTableMap { + points_v: 1, + points_h: 1, + points_n: 1, + spacing_v: 1.0, + spacing_h: 1.0, + origin_v: 0.0, + origin_h: 0.0, + weights: [0.2, 0.2, 0.2, 0.2, 0.2], + gamma: 1.0, + gains: vec![2.0], + }); + let with_map = lightcraft_raw::write_dng(&raw, &Default::default()).unwrap(); + assert!(lightcraft_raw::decode(&with_map).unwrap().color.profile.gain_table_map.is_some(), "the map is kept"); + let (after, _) = load_bytes(&with_map, 64).unwrap(); + assert_eq!(before.data, after.data); + } + #[test] fn unsupported_raw_falls_back_to_embedded_preview() { let b = cr3_with_preview(48, 32); diff --git a/crates/engine/src/lib.rs b/crates/engine/src/lib.rs index f26b4ca7..c3fcdd33 100644 --- a/crates/engine/src/lib.rs +++ b/crates/engine/src/lib.rs @@ -14,11 +14,17 @@ pub mod availability; mod camera_preview; pub mod camera_profiles; pub mod cmd; +pub mod config; pub mod crs; pub mod crs_masks; pub mod demo; pub mod devices; pub mod export; +pub mod face_download; +#[cfg(not(target_arch = "wasm32"))] +mod faces_index; +#[cfg(not(target_arch = "wasm32"))] +mod faces_worker; pub mod files; pub mod fonts; pub mod guard; @@ -157,6 +163,15 @@ pub struct Session { /// Set by a command whose change must not rewrite the photo's XMP sidecar even with auto-write on /// (a catalog-only edit of data the sidecar writer does not emit); consumed when the command ends. pub(crate) skip_auto_write: bool, + /// Where the host keeps face models (one folder each); `None` where there is no file system (the web). + pub face_models_dir: Option, + /// Face model downloads started this session (the staged files wait in `/.downloads`). + pub face_downloads: face_download::Downloads, + /// The user's own list of models to download (`catalog.json` in the models folder), as last read. + pub(crate) face_catalog: lightcraft_faces::catalog::Catalog, + /// The loaded recognition model and the face embeddings made with it. + #[cfg(not(target_arch = "wasm32"))] + pub(crate) faces: faces_index::FacesState, /// Copied develop settings (partial JSON) for Paste. pub clipboard: Option, /// The folder on disk the [`LibrarySource::Folder`] view browses. @@ -263,6 +278,11 @@ impl Session { redo: Vec::new(), interaction: None, skip_auto_write: false, + face_models_dir: None, + face_downloads: Default::default(), + face_catalog: Default::default(), + #[cfg(not(target_arch = "wasm32"))] + faces: Default::default(), clipboard: None, meta_clipboard: None, browse: None, @@ -403,6 +423,27 @@ impl Session { Ok(()) } + /// The box to cut a face's picture from: the detector's, when the scan has looked at the face (every face is then shown + /// equally close, however loosely or tightly its own region was drawn), else the region's own box. + pub fn face_view(&self, id: PhotoId, rect: lightcraft_geom::Rect) -> lightcraft_geom::Rect { + #[cfg(not(target_arch = "wasm32"))] + if let Some(v) = self.faces.index.view(id.0, &rect) { + return v; + } + let _ = id; + rect + } + + /// Apply an op that is LightCraft's own bookkeeping rather than something the user did (faces found by the background + /// scan): journaled like any op, but not an undo step, and it leaves the redo stack alone. + #[cfg(not(target_arch = "wasm32"))] + pub(crate) fn apply_system(&mut self, op: Op) -> Result<()> { + let fwd = op.clone(); + self.catalog.apply(op)?; + self.pending_log.push(fwd); + Ok(()) + } + /// Leaving photo `id` after editing it: keep its settings as an automatic version (when they /// differ from its latest version; at most [`AUTO_VERSIONS`] auto versions, oldest dropped). /// Saved with the library but not an undo step. @@ -785,6 +826,11 @@ mod tests_color; #[cfg(test)] mod tests_export; #[cfg(test)] +mod tests_face_models; +#[cfg(test)] +#[cfg(not(target_arch = "wasm32"))] +mod tests_face_recognize; +#[cfg(test)] mod tests_folders; #[cfg(test)] mod tests_forget_local; diff --git a/crates/engine/src/library.rs b/crates/engine/src/library.rs index c1b91c4c..d01c1b72 100644 --- a/crates/engine/src/library.rs +++ b/crates/engine/src/library.rs @@ -297,6 +297,9 @@ impl Session { { log::error!("library: {e}"); } + // what was learned about the old library's faces is saved, and none of it carries over (photo ids are per library) + #[cfg(not(target_arch = "wasm32"))] + self.faces.library_changed(); let (mut journal, mut catalog, report) = Journal::open(catalog)?; // A loaded catalog counts revisions from 0, like every other one. Caches (sidebar counts, // keyword tree, the grid's list…) are keyed on the revision, so give each library loaded diff --git a/crates/engine/src/media.rs b/crates/engine/src/media.rs index 1ed89e47..e3ed5815 100644 --- a/crates/engine/src/media.rs +++ b/crates/engine/src/media.rs @@ -62,7 +62,7 @@ impl SettingsHashes { } /// Bump when the pipeline's output changes, to invalidate cached thumbnails. -pub const RENDER_CACHE_VERSION: u64 = 15; +pub const RENDER_CACHE_VERSION: u64 = 16; /// Thumbnails render at one of these long edges (so window/cell size changes reuse the cache). pub const THUMB_SIZES: [usize; 4] = [128, 256, 384, 512]; @@ -143,7 +143,7 @@ impl DecodedSource { /// What to render these pixels against: the decoder's facts, else `header` (the catalog's) /// with any stored camera tone curve. pub fn info_or(&self, header: SourceInfo) -> SourceInfo { - self.info.unwrap_or(SourceInfo { camera_tone: self.camera_tone.or(header.camera_tone), ..header }) + self.info.clone().unwrap_or_else(|| SourceInfo { camera_tone: self.camera_tone.or(header.camera_tone), ..header }) } } @@ -253,8 +253,9 @@ fn rendered_budget(share: usize) -> usize { RENDERED_MEM_BYTES.min(share / 4) } -fn source_bytes(img: &Rgb32f) -> usize { - img.data.len() * 12 + 64 + std::mem::size_of::() +fn source_bytes(s: &DecodedSource) -> usize { + let mattes = s.info.as_ref().and_then(|i| i.mattes.as_ref()).map_or(0, |m| m.bytes()); + s.image.data.len() * 12 + 64 + std::mem::size_of::() + mattes } impl MediaCache { @@ -319,8 +320,8 @@ impl MediaCache { let Some((tick, which)) = candidates.into_iter().flatten().min() else { break }; let freed = match which { 0 => self.thumbs.pop_oldest(), - 1 => self.previews.iter().position(|e| e.2 == tick).map(|i| source_bytes(&self.previews.remove(i).1.image)), - 2 => self.full.take().map(|e| source_bytes(&e.1.image)), + 1 => self.previews.iter().position(|e| e.2 == tick).map(|i| source_bytes(&self.previews.remove(i).1)), + 2 => self.full.take().map(|e| source_bytes(&e.1)), _ => self.rendered.evict_oldest(), }; match freed { @@ -372,7 +373,7 @@ impl MediaCache { let tick = lightcraft_preview::next_tick(); match level { SourceLevel::Thumb => { - let cost = source_bytes(&img.image); + let cost = source_bytes(&img); self.thumbs.insert(id, img, cost); } SourceLevel::Preview => { @@ -415,8 +416,8 @@ impl MediaCache { /// Decoded sources held: (thumbnail level, preview level, full size). pub fn usage(&self) -> (crate::memory::Usage, crate::memory::Usage, crate::memory::Usage) { use crate::memory::Usage; - let previews = Usage::new(self.previews.len(), self.previews.iter().map(|e| source_bytes(&e.1.image)).sum()); - let full = self.full.as_ref().map(|e| Usage::new(1, source_bytes(&e.1.image))).unwrap_or_default(); + let previews = Usage::new(self.previews.len(), self.previews.iter().map(|e| source_bytes(&e.1)).sum()); + let full = self.full.as_ref().map(|e| Usage::new(1, source_bytes(&e.1))).unwrap_or_default(); (Usage::new(self.thumbs.len(), self.thumbs.cost()), previews, full) } @@ -979,6 +980,17 @@ impl crate::Session { } } + /// The camera's embedded preview of a raw file, for the face scan. Faces are read from the unedited picture, so the + /// photo's edits do not matter, and the camera's own JPEG (oriented, display-ready) costs a fraction of decoding the + /// raw. `None` for anything else, and where the host installed no loader. + #[cfg(not(target_arch = "wasm32"))] + pub(crate) fn scan_preview_of(&self, p: &Photo) -> Option<(String, PreviewLoader)> { + match (&p.source, &self.media.preview_loader) { + (Source::File { path }, Some(l)) if p.kind == MediaKind::Raw => Some((path.clone(), l.clone())), + _ => None, + } + } + /// Something to show in the loupe right away for `id` (see [`QuickJob`]): its cached view /// render, else the embedded preview of an unedited raw, else a cached or fresh thumbnail. pub fn quick_view_job(&mut self, id: PhotoId, max_edge: usize, apply_crop: bool) -> Option { @@ -1210,20 +1222,21 @@ mod tests { let mut s = crate::Session::with_demo(); let id = s.active().unwrap(); let mut job = s.render_job(id, 64, 64, false, true).unwrap(); + let loaded_info = info.clone(); job.source = SourceRef::File { path: "synthetic.arw".into(), max_edge: 64, loader: Some(Arc::new(move |_, _| { let mut image = Rgb32f::new(64, 64); image.data.fill([0.1; 3]); - Ok((image, info)) + Ok((image, loaded_info.clone())) })), fallback: None, }; job.info = SourceInfo::default(); // Header facts cannot override decoder facts. job.settings = Arc::new(DevelopSettings::default()); let r = job.clone().run(); - assert_eq!(r.loaded.as_ref().unwrap().info, Some(info)); + assert_eq!(r.loaded.as_ref().unwrap().info.as_ref(), Some(&info)); let expected = lightcraft_pipeline::render(&r.loaded.as_ref().unwrap().image, &info, &job.settings, &job.request); assert_eq!(r.rendered.as_ref().unwrap().image.data, expected.image.data); s.accept(&r); diff --git a/crates/engine/src/tests_face_models.rs b/crates/engine/src/tests_face_models.rs new file mode 100644 index 00000000..c71f7afe --- /dev/null +++ b/crates/engine/src/tests_face_models.rs @@ -0,0 +1,586 @@ +//! The `faces.*` model commands: list, inspect, install behind the licence, select, remove, the on/off +//! setting, and the refusals (no folder, licence not accepted, unsupported or hostile files and ids). + +use std::path::PathBuf; + +use serde_json::{Value, json}; + +use crate::Session; + +fn temp(name: &str) -> PathBuf { + let d = std::env::temp_dir().join(format!("lc-facemodels-{name}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&d); + std::fs::create_dir_all(&d).unwrap(); + d +} + +fn session(dir: &std::path::Path) -> Session { + let mut s = Session::new(); + s.face_models_dir = Some(dir.join("models")); + s +} + +fn model_file(dir: &std::path::Path, name: &str, dim: u64) -> String { + let p = dir.join(name); + std::fs::write(&p, lightcraft_faces::synthetic::tiny_embedder_model(dim)).unwrap(); + p.to_string_lossy().into_owned() +} + +fn find<'a>(list: &'a Value, id: &str) -> &'a Value { + list["models"].as_array().unwrap().iter().find(|m| m["id"] == id).unwrap_or(&Value::Null) +} + +#[test] +fn the_list_knows_the_known_models_and_is_honest_about_the_runtime() { + let d = temp("list"); + let mut s = session(&d); + let l = s.execute("faces.models.list", &json!({})).unwrap(); + assert_eq!(l["enabled"], false); + assert_eq!(l["runtime"], cfg!(not(target_arch = "wasm32"))); + assert_eq!(find(&l, "yunet-2023mar")["installed"], false, "nothing ships with LightCraft"); + assert_eq!(find(&l, "auraface-v1")["installed"], false); + assert_eq!(find(&l, "auraface-v1")["licence"]["commercial"], "yes"); + assert_eq!(find(&l, "sface-2021dec")["licence"]["commercial"], "unknown"); + // speed is a ratio to a ResNet-100 model, in words, for models that are not installed yet too; never milliseconds + assert_eq!(find(&l, "sface-2021dec")["speedText"], "6.9× faster than a ResNet-100 model"); + assert_eq!(find(&l, "auraface-v1")["speedText"], "Same speed as a ResNet-100 model"); + assert_eq!(find(&l, "yunet-2023mar")["speedText"], Value::Null, "nobody timed the detector against it"); + // a build with no folder lists too, and cannot install + let mut web = Session::new(); + assert!(web.execute("faces.models.list", &json!({})).is_ok()); + let f = model_file(&d, "m.onnx", 512); + assert!(web.execute("faces.models.install", &json!({"path": f, "acknowledged": true})).is_err()); + assert!(web.execute("faces.enable", &json!({"enabled": true})).is_err()); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn inspect_describes_a_file_and_installs_nothing() { + let d = temp("inspect"); + let mut s = session(&d); + let f = model_file(&d, "My Model.onnx", 512); + let r = s.execute("faces.models.inspect", &json!({"path": f})).unwrap(); + assert_eq!(r["kind"], "draft"); + assert_eq!(r["model"]["licence"]["commercial"], "unknown"); + assert_eq!(r["model"]["known"], false); + assert!(r["assumptions"].as_array().unwrap().iter().any(|a| a.as_str().unwrap().contains("127.5"))); + assert_eq!(r["alreadyInstalled"], false); + assert!(!d.join("models").exists(), "inspecting writes nothing"); + // not a model: a plain answer, not an error + let junk = d.join("junk.onnx"); + std::fs::write(&junk, b"this is not a model").unwrap(); + let r = s.execute("faces.models.inspect", &json!({"path": junk.to_string_lossy()})).unwrap(); + assert_eq!(r["kind"], "unsupported"); + assert!(r["reason"].as_str().unwrap().len() > 5); + // missing and empty files and folders are errors + assert!(s.execute("faces.models.inspect", &json!({"path": d.join("nope.onnx").to_string_lossy()})).is_err()); + assert!(s.execute("faces.models.inspect", &json!({"path": d.to_string_lossy()})).is_err()); + let empty = d.join("empty.onnx"); + std::fs::write(&empty, b"").unwrap(); + assert!(s.execute("faces.models.inspect", &json!({"path": empty.to_string_lossy()})).is_err()); + assert!(s.execute("faces.models.inspect", &json!({})).is_err()); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn install_needs_the_licence_accepted_then_select_enable_and_remove_work() { + let d = temp("flow"); + let mut s = session(&d); + let f = model_file(&d, "Mine.onnx", 512); + assert!(s.execute("faces.models.install", &json!({"path": f})).is_err(), "no acknowledgement"); + assert!(s.execute("faces.models.install", &json!({"path": f, "acknowledged": false})).is_err()); + assert!(s.execute("faces.models.install", &json!({"path": f, "acknowledged": "yes"})).is_err(), "must be the boolean true"); + assert!(!d.join("models").exists(), "refused installs leave nothing behind"); + + let r = s.execute("faces.models.install", &json!({"path": f, "acknowledged": true})).unwrap(); + let id = r["installed"]["id"].as_str().unwrap().to_string(); + assert!(id.starts_with("custom-mine-")); + let home = d.join("models").join(&id); + assert!(home.join("model.onnx").is_file() && home.join("face-model.json").is_file() && home.join("installed.json").is_file()); + assert!(!home.join("model.onnx.part").exists()); + assert_eq!(std::fs::read(home.join("model.onnx")).unwrap(), std::fs::read(&f).unwrap()); + let rec: Value = serde_json::from_slice(&std::fs::read(home.join("installed.json")).unwrap()).unwrap(); + assert_eq!(rec["commercial"], "unknown"); + assert_eq!(rec["fileName"], "Mine.onnx"); + + // listed as installed and custom; installing again is harmless + let l = s.execute("faces.models.list", &json!({})).unwrap(); + assert_eq!(find(&l, &id)["installed"], true); + assert_eq!(find(&l, &id)["known"], false); + s.execute("faces.models.install", &json!({"path": f, "acknowledged": true})).unwrap(); + assert_eq!(s.execute("faces.models.inspect", &json!({"path": f})).unwrap()["alreadyInstalled"], true); + + // choose it; the choice and the on/off setting survive a new session on the same folder + assert!(s.execute("faces.models.select", &json!({"id": "auraface-v1"})).is_err(), "not installed"); + assert!(s.execute("faces.models.select", &json!({"id": "yunet-2023mar"})).is_err(), "a detector is not a recogniser"); + s.execute("faces.models.select", &json!({"id": id})).unwrap(); + assert_eq!(s.execute("faces.enable", &json!({"enabled": true})).unwrap()["enabled"], true); + let mut again = session(&d); + let l = again.execute("faces.models.list", &json!({})).unwrap(); + assert_eq!((l["enabled"].clone(), l["embedder"].clone()), (json!(true), json!(id))); + assert_eq!(again.execute("faces.enable", &json!({})).unwrap()["enabled"], false, "toggles when omitted"); + + // removing it clears the choice + assert_eq!(s.execute("faces.models.remove", &json!({"id": id})).unwrap()["removed"], id); + assert!(!home.exists()); + assert_eq!(s.execute("faces.models.list", &json!({})).unwrap()["embedder"], Value::Null); + assert!(s.execute("faces.models.remove", &json!({"id": id})).is_err(), "already gone"); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn hostile_ids_unsupported_models_and_odd_folders_are_handled() { + let d = temp("hostile"); + let mut s = session(&d); + let models = d.join("models"); + std::fs::create_dir_all(&models).unwrap(); + // something that is not a model sits in the folder: it is never listed and never deleted + let precious = models.join("precious"); + std::fs::create_dir_all(&precious).unwrap(); + std::fs::write(precious.join("keep.txt"), b"mine").unwrap(); + for id in ["precious", "..", ".", "", "../x", "a/b", r"a\b", ".hidden", "yunet-2023mar", "UPPER"] { + assert!(s.execute("faces.models.remove", &json!({"id": id})).is_err(), "{id:?}"); + } + assert!(precious.join("keep.txt").is_file()); + assert!(d.join("models").is_dir()); + // a folder whose manifest lies about its id is ignored + let liar = models.join("liar"); + std::fs::create_dir_all(&liar).unwrap(); + std::fs::write(liar.join("model.onnx"), lightcraft_faces::synthetic::embedder_model(64)).unwrap(); + let mut m = lightcraft_faces::known::auraface(); + m.id = "someone-else".into(); + std::fs::write(liar.join("face-model.json"), serde_json::to_vec(&m).unwrap()).unwrap(); + std::fs::write(models.join("settings.json"), b"{ not json").unwrap(); + let l = s.execute("faces.models.list", &json!({})).unwrap(); + assert!(find(&l, "someone-else").is_null() && find(&l, "liar").is_null()); + assert_eq!(l["enabled"], false, "a corrupt settings file is ignored"); + // not-a-model and wrong-shaped files cannot be installed + let junk = d.join("junk.onnx"); + std::fs::write(&junk, b"nonsense").unwrap(); + let e = s.execute("faces.models.install", &json!({"path": junk.to_string_lossy(), "acknowledged": true})).unwrap_err().to_string(); + assert!(e.contains("cannot be used yet"), "{e}"); + let tokens = d.join("tokens.onnx"); + { + use lightcraft_faces::synthetic::{len_field, value_info_bytes}; + let mut graph = Vec::new(); + len_field(11, &value_info_bytes("pixels", 1, &[Err("N"), Ok(3), Ok(224), Ok(224)]), &mut graph); + len_field(12, &value_info_bytes("tokens", 1, &[Err("N"), Ok(257), Ok(384)]), &mut graph); + let mut model = Vec::new(); + len_field(7, &graph, &mut model); + std::fs::write(&tokens, model).unwrap(); + } + assert!(s.execute("faces.models.install", &json!({"path": tokens.to_string_lossy(), "acknowledged": true})).is_err()); + assert!(s.execute("faces.models.select", &json!({"id": 5})).is_err()); + assert!(s.execute("faces.models.select", &json!({"id": ""})).is_err()); + let _ = std::fs::remove_dir_all(&d); +} + +/// A demo session with the real YuNet installed, for the tests that need it: `LC_YUNET_MODEL=` (what +/// Settings > Faces downloads; nothing is committed). Without it those tests say so and pass. +fn demo_with_yunet() -> Option { + let Some(file) = std::env::var_os("LC_YUNET_MODEL") else { + eprintln!("LC_YUNET_MODEL is not set: skipping"); + return None; + }; + let dir = temp("yunet"); + let mut s = Session::with_demo(); + s.face_models_dir = Some(dir.join("models")); + let r = s.execute("faces.models.install", &json!({"path": file.to_string_lossy(), "acknowledged": true})); + assert!(r.is_ok(), "{r:?}"); + Some(s) +} + +#[test] +fn installing_the_detector_does_not_choose_it_as_the_recogniser_or_switch_recognition_on() { + let Some(mut s) = demo_with_yunet() else { return }; + let l = s.execute("faces.models.list", &json!({})).unwrap(); + assert_eq!(find(&l, "yunet-2023mar")["installed"], true); + assert_eq!(l["enabled"], false, "{l}"); + assert!(l["embedder"].is_null(), "{l}"); + assert!(s.execute("faces.models.test", &json!({"id": "yunet-2023mar"})).is_err(), "a detector has no recognition self-test"); +} + +#[test] +fn detect_without_the_model_says_where_to_get_it() { + let d = temp("nodetector"); + let mut s = Session::with_demo(); + // no folder for models at all, then a folder that has none yet + assert!(s.execute("faces.detect", &json!({})).is_err()); + s.face_models_dir = Some(d.join("models")); + let e = s.execute("faces.detect", &json!({})).unwrap_err().to_string(); + assert!(e.contains("Settings > Faces"), "{e}"); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn the_detector_and_the_recognisers_can_be_downloaded_and_the_download_is_refused_otherwise() { + let d = temp("dl-refusals"); + let mut s = session(&d); + let l = s.execute("faces.models.list", &json!({})).unwrap(); + for (id, host) in [("yunet-2023mar", "github.com"), ("sface-2021dec", "github.com"), ("auraface-v1", "huggingface.co")] { + assert_eq!(find(&l, id)["downloadHost"], host, "{id}"); + } + // without the user's acceptance, for an id nobody knows, and without a folder: nothing starts + assert!(s.execute("faces.models.download", &json!({"id": "yunet-2023mar"})).is_err()); + assert!(s.execute("faces.models.download", &json!({"id": "yunet-2023mar", "acknowledged": false})).is_err()); + for id in ["nope", "", "../yunet-2023mar"] { + assert!(s.execute("faces.models.download", &json!({"id": id, "acknowledged": true})).is_err(), "{id}"); + } + assert!(Session::new().execute("faces.models.download", &json!({"id": "yunet-2023mar", "acknowledged": true})).is_err()); + assert_eq!(s.execute("faces.models.downloads", &json!({})).unwrap()["downloads"], json!([])); + assert!(!d.join("models").join(".downloads").exists(), "a refused download touches nothing"); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn a_download_that_has_arrived_is_installed_by_itself_and_its_staged_file_goes() { + let d = temp("dl-arrived"); + let mut s = session(&d); + let staging = d.join("models").join(".downloads"); + std::fs::create_dir_all(&staging).unwrap(); + let staged = staging.join("model.onnx"); + std::fs::write(&staged, lightcraft_faces::synthetic::tiny_embedder_model(512)).unwrap(); + let sha = lightcraft_faces::hash::sha256_file(&staged).unwrap(); + s.face_downloads.arrived("my-model", staged.clone(), &sha); + + let r = s.execute("faces.models.downloads", &json!({})).unwrap(); + let row = &r["downloads"][0]; + assert_eq!((row["id"].as_str(), row["state"].as_str()), (Some("my-model"), Some("installed")), "{r}"); + assert!(!staged.exists(), "the staged file was moved into the model's folder"); + let l = s.execute("faces.models.list", &json!({})).unwrap(); + assert_eq!(l["models"].as_array().unwrap().iter().filter(|m| m["installed"] == true).count(), 1); + // announced once, cleared on request + assert_eq!(s.execute("faces.models.downloadCancel", &json!({"id": "my-model"})).unwrap()["discarded"], true); + assert_eq!(s.execute("faces.models.downloads", &json!({})).unwrap()["downloads"], json!([])); + assert_eq!(s.execute("faces.models.downloadCancel", &json!({"id": "my-model"})).unwrap()["discarded"], false); + + // a file that cannot be installed is thrown away with a reason + std::fs::write(&staged, b"this is not a model").unwrap(); + let sha = lightcraft_faces::hash::sha256_file(&staged).unwrap(); + s.face_downloads.arrived("bad", staged.clone(), &sha); + let r = s.execute("faces.models.downloads", &json!({})).unwrap(); + assert_eq!(r["downloads"][0]["state"], "failed", "{r}"); + assert!(!staged.exists()); + let _ = std::fs::remove_dir_all(&d); +} + +/// `faces.detect` on photos without faces: it reports nothing, and a new run replaces earlier detections but +/// never regions that came from XMP or were named; one undo step restores everything. +#[test] +fn detect_replaces_only_earlier_detections_and_is_one_undo_step() { + use lightcraft_catalog::Op; + use lightcraft_meta::{Region, RegionKind}; + let region = |name: Option<&str>, description: Option<&str>, x: f64| Region { + rect: lightcraft_geom::Rect { x0: x, y0: 0.2, x1: x + 0.2, y1: 0.5 }, + kind: RegionKind::Face, + name: name.map(str::to_string), + description: description.map(str::to_string), + }; + let Some(mut s) = demo_with_yunet() else { return }; + let id = s.active().unwrap(); + let mut meta = s.catalog.photo(id).unwrap().meta.clone(); + meta.regions = + vec![region(Some("Jane Doe"), None, 0.1), region(None, Some("Detected by YuNet 2023mar"), 0.5), region(None, Some("Drawn by hand"), 0.7)]; + s.commit("setup", Op::SetMeta { id, meta: Box::new(meta) }).unwrap(); + let undo_before = s.undo.len(); + + // a dry run reports and changes nothing + let r = s.execute("faces.detect", &json!({"apply": false})).unwrap(); + assert_eq!(r["applied"], false); + assert_eq!(s.catalog.photo(id).unwrap().meta.regions.len(), 3); + assert_eq!(s.undo.len(), undo_before); + + // a real run (the default) drops the earlier detection (the demo has no faces) and keeps the others + let r = s.execute("faces.detect", &json!({})).unwrap(); + assert_eq!(r["applied"], true); + assert_eq!(r["detector"], "YuNet 2023mar"); + let photo = r["photos"].as_array().unwrap().iter().find(|p| p["id"] == id.0).unwrap(); + assert_eq!(photo["faces"].as_array().unwrap().len(), 0, "the procedural demo photos have no faces"); + let names: Vec<_> = s.catalog.photo(id).unwrap().meta.regions.iter().map(|r| (r.name.clone(), r.description.clone())).collect(); + assert_eq!(names, vec![(Some("Jane Doe".to_string()), None), (None, Some("Drawn by hand".to_string()))]); + assert_eq!(s.undo.len(), undo_before + 1, "one undo step for the whole run"); + s.execute("edit.undo", &json!({})).unwrap(); + assert_eq!(s.catalog.photo(id).unwrap().meta.regions.len(), 3, "undo brings the detection back"); + + // odd parameters are tolerated, and nothing is detected without a selection + for p in [json!({"score": 5}), json!({"score": -1}), json!({"nmsIou": 99}), json!({"maxFaces": 0}), json!({"score": "x"})] { + assert!(s.execute("faces.detect", &p).is_ok(), "{p}"); + } + let mut empty = Session::new(); + assert!(empty.execute("faces.detect", &json!({})).is_err()); +} + +/// Opt-in, needs the internet: `cargo test -p lightcraft-engine real_download -- --ignored --nocapture`. Downloads YuNet +/// from its pinned address the way the app does, waits for it to be installed, then finds faces with it. +#[test] +#[ignore = "downloads from github.com"] +fn real_download_of_yunet_installs_it_and_detects() { + let d = temp("real-download"); + let mut s = Session::with_demo(); + s.face_models_dir = Some(d.join("models")); + let started = s.execute("faces.models.download", &json!({"id": "yunet-2023mar", "acknowledged": true})).unwrap(); + assert_eq!(started["from"], "github.com"); + let t0 = std::time::Instant::now(); + loop { + let r = s.execute("faces.models.downloads", &json!({})).unwrap(); + let row = r["downloads"][0].clone(); + match row["state"].as_str() { + Some("installed") => break, + Some("running") | Some("done") => {} + _ => panic!("{r}"), + } + assert!(t0.elapsed() < std::time::Duration::from_secs(120), "{r}"); + std::thread::sleep(std::time::Duration::from_millis(100)); + } + println!("downloaded and installed in {:?}", t0.elapsed()); + let l = s.execute("faces.models.list", &json!({})).unwrap(); + assert_eq!(find(&l, "yunet-2023mar")["installed"], true); + let r = s.execute("faces.detect", &json!({"apply": false})).unwrap(); + assert_eq!(r["detector"], "YuNet 2023mar"); + // asking again once installed is refused + assert!(s.execute("faces.models.download", &json!({"id": "yunet-2023mar", "acknowledged": true})).is_err()); + let _ = std::fs::remove_dir_all(&d); +} + +/// Opt-in, needs the internet: the recognisers are big, so this only checks that bytes start to arrive (through the +/// hosts' redirects) and that cancelling leaves nothing behind. +#[test] +#[ignore = "downloads from github.com and huggingface.co"] +fn real_downloads_of_the_recognisers_start_and_can_be_cancelled() { + let d = temp("real-start"); + let mut s = Session::with_demo(); + s.face_models_dir = Some(d.join("models")); + for id in ["sface-2021dec", "auraface-v1"] { + s.execute("faces.models.download", &json!({"id": id, "acknowledged": true})).unwrap(); + let t0 = std::time::Instant::now(); + loop { + let r = s.execute("faces.models.downloads", &json!({})).unwrap(); + let row = r["downloads"].as_array().unwrap().iter().find(|x| x["id"] == id).unwrap().clone(); + if row["state"] == "running" && row["bytes"].as_u64().unwrap_or(0) > 100_000 { + println!("{id}: {} of {} bytes after {:?}", row["bytes"], row["total"], t0.elapsed()); + break; + } + assert!(row["state"] == "running", "{r}"); + assert!(t0.elapsed() < std::time::Duration::from_secs(60), "{r}"); + std::thread::sleep(std::time::Duration::from_millis(50)); + } + s.execute("faces.models.downloadCancel", &json!({"id": id})).unwrap(); + } + drop(s); + let staging = d.join("models").join(".downloads"); + let left: Vec<_> = std::fs::read_dir(&staging).map(|r| r.flatten().map(|e| e.file_name()).collect()).unwrap_or_default(); + println!("left in staging: {left:?}"); + let _ = std::fs::remove_dir_all(&d); +} + +/// With the recognition runtime: installing runs the model's self-test, a model that does not work is refused +/// and leaves nothing behind, and `faces.models.test` runs it again. +#[cfg(not(target_arch = "wasm32"))] +#[test] +fn installing_runs_the_self_test_and_refuses_models_that_do_not_work() { + let d = temp("selftest"); + let mut s = session(&d); + // a working model passes; its record keeps the result + let f = model_file(&d, "Works.onnx", 64); + let r = s.execute("faces.models.install", &json!({"path": f, "acknowledged": true})).unwrap(); + let id = r["installed"]["id"].as_str().unwrap().to_string(); + assert_eq!(r["installed"]["accepted"]["selfTest"]["ok"], true, "{r}"); + assert_eq!(r["installed"]["accepted"]["selfTest"]["dimension"], 64); + let t = s.execute("faces.models.test", &json!({"id": id})).unwrap(); + assert_eq!(t["ok"], true, "{t}"); + assert!(t["result"]["embedMs"].as_f64().unwrap() >= 0.0); + assert!(s.execute("faces.models.test", &json!({"id": "nope"})).is_err()); + assert!(s.execute("faces.models.test", &json!({})).is_err()); + // a graph with no layers passes the shape check but cannot run: refused, and nothing is left on disk + let broken = d.join("Broken.onnx"); + std::fs::write(&broken, lightcraft_faces::synthetic::embedder_model(64)).unwrap(); + let e = s.execute("faces.models.install", &json!({"path": broken.to_string_lossy(), "acknowledged": true})).unwrap_err().to_string(); + assert!(e.contains("could not be loaded") || e.contains("self-test"), "{e}"); + let leftovers: Vec<_> = std::fs::read_dir(d.join("models")) + .unwrap() + .flatten() + .map(|e| e.file_name().to_string_lossy().into_owned()) + .filter(|n| n != "settings.json") + .collect(); + assert!(leftovers.len() == 1 && leftovers[0].starts_with("custom-works-"), "only the working model remains: {leftovers:?}"); + let _ = std::fs::remove_dir_all(&d); +} + +#[cfg(not(target_arch = "wasm32"))] +#[test] +fn failed_replacement_keeps_the_installed_model_and_records() { + let d = temp("replacement"); + let models = d.join("models"); + std::fs::create_dir_all(&models).unwrap(); + let good = model_file(&d, "Good.onnx", 64); + let write_catalog = |path: &std::path::Path| { + let catalog = json!({"models": [{ + "id": "replacement", "name": "Replacement", "version": "1", "role": "embedder", + "url": "https://models.example.org/weights.onnx", + "sha256": lightcraft_faces::hash::sha256_file(path).unwrap(), + "sizeBytes": std::fs::metadata(path).unwrap().len(), + "licence": {"name": "MIT", "commercial": "yes"}, + "output": {"kind": "embedding", "dim": 64} + }]}); + std::fs::write(models.join("catalog.json"), catalog.to_string()).unwrap(); + }; + write_catalog(std::path::Path::new(&good)); + let mut s = session(&d); + s.execute("faces.models.install", &json!({"path": good, "acknowledged": true})).unwrap(); + let home = models.join("replacement"); + let files = ["model.onnx", "face-model.json", "installed.json"]; + let before: Vec<_> = files.iter().map(|name| std::fs::read(home.join(name)).unwrap()).collect(); + let settings = std::fs::read(models.join("settings.json")).unwrap(); + let bad = d.join("Bad.onnx"); + std::fs::write(&bad, lightcraft_faces::synthetic::embedder_model(64)).unwrap(); + write_catalog(&bad); + let mut s = session(&d); + assert!(s.execute("faces.models.install", &json!({"path": bad, "acknowledged": true})).is_err()); + for (name, bytes) in files.iter().zip(before) { + assert_eq!(std::fs::read(home.join(name)).unwrap(), bytes, "{name} was replaced"); + } + assert_eq!(std::fs::read(models.join("settings.json")).unwrap(), settings); + assert!(!home.join("model.onnx.part").exists()); + let _ = std::fs::remove_dir_all(d); +} + +fn settings_of(d: &std::path::Path) -> Value { + serde_json::from_slice(&std::fs::read(d.join("models").join("settings.json")).unwrap()).unwrap() +} + +#[test] +fn installing_a_model_chooses_it_and_switches_recognition_on() { + let d = temp("activate"); + let mut s = session(&d); + let a = s.execute("faces.models.install", &json!({"path": model_file(&d, "First.onnx", 64), "acknowledged": true})).unwrap(); + let a_id = a["installed"]["id"].as_str().unwrap().to_string(); + assert_eq!(a["active"], json!({"embedder": a_id, "enabled": true})); + assert_eq!(a["installed"]["selected"], true); + assert_eq!(settings_of(&d), json!({"enabled": true, "embedder": a_id})); + // a second model takes over (the newest is the one the user just asked for); the first stays installed + let b = s.execute("faces.models.install", &json!({"path": model_file(&d, "Second.onnx", 96), "acknowledged": true})).unwrap(); + let b_id = b["installed"]["id"].as_str().unwrap().to_string(); + assert_ne!(a_id, b_id); + assert_eq!(b["active"]["embedder"], json!(b_id)); + let l = s.execute("faces.models.list", &json!({})).unwrap(); + assert_eq!((find(&l, &a_id)["installed"].clone(), find(&l, &a_id)["selected"].clone()), (json!(true), json!(false))); + assert_eq!(find(&l, &b_id)["selected"], true); + // switching back is a choice, and a model added with `activate: false` leaves the choice alone + s.execute("faces.models.select", &json!({"id": a_id})).unwrap(); + let c = s.execute("faces.models.install", &json!({"path": model_file(&d, "Third.onnx", 128), "acknowledged": true, "activate": false})).unwrap(); + let c_id = c["installed"]["id"].as_str().unwrap().to_string(); + assert_eq!(c["active"]["embedder"], json!(a_id)); + assert_eq!(c["installed"]["selected"], false); + // a switched-off recognition is switched back on by the next install, not by a refused one + s.execute("faces.enable", &json!({"enabled": false})).unwrap(); + assert!(s.execute("faces.models.install", &json!({"path": model_file(&d, "Nope.onnx", 64)})).is_err()); + assert_eq!(settings_of(&d)["enabled"], false); + s.execute("faces.models.install", &json!({"path": model_file(&d, "Fourth.onnx", 64), "acknowledged": true})).unwrap(); + assert_eq!(settings_of(&d)["enabled"], true); + // removing the model in use hands over to another installed one instead of leaving nothing chosen + let in_use = settings_of(&d)["embedder"].as_str().unwrap().to_string(); + let r = s.execute("faces.models.remove", &json!({"id": in_use})).unwrap(); + let next = r["embedder"].as_str().unwrap().to_string(); + assert_ne!(next, in_use); + assert!([&a_id, &b_id, &c_id].contains(&&next) || next.starts_with("custom-"), "{next}"); + assert_eq!(settings_of(&d)["embedder"], json!(next)); + // removing a model that is not in use changes nothing + let other = [&a_id, &b_id, &c_id].into_iter().find(|i| **i != next && **i != in_use).unwrap().clone(); + s.execute("faces.models.remove", &json!({"id": other})).unwrap(); + assert_eq!(settings_of(&d)["embedder"], json!(next)); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn a_download_that_arrives_is_installed_and_switched_on_by_itself() { + use crate::face_download::{STAGING, State}; + let d = temp("finish"); + let mut s = session(&d); + let staging = d.join("models").join(STAGING); + std::fs::create_dir_all(&staging).unwrap(); + // what the download thread leaves when it succeeds: a verified file in the staging folder and a finished job + let staged = model_file(&staging, "arrived.onnx", 64); + let sha = lightcraft_faces::hash::sha256_file(std::path::Path::new(&staged)).unwrap(); + s.face_downloads.set_outcome("fake-model", State::Done { path: staged.clone().into(), sha256: sha }); + let r = s.execute("faces.models.downloads", &json!({})).unwrap(); + assert_eq!(r["downloads"][0]["id"], "fake-model"); + assert_eq!(r["downloads"][0]["state"], "installed", "{r}"); + let st = settings_of(&d); + assert_eq!(st["enabled"], true); + assert!(st["embedder"].as_str().unwrap().starts_with("custom-arrived-"), "{st}"); + assert!(!std::path::Path::new(&staged).exists(), "the staged file was moved into place"); + // the record stays until it is cleared, and is cleared once + assert_eq!(s.execute("faces.models.downloads", &json!({})).unwrap()["downloads"][0]["state"], "installed"); + assert_eq!(s.execute("faces.models.downloadCancel", &json!({"id": "fake-model"})).unwrap()["discarded"], true); + assert_eq!(s.execute("faces.models.downloads", &json!({})).unwrap()["downloads"], json!([])); + + // a download that turns out not to be a usable model is thrown away, and the reason is shown + let junk = staging.join("junk.onnx"); + std::fs::write(&junk, b"definitely not a model").unwrap(); + let sha = lightcraft_faces::hash::sha256_file(&junk).unwrap(); + s.face_downloads.set_outcome("fake-junk", State::Done { path: junk.clone(), sha256: sha }); + let r = s.execute("faces.models.downloads", &json!({})).unwrap(); + assert_eq!(r["downloads"][0]["state"], "failed", "{r}"); + assert!(r["downloads"][0]["error"].as_str().unwrap().contains("cannot be used"), "{r}"); + assert!(!junk.exists()); + assert_eq!(settings_of(&d)["embedder"], st["embedder"], "a failed download leaves the choice alone"); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn the_users_own_catalog_adds_models_with_the_same_download_and_install() { + use crate::face_download::State; + let d = temp("catalog"); + let mut s = session(&d); + let models = d.join("models"); + std::fs::create_dir_all(&models).unwrap(); + // a model file, and a catalog that lists it by its hash (plus an entry that cannot be used) + let file = model_file(&d, "weights.onnx", 64); + let (sha, size) = (lightcraft_faces::hash::sha256_file(std::path::Path::new(&file)).unwrap(), std::fs::metadata(&file).unwrap().len()); + let catalog = json!({"models": [ + {"id": "my-research-model", "name": "My research model", "version": "1", "role": "embedder", "url": "https://models.example.org/w/weights.onnx", + "sha256": sha, "sizeBytes": size, "licence": {"name": "Research only", "commercial": "no", "notice": "Not for commercial use."}, + "provenance": "Some dataset", "output": {"kind": "embedding", "dim": 64}, "thresholds": {"matchCosine": 0.42}}, + {"id": "broken", "name": "No hash", "version": "1", "role": "embedder", "url": "https://models.example.org/x.onnx", "output": {"kind": "embedding", "dim": 64}}, + ]}); + std::fs::write(models.join("catalog.json"), catalog.to_string()).unwrap(); + // it is listed like a built-in model, with its own terms, a download address, and the problem with the other entry + let l = s.execute("faces.models.list", &json!({})).unwrap(); + let m = find(&l, "my-research-model"); + assert_eq!( + (m["fromCatalog"].clone(), m["downloadHost"].clone(), m["licence"]["commercial"].clone()), + (json!(true), json!("models.example.org"), json!("no")) + ); + assert_eq!(m["installed"], false); + assert_eq!(find(&l, "broken"), &Value::Null); + assert_eq!(l["catalog"]["models"], 1); + assert!(l["catalog"]["errors"][0].as_str().unwrap().contains("broken"), "{}", l["catalog"]); + // nothing is fetched without the terms being accepted + assert!(s.execute("faces.models.download", &json!({"id": "my-research-model"})).is_err()); + // a file that matches is recognised as that model (not as a draft), so its terms and settings are used + let seen = s.execute("faces.models.inspect", &json!({"path": file})).unwrap(); + assert_eq!((seen["kind"].clone(), seen["model"]["id"].clone()), (json!("known"), json!("my-research-model"))); + // a finished download of it installs under the catalog's id, with the catalog's licence record + let staging = models.join(crate::face_download::STAGING); + std::fs::create_dir_all(&staging).unwrap(); + let staged = staging.join("weights.onnx"); + std::fs::copy(&file, &staged).unwrap(); + s.face_downloads.set_outcome("my-research-model", State::Done { path: staged, sha256: sha }); + let r = s.execute("faces.models.downloads", &json!({})).unwrap(); + assert_eq!(r["downloads"][0]["state"], "installed", "{r}"); + let st = settings_of(&d); + assert_eq!((st["embedder"].clone(), st["enabled"].clone()), (json!("my-research-model"), json!(true))); + let l = s.execute("faces.models.list", &json!({})).unwrap(); + assert_eq!( + (find(&l, "my-research-model")["installed"].clone(), find(&l, "my-research-model")["accepted"]["commercial"].clone()), + (json!(true), json!("no")) + ); + // a broken or hostile catalog never stops the list + for bytes in [&b"{"[..], b"[]", &vec![b'x'; 400_000]] { + std::fs::write(models.join("catalog.json"), bytes).unwrap(); + let l = s.execute("faces.models.list", &json!({})).unwrap(); + assert!(l["catalog"]["errors"].as_array().is_some_and(|e| !e.is_empty()), "{}", l["catalog"]); + assert!(find(&l, "sface-2021dec")["id"] == "sface-2021dec"); + } + let _ = std::fs::remove_dir_all(&d); +} diff --git a/crates/engine/src/tests_face_recognize.rs b/crates/engine/src/tests_face_recognize.rs new file mode 100644 index 00000000..2fd53c67 --- /dev/null +++ b/crates/engine/src/tests_face_recognize.rs @@ -0,0 +1,489 @@ +//! Recognition: indexing faces, suggesting names, naming, with a tiny working network standing in for a real +//! model (its embedding depends on the picture's colours, which is enough to test the machinery; how well real +//! models recognise people is measured separately, on real photos). + +use std::path::PathBuf; + +use lightcraft_catalog::{Op, PhotoId}; +use lightcraft_geom::Rect; +use lightcraft_meta::{Region, RegionKind}; +use serde_json::{Value, json}; + +use crate::Session; + +pub(crate) fn temp(name: &str) -> PathBuf { + let d = std::env::temp_dir().join(format!("lc-facerec-{name}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&d); + std::fs::create_dir_all(&d).unwrap(); + d +} + +/// A demo session with a tiny model installed and chosen. +pub(crate) fn setup(d: &std::path::Path, dim: u64) -> (Session, String) { + let mut s = Session::with_demo(); + s.face_models_dir = Some(d.join("models")); + let model = d.join(format!("tiny{dim}.onnx")); + std::fs::write(&model, lightcraft_faces::synthetic::tiny_embedder_model(dim)).unwrap(); + let r = s.execute("faces.models.install", &json!({"path": model.to_string_lossy(), "acknowledged": true})).unwrap(); + let id = r["installed"]["id"].as_str().unwrap().to_string(); + s.execute("faces.models.select", &json!({"id": id})).unwrap(); + (s, id) +} + +/// Installs the real YuNet from `LC_YUNET_MODEL=` (what Settings > Faces downloads; nothing is +/// committed) and says whether it did: the parts of a test that need to find faces run only then. +pub(crate) fn install_yunet(s: &mut Session) -> bool { + let Some(file) = std::env::var_os("LC_YUNET_MODEL") else { + eprintln!("LC_YUNET_MODEL is not set: skipping the face-finding part"); + return false; + }; + let r = s.execute("faces.models.install", &json!({"path": file.to_string_lossy(), "acknowledged": true})); + assert!(r.is_ok(), "{r:?}"); + true +} + +pub(crate) fn region(x: f64, name: Option<&str>) -> Region { + Region { rect: Rect { x0: x, y0: 0.2, x1: x + 0.25, y1: 0.6 }, kind: RegionKind::Face, name: name.map(str::to_string), description: None } +} + +pub(crate) fn set_regions(s: &mut Session, id: PhotoId, regions: Vec) { + let mut meta = s.catalog.photo(id).unwrap().meta.clone(); + meta.regions = regions; + s.commit("setup", Op::SetMeta { id, meta: Box::new(meta) }).unwrap(); +} + +pub(crate) fn two_photos(s: &Session) -> (PhotoId, PhotoId) { + let mut ids = s.catalog.photos().map(|p| p.id).collect::>(); + ids.sort(); + (ids[0], ids[1]) +} + +#[test] +fn indexing_embeds_each_face_once_and_reports_progress() { + let d = temp("index"); + let (mut s, _) = setup(&d, 64); + let (a, b) = two_photos(&s); + set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, None)]); + set_regions(&mut s, b, vec![region(0.3, None)]); + // only a report + let r = s.execute("faces.index", &json!({"budgetMs": 0})).unwrap(); + assert_eq!((r["embedded"].clone(), r["photosDone"].clone(), r["pendingPhotos"].clone()), (json!(0), json!(0), json!(2)), "{r}"); + // a generous budget does it all, the named photo first + let r = s.execute("faces.index", &json!({"budgetMs": 60_000})).unwrap(); + assert_eq!((r["embedded"].clone(), r["pendingPhotos"].clone(), r["indexedFaces"].clone()), (json!(3), json!(0), json!(3)), "{r}"); + // nothing is embedded twice + let r = s.execute("faces.index", &json!({"budgetMs": 60_000})).unwrap(); + assert_eq!((r["embedded"].clone(), r["photosDone"].clone()), (json!(0), json!(0))); + // a small budget still makes progress (at least one photo per call), and `ids` go first + let c = s.catalog.photos().map(|p| p.id).find(|id| *id != a && *id != b).unwrap(); + set_regions(&mut s, c, vec![region(0.2, None)]); + set_regions(&mut s, b, vec![region(0.3, None), region(0.6, None)]); + let r = s.execute("faces.index", &json!({"budgetMs": 1, "ids": [c.0]})).unwrap(); + assert_eq!(r["photosDone"], 1, "{r}"); + assert_eq!(r["pendingPhotos"], 1, "{r}"); + assert!(s.faces.index.contains(c.0, ®ion(0.2, None).rect), "the photo asked for first was done first"); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn suggestions_come_from_the_named_faces_and_are_never_applied() { + let d = temp("suggest"); + let (mut s, _) = setup(&d, 64); + let (a, b) = two_photos(&s); + set_regions(&mut s, a, vec![region(0.1, Some("Ann"))]); + set_regions(&mut s, b, vec![region(0.1, None), region(0.5, Some("Pet?")), region(0.7, None)]); + // thresholds that accept anything: Ann is the only person with a named face except "Pet?" + let r = s.execute("faces.suggest", &json!({"ids": [b.0], "threshold": -1.0, "margin": -2.0, "budgetMs": 60_000})).unwrap(); + assert_eq!(r["galleryFaces"], 2, "{r}"); + let faces = r["photos"][0]["faces"].as_array().unwrap(); + assert_eq!(faces.len(), 2, "only the unnamed faces get suggestions: {r}"); + assert!(faces.iter().all(|f| f["suggestion"]["name"].is_string()), "{r}"); + assert_eq!(faces[0]["index"], 0); + assert_eq!(faces[1]["index"], 2); + assert!(faces[0]["candidates"].as_array().unwrap().len() == 2); + // an impossible threshold suggests nobody but still lists the candidates + let r = s.execute("faces.suggest", &json!({"ids": [b.0], "threshold": 2.0, "budgetMs": 60_000})).unwrap(); + assert!(r["photos"][0]["faces"][0]["suggestion"].is_null()); + assert!(!r["photos"][0]["faces"][0]["candidates"].as_array().unwrap().is_empty()); + // nothing was named by suggesting + assert!(s.catalog.photo(b).unwrap().meta.regions.iter().filter(|r| r.name.is_none()).count() == 2); + // pets and other kinds are not people + let mut meta = s.catalog.photo(b).unwrap().meta.clone(); + meta.regions[1].kind = RegionKind::Pet; + s.commit("pet", Op::SetMeta { id: b, meta: Box::new(meta) }).unwrap(); + let r = s.execute("faces.suggest", &json!({"ids": [b.0], "threshold": -1.0, "margin": -2.0, "budgetMs": 60_000})).unwrap(); + assert_eq!(r["galleryFaces"], 1, "{r}"); + // odd parameters are tolerated + for p in [json!({"threshold": "x"}), json!({"margin": null}), json!({"budgetMs": 0}), json!({"ids": []}), json!({"ids": [999999]})] { + assert!(s.execute("faces.suggest", &p).is_ok(), "{p}"); + } + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn naming_a_face_is_undoable_and_makes_it_yours() { + let d = temp("name"); + let (mut s, _) = setup(&d, 64); + let (a, _) = two_photos(&s); + let mut detected = region(0.1, None); + detected.description = Some("Detected by YuNet 2023mar".into()); + set_regions(&mut s, a, vec![detected, region(0.5, None)]); + let undo = s.undo.len(); + let r = s.execute("faces.setName", &json!({"id": a.0, "index": 0, "name": " Ann "})).unwrap(); + assert_eq!(r["name"], "Ann"); + let regions = s.catalog.photo(a).unwrap().meta.regions.clone(); + assert_eq!((regions[0].name.as_deref(), regions[0].description.as_deref()), (Some("Ann"), None), "a named detection loses its detected marker"); + assert_eq!(s.undo.len(), undo + 1); + // a new detection run now leaves it alone + if install_yunet(&mut s) { + s.execute("faces.detect", &json!({"id": a.0})).unwrap(); + assert_eq!(s.catalog.photo(a).unwrap().meta.regions.iter().filter(|r| r.name.as_deref() == Some("Ann")).count(), 1); + } + // clearing, undo + s.execute("faces.setName", &json!({"id": a.0, "index": 0, "name": null})).unwrap(); + assert_eq!(s.catalog.photo(a).unwrap().meta.regions[0].name, None); + s.execute("edit.undo", &json!({})).unwrap(); + assert_eq!(s.catalog.photo(a).unwrap().meta.regions[0].name.as_deref(), Some("Ann")); + // bad input + for p in [ + json!({"id": a.0}), + json!({"id": a.0, "index": 9, "name": "x"}), + json!({"id": a.0, "index": 0, "name": 5}), + json!({"id": a.0, "index": 0, "name": "a\u{7}b"}), + json!({"id": a.0, "index": 0, "name": "x".repeat(201)}), + json!({"id": 999999, "index": 0, "name": "x"}), + ] { + assert!(s.execute("faces.setName", &p).is_err(), "{p}"); + } + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn embeddings_are_cached_beside_the_library_and_reset_for_another_model() { + let d = temp("persist"); + let lib = d.join("lib"); + let models = d.join("models"); + let install = |s: &mut Session, dim: u64| -> String { + let model = d.join(format!("tiny{dim}.onnx")); + std::fs::write(&model, lightcraft_faces::synthetic::tiny_embedder_model(dim)).unwrap(); + let r = s.execute("faces.models.install", &json!({"path": model.to_string_lossy(), "acknowledged": true})).unwrap(); + let id = r["installed"]["id"].as_str().unwrap().to_string(); + s.execute("faces.models.select", &json!({"id": id})).unwrap(); + id + }; + let mut s = Session::new().with_fs(); + s.face_models_dir = Some(models.clone()); + s.open_library(&lib, true).unwrap(); + let (a, _) = two_photos(&s); + set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, None)]); + install(&mut s, 64); + s.execute("faces.index", &json!({"budgetMs": 60_000})).unwrap(); + assert!(lib.join(format!("face-embeddings-{}.bin", s.faces.index.tag.split('@').next().unwrap())).is_file()); + s.close_library().unwrap(); + drop(s); + + // a new session finds the embeddings on disk and has nothing left to do + let mut s2 = Session::new().with_fs(); + s2.face_models_dir = Some(models); + s2.open_library(&lib, true).unwrap(); + let r = s2.execute("faces.index", &json!({"budgetMs": 60_000})).unwrap(); + assert_eq!((r["embedded"].clone(), r["indexedFaces"].clone(), r["pendingPhotos"].clone()), (json!(0), json!(2), json!(0)), "{r}"); + // choosing another model starts over: its faces are different numbers, the old ones are not used + install(&mut s2, 32); + let r = s2.execute("faces.index", &json!({"budgetMs": 0})).unwrap(); + assert_eq!((r["indexedFaces"].clone(), r["pendingPhotos"].clone()), (json!(0), json!(1)), "{r}"); + let r = s2.execute("faces.index", &json!({"budgetMs": 60_000})).unwrap(); + assert_eq!(r["embedded"], 2); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn without_a_chosen_model_recognition_says_so() { + let d = temp("nomodel"); + let mut s = Session::with_demo(); + assert!(s.execute("faces.index", &json!({})).is_err(), "no models folder"); + s.face_models_dir = Some(d.join("models")); + for cmd in ["faces.index", "faces.suggest"] { + let e = s.execute(cmd, &json!({})).unwrap_err().to_string(); + assert!(e.contains("no recognition model is chosen"), "{cmd}: {e}"); + } + let _: Value = json!(null); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn the_background_pump_indexes_by_itself_once_recognition_is_on() { + let d = temp("pump"); + let (mut s, _) = setup(&d, 64); + let (a, b) = two_photos(&s); + set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, None)]); + set_regions(&mut s, b, vec![region(0.3, None)]); + // installing the model switched recognition on; switched off, the pump does nothing, and says so + assert_eq!(s.execute("faces.enable", &json!({"enabled": false})).unwrap()["enabled"], false); + assert_eq!(s.execute("faces.pump", &json!({})).unwrap()["active"], false); + assert_eq!(s.faces.index.len(), 0); + s.execute("faces.enable", &json!({"enabled": true})).unwrap(); + // `enabled` is looked up at most once a second, so the first pump may still see "off" + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(60); + let mut last = json!(null); + while std::time::Instant::now() < deadline { + last = s.execute("faces.pump", &json!({})).unwrap(); + if last["active"] == true && last["pendingPhotos"] == 0 && last["indexedFaces"] == 3 { + break; + } + std::thread::sleep(std::time::Duration::from_millis(50)); + } + assert_eq!((last["active"].clone(), last["pendingPhotos"].clone(), last["indexedFaces"].clone()), (json!(true), json!(0), json!(3)), "{last}"); + // a new face is picked up (the catalog changed), and suggestions use what the pump embedded + set_regions(&mut s, b, vec![region(0.3, None), region(0.6, None)]); + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(60); + while std::time::Instant::now() < deadline && s.faces.index.len() < 4 { + s.execute("faces.pump", &json!({})).unwrap(); + std::thread::sleep(std::time::Duration::from_millis(50)); + } + assert_eq!(s.faces.index.len(), 4); + let r = s.execute("faces.suggest", &json!({"ids": [b.0], "threshold": -1.0, "margin": -2.0, "budgetMs": 0})).unwrap(); + assert_eq!(r["photos"][0]["faces"].as_array().unwrap().len(), 2); + assert!(r["photos"][0]["faces"][0]["suggestion"]["name"].is_string(), "{r}"); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn the_background_pump_saves_what_it_embeds_for_the_next_launch() { + let d = temp("pumpsave"); + let (lib, models) = (d.join("lib"), d.join("models")); + let mut s = Session::new().with_fs(); + s.face_models_dir = Some(models); + s.open_library(&lib, true).unwrap(); + let model = d.join("tiny64.onnx"); + std::fs::write(&model, lightcraft_faces::synthetic::tiny_embedder_model(64)).unwrap(); + // installing switches recognition on: nothing else is asked for + s.execute("faces.models.install", &json!({"path": model.to_string_lossy(), "acknowledged": true})).unwrap(); + // without the detector (not bundled) photos without faces are left unsearched + let searching = install_yunet(&mut s); + let (a, b) = two_photos(&s); + set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, None)]); + set_regions(&mut s, b, vec![region(0.3, None)]); + let photos = s.catalog.photos().count(); + let searched_photos = if searching { photos - 2 } else { 0 }; + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(120); + let mut last = json!(null); + while std::time::Instant::now() < deadline { + last = s.execute("faces.pump", &json!({})).unwrap(); + if last["pendingPhotos"] == 0 && s.faces.index.len() == 3 && s.faces.scanned.len() == searched_photos { + break; + } + std::thread::sleep(std::time::Duration::from_millis(50)); + } + assert_eq!((s.faces.index.len(), s.faces.scanned.len()), (3, searched_photos), "{last}"); + let (tag, path) = (s.faces.index.tag.clone(), s.faces.index.path.clone().unwrap()); + assert!(path.file_name().unwrap().to_string_lossy().starts_with("face-embeddings-custom-"), "one cache file per model: {}", path.display()); + // the first batch is written at once; whatever is left is written when the session ends, so a launch never redoes + // work: neither the embeddings nor the list of photos already searched for faces + s.close_library().unwrap(); + drop(s); + let mut again = crate::faces_index::Index::default(); + again.reset(&tag, 64, Some(path)); + assert_eq!(again.len(), 3); + let mut searched = crate::faces_index::Scanned::default(); + searched.reset(Some(lib.join("face-scanned.bin"))); + assert_eq!(searched.len(), searched_photos, "every photo without faces is remembered as searched"); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn faces_person_lists_a_persons_faces_and_looks_for_more_only_while_recognition_runs() { + let d = temp("person"); + let (mut s, _) = setup(&d, 64); + s.execute("faces.index", &json!({"budgetMs": 0})).unwrap(); + let (a, b) = two_photos(&s); + set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, Some("Bob"))]); + set_regions(&mut s, b, vec![region(0.1, None), region(0.5, Some("ann"))]); + let unit = |parts: &[(usize, f32)]| { + let mut v = vec![0.0f32; 64]; + for (i, x) in parts { + v[*i] = *x; + } + let n = v.iter().map(|x| x * x).sum::().sqrt(); + v.iter().map(|x| x / n).collect::>() + }; + for (id, x, v) in [(a, 0.1, unit(&[(0, 1.0)])), (a, 0.5, unit(&[(1, 1.0)])), (b, 0.1, unit(&[(0, 0.9), (1, 0.1)])), (b, 0.5, unit(&[(0, 1.0)]))] { + s.faces.index.insert(id.0, ®ion(x, None).rect, v); + } + // the name is matched without regard to case or surrounding spaces; every face is listed, wherever it is + let r = s.execute("faces.person", &json!({"name": " ANN "})).unwrap(); + assert_eq!((r["name"].clone(), r["total"].clone()), (json!("Ann"), json!(2))); + let mut at: Vec<(u64, u64)> = + r["confirmed"].as_array().unwrap().iter().map(|f| (f["photo"].as_u64().unwrap(), f["index"].as_u64().unwrap())).collect(); + at.sort(); + let mut want = vec![(a.0, 0), (b.0, 1)]; + want.sort(); + assert_eq!(at, want); + assert!(r["confirmed"].as_array().unwrap().iter().any(|f| f["rect"]["x0"] == 0.1 && f["rect"]["y1"] == 0.6)); + // recognition is on (installing the model switched it on): the unnamed face that looks like Ann is offered + assert_eq!(r["ready"], true); + let more = r["more"].as_array().unwrap(); + assert_eq!((more.len(), more[0]["photo"].clone(), more[0]["index"].clone()), (1, json!(b.0), json!(0)), "{r}"); + assert!(more[0]["score"].as_f64().unwrap() > 0.9); + // confirming it moves it from one list to the other (and is one undo step) + s.execute("faces.setName", &json!({"id": b.0, "index": 0, "name": "Ann"})).unwrap(); + let r = s.execute("faces.person", &json!({"name": "Ann"})).unwrap(); + assert_eq!((r["total"].clone(), r["more"].as_array().unwrap().len()), (json!(3), 0)); + // switched off, the faces are still listed but nothing is searched for + s.execute("faces.enable", &json!({"enabled": false})).unwrap(); + let r = s.execute("faces.person", &json!({"name": "Ann"})).unwrap(); + assert_eq!((r["total"].clone(), r["ready"].clone(), r["more"].as_array().unwrap().len()), (json!(3), json!(false), 0)); + // nobody by that name, and bad input, are plain answers and errors + let r = s.execute("faces.person", &json!({"name": "Nobody"})).unwrap(); + assert_eq!((r["total"].clone(), r["confirmed"].as_array().unwrap().len(), r["ready"].clone()), (json!(0), 0, json!(false))); + for bad in [json!({}), json!({"name": ""}), json!({"name": " "}), json!({"name": 7})] { + assert!(s.execute("faces.person", &bad).is_err(), "{bad}"); + } + assert!(s.execute("faces.person", &json!({"name": "Ann", "more": u64::MAX})).is_ok()); + let _ = std::fs::remove_dir_all(&d); +} + +fn unit64(parts: &[(usize, f32)]) -> Vec { + let mut v = vec![0.0f32; 64]; + for (i, x) in parts { + v[*i] = *x; + } + let n = v.iter().map(|x| x * x).sum::().sqrt(); + v.iter().map(|x| x / n).collect() +} + +#[test] +fn many_faces_are_named_in_one_undo_step() { + let d = temp("namemany"); + let (mut s, _) = setup(&d, 64); + let (a, b) = two_photos(&s); + let detected = |x: f64| Region { description: Some("Detected by YuNet 2023mar".into()), ..region(x, None) }; + set_regions(&mut s, a, vec![detected(0.1), region(0.5, None)]); + set_regions(&mut s, b, vec![detected(0.3)]); + let undo = s.undo.len(); + let r = s + .execute( + "faces.nameFaces", + &json!({"faces": [{"photo": a.0, "index": 0}, {"photo": a.0, "index": 1}, {"photo": b.0, "index": 0}], "name": " Jane Doe "}), + ) + .unwrap(); + assert_eq!((r["named"].clone(), r["photos"].clone()), (json!(3), json!(2))); + assert_eq!(s.undo.len(), undo + 1, "one step for the whole group"); + for (id, i) in [(a, 0), (a, 1), (b, 0)] { + let reg = &s.catalog.photo(id).unwrap().meta.regions[i]; + assert_eq!(reg.name.as_deref(), Some("Jane Doe")); + assert!(reg.description.is_none() || !reg.description.as_deref().unwrap().starts_with("Detected by"), "a named face is the user's now"); + } + // undo restores every one of them together + s.execute("edit.undo", &json!({})).unwrap(); + assert!(s.catalog.photo(a).unwrap().meta.regions.iter().chain(s.catalog.photo(b).unwrap().meta.regions.iter()).all(|r| r.name.is_none())); + // an empty name clears; faces that do not exist are skipped, and with none left it is an error + s.execute("faces.nameFaces", &json!({"faces": [{"photo": a.0, "index": 0}], "name": "Ann"})).unwrap(); + let r = s + .execute( + "faces.nameFaces", + &json!({"faces": [{"photo": a.0, "index": 0}, {"photo": a.0, "index": 99}, {"photo": 987_654, "index": 0}], "name": null}), + ) + .unwrap(); + assert_eq!(r["named"], 1); + assert!(s.catalog.photo(a).unwrap().meta.regions[0].name.is_none()); + for bad in [ + json!({"name": "x"}), + json!({"faces": "x", "name": "x"}), + json!({"faces": [{"photo": a.0}], "name": "x"}), + json!({"faces": [{"photo": 987_654, "index": 0}], "name": "x"}), + json!({"faces": [{"photo": a.0, "index": 0}], "name": 7}), + json!({"faces": [{"photo": a.0, "index": 0}], "name": "bad\u{7}name"}), + json!({"faces": vec![json!({"photo": a.0, "index": 0}); 5001], "name": "x"}), + ] { + assert!(s.execute("faces.nameFaces", &bad).is_err(), "{bad}"); + } + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn unnamed_faces_come_with_look_alikes_together_and_the_name_the_named_ones_suggest() { + let d = temp("unnamed"); + let (mut s, _) = setup(&d, 64); + s.execute("faces.index", &json!({"budgetMs": 0})).unwrap(); + let (a, b) = two_photos(&s); + let c = s.catalog.photos().map(|p| p.id).find(|id| *id != a && *id != b).unwrap(); + set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, None), region(0.7, None)]); + set_regions(&mut s, b, vec![region(0.1, None), region(0.5, None)]); + set_regions(&mut s, c, vec![region(0.3, None)]); + // three faces like Ann, two like Bob (unnamed), one unembedded + let put = |s: &mut Session, id: lightcraft_catalog::PhotoId, x: f64, v: Vec| s.faces.index.insert(id.0, ®ion(x, None).rect, v); + put(&mut s, a, 0.1, unit64(&[(0, 1.0)])); // Ann, named + put(&mut s, a, 0.5, unit64(&[(1, 1.0)])); // like Bob + put(&mut s, b, 0.1, unit64(&[(0, 0.97), (2, 0.2)])); // like Ann + put(&mut s, b, 0.5, unit64(&[(1, 0.98), (2, 0.2)])); // like Bob + put(&mut s, c, 0.3, unit64(&[(0, 0.95), (3, 0.3)])); // like Ann + // (a's last face, at 0.7, is not embedded yet) + let r = s.execute("faces.unnamed", &json!({})).unwrap(); + assert_eq!((r["total"].clone(), r["ordered"].clone(), r["ready"].clone()), (json!(5), json!(true), json!(true)), "{r}"); + let faces = r["faces"].as_array().unwrap(); + let at = |f: &Value| (f["photo"].as_u64().unwrap(), f["index"].as_u64().unwrap()); + let order: Vec<(u64, u64)> = faces.iter().map(at).collect(); + // look-alikes are neighbours: the two like Bob are adjacent, and the three like Ann are adjacent; the unembedded one is last + let pos = |x: (u64, u64)| order.iter().position(|o| *o == x).unwrap(); + assert_eq!(pos((a.0, 1)).abs_diff(pos((b.0, 1))), 1, "{order:?}"); + assert_eq!(order.last(), Some(&(a.0, 2)), "{order:?}"); + // a face the named ones recognise carries the name they suggest + let like_ann = faces.iter().find(|f| at(f) == (b.0, 0)).unwrap(); + assert_eq!(like_ann["suggestion"]["name"], "Ann", "{like_ann}"); + assert!(faces.iter().find(|f| at(f) == (b.0, 1)).unwrap()["suggestion"].is_null(), "Bob is nobody the named faces know"); + // each face says which box to show it by; with no scan view it is its own box + for k in ["x0", "y0", "x1", "y1"] { + assert!((like_ann["view"][k].as_f64().unwrap() - like_ann["rect"][k].as_f64().unwrap()).abs() < 1e-3, "{like_ann}"); + } + // a limit cuts the list, not the count + let r = s.execute("faces.unnamed", &json!({"limit": 2})).unwrap(); + assert_eq!((r["total"].clone(), r["faces"].as_array().unwrap().len()), (json!(5), 2)); + // switched off, the faces are still listed, in photo order, without suggestions + s.execute("faces.enable", &json!({"enabled": false})).unwrap(); + let r = s.execute("faces.unnamed", &json!({})).unwrap(); + assert_eq!((r["total"].clone(), r["ordered"].clone(), r["ready"].clone()), (json!(5), json!(false), json!(false))); + assert!(r["faces"].as_array().unwrap().iter().all(|f| f["suggestion"].is_null())); + for odd in [json!({"limit": 0}), json!({"limit": "x"}), json!({"limit": u64::MAX})] { + assert!(s.execute("faces.unnamed", &odd).is_ok(), "{odd}"); + } + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn every_scanned_face_is_shown_by_the_same_kind_of_box() { + let d = temp("view"); + let (mut s, _) = setup(&d, 64); + s.execute("faces.index", &json!({"budgetMs": 0})).unwrap(); + let (a, _) = two_photos(&s); + // a loosely drawn box (from some other tool) and the detector's tighter box of the same face + let loose = lightcraft_geom::Rect { x0: 0.2, y0: 0.1, x1: 0.8, y1: 0.9 }; + let tight = lightcraft_geom::Rect { x0: 0.35, y0: 0.25, x1: 0.65, y1: 0.65 }; + let mut meta = s.catalog.photo(a).unwrap().meta.clone(); + meta.regions = vec![Region { rect: loose, kind: RegionKind::Face, name: Some("Ann".into()), description: None }]; + s.commit("setup", Op::SetMeta { id: a, meta: Box::new(meta) }).unwrap(); + // before the scan has looked at it, the face is shown by its own box + assert_eq!(s.face_view(a, loose), loose); + s.faces.index.insert_with_view(a.0, &loose, unit64(&[(0, 1.0)]), &tight); + let shown = s.face_view(a, loose); + assert!((shown.x0 - tight.x0).abs() < 1e-3 && (shown.y1 - tight.y1).abs() < 1e-3, "shown by the detector's box once known: {shown:?}"); + // a face nobody has looked at, or in another photo, is shown by its own box + let other = lightcraft_geom::Rect { x0: 0.1, y0: 0.1, x1: 0.2, y1: 0.2 }; + assert_eq!(s.face_view(a, other), other); + // the page of the person uses it too + let r = s.execute("faces.person", &json!({"name": "Ann"})).unwrap(); + let view = &r["confirmed"][0]["view"]; + assert!((view["x0"].as_f64().unwrap() - tight.x0).abs() < 1e-3, "{r}"); + // and it survives a restart: the index file keeps it beside the embedding + let (tag, path) = (s.faces.index.tag.clone(), d.join("view-cache.bin")); + let mut written = crate::faces_index::Index::default(); + written.reset(&tag, 64, Some(path.clone())); + written.insert_with_view(a.0, &loose, unit64(&[(0, 1.0)]), &tight); + written.save().unwrap(); + drop(written); + let mut again = crate::faces_index::Index::default(); + again.reset(&tag, 64, Some(path)); + let v = again.view(a.0, &loose).expect("the face is remembered"); + assert!((v.x0 - tight.x0).abs() < 1e-3 && (v.y0 - tight.y0).abs() < 1e-3, "{v:?}"); + let _ = std::fs::remove_dir_all(&d); +} diff --git a/crates/faces/Cargo.toml b/crates/faces/Cargo.toml new file mode 100644 index 00000000..112967ce --- /dev/null +++ b/crates/faces/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "lightcraft-faces" +version.workspace = true +edition.workspace = true +license.workspace = true +description = "LightCraft faces (L1): bounded ONNX CPU inference, face-model manifests, validation and detection" + +[lints] +workspace = true + +[features] +default = [] + +[dependencies] +ndarray = { version = "=0.17.2", default-features = false, features = ["std"] } +# Rust intrinsics only, CPU-dispatched; no BLAS, cc/assembly or nested threading. +matrixmultiply = { version = "=0.3.11", default-features = false, features = ["std", "avx512"] } +web-time = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +sha2 = { workspace = true } +thiserror = { workspace = true } + +[dev-dependencies] +proptest = { workspace = true } diff --git a/crates/faces/src/align.rs b/crates/faces/src/align.rs new file mode 100644 index 00000000..4982325a --- /dev/null +++ b/crates/faces/src/align.rs @@ -0,0 +1,273 @@ +//! Cutting a face out of a photo the way recognition models expect it. +//! +//! With the five landmarks the detector finds (eyes, nose tip, mouth corners), the face is turned and scaled +//! (a similarity transform, found by least squares) so those points land on the standard 112 × 112 +//! ArcFace template. Without landmarks (a box from an XMP sidecar) a padded square around the box is used +//! instead. Both warp with supersampling, so shrinking a large face does not alias. +//! +//! Pure geometry on 8-bit RGB: no image types, no panics (every access is checked). + +/// The five template points of a 112 × 112 aligned face, image-left eye first: eye, eye, nose tip, mouth +/// corner, mouth corner. (The standard ArcFace template.) +pub const TEMPLATE_112: [(f32, f32); 5] = [(38.2946, 51.6963), (73.5318, 51.5014), (56.0252, 71.7366), (41.5493, 92.3655), (70.7299, 92.2041)]; + +/// A borrowed 8-bit RGB picture. +#[derive(Clone, Copy)] +pub struct Rgb<'a> { + pub data: &'a [u8], + pub width: usize, + pub height: usize, +} + +impl Rgb<'_> { + fn valid(&self) -> bool { + self.width > 0 && self.height > 0 && self.width.checked_mul(self.height).and_then(|n| n.checked_mul(3)) == Some(self.data.len()) + } + + /// The pixel at integer coordinates, clamped to the picture. + fn at(&self, x: isize, y: isize) -> [f32; 3] { + let (x, y) = (x.clamp(0, self.width as isize - 1) as usize, y.clamp(0, self.height as isize - 1) as usize); + let i = (y * self.width + x) * 3; + match self.data.get(i..i + 3) { + Some(p) => [f32::from(p[0]), f32::from(p[1]), f32::from(p[2])], + None => [0.0; 3], + } + } + + /// Bilinear sample at a fractional position (pixel centres at integers). + fn sample(&self, x: f32, y: f32) -> [f32; 3] { + let (x0, y0) = (x.floor(), y.floor()); + let (fx, fy) = (x - x0, y - y0); + let (x0, y0) = (x0 as isize, y0 as isize); + let (a, b, c, d) = (self.at(x0, y0), self.at(x0 + 1, y0), self.at(x0, y0 + 1), self.at(x0 + 1, y0 + 1)); + let mut out = [0.0; 3]; + for k in 0..3 { + let top = a[k] * (1.0 - fx) + b[k] * fx; + let bottom = c[k] * (1.0 - fx) + d[k] * fx; + out[k] = top * (1.0 - fy) + bottom * fy; + } + out + } +} + +/// `q = a · p + t` with complex `a` (scale and rotation) and `t`, mapping source positions to output positions. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Similarity { + pub a: (f32, f32), + pub t: (f32, f32), +} + +impl Similarity { + pub fn scale(&self) -> f32 { + self.a.0.hypot(self.a.1) + } + + /// Where an output position `q` comes from in the source. + fn invert(&self, q: (f32, f32)) -> Option<(f32, f32)> { + let n = self.a.0 * self.a.0 + self.a.1 * self.a.1; + if !(n.is_finite() && n > 1e-12) { + return None; + } + let (dx, dy) = (q.0 - self.t.0, q.1 - self.t.1); + Some(((dx * self.a.0 + dy * self.a.1) / n, (dy * self.a.0 - dx * self.a.1) / n)) + } +} + +/// The least-squares similarity transform taking `from[i]` to `to[i]` (no reflection). `None` when the points +/// do not determine one (all the same point, or non-finite). +pub fn fit_similarity(from: &[(f32, f32)], to: &[(f32, f32)]) -> Option { + if from.len() != to.len() || from.len() < 2 { + return None; + } + let n = from.len() as f32; + let mean = |pts: &[(f32, f32)]| (pts.iter().map(|p| p.0).sum::() / n, pts.iter().map(|p| p.1).sum::() / n); + let (mp, mq) = (mean(from), mean(to)); + // a = Σ conj(p') q' / Σ |p'|² (complex numbers) + let (mut re, mut im, mut norm) = (0.0f32, 0.0f32, 0.0f32); + for (p, q) in from.iter().zip(to) { + let (px, py, qx, qy) = (p.0 - mp.0, p.1 - mp.1, q.0 - mq.0, q.1 - mq.1); + re += px * qx + py * qy; + im += px * qy - py * qx; + norm += px * px + py * py; + } + if !(norm.is_finite() && norm > 1e-9) { + return None; + } + let a = (re / norm, im / norm); + let t = (mq.0 - (a.0 * mp.0 - a.1 * mp.1), mq.1 - (a.0 * mp.1 + a.1 * mp.0)); + (a.0.is_finite() && a.1.is_finite() && t.0.is_finite() && t.1.is_finite()).then_some(Similarity { a, t }) +} + +/// Fill an `out_w` × `out_h` picture by taking, for each output pixel, the source under it, averaging +/// enough sub-samples that shrinking does not alias. +fn warp(img: &Rgb, m: &Similarity, out_w: usize, out_h: usize) -> Option> { + if !img.valid() || out_w == 0 || out_h == 0 || out_w > 4096 || out_h > 4096 { + return None; + } + if ![m.a.0, m.a.1, m.t.0, m.t.1].iter().all(|v| v.is_finite()) { + return None; + } + // how many source pixels one output pixel spans, per side + let span = 1.0 / m.scale(); + let n = if span.is_finite() { (span.ceil() as usize).clamp(1, 4) } else { 1 }; + let mut out = vec![0u8; out_w * out_h * 3]; + for (i, px) in out.as_chunks_mut::<3>().0.iter_mut().enumerate() { + let (ox, oy) = ((i % out_w) as f32, (i / out_w) as f32); + let mut acc = [0.0f32; 3]; + for sy in 0..n { + for sx in 0..n { + let q = (ox + (sx as f32 + 0.5) / n as f32 - 0.5, oy + (sy as f32 + 0.5) / n as f32 - 0.5); + let (px_, py_) = m.invert(q)?; + let s = img.sample(px_, py_); + acc.iter_mut().zip(s).for_each(|(a, v)| *a += v); + } + } + let k = (n * n) as f32; + for (o, a) in px.iter_mut().zip(acc) { + *o = (a / k).round().clamp(0.0, 255.0) as u8; + } + } + Some(out) +} + +/// The face aligned to the ArcFace template, `out_w` × `out_h` (the template is for 112 × 112 and is scaled for +/// other sizes). `landmarks` are in pixels of `img`. `None` when the landmarks cannot define a transform. +pub fn align_to_template(img: &Rgb, landmarks: &[(f32, f32); 5], out_w: usize, out_h: usize) -> Option> { + let (sx, sy) = (out_w as f32 / 112.0, out_h as f32 / 112.0); + let template: Vec<(f32, f32)> = TEMPLATE_112.iter().map(|(x, y)| (x * sx, y * sy)).collect(); + let m = fit_similarity(landmarks, &template)?; + // a face smaller than a few pixels, or absurdly large, is not a face we can use + let scale = m.scale(); + if !(scale.is_finite() && (0.01..=100.0).contains(&scale)) { + return None; + } + warp(img, &m, out_w, out_h) +} + +/// A square around a face box (pixels of `img`), `pad` times its larger side, as `out_w` × `out_h`. +pub fn crop_box(img: &Rgb, x0: f32, y0: f32, x1: f32, y1: f32, pad: f32, out_w: usize, out_h: usize) -> Option> { + if ![x0, y0, x1, y1, pad].iter().all(|v| v.is_finite()) { + return None; + } + let side = (x1 - x0).abs().max((y1 - y0).abs()) * pad; + if !(side.is_finite() && side >= 2.0) { + return None; + } + let (cx, cy) = ((x0 + x1) / 2.0, (y0 + y1) / 2.0); + let s = out_w.min(out_h) as f32 / side; + let m = Similarity { a: (s, 0.0), t: (out_w as f32 / 2.0 - s * cx, out_h as f32 / 2.0 - s * cy) }; + warp(img, &m, out_w, out_h) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn close(a: (f32, f32), b: (f32, f32), tol: f32) -> bool { + (a.0 - b.0).abs() <= tol && (a.1 - b.1).abs() <= tol + } + + fn apply(m: &Similarity, p: (f32, f32)) -> (f32, f32) { + (m.a.0 * p.0 - m.a.1 * p.1 + m.t.0, m.a.0 * p.1 + m.a.1 * p.0 + m.t.1) + } + + #[test] + fn a_known_similarity_is_recovered() { + // rotate 20 degrees, scale 0.7, translate + let (s, th) = (0.7f32, 20f32.to_radians()); + let truth = Similarity { a: (s * th.cos(), s * th.sin()), t: (13.0, -5.0) }; + let from: Vec<(f32, f32)> = TEMPLATE_112.to_vec(); + let to: Vec<(f32, f32)> = from.iter().map(|p| apply(&truth, *p)).collect(); + let m = fit_similarity(&from, &to).unwrap(); + assert!(close(m.a, truth.a, 1e-4) && close(m.t, truth.t, 1e-3), "{m:?} vs {truth:?}"); + // noisy points still give the best fit (the residual is small) + let noisy: Vec<(f32, f32)> = to.iter().enumerate().map(|(i, p)| (p.0 + (i as f32 - 2.0) * 0.2, p.1 - (i as f32 - 2.0) * 0.1)).collect(); + let m2 = fit_similarity(&from, &noisy).unwrap(); + assert!(from.iter().zip(&noisy).all(|(p, q)| close(apply(&m2, *p), *q, 1.0))); + } + + #[test] + fn degenerate_points_give_nothing() { + let same = vec![(5.0, 5.0); 5]; + assert!(fit_similarity(&same, &TEMPLATE_112).is_none()); + assert!(fit_similarity(&TEMPLATE_112, &TEMPLATE_112[..3]).is_none()); + assert!(fit_similarity(&[(0.0, 0.0)], &[(1.0, 1.0)]).is_none()); + let nan = vec![(f32::NAN, 0.0); 5]; + assert!(fit_similarity(&nan, &TEMPLATE_112).is_none()); + let inf = vec![(f32::INFINITY, 0.0), (1.0, 1.0), (2.0, 2.0), (3.0, 3.0), (4.0, 4.0)]; + assert!(fit_similarity(&inf, &TEMPLATE_112).is_none()); + } + + /// A picture with a bright dot at each of five landmarks, the landmarks being the template moved by a known + /// similarity: aligning must put the dots back on the template. + #[test] + fn alignment_puts_the_landmarks_on_the_template() { + let (w, h) = (400usize, 300usize); + let (s, th) = (1.7f32, -15f32.to_radians()); + let place = Similarity { a: (s * th.cos(), s * th.sin()), t: (150.0, 60.0) }; + let landmarks: Vec<(f32, f32)> = TEMPLATE_112.iter().map(|p| apply(&place, *p)).collect(); + let mut data = vec![0u8; w * h * 3]; + for l in &landmarks { + for dy in -3i32..=3 { + for dx in -3i32..=3 { + let (x, y) = ((l.0 as i32 + dx) as usize, (l.1 as i32 + dy) as usize); + if x < w && y < h { + data[(y * w + x) * 3] = 255; + } + } + } + } + let img = Rgb { data: &data, width: w, height: h }; + let lm: [(f32, f32); 5] = [landmarks[0], landmarks[1], landmarks[2], landmarks[3], landmarks[4]]; + let out = align_to_template(&img, &lm, 112, 112).unwrap(); + assert_eq!(out.len(), 112 * 112 * 3); + // the red channel's centre of mass near each template point is on the point + for t in TEMPLATE_112 { + let (mut sx, mut sy, mut sw) = (0.0f32, 0.0f32, 0.0f32); + for y in (t.1 as usize).saturating_sub(8)..(t.1 as usize + 9).min(112) { + for x in (t.0 as usize).saturating_sub(8)..(t.0 as usize + 9).min(112) { + let v = f32::from(out[(y * 112 + x) * 3]); + sx += v * x as f32; + sy += v * y as f32; + sw += v; + } + } + assert!(sw > 100.0, "no dot near {t:?}"); + assert!(close((sx / sw, sy / sw), t, 1.5), "dot at {:?}, template {t:?}", (sx / sw, sy / sw)); + } + } + + #[test] + fn box_crops_are_centred_and_bad_inputs_are_refused() { + let (w, h) = (200usize, 100usize); + let mut data = vec![0u8; w * h * 3]; + for y in 40..60 { + for x in 90..110 { + data[(y * w + x) * 3 + 1] = 255; + } + } + let img = Rgb { data: &data, width: w, height: h }; + let out = crop_box(&img, 90.0, 40.0, 110.0, 60.0, 2.0, 64, 64).unwrap(); + // the 20 px block fills the middle half of the 64 px output + let at = |x: usize, y: usize| out[(y * 64 + x) * 3 + 1]; + assert!(at(32, 32) > 250 && at(20, 20) > 250 && at(44, 44) > 250 && at(5, 5) < 5 && at(60, 32) < 5); + assert!(crop_box(&img, 5.0, 5.0, 5.0, 5.0, 2.0, 64, 64).is_none(), "an empty box"); + assert!(crop_box(&img, f32::NAN, 0.0, 10.0, 10.0, 2.0, 64, 64).is_none()); + assert!(crop_box(&img, 0.0, 0.0, 10.0, 10.0, 2.0, 0, 64).is_none()); + assert!(crop_box(&img, 0.0, 0.0, 10.0, 10.0, 2.0, 100_000, 64).is_none()); + let bad = Rgb { data: &data[..10], width: w, height: h }; + assert!(crop_box(&bad, 0.0, 0.0, 10.0, 10.0, 2.0, 64, 64).is_none(), "wrong data length"); + } + + #[test] + fn faces_off_the_edge_and_tiny_or_huge_landmarks_do_not_panic() { + let data = vec![80u8; 50 * 50 * 3]; + let img = Rgb { data: &data, width: 50, height: 50 }; + // a box mostly outside the picture is padded by clamping, not an error + assert!(crop_box(&img, -30.0, -30.0, 20.0, 20.0, 1.5, 32, 32).is_some()); + let far: [(f32, f32); 5] = [(1e6, 1e6), (1e6 + 5.0, 1e6), (1e6 + 2.0, 1e6 + 3.0), (1e6, 1e6 + 6.0), (1e6 + 5.0, 1e6 + 6.0)]; + let _ = align_to_template(&img, &far, 112, 112); + let tiny: [(f32, f32); 5] = [(1.0, 1.0), (1.01, 1.0), (1.0, 1.01), (1.01, 1.01), (1.005, 1.005)]; + assert!(align_to_template(&img, &tiny, 112, 112).is_none(), "landmarks a hundredth of a pixel apart"); + } +} diff --git a/crates/faces/src/catalog.rs b/crates/faces/src/catalog.rs new file mode 100644 index 00000000..07a20622 --- /dev/null +++ b/crates/faces/src/catalog.rs @@ -0,0 +1,228 @@ +//! Your own list of models to download: `catalog.json` in the models folder. +//! +//! LightCraft's built-in list only holds models whose terms let the project point at them. Anything else (a stronger +//! recogniser whose weights are for research use, a model you trained, a mirror you trust) can be added by you: each entry +//! is a model manifest (the same fields as a `face-model.json`) plus the address to fetch it from, and it then gets the same +//! Download button as a built-in one, with its licence notice shown before anything is fetched. LightCraft ships and links to +//! none of them; what the file says, and whether you may use the weights, is yours to check. +//! +//! ```json +//! { "models": [ { +//! "id": "my-recogniser", "name": "My recogniser", "version": "1", "role": "embedder", +//! "url": "https://example.org/weights/recogniser.onnx", +//! "sha256": "…64 lowercase hex digits…", "sizeBytes": 166000000, +//! "licence": { "name": "Research use only", "commercial": "no", "notice": "Not for commercial use." }, +//! "provenance": "Trained on …", +//! "output": { "kind": "embedding", "dim": 512 }, +//! "thresholds": { "matchCosine": 0.4 }, +//! "speed": 1.7 +//! } ] } +//! ``` +//! +//! `speed` (optional) is how many times faster the model is than a ResNet-100, shown in Settings instead of timings. +//! `input` may be left out (112 × 112 RGB, `(x − 127.5) / 127.5`, the ArcFace convention). The file is untrusted: it is read +//! with a size limit, every entry is validated, and a bad entry is reported and skipped, never a reason to refuse the rest. +//! The `sha256` and `sizeBytes` are required, since they are what a downloaded file is checked against. + +use serde::Deserialize; +use serde_json::Value; + +use crate::known::{self, Download}; +use crate::manifest::{self, ModelManifest, Role}; + +/// Largest `catalog.json` we read. +pub const MAX_CATALOG_BYTES: usize = 256 * 1024; +/// Most entries we take from it. +pub const MAX_ENTRIES: usize = 50; + +/// One model from the catalog: its manifest, where to fetch it and what to call the file meanwhile. +#[derive(Clone, Debug, PartialEq)] +pub struct Entry { + pub manifest: ModelManifest, + pub url: String, + pub file_name: String, +} + +impl Entry { + /// What the download machinery needs. (`sha256` and `size_bytes` are present: [`parse`] requires them.) + pub fn download(&self) -> Option { + Some(Download { + id: self.manifest.id.clone(), + url: self.url.clone(), + file_name: self.file_name.clone(), + size_bytes: self.manifest.size_bytes?, + sha256: self.manifest.sha256.clone()?, + }) + } +} + +/// What a catalog file holds. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct Catalog { + pub entries: Vec, + /// One plain sentence for each entry (or the file) that could not be used. + pub errors: Vec, +} + +#[derive(Deserialize)] +struct Raw { + url: String, + #[serde(flatten)] + manifest: ModelManifest, +} + +/// A file name for the download, from the address: letters, digits, `.`, `_`, `-`, ending in `.onnx`. +fn file_name_of(url: &str) -> String { + let last = url.split(['?', '#']).next().unwrap_or("").rsplit('/').next().unwrap_or(""); + let clean: String = last.chars().filter(|c| c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '-')).take(80).collect(); + if clean.ends_with(".onnx") && clean.len() > 5 && !clean.starts_with('.') { clean } else { "model.onnx".to_string() } +} + +fn entry(v: &Value) -> Result { + let raw: Raw = serde_json::from_value(v.clone()).map_err(|e| e.to_string())?; + manifest::validate(&raw.manifest).map_err(|e| e.to_string())?; + let m = &raw.manifest; + if m.role != Role::Embedder { + return Err("only face recognition models (role \"embedder\") can be downloaded".into()); + } + if m.sha256.is_none() || m.size_bytes.is_none() { + return Err("needs `sha256` and `sizeBytes`: a download is checked against them".into()); + } + if !(raw.url.len() <= 500 && raw.url.starts_with("https://") && !raw.url.chars().any(|c| c.is_control() || c.is_whitespace())) { + return Err("`url` must be an https:// address of at most 500 characters, without spaces".into()); + } + if known::all().iter().any(|k| k.id == m.id) { + return Err("that id belongs to a model LightCraft already knows".into()); + } + if m.sha256.as_deref().is_some_and(|h| known::lookup(h).is_some()) { + return Err("LightCraft already knows a model with that file".into()); + } + Ok(Entry { file_name: file_name_of(&raw.url), url: raw.url, manifest: raw.manifest }) +} + +/// Read a catalog file's bytes: the usable entries, and why the others were left out. +pub fn parse(bytes: &[u8]) -> Catalog { + let mut out = Catalog::default(); + if bytes.len() > MAX_CATALOG_BYTES { + out.errors.push(format!("catalog.json is larger than {} KB: not read", MAX_CATALOG_BYTES / 1024)); + return out; + } + // Notepad and older PowerShell begin a UTF-8 file with a byte-order mark + let bytes = bytes.strip_prefix(&[0xEF, 0xBB, 0xBF]).unwrap_or(bytes); + let value: Value = match serde_json::from_slice(bytes) { + Ok(v) => v, + Err(e) => { + out.errors.push(format!("catalog.json is not valid JSON: {e}")); + return out; + } + }; + let Some(models) = value.get("models").and_then(Value::as_array) else { + out.errors.push("catalog.json needs a \"models\" list".into()); + return out; + }; + if models.len() > MAX_ENTRIES { + out.errors.push(format!("only the first {MAX_ENTRIES} models of the catalog are read")); + } + for (i, v) in models.iter().take(MAX_ENTRIES).enumerate() { + let label = v + .get("id") + .and_then(Value::as_str) + .map_or_else(|| format!("model {}", i + 1), |id| format!("`{}`", id.chars().take(40).collect::())); + match entry(v) { + Ok(e) if out.entries.iter().any(|x| x.manifest.id == e.manifest.id) => { + out.errors.push(format!("{label}: listed twice, the first is used")) + } + Ok(e) => out.entries.push(e), + Err(why) => out.errors.push(format!("{label}: {why}")), + } + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn good(id: &str) -> Value { + json!({ + "id": id, "name": "Test model", "version": "1", "role": "embedder", + "url": "https://example.org/w/recogniser-r50.onnx?download=1", + "sha256": "ab".repeat(32), "sizeBytes": 1234, + "licence": {"name": "Research only", "commercial": "no", "notice": "Not for commercial use."}, + "output": {"kind": "embedding", "dim": 512}, "thresholds": {"matchCosine": 0.4}, "speed": 1.7, + }) + } + + fn cat(models: Vec) -> Catalog { + parse(json!({"models": models}).to_string().as_bytes()) + } + + #[test] + fn a_manifest_with_an_address_is_a_downloadable_model() { + let c = cat(vec![good("my-model")]); + assert_eq!(c.errors, Vec::::new()); + let e = &c.entries[0]; + assert_eq!((e.manifest.id.as_str(), e.file_name.as_str()), ("my-model", "recogniser-r50.onnx"), "the file name comes from the address"); + assert_eq!(e.manifest.licence.commercial, crate::Commercial::No, "a non-commercial model is the user's to add"); + let d = e.download().unwrap(); + assert_eq!((d.size_bytes, d.sha256.len(), d.url.as_str()), (1234, 64, "https://example.org/w/recogniser-r50.onnx?download=1")); + assert_eq!(e.manifest.input, crate::InputSpec::default(), "the input defaults to the ArcFace convention"); + assert_eq!(e.manifest.speed, Some(1.7), "the speed ratio is the user's to state"); + let mut nonsense = good("nonsense-speed"); + nonsense["speed"] = json!(-3); + let c = cat(vec![nonsense]); + assert!(c.entries.is_empty() && c.errors.iter().any(|e| e.contains("speed")), "{:?}", c.errors); + } + + #[test] + fn a_file_saved_with_a_byte_order_mark_still_reads() { + let mut bytes = vec![0xEF, 0xBB, 0xBF]; + bytes.extend_from_slice(json!({"models": [good("with-bom")]}).to_string().as_bytes()); + let c = parse(&bytes); + assert_eq!((c.entries.len(), c.errors.len()), (1, 0), "{:?}", c.errors); + } + + #[test] + fn bad_entries_are_reported_and_the_rest_still_load() { + let mut no_hash = good("no-hash"); + no_hash.as_object_mut().unwrap().remove("sha256"); + let mut http = good("plain-http"); + http["url"] = json!("http://example.org/m.onnx"); + let mut spaces = good("spaces"); + spaces["url"] = json!("https://example.org/a b.onnx"); + let mut detector = good("a-detector"); + detector["role"] = json!("detector"); + detector["output"] = json!({"kind": "detector", "decoder": "yunet-v2"}); + let mut bad_id = good("Bad/Id"); + bad_id["id"] = json!("../escape"); + let builtin = good("sface-2021dec"); + let mut known_file = good("known-file"); + known_file["sha256"] = json!(known::SFACE_SHA256); + let c = cat(vec![good("fine"), no_hash, http, spaces, detector, bad_id, builtin, known_file, good("fine"), json!(7), json!({"id": "x"})]); + assert_eq!(c.entries.iter().map(|e| e.manifest.id.as_str()).collect::>(), ["fine"]); + assert_eq!(c.errors.len(), 10, "{:#?}", c.errors); + assert!(c.errors.iter().any(|e| e.contains("`no-hash`") && e.contains("sha256"))); + assert!(c.errors.iter().any(|e| e.contains("https"))); + assert!(c.errors.iter().any(|e| e.contains("already knows"))); + assert!(c.errors.iter().any(|e| e.contains("listed twice"))); + } + + #[test] + fn a_hostile_or_broken_file_is_an_error_message_not_a_crash() { + for bytes in [&b""[..], b"not json", b"[]", b"{}", b"{\"models\": 5}", b"\xff\xfe\x00", &vec![b' '; MAX_CATALOG_BYTES + 1]] { + let c = parse(bytes); + assert!(c.entries.is_empty() && !c.errors.is_empty(), "{:?}", String::from_utf8_lossy(&bytes[..bytes.len().min(20)])); + } + // more entries than allowed: the extra are ignored, with a note + let many: Vec = (0..MAX_ENTRIES + 20).map(|i| good(&format!("m{i}"))).collect(); + let c = cat(many); + assert_eq!(c.entries.len(), MAX_ENTRIES); + assert!(c.errors.iter().any(|e| e.contains("first"))); + // file names from odd addresses are safe + for url in ["https://x.org/a/b/c.onnx", "https://x.org/", "https://x.org/..%2f..%2fevil.onnx", "https://x.org/q.bin", "https://x.org/.onnx"] { + let got = file_name_of(url); + assert!(!got.contains(['/', '\\']) && got.ends_with(".onnx") && !got.starts_with('.'), "{url} -> {got}"); + } + } +} diff --git a/crates/faces/src/graph.rs b/crates/faces/src/graph.rs new file mode 100644 index 00000000..a5a22d26 --- /dev/null +++ b/crates/faces/src/graph.rs @@ -0,0 +1,442 @@ +//! A full, bounded read of a small ONNX model: its nodes, attributes and weights. +//! +//! The detector and installed recognisers share this reader and [`crate::net`]. Files are capped at +//! 512 MiB, with at most 64 million weight elements. The bytes are still treated as hostile: the +//! number of nodes, weight tensors and elements is capped, every length is checked against its parent, +//! and a weight blob whose size disagrees with its shape is an error. + +use std::collections::HashMap; +use std::io::Cursor; + +use crate::onnx::{ProbeError, Walker, Wire}; + +type Result = std::result::Result; + +const MAX_NODES: usize = 20_000; +const MAX_TENSORS: usize = 20_000; +/// Elements across all weight tensors (4 bytes each as float32). +const MAX_ELEMENTS: usize = 64 * 1024 * 1024; +const MAX_NAME: usize = 4096; +const MAX_LIST: usize = 1_000_000; + +#[derive(Clone, Debug, PartialEq)] +pub enum Attr { + Int(i64), + Float(f32), + Str(String), + Ints(Vec), + Floats(Vec), + /// A kind we do not read (a tensor, a sub-graph…). + Other, +} + +#[derive(Clone, Debug)] +pub struct Node { + pub op: String, + pub inputs: Vec, + pub outputs: Vec, + pub attrs: HashMap, +} + +impl Node { + pub fn int(&self, name: &str) -> Option { + match self.attrs.get(name) { + Some(Attr::Int(v)) => Some(*v), + _ => None, + } + } + pub fn ints(&self, name: &str) -> Option<&[i64]> { + match self.attrs.get(name) { + Some(Attr::Ints(v)) => Some(v), + _ => None, + } + } + pub fn string(&self, name: &str) -> Option<&str> { + match self.attrs.get(name) { + Some(Attr::Str(v)) => Some(v), + _ => None, + } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub enum Data { + F32(Vec), + I64(Vec), +} + +/// A weight tensor. +#[derive(Clone, Debug, PartialEq)] +pub struct Weight { + pub dims: Vec, + pub data: Data, +} + +#[derive(Clone, Debug)] +pub struct Graph { + pub nodes: Vec, + pub weights: HashMap, + /// The inputs the caller feeds (weights some exporters also list as inputs are left out), and their shapes + /// (0 for a dynamic size). + pub inputs: Vec<(String, Vec)>, + pub outputs: Vec, +} + +fn malformed(why: &'static str) -> Result { + Err(ProbeError::Malformed(why)) +} + +fn f32s(raw: &[u8]) -> Vec { + raw.as_chunks::<4>().0.iter().map(|b| f32::from_le_bytes(*b)).collect() +} + +fn i64s(raw: &[u8]) -> Vec { + raw.as_chunks::<8>().0.iter().map(|b| i64::from_le_bytes(*b)).collect() +} + +/// A list of varints in a packed field (`start..end`), appended to `out`. +fn packed_varints(w: &mut Walker>, end: u64, out: &mut Vec) -> Result<()> { + while w.pos()? < end { + if out.len() >= MAX_LIST { + return malformed("a list is too long"); + } + out.push(w.varint()?); + if w.pos()? > end { + return malformed("a packed number runs past its field"); + } + } + Ok(()) +} + +fn attribute(w: &mut Walker>, end: u64) -> Result<(String, Attr)> { + let (mut name, mut kind) = (String::new(), 0u64); + let (mut i, mut f, mut s) = (0i64, 0f32, String::new()); + let (mut ints, mut floats): (Vec, Vec) = (Vec::new(), Vec::new()); + while let Some((n, f_)) = w.next(end)? { + match (n, f_) { + (1, Wire::Len { start, end: e }) => { + name = w.string(start, e)?; + w.seek_to(e)?; + } + (2, Wire::Fixed32(v)) => f = f32::from_bits(v), + (3, Wire::Varint(v)) => i = v as i64, + (4, Wire::Len { start, end: e }) => { + s = w.string(start, e)?; + w.seek_to(e)?; + } + (7, Wire::Len { end: e, .. }) => { + // packed floats + let at = w.pos()?; + let raw = w.bytes(at, e, (MAX_LIST * 4) as u64)?; + if !raw.len().is_multiple_of(4) || floats.len().saturating_add(raw.len() / 4) > MAX_LIST { + return malformed("an attribute float list is invalid or too long"); + } + floats.extend(f32s(&raw)); + } + (7, Wire::Fixed32(v)) => { + if floats.len() >= MAX_LIST { + return malformed("an attribute float list is too long"); + } + floats.push(f32::from_bits(v)); + } + (8, Wire::Len { end: e, .. }) => { + packed_varints(w, e, &mut ints)?; + w.seek_to(e)?; + } + (8, Wire::Varint(v)) => { + if ints.len() >= MAX_LIST { + return malformed("a list is too long"); + } + ints.push(v); + } + (20, Wire::Varint(v)) => kind = v, + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + let ints = ints.into_iter().map(|v| v as i64).collect(); + let value = match kind { + 1 => Attr::Float(f), + 2 => Attr::Int(i), + 3 => Attr::Str(s), + 6 => Attr::Floats(floats), + 7 => Attr::Ints(ints), + _ => Attr::Other, + }; + Ok((name, value)) +} + +fn node(w: &mut Walker>, end: u64) -> Result { + let mut n = Node { op: String::new(), inputs: Vec::new(), outputs: Vec::new(), attrs: HashMap::new() }; + while let Some((num, f)) = w.next(end)? { + match (num, f) { + (1, Wire::Len { start, end: e }) => { + n.inputs.push(w.string(start, e)?); + w.seek_to(e)?; + } + (2, Wire::Len { start, end: e }) => { + n.outputs.push(w.string(start, e)?); + w.seek_to(e)?; + } + (4, Wire::Len { start, end: e }) => { + n.op = w.string(start, e)?; + w.seek_to(e)?; + } + (5, Wire::Len { end: e, .. }) => { + let (k, v) = attribute(w, e)?; + if k.is_empty() || n.attrs.insert(k, v).is_some() { + return malformed("an attribute name is empty or repeated"); + } + w.seek_to(e)?; + } + (7, Wire::Len { start, end: e }) => { + let domain = w.string(start, e)?; + if !domain.is_empty() && domain != "ai.onnx" { + return malformed("custom operator domains are unsupported"); + } + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + if n.inputs.len() > 64 || n.outputs.len() > 64 || n.attrs.len() > 64 { + return malformed("a node is too large"); + } + } + Ok(n) +} + +/// A weight tensor (`TensorProto`): its name, shape and float32 or int64 values. +fn tensor(w: &mut Walker>, end: u64, budget: &mut usize) -> Result<(String, Option)> { + let (mut name, mut dtype) = (String::new(), 0u64); + let mut dims: Vec = Vec::new(); + let mut raw: Option<&[u8]> = None; + let (mut float_data, mut int_data): (Vec, Vec) = (Vec::new(), Vec::new()); + let max_bytes = (*budget as u64).saturating_mul(8); + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Len { end: e, .. }) => { + packed_varints(w, e, &mut dims)?; + w.seek_to(e)?; + } + (1, Wire::Varint(v)) => dims.push(v), + (2, Wire::Varint(v)) => dtype = v, + (4, Wire::Len { end: e, .. }) => { + let at = w.pos()?; + let bytes = w.bytes(at, e, max_bytes)?; + if !bytes.len().is_multiple_of(4) || float_data.len().saturating_add(bytes.len() / 4) > *budget { + return malformed("invalid or excessive packed float weights"); + } + float_data.extend(f32s(&bytes)); + } + (4, Wire::Fixed32(v)) => { + if float_data.len() >= *budget { + return malformed("too many float weights"); + } + float_data.push(f32::from_bits(v)); + } + (7, Wire::Varint(v)) => { + if int_data.len() >= *budget { + return malformed("too many integer weights"); + } + int_data.push(v); + } + (7, Wire::Len { end: e, .. }) => { + packed_varints(w, e, &mut int_data)?; + w.seek_to(e)?; + } + (8, Wire::Len { start, end: e }) => { + name = w.string(start, e)?; + w.seek_to(e)?; + } + (9, Wire::Len { start, end: e }) => { + if raw.is_some() || e.saturating_sub(start) > max_bytes { + return malformed("repeated or excessive raw weights"); + } + let from = usize::try_from(start).map_err(|_| ProbeError::Malformed("weight offset overflow"))?; + let to = usize::try_from(e).map_err(|_| ProbeError::Malformed("weight offset overflow"))?; + raw = Some((*w.r.get_ref()).get(from..to).ok_or(ProbeError::Malformed("short raw weights"))?); + w.seek_to(e)?; + } + (13, Wire::Len { .. }) | (14, Wire::Varint(1..)) => return malformed("external tensor data is unsupported"), + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + if dims.len() > 8 { + return malformed("a tensor has too many dimensions"); + } + } + let mut elements: usize = 1; + let mut shape = Vec::new(); + for d in &dims { + let d = usize::try_from(*d).map_err(|_| ProbeError::Malformed("a tensor dimension is too large"))?; + elements = elements.checked_mul(d).ok_or(ProbeError::Malformed("a tensor is too large"))?; + shape.push(d); + } + if elements > *budget { + return malformed("the weights are too large"); + } + *budget -= elements; + if raw.is_some() && (!float_data.is_empty() || !int_data.is_empty()) { + return malformed("ambiguous raw and typed weights"); + } + let data = match dtype { + 1 => { + let v = if let Some(raw) = raw { + if elements.checked_mul(4) != Some(raw.len()) { + return malformed("float weight bytes disagree with shape"); + } + f32s(raw) + } else { + float_data + }; + if v.len() != elements { + return malformed("a weight tensor's size disagrees with its shape"); + } + Some(Data::F32(v)) + } + 7 => { + let v = if let Some(raw) = raw { + if elements.checked_mul(8) != Some(raw.len()) { + return malformed("integer weight bytes disagree with shape"); + } + i64s(raw) + } else { + int_data.into_iter().map(|v| v as i64).collect() + }; + if v.len() != elements { + return malformed("a weight tensor's size disagrees with its shape"); + } + Some(Data::I64(v)) + } + _ => return malformed("only float32 and int64 initializers are supported"), + }; + if name.is_empty() { + return malformed("a weight name is empty"); + } + Ok((name, data.map(|data| Weight { dims: shape, data }))) +} + +fn value_info(w: &mut Walker>, end: u64) -> Result<(String, Vec)> { + let (mut name, mut shape) = (String::new(), Vec::new()); + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Len { start, end: e }) => { + name = w.string(start, e)?; + w.seek_to(e)?; + } + (2, Wire::Len { end: e, .. }) => { + // TypeProto > tensor_type (1) > shape (2) > dim (1) > dim_value (1) + let mut info = crate::onnx::TensorInfo { name: String::new(), elem_type: 0, shape: Vec::new() }; + crate::onnx::type_proto(w, e, &mut info)?; + shape = info + .shape + .iter() + .map(|d| match d { + crate::onnx::Dim::Fixed(v) => { + usize::try_from(*v).map_err(|_| ProbeError::Malformed("a declared dimension exceeds this platform")) + } + crate::onnx::Dim::Dynamic(_) => Ok(0), + }) + .collect::>()?; + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + if name.len() > MAX_NAME { + return malformed("a name is too long"); + } + Ok((name, shape)) +} + +/// Read a bounded float32 ONNX inference model. +pub fn load(bytes: &[u8]) -> Result { + let total = bytes.len() as u64; + if total > 512 * 1024 * 1024 { + return malformed("the model exceeds 512 MiB"); + } + let info = crate::onnx::probe_reader(&mut Cursor::new(bytes))?; + if info.opset.is_none_or(|v| !(10..=17).contains(&v)) { + return malformed("only ONNX default-domain opsets 10 through 17 are supported"); + } + if info.inputs.iter().chain(&info.outputs).any(|v| v.shape.iter().any(|d| matches!(d, crate::onnx::Dim::Fixed(0)))) { + return malformed("a declared input/output has a zero dimension"); + } + if info.inputs.iter().chain(&info.outputs).any(|v| v.elem_type != 1) { + return malformed("model inputs and outputs must be float32"); + } + if total == 0 { + return malformed("the file is empty"); + } + let mut cursor = Cursor::new(bytes); + let mut w = Walker { r: &mut cursor, total, fields: 0 }; + let (mut nodes, mut weights, mut inputs, mut outputs) = (Vec::new(), HashMap::new(), Vec::new(), Vec::new()); + let mut budget = MAX_ELEMENTS; + let mut saw_graph = false; + while let Some((n, f)) = w.next(total)? { + match (n, f) { + (7, Wire::Len { end, .. }) => { + if saw_graph { + return malformed("the model contains multiple graphs"); + } + saw_graph = true; + while let Some((gn, gf)) = w.next(end)? { + match (gn, gf) { + (1, Wire::Len { end: e, .. }) => { + if nodes.len() >= MAX_NODES { + return malformed("too many nodes"); + } + nodes.push(node(&mut w, e)?); + w.seek_to(e)?; + } + (5, Wire::Len { end: e, .. }) => { + if weights.len() >= MAX_TENSORS { + return malformed("too many weight tensors"); + } + let (name, wt) = tensor(&mut w, e, &mut budget)?; + if let Some(wt) = wt + && weights.insert(name, wt).is_some() + { + return malformed("a weight is defined twice"); + } + w.seek_to(e)?; + } + (11, Wire::Len { end: e, .. }) => { + if inputs.len() >= MAX_TENSORS { + return malformed("too many graph inputs"); + } + inputs.push(value_info(&mut w, e)?); + w.seek_to(e)?; + } + (12, Wire::Len { end: e, .. }) => { + if outputs.len() >= 64 { + return malformed("too many graph outputs"); + } + outputs.push(value_info(&mut w, e)?.0); + w.seek_to(e)?; + } + (15, Wire::Len { .. }) => return malformed("sparse initializers are unsupported"), + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + w.seek_to(end)?; + } + (_, Wire::Len { end, .. }) => w.seek_to(end)?, + _ => {} + } + } + if !saw_graph || nodes.is_empty() { + return malformed("no graph"); + } + inputs.retain(|(name, _)| !weights.contains_key(name)); + if inputs.is_empty() || outputs.is_empty() { + return malformed("the graph has no inputs or no outputs"); + } + Ok(Graph { nodes, weights, inputs, outputs }) +} + +#[cfg(test)] +mod tests; diff --git a/crates/faces/src/graph/tests.rs b/crates/faces/src/graph/tests.rs new file mode 100644 index 00000000..38635459 --- /dev/null +++ b/crates/faces/src/graph/tests.rs @@ -0,0 +1,117 @@ +use super::*; +use crate::synthetic::{len_field, value_info_bytes, varint, varint_field}; +fn node() -> Vec { + let mut n = Vec::new(); + for (id, v) in [(1, "x"), (2, "y"), (4, "Identity")] { + len_field(id, v.as_bytes(), &mut n); + } + n +} +fn weight(raw: &[u8]) -> Vec { + let mut t = Vec::new(); + varint_field(1, 1, &mut t); + varint_field(2, 1, &mut t); + len_field(8, b"w", &mut t); + len_field(9, raw, &mut t); + t +} +fn model(node: &[u8], weights: &[Vec]) -> Vec { + let mut g = Vec::new(); + len_field(1, node, &mut g); + for t in weights { + len_field(5, t, &mut g); + } + for (id, name) in [(11, "x"), (12, "y")] { + len_field(id, &value_info_bytes(name, 1, &[Ok(1), Ok(3), Ok(112), Ok(112)]), &mut g); + } + let mut ops = Vec::new(); + varint_field(2, 11, &mut ops); + let mut out = Vec::new(); + len_field(7, &g, &mut out); + len_field(8, &ops, &mut out); + out +} +#[test] +fn raw_and_typed_weights_are_exact_not_truncated_or_ambiguous() { + assert!(load(&model(&node(), &[weight(&1f32.to_le_bytes())])).is_ok()); + for raw in [vec![0; 3], vec![0; 5], vec![0; 8]] { + assert!(load(&model(&node(), &[weight(&raw)])).is_err()); + } + let mut t = weight(&1f32.to_le_bytes()); + varint(4 << 3 | 5, &mut t); + t.extend_from_slice(&1f32.to_le_bytes()); + assert!(load(&model(&node(), &[t])).is_err()); + let mut t = weight(&1f32.to_le_bytes()); + len_field(9, &1f32.to_le_bytes(), &mut t); + assert!(load(&model(&node(), &[t])).is_err()); + // The protobuf float_data form must work as well as raw_data. + let mut t = Vec::new(); + varint_field(1, 1, &mut t); + varint_field(2, 1, &mut t); + len_field(8, b"w", &mut t); + varint(4 << 3 | 5, &mut t); + t.extend_from_slice(&1f32.to_le_bytes()); + assert!(load(&model(&node(), &[t])).is_ok()); +} +#[test] +fn external_sparse_duplicate_custom_domain_and_wrong_dtype_are_errors() { + let t = weight(&1f32.to_le_bytes()); + assert!(load(&model(&node(), &[t.clone(), t.clone()])).is_err()); + for field in [13, 14] { + let mut bad = t.clone(); + if field == 13 { + len_field(field, b"external", &mut bad); + } else { + varint_field(field, 1, &mut bad); + } + assert!(load(&model(&node(), &[bad])).is_err()); + } + let mut bad = t.clone(); + varint_field(2, 10, &mut bad); + assert!(load(&model(&node(), &[bad])).is_err()); + let mut n = node(); + len_field(7, b"vendor.example", &mut n); + assert!(load(&model(&n, &[])).is_err()); + let mut attr = Vec::new(); + len_field(1, b"ratio", &mut attr); + varint_field(20, 1, &mut attr); + let mut n = node(); + len_field(5, &attr, &mut n); + len_field(5, &attr, &mut n); + assert!(load(&model(&n, &[])).is_err()); + let mut g = Vec::new(); + len_field(15, b"sparse", &mut g); + let mut bad = model(&node(), &[]); + len_field(7, &g, &mut bad); + assert!(load(&bad).is_err()); +} +#[test] +fn overflowing_and_cross_boundary_varints_long_strings_and_bad_shapes_are_errors() { + assert!(load(&[0x08, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x02]).is_err()); + let mut t = Vec::new(); + len_field(1, &[0x80], &mut t); + varint_field(2, 1, &mut t); + len_field(8, b"w", &mut t); + assert!(load(&model(&node(), &[t])).is_err()); + let mut n = node(); + len_field(1, &vec![b'x'; 4097], &mut n); + assert!(load(&model(&n, &[])).is_err()); + let mut n = node(); + len_field(1, &[0xff], &mut n); + assert!(load(&model(&n, &[])).is_err()); + let mut t = Vec::new(); + varint_field(1, u64::MAX, &mut t); + varint_field(1, 8, &mut t); + varint_field(2, 1, &mut t); + len_field(8, b"w", &mut t); + assert!(load(&model(&node(), &[t])).is_err()); + let mut bad = model(&node(), &[]); + let mut ops = Vec::new(); + varint_field(2, 999, &mut ops); + len_field(8, &ops, &mut bad); + assert!(load(&bad).is_err()); +} +proptest::proptest! { + #[test] + fn arbitrary_graph_bytes_do_not_panic(bytes in proptest::collection::vec(proptest::num::u8::ANY,0..4096)) { let _=load(&bytes); } +} diff --git a/crates/faces/src/hash.rs b/crates/faces/src/hash.rs new file mode 100644 index 00000000..50f83be3 --- /dev/null +++ b/crates/faces/src/hash.rs @@ -0,0 +1,46 @@ +//! SHA-256 of a model file, streamed (a model can be hundreds of megabytes). + +use std::io::{self, Read}; +use std::path::Path; + +use sha2::{Digest, Sha256}; + +/// Lowercase hex SHA-256 of everything `r` yields. +pub fn sha256_reader(mut r: R) -> io::Result { + let mut h = Sha256::new(); + let mut buf = vec![0u8; 1 << 20]; + loop { + let n = r.read(&mut buf)?; + if n == 0 { + break; + } + h.update(buf.get(..n).unwrap_or(&[])); + } + Ok(hex(&h.finalize())) +} + +/// Lowercase hex SHA-256 of a file. +pub fn sha256_file(path: &Path) -> io::Result { + sha256_reader(io::BufReader::new(std::fs::File::open(path)?)) +} + +/// Lowercase hex SHA-256 of `bytes`. +pub fn sha256_hex(bytes: &[u8]) -> String { + hex(&Sha256::digest(bytes)) +} + +fn hex(bytes: &[u8]) -> String { + bytes.iter().map(|b| format!("{b:02x}")).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn known_digests() { + assert_eq!(sha256_hex(b""), "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"); + assert_eq!(sha256_hex(b"abc"), "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"); + assert_eq!(sha256_reader(&b"abc"[..]).unwrap(), sha256_hex(b"abc")); + } +} diff --git a/crates/faces/src/known.rs b/crates/faces/src/known.rs new file mode 100644 index 00000000..f1ff541c --- /dev/null +++ b/crates/faces/src/known.rs @@ -0,0 +1,223 @@ +//! The models LightCraft recognises by their SHA-256, with what is honestly known about each. +//! +//! The licence and provenance texts here are what the user reads before enabling a model, so they say what +//! is *not* known too. Weights are never part of LightCraft: every model is an opt-in the user installs. + +use crate::manifest::{Colour, Commercial, InputSpec, Licence, ModelManifest, OutputSpec, Resize, Role, Thresholds}; + +/// Detector output decoders built into LightCraft. +pub const DECODERS: &[&str] = &["yunet-v2"]; + +/// The id of the face detector LightCraft runs (see [`yunet`]). +pub const YUNET_ID: &str = "yunet-2023mar"; +pub const YUNET_SHA256: &str = "8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4"; +pub const SFACE_SHA256: &str = "0ba9fbfa01b5270c96627c4ef784da859931e02f04419c829e83484087c34e79"; +pub const AURAFACE_SHA256: &str = "a7933ea5330113b01c9b60351d8f4c33003f145d8470ac5f0e52ee2effe25c60"; + +fn licence(name: &str, commercial: Commercial, url: &str, notice: &str) -> Licence { + Licence { name: name.into(), commercial, url: Some(url.into()), notice: notice.into() } +} + +/// YuNet 2023mar (OpenCV Zoo): a small face detector. +pub fn yunet() -> ModelManifest { + ModelManifest { + id: YUNET_ID.into(), + name: "YuNet (face detector)".into(), + version: "2023mar".into(), + role: Role::Detector, + licence: licence( + "MIT", + Commercial::Yes, + "https://github.com/opencv/opencv_zoo/blob/main/models/face_detection_yunet/LICENSE", + "MIT licence, copyright Shiqi Yu. A separate opt-in download.", + ), + source: Some("https://github.com/opencv/opencv_zoo/tree/main/models/face_detection_yunet".into()), + sha256: Some(YUNET_SHA256.into()), + size_bytes: Some(232_589), + provenance: + "Trained on WIDER FACE, whose terms are not settled by the weights' MIT licence. A detector: it finds faces, it does not identify anyone." + .into(), + input: InputSpec { width: 640, height: 640, colour: Colour::Bgr, mean: [0.0; 3], std: [1.0; 3], resize: Resize::Letterbox }, + output: OutputSpec::Detector { decoder: "yunet-v2".into() }, + thresholds: Thresholds { match_cosine: None, score: Some(0.6), nms_iou: Some(0.3) }, + speed: None, + } +} + +/// SFace 2021dec (OpenCV Zoo). Its match threshold is a starting point: on 755 named faces (busts, matched across +/// shots more than five seconds apart) 0.55 gave 97.5% right suggestions for 88% of faces; `faces.evaluate` checks it on +/// your own photos. +pub fn sface() -> ModelManifest { + ModelManifest { + id: "sface-2021dec".into(), + name: "SFace (face recogniser)".into(), + version: "2021dec".into(), + role: Role::Embedder, + licence: licence( + "Apache-2.0 (as labelled by OpenCV Zoo)", + Commercial::Unknown, + "https://github.com/opencv/opencv_zoo/tree/main/models/face_recognition_sface", + "Labelled Apache-2.0, but what it was trained on is not documented, and two questions about commercial use (opencv_zoo issues 313 and 318) are unanswered. Fine to try for yourself; do not redistribute.", + ), + source: Some("https://github.com/opencv/opencv_zoo/tree/main/models/face_recognition_sface".into()), + sha256: Some(SFACE_SHA256.into()), + size_bytes: Some(38_696_353), + provenance: "Undocumented. The original SFace repository mentions CASIA-WebFace, VGGFace2 and MS1MV2.".into(), + input: InputSpec { width: 112, height: 112, colour: Colour::Rgb, mean: [0.0; 3], std: [1.0; 3], resize: Resize::Stretch }, + output: OutputSpec::Embedding { dim: 128 }, + thresholds: Thresholds { match_cosine: Some(0.55), ..Thresholds::default() }, + speed: Some(6.9), + } +} + +/// AuraFace v1 `glintr100` (fal.ai). Starting-point threshold: on the same 755 faces 0.40 gave 96.6% right suggestions +/// for 56% of faces (0.30: 97.3% for 85%). +pub fn auraface() -> ModelManifest { + ModelManifest { + id: "auraface-v1".into(), + name: "AuraFace v1 (face recogniser)".into(), + version: "1".into(), + role: Role::Embedder, + licence: licence( + "Apache-2.0", + Commercial::Yes, + "https://huggingface.co/fal/AuraFace-v1", + "Apache-2.0. Its training data is described only as a commercial dataset, with no dataset named and no consent statement, and it covers some ethnicities less well. About 261 MB, and slow on a CPU.", + ), + source: Some("https://huggingface.co/fal/AuraFace-v1".into()), + sha256: Some(AURAFACE_SHA256.into()), + size_bytes: Some(260_694_151), + provenance: "Undisclosed: \"a commercial dataset comprising face images from various sources\".".into(), + input: InputSpec::default(), + output: OutputSpec::Embedding { dim: 512 }, + thresholds: Thresholds { match_cosine: Some(0.40), ..Thresholds::default() }, + speed: Some(1.0), + } +} + +/// Every model we know by hash. +pub fn all() -> Vec { + vec![yunet(), sface(), auraface()] +} + +/// The known model whose file has this SHA-256 (lowercase hex). +pub fn lookup(sha256: &str) -> Option { + all().into_iter().find(|m| m.sha256.as_deref() == Some(sha256)) +} + +/// Where the user's "Download" button fetches a model from. The URL is pinned to a commit of the model's own +/// repository, and what arrives is checked against the manifest's size and SHA-256 before anything uses it. +#[derive(Clone, Debug, PartialEq)] +pub struct Download { + pub id: String, + pub url: String, + /// The name the file is kept under while it waits for the user to accept its terms. + pub file_name: String, + pub size_bytes: u64, + pub sha256: String, +} + +/// Pinned download addresses by model id. Only models whose terms allow us to point at them are here: a model +/// that is marked non-commercial is never offered for download (the user brings the file themselves). +const SOURCES: &[(&str, &str, &str)] = &[ + ( + YUNET_ID, + "https://github.com/opencv/opencv_zoo/raw/25f423d0e04c31a17254620e58febd7386da523b/models/face_detection_yunet/face_detection_yunet_2023mar.onnx", + "face_detection_yunet_2023mar.onnx", + ), + ( + "sface-2021dec", + "https://github.com/opencv/opencv_zoo/raw/25f423d0e04c31a17254620e58febd7386da523b/models/face_recognition_sface/face_recognition_sface_2021dec.onnx", + "face_recognition_sface_2021dec.onnx", + ), + ("auraface-v1", "https://huggingface.co/fal/AuraFace-v1/resolve/af6d057c9b0ec4071d4c49c80e3539258798b609/glintr100.onnx", "glintr100.onnx"), +]; + +/// How to download the model with this id, if LightCraft offers to. +pub fn download(id: &str) -> Option { + let (_, url, file_name) = SOURCES.iter().find(|(i, _, _)| *i == id)?; + let m = all().into_iter().find(|m| m.id == id)?; + if m.licence.commercial == Commercial::No { + return None; + } + Some(Download { id: m.id, url: (*url).into(), file_name: (*file_name).into(), size_bytes: m.size_bytes?, sha256: m.sha256? }) +} + +/// "github.com" for a download address: what the user is told the file comes from. +pub fn host(url: &str) -> &str { + url.strip_prefix("https://").and_then(|r| r.split('/').next()).unwrap_or("") +} + +impl Download { + /// The address without the file name (the downloader's "mirror": it fetches `/`). + pub fn base(&self) -> &str { + self.url.strip_suffix(self.file_name.as_str()).map_or(self.url.as_str(), |b| b.trim_end_matches('/')) + } + + /// The site the file comes from, as the user is told. + pub fn host(&self) -> &str { + host(&self.url) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::manifest::validate; + + #[test] + fn known_models_are_valid_and_unique() { + let all = all(); + for m in &all { + assert_eq!(validate(m), Ok(()), "{}", m.id); + } + let mut ids: Vec<_> = all.iter().map(|m| &m.id).collect(); + ids.sort(); + ids.dedup(); + assert_eq!(ids.len(), all.len(), "unique ids"); + let mut hashes: Vec<_> = all.iter().filter_map(|m| m.sha256.as_deref()).collect(); + hashes.sort(); + hashes.dedup(); + assert_eq!(hashes.len(), all.len(), "unique hashes"); + } + + #[test] + fn lookup_finds_by_hash() { + assert_eq!(lookup(AURAFACE_SHA256).map(|m| m.id), Some("auraface-v1".to_string())); + assert_eq!(lookup(YUNET_SHA256).map(|m| m.role), Some(Role::Detector)); + assert!(lookup(&"0".repeat(64)).is_none()); + assert!(lookup("").is_none()); + } + + #[test] + fn downloads_are_pinned_checked_and_only_for_models_we_may_point_at() { + for m in all() { + let Some(d) = download(&m.id) else { continue }; + assert!(d.url.starts_with("https://"), "{}: https only", d.id); + assert!(!d.url.contains("/main/") && !d.url.contains("/master/"), "{}: pinned to a commit, not a branch", d.id); + assert!(d.file_name.ends_with(".onnx") && !d.file_name.contains(['/', '\\']), "{}", d.file_name); + assert!(d.url.ends_with(&format!("/{}", d.file_name)), "{}: the address ends in the file's name", d.id); + assert!(d.base().starts_with("https://") && !d.base().ends_with('/') && !d.base().ends_with(".onnx"), "{}", d.base()); + assert_eq!((Some(d.size_bytes), Some(d.sha256.as_str())), (m.size_bytes, m.sha256.as_deref()), "{}", d.id); + assert!(m.licence.commercial != Commercial::No, "{} is non-commercial and must not be offered", d.id); + } + assert_eq!(download("yunet-2023mar").map(|d| d.host().to_string()).as_deref(), Some("github.com")); + assert_eq!(download("sface-2021dec").map(|d| host(&d.url).to_string()).as_deref(), Some("github.com")); + assert_eq!(download("auraface-v1").map(|d| host(&d.url).to_string()).as_deref(), Some("huggingface.co")); + // unknown ids and path tricks have nothing to download + for id in ["nope", "", "../sface-2021dec", "SFACE-2021DEC"] { + assert!(download(id).is_none(), "{id}"); + } + // every address in the table belongs to a known model + for (id, _, _) in SOURCES { + assert!(all().iter().any(|m| m.id == *id), "{id}"); + } + } + + #[test] + fn unresolved_models_say_so() { + assert_eq!(sface().licence.commercial, Commercial::Unknown); + assert!(sface().provenance.to_lowercase().contains("undocumented")); + assert!(auraface().provenance.to_lowercase().contains("undisclosed")); + } +} diff --git a/crates/faces/src/lib.rs b/crates/faces/src/lib.rs new file mode 100644 index 00000000..b4b7066d --- /dev/null +++ b/crates/faces/src/lib.rs @@ -0,0 +1,26 @@ +//! Face models for LightCraft: what a model is ([`ModelManifest`]), whether to trust a description of one +//! ([`manifest::validate`]), the models we know by hash ([`known`]), and a probe that reads an unknown +//! `.onnx` file's input and output shapes without running it ([`onnx`]), so installing a model of your +//! own can be one drop and one licence prompt. +//! +//! Everything here treats model files and manifests as hostile input: sizes are capped, numbers must be +//! finite and sane, and a malformed file is an error, never a panic. Detection and recognition run +//! on the same checked, single-thread pure-Rust CPU interpreter. +#![forbid(unsafe_code)] +#![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)] + +pub mod align; +pub mod catalog; +pub mod graph; +pub mod hash; +pub mod known; +pub mod manifest; +pub mod matching; +pub mod net; +pub mod onnx; +pub mod runtime; +pub mod suggest; +pub mod synthetic; +pub mod yunet; + +pub use manifest::{Colour, Commercial, InputSpec, Licence, ManifestError, ModelManifest, OutputSpec, Resize, Role, Thresholds}; diff --git a/crates/faces/src/manifest.rs b/crates/faces/src/manifest.rs new file mode 100644 index 00000000..d3d99501 --- /dev/null +++ b/crates/faces/src/manifest.rs @@ -0,0 +1,386 @@ +//! A face model's manifest: what the model is, what it needs and who may use it, as plain data. +//! +//! A manifest comes from the known-model table, from a suggestion made for a file the user dropped in, or +//! from a `face-model.json` next to a model. All of those are untrusted: [`validate`] must pass before a +//! manifest is used, and [`parse`] applies it after reading JSON of bounded size. + +use serde::{Deserialize, Serialize}; + +/// Largest `face-model.json` we read. +pub const MAX_MANIFEST_BYTES: usize = 64 * 1024; +/// Largest model file we accept (bytes). +pub const MAX_MODEL_BYTES: u64 = 512 * 1024 * 1024; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum Role { + /// Finds faces (boxes, and landmarks when the model has them). + Detector, + /// Turns an aligned face into a vector whose distance says whether two faces are the same person. + Embedder, +} + +/// Whether the weights may be used commercially. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum Commercial { + Yes, + /// Research or personal use only. Never bundled or redistributed by LightCraft. + No, + /// The licence or the training data is unclear: the user is told so before installing. + Unknown, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub struct Licence { + /// SPDX id or a short name ("MIT", "Apache-2.0", "InsightFace non-commercial research"). + pub name: String, + pub commercial: Commercial, + pub url: Option, + /// Plain-language terms shown before the model is enabled. + pub notice: String, +} + +impl Default for Licence { + fn default() -> Self { + Self { name: String::new(), commercial: Commercial::Unknown, url: None, notice: String::new() } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum Colour { + Rgb, + Bgr, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum Resize { + /// Stretch the crop to the input size (aligned face crops). + Stretch, + /// Keep the aspect ratio and pad (whole photos for a detector). + Letterbox, +} + +/// How an image becomes the model's input: NCHW float, `(pixel − mean) / std` per channel, pixel in 0..255. +#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub struct InputSpec { + pub width: u32, + pub height: u32, + pub colour: Colour, + pub mean: [f32; 3], + pub std: [f32; 3], + pub resize: Resize, +} + +impl Default for InputSpec { + /// ArcFace-style: 112 × 112, RGB, `(x − 127.5) / 127.5`. + fn default() -> Self { + Self { width: 112, height: 112, colour: Colour::Rgb, mean: [127.5; 3], std: [127.5; 3], resize: Resize::Stretch } + } +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "camelCase")] +pub enum OutputSpec { + /// One vector of `dim` numbers per face. + Embedding { dim: u32 }, + /// A detector whose outputs a built-in decoder understands (`"yunet-v2"`). + Detector { decoder: String }, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub struct Thresholds { + /// Embedders: cosine similarity at or above which two faces are suggested as a match. + pub match_cosine: Option, + /// Detectors: lowest detection score kept. + pub score: Option, + /// Detectors: overlap above which two boxes are one face. + pub nms_iou: Option, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ModelManifest { + /// Lowercase letters, digits, `.` `_` `-`. + pub id: String, + pub name: String, + pub version: String, + pub role: Role, + #[serde(default)] + pub licence: Licence, + /// Where the weights come from (an `https://` page or file), for the user's information. + #[serde(default)] + pub source: Option, + /// Lowercase hex SHA-256 of the `.onnx` file, when known. + #[serde(default)] + pub sha256: Option, + #[serde(default)] + pub size_bytes: Option, + /// What the model was trained on, in plain words; says "undisclosed" when it is. + #[serde(default)] + pub provenance: String, + #[serde(default)] + pub input: InputSpec, + pub output: OutputSpec, + #[serde(default)] + pub thresholds: Thresholds, + /// How fast it is as a multiple of a ResNet-100 model's speed (`1.0`): `2.0` is twice as fast. Shown instead of + /// timings, which depend on the computer; leave it out when you have not measured it. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub speed: Option, +} + +/// The reference model the `speed` ratios are measured against. +pub const SPEED_REFERENCE: &str = "a ResNet-100 model"; + +/// A speed ratio in plain words, as times faster than the reference ("4.6× faster than a ResNet-100 model"). +pub fn describe_speed(speed: f32) -> String { + if !(speed.is_finite() && speed > 0.0) { + return String::new(); + } + let times = |x: f32| { + let s = format!("{x:.1}"); + s.strip_suffix(".0").map_or(s.clone(), str::to_string) + }; + if speed >= 1.15 { + format!("{}× faster than {SPEED_REFERENCE}", times(speed)) + } else if speed > 0.87 { + format!("Same speed as {SPEED_REFERENCE}") + } else { + format!("{}× slower than {SPEED_REFERENCE}", times(1.0 / speed)) + } +} + +#[derive(Debug, PartialEq, thiserror::Error)] +pub enum ManifestError { + #[error("`{0}` is not valid: {1}")] + Field(&'static str, String), + #[error("not a face-model manifest: {0}")] + Parse(String), + #[error("a manifest may be at most {MAX_MANIFEST_BYTES} bytes")] + TooLarge, +} + +fn bad(field: &'static str, why: impl Into) -> Result { + Err(ManifestError::Field(field, why.into())) +} + +/// A model id is a folder name: 1 to 64 characters, starting with a lowercase letter or digit, then +/// lowercase letters, digits, `.`, `_` or `-`. So it is never `.`, `..`, hidden, or a path. +pub fn valid_id(id: &str) -> bool { + let mut bytes = id.bytes(); + let first_ok = bytes.next().is_some_and(|b| b.is_ascii_lowercase() || b.is_ascii_digit()); + first_ok && id.len() <= 64 && bytes.all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || matches!(b, b'.' | b'_' | b'-')) +} + +fn text(field: &'static str, s: &str, min: usize, max: usize) -> Result<(), ManifestError> { + let n = s.chars().count(); + if n < min || n > max { + return bad(field, format!("must be {min} to {max} characters")); + } + if s.chars().any(|c| c.is_control() && c != '\n') { + return bad(field, "must not contain control characters"); + } + Ok(()) +} + +fn url(field: &'static str, u: &Option) -> Result<(), ManifestError> { + match u { + None => Ok(()), + Some(u) if u.len() <= 500 && u.starts_with("https://") && !u.chars().any(|c| c.is_control() || c.is_whitespace()) => Ok(()), + Some(_) => bad(field, "must be an https:// address of at most 500 characters"), + } +} + +fn ratio(field: &'static str, v: Option, lo: f32, hi: f32) -> Result<(), ManifestError> { + match v { + Some(v) if !(v.is_finite() && (lo..=hi).contains(&v)) => bad(field, format!("must be a number from {lo} to {hi}")), + _ => Ok(()), + } +} + +/// Check a manifest from any source. Everything the rest of LightCraft relies on (sizes it will allocate, +/// numbers it will divide by, text it will show) is bounded here. +pub fn validate(m: &ModelManifest) -> Result<(), ManifestError> { + if !valid_id(&m.id) { + return bad("id", "must be 1 to 64 characters: lowercase letters and digits, then also '.', '_' or '-'"); + } + text("name", &m.name, 1, 120)?; + text("version", &m.version, 1, 40)?; + text("licence.name", &m.licence.name, 0, 80)?; + text("licence.notice", &m.licence.notice, 0, 800)?; + text("provenance", &m.provenance, 0, 800)?; + url("licence.url", &m.licence.url)?; + url("source", &m.source)?; + if let Some(h) = &m.sha256 + && !(h.len() == 64 && h.bytes().all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b))) + { + return bad("sha256", "must be 64 lowercase hexadecimal digits"); + } + if m.size_bytes.is_some_and(|s| s == 0 || s > MAX_MODEL_BYTES) { + return bad("sizeBytes", format!("must be 1 to {MAX_MODEL_BYTES}")); + } + let i = &m.input; + if !(16..=4096).contains(&i.width) || !(16..=4096).contains(&i.height) { + return bad("input.width/height", "must be 16 to 4096"); + } + if i.mean.iter().any(|v| !v.is_finite() || v.abs() > 1000.0) { + return bad("input.mean", "must be finite numbers within ±1000"); + } + if i.std.iter().any(|v| !v.is_finite() || v.abs() < 1e-3 || v.abs() > 1000.0) { + return bad("input.std", "must be finite, not zero, within ±1000"); + } + match &m.output { + OutputSpec::Embedding { dim } => { + if m.role != Role::Embedder { + return bad("output", "an embedding output needs the embedder role"); + } + if !(1..=4096).contains(dim) { + return bad("output.dim", "must be 1 to 4096"); + } + } + OutputSpec::Detector { decoder } => { + if m.role != Role::Detector { + return bad("output", "a detector output needs the detector role"); + } + if !crate::known::DECODERS.contains(&decoder.as_str()) { + return bad("output.decoder", format!("unknown decoder `{}`", decoder.chars().take(40).collect::())); + } + } + } + ratio("thresholds.matchCosine", m.thresholds.match_cosine, -1.0, 1.0)?; + ratio("thresholds.score", m.thresholds.score, 0.0, 1.0)?; + ratio("thresholds.nmsIou", m.thresholds.nms_iou, 0.0, 1.0)?; + ratio("speed", m.speed, 0.01, 1000.0)?; + Ok(()) +} + +/// Read a `face-model.json` (at most [`MAX_MANIFEST_BYTES`]) and validate it. +pub fn parse(bytes: &[u8]) -> Result { + if bytes.len() > MAX_MANIFEST_BYTES { + return Err(ManifestError::TooLarge); + } + let m: ModelManifest = serde_json::from_slice(bytes).map_err(|e| ManifestError::Parse(e.to_string()))?; + validate(&m)?; + Ok(m) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn good() -> ModelManifest { + ModelManifest { + id: "example-r50".into(), + name: "Example R50".into(), + version: "1".into(), + role: Role::Embedder, + licence: Licence { + name: "MIT".into(), + commercial: Commercial::Yes, + url: Some("https://example.org/licence".into()), + notice: String::new(), + }, + source: Some("https://example.org/model.onnx".into()), + sha256: Some("a".repeat(64)), + size_bytes: Some(1234), + provenance: "synthetic".into(), + input: InputSpec::default(), + output: OutputSpec::Embedding { dim: 512 }, + thresholds: Thresholds { match_cosine: Some(0.4), ..Default::default() }, + speed: Some(1.7), + } + } + + #[test] + fn speed_reads_as_times_faster_than_the_reference() { + assert_eq!(describe_speed(7.0), "7× faster than a ResNet-100 model"); + assert_eq!(describe_speed(4.6), "4.6× faster than a ResNet-100 model"); + assert_eq!(describe_speed(1.7), "1.7× faster than a ResNet-100 model"); + assert_eq!(describe_speed(1.0), "Same speed as a ResNet-100 model"); + assert_eq!(describe_speed(0.5), "2× slower than a ResNet-100 model"); + for odd in [0.0, -1.0, f32::NAN, f32::INFINITY] { + assert_eq!(describe_speed(odd), "", "{odd}"); + } + } + + #[test] + fn a_good_manifest_passes_and_round_trips() { + let m = good(); + assert_eq!(validate(&m), Ok(())); + let json = serde_json::to_vec(&m).unwrap(); + assert_eq!(parse(&json).unwrap(), m); + } + + #[test] + fn every_field_is_bounded() { + let cases: Vec<(&str, Box)> = vec![ + ("empty id", Box::new(|m| m.id = String::new())), + ("uppercase id", Box::new(|m| m.id = "Bad".into())), + ("path in id", Box::new(|m| m.id = "../etc".into())), + ("dot id", Box::new(|m| m.id = ".".into())), + ("dotdot id", Box::new(|m| m.id = "..".into())), + ("hidden id", Box::new(|m| m.id = ".hidden".into())), + ("slash id", Box::new(|m| m.id = "a/b".into())), + ("backslash id", Box::new(|m| m.id = r"a\b".into())), + ("long id", Box::new(|m| m.id = "a".repeat(65))), + ("empty name", Box::new(|m| m.name = String::new())), + ("control in name", Box::new(|m| m.name = "a\u{7}b".into())), + ("huge notice", Box::new(|m| m.licence.notice = "x".repeat(801))), + ("http source", Box::new(|m| m.source = Some("http://example.org/m.onnx".into()))), + ("file source", Box::new(|m| m.source = Some("file:///etc/passwd".into()))), + ("space in url", Box::new(|m| m.source = Some("https://a b".into()))), + ("short sha", Box::new(|m| m.sha256 = Some("abc".into()))), + ("uppercase sha", Box::new(|m| m.sha256 = Some("A".repeat(64)))), + ("zero size", Box::new(|m| m.size_bytes = Some(0))), + ("huge size", Box::new(|m| m.size_bytes = Some(u64::MAX))), + ("tiny input", Box::new(|m| m.input.width = 4)), + ("huge input", Box::new(|m| m.input.height = 1 << 20)), + ("nan mean", Box::new(|m| m.input.mean[1] = f32::NAN)), + ("inf mean", Box::new(|m| m.input.mean[0] = f32::INFINITY)), + ("zero std", Box::new(|m| m.input.std[2] = 0.0)), + ("nan std", Box::new(|m| m.input.std[0] = f32::NAN)), + ("zero dim", Box::new(|m| m.output = OutputSpec::Embedding { dim: 0 })), + ("huge dim", Box::new(|m| m.output = OutputSpec::Embedding { dim: 10_000_000 })), + ("role mismatch", Box::new(|m| m.role = Role::Detector)), + ( + "unknown decoder", + Box::new(|m| { + m.role = Role::Detector; + m.output = OutputSpec::Detector { decoder: "made-up".into() }; + }), + ), + ("nan threshold", Box::new(|m| m.thresholds.match_cosine = Some(f32::NAN))), + ("threshold range", Box::new(|m| m.thresholds.score = Some(1.5))), + ("nan speed", Box::new(|m| m.speed = Some(f32::NAN))), + ("zero speed", Box::new(|m| m.speed = Some(0.0))), + ("huge speed", Box::new(|m| m.speed = Some(1e9))), + ]; + for (what, f) in cases { + let mut m = good(); + f(&mut m); + assert!(validate(&m).is_err(), "{what} must be rejected"); + } + } + + #[test] + fn hostile_json_is_an_error() { + assert!(matches!(parse(&vec![b' '; MAX_MANIFEST_BYTES + 1]), Err(ManifestError::TooLarge))); + for bytes in [ + &b""[..], + b"null", + b"[]", + b"{", + b"{\"id\": 1}", + b"\xff\xfe\x00", + b"{\"id\":\"x\",\"name\":\"n\",\"version\":\"1\",\"role\":\"embedder\"}", + ] { + assert!(parse(bytes).is_err(), "{:?}", String::from_utf8_lossy(bytes)); + } + } +} diff --git a/crates/faces/src/matching.rs b/crates/faces/src/matching.rs new file mode 100644 index 00000000..de50ad3e --- /dev/null +++ b/crates/faces/src/matching.rs @@ -0,0 +1,128 @@ +//! Who does this face look like? Rank the people in a gallery of named faces by how close their faces are, +//! and suggest one only when the match is strong *and* clearly ahead of the runner-up. +//! +//! Embeddings are unit vectors, so closeness is the dot product (cosine similarity). A person's score is the +//! mean of their best few faces (one stray, wrongly named face cannot make a stranger look like them). The +//! result is only ever a suggestion for the user to confirm: it is better to leave a face unnamed than to +//! name it wrongly. + +use std::cmp::Ordering; + +/// How many of a person's best matching faces are averaged into their score. +pub const TOP_FACES: usize = 2; + +#[derive(Clone, Debug, PartialEq)] +pub struct Ranked { + pub name: String, + pub score: f32, + /// How many of their faces the score is based on. + pub faces: usize, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Suggestion { + pub name: String, + pub score: f32, + /// The next best person, and how far behind. + pub runner_up: Option<(String, f32)>, +} + +/// Cosine similarity of two unit vectors. `None` when their lengths differ or either has a non-finite number. +pub fn cosine(a: &[f32], b: &[f32]) -> Option { + if a.len() != b.len() || a.is_empty() { + return None; + } + let d: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum(); + d.is_finite().then_some(d) +} + +/// Rank the people in `gallery` (a name and one face's embedding per entry; a person appears once per face), +/// best first. Names are compared case-insensitively and shown as first seen. +pub fn rank(query: &[f32], gallery: &[(&str, &[f32])]) -> Vec { + let mut people: Vec<(String, String, Vec)> = Vec::new(); // (key, shown name, similarities) + for (name, emb) in gallery { + let Some(c) = cosine(query, emb) else { continue }; + let key = name.to_lowercase(); + match people.iter_mut().find(|p| p.0 == key) { + Some(p) => p.2.push(c), + None => people.push((key, (*name).to_string(), vec![c])), + } + } + let mut out: Vec = people + .into_iter() + .map(|(_, name, mut sims)| { + sims.sort_by(|a, b| b.partial_cmp(a).unwrap_or(Ordering::Equal)); + let top = sims.get(..sims.len().min(TOP_FACES)).unwrap_or(&[]); + Ranked { name, score: top.iter().sum::() / top.len().max(1) as f32, faces: sims.len() } + }) + .collect(); + out.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(Ordering::Equal).then_with(|| a.name.cmp(&b.name))); + out +} + +/// The best person if they pass both bars: a score of at least `threshold`, and at least `margin` ahead of the +/// next person (when there is one). +pub fn suggest(ranked: &[Ranked], threshold: f32, margin: f32) -> Option { + let best = ranked.first()?; + let next = ranked.get(1); + if best.score.is_nan() || best.score < threshold || next.is_some_and(|n| best.score - n.score < margin) { + return None; + } + Some(Suggestion { name: best.name.clone(), score: best.score, runner_up: next.map(|n| (n.name.clone(), n.score)) }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn unit(v: &[f32]) -> Vec { + let n = v.iter().map(|x| x * x).sum::().sqrt(); + v.iter().map(|x| x / n).collect() + } + + #[test] + fn people_are_ranked_by_their_closest_faces() { + let (a1, a2, b1) = (unit(&[1.0, 0.1, 0.0]), unit(&[0.9, 0.2, 0.0]), unit(&[0.0, 1.0, 0.2])); + let gallery: Vec<(&str, &[f32])> = vec![("Ann", &a1), ("ann", &a2), ("Bob", &b1)]; + let ranked = rank(&unit(&[1.0, 0.0, 0.0]), &gallery); + assert_eq!(ranked.iter().map(|r| r.name.as_str()).collect::>(), vec!["Ann", "Bob"], "case-insensitive, shown as first seen"); + assert_eq!(ranked[0].faces, 2); + assert!(ranked[0].score > 0.95 && ranked[1].score < 0.2); + // a stray, wrongly named face cannot carry a person: Cy has one far face and one good one + let (c_good, c_bad) = (unit(&[1.0, 0.0, 0.0]), unit(&[0.0, 0.0, 1.0])); + let g2: Vec<(&str, &[f32])> = vec![("Cy", &c_good), ("Cy", &c_bad), ("Di", &a1), ("Di", &a2)]; + let r2 = rank(&unit(&[1.0, 0.0, 0.0]), &g2); + assert_eq!(r2[0].name, "Di", "{r2:?}"); + } + + #[test] + fn a_suggestion_needs_a_strong_match_clearly_ahead() { + let r = |name: &str, score: f32| Ranked { name: name.into(), score, faces: 1 }; + // strong and clear + let s = suggest(&[r("Ann", 0.62), r("Bob", 0.20)], 0.45, 0.08).unwrap(); + assert_eq!((s.name.as_str(), s.runner_up.as_ref().map(|x| x.0.as_str())), ("Ann", Some("Bob"))); + // strong but too close to the runner-up: nobody + assert!(suggest(&[r("Ann", 0.62), r("Bob", 0.58)], 0.45, 0.08).is_none()); + // clear but not strong enough + assert!(suggest(&[r("Ann", 0.40), r("Bob", 0.10)], 0.45, 0.08).is_none()); + // a single person has no runner-up to be ahead of + assert!(suggest(&[r("Ann", 0.5)], 0.45, 0.08).is_some()); + assert!(suggest(&[], 0.45, 0.08).is_none()); + assert!(suggest(&[r("Ann", f32::NAN)], 0.45, 0.08).is_none()); + } + + #[test] + fn mismatched_or_non_finite_embeddings_are_skipped_not_panics() { + assert!(cosine(&[1.0], &[1.0, 0.0]).is_none()); + assert!(cosine(&[], &[]).is_none()); + assert!(cosine(&[f32::NAN], &[1.0]).is_none()); + let good = unit(&[1.0, 0.0]); + let bad = [f32::NAN, 0.0]; + let short = [1.0f32]; + let gallery: Vec<(&str, &[f32])> = vec![("Ann", &good), ("Nan", &bad), ("Short", &short)]; + let ranked = rank(&good, &gallery); + assert_eq!(ranked.len(), 1); + assert!(rank(&[], &gallery).is_empty()); + assert!(rank(&good, &[]).is_empty()); + } +} diff --git a/crates/faces/src/net.rs b/crates/faces/src/net.rs new file mode 100644 index 00000000..6222314f --- /dev/null +++ b/crates/faces/src/net.rs @@ -0,0 +1,570 @@ +//! A small CPU interpreter for the convolutional networks LightCraft bundles (the YuNet face detector). +//! +//! It runs checked convolution, affine/activation, broadcast, pooling, matrix and shape operators on +//! float32 tensors. Unsupported semantics fail at load time. Dense convolutions use bounded im2col and +//! single-thread GEMM; depthwise convolutions use row-wise Rust loops. Every file-derived shape is checked. +//! +//! The same checked float32 engine runs YuNet and installed SFace/AuraFace recognisers. +//! Unknown operators, attributes and shapes are rejected before recognition starts. + +use std::collections::HashMap; +use std::rc::Rc; + +mod kernels; +mod shape; + +use crate::graph::{Data, Graph, Node}; + +/// Elements of any one tensor the interpreter will allocate. +const MAX_TENSOR: usize = 64 * 1024 * 1024; + +#[derive(Debug, thiserror::Error, PartialEq)] +pub enum NetError { + #[error("this model uses something the built-in runtime does not support: {0}")] + Unsupported(String), + #[error("the model does not fit together: {0}")] + Shape(String), +} + +type Result = std::result::Result; + +fn shape_err(why: impl Into) -> Result { + Err(NetError::Shape(why.into())) +} + +fn unsupported(why: impl Into) -> Result { + Err(NetError::Unsupported(why.into())) +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Tensor { + pub shape: Vec, + pub data: Vec, +} + +impl Tensor { + pub fn new(shape: Vec, data: Vec) -> Result { + let n = shape.iter().try_fold(1usize, |a, d| a.checked_mul(*d)); + if shape.len() > 8 || shape.contains(&0) || n.is_none_or(|n| n > MAX_TENSOR) || n != Some(data.len()) || data.iter().any(|v| !v.is_finite()) { + return shape_err(format!("{} values do not fill the shape {shape:?}", data.len())); + } + Ok(Tensor { shape, data }) + } + + fn zeros(shape: Vec) -> Result { + match shape.iter().try_fold(1usize, |a, d| a.checked_mul(*d)) { + Some(n) if n > 0 && n <= MAX_TENSOR && shape.len() <= 8 && !shape.contains(&0) => { + let mut data = Vec::new(); + data.try_reserve_exact(n).map_err(|_| NetError::Shape("not enough memory for a model tensor".into()))?; + data.resize(n, 0.0); + Ok(Tensor { shape, data }) + } + _ => shape_err("a tensor would be too large"), + } + } +} + +/// A loaded network, ready to run. +pub struct Net { + graph: Graph, +} + +/// Validate semantics rather than silently ignoring unknown or mistyped attributes. +fn check(node: &Node, g: &Graph) -> Result<()> { + shape::check(node, g) +} + +impl Net { + pub fn new(graph: Graph) -> Result { + shape::check_graph(&graph)?; + for n in &graph.nodes { + check(n, &graph)?; + } + Ok(Net { graph }) + } + + /// The shape the model declares for its input (0 for a dynamic size). + pub fn input_shape(&self) -> Option<&[usize]> { + self.graph.inputs.first().map(|(_, s)| s.as_slice()) + } + + /// Infer and check every intermediate without allocating its pixel data. + /// Recognition calls this at load time with the manifest's fixed input. + pub fn validate_input(&self, shape: &[usize]) -> Result>> { + shape::infer(&self.graph, shape) + } + + /// Run the network on `input` (its single input); returns each output by name. + pub fn run(&self, input: Tensor) -> Result> { + shape::tensor(&input.shape, input.data.len(), &input.data)?; + self.validate_input(&input.shape)?; + let mut scratch = kernels::Scratch::default(); + let profile = std::env::var_os("LIGHTCRAFT_PROFILE").is_some(); + let g = &self.graph; + let Some((input_name, _)) = g.inputs.first() else { return shape_err("the model has no input") }; + // when each value is last read, so big intermediates are dropped as soon as nothing needs them + let mut last_use: HashMap<&str, usize> = HashMap::new(); + for (i, n) in g.nodes.iter().enumerate() { + for name in &n.inputs { + last_use.insert(name.as_str(), i); + } + } + for o in &g.outputs { + last_use.insert(o.as_str(), usize::MAX); + } + let mut values: HashMap> = HashMap::new(); + values.insert(input_name.clone(), Rc::new(input)); + for (i, node) in g.nodes.iter().enumerate() { + let get = |idx: usize| -> Result> { + let name = node.inputs.get(idx).ok_or_else(|| NetError::Shape(format!("{} is missing an input", node.op)))?; + if let Some(value) = values.get(name) { + return Ok(value.clone()); + } + if let Some(crate::graph::Weight { dims, data: Data::F32(data) }) = g.weights.get(name) { + return Ok(Rc::new(Tensor::new(dims.clone(), data.clone())?)); + } + values.get(name).cloned().ok_or_else(|| NetError::Shape(format!("`{name}` is read before it is computed"))) + }; + let started = web_time::Instant::now(); + let out = match node.op.as_str() { + "Conv" => kernels::conv(node, &*get(0)?, g, &mut scratch)?, + "Relu" => map(&*get(0)?, |v| v.max(0.0)), + "Sigmoid" => map(&*get(0)?, |v| 1.0 / (1.0 + (-v).exp())), + "Add" | "Sub" | "Mul" => kernels::binary(node.op.as_str(), &*get(0)?, &*get(1)?)?, + "PRelu" => kernels::prelu(node, &*get(0)?, g)?, + "BatchNormalization" => kernels::batch_norm(node, &*get(0)?, g)?, + "Gemm" => kernels::gemm(node, &*get(0)?, g)?, + "Flatten" => kernels::flatten(node, &*get(0)?)?, + "Dropout" | "Identity" => (*get(0)?).clone(), + "GlobalAveragePool" => kernels::global_average(&*get(0)?)?, + "MaxPool" => max_pool(node, &*get(0)?)?, + "Resize" => resize(node, &*get(0)?, g)?, + "Transpose" => transpose(node, &*get(0)?)?, + "Reshape" => reshape(node, &*get(0)?, g)?, + other => return unsupported(format!("the operator `{other}`")), + }; + if out.data.iter().any(|v| !v.is_finite()) { + return shape_err(format!("{} produced non-finite values", node.op)); + } + if profile { + eprintln!("[faces-cpu] layer {i} {} {:?}: {:.3} ms", node.op, out.shape, started.elapsed().as_secs_f64() * 1000.0); + } + let Some(name) = node.outputs.first() else { return shape_err("a node has no output") }; + values.insert(name.clone(), Rc::new(out)); + for name in &node.inputs { + if last_use.get(name.as_str()) == Some(&i) { + values.remove(name); + } + } + } + g.outputs + .iter() + .map(|name| { + let t = values.remove(name).ok_or_else(|| NetError::Shape(format!("output `{name}` was never computed")))?; + let t = Rc::try_unwrap(t).map_err(|_| NetError::Shape(format!("output `{name}` is still shared")))?; + Ok((name.clone(), t)) + }) + .collect() + } +} + +fn map(x: &Tensor, f: impl Fn(f32) -> f32) -> Tensor { + Tensor { shape: x.shape.clone(), data: x.data.iter().map(|v| f(*v)).collect() } +} + +/// `[1, c, h, w]` → `(c, h, w)`. +fn chw(x: &Tensor) -> Result<(usize, usize, usize)> { + match x.shape.as_slice() { + [1, c, h, w] => Ok((*c, *h, *w)), + s => unsupported(format!("a tensor of shape {s:?} (only [1, channels, height, width])")), + } +} + +/// `pads` of an ONNX node as (top, left, bottom, right). +fn pads(node: &Node) -> Result<(usize, usize, usize, usize)> { + match node.ints("pads") { + None => Ok((0, 0, 0, 0)), + Some([t, l, b, r]) => { + let u = |v: &i64| usize::try_from(*v).map_err(|_| NetError::Shape("a negative pad".into())); + Ok((u(t)?, u(l)?, u(b)?, u(r)?)) + } + Some(_) => unsupported("pads that are not [top, left, bottom, right]"), + } +} + +fn pair(node: &Node, name: &str, default: usize) -> Result<(usize, usize)> { + match node.ints(name) { + None => Ok((default, default)), + Some([a, b]) => match (usize::try_from(*a), usize::try_from(*b)) { + (Ok(a), Ok(b)) if a >= 1 && b >= 1 => Ok((a, b)), + _ => shape_err(format!("`{name}` must be positive")), + }, + Some(_) => unsupported(format!("`{name}` that is not two numbers")), + } +} + +fn conv_direct(node: &Node, x: &Tensor, g: &Graph) -> Result { + let (cin, h, w) = chw(x)?; + let weight = node.inputs.get(1).and_then(|n| g.weights.get(n)).ok_or_else(|| NetError::Shape("a Conv lost its weights".into()))?; + let Data::F32(wd) = &weight.data else { return unsupported("non-float Conv weights") }; + let [cout, cin_g, kh, kw] = weight.dims.as_slice() else { return unsupported("a Conv that is not 2-D") }; + let (cout, cin_g, kh, kw) = (*cout, *cin_g, *kh, *kw); + let groups = + usize::try_from(node.int("group").unwrap_or(1)).ok().filter(|g| *g >= 1).ok_or_else(|| NetError::Shape("a bad Conv group".into()))?; + if groups.checked_mul(cin_g) != Some(cin) || cout % groups != 0 || kh == 0 || kw == 0 { + return shape_err("the Conv weights do not match its input"); + } + let bias: Option<&[f32]> = match node.inputs.get(2).filter(|b| !b.is_empty()).and_then(|b| g.weights.get(b)) { + Some(b) => match &b.data { + Data::F32(v) if v.len() == cout => Some(v), + _ => return shape_err("the Conv bias does not match"), + }, + None => None, + }; + let (sh, sw) = pair(node, "strides", 1)?; + let (pt, pl, pb, pr) = pads(node)?; + let ph = h.checked_add(pt).and_then(|v| v.checked_add(pb)).ok_or_else(|| NetError::Shape("Conv padding overflow".into()))?; + let pw = w.checked_add(pl).and_then(|v| v.checked_add(pr)).ok_or_else(|| NetError::Shape("Conv padding overflow".into()))?; + if ph < kh || pw < kw { + return shape_err("the Conv kernel is larger than its input"); + } + let (oh, ow) = ((ph - kh) / sh + 1, (pw - kw) / sw + 1); + let mut out = Tensor::zeros(vec![1, cout, oh, ow])?; + let plane = oh * ow; + let in_planes: Vec<&[f32]> = x.data.chunks_exact((h * w).max(1)).collect(); + let cout_g = cout / groups; + for (co, out_plane) in out.data.chunks_exact_mut(plane.max(1)).enumerate() { + if let Some(b) = bias.and_then(|b| b.get(co)) { + out_plane.iter_mut().for_each(|o| *o = *b); + } + let group = co / cout_g; + for ci_l in 0..cin_g { + let Some(in_plane) = in_planes.get(group * cin_g + ci_l) else { return shape_err("a Conv input channel is missing") }; + let base = (co * cin_g + ci_l) * kh * kw; + let Some(kernel) = wd.get(base..base + kh * kw) else { return shape_err("the Conv weights are too short") }; + for (k, wv) in kernel.iter().enumerate() { + if *wv == 0.0 { + continue; + } + let (ky, kx) = (k / kw, k % kw); + for oy in 0..oh { + // input row for this output row and kernel row (in padded coordinates minus the pad) + let Some(iy) = (oy * sh + ky).checked_sub(pt).filter(|iy| *iy < h) else { continue }; + let in_row = in_plane.get(iy * w..(iy + 1) * w).unwrap_or(&[]); + let out_row = out_plane.get_mut(oy * ow..(oy + 1) * ow).unwrap_or(&mut []); + if sw == 1 { + // ox runs over outputs whose input column ox + kx - pl lies inside the row + let lo = pl.saturating_sub(kx); + let hi = ow.min((w + pl).saturating_sub(kx)); + if lo < hi { + let src = in_row.get(lo + kx - pl..hi + kx - pl).unwrap_or(&[]); + let dst = out_row.get_mut(lo..hi).unwrap_or(&mut []); + dst.iter_mut().zip(src).for_each(|(o, i)| *o += wv * i); + } + } else { + for (ox, o) in out_row.iter_mut().enumerate() { + if let Some(ix) = (ox * sw + kx).checked_sub(pl).filter(|ix| *ix < w) + && let Some(i) = in_row.get(ix) + { + *o += wv * i; + } + } + } + } + } + } + } + Ok(out) +} + +fn max_pool(node: &Node, x: &Tensor) -> Result { + let (c, h, w) = chw(x)?; + let (kh, kw) = pair(node, "kernel_shape", 1)?; + let (sh, sw) = pair(node, "strides", 1)?; + let (pt, pl, pb, pr) = pads(node)?; + if (pt, pl, pb, pr) != (0, 0, 0, 0) { + return unsupported("a padded MaxPool"); + } + if h < kh || w < kw { + return shape_err("the MaxPool window is larger than its input"); + } + let (oh, ow) = ((h - kh) / sh + 1, (w - kw) / sw + 1); + let mut out = Tensor::zeros(vec![1, c, oh, ow])?; + for (plane_in, plane_out) in x.data.chunks_exact((h * w).max(1)).zip(out.data.chunks_exact_mut((oh * ow).max(1))) { + for oy in 0..oh { + for ox in 0..ow { + let mut m = f32::NEG_INFINITY; + for ky in 0..kh { + let row = plane_in.get((oy * sh + ky) * w..(oy * sh + ky + 1) * w).unwrap_or(&[]); + for kx in 0..kw { + if let Some(v) = row.get(ox * sw + kx) { + m = m.max(*v); + } + } + } + if let Some(o) = plane_out.get_mut(oy * ow + ox) { + *o = m; + } + } + } + } + Ok(out) +} + +/// Nearest-neighbour upscale by whole numbers (YuNet's feature pyramid doubles each map). +fn resize(node: &Node, x: &Tensor, g: &Graph) -> Result { + let (c, h, w) = chw(x)?; + let scales = match node.inputs.get(2).and_then(|n| g.weights.get(n)).map(|w| &w.data) { + Some(Data::F32(s)) if s.len() == 4 => s, + _ => return unsupported("a Resize whose scales are not four floats"), + }; + let whole = |v: f32| ((1.0..=16.0).contains(&v) && v.fract() == 0.0).then_some(v as usize); + let (Some(sh), Some(sw)) = (whole(scales.get(2).copied().unwrap_or(0.0)), whole(scales.get(3).copied().unwrap_or(0.0))) else { + return unsupported("a Resize by a scale that is not a whole number from 1 to 16"); + }; + if scales.first() != Some(&1.0) || scales.get(1) != Some(&1.0) { + return unsupported("a Resize that scales channels"); + } + let (oh, ow) = (h * sh, w * sw); + let mut out = Tensor::zeros(vec![1, c, oh, ow])?; + for (plane_in, plane_out) in x.data.chunks_exact((h * w).max(1)).zip(out.data.chunks_exact_mut((oh * ow).max(1))) { + for (oy, row_out) in plane_out.chunks_exact_mut(ow.max(1)).enumerate() { + let row_in = plane_in.get((oy / sh) * w..(oy / sh + 1) * w).unwrap_or(&[]); + for (ox, o) in row_out.iter_mut().enumerate() { + *o = row_in.get(ox / sw).copied().unwrap_or(0.0); + } + } + } + Ok(out) +} + +fn transpose(node: &Node, x: &Tensor) -> Result { + let default_perm: Vec = (0..x.shape.len() as i64).rev().collect(); + let perm: Vec = node + .ints("perm") + .unwrap_or(&default_perm) + .iter() + .map(|p| usize::try_from(*p).map_err(|_| NetError::Shape("a negative permutation".into()))) + .collect::>()?; + let rank = x.shape.len(); + let mut seen = vec![false; rank]; + if perm.len() != rank || perm.iter().any(|p| *p >= rank || std::mem::replace(seen.get_mut(*p).unwrap_or(&mut true), true)) { + return shape_err("the Transpose permutation does not match the tensor"); + } + let out_shape: Vec = perm.iter().filter_map(|p| x.shape.get(*p).copied()).collect(); + // strides of the input, in elements + let mut in_strides = vec![1usize; rank]; + for i in (0..rank.saturating_sub(1)).rev() { + let v = in_strides.get(i + 1).copied().unwrap_or(1) * x.shape.get(i + 1).copied().unwrap_or(1); + *in_strides.get_mut(i).ok_or_else(|| NetError::Shape("missing transpose stride".into()))? = v; + } + let strides_for_out: Vec = perm.iter().filter_map(|p| in_strides.get(*p).copied()).collect(); + let mut out = Tensor::zeros(out_shape.clone())?; + let mut counter = vec![0usize; rank]; + let mut src = 0usize; + for o in out.data.iter_mut() { + *o = x.data.get(src).copied().unwrap_or(0.0); + // advance the multi-index of the output and the matching input offset + for d in (0..rank).rev() { + let (Some(c), Some(len), Some(stride)) = (counter.get_mut(d), out_shape.get(d), strides_for_out.get(d)) else { break }; + *c += 1; + src += stride; + if *c < *len { + break; + } + src -= stride * *len; + *c = 0; + } + } + Ok(out) +} + +fn reshape(node: &Node, x: &Tensor, g: &Graph) -> Result { + let Some(Data::I64(dims)) = node.inputs.get(1).and_then(|n| g.weights.get(n)).map(|w| &w.data) else { + return unsupported("a Reshape with a computed shape"); + }; + let total = x.data.len(); + let mut shape: Vec = Vec::with_capacity(dims.len()); + let mut infer = None; + for (i, d) in dims.iter().enumerate() { + match *d { + -1 if infer.is_none() => { + infer = Some(i); + shape.push(1); + } + 0 => shape.push(x.shape.get(i).copied().ok_or_else(|| NetError::Shape("a Reshape copies a missing dimension".into()))?), + d if d > 0 => shape.push(usize::try_from(d).map_err(|_| NetError::Shape("a Reshape dimension is too large".into()))?), + _ => return shape_err("a bad Reshape dimension"), + } + } + let known = shape.iter().try_fold(1usize, |a, d| a.checked_mul(*d)).filter(|k| *k > 0); + if let (Some(i), Some(known)) = (infer, known) { + if !total.is_multiple_of(known) { + return shape_err("the Reshape does not divide evenly"); + } + if let Some(s) = shape.get_mut(i) { + *s = total / known; + } + } + Tensor::new(shape, x.data.clone()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::graph::{Attr, Weight}; + + fn node(op: &str, inputs: &[&str], attrs: &[(&str, Attr)]) -> Node { + Node { + op: op.into(), + inputs: inputs.iter().map(|s| s.to_string()).collect(), + outputs: vec!["y".into()], + attrs: attrs.iter().map(|(k, v)| (k.to_string(), v.clone())).collect(), + } + } + fn graph(nodes: Vec, weights: Vec<(&str, Weight)>) -> Graph { + Graph { + nodes, + weights: weights.into_iter().map(|(k, v)| (k.to_string(), v)).collect(), + inputs: vec![("x".into(), vec![1, 1, 3, 3])], + outputs: vec!["y".into()], + } + } + fn w(dims: &[usize], v: Vec) -> Weight { + Weight { dims: dims.to_vec(), data: Data::F32(v) } + } + fn input(c: usize, h: usize, wd: usize, f: impl Fn(usize) -> f32) -> Tensor { + Tensor::new(vec![1, c, h, wd], (0..c * h * wd).map(f).collect()).unwrap() + } + fn run(g: Graph, x: Tensor) -> Result { + Ok(Net::new(g)?.run(x)?.remove(0).1) + } + + /// A direct, slow reference convolution to check the fast one against. + fn reference_conv(x: &Tensor, wd: &[f32], dims: [usize; 4], bias: &[f32], s: usize, p: usize, groups: usize) -> Tensor { + let (c, h, wi) = (x.shape[1], x.shape[2], x.shape[3]); + let [co_n, cin_g, kh, kw] = dims; + let (oh, ow) = ((h + 2 * p - kh) / s + 1, (wi + 2 * p - kw) / s + 1); + let mut out = vec![0.0f32; co_n * oh * ow]; + for co in 0..co_n { + let g = co / (co_n / groups); + for oy in 0..oh { + for ox in 0..ow { + let mut acc = bias[co]; + for cl in 0..cin_g { + for ky in 0..kh { + for kx in 0..kw { + let (iy, ix) = ((oy * s + ky) as isize - p as isize, (ox * s + kx) as isize - p as isize); + if iy >= 0 && ix >= 0 && (iy as usize) < h && (ix as usize) < wi { + acc += + x.data[((g * cin_g + cl) * h + iy as usize) * wi + ix as usize] * wd[((co * cin_g + cl) * kh + ky) * kw + kx]; + } + } + } + } + out[(co * oh + oy) * ow + ox] = acc; + } + } + } + let _ = c; + Tensor::new(vec![1, co_n, oh, ow], out).unwrap() + } + + #[test] + fn convolutions_match_a_slow_reference() { + // (cin, cout, groups, kernel, stride, pad, h, w) + for (cin, cout, groups, k, s, p, h, wi) in [ + (3, 4, 1, 3, 1, 1, 7, 9), + (3, 4, 1, 3, 2, 1, 8, 8), + (4, 4, 4, 3, 1, 1, 6, 5), + (2, 6, 1, 1, 1, 0, 5, 5), + (4, 8, 2, 3, 1, 1, 5, 6), + (1, 1, 1, 3, 1, 0, 3, 3), + ] { + let cin_g = cin / groups; + let wd: Vec = (0..cout * cin_g * k * k).map(|i| ((i * 37 % 11) as f32 - 5.0) * 0.1).collect(); + let bias: Vec = (0..cout).map(|i| i as f32 * 0.25 - 0.5).collect(); + let x = input(cin, h, wi, |i| ((i * 13 % 17) as f32 - 8.0) * 0.2); + let attrs = [("group", Attr::Int(groups as i64)), ("strides", Attr::Ints(vec![s as i64; 2])), ("pads", Attr::Ints(vec![p as i64; 4]))]; + let g = graph( + vec![node("Conv", &["x", "w", "b"], &attrs)], + vec![("w", w(&[cout, cin_g, k, k], wd.clone())), ("b", w(&[cout], bias.clone()))], + ); + let got = run(g, x.clone()).unwrap(); + let want = reference_conv(&x, &wd, [cout, cin_g, k, k], &bias, s, p, groups); + assert_eq!(got.shape, want.shape, "{cin} {cout} {groups} {k} {s} {p}"); + for (a, b) in got.data.iter().zip(&want.data) { + assert!((a - b).abs() < 1e-4, "{cin} {cout} {groups} {k} {s} {p}: {a} vs {b}"); + } + } + } + + #[test] + fn the_small_operators() { + let x = input(1, 4, 4, |i| i as f32 - 6.0); + // Relu, Sigmoid + let r = run(graph(vec![node("Relu", &["x"], &[])], vec![]), x.clone()).unwrap(); + assert!(r.data.iter().all(|v| *v >= 0.0) && r.data.get(7) == Some(&1.0)); + let s = run(graph(vec![node("Sigmoid", &["x"], &[])], vec![]), x.clone()).unwrap(); + assert!((s.data[6] - 0.5).abs() < 1e-6 && s.data.iter().all(|v| (0.0..=1.0).contains(v))); + // MaxPool 2x2/2 + let m = run( + graph(vec![node("MaxPool", &["x"], &[("kernel_shape", Attr::Ints(vec![2, 2])), ("strides", Attr::Ints(vec![2, 2]))])], vec![]), + x.clone(), + ) + .unwrap(); + assert_eq!((m.shape.clone(), m.data.clone()), (vec![1, 1, 2, 2], vec![-1.0, 1.0, 7.0, 9.0])); + // Resize x2 nearest + let z = run(graph(vec![node("Resize", &["x", "", "s"], &[])], vec![("s", w(&[4], vec![1.0, 1.0, 2.0, 2.0]))]), input(1, 2, 2, |i| i as f32)) + .unwrap(); + assert_eq!(z.shape, vec![1, 1, 4, 4]); + assert_eq!(&z.data[..8], &[0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0]); + // Add + let a = run(graph(vec![node("Add", &["x", "x"], &[])], vec![]), x.clone()).unwrap(); + assert_eq!(a.data[5], 2.0 * x.data[5]); + // Transpose [0,2,3,1] then Reshape [1,-1,2] + let t = run(graph(vec![node("Transpose", &["x"], &[("perm", Attr::Ints(vec![0, 2, 3, 1]))])], vec![]), input(2, 2, 3, |i| i as f32)).unwrap(); + assert_eq!(t.shape, vec![1, 2, 3, 2]); + assert_eq!(t.data, vec![0.0, 6.0, 1.0, 7.0, 2.0, 8.0, 3.0, 9.0, 4.0, 10.0, 5.0, 11.0]); + let rs = run( + graph(vec![node("Reshape", &["x", "s"], &[])], vec![("s", Weight { dims: vec![3], data: Data::I64(vec![1, -1, 2]) })]), + input(2, 2, 3, |i| i as f32), + ) + .unwrap(); + assert_eq!(rs.shape, vec![1, 6, 2]); + } + + #[test] + fn unsupported_and_inconsistent_models_are_errors() { + let x = input(1, 3, 3, |i| i as f32); + for (what, g) in [ + ("unknown op", graph(vec![node("Softmax", &["x"], &[])], vec![])), + ( + "dilated conv", + graph(vec![node("Conv", &["x", "w"], &[("dilations", Attr::Ints(vec![2, 2]))])], vec![("w", w(&[1, 1, 1, 1], vec![1.0]))]), + ), + ("conv without weights", graph(vec![node("Conv", &["x", "w"], &[])], vec![])), + ("channels do not match", graph(vec![node("Conv", &["x", "w"], &[])], vec![("w", w(&[1, 5, 1, 1], vec![1.0; 5]))])), + ("kernel too big", graph(vec![node("Conv", &["x", "w"], &[])], vec![("w", w(&[1, 1, 5, 5], vec![1.0; 25]))])), + ("zero group", graph(vec![node("Conv", &["x", "w"], &[("group", Attr::Int(0))])], vec![("w", w(&[1, 1, 1, 1], vec![1.0]))])), + ("zero stride", graph(vec![node("Conv", &["x", "w"], &[("strides", Attr::Ints(vec![0, 0]))])], vec![("w", w(&[1, 1, 1, 1], vec![1.0]))])), + ( + "negative pad", + graph(vec![node("Conv", &["x", "w"], &[("pads", Attr::Ints(vec![-1, 0, 0, 0]))])], vec![("w", w(&[1, 1, 1, 1], vec![1.0]))]), + ), + ("bad perm", graph(vec![node("Transpose", &["x"], &[("perm", Attr::Ints(vec![0, 0, 1, 2]))])], vec![])), + ( + "huge reshape", + graph(vec![node("Reshape", &["x", "s"], &[])], vec![("s", Weight { dims: vec![2], data: Data::I64(vec![i64::MAX, 1]) })]), + ), + ("reshape mismatch", graph(vec![node("Reshape", &["x", "s"], &[])], vec![("s", Weight { dims: vec![2], data: Data::I64(vec![4, 4]) })])), + ("non-integer resize", graph(vec![node("Resize", &["x", "", "s"], &[])], vec![("s", w(&[4], vec![1.0, 1.0, 1.5, 1.5]))])), + ("add that broadcasts", graph(vec![node("Add", &["x", "c"], &[])], vec![])), + ] { + assert!(run(g, x.clone()).is_err(), "{what}"); + } + } +} + +#[cfg(test)] +mod recognition_tests; diff --git a/crates/faces/src/net/kernels.rs b/crates/faces/src/net/kernels.rs new file mode 100644 index 00000000..73cc72b9 --- /dev/null +++ b/crates/faces/src/net/kernels.rs @@ -0,0 +1,306 @@ +//! Safe CPU kernels shared by detection and recognition. No inner thread pool. +use super::*; +use ndarray::{ArrayView2, ArrayViewMut2, linalg::general_mat_mul}; +const BLOCK: usize = 256; +#[derive(Default)] +pub(super) struct Scratch { + lower: Vec, + product: Vec, +} +fn storage(v: &mut Vec, n: usize) -> Result<&mut [f32]> { + if n > MAX_TENSOR { + return shape_err("convolution scratch exceeds the limit"); + } + if v.len() < n { + v.try_reserve_exact(n - v.len()).map_err(|_| NetError::Shape("not enough memory for convolution scratch".into()))?; + v.resize(n, 0.0); + } + v.get_mut(..n).ok_or_else(|| NetError::Shape("missing convolution scratch".into())) +} +/// The safe matrix views assert lengths, so validate exact slices before creating any view. +fn multiply(w: &[f32], x: &[f32], y: &mut [f32], m: usize, k: usize, n: usize) -> Result<()> { + if m == 0 || k == 0 || n == 0 || m.checked_mul(k) != Some(w.len()) || k.checked_mul(n) != Some(x.len()) || m.checked_mul(n) != Some(y.len()) { + return shape_err("invalid GEMM dimensions"); + } + let a = ArrayView2::from_shape((m, k), w).map_err(|e| NetError::Shape(e.to_string()))?; + let b = ArrayView2::from_shape((k, n), x).map_err(|e| NetError::Shape(e.to_string()))?; + let mut c = ArrayViewMut2::from_shape((m, n), y).map_err(|e| NetError::Shape(e.to_string()))?; + general_mat_mul(1.0, &a, &b, 0.0, &mut c); + Ok(()) +} +/// For equal-width, unit-stride maps a shifted plane is one contiguous copy. Only horizontal +/// border stripes need masking; clipping the flat source range also removes vertical padding. +fn shifted_plane(plane: &[f32], width: usize, ky: usize, kx: usize, pt: usize, pl: usize, start: usize, row: &mut [f32]) -> Result<()> { + let end = start + row.len(); + let delta = (ky as i128 - pt as i128) * width as i128 + kx as i128 - pl as i128; + let lo = (-delta).clamp(start as i128, end as i128); + let hi = (plane.len() as i128 - delta).clamp(start as i128, end as i128); + row.fill(0.0); + if lo < hi { + let from = usize::try_from(lo + delta).map_err(|_| NetError::Shape("shifted source overflow".into()))?; + let at = usize::try_from(lo - start as i128).map_err(|_| NetError::Shape("shifted destination overflow".into()))?; + let len = usize::try_from(hi - lo).map_err(|_| NetError::Shape("shifted length overflow".into()))?; + row.get_mut(at..at + len) + .ok_or_else(|| NetError::Shape("short shifted destination".into()))? + .copy_from_slice(plane.get(from..from + len).ok_or_else(|| NetError::Shape("short shifted source".into()))?); + } + let left = pl.saturating_sub(kx).min(width); + let right = (width + pl).saturating_sub(kx).min(width); + if left > 0 || right < width { + for y in start / width..=(end - 1) / width { + for (a, b) in [(y * width, y * width + left), (y * width + right, (y + 1) * width)] { + let lo = a.max(start); + let hi = b.min(end); + if lo < hi { + row.get_mut(lo - start..hi - start).ok_or_else(|| NetError::Shape("short border stripe".into()))?.fill(0.0); + } + } + } + } + Ok(()) +} +pub(super) fn conv(n: &Node, x: &Tensor, g: &Graph, scratch: &mut Scratch) -> Result { + let (cin, h, w) = chw(x)?; + let wt = shape::weight(n, g, 1)?; + let wd = shape::floats(n, g, 1)?; + let [cout, cg, kh, kw] = wt.dims.as_slice() else { + return shape_err("invalid convolution weights"); + }; + let (cout, cg, kh, kw) = (*cout, *cg, *kh, *kw); + let groups = usize::try_from(n.int("group").unwrap_or(1)).map_err(|_| NetError::Shape("invalid group".into()))?; + if cg == 1 { + return conv_direct(n, x, g); + } + let (sh, sw) = pair(n, "strides", 1)?; + let (pt, pl, pb, pr) = pads(n)?; + let ph = h.checked_add(pt).and_then(|v| v.checked_add(pb)).ok_or_else(|| NetError::Shape("padding overflow".into()))?; + let pw = w.checked_add(pl).and_then(|v| v.checked_add(pr)).ok_or_else(|| NetError::Shape("padding overflow".into()))?; + if ph < kh || pw < kw || groups == 0 || cin != groups * cg { + return shape_err("invalid convolution shape"); + } + let (oh, ow) = ((ph - kh) / sh + 1, (pw - kw) / sw + 1); + let mut out = Tensor::zeros(vec![1, cout, oh, ow])?; + let pixels = oh * ow; + let k = cg * kh * kw; + let m = cout / groups; + let bias = if n.inputs.get(2).is_some_and(|s| !s.is_empty()) { Some(shape::floats(n, g, 2)?) } else { None }; + for group in 0..groups { + let weights = wd.get(group * m * k..(group + 1) * m * k).ok_or_else(|| NetError::Shape("short convolution weights".into()))?; + // A pointwise convolution is already a matrix: no lowering, copy or spatial blocking needed. + if kh == 1 && kw == 1 && sh == 1 && sw == 1 && (pt, pl, pb, pr) == (0, 0, 0, 0) { + let input = x.data.get(group * cg * pixels..(group + 1) * cg * pixels).ok_or_else(|| NetError::Shape("short pointwise input".into()))?; + let output = + out.data.get_mut(group * m * pixels..(group + 1) * m * pixels).ok_or_else(|| NetError::Shape("short pointwise output".into()))?; + multiply(weights, input, output, m, cg, pixels)?; + for (co, plane) in output.chunks_exact_mut(pixels).enumerate() { + let b = bias.and_then(|v| v.get(group * m + co)).copied().unwrap_or(0.0); + if b != 0.0 { + for v in plane { + *v += b; + } + } + } + continue; + } + for start in (0..pixels).step_by(BLOCK) { + let len = BLOCK.min(pixels - start); + let lower = storage(&mut scratch.lower, k * len)?; + for cl in 0..cg { + let plane = x + .data + .get((group * cg + cl) * h * w..(group * cg + cl + 1) * h * w) + .ok_or_else(|| NetError::Shape("short convolution input".into()))?; + for ky in 0..kh { + for kx in 0..kw { + let row = lower + .get_mut((cl * kh * kw + ky * kw + kx) * len..(cl * kh * kw + ky * kw + kx + 1) * len) + .ok_or_else(|| NetError::Shape("short im2col row".into()))?; + if sh == 1 && sw == 1 && ow == w { + shifted_plane( + plane, + w, + ky, + kx, + pt, + pl, + start, + row.get_mut(..len).ok_or_else(|| NetError::Shape("short lower panel".into()))?, + )?; + continue; + } + let mut at = 0; + while at < len { + let p = start + at; + let (oy, ox) = (p / ow, p % ow); + let span = (ow - ox).min(len - at); + if let Some(yy) = (oy * sh + ky).checked_sub(pt).filter(|&v| v < h) { + let lo = ox.max(pl.saturating_sub(kx).div_ceil(sw)); + let hi = (ox + span).min((w + pl).saturating_sub(kx).div_ceil(sw)); + if lo < hi { + let from = yy * w + lo * sw + kx - pl; + let to = at + lo - ox; + row.get_mut(at..to).ok_or_else(|| NetError::Shape("short im2col padding".into()))?.fill(0.0); + let src = plane + .get(from..from + (hi - lo - 1) * sw + 1) + .ok_or_else(|| NetError::Shape("short im2col source row".into()))?; + let dst = row.get_mut(to..to + hi - lo).ok_or_else(|| NetError::Shape("short im2col row".into()))?; + if sw == 1 { + dst.copy_from_slice(src); + } else { + for (o, &v) in dst.iter_mut().zip(src.iter().step_by(sw)) { + *o = v; + } + } + row.get_mut(to + hi - lo..at + span).ok_or_else(|| NetError::Shape("short im2col padding".into()))?.fill(0.0); + } else { + row.get_mut(at..at + span).ok_or_else(|| NetError::Shape("short im2col row".into()))?.fill(0.0); + } + } else { + row.get_mut(at..at + span).ok_or_else(|| NetError::Shape("short im2col row".into()))?.fill(0.0); + } + at += span; + } + } + } + } + let product = storage(&mut scratch.product, m * len)?; + multiply(weights, lower, product, m, k, len)?; + for (co, row) in product.chunks_exact(len).enumerate() { + let channel = group * m + co; + let b = bias.and_then(|v| v.get(channel)).copied().unwrap_or(0.0); + let dest = out + .data + .get_mut(channel * pixels + start..channel * pixels + start + len) + .ok_or_else(|| NetError::Shape("short convolution output".into()))?; + for (to, &value) in dest.iter_mut().zip(row) { + *to = value + b; + } + } + } + } + Ok(out) +} +/// Flat offset under ONNX's trailing-dimension broadcasting; shapes are checked before execution. +fn broadcast_offset(mut at: usize, out: &[usize], input: &[usize]) -> usize { + let mut offset = 0; + let mut stride = 1; + for (axis, &size) in out.iter().enumerate().rev() { + let coordinate = at % size; + at /= size; + if let Some(&d) = axis.checked_sub(out.len() - input.len()).and_then(|i| input.get(i)) { + if d != 1 { + offset += coordinate * stride; + } + stride *= d; + } + } + offset +} +pub(super) fn binary(op: &str, a: &Tensor, b: &Tensor) -> Result { + let dims = shape::broadcast(&a.shape, &b.shape)?; + let mut out = Tensor::zeros(dims)?; + let f = |x: f32, y: f32| match op { + "Add" => x + y, + "Sub" => x - y, + _ => x * y, + }; + if a.shape == b.shape { + for ((o, &x), &y) in out.data.iter_mut().zip(&a.data).zip(&b.data) { + *o = f(x, y); + } + } else if b.data.len() == 1 && a.shape == out.shape { + let y = *b.data.first().ok_or_else(|| NetError::Shape("empty scalar".into()))?; + for (o, &x) in out.data.iter_mut().zip(&a.data) { + *o = f(x, y); + } + } else { + for (i, o) in out.data.iter_mut().enumerate() { + let x = a.data.get(broadcast_offset(i, &out.shape, &a.shape)).ok_or_else(|| NetError::Shape("bad broadcast A".into()))?; + let y = b.data.get(broadcast_offset(i, &out.shape, &b.shape)).ok_or_else(|| NetError::Shape("bad broadcast B".into()))?; + *o = f(*x, *y); + } + } + Ok(out) +} +pub(super) fn prelu(n: &Node, x: &Tensor, g: &Graph) -> Result { + let wt = shape::weight(n, g, 1)?; + let slope = shape::floats(n, g, 1)?; + let mut out = Tensor::zeros(x.shape.clone())?; + let channels = x.shape.get(1).copied().unwrap_or(1); + let mut expanded = vec![1; x.shape.len().saturating_sub(wt.dims.len())]; + expanded.extend_from_slice(&wt.dims); + let channel_slope = expanded.iter().enumerate().all(|(i, &d)| d == 1 || i == 1) && slope.len() == channels; + if channel_slope || slope.len() == 1 { + let plane = shape::elements(x.shape.get(2..).unwrap_or(&[]))?; + for (c, (src, dst)) in x.data.chunks_exact(plane).zip(out.data.chunks_exact_mut(plane)).enumerate() { + let alpha = *slope.get(if slope.len() == 1 { 0 } else { c % channels }).ok_or_else(|| NetError::Shape("short PRelu slope".into()))?; + for (to, &v) in dst.iter_mut().zip(src) { + *to = if v < 0.0 { v * alpha } else { v }; + } + } + } else { + for (i, (o, &v)) in out.data.iter_mut().zip(&x.data).enumerate() { + let alpha = *slope.get(broadcast_offset(i, &x.shape, &wt.dims)).ok_or_else(|| NetError::Shape("bad PRelu broadcast".into()))?; + *o = if v < 0.0 { v * alpha } else { v }; + } + } + Ok(out) +} +pub(super) fn batch_norm(n: &Node, x: &Tensor, g: &Graph) -> Result { + let channels = *x.shape.get(1).ok_or_else(|| NetError::Shape("BatchNormalization rank".into()))?; + let (scale, bias, mean, var) = (shape::floats(n, g, 1)?, shape::floats(n, g, 2)?, shape::floats(n, g, 3)?, shape::floats(n, g, 4)?); + let eps = shape::float(n, "epsilon", 1e-5); + let plane = shape::elements(x.shape.get(2..).unwrap_or(&[]))?; + let mut out = Tensor::zeros(x.shape.clone())?; + for (j, (src, dst)) in x.data.chunks_exact(plane).zip(out.data.chunks_exact_mut(plane)).enumerate() { + let c = j % channels; + let coefficient = *scale.get(c).ok_or_else(|| NetError::Shape("short BN scale".into()))? + / (var.get(c).copied().ok_or_else(|| NetError::Shape("short BN variance".into()))? + eps).sqrt(); + let offset = bias.get(c).copied().ok_or_else(|| NetError::Shape("short BN bias".into()))? + - mean.get(c).copied().ok_or_else(|| NetError::Shape("short BN mean".into()))? * coefficient; + for (to, &value) in dst.iter_mut().zip(src) { + *to = value * coefficient + offset; + } + } + Ok(out) +} +pub(super) fn gemm(n: &Node, a: &Tensor, g: &Graph) -> Result { + let wt = shape::weight(n, g, 1)?; + let b = shape::floats(n, g, 1)?; + let [ar, ac] = a.shape.as_slice() else { + return shape_err("Gemm A is not a matrix"); + }; + let [br, bc] = wt.dims.as_slice() else { + return shape_err("Gemm B is not a matrix"); + }; + if ar.checked_mul(*ac) != Some(a.data.len()) || br.checked_mul(*bc) != Some(b.len()) { + return shape_err("invalid matrix buffers"); + } + let av = ArrayView2::from_shape((*ar, *ac), a.data.as_slice()).map_err(|e| NetError::Shape(e.to_string()))?; + let bv = ArrayView2::from_shape((*br, *bc), b).map_err(|e| NetError::Shape(e.to_string()))?; + let av = if n.int("transA").unwrap_or(0) != 0 { av.reversed_axes() } else { av }; + let bv = if n.int("transB").unwrap_or(0) != 0 { bv.reversed_axes() } else { bv }; + if av.ncols() != bv.nrows() { + return shape_err("incompatible matrices"); + } + let mut out = Tensor::zeros(vec![av.nrows(), bv.ncols()])?; + let mut view = ArrayViewMut2::from_shape((av.nrows(), bv.ncols()), out.data.as_mut_slice()).map_err(|e| NetError::Shape(e.to_string()))?; + general_mat_mul(shape::float(n, "alpha", 1.0), &av, &bv, 0.0, &mut view); + if n.inputs.get(2).is_some_and(|s| !s.is_empty()) { + let wt = shape::weight(n, g, 2)?; + let bias = shape::floats(n, g, 2)?; + let beta = shape::float(n, "beta", 1.0); + for (i, to) in out.data.iter_mut().enumerate() { + *to += beta * bias.get(broadcast_offset(i, &out.shape, &wt.dims)).copied().ok_or_else(|| NetError::Shape("short matrix bias".into()))?; + } + } + Ok(out) +} +pub(super) fn flatten(n: &Node, x: &Tensor) -> Result { + let axis = shape::axis(n, x.shape.len())?; + Tensor::new(vec![shape::elements(x.shape.get(..axis).unwrap_or(&[]))?, shape::elements(x.shape.get(axis..).unwrap_or(&[]))?], x.data.clone()) +} +pub(super) fn global_average(x: &Tensor) -> Result { + let (c, h, w) = chw(x)?; + Tensor::new(vec![1, c, 1, 1], x.data.chunks_exact(h * w).map(|v| v.iter().sum::() / (h * w) as f32).collect()) +} diff --git a/crates/faces/src/net/recognition_tests.rs b/crates/faces/src/net/recognition_tests.rs new file mode 100644 index 00000000..e245bb33 --- /dev/null +++ b/crates/faces/src/net/recognition_tests.rs @@ -0,0 +1,337 @@ +use super::*; +use crate::graph::{Attr, Weight}; +fn weight(shape: &[usize], values: Vec) -> Weight { + Weight { dims: shape.to_vec(), data: Data::F32(values) } +} +fn graph(op: &str, extra: &[&str], weights: Vec<(&str, Weight)>, attrs: Vec<(&str, Attr)>) -> Graph { + Graph { + nodes: vec![Node { + op: op.into(), + inputs: std::iter::once("x").chain(extra.iter().copied()).map(str::to_string).collect(), + outputs: vec!["y".into()], + attrs: attrs.into_iter().map(|(s, a)| (s.into(), a)).collect(), + }], + weights: weights.into_iter().map(|(s, w)| (s.into(), w)).collect(), + inputs: vec![("x".into(), vec![1, 3, 2, 2])], + outputs: vec!["y".into()], + } +} +fn execute(g: Graph, shape: &[usize], values: Vec) -> Tensor { + Net::new(g).unwrap().run(Tensor::new(shape.to_vec(), values).unwrap()).unwrap().remove(0).1 +} +#[test] +fn batch_norm_matches_the_inference_formula_for_images_and_matrices() { + for shape in [vec![1, 3, 2, 5], vec![2, 3]] { + let values: Vec<_> = (0..shape.iter().product()).map(|i| i as f32 * 0.07 - 0.8).collect(); + let g = graph( + "BatchNormalization", + &["scale", "bias", "mean", "var"], + vec![ + ("scale", weight(&[3], vec![0.7, -1.1, 2.3])), + ("bias", weight(&[3], vec![-0.4, 0.6, 0.2])), + ("mean", weight(&[3], vec![0.5, -0.3, 0.9])), + ("var", weight(&[3], vec![0.2, 0.0, 1.4])), + ], + vec![("epsilon", Attr::Float(1e-3))], + ); + let out = execute(g.clone(), &shape, values.clone()); + let spatial: usize = shape.get(2..).unwrap_or(&[]).iter().product(); + for (i, (&value, &got)) in values.iter().zip(&out.data).enumerate() { + let c = (i / spatial) % 3; + let [s, b, m, v] = ["scale", "bias", "mean", "var"].map(|name| { + let Data::F32(v) = &g.weights[name].data else { panic!() }; + v[c] + }); + let want = (value - m) * s / (v + 1e-3).sqrt() + b; + assert!((got - want).abs() < 1e-5); + } + } +} +#[test] +fn prelu_obeys_trailing_broadcasting_including_spatial_slopes() { + for (shape, slopes) in + [(vec![3, 1, 1], vec![0.1, 0.2, 0.3]), (vec![1, 3, 1, 1], vec![0.1, 0.2, 0.3]), (vec![1], vec![0.3]), (vec![2, 2], vec![0.1, 0.2, 0.3, 0.4])] + { + let values: Vec<_> = (0..12).map(|i| i as f32 - 6.0).collect(); + let out = execute(graph("PRelu", &["s"], vec![("s", weight(&shape, slopes.clone()))], vec![]), &[1, 3, 2, 2], values.clone()); + for (i, (&v, &got)) in values.iter().zip(&out.data).enumerate() { + let alpha = if slopes.len() == 1 { + slopes[0] + } else if slopes.len() == 3 { + slopes[i / 4] + } else { + slopes[i % 4] + }; + assert_eq!(got, if v < 0.0 { v * alpha } else { v }); + } + } +} +#[test] +fn gemm_transpositions_scalars_vectors_matrices_and_optional_bias() { + for ta in [0, 1] { + for tb in [0, 1] { + for bias_shape in [vec![], vec![2], vec![1, 2], vec![2, 1], vec![2, 2]] { + let (ash, bsh) = (if ta == 0 { vec![2, 3] } else { vec![3, 2] }, if tb == 0 { vec![3, 2] } else { vec![2, 3] }); + let a: Vec<_> = (0..6).map(|i| i as f32 * 0.4 - 0.8).collect(); + let b: Vec<_> = (0..6).map(|i| i as f32 * 0.1 - 0.3).collect(); + let bias: Vec<_> = (0..bias_shape.iter().product()).map(|i| i as f32 * 0.25 + 0.3).collect(); + let g = graph( + "Gemm", + &["w", "b"], + vec![("w", weight(&bsh, b.clone())), ("b", weight(&bias_shape, bias.clone()))], + vec![("transA", Attr::Int(ta)), ("transB", Attr::Int(tb)), ("alpha", Attr::Float(0.7)), ("beta", Attr::Float(-0.4))], + ); + let out = execute(g, &ash, a.clone()); + assert_eq!(out.shape, vec![2, 2]); + for i in 0..2 { + for j in 0..2 { + let sum = (0..3) + .map(|k| a[if ta == 0 { i * 3 + k } else { k * 2 + i }] * b[if tb == 0 { k * 2 + j } else { j * 3 + k }]) + .sum::(); + let off = match bias_shape.as_slice() { + [] => 0, + [2] | [1, 2] => j, + [2, 1] => i, + _ => i * 2 + j, + }; + assert!((out.data[i * 2 + j] - (0.7 * sum - 0.4 * bias[off])).abs() < 1e-6); + } + } + } + } + } + let out = execute(graph("Gemm", &["w"], vec![("w", weight(&[2, 1], vec![2.0, 3.0]))], vec![]), &[1, 2], vec![4.0, 5.0]); + assert_eq!(out.data, vec![23.0]); +} +#[test] +fn arithmetic_broadcasts_and_keeps_operand_order() { + for op in ["Add", "Sub", "Mul"] { + let out = + execute(graph(op, &["w"], vec![("w", weight(&[3, 1], vec![0.5, 1.5, 2.5]))], vec![]), &[1, 3, 2], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + for (i, &got) in out.data.iter().enumerate() { + let x = (i + 1) as f32; + let y = (i / 2) as f32 + 0.5; + let want = match op { + "Add" => x + y, + "Sub" => x - y, + _ => x * y, + }; + assert_eq!(got, want); + } + } +} +#[test] +fn flatten_negative_axis_and_default_transpose_are_correct() { + let out = execute(graph("Flatten", &[], vec![], vec![("axis", Attr::Int(-1))]), &[1, 3, 2, 5], (0..30).map(|i| i as f32).collect()); + assert_eq!(out.shape, vec![6, 5]); + let out = execute(graph("Transpose", &[], vec![], vec![]), &[2, 3], (0..6).map(|i| i as f32).collect()); + assert_eq!(out.shape, vec![3, 2]); + assert_eq!(out.data, vec![0.0, 3.0, 1.0, 4.0, 2.0, 5.0]); +} +#[test] +fn all_recognition_operators_run_on_seeded_synthetic_networks() { + let input = Tensor::new(vec![1, 3, 112, 112], (0..3 * 112 * 112).map(|i| (i * 37 % 256) as f32).collect()).unwrap(); + for seed in 0..20 { + let net = Net::new(crate::graph::load(&crate::synthetic::recognition_model(seed)).unwrap()).unwrap(); + let a = net.run(input.clone()).unwrap(); + let b = net.run(input.clone()).unwrap(); + assert_eq!(a, b); + assert_eq!(a[0].1.shape, vec![1, 8]); + } +} +#[test] +fn hostile_attributes_weights_shapes_and_public_tensors_return_errors() { + for (name, attr) in [ + ("strides", Attr::Ints(vec![0, 1])), + ("pads", Attr::Ints(vec![i64::MAX; 4])), + ("group", Attr::Int(0)), + ("kernel_shape", Attr::Ints(vec![7, 7])), + ("strides", Attr::Float(1.0)), + ("new_attribute", Attr::Int(1)), + ] { + let g = graph("Conv", &["w"], vec![("w", weight(&[3, 3, 1, 1], vec![0.1; 9]))], vec![(name, attr)]); + let result = Net::new(g).and_then(|n| n.validate_input(&[1, 3, 3, 3])); + assert!(result.is_err(), "{name}"); + } + for op in ["Loop", "MatMul", "If", "Concat", "Unknown"] { + assert!(Net::new(graph(op, &[], vec![], vec![])).is_err()); + } + for value in [f32::NAN, f32::INFINITY] { + assert!(Net::new(graph("PRelu", &["w"], vec![("w", weight(&[1], vec![value]))], vec![])).is_err()); + } + let net = Net::new(graph("Relu", &[], vec![], vec![])).unwrap(); + for input in [ + Tensor { shape: vec![usize::MAX, 2], data: vec![] }, + Tensor { shape: vec![0, 3], data: vec![] }, + Tensor { shape: vec![1], data: vec![f32::NAN] }, + Tensor { shape: vec![1; 9], data: vec![1.0] }, + Tensor { shape: vec![2], data: vec![1.0] }, + ] { + assert!(net.run(input).is_err()); + } + let mut g = graph("Relu", &[], vec![], vec![]); + g.nodes.push(g.nodes[0].clone()); + assert!(Net::new(g).is_err()); + let mut g = graph("Relu", &[], vec![], vec![]); + g.nodes[0].inputs[0] = "future".into(); + assert!(Net::new(g).is_err()); +} +#[test] +fn malformed_normalization_prelu_matrix_dropout_and_axes_are_errors() { + let bn = || { + graph( + "BatchNormalization", + &["s", "b", "m", "v"], + vec![ + ("s", weight(&[3], vec![1.0; 3])), + ("b", weight(&[3], vec![0.0; 3])), + ("m", weight(&[3], vec![0.0; 3])), + ("v", weight(&[3], vec![1.0; 3])), + ], + vec![], + ) + }; + for mutation in 0..4 { + let mut g = bn(); + match mutation { + 0 => { + g.nodes[0].attrs.insert("training_mode".into(), Attr::Int(1)); + } + 1 => { + g.nodes[0].attrs.insert("epsilon".into(), Attr::Float(0.0)); + } + 2 => { + g.weights.insert("v".into(), weight(&[3], vec![-1.0; 3])); + } + _ => { + g.weights.insert("b".into(), weight(&[1, 3], vec![0.0; 3])); + } + }; + assert!(Net::new(g).and_then(|n| n.validate_input(&[1, 3, 2, 2])).is_err()); + } + for g in [ + graph("PRelu", &["s"], vec![("s", weight(&[3], vec![0.1; 3]))], vec![]), + graph("Gemm", &["w"], vec![("w", weight(&[4, 2], vec![0.1; 8]))], vec![]), + graph("Flatten", &[], vec![], vec![("axis", Attr::Int(i64::MIN))]), + graph("Dropout", &[], vec![], vec![("ratio", Attr::Float(1.0))]), + graph("Transpose", &[], vec![], vec![("perm", Attr::Ints(vec![0, 1, 1, 3]))]), + ] { + assert!(Net::new(g).and_then(|n| n.validate_input(&[1, 3, 2, 2])).is_err()); + } +} +#[test] +fn arithmetic_overflow_is_reported_and_next_run_succeeds() { + let net = Net::new(graph("Mul", &["w"], vec![("w", weight(&[1], vec![f32::MAX]))], vec![])).unwrap(); + assert!(net.run(Tensor::new(vec![1], vec![2.0]).unwrap()).is_err()); + assert!(net.run(Tensor::new(vec![1], vec![0.0]).unwrap()).is_ok()); +} + +#[test] +fn multiple_batches_prelu_keeps_channels_independent() { + let out = execute(graph("PRelu", &["s"], vec![("s", weight(&[3, 1, 1], vec![0.1, 0.2, 0.3]))], vec![]), &[2, 3, 2, 2], vec![-1.0; 24]); + for (i, &v) in out.data.iter().enumerate() { + assert_eq!(v, -[0.1, 0.2, 0.3][i / 4 % 3]); + } +} +#[test] +fn integer_activations_and_constant_outputs_fail_at_load() { + let integer = Weight { dims: vec![1], data: Data::I64(vec![1]) }; + assert!(Net::new(graph("Add", &["w"], vec![("w", integer.clone())], vec![])).is_err()); + let mut g = graph("Relu", &[], vec![("w", integer)], vec![]); + g.outputs = vec!["w".into()]; + assert!(Net::new(g).is_err()); +} +#[test] +fn normalisation_coefficient_and_working_memory_overflows_fail_before_running() { + let g = graph( + "BatchNormalization", + &["s", "b", "m", "v"], + vec![ + ("s", weight(&[3], vec![f32::MAX; 3])), + ("b", weight(&[3], vec![0.0; 3])), + ("m", weight(&[3], vec![f32::MAX; 3])), + ("v", weight(&[3], vec![0.0; 3])), + ], + vec![], + ); + assert!(Net::new(g).unwrap().validate_input(&[1, 3, 2, 2]).is_err()); + let g = graph("Conv", &["w"], vec![("w", weight(&[4096, 3, 1, 1], vec![0.1; 4096 * 3]))], vec![]); + assert!(Net::new(g).unwrap().validate_input(&[1, 3, 4096, 4096]).is_err()); +} + +#[test] +fn empty_broadcast_operands_are_rejected_before_allocation() { + for op in ["Add", "Sub", "Mul", "PRelu"] { + let g = graph(op, &["w"], vec![("w", weight(&[0, 3, 1, 1], vec![]))], vec![]); + assert!(Net::new(g).and_then(|n| n.validate_input(&[1, 3, 2, 2])).is_err()); + } +} + +#[test] +fn odd_maps_cross_panel_boundaries_and_reused_padding_is_zero() { + for (c, co, groups) in [(3, 11, 1), (6, 10, 2)] { + let weights: Vec<_> = (0..co * (c / groups) * 9).map(|i| (i as f32 % 19.0 - 9.0) * 0.013).collect(); + let g = graph( + "Conv", + &["w", "b"], + vec![("w", weight(&[co, c / groups, 3, 3], weights)), ("b", weight(&[co], vec![0.1; co]))], + vec![("group", Attr::Int(groups as i64)), ("strides", Attr::Ints(vec![1, 1])), ("pads", Attr::Ints(vec![2, 1, 0, 1]))], + ); + let mut scratch = kernels::Scratch::default(); + for zero in [false, true, false] { + let input = Tensor::new(vec![1, c, 23, 29], (0..c * 23 * 29).map(|i| if zero { 0.0 } else { (i as f32 % 29.0 - 14.0) * 0.12 }).collect()) + .unwrap(); + let got = kernels::conv(&g.nodes[0], &input, &g, &mut scratch).unwrap(); + let want = conv_direct(&g.nodes[0], &input, &g).unwrap(); + assert_eq!(got.shape, want.shape); + for (&a, &b) in got.data.iter().zip(&want.data) { + assert!((a - b).abs() < 1e-5, "{a} vs {b}"); + } + } + } +} + +#[test] +fn retained_convolution_scratch_counts_towards_the_working_limit() { + let network = |output_channels: usize| { + let pads = vec![("pads", Attr::Ints(vec![1, 0, 0, 0]))]; + let mut g = graph("Conv", &["w"], vec![("w", weight(&[2, 130_000, 2, 1], vec![0.0; 2 * 130_000 * 2]))], pads.clone()); + g.inputs[0].1 = vec![1, 130_000, 2, 256]; + g.weights.insert("v".into(), weight(&[output_channels, 2, 2, 1], vec![0.0; output_channels * 4])); + g.nodes.push(Node { + op: "Conv".into(), + inputs: vec!["y".into(), "v".into()], + outputs: vec!["z".into()], + attrs: pads.into_iter().map(|(s, a)| (s.into(), a)).collect(), + }); + g.outputs = vec!["z".into()]; + Net::new(g).unwrap() + }; + // The first layer retains a large im2col buffer; the second grows the product buffer. + // Both layers individually fit, but keeping both scratch peaks would exceed 512 MiB. + let shape = [1, 130_000, 2, 256]; + assert!(network(131_072).validate_input(&shape).is_err()); + assert!(network(65_536).validate_input(&shape).is_ok()); +} + +#[test] +fn constant_input_copies_count_towards_the_working_limit() { + let network = |retained_outputs: usize| { + let dims = [1, 1, 2048, 2048]; + let mut g = graph("Identity", &[], vec![("w", weight(&dims, vec![0.0; 2048 * 2048]))], vec![]); + g.inputs[0].1 = dims.to_vec(); + g.nodes.clear(); + g.outputs.clear(); + for i in 0..retained_outputs { + let name = format!("retained-{i}"); + g.nodes.push(Node { op: "Identity".into(), inputs: vec!["x".into()], outputs: vec![name.clone()], attrs: HashMap::new() }); + g.outputs.push(name); + } + g.nodes.push(Node { op: "Add".into(), inputs: vec!["x".into(), "w".into()], outputs: vec!["z".into()], attrs: HashMap::new() }); + g.outputs.push("z".into()); + Net::new(g).unwrap() + }; + // Infer only: no pixel tensors are allocated. The constant's temporary copy must be counted. + assert!(network(30).validate_input(&[1, 1, 2048, 2048]).is_err()); + assert!(network(29).validate_input(&[1, 1, 2048, 2048]).is_ok()); +} diff --git a/crates/faces/src/net/shape.rs b/crates/faces/src/net/shape.rs new file mode 100644 index 00000000..d76271bc --- /dev/null +++ b/crates/faces/src/net/shape.rs @@ -0,0 +1,496 @@ +//! Bounded, allocation-free (apart from small shape lists) graph checking. +use super::*; +use crate::graph::{Attr, Weight}; + +pub(super) fn elements(s: &[usize]) -> Result { + if s.len() > 8 || s.contains(&0) { + return shape_err("a tensor has a zero dimension or excessive rank"); + } + s.iter() + .try_fold(1usize, |n, &d| n.checked_mul(d)) + .filter(|&n| n <= MAX_TENSOR) + .ok_or_else(|| NetError::Shape("a tensor exceeds the element limit".into())) +} +pub(super) fn tensor(s: &[usize], n: usize, data: &[f32]) -> Result<()> { + if elements(s)? != n || data.iter().any(|v| !v.is_finite()) { + return shape_err("invalid tensor size or non-finite values"); + } + Ok(()) +} +pub(super) fn float(n: &Node, name: &str, default: f32) -> f32 { + match n.attrs.get(name) { + Some(Attr::Float(v)) => *v, + _ => default, + } +} +pub(super) fn weight<'a>(n: &Node, g: &'a Graph, i: usize) -> Result<&'a Weight> { + n.inputs.get(i).and_then(|s| g.weights.get(s)).ok_or_else(|| NetError::Unsupported(format!("{} needs constant input {i}", n.op))) +} +pub(super) fn floats<'a>(n: &Node, g: &'a Graph, i: usize) -> Result<&'a [f32]> { + match &weight(n, g, i)?.data { + Data::F32(v) => Ok(v), + _ => unsupported(format!("{} needs float32 weights", n.op)), + } +} +pub(super) fn check_graph(g: &Graph) -> Result<()> { + if g.inputs.len() != 1 || g.outputs.is_empty() || g.outputs.len() > 64 || g.nodes.len() > 20_000 || g.weights.len() > 20_000 { + return unsupported("one input and at most 64 outputs are required"); + } + let mut known = std::collections::HashSet::new(); + let mut total_weights = 0usize; + for (name, w) in &g.weights { + let count = w.dims.iter().try_fold(1usize, |n, &d| n.checked_mul(d)); + let len = match &w.data { + Data::F32(v) => { + if v.iter().any(|v| !v.is_finite()) { + return shape_err("non-finite weights"); + } + v.len() + } + Data::I64(v) => v.len(), + }; + if name.is_empty() + || name.len() > 4096 + || w.dims.len() > 8 + || w.dims.iter().any(|&d| d > MAX_TENSOR) + || count != Some(len) + || len > MAX_TENSOR + { + return shape_err("invalid weight name, shape or size"); + } + total_weights = total_weights + .checked_add(len) + .filter(|&v| v <= MAX_TENSOR) + .ok_or_else(|| NetError::Shape("combined weights exceed 64 Mi elements".into()))?; + known.insert(name.as_str()); + } + let input = g.inputs.first().ok_or_else(|| NetError::Shape("missing input".into()))?; + if input.0.is_empty() || !known.insert(&input.0) { + return shape_err("invalid or repeated input name"); + } + for n in &g.nodes { + if n.inputs.len() > 64 || n.outputs.len() != 1 || n.outputs.first().is_none_or(|s| s.is_empty() || s.len() > 4096) { + return unsupported(format!("{} requires one named output", n.op)); + } + if n.inputs.iter().any(|s| !s.is_empty() && !known.contains(s.as_str())) { + return shape_err(format!("{} reads an unknown or forward value", n.op)); + } + if !known.insert(n.outputs.first().ok_or_else(|| NetError::Shape("missing output".into()))?) { + return shape_err("a graph value is defined twice"); + } + } + if g.outputs.iter().any(|s| g.weights.contains_key(s)) { + return unsupported("constant graph outputs"); + } + let mut output_names = std::collections::HashSet::new(); + if g.outputs.iter().any(|s| !known.contains(s.as_str()) || !output_names.insert(s)) { + return shape_err("unknown or repeated graph output"); + } + Ok(()) +} +pub(super) fn check(n: &Node, g: &Graph) -> Result<()> { + let (min, max, attrs): (usize, usize, &[(&str, u8)]) = match n.op.as_str() { + "Conv" => (2, 3, &[("group", 0), ("kernel_shape", 2), ("strides", 2), ("pads", 2), ("dilations", 2), ("auto_pad", 3)]), + "BatchNormalization" => (5, 5, &[("epsilon", 1), ("momentum", 1), ("training_mode", 0)]), + "PRelu" | "Add" | "Sub" | "Mul" => (2, 2, &[]), + "Relu" | "Sigmoid" | "GlobalAveragePool" | "Identity" => (1, 1, &[]), + "Dropout" => (1, 1, &[("ratio", 1), ("seed", 0)]), + "Flatten" => (1, 1, &[("axis", 0)]), + "Gemm" => (2, 3, &[("alpha", 1), ("beta", 1), ("transA", 0), ("transB", 0)]), + "MaxPool" => { + (1, 1, &[("kernel_shape", 2), ("strides", 2), ("pads", 2), ("dilations", 2), ("ceil_mode", 0), ("auto_pad", 3), ("storage_order", 0)]) + } + "Resize" => ( + 3, + 4, + &[ + ("mode", 3), + ("coordinate_transformation_mode", 3), + ("nearest_mode", 3), + ("cubic_coeff_a", 1), + ("exclude_outside", 0), + ("extrapolation_value", 1), + ], + ), + "Transpose" => (1, 1, &[("perm", 2)]), + "Reshape" => (2, 2, &[("allowzero", 0)]), + op => return unsupported(format!("the operator `{op}`")), + }; + if !(min..=max).contains(&n.inputs.len()) || n.inputs.first().is_none_or(String::is_empty) { + return shape_err(format!("{} has invalid inputs", n.op)); + } + for (name, value) in &n.attrs { + let actual = match value { + Attr::Int(_) => 0, + Attr::Float(v) if v.is_finite() => 1, + Attr::Ints(_) => 2, + Attr::Str(_) => 3, + _ => 255, + }; + if attrs.iter().find(|(key, _)| *key == name).is_none_or(|(_, kind)| *kind != actual) { + return unsupported(format!("{} attribute `{name}` (or its type)", n.op)); + } + } + for (i, name) in n.inputs.iter().enumerate() { + if let Some(w) = g.weights.get(name) + && matches!(w.data, Data::I64(_)) + && !(n.op == "Reshape" && i == 1) + { + return unsupported(format!("{} needs float32 input {i}", n.op)); + } + } + match n.op.as_str() { + "Conv" => { + let w = weight(n, g, 1)?; + floats(n, g, 1)?; + if w.dims.len() != 4 || w.dims.contains(&0) { + return shape_err("Conv needs nonempty 2-D weights"); + } + if n.ints("kernel_shape").is_some_and(|v| { + v.iter().map(|&d| usize::try_from(d).ok()).collect::>>().as_deref() != Some(w.dims.get(2..).unwrap_or(&[])) + }) { + return shape_err("Conv kernel_shape differs from its weights"); + } + if n.ints("dilations").is_some_and(|d| d != [1, 1]) || n.string("auto_pad").is_some_and(|s| s != "NOTSET") { + return unsupported("Conv dilation or automatic padding"); + } + let (sh, sw) = pair(n, "strides", 1)?; + let (pt, pl, pb, pr) = pads(n)?; + if [sh, sw, pt, pl, pb, pr].iter().any(|&v| v > 65536) { + return shape_err("Conv stride or padding exceeds 65536"); + } + if n.int("group").unwrap_or(1) < 1 { + return shape_err("Conv group must be positive"); + } + if n.inputs.get(2).is_some_and(|s| !s.is_empty()) { + floats(n, g, 2)?; + } + } + "BatchNormalization" => { + if n.int("training_mode").unwrap_or(0) != 0 { + return unsupported("training BatchNormalization"); + } + if float(n, "epsilon", 1e-5) <= 0.0 { + return shape_err("BatchNormalization epsilon must be positive"); + } + for i in 1..5 { + floats(n, g, i)?; + } + } + "PRelu" => { + floats(n, g, 1)?; + } + "Gemm" => { + floats(n, g, 1)?; + if n.inputs.get(2).is_some_and(|s| !s.is_empty()) { + floats(n, g, 2)?; + } + } + "Dropout" => { + if !(0.0..1.0).contains(&float(n, "ratio", 0.5)) { + return shape_err("Dropout ratio is invalid"); + } + } + "MaxPool" => { + pair(n, "kernel_shape", 1)?; + pair(n, "strides", 1)?; + if n.ints("kernel_shape").is_none() + || n.int("ceil_mode").unwrap_or(0) != 0 + || n.int("storage_order").unwrap_or(0) != 0 + || n.ints("dilations").is_some_and(|v| v != [1, 1]) + || pads(n)? != (0, 0, 0, 0) + || n.string("auto_pad").is_some_and(|s| s != "NOTSET") + { + return unsupported("this MaxPool padding, dilation or ceil mode"); + } + } + "Resize" => { + floats(n, g, 2)?; + if n.string("mode").unwrap_or("nearest") != "nearest" || n.inputs.get(3).is_some_and(|s| !s.is_empty()) { + return unsupported("Resize requires nearest constant scales"); + } + if n.int("exclude_outside").unwrap_or(0) != 0 || float(n, "extrapolation_value", 0.0) != 0.0 { + return unsupported("Resize excludes or extrapolates outside values"); + } + let coord = n.string("coordinate_transformation_mode").unwrap_or("half_pixel"); + let nearest = n.string("nearest_mode").unwrap_or("round_prefer_floor"); + if !matches!((coord, nearest), ("asymmetric", "floor") | ("half_pixel", "round_prefer_floor")) { + return unsupported("this Resize coordinate/nearest mode"); + } + if n.inputs.get(1).is_some_and(|s| !s.is_empty()) && !floats(n, g, 1)?.is_empty() { + return unsupported("Resize with an ROI"); + } + } + "Reshape" => { + if n.int("allowzero").unwrap_or(0) != 0 { + return unsupported("Reshape allowzero"); + } + if !matches!(weight(n, g, 1)?.data, Data::I64(_)) { + return unsupported("Reshape needs a constant int64 shape"); + } + } + _ => {} + } + Ok(()) +} +pub(super) fn broadcast(a: &[usize], b: &[usize]) -> Result> { + elements(a)?; + elements(b)?; + let rank = a.len().max(b.len()); + let mut out = vec![1; rank]; + for (i, o) in out.iter_mut().enumerate() { + let aa = i.checked_sub(rank - a.len()).and_then(|j| a.get(j)).copied().unwrap_or(1); + let bb = i.checked_sub(rank - b.len()).and_then(|j| b.get(j)).copied().unwrap_or(1); + if aa != bb && aa != 1 && bb != 1 { + return shape_err("incompatible broadcast shapes"); + } + *o = aa.max(bb); + } + elements(&out)?; + Ok(out) +} +pub(super) fn axis(n: &Node, rank: usize) -> Result { + let raw = n.int("axis").unwrap_or(1); + let raw = if raw < 0 { raw.checked_add(rank as i64).ok_or_else(|| NetError::Shape("axis overflow".into()))? } else { raw }; + usize::try_from(raw).ok().filter(|&a| a <= rank).ok_or_else(|| NetError::Shape("Flatten axis is outside its rank".into())) +} +fn image(s: &[usize]) -> Result<(usize, usize, usize)> { + match s { + [1, c, h, w] => Ok((*c, *h, *w)), + _ => unsupported("a spatial operator needs [1,C,H,W]"), + } +} +pub(super) fn infer(g: &Graph, input: &[usize]) -> Result>> { + elements(input)?; + let mut shapes: HashMap<&str, Vec> = g.weights.iter().map(|(n, w)| (n.as_str(), w.dims.clone())).collect(); + shapes.insert(g.inputs.first().ok_or_else(|| NetError::Shape("no input".into()))?.0.as_str(), input.to_vec()); + let mut live = elements(input)?; + // Scratch grows independently and remains allocated across all subsequent layers. + let mut scratch_lower = 0; + let mut scratch_product = 0; + let mut last = HashMap::new(); + for (i, n) in g.nodes.iter().enumerate() { + for s in &n.inputs { + last.insert(s.as_str(), i); + } + } + for s in &g.outputs { + last.insert(s.as_str(), usize::MAX); + } + for (i, n) in g.nodes.iter().enumerate() { + let a = n.inputs.first().and_then(|s| shapes.get(s.as_str())).ok_or_else(|| NetError::Shape("missing input shape".into()))?; + elements(a)?; + let out = match n.op.as_str() { + "Conv" => { + let (cin, h, w) = image(a)?; + let dims = &weight(n, g, 1)?.dims; + let [co, cg, kh, kw] = dims.as_slice() else { + return shape_err("Conv dimensions"); + }; + let groups = usize::try_from(n.int("group").unwrap_or(1)).map_err(|_| NetError::Shape("Conv group".into()))?; + if groups.checked_mul(*cg) != Some(cin) || !co.is_multiple_of(groups) { + return shape_err("Conv channels do not match weights/groups"); + } + if n.inputs.get(2).is_some_and(|s| !s.is_empty()) && weight(n, g, 2)?.dims != [*co] { + return shape_err("Conv bias is not a channel vector"); + } + let (sh, sw) = pair(n, "strides", 1)?; + let (pt, pl, pb, pr) = pads(n)?; + let padded = |size: usize, p: usize, q: usize, k: usize, s: usize| { + size.checked_add(p) + .and_then(|v| v.checked_add(q)) + .and_then(|v| v.checked_sub(k)) + .and_then(|v| (v / s).checked_add(1)) + .ok_or_else(|| NetError::Shape("Conv padding overflow or oversized kernel".into())) + }; + vec![1, *co, padded(h, pt, pb, *kh, sh)?, padded(w, pl, pr, *kw, sw)?] + } + "BatchNormalization" => { + let channels = *a.get(1).ok_or_else(|| NetError::Shape("BatchNormalization rank must be at least two".into()))?; + for j in 1..5 { + if weight(n, g, j)?.dims != [channels] { + return shape_err("BatchNormalization parameters must be channel vectors"); + } + } + if floats(n, g, 4)?.iter().any(|v| *v < 0.0) { + return shape_err("BatchNormalization has negative variance"); + } + let eps = float(n, "epsilon", 1e-5); + let (s, b, m, v) = (floats(n, g, 1)?, floats(n, g, 2)?, floats(n, g, 3)?, floats(n, g, 4)?); + for (((&s, &b), &m), &v) in s.iter().zip(b).zip(m).zip(v) { + let variance = v + eps; + let coefficient = s / variance.sqrt(); + let offset = b - m * coefficient; + if !variance.is_finite() || !coefficient.is_finite() || !offset.is_finite() { + return shape_err("BatchNormalization coefficients overflow"); + } + } + + a.clone() + } + "PRelu" => { + if broadcast(a, &weight(n, g, 1)?.dims)? != *a { + return shape_err("PRelu slope would enlarge its input"); + } + a.clone() + } + "Add" | "Sub" | "Mul" => { + let b = n.inputs.get(1).and_then(|s| shapes.get(s.as_str())).ok_or_else(|| NetError::Shape("missing second input".into()))?; + broadcast(a, b)? + } + "Gemm" => { + let b = &weight(n, g, 1)?.dims; + let [ar, ac] = a.as_slice() else { + return shape_err("Gemm needs matrix A"); + }; + let [br, bc] = b.as_slice() else { + return shape_err("Gemm needs matrix B"); + }; + let (m, k) = if n.int("transA").unwrap_or(0) != 0 { (*ac, *ar) } else { (*ar, *ac) }; + let (bk, cols) = if n.int("transB").unwrap_or(0) != 0 { (*bc, *br) } else { (*br, *bc) }; + if k != bk { + return shape_err("Gemm inner dimensions differ"); + } + let s = vec![m, cols]; + if n.inputs.get(2).is_some_and(|s| !s.is_empty()) && broadcast(&s, &weight(n, g, 2)?.dims)? != s { + return shape_err("Gemm bias would enlarge the output"); + } + s + } + "Flatten" => { + let at = axis(n, a.len())?; + vec![ + elements(a.get(..at).ok_or_else(|| NetError::Shape("bad axis".into()))?)?, + elements(a.get(at..).ok_or_else(|| NetError::Shape("bad axis".into()))?)?, + ] + } + "GlobalAveragePool" => { + let (c, _, _) = image(a)?; + vec![1, c, 1, 1] + } + "MaxPool" => { + let (c, h, w) = image(a)?; + let (kh, kw) = pair(n, "kernel_shape", 1)?; + let (sh, sw) = pair(n, "strides", 1)?; + if h < kh || w < kw { + return shape_err("MaxPool kernel exceeds input"); + } + vec![1, c, (h - kh) / sh + 1, (w - kw) / sw + 1] + } + "Resize" => { + let (c, h, w) = image(a)?; + let s = floats(n, g, 2)?; + if s.len() != 4 + || s.first() != Some(&1.0) + || s.get(1) != Some(&1.0) + || s.iter().any(|&v| !(1.0..=16.0).contains(&v) || v.fract() != 0.0) + { + return unsupported("Resize supports whole spatial scales from 1 to 16"); + } + vec![ + 1, + c, + h.checked_mul(s.get(2).copied().unwrap_or(0.0) as usize).ok_or_else(|| NetError::Shape("Resize overflow".into()))?, + w.checked_mul(s.get(3).copied().unwrap_or(0.0) as usize).ok_or_else(|| NetError::Shape("Resize overflow".into()))?, + ] + } + "Transpose" => { + let perm: Vec = n.ints("perm").map_or_else(|| (0..a.len() as i64).rev().collect(), Vec::from); + let mut seen = std::collections::HashSet::new(); + let mut s = Vec::new(); + if perm.len() != a.len() { + return shape_err("Transpose rank mismatch"); + } + for d in perm { + let d = usize::try_from(d).map_err(|_| NetError::Shape("negative permutation".into()))?; + if !seen.insert(d) { + return shape_err("repeated permutation axis"); + } + s.push(*a.get(d).ok_or_else(|| NetError::Shape("bad permutation axis".into()))?); + } + s + } + "Reshape" => { + let Data::I64(dims) = &weight(n, g, 1)?.data else { + return shape_err("Reshape datatype"); + }; + if dims.len() > 8 { + return shape_err("Reshape rank exceeds eight"); + } + let total = elements(a)?; + let mut s = Vec::new(); + let mut infer = None; + for (j, &v) in dims.iter().enumerate() { + s.push(match v { + -1 if infer.is_none() => { + infer = Some(j); + 1 + } + 0 => *a.get(j).ok_or_else(|| NetError::Shape("Reshape copies missing dimension".into()))?, + v if v > 0 => usize::try_from(v).map_err(|_| NetError::Shape("Reshape dimension overflow".into()))?, + _ => return shape_err("invalid Reshape dimension"), + }); + } + let known = elements(&s)?; + if let Some(j) = infer { + if !total.is_multiple_of(known) { + return shape_err("Reshape does not divide evenly"); + } + *s.get_mut(j).ok_or_else(|| NetError::Shape("missing inferred dimension".into()))? = total / known; + } + if elements(&s)? != total { + return shape_err("Reshape changes element count"); + } + s + } + _ => a.clone(), + }; + let out_count = elements(&out)?; + let mut copied_constants = 0usize; + for (j, name) in n.inputs.iter().enumerate() { + // Execution borrows operator weights, but get(0) and binary get(1) materialize constants. + if (j == 0 || (j == 1 && matches!(n.op.as_str(), "Add" | "Sub" | "Mul"))) + && let Some(w) = g.weights.get(name) + { + copied_constants = + copied_constants.checked_add(elements(&w.dims)?).ok_or_else(|| NetError::Shape("constant copies overflow".into()))?; + } + } + if n.op == "Conv" { + let wd = &weight(n, g, 1)?.dims; + let [co, cg, kh, kw] = wd.as_slice() else { + return shape_err("Conv rank"); + }; + let block = 256.min(out.get(2).copied().unwrap_or(1) * out.get(3).copied().unwrap_or(1)); + let lower = cg + .checked_mul(*kh) + .and_then(|v| v.checked_mul(*kw)) + .and_then(|v| v.checked_mul(block)) + .filter(|&v| v <= MAX_TENSOR) + .ok_or_else(|| NetError::Shape("convolution scratch exceeds the limit".into()))?; + scratch_lower = scratch_lower.max(lower); + scratch_product = scratch_product.max(co * block); + } + let scratch = scratch_lower + scratch_product; + if live + .checked_add(out_count) + .and_then(|v| v.checked_add(scratch)) + .and_then(|v| v.checked_add(copied_constants)) + .is_none_or(|v| v > 128 * 1024 * 1024) + { + return shape_err("tensors and scratch exceed the 512 MiB working limit"); + } + live = live + .checked_add(out_count) + .filter(|&n| n <= 128 * 1024 * 1024) + .ok_or_else(|| NetError::Shape("live tensors exceed the 512 MiB working limit".into()))?; + shapes.insert(n.outputs.first().ok_or_else(|| NetError::Shape("no output".into()))?, out); + let mut removed = std::collections::HashSet::new(); + for name in &n.inputs { + if last.get(name.as_str()) == Some(&i) && !g.weights.contains_key(name) && removed.insert(name) { + live = live.saturating_sub(shapes.get(name.as_str()).map_or(0, |s| elements(s).unwrap_or(0))); + } + } + } + g.outputs.iter().map(|s| shapes.get(s.as_str()).cloned().ok_or_else(|| NetError::Shape("missing output".into()))).collect() +} diff --git a/crates/faces/src/onnx.rs b/crates/faces/src/onnx.rs new file mode 100644 index 00000000..453be241 --- /dev/null +++ b/crates/faces/src/onnx.rs @@ -0,0 +1,514 @@ +//! A streaming probe of an ONNX file: the graph's inputs and outputs (name, element type, shape), without +//! running or even fully loading the model. +//! +//! ONNX is protobuf. We read only the few fields we need (the producer, the opset, and the graph's input, +//! output and initializer names) and seek past everything else, so a 260 MB model costs a few kilobytes and +//! a few milliseconds. The file is hostile input: every length is checked against its parent message and +//! the file size, the number of fields read is capped, strings are bounded, and nothing indexes or unwraps. + +use std::collections::HashSet; +use std::io::{self, Read, Seek, SeekFrom}; +use std::path::Path; + +const MAX_FIELDS: u32 = 5_000_000; +const MAX_STRING: u64 = 4096; +const MAX_DIMS: usize = 16; +const MAX_TENSORS: usize = 64; +/// Older exporters list every weight as a graph input too; that many are read before the weights are filtered out. +const MAX_LISTED_INPUTS: usize = 200_000; +const MAX_INITIALIZERS: usize = 1_000_000; + +#[derive(Clone, Debug, PartialEq)] +pub enum Dim { + Fixed(u64), + /// A named or unknown size (typically the batch). + Dynamic(String), +} + +#[derive(Clone, Debug, PartialEq)] +pub struct TensorInfo { + pub name: String, + /// ONNX element type (1 = float32). + pub elem_type: u32, + pub shape: Vec, +} + +#[derive(Clone, Debug, Default, PartialEq)] +pub struct OnnxInfo { + pub producer: String, + /// The default-domain opset version. + pub opset: Option, + /// Real inputs only (weights that older exporters also list as inputs are left out). + pub inputs: Vec, + pub outputs: Vec, +} + +#[derive(Debug, thiserror::Error)] +pub enum ProbeError { + #[error("not an ONNX model: {0}")] + Malformed(&'static str), + #[error("could not read the model: {0}")] + Io(#[from] io::Error), +} + +pub(crate) type Result = std::result::Result; + +pub(crate) enum Wire { + Varint(u64), + /// A length-delimited field: its bytes are `start..end` and the reader sits at `start`. + Len { + start: u64, + end: u64, + }, + Fixed32(u32), + #[allow(dead_code)] + Fixed64(u64), +} + +pub(crate) struct Walker<'a, R> { + pub(crate) r: &'a mut R, + pub(crate) total: u64, + pub(crate) fields: u32, +} + +impl Walker<'_, R> { + pub(crate) fn pos(&mut self) -> Result { + Ok(self.r.stream_position()?) + } + + pub(crate) fn seek_to(&mut self, at: u64) -> Result<()> { + self.r.seek(SeekFrom::Start(at))?; + Ok(()) + } + + pub(crate) fn byte(&mut self) -> Result { + let mut b = [0u8; 1]; + self.r + .read_exact(&mut b) + .map_err(|e| if e.kind() == io::ErrorKind::UnexpectedEof { ProbeError::Malformed("file ends inside a field") } else { e.into() })?; + let [x] = b; + Ok(x) + } + + pub(crate) fn varint(&mut self) -> Result { + let mut v = 0u64; + for shift in (0..70u32).step_by(7) { + let b = self.byte()?; + if shift == 63 && b > 1 { + return Err(ProbeError::Malformed("number overflows 64 bits")); + } + v |= u64::from(b & 0x7f) << shift; + if b & 0x80 == 0 { + return Ok(v); + } + } + Err(ProbeError::Malformed("number too long")) + } + + /// The next field of the message that ends at `end`, or `None` at its end. + pub(crate) fn next(&mut self, end: u64) -> Result> { + let here = self.pos()?; + if here >= end { + return Ok(None); + } + self.fields += 1; + if self.fields > MAX_FIELDS { + return Err(ProbeError::Malformed("too many fields")); + } + let tag = self.varint()?; + let number = u32::try_from(tag >> 3).map_err(|_| ProbeError::Malformed("field number too large"))?; + if number == 0 || number > 0x1fff_ffff { + return Err(ProbeError::Malformed("invalid field number")); + } + let wire = match tag & 7 { + 0 => Wire::Varint(self.varint()?), + 1 => { + let at = self.pos()?.checked_add(8).ok_or(ProbeError::Malformed("offset overflow"))?; + if at > end || at > self.total { + return Err(ProbeError::Malformed("field runs past its message")); + } + let mut b = [0u8; 8]; + self.r.read_exact(&mut b).map_err(|_| ProbeError::Malformed("file ends inside a field"))?; + Wire::Fixed64(u64::from_le_bytes(b)) + } + 5 => { + let at = self.pos()?.checked_add(4).ok_or(ProbeError::Malformed("offset overflow"))?; + if at > end || at > self.total { + return Err(ProbeError::Malformed("field runs past its message")); + } + let mut b = [0u8; 4]; + self.r.read_exact(&mut b).map_err(|_| ProbeError::Malformed("file ends inside a field"))?; + Wire::Fixed32(u32::from_le_bytes(b)) + } + 2 => { + let len = self.varint()?; + let start = self.pos()?; + let stop = start.checked_add(len).ok_or(ProbeError::Malformed("offset overflow"))?; + if stop > end || stop > self.total { + return Err(ProbeError::Malformed("field runs past its message")); + } + Wire::Len { start, end: stop } + } + _ => return Err(ProbeError::Malformed("unsupported field type")), + }; + if self.pos()? > end { + return Err(ProbeError::Malformed("field runs past its message")); + } + Ok(Some((number, wire))) + } + + /// A bounded UTF-8 string occupying `start..end`; never silently truncate graph names. + pub(crate) fn string(&mut self, start: u64, end: u64) -> Result { + let len = end.saturating_sub(start); + if len > MAX_STRING { + return Err(ProbeError::Malformed("a string exceeds 4096 bytes")); + } + let mut buf = Vec::new(); + self.r.by_ref().take(len).read_to_end(&mut buf)?; + String::from_utf8(buf).map_err(|_| ProbeError::Malformed("a string is not UTF-8")) + } + + /// The bytes of `start..end`, which must be at most `max` long. + pub(crate) fn bytes(&mut self, start: u64, end: u64, max: u64) -> Result> { + let len = end.saturating_sub(start); + if len > max { + return Err(ProbeError::Malformed("a field is larger than allowed")); + } + let mut buf = Vec::new(); + self.r.by_ref().take(len).read_to_end(&mut buf)?; + if buf.len() as u64 != len { + return Err(ProbeError::Malformed("file ends inside a field")); + } + Ok(buf) + } +} + +fn dim(w: &mut Walker, end: u64) -> Result> { + let mut out = None; + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Varint(v)) => out = Some(Dim::Fixed(v)), + (2, Wire::Len { start, end: e }) => { + out = Some(Dim::Dynamic(w.string(start, e)?)); + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(out) +} + +fn shape(w: &mut Walker, end: u64) -> Result> { + let mut dims = Vec::new(); + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Len { end: e, .. }) => { + if dims.len() >= MAX_DIMS { + return Err(ProbeError::Malformed("too many dimensions")); + } + dims.push(dim(w, e)?.unwrap_or(Dim::Dynamic(String::new()))); + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(dims) +} + +fn tensor_type(w: &mut Walker, end: u64, info: &mut TensorInfo) -> Result<()> { + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Varint(v)) => info.elem_type = u32::try_from(v).unwrap_or(0), + (2, Wire::Len { end: e, .. }) => { + info.shape = shape(w, e)?; + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(()) +} + +pub(crate) fn type_proto(w: &mut Walker, end: u64, info: &mut TensorInfo) -> Result<()> { + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Len { end: e, .. }) => { + tensor_type(w, e, info)?; + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(()) +} + +fn value_info(w: &mut Walker, end: u64) -> Result { + let mut info = TensorInfo { name: String::new(), elem_type: 0, shape: Vec::new() }; + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Len { start, end: e }) => { + info.name = w.string(start, e)?; + w.seek_to(e)?; + } + (2, Wire::Len { end: e, .. }) => { + type_proto(w, e, &mut info)?; + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(info) +} + +/// The `name` (field 8) of a weight tensor; its bulk data is skipped, never read. +fn tensor_name(w: &mut Walker, end: u64) -> Result> { + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (8, Wire::Len { start, end: e }) => return Ok(Some(w.string(start, e)?)), + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(None) +} + +fn graph(w: &mut Walker, end: u64, info: &mut OnnxInfo, weights: &mut HashSet) -> Result<()> { + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (5, Wire::Len { end: e, .. }) => { + if weights.len() < MAX_INITIALIZERS + && let Some(name) = tensor_name(w, e)? + { + weights.insert(name); + } + w.seek_to(e)?; + } + (11, Wire::Len { end: e, .. }) => { + if info.inputs.len() >= MAX_LISTED_INPUTS { + return Err(ProbeError::Malformed("too many inputs")); + } + info.inputs.push(value_info(w, e)?); + w.seek_to(e)?; + } + (12, Wire::Len { end: e, .. }) => { + if info.outputs.len() >= MAX_TENSORS { + return Err(ProbeError::Malformed("too many outputs")); + } + info.outputs.push(value_info(w, e)?); + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(()) +} + +fn opset(w: &mut Walker, end: u64) -> Result> { + let (mut domain, mut version) = (String::new(), None); + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Len { start, end: e }) => { + domain = w.string(start, e)?; + w.seek_to(e)?; + } + (2, Wire::Varint(v)) => version = Some(v), + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(if domain.is_empty() || domain == "ai.onnx" { version } else { None }) +} + +/// Read an ONNX model's graph inputs and outputs from any seekable source. +pub fn probe_reader(r: &mut R) -> Result { + let total = r.seek(SeekFrom::End(0))?; + if total > crate::manifest::MAX_MODEL_BYTES { + return Err(ProbeError::Malformed("the model exceeds 512 MiB")); + } + r.seek(SeekFrom::Start(0))?; + if total == 0 { + return Err(ProbeError::Malformed("the file is empty")); + } + let mut w = Walker { r, total, fields: 0 }; + let (mut info, mut weights, mut saw_graph) = (OnnxInfo::default(), HashSet::new(), false); + while let Some((n, f)) = w.next(total)? { + match (n, f) { + (2, Wire::Len { start, end }) => { + info.producer = w.string(start, end)?; + w.seek_to(end)?; + } + (7, Wire::Len { end, .. }) => { + saw_graph = true; + graph(&mut w, end, &mut info, &mut weights)?; + w.seek_to(end)?; + } + (8, Wire::Len { end, .. }) => { + if let Some(v) = opset(&mut w, end)? { + info.opset = Some(v); + } + w.seek_to(end)?; + } + (_, Wire::Len { end, .. }) => w.seek_to(end)?, + _ => {} + } + } + if !saw_graph { + return Err(ProbeError::Malformed("no graph")); + } + info.inputs.retain(|t| !weights.contains(&t.name)); + if info.inputs.len() > MAX_TENSORS { + return Err(ProbeError::Malformed("too many inputs")); + } + if info.inputs.is_empty() || info.outputs.is_empty() { + return Err(ProbeError::Malformed("the graph has no inputs or no outputs")); + } + Ok(info) +} + +/// Read an ONNX file's graph inputs and outputs. +pub fn probe_path(path: &Path) -> Result { + probe_reader(&mut io::BufReader::new(std::fs::File::open(path)?)) +} + +#[cfg(test)] +pub(crate) mod tests { + use super::*; + use std::io::Cursor; + + pub(crate) use crate::synthetic::embedder_model as model_bytes; + use crate::synthetic::{len_field, value_info_bytes, varint}; + + #[test] + fn reads_inputs_outputs_and_skips_weights() { + let info = probe_reader(&mut Cursor::new(model_bytes(512))).unwrap(); + assert_eq!(info.producer, "pytorch"); + assert_eq!(info.opset, Some(17)); + assert_eq!(info.inputs.len(), 1, "the weight listed as an input is left out"); + assert_eq!(info.inputs[0].name, "data"); + assert_eq!(info.inputs[0].elem_type, 1); + assert_eq!(info.inputs[0].shape, vec![Dim::Dynamic("N".into()), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]); + assert_eq!(info.outputs[0].shape, vec![Dim::Dynamic("N".into()), Dim::Fixed(512)]); + } + + #[test] + fn old_style_files_listing_hundreds_of_weights_as_inputs_still_probe() { + let mut graph = Vec::new(); + len_field(11, &value_info_bytes("data", 1, &[Err("N"), Ok(3), Ok(112), Ok(112)]), &mut graph); + for i in 0..500 { + let name = format!("weight{i}"); + len_field(11, &value_info_bytes(&name, 1, &[Ok(4)]), &mut graph); + let mut tensor = Vec::new(); + len_field(8, name.as_bytes(), &mut tensor); + len_field(5, &tensor, &mut graph); + } + len_field(12, &value_info_bytes("emb", 1, &[Err("N"), Ok(128)]), &mut graph); + let mut m = Vec::new(); + len_field(7, &graph, &mut m); + let info = probe_reader(&mut Cursor::new(m)).unwrap(); + assert_eq!(info.inputs.len(), 1); + assert_eq!(info.inputs[0].name, "data"); + } + + #[test] + fn many_real_inputs_are_refused() { + let mut graph = Vec::new(); + for i in 0..70 { + len_field(11, &value_info_bytes(&format!("in{i}"), 1, &[Ok(4)]), &mut graph); + } + len_field(12, &value_info_bytes("emb", 1, &[Ok(4)]), &mut graph); + let mut m = Vec::new(); + len_field(7, &graph, &mut m); + assert!(probe_reader(&mut Cursor::new(m)).is_err()); + } + + #[test] + fn truncated_at_every_length_never_panics() { + let bytes = model_bytes(128); + for n in 0..bytes.len() { + let _ = probe_reader(&mut Cursor::new(bytes[..n].to_vec())); + } + assert!(probe_reader(&mut Cursor::new(bytes)).is_ok()); + } + + #[test] + fn hostile_files_are_errors() { + let mut lying = Vec::new(); + varint(7 << 3 | 2, &mut lying); + varint(1 << 40, &mut lying); // a graph 1 TB long in a 10-byte file + let mut overflow = Vec::new(); + varint(7 << 3 | 2, &mut overflow); + varint(u64::MAX, &mut overflow); + let mut long_number = vec![0x08]; + long_number.extend_from_slice(&[0xff; 11]); + for (what, bytes) in [ + ("empty", Vec::new()), + ("zeros", vec![0u8; 64]), + ("text", b"this is not an onnx model at all".to_vec()), + ("lying length", lying), + ("overflowing length", overflow), + ("endless number", long_number), + ("no graph", { + let mut m = Vec::new(); + len_field(2, b"x", &mut m); + m + }), + ("graph without io", { + let mut m = Vec::new(); + len_field(7, &[], &mut m); + m + }), + ] { + assert!(probe_reader(&mut Cursor::new(bytes)).is_err(), "{what}"); + } + } + + #[test] + fn too_many_dimensions_or_tensors_are_errors() { + let dims: Vec> = (0..40).map(Ok).collect(); + let mut graph = Vec::new(); + len_field(11, &value_info_bytes("data", 1, &dims), &mut graph); + len_field(12, &value_info_bytes("out", 1, &[Ok(1)]), &mut graph); + let mut m = Vec::new(); + len_field(7, &graph, &mut m); + assert!(probe_reader(&mut Cursor::new(m)).is_err()); + } + + proptest::proptest! { + #[test] + fn arbitrary_bytes_never_panic(bytes in proptest::collection::vec(proptest::prelude::any::(), 0..400)) { + let _ = probe_reader(&mut Cursor::new(bytes)); + } + + #[test] + fn a_valid_model_with_flipped_bytes_never_panics(flips in proptest::collection::vec((0usize..400, proptest::prelude::any::()), 1..6)) { + let mut bytes = model_bytes(512); + for (at, v) in flips { + if let Some(b) = bytes.get_mut(at) { + *b = v; + } + } + let _ = probe_reader(&mut Cursor::new(bytes)); + } + } + #[test] + fn oversized_models_are_rejected_without_reading_or_allocating() { + struct Oversized; + impl std::io::Read for Oversized { + fn read(&mut self, _: &mut [u8]) -> std::io::Result { + panic!("an oversized model must not be read") + } + } + impl std::io::Seek for Oversized { + fn seek(&mut self, _: std::io::SeekFrom) -> std::io::Result { + Ok(crate::manifest::MAX_MODEL_BYTES + 1) + } + } + assert!(probe_reader(&mut Oversized).is_err()); + } +} diff --git a/crates/faces/src/runtime.rs b/crates/faces/src/runtime.rs new file mode 100644 index 00000000..bbb1457e --- /dev/null +++ b/crates/faces/src/runtime.rs @@ -0,0 +1,352 @@ +//! Installed recognition models on LightCraft's own checked, single-thread float32 CPU engine. +//! The detector and recognisers share graph validation and convolution kernels. + +use std::panic::{AssertUnwindSafe, catch_unwind}; +use std::path::Path; +use std::sync::Arc; +use web_time::Instant; + +use crate::{ + graph, + net::{Net, Tensor}, +}; +use serde::Serialize; + +use crate::manifest::{Colour, ModelManifest, OutputSpec}; + +#[derive(Debug, thiserror::Error)] +pub enum RuntimeError { + #[error("the model could not be loaded: {0}")] + Load(String), + #[error("the model failed while running: {0}")] + Run(String), + #[error("{0}")] + Input(String), +} + +/// What `self_test` found. +#[derive(Clone, Debug, PartialEq, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct SelfTest { + pub ok: bool, + /// Milliseconds to load and optimise the model. + pub load_ms: f64, + /// Median milliseconds to embed one face. + pub embed_ms: f64, + pub dimension: usize, + /// Each check: what was checked and whether it held. + pub checks: Vec<(String, bool)>, +} + +type Plan = Net; + +/// A loaded recognition model. +#[derive(Clone)] +pub struct Embedder { + plan: Arc, + manifest: ModelManifest, + width: usize, + height: usize, + dim: usize, + load_ms: f64, +} + +fn guarded(what: &str, f: impl FnOnce() -> Result) -> Result { + match catch_unwind(AssertUnwindSafe(f)) { + Ok(r) => r, + Err(_) => Err(RuntimeError::Run(format!("the runtime gave up on this model while {what}"))), + } +} + +impl Embedder { + /// Load the `.onnx` at `path` as described by `manifest` (an embedder). + pub fn load(path: &Path, manifest: &ModelManifest) -> Result { + let OutputSpec::Embedding { dim } = manifest.output else { + return Err(RuntimeError::Load("this is not a recognition model".into())); + }; + let (w, h) = (manifest.input.width as usize, manifest.input.height as usize); + let started = Instant::now(); + crate::manifest::validate(manifest).map_err(|e| RuntimeError::Load(e.to_string()))?; + let mut file = std::fs::File::open(path).map_err(|e| RuntimeError::Load(e.to_string()))?; + let length = file.metadata().map_err(|e| RuntimeError::Load(e.to_string()))?.len(); + if length == 0 || length > 512 * 1024 * 1024 { + return Err(RuntimeError::Load("the model file is empty or exceeds 512 MiB".into())); + } + let mut bytes = Vec::new(); + bytes + .try_reserve_exact(usize::try_from(length).map_err(|_| RuntimeError::Load("model size overflow".into()))?) + .map_err(|_| RuntimeError::Load("not enough memory for the model".into()))?; + use std::io::Read; + (&mut file).take(512 * 1024 * 1024 + 1).read_to_end(&mut bytes).map_err(|e| RuntimeError::Load(e.to_string()))?; + let graph = graph::load(&bytes).map_err(|e| RuntimeError::Load(e.to_string()))?; + let input = graph.inputs.first().ok_or_else(|| RuntimeError::Load("no input".into()))?; + let want = [1, 3, h, w]; + if input.1.len() != 4 || input.1.iter().zip(want).any(|(&d, w)| d != 0 && d != w) { + return Err(RuntimeError::Load("the graph input differs from the manifest's [1,3,H,W]".into())); + } + let plan = Net::new(graph).map_err(|e| RuntimeError::Load(e.to_string()))?; + let shapes = plan.validate_input(&want).map_err(|e| RuntimeError::Load(e.to_string()))?; + if shapes.len() != 1 || shapes.first().is_none_or(|s| s.first() != Some(&1) || s.iter().product::() != dim as usize) { + return Err(RuntimeError::Load("the graph output differs from the manifest's embedding dimension".into())); + } + Ok(Embedder { + plan: Arc::new(plan), + manifest: manifest.clone(), + width: w, + height: h, + dim: dim as usize, + load_ms: started.elapsed().as_secs_f64() * 1000.0, + }) + } + + /// The size of the aligned face picture the model wants (width, height). + pub fn input_size(&self) -> (usize, usize) { + (self.width, self.height) + } + + pub fn manifest(&self) -> &ModelManifest { + &self.manifest + } + + /// The face as a unit-length vector. `rgb` is an aligned face at exactly [`Self::input_size`], 8-bit RGB. + pub fn embed(&self, rgb: &[u8]) -> Result, RuntimeError> { + if rgb.len() != self.width * self.height * 3 { + return Err(RuntimeError::Input("the face picture is not the size the model wants".into())); + } + let (mean, std, bgr) = (self.manifest.input.mean, self.manifest.input.std, self.manifest.input.colour == Colour::Bgr); + let plane = self.width * self.height; + let mut data = vec![0.0f32; 3 * plane]; + for (i, px) in rgb.as_chunks::<3>().0.iter().enumerate() { + for c in 0..3 { + let src = if bgr { 2 - c } else { c }; + if let (Some(v), Some(m), Some(s), Some(o)) = (px.get(src), mean.get(c), std.get(c), data.get_mut(c * plane + i)) { + *o = (f32::from(*v) - m) / s; + } + } + } + let out = guarded("running it", || { + let input = Tensor::new(vec![1, 3, self.height, self.width], data).map_err(|e| RuntimeError::Input(e.to_string()))?; + let mut result = self.plan.run(input).map_err(|e| RuntimeError::Run(e.to_string()))?; + if result.len() != 1 { + return Err(RuntimeError::Run("the model must give one embedding".into())); + } + Ok(result.pop().ok_or_else(|| RuntimeError::Run("no output".into()))?.1.data) + })?; + if out.len() != self.dim { + return Err(RuntimeError::Run(format!("the model gave {} numbers per face, not the {} its description says", out.len(), self.dim))); + } + if out.iter().any(|v| !v.is_finite()) { + return Err(RuntimeError::Run("the model gave numbers that are not finite".into())); + } + let norm = out.iter().map(|&v| f64::from(v).powi(2)).sum::().sqrt(); + if !norm.is_finite() || norm < 1e-9 { + return Err(RuntimeError::Run("the model gave an all-zero face".into())); + } + Ok(out.into_iter().map(|v| (f64::from(v) / norm) as f32).collect()) + } + + /// Check the model does something sensible before it is ever used on a photo: the right number of finite + /// numbers, the same picture twice gives the same face, and two different pictures give different faces. + pub fn self_test(&self) -> SelfTest { + let mut checks: Vec<(String, bool)> = Vec::new(); + let n = self.width * self.height * 3; + // two deterministic, different test pictures: noise, and a diagonal gradient with a colour cast + let mut seed = 0x2545_f491u32; + let noise: Vec = (0..n) + .map(|_| { + seed ^= seed << 13; + seed ^= seed >> 17; + seed ^= seed << 5; + (seed >> 8) as u8 + }) + .collect(); + let gradient: Vec = (0..n) + .map(|i| (((i / 3) % self.width + (i / 3) / self.width) * 255 / (self.width + self.height) + (i % 3) * 40).min(255) as u8) + .collect(); + let mut times = Vec::new(); + let mut run = |img: &[u8]| { + let t = Instant::now(); + let r = self.embed(img); + times.push(t.elapsed().as_secs_f64() * 1000.0); + r + }; + let (a, a2, b) = (run(&noise), run(&noise), run(&gradient)); + let dot = |x: &[f32], y: &[f32]| x.iter().zip(y).map(|(p, q)| p * q).sum::(); + match (&a, &a2, &b) { + (Ok(a), Ok(a2), Ok(b)) => { + checks.push((format!("gives {} finite numbers per face", a.len()), true)); + checks.push(("gives the same face for the same picture".into(), dot(a, a2) > 0.9999)); + checks.push(("gives different faces for different pictures".into(), dot(a, b) < 0.9995)); + } + (Err(e), _, _) | (_, Err(e), _) | (_, _, Err(e)) => checks.push((format!("runs: {e}"), false)), + } + times.sort_by(|x, y| x.total_cmp(y)); + SelfTest { + ok: checks.iter().all(|(_, ok)| *ok), + load_ms: self.load_ms, + embed_ms: times.get(times.len() / 2).copied().unwrap_or(0.0), + dimension: self.dim, + checks, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::manifest::{InputSpec, Licence, OutputSpec, Role, Thresholds}; + + fn manifest(dim: u32) -> ModelManifest { + ModelManifest { + id: "tiny".into(), + name: "Tiny".into(), + version: "1".into(), + role: Role::Embedder, + licence: Licence::default(), + source: None, + sha256: None, + size_bytes: None, + provenance: String::new(), + input: InputSpec::default(), + output: OutputSpec::Embedding { dim }, + thresholds: Thresholds::default(), + speed: None, + } + } + + fn temp_model(name: &str, bytes: &[u8]) -> std::path::PathBuf { + let p = std::env::temp_dir().join(format!("lc-faces-runtime-{name}-{}.onnx", std::process::id())); + std::fs::write(&p, bytes).unwrap(); + p + } + + #[test] + fn a_working_model_loads_embeds_and_passes_the_self_test() { + let p = temp_model("ok", &crate::synthetic::tiny_embedder_model(16)); + let e = Embedder::load(&p, &manifest(16)).unwrap(); + assert_eq!(e.input_size(), (112, 112)); + let v = e.embed(&vec![100u8; 112 * 112 * 3]).unwrap(); + assert_eq!(v.len(), 16); + assert!((v.iter().map(|x| x * x).sum::() - 1.0).abs() < 1e-4, "unit length"); + let t = e.self_test(); + assert!(t.ok, "{t:?}"); + assert_eq!(t.dimension, 16); + let _ = std::fs::remove_file(p); + } + + #[test] + fn wrong_sizes_wrong_dimensions_and_broken_files_are_errors() { + let p = temp_model("dim", &crate::synthetic::tiny_embedder_model(16)); + let e = Embedder::load(&p, &manifest(16)).unwrap(); + assert!(e.embed(&[0u8; 10]).is_err(), "wrong picture size"); + // the manifest promises 32 numbers, the model gives 16 + assert!(Embedder::load(&p, &manifest(32)).is_err()); + // not a model, an empty file, a graph with no layers, a detector-shaped manifest + for (name, bytes) in [("junk", b"nonsense".to_vec()), ("empty", Vec::new()), ("nolayers", crate::synthetic::embedder_model(16))] { + let q = temp_model(name, &bytes); + assert!(Embedder::load(&q, &manifest(16)).is_err(), "{name}"); + let _ = std::fs::remove_file(q); + } + let mut detector = manifest(16); + detector.output = OutputSpec::Detector { decoder: "yunet-v2".into() }; + assert!(Embedder::load(&p, &detector).is_err()); + assert!(Embedder::load(&std::env::temp_dir().join("does-not-exist.onnx"), &manifest(16)).is_err()); + let _ = std::fs::remove_file(p); + } + + #[test] + fn truncated_and_corrupted_models_never_panic() { + let good = crate::synthetic::tiny_embedder_model(16); + for n in (0..good.len()).step_by(7) { + let q = temp_model("trunc", &good[..n]); + let _ = Embedder::load(&q, &manifest(16)); + let _ = std::fs::remove_file(q); + } + for at in (0..good.len()).step_by(11) { + let mut bad = good.clone(); + if let Some(b) = bad.get_mut(at) { + *b = b.wrapping_add(0x55); + } + let q = temp_model("flip", &bad); + if let Ok(e) = Embedder::load(&q, &manifest(16)) { + let _ = e.embed(&vec![7u8; 112 * 112 * 3]); + } + let _ = std::fs::remove_file(q); + } + } + + /// Opt-in: `LC_FACE_MODELS= cargo test -p lightcraft-faces -- --ignored --nocapture` + /// self-tests every known recognition model found there. + #[test] + #[ignore = "needs real model files: set LC_FACE_MODELS"] + fn real_models_pass_the_self_test() { + let Some(dir) = std::env::var_os("LC_FACE_MODELS") else { return }; + for entry in std::fs::read_dir(dir).unwrap().flatten() { + let path = entry.path(); + let Ok(sha) = crate::hash::sha256_file(&path) else { continue }; + let Some(m) = crate::known::lookup(&sha).filter(|m| m.role == Role::Embedder) else { continue }; + let e = Embedder::load(&path, &m).unwrap(); + let t = e.self_test(); + println!("{}: ok {} | load {:.0} ms | {:.0} ms per face | {} numbers | {:?}", m.id, t.ok, t.load_ms, t.embed_ms, t.dimension, t.checks); + assert!(t.ok); + } + } + + #[test] + fn bgr_changes_the_input() { + let p = temp_model("bgr", &crate::synthetic::tiny_embedder_model(16)); + let rgb = Embedder::load(&p, &manifest(16)).unwrap(); + let mut m = manifest(16); + m.input.colour = crate::manifest::Colour::Bgr; + let bgr = Embedder::load(&p, &m).unwrap(); + // a picture with different channel levels: swapping the order must change the embedding + let img: Vec = (0..112 * 112).flat_map(|_| [200u8, 100, 20]).collect(); + let (a, b) = (rgb.embed(&img).unwrap(), bgr.embed(&img).unwrap()); + let cos: f32 = a.iter().zip(&b).map(|(x, y)| x * y).sum(); + assert!(cos < 0.9999, "{cos}"); + let _ = std::fs::remove_file(p); + } + #[test] + fn large_finite_embeddings_have_a_finite_unit_norm() { + let p = temp_model("large-norm", &crate::synthetic::norm_overflow_model()); + let e = Embedder::load(&p, &manifest(3)).unwrap(); + let out = e.embed(&vec![255; 112 * 112 * 3]).unwrap(); + assert!(out.iter().all(|v| v.is_finite())); + assert!((out.iter().map(|v| v * v).sum::() - 1.0).abs() < 1e-6); + let _ = std::fs::remove_file(p); + } + #[test] + fn invalid_manifest_numbers_are_rejected_before_preprocessing() { + let p = temp_model("invalid-manifest", &crate::synthetic::tiny_embedder_model(16)); + for case in 0..5 { + let mut m = manifest(16); + match case { + 0 => m.input.std = [0.0; 3], + 1 => m.input.mean = [f32::NAN; 3], + 2 => m.input.width = 0, + 3 => m.input.width = u32::MAX, + _ => m.output = OutputSpec::Embedding { dim: 0 }, + }; + assert!(Embedder::load(&p, &m).is_err()); + } + let _ = std::fs::remove_file(p); + } + #[test] + fn concurrent_faces_share_weights_and_give_deterministic_embeddings() { + let p = temp_model("concurrent", &crate::synthetic::tiny_embedder_model(16)); + let e = Embedder::load(&p, &manifest(16)).unwrap(); + let input = vec![123; 112 * 112 * 3]; + let want = e.embed(&input).unwrap(); + let threads: Vec<_> = (0..8) + .map(|_| { + let e = e.clone(); + let input = input.clone(); + std::thread::spawn(move || e.embed(&input)) + }) + .collect(); + for thread in threads { + assert_eq!(thread.join().unwrap().unwrap(), want); + } + let _ = std::fs::remove_file(p); + } +} diff --git a/crates/faces/src/suggest.rs b/crates/faces/src/suggest.rs new file mode 100644 index 00000000..377452e4 --- /dev/null +++ b/crates/faces/src/suggest.rs @@ -0,0 +1,237 @@ +//! What to do with a model file the user drops in: recognise it by hash, or look at its shape and propose a +//! manifest the user confirms, or say plainly why it cannot be used yet. + +use crate::manifest::{Commercial, InputSpec, Licence, ModelManifest, OutputSpec, Role, Thresholds}; +use crate::onnx::{Dim, OnnxInfo, TensorInfo}; + +/// ONNX element type for float32. +const FLOAT32: u32 = 1; + +#[derive(Clone, Debug, PartialEq)] +pub enum Suggestion { + /// A model we know by its SHA-256: its manifest as it is. + Known(ModelManifest), + /// An unknown model shaped like a face embedder. `assumptions` is what we guessed, in plain words. + Draft { manifest: ModelManifest, assumptions: Vec }, + /// Not usable yet, and why. + Unsupported(String), +} + +fn fixed(d: &Dim) -> Option { + match d { + Dim::Fixed(v) => Some(*v), + Dim::Dynamic(_) => None, + } +} + +/// An id made of the file name and the start of the hash: lowercase letters, digits, `.` `_` `-`. +fn make_id(file_name: &str, sha256: &str) -> String { + let stem = file_name.rsplit_once('.').map_or(file_name, |(s, _)| s); + let clean: String = stem.chars().map(|c| if c.is_ascii_alphanumeric() { c.to_ascii_lowercase() } else { '-' }).collect(); + let clean: String = clean.trim_matches('-').chars().take(40).collect(); + let short: String = sha256.chars().take(8).collect(); + match (clean.is_empty(), short.is_empty()) { + (true, true) => "custom".into(), + (true, false) => format!("custom-{short}"), + (false, true) => format!("custom-{clean}"), + (false, false) => format!("custom-{clean}-{short}"), + } +} + +fn display_name(file_name: &str) -> String { + let stem = file_name.rsplit_once('.').map_or(file_name, |(s, _)| s); + let name: String = stem.chars().filter(|c| !c.is_control()).take(60).collect(); + if name.trim().is_empty() { "Custom face model".into() } else { format!("{} (your model)", name.trim()) } +} + +/// The input of a face embedder: float32, `[batch, 3, height, width]`. +fn face_input(t: &TensorInfo) -> Option<(u32, u32)> { + if t.elem_type != FLOAT32 || t.shape.len() != 4 || t.shape.get(1).and_then(fixed) != Some(3) { + return None; + } + let size = |d: Option<&Dim>| match d.and_then(fixed) { + Some(v) if (16..=1024).contains(&v) => u32::try_from(v).ok(), + Some(_) => None, + None => Some(112), // a dynamic size: the common face-crop size + }; + Some((size(t.shape.get(3))?, size(t.shape.get(2))?)) +} + +/// Propose what to do with a probed model. `sha256` and `size` describe the file, `file_name` names it. +pub fn suggest(info: &OnnxInfo, sha256: &str, size: u64, file_name: &str) -> Suggestion { + suggest_with(info, sha256, size, file_name, &[]) +} + +/// [`suggest`], also recognising the models in `extra` (the user's own catalog) by their SHA-256. +pub fn suggest_with(info: &OnnxInfo, sha256: &str, size: u64, file_name: &str, extra: &[ModelManifest]) -> Suggestion { + if let Some(m) = crate::known::lookup(sha256).or_else(|| extra.iter().find(|m| m.sha256.as_deref() == Some(sha256)).cloned()) { + return Suggestion::Known(m); + } + let [input] = info.inputs.as_slice() else { + return Suggestion::Unsupported(format!("the model has {} inputs; face models here take one image input", info.inputs.len())); + }; + let Some((width, height)) = face_input(input) else { + return Suggestion::Unsupported("its input is not a float image of shape [batch, 3, height, width]".into()); + }; + let [output] = info.outputs.as_slice() else { + return Suggestion::Unsupported( + "it has several outputs, like a face detector; only models that give one vector per face are supported for now".into(), + ); + }; + if output.elem_type != FLOAT32 || output.shape.len() != 2 { + return Suggestion::Unsupported("its output is not one vector per face (shape [batch, dimension])".into()); + } + let dim = match output.shape.get(1).and_then(fixed).map(u32::try_from) { + Some(Ok(d)) if (8..=4096).contains(&d) => d, + _ => return Suggestion::Unsupported("its output size is not a fixed number between 8 and 4096".into()), + }; + let manifest = ModelManifest { + id: make_id(file_name, sha256), + name: display_name(file_name), + version: "1".into(), + role: Role::Embedder, + licence: Licence { + name: "Unknown".into(), + commercial: Commercial::Unknown, + url: None, + notice: "LightCraft does not know where this model came from or what its licence allows. Use it only if its licence allows what you are doing, and do not redistribute it.".into(), + }, + source: None, + sha256: Some(sha256.to_string()).filter(|h| h.len() == 64), + size_bytes: Some(size).filter(|s| *s > 0), + provenance: "Unknown.".into(), + input: InputSpec { width, height, ..InputSpec::default() }, + output: OutputSpec::Embedding { dim }, + thresholds: Thresholds::default(), + speed: None, + }; + let assumptions = vec![ + format!("Faces are aligned crops of {width} × {height} pixels."), + "Pixels are RGB, normalised as (x − 127.5) / 127.5: the usual ArcFace convention, which InsightFace models use.".into(), + format!("Each face becomes a vector of {dim} numbers; two faces match when their cosine similarity is high."), + ]; + Suggestion::Draft { manifest, assumptions } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::manifest::validate; + use crate::onnx::probe_reader; + use std::io::Cursor; + + fn tensor(elem: u32, shape: &[Dim]) -> TensorInfo { + TensorInfo { name: "t".into(), elem_type: elem, shape: shape.to_vec() } + } + fn dynamic() -> Dim { + Dim::Dynamic("N".into()) + } + fn info(input: TensorInfo, output: TensorInfo) -> OnnxInfo { + OnnxInfo { producer: "x".into(), opset: Some(17), inputs: vec![input], outputs: vec![output] } + } + const SHA: &str = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"; + + #[test] + fn a_known_hash_wins() { + let i = info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(10), Dim::Fixed(10)]), tensor(1, &[dynamic(), Dim::Fixed(5)])); + assert!(matches!(suggest(&i, crate::known::AURAFACE_SHA256, 1, "whatever.onnx"), Suggestion::Known(m) if m.id == "auraface-v1")); + } + + #[test] + fn an_unknown_embedder_gets_a_draft_that_validates() { + let i = info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(512)])); + let Suggestion::Draft { manifest, assumptions } = suggest(&i, SHA, 5000, "My Model (r50).onnx") else { unreachable_in_test() }; + assert_eq!(validate(&manifest), Ok(())); + assert_eq!(manifest.id, "custom-my-model--r50-01234567"); + assert_eq!(manifest.output, OutputSpec::Embedding { dim: 512 }); + assert_eq!((manifest.input.width, manifest.input.height), (112, 112)); + assert_eq!(manifest.licence.commercial, Commercial::Unknown); + assert!(assumptions.iter().any(|a| a.contains("127.5"))); + } + + #[allow(clippy::panic)] + fn unreachable_in_test() -> ! { + panic!("expected a draft"); + } + + #[test] + fn dynamic_sizes_default_and_odd_shapes_are_refused() { + let ok = info(tensor(1, &[dynamic(), Dim::Fixed(3), dynamic(), dynamic()]), tensor(1, &[dynamic(), Dim::Fixed(128)])); + assert!(matches!(suggest(&ok, SHA, 1, "a.onnx"), Suggestion::Draft { manifest, .. } if manifest.input.width == 112)); + let refused = [ + info(tensor(1, &[dynamic(), Dim::Fixed(1), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(128)])), // grey + info(tensor(7, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(128)])), // int64 input + info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(128)])), // rank 3 + info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(1_000_000), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(128)])), + info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(257), Dim::Fixed(384)])), // tokens + info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), dynamic()])), + info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(2)])), + info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(1_000_000)])), + ]; + for r in refused { + assert!(matches!(suggest(&r, SHA, 1, "a.onnx"), Suggestion::Unsupported(_)), "{r:?}"); + } + let mut two = ok.clone(); + two.outputs.push(two.outputs[0].clone()); + assert!(matches!(suggest(&two, SHA, 1, "a.onnx"), Suggestion::Unsupported(_)), "several outputs"); + } + + #[test] + fn file_names_cannot_break_the_manifest() { + let i = info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(512)])); + for name in ["", ".onnx", "../../etc/passwd", "\u{7}\u{0}x.onnx", &"x".repeat(5000), "ÅÄÖ 🙂.onnx", "UPPER lower.ONNX"] { + let Suggestion::Draft { manifest, .. } = suggest(&i, SHA, 1, name) else { unreachable_in_test() }; + assert_eq!(validate(&manifest), Ok(()), "{name:?}"); + } + // and with no usable hash at all + let Suggestion::Draft { manifest, .. } = suggest(&i, "", 0, "x.onnx") else { unreachable_in_test() }; + assert_eq!(validate(&manifest), Ok(())); + assert!(manifest.sha256.is_none() && manifest.size_bytes.is_none()); + } + + /// Opt-in: `LC_FACE_MODELS=/folder/with/onnx/files cargo test -p lightcraft-faces -- --ignored --nocapture` + /// probes real models (never committed) and prints what would be suggested, with timings. + #[test] + #[ignore = "needs real model files: set LC_FACE_MODELS"] + fn probing_real_models() { + let Some(dir) = std::env::var_os("LC_FACE_MODELS") else { return }; + for entry in std::fs::read_dir(dir).unwrap().flatten() { + let path = entry.path(); + if path.extension().is_none_or(|e| e != "onnx") { + continue; + } + let t = std::time::Instant::now(); + let sha = crate::hash::sha256_file(&path).unwrap(); + let hashed = t.elapsed(); + let t = std::time::Instant::now(); + let info = crate::onnx::probe_path(&path); + let probed = t.elapsed(); + let size = std::fs::metadata(&path).unwrap().len(); + let name = path.file_name().unwrap().to_string_lossy().into_owned(); + let verdict = match &info { + Ok(i) => match suggest(i, &sha, size, &name) { + Suggestion::Known(m) => format!("KNOWN {} ({:?})", m.id, m.licence.commercial), + Suggestion::Draft { manifest, .. } => format!("DRAFT {} {:?}", manifest.id, manifest.output), + Suggestion::Unsupported(why) => format!("UNSUPPORTED: {why}"), + }, + Err(e) => format!("PROBE ERROR: {e}"), + }; + println!("{name}: {} MB | hash {hashed:?} | probe {probed:?} | {verdict}", size / 1_048_576); + if let Ok(i) = info { + println!( + " in {:?} + out {:?}", + i.inputs, i.outputs + ); + } + } + } + + #[test] + fn probing_then_suggesting_end_to_end() { + let probed = probe_reader(&mut Cursor::new(crate::onnx::tests::model_bytes(512))).unwrap(); + assert!( + matches!(suggest(&probed, SHA, 9, "m.onnx"), Suggestion::Draft { manifest, .. } if manifest.output == OutputSpec::Embedding { dim: 512 }) + ); + } +} diff --git a/crates/faces/src/synthetic.rs b/crates/faces/src/synthetic.rs new file mode 100644 index 00000000..61010a24 --- /dev/null +++ b/crates/faces/src/synthetic.rs @@ -0,0 +1,250 @@ +//! A minimal, valid ONNX model made in memory: for tests, and later for the install self-check. It is not a +//! working network (it has no layers), only a graph with the right inputs, outputs and a weight blob. + +pub(crate) fn varint(mut v: u64, out: &mut Vec) { + loop { + let b = (v & 0x7f) as u8; + v >>= 7; + if v == 0 { + out.push(b); + return; + } + out.push(b | 0x80); + } +} + +pub fn len_field(n: u64, body: &[u8], out: &mut Vec) { + varint(n << 3 | 2, out); + varint(body.len() as u64, out); + out.extend_from_slice(body); +} + +pub(crate) fn varint_field(n: u64, v: u64, out: &mut Vec) { + varint(n << 3, out); + varint(v, out); +} + +/// A graph input or output: `Ok(n)` is a fixed dimension, `Err(name)` a named (dynamic) one. +pub fn value_info_bytes(name: &str, elem: u64, dims: &[std::result::Result]) -> Vec { + let mut shape = Vec::new(); + for d in dims { + let mut dim = Vec::new(); + match d { + Ok(v) => varint_field(1, *v, &mut dim), + Err(p) => len_field(2, p.as_bytes(), &mut dim), + } + len_field(1, &dim, &mut shape); + } + let mut tensor = Vec::new(); + varint_field(1, elem, &mut tensor); + len_field(2, &shape, &mut tensor); + let mut ty = Vec::new(); + len_field(1, &tensor, &mut ty); + let mut vi = Vec::new(); + len_field(1, name.as_bytes(), &mut vi); + len_field(2, &ty, &mut vi); + vi +} + +/// A small valid model shaped like a face embedder: float input `data` `[N, 3, 112, 112]`, an old-style +/// weight also listed as an input, a weight with bulk data, float output `emb` `[N, output_dim]`, opset 17. +pub fn embedder_model(output_dim: u64) -> Vec { + let mut graph = Vec::new(); + len_field(11, &value_info_bytes("data", 1, &[Err("N"), Ok(3), Ok(112), Ok(112)]), &mut graph); + len_field(11, &value_info_bytes("w", 1, &[Ok(4)]), &mut graph); + let mut tensor = Vec::new(); + varint_field(1, 4, &mut tensor); + len_field(8, b"w", &mut tensor); + len_field(9, &[7u8; 5000], &mut tensor); + len_field(5, &tensor, &mut graph); + len_field(12, &value_info_bytes("emb", 1, &[Err("N"), Ok(output_dim)]), &mut graph); + let mut ops = Vec::new(); + len_field(1, b"", &mut ops); + varint_field(2, 17, &mut ops); + let mut model = Vec::new(); + varint_field(1, 8, &mut model); + len_field(2, b"pytorch", &mut model); + len_field(7, &graph, &mut model); + len_field(8, &ops, &mut model); + model +} + +/// A node attribute for [`node_bytes`]. +enum Attribute<'a> { + Int(i64), + Ints(&'a [i64]), + Float(f32), +} + +fn node_bytes(op: &str, inputs: &[&str], output: &str, attrs: &[(&str, Attribute)]) -> Vec { + let mut n = Vec::new(); + for i in inputs { + len_field(1, i.as_bytes(), &mut n); + } + len_field(2, output.as_bytes(), &mut n); + len_field(4, op.as_bytes(), &mut n); + for (name, value) in attrs { + let mut a = Vec::new(); + len_field(1, name.as_bytes(), &mut a); + match value { + Attribute::Float(v) => { + varint(2 << 3 | 5, &mut a); + a.extend_from_slice(&v.to_le_bytes()); + varint_field(20, 1, &mut a); + } + Attribute::Int(v) => { + varint_field(3, *v as u64, &mut a); + varint_field(20, 2, &mut a); // INT + } + Attribute::Ints(values) => { + for v in *values { + varint_field(8, *v as u64, &mut a); + } + varint_field(20, 7, &mut a); // INTS + } + } + len_field(5, &a, &mut n); + } + n +} + +fn float_tensor_bytes(name: &str, dims: &[u64], values: &[f32]) -> Vec { + let mut t = Vec::new(); + for d in dims { + varint_field(1, *d, &mut t); + } + varint_field(2, 1, &mut t); // FLOAT + len_field(8, name.as_bytes(), &mut t); + let raw: Vec = values.iter().flat_map(|v| v.to_le_bytes()).collect(); + len_field(9, &raw, &mut t); + t +} + +/// A tiny working network shaped like a face embedder, for tests of the runtime: input `data` +/// `[1, 3, 112, 112]`, a 1 × 1 convolution to `dim` channels, global average pooling and flattening to +/// `[1, dim]`. Its output depends on the picture's per-channel brightness, which is all the tests need. +pub fn tiny_embedder_model(dim: u64) -> Vec { + let mut graph = Vec::new(); + len_field(1, &node_bytes("Conv", &["data", "w", "b"], "c", &[("kernel_shape", Attribute::Ints(&[1, 1]))]), &mut graph); + len_field(1, &node_bytes("GlobalAveragePool", &["c"], "p", &[]), &mut graph); + len_field(1, &node_bytes("Flatten", &["p"], "emb", &[("axis", Attribute::Int(1))]), &mut graph); + let w: Vec = (0..dim * 3).map(|i| ((i * 37 % 19) as f32 - 9.0) / 20.0).collect(); + let b: Vec = (0..dim).map(|i| (i % 5) as f32 / 10.0 - 0.2).collect(); + len_field(5, &float_tensor_bytes("w", &[dim, 3, 1, 1], &w), &mut graph); + len_field(5, &float_tensor_bytes("b", &[dim], &b), &mut graph); + len_field(11, &value_info_bytes("data", 1, &[Ok(1), Ok(3), Ok(112), Ok(112)]), &mut graph); + len_field(12, &value_info_bytes("emb", 1, &[Ok(1), Ok(dim)]), &mut graph); + let mut ops = Vec::new(); + len_field(1, b"", &mut ops); + varint_field(2, 13, &mut ops); + let mut model = Vec::new(); + varint_field(1, 8, &mut model); + len_field(7, &graph, &mut model); + len_field(8, &ops, &mut model); + model +} + +/// A deterministic small recognition graph using all SFace/AuraFace operators (no downloaded weights). +/// Intended for runtime regression and independent reference comparisons. +pub fn recognition_model(mut seed: u32) -> Vec { + let mut random = || { + seed = seed.wrapping_mul(1664525).wrapping_add(1013904223); + ((seed >> 16) as f32 / 65535.0 - 0.5) * 0.2 + }; + let mut g = Vec::new(); + for (op, inputs, out, attrs) in [ + ("Sub", vec!["data", "center"], "centered", vec![]), + ("Mul", vec!["centered", "scale"], "scaled", vec![]), + ("Conv", vec!["scaled", "conv.w", "conv.b"], "conv", vec![("pads", Attribute::Ints(&[1, 1, 1, 1])), ("strides", Attribute::Ints(&[2, 2]))]), + ("BatchNormalization", vec!["conv", "bn.s", "bn.b", "bn.m", "bn.v"], "bn", vec![("epsilon", Attribute::Float(1e-5))]), + ("PRelu", vec!["bn", "slope"], "act", vec![]), + ("Add", vec!["act", "bn"], "skip", vec![]), + ("Dropout", vec!["skip"], "drop", vec![("ratio", Attribute::Float(0.4))]), + ("GlobalAveragePool", vec!["drop"], "pool", vec![]), + ("Flatten", vec!["pool"], "flat", vec![]), + ("Gemm", vec!["flat", "fc.w", "fc.b"], "emb", vec![("transB", Attribute::Int(1))]), + ] { + len_field(1, &node_bytes(op, &inputs, out, &attrs), &mut g); + } + for (name, dims, data) in [ + ("center", vec![1], vec![127.5]), + ("scale", vec![1], vec![1.0 / 127.5]), + ("conv.w", vec![5, 3, 3, 3], (0..135).map(|_| random()).collect()), + ("conv.b", vec![5], (0..5).map(|_| random()).collect()), + ("bn.s", vec![5], (0..5).map(|_| 1.0 + random()).collect()), + ("bn.b", vec![5], (0..5).map(|_| random()).collect()), + ("bn.m", vec![5], (0..5).map(|_| random()).collect()), + ("bn.v", vec![5], (0..5).map(|_| 1.0 + random()).collect()), + ("slope", vec![5, 1, 1], (0..5).map(|_| 0.2 + random()).collect()), + ("fc.w", vec![8, 5], (0..40).map(|_| random()).collect()), + ("fc.b", vec![8], (0..8).map(|_| random()).collect()), + ] { + len_field(5, &float_tensor_bytes(name, &dims, &data), &mut g); + } + len_field(11, &value_info_bytes("data", 1, &[Ok(1), Ok(3), Ok(112), Ok(112)]), &mut g); + len_field(12, &value_info_bytes("emb", 1, &[Ok(1), Ok(8)]), &mut g); + let mut ops = Vec::new(); + varint_field(2, 11, &mut ops); + let mut model = Vec::new(); + varint_field(1, 8, &mut model); + len_field(7, &g, &mut model); + len_field(8, &ops, &mut model); + model +} + +#[cfg(test)] +pub(crate) fn norm_overflow_model() -> Vec { + let mut graph = Vec::new(); + len_field(1, &node_bytes("Flatten", &["data"], "flat", &[]), &mut graph); + len_field(1, &node_bytes("Gemm", &["flat", "w"], "emb", &[("transB", Attribute::Int(1))]), &mut graph); + let plane = 112 * 112; + let mut weights = vec![0.0; 3 * 3 * plane]; + for c in 0..3 { + weights[c * 3 * plane + c * plane] = 5e37; + } + len_field(5, &float_tensor_bytes("w", &[3, (3 * plane) as u64], &weights), &mut graph); + len_field(11, &value_info_bytes("data", 1, &[Ok(1), Ok(3), Ok(112), Ok(112)]), &mut graph); + len_field(12, &value_info_bytes("emb", 1, &[Ok(1), Ok(3)]), &mut graph); + let mut ops = Vec::new(); + varint_field(2, 11, &mut ops); + let mut model = Vec::new(); + len_field(7, &graph, &mut model); + len_field(8, &ops, &mut model); + model +} + +/// A small YuNet-shaped stand-in that reports no faces, independent of the real downloaded model. +pub fn tiny_detector_model() -> Vec { + let mut graph = Vec::new(); + len_field(11, &value_info_bytes("input", 1, &[Ok(1), Ok(3), Ok(64), Ok(64)]), &mut graph); + for stride in [8u64, 16, 32] { + let cells = (64 / stride) * (64 / stride); + for (kind, components) in [("cls", 1u64), ("obj", 1), ("bbox", 4), ("kps", 10)] { + let name = format!("{kind}_{stride}"); + let weight = format!("weight_{name}"); + let values = vec![0.0; (cells * components) as usize]; + len_field(5, &float_tensor_bytes(&weight, &[1, cells, components], &values), &mut graph); + len_field(1, &node_bytes("Identity", &[&weight], &name, &[]), &mut graph); + len_field(12, &value_info_bytes(&name, 1, &[Ok(1), Ok(cells), Ok(components)]), &mut graph); + } + } + let mut model = Vec::new(); + varint_field(1, 8, &mut model); + len_field(7, &graph, &mut model); + let mut ops = Vec::new(); + varint_field(2, 13, &mut ops); + len_field(8, &ops, &mut model); + model +} + +/// Manifest for the stand-in, for model-install and background-scan tests only. +pub fn tiny_detector_manifest() -> crate::manifest::ModelManifest { + let bytes = tiny_detector_model(); + let mut m = crate::known::yunet(); + m.name = "Synthetic detector (tests only)".into(); + m.input.width = 64; + m.input.height = 64; + m.sha256 = Some(crate::hash::sha256_hex(&bytes)); + m.size_bytes = Some(bytes.len() as u64); + m +} diff --git a/crates/faces/src/yunet.rs b/crates/faces/src/yunet.rs new file mode 100644 index 00000000..b77f318b --- /dev/null +++ b/crates/faces/src/yunet.rs @@ -0,0 +1,370 @@ +//! The YuNet face detector (232 KB, MIT, downloaded by the user in Settings ▸ Faces), run by LightCraft's own small +//! interpreter ([`crate::net`]): no extra runtime. +//! +//! A photo is letterboxed (aspect ratio kept, padded at the right and bottom) into the network's square +//! input, the network gives a score, a box and five landmarks for every cell of three grids (strides 8, 16 and +//! 32), and the cells above the score threshold are decoded, de-duplicated by non-maximum suppression and +//! mapped back to the photo as fractions of its width and height. +//! +//! The pixels come from the caller as plain 8-bit RGB; this module has no image types of its own. + +use crate::graph; +use crate::net::{Net, NetError, Tensor}; + +const STRIDES: [usize; 3] = [8, 16, 32]; +/// Largest photo side we accept. +const MAX_SIDE: usize = 65_536; +/// Candidate cells kept before suppression, so a hostile or noisy output cannot make it slow. +const MAX_CANDIDATES: usize = 4000; + +#[derive(Debug, thiserror::Error)] +pub enum DetectError { + #[error("the detector model is not usable: {0}")] + Model(String), + #[error("{0}")] + Input(String), +} + +impl From for DetectError { + fn from(e: NetError) -> Self { + DetectError::Model(e.to_string()) + } +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Options { + /// Lowest detection score kept (0 to 1). + pub score: f32, + /// Boxes overlapping a better one by more than this (intersection over union) are dropped. + pub nms_iou: f32, + /// Most faces returned. + pub max_faces: usize, +} + +impl Default for Options { + fn default() -> Self { + Options { score: 0.6, nms_iou: 0.3, max_faces: 200 } + } +} + +/// A detected face, in fractions of the photo's width and height (0 to 1, y down). +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Face { + pub x0: f32, + pub y0: f32, + pub x1: f32, + pub y1: f32, + pub score: f32, + /// Right eye, left eye, nose tip, right and left mouth corner, as the detector names them (from the + /// subject's point of view, so the "right eye" is on the left of the picture). May lie slightly outside 0 to 1. + pub landmarks: [(f32, f32); 5], +} + +pub struct Detector { + net: Net, + /// Side of the square network input. + side: usize, +} + +impl Detector { + /// A detector from YuNet-architecture ONNX bytes (a retrained model of the same shape works too). + pub fn new(onnx: &[u8]) -> Result { + let g = graph::load(onnx).map_err(|e| DetectError::Model(e.to_string()))?; + let outputs = g.outputs.clone(); + let net = Net::new(g)?; + let side = match net.input_shape() { + Some(&[1, 3, h, w]) if h == w && (64..=4096).contains(&h) && h % 32 == 0 => h, + _ => return Err(DetectError::Model("the model's input is not a fixed square image [1, 3, N, N] with N a multiple of 32".into())), + }; + let shapes = net.validate_input(&[1, 3, side, side])?; + for stride in STRIDES { + for (kind, components) in [("cls", 1), ("obj", 1), ("bbox", 4), ("kps", 10)] { + let name = format!("{kind}_{stride}"); + let index = outputs.iter().position(|n| n == &name).ok_or_else(|| DetectError::Model(format!("the model has no output `{name}`")))?; + let count = shapes.get(index).and_then(|s| s.iter().try_fold(1usize, |n, d| n.checked_mul(*d))); + if count != Some((side / stride) * (side / stride) * components) { + return Err(DetectError::Model(format!("the output `{name}` has the wrong shape"))); + } + } + } + Ok(Detector { net, side }) + } + + /// The square input side, used to check an installed manifest. + pub fn input_side(&self) -> usize { + self.side + } + + /// Find faces in an upright `width` × `height` photo of 8-bit RGB pixels (`rgb` is `width * height * 3` bytes). + pub fn detect(&self, rgb: &[u8], width: usize, height: usize, opts: &Options) -> Result, DetectError> { + if width == 0 || height == 0 || width > MAX_SIDE || height > MAX_SIDE { + return Err(DetectError::Input(format!("a {width} × {height} picture cannot be searched"))); + } + if width.checked_mul(height).and_then(|n| n.checked_mul(3)) != Some(rgb.len()) { + return Err(DetectError::Input("the pixel data does not match the picture's size".into())); + } + let opts = Options { + score: if opts.score.is_finite() { opts.score.clamp(0.0, 1.0) } else { 0.6 }, + nms_iou: if opts.nms_iou.is_finite() { opts.nms_iou.clamp(0.0, 1.0) } else { 0.3 }, + max_faces: opts.max_faces.min(10_000), + }; + let scale = (self.side as f32 / width as f32).min(self.side as f32 / height as f32); + let input = self.letterbox(rgb, width, height, scale)?; + let outputs = self.net.run(input)?; + let find = |name: String| { + outputs.iter().find(|(n, _)| *n == name).map(|(_, t)| t).ok_or_else(|| DetectError::Model(format!("the model has no output `{name}`"))) + }; + let mut candidates: Vec = Vec::new(); + for stride in STRIDES { + let (cls, obj) = (find(format!("cls_{stride}"))?, find(format!("obj_{stride}"))?); + let (bbox, kps) = (find(format!("bbox_{stride}"))?, find(format!("kps_{stride}"))?); + let cols = self.side / stride; + let cells = cols * cols; + if cls.data.len() != cells || obj.data.len() != cells || bbox.data.len() != cells * 4 || kps.data.len() != cells * 10 { + return Err(DetectError::Model(format!("the outputs for stride {stride} have the wrong size"))); + } + for (i, (c, o)) in cls.data.iter().zip(&obj.data).enumerate() { + let score = (c.clamp(0.0, 1.0) * o.clamp(0.0, 1.0)).sqrt(); + if score.is_nan() || score < opts.score { + continue; + } + let (row, col) = ((i / cols) as f32, (i % cols) as f32); + let (Some(b), Some(k)) = (bbox.data.get(i * 4..i * 4 + 4), kps.data.get(i * 10..i * 10 + 10)) else { continue }; + let s = stride as f32; + let (cx, cy) = ((col + b[0]) * s, (row + b[1]) * s); + let (w, h) = (b[2].clamp(-10.0, 10.0).exp() * s, b[3].clamp(-10.0, 10.0).exp() * s); + let mut landmarks = [(0.0, 0.0); 5]; + for (n, l) in landmarks.iter_mut().enumerate() { + let (lx, ly) = (k.get(n * 2).copied().unwrap_or(0.0), k.get(n * 2 + 1).copied().unwrap_or(0.0)); + *l = (((lx + col) * s) / scale / width as f32, ((ly + row) * s) / scale / height as f32); + } + candidates.push(Face { + x0: (cx - w / 2.0) / scale / width as f32, + y0: (cy - h / 2.0) / scale / height as f32, + x1: (cx + w / 2.0) / scale / width as f32, + y1: (cy + h / 2.0) / scale / height as f32, + score, + landmarks, + }); + } + } + candidates.sort_by(|a, b| b.score.total_cmp(&a.score)); + candidates.truncate(MAX_CANDIDATES); + let mut kept: Vec = Vec::new(); + for c in candidates { + if kept.len() >= opts.max_faces { + break; + } + // the box as it lies inside the photo + let clipped = Face { x0: c.x0.clamp(0.0, 1.0), y0: c.y0.clamp(0.0, 1.0), x1: c.x1.clamp(0.0, 1.0), y1: c.y1.clamp(0.0, 1.0), ..c }; + if !(clipped.x1 - clipped.x0 > 1e-4 && clipped.y1 - clipped.y0 > 1e-4) || kept.iter().any(|k| iou(k, &clipped) > opts.nms_iou) { + continue; + } + kept.push(clipped); + } + Ok(kept) + } + + /// The photo scaled by `scale` into the top left of a zero-filled square, as BGR 0..255 (the model's own input). + fn letterbox(&self, rgb: &[u8], width: usize, height: usize, scale: f32) -> Result { + let side = self.side; + let nw = ((width as f32 * scale).round() as usize).clamp(1, side); + let nh = ((height as f32 * scale).round() as usize).clamp(1, side); + let (xt, yt) = (taps(width, nw), taps(height, nh)); + // horizontal pass: every source row, nw columns, three channels + let mut mid = vec![0.0f32; height * nw * 3]; + for (src_row, mid_row) in rgb.chunks_exact(width * 3).zip(mid.chunks_exact_mut(nw * 3)) { + for (x, tap) in xt.iter().enumerate() { + let mut acc = [0.0f32; 3]; + for (j, wt) in tap.weights.iter().enumerate() { + if let Some(px) = src_row.get((tap.start + j) * 3..(tap.start + j) * 3 + 3) { + acc[0] += wt * f32::from(px[0]); + acc[1] += wt * f32::from(px[1]); + acc[2] += wt * f32::from(px[2]); + } + } + if let Some(o) = mid_row.get_mut(x * 3..x * 3 + 3) { + o.copy_from_slice(&acc); + } + } + } + // vertical pass into the planes B, G, R + let plane = side * side; + let mut data = vec![0.0f32; 3 * plane]; + for (y, tap) in yt.iter().enumerate() { + for x in 0..nw { + let mut acc = [0.0f32; 3]; + for (j, wt) in tap.weights.iter().enumerate() { + if let Some(px) = mid.get(((tap.start + j) * nw + x) * 3..((tap.start + j) * nw + x) * 3 + 3) { + acc[0] += wt * px[0]; + acc[1] += wt * px[1]; + acc[2] += wt * px[2]; + } + } + let at = y * side + x; + for (c, v) in [acc[2], acc[1], acc[0]].into_iter().enumerate() { + if let Some(o) = data.get_mut(c * plane + at) { + *o = v; + } + } + } + } + Ok(Tensor::new(vec![1, 3, side, side], data)?) + } +} + +struct Tap { + start: usize, + weights: Vec, +} + +/// For each of `dst` output positions, which `src` positions to average and with what weights: exact area +/// averaging when shrinking, linear interpolation when enlarging. +fn taps(src: usize, dst: usize) -> Vec { + let r = src as f32 / dst as f32; + (0..dst) + .map(|i| { + if r >= 1.0 { + let (a, b) = (i as f32 * r, (i as f32 + 1.0) * r); + let first = a.floor() as usize; + let last = ((b.ceil() as usize).min(src)).max(first + 1); + let weights: Vec = (first..last).map(|j| ((j as f32 + 1.0).min(b) - (j as f32).max(a)).max(0.0) / r).collect(); + Tap { start: first.min(src - 1), weights } + } else { + let c = ((i as f32 + 0.5) * r - 0.5).max(0.0); + let lo = (c.floor() as usize).min(src - 1); + let hi = (lo + 1).min(src - 1); + let f = c - lo as f32; + if hi == lo { Tap { start: lo, weights: vec![1.0] } } else { Tap { start: lo, weights: vec![1.0 - f, f] } } + } + }) + .collect() +} + +fn iou(a: &Face, b: &Face) -> f32 { + let (w, h) = ((a.x1.min(b.x1) - a.x0.max(b.x0)).max(0.0), (a.y1.min(b.y1) - a.y0.max(b.y0)).max(0.0)); + let inter = w * h; + let union = (a.x1 - a.x0) * (a.y1 - a.y0) + (b.x1 - b.x0) * (b.y1 - b.y0) - inter; + if union > 0.0 { inter / union } else { 0.0 } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The real model, for the tests that need it: `LC_YUNET_MODEL=` (the file + /// Settings ▸ Faces downloads; nothing is committed). Without it those tests say so and pass. + fn model() -> Option> { + let path = std::env::var_os("LC_YUNET_MODEL")?; + let bytes = std::fs::read(path).ok()?; + assert_eq!(crate::hash::sha256_hex(&bytes), crate::known::YUNET_SHA256, "LC_YUNET_MODEL is not YuNet 2023mar"); + Some(bytes) + } + + fn detector() -> Option { + let Some(bytes) = model() else { + eprintln!("LC_YUNET_MODEL is not set: skipping"); + return None; + }; + Some(Detector::new(&bytes).unwrap()) + } + + #[test] + fn the_model_loads_and_has_a_square_input() { + let Some(d) = detector() else { return }; + assert_eq!(d.side, 640); + } + + #[test] + fn a_blank_picture_has_no_faces_and_odd_sizes_are_handled() { + let Some(d) = detector() else { return }; + for (w, h) in [(640, 480), (100, 700), (1, 1), (37, 41), (1000, 20)] { + let rgb = vec![128u8; w * h * 3]; + let faces = d.detect(&rgb, w, h, &Options::default()).unwrap(); + assert!(faces.is_empty(), "{w}×{h}: {faces:?}"); + } + } + + #[test] + fn bad_input_is_an_error_not_a_panic() { + let d = Detector::new(&crate::synthetic::tiny_detector_model()).unwrap(); + let o = Options::default(); + assert!(d.detect(&[], 0, 0, &o).is_err()); + assert!(d.detect(&[0; 12], 2, 3, &o).is_err(), "wrong byte count"); + assert!(d.detect(&[0; 3], usize::MAX, 2, &o).is_err(), "overflowing size"); + assert!(d.detect(&[0; 3], MAX_SIDE + 1, 1, &o).is_err()); + let nan = Options { score: f32::NAN, nms_iou: f32::INFINITY, max_faces: usize::MAX }; + assert!(d.detect(&vec![0u8; 12 * 12 * 3], 12, 12, &nan).is_ok(), "non-finite options fall back to defaults"); + } + + #[test] + fn missing_decoder_outputs_are_refused_at_load_time() { + let mut bytes = crate::synthetic::tiny_detector_model(); + let mut i = 0; + while i + 5 <= bytes.len() { + if &bytes[i..i + 5] == b"cls_8" { + bytes[i..i + 5].copy_from_slice(b"bad_8"); + } + i += 1; + } + let error = Detector::new(&bytes).err().unwrap().to_string(); + assert!(error.contains("cls_8"), "{error}"); + } + + #[test] + fn a_model_that_is_not_yunet_shaped_is_refused() { + let d = Detector::new(&crate::synthetic::tiny_detector_model()).unwrap(); + assert!(d.detect(&[128; 3], 1, 1, &Options::default()).unwrap().is_empty()); + assert!(Detector::new(b"nonsense").is_err()); + assert!(Detector::new(&crate::synthetic::embedder_model(512)).is_err()); + if let Some(bytes) = model() { + for n in [0, 1, 100, bytes.len() / 2, bytes.len() - 1] { + assert!(Detector::new(&bytes[..n]).is_err(), "truncated at {n}"); + } + } + } + + /// Opt-in: `LC_YUNET_REF= cargo test -p lightcraft-faces -- --ignored --nocapture` compares every output + /// of the network with ONNX Runtime's on the same input (`input.bin`, `out_.bin`, raw little-endian f32). + #[test] + #[ignore = "needs reference files: set LC_YUNET_REF"] + fn the_network_matches_onnx_runtime() { + let Some(dir) = std::env::var_os("LC_YUNET_REF") else { return }; + let dir = std::path::PathBuf::from(dir); + let floats = + |name: &str| -> Vec { std::fs::read(dir.join(name)).unwrap().as_chunks::<4>().0.iter().map(|b| f32::from_le_bytes(*b)).collect() }; + let Some(d) = detector() else { return }; + let started = std::time::Instant::now(); + let outs = d.net.run(Tensor::new(vec![1, 3, 640, 640], floats("input.bin")).unwrap()).unwrap(); + println!("network run: {:?}", started.elapsed()); + assert_eq!(outs.len(), 12); + for (name, t) in outs { + let want = floats(&format!("out_{name}.bin")); + assert_eq!(t.data.len(), want.len(), "{name}"); + let worst = t.data.iter().zip(&want).map(|(a, b)| (a - b).abs()).fold(0.0f32, f32::max); + println!("{name}: {} values, worst difference {worst:e}", want.len()); + assert!(worst < 2e-3, "{name} differs by {worst}"); + } + } + + #[test] + fn taps_cover_the_source_and_sum_to_one() { + for (src, dst) in [(6000, 640), (640, 640), (37, 640), (3, 5), (1, 1), (1, 640), (641, 640)] { + let t = taps(src, dst); + assert_eq!(t.len(), dst); + for tap in &t { + assert!(tap.start + tap.weights.len() <= src, "{src}→{dst}"); + let sum: f32 = tap.weights.iter().sum(); + assert!((sum - 1.0).abs() < 1e-3, "{src}→{dst}: {sum}"); + } + } + } + + #[test] + fn overlap_is_symmetric() { + let f = |x0, y0, x1, y1| Face { x0, y0, x1, y1, score: 1.0, landmarks: [(0.0, 0.0); 5] }; + let (a, b) = (f(0.0, 0.0, 0.5, 0.5), f(0.25, 0.25, 0.75, 0.75)); + assert!((iou(&a, &b) - iou(&b, &a)).abs() < 1e-6 && (iou(&a, &b) - 0.0625 / 0.4375).abs() < 1e-5); + assert_eq!(iou(&a, &f(0.6, 0.6, 0.9, 0.9)), 0.0); + } +} diff --git a/crates/gpu/src/render.rs b/crates/gpu/src/render.rs index 60ffd806..7e7ac8a3 100644 --- a/crates/gpu/src/render.rs +++ b/crates/gpu/src/render.rs @@ -717,7 +717,7 @@ pub fn render( Err(m) => { let img = host.lin.get_or_insert_with(|| Arc::new(cx.read_rgb(&lin, w, h))).clone(); let l = host.log_l.get_or_insert_with(|| Arc::new(cx.read_plane(&prep.log_l, w, h))).clone(); - lightcraft_pipeline::masks::evaluate_one(m, &plan.frame, w, h, &img, &l, s.light.exposure as f32) + lightcraft_pipeline::masks::evaluate_one(m, &plan.frame, w, h, &img, &l, s.light.exposure as f32, plan.mattes.as_deref()) } }); lightcraft_pipeline::visualize::apply(&mut image, req.overlay, &plan, overlay_mask.as_ref()); @@ -1018,7 +1018,7 @@ fn masks(cx: &mut Cx<'_>, lin: &Buf, prep: &Prep, plan: &Plan<'_>, host: &mut Ho // no kernel: evaluate on the CPU let img = host.lin.get_or_insert_with(|| Arc::new(cx.read_rgb(lin, w, h))).clone(); let l = host.log_l.get_or_insert_with(|| Arc::new(cx.read_plane(&prep.log_l, w, h))).clone(); - let mut v = lightcraft_pipeline::masks::shape_alpha(&comp.shape, frame, w, h, &img, &l, ev); + let mut v = lightcraft_pipeline::masks::shape_alpha(&comp.shape, frame, w, h, &img, &l, ev, plan.mattes.as_deref()); if comp.invert { v.data.iter_mut().for_each(|x| *x = 1.0 - *x); } diff --git a/crates/heif/Cargo.toml b/crates/heif/Cargo.toml new file mode 100644 index 00000000..9a12bb2e --- /dev/null +++ b/crates/heif/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "lightcraft-heif" +description = "Optional HEIF/HEIC still-image decoder for LightCraft (a thin, panic-guarded wrapper around heic-rs)" +version.workspace = true +edition.workspace = true +license.workspace = true + +[dependencies] +# Pure Rust, no unsafe. Pinned exactly: a young, single-owner crate that parses untrusted files; +# moving the pin is a reviewed change (re-run the heif tests and the codecs corpus tests). +# Mirrors PhotoCraft's photocraft-heif (same team, same decision): heic-rs was chosen over the +# other permissive candidates (heif-oxide lacks 4:2:2; heifer was published the same week). +heic-rs = { version = "=0.1.1", default-features = false, features = ["std"] } + +# The tiles of a HEIC photo decode in parallel on native; wasm stays single-threaded. +[target.'cfg(not(target_arch = "wasm32"))'.dependencies] +heic-rs = { version = "=0.1.1", default-features = false, features = ["std", "parallel"] } + +[lints] +workspace = true diff --git a/crates/heif/README.md b/crates/heif/README.md new file mode 100644 index 00000000..e008a01c --- /dev/null +++ b/crates/heif/README.md @@ -0,0 +1,33 @@ +# lightcraft-heif + +The optional HEIF / HEIC decoder of LightCraft: iPhone and Mac photos. A thin, panic-guarded wrapper +around [`heic-rs`](https://github.com/tbraun96/heic-rs), a pure-Rust HEVC still-picture decoder +(no `unsafe`, MIT OR Apache-2.0): single pictures and grid-tiled photos, 8- and 10-bit (10-bit +decodes to 16-bit), alpha auxiliary images, the container's rotation/mirror/crop, ICC, EXIF and XMP. +Read-only: writing would need an HEVC encoder. + +```rust +let info = lightcraft_heif::probe(&bytes)?; // size, depth, alpha: check limits first +let img = lightcraft_heif::decode(&bytes, &lightcraft_heif::Options::default())?; +// img.width, img.height, img.has_alpha, img.sixteen_bit, img.data (RGB/RGBA), img.icc, img.exif, img.xmp +``` + +Errors are `Error::Unsupported` (image sequences, overlays, …), `Error::Limit` or +`Error::Malformed`. It never panics: every call into heic-rs runs under `catch_unwind`, and a panic +inside it (PhotoCraft's fuzzing found two in 0.1.1) becomes `Error::Malformed`. The crate is a port +of PhotoCraft's `photocraft-heif` — the sibling apps deliberately share one decoder, one pin and +one set of regression fixtures. + +## Why it is a separate, optional crate + +- **A young decoder.** heic-rs is new and has a single maintainer. It is pinned exactly + (`=0.1.1`); moving the pin is a reviewed change (re-run the crate tests and + `cargo test -p lightcraft-codecs --features heif`). +- **HEVC patents are a distributor's call.** HEVC is patent-encumbered in some jurisdictions, so + whether a build includes an HEVC decoder is a build-time choice. `lightcraft-codecs` uses this + crate only behind its non-default `heif` feature; the apps forward it + (`cargo build -p lightcraft --features heif`). Without it, HEIC files are still recognised and + opening one is a clear "HEIC/HEIF support isn't included in this build" error — the same policy + as PhotoCraft. + +The crate depends on no other LightCraft crate; `codecs` → `heif` is the only edge between them. diff --git a/crates/heif/src/lib.rs b/crates/heif/src/lib.rs new file mode 100644 index 00000000..7b19b9b9 --- /dev/null +++ b/crates/heif/src/lib.rs @@ -0,0 +1,229 @@ +//! `lightcraft-heif`: the optional HEIF / HEIC (iPhone and Mac photo) decoder. +//! +//! A thin wrapper around `heic-rs`, a pure-Rust HEVC still-picture decoder: single pictures and +//! grid-tiled photos, 8- and 10-bit (10-bit decodes to 16-bit), alpha auxiliary images, ICC, +//! EXIF and XMP. Its API is plain data ([`Info`], [`Decoded`], [`Error`]) so the crate knows +//! nothing about the rest of LightCraft; `lightcraft-codecs` adapts it behind its `heif` feature. +//! A port of PhotoCraft's `photocraft-heif` (same team, same licence) so the sibling apps share +//! one decision, one dependency pin and one set of regression fixtures. +//! +//! HEIF records orientation in the container (`irot`/`imir`, plus a `clap` crop), not in EXIF. +//! [`Options::apply_transforms`] applies them; with it off, the pixels come back as coded. +//! +//! Never panics: heic-rs is young, so every call into it runs under `catch_unwind` and a panic +//! inside it becomes [`Error::Malformed`]. + +#![forbid(unsafe_code)] +#![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)] + +use std::borrow::Cow; + +use heic_rs::PixelLayout; + +/// What a HEIF file declares, read from the container alone (no pixel is decoded). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Info { + /// Size after the container transforms (rotation, mirror, crop): what a viewer shows. + pub width: u32, + pub height: u32, + /// Size as coded, before the transforms. + pub coded_width: u32, + pub coded_height: u32, + /// The primary image has an alpha auxiliary image. + pub has_alpha: bool, + /// More than 8 bits per sample: it decodes to 16-bit. + pub sixteen_bit: bool, +} + +/// Decode settings. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Options { + /// Refuse images with more pixels than this (checked before decoding). + pub max_pixels: u64, + /// Apply the container's rotation, mirror and crop (`irot`/`imir`/`clap`). + pub apply_transforms: bool, +} + +impl Default for Options { + fn default() -> Self { + Options { max_pixels: 1 << 28, apply_transforms: true } + } +} + +/// A decoded image: interleaved RGB or RGBA samples, row-major, no padding. 16-bit samples are +/// native-endian `u16`s. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Decoded { + pub width: u32, + pub height: u32, + /// Four channels (RGBA) instead of three (RGB). + pub has_alpha: bool, + /// Two bytes per sample instead of one. + pub sixteen_bit: bool, + pub data: Vec, + /// The primary image's ICC profile. + pub icc: Option>, + /// EXIF as a TIFF structure (without HEIF's offset header). Its Orientation tag only mirrors + /// what the container says; it is returned verbatim. + pub exif: Option>, + /// The XMP packet. + pub xmp: Option, +} + +/// Why a file could not be decoded. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Error { + /// A valid file using something this decoder does not handle (image sequences, overlays…). + Unsupported(String), + /// The image is larger than [`Options::max_pixels`] or holds an oversized box. + Limit(String), + /// Broken or truncated data (or a decoder bug, reported the same way). + Malformed(String), +} + +impl std::fmt::Display for Error { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Error::Unsupported(m) => write!(f, "unsupported HEIF: {m}"), + Error::Limit(m) => write!(f, "limit exceeded: {m}"), + Error::Malformed(m) => write!(f, "malformed HEIF: {m}"), + } + } +} + +impl std::error::Error for Error {} + +const SEQUENCE: &str = "this file holds an image sequence; only HEIF still images can be opened"; + +fn err(e: heic_rs::Error) -> Error { + match e { + heic_rs::Error::Unsupported(what) => Error::Unsupported(what.to_string()), + // The primary item is neither an HEVC picture nor a grid of them. + heic_rs::Error::MissingBox("hvcC") => Error::Unsupported("the image is an overlay, an identity derivation or not HEVC-coded".into()), + heic_rs::Error::PixelLimit { .. } | heic_rs::Error::BoxTooLarge { .. } => Error::Limit(e.to_string()), + e => Error::Malformed(e.to_string()), + } +} + +/// Runs `f`, turning a heic-rs panic into an error. Fuzzing in the sibling app found two +/// out-of-range slices in heic-rs 0.1.1 on malformed files: in its box parser (`boxes.rs:130`) +/// and during reconstruction (`hevc/decode/recon.rs:70`). +fn guarded(f: impl FnOnce() -> Result) -> Result { + std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)).unwrap_or_else(|_| Err(Error::Malformed("the HEIF decoder failed on this file".into()))) +} + +/// Reads the container's declared size, depth and alpha without decoding pixels, so callers can +/// check their limits first. +pub fn probe(bytes: &[u8]) -> Result { + guarded(|| probe_unguarded(bytes)) +} + +fn probe_unguarded(bytes: &[u8]) -> Result { + let info = heic_rs::probe(bytes).map_err(|e| match e { + // No `meta` box: no still image, only a `moov` image sequence (a video track). + heic_rs::Error::MissingBox("meta") => Error::Unsupported(SEQUENCE.into()), + e => err(e), + })?; + Ok(Info { + width: info.width, + height: info.height, + coded_width: info.coded_width, + coded_height: info.coded_height, + has_alpha: info.has_alpha, + sixteen_bit: info.bit_depth > 8, + }) +} + +/// Decodes the primary image, with its metadata. +pub fn decode(bytes: &[u8], options: &Options) -> Result { + guarded(|| decode_unguarded(bytes, options)) +} + +fn decode_unguarded(bytes: &[u8], options: &Options) -> Result { + let info = probe_unguarded(bytes)?; + let pixel_layout = match (info.has_alpha, info.sixteen_bit) { + (true, true) => PixelLayout::Rgba16, + (true, false) => PixelLayout::Rgba8, + (false, true) => PixelLayout::Rgb16, + (false, false) => PixelLayout::Rgb8, + }; + let raw = heic_rs::DecodeOptions::default() + .with_layout(pixel_layout) + .with_max_pixels(Some(options.max_pixels)) + .with_transforms(options.apply_transforms) + .with_alpha(info.has_alpha); + let decoded = heic_rs::decode(bytes, &raw).map_err(err)?; + let (icc, exif, xmp) = metadata(bytes); + Ok(Decoded { + width: decoded.width, + height: decoded.height, + has_alpha: info.has_alpha, + sixteen_bit: info.sixteen_bit, + data: decoded.data, + icc, + exif, + xmp, + }) +} + +/// The primary image's ICC profile, EXIF (TIFF structure) and XMP packet. Metadata is a +/// courtesy: a broken metadata item never fails a decode whose pixels came out fine. +fn metadata(bytes: &[u8]) -> (Option>, Option>, Option) { + let Ok(ctx) = heic_rs::context::Context::open(bytes) else { + return (None, None, None); + }; + let id = ctx.meta.primary; + let icc = ctx.props(id).ok().and_then(|p| ctx.icc(&p).map(<[u8]>::to_vec)); + let exif = ctx.exif(id).ok().flatten().map(<[u8]>::to_vec); + let xmp = ctx.xmp_item(id).and_then(|x| ctx.item_data(x).ok()).and_then(|d| match d { + Cow::Borrowed(b) => String::from_utf8(b.to_vec()).ok(), + Cow::Owned(b) => String::from_utf8(b).ok(), + }); + (icc, exif, xmp.map(|x| x.trim_end_matches('\0').to_string())) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Found by PhotoCraft's `decode_heif` fuzz target: a malformed box makes heic-rs 0.1.1 slice + /// out of range (`boxes.rs:130`). It must be an error, not a crash. + const HEIC_RS_BOX_PANIC: [u8; 72] = [ + 0x00, 0x00, 0x00, 0x24, 0x66, 0x74, 0x79, 0x70, 0x68, 0x65, 0x69, 0x63, 0x00, 0x00, 0x00, 0x00, 0x6d, 0x69, 0x66, 0x31, 0x4d, 0x69, 0x50, + 0x72, 0x6d, 0x69, 0x61, 0x66, 0x4d, 0x69, 0x48, 0x42, 0x72, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x6d, 0x65, 0x74, 0x61, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x75, 0x75, 0x69, 0x64, 0x00, 0x00, 0x00, 0x00, 0x03, 0x08, 0x08, 0x08, 0x00, 0x03, 0x01, 0x03, 0x70, + 0x00, 0xa8, 0x00, + ]; + + #[test] + fn a_heic_rs_panic_is_a_malformed_error() { + for apply_transforms in [true, false] { + let r = decode(&HEIC_RS_BOX_PANIC, &Options { apply_transforms, ..Default::default() }); + assert!(matches!(r, Err(Error::Malformed(_))), "{r:?}"); + } + } + + #[test] + fn image_sequence_is_unsupported_with_a_clear_message() { + // An ftyp and nothing else: a HEIF image sequence would have a moov here instead of a meta. + let bytes = b"\0\0\0\x18ftypmsf1\0\0\0\0msf1hevc"; + for r in [probe(bytes).map(|_| ()), decode(bytes, &Options::default()).map(|_| ())] { + assert!(matches!(&r, Err(Error::Unsupported(m)) if m.contains("sequence")), "{r:?}"); + } + } + + #[test] + fn garbage_is_an_error() { + for bytes in [&b""[..], b"\0\0\0\x18ftypheic\0\0\0\0mif1heic", &[0xFF; 64]] { + assert!(probe(bytes).is_err()); + assert!(decode(bytes, &Options::default()).is_err()); + } + } + + #[test] + fn errors_display_their_kind() { + assert!(Error::Unsupported("x".into()).to_string().contains("unsupported")); + assert!(Error::Limit("x".into()).to_string().contains("limit")); + assert!(Error::Malformed("x".into()).to_string().contains("malformed")); + } +} diff --git a/crates/mcp/src/headless.rs b/crates/mcp/src/headless.rs index bdebec60..7347d33b 100644 --- a/crates/mcp/src/headless.rs +++ b/crates/mcp/src/headless.rs @@ -31,7 +31,7 @@ impl Drop for Headless { impl Default for Headless { fn default() -> Self { - Self::new(Session::new().with_fs()) + Self::new(Session::new().with_fs().with_default_face_models()) } } @@ -42,7 +42,7 @@ impl Headless { /// A headless session with the procedurally generated demo library. pub fn demo() -> Self { - Self::new(Session::with_demo().with_fs()) + Self::new(Session::with_demo().with_fs().with_default_face_models()) } fn photo_or_active(&self, p: &Value) -> Result { diff --git a/crates/merge/src/output.rs b/crates/merge/src/output.rs index 442fa2ee..e020ac45 100644 --- a/crates/merge/src/output.rs +++ b/crates/merge/src/output.rs @@ -29,10 +29,16 @@ pub enum DngSamples { U16, } -/// The DNG colour tags for frames of `color`. +/// The DNG colour tags for frames of `color`. A source's gain table map is dropped: it is +/// positioned on that one capture's active area, which a merged image (aligned, oriented, +/// stitched, cropped) no longer shares. pub fn color_data(color: &FrameColor) -> ColorData { match color { - FrameColor::Camera(c) => (**c).clone(), + FrameColor::Camera(c) => { + let mut c = (**c).clone(); + c.profile.gain_table_map = None; + c + } FrameColor::Linear { to_xyz_d50 } => { let to_xyz_d65 = bradford(D50, D65).mul(to_xyz_d50); let cm = to_xyz_d65.inverse().unwrap_or(Mat3::IDENTITY); @@ -45,6 +51,7 @@ pub fn color_data(color: &FrameColor) -> ColorData { as_shot_neutral: Some([1.0, 1.0, 1.0]), as_shot_white_xy: None, baseline_exposure: 0.0, + baseline_sharpness: None, profile: Default::default(), } } diff --git a/crates/pipeline/Cargo.toml b/crates/pipeline/Cargo.toml index 27d77816..c8219f02 100644 --- a/crates/pipeline/Cargo.toml +++ b/crates/pipeline/Cargo.toml @@ -13,6 +13,7 @@ lightcraft-geom = { workspace = true } lightcraft-color = { workspace = true } lightcraft-raster = { workspace = true } lightcraft-develop = { workspace = true } +lightcraft-raw = { workspace = true } rayon = { workspace = true } serde = { workspace = true } diff --git a/crates/pipeline/src/auto.rs b/crates/pipeline/src/auto.rs index 45100a1e..3455fff2 100644 --- a/crates/pipeline/src/auto.rs +++ b/crates/pipeline/src/auto.rs @@ -120,10 +120,17 @@ pub fn auto_wb(src: &Rgb32f, info: &SourceInfo) -> (f64, f64) { return (info.as_shot_temp, info.as_shot_tint); } let avg = acc.map(|v| v / wsum); - let xyz = REC2020.to_xyz().apply(avg); - let shot = lightcraft_color::cct::temp_tint_to_xy(info.as_shot_temp, info.as_shot_tint); - let seen = bradford(REC2020.white, shot).apply(xyz); - let (t, tint) = xy_to_temp_tint(Xy::from_xyz(seen)); + // the white that makes `avg` neutral, through the same model the white balance renders with + // (`local::wb_matrix_for`): the camera's own colour model, else a Bradford adaptation + let camera = info.camera_color.as_deref().filter(|_| info.raw && !info.relative_wb); + let seen = match camera.and_then(|cc| lightcraft_raw::color::neutral_white(&cc.tags, cc.developed_for, avg)) { + Some(white) => white, + None => { + let shot = lightcraft_color::cct::temp_tint_to_xy(info.as_shot_temp, info.as_shot_tint); + Xy::from_xyz(bradford(REC2020.white, shot).apply(REC2020.to_xyz().apply(avg))) + } + }; + let (t, tint) = xy_to_temp_tint(seen); (t.clamp(2000.0, 50000.0).round(), tint.clamp(-150.0, 150.0).round()) } diff --git a/crates/pipeline/src/lib.rs b/crates/pipeline/src/lib.rs index 057fb42f..84f5bd98 100644 --- a/crates/pipeline/src/lib.rs +++ b/crates/pipeline/src/lib.rs @@ -54,8 +54,19 @@ use lightcraft_raster::{Histogram, Plane, Rgb32f, Rgba8, par_rows}; pub use tone::ToneMap; +/// A raw source's own colour model, for white balance: like Lightroom, a white balance +/// re-evaluates the camera's colour matrices at the chosen white (camera-space white balance, DNG +/// spec ch. 6) instead of adapting the as-shot rendering ([`local::wb_matrix_for`]). +#[derive(Debug, PartialEq)] +pub struct CameraColor { + /// The file's colour tags (its profile look left out: it is applied at load). + pub tags: lightcraft_raw::ColorData, + /// The white the source pixels were developed for. + pub developed_for: lightcraft_color::Xy, +} + /// Facts about the source the settings are interpreted against. -#[derive(Clone, Copy, Debug, PartialEq)] +#[derive(Clone, Debug, PartialEq)] pub struct SourceInfo { /// Lens corrections embedded in the file (DNG opcodes), relative to the EXIF-oriented source. pub lens: Option, @@ -66,12 +77,26 @@ pub struct SourceInfo { pub as_shot_tint: f64, /// No measured camera illuminant: WB adjustments are relative to the camera's rendered look. pub relative_wb: bool, + /// The camera's colour model (raw files with a colour matrix); `None`: white balance adapts + /// the developed pixels (Bradford, linear Rec.2020). + pub camera_color: Option>, pub camera_tone: Option, + /// Segmentation mattes stored in the file (DNG semantic masks): AI masks use them. + pub mattes: Option>, } impl Default for SourceInfo { fn default() -> Self { - Self { raw: false, as_shot_temp: 6500.0, as_shot_tint: 0.0, lens: None, relative_wb: false, camera_tone: None } + Self { + raw: false, + as_shot_temp: 6500.0, + as_shot_tint: 0.0, + lens: None, + relative_wb: false, + camera_color: None, + camera_tone: None, + mattes: None, + } } } @@ -332,6 +357,8 @@ pub struct Plan<'a> { /// [`spots::window_for_reads`]); this is the requested window inside the rendered one, which /// the result is cut to. pub keep: Option, + /// The source's segmentation mattes ([`SourceInfo::mattes`]). + pub mattes: Option>, } /// Long edge of the small render a windowed render estimates the whole frame's airlight from. @@ -392,7 +419,7 @@ pub fn plan<'a>(src: &Rgb32f, info: &SourceInfo, s: &'a DevelopSettings, req: &R [d.nr_luminance, d.nr_detail, d.nr_color, d.nr_color_detail, d.nr_color_smoothness].map(f64::to_bits), src_long, )); - Plan { settings, frame, w, h, px_per_long, src_long, geo, lin_key, eyes, fixed_air, keep } + Plan { settings, frame, w, h, px_per_long, src_long, geo, lin_key, eyes, fixed_air, keep, mattes: info.mattes.clone() } } /// Whether the scene-linear stage needs work only the CPU does (defringe, spot removal). @@ -474,7 +501,7 @@ fn render_impl(src: Src<'_>, info: &SourceInfo, s: &DevelopSettings, req: &Rende Some(p) if p.key == lin_key => p, _ => local::Planes { key: lin_key, ..Default::default() }, }; - let mut prep = local::prepare(lin.clone(), s, frame, px_per_long, req.quality, &mut planes); + let mut prep = local::prepare(lin.clone(), s, frame, px_per_long, req.quality, &mut planes, plan.mattes.as_deref()); if let Some(a) = plan.fixed_air { prep.air = a; } @@ -543,7 +570,7 @@ fn overlay_alpha(o: Overlay, plan: &Plan<'_>, prep: &Prepared) -> Option return Some(e.alpha.clone()); } let ev = plan.settings.light.exposure as f32; - Some(masks::evaluate_one(m, &plan.frame, plan.w, plan.h, &prep.img, &prep.log_l, ev)) + Some(masks::evaluate_one(m, &plan.frame, plan.w, plan.h, &prep.img, &prep.log_l, ev, plan.mattes.as_deref())) } /// Convenience: render a before/after pair side by side is up to the UI; this renders "before" diff --git a/crates/pipeline/src/local.rs b/crates/pipeline/src/local.rs index 61ec8482..fd7fefd3 100644 --- a/crates/pipeline/src/local.rs +++ b/crates/pipeline/src/local.rs @@ -22,13 +22,21 @@ pub fn white_balance(img: &mut Rgb32f, info: &SourceInfo, s: &DevelopSettings) { wb_gain(img, info, s, 1.0); } -/// The white-balance matrix (linear Rec.2020, luminance-preserving) for the settings, or `None` -/// when the as-shot white is kept. +/// The white-balance matrix (linear Rec.2020) for the settings, or `None` when the as-shot white +/// is kept. A raw source with its camera colour model ([`crate::CameraColor`]) is re-developed for +/// the new white in camera space, as Lightroom does (so a neutral under the chosen white renders +/// neutral and saturated colours move as the camera's matrices say); other sources are adapted +/// with Bradford, luminance-preserving. pub fn wb_matrix_for(info: &SourceInfo, s: &DevelopSettings) -> Option<[[f32; 3]; 3]> { let (t, tint) = effective_wb(info, s); if (t - info.as_shot_temp).abs() < 1e-6 && (tint - info.as_shot_tint).abs() < 1e-6 { return None; } + if let Some(cc) = info.camera_color.as_deref().filter(|_| info.raw && !info.relative_wb) + && let Some(m) = lightcraft_raw::color::rebalance(&cc.tags, cc.developed_for, temp_tint_to_xy(t, tint)) + { + return Some(m.to_f32()); + } let set = wb_matrix(&REC2020, temp_tint_to_xy(t, tint)); let shot = wb_matrix(&REC2020, temp_tint_to_xy(info.as_shot_temp, info.as_shot_tint)); let m = set.mul(&shot.inverse().unwrap_or(lightcraft_color::Mat3::IDENTITY)); @@ -279,7 +287,16 @@ pub fn plane_sigmas(s: &DevelopSettings, px_per_long: f64, q: Quality) -> PlaneS /// The spatial planes the per-pixel stage needs for `img` (white-balanced, before exposure), /// reusing whatever `planes` already holds for it. Missing planes are computed side by side (each /// one alone scales poorly: the guided filters work on small subsampled grids). -pub(crate) fn prepare(img: Arc, s: &DevelopSettings, frame: &Frame, px_per_long: f64, q: Quality, planes: &mut Planes) -> Prepared { +#[allow(clippy::too_many_arguments)] +pub(crate) fn prepare( + img: Arc, + s: &DevelopSettings, + frame: &Frame, + px_per_long: f64, + q: Quality, + planes: &mut Planes, + mattes: Option<&masks::Mattes>, +) -> Prepared { let log_l = planes.log_l.get_or_insert_with(|| Arc::new(timed("log_l", || img.map(log_lum)))).clone(); let PlaneSigmas { base: base_sigma, clarity: clarity_sigma, texture: texture_sigma, dark: dark_sigma, chroma: chroma_sigma } = plane_sigmas(s, px_per_long, q); @@ -324,7 +341,7 @@ pub(crate) fn prepare(img: Arc, s: &DevelopSettings, frame: &Frame, px_p None => (None, 1.0), }; let ev = s.light.exposure as f32; - let masks = timed("masks", || masks::evaluate(&s.masks, frame, img.width, img.height, &img, &log_l, ev)); + let masks = timed("masks", || masks::evaluate(&s.masks, frame, img.width, img.height, &img, &log_l, ev, mattes)); Prepared { img, log_l, base, clarity_blur, texture_blur, dark, chroma_blur, air, masks, px_per_long } } diff --git a/crates/pipeline/src/masks.rs b/crates/pipeline/src/masks.rs index aae261c5..4777c878 100644 --- a/crates/pipeline/src/masks.rs +++ b/crates/pipeline/src/masks.rs @@ -2,14 +2,96 @@ //! //! Shapes are defined in normalized oriented-image coordinates; sizes are in "long-edge units". //! Components combine with Add (max), Subtract (a·(1−c)) and Intersect (a·c). +//! +//! AI shapes (Sky, Subject, Background, People) use the segmentation [`Mattes`] the source carries +//! (DNG semantic masks, e.g. iPhone ProRAW's) when it has a matching one, else a heuristic. use lightcraft_develop::{BrushStroke, LocalAdjustments, Mask, MaskOp, MaskShape, SegMask}; -use lightcraft_geom::Point; -use lightcraft_raster::{Plane, Rgb32f}; +use lightcraft_geom::{Orientation, Point}; +use lightcraft_raster::{Image, Plane, Rgb32f}; use crate::for_rows; use crate::geometry::Frame; +/// What an embedded segmentation matte selects. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum MatteKind { + Sky, + /// A whole person (all people in the photo). + Person, + Skin, + Hair, + Teeth, + Glasses, +} + +/// Segmentation mattes that came with the source: 0..255 alpha planes, each covering the whole +/// (EXIF-oriented) source image at its own resolution. Several mattes of one kind (e.g. one per +/// person) are combined with max. +#[derive(Clone, Default, PartialEq)] +pub struct Mattes { + mattes: Vec<(MatteKind, Image)>, +} + +impl std::fmt::Debug for Mattes { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_list().entries(self.mattes.iter().map(|(k, m)| (k, m.width, m.height))).finish() + } +} + +impl Mattes { + /// Add a matte (empty ones are ignored). + pub fn push(&mut self, kind: MatteKind, matte: Image) { + if matte.width > 0 && matte.height > 0 && matte.data.len() == matte.width * matte.height { + self.mattes.push((kind, matte)); + } + } + + pub fn is_empty(&self) -> bool { + self.mattes.is_empty() + } + + /// Memory held (bytes). + pub fn bytes(&self) -> usize { + self.mattes.iter().map(|(_, m)| m.data.len()).sum() + } + + fn of(&self, kinds: &[MatteKind]) -> Vec<&Image> { + self.mattes.iter().filter(|(k, _)| kinds.contains(k)).map(|(_, m)| m).collect() + } + + /// The mattes a shape uses instead of its heuristic (empty: none fits). + fn for_shape(&self, shape: &MaskShape) -> Vec<&Image> { + use MatteKind::*; + // the whole person: a person matte, else the union of the parts + let person = || { + let p = self.of(&[Person]); + if p.is_empty() { self.of(&[Skin, Hair, Teeth, Glasses]) } else { p } + }; + match shape { + MaskShape::Sky => self.of(&[Sky]), + MaskShape::Subject => person(), + MaskShape::People { parts, .. } => { + let wanted: Vec = parts + .iter() + .filter_map(|p| { + let p = p.to_ascii_lowercase(); + [("skin", Skin), ("hair", Hair), ("teeth", Teeth), ("glasses", Glasses)] + .into_iter() + .find(|(n, _)| p.contains(n)) + .map(|(_, k)| k) + }) + .collect(); + let whole = + parts.is_empty() || parts.iter().any(|p| matches!(p.to_ascii_lowercase().as_str(), "person" | "entire person" | "entireperson")); + // parts without a matte (lips, clothes, …) keep the heuristic + if whole { person() } else { self.of(&wanted) } + } + _ => Vec::new(), + } + } +} + pub struct Evaluated { /// The mask's id ([`Mask::id`]). pub id: u32, @@ -23,22 +105,24 @@ fn smooth(e0: f32, e1: f32, x: f32) -> f32 { t * t * (3.0 - 2.0 * t) } -/// The visible masks with components, evaluated in order. -pub fn evaluate(masks: &[Mask], frame: &Frame, w: usize, h: usize, img: &Rgb32f, log_l: &Plane, ev: f32) -> Vec { +/// The visible masks with components, evaluated in order (`mattes`: the source's, see [`Mattes`]). +#[allow(clippy::too_many_arguments)] +pub fn evaluate(masks: &[Mask], frame: &Frame, w: usize, h: usize, img: &Rgb32f, log_l: &Plane, ev: f32, mattes: Option<&Mattes>) -> Vec { masks .iter() .filter(|m| m.visible && !m.components.is_empty()) - .map(|m| Evaluated { id: m.id, alpha: evaluate_one(m, frame, w, h, img, log_l, ev), adjust: m.adjust }) + .map(|m| Evaluated { id: m.id, alpha: evaluate_one(m, frame, w, h, img, log_l, ev, mattes), adjust: m.adjust }) .collect() } /// The alpha plane of mask `m` (whether visible or not): its components combined, inverted and /// scaled by its amount. -pub fn evaluate_one(m: &Mask, frame: &Frame, w: usize, h: usize, img: &Rgb32f, log_l: &Plane, ev: f32) -> Plane { +#[allow(clippy::too_many_arguments)] +pub fn evaluate_one(m: &Mask, frame: &Frame, w: usize, h: usize, img: &Rgb32f, log_l: &Plane, ev: f32, mattes: Option<&Mattes>) -> Plane { let mut alpha = Plane::new(w, h); let mut first = true; for comp in &m.components { - let mut c = shape_alpha(&comp.shape, frame, w, h, img, log_l, ev); + let mut c = shape_alpha(&comp.shape, frame, w, h, img, log_l, ev, mattes); if comp.invert { c.data.iter_mut().for_each(|v| *v = 1.0 - *v); } @@ -88,7 +172,13 @@ fn for_each_pos(frame: &Frame, w: usize, h: usize, out: &mut Plane, f: impl Fn(P }); } -pub fn shape_alpha(shape: &MaskShape, frame: &Frame, w: usize, h: usize, img: &Rgb32f, log_l: &Plane, ev: f32) -> Plane { +/// The alpha plane of one component shape. AI shapes use the source's matching `mattes`, if any. +#[allow(clippy::too_many_arguments)] +pub fn shape_alpha(shape: &MaskShape, frame: &Frame, w: usize, h: usize, img: &Rgb32f, log_l: &Plane, ev: f32, mattes: Option<&Mattes>) -> Plane { + let fitting = mattes.map(|m| m.for_shape(shape)).unwrap_or_default(); + if !fitting.is_empty() { + return matte_alpha(&fitting, frame, w, h); + } let mut out = Plane::new(w, h); let to_long = |p: Point| frame.norm_to_long(p); // `img`/`log_l` are before exposure: range masks select on the exposed values. @@ -180,7 +270,7 @@ pub fn shape_alpha(shape: &MaskShape, frame: &Frame, w: usize, h: usize, img: &R smooth_plane(&mut out, 0.015 * frame_px(frame, w)); } MaskShape::Background => { - let mut s = shape_alpha(&MaskShape::Subject, frame, w, h, img, log_l, ev); + let mut s = shape_alpha(&MaskShape::Subject, frame, w, h, img, log_l, ev, mattes); s.data.iter_mut().for_each(|v| *v = 1.0 - *v); out = s; } @@ -252,6 +342,71 @@ fn smooth_plane(p: &mut Plane, sigma: f32) { *p = lightcraft_raster::blur::gaussian(p, sigma.max(0.5)); } +/// The union (max) of `mattes` (0..255, over the EXIF-oriented source) at output resolution, +/// sampled like the image: user orientation, lens / perspective warp, crop and flips; area-averaged +/// first when the output is much smaller than the matte, bilinear otherwise. +fn matte_alpha(mattes: &[&Image], frame: &Frame, w: usize, h: usize) -> Plane { + let mut out = Plane::new(w, h); + let o2t = frame.out_to_oriented(w, h); + let ppl = frame.px_per_long(w).max(1e-9); + for m in mattes { + let oriented = + if frame.orient == Orientation::Normal { std::borrow::Cow::Borrowed(*m) } else { std::borrow::Cow::Owned(m.oriented(frame.orient)) }; + // matte px per output px + let k = oriented.width as f64 / frame.ow * frame.ow.max(frame.oh) / ppl; + let base = box_down(&oriented, if k >= 2.0 { (k.floor() as usize).min(64) } else { 1 }); + let (sx, sy) = ((base.width as f64 / frame.ow) as f32, (base.height as f64 / frame.oh) as f32); + for_rows(&mut out.data, w, |y, row| { + for (x, v) in row.iter_mut().enumerate() { + let s = match &frame.warp { + Some(wp) => { + let Some((_, s)) = wp.frame(&o2t, x, y) else { + continue; + }; + s + } + None => o2t.apply(Point::new(x as f64 + 0.5, y as f64 + 0.5)), + }; + *v = v.max(bilinear(&base, s.x as f32 * sx, s.y as f32 * sy) / 255.0); + } + }); + } + out +} + +/// `m` area-averaged by `f × f` (0..255 values kept); `f == 1` converts only. +fn box_down(m: &Image, f: usize) -> Plane { + let f = f.max(1); + let (bw, bh) = (m.width.div_ceil(f).max(1), m.height.div_ceil(f).max(1)); + let mut out = Plane::new(bw, bh); + for_rows(&mut out.data, bw, |by, row| { + let rows = by * f..((by + 1) * f).min(m.height); + for (bx, v) in row.iter_mut().enumerate() { + let cols = bx * f..((bx + 1) * f).min(m.width); + let (mut sum, mut n) = (0u32, 0u32); + for y in rows.clone() { + for x in cols.clone() { + sum += m.data.get(y * m.width + x).copied().unwrap_or(0) as u32; + n += 1; + } + } + *v = sum as f32 / n.max(1) as f32; + } + }); + out +} + +/// Bilinear sample at continuous pixel coordinates (pixel centres at +0.5), edges extended. +fn bilinear(p: &Plane, x: f32, y: f32) -> f32 { + let (fx, fy) = (x - 0.5, y - 0.5); + let (x0, y0) = (fx.floor(), fy.floor()); + let (tx, ty) = (fx - x0, fy - y0); + let (x0, y0) = (x0 as isize, y0 as isize); + let a = p.get_clamped(x0, y0) + (p.get_clamped(x0 + 1, y0) - p.get_clamped(x0, y0)) * tx; + let b = p.get_clamped(x0, y0 + 1) + (p.get_clamped(x0 + 1, y0 + 1) - p.get_clamped(x0, y0 + 1)) * tx; + a + (b - a) * ty +} + /// A rough display mapping for colour picking (so samples taken on screen match). pub(crate) fn tonemap_for_select(c: [f32; 3]) -> [f32; 3] { c.map(|v| v / (1.0 + v)) @@ -436,7 +591,7 @@ mod tests { let f = frame(100, 50); let img = Rgb32f::new(100, 50); let l = Plane::new(100, 50); - let a = shape_alpha(&MaskShape::Linear { start: Point::new(0.0, 0.0), end: Point::new(1.0, 0.0) }, &f, 100, 50, &img, &l, 0.0); + let a = shape_alpha(&MaskShape::Linear { start: Point::new(0.0, 0.0), end: Point::new(1.0, 0.0) }, &f, 100, 50, &img, &l, 0.0, None); assert!(a.get(1, 25) > 0.99 && a.get(98, 25) < 0.01); assert!((a.get(50, 25) - 0.5).abs() < 0.05); } @@ -452,7 +607,7 @@ mod tests { MaskShape::Prompt { text: "left".into(), seg: Some(seg.clone()), detail: vec![], edge: 0.0 }, ] { for (w, h) in [(40, 20), (400, 200)] { - let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0); + let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0, None); assert!(a.get(1, h / 2) > 0.99 && a.get(w - 2, h / 2) < 0.01, "{w}×{h}"); if w >= 400 { let soft = (0..w).filter(|x| (0.05..0.95).contains(&a.get(*x, h / 2))).count(); @@ -465,7 +620,7 @@ mod tests { d.crop.geometry.rect = lightcraft_geom::Rect::from_center(Point::new(0.75, 0.5), 0.5, 1.0); let f = Frame::new(40, 20, &d, true); let shape = MaskShape::Prompt { text: "left".into(), seg: Some(seg), detail: vec![], edge: 0.0 }; - let a = shape_alpha(&shape, &f, 20, 20, &Rgb32f::new(20, 20), &Plane::new(20, 20), 0.0); + let a = shape_alpha(&shape, &f, 20, 20, &Rgb32f::new(20, 20), &Plane::new(20, 20), 0.0, None); assert!(a.data.iter().all(|v| *v < 0.05)); } @@ -478,7 +633,7 @@ mod tests { let patch = SegMask::from_logits_in(side, &l, [0.5, 0.0, 1.0, 1.0]); let shape = MaskShape::Object { hint: vec![Point::new(0.6, 0.5)], exclude: vec![], seg: Some(coarse), detail: vec![patch], edge: 0.0 }; let (w, h) = (200, 100); - let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0); + let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0, None); assert!(a.get(20, 50) < 0.01, "outside the patch: the coarse mask"); assert!(a.get(120, 50) > 0.99, "inside the patch, its selected half"); assert!(a.get(190, 50) < 0.01, "inside the patch, its unselected half"); @@ -494,7 +649,7 @@ mod tests { let shape = MaskShape::Object { hint: vec![Point::new(0.2, 0.5)], exclude: vec![], seg: None, detail: vec![patch; 20_000], edge: 0.0 }; let (w, h) = (64, 32); let t = std::time::Instant::now(); - let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0); + let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0, None); assert!(a.get(5, 16) > 0.99 && a.get(60, 16) < 0.01); assert!(t.elapsed() < std::time::Duration::from_secs(20), "{:?}", t.elapsed()); } @@ -507,7 +662,7 @@ mod tests { let (w, h) = (400, 40); let soft_px = |edge: f64| { let shape = MaskShape::Prompt { text: "x".into(), seg: Some(seg.clone()), detail: vec![], edge }; - let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0); + let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0, None); assert!(a.get(2, h / 2) > 0.9 && a.get(w - 3, h / 2) < 0.1, "still the same selection at {edge}"); (0..w).filter(|x| (0.1..0.9).contains(&a.get(*x, h / 2))).count() }; @@ -524,7 +679,7 @@ mod tests { MaskShape::Prompt { text: "sky".into(), seg: None, detail: vec![], edge: 0.0 }, MaskShape::Prompt { text: "sky".into(), seg: Some(SegMask { side: 4, data: "damaged!".into(), rect: None }), detail: vec![], edge: 0.0 }, ] { - assert!(shape_alpha(&shape, &f, 30, 30, &img, &l, 0.0).data.iter().all(|v| *v == 0.0)); + assert!(shape_alpha(&shape, &f, 30, 30, &img, &l, 0.0, None).data.iter().all(|v| *v == 0.0)); } } @@ -534,10 +689,10 @@ mod tests { let img = Rgb32f::new(100, 100); let l = Plane::new(100, 100); let shape = MaskShape::Radial { center: Point::new(0.5, 0.5), rx: 0.2, ry: 0.2, angle: 0.0, feather: 20.0, invert: false }; - let a = shape_alpha(&shape, &f, 100, 100, &img, &l, 0.0); + let a = shape_alpha(&shape, &f, 100, 100, &img, &l, 0.0, None); assert!(a.get(50, 50) > 0.99 && a.get(5, 5) < 0.01); let m = Mask { components: vec![MaskComponent { name: None, op: MaskOp::Add, invert: true, shape }], ..Default::default() }; - let e = evaluate(&[m], &f, 100, 100, &img, &l, 0.0); + let e = evaluate(&[m], &f, 100, 100, &img, &l, 0.0, None); assert!(e[0].alpha.get(50, 50) < 0.01); } @@ -552,7 +707,7 @@ mod tests { components: vec![MaskComponent { name: None, op: MaskOp::Add, invert: false, shape: MaskShape::Brush { strokes: vec![stroke, erase] } }], ..Default::default() }; - let e = evaluate(&[m], &f, 200, 100, &img, &l, 0.0); + let e = evaluate(&[m], &f, 200, 100, &img, &l, 0.0, None); let a = &e[0].alpha; assert!(a.get(40, 50) > 0.9, "{}", a.get(40, 50)); assert!(a.get(100, 50) < 0.05, "erased centre {}", a.get(100, 50)); @@ -574,8 +729,8 @@ mod tests { ..Default::default() }; let comp = |auto| MaskShape::Brush { strokes: vec![stroke(auto)] }; - let plain = shape_alpha(&comp(false), &f, w, h, &img, &l, 0.0); - let auto = shape_alpha(&comp(true), &f, w, h, &img, &l, 0.0); + let plain = shape_alpha(&comp(false), &f, w, h, &img, &l, 0.0, None); + let auto = shape_alpha(&comp(true), &f, w, h, &img, &l, 0.0, None); // without Auto Mask the brush spills over the edge; with it, it stays on the dark side assert!(plain.get(105, 50) > 0.9, "{}", plain.get(105, 50)); assert!(auto.get(105, 50) < 0.05, "spill {}", auto.get(105, 50)); @@ -585,8 +740,8 @@ mod tests { // on a flat area Auto Mask paints like the plain brush let flat = Rgb32f::from_fn(w, h, |_, _| [0.2; 3]); let fl = flat.map(crate::local::log_lum); - let a = shape_alpha(&comp(true), &f, w, h, &flat, &fl, 0.0); - let b = shape_alpha(&comp(false), &f, w, h, &flat, &fl, 0.0); + let a = shape_alpha(&comp(true), &f, w, h, &flat, &fl, 0.0, None); + let b = shape_alpha(&comp(false), &f, w, h, &flat, &fl, 0.0, None); assert!((a.get(94, 50) - b.get(94, 50)).abs() < 0.02, "{} vs {}", a.get(94, 50), b.get(94, 50)); } @@ -609,7 +764,7 @@ mod tests { }; let refined = Mask { refine: 100.0, ..soft.clone() }; let step = |m: &Mask| { - let a = evaluate_one(m, &f, w, h, &img, &log_l, 0.0); + let a = evaluate_one(m, &f, w, h, &img, &log_l, 0.0, None); a.data[40 * w + 63] - a.data[40 * w + 56] }; let (plain, sharp) = (step(&soft), step(&refined)); @@ -640,7 +795,7 @@ mod tests { } let alpha = |refine| { let shape = MaskShape::ColorRange { samples: vec![sample], refine }; - let a = shape_alpha(&shape, &f, w, h, &img, &log_l, 0.0); + let a = shape_alpha(&shape, &f, w, h, &img, &log_l, 0.0, None); (a.get(30, 20), a.get(90, 20)) }; for refine in [0.0, 50.0, 85.0] { @@ -653,4 +808,63 @@ mod tests { let (l, r) = alpha(100.0); assert!(l > 0.95 && r < 0.3, "refine 100: {l} vs {r}"); } + + /// A half-resolution matte selecting the left `frac` of a `w × h` source. + fn left_matte(w: usize, h: usize, frac: f64) -> Image { + Image::from_fn(w / 2, h / 2, |x, _| if (x as f64 + 0.5) < frac * (w / 2) as f64 { 255 } else { 0 }) + } + + /// The Sky mask follows the source's sky matte where it has one (here: the bottom-left of a + /// uniformly bright blue photo, where the heuristic sees no sky, and not its top right, where + /// it does), through user orientation, at any output size; without one it keeps the heuristic. + #[test] + fn sky_uses_the_sources_matte() { + let (w, h) = (200, 100); + let img = Rgb32f::filled(w, h, [0.5, 0.7, 1.4]); + let l = img.map(crate::local::log_lum); + let mut mattes = Mattes::default(); + mattes.push(MatteKind::Sky, left_matte(w, h, 0.5)); + let f = frame(w, h); + let heuristic = shape_alpha(&MaskShape::Sky, &f, w, h, &img, &l, 0.0, None); + let matte = shape_alpha(&MaskShape::Sky, &f, w, h, &img, &l, 0.0, Some(&mattes)); + assert!(heuristic.get(150, 5) > 0.9 && heuristic.get(25, 90) < 0.1, "{} {}", heuristic.get(150, 5), heuristic.get(25, 90)); + assert!(matte.get(150, 5) < 0.01 && matte.get(25, 90) > 0.99, "{} {}", matte.get(150, 5), matte.get(25, 90)); + // mattes of other kinds leave the heuristic alone + let mut hair = Mattes::default(); + hair.push(MatteKind::Hair, left_matte(w, h, 0.5)); + assert_eq!(shape_alpha(&MaskShape::Sky, &f, w, h, &img, &l, 0.0, Some(&hair)), heuristic); + // rotated 90° clockwise, the source's left half is the top half; a preview and a larger + // render agree + let s = lightcraft_develop::DevelopSettings { orientation: Orientation::Rotate90, ..Default::default() }; + let f = Frame::new(w, h, &s, true); + for (ow, oh) in [(100, 200), (25, 50)] { + let small = Rgb32f::filled(ow, oh, [0.5, 0.7, 1.4]); + let a = shape_alpha(&MaskShape::Sky, &f, ow, oh, &small, &small.map(crate::local::log_lum), 0.0, Some(&mattes)); + let at = |x: f64, y: f64| a.get((x * ow as f64) as usize, (y * oh as f64) as usize); + assert!(at(0.5, 0.2) > 0.99 && at(0.5, 0.8) < 0.01, "{ow}×{oh}: {} {}", at(0.5, 0.2), at(0.5, 0.8)); + } + } + + /// Subject, Background and People use person mattes: the whole person is the union of the + /// parts; a People part without a matte keeps the heuristic. + #[test] + fn people_and_subject_use_person_mattes() { + let (w, h) = (200, 100); + let img = Rgb32f::filled(w, h, [0.2; 3]); + let l = img.map(crate::local::log_lum); + let mut mattes = Mattes::default(); + mattes.push(MatteKind::Skin, left_matte(w, h, 0.25)); + let right: Image = Image::from_fn(w / 2, h / 2, |x, _| if x >= 75 { 255 } else { 0 }); + mattes.push(MatteKind::Hair, right); + let f = frame(w, h); + let eval = |shape: &MaskShape| shape_alpha(shape, &f, w, h, &img, &l, 0.0, Some(&mattes)); + let subject = eval(&MaskShape::Subject); + assert!(subject.get(10, 50) > 0.99 && subject.get(190, 50) > 0.99 && subject.get(100, 50) < 0.01); + let background = eval(&MaskShape::Background); + assert!(background.get(100, 50) > 0.99 && background.get(10, 50) < 0.01); + let hair = eval(&MaskShape::People { person: 0, parts: vec!["Hair".into()] }); + assert!(hair.get(190, 50) > 0.99 && hair.get(10, 50) < 0.01); + let lips = MaskShape::People { person: 0, parts: vec!["Lips".into()] }; + assert_eq!(eval(&lips), shape_alpha(&lips, &f, w, h, &img, &l, 0.0, None)); + } } diff --git a/crates/pipeline/src/tests.rs b/crates/pipeline/src/tests.rs index 697a0f97..f572b21d 100644 --- a/crates/pipeline/src/tests.rs +++ b/crates/pipeline/src/tests.rs @@ -369,6 +369,8 @@ fn tint_negative_is_green_and_positive_is_magenta() { SourceInfo::default(), SourceInfo { raw: true, relative_wb: true, ..Default::default() }, SourceInfo { raw: true, as_shot_temp: 4200.0, as_shot_tint: 15.0, ..Default::default() }, + // camera-space white balance: the same sign through the camera's own colour model + SourceInfo { raw: true, as_shot_temp: 4200.0, as_shot_tint: 15.0, camera_color: Some(camera_color(4200.0, 15.0)), ..Default::default() }, ] { let mut s = DevelopSettings::default(); let neutral = render(&src, &info, &s, &RenderRequest::fit(16, 16)).image.get(8, 8); @@ -384,3 +386,52 @@ fn tint_negative_is_green_and_positive_is_magenta() { } } } + +/// A dual-illuminant (A / D65) camera colour model, its pixels developed for `temp` / `tint`. +fn camera_color(temp: f64, tint: f64) -> std::sync::Arc { + use lightcraft_color::Mat3; + let tags = lightcraft_raw::ColorData { + illuminant: [17, 21], + color_matrix: [ + Some(Mat3([[0.9, 0.2, -0.15], [-0.3, 1.25, 0.08], [0.02, -0.12, 0.85]])), + Some(Mat3([[0.7, 0.3, -0.1], [-0.35, 1.3, 0.1], [0.05, -0.2, 1.0]])), + ], + ..Default::default() + }; + std::sync::Arc::new(crate::CameraColor { tags, developed_for: lightcraft_color::cct::temp_tint_to_xy(temp, tint) }) +} + +/// Camera-space white balance: a raw source with its colour model is re-developed for the chosen +/// white, so a grey card lit by that white renders neutral. Without the model, or with a relative +/// white balance, the developed pixels are adapted (Bradford, neutral luminance kept). +#[test] +fn custom_white_balance_redevelops_in_camera_space() { + use crate::local::wb_matrix_for; + use lightcraft_color::cct::temp_tint_to_xy; + use lightcraft_develop::WbMode; + let cc = camera_color(5500.0, 0.0); + let info = SourceInfo { raw: true, as_shot_temp: 5500.0, as_shot_tint: 0.0, camera_color: Some(cc.clone()), ..Default::default() }; + let mut s = DevelopSettings::default(); + (s.wb.mode, s.wb.temp, s.wb.tint) = (WbMode::Custom, 3200.0, 10.0); + // a grey card lit by 3200 K / +10, as the camera records it, developed for the as-shot white + let a = lightcraft_raw::color::camera_transform_of(&cc.tags, cc.developed_for); + let n = lightcraft_raw::color::camera_neutral(&cc.tags, temp_tint_to_xy(3200.0, 10.0)); + let card = a.matrix.apply(std::array::from_fn(|i| n[i] * a.wb[i] as f64)).map(|v| v as f32); + let apply = |m: [[f32; 3]; 3], c: [f32; 3]| -> [f32; 3] { std::array::from_fn(|r| m[r][0] * c[0] + m[r][1] * c[1] + m[r][2] * c[2]) }; + let m = wb_matrix_for(&info, &s).unwrap(); + let out = apply(m, card); + assert!(out.iter().all(|v| (v / out[1] - 1.0).abs() < 1e-4), "{card:?} -> {out:?}"); + // the adaptation path for sources without a colour model, and for relative white balance + let luma = |c: [f32; 3]| lightcraft_color::luminance_2020(c); + for other in [SourceInfo { camera_color: None, ..info.clone() }, SourceInfo { relative_wb: true, ..info.clone() }] { + let b = wb_matrix_for(&other, &s).unwrap(); + assert_ne!(b, m); + assert!((luma(apply(b, [1.0; 3])) - 1.0).abs() < 1e-4); + } + // the white-balance picker / auto white balance invert the same model: the card names its white + let picked = crate::auto::auto_wb(&Rgb32f::filled(4, 4, card), &info); + assert_eq!(picked, (3200.0, 10.0)); + // the as-shot white needs nothing + s.wb.mode = WbMode::AsShot; + assert!(wb_matrix_for(&info, &s).is_none()); +} diff --git a/crates/raw/Cargo.toml b/crates/raw/Cargo.toml index a40fcc43..2ca9ef24 100644 --- a/crates/raw/Cargo.toml +++ b/crates/raw/Cargo.toml @@ -20,9 +20,16 @@ rayon = { workspace = true } miniz_oxide = { workspace = true } zune-jpeg = { workspace = true } zune-core = { workspace = true } +jxl-oxide = { workspace = true, optional = true } + +[features] +default = ["jxl"] +## JPEG XL compressed DNGs (DNG 1.7, compression 52546) via jxl-oxide (pure Rust, MIT OR Apache-2.0). +jxl = ["dep:jxl-oxide"] [dev-dependencies] jpeg-encoder = { workspace = true } proptest = { workspace = true } +zune-jpegxl = { workspace = true } # Pinned CC0 corpus identities; not a product dependency. sha2 = { version = "0.10", default-features = false, features = ["std"] } diff --git a/crates/raw/src/color.rs b/crates/raw/src/color.rs index 050d21ee..d738bb1f 100644 --- a/crates/raw/src/color.rs +++ b/crates/raw/src/color.rs @@ -181,7 +181,11 @@ pub struct CameraTransform { /// Camera transform for white `wb_xy`: `rec2020 = matrix · (wb ⊙ camera)`. pub fn camera_transform(raw: &RawImage, wb_xy: Xy) -> CameraTransform { - let color = &raw.color; + camera_transform_of(&raw.color, wb_xy) +} + +/// [`camera_transform`] from the colour tags alone. +pub fn camera_transform_of(color: &ColorData, wb_xy: Xy) -> CameraTransform { let fallback = !has_matrix(color); let wb = wb_multipliers(color, wb_xy); let to_d50 = camera_to_xyz_d50(color, wb_xy); @@ -192,6 +196,34 @@ pub fn camera_transform(raw: &RawImage, wb_xy: Xy) -> CameraTransform { CameraTransform { matrix, wb: wb.map(|v| v as f32), white_xy: wb_xy, matrix_is_fallback: fallback, baseline_exposure: color.baseline_exposure } } +/// Linear Rec.2020 → linear Rec.2020 matrix that re-develops pixels developed with the camera +/// transform for white `from` as if developed for white `to`: the camera's own colour model is +/// evaluated at the new white (camera-space white balance, DNG spec ch. 6), as Lightroom does, so +/// a neutral lit by `to` comes out neutral. `None` when the tags have no colour matrix or the +/// transform is singular. +pub fn rebalance(color: &ColorData, from: Xy, to: Xy) -> Option { + if !has_matrix(color) { + return None; + } + let (a, b) = (camera_transform_of(color, from), camera_transform_of(color, to)); + let ratio = Mat3::diag(f64::from(b.wb[0] / a.wb[0]), f64::from(b.wb[1] / a.wb[1]), f64::from(b.wb[2] / a.wb[2])); + Some(b.matrix.mul(&ratio).mul(&a.matrix.inverse()?)) +} + +/// The white under which `rgb` (linear Rec.2020, developed with the camera transform for white +/// `developed_for`) is a neutral: the white [`rebalance`] must re-develop for to make it neutral +/// (white-balance picker, auto white balance). `None` without a colour matrix, or when `rgb` maps +/// to a non-positive camera value. +pub fn neutral_white(color: &ColorData, developed_for: Xy, rgb: [f64; 3]) -> Option { + if !has_matrix(color) { + return None; + } + let a = camera_transform_of(color, developed_for); + let balanced = a.matrix.inverse()?.apply(rgb); + let camera: [f64; 3] = std::array::from_fn(|i| balanced[i] / f64::from(a.wb[i])); + camera.iter().all(|v| v.is_finite() && *v > 0.0).then(|| neutral_to_xy(color, camera)) +} + /// `(matrix, multipliers)` with `rec2020 = matrix · (multipliers ⊙ camera)`, linear Rec.2020 D65. pub fn camera_to_rec2020(raw: &RawImage, wb_xy: Xy) -> (Mat3, [f32; 3]) { let t = camera_transform(raw, wb_xy); @@ -373,6 +405,30 @@ mod tests { assert_eq!(illuminant_temperature(255), None); } + /// Camera-space white balance: re-developing for another white turns a neutral lit by that + /// white neutral (the camera's own matrices, interpolated for the new white, say how); the same + /// white is the identity; without a colour matrix there is nothing to re-evaluate. + #[test] + fn rebalance_redevelops_for_the_new_white() { + let color = ColorData { illuminant: [17, 21], color_matrix: [Some(cam_a()), Some(cam_b())], ..Default::default() }; + let (from, to) = (cct::temp_tint_to_xy(5500.0, 0.0), cct::temp_tint_to_xy(3200.0, 10.0)); + let a = camera_transform_of(&color, from); + // a grey card lit by `to`, as the camera records it, developed for `from`: warm + let n = camera_neutral(&color, to); + let developed = a.matrix.apply(std::array::from_fn(|i| n[i] * a.wb[i] as f64)); + assert!(developed[0] > 1.2 * developed[2], "{developed:?}"); + let out = rebalance(&color, from, to).unwrap().apply(developed); + assert!(close(out.map(|v| v / out[1]), [1.0; 3], 1e-6), "{out:?}"); + let same = rebalance(&color, from, from).unwrap(); + assert!((0..3).all(|r| close(same.0[r], Mat3::IDENTITY.0[r], 1e-9)), "{same:?}"); + assert!(rebalance(&ColorData::default(), from, to).is_none()); + // and back: the developed grey card names the white it was lit by + let white = neutral_white(&color, from, developed).unwrap(); + assert!((white.x - to.x).abs() < 1e-6 && (white.y - to.y).abs() < 1e-6, "{white:?} vs {to:?}"); + assert!(neutral_white(&color, from, [-1.0, 0.5, 0.5]).is_none()); + assert!(neutral_white(&ColorData::default(), from, developed).is_none()); + } + #[test] fn grey_world_balances() { let img = lightcraft_raster::Rgb32f::from_fn(20, 20, |x, _| { diff --git a/crates/raw/src/dng.rs b/crates/raw/src/dng.rs index 1e80c0b3..63ddf91a 100644 --- a/crates/raw/src/dng.rs +++ b/crates/raw/src/dng.rs @@ -1,13 +1,14 @@ //! DNG (Adobe Digital Negative Specification 1.7): raw IFD selection, pixel data (via [`crate::tiffraw`]), //! linearization, black/white levels, active area, default crop, CFA description, colour tags, opcode lists. +use crate::gaintable::GainTableMap; use crate::profile::{HsvTable, ProfileLook, ToneCurve}; use crate::tiffraw::{Packing, read_image_in}; use crate::{BlackLevel, Cfa, ColorData, Mat3, Mode, RawData, RawError, RawFormat, RawImage, Rect, Result, opcodes}; use lightcraft_color::Xy; use lightcraft_geom::Orientation; use lightcraft_tiff::tags::{self as t, photometric}; -use lightcraft_tiff::{Ifd, Tiff}; +use lightcraft_tiff::{ByteOrder, Ifd, Tiff}; /// The main raw IFD: full-resolution (NewSubfileType 0) CFA or LinearRaw image with the most pixels. pub(crate) fn raw_ifd(tiff: &Tiff) -> Option<&Ifd> { @@ -33,7 +34,8 @@ fn vec3(v: Option>) -> Option<[f64; 3]> { } /// Colour tags: DNG puts them in IFD0, but some writers use the raw IFD; prefer the raw IFD. -pub(crate) fn color_data(ifd0: &Ifd, raw: &Ifd) -> ColorData { +/// `order` is the file's byte order (for the binary gain-table-map tags). +pub(crate) fn color_data(ifd0: &Ifd, raw: &Ifd, order: ByteOrder) -> ColorData { let get = |tag: u16| raw.f64s(tag).or_else(|| ifd0.f64s(tag)); let geti = |tag: u16| raw.u16(tag).or_else(|| ifd0.u16(tag)); ColorData { @@ -46,13 +48,16 @@ pub(crate) fn color_data(ifd0: &Ifd, raw: &Ifd) -> ColorData { as_shot_white_xy: get(t::AS_SHOT_WHITE_XY).filter(|v| v.len() == 2 && v[0] > 0.0 && v[1] > 0.0).map(|v| Xy::new(v[0], v[1])), baseline_exposure: get(t::BASELINE_EXPOSURE).and_then(|v| v.first().copied()).filter(|v| v.is_finite()).unwrap_or(0.0) + get(t::BASELINE_EXPOSURE_OFFSET).and_then(|v| v.first().copied()).filter(|v| v.is_finite()).unwrap_or(0.0), - profile: profile_look(ifd0, raw), + // DNG 1.7 BaselineSharpness (IFD 0 or the enhanced IFD); implausible values are ignored + baseline_sharpness: get(t::BASELINE_SHARPNESS).and_then(|v| v.first().copied()).filter(|v| v.is_finite() && *v > 0.0 && *v <= 16.0), + profile: profile_look(ifd0, raw, order), } } -/// The profile look tags (`ProfileHueSatMap*`, `ProfileLookTable*`, `ProfileToneCurve`); malformed -/// ones are ignored. Like the colour tags, read from the raw IFD first, else IFD 0. -pub(crate) fn profile_look(ifd0: &Ifd, raw: &Ifd) -> ProfileLook { +/// The profile look tags (`ProfileHueSatMap*`, `ProfileLookTable*`, `ProfileToneCurve`, +/// `ProfileGainTableMap*`); malformed ones are ignored. Like the colour tags, read from the raw IFD +/// first, else IFD 0. +pub(crate) fn profile_look(ifd0: &Ifd, raw: &Ifd, order: ByteOrder) -> ProfileLook { let pick = |tag: u16| if raw.contains(tag) { raw } else { ifd0 }; let table = |dims: u16, data: u16, enc: u16| { let ifd = pick(dims); @@ -67,9 +72,22 @@ pub(crate) fn profile_look(ifd0: &Ifd, raw: &Ifd) -> ProfileLook { ], look_table: table(t::PROFILE_LOOK_TABLE_DIMS, t::PROFILE_LOOK_TABLE_DATA, t::PROFILE_LOOK_TABLE_ENCODING), tone_curve: pick(t::PROFILE_TONE_CURVE).f64s(t::PROFILE_TONE_CURVE).and_then(|v| ToneCurve::from_tag(&v)), + gain_table_map: gain_table_map(ifd0, raw, order), } } +/// The gain table map (DNG 1.7 precedence): `ProfileGainTableMap2` from IFD 0 (the embedded +/// camera profile), else from the raw IFD; else `ProfileGainTableMap` from the raw IFD (where DNG +/// 1.6 put it, and where Apple writes it), else from IFD 0 (where DNG 1.7 says it was meant to be). +/// A malformed tag falls through to the next one. +fn gain_table_map(ifd0: &Ifd, raw: &Ifd, order: ByteOrder) -> Option { + let parse = |ifd: &Ifd, tag: u16| ifd.bytes(tag).and_then(|b| GainTableMap::parse(b, order, tag == t::PROFILE_GAIN_TABLE_MAP_2)); + parse(ifd0, t::PROFILE_GAIN_TABLE_MAP_2) + .or_else(|| parse(raw, t::PROFILE_GAIN_TABLE_MAP_2)) + .or_else(|| parse(raw, t::PROFILE_GAIN_TABLE_MAP)) + .or_else(|| parse(ifd0, t::PROFILE_GAIN_TABLE_MAP)) +} + pub(crate) fn decode(bytes: &[u8], mode: Mode) -> Result { let tiff = Tiff::parse(bytes)?; let ifd0 = &tiff.ifds[0]; @@ -79,9 +97,6 @@ pub(crate) fn decode(bytes: &[u8], mode: Mode) -> Result { if !(1..=4).contains(&cpp) { return Err(RawError::Unsupported(format!("{cpp} samples per pixel"))); } - if info.compression == t::compression::JPEG_XL { - return Err(RawError::Unsupported("JPEG XL DNG".into())); - } let mut data = read_image_in(mode, bytes, &info, tiff.order, Packing::Msb)?; let bits = info.bits() as u32; @@ -175,7 +190,7 @@ pub(crate) fn decode(bytes: &[u8], mode: Mode) -> Result { active_area, crop, orientation: Orientation::from_exif(ifd0.u16(t::ORIENTATION).unwrap_or(1)), - color: color_data(ifd0, raw), + color: color_data(ifd0, raw, tiff.order), wb_multipliers: None, linearized, opcodes, diff --git a/crates/raw/src/dngwrite.rs b/crates/raw/src/dngwrite.rs index bd46a2c5..2f04f016 100644 --- a/crates/raw/src/dngwrite.rs +++ b/crates/raw/src/dngwrite.rs @@ -182,6 +182,9 @@ pub fn write_dng(raw: &RawImage, opts: &DngWriteOptions) -> Result> { ifd.set(t::AS_SHOT_NEUTRAL, rat_vec(&m.map(|v| 1.0 / v as f64))); } ifd.set(t::BASELINE_EXPOSURE, Value::SRational(vec![srational(c.baseline_exposure)])); + if let Some(v) = c.baseline_sharpness { + ifd.set(t::BASELINE_SHARPNESS, rat_vec(&[v])); + } // the source's own profile look travels with its data (conversions, smart previews, merges) let p = &c.profile; let table_tags = |ifd: &mut IfdBuilder, table: &crate::profile::HsvTable, dims: u16, enc: u16| { @@ -214,6 +217,11 @@ pub fn write_dng(raw: &RawImage, opts: &DngWriteOptions) -> Result> { if let Some(curve) = &p.tone_curve { ifd.set(t::PROFILE_TONE_CURVE, Value::Float(curve.points.iter().flatten().copied().collect())); } + if let Some(map) = &p.gain_table_map { + // the raw IFD is IFD 0 here, where both tag versions are valid + let (version2, bytes) = map.to_bytes(opts.order); + ifd.set(if version2 { t::PROFILE_GAIN_TABLE_MAP_2 } else { t::PROFILE_GAIN_TABLE_MAP }, Value::Undefined(bytes)); + } for (list, tag) in [(&raw.opcodes.list1, t::OPCODE_LIST_1), (&raw.opcodes.list2, t::OPCODE_LIST_2), (&raw.opcodes.list3, t::OPCODE_LIST_3)] { if !list.is_empty() { ifd.set(tag, Value::Undefined(opcodes::write_list(list))); diff --git a/crates/raw/src/gaintable.rs b/crates/raw/src/gaintable.rs new file mode 100644 index 00000000..750a1c70 --- /dev/null +++ b/crates/raw/src/gaintable.rs @@ -0,0 +1,325 @@ +//! DNG `ProfileGainTableMap` (tag 52525, DNG 1.6) and `ProfileGainTableMap2` (tag 52544, DNG 1.7): +//! a coarse grid of 1D gain tables carrying the maker's local tone mapping. Apple ProRAW has one. +//! +//! We read and write the tag (a DNG export keeps it) but don't render it: Lightroom Classic +//! renders Apple ProRAW without the map (see [`crate::profile`]), and so do we. +//! +//! - The grid has `MapPointsV × MapPointsH` tables of `MapPointsN` gains. Its origin and spacing +//! are relative to the active area (1.0 = the active area's height or width). +//! - The table input is the dot product of (R, G, B, min, max) with the tag's five weights; +//! version 2 then raises it to the tag's `Gamma`. +//! - Version 2 stores gains as u8 or u16 (mapped linearly onto `GainMin..=GainMax`), f16 or f32. +//! It takes precedence over version 1 when both are present. +//! +//! Malformed tags are ignored: a wrong size, zero or huge dimensions, non-finite or negative gains, +//! non-positive spacing, or gamma outside 0.25..=4. +//! +//! Byte order: the tag's fields are TIFF types (LONG, DOUBLE, FLOAT, …), and neither DNG 1.6 nor +//! 1.7.1 gives this tag a byte order of its own, as they do for opcode lists and +//! `OriginalRawFileData` (always big-endian). What applies is the general rule of the DNG 1.6 +//! specification ("DNG Format Overview" → "Byte Order", worded the same in 1.7.1): +//! +//! > DNG readers are required to support either byte order, even for files from a particular +//! > camera model. +//! +//! (, "DNG Specification".) So we read +//! the fields in the file's byte order, and fall back to big-endian for writers that followed the +//! opcode-list convention. A read in the wrong order can't pass for a valid map: one of +//! `MapPointsV` / `MapPointsH` is below 256 (their product is at most [`MAX_TABLES`]), so swapped +//! it is at least 2^24 and the size checks reject it. We write in the file's byte order. + +use lightcraft_tiff::ByteOrder; +use serde::{Deserialize, Serialize}; + +/// Upper bound on `MapPointsV × MapPointsH` (Apple uses 6 × 8). +const MAX_TABLES: usize = 1 << 12; +/// Upper bound on the total number of gains (Apple: 6 × 8 × 257 = 12 336). +const MAX_GAINS: usize = 1 << 22; + +/// A parsed `ProfileGainTableMap` / `ProfileGainTableMap2`. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct GainTableMap { + /// Tables vertically (`MapPointsV`) and horizontally (`MapPointsH`), and points per table (`MapPointsN`). + pub points_v: usize, + pub points_h: usize, + pub points_n: usize, + /// Grid spacing and origin, relative to the active area's height (`_v`) and width (`_h`). + pub spacing_v: f64, + pub spacing_h: f64, + pub origin_v: f64, + pub origin_h: f64, + /// Weights of (R, G, B, min, max) for the table input. + pub weights: [f32; 5], + /// Exponent for the table input (1.0 for the version-1 tag). + pub gamma: f32, + /// `points_v × points_h × points_n` gains, row-major with the table points innermost. Integer + /// storage has already been mapped onto `GainMin..=GainMax`. Mapping before interpolating is + /// the same as mapping after, because the mapping is linear. + pub gains: Vec, +} + +impl GainTableMap { + /// Parse the tag's bytes from a file in byte order `order` (`version2`: `ProfileGainTableMap2`): + /// in that order, else big-endian (see the module docs). `None` when malformed. + pub fn parse(b: &[u8], order: ByteOrder, version2: bool) -> Option { + Self::parse_in(b, order, version2).or_else(|| if order == ByteOrder::Big { None } else { Self::parse_in(b, ByteOrder::Big, version2) }) + } + + /// Parse the tag's bytes with every multi-byte field in `order`. + fn parse_in(b: &[u8], order: ByteOrder, version2: bool) -> Option { + let take = |o: usize, n: usize| b.get(o..o.checked_add(n)?); + let u32_at = |o: usize| take(o, 4).and_then(|s| s.try_into().ok()).map(|a: [u8; 4]| order.u32(a)); + let f32_at = |o: usize| u32_at(o).map(f32::from_bits); + let f64_at = |o: usize| take(o, 8).and_then(|s| s.try_into().ok()).map(|a: [u8; 8]| f64::from_bits(order.u64(a))); + let (points_v, points_h) = (u32_at(0)? as usize, u32_at(4)? as usize); + let (spacing_v, spacing_h, origin_v, origin_h) = (f64_at(8)?, f64_at(16)?, f64_at(24)?, f64_at(32)?); + let points_n = u32_at(40)? as usize; + let mut weights = [0f32; 5]; + for (i, w) in weights.iter_mut().enumerate() { + *w = f32_at(44 + 4 * i)?; + } + let tables = points_v.checked_mul(points_h).filter(|&t| t > 0 && t <= MAX_TABLES)?; + let count = tables.checked_mul(points_n).filter(|&c| c > 0 && c <= MAX_GAINS)?; + let spacing_ok = |s: f64, points: usize| s.is_finite() && (s > 0.0 || points == 1); + if !spacing_ok(spacing_v, points_v) + || !spacing_ok(spacing_h, points_h) + || !origin_v.is_finite() + || !origin_h.is_finite() + || !weights.iter().all(|w| w.is_finite()) + { + return None; + } + let (header, data_type, gamma, gain_min, gain_max) = + if version2 { (80, u32_at(64)?, f32_at(68)?, f32_at(72)?, f32_at(76)?) } else { (64, 3, 1.0, 0.0, 1.0) }; + if !(0.25..=4.0).contains(&gamma) { + return None; + } + let width = match data_type { + 0 => 1, + 1 | 2 => 2, + 3 => 4, + _ => return None, + }; + if b.len() != count.checked_mul(width)?.checked_add(header)? { + return None; + } + let data = b.get(header..)?; + let scaled = |q: f32, max: f32| gain_min + q / max * (gain_max - gain_min); + let gains: Vec = match data_type { + 0 => data.iter().map(|&q| scaled(q as f32, 255.0)).collect(), + 1 => data.as_chunks::<2>().0.iter().map(|c| scaled(order.u16(*c) as f32, 65535.0)).collect(), + 2 => data.as_chunks::<2>().0.iter().map(|c| crate::unpack::f16_to_f32(order.u16(*c))).collect(), + _ => data.as_chunks::<4>().0.iter().map(|c| f32::from_bits(order.u32(*c))).collect(), + }; + if gains.len() != count || !gains.iter().all(|g| g.is_finite() && *g >= 0.0) { + return None; + } + Some(GainTableMap { points_v, points_h, points_n, spacing_v, spacing_h, origin_v, origin_h, weights, gamma, gains }) + } + + /// The tag's bytes in `order`: version 1 (DNG 1.6, f32 gains) when the gamma is 1, else + /// version 2 with f32 gains. Returns `(version2, bytes)`. + pub fn to_bytes(&self, order: ByteOrder) -> (bool, Vec) { + let version2 = self.gamma != 1.0; + let mut out = Vec::with_capacity(80 + 4 * self.gains.len()); + let u32b = |v: u32| match order { + ByteOrder::Big => v.to_be_bytes(), + ByteOrder::Little => v.to_le_bytes(), + }; + let f64b = |v: f64| match order { + ByteOrder::Big => v.to_be_bytes(), + ByteOrder::Little => v.to_le_bytes(), + }; + let dim = |v: usize| u32b(u32::try_from(v).unwrap_or(u32::MAX)); + out.extend_from_slice(&dim(self.points_v)); + out.extend_from_slice(&dim(self.points_h)); + for v in [self.spacing_v, self.spacing_h, self.origin_v, self.origin_h] { + out.extend_from_slice(&f64b(v)); + } + out.extend_from_slice(&dim(self.points_n)); + for w in self.weights { + out.extend_from_slice(&u32b(w.to_bits())); + } + if version2 { + // DataType 3 (f32); GainMin/GainMax are ignored for float data + for v in [3, self.gamma.to_bits(), 0f32.to_bits(), 1f32.to_bits()] { + out.extend_from_slice(&u32b(v)); + } + } + for g in &self.gains { + out.extend_from_slice(&u32b(g.to_bits())); + } + (version2, out) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Tag bytes: the header fields (`v2`: data type, gamma, gain min, gain max), then `data`. + fn tag( + order: ByteOrder, + dims: [u32; 3], + spacing: [f64; 2], + origin: [f64; 2], + weights: [f32; 5], + v2: Option<(u32, f32, f32, f32)>, + data: &[u8], + ) -> Vec { + let be = order == ByteOrder::Big; + let u = |v: u32| if be { v.to_be_bytes() } else { v.to_le_bytes() }; + let d = |v: f64| if be { v.to_be_bytes() } else { v.to_le_bytes() }; + let mut b = Vec::new(); + b.extend_from_slice(&u(dims[0])); + b.extend_from_slice(&u(dims[1])); + for v in [spacing[0], spacing[1], origin[0], origin[1]] { + b.extend_from_slice(&d(v)); + } + b.extend_from_slice(&u(dims[2])); + for w in weights { + b.extend_from_slice(&u(w.to_bits())); + } + if let Some((dt, gamma, min, max)) = v2 { + for v in [dt, gamma.to_bits(), min.to_bits(), max.to_bits()] { + b.extend_from_slice(&u(v)); + } + } + b.extend_from_slice(data); + b + } + + fn floats(order: ByteOrder, g: &[f32]) -> Vec { + g.iter().flat_map(|v| if order == ByteOrder::Big { v.to_be_bytes() } else { v.to_le_bytes() }).collect() + } + + const GREEN: [f32; 5] = [0.0, 1.0, 0.0, 0.0, 0.0]; + + /// 2 × 2 tables of N = 4 points: table (r, c) = base · [1, 2, 3, 4] with base 1, 2, 3, 4. + fn grid() -> Vec { + (0..4).flat_map(|t| (1..=4).map(move |i| (t + 1) as f32 * i as f32)).collect() + } + + #[test] + fn version_2_storage_types_and_gamma() { + let order = ByteOrder::Little; + let v2 = |dt: u32, gamma: f32, data: &[u8]| { + GainTableMap::parse(&tag(order, [1, 1, 4], [1.0, 1.0], [0.0, 0.0], GREEN, Some((dt, gamma, 0.5, 2.5)), data), order, true) + }; + // u8: 0 → GainMin, 255 → GainMax + let m = v2(0, 1.0, &[0, 51, 255, 255]).unwrap(); + let want = [0.5, 0.9, 2.5, 2.5]; + assert!(m.gains.iter().zip(want).all(|(g, w)| (g - w).abs() < 1e-6), "{:?}", m.gains); + // u16 + let m = v2(1, 1.0, &[0, 0, 0xff, 0xff, 0, 0, 0xff, 0xff]).unwrap(); + assert_eq!(m.gains, vec![0.5, 2.5, 0.5, 2.5]); + // f16 ignores GainMin/GainMax: 1.0, 2.0, 0.5, 0 + let m = v2(2, 1.0, &[0x00, 0x3c, 0x00, 0x40, 0x00, 0x38, 0, 0]).unwrap(); + assert_eq!(m.gains, vec![1.0, 2.0, 0.5, 0.0]); + // f32 with gamma 2 + let m = v2(3, 2.0, &floats(order, &[1.0, 2.0, 3.0, 4.0])).unwrap(); + assert_eq!((m.gains.clone(), m.gamma), (vec![1.0, 2.0, 3.0, 4.0], 2.0)); + // unknown data type, gamma out of range, size for another type: ignored + assert!(v2(4, 1.0, &[0; 4]).is_none()); + assert!(v2(0, 0.1, &[0; 4]).is_none()); + assert!(v2(0, 5.0, &[0; 4]).is_none()); + assert!(v2(1, 1.0, &[0; 4]).is_none()); + // negative integer range + let neg = tag(order, [1, 1, 4], [1.0, 1.0], [0.0, 0.0], GREEN, Some((0, 1.0, -1.0, 1.0)), &[0; 4]); + assert!(GainTableMap::parse(&neg, order, true).is_none()); + } + + #[test] + fn malformed_tags_are_rejected() { + let order = ByteOrder::Big; + let good = tag(order, [2, 2, 4], [0.5, 0.5], [0.25, 0.25], GREEN, None, &floats(order, &grid())); + assert!(GainTableMap::parse(&good, order, false).is_some()); + // every truncation (and an extra byte) + for n in 0..good.len() { + assert!(GainTableMap::parse(&good[..n], order, false).is_none(), "{n} bytes"); + } + let mut long = good.clone(); + long.push(0); + assert!(GainTableMap::parse(&long, order, false).is_none()); + // the wrong byte order reads absurd dimensions + assert!(GainTableMap::parse_in(&good, ByteOrder::Little, false).is_none()); + // read as version 2 the header is too short for the data + assert!(GainTableMap::parse(&good, order, true).is_none()); + let bad = |dims: [u32; 3], spacing: [f64; 2], origin: [f64; 2], gains: &[f32]| { + GainTableMap::parse(&tag(order, dims, spacing, origin, GREEN, None, &floats(order, gains)), order, false) + }; + let mut g = grid(); + g[5] = f32::NAN; + assert!(bad([2, 2, 4], [0.5, 0.5], [0.0, 0.0], &g).is_none(), "NaN gain"); + g[5] = -1.0; + assert!(bad([2, 2, 4], [0.5, 0.5], [0.0, 0.0], &g).is_none(), "negative gain"); + g[5] = f32::INFINITY; + assert!(bad([2, 2, 4], [0.5, 0.5], [0.0, 0.0], &g).is_none(), "infinite gain"); + assert!(bad([2, 2, 4], [0.0, 0.5], [0.0, 0.0], &grid()).is_none(), "zero spacing"); + assert!(bad([2, 2, 4], [0.5, f64::NAN], [0.0, 0.0], &grid()).is_none(), "NaN spacing"); + assert!(bad([2, 2, 4], [0.5, 0.5], [f64::INFINITY, 0.0], &grid()).is_none(), "infinite origin"); + assert!(bad([0, 2, 4], [0.5, 0.5], [0.0, 0.0], &[]).is_none(), "no rows"); + assert!(bad([1, 1, 0], [0.5, 0.5], [0.0, 0.0], &[]).is_none(), "no points"); + // huge dimensions (whose product overflows or exceeds the cap) are rejected before allocating + assert!(bad([u32::MAX, u32::MAX, u32::MAX], [0.5, 0.5], [0.0, 0.0], &grid()).is_none()); + assert!(bad([4096, 4096, 1], [0.5, 0.5], [0.0, 0.0], &grid()).is_none()); + assert!(bad([1, 1, u32::MAX], [0.5, 0.5], [0.0, 0.0], &grid()).is_none()); + // a single table needs no spacing + assert!(bad([1, 1, 2], [0.0, 0.0], [0.0, 0.0], &[1.0, 2.0]).is_some()); + let mut w = GREEN; + w[3] = f32::NAN; + assert!(GainTableMap::parse(&tag(order, [1, 1, 1], [1.0, 1.0], [0.0, 0.0], w, None, &floats(order, &[1.0])), order, false).is_none()); + } + + /// The payload is read in the file's byte order, else big-endian; a little-endian payload in + /// a big-endian file is malformed. + #[test] + fn byte_order_is_the_files_with_a_big_endian_fallback() { + use ByteOrder::{Big, Little}; + let v1 = |o: ByteOrder| tag(o, [2, 2, 4], [0.5, 0.5], [0.25, 0.25], GREEN, None, &floats(o, &grid())); + let v2 = |o: ByteOrder| { + let mut gains = Vec::new(); + for q in [0u16, 0x1234, 0xffff, 0x8000] { + gains.extend_from_slice(&if o == Big { q.to_be_bytes() } else { q.to_le_bytes() }); + } + tag(o, [1, 1, 4], [1.0, 1.0], [0.0, 0.0], GREEN, Some((1, 2.0, 0.5, 2.5)), &gains) + }; + let (want1, want2) = (GainTableMap::parse_in(&v1(Big), Big, false).unwrap(), GainTableMap::parse_in(&v2(Big), Big, true).unwrap()); + assert_eq!(want1.gains, grid()); + assert_eq!((want2.gains[0], want2.gains[2], want2.gamma), (0.5, 2.5, 2.0)); + // in the file's order + assert_eq!(GainTableMap::parse(&v1(Little), Little, false).as_ref(), Some(&want1)); + assert_eq!(GainTableMap::parse(&v2(Little), Little, true).as_ref(), Some(&want2)); + assert_eq!(GainTableMap::parse(&v1(Big), Big, false).as_ref(), Some(&want1)); + // a big-endian payload in a little-endian file: the fallback + assert_eq!(GainTableMap::parse(&v1(Big), Little, false).as_ref(), Some(&want1)); + assert_eq!(GainTableMap::parse(&v2(Big), Little, true).as_ref(), Some(&want2)); + // no little-endian fallback in a big-endian file + assert!(GainTableMap::parse(&v1(Little), Big, false).is_none()); + assert!(GainTableMap::parse(&v2(Little), Big, true).is_none()); + // whatever the grid, a read in the wrong order fails the size checks + for dims in [[1, 1, 1], [6, 8, 257], [48, 64, 2], [255, 16, 1], [256, 16, 1], [1, 4096, 1]] { + for (o, wrong) in [(Big, Little), (Little, Big)] { + let g = vec![1.0; dims.iter().product::() as usize]; + let b = tag(o, dims, [0.5, 0.5], [0.0, 0.0], GREEN, None, &floats(o, &g)); + assert!(GainTableMap::parse_in(&b, o, false).is_some(), "{dims:?}"); + assert!(GainTableMap::parse_in(&b, wrong, false).is_none(), "{dims:?} read {wrong:?}"); + } + } + } + + #[test] + fn writes_what_it_reads() { + for order in [ByteOrder::Big, ByteOrder::Little] { + let m = + GainTableMap::parse(&tag(order, [2, 2, 4], [0.5, 0.5], [0.25, 0.25], GREEN, None, &floats(order, &grid())), order, false).unwrap(); + let (v2, b) = m.to_bytes(order); + assert!(!v2); + assert_eq!(GainTableMap::parse(&b, order, false).unwrap(), m); + let g = GainTableMap { gamma: 2.0, ..m }; + let (v2, b) = g.to_bytes(order); + assert!(v2); + assert_eq!(GainTableMap::parse(&b, order, true).unwrap(), g); + } + } +} diff --git a/crates/raw/src/jxl.rs b/crates/raw/src/jxl.rs new file mode 100644 index 00000000..9e5bfc5f --- /dev/null +++ b/crates/raw/src/jxl.rs @@ -0,0 +1,121 @@ +//! JPEG XL tiles and strips (DNG 1.7 compression 52546), decoded with `jxl-oxide` (pure Rust). +//! +//! Per the DNG 1.7 specification a JPEG XL chunk holds N-bit unsigned integers (8 ≤ N ≤ 16) or 16-bit +//! floats, in 1 or 3 planes, as either a bare codestream or an ISO-BMFF container; jxl-oxide reads both. +//! +//! jxl-oxide renders every channel as `f32`: integer samples come out as `value / (2^bits − 1)` with +//! `bits` the codestream's own declared bit depth, so scaling back by that maximum and rounding gives the +//! stored integers exactly (an `f32` holds `v / 65535` to far better than half a code value) — lossless +//! tiles decode bit-exactly, whatever bit depth the encoder declared relative to the IFD's +//! `BitsPerSample`. Float samples are returned as decoded. +//! +//! We never request an output colour encoding, so jxl-oxide renders in the codestream's *original* +//! colour encoding: for lossless (non-XYB) data that is the identity, keeping the values in the raw IFD's +//! sample space. Lossy (XYB) tiles are converted from XYB back to that signalled encoding by jxl-oxide; +//! [INFERENCE] DNG writers signal a linear encoding matching the raw data there, so the result is the raw +//! sample space up to the codec's loss — not verified against a real lossy JXL DNG. + +use crate::tiffraw::ChunkPx; +use crate::{RawError, Result}; +use jxl_oxide::frame::Encoding; +use jxl_oxide::image::BitDepth; +use jxl_oxide::{AllocTracker, JxlImage, JxlThreadPool}; + +/// Decoder memory (bytes) allowed per sample of a tile. jxl-oxide 0.12 peaks at about three 32-bit +/// buffers per sample: measured ≤ 10.3 bytes for lossless (modular) tiles and ≤ 13.5 bytes per +/// colour sample for lossy (VarDCT) ones, photon noise included. +const BYTES_PER_SAMPLE: usize = 16; +/// Fixed decoder overhead on top (headers, entropy-coding tables, group metadata): under 10 KiB +/// on a 16 × 16 tile. +const MARGIN_BYTES: usize = 64 << 10; + +fn corrupt(e: impl std::fmt::Display) -> RawError { + RawError::Corrupt(format!("JPEG XL tile: {e}")) +} + +/// Parse a chunk's headers (no pixel decoding) and check them against the chunk it must fill: +/// nominal `cw × ch`, or the part of it inside the image `vw × vh` (edge tiles), `cpp` colour planes, +/// a sample type matching the IFD. Returns the decoder and the codestream's integer maximum (0 for float). +#[allow(clippy::too_many_arguments)] +fn open(src: &[u8], cw: usize, ch: usize, vw: usize, vh: usize, cpp: usize, bits: u32, float: bool) -> Result<(JxlImage, f32)> { + if !matches!(cpp, 1 | 3) { + return Err(RawError::Unsupported(format!("JPEG XL with {cpp} samples per pixel"))); + } + if float && bits != 16 || !float && !(8..=16).contains(&bits) { + return Err(RawError::Unsupported(format!("JPEG XL with {bits}-bit {} samples", if float { "float" } else { "integer" }))); + } + // The decoder's memory is bounded by the tile itself: a codestream that needs more (extra + // channels, layers, …) fails to decode instead of allocating what its header asks for. + let pixels = cw.saturating_mul(ch); + let budget = |planes: usize| pixels.saturating_mul(planes).saturating_mul(BYTES_PER_SAMPLE); + let tracker = AllocTracker::with_limit(budget(cpp).saturating_add(MARGIN_BYTES)); + let img = JxlImage::builder() + .pool(JxlThreadPool::none()) // chunks are already decoded in parallel + .alloc_tracker(tracker.clone()) + .read(std::io::Cursor::new(src)) + .map_err(corrupt)?; + // lossy (VarDCT) decoding always works on three colour channels, even for a grey tile + if cpp == 1 && img.frame_header(0).is_some_and(|f| f.encoding == Encoding::VarDct) { + tracker.expand_limit(budget(2)); + } + let header = img.image_header(); + let (jw, jh) = (img.width() as usize, img.height() as usize); + if (jw, jh) != (cw, ch) && (jw, jh) != (vw, vh) { + return Err(RawError::Corrupt(format!("JPEG XL tile is {jw}×{jh}, expected {cw}×{ch} (or {vw}×{vh} at the image edge)"))); + } + if header.metadata.orientation != 1 { + return Err(RawError::Corrupt("JPEG XL tile with a non-identity orientation".into())); + } + let planes = if header.metadata.grayscale() { 1 } else { 3 }; + if planes != cpp { + return Err(RawError::Corrupt(format!("JPEG XL tile has {planes} colour planes, the IFD {cpp} samples per pixel"))); + } + let max = match header.metadata.bit_depth { + BitDepth::IntegerSample { bits_per_sample } if !float && (1..=16).contains(&bits_per_sample) => ((1u32 << bits_per_sample) - 1) as f32, + BitDepth::FloatSample { .. } if float => 0.0, + d => return Err(RawError::Corrupt(format!("JPEG XL tile sample type {d:?} does not match the IFD"))), + }; + Ok((img, max)) +} + +/// Check the first chunk's headers ([`open`]) without decoding pixels (header-only probing). +#[allow(clippy::too_many_arguments)] +pub(crate) fn check(src: &[u8], cw: usize, ch: usize, vw: usize, vh: usize, cpp: usize, bits: u32, float: bool) -> Result<()> { + open(src, cw, ch, vw, vh, cpp, bits, float).map(|_| ()) +} + +/// Decode one chunk into `cw × ch × cpp` samples (row-major, interleaved); a chunk coded at its +/// in-image size `vw × vh` fills the top-left of that, the rest stays 0. +#[allow(clippy::too_many_arguments)] +pub(crate) fn decode(src: &[u8], cw: usize, ch: usize, vw: usize, vh: usize, cpp: usize, bits: u32, float: bool) -> Result { + let (img, max) = open(src, cw, ch, vw, vh, cpp, bits, float)?; + if img.num_loaded_keyframes() == 0 { + return Err(corrupt("no complete frame")); + } + let render = img.render_frame(0).map_err(corrupt)?; + let fb = render.image_all_channels(); + let (jw, jh, fc) = (fb.width(), fb.height(), fb.channels()); + if jw == 0 || jh == 0 || jw > cw || jh > ch || fc < cpp { + return Err(corrupt(format!("rendered {jw}×{jh}×{fc}, expected {cw}×{ch}×{cpp}"))); + } + let n = cw.saturating_mul(ch).saturating_mul(cpp); + let rows = fb.buf().chunks_exact(jw * fc); + Ok(if float { + let mut out = vec![0f32; n]; + for (dst, src) in out.chunks_exact_mut(cw * cpp).zip(rows) { + for (d, s) in dst.chunks_exact_mut(cpp).zip(src.chunks_exact(fc)) { + d.iter_mut().zip(s).for_each(|(d, s)| *d = *s); + } + } + ChunkPx::F32(out) + } else { + let mut out = vec![0u16; n]; + for (dst, src) in out.chunks_exact_mut(cw * cpp).zip(rows) { + for (d, s) in dst.chunks_exact_mut(cpp).zip(src.chunks_exact(fc)) { + // `as` saturates (NaN → 0): lossy overshoot clips to the integer range + d.iter_mut().zip(s).for_each(|(d, s)| *d = (s * max).round() as u16); + } + } + ChunkPx::U16(out) + }) +} diff --git a/crates/raw/src/lib.rs b/crates/raw/src/lib.rs index d63e5254..6732ad92 100644 --- a/crates/raw/src/lib.rs +++ b/crates/raw/src/lib.rs @@ -2,16 +2,18 @@ //! //! - [`probe`] recognises raw containers; [`decode`] turns a file into a [`RawImage`] (sensor data + everything //! needed to render it: CFA, black/white levels, active area, default crop, orientation, DNG colour tags, -//! opcode lists, [`Metadata`]); [`embedded_preview`] returns the largest embedded JPEG. +//! opcode lists, [`Metadata`]); [`embedded_preview`] returns the largest embedded JPEG (or DNG 1.7 JPEG XL) preview; +//! [`semantic_masks`] reads a DNG's semantic masks (segmentation mattes, e.g. iPhone ProRAW's sky matte). //! - [`RawImage::normalized`] subtracts black, scales white to 1.0 and crops to the active area (applying DNG //! `OpcodeList1`/`OpcodeList2`); [`demosaic`] turns CFA data into camera-RGB [`Rgb32f`]; //! [`RawImage::develop`] does all of it plus `OpcodeList3` and the default crop; [`RawImage::develop_binned`] //! produces the same at 1/k of the size straight from the mosaic (previews, thumbnails). //! - [`color`] implements the DNG colour model (dual-illuminant interpolation, forward matrices, white balance) //! and produces camera → linear Rec.2020 D65 matrices; [`profile`] reads and applies a DNG's own profile -//! look tables and tone curve. +//! look tables and tone curve, and [`gaintable`] its gain table map (Apple ProRAW's local tone mapping). //! -//! Formats: DNG (uncompressed, lossless JPEG, lossy JPEG (Smart Previews), Deflate incl. floating point, tiled/stripped, CFA and LinearRaw), +//! Formats: DNG (uncompressed, lossless JPEG, lossy JPEG (Smart Previews), Deflate incl. floating point, JPEG XL (DNG 1.7, +//! `jxl` feature, on by default), tiled/stripped, CFA and LinearRaw), //! Canon CR2 / CR3 (lossless CRX Bayer and version 0x100/0x200 C-RAW), Nikon NEF/NRW (uncompressed, Huffman lossless / lossy compressed), Sony ARW (uncompressed, ARW2, lossless), Fujifilm RAF (uncompressed Bayer //! and X-Trans, lossless and lossy compressed), Panasonic RW2 / Leica RWL / Panasonic RAW (every raw format: compressed 4 and 6, the prefix-coded strips of 8, //! packed 2/5/7, the 16-bit words of the oldest bodies), Pentax PEF (uncompressed, Huffman), Olympus ORF (uncompressed). @@ -27,11 +29,15 @@ pub mod color; pub mod demosaic; mod dng; pub mod dngwrite; +pub mod gaintable; pub mod highlight; +#[cfg(feature = "jxl")] +mod jxl; pub mod ljpeg; pub mod opcodes; mod preview; pub mod profile; +pub mod semantic; mod tiffraw; mod unpack; mod vendor; @@ -44,6 +50,7 @@ pub use lightcraft_meta::Metadata; pub use lightcraft_raster::Rgb32f; pub use opcodes::{Opcode, OpcodeLists}; pub use preview::{PreviewColorSpace, embedded_preview, embedded_preview_color_space}; +pub use semantic::{SemanticMask, semantic_masks}; use lightcraft_color::Xy; use lightcraft_tiff::{Tiff, TiffError}; @@ -509,8 +516,13 @@ pub struct ColorData { pub as_shot_white_xy: Option, /// EV to add for a "normal" rendering (`BaselineExposure` + `BaselineExposureOffset`). pub baseline_exposure: f64, + /// `BaselineSharpness`: sharpening relative to a reference camera (`None`: the DNG default, + /// 1). Read and written back by DNG export; rendering doesn't use it yet. + #[serde(default)] + pub baseline_sharpness: Option, /// The file's own camera-profile look (`ProfileHueSatMap*`, `ProfileLookTable*`, - /// `ProfileToneCurve`), applied by [`color`]'s users at render time. + /// `ProfileToneCurve`), applied by [`color`]'s users at render time, and its + /// `ProfileGainTableMap*`, kept (a DNG export writes it back) but not rendered. #[serde(default)] pub profile: profile::ProfileLook, } diff --git a/crates/raw/src/preview.rs b/crates/raw/src/preview.rs index 6ec326ea..371b5e4a 100644 --- a/crates/raw/src/preview.rs +++ b/crates/raw/src/preview.rs @@ -1,5 +1,7 @@ //! Embedded preview extraction: the largest baseline/progressive JPEG stored in a TIFF-based raw (IFD strips -//! with JPEG compression, `JPEGInterchangeFormat` pointers in any IFD, Nikon/others' maker-note preview IFDs). +//! with JPEG compression, `JPEGInterchangeFormat` pointers in any IFD, Nikon/others' maker-note preview IFDs), +//! or a DNG 1.7 JPEG XL preview IFD stored as a single tile/strip (a standalone `.jxl` file, as the spec +//! recommends for previews). use lightcraft_tiff::image::chunk_bytes; use lightcraft_tiff::tags as t; @@ -56,6 +58,11 @@ fn is_dct_jpeg(b: &[u8]) -> bool { false } +/// Whether `b` is a JPEG XL file: a bare codestream or the ISO-BMFF container's signature box. +fn is_jxl(b: &[u8]) -> bool { + b.starts_with(&[0xff, 0x0a]) || b.starts_with(&[0, 0, 0, 0x0c, b'J', b'X', b'L', b' ', 0x0d, 0x0a, 0x87, 0x0a]) +} + fn candidates<'a>(data: &'a [u8], ifd: &Ifd, base: u64, out: &mut Vec<&'a [u8]>) { // whole JPEG files stored as an undefined-type tag value (e.g. Panasonic `JpgFromRaw` 0x002e) for e in &ifd.entries { @@ -83,9 +90,20 @@ fn candidates<'a>(data: &'a [u8], ifd: &Ifd, base: u64, out: &mut Vec<&'a [u8]>) out.push(s); } } + // a rendered (RGB or grey, never CFA/LinearRaw/mask) JPEG XL preview in one chunk + if ifd.u16(t::COMPRESSION) == Some(t::compression::JPEG_XL) + && matches!(ifd.u16(t::PHOTOMETRIC), Some(t::photometric::BLACK_IS_ZERO | t::photometric::RGB)) + && let Ok(info) = ifd.image() + && let [chunk] = info.chunks(data.len() as u64).as_slice() + && let Some(s) = chunk_bytes(data, chunk) + && is_jxl(s) + { + out.push(s); + } } -/// The largest embedded JPEG preview, if any. +/// The largest embedded preview, if any: a JPEG, or (DNG 1.7) a JPEG XL file — both decode with +/// `lightcraft_codecs::decode`. pub fn embedded_preview(bytes: &[u8]) -> Option> { if bytes.starts_with(b"FUJIFILMCCD-RAW") { let j = crate::vendor::raf::header(bytes).ok()?.jpeg?; @@ -121,7 +139,11 @@ pub fn embedded_preview(bytes: &[u8]) -> Option> { if let Some(p) = crate::vendor::orf::preview(bytes) { found.push(p); } - let best = found.into_iter().filter(|s| is_dct_jpeg(s)).max_by_key(|s| s.len()).map(|s| trim_eoi(s).to_vec()); + let best = found + .into_iter() + .filter(|s| is_dct_jpeg(s) || is_jxl(s)) + .max_by_key(|s| s.len()) + .map(|s| if is_jxl(s) { s.to_vec() } else { trim_eoi(s).to_vec() }); // a TIFF raw we can't decode whose preview no tag points to (Leaf MOS, Epson ERF) best.or_else(|| format.filter(|f| !f.is_supported()).and_then(|_| scan_for_jpeg(bytes))) } diff --git a/crates/raw/src/profile.rs b/crates/raw/src/profile.rs index 8b9761ab..a21444c4 100644 --- a/crates/raw/src/profile.rs +++ b/crates/raw/src/profile.rs @@ -1,5 +1,6 @@ //! DNG camera-profile look data carried by the file itself (Adobe DNG Specification 1.6/1.7, -//! chapter 6 and the `ProfileHueSatMap*`, `ProfileLookTable*` and `ProfileToneCurve` tags). +//! chapter 6 and the `ProfileHueSatMap*`, `ProfileLookTable*`, `ProfileToneCurve` and +//! `ProfileGainTableMap*` tags). //! //! - A hue/saturation/value table holds, for a grid of (value, hue, saturation) inputs, a hue shift //! in degrees, a saturation scale and a value scale (value-major, hue-middle, saturation-minor). @@ -10,11 +11,17 @@ //! - `ProfileHueSatMapData1/2` follow the colour matrices (interpolated by the same illuminant //! weight); `ProfileLookTableData` comes after exposure compensation and before any tone curve. //! - `ProfileToneCurve`: (input, output) pairs in linear gamma from (0, 0) to (1, 1). +//! - `ProfileGainTableMap*` ([`crate::gaintable`]) is read and kept (a DNG export writes it back) +//! but not rendered: Lightroom Classic renders Apple ProRAW without it (measured on the iPhone +//! 12 Pro corpus file against Lightroom Classic 15.6: with the map the default render is far +//! brighter than Lightroom's, mean ΔE00 19; without it Lightroom's output is a tight function of +//! the hue/saturation map, exposure compensation and tone curve alone). //! //! The data is read from the user's own file at run time; nothing here ships profile data. //! Values above 1.0 (highlight headroom) keep their headroom: the table's value scale only ever //! lowers them, so a lookup can't clip what the scene-referred pipeline still needs. +use crate::gaintable::GainTableMap; use lightcraft_color::transfer::{linear_to_srgb, srgb_to_linear}; use lightcraft_color::{D50, D65, Mat3, PROPHOTO, REC2020, bradford}; use serde::{Deserialize, Serialize}; @@ -207,11 +214,14 @@ pub struct ProfileLook { pub look_table: Option, /// `ProfileToneCurve`. pub tone_curve: Option, + /// `ProfileGainTableMap2`, else `ProfileGainTableMap`. + #[serde(default)] + pub gain_table_map: Option, } impl ProfileLook { pub fn is_empty(&self) -> bool { - self.hue_sat_map.iter().all(Option::is_none) && self.look_table.is_none() && self.tone_curve.is_none() + self.hue_sat_map.iter().all(Option::is_none) && self.look_table.is_none() && self.tone_curve.is_none() && self.gain_table_map.is_none() } /// The hue/saturation map for calibration-1 weight `g` (one table: used for any white). @@ -360,7 +370,7 @@ mod tests { let m = HsvTable::blend(&a, &b, 0.25).unwrap(); assert_eq!(m.data[1], [7.5, 1.75, 1.0]); assert!(HsvTable::blend(&a, &identity(3, 2, 1), 0.5).is_none()); - let look = ProfileLook { hue_sat_map: [Some(a.clone()), None], look_table: Some(a), tone_curve: None }; + let look = ProfileLook { hue_sat_map: [Some(a.clone()), None], look_table: Some(a), tone_curve: None, gain_table_map: None }; let t = ProfileTables::new(&look, 0.5).unwrap(); let c = [0.3, 0.2, 0.1]; assert!(close(t.apply(c, 2.0), [0.6, 0.4, 0.2], 1e-4), "{:?}", t.apply(c, 2.0)); diff --git a/crates/raw/src/semantic.rs b/crates/raw/src/semantic.rs new file mode 100644 index 00000000..e9af87c0 --- /dev/null +++ b/crates/raw/src/semantic.rs @@ -0,0 +1,136 @@ +//! DNG semantic masks (DNG 1.6+, chapter "Semantic Masks"): single-channel segmentation mattes +//! stored in their own IFDs (NewSubFileType 10004.H, PhotometricInterpretation 52527) with a +//! `SemanticName`, an optional `SemanticInstanceID` and an optional `MaskSubArea`. iPhone ProRAW +//! stores e.g. a sky matte this way. +//! +//! [`semantic_masks`] reads them separately from [`crate::decode`] (only callers that use them +//! pay for them); malformed or unsupported mask IFDs are skipped. + +use crate::Rect; +use crate::tiffraw::{Packing, read_image}; +use crate::{RawData, Result}; +use lightcraft_raster::Image; +use lightcraft_tiff::Tiff; +use lightcraft_tiff::tags::{self as t, photometric}; + +/// At most this many mask IFDs are read from one file. +pub const MAX_MASKS: usize = 16; +/// Masks with more samples than this are skipped (a mask is a low-resolution matte; this is +/// 8192 × 4096). +pub const MAX_MASK_SAMPLES: usize = 1 << 25; + +/// Where a cropped mask sits in the uncropped mask (DNG `MaskSubArea`). The uncropped mask +/// (`full_width × full_height`) spans the main image's active area. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct MaskSubArea { + pub top: u32, + pub left: u32, + pub full_width: u32, + pub full_height: u32, +} + +/// One semantic mask: `width × height` values, 0 = not part of the attribute, 65535 = fully part. +#[derive(Clone, Debug, PartialEq)] +pub struct SemanticMask { + /// `SemanticName`, e.g. `urn:com:apple:photo:2020:aux:semanticskymatte`. + pub name: String, + pub instance_id: Option, + pub width: usize, + pub height: usize, + /// `None`: the mask covers the whole active area. + pub sub_area: Option, + pub data: Vec, +} + +/// The semantic masks of a DNG, in file order (none for other files). Only the first +/// [`MAX_MASKS`] mask IFDs are looked at, malformed ones included, so a file full of mask IFDs +/// costs at most that many decodes. +pub fn semantic_masks(bytes: &[u8]) -> Vec { + let Ok(tiff) = Tiff::parse(bytes) else { return Vec::new() }; + tiff.all_ifds() + .into_iter() + .filter(|i| i.u16(t::PHOTOMETRIC) == Some(photometric::MASK) || i.u32(t::NEW_SUBFILE_TYPE) == Some(t::SUBFILE_SEMANTIC_MASK)) + .take(MAX_MASKS) + .filter_map(|ifd| read_mask(bytes, &tiff, ifd).ok().flatten()) + .collect() +} + +fn read_mask(bytes: &[u8], tiff: &Tiff, ifd: &lightcraft_tiff::Ifd) -> Result> { + let Some(name) = ifd.string(t::SEMANTIC_NAME).filter(|n| !n.is_empty()) else { return Ok(None) }; + let info = ifd.image()?; + let (w, h) = (info.width as usize, info.height as usize); + if info.samples_per_pixel != 1 || w.checked_mul(h).is_none_or(|n| n > MAX_MASK_SAMPLES) { + return Ok(None); + } + let bits = info.bits() as u32; + let data: Vec = match read_image(bytes, &info, tiff.order, Packing::Msb)? { + RawData::U16(v) if (1..=16).contains(&bits) => { + let max = ((1u32 << bits) - 1) as f32; + if bits == 16 { v } else { v.into_iter().map(|s| (s as f32 / max * 65535.0).round().min(65535.0) as u16).collect() } + } + RawData::F32(v) => v.into_iter().map(|s| if s.is_finite() { (s.clamp(0.0, 1.0) * 65535.0).round() as u16 } else { 0 }).collect(), + RawData::U16(_) => return Ok(None), + }; + if data.len() != w * h { + return Ok(None); + } + // the tag is ignored when the stored mask does not fit inside the uncropped one + let sub_area = match ifd.u64s(t::MASK_SUB_AREA).as_deref() { + Some(&[top, left, fw, fh]) => { + let fits = fw > 0 && fh > 0 && left.saturating_add(w as u64) <= fw && top.saturating_add(h as u64) <= fh && fw.max(fh) <= u32::MAX as u64; + fits.then_some(MaskSubArea { top: top as u32, left: left as u32, full_width: fw as u32, full_height: fh as u32 }) + } + _ => None, + }; + Ok(Some(SemanticMask { name, instance_id: ifd.string(t::SEMANTIC_INSTANCE_ID).filter(|s| !s.is_empty()), width: w, height: h, sub_area, data })) +} + +impl SemanticMask { + /// The mask value (0..1) at a continuous position of the uncropped mask (pixel centres at + /// +0.5); 0 outside the stored (cropped) mask. + pub fn value_at(&self, x: f64, y: f64) -> f32 { + let (left, top) = self.sub_area.map_or((0.0, 0.0), |a| (a.left as f64, a.top as f64)); + let (x, y) = (x - left, y - top); + if self.width == 0 || self.height == 0 || !(x >= 0.0 && y >= 0.0 && x <= self.width as f64 && y <= self.height as f64) { + return 0.0; + } + let (fx, fy) = ((x - 0.5).max(0.0), (y - 0.5).max(0.0)); + let (x0, y0) = (fx.floor() as usize, fy.floor() as usize); + let (tx, ty) = ((fx - x0 as f64) as f32, (fy - y0 as f64) as f32); + let at = |x: usize, y: usize| { + let (x, y) = (x.min(self.width - 1), y.min(self.height - 1)); + self.data.get(y * self.width + x).copied().unwrap_or(0) as f32 / 65535.0 + }; + let top = at(x0, y0) + (at(x0 + 1, y0) - at(x0, y0)) * tx; + let bottom = at(x0, y0 + 1) + (at(x0 + 1, y0 + 1) - at(x0, y0 + 1)) * tx; + top + (bottom - top) * ty + } + + /// The mask over the developed image ([`crate::RawImage::develop`]: the default `crop` of the + /// `active_area`, before orientation), at the mask's own resolution, as 0..255 alpha. `None` + /// for an empty crop. + pub fn developed(&self, active_area: Rect, crop: Rect) -> Option> { + let c = crop.clipped(active_area.width, active_area.height); + if c.width == 0 || c.height == 0 || self.width == 0 || self.height == 0 { + return None; + } + let (fw, fh) = self.sub_area.map_or((self.width, self.height), |a| (a.full_width as usize, a.full_height as usize)); + // uncropped-mask px per active-area px + let (sx, sy) = (fw as f64 / active_area.width as f64, fh as f64 / active_area.height as f64); + let ow = ((c.width as f64 * sx).round() as usize).clamp(1, self.width.max(fw)); + let oh = ((c.height as f64 * sy).round() as usize).clamp(1, self.height.max(fh)); + if ow.checked_mul(oh).is_none_or(|n| n > MAX_MASK_SAMPLES) { + return None; + } + let (kx, ky) = (c.width as f64 / ow as f64, c.height as f64 / oh as f64); + let mut out = Image::new(ow, oh); + lightcraft_raster::par_rows(&mut out.data, ow, |y, row| { + let ay = c.y as f64 + (y as f64 + 0.5) * ky; + for (x, v) in row.iter_mut().enumerate() { + let ax = c.x as f64 + (x as f64 + 0.5) * kx; + *v = (self.value_at(ax * sx, ay * sy) * 255.0).round() as u8; + } + }); + Some(out) + } +} diff --git a/crates/raw/src/tiffraw.rs b/crates/raw/src/tiffraw.rs index 200dd391..c28b8b64 100644 --- a/crates/raw/src/tiffraw.rs +++ b/crates/raw/src/tiffraw.rs @@ -1,4 +1,4 @@ -//! Reading the pixel data of a TIFF IFD (strips or tiles; uncompressed, lossless JPEG, Deflate), in parallel. +//! Reading the pixel data of a TIFF IFD (strips or tiles; uncompressed, lossless/lossy JPEG, Deflate, JPEG XL), in parallel. use crate::unpack::*; use crate::{MAX_SAMPLES, RawData, RawError, Result, ljpeg}; @@ -16,9 +16,12 @@ pub enum Packing { Lsb, /// Each sample stored in 16 bits in file byte order, regardless of `bits`. Word16, + /// MSB-first bit stream inside 32-bit words stored in file byte order ([`unpack_words32_msb`]); rows start on a + /// word boundary. + Words32Msb, } -enum ChunkPx { +pub(crate) enum ChunkPx { U16(Vec), F32(Vec), } @@ -61,11 +64,24 @@ pub(crate) fn read_image_in(mode: crate::Mode, data: &[u8], info: &ImageInfo, or crate::Mode::Header => { check_image(data, info)?; // a lossless-JPEG layout the decoder rejects (subsampled components, e.g. Sony's - // lossless M/S sizes) must fail here too, as the full decode will - if info.compression == 7 - && let Some(src) = info.chunks(data.len() as u64).first().and_then(|c| chunk_bytes(data, c)) + // lossless M/S sizes) must fail here too, as the full decode will; likewise a JPEG XL + // chunk whose headers don't match the IFD + let first = info.chunks(data.len() as u64).first().copied(); + if let Some(c) = first + && let Some(src) = chunk_bytes(data, &c) { - ljpeg::frame_info(src)?; + match info.compression { + 7 => { + ljpeg::frame_info(src)?; + } + comp::JPEG_XL => { + let planar = info.planar == 2 && info.samples_per_pixel > 1; + let cpp = if planar { 1 } else { info.samples_per_pixel as usize }; + let (cw, ch, vw, vh) = chunk_dims(info, &c); + jxl_chunk_check(src, cw, ch, vw, vh, cpp, info.bits() as u32, info.sample_format == 3)?; + } + _ => {} + } } Ok(if info.sample_format == 3 { RawData::F32(Vec::new()) } else { RawData::U16(Vec::new()) }) } @@ -138,7 +154,7 @@ pub fn read_image(data: &[u8], info: &ImageInfo, order: ByteOrder, packing: Pack #[allow(clippy::too_many_arguments)] fn decode_chunk(data: &[u8], info: &ImageInfo, c: &Chunk, order: ByteOrder, packing: Packing, cpp: usize, bits: u32, float: bool) -> Result { let src = chunk_bytes(data, c).ok_or_else(|| RawError::Corrupt("chunk offset past end of file".into()))?; - let (cw, ch) = (c.width as usize, c.height as usize); + let (cw, ch, vw, vh) = chunk_dims(info, c); let n = cw.checked_mul(ch).and_then(|v| v.checked_mul(cpp)).ok_or(RawError::Limit("chunk too large"))?; if n > MAX_SAMPLES { return Err(RawError::Limit("chunk too large")); @@ -176,10 +192,34 @@ fn decode_chunk(data: &[u8], info: &ImageInfo, c: &Chunk, order: ByteOrder, pack let raw = inflate(src, row_bytes * ch)?; unpack_chunk(&raw, order, packing, cw, ch, cpp, bits, float, info.predictor) } + comp::JPEG_XL => jxl_chunk(src, cw, ch, vw, vh, cpp, bits, float), other => Err(RawError::Unsupported(format!("TIFF compression {other}"))), } } +/// A chunk's nominal size and the part of it inside the image (smaller for edge tiles). +fn chunk_dims(info: &ImageInfo, c: &Chunk) -> (usize, usize, usize, usize) { + let (cw, ch) = (c.width as usize, c.height as usize); + let vw = cw.min(info.width.saturating_sub(c.x) as usize); + let vh = ch.min(info.height.saturating_sub(c.y) as usize); + (cw, ch, vw, vh) +} + +#[cfg(feature = "jxl")] +use crate::jxl::{check as jxl_chunk_check, decode as jxl_chunk}; + +#[cfg(not(feature = "jxl"))] +#[allow(clippy::too_many_arguments)] +fn jxl_chunk(_: &[u8], _: usize, _: usize, _: usize, _: usize, _: usize, _: u32, _: bool) -> Result { + Err(RawError::Unsupported("JPEG XL DNG (built without the `jxl` feature)".into())) +} + +#[cfg(not(feature = "jxl"))] +#[allow(clippy::too_many_arguments)] +fn jxl_chunk_check(src: &[u8], cw: usize, ch: usize, vw: usize, vh: usize, cpp: usize, bits: u32, float: bool) -> Result<()> { + jxl_chunk(src, cw, ch, vw, vh, cpp, bits, float).map(|_| ()) +} + /// Decode one baseline JPEG chunk to 8-bit samples: (samples, width, height, channels). fn lossy_jpeg(src: &[u8], cpp: usize) -> Result<(Vec, usize, usize, usize)> { use zune_core::bytestream::ZCursor; @@ -263,6 +303,7 @@ fn unpack_chunk( (16, _) | (_, Packing::Word16) => (row_n * 2, 1), (_, Packing::Msb) => ((row_n * bits as usize).div_ceil(8), 2), (_, Packing::Lsb) => ((row_n * bits as usize).div_ceil(8), 3), + (_, Packing::Words32Msb) => ((row_n * bits as usize).div_ceil(32) * 4, 4), }; for y in 0..ch { let s = src.get(y * row_bytes..).unwrap_or(&[]); @@ -272,6 +313,7 @@ fn unpack_chunk( 0 => row.iter_mut().zip(s).for_each(|(o, &b)| *o = b as u16), 1 => read_u16s(s, order, row), 2 => unpack_msb(s, bits, row), + 4 => unpack_words32_msb(s, order, bits, row), _ => unpack_lsb(s, bits, row), } if factor > 0 { diff --git a/crates/raw/src/unpack.rs b/crates/raw/src/unpack.rs index 15451548..77c69c00 100644 --- a/crates/raw/src/unpack.rs +++ b/crates/raw/src/unpack.rs @@ -40,6 +40,19 @@ pub fn unpack_lsb(src: &[u8], bits: u32, out: &mut [u16]) { } } +/// Unpack MSB-first packed samples of `bits` from a stream of 32-bit words stored in `order` (a word is read as a +/// number first, then its bits are taken from the top): the layout of some cameras' uncompressed 12- and 14-bit +/// data, which is not the byte-wise MSB stream of [`unpack_msb`] unless `order` is big-endian. A short last word +/// is padded with zeros. +pub fn unpack_words32_msb(src: &[u8], order: ByteOrder, bits: u32, out: &mut [u16]) { + let words = src.chunks(4).flat_map(|c| { + let mut w = [0u8; 4]; + w[..c.len()].copy_from_slice(c); + order.u32(w).to_be_bytes() + }); + unpack_msb(&words.collect::>(), bits, out); +} + pub fn read_u16s(src: &[u8], order: ByteOrder, out: &mut [u16]) { for (o, c) in out.iter_mut().zip(src.as_chunks::<2>().0) { *o = order.u16([c[0], c[1]]); @@ -128,6 +141,30 @@ pub fn undo_float_predictor(row: &mut [u8], n: usize, bytes_per: usize, stride: mod tests { use super::*; + #[test] + fn words32_are_read_as_numbers_then_msb_first() { + // eight 12-bit samples occupy 96 bits = three 32-bit words, filled from the top + let samples: [u16; 8] = [0x123, 0x456, 0x789, 0xabc, 0xdef, 0x012, 0x345, 0x678]; + let stream = samples.iter().fold(0u128, |acc, &v| (acc << 12) | v as u128); + let words = [(stream >> 64) as u32, (stream >> 32) as u32, stream as u32]; + let le: Vec = words.iter().flat_map(|w| w.to_le_bytes()).collect(); + let be: Vec = words.iter().flat_map(|w| w.to_be_bytes()).collect(); + let mut o = [0u16; 8]; + unpack_words32_msb(&le, ByteOrder::Little, 12, &mut o); + assert_eq!(o, samples); + o = [0; 8]; + unpack_words32_msb(&be, ByteOrder::Big, 12, &mut o); + assert_eq!(o, samples); + // read as a plain byte stream the little-endian words give different samples + o = [0; 8]; + unpack_msb(&le, 12, &mut o); + assert_ne!(o, samples); + // a short last word is zero padded + let mut o = [9u16; 2]; + unpack_words32_msb(&[0x34, 0x12], ByteOrder::Little, 12, &mut o); + assert_eq!(o, [0x000, 0x012]); // the word 0x00001234, top bits first + } + #[test] fn packing() { // 12-bit MSB: 0xABC, 0xDEF -> AB CD EF diff --git a/crates/raw/src/vendor/nef.rs b/crates/raw/src/vendor/nef.rs index 8d1898e0..5ca66af4 100644 --- a/crates/raw/src/vendor/nef.rs +++ b/crates/raw/src/vendor/nef.rs @@ -18,6 +18,25 @@ const WB_RB_LEVELS: u16 = 0x000c; const BLACK_LEVEL: u16 = 0x003d; const LINEARIZATION_TABLE: u16 = 0x0096; const CROP_AREA: u16 = 0x0045; +/// Maker note `0x0014`, the "NRW" data block of Coolpix raws: `"NRW "` + a 4-character version, then fields. In +/// version `0104` the `u32` at byte `0x20` is the black level in sample units (200 in the four 12-bit files here +/// that have 200 as their darkest samples, 0 in the one whose samples reach 0). +const NRW_DATA: u16 = 0x0014; +const NRW_BLACK_AT: usize = 0x20; + +/// The NRW data block, when the maker note has one. +fn nrw_data(mn: Option<&makernote::MakerNote>) -> Option<(&[u8], ByteOrder)> { + let m = mn?; + m.ifd.bytes(NRW_DATA).filter(|b| b.starts_with(b"NRW ")).map(|b| (b, m.order)) +} + +/// Black level from the NRW data block, bounded to a quarter of the sample range so an unknown version's field +/// can't turn into a large offset. +fn nrw_black(mn: Option<&makernote::MakerNote>, bits: u32) -> Option { + let (b, order) = nrw_data(mn)?; + let v = order.u32(b.get(NRW_BLACK_AT..NRW_BLACK_AT + 4)?.try_into().ok()?); + (u64::from(v) < (1u64 << bits.min(16)) / 4).then_some(v as f32) +} /// Nikon CropArea is `[left, top, width, height]` in sensor pixels (maker-note tag documentation). /// Keep the active area/CFA origin unchanged: the default crop is applied after demosaicing. @@ -65,8 +84,11 @@ pub(crate) fn decode(bytes: &[u8]) -> Result { let make = ifd0.string(t::MAKE).unwrap_or_default(); let mn = tiff.exif().and_then(|e| e.get(t::MAKER_NOTE)).and_then(|e| makernote::parse_makernote(bytes, e.offset, e.count() as u64, tiff.order, &make)); - let data = - if info.compression == t::compression::NIKON { compressed(bytes, &info, mn.as_ref())? } else { uncompressed(bytes, &info, tiff.order)? }; + let data = if info.compression == t::compression::NIKON { + compressed(bytes, &info, mn.as_ref())? + } else { + uncompressed(bytes, &info, tiff.order, mn.as_ref())? + }; let RawData::U16(ref samples) = data else { return Err(RawError::Unsupported("float NEF".into())) }; // Maker note 0x003d is in 14-bit units whatever the sample depth: 12-bit files from the D750, D780, D850, @@ -80,7 +102,7 @@ pub(crate) fn decode(bytes: &[u8]) -> Result { values: [a, b, c, d].iter().map(|v| (**v * black_scale) as f32).collect(), ..Default::default() }, - _ => BlackLevel::uniform(0.0), + _ => BlackLevel::uniform(nrw_black(mn.as_ref(), bits).unwrap_or(0.0)), }; let wb = mn .as_ref() @@ -120,8 +142,10 @@ pub(crate) fn decode(bytes: &[u8]) -> Result { Ok(img) } -/// Uncompressed strips: 16-bit words or 12-bit MSB-packed, told apart by the strip size. -fn uncompressed(bytes: &[u8], info: &ImageInfo, order: ByteOrder) -> Result { +/// Uncompressed strips: 16-bit words or 12-bit packed, told apart by the strip size. Packed 12-bit data is an MSB-first +/// byte stream, except in files with an NRW data block (Coolpix), where it is MSB-first inside little-endian 32-bit +/// words (the rows are whole words). +fn uncompressed(bytes: &[u8], info: &ImageInfo, order: ByteOrder, mn: Option<&makernote::MakerNote>) -> Result { let (w, h) = (info.width as usize, info.height as usize); let bits = info.bits() as u32; let row_samples = w * info.samples_per_pixel as usize; @@ -131,7 +155,7 @@ fn uncompressed(bytes: &[u8], info: &ImageInfo, order: ByteOrder) -> Result= row_samples * 2 { Packing::Word16 } else if bits == 12 && bytes_per_row * 8 >= row_samples * 12 && bytes_per_row * 8 < row_samples * 13 { - Packing::Msb + if nrw_data(mn).is_some() && bytes_per_row.is_multiple_of(4) { Packing::Words32Msb } else { Packing::Msb } } else if info.compression == 1 && bits != 8 && bits != 16 { return Err(RawError::Unsupported(format!("NEF uncompressed packing ({bytes_per_row} bytes per {row_samples}-sample row)"))); } else { @@ -171,6 +195,12 @@ mod tests { /// [`nef`] with an optional maker note (`Nikon\0` v2 header + embedded big-endian TIFF holding `note`). fn nef_with_note(compression: u16, bits: u16, strips: Vec>, w: u32, h: u32, rps: u32, note: Option) -> Vec { + nef_in(ByteOrder::Big, compression, bits, strips, w, h, rps, note) + } + + /// [`nef_with_note`] in either byte order (file and maker note). + #[allow(clippy::too_many_arguments)] + fn nef_in(order: ByteOrder, compression: u16, bits: u16, strips: Vec>, w: u32, h: u32, rps: u32, note: Option) -> Vec { let mut raw = IfdBuilder::new(); raw.set(t::NEW_SUBFILE_TYPE, Value::Long(vec![0])); raw.set(t::IMAGE_WIDTH, Value::Long(vec![w])); @@ -186,13 +216,13 @@ mod tests { ifd0.set(t::MODEL, Value::Ascii("NIKON TEST".into())); if let Some(mn) = note { let mut bytes = b"Nikon\0\x02\x10\0\0".to_vec(); - bytes.extend(TiffWriter::new(ByteOrder::Big, false).write(&[mn]).unwrap()); + bytes.extend(TiffWriter::new(order, false).write(&[mn]).unwrap()); let mut exif = IfdBuilder::new(); exif.set(t::MAKER_NOTE, Value::Undefined(bytes)); ifd0.set_child(t::EXIF_IFD, exif); } ifd0.add_sub_ifd(raw); - TiffWriter::new(ByteOrder::Big, false).write(&[ifd0]).unwrap() + TiffWriter::new(order, false).write(&[ifd0]).unwrap() } #[test] @@ -266,6 +296,61 @@ mod tests { assert_eq!(crate::decode(&bytes).unwrap().data, RawData::U16(px)); } + /// An NRW data block (`"NRW 0104"`, fields from byte 8) with `black` at byte `0x20`, as a maker note. + fn nrw_note(black: u32) -> IfdBuilder { + let mut block = b"NRW 0104".to_vec(); + block.resize(0x20, 0); + block.extend_from_slice(&black.to_le_bytes()); + block.resize(0x40, 0); + let mut mn = IfdBuilder::new(); + mn.set(NRW_DATA, Value::Undefined(block)); + mn + } + + /// 12-bit samples as an MSB-first bit stream inside 32-bit words in `order` (rows are whole words). + fn words32(px: &[u16], w: usize, order: ByteOrder) -> Vec { + let mut out = Vec::new(); + for row in px.chunks(w) { + let stream = row.iter().fold(Vec::new(), |mut bits: Vec, v| { + bits.extend((0..12).rev().map(|b| v >> b & 1 == 1)); + bits + }); + for word in stream.chunks(32) { + let v = word.iter().enumerate().fold(0u32, |a, (i, b)| a | (*b as u32) << (31 - i)); + out.extend_from_slice(&match order { + ByteOrder::Little => v.to_le_bytes(), + ByteOrder::Big => v.to_be_bytes(), + }); + } + } + out + } + + #[test] + fn nrw_block_selects_32_bit_word_packing_and_supplies_the_black_level() { + let (w, h) = (16usize, 6usize); + let px: Vec = (0..w * h).map(|i| (200 + i * 131 % 3800) as u16).collect(); + let packed = words32(&px, w, ByteOrder::Little); + assert_eq!(packed.len(), w * h * 3 / 2); + let decode = |note: Option| { + crate::decode(&nef_in(ByteOrder::Little, 1, 12, vec![packed.clone()], w as u32, h as u32, h as u32, note)).unwrap() + }; + let r = decode(Some(nrw_note(200))); + assert_eq!(r.data, RawData::U16(px.clone())); + assert_eq!(r.black.values, vec![200.0]); + assert_eq!(decode(Some(nrw_note(0))).black.values, vec![0.0]); + // an implausible field is not a black level + assert_eq!(decode(Some(nrw_note(60000))).black.values, vec![0.0]); + // without the block the same bytes are an MSB-first byte stream + let plain = decode(None); + assert_ne!(plain.data, RawData::U16(px)); + assert_eq!(plain.black.values, vec![0.0]); + // the maker-note black level (14-bit units) wins when present + let mut both = nrw_note(200); + both.set(BLACK_LEVEL, Value::Short(vec![400, 404, 408, 412])); + assert_eq!(decode(Some(both)).black.values, vec![100.0, 101.0, 102.0, 103.0]); + } + #[test] fn compressed_without_table_is_unsupported() { let bytes = nef(34713, 14, vec![vec![0; 64]], 8, 8, 8); diff --git a/crates/raw/src/vendor/raf.rs b/crates/raw/src/vendor/raf.rs index 9d950121..61205a3f 100644 --- a/crates/raw/src/vendor/raf.rs +++ b/crates/raw/src/vendor/raf.rs @@ -12,6 +12,8 @@ //! IFD; offsets in it are relative to the raw block. //! - Sample packing, found by testing candidate bit orders for the smoothest image: 16-bit little-endian words; //! 12-bit LSB-first bit stream; 14-bit (and other depths) stored as little-endian 32-bit words read MSB-first. +//! Raw-IFD flag `0xf006` = 1 selects the 32-bit-word packing at any depth: 12-bit files with the flag set +//! (FinePix S1, X-A10, SL1000) are word-packed like the 14-bit ones, the 12-bit files without it are LSB-first. //! `0xf005`, when non-zero, is the row stride in 32-bit words. //! - CFA: the X-Trans layout record (36 bytes, 0 = R, 1 = G, 2 = B) is stored in reverse order and, reversed, //! is anchored at raw pixel (0, 0) — verified on X-Trans I and III samples by minimising the difference @@ -20,7 +22,7 @@ //! - Striped predictive compression is decoded by `rafc`; see `docs/raf-compression.md`. use super::white_from_data; -use crate::unpack::{unpack_lsb, unpack_msb}; +use crate::unpack::{unpack_lsb, unpack_words32_msb}; use crate::{BlackLevel, Cfa, ColorData, Mode, OpcodeLists, RawData, RawError, RawFormat, RawImage, Rect, Result}; use lightcraft_geom::Orientation; use lightcraft_tiff::{ByteOrder, Ifd, Tiff}; @@ -29,6 +31,9 @@ use rayon::prelude::*; const CROP_TOP_LEFT: u16 = 0x0110; const CROPPED_SIZE: u16 = 0x0111; const XTRANS_LAYOUT: u16 = 0x0131; +/// Raw-IFD flag `0xf006`: 1 in every 12- and 14-bit file here whose samples are packed into 32-bit words (the 12-bit +/// ones are otherwise an LSB-first stream), 0 where they are an LSB-first stream or 16-bit words. +const PACKED_WORDS: u16 = 0xf006; /// The fixed-position header pointers. pub(crate) struct Header<'a> { @@ -95,22 +100,14 @@ fn cfa(h: &Header) -> Cfa { } } -/// Unpack one row of `bits`-bit samples in Fujifilm's packing. -pub(crate) fn unpack_row(src: &[u8], bits: u32, out: &mut [u16]) { +/// Unpack one row of `bits`-bit samples in Fujifilm's packing. `words32`: the raw IFD's `0xf006` flag is set, the +/// samples are MSB-first in little-endian 32-bit words whatever the depth; otherwise 16-bit words, a 12-bit +/// LSB-first bit stream, or (other depths) the same 32-bit words. +pub(crate) fn unpack_row(src: &[u8], bits: u32, words32: bool, out: &mut [u16]) { match bits { 16 => out.iter_mut().zip(src.as_chunks::<2>().0).for_each(|(o, c)| *o = u16::from_le_bytes([c[0], c[1]])), - 12 => unpack_lsb(src, 12, out), - _ => { - let swapped: Vec = src - .chunks(4) - .flat_map(|c| { - let mut w = [0u8; 4]; - w[..c.len()].copy_from_slice(c); - [w[3], w[2], w[1], w[0]] - }) - .collect(); - unpack_msb(&swapped, bits, out) - } + 12 if !words32 => unpack_lsb(src, 12, out), + _ => unpack_words32_msb(src, ByteOrder::Little, bits, out), } } @@ -147,6 +144,7 @@ pub(crate) fn decode(bytes: &[u8], mode: Mode) -> Result { } data = super::rafc::decode(src, w, hgt, bits, cfa.width == 6, mode)?; } else { + let words32 = ifd.u64(PACKED_WORDS).is_some_and(|v| v != 0); // Samples may use 16-bit words (e.g. 14-bit X-T20); otherwise they are packed at `bits`. let store = if len as u64 >= (n as u64) * 2 { 16 } else { bits }; let packed_row = w.checked_mul(store as usize).map(|n| n.div_ceil(8)).ok_or(RawError::Limit("RAF row too large"))?; @@ -163,7 +161,7 @@ pub(crate) fn decode(bytes: &[u8], mode: Mode) -> Result { data.par_chunks_mut(w).enumerate().for_each(|(y, row)| { // `required <= src.len()` proves these row offsets fit, including the final row. let s = src.get(y * stride..y * stride + packed_row).unwrap_or(&[]); - unpack_row(s, store, row); + unpack_row(s, store, words32, row); }); } } @@ -220,19 +218,25 @@ pub(super) mod tests { /// Build a minimal RAF: header, records, raw block (LE TIFF whose IFD0 points to a 0xf000 raw IFD). pub(crate) fn raf(w: u32, h: u32, bits: u32, strip: Vec, layout: Option<[u8; 36]>, jpeg: &[u8]) -> Vec { + raf_flagged(w, h, bits, strip, layout, jpeg, 0) + } + + /// [`raf`] with the raw IFD's `0xf006` flag set to `words32`. + pub(crate) fn raf_flagged(w: u32, h: u32, bits: u32, strip: Vec, layout: Option<[u8; 36]>, jpeg: &[u8], words32: u32) -> Vec { let mut block = b"II*\0\x08\0\0\0".to_vec(); let sub_off: u32 = 8 + 2 + 12 + 4; block.extend_from_slice(&1u16.to_le_bytes()); block.extend_from_slice(&[0x00, 0xf0, 13, 0, 1, 0, 0, 0]); block.extend_from_slice(&sub_off.to_le_bytes()); block.extend_from_slice(&0u32.to_le_bytes()); - let n = 7u32; + let n = 8u32; let wb_off = sub_off + 2 + 12 * n + 4; let strip_off = wb_off + 12; - let entries: [(u16, u32, u32); 7] = [ + let entries: [(u16, u32, u32); 8] = [ (0xf001, 1, w), (0xf002, 1, h), (0xf003, 1, bits), + (0xf006, 1, words32), (0xf007, 1, strip_off), (0xf008, 1, strip.len() as u32), (0xf00a, 1, 64), @@ -345,14 +349,39 @@ pub(super) mod tests { assert!(matches!(crate::decode(&bytes), Err(RawError::Unsupported(_)))); } + /// Raw IFD flag `0xf006` selects the 32-bit-word packing at 12 bits; without it 12 bits are an LSB-first stream. + #[test] + fn flag_selects_word_packed_12_bit() { + let (w, h) = (16u32, 6u32); + let px: Vec = (0..w * h).map(|i| ((i * 977) % 4096) as u16).collect(); + // the same 12-bit samples as little-endian 32-bit words (what the 14-bit branch of `pack` writes) + let mut stream: Vec = Vec::new(); + let (mut acc, mut n) = (0u64, 0u32); + for &v in &px { + acc = (acc << 12) | v as u64; + n += 12; + while n >= 8 { + stream.push((acc >> (n - 8)) as u8); + n -= 8; + } + } + let le: Vec = stream.chunks(4).flat_map(|c| [c[3], c[2], c[1], c[0]]).collect(); + let decode = |bytes: Vec| crate::decode(&bytes).unwrap().data; + assert_eq!(decode(raf_flagged(w, h, 12, le.clone(), None, b"", 1)), RawData::U16(px.clone())); + // flag clear: the same bytes are read as the LSB-first stream, which gives other samples + assert_ne!(decode(raf_flagged(w, h, 12, le, None, b"", 0)), RawData::U16(px.clone())); + // flag clear: the LSB-first stream still decodes + assert_eq!(decode(raf_flagged(w, h, 12, pack(&px, 12), None, b"", 0)), RawData::U16(px)); + } + #[test] fn overflowing_row_stride_is_an_error() { let mut bytes = raf(16, 6, 16, vec![0; 192], None, &[]); let off = be32(&bytes, 100).unwrap(); // Replace WB with a LONG8 stride, pointing to the existing out-of-line WB storage. - bytes[off + 100..off + 104].copy_from_slice(&[0x05, 0xf0, 16, 0]); - bytes[off + 104..off + 108].copy_from_slice(&1u32.to_le_bytes()); - bytes[off + 116..off + 124].copy_from_slice(&u64::MAX.to_le_bytes()); + bytes[off + 112..off + 116].copy_from_slice(&[0x05, 0xf0, 16, 0]); + bytes[off + 116..off + 120].copy_from_slice(&1u32.to_le_bytes()); + bytes[off + 128..off + 136].copy_from_slice(&u64::MAX.to_le_bytes()); assert!(matches!(crate::decode(&bytes), Err(RawError::Limit(_)))); } } diff --git a/crates/raw/tests/corpus.rs b/crates/raw/tests/corpus.rs index b3c0ade9..53193f1f 100644 --- a/crates/raw/tests/corpus.rs +++ b/crates/raw/tests/corpus.rs @@ -43,7 +43,9 @@ fn corpus_raw_decodes() { // DNG previews are optional (and some carry only an uncompressed RGB thumbnail); vendor raws embed a JPEG, // except the Panasonic `.RAW` files of 2005–2007 and Canon's explicit HEVC preview tracks. if let Some(p) = &preview { - assert!(p.starts_with(&[0xff, 0xd8]) && p.ends_with(&[0xff, 0xd9]), "{name}: preview is not a JPEG"); + let jpeg = p.starts_with(&[0xff, 0xd8]) && p.ends_with(&[0xff, 0xd9]); + let jxl = p.starts_with(&[0xff, 0x0a]) || p.starts_with(b"\0\0\0\x0cJXL "); + assert!(jpeg || jxl, "{name}: preview is neither a JPEG nor a JPEG XL file"); } else { let hevc_preview = fmt == RawFormat::Cr3 && lightcraft_meta::cr3::parse_cr3(&bytes) @@ -250,9 +252,11 @@ fn corpus_nef_12_bit_black_level_matches_14_bit() { } /// Issue #138: DNGs converted by Adobe software carry their camera profile's hue/saturation map and -/// look table; we read them (and render with them). Camera-written DNGs here carry none. +/// look table; we read them (and render with them). Apple ProRAW carries a tone curve and a gain +/// table map (its local tone mapping, read and kept but not rendered, as in Lightroom Classic). +/// Other camera-written DNGs here carry none. #[test] -fn corpus_adobe_dngs_carry_profile_looks() { +fn corpus_dngs_carry_profile_looks() { let dir = corpus_root().join("raw"); let Ok(rd) = std::fs::read_dir(&dir) else { eprintln!("skip: {} absent", dir.display()); @@ -276,14 +280,51 @@ fn corpus_adobe_dngs_carry_profile_looks() { let g = t.apply([0.18; 3], 1.0); assert!(g.iter().all(|v| (v - g[0]).abs() < 0.01 * g[0].max(0.01)), "{name}: grey → {g:?}"); } + if name.starts_with("dng-apple-") { + seen += 1; + let map = look.gain_table_map.as_ref().unwrap_or_else(|| panic!("{name}: no gain table map")); + assert!(map.points_v > 1 && map.points_h > 1 && map.points_n > 1, "{name}"); + assert!(look.tone_curve.is_some(), "{name}: no tone curve"); + // the map lifts dark tones (Apple's local tone mapping), most at the darkest input + let first = &map.gains[..map.points_n]; + assert!(first[0] > 1.5 && first[0] >= first[map.points_n - 1], "{name}: first table {first:?}"); + } eprintln!( - "{name:44} profile look: hsm {} look {} tone {}", + "{name:44} profile look: hsm {} look {} tone {} gain map {}", look.hue_sat_map[0].is_some(), look.look_table.is_some(), - look.tone_curve.is_some() + look.tone_curve.is_some(), + look.gain_table_map.is_some() ); } - eprintln!("{seen} Adobe-converted DNGs checked"); + eprintln!("{seen} DNGs with profile looks checked"); +} + +/// iPhone ProRAW (`dng-apple-iphone12pro-proraw.dng`, iPhone 12 Pro, CC0 from raw.pixls.us, see +/// `cargo xtask corpus`) carries a lossy-JPEG sky matte at half resolution as a DNG semantic mask: +/// it decodes, and it is a sky. The matte is stored like the raw (landscape; the photo is shown +/// rotated 90° clockwise), so the sky is on its left and the lake and shore on its right. +#[test] +fn corpus_proraw_sky_matte() { + let path = corpus_root().join("raw/dng-apple-iphone12pro-proraw.dng"); + let Ok(bytes) = std::fs::read(&path) else { + eprintln!("skip: {} absent", path.display()); + return; + }; + let t0 = Instant::now(); + let masks = lightcraft_raw::semantic_masks(&bytes); + eprintln!("semantic masks read in {:.1} ms", t0.elapsed().as_secs_f64() * 1e3); + let names: Vec<&str> = masks.iter().map(|m| m.name.as_str()).collect(); + assert_eq!(names, ["urn:com:apple:photo:2020:aux:semanticskymatte"]); + let m = &masks[0]; + assert_eq!((m.width, m.height, m.sub_area), (2016, 1512, None)); + let cols = |x0: usize, x1: usize| { + let sum: f64 = (0..m.height).flat_map(|y| m.data[y * m.width + x0..y * m.width + x1].iter()).map(|&v| v as f64 / 65535.0).sum(); + sum / ((x1 - x0) * m.height) as f64 + }; + let (left, right) = (cols(0, 100), cols(1916, 2016)); + eprintln!("sky matte: left columns {left:.3}, right columns {right:.3}"); + assert!(left > 0.5 && right < 0.02, "left {left}, right {right}"); } /// Issue #148: Sony ARWs from before ~2017 carry no plain white-balance, black-level or crop tags in the raw IFD. diff --git a/crates/raw/tests/dng.rs b/crates/raw/tests/dng.rs index 0d38ebb4..bba5ac81 100644 --- a/crates/raw/tests/dng.rs +++ b/crates/raw/tests/dng.rs @@ -39,6 +39,7 @@ fn synthetic(w: usize, h: usize, cfa: Option, cpp: usize) -> RawImage { ], as_shot_neutral: Some([0.5, 1.0, 0.7]), baseline_exposure: 0.35, + baseline_sharpness: Some(1.5), ..Default::default() }; let metadata = lightcraft_meta::Metadata { make: Some("Synth".into()), model: Some("Cam 1".into()), rating: Some(3), ..Default::default() }; @@ -84,6 +85,7 @@ fn assert_same(a: &RawImage, b: &RawImage) { let (n1, n2) = (a.color.as_shot_neutral.unwrap(), b.color.as_shot_neutral.unwrap()); assert!((0..3).all(|i| (n1[i] - n2[i]).abs() < 1e-6)); assert!((a.color.baseline_exposure - b.color.baseline_exposure).abs() < 1e-6); + assert!(b.color.baseline_sharpness.is_some_and(|v| (v - 1.5).abs() < 1e-6), "{:?}", b.color.baseline_sharpness); assert_eq!(b.metadata.make.as_deref(), Some("Synth")); assert_eq!(b.metadata.rating, Some(3)); } @@ -548,3 +550,422 @@ fn malformed_profile_look_tags_are_ignored() { let plain = decode(&profile_dng(ByteOrder::Little, |_| {})).unwrap(); assert!(plain.color.profile.is_empty()); } + +/// Lossless JPEG XL codestream of `px` (`w × h × cpp` interleaved), 16-bit unless `eight`. +fn jxl_encode(px: &[u16], w: usize, h: usize, cpp: usize, eight: bool) -> Vec { + use zune_core::bit_depth::BitDepth; + use zune_core::colorspace::ColorSpace; + use zune_core::options::EncoderOptions; + let cs = if cpp == 1 { ColorSpace::Luma } else { ColorSpace::RGB }; + let (bytes, depth): (Vec, _) = if eight { + (px.iter().map(|&v| v as u8).collect(), BitDepth::Eight) + } else { + (px.iter().flat_map(|v| v.to_ne_bytes()).collect(), BitDepth::Sixteen) + }; + let mut out = Vec::new(); + zune_jpegxl::JxlSimpleEncoder::new(&bytes, EncoderOptions::new(w, h, cs, depth)).encode(&mut out).unwrap(); + out +} + +/// A bare codestream wrapped in the ISO-BMFF container (signature, `ftyp`, `jxlc`). +fn jxl_container(cs: &[u8]) -> Vec { + let mut b = vec![0, 0, 0, 0x0c, b'J', b'X', b'L', b' ', 0x0d, 0x0a, 0x87, 0x0a, 0, 0, 0, 0x14]; + b.extend_from_slice(b"ftypjxl \0\0\0\0jxl "); + b.extend_from_slice(&((cs.len() + 8) as u32).to_be_bytes()); + b.extend_from_slice(b"jxlc"); + b.extend_from_slice(cs); + b +} + +/// A tiled JPEG XL raw IFD of `px`. Edge tiles alternate between being coded at their in-image size +/// and padded to the nominal tile size (with junk the decoder must drop); tiles alternate between bare +/// codestreams and containers. +fn jxl_tiled(px: &[u16], w: usize, h: usize, cpp: usize, bits: u16, cfa: bool, tile: (usize, usize)) -> IfdBuilder { + let (tw, th) = tile; + let mut tiles = Vec::new(); + for ty in 0..h.div_ceil(th) { + for tx in 0..w.div_ceil(tw) { + let i = tiles.len(); + let (vw, vh) = (tw.min(w - tx * tw), th.min(h - ty * th)); + let (cw, ch) = if i % 2 == 0 { (tw, th) } else { (vw, vh) }; + let mut t = vec![bits.min(8) * 7; cw * ch * cpp]; + for y in 0..vh { + for x in 0..vw { + for s in 0..cpp { + t[(y * cw + x) * cpp + s] = px[((ty * th + y) * w + tx * tw + x) * cpp + s]; + } + } + } + let cs = jxl_encode(&t, cw, ch, cpp, bits == 8); + tiles.push(if i % 3 == 1 { jxl_container(&cs) } else { cs }); + } + } + let mut raw = base_ifd(w, h, bits, cpp, cfa); + raw.set(t::COMPRESSION, Value::Short(vec![compression::JPEG_XL])); + raw.set_image(ImageData::Tiles { tile_width: tw as u32, tile_height: th as u32, tiles }); + raw +} + +#[test] +fn jpeg_xl_tiles_decode_bit_exactly() { + let (w, h) = (37, 29); + // LinearRaw 3 × 16-bit, the full code range + let px = pattern(w, h, 3, 16); + let bytes = dng(jxl_tiled(&px, w, h, 3, 16, false, (16, 16)), ByteOrder::Big); + let img = decode(&bytes).unwrap(); + assert_eq!((img.width, img.height, img.cpp), (w, h, 3)); + assert_eq!(img.data, RawData::U16(px)); + assert_eq!(lightcraft_raw::probe_info(&bytes).unwrap(), img.info()); + // CFA 1 × 8/12/14-bit + for bits in [8u32, 12, 14] { + let px = pattern(w, h, 1, bits); + let img = decode(&dng(jxl_tiled(&px, w, h, 1, bits as u16, true, (16, 8)), ByteOrder::Little)).unwrap(); + assert_eq!(img.data, RawData::U16(px), "{bits} bits"); + assert_eq!(img.white, vec![((1u32 << bits) - 1) as f32]); + assert_eq!(img.cfa.unwrap().name(), "GRBG"); + } + // strips: a single JPEG XL strip per 8 rows (the last one shorter) + let px = pattern(w, h, 1, 12); + let strips = px.chunks(w * 8).map(|c| jxl_encode(c, w, c.len() / w, 1, false)).collect(); + let mut raw = base_ifd(w, h, 12, 1, true); + raw.set(t::COMPRESSION, Value::Short(vec![compression::JPEG_XL])); + raw.set_image(ImageData::Strips { rows_per_strip: 8, strips }); + assert_eq!(decode(&dng(raw, ByteOrder::Big)).unwrap().data, RawData::U16(px)); +} + +#[test] +fn malformed_jpeg_xl_tiles_are_errors() { + use lightcraft_raw::RawError; + let (w, h) = (16, 16); + let px = pattern(w, h, 1, 12); + let good = jxl_encode(&px, w, h, 1, false); + let one = |tile: Vec, cpp: usize| { + let mut raw = base_ifd(w, h, 12, cpp, cpp == 1); + raw.set(t::COMPRESSION, Value::Short(vec![compression::JPEG_XL])); + raw.set_image(ImageData::Tiles { tile_width: w as u32, tile_height: h as u32, tiles: vec![tile] }); + dng(raw, ByteOrder::Little) + }; + // garbage, every truncation, and flipped bytes: errors (or, for a truncation that only drops + // trailing padding, the exact samples), never panics + let mut garbage = vec![0xff, 0x0a]; + garbage.extend((0..400u32).map(|i| (i.wrapping_mul(2654435761) >> 13) as u8)); + assert!(matches!(decode(&one(garbage, 1)), Err(RawError::Corrupt(_)))); + for n in 0..good.len() { + let bytes = one(good[..n].to_vec(), 1); + if let Ok(img) = decode(&bytes) { + assert_eq!(img.data, RawData::U16(px.clone()), "truncated to {n}"); + } + let _ = lightcraft_raw::probe_info(&bytes); + } + for i in 0..good.len() { + let mut b = good.clone(); + b[i] ^= 0x5a; + let _ = decode(&one(b, 1)); + } + // a tile of the wrong size, or the wrong plane count for the IFD + let small = jxl_encode(&px[..w * 8], w, 8, 1, false); + assert!(matches!(decode(&one(small, 1)), Err(RawError::Corrupt(_)))); + assert!(matches!(decode(&one(good.clone(), 3)), Err(RawError::Corrupt(_)))); + assert!(matches!(lightcraft_raw::probe_info(&one(good, 3)), Err(RawError::Corrupt(_)))); +} + +/// The decoder's memory is bounded by the tile: a grey tile that also codes an alpha channel needs +/// about twice what one plane may use, and is an error rather than an allocation the header asks +/// for. The same samples without the alpha channel decode. +#[test] +fn jpeg_xl_tiles_beyond_their_decode_budget_are_errors() { + use lightcraft_raw::RawError; + use zune_core::bit_depth::BitDepth; + use zune_core::colorspace::ColorSpace; + use zune_core::options::EncoderOptions; + let (w, h) = (512, 512); + let grey: Vec = (0..w * h).map(|i| ((i % w) * 97 + (i / w) * 31) as u16 & 0x0fff).collect(); + let with_alpha: Vec = grey.iter().flat_map(|&v| [v, 4095 - v]).flat_map(u16::to_ne_bytes).collect(); + let mut tile = Vec::new(); + zune_jpegxl::JxlSimpleEncoder::new(&with_alpha, EncoderOptions::new(w, h, ColorSpace::LumaA, BitDepth::Sixteen)).encode(&mut tile).unwrap(); + let one = |tile: Vec| { + let mut raw = base_ifd(w, h, 12, 1, true); + raw.set(t::COMPRESSION, Value::Short(vec![compression::JPEG_XL])); + raw.set_image(ImageData::Tiles { tile_width: w as u32, tile_height: h as u32, tiles: vec![tile] }); + dng(raw, ByteOrder::Little) + }; + assert!(matches!(decode(&one(tile)), Err(RawError::Corrupt(_)))); + assert_eq!(decode(&one(jxl_encode(&grey, w, h, 1, false))).unwrap().data, RawData::U16(grey)); +} + +#[test] +fn jpeg_xl_preview_ifd_is_offered() { + let raw = synthetic(16, 16, Some(Cfa::bayer("RGGB").unwrap()), 1); + let RawData::U16(px) = &raw.data else { unreachable!() }; + let mut sub = base_ifd(16, 16, 14, 1, true); + sub.set(t::COMPRESSION, Value::Short(vec![compression::JPEG_XL])); + sub.set_image(ImageData::Tiles { tile_width: 16, tile_height: 16, tiles: vec![jxl_encode(px, 16, 16, 1, false)] }); + let preview = jxl_container(&jxl_encode(&[200; 4 * 2 * 3], 4, 2, 3, true)); + let bytes = dng_with(sub, ByteOrder::Big, |ifd0| { + ifd0.set(t::COMPRESSION, Value::Short(vec![compression::JPEG_XL])); + ifd0.set_image(ImageData::Strips { rows_per_strip: 2, strips: vec![preview.clone()] }); + }); + // the rendered RGB preview is offered whole; the (single-tile) CFA raw never is + assert_eq!(embedded_preview(&bytes).unwrap(), preview); + assert_eq!(decode(&bytes).unwrap().data, raw.data); +} + +/// `ProfileGainTableMap` / `ProfileGainTableMap2` bytes in `order`: a `dims` = (V, H, N) grid with +/// spacing 0.5 and origin 0.25, green-only input weights, `v2` = (data type, gamma, gain min, +/// gain max), then `data` (the stored gains, already in `order`). +fn gain_map_tag(order: ByteOrder, dims: [u32; 3], v2: Option<(u32, f32, f32, f32)>, data: &[u8]) -> Vec { + let be = order == ByteOrder::Big; + let u = |v: u32| if be { v.to_be_bytes() } else { v.to_le_bytes() }; + let mut b = Vec::new(); + b.extend_from_slice(&u(dims[0])); + b.extend_from_slice(&u(dims[1])); + for v in [0.5f64, 0.5, 0.25, 0.25] { + b.extend_from_slice(&if be { v.to_be_bytes() } else { v.to_le_bytes() }); + } + b.extend_from_slice(&u(dims[2])); + for w in [0.0f32, 1.0, 0.0, 0.0, 0.0] { + b.extend_from_slice(&u(w.to_bits())); + } + if let Some((dt, gamma, min, max)) = v2 { + for v in [dt, gamma.to_bits(), min.to_bits(), max.to_bits()] { + b.extend_from_slice(&u(v)); + } + } + b.extend_from_slice(data); + b +} + +fn f32_bytes(v: &[f32], order: ByteOrder) -> Vec { + v.iter().flat_map(|f| if order == ByteOrder::Big { f.to_be_bytes() } else { f.to_le_bytes() }).collect() +} + +/// [`profile_dng`] with binary tags in the raw IFD (where Apple writes `ProfileGainTableMap`) and +/// in IFD 0. +fn gain_map_dng(order: ByteOrder, raw_tags: &[(u16, &[u8])], ifd0_tags: &[(u16, &[u8])]) -> Vec { + let (w, h) = (8, 6); + let px = pattern(w, h, 3, 16); + let mut raw = base_ifd(w, h, 16, 3, false); + raw.set(t::COMPRESSION, Value::Short(vec![1])); + raw.set_image(ImageData::Strips { rows_per_strip: h as u32, strips: vec![u16_bytes(&px, order)] }); + for (tag, b) in raw_tags { + raw.set(*tag, Value::Undefined(b.to_vec())); + } + dng_with(raw, order, |ifd0| { + for (tag, b) in ifd0_tags { + ifd0.set(*tag, Value::Undefined(b.to_vec())); + } + }) +} + +#[test] +fn gain_table_maps_follow_dng_precedence() { + const V1: u16 = t::PROFILE_GAIN_TABLE_MAP; + const V2: u16 = t::PROFILE_GAIN_TABLE_MAP_2; + for order in [ByteOrder::Big, ByteOrder::Little] { + let v1 = gain_map_tag(order, [2, 2, 4], None, &f32_bytes(&(1..=16).map(|i| i as f32 * 0.25).collect::>(), order)); + // Apple's layout: version 1 in the raw IFD, in the file's byte order + let apple = gain_map_dng(order, &[(V1, &v1[..])], &[]); + let img = decode(&apple).unwrap(); + let map = img.color.profile.gain_table_map.clone().expect("gain table map"); + assert_eq!((map.points_v, map.points_h, map.points_n, map.gamma), (2, 2, 4, 1.0)); + assert_eq!((map.spacing_v, map.origin_h, map.weights[1]), (0.5, 0.25, 1.0)); + assert_eq!(map.gains[5], 1.5); + assert_eq!(lightcraft_raw::probe_info(&apple).unwrap().color.profile, img.color.profile); + // our DNG writer keeps it + let again = decode(&write_dng(&img, &DngWriteOptions::default()).unwrap()).unwrap(); + assert_eq!(again.color.profile.gain_table_map.as_ref(), Some(&map)); + // version 1 in IFD 0 (where DNG 1.7 says it was meant to go) is read too + let in_ifd0 = decode(&gain_map_dng(order, &[], &[(V1, &v1[..])])).unwrap(); + assert_eq!(in_ifd0.color.profile.gain_table_map.as_ref(), Some(&map)); + // version 2 in IFD 0 supersedes version 1: 8-bit gains over 0.5..=1.5, gamma 2 + let v2 = gain_map_tag(order, [1, 1, 2], Some((0, 2.0, 0.5, 1.5)), &[0, 255]); + let img = decode(&gain_map_dng(order, &[(V1, &v1[..])], &[(V2, &v2[..])])).unwrap(); + let map2 = img.color.profile.gain_table_map.clone().unwrap(); + assert_eq!((map2.points_n, map2.gamma, map2.gains.clone()), (2, 2.0, vec![0.5, 1.5])); + // a gamma ≠ 1 survives our writer (as version 2) + let again = decode(&write_dng(&img, &DngWriteOptions::default()).unwrap()).unwrap(); + assert_eq!(again.color.profile.gain_table_map, Some(map2)); + // a malformed version 2 (gamma outside 0.25..=4) falls back to version 1 + let bad_v2 = gain_map_tag(order, [1, 1, 2], Some((0, 9.0, 0.5, 1.5)), &[0, 255]); + assert_eq!(decode(&gain_map_dng(order, &[(V1, &v1[..])], &[(V2, &bad_v2[..])])).unwrap().color.profile.gain_table_map, Some(map)); + } +} + +#[test] +fn malformed_gain_table_maps_are_ignored() { + let order = ByteOrder::Little; + let ok = gain_map_tag(order, [1, 1, 2], None, &f32_bytes(&[1.0, 2.0], order)); + let cases: Vec<(&str, Vec, bool)> = vec![ + ("valid", ok.clone(), false), + ("empty", vec![], false), + ("header only", ok[..64].to_vec(), false), + ("one byte short", ok[..ok.len() - 1].to_vec(), false), + ("NaN gain", gain_map_tag(order, [1, 1, 2], None, &f32_bytes(&[1.0, f32::NAN], order)), false), + ("negative gain", gain_map_tag(order, [1, 1, 2], None, &f32_bytes(&[1.0, -2.0], order)), false), + ("huge dimensions", gain_map_tag(order, [u32::MAX, u32::MAX, u32::MAX], None, &f32_bytes(&[1.0, 2.0], order)), false), + ("zero points", gain_map_tag(order, [1, 0, 2], None, &[]), false), + ("unknown data type", gain_map_tag(order, [1, 1, 2], Some((9, 1.0, 0.0, 1.0)), &[0, 0]), true), + ("NaN gain range", gain_map_tag(order, [1, 1, 2], Some((0, 1.0, f32::NAN, 1.0)), &[0, 255]), true), + ]; + for (what, bytes, v2) in cases { + let tag = if v2 { t::PROFILE_GAIN_TABLE_MAP_2 } else { t::PROFILE_GAIN_TABLE_MAP }; + let img = decode(&gain_map_dng(order, &[(tag, &bytes[..])], &[])).unwrap_or_else(|e| panic!("{what}: {e}")); + assert_eq!(img.color.profile.gain_table_map.is_some(), what == "valid", "{what}"); + } +} + +/// The map is read in the file's byte order, else big-endian (`lightcraft_raw::gaintable`): a +/// big-endian payload in a little-endian file is read, a little-endian one in a big-endian file +/// is not. +#[test] +fn gain_table_map_byte_orders() { + use ByteOrder::{Big, Little}; + let payload = |o: ByteOrder| gain_map_tag(o, [1, 1, 2], None, &f32_bytes(&[1.0, 2.0], o)); + for (file, data, read) in [(Little, Little, true), (Little, Big, true), (Big, Big, true), (Big, Little, false)] { + let img = decode(&gain_map_dng(file, &[(t::PROFILE_GAIN_TABLE_MAP, &payload(data)[..])], &[])).unwrap(); + assert_eq!(img.color.profile.gain_table_map.map(|m| m.gains), read.then(|| vec![1.0, 2.0]), "{data:?} payload, {file:?} file"); + } +} + +/// A semantic mask IFD (8-bit, uncompressed) named `name`. +fn mask_ifd(w: usize, h: usize, name: &str, px: Vec) -> IfdBuilder { + let mut m = IfdBuilder::new(); + m.set(t::NEW_SUBFILE_TYPE, Value::Long(vec![t::SUBFILE_SEMANTIC_MASK])); + m.set(t::IMAGE_WIDTH, Value::Long(vec![w as u32])); + m.set(t::IMAGE_LENGTH, Value::Long(vec![h as u32])); + m.set(t::BITS_PER_SAMPLE, Value::Short(vec![8])); + m.set(t::SAMPLES_PER_PIXEL, Value::Short(vec![1])); + m.set(t::PHOTOMETRIC, Value::Short(vec![photometric::MASK])); + m.set(t::COMPRESSION, Value::Short(vec![1])); + if !name.is_empty() { + m.set(t::SEMANTIC_NAME, Value::Ascii(name.into())); + } + m.set_image(ImageData::Strips { rows_per_strip: h as u32, strips: vec![px] }); + m +} + +/// Semantic masks: a cropped mask is placed by its `MaskSubArea` (zero outside it), mapped onto the +/// developed image (default crop of the active area), and doesn't disturb the raw decode. +#[test] +fn semantic_masks_are_read_and_placed() { + use lightcraft_raw::semantic::MaskSubArea; + for order in [ByteOrder::Little, ByteOrder::Big] { + let bytes = profile_dng(order, |ifd0| { + // a 4 × 2 crop at (top 1, left 2) of an 8 × 4 mask over the 8 × 6 image + let mut sky = mask_ifd(4, 2, "sky", vec![10, 20, 30, 40, 50, 60, 70, 255]); + sky.set(t::MASK_SUB_AREA, Value::Long(vec![1, 2, 8, 4])); + ifd0.add_sub_ifd(sky); + let mut skin = mask_ifd(2, 2, "skin", vec![0, 255, 255, 0]); + skin.set(t::SEMANTIC_INSTANCE_ID, Value::Ascii("person_a".into())); + ifd0.add_sub_ifd(skin); + }); + let raw = decode(&bytes).unwrap(); + assert_eq!((raw.width, raw.height, raw.cpp), (8, 6, 3), "the masks are not the raw image"); + let masks = lightcraft_raw::semantic_masks(&bytes); + assert_eq!(masks.len(), 2); + let (sky, skin) = (&masks[0], &masks[1]); + assert_eq!((sky.name.as_str(), sky.instance_id.as_deref(), sky.width, sky.height), ("sky", None, 4, 2)); + assert_eq!(sky.sub_area, Some(MaskSubArea { top: 1, left: 2, full_width: 8, full_height: 4 })); + assert_eq!((sky.data[0], sky.data[7]), (10 * 257, 65535)); + assert_eq!((skin.name.as_str(), skin.instance_id.as_deref(), skin.sub_area), ("skin", Some("person_a"), None)); + // positions are in uncropped-mask pixels: zero outside the stored crop + assert!((sky.value_at(2.5, 1.5) - 10.0 / 255.0).abs() < 1e-4); + assert!((sky.value_at(5.5, 2.5) - 1.0).abs() < 1e-4); + assert_eq!((sky.value_at(1.5, 1.5), sky.value_at(2.5, 0.5), sky.value_at(6.5, 1.5)), (0.0, 0.0, 0.0)); + // over the whole image: the uncropped mask's 8 × 4 grid + let full = sky.developed(Rect::new(0, 0, 8, 6), Rect::new(0, 0, 8, 6)).unwrap(); + assert_eq!((full.width, full.height), (8, 4)); + assert_eq!((full.get(2, 1), full.get(5, 2), full.get(1, 1), full.get(6, 2), full.get(2, 0)), (10, 255, 0, 0, 0)); + // a default crop keeping columns 2..6 of the image: the stored crop fills it + let cropped = sky.developed(Rect::new(0, 0, 8, 6), Rect::new(2, 0, 4, 6)).unwrap(); + assert_eq!((cropped.width, cropped.height), (4, 4)); + assert_eq!((cropped.get(0, 1), cropped.get(3, 2), cropped.get(0, 0)), (10, 255, 0)); + } +} + +/// Malformed semantic masks are skipped (never a panic), and so is the `MaskSubArea` of a mask it +/// doesn't fit; the raw image still decodes. +#[test] +fn malformed_semantic_masks_are_skipped() { + let bytes = profile_dng(ByteOrder::Little, |ifd0| { + ifd0.add_sub_ifd(mask_ifd(2, 2, "", vec![1, 2, 3, 4])); // no SemanticName + let mut rgb = mask_ifd(2, 2, "rgb", vec![0; 12]); + rgb.set(t::SAMPLES_PER_PIXEL, Value::Short(vec![3])); + rgb.set(t::BITS_PER_SAMPLE, Value::Short(vec![8, 8, 8])); + ifd0.add_sub_ifd(rgb); + let mut huge = mask_ifd(2, 2, "huge", vec![0; 4]); + huge.set(t::IMAGE_WIDTH, Value::Long(vec![100_000])); + huge.set(t::IMAGE_LENGTH, Value::Long(vec![100_000])); + ifd0.add_sub_ifd(huge); + let mut tiny = mask_ifd(2, 2, "tiny data", vec![0; 4]); + tiny.set(t::IMAGE_WIDTH, Value::Long(vec![4000])); + tiny.set(t::IMAGE_LENGTH, Value::Long(vec![3000])); + ifd0.add_sub_ifd(tiny); + let mut jxl = mask_ifd(2, 2, "unknown coding", vec![0; 4]); + jxl.set(t::COMPRESSION, Value::Short(vec![60_000])); + ifd0.add_sub_ifd(jxl); + let mut zero = mask_ifd(2, 2, "zero size", vec![0; 4]); + zero.set(t::IMAGE_WIDTH, Value::Long(vec![0])); + ifd0.add_sub_ifd(zero); + let mut short = mask_ifd(4, 4, "short strip", vec![200; 3]); + short.set(t::MASK_SUB_AREA, Value::Long(vec![u32::MAX, u32::MAX, u32::MAX, u32::MAX])); + ifd0.add_sub_ifd(short); + let mut outside = mask_ifd(2, 2, "sub area too small", vec![255; 4]); + outside.set(t::MASK_SUB_AREA, Value::Long(vec![3, 0, 2, 4])); + ifd0.add_sub_ifd(outside); + }); + assert_eq!(decode(&bytes).unwrap().cpp, 3); + let masks = lightcraft_raw::semantic_masks(&bytes); + let names: Vec<&str> = masks.iter().map(|m| m.name.as_str()).collect(); + assert_eq!(names, ["short strip", "sub area too small"]); + assert!(masks.iter().all(|m| m.sub_area.is_none()), "sub areas the mask doesn't fit are ignored"); + for m in &masks { + let d = m.developed(Rect::new(0, 0, 8, 6), Rect::new(0, 0, 8, 6)).unwrap(); + assert_eq!((d.width, d.height), (m.width, m.height)); + assert!(m.developed(Rect::new(0, 0, 8, 6), Rect::new(8, 6, 0, 0)).is_none()); + } + // not a TIFF at all + assert!(lightcraft_raw::semantic_masks(b"not a dng").is_empty()); +} + +/// Only the first `MAX_MASKS` mask IFDs are decoded, malformed ones included, so a file full of +/// mask IFDs costs a bounded number of decodes: here the first `MAX_MASKS` can't be decoded, and +/// the valid masks after them are never looked at. +#[test] +fn semantic_mask_ifds_past_the_cap_are_not_decoded() { + use lightcraft_raw::semantic::MAX_MASKS; + let with = |undecodable: usize, valid: usize| { + profile_dng(ByteOrder::Little, |ifd0| { + for i in 0..undecodable { + let mut m = mask_ifd(2, 2, &format!("undecodable {i}"), vec![0; 4]); + m.set(t::COMPRESSION, Value::Short(vec![60_000])); + ifd0.add_sub_ifd(m); + } + for i in 0..valid { + ifd0.add_sub_ifd(mask_ifd(2, 2, &format!("valid {i}"), vec![255; 4])); + } + }) + }; + let bytes = with(MAX_MASKS, 3 * MAX_MASKS); + assert_eq!(decode(&bytes).unwrap().cpp, 3); + assert!(lightcraft_raw::semantic_masks(&bytes).is_empty()); + // one fewer undecodable mask leaves room for exactly one valid one + let names: Vec = lightcraft_raw::semantic_masks(&with(MAX_MASKS - 1, 3 * MAX_MASKS)).into_iter().map(|m| m.name).collect(); + assert_eq!(names, ["valid 0"]); +} + +/// `BaselineSharpness` is read from IFD 0; a zero, negative or absurd value is ignored. +#[test] +fn baseline_sharpness_is_read_when_plausible() { + for (value, want) in [((3, 2), Some(1.5)), ((0, 1), None), ((1, 0), None), ((100, 1), None)] { + let bytes = profile_dng(ByteOrder::Big, |ifd0| { + ifd0.set(t::BASELINE_SHARPNESS, Value::Rational(vec![value])); + }); + assert_eq!(decode(&bytes).unwrap().color.baseline_sharpness, want, "{value:?}"); + } + let bytes = profile_dng(ByteOrder::Big, |ifd0| { + ifd0.set(t::BASELINE_SHARPNESS, Value::SRational(vec![(-1, 1)])); + }); + assert_eq!(decode(&bytes).unwrap().color.baseline_sharpness, None); + assert_eq!(decode(&profile_dng(ByteOrder::Big, |_| {})).unwrap().color.baseline_sharpness, None); +} diff --git a/crates/tiff/src/tags.rs b/crates/tiff/src/tags.rs index aa24ec87..5771316a 100644 --- a/crates/tiff/src/tags.rs +++ b/crates/tiff/src/tags.rs @@ -172,6 +172,7 @@ pub const OPCODE_LIST_2: u16 = 51009; pub const OPCODE_LIST_3: u16 = 51022; pub const NOISE_PROFILE: u16 = 51041; pub const DEFAULT_USER_CROP: u16 = 51125; +pub const PROFILE_GAIN_TABLE_MAP: u16 = 52525; pub const CALIBRATION_ILLUMINANT_3: u16 = 52529; pub const CAMERA_CALIBRATION_3: u16 = 52530; pub const COLOR_MATRIX_3: u16 = 52531; @@ -179,6 +180,13 @@ pub const FORWARD_MATRIX_3: u16 = 52532; pub const ILLUMINANT_DATA_1: u16 = 52533; pub const ILLUMINANT_DATA_2: u16 = 52534; pub const ILLUMINANT_DATA_3: u16 = 52535; +// DNG 1.6 semantic masks (spec chapter 4, "Semantic Masks") +pub const SEMANTIC_NAME: u16 = 52526; +pub const SEMANTIC_INSTANCE_ID: u16 = 52528; +pub const MASK_SUB_AREA: u16 = 52536; +/// NewSubFileType of a semantic mask IFD. +pub const SUBFILE_SEMANTIC_MASK: u32 = 0x10004; +pub const PROFILE_GAIN_TABLE_MAP_2: u16 = 52544; // --- Compression codes (TIFF 6.0, DNG, and widely used registrations) pub mod compression { @@ -206,4 +214,6 @@ pub mod photometric { pub const YCBCR: u16 = 6; pub const CFA: u16 = 32803; pub const LINEAR_RAW: u16 = 34892; + /// DNG 1.6 semantic mask. + pub const MASK: u16 = 52527; } diff --git a/crates/ui-egui/Cargo.toml b/crates/ui-egui/Cargo.toml index 49bdfc90..96779e31 100644 --- a/crates/ui-egui/Cargo.toml +++ b/crates/ui-egui/Cargo.toml @@ -26,6 +26,7 @@ log = { workspace = true } [dev-dependencies] lightcraft-codecs = { workspace = true } +lightcraft-faces = { workspace = true } [build-dependencies] serde_json = { workspace = true } diff --git a/crates/ui-egui/locales/de.json b/crates/ui-egui/locales/de.json index 2a472186..520c73f9 100644 --- a/crates/ui-egui/locales/de.json +++ b/crates/ui-egui/locales/de.json @@ -1797,5 +1797,6 @@ "How to install the model": "So installierst du das Modell", "Install the SAM 3 Model": "SAM-3-Modell installieren", "This build can't download the model yet. To install it by hand, put model.safetensors, vocab.json and merges.txt from Meta's facebook/sam3 in the folder above: Object and Describe work as soon as they are there.": "Diese Version kann das Modell noch nicht herunterladen. Zum Installieren von Hand lege model.safetensors, vocab.json und merges.txt aus Metas facebook/sam3 in den obigen Ordner: Objekt und Beschreiben funktionieren, sobald sie dort liegen.", - "Open Log Folder": "Protokollordner öffnen" + "Open Log Folder": "Protokollordner öffnen", + "Back to this person": "Back to this person" } diff --git a/crates/ui-egui/locales/es.json b/crates/ui-egui/locales/es.json index c0986c19..1c8b0013 100644 --- a/crates/ui-egui/locales/es.json +++ b/crates/ui-egui/locales/es.json @@ -1642,5 +1642,6 @@ "Captured from": "Capturadas desde", "Captured until": "Capturadas hasta", "Only": "Solo", - "Open Log Folder": "Abrir carpeta de registros" + "Open Log Folder": "Abrir carpeta de registros", + "Back to this person": "Back to this person" } diff --git a/crates/ui-egui/locales/ja.json b/crates/ui-egui/locales/ja.json index 31c61a7a..afc2fe44 100644 --- a/crates/ui-egui/locales/ja.json +++ b/crates/ui-egui/locales/ja.json @@ -1642,5 +1642,6 @@ "How to install the model": "モデルのインストール方法", "Install the SAM 3 Model": "SAM 3モデルのインストール", "This build can't download the model yet. To install it by hand, put model.safetensors, vocab.json and merges.txt from Meta's facebook/sam3 in the folder above: Object and Describe work as soon as they are there.": "このビルドではまだモデルをダウンロードできません。手動でインストールするには、Meta の facebook/sam3 の model.safetensors、vocab.json、merges.txt を上のフォルダに置いてください。置いた時点で「オブジェクト」と「説明で選択」が使えるようになります。", - "Open Log Folder": "ログフォルダーを開く" + "Open Log Folder": "ログフォルダーを開く", + "Back to this person": "Back to this person" } diff --git a/crates/ui-egui/locales/pt-br.json b/crates/ui-egui/locales/pt-br.json index 03ae0773..be7d22ba 100644 --- a/crates/ui-egui/locales/pt-br.json +++ b/crates/ui-egui/locales/pt-br.json @@ -1642,5 +1642,6 @@ "≤ at most": "≤ no máximo", "≥ at least": "≥ no mínimo", "日本語": "日本語", - "繁體中文（台灣）": "繁體中文（台灣）" + "繁體中文（台灣）": "繁體中文（台灣）", + "Back to this person": "Back to this person" } diff --git a/crates/ui-egui/locales/ru.json b/crates/ui-egui/locales/ru.json index 530cdfcb..f1d33491 100644 --- a/crates/ui-egui/locales/ru.json +++ b/crates/ui-egui/locales/ru.json @@ -1312,5 +1312,6 @@ "of the following rules:": "из следующих условий:", "saved in": "сохранено в", "this photo": "это фото", - "− harder border · + softer border": "− жёстче край · + мягче край" + "− harder border · + softer border": "− жёстче край · + мягче край", + "Back to this person": "Back to this person" } diff --git a/crates/ui-egui/locales/zh-hans.json b/crates/ui-egui/locales/zh-hans.json index cea313d0..3d88dcc2 100644 --- a/crates/ui-egui/locales/zh-hans.json +++ b/crates/ui-egui/locales/zh-hans.json @@ -1642,5 +1642,6 @@ "How to install the model": "如何安装模型", "Install the SAM 3 Model": "安装 SAM 3 模型", "This build can't download the model yet. To install it by hand, put model.safetensors, vocab.json and merges.txt from Meta's facebook/sam3 in the folder above: Object and Describe work as soon as they are there.": "此版本尚无法下载模型。如需手动安装，请将 Meta 的 facebook/sam3 中的 model.safetensors、vocab.json 和 merges.txt 放入上方的文件夹：放入后即可使用“对象”和“描述”。", - "Open Log Folder": "打开日志文件夹" + "Open Log Folder": "打开日志文件夹", + "Back to this person": "Back to this person" } diff --git a/crates/ui-egui/locales/zh-hant.json b/crates/ui-egui/locales/zh-hant.json index 597ecbf9..085ccadc 100644 --- a/crates/ui-egui/locales/zh-hant.json +++ b/crates/ui-egui/locales/zh-hant.json @@ -1642,5 +1642,6 @@ "How to install the model": "如何安裝模型", "Install the SAM 3 Model": "安裝 SAM 3 模型", "This build can't download the model yet. To install it by hand, put model.safetensors, vocab.json and merges.txt from Meta's facebook/sam3 in the folder above: Object and Describe work as soon as they are there.": "此版本尚無法下載模型。如需手動安裝，請將 Meta 的 facebook/sam3 中的 model.safetensors、vocab.json 和 merges.txt 放入上方的資料夾：放入後即可使用「物件」與「描述」。", - "Open Log Folder": "開啟記錄資料夾" + "Open Log Folder": "開啟記錄資料夾", + "Back to this person": "Back to this person" } diff --git a/crates/ui-egui/src/control.rs b/crates/ui-egui/src/control.rs index cfcb287f..dcd3d1f8 100644 --- a/crates/ui-egui/src/control.rs +++ b/crates/ui-egui/src/control.rs @@ -105,6 +105,8 @@ pub fn inspect(app: &LightcraftApp, ctx: &egui::Context) -> Value { "status": app.ui.status, "notices": app.notices, "quitPrompt": app.quit_prompt, + // how hard the background face scan may work right now, and why it is judged so (see `panels::faces::scan_pace`) + "faceScan": {"pace": app.caches.faces_pace, "focused": app.caches.faces_in_front, "pending": app.caches.faces_pending, "indexed": app.caches.faces_indexed}, "unsaved": app.session.unsaved().map(|(n, e)| json!({"ops": n, "error": e})), "libraryProblem": app.library_problem.as_ref().map(crate::panels::library_problem::LibraryProblem::to_json), "scan": app.scan.as_ref().map(crate::import::ScanTask::status), @@ -323,8 +325,18 @@ pub fn handle(app: &mut LightcraftApp, ctx: &egui::Context, req: &ControlRequest ok(Value::Null) } "ui.resize" => { + for name in ["width", "height"] { + if p.get(name).is_some() && f(name).is_none() { + return err(format!("ui.resize: {name} must be a number")); + } + } let w = f("width").unwrap_or(1600.0) as f32; let h = f("height").unwrap_or(1000.0) as f32; + let zoom = ctx.zoom_factor(); + let native_scale = ctx.input(|i| i.viewport().native_pixels_per_point).unwrap_or_else(|| ctx.pixels_per_point() / zoom); + if let Err(e) = crate::headless::resized_viewport([w, h], native_scale, zoom) { + return err(e); + } ctx.send_viewport_cmd(egui::ViewportCommand::InnerSize(egui::vec2(w, h))); ok(Value::Null) } diff --git a/crates/ui-egui/src/headless.rs b/crates/ui-egui/src/headless.rs index 206d75b1..77aa77b3 100644 --- a/crates/ui-egui/src/headless.rs +++ b/crates/ui-egui/src/headless.rs @@ -28,13 +28,56 @@ const CLEAR: Color32 = Color32::from_rgb(12, 12, 12); /// Simulated frame interval (deterministic animation time). const FRAME_DT: f64 = 1.0 / 60.0; +/// Validate the scaled screenshot dimensions before layout or pixel allocation. +pub fn viewport_pixels(size: [f32; 2], pixels_per_point: f32) -> Result<[usize; 2], String> { + if size.iter().any(|edge| *edge < 1.0) { + return Err("viewport dimensions must be at least one logical point".into()); + } + native_viewport_pixels(size, pixels_per_point) +} + +fn native_viewport_pixels(size: [f32; 2], pixels_per_point: f32) -> Result<[usize; 2], String> { + let [width, height] = size; + if !width.is_finite() || !height.is_finite() || width <= 0.0 || height <= 0.0 || !pixels_per_point.is_finite() || pixels_per_point <= 0.0 { + return Err("viewport dimensions and scale must be finite and positive".into()); + } + let axis = |value: f32| { + // Match egui/scissor rounding and the existing screenshot dimension contract. + let pixels = (value * pixels_per_point).round(); + if !(1.0..=softpaint::MAX_IMAGE_EDGE as f32).contains(&pixels) { + return Err(format!("scaled viewport edges must be between 1 and {} pixels", softpaint::MAX_IMAGE_EDGE)); + } + Ok(pixels as usize) + }; + let (w, h) = (axis(width)?, axis(height)?); + if w.checked_mul(h).is_none_or(|pixels| pixels > softpaint::MAX_IMAGE_PIXELS) { + return Err(format!("scaled viewport exceeds {} pixels", softpaint::MAX_IMAGE_PIXELS)); + } + Ok([w, h]) +} + +/// Convert an effective-point resize to the native viewport before validating pixel bounds. +/// Sub-point native edges after zooming out are valid when they round to at least one pixel. +pub fn resized_viewport(size: [f32; 2], native_pixels_per_point: f32, zoom: f32) -> Result { + if size.iter().any(|edge| !edge.is_finite() || *edge < 1.0) || !zoom.is_finite() || zoom <= 0.0 { + return Err("resize dimensions must be at least one finite logical point with a positive zoom".into()); + } + let native_size = egui::Vec2::from(size) * zoom; + native_viewport_pixels([native_size.x, native_size.y], native_pixels_per_point)?; + Ok(native_size) +} + +fn bounded_viewport(size: egui::Vec2, scale: f32) -> (egui::Vec2, f32) { + if native_viewport_pixels([size.x, size.y], scale).is_ok() { (size, scale) } else { (egui::vec2(1600.0, 1000.0), 1.0) } +} + /// An offscreen egui context with our fonts and theme, and a CPU mirror of its textures. pub struct HeadlessView { pub ctx: egui::Context, pub textures: TextureStore, pub(crate) shapes: Vec, pixels_per_point: f32, - size: egui::Vec2, + size: [usize; 2], frames: u64, } @@ -49,11 +92,12 @@ impl HeadlessView { let ctx = egui::Context::default(); crate::theme::install_fonts(&ctx); crate::theme::apply(&ctx); - HeadlessView { ctx, textures: TextureStore::default(), shapes: vec![], pixels_per_point: 1.0, size: egui::vec2(1600.0, 1000.0), frames: 0 } + HeadlessView { ctx, textures: TextureStore::default(), shapes: vec![], pixels_per_point: 1.0, size: [1600, 1000], frames: 0 } } /// Input for one frame of a `size` (points) viewport at `pixels_per_point`. pub fn raw_input(size: egui::Vec2, pixels_per_point: f32, time: f64, events: Vec) -> RawInput { + let (size, pixels_per_point) = bounded_viewport(size, pixels_per_point); let rect = egui::Rect::from_min_size(egui::Pos2::ZERO, size); let mut raw = RawInput { screen_rect: Some(rect), @@ -73,9 +117,20 @@ impl HeadlessView { } /// Run one frame; keeps the shapes for [`Self::paint`]. Returns the root viewport's commands. - pub fn run(&mut self, raw: RawInput, run_ui: impl FnMut(&mut egui::Ui)) -> Vec { - if let Some(r) = raw.screen_rect { - self.size = r.size(); + pub fn run(&mut self, mut raw: RawInput, run_ui: impl FnMut(&mut egui::Ui)) -> Vec { + if let Some(rect) = raw.screen_rect { + let native_scale = raw.viewports.get(&ViewportId::ROOT).and_then(|info| info.native_pixels_per_point).unwrap_or(1.0); + self.size = native_viewport_pixels([rect.width(), rect.height()], native_scale).unwrap_or([0, 0]); + // The host viewport is measured before UI zoom. Egui lays out in zoomed points; + // retain the host's physical size while allowing text and controls to grow. + let zoom = self.ctx.zoom_factor(); + if zoom.is_finite() && zoom > 0.0 { + raw.screen_rect = Some(rect / zoom); + if let Some(info) = raw.viewports.get_mut(&ViewportId::ROOT) { + info.inner_rect = info.inner_rect.map(|rect| rect / zoom); + info.outer_rect = info.outer_rect.map(|rect| rect / zoom); + } + } } let mut out = self.ctx.run_ui(raw, run_ui); self.frames += 1; @@ -92,11 +147,17 @@ impl HeadlessView { /// Size in pixels of the last frame. pub fn size_px(&self) -> [usize; 2] { - [(self.size.x * self.pixels_per_point).round() as usize, (self.size.y * self.pixels_per_point).round() as usize] + self.size } /// Rasterize the last frame. `extra` textures (by id) take precedence over the context's own. pub fn paint(&self, extra: &HashMap) -> ColorImage { + if viewport_pixels([self.size[0] as f32, self.size[1] as f32], 1.0).is_err() + || !self.pixels_per_point.is_finite() + || self.pixels_per_point <= 0.0 + { + return ColorImage::new([0, 0], Vec::new()); + } let prims = self.ctx.tessellate(self.shapes.clone(), self.pixels_per_point); softpaint::paint(&prims, &Layered { over: extra, base: &self.textures }, self.size_px(), self.pixels_per_point, CLEAR) } @@ -125,7 +186,14 @@ pub struct Headless { impl Headless { /// Wrap `app` (its control channel is replaced by the driver's). + /// Invalid dimensions fall back to 1600×1000 at scale 1; external callers can use + /// [`viewport_pixels`] to reject them before constructing the session. pub fn new(app: LightcraftApp, size: [f32; 2], pixels_per_point: f32) -> Self { + let (size, pixels_per_point) = if viewport_pixels(size, pixels_per_point).is_ok() { + (egui::Vec2::from(size), pixels_per_point) + } else { + (egui::vec2(1600.0, 1000.0), 1.0) + }; let (tx, rx) = std::sync::mpsc::channel(); let mut app = app.with_control(rx); app.headless_host = true; @@ -133,7 +201,7 @@ impl Headless { app, view: HeadlessView::new(), max_texture_side: 16384, - size: egui::vec2(size[0], size[1]), + size, pixels_per_point, time: 0.0, frames: 0, @@ -156,6 +224,7 @@ impl Headless { /// Run one frame. pub fn step(&mut self) { + (self.size, self.pixels_per_point) = bounded_viewport(self.size, self.pixels_per_point); let mut raw = HeadlessView::raw_input(self.size, self.pixels_per_point, self.time, std::mem::take(&mut self.events)); raw.viewports.entry(ViewportId::ROOT).or_default().maximized = Some(self.window_maximized); raw.max_texture_side = Some(self.max_texture_side); @@ -173,7 +242,12 @@ impl Headless { let image = Arc::new(self.paint()); self.events.push(egui::Event::Screenshot { viewport_id: ViewportId::ROOT, user_data, image }); } - ViewportCommand::InnerSize(s) if s.x >= 1.0 && s.y >= 1.0 => self.size = s, + ViewportCommand::InnerSize(size) => { + // Viewport commands use current egui points, which include UI zoom. + if let Ok(native_size) = resized_viewport([size.x, size.y], self.pixels_per_point, self.view.ctx.zoom_factor()) { + self.size = native_size; + } + } ViewportCommand::Close => self.quit = true, ViewportCommand::Maximized(on) => { self.window_maximized = on; @@ -293,6 +367,151 @@ impl Headless { mod tests { use super::*; + #[test] + fn screenshot_dimensions_are_checked_before_layout() { + assert_eq!(viewport_pixels([1600.0, 1000.0], 2.0).unwrap(), [3200, 2000]); + assert_eq!(viewport_pixels([345.0, 200.0], 0.9).unwrap(), [311, 180]); + assert_eq!(viewport_pixels([8000.0, 8000.0], 1.0).unwrap(), [8000, 8000]); + for (size, scale) in [ + ([f32::NAN, 10.0], 1.0), + ([10.0, f32::INFINITY], 1.0), + ([10.0, 10.0], f32::INFINITY), + ([10.0, 10.0], f32::NAN), + ([10.0, 10.0], 0.0), + ([0.0, 10.0], 1.0), + ([10.0, 10.0], 0.01), + ([9000.0, 9000.0], 1.0), + ([9000.0, 100.0], 2.0), + ] { + assert!(viewport_pixels(size, scale).is_err(), "{size:?} at {scale}"); + } + let raw = HeadlessView::raw_input(egui::vec2(f32::INFINITY, 10.0), 1.0, 0.0, vec![]); + assert_eq!(raw.screen_rect.unwrap().size(), egui::vec2(1600.0, 1000.0)); + } + + #[test] + fn ui_zoom_preserves_host_pixels_and_scales_layout() { + for (size, scale, expected) in + [(egui::vec2(400.0, 240.0), 1.0, [400, 240]), (egui::vec2(345.0, 200.0), 0.9, [311, 180]), (egui::vec2(400.0, 240.0), 2.0, [800, 480])] + { + let mut view = HeadlessView::new(); + for frame in 0..9 { + let key = match frame { + 1 => Some(egui::Key::Plus), + 5 => Some(egui::Key::Minus), + _ => None, + }; + let mut raw = HeadlessView::raw_input(size, scale, frame as f64 / 60.0, vec![]); + if let Some(key) = key { + let modifiers = egui::Modifiers { command: true, ..Default::default() }; + raw.events.push(egui::Event::ModifiersChanged(modifiers)); + raw.events.extend(vec![ + egui::Event::Key { key, physical_key: None, pressed: true, repeat: false, modifiers }, + egui::Event::Key { key, physical_key: None, pressed: false, repeat: false, modifiers }, + ]); + } + view.run(raw, |ui| { + ui.label("UI zoom changes layout, keeping the host viewport size"); + }); + assert_eq!(view.size_px(), expected, "frame {frame} at scale {scale}"); + assert_eq!(view.paint(&HashMap::new()).size, expected); + if frame == 4 { + assert!(view.ctx.zoom_factor() > 1.0); + assert!(view.ctx.content_rect().width() < size.x); + } + if frame == 8 { + assert!(view.ctx.zoom_factor() < 1.01); + } + } + } + } + + #[test] + fn shadow_screenshot_keeps_the_zoomed_hosts_size_and_scale() { + let mut app = LightcraftApp::new(lightcraft_engine::Session::new(), crate::Services::default()); + app.ui.settings.gpu = false; + let mut h = Headless::new(app, [480.0, 320.0], 2.0); + let t = Duration::from_secs(5); + assert_eq!(h.request("ui.key", json!({"key": "Plus", "cmd": true}), t)["ok"], true); + assert!(h.view.ctx.zoom_factor() > 1.0); + let main = h.view.ctx.clone(); + let screenshot = h.app.headless_screenshot(&main, true).unwrap(); + assert_eq!(screenshot.size, [960, 640]); + let shadow = h.app.shadow.as_ref().unwrap(); + assert!((shadow.ctx.pixels_per_point() - main.pixels_per_point()).abs() < 0.001); + assert!((shadow.ctx.content_rect().width() - main.content_rect().width()).abs() < 0.1); + } + + #[test] + fn control_resize_uses_effective_points_after_ui_zoom() { + for scale in [0.9, 1.0, 2.0] { + let mut app = LightcraftApp::new(lightcraft_engine::Session::new(), crate::Services::default()); + app.ui.settings.gpu = false; + let mut h = Headless::new(app, [480.0, 320.0], scale); + let t = Duration::from_secs(5); + assert_eq!(h.request("ui.key", json!({"key": "Plus", "cmd": true}), t)["ok"], true); + let native_size = resized_viewport([320.0, 240.0], scale, h.view.ctx.zoom_factor()).unwrap(); + let expected = native_viewport_pixels([native_size.x, native_size.y], scale).unwrap(); + assert_eq!(h.request("ui.resize", json!({"width": 320, "height": 240}), t)["ok"], true); + assert_eq!(h.view.size_px(), expected); + assert!((h.view.ctx.content_rect().width() - 320.0).abs() < 0.1); + let main = h.view.ctx.clone(); + assert_eq!(h.app.headless_screenshot(&main, true).unwrap().size, expected); + } + } + + #[test] + fn control_resize_keeps_a_minimum_width_after_zooming_out() { + let mut app = LightcraftApp::new(lightcraft_engine::Session::new(), crate::Services::default()); + app.ui.settings.gpu = false; + let mut h = Headless::new(app, [480.0, 320.0], 1.0); + h.view.ctx.set_zoom_factor(0.5); + let t = Duration::from_secs(5); + assert_eq!(h.request("ui.resize", json!({"width": 1, "height": 120}), t)["ok"], true); + assert_eq!(h.size, egui::vec2(0.5, 60.0)); + assert_eq!(h.view.size_px(), [1, 60]); + assert!((h.view.ctx.content_rect().width() - 1.0).abs() < 0.01); + let main = h.view.ctx.clone(); + assert_eq!(h.app.headless_screenshot(&main, true).unwrap().size, [1, 60]); + } + + #[test] + fn control_resize_rounds_in_native_points_after_a_fractional_zoom_out() { + let mut app = LightcraftApp::new(lightcraft_engine::Session::new(), crate::Services::default()); + app.ui.settings.gpu = false; + let mut h = Headless::new(app, [480.0, 320.0], 0.9); + let t = Duration::from_secs(5); + assert_eq!(h.request("ui.key", json!({"key": "Minus", "cmd": true}), t)["ok"], true); + assert!((h.view.ctx.zoom_factor() - 0.9).abs() < 0.001); + assert_eq!(h.request("ui.resize", json!({"width": 150, "height": 150}), t)["ok"], true); + assert_eq!(h.size, egui::vec2(135.0, 135.0)); + assert_eq!(h.view.size_px(), [122, 122]); + assert!((h.view.ctx.content_rect().width() - 150.0).abs() < 0.01); + let main = h.view.ctx.clone(); + assert_eq!(h.app.headless_screenshot(&main, true).unwrap().size, [122, 122]); + } + + #[test] + fn invalid_control_resize_preserves_the_headless_viewport() { + let app = LightcraftApp::new(lightcraft_engine::Session::new(), crate::Services::default()); + let mut h = Headless::new(app, [480.0, 320.0], 2.0); + h.app.ui.settings.gpu = false; + h.step(); + let size = h.size; + let rejected = h.request("ui.resize", json!({"width": 9000, "height": 8000}), Duration::from_secs(5)); + assert_eq!(rejected["ok"], false, "{rejected}"); + assert_eq!(h.size, size); + for bad in [json!({"width": "320", "height": 240}), json!({"width": 320, "height": null}), json!({"width": true})] { + let rejected = h.request("ui.resize", bad, Duration::from_secs(5)); + assert_eq!(rejected["ok"], false, "{rejected}"); + assert_eq!(h.size, size); + } + let accepted = h.request("ui.resize", json!({"width": 320, "height": 240}), Duration::from_secs(5)); + assert_eq!(accepted["ok"], true, "{accepted}"); + assert_eq!(h.size, egui::vec2(320.0, 240.0)); + assert_eq!(h.view.size_px(), [640, 480]); + } + /// Generous: renders are slow when the machine is loaded (parallel builds), and a timed-out /// settle would show a half-rendered UI. const SETTLE: Duration = Duration::from_secs(120); @@ -401,6 +620,494 @@ mod tests { h.settle(SETTLE); } + /// Naming a face in the loupe: point at an unnamed face, click its "Add name" label, type, Enter. Needs the + /// recognition runtime only for the background indexer, so the naming itself is tested whatever the build. + #[test] + fn naming_a_face_in_the_loupe() { + use crate::state::NameEdit; + use lightcraft_catalog::Op; + let mut h = demo([1400.0, 900.0]); + let t = Duration::from_secs(10); + let id = h.app.session.active().unwrap(); + let mut meta = h.app.session.catalog.photo(id).unwrap().meta.clone(); + meta.regions = vec![lightcraft_meta::Region { + rect: lightcraft_geom::Rect { x0: 0.3, y0: 0.2, x1: 0.55, y1: 0.6 }, + kind: lightcraft_meta::RegionKind::Face, + name: None, + description: None, + }]; + h.app.session.commit("setup", Op::SetMeta { id, meta: Box::new(meta) }).unwrap(); + h.request("ui.set", json!({"view": "detail", "right": "none"}), t); + h.settle(SETTLE); + // pointing at the box shows the invitation; clicking it opens the name box + h.request("ui.pointer", json!({"events": [{"kind": "move", "x": 0.42, "y": 0.4}]}), t); + h.step(); + h.step(); + let r = h.request("ui.clickWidget", json!({"id": "regionLabel:0"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + assert!(matches!(&h.app.ui.name_edit, Some(NameEdit { index: 0, .. })), "{:?}", h.app.ui.name_edit); + // Escape closes it without naming anything + let r = h.request("ui.key", json!({"key": "escape"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + h.step(); + assert!(h.app.ui.name_edit.is_none(), "{:?}", h.app.ui.name_edit); + assert_eq!(h.app.session.catalog.photo(id).unwrap().meta.regions[0].name, None); + // open it again, type a name, Enter: the face is named, in one undo step + h.request("ui.pointer", json!({"events": [{"kind": "move", "x": 0.42, "y": 0.4}]}), t); + h.step(); + h.step(); + h.request("ui.clickWidget", json!({"id": "regionLabel:0"}), t); + h.step(); + h.step(); + let undo = h.app.session.undo.len(); + h.request("ui.text", json!({"text": "Ann"}), t); + h.step(); + h.request("ui.key", json!({"key": "enter"}), t); + h.step(); + h.step(); + assert_eq!(h.app.session.catalog.photo(id).unwrap().meta.regions[0].name.as_deref(), Some("Ann")); + assert_eq!(h.app.session.undo.len(), undo + 1); + assert!(h.app.ui.name_edit.is_none()); + } + + /// Adding a face model: the dialog shows the file's terms, the model is installed only once they are + /// accepted, and a file LightCraft cannot use only says why. + #[test] + fn adding_a_face_model_shows_its_terms_and_installs_only_once_accepted() { + use crate::state::Dialog; + let mut h = demo([1400.0, 900.0]); + let t = Duration::from_secs(10); + let dir = std::env::temp_dir().join(format!("lc-ui-facemodel-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + h.app.session.face_models_dir = Some(dir.join("models")); + let model = dir.join("Mine.onnx"); + std::fs::write(&model, lightcraft_faces::synthetic::tiny_embedder_model(512)).unwrap(); + let open = |h: &mut Headless, path: &std::path::Path| { + h.request("engine.execute", json!({"command": "dialog.faceModel", "params": {"path": path.to_string_lossy()}}), t) + }; + + // the dialog opens for the chosen file, with nothing accepted and nothing installed + assert_eq!(open(&mut h, &model)["ok"], true); + h.settle(SETTLE); + h.step(); + let Some(dlg) = h.app.ui.dialog.clone() else { panic!("no dialog") }; + let Dialog::FaceModel { info, accepted, .. } = &dlg else { panic!("{dlg:?}") }; + assert!(!accepted); + assert_eq!(info["kind"], "draft"); + assert!(!dir.join("models").exists()); + // OK does nothing yet + assert!(crate::panels::dialogs::confirm_dialog(&mut h.app, &dlg).is_err()); + assert!(!dir.join("models").exists()); + // ticking the box and confirming installs it + let r = h.request("ui.clickWidget", json!({"id": "check:faceModel.accept"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + assert!(matches!(&h.app.ui.dialog, Some(Dialog::FaceModel { accepted: true, .. }))); + let r = h.request("ui.dialog.confirm", json!({}), t); + assert_eq!(r["ok"], true, "{r}"); + let listed = h.app.run("faces.models.list", json!({})).unwrap(); + assert!(listed["models"].as_array().unwrap().iter().any(|m| m["installed"] == true && m["known"] == false), "{listed}"); + // installed, in use and recognition on: back in Settings ▸ Faces, where the scan can be watched + assert_eq!(h.app.ui.dialog, Some(Dialog::Settings { tab: "faces".into() })); + assert_eq!(listed["enabled"], true); + h.app.ui.dialog = None; + + // a file that is not a usable model says why and offers no install + let junk = dir.join("junk.onnx"); + std::fs::write(&junk, b"not a model").unwrap(); + assert_eq!(open(&mut h, &junk)["ok"], true); + h.step(); + let Some(Dialog::FaceModel { info, .. }) = h.app.ui.dialog.clone() else { panic!("no dialog") }; + assert_eq!(info["kind"], "unsupported"); + // a path that does not exist is an error, not a dialog + h.app.ui.dialog = None; + assert_eq!(open(&mut h, &dir.join("nope.onnx"))["ok"], false); + assert!(h.app.ui.dialog.is_none()); + let _ = std::fs::remove_dir_all(&dir); + } + + /// Photo > Detect Faces without the detector offers the download: the terms first, nothing fetched before they are + /// accepted, and the model's row in Settings has its Download button. + #[test] + fn detect_faces_without_the_detector_offers_the_download_and_fetches_nothing_before_acceptance() { + use crate::state::Dialog; + let mut h = demo([1400.0, 900.0]); + let dir = std::env::temp_dir().join(format!("lc-ui-detector-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + h.app.session.face_models_dir = Some(dir.join("models")); + + let r = h.app.run("faces.detect", json!({})); + assert!(r.is_err(), "no detector yet"); + let Some(dlg) = h.app.ui.dialog.clone() else { panic!("the download was not offered") }; + let Dialog::FaceModel { info, accepted, .. } = &dlg else { panic!("{dlg:?}") }; + assert!(!accepted); + assert_eq!((info["download"].as_str(), info["host"].as_str()), (Some("yunet-2023mar"), Some("github.com"))); + assert_eq!(info["model"]["licence"]["name"], "MIT"); + // OK without accepting the terms starts nothing + assert!(crate::panels::dialogs::confirm_dialog(&mut h.app, &dlg).is_err()); + assert!(h.app.session.face_downloads.snapshot().is_empty()); + assert!(!dir.join("models").join(".downloads").exists()); + assert!(h.app.caches.faces_dl_watch.is_empty()); + + // Settings > Faces shows the detector with a Download button + h.app.ui.dialog = Some(Dialog::Settings { tab: "faces".into() }); + h.step(); + // (the first request switches the widget list on, the next frame fills it) + h.request("ui.widgets", json!({}), Duration::from_secs(10)); + h.step(); + let w = h.request("ui.widgets", json!({"filter": "faces:"}), Duration::from_secs(10)); + assert!(w.to_string().contains("faces:download:yunet-2023mar"), "{w}"); + let _ = std::fs::remove_dir_all(&dir); + } + + /// Pressing Download in Settings shows the model's terms first and fetches nothing until they are accepted; a build + /// that cannot run recognition models offers no download. The licence dialog replaces Settings, which it was opened from. + #[test] + fn download_shows_the_terms_first_and_fetches_nothing_until_accepted() { + use crate::state::Dialog; + let mut h = demo([1400.0, 900.0]); + let t = Duration::from_secs(10); + let dir = std::env::temp_dir().join(format!("lc-ui-facedl-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + h.app.session.face_models_dir = Some(dir.join("models")); + h.request("engine.execute", json!({"command": "app.settings", "params": {"tab": "faces"}}), t); + h.settle(SETTLE); + h.step(); + let runtime = h.app.session.execute("faces.models.list", &json!({})).unwrap()["runtime"] == true; + let r = h.request("ui.clickWidget", json!({"id": "faces:download:sface-2021dec"}), t); + if !runtime { + assert_eq!(r["ok"], false, "a build that cannot run the model offers no download: {r}"); + let _ = std::fs::remove_dir_all(&dir); + return; + } + assert_eq!(r["ok"], true, "{r}"); + h.step(); + h.step(); + let Some(dlg) = h.app.ui.dialog.clone() else { panic!("no dialog") }; + let Dialog::FaceModel { path, info, accepted } = &dlg else { panic!("the terms did not replace Settings: {dlg:?}") }; + assert_eq!((info["download"].as_str(), *accepted, path.as_str()), (Some("sface-2021dec"), false, "")); + assert_eq!(info["model"]["licence"]["commercial"], "unknown"); + // OK does nothing until the terms are accepted, and nothing has been fetched + assert!(crate::panels::dialogs::confirm_dialog(&mut h.app, &dlg).is_err()); + assert_eq!(h.app.session.execute("faces.models.downloads", &json!({})).unwrap()["downloads"], json!([])); + // cancelling leaves it that way + h.request("ui.key", json!({"key": "escape"}), t); + h.step(); + assert!(h.app.ui.dialog.is_none()); + assert_eq!(h.app.session.execute("faces.models.downloads", &json!({})).unwrap()["downloads"], json!([])); + let _ = std::fs::remove_dir_all(&dir); + } + + /// Until face recognition is set up, the People view and the loupe's name box say so and offer the next step: with no + /// model, a button that opens Settings ▸ Faces; with a model installed but recognition off, one that switches it on. + #[test] + fn people_and_the_name_box_offer_to_set_face_recognition_up() { + use crate::state::{Dialog, ViewMode}; + use lightcraft_catalog::Op; + let mut h = demo([1400.0, 900.0]); + let t = Duration::from_secs(10); + let dir = std::env::temp_dir().join(format!("lc-ui-facesetup-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + h.app.session.face_models_dir = Some(dir.join("models")); + let runtime = h.app.session.execute("faces.models.list", &json!({})).unwrap()["runtime"] == true; + let id = h.app.session.active().unwrap(); + let mut meta = h.app.session.catalog.photo(id).unwrap().meta.clone(); + meta.regions = vec![lightcraft_meta::Region { + rect: lightcraft_geom::Rect { x0: 0.3, y0: 0.2, x1: 0.55, y1: 0.6 }, + kind: lightcraft_meta::RegionKind::Face, + name: None, + description: None, + }]; + h.app.session.commit("setup", Op::SetMeta { id, meta: Box::new(meta) }).unwrap(); + let offers = |h: &mut Headless| h.request("ui.widgets", json!({}), t).to_string().contains("\"faces:setup\""); + + h.request("engine.execute", json!({"command": "view.people"}), t); + h.settle(SETTLE); + h.step(); + assert_eq!(h.app.ui.view, ViewMode::People); + if !runtime { + assert!(!offers(&mut h), "a build that cannot run recognition offers nothing"); + let _ = std::fs::remove_dir_all(&dir); + return; + } + // no model: the offer is there, and one click lands in Settings ▸ Faces + assert!(offers(&mut h)); + let r = h.request("ui.clickWidget", json!({"id": "faces:setup"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + h.step(); + assert_eq!(h.app.ui.dialog, Some(Dialog::Settings { tab: "faces".into() })); + h.app.ui.dialog = None; + + // the loupe's name box offers it too, and its button leaves for Settings with the box closed + h.request("ui.set", json!({"view": "detail", "right": "none"}), t); + h.settle(SETTLE); + h.request("ui.pointer", json!({"events": [{"kind": "move", "x": 0.42, "y": 0.4}]}), t); + h.step(); + h.step(); + h.request("ui.clickWidget", json!({"id": "regionLabel:0"}), t); + h.step(); + h.step(); + assert!(h.app.ui.name_edit.is_some() && offers(&mut h), "the name box offers the setup"); + let r = h.request("ui.clickWidget", json!({"id": "faces:setup"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + h.step(); + assert!(h.app.ui.name_edit.is_none()); + assert_eq!(h.app.ui.dialog, Some(Dialog::Settings { tab: "faces".into() })); + h.app.ui.dialog = None; + + // a model installed but recognition off: "Turn on" does it on the spot + let model = dir.join("Mine.onnx"); + std::fs::write(&model, lightcraft_faces::synthetic::tiny_embedder_model(512)).unwrap(); + let installed = h.app.run("faces.models.install", json!({"path": model.to_string_lossy(), "acknowledged": true, "activate": false})).unwrap(); + h.app.run("faces.models.select", json!({"id": installed["installed"]["id"]})).unwrap(); + h.app.caches.faces_epoch += 1; + h.request("engine.execute", json!({"command": "view.people"}), t); + h.settle(SETTLE); + h.step(); + assert!(offers(&mut h)); + assert_eq!(h.app.session.execute("faces.models.list", &json!({})).unwrap()["enabled"], false); + let r = h.request("ui.clickWidget", json!({"id": "faces:setup"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + h.step(); + assert_eq!(h.app.session.execute("faces.models.list", &json!({})).unwrap()["enabled"], true); + assert_eq!(h.app.ui.dialog, None, "nothing to go to Settings for"); + h.step(); + assert!(!offers(&mut h), "once it is on the offer is gone"); + h.settle(SETTLE); + let _ = std::fs::remove_dir_all(&dir); + } + + /// A click on a person's card opens their page, which shows only cropped faces (one per face, not per photo); a click + /// on one opens its photo; Back, the People button and Escape return to everyone. + #[test] + fn a_persons_page_shows_their_faces_and_back_returns_to_everyone() { + use crate::state::ViewMode; + use lightcraft_catalog::Op; + let mut h = demo([1400.0, 900.0]); + let t = Duration::from_secs(10); + let ids: Vec<_> = h.app.session.catalog.photos().map(|p| p.id).take(3).collect(); + let face = |x: f64, name: &str| lightcraft_meta::Region { + rect: lightcraft_geom::Rect { x0: x, y0: 0.2, x1: x + 0.25, y1: 0.6 }, + kind: lightcraft_meta::RegionKind::Face, + name: Some(name.to_string()), + description: None, + }; + // Jane Doe twice in the first photo and once in the second; John Roe in the third + for (id, regions) in [ + (ids[0], vec![face(0.1, "Jane Doe"), face(0.5, "jane doe")]), + (ids[1], vec![face(0.3, "Jane Doe")]), + (ids[2], vec![face(0.3, "John Roe")]), + ] { + let mut meta = h.app.session.catalog.photo(id).unwrap().meta.clone(); + meta.regions = regions; + h.app.session.commit("setup", Op::SetMeta { id, meta: Box::new(meta) }).unwrap(); + } + let open_people = |h: &mut Headless| { + h.request("engine.execute", json!({"command": "view.people"}), t); + h.settle(SETTLE); + h.step(); + }; + open_people(&mut h); + assert_eq!((h.app.ui.view, h.app.ui.person_page.clone()), (ViewMode::People, None)); + // the card opens the page: a face tile for each of the three faces, not a tile per photo + let r = h.request("ui.clickWidget", json!({"id": "person:Jane Doe"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + h.step(); + assert_eq!((h.app.ui.view, h.app.ui.person_page.as_deref()), (ViewMode::People, Some("Jane Doe"))); + let widgets = h.request("ui.widgets", json!({}), t).to_string(); + for (id, index) in [(ids[0], 0), (ids[0], 1), (ids[1], 0)] { + assert!(widgets.contains(&format!("\"person-face:{}:{index}\"", id.0)), "a tile for each face: {widgets}"); + } + assert!(!widgets.contains(&format!("\"person-face:{}:0\"", ids[2].0)), "someone else's face is not on the page"); + // a face opens its photo in the detail view + h.request("ui.clickWidget", json!({"id": format!("person-face:{}:0", ids[1].0)}), t); + h.step(); + h.step(); + assert_eq!((h.app.ui.view, h.app.session.active()), (ViewMode::Detail, Some(ids[1]))); + // Escape from that photo goes back to the page, not to the grid + h.request("ui.key", json!({"key": "escape"}), t); + h.step(); + assert_eq!((h.app.ui.view, h.app.ui.person_page.as_deref()), (ViewMode::People, Some("Jane Doe"))); + // the People button from a photo shows everyone, with the person last left one click away next to the title + h.request("ui.clickWidget", json!({"id": format!("person-face:{}:0", ids[1].0)}), t); + h.step(); + h.step(); + assert_eq!(h.app.ui.view, ViewMode::Detail); + open_people(&mut h); + assert_eq!((h.app.ui.view, h.app.ui.person_page.clone()), (ViewMode::People, None)); + h.step(); + let r = h.request("ui.clickWidget", json!({"id": "people:last"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + h.step(); + assert_eq!((h.app.ui.view, h.app.ui.person_page.as_deref()), (ViewMode::People, Some("Jane Doe"))); + // a photo reached any other way goes back to the grid on Escape + h.request("engine.execute", json!({"command": "view.detail"}), t); + h.request("ui.key", json!({"key": "escape"}), t); + h.step(); + assert_eq!(h.app.ui.view, ViewMode::PhotoGrid); + h.request("engine.execute", json!({"command": "view.person", "params": {"name": "Jane Doe"}}), t); + h.step(); + let r = h.request("ui.clickWidget", json!({"id": "person:back"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + h.step(); + assert_eq!(h.app.ui.person_page, None); + h.request("engine.execute", json!({"command": "view.person", "params": {"name": "Jane Doe"}}), t); + h.request("ui.key", json!({"key": "escape"}), t); + h.step(); + assert_eq!((h.app.ui.view, h.app.ui.person_page.clone()), (ViewMode::People, None), "Escape goes back to everyone"); + h.request("engine.execute", json!({"command": "view.person", "params": {"name": "John Roe"}}), t); + open_people(&mut h); + assert_eq!(h.app.ui.person_page, None, "the People button shows everyone"); + // a missing or empty name is refused + assert_eq!(h.request("engine.execute", json!({"command": "view.person", "params": {}}), t)["ok"], false); + assert_eq!(h.request("engine.execute", json!({"command": "view.person", "params": {"name": " "}}), t)["ok"], false); + } + + /// The People view lists the unnamed faces below the named people: select some (a click each, or all), type a name, + /// press Enter, and they are all named at once, in one undo step; the new person appears among the named. + #[test] + fn unnamed_faces_are_selected_and_named_together() { + use lightcraft_catalog::Op; + let mut h = demo([1400.0, 900.0]); + let t = Duration::from_secs(10); + let ids: Vec<_> = h.app.session.catalog.photos().map(|p| p.id).take(3).collect(); + let face = |x: f64, name: Option<&str>| lightcraft_meta::Region { + rect: lightcraft_geom::Rect { x0: x, y0: 0.2, x1: x + 0.2, y1: 0.55 }, + kind: lightcraft_meta::RegionKind::Face, + name: name.map(str::to_string), + description: None, + }; + for (id, regions) in [ + (ids[0], vec![face(0.1, Some("Jane Doe")), face(0.5, None)]), + (ids[1], vec![face(0.3, None), face(0.6, None)]), + (ids[2], vec![face(0.3, None)]), + ] { + let mut meta = h.app.session.catalog.photo(id).unwrap().meta.clone(); + meta.regions = regions; + h.app.session.commit("setup", Op::SetMeta { id, meta: Box::new(meta) }).unwrap(); + } + h.request("engine.execute", json!({"command": "view.people"}), t); + h.settle(SETTLE); + h.step(); + h.step(); + let tiles = |h: &mut Headless| h.request("ui.widgets", json!({"filter": "unnamed-face:"}), t)["result"].as_array().map_or(0, Vec::len); + // the named person is a card, and the four unnamed faces are tiles below + let r = h.request("ui.clickWidget", json!({"id": "person:Jane Doe"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.request("engine.execute", json!({"command": "view.people"}), t); + h.step(); + h.step(); + assert_eq!(tiles(&mut h), 4); + // select two (nothing is named by selecting), and the naming bar appears + for (id, i) in [(ids[1], 0), (ids[2], 0)] { + let r = h.request("ui.clickWidget", json!({"id": format!("unnamed-face:{}:{i}", id.0)}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + h.step(); + } + assert_eq!(h.app.ui.unnamed_selected.len(), 2); + let bar = h.request("ui.widgets", json!({"filter": "unnamed:"}), t).to_string(); + assert!(bar.contains("unnamed:name") && bar.contains("unnamed:clear"), "{bar}"); + assert!(h.request("ui.widgets", json!({"filter": "field:unnamedName"}), t).to_string().contains("field:unnamedName")); + // type a name and press Enter: both faces are named, in one undo step, and the selection is gone + let undo = h.app.session.undo.len(); + h.request("ui.text", json!({"text": "Ann Example"}), t); + h.step(); + h.request("ui.key", json!({"key": "enter"}), t); + h.step(); + h.step(); + let named = |h: &Headless, id: lightcraft_catalog::PhotoId, i: usize| h.app.session.catalog.photo(id).unwrap().meta.regions[i].name.clone(); + assert_eq!((named(&h, ids[1], 0), named(&h, ids[2], 0)), (Some("Ann Example".to_string()), Some("Ann Example".to_string()))); + assert_eq!(named(&h, ids[1], 1), None, "a face that was not selected stays unnamed"); + assert_eq!(h.app.session.undo.len(), undo + 1, "one step for both"); + assert!(h.app.ui.unnamed_selected.is_empty() && h.app.ui.unnamed_name.is_empty()); + h.settle(SETTLE); + h.step(); + assert_eq!(tiles(&mut h), 2, "the named faces left the unnamed list"); + assert_eq!(h.app.session.catalog.people().len(), 2, "and the new person is among the named"); + // select all, then clear: nothing is named + h.request("ui.clickWidget", json!({"id": "unnamed:selectAll"}), t); + h.step(); + h.step(); + assert_eq!(h.app.ui.unnamed_selected.len(), 2); + h.request("ui.clickWidget", json!({"id": "unnamed:clear"}), t); + h.step(); + assert!(h.app.ui.unnamed_selected.is_empty()); + // undo gives the two faces back + h.request("engine.execute", json!({"command": "edit.undo"}), t); + h.step(); + assert_eq!(named(&h, ids[1], 0), None); + } + + /// A screenful of faces larger than the picture cache's usual budget (96) is all kept: with a fixed budget the same few + /// tiles were evicted and re-requested every frame and stayed blank. + #[test] + fn a_screenful_of_small_faces_is_not_evicted_and_left_blank() { + use lightcraft_catalog::Op; + let mut h = demo([1500.0, 1000.0]); + let t = Duration::from_secs(20); + let ids: Vec<_> = h.app.session.catalog.photos().map(|p| p.id).collect(); + // seven unnamed faces in every photo of the demo library + for id in &ids { + let mut meta = h.app.session.catalog.photo(*id).unwrap().meta.clone(); + meta.regions = (0..7) + .map(|k| lightcraft_meta::Region { + rect: lightcraft_geom::Rect { x0: 0.05 + 0.12 * k as f64, y0: 0.2, x1: 0.15 + 0.12 * k as f64, y1: 0.45 }, + kind: lightcraft_meta::RegionKind::Face, + name: None, + description: None, + }) + .collect(); + h.app.session.commit("setup", Op::SetMeta { id: *id, meta: Box::new(meta) }).unwrap(); + } + // the smallest faces, so a lot of them fit on the screen + h.request("ui.set", json!({"thumbSize": 90.0, "view": "people"}), t); + h.settle(SETTLE); + h.step(); + let drawn = h.request("ui.widgets", json!({"filter": "unnamed-face:"}), t)["result"].as_array().map_or(0, Vec::len); + assert!(drawn > 96, "the test needs more faces on screen than the usual budget, got {drawn}"); + // the pictures are made and kept, well beyond the old budget (a few demo photos share a scene, so two faces can share a + // picture: the count is a little under the number of tiles) + let started = Instant::now(); + while h.app.renderer.variant_textures() < drawn - 20 && started.elapsed() < Duration::from_secs(20) { + h.step(); + std::thread::sleep(Duration::from_millis(20)); + } + let kept = h.app.renderer.variant_textures(); + assert!(kept > 96 + 30, "{kept} pictures kept for {drawn} faces on screen"); + // and they stay: many more frames, and none is evicted to be asked for again + for _ in 0..30 { + h.step(); + } + assert!(h.app.renderer.variant_textures() >= kept, "{} pictures kept after more frames, {kept} before", h.app.renderer.variant_textures()); + } + + /// `ui.inspect` says how hard the face scan is allowed to work and whether the window counts as in front, so a slow + /// scan can be told from a stuck one. + #[test] + fn inspect_reports_the_face_scan_pace() { + let mut h = demo([1000.0, 700.0]); + let t = Duration::from_secs(10); + h.settle(SETTLE); + h.step(); + let r = h.request("ui.inspect", json!({}), t); + let scan = &r["result"]["faceScan"]; + assert!(["pause", "light", "normal", "full"].contains(&scan["pace"].as_str().unwrap_or("")), "{scan}"); + assert!(scan["focused"].is_boolean() && scan["pending"].is_u64() && scan["indexed"].is_u64(), "{scan}"); + } + /// Profile browser: live variant thumbnails, hover previews in the loupe without touching the /// photo or its history, click applies, the star toggles the favourite. #[test] diff --git a/crates/ui-egui/src/lib.rs b/crates/ui-egui/src/lib.rs index aeab2794..ea82b0d7 100644 --- a/crates/ui-egui/src/lib.rs +++ b/crates/ui-egui/src/lib.rs @@ -106,6 +106,8 @@ pub struct Services { pub picker: Option, /// Show an open dialog for photos; returns paths. pub pick_files: Option, + /// Open dialog for a face model (`.onnx`; Settings ▸ Faces ▸ Add a model file…). + pub pick_model_file: Option, /// Open dialog for preset files (`.lcpreset`, `.xmp`, `.lrtemplate`, `.zip`, `.dng`, Luminar `.lmp` / `.mplumpack`). pub pick_preset_files: Option, /// Open dialog for a GPS track log (`.gpx`; Photo ▸ Auto-Tag from Tracklog…). @@ -342,9 +344,17 @@ impl LightcraftApp { let now = self.tasks.repaint.as_ref().map(|ctx| ctx.input(|i| i.time)).unwrap_or(self.last_time); self.ui.toast = Some((text, now + TOAST_SECONDS, label)); } - if let Err(e) = &r { - log::warn!("{id}: {e}"); - self.ui.status = e.clone(); + match &r { + Err(e) => { + log::warn!("{id}: {e}"); + self.ui.status = e.clone(); + // no detector yet: offer it (its terms first) instead of only saying so + if id == "faces.detect" && e.contains("Settings > Faces") { + panels::faces::offer_detector(self); + } + } + Ok(v) if id == "faces.detect" => self.ui.status = detect_summary(v), + Ok(_) => {} } r } @@ -606,7 +616,10 @@ impl LightcraftApp { let mut view = self.shadow.take().unwrap_or_default(); let size = main.input(|i| i.content_rect()).size(); let size = if size.x >= 1.0 && size.y >= 1.0 { size } else { egui::vec2(1600.0, 1000.0) }; - let ppp = main.pixels_per_point(); + // Feed native viewport points; the shadow context applies the copied UI zoom. + let zoom = main.zoom_factor(); + let size = size * zoom; + let ppp = main.pixels_per_point() / zoom; let time = main.input(|i| i.time); if view.frames() == 0 { // warm-up pass: activates our fonts (pending font definitions live in `Memory`, which @@ -708,6 +721,7 @@ impl LightcraftApp { self.preview_build_status(ctx); self.save_status(ctx); self.slideshow_tick(ctx); + panels::faces::pump(self, ctx); // back from an external editor: pick up the files it saved let focused = ctx.input(|i| i.focused); if focused && !self.ui.was_focused && !self.ui.external_edits.is_empty() { @@ -751,6 +765,7 @@ impl LightcraftApp { ctx.request_repaint_after(std::time::Duration::from_secs(3)); } self.session.persist_if_dirty(); + panels::faces::pump(self, ctx); self.collect_screenshots(ctx); self.issue_screenshots(ctx); if self.fonts_ready { @@ -761,6 +776,13 @@ impl LightcraftApp { let dropped: Vec = ctx.input(|i| i.raw.dropped_files.iter().map(|f| f.path().to_string_lossy().to_string()).filter(|p| !p.is_empty()).collect()); // preset files import as presets, everything else as photos + // a dropped face model opens its licence dialog; presets and photos as before + let (models, dropped): (Vec, Vec) = dropped.into_iter().partition(|p| is_model_file(p)); + if let Some(model) = models.first() + && let Err(e) = panels::faces::open_dialog(self, model) + { + self.toast(ctx, e); + } let (presets, photos): (Vec, Vec) = dropped.into_iter().partition(|p| is_preset_file(p)); if !presets.is_empty() { let _ = self.run("file.importPresets", serde_json::json!({"paths": presets})); @@ -956,6 +978,25 @@ pub fn is_bw(d: &lightcraft_develop::DevelopSettings) -> bool { d.treatment == lightcraft_develop::Treatment::Bw || d.profile.id == "lc.mono" || d.profile.id.starts_with("lc.bw.") } +/// "Found 3 faces in 2 photos" for a `faces.detect` result. +fn detect_summary(v: &Value) -> String { + let photos = v["photos"].as_array().map(Vec::as_slice).unwrap_or_default(); + let counts: Vec = photos.iter().map(|p| p["faces"].as_array().map_or(0, Vec::len)).collect(); + let (faces, with) = (counts.iter().sum::(), counts.iter().filter(|n| **n > 0).count()); + let plural = |n: usize, w: &str| format!("{n} {w}{}", if n == 1 { "" } else { "s" }); + match (faces, photos.len()) { + (0, 1) => "No faces found".into(), + (0, n) => format!("No faces found in {}", plural(n, "photo")), + (f, 1) => format!("Found {}", plural(f, "face")), + (f, _) => format!("Found {} in {}", plural(f, "face"), plural(with, "photo")), + } +} + +/// A dropped face model file (`.onnx`). +pub fn is_model_file(path: &str) -> bool { + std::path::Path::new(path).extension().is_some_and(|e| e.eq_ignore_ascii_case("onnx")) +} + /// Files dropped on the window that are presets rather than photos. pub fn is_preset_file(path: &str) -> bool { let ext = std::path::Path::new(path).extension().map(|e| e.to_string_lossy().to_ascii_lowercase()).unwrap_or_default(); @@ -964,6 +1005,28 @@ pub fn is_preset_file(path: &str) -> bool { #[cfg(test)] mod drop_tests { + #[test] + fn detection_results_are_summarised_in_words() { + let r = + |counts: &[usize]| serde_json::json!({"photos": counts.iter().map(|n| serde_json::json!({"faces": vec![0; *n]})).collect::>()}); + assert_eq!(super::detect_summary(&r(&[0])), "No faces found"); + assert_eq!(super::detect_summary(&r(&[1])), "Found 1 face"); + assert_eq!(super::detect_summary(&r(&[3])), "Found 3 faces"); + assert_eq!(super::detect_summary(&r(&[0, 0])), "No faces found in 2 photos"); + assert_eq!(super::detect_summary(&r(&[2, 0, 1])), "Found 3 faces in 2 photos"); + assert_eq!(super::detect_summary(&serde_json::json!({})), "No faces found in 0 photos"); + } + + #[test] + fn dropped_face_models_are_told_apart() { + for p in ["/a/model.onnx", "/a/dir/M.ONNX"] { + assert!(super::is_model_file(p), "{p}"); + } + for p in ["/a/model.onnx.jpg", "/a/onnx", "/a/b.png"] { + assert!(!super::is_model_file(p), "{p}"); + } + } + #[test] fn dropped_presets_are_told_apart_from_photos() { for p in ["/a/Look.lrtemplate", "/a/b.XMP", "/a/pack.zip", "/a/x.lcpreset", "/a/Magic Hour.mplumpack", "/a/Pop.lmp", "/a/Bundle.LMP"] { @@ -988,6 +1051,31 @@ pub struct Caches { album_counts: Option<(u64, std::sync::Arc>)>, /// How often the album counts were recomputed (tests check that unchanged frames don't). pub album_count_scans: usize, + /// Bumped when a face model is installed, removed or chosen, so Settings re-reads the list at once. + pub faces_epoch: u64, + /// Face model downloads the user started, followed until installed (see `panels::faces::pump`). + pub faces_dl_watch: Vec, + /// Whether the background face indexer is running, how many faces it has embedded, how many photos are left, and + /// when to ask it again. + pub faces_active: bool, + pub faces_indexed: u64, + pub faces_pending: u64, + pub faces_next_pump: f64, + /// When the user last dragged, typed or scrolled, and when they last moved the pointer (egui time). + pub last_input: f64, + /// The pace the face scan was last given, and whether the window was in front then (for `ui.inspect`). + pub faces_pace: &'static str, + pub faces_in_front: bool, + pub last_move: f64, + /// Name suggestions for the photo in the loupe: (photo, catalog revision, faces indexed, suggestions by region). + pub face_hints: Option<(u64, u64, u64, std::sync::Arc)>, + /// The open person page: (name, catalog revision, faces indexed, when it was asked for, the page). + pub person_page: Option<(String, u64, u64, f64, std::sync::Arc)>, + /// The unnamed faces: (catalog revision, faces indexed, when it was asked for, the list). + pub unnamed: Option<(u64, u64, f64, std::sync::Arc)>, + /// The most photos the face scan has had left at once since it last finished (the progress bar's whole). + pub faces_peak: u64, + person_names: Option<(u64, std::sync::Arc>)>, /// The grid's date runs, layout and indexes (by the visible list's generation). pub grid: panels::grid::GridCache, /// What the grid did on its frames (benchmarks and tests check unchanged frames stay cheap). @@ -1076,6 +1164,17 @@ impl Caches { self.counts = Some((cat.revision, c)); c } + /// Everyone named on a face in the library (for completing a name as it is typed). + pub fn person_names(&mut self, cat: &lightcraft_catalog::Catalog) -> std::sync::Arc> { + match &self.person_names { + Some((r, v)) if *r == cat.revision => v.clone(), + _ => { + let v = std::sync::Arc::new(cat.people().into_iter().map(|p| p.name).collect::>()); + self.person_names = Some((cat.revision, v.clone())); + v + } + } + } /// The By Date tree. pub fn date_groups(&mut self, cat: &lightcraft_catalog::Catalog) -> std::sync::Arc> { match &self.date_groups { diff --git a/crates/ui-egui/src/menubar.rs b/crates/ui-egui/src/menubar.rs index 6ecdb169..7d964856 100644 --- a/crates/ui-egui/src/menubar.rs +++ b/crates/ui-egui/src/menubar.rs @@ -753,10 +753,29 @@ pub fn show_in_window(app: &mut LightcraftApp, ui: &mut egui::Ui, max_width: f32 if total <= max_width { let saved = ui.spacing().item_spacing.x; ui.spacing_mut().item_spacing.x = TITLE_GAP; - for (title, items) in &bar { - let r = ui.add(egui::Button::new(egui::RichText::new(crate::i18n::tr(title)).font(font.clone()).color(t.text_label)).frame(false)); - crate::widgets::register(ui.ctx(), format!("menu:{title}"), r.rect); - egui::Popup::menu(&r).show(|ui| crate::menu_level::level(ui, 1, bar_bottom, |ui| nodes_ui(ui, items, mac, &mut clicked, 1, bar_bottom))); + // All titles first, then their popups, so a title the pointer moves onto can take over + // from the open menu before either is drawn (no frame with two menus or none). + let titles: Vec = bar + .iter() + .map(|(title, _)| { + let r = ui.add(egui::Button::new(egui::RichText::new(crate::i18n::tr(title)).font(font.clone()).color(t.text_label)).frame(false)); + crate::widgets::register(ui.ctx(), format!("menu:{title}"), r.rect); + r + }) + .collect(); + let confirmed = switch_open_title(ui.ctx(), &titles); + for ((_, items), r) in bar.iter().zip(&titles) { + // (opening and closing by click is done in `switch_open_title`) + // (the click confirming a menu a hover opened must not close it again) + let close = if confirmed == Some(egui::Popup::default_response_id(r)) { + egui::PopupCloseBehavior::IgnoreClicks + } else { + egui::PopupCloseBehavior::CloseOnClick + }; + egui::Popup::menu(r) + .close_behavior(close) + .open_memory(None) + .show(|ui| crate::menu_level::level(ui, 1, bar_bottom, |ui| nodes_ui(ui, items, mac, &mut clicked, 1, bar_bottom))); } ui.spacing_mut().item_spacing.x = saved; } else { @@ -782,6 +801,47 @@ pub fn show_in_window(app: &mut LightcraftApp, ui: &mut egui::Ui, max_width: f32 ui.cursor().left() - start } +/// Native menu-bar behaviour for the top-level titles: while one menu is open, resting the pointer +/// on another title closes it and opens that one, and a click on another title opens it in that +/// same click (instead of only closing the open menu as a click outside). Without an open menu +/// hovering does nothing; clicking a title toggles its menu. egui keeps one popup open at a time. +/// Returns the menu whose title was just clicked right after a hover opened it. +fn switch_open_title(ctx: &egui::Context, titles: &[egui::Response]) -> Option { + let popup = egui::Popup::default_response_id; + let key = egui::Id::new("lc-menubar-hover-opened"); + let open = titles.iter().position(|r| egui::Popup::is_id_open(ctx, popup(r))); + // the menu a hover opened, until it closes or is clicked + let mut confirmed = None; + let mut hover_opened: Option = ctx.data(|d| d.get_temp(key)).filter(|id| open.is_some_and(|o| popup(&titles[o]) == *id)); + for (i, r) in titles.iter().enumerate() { + if r.clicked() { + // the click that follows a hover switch confirms the menu the hover opened + if hover_opened == Some(popup(r)) { + confirmed = Some(popup(r)); + } else { + egui::Popup::toggle_id(ctx, popup(r)); + } + hover_opened = None; + } else if r.hovered() && open.is_some_and(|o| o != i) { + egui::Popup::open_id(ctx, popup(r)); + hover_opened = Some(popup(r)); + // a pointer move alone doesn't schedule the frame that draws the new menu + ctx.request_repaint(); + } + } + // (leaving the titles ends it: a later click on the open title closes the menu) + let hover_opened = hover_opened.filter(|_| titles.iter().any(|r| r.hovered())); + ctx.data_mut(|d| match hover_opened { + Some(id) => { + d.insert_temp(key, id); + } + None => { + d.remove::(key); + } + }); + confirmed +} + /// A submenu row showing `text` at `depth`, whose rows `children` draws one level deeper, /// bounded like every level. The row is registered as `menusub: