Merge pull request #439 from storytold/claude/integrate-batch-14
Integrate #429, #426, #422, #395, #172, #173, #174, #335, #427
This commit is contained in:
commit
c4c96ecabb
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
11
ROADMAP.md
11
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
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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<Value> = 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::<f64>().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<Value> = 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<Value> = 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());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -209,6 +209,8 @@ fn services(originals: Originals, backend: Option<Backend>, 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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 }
|
||||
|
|
|
|||
|
|
@ -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 (`<config>/models`).
|
||||
fn open_session(in_memory: bool, dir: Option<std::path::PathBuf>, seed_demo: bool) -> (Session, Option<LibraryProblem>) {
|
||||
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<std::path::PathBuf>, seed_demo: bool) -> (Session, Option<LibraryProblem>) {
|
||||
fn open_library_session(in_memory: bool, dir: Option<std::path::PathBuf>, seed_demo: bool) -> (Session, Option<LibraryProblem>) {
|
||||
if in_memory {
|
||||
return (if seed_demo { Session::with_demo() } else { Session::new() }.with_fs().with_system_clock(), None);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -21,3 +21,6 @@
|
|||
real_name = "Brandon Thomas"
|
||||
display_name = "Brandon Thomas"
|
||||
sync_github_name = false
|
||||
|
||||
[people.sanyamsingh001]
|
||||
display_name = "sanyamsingh001"
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
||||
|
|
|
|||
|
|
@ -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<Decoded> {
|
||||
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<Decoded> {
|
||||
// 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:?}");
|
||||
}
|
||||
}
|
||||
|
|
@ -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")),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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<Arc<Detector>, String> {
|
||||
static LOADED: Mutex<Option<(PathBuf, Arc<Detector>)>> = 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::<Vec<_>>(),
|
||||
})
|
||||
}
|
||||
|
||||
/// 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<u8>, 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<u8> = 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<Value> {
|
||||
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::<Vec<_>>()}));
|
||||
}
|
||||
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<CommandSpec> {
|
||||
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), &[]));
|
||||
}
|
||||
}
|
||||
|
|
@ -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<String>,
|
||||
}
|
||||
|
||||
fn fail(what: &str, e: impl std::fmt::Display) -> EngineError {
|
||||
EngineError::Other(format!("{what}: {e}"))
|
||||
}
|
||||
|
||||
fn models_dir(s: &Session, cmd: &str) -> Result<PathBuf> {
|
||||
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<Vec<u8>> {
|
||||
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<Installed> {
|
||||
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<ModelManifest> {
|
||||
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> {
|
||||
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<Inspection> {
|
||||
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<Value> {
|
||||
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<Value> = 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<Value> {
|
||||
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<Value> {
|
||||
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<Value> {
|
||||
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<Value> {
|
||||
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<Value> {
|
||||
finish_downloads(s);
|
||||
let all: Vec<Value> = 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<Value> {
|
||||
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<Value, String> {
|
||||
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<Value, String> {
|
||||
Ok(Value::Null)
|
||||
}
|
||||
|
||||
/// `faces.models.test {id}`: run an installed model's self-test again.
|
||||
fn test(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
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<Value> {
|
||||
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<Value> {
|
||||
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<Value> {
|
||||
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<CommandSpec> {
|
||||
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),
|
||||
]
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
|
|
@ -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
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<PathBuf> {
|
||||
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 `<config>/models`.
|
||||
pub fn default_face_models_dir() -> Option<PathBuf> {
|
||||
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
|
||||
}
|
||||
}
|
||||
|
|
@ -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<State>,
|
||||
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<String, Arc<Shared>>,
|
||||
}
|
||||
|
||||
impl Downloads {
|
||||
/// Start fetching `spec` into `<models dir>/.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<u8>) -> 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<u8> = (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);
|
||||
}
|
||||
}
|
||||
|
|
@ -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-<model id>.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 (`<model id>@<sha256>`), 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 {
|
||||
/// `<model id>@<sha256>` of the model the embeddings belong to.
|
||||
pub tag: String,
|
||||
pub dim: usize,
|
||||
entries: HashMap<Key, Entry>,
|
||||
unsaved: Vec<(Key, Entry)>,
|
||||
/// Faces that could not be embedded this session (so they are not retried on every call).
|
||||
skipped: std::collections::HashSet<Key>,
|
||||
/// The cache file on disk belongs to this model, so new records can be appended to it.
|
||||
file_matches: bool,
|
||||
pub path: Option<PathBuf>,
|
||||
}
|
||||
|
||||
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<Rect> {
|
||||
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<f32>) {
|
||||
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<f32>, 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<PathBuf>) {
|
||||
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<f32> = 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<u8> {
|
||||
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<u8> {
|
||||
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<u64>,
|
||||
unsaved: Vec<u64>,
|
||||
path: Option<PathBuf>,
|
||||
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<PathBuf>) {
|
||||
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<u8> {
|
||||
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<u8> = 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 {
|
||||
/// `<model id>@<sha256>` and the loaded model.
|
||||
pub embedder: Option<(String, Arc<lightcraft_faces::runtime::Embedder>)>,
|
||||
pub index: Index,
|
||||
/// When the models folder was last looked at (the per-frame pump trusts a look for a second).
|
||||
pub checked: Option<std::time::Instant>,
|
||||
/// Whether recognition was switched on when last looked up, and when.
|
||||
pub enabled_seen: Option<(std::time::Instant, bool)>,
|
||||
pub worker: Option<crate::faces_worker::Worker>,
|
||||
/// After a failure to load the chosen model, the background pump leaves it alone until then.
|
||||
pub retry_at: Option<std::time::Instant>,
|
||||
/// Photos the worker is busy with.
|
||||
pub in_flight: std::collections::HashSet<lightcraft_catalog::PhotoId>,
|
||||
/// Photos still to hand to the worker, and the catalog revision this list was made at.
|
||||
pub queue: Vec<lightcraft_catalog::PhotoId>,
|
||||
pub queue_stamp: Option<u64>,
|
||||
/// When the index was last written to its cache file by the background pump.
|
||||
pub saved_at: Option<std::time::Instant>,
|
||||
/// 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<u64>,
|
||||
}
|
||||
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
|
@ -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>) -> 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<usize> {
|
||||
std::env::var("LIGHTCRAFT_FACE_THREADS").ok().and_then(|v| v.trim().parse::<usize>().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>) -> 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 {
|
||||
/// `<model id>@<sha256>` 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<Rect>,
|
||||
/// 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<Arc<Detector>>,
|
||||
}
|
||||
|
||||
/// 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<Vec<f32>>, 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<Vec<Rect>>,
|
||||
}
|
||||
|
||||
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<Rect> {
|
||||
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<Rect>| todo.into_iter().map(|r| (r, None, r)).collect::<Vec<_>>();
|
||||
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<u8> = 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<Vec<f32>>, 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<Embedder>, 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<usize>,
|
||||
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<Option<rayon::ThreadPool>> = std::sync::OnceLock::new();
|
||||
static NORMAL: std::sync::OnceLock<Option<rayon::ThreadPool>> = std::sync::OnceLock::new();
|
||||
static FULL: std::sync::OnceLock<Option<rayon::ThreadPool>> = 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<Job>,
|
||||
done: Receiver<Done>,
|
||||
limit: Arc<Limit>,
|
||||
}
|
||||
|
||||
impl Worker {
|
||||
pub fn start(threads: usize) -> Option<Worker> {
|
||||
let (jobs, job_rx) = channel::<Job>();
|
||||
let (done_tx, done) = channel::<Done>();
|
||||
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<Embedder>, prepared: Prepared) -> bool {
|
||||
self.jobs.send((embedder, prepared)).is_ok()
|
||||
}
|
||||
|
||||
/// Everything finished since the last call.
|
||||
pub fn finished(&self) -> Vec<Done> {
|
||||
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:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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<Arc<lightcraft_pipeline::masks::Mattes>> {
|
||||
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:<year>:aux:<type>`. Other names: `None`.
|
||||
fn matte_kind(name: &str) -> Option<lightcraft_pipeline::masks::MatteKind> {
|
||||
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<u8>| {
|
||||
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);
|
||||
|
|
|
|||
|
|
@ -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<std::path::PathBuf>,
|
||||
/// Face model downloads started this session (the staged files wait in `<face_models_dir>/.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<Value>,
|
||||
/// 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;
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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::<SourceInfo>()
|
||||
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::<SourceInfo>() + 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<QuickJob> {
|
||||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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=<the .onnx file>` (what
|
||||
/// Settings > Faces downloads; nothing is committed). Without it those tests say so and pass.
|
||||
fn demo_with_yunet() -> Option<Session> {
|
||||
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);
|
||||
}
|
||||
|
|
@ -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=<the .onnx file>` (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<Region>) {
|
||||
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::<Vec<_>>();
|
||||
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::<f32>().sqrt();
|
||||
v.iter().map(|x| x / n).collect::<Vec<f32>>()
|
||||
};
|
||||
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<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::<f32>().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<f32>| 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);
|
||||
}
|
||||
|
|
@ -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 }
|
||||
|
|
@ -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<Similarity> {
|
||||
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::<f32>() / n, pts.iter().map(|p| p.1).sum::<f32>() / 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<Vec<u8>> {
|
||||
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<Vec<u8>> {
|
||||
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<Vec<u8>> {
|
||||
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");
|
||||
}
|
||||
}
|
||||
|
|
@ -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<Download> {
|
||||
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<Entry>,
|
||||
/// One plain sentence for each entry (or the file) that could not be used.
|
||||
pub errors: Vec<String>,
|
||||
}
|
||||
|
||||
#[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<Entry, String> {
|
||||
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::<String>()));
|
||||
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<Value>) -> 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::<String>::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::<Vec<_>>(), ["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<Value> = (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}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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<T> = std::result::Result<T, ProbeError>;
|
||||
|
||||
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<i64>),
|
||||
Floats(Vec<f32>),
|
||||
/// A kind we do not read (a tensor, a sub-graph…).
|
||||
Other,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Node {
|
||||
pub op: String,
|
||||
pub inputs: Vec<String>,
|
||||
pub outputs: Vec<String>,
|
||||
pub attrs: HashMap<String, Attr>,
|
||||
}
|
||||
|
||||
impl Node {
|
||||
pub fn int(&self, name: &str) -> Option<i64> {
|
||||
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<f32>),
|
||||
I64(Vec<i64>),
|
||||
}
|
||||
|
||||
/// A weight tensor.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct Weight {
|
||||
pub dims: Vec<usize>,
|
||||
pub data: Data,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Graph {
|
||||
pub nodes: Vec<Node>,
|
||||
pub weights: HashMap<String, Weight>,
|
||||
/// 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<usize>)>,
|
||||
pub outputs: Vec<String>,
|
||||
}
|
||||
|
||||
fn malformed<T>(why: &'static str) -> Result<T> {
|
||||
Err(ProbeError::Malformed(why))
|
||||
}
|
||||
|
||||
fn f32s(raw: &[u8]) -> Vec<f32> {
|
||||
raw.as_chunks::<4>().0.iter().map(|b| f32::from_le_bytes(*b)).collect()
|
||||
}
|
||||
|
||||
fn i64s(raw: &[u8]) -> Vec<i64> {
|
||||
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<Cursor<&[u8]>>, end: u64, out: &mut Vec<u64>) -> 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<Cursor<&[u8]>>, 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<u64>, Vec<f32>) = (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<Cursor<&[u8]>>, end: u64) -> Result<Node> {
|
||||
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<Cursor<&[u8]>>, end: u64, budget: &mut usize) -> Result<(String, Option<Weight>)> {
|
||||
let (mut name, mut dtype) = (String::new(), 0u64);
|
||||
let mut dims: Vec<u64> = Vec::new();
|
||||
let mut raw: Option<&[u8]> = None;
|
||||
let (mut float_data, mut int_data): (Vec<f32>, Vec<u64>) = (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<Cursor<&[u8]>>, end: u64) -> Result<(String, Vec<usize>)> {
|
||||
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::<Result<_>>()?;
|
||||
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<Graph> {
|
||||
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;
|
||||
|
|
@ -0,0 +1,117 @@
|
|||
use super::*;
|
||||
use crate::synthetic::{len_field, value_info_bytes, varint, varint_field};
|
||||
fn node() -> Vec<u8> {
|
||||
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<u8> {
|
||||
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<u8>]) -> Vec<u8> {
|
||||
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); }
|
||||
}
|
||||
|
|
@ -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<R: Read>(mut r: R) -> io::Result<String> {
|
||||
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<String> {
|
||||
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"));
|
||||
}
|
||||
}
|
||||
|
|
@ -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<ModelManifest> {
|
||||
vec![yunet(), sface(), auraface()]
|
||||
}
|
||||
|
||||
/// The known model whose file has this SHA-256 (lowercase hex).
|
||||
pub fn lookup(sha256: &str) -> Option<ModelManifest> {
|
||||
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<Download> {
|
||||
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 `<mirror>/<file name>`).
|
||||
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"));
|
||||
}
|
||||
}
|
||||
|
|
@ -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};
|
||||
|
|
@ -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<String>,
|
||||
/// 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<f32>,
|
||||
/// Detectors: lowest detection score kept.
|
||||
pub score: Option<f32>,
|
||||
/// Detectors: overlap above which two boxes are one face.
|
||||
pub nms_iou: Option<f32>,
|
||||
}
|
||||
|
||||
#[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<String>,
|
||||
/// Lowercase hex SHA-256 of the `.onnx` file, when known.
|
||||
#[serde(default)]
|
||||
pub sha256: Option<String>,
|
||||
#[serde(default)]
|
||||
pub size_bytes: Option<u64>,
|
||||
/// 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<f32>,
|
||||
}
|
||||
|
||||
/// 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<T>(field: &'static str, why: impl Into<String>) -> Result<T, ManifestError> {
|
||||
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<String>) -> 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<f32>, 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::<String>()));
|
||||
}
|
||||
}
|
||||
}
|
||||
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<ModelManifest, ManifestError> {
|
||||
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<dyn Fn(&mut ModelManifest)>)> = 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));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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<f32> {
|
||||
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<Ranked> {
|
||||
let mut people: Vec<(String, String, Vec<f32>)> = 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<Ranked> = 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::<f32>() / 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<Suggestion> {
|
||||
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<f32> {
|
||||
let n = v.iter().map(|x| x * x).sum::<f32>().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<_>>(), 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());
|
||||
}
|
||||
}
|
||||
|
|
@ -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<T> = std::result::Result<T, NetError>;
|
||||
|
||||
fn shape_err<T>(why: impl Into<String>) -> Result<T> {
|
||||
Err(NetError::Shape(why.into()))
|
||||
}
|
||||
|
||||
fn unsupported<T>(why: impl Into<String>) -> Result<T> {
|
||||
Err(NetError::Unsupported(why.into()))
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct Tensor {
|
||||
pub shape: Vec<usize>,
|
||||
pub data: Vec<f32>,
|
||||
}
|
||||
|
||||
impl Tensor {
|
||||
pub fn new(shape: Vec<usize>, data: Vec<f32>) -> Result<Tensor> {
|
||||
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<usize>) -> Result<Tensor> {
|
||||
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<Net> {
|
||||
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<Vec<Vec<usize>>> {
|
||||
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<Vec<(String, Tensor)>> {
|
||||
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<String, Rc<Tensor>> = HashMap::new();
|
||||
values.insert(input_name.clone(), Rc::new(input));
|
||||
for (i, node) in g.nodes.iter().enumerate() {
|
||||
let get = |idx: usize| -> Result<Rc<Tensor>> {
|
||||
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<Tensor> {
|
||||
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<Tensor> {
|
||||
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<Tensor> {
|
||||
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<Tensor> {
|
||||
let default_perm: Vec<i64> = (0..x.shape.len() as i64).rev().collect();
|
||||
let perm: Vec<usize> = node
|
||||
.ints("perm")
|
||||
.unwrap_or(&default_perm)
|
||||
.iter()
|
||||
.map(|p| usize::try_from(*p).map_err(|_| NetError::Shape("a negative permutation".into())))
|
||||
.collect::<Result<_>>()?;
|
||||
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<usize> = 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<usize> = 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<Tensor> {
|
||||
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<usize> = 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<Node>, 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<f32>) -> 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<Tensor> {
|
||||
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<f32> = (0..cout * cin_g * k * k).map(|i| ((i * 37 % 11) as f32 - 5.0) * 0.1).collect();
|
||||
let bias: Vec<f32> = (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;
|
||||
|
|
@ -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<f32>,
|
||||
product: Vec<f32>,
|
||||
}
|
||||
fn storage(v: &mut Vec<f32>, 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<Tensor> {
|
||||
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<Tensor> {
|
||||
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<Tensor> {
|
||||
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<Tensor> {
|
||||
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<Tensor> {
|
||||
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<Tensor> {
|
||||
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<Tensor> {
|
||||
let (c, h, w) = chw(x)?;
|
||||
Tensor::new(vec![1, c, 1, 1], x.data.chunks_exact(h * w).map(|v| v.iter().sum::<f32>() / (h * w) as f32).collect())
|
||||
}
|
||||
|
|
@ -0,0 +1,337 @@
|
|||
use super::*;
|
||||
use crate::graph::{Attr, Weight};
|
||||
fn weight(shape: &[usize], values: Vec<f32>) -> 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<f32>) -> 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::<f32>();
|
||||
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());
|
||||
}
|
||||
|
|
@ -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<usize> {
|
||||
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::<Option<Vec<_>>>().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<Vec<usize>> {
|
||||
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<usize> {
|
||||
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<Vec<Vec<usize>>> {
|
||||
elements(input)?;
|
||||
let mut shapes: HashMap<&str, Vec<usize>> = 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<i64> = 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()
|
||||
}
|
||||
|
|
@ -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<Dim>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq)]
|
||||
pub struct OnnxInfo {
|
||||
pub producer: String,
|
||||
/// The default-domain opset version.
|
||||
pub opset: Option<u64>,
|
||||
/// Real inputs only (weights that older exporters also list as inputs are left out).
|
||||
pub inputs: Vec<TensorInfo>,
|
||||
pub outputs: Vec<TensorInfo>,
|
||||
}
|
||||
|
||||
#[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<T> = std::result::Result<T, ProbeError>;
|
||||
|
||||
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<R: Read + Seek> Walker<'_, R> {
|
||||
pub(crate) fn pos(&mut self) -> Result<u64> {
|
||||
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<u8> {
|
||||
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<u64> {
|
||||
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<Option<(u32, Wire)>> {
|
||||
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<String> {
|
||||
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<Vec<u8>> {
|
||||
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<R: Read + Seek>(w: &mut Walker<R>, end: u64) -> Result<Option<Dim>> {
|
||||
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<R: Read + Seek>(w: &mut Walker<R>, end: u64) -> Result<Vec<Dim>> {
|
||||
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<R: Read + Seek>(w: &mut Walker<R>, 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<R: Read + Seek>(w: &mut Walker<R>, 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<R: Read + Seek>(w: &mut Walker<R>, end: u64) -> Result<TensorInfo> {
|
||||
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<R: Read + Seek>(w: &mut Walker<R>, end: u64) -> Result<Option<String>> {
|
||||
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<R: Read + Seek>(w: &mut Walker<R>, end: u64, info: &mut OnnxInfo, weights: &mut HashSet<String>) -> 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<R: Read + Seek>(w: &mut Walker<R>, end: u64) -> Result<Option<u64>> {
|
||||
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: Read + Seek>(r: &mut R) -> Result<OnnxInfo> {
|
||||
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<OnnxInfo> {
|
||||
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<std::result::Result<u64, &str>> = (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::<u8>(), 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::<u8>()), 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<usize> {
|
||||
panic!("an oversized model must not be read")
|
||||
}
|
||||
}
|
||||
impl std::io::Seek for Oversized {
|
||||
fn seek(&mut self, _: std::io::SeekFrom) -> std::io::Result<u64> {
|
||||
Ok(crate::manifest::MAX_MODEL_BYTES + 1)
|
||||
}
|
||||
}
|
||||
assert!(probe_reader(&mut Oversized).is_err());
|
||||
}
|
||||
}
|
||||
|
|
@ -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<Plan>,
|
||||
manifest: ModelManifest,
|
||||
width: usize,
|
||||
height: usize,
|
||||
dim: usize,
|
||||
load_ms: f64,
|
||||
}
|
||||
|
||||
fn guarded<T>(what: &str, f: impl FnOnce() -> Result<T, RuntimeError>) -> Result<T, RuntimeError> {
|
||||
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<Embedder, RuntimeError> {
|
||||
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::<usize>() != 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<Vec<f32>, 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::<f64>().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<u8> = (0..n)
|
||||
.map(|_| {
|
||||
seed ^= seed << 13;
|
||||
seed ^= seed >> 17;
|
||||
seed ^= seed << 5;
|
||||
(seed >> 8) as u8
|
||||
})
|
||||
.collect();
|
||||
let gradient: Vec<u8> = (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::<f32>();
|
||||
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::<f32>() - 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=<folder> 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<u8> = (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::<f32>() - 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);
|
||||
}
|
||||
}
|
||||
|
|
@ -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<String> },
|
||||
/// Not usable yet, and why.
|
||||
Unsupported(String),
|
||||
}
|
||||
|
||||
fn fixed(d: &Dim) -> Option<u64> {
|
||||
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 })
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -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<u8>) {
|
||||
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<u8>) {
|
||||
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<u8>) {
|
||||
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<u64, &str>]) -> Vec<u8> {
|
||||
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<u8> {
|
||||
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<u8> {
|
||||
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<u8> {
|
||||
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<u8> = 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<u8> {
|
||||
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<f32> = (0..dim * 3).map(|i| ((i * 37 % 19) as f32 - 9.0) / 20.0).collect();
|
||||
let b: Vec<f32> = (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<u8> {
|
||||
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<u8> {
|
||||
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<u8> {
|
||||
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
|
||||
}
|
||||
|
|
@ -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<NetError> 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<Detector, DetectError> {
|
||||
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<Vec<Face>, 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<Face> = 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<Face> = 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<Tensor, DetectError> {
|
||||
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<f32>,
|
||||
}
|
||||
|
||||
/// 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<Tap> {
|
||||
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<f32> = (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=<face_detection_yunet_2023mar.onnx>` (the file
|
||||
/// Settings ▸ Faces downloads; nothing is committed). Without it those tests say so and pass.
|
||||
fn model() -> Option<Vec<u8>> {
|
||||
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<Detector> {
|
||||
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=<folder> cargo test -p lightcraft-faces -- --ignored --nocapture` compares every output
|
||||
/// of the network with ONNX Runtime's on the same input (`input.bin`, `out_<name>.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<f32> { 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);
|
||||
}
|
||||
}
|
||||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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.
|
||||
|
|
@ -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<u8>,
|
||||
/// The primary image's ICC profile.
|
||||
pub icc: Option<Vec<u8>>,
|
||||
/// 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<Vec<u8>>,
|
||||
/// The XMP packet.
|
||||
pub xmp: Option<String>,
|
||||
}
|
||||
|
||||
/// 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<T>(f: impl FnOnce() -> Result<T, Error>) -> Result<T, Error> {
|
||||
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<Info, Error> {
|
||||
guarded(|| probe_unguarded(bytes))
|
||||
}
|
||||
|
||||
fn probe_unguarded(bytes: &[u8]) -> Result<Info, Error> {
|
||||
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<Decoded, Error> {
|
||||
guarded(|| decode_unguarded(bytes, options))
|
||||
}
|
||||
|
||||
fn decode_unguarded(bytes: &[u8], options: &Options) -> Result<Decoded, Error> {
|
||||
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<Vec<u8>>, Option<Vec<u8>>, Option<String>) {
|
||||
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"));
|
||||
}
|
||||
}
|
||||
|
|
@ -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<PhotoId, String> {
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 }
|
||||
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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<lightcraft_develop::EmbeddedLens>,
|
||||
|
|
@ -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<Arc<CameraColor>>,
|
||||
pub camera_tone: Option<tone::CameraTone>,
|
||||
/// Segmentation mattes stored in the file (DNG semantic masks): AI masks use them.
|
||||
pub mattes: Option<Arc<masks::Mattes>>,
|
||||
}
|
||||
|
||||
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<PixelWindow>,
|
||||
/// The source's segmentation mattes ([`SourceInfo::mattes`]).
|
||||
pub mattes: Option<Arc<masks::Mattes>>,
|
||||
}
|
||||
|
||||
/// 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<Plane>
|
|||
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"
|
||||
|
|
|
|||
|
|
@ -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<Rgb32f>, 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<Rgb32f>,
|
||||
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<Rgb32f>, 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 }
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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<u8>)>,
|
||||
}
|
||||
|
||||
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<u8>) {
|
||||
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<u8>> {
|
||||
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<u8>> {
|
||||
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<MatteKind> = 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<Evaluated> {
|
||||
/// 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<Evaluated> {
|
||||
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<u8>], 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<u8>, 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<u8> {
|
||||
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<u8> = 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));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<crate::CameraColor> {
|
||||
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());
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
|
|
@ -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<Mat3> {
|
||||
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<Xy> {
|
||||
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, _| {
|
||||
|
|
|
|||
|
|
@ -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<Vec<f64>>) -> 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<GainTableMap> {
|
||||
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<RawImage> {
|
||||
let tiff = Tiff::parse(bytes)?;
|
||||
let ifd0 = &tiff.ifds[0];
|
||||
|
|
@ -79,9 +97,6 @@ pub(crate) fn decode(bytes: &[u8], mode: Mode) -> Result<RawImage> {
|
|||
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<RawImage> {
|
|||
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,
|
||||
|
|
|
|||
|
|
@ -182,6 +182,9 @@ pub fn write_dng(raw: &RawImage, opts: &DngWriteOptions) -> Result<Vec<u8>> {
|
|||
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<Vec<u8>> {
|
|||
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)));
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
//!
|
||||
//! (<https://helpx.adobe.com/camera-raw/digital-negative.html>, "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<f32>,
|
||||
}
|
||||
|
||||
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<GainTableMap> {
|
||||
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<GainTableMap> {
|
||||
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<f32> = 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<u8>) {
|
||||
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<u8> {
|
||||
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<u8> {
|
||||
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<f32> {
|
||||
(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::<u32>() 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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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<ChunkPx> {
|
||||
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)
|
||||
})
|
||||
}
|
||||
|
|
@ -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<Xy>,
|
||||
/// 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<f64>,
|
||||
/// 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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<Vec<u8>> {
|
||||
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<Vec<u8>> {
|
|||
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)))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<HsvTable>,
|
||||
/// `ProfileToneCurve`.
|
||||
pub tone_curve: Option<ToneCurve>,
|
||||
/// `ProfileGainTableMap2`, else `ProfileGainTableMap`.
|
||||
#[serde(default)]
|
||||
pub gain_table_map: Option<GainTableMap>,
|
||||
}
|
||||
|
||||
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));
|
||||
|
|
|
|||
|
|
@ -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<String>,
|
||||
pub width: usize,
|
||||
pub height: usize,
|
||||
/// `None`: the mask covers the whole active area.
|
||||
pub sub_area: Option<MaskSubArea>,
|
||||
pub data: Vec<u16>,
|
||||
}
|
||||
|
||||
/// 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<SemanticMask> {
|
||||
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<Option<SemanticMask>> {
|
||||
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<u16> = 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<Image<u8>> {
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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<u16>),
|
||||
F32(Vec<f32>),
|
||||
}
|
||||
|
|
@ -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<ChunkPx> {
|
||||
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<ChunkPx> {
|
||||
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<u8>, 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 {
|
||||
|
|
|
|||
|
|
@ -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::<Vec<u8>>(), 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<u8> = words.iter().flat_map(|w| w.to_le_bytes()).collect();
|
||||
let be: Vec<u8> = 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
|
||||
|
|
|
|||
|
|
@ -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<f32> {
|
||||
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<RawImage> {
|
|||
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<RawImage> {
|
|||
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<RawImage> {
|
|||
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<RawData> {
|
||||
/// 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<RawData> {
|
||||
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<RawD
|
|||
let packing = if bytes_per_row >= 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<Vec<u8>>, w: u32, h: u32, rps: u32, note: Option<IfdBuilder>) -> Vec<u8> {
|
||||
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<Vec<u8>>, w: u32, h: u32, rps: u32, note: Option<IfdBuilder>) -> Vec<u8> {
|
||||
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<u8> {
|
||||
let mut out = Vec::new();
|
||||
for row in px.chunks(w) {
|
||||
let stream = row.iter().fold(Vec::new(), |mut bits: Vec<bool>, 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<u16> = (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<IfdBuilder>| {
|
||||
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);
|
||||
|
|
|
|||
|
|
@ -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<u8> = 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<RawImage> {
|
|||
}
|
||||
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<RawImage> {
|
|||
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<u8>, layout: Option<[u8; 36]>, jpeg: &[u8]) -> Vec<u8> {
|
||||
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<u8>, layout: Option<[u8; 36]>, jpeg: &[u8], words32: u32) -> Vec<u8> {
|
||||
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<u16> = (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<u8> = 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<u8> = stream.chunks(4).flat_map(|c| [c[3], c[2], c[1], c[0]]).collect();
|
||||
let decode = |bytes: Vec<u8>| 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(_))));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ fn synthetic(w: usize, h: usize, cfa: Option<Cfa>, 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<u8> {
|
||||
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<u8>, _) = 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<u8> {
|
||||
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<u8>, 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<u16> = (0..w * h).map(|i| ((i % w) * 97 + (i / w) * 31) as u16 & 0x0fff).collect();
|
||||
let with_alpha: Vec<u8> = 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<u8>| {
|
||||
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<u8> {
|
||||
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<u8> {
|
||||
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<u8> {
|
||||
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::<Vec<_>>(), 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<u8>, 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<u8>) -> 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<String> = 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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ log = { workspace = true }
|
|||
|
||||
[dev-dependencies]
|
||||
lightcraft-codecs = { workspace = true }
|
||||
lightcraft-faces = { workspace = true }
|
||||
|
||||
[build-dependencies]
|
||||
serde_json = { workspace = true }
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1642,5 +1642,6 @@
|
|||
"≤ at most": "≤ no máximo",
|
||||
"≥ at least": "≥ no mínimo",
|
||||
"日本語": "日本語",
|
||||
"繁體中文(台灣)": "繁體中文(台灣)"
|
||||
"繁體中文(台灣)": "繁體中文(台灣)",
|
||||
"Back to this person": "Back to this person"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<egui::Vec2, String> {
|
||||
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<egui::epaint::ClippedShape>,
|
||||
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<egui::Event>) -> 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<ViewportCommand> {
|
||||
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<ViewportCommand> {
|
||||
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<TextureId, CpuTexture>) -> 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]
|
||||
|
|
|
|||
|
|
@ -106,6 +106,8 @@ pub struct Services {
|
|||
pub picker: Option<pick::Picker>,
|
||||
/// Show an open dialog for photos; returns paths.
|
||||
pub pick_files: Option<PickFiles>,
|
||||
/// Open dialog for a face model (`.onnx`; Settings ▸ Faces ▸ Add a model file…).
|
||||
pub pick_model_file: Option<PickFiles>,
|
||||
/// Open dialog for preset files (`.lcpreset`, `.xmp`, `.lrtemplate`, `.zip`, `.dng`, Luminar `.lmp` / `.mplumpack`).
|
||||
pub pick_preset_files: Option<PickFiles>,
|
||||
/// 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<String> =
|
||||
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<String>, Vec<String>) = 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<String>, Vec<String>) = 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<usize> = photos.iter().map(|p| p["faces"].as_array().map_or(0, Vec::len)).collect();
|
||||
let (faces, with) = (counts.iter().sum::<usize>(), 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::<Vec<_>>()});
|
||||
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<std::collections::HashMap<lightcraft_catalog::AlbumId, usize>>)>,
|
||||
/// 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<String>,
|
||||
/// 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<panels::faces::Hints>)>,
|
||||
/// 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<panels::person::PersonPage>)>,
|
||||
/// The unnamed faces: (catalog revision, faces indexed, when it was asked for, the list).
|
||||
pub unnamed: Option<(u64, u64, f64, std::sync::Arc<panels::unnamed::Unnamed>)>,
|
||||
/// 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<Vec<String>>)>,
|
||||
/// 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<Vec<String>> {
|
||||
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::<Vec<_>>());
|
||||
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<Vec<lightcraft_catalog::DateGroup>> {
|
||||
match &self.date_groups {
|
||||
|
|
|
|||
|
|
@ -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<egui::Response> = 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<egui::Id> {
|
||||
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<egui::Id> = 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::<egui::Id>(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:<label>` (the untranslated
|
||||
/// label), and while its submenu is open it is the anchor the next frame's room is measured from.
|
||||
|
|
|
|||
|
|
@ -56,6 +56,7 @@ pub const UI_COMMANDS: &[UiCommand] = &[
|
|||
("view.compare", "Compare", Some("Shift+C"), "View"),
|
||||
("view.survey", "Survey", Some("N"), "View"),
|
||||
("view.people", "People", None, "View"),
|
||||
("view.person", "Show Person", None, ""),
|
||||
("view.faceBoxes", "Face Boxes", None, "View"),
|
||||
("view.reference", "Reference View", Some("Shift+R"), "View"),
|
||||
("photo.setReference", "Set as Reference Photo", None, ""),
|
||||
|
|
@ -135,6 +136,7 @@ pub const UI_COMMANDS: &[UiCommand] = &[
|
|||
("tool.keywordPainter", "Keyword Painter", None, ""),
|
||||
("view.gridInfo", "Grid Info", None, ""),
|
||||
("dialog.allMetadata", "All Metadata…", None, "Photo"),
|
||||
("dialog.faceModel", "Add Face Model…", None, ""),
|
||||
("dialog.newSmartAlbum", "New Smart Album from Filter…", Some("Cmd+Alt+N"), "File"),
|
||||
("dialog.createPreset", "Create Preset…", Some("Cmd+Shift+P"), "Photo"),
|
||||
("dialog.autoStack", "Auto-Stack by Capture Time…", None, "Photo>Stack"),
|
||||
|
|
@ -316,9 +318,24 @@ pub fn run_ui_command(app: &mut LightcraftApp, id: &str, p: &Value) -> Option<Re
|
|||
Ok(json!({"show": app.ui.face_boxes}))
|
||||
}
|
||||
"view.people" => {
|
||||
// Everyone; the page just left stays one click away (the chip next to the title)
|
||||
if let Some(name) = app.ui.person_page.take() {
|
||||
app.ui.last_person = Some(name);
|
||||
}
|
||||
app.ui.person_from = None;
|
||||
app.ui.view = ViewMode::People;
|
||||
Ok(json!({"people": app.session.catalog.people().len()}))
|
||||
}
|
||||
"view.person" => {
|
||||
// {name}: one person's page in the People view: their faces, and the faces that look like them
|
||||
let Some(name) = p.get("name").and_then(Value::as_str).map(str::trim).filter(|n| !n.is_empty()) else {
|
||||
return Some(Err("view.person: missing `name`".into()));
|
||||
};
|
||||
app.ui.view = ViewMode::People;
|
||||
app.ui.person_from = None;
|
||||
app.ui.person_page = Some(name.to_string());
|
||||
Ok(json!({"person": name}))
|
||||
}
|
||||
"view.survey" => {
|
||||
app.ui.view = ViewMode::Survey;
|
||||
Ok(json!({"photos": crate::panels::compare::survey_photos(app).len()}))
|
||||
|
|
@ -364,10 +381,19 @@ pub fn run_ui_command(app: &mut LightcraftApp, id: &str, p: &Value) -> Option<Re
|
|||
app.ui.slideshow = None;
|
||||
} else if !app.ui.tool.is_empty() {
|
||||
app.ui.tool.clear();
|
||||
} else if app.ui.view == ViewMode::People && app.ui.person_page.is_some() {
|
||||
app.ui.person_page = None;
|
||||
} else if matches!(app.ui.view, ViewMode::Compare | ViewMode::Survey) {
|
||||
app.ui.view = ViewMode::Detail;
|
||||
} else if app.ui.view == ViewMode::Detail {
|
||||
app.ui.view = ViewMode::PhotoGrid;
|
||||
// a photo opened from a person's page goes back to that page, any other to the grid
|
||||
match app.ui.person_from.take() {
|
||||
Some((name, photo)) if app.session.active().is_some_and(|a| a.0 == photo) => {
|
||||
app.ui.view = ViewMode::People;
|
||||
app.ui.person_page = Some(name);
|
||||
}
|
||||
_ => app.ui.view = ViewMode::PhotoGrid,
|
||||
}
|
||||
}
|
||||
Ok(Value::Null)
|
||||
}
|
||||
|
|
@ -433,7 +459,7 @@ pub fn run_ui_command(app: &mut LightcraftApp, id: &str, p: &Value) -> Option<Re
|
|||
"app.settings" => {
|
||||
let tab = p.get("tab").and_then(Value::as_str).unwrap_or("general");
|
||||
if !crate::panels::settings::TABS.iter().any(|(id, _)| *id == tab) {
|
||||
return Some(Err(format!("unknown settings tab `{tab}` (general|import|performance|interface)")));
|
||||
return Some(Err(format!("unknown settings tab `{tab}` (general|import|performance|interface|faces)")));
|
||||
}
|
||||
app.ui.dialog = Some(Dialog::Settings { tab: tab.into() });
|
||||
Ok(Value::Null)
|
||||
|
|
@ -906,6 +932,16 @@ pub fn run_ui_command(app: &mut LightcraftApp, id: &str, p: &Value) -> Option<Re
|
|||
"merge.panoramaLast" => crate::merge::start_last(app, "merge.panorama"),
|
||||
"merge.hdrPanoramaLast" => crate::merge::start_last(app, "merge.hdrPanorama"),
|
||||
"dialog.mergeHdr" => crate::merge::open(app, "merge.hdr"),
|
||||
"dialog.faceModel" => {
|
||||
let path = match p.get("path").and_then(Value::as_str) {
|
||||
Some(x) => Some(x.to_string()),
|
||||
None => app.services.pick_model_file.as_mut().and_then(|f| f().into_iter().next()),
|
||||
};
|
||||
match path {
|
||||
Some(path) => crate::panels::faces::open_dialog(app, &path),
|
||||
None => Ok(Value::Null),
|
||||
}
|
||||
}
|
||||
"dialog.mergePanorama" => crate::merge::open(app, "merge.panorama"),
|
||||
"dialog.mergeHdrPanorama" => crate::merge::open(app, "merge.hdrPanorama"),
|
||||
"app.about" => {
|
||||
|
|
|
|||
|
|
@ -631,10 +631,26 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui) {
|
|||
filter_pill(app, ui, canvas);
|
||||
}
|
||||
if app.ui.face_boxes {
|
||||
match region_overlay(ui, &p, &map, &photo, d.orientation) {
|
||||
// a name box left open on another photo is dropped
|
||||
if app.ui.name_edit.as_ref().is_some_and(|e| e.photo != id.0) {
|
||||
app.ui.name_edit = None;
|
||||
}
|
||||
let hints = super::faces::hints_for(app, id.0);
|
||||
let people = if app.ui.name_edit.is_some() { app.caches.person_names(&app.session.catalog) } else { Default::default() };
|
||||
// while the name box is open and recognition is not set up, the box offers the next step
|
||||
let setup = if app.ui.name_edit.is_some() { super::faces::setup(app, ui.ctx()) } else { super::faces::Setup::Running };
|
||||
match region_overlay(ui, &p, &map, &photo, d.orientation, hints.as_deref(), &people, &mut app.ui.name_edit, setup) {
|
||||
Some(RegionEdit::SetUp) => {
|
||||
app.ui.name_edit = None;
|
||||
super::faces::take_step(app, ui.ctx(), setup);
|
||||
}
|
||||
Some(RegionEdit::Remove(index)) => {
|
||||
let _ = app.run("photo.removeRegion", json!({"id": id.0, "index": index}));
|
||||
}
|
||||
Some(RegionEdit::Name(index, name)) => {
|
||||
app.ui.name_edit = None;
|
||||
let _ = app.run("faces.setName", json!({"id": id.0, "index": index, "name": name}));
|
||||
}
|
||||
Some(RegionEdit::Resize(index, r)) => {
|
||||
let _ = app.run("photo.setRegion", json!({"id": id.0, "index": index, "rect": {"x0": r.x0, "y0": r.y0, "x1": r.x1, "y1": r.y1}}));
|
||||
}
|
||||
|
|
@ -740,6 +756,10 @@ enum RegionEdit {
|
|||
Remove(usize),
|
||||
/// A handle drag ended: the region's new box (normalized, upright frame).
|
||||
Resize(usize, lightcraft_geom::Rect),
|
||||
/// A name was typed or picked for the face.
|
||||
Name(usize, String),
|
||||
/// The name box's offer to set up face recognition was pressed.
|
||||
SetUp,
|
||||
}
|
||||
|
||||
/// Handles of a box: (x, y) as fractions of its width and height, and the cursor they show.
|
||||
|
|
@ -765,6 +785,10 @@ fn region_overlay(
|
|||
map: &CanvasMap,
|
||||
photo: &lightcraft_catalog::Photo,
|
||||
orient: lightcraft_geom::Orientation,
|
||||
hints: Option<&super::faces::Hints>,
|
||||
people: &[String],
|
||||
editing: &mut Option<crate::state::NameEdit>,
|
||||
setup: super::faces::Setup,
|
||||
) -> Option<RegionEdit> {
|
||||
let t = Tokens::get(p.ctx());
|
||||
let clip = p.clip_rect();
|
||||
|
|
@ -840,9 +864,20 @@ fn region_overlay(
|
|||
edit = Some(RegionEdit::Remove(index));
|
||||
}
|
||||
}
|
||||
let Some(name) = &r.name else { continue };
|
||||
// what the label says: the name, a guess to confirm ("Jane Doe?"), or on hover an invitation to name the face
|
||||
let is_face = r.kind == lightcraft_meta::RegionKind::Face;
|
||||
let hint = hints.and_then(|h| h.by_index.get(&index));
|
||||
let hovered = pointer.is_some_and(|h| rect.expand(3.0).contains(h));
|
||||
let editing_this = editing.as_ref().is_some_and(|e| e.photo == photo.id.0 && e.index == index);
|
||||
let (text, text_color, fill_alpha, invitation) = match (&r.name, hint.and_then(|h| h.suggestion.as_ref())) {
|
||||
(Some(n), _) => (n.clone(), Color32::from_gray(225), 235, false),
|
||||
(None, Some((n, _))) if is_face => (format!("{n}?"), Color32::from_gray(185), 200, false),
|
||||
(None, None) if is_face && live.is_none() => ("Add name".to_string(), Color32::from_gray(150), 190, true),
|
||||
_ => continue,
|
||||
};
|
||||
// the name in a dark label with a caret, centred above the box (below it when there is no room)
|
||||
let g = p.layout_no_wrap(name.clone(), t.font(13.0), Color32::from_gray(225));
|
||||
let name = &text;
|
||||
let g = p.layout_no_wrap(name.clone(), t.font(13.0), text_color);
|
||||
let (pad, caret) = (vec2(14.0, 7.0), 5.0);
|
||||
let size = g.size() + pad * 2.0;
|
||||
let above = rect.top() - caret - size.y >= clip.top();
|
||||
|
|
@ -853,12 +888,34 @@ fn region_overlay(
|
|||
};
|
||||
let left = (rect.center().x - size.x / 2.0).clamp(clip.left(), (clip.right() - size.x).max(clip.left()));
|
||||
let label = Rect::from_min_size(pos2(left, top), size);
|
||||
let fill = Color32::from_rgba_unmultiplied(56, 56, 56, 235);
|
||||
// the invitation shows while the pointer is on the box *or on the label* (and the gap between), so it can be reached
|
||||
if invitation && !(hovered || editing_this || pointer.is_some_and(|h| rect.union(label).expand(3.0).contains(h))) {
|
||||
continue;
|
||||
}
|
||||
let fill = Color32::from_rgba_unmultiplied(56, 56, 56, fill_alpha);
|
||||
p.rect_filled(label, 3.0, fill);
|
||||
p.rect_stroke(label, 3.0, Stroke::new(1.0, Color32::from_black_alpha(160)), StrokeKind::Inside);
|
||||
let cx = rect.center().x.clamp(label.left() + caret + 4.0, label.right() - caret - 4.0);
|
||||
p.add(egui::Shape::convex_polygon(vec![pos2(cx - caret, base), pos2(cx + caret, base), pos2(cx, tip)], fill, Stroke::NONE));
|
||||
p.galley(label.min + pad, g, Color32::from_gray(225));
|
||||
p.galley(label.min + pad, g, text_color);
|
||||
if is_face && live.is_none() {
|
||||
// a click on the label opens the name box, filled with the name or the guess
|
||||
let resp = ui.interact(label, egui::Id::new(("region-label", index)), Sense::click());
|
||||
register(ui.ctx(), format!("regionLabel:{index}"), label);
|
||||
if resp.on_hover_text("Click to name this face").clicked() {
|
||||
let start = r.name.clone().or_else(|| hint.and_then(|h| h.suggestion.as_ref().map(|(n, _)| n.clone()))).unwrap_or_default();
|
||||
*editing = Some(crate::state::NameEdit { photo: photo.id.0, index, text: start, fresh: true });
|
||||
}
|
||||
}
|
||||
if editing_this && let Some(e) = editing.as_mut() {
|
||||
let candidates = hint.map(|h| h.candidates.as_slice()).unwrap_or(&[]);
|
||||
match super::faces::name_editor(ui.ctx(), pos2(label.left(), rect.bottom() + 8.0), e, candidates, people, setup) {
|
||||
super::faces::Editor::Submit(name) => edit = Some(RegionEdit::Name(index, name)),
|
||||
super::faces::Editor::Setup => edit = Some(RegionEdit::SetUp),
|
||||
super::faces::Editor::Cancel => *editing = None,
|
||||
super::faces::Editor::Open => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
edit
|
||||
}
|
||||
|
|
|
|||
|
|
@ -62,6 +62,7 @@ pub const WHATS_NEW: &str = include_str!("../../../../docs/whats-new.md");
|
|||
|
||||
pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) {
|
||||
let Some(mut dlg) = app.ui.dialog.clone() else { return };
|
||||
let at_start = dlg.clone();
|
||||
let t = Tokens::get(ctx);
|
||||
let screen = ctx.content_rect();
|
||||
// The backdrop is an area below the dialog window (a bare `Middle` layer painter would be
|
||||
|
|
@ -85,6 +86,8 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) {
|
|||
Dialog::NewSmartAlbum { .. } => "Create Smart Album",
|
||||
Dialog::AllMetadata { .. } => "All Metadata",
|
||||
Dialog::SystemInfo { .. } => "System Info",
|
||||
Dialog::FaceModel { info, .. } if info["download"].is_string() => "Download Face Model",
|
||||
Dialog::FaceModel { .. } => "Add Face Model",
|
||||
Dialog::WhatsNew => "What's New",
|
||||
Dialog::Cull { .. } => "Assisted Culling",
|
||||
Dialog::SmartRules { id: None, .. } => "New Smart Album",
|
||||
|
|
@ -123,6 +126,7 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) {
|
|||
Dialog::Import { .. } => crate::import::DIALOG_SIZE[0],
|
||||
Dialog::SmartRules { .. } => 680.0,
|
||||
Dialog::AllMetadata { .. } => 620.0,
|
||||
Dialog::FaceModel { .. } => 460.0,
|
||||
_ => 380.0,
|
||||
})
|
||||
.show(ctx, |ui| {
|
||||
|
|
@ -174,6 +178,7 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) {
|
|||
}
|
||||
});
|
||||
}
|
||||
Dialog::FaceModel { info, accepted, .. } => crate::panels::faces::model_dialog(app, ui, &t, info, accepted),
|
||||
Dialog::SystemInfo { rows } => {
|
||||
egui::Grid::new("sysinfo").num_columns(2).spacing([16.0, 4.0]).striped(true).show(ui, |ui| {
|
||||
for (k, v) in rows.iter() {
|
||||
|
|
@ -784,7 +789,9 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) {
|
|||
}
|
||||
ui.add_space(4.0);
|
||||
ui.horizontal(|ui| {
|
||||
let informational = matches!(dlg, Dialog::About | Dialog::Shortcuts | Dialog::Settings { .. });
|
||||
// a model file LightCraft cannot use has nothing to confirm
|
||||
let unusable_model = matches!(&dlg, Dialog::FaceModel { info, .. } if info["kind"] == "unsupported");
|
||||
let informational = unusable_model || matches!(dlg, Dialog::About | Dialog::Shortcuts | Dialog::Settings { .. });
|
||||
let sam = &app.session.segmenter;
|
||||
let (sam_installed, sam_running, sam_failed) = (sam.installed(), sam.download_status().running, sam.download_status().error.is_some());
|
||||
// no download location in this build: nothing to offer but the manual install
|
||||
|
|
@ -810,6 +817,8 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) {
|
|||
add_label.as_str()
|
||||
}
|
||||
Dialog::Merge { .. } => "Merge",
|
||||
Dialog::FaceModel { info, .. } if !informational && info["download"].is_string() => "Download",
|
||||
Dialog::FaceModel { .. } if !informational => "Install",
|
||||
Dialog::ConfirmDelete { .. } => "Delete",
|
||||
Dialog::RemoveFolder { .. } => "Remove",
|
||||
Dialog::SamModel { then: Some(_), .. } if sam_installed => "Continue",
|
||||
|
|
@ -819,7 +828,9 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) {
|
|||
_ if informational => "Close",
|
||||
_ => "OK",
|
||||
};
|
||||
let r = (!ok.is_empty()).then(|| ui.button(crate::i18n::tr(ok)));
|
||||
// installing waits for the licence to be accepted
|
||||
let can_confirm = informational || !matches!(&dlg, Dialog::FaceModel { accepted: false, .. });
|
||||
let r = (!ok.is_empty()).then(|| ui.add_enabled(can_confirm, egui::Button::new(crate::i18n::tr(ok))));
|
||||
if let Some(r) = &r {
|
||||
crate::widgets::register(ui.ctx(), "button:dialogOk", r.rect);
|
||||
}
|
||||
|
|
@ -862,6 +873,10 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) {
|
|||
_ => close = true,
|
||||
}
|
||||
}
|
||||
// a dialog opened from inside this one (a button in Settings that shows a licence) takes its place
|
||||
if app.ui.dialog.as_ref().is_some_and(|now| *now != at_start) {
|
||||
return;
|
||||
}
|
||||
app.ui.dialog = if close { None } else { Some(dlg) };
|
||||
}
|
||||
|
||||
|
|
@ -1025,6 +1040,7 @@ pub fn confirm_dialog(app: &mut LightcraftApp, dlg: &Dialog) -> Result<serde_jso
|
|||
Dialog::RenameKeyword { from, to } => app.run("keyword.rename", json!({"from": from, "to": to})),
|
||||
Dialog::MergeKeywords { from, into } => app.run("keyword.merge", json!({"from": from, "into": into})),
|
||||
Dialog::AutoStack { gap } => app.run("stack.auto", json!({"gap": gap})),
|
||||
Dialog::FaceModel { path, info, accepted } => crate::panels::faces::install(app, path, info, *accepted),
|
||||
Dialog::AllMetadata { .. } | Dialog::SystemInfo { .. } | Dialog::WhatsNew => Ok(serde_json::Value::Null),
|
||||
Dialog::Cull { reject_below, pick_best } => {
|
||||
let mut p = json!({"pickBest": pick_best});
|
||||
|
|
|
|||
|
|
@ -0,0 +1,779 @@
|
|||
//! Faces: the Settings ▸ Faces tab (the on/off toggle and the list of face models) and the dialog that
|
||||
//! shows a model file's licence terms before it is installed.
|
||||
//!
|
||||
//! Adding a model is: pick or drop a `.onnx` file → the engine inspects it (`faces.models.inspect`) →
|
||||
//! this dialog shows what it is and its terms → the user ticks "I accept" → `faces.models.install`.
|
||||
//! LightCraft fetches a model only when the user presses Download on a model it has a pinned address for: the
|
||||
//! same dialog shows its terms first, and accepting them downloads it (`faces.models.download`: checked against its
|
||||
//! size and SHA-256), installs it, starts using it and switches recognition on, with no further question. "Open
|
||||
//! page" opens the model's own page in the browser for anything else.
|
||||
|
||||
use egui::RichText;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::settings::{check, heading, hint};
|
||||
use crate::LightcraftApp;
|
||||
use crate::theme::Tokens;
|
||||
use crate::widgets::register;
|
||||
|
||||
/// How hard the background scan may work (the engine's `faces.pump` pace), from what the user is doing: nothing new
|
||||
/// while they drag, type or scroll; one photo at a time while they are around or the window is minimized; half the
|
||||
/// machine once they have been idle for a few seconds, or while another app has the keyboard but this window is still
|
||||
/// on screen; most of it while they are looking at the scan's progress.
|
||||
fn scan_pace(app: &LightcraftApp, ctx: &egui::Context, now: f64) -> (&'static str, bool) {
|
||||
let (focused, minimized) = ctx.input(|i| (i.focused, i.raw.viewports.get(&i.raw.viewport_id).and_then(|v| v.minimized).unwrap_or(false)));
|
||||
let watching =
|
||||
matches!(&app.ui.dialog, Some(crate::state::Dialog::Settings { tab }) if tab == "faces") || app.ui.view == crate::state::ViewMode::People;
|
||||
(pace_for(focused, minimized, now - app.caches.last_input, now - app.caches.last_move, watching), focused && !minimized)
|
||||
}
|
||||
|
||||
/// The pace for a window that has the keyboard (or not) and is minimized (or not), `worked` seconds after the user last
|
||||
/// dragged, typed or scrolled and `moved` seconds after they last moved the pointer, while they are (or are not) looking
|
||||
/// at the scan's progress.
|
||||
fn pace_for(focused: bool, minimized: bool, worked: f64, moved: f64, watching: bool) -> &'static str {
|
||||
if minimized {
|
||||
"light"
|
||||
} else if !focused {
|
||||
// still on screen, but the user is working in another app: no input here to get in the way of
|
||||
"normal"
|
||||
} else if worked < 0.4 {
|
||||
"pause"
|
||||
} else if moved < 3.0 {
|
||||
"light"
|
||||
} else if watching {
|
||||
"full"
|
||||
} else {
|
||||
"normal"
|
||||
}
|
||||
}
|
||||
|
||||
/// How long a read of the model list is reused (it is a few small files, but not for every frame).
|
||||
const REFRESH_SECS: f64 = 1.5;
|
||||
|
||||
fn list_id() -> egui::Id {
|
||||
egui::Id::new("faces-model-list")
|
||||
}
|
||||
|
||||
/// The engine's model list, re-read at most every [`REFRESH_SECS`], or at once after an action changed it.
|
||||
fn models(app: &mut LightcraftApp, ctx: &egui::Context) -> Value {
|
||||
let now = ctx.input(|i| i.time);
|
||||
let epoch = app.caches.faces_epoch;
|
||||
if let Some((e, at, v)) = ctx.data(|d| d.get_temp::<(u64, f64, Value)>(list_id()))
|
||||
&& e == epoch
|
||||
&& now - at < REFRESH_SECS
|
||||
{
|
||||
return v;
|
||||
}
|
||||
let v = app.run("faces.models.list", json!({})).unwrap_or(Value::Null);
|
||||
ctx.data_mut(|d| d.insert_temp(list_id(), (epoch, now, v.clone())));
|
||||
v
|
||||
}
|
||||
|
||||
/// What stands between the user and name suggestions, for the prompts where faces are shown or named.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum Setup {
|
||||
/// Recognition is on (or about to start): nothing to offer.
|
||||
Running,
|
||||
/// A model is installed and chosen; only the switch is off.
|
||||
TurnOn,
|
||||
/// No model to use yet: Settings ▸ Faces has the download.
|
||||
GetModel,
|
||||
/// This build cannot keep or run models: nothing to offer.
|
||||
Unavailable,
|
||||
}
|
||||
|
||||
/// Where the user stands with face recognition. Cheap: the model list is read at most every 1.5 s, and not at all while
|
||||
/// recognition runs.
|
||||
pub fn setup(app: &mut LightcraftApp, ctx: &egui::Context) -> Setup {
|
||||
if app.caches.faces_active {
|
||||
return Setup::Running;
|
||||
}
|
||||
let list = models(app, ctx);
|
||||
if list["dir"].is_null() || list["runtime"].as_bool() != Some(true) {
|
||||
return Setup::Unavailable;
|
||||
}
|
||||
let chosen = list["embedder"].as_str();
|
||||
let usable = chosen.is_some_and(|id| list["models"].as_array().into_iter().flatten().any(|m| m["id"] == id && m["installed"] == true));
|
||||
match (usable, list["enabled"].as_bool() == Some(true)) {
|
||||
(true, true) => Setup::Running,
|
||||
(true, false) => Setup::TurnOn,
|
||||
(false, _) => Setup::GetModel,
|
||||
}
|
||||
}
|
||||
|
||||
impl Setup {
|
||||
/// The sentence and the button label that offer the next step; none when there is nothing to do.
|
||||
pub fn prompt(self) -> Option<(&'static str, &'static str)> {
|
||||
match self {
|
||||
Setup::TurnOn => Some(("Face recognition is off. Turn it on to get name suggestions and see look-alike faces together.", "Turn on")),
|
||||
Setup::GetModel => {
|
||||
Some(("Face recognition needs a model to suggest names: a one-time download, in Settings.", "Set up face recognition"))
|
||||
}
|
||||
Setup::Running | Setup::Unavailable => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// The same offer as one line of a small box (the name box in the loupe).
|
||||
pub fn short(self) -> Option<&'static str> {
|
||||
match self {
|
||||
Setup::TurnOn => Some("Turn on name suggestions"),
|
||||
Setup::GetModel => Some("Set up name suggestions…"),
|
||||
Setup::Running | Setup::Unavailable => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The prompt's button: switch recognition on, or open Settings ▸ Faces, where the model is one click away.
|
||||
pub fn take_step(app: &mut LightcraftApp, ctx: &egui::Context, step: Setup) {
|
||||
match step {
|
||||
Setup::TurnOn => {
|
||||
if app.run("faces.enable", json!({"enabled": true})).is_ok() {
|
||||
app.caches.faces_epoch += 1;
|
||||
app.ui.status = "Face recognition is on".into();
|
||||
app.toast(ctx, "Face recognition is on");
|
||||
}
|
||||
}
|
||||
Setup::GetModel => {
|
||||
let _ = app.run("app.settings", json!({"tab": "faces"}));
|
||||
}
|
||||
Setup::Running | Setup::Unavailable => {}
|
||||
}
|
||||
}
|
||||
|
||||
/// A slim line under a view's heading offering to set recognition up, shown only while it is not set up.
|
||||
pub fn setup_banner(app: &mut LightcraftApp, ui: &mut egui::Ui) {
|
||||
let step = setup(app, ui.ctx());
|
||||
let Some((message, button)) = step.prompt() else { return };
|
||||
let t = Tokens::get(ui.ctx());
|
||||
egui::Frame::NONE.inner_margin(egui::Margin::symmetric(12, 6)).show(ui, |ui| {
|
||||
ui.horizontal_wrapped(|ui| {
|
||||
ui.label(RichText::new(message).font(t.font(12.5)).color(t.text_dim));
|
||||
let r = ui.button(button);
|
||||
register(ui.ctx(), "faces:setup", r.rect);
|
||||
if r.clicked() {
|
||||
take_step(app, ui.ctx(), step);
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/// A download-style size: decimal megabytes, as file managers and model pages show them.
|
||||
fn mb(bytes: Option<u64>) -> String {
|
||||
match bytes {
|
||||
Some(b) if b >= 10_000_000 => format!("{} MB", (b as f64 / 1e6).round() as u64),
|
||||
Some(b) if b >= 1_000_000 => format!("{:.1} MB", b as f64 / 1e6),
|
||||
Some(b) => format!("{} KB", (b as f64 / 1e3).round().max(1.0) as u64),
|
||||
None => String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// "its input is not an image" → "Its input is not an image."
|
||||
fn sentence(s: &str) -> String {
|
||||
let s = s.trim();
|
||||
let mut c = s.chars();
|
||||
let mut out: String = c.next().map(|f| f.to_uppercase().collect()).unwrap_or_default();
|
||||
out.push_str(c.as_str());
|
||||
if !out.is_empty() && !out.ends_with(['.', '!', '?']) {
|
||||
out.push('.');
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// "Apache-2.0 · commercial use allowed"
|
||||
fn licence_line(m: &Value) -> String {
|
||||
let name = m["licence"]["name"].as_str().filter(|s| !s.is_empty()).unwrap_or("Unknown licence");
|
||||
let terms = match m["licence"]["commercial"].as_str() {
|
||||
Some("yes") => "commercial use allowed",
|
||||
Some("no") => "non-commercial use only",
|
||||
_ => "terms unclear",
|
||||
};
|
||||
format!("{name} · {terms}")
|
||||
}
|
||||
|
||||
fn open_page(app: &mut LightcraftApp, url: &str) {
|
||||
if let Some(f) = app.services.open_url.as_mut() {
|
||||
let _ = f(url);
|
||||
}
|
||||
}
|
||||
|
||||
/// The Settings ▸ Faces tab.
|
||||
pub fn settings_tab(app: &mut LightcraftApp, ui: &mut egui::Ui, t: &Tokens) {
|
||||
let list = models(app, ui.ctx());
|
||||
heading(ui, t, "Face recognition");
|
||||
if list["dir"].is_null() {
|
||||
hint(ui, t, "This build has nowhere to keep face models, so the model list is not available. (The desktop app has.)");
|
||||
return;
|
||||
}
|
||||
let mut on = list["enabled"].as_bool().unwrap_or(false);
|
||||
if check(ui, "settings.facesEnabled", &mut on, "Recognise faces (experimental)") {
|
||||
let _ = app.run("faces.enable", json!({"enabled": on}));
|
||||
app.caches.faces_epoch += 1;
|
||||
}
|
||||
hint(ui, t, "Suggests who is in a photo from the faces you have named. Everything stays on your computer.");
|
||||
let runtime = list["runtime"].as_bool() == Some(true);
|
||||
if !runtime {
|
||||
hint(ui, t, "This build cannot run recognition models: they can be added and chosen, not used.");
|
||||
} else if on && list["embedder"].is_null() {
|
||||
hint(ui, t, "Choose a model below (Use) to start.");
|
||||
} else if on && app.caches.faces_active {
|
||||
scan_progress(app, ui, t);
|
||||
}
|
||||
let all: Vec<Value> = list["models"].as_array().cloned().unwrap_or_default();
|
||||
let downloads =
|
||||
app.session.execute("faces.models.downloads", &json!({})).map(|v| v["downloads"].as_array().cloned().unwrap_or_default()).unwrap_or_default();
|
||||
heading(ui, t, "Face detector");
|
||||
for m in all.iter().filter(|m| m["role"] == "detector") {
|
||||
model_row(app, ui, t, m, downloads.iter().find(|d| d["id"] == m["id"]), runtime);
|
||||
}
|
||||
if !all.iter().any(|m| m["role"] == "detector" && m["installed"] == true) {
|
||||
hint(ui, t, "Detect Faces (Photo menu) needs it: a small one-time download.");
|
||||
}
|
||||
heading(ui, t, "Recognition models");
|
||||
for m in all.iter().filter(|m| m["role"] == "embedder") {
|
||||
model_row(app, ui, t, m, downloads.iter().find(|d| d["id"] == m["id"]), runtime);
|
||||
}
|
||||
ui.add_space(6.0);
|
||||
ui.horizontal(|ui| {
|
||||
let can = app.services.pick_model_file.is_some();
|
||||
let r = ui.add_enabled(can, egui::Button::new("Add a model file…"));
|
||||
register(ui.ctx(), "faces:addModel", r.rect);
|
||||
if r.clicked() {
|
||||
let _ = app.run("dialog.faceModel", json!({}));
|
||||
}
|
||||
ui.label(RichText::new("or drop a .onnx file on the window").color(t.text_dim));
|
||||
});
|
||||
// what is wrong with the user's own catalog (catalog.json in the models folder), if anything
|
||||
for e in list["catalog"]["errors"].as_array().into_iter().flatten().filter_map(Value::as_str).take(3) {
|
||||
ui.add(egui::Label::new(RichText::new(format!("catalog.json: {e}")).font(t.font(11.5)).color(t.caution)).wrap());
|
||||
}
|
||||
if let Some(d) = all.iter().find(|m| m["role"] == "detector") {
|
||||
ui.add_space(4.0);
|
||||
hint(ui, t, &format!("Faces are found by {} ({}, included).", d["name"].as_str().unwrap_or("the detector"), mb(d["sizeBytes"].as_u64())));
|
||||
}
|
||||
}
|
||||
|
||||
/// How the scan is going: a bar while photos are left, then how many faces it learned.
|
||||
fn scan_progress(app: &mut LightcraftApp, ui: &mut egui::Ui, t: &Tokens) {
|
||||
let (left, peak, faces) = (app.caches.faces_pending, app.caches.faces_peak.max(1), app.caches.faces_indexed);
|
||||
if left == 0 {
|
||||
hint(ui, t, &format!("{faces} faces learned."));
|
||||
return;
|
||||
}
|
||||
let text = format!("Scanning photos for faces: {left} left · {faces} faces so far");
|
||||
let bar = ui.add(
|
||||
egui::ProgressBar::new((1.0 - left as f32 / peak as f32).clamp(0.0, 1.0)).desired_width(380.0).text(RichText::new(text).font(t.font(11.5))),
|
||||
);
|
||||
register(ui.ctx(), "faces:scanProgress", bar.rect);
|
||||
}
|
||||
|
||||
fn model_row(app: &mut LightcraftApp, ui: &mut egui::Ui, t: &Tokens, m: &Value, dl: Option<&Value>, can_run: bool) {
|
||||
let id = m["id"].as_str().unwrap_or("").to_string();
|
||||
let detector = m["role"] == "detector";
|
||||
let (installed, selected) = (m["installed"].as_bool() == Some(true), m["selected"].as_bool() == Some(true));
|
||||
let dl_state = if installed { None } else { dl.and_then(|d| d["state"].as_str()) };
|
||||
// a build that cannot run recognition models has no use for a download
|
||||
let host = if can_run { m["downloadHost"].as_str().unwrap_or("").to_string() } else { String::new() };
|
||||
egui::Frame::NONE.inner_margin(egui::Margin::symmetric(0, 3)).show(ui, |ui| {
|
||||
ui.horizontal(|ui| {
|
||||
ui.vertical(|ui| {
|
||||
ui.horizontal(|ui| {
|
||||
ui.label(RichText::new(m["name"].as_str().unwrap_or("")).color(t.text));
|
||||
let size = mb(m["sizeBytes"].as_u64());
|
||||
if !size.is_empty() {
|
||||
ui.label(RichText::new(size).color(t.text_dim));
|
||||
}
|
||||
if selected {
|
||||
ui.label(RichText::new("in use").color(t.accent));
|
||||
} else if installed {
|
||||
ui.label(RichText::new("installed").color(t.text_dim));
|
||||
}
|
||||
});
|
||||
ui.label(RichText::new(licence_line(m)).font(t.font(11.5)).color(if m["licence"]["commercial"] == "yes" {
|
||||
t.text_dim
|
||||
} else {
|
||||
t.caution
|
||||
}));
|
||||
// how fast it is, as times faster than a ResNet-100 model (milliseconds would depend on the computer),
|
||||
// with whether it passed its test when it has been installed
|
||||
let speed = m["speedText"].as_str().filter(|s| !s.is_empty());
|
||||
let (line, ok) = match m["accepted"]["selfTest"].as_object() {
|
||||
Some(test) if test.get("ok").and_then(Value::as_bool) == Some(true) => {
|
||||
(Some(speed.map_or("Works".to_string(), |s| format!("Works · {s}"))), true)
|
||||
}
|
||||
Some(_) => (Some("Failed its last test".to_string()), false),
|
||||
None => (speed.map(str::to_string), true),
|
||||
};
|
||||
if let Some(line) = line {
|
||||
ui.label(RichText::new(line).font(t.font(11.5)).color(if ok { t.text_dim } else { t.caution }));
|
||||
}
|
||||
match (dl_state, dl) {
|
||||
(Some("running"), Some(d)) => {
|
||||
let (got, total) = (d["bytes"].as_u64().unwrap_or(0), d["total"].as_u64().unwrap_or(0));
|
||||
let frac = if total > 0 { (got as f32 / total as f32).clamp(0.0, 1.0) } else { 0.0 };
|
||||
let text = if total > 0 && got >= total {
|
||||
"Checking…".to_string()
|
||||
} else {
|
||||
format!("{} of {} from {host}", mb(Some(got)), mb(Some(total)))
|
||||
};
|
||||
let bar = ui.add(egui::ProgressBar::new(frac).desired_width(260.0).text(RichText::new(text).font(t.font(11.5))));
|
||||
register(ui.ctx(), format!("faces:progress:{id}"), bar.rect);
|
||||
}
|
||||
(Some("done"), _) => {
|
||||
ui.label(RichText::new("Downloaded and checked. Installing…").font(t.font(11.5)).color(t.text_dim));
|
||||
}
|
||||
(Some("failed"), Some(d)) => {
|
||||
let why = sentence(d["error"].as_str().unwrap_or("The download failed"));
|
||||
ui.add(egui::Label::new(RichText::new(why).font(t.font(11.5)).color(t.caution)).wrap());
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
});
|
||||
ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| {
|
||||
if installed {
|
||||
let r = ui.button("Remove");
|
||||
register(ui.ctx(), format!("faces:remove:{id}"), r.rect);
|
||||
if r.clicked() {
|
||||
let _ = app.run("faces.models.remove", json!({"id": id}));
|
||||
app.caches.faces_epoch += 1;
|
||||
}
|
||||
if !selected && !detector {
|
||||
let r = ui.button("Use");
|
||||
register(ui.ctx(), format!("faces:use:{id}"), r.rect);
|
||||
if r.clicked() {
|
||||
let _ = app.run("faces.models.select", json!({"id": id}));
|
||||
app.caches.faces_epoch += 1;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
match dl_state {
|
||||
Some("running") => {
|
||||
let r = ui.button("Cancel");
|
||||
register(ui.ctx(), format!("faces:cancelDownload:{id}"), r.rect);
|
||||
if r.clicked() {
|
||||
let _ = app.run("faces.models.downloadCancel", json!({"id": id}));
|
||||
app.caches.faces_dl_watch.retain(|w| w != &id);
|
||||
}
|
||||
}
|
||||
// the engine installs it within a moment; nothing to press
|
||||
Some("done") => {}
|
||||
_ => {
|
||||
if let Some(url) = m["source"].as_str() {
|
||||
let r = ui
|
||||
.button("Open page")
|
||||
.on_hover_text("Opens the model's own page in your browser, to read about it or get the file yourself.");
|
||||
register(ui.ctx(), format!("faces:get:{id}"), r.rect);
|
||||
if r.clicked() {
|
||||
open_page(app, url);
|
||||
}
|
||||
}
|
||||
if !host.is_empty() {
|
||||
let label = if dl_state == Some("failed") { "Try again" } else { "Download" };
|
||||
let tip = format!(
|
||||
"Shows the model's terms, then downloads {} from {host} and installs it{}. Nothing else is sent.",
|
||||
mb(m["sizeBytes"].as_u64()),
|
||||
if detector { "" } else { ", starts using it" }
|
||||
);
|
||||
let r = ui.button(label).on_hover_text(tip);
|
||||
register(ui.ctx(), format!("faces:download:{id}"), r.rect);
|
||||
if r.clicked() {
|
||||
open_download_dialog(app, m);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/// The body of the "Add Face Model" dialog: what the file is, its terms, and the accept box.
|
||||
/// `info` is the engine's `faces.models.inspect` answer.
|
||||
pub fn model_dialog(app: &mut LightcraftApp, ui: &mut egui::Ui, t: &Tokens, info: &Value, accepted: &mut bool) {
|
||||
let kind = info["kind"].as_str().unwrap_or("unsupported");
|
||||
let file = info["fileName"].as_str().unwrap_or("");
|
||||
if kind == "unsupported" {
|
||||
ui.label(RichText::new("This file cannot be used yet").font(t.semibold(13.5)).color(t.caution));
|
||||
ui.add(
|
||||
egui::Label::new(
|
||||
RichText::new(sentence(info["reason"].as_str().unwrap_or("It is not a face recognition model LightCraft understands")))
|
||||
.color(t.text_label),
|
||||
)
|
||||
.wrap(),
|
||||
);
|
||||
ui.label(RichText::new(format!("{file} · {}", mb(info["sizeBytes"].as_u64()))).color(t.text_dim));
|
||||
return;
|
||||
}
|
||||
let m = &info["model"];
|
||||
ui.horizontal(|ui| {
|
||||
ui.label(RichText::new(m["name"].as_str().unwrap_or(file)).font(t.semibold(14.0)).color(t.text));
|
||||
ui.label(RichText::new(mb(info["sizeBytes"].as_u64())).color(t.text_dim));
|
||||
});
|
||||
if info["alreadyInstalled"] == true {
|
||||
ui.label(RichText::new("Already installed. Installing again is harmless.").color(t.text_dim));
|
||||
}
|
||||
let download = info["download"].is_string();
|
||||
let detector = m["role"] == "detector";
|
||||
if download {
|
||||
let host = info["host"].as_str().unwrap_or("its own repository");
|
||||
ui.add(
|
||||
egui::Label::new(
|
||||
RichText::new(format!(
|
||||
"Downloads from {host}. Once it has arrived and checked out, LightCraft installs it{}.",
|
||||
if detector { "" } else { ", starts using it and turns face recognition on" }
|
||||
))
|
||||
.color(t.text_label),
|
||||
)
|
||||
.wrap(),
|
||||
);
|
||||
}
|
||||
let commercial = m["licence"]["commercial"].as_str().unwrap_or("unknown");
|
||||
ui.label(RichText::new(licence_line(m)).font(t.semibold(12.5)).color(if commercial == "yes" { t.text } else { t.caution }));
|
||||
let notice = m["licence"]["notice"].as_str().unwrap_or("");
|
||||
if !notice.is_empty() {
|
||||
ui.add(egui::Label::new(RichText::new(notice).color(t.text_label)).wrap());
|
||||
}
|
||||
let provenance = m["provenance"].as_str().unwrap_or("");
|
||||
if !provenance.is_empty() {
|
||||
ui.add(egui::Label::new(RichText::new(format!("Trained on: {provenance}")).color(t.text_dim)).wrap());
|
||||
}
|
||||
if kind == "draft" {
|
||||
ui.add_space(2.0);
|
||||
ui.label(RichText::new("LightCraft does not know this model, so it assumed:").color(t.text_label));
|
||||
for a in info["assumptions"].as_array().into_iter().flatten().filter_map(Value::as_str) {
|
||||
ui.add(egui::Label::new(RichText::new(format!("• {a}")).font(t.font(12.0)).color(t.text_dim)).wrap());
|
||||
}
|
||||
}
|
||||
if let Some(url) = m["source"].as_str().or(m["licence"]["url"].as_str()) {
|
||||
let r = ui.link("Open the model's page");
|
||||
register(ui.ctx(), "faceModel:page", r.rect);
|
||||
if r.clicked() {
|
||||
open_page(app, url);
|
||||
}
|
||||
}
|
||||
ui.add_space(4.0);
|
||||
check(ui, "faceModel.accept", accepted, "I have read these terms and accept them for my own use");
|
||||
let closing = if download {
|
||||
"Nothing is fetched until you accept. The model is kept on this computer only; LightCraft never uploads or shares it."
|
||||
} else {
|
||||
"The model is kept on this computer only. LightCraft never uploads or shares it."
|
||||
};
|
||||
ui.label(RichText::new(closing).font(t.font(11.5)).color(t.text_dim));
|
||||
}
|
||||
|
||||
/// Pressing Download: the model's terms first, in the same dialog as a model file. Accepting them starts the download;
|
||||
/// nothing is fetched before.
|
||||
pub fn open_download_dialog(app: &mut LightcraftApp, m: &Value) {
|
||||
let info = json!({
|
||||
"kind": "known",
|
||||
"fileName": m["id"],
|
||||
"sizeBytes": m["sizeBytes"],
|
||||
"model": m,
|
||||
"alreadyInstalled": false,
|
||||
"download": m["id"],
|
||||
"host": m["downloadHost"],
|
||||
"assumptions": [],
|
||||
"reason": null,
|
||||
});
|
||||
app.ui.dialog = Some(crate::state::Dialog::FaceModel { path: String::new(), info, accepted: false });
|
||||
}
|
||||
|
||||
/// Photo ▸ Detect Faces without the detector: offer to download it (its terms first), the same as in Settings.
|
||||
pub fn offer_detector(app: &mut LightcraftApp) {
|
||||
let list = app.session.execute("faces.models.list", &json!({})).unwrap_or(Value::Null);
|
||||
if let Some(m) = list["models"].as_array().and_then(|a| a.iter().find(|m| m["role"] == "detector" && m["installed"] != true)) {
|
||||
open_download_dialog(app, m);
|
||||
}
|
||||
}
|
||||
|
||||
/// Inspect `path` and open the dialog for it (also what dropping a `.onnx` file on the window does).
|
||||
pub fn open_dialog(app: &mut LightcraftApp, path: &str) -> Result<Value, String> {
|
||||
let info = app.run("faces.models.inspect", json!({"path": path}))?;
|
||||
app.ui.dialog = Some(crate::state::Dialog::FaceModel { path: path.to_string(), info, accepted: false });
|
||||
Ok(Value::Null)
|
||||
}
|
||||
|
||||
/// The dialog's OK (the engine refuses without the acceptance): for a download, start it and watch it; for a file,
|
||||
/// install it. Either way the model ends up installed, in use, with recognition on.
|
||||
pub fn install(app: &mut LightcraftApp, path: &str, info: &Value, accepted: bool) -> Result<Value, String> {
|
||||
if !accepted {
|
||||
return Err("Tick the box to accept the model's terms first".into());
|
||||
}
|
||||
app.caches.faces_epoch += 1;
|
||||
let back_to_settings = |app: &mut LightcraftApp| app.ui.dialog = Some(crate::state::Dialog::Settings { tab: "faces".into() });
|
||||
if let Some(id) = info["download"].as_str() {
|
||||
let r = app.run("faces.models.download", json!({"id": id, "acknowledged": true}));
|
||||
if r.is_ok() {
|
||||
app.caches.faces_dl_watch.push(id.to_string());
|
||||
// the progress is in Settings: stay there
|
||||
back_to_settings(app);
|
||||
}
|
||||
return r;
|
||||
}
|
||||
let r = app.run("faces.models.install", json!({"path": path, "acknowledged": true}));
|
||||
if let Ok(v) = &r {
|
||||
let name = v["installed"]["name"].as_str().unwrap_or("The model");
|
||||
app.ui.status = if v["installed"]["role"] == "detector" {
|
||||
format!("{name} is installed")
|
||||
} else {
|
||||
format!("{name} is installed and in use: face recognition is on")
|
||||
};
|
||||
back_to_settings(app);
|
||||
}
|
||||
r
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------------------------------
|
||||
// Naming faces in the loupe
|
||||
|
||||
/// A suggestion for one unnamed face: the best guess (when it passes the bar) and the closest few people.
|
||||
#[derive(Clone, Debug, Default, PartialEq)]
|
||||
pub struct Hint {
|
||||
pub suggestion: Option<(String, f32)>,
|
||||
pub candidates: Vec<(String, f32)>,
|
||||
}
|
||||
|
||||
/// Suggestions for the unnamed faces of one photo, by region index.
|
||||
#[derive(Clone, Debug, Default, PartialEq)]
|
||||
pub struct Hints {
|
||||
pub by_index: std::collections::HashMap<usize, Hint>,
|
||||
}
|
||||
|
||||
/// Suggestions for `photo`'s unnamed faces from what is already indexed (nothing is embedded for this: the background
|
||||
/// indexer does that). `None` while recognition is off. Asked again only when the catalog or the index changed.
|
||||
pub fn hints_for(app: &mut LightcraftApp, photo: u64) -> Option<std::sync::Arc<Hints>> {
|
||||
if !app.caches.faces_active {
|
||||
return None;
|
||||
}
|
||||
let (rev, indexed) = (app.session.catalog.revision, app.caches.faces_indexed);
|
||||
if let Some((p, r, i, h)) = &app.caches.face_hints
|
||||
&& (*p, *r, *i) == (photo, rev, indexed)
|
||||
{
|
||||
return Some(h.clone());
|
||||
}
|
||||
let answer = app.session.execute("faces.suggest", &json!({"ids": [photo], "budgetMs": 0})).ok()?;
|
||||
let pair = |v: &Value| Some((v["name"].as_str()?.to_string(), v["score"].as_f64()? as f32));
|
||||
let mut hints = Hints::default();
|
||||
for f in answer["photos"][0]["faces"].as_array().into_iter().flatten() {
|
||||
let Some(index) = f["index"].as_u64().and_then(|i| usize::try_from(i).ok()) else { continue };
|
||||
hints.by_index.insert(
|
||||
index,
|
||||
Hint { suggestion: pair(&f["suggestion"]), candidates: f["candidates"].as_array().into_iter().flatten().filter_map(pair).collect() },
|
||||
);
|
||||
}
|
||||
let hints = std::sync::Arc::new(hints);
|
||||
app.caches.face_hints = Some((photo, rev, indexed, hints.clone()));
|
||||
Some(hints)
|
||||
}
|
||||
|
||||
/// What the name box did this frame.
|
||||
pub enum Editor {
|
||||
Open,
|
||||
Submit(String),
|
||||
Cancel,
|
||||
/// The offer to set up recognition was pressed.
|
||||
Setup,
|
||||
}
|
||||
|
||||
/// The inline name box under a face: type a name (completed from the people already named), pick one of the
|
||||
/// suggestions, Enter to confirm, Escape to cancel. While recognition is not set up it offers the next step (`setup`).
|
||||
pub fn name_editor(
|
||||
ctx: &egui::Context,
|
||||
at: egui::Pos2,
|
||||
edit: &mut crate::state::NameEdit,
|
||||
candidates: &[(String, f32)],
|
||||
people: &[String],
|
||||
setup: Setup,
|
||||
) -> Editor {
|
||||
let t = Tokens::get(ctx);
|
||||
let mut outcome = Editor::Open;
|
||||
let opening = edit.fresh;
|
||||
let shown = egui::Area::new(egui::Id::new("face-name-editor")).order(egui::Order::Foreground).fixed_pos(at).constrain(true).show(ctx, |ui| {
|
||||
egui::Frame::popup(ui.style()).show(ui, |ui| {
|
||||
ui.set_min_width(230.0);
|
||||
let r = ui.add(egui::TextEdit::singleline(&mut edit.text).hint_text("Name").desired_width(220.0));
|
||||
register(ui.ctx(), "field:faceName", r.rect);
|
||||
if edit.fresh {
|
||||
r.request_focus();
|
||||
edit.fresh = false;
|
||||
}
|
||||
if r.lost_focus() && ui.input(|i| i.key_pressed(egui::Key::Enter)) {
|
||||
outcome = Editor::Submit(edit.text.trim().to_string());
|
||||
}
|
||||
let typed = edit.text.trim().to_lowercase();
|
||||
let mut offered: Vec<(String, Option<f32>)> = candidates.iter().take(3).map(|(n, s)| (n.clone(), Some(*s))).collect();
|
||||
if !typed.is_empty() {
|
||||
for p in people.iter().filter(|p| p.to_lowercase().starts_with(&typed)) {
|
||||
if offered.len() >= 7 {
|
||||
break;
|
||||
}
|
||||
if !offered.iter().any(|(n, _)| n.eq_ignore_ascii_case(p)) {
|
||||
offered.push((p.clone(), None));
|
||||
}
|
||||
}
|
||||
}
|
||||
for (name, score) in offered {
|
||||
let label = match score {
|
||||
Some(s) => format!("{name} {:.0}%", (s * 100.0).max(0.0)),
|
||||
None => name.clone(),
|
||||
};
|
||||
let b = ui.add(egui::Button::new(RichText::new(label).color(t.text_label)).frame(false).min_size(egui::vec2(220.0, 0.0)));
|
||||
register(ui.ctx(), format!("faceName:pick:{name}"), b.rect);
|
||||
if b.clicked() {
|
||||
outcome = Editor::Submit(name);
|
||||
}
|
||||
}
|
||||
if let Some(label) = setup.short() {
|
||||
let b = ui.add(egui::Button::new(RichText::new(label).font(t.font(12.0)).color(t.accent)).frame(false));
|
||||
register(ui.ctx(), "faces:setup", b.rect);
|
||||
if b.clicked() {
|
||||
outcome = Editor::Setup;
|
||||
}
|
||||
}
|
||||
ui.label(RichText::new("Enter to confirm, Esc to cancel").font(t.font(11.0)).color(t.text_dim));
|
||||
});
|
||||
});
|
||||
if ctx.input(|i| i.key_pressed(egui::Key::Escape)) {
|
||||
outcome = Editor::Cancel;
|
||||
}
|
||||
// a click anywhere else closes it (not on the frame it opened, which is the click that opened it)
|
||||
if !opening && ctx.input(|i| i.pointer.any_pressed()) && ctx.input(|i| i.pointer.interact_pos()).is_some_and(|p| !shown.response.rect.contains(p))
|
||||
{
|
||||
outcome = Editor::Cancel;
|
||||
}
|
||||
outcome
|
||||
}
|
||||
|
||||
/// Called every frame: keeps the background face indexer going while recognition is on, and notes whether it is
|
||||
/// running, how many faces it has done and how many photos are left (the loupe's suggestions depend on it).
|
||||
pub fn pump(app: &mut LightcraftApp, ctx: &egui::Context) {
|
||||
watch_downloads(app, ctx);
|
||||
let now = ctx.input(|i| i.time);
|
||||
let (working, moving) = ctx.input(|i| {
|
||||
let key_or_scroll = i.events.iter().any(|e| matches!(e, egui::Event::Key { .. } | egui::Event::Text(_) | egui::Event::MouseWheel { .. }));
|
||||
(i.pointer.any_down() || key_or_scroll, i.pointer.is_moving())
|
||||
});
|
||||
if working {
|
||||
app.caches.last_input = now;
|
||||
}
|
||||
if moving || working {
|
||||
app.caches.last_move = now;
|
||||
}
|
||||
if now < app.caches.faces_next_pump {
|
||||
return;
|
||||
}
|
||||
let (pace, in_front) = scan_pace(app, ctx, now);
|
||||
app.caches.faces_pace = pace;
|
||||
app.caches.faces_in_front = in_front;
|
||||
let Ok(v) = app.session.execute("faces.pump", &json!({"pace": pace})) else { return };
|
||||
app.caches.faces_active = v["active"] == true;
|
||||
app.caches.faces_indexed = v["indexedFaces"].as_u64().unwrap_or(0);
|
||||
app.caches.faces_pending = v["pendingPhotos"].as_u64().unwrap_or(0);
|
||||
// the progress bar's whole: the most photos that were left at once since the scan last finished
|
||||
app.caches.faces_peak = if app.caches.faces_pending == 0 { 0 } else { app.caches.faces_peak.max(app.caches.faces_pending) };
|
||||
// Frames are drawn only when something asks for one, so nothing here wakes the window needlessly: while there is work it
|
||||
// asks to be called again in 50 ms (the progress in Settings, the next photos for the workers); with recognition on and
|
||||
// nothing to do it looks again every few seconds (a few wake-ups a minute); with recognition off it asks for nothing
|
||||
// (and is called on every frame, which is cheap: whatever makes a frame, such as switching recognition on, is seen at once).
|
||||
let busy = app.caches.faces_active && (app.caches.faces_pending > 0 || v["inFlight"].as_u64().unwrap_or(0) > 0);
|
||||
let wait = match (busy, app.caches.faces_active) {
|
||||
(true, _) => Some(0.05),
|
||||
(false, true) => Some(5.0),
|
||||
(false, false) => None,
|
||||
};
|
||||
app.caches.faces_next_pump = now
|
||||
+ if busy {
|
||||
0.05
|
||||
} else if app.caches.faces_active {
|
||||
1.0
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
if let Some(secs) = wait {
|
||||
ctx.request_repaint_after(std::time::Duration::from_secs_f64(secs));
|
||||
}
|
||||
}
|
||||
|
||||
/// Sort the watched downloads: those that have been installed (to announce), and those still to watch (running, or
|
||||
/// arrived and being installed). A failed or cancelled one is dropped: the model's row in Settings says what happened.
|
||||
fn due(rows: &[Value], watched: Vec<String>) -> (Vec<String>, Vec<String>) {
|
||||
let (mut installed, mut keep) = (Vec::new(), Vec::new());
|
||||
for id in watched {
|
||||
match rows.iter().find(|r| r["id"] == id.as_str()).and_then(|r| r["state"].as_str()) {
|
||||
Some("running" | "done") => keep.push(id),
|
||||
Some("installed") => installed.push(id),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
(installed, keep)
|
||||
}
|
||||
|
||||
/// Follows the downloads the user started: keeps redrawing while one runs (the progress bar), and says so when a
|
||||
/// model has been installed, is in use and recognition is on (the engine does all of that by itself).
|
||||
fn watch_downloads(app: &mut LightcraftApp, ctx: &egui::Context) {
|
||||
if app.caches.faces_dl_watch.is_empty() {
|
||||
return;
|
||||
}
|
||||
let Ok(v) = app.session.execute("faces.models.downloads", &json!({})) else {
|
||||
app.caches.faces_dl_watch.clear();
|
||||
return;
|
||||
};
|
||||
let rows: Vec<Value> = v["downloads"].as_array().cloned().unwrap_or_default();
|
||||
let (installed, keep) = due(&rows, std::mem::take(&mut app.caches.faces_dl_watch));
|
||||
app.caches.faces_dl_watch = keep;
|
||||
for id in installed {
|
||||
let name = app
|
||||
.session
|
||||
.execute("faces.models.list", &json!({}))
|
||||
.ok()
|
||||
.and_then(|l| {
|
||||
l["models"].as_array().and_then(|a| a.iter().find(|m| m["id"] == id.as_str()).and_then(|m| m["name"].as_str().map(str::to_string)))
|
||||
})
|
||||
.unwrap_or_else(|| id.clone());
|
||||
let _ = app.session.execute("faces.models.downloadCancel", &json!({"id": id}));
|
||||
app.caches.faces_epoch += 1;
|
||||
let detector = app
|
||||
.session
|
||||
.execute("faces.models.list", &json!({}))
|
||||
.ok()
|
||||
.is_some_and(|l| l["models"].as_array().is_some_and(|a| a.iter().any(|m| m["id"] == id.as_str() && m["role"] == "detector")));
|
||||
let text = if detector { format!("{name} is installed") } else { format!("{name} is installed and in use: face recognition is on") };
|
||||
app.ui.status = text.clone();
|
||||
app.toast(ctx, text);
|
||||
}
|
||||
ctx.request_repaint_after(std::time::Duration::from_millis(120));
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn the_scan_works_hard_only_when_nobody_is_in_the_way() {
|
||||
// (focused, minimized, seconds since a drag/key/scroll, seconds since the pointer moved, watching the progress)
|
||||
assert_eq!(pace_for(true, false, 0.1, 0.1, true), "pause", "dragging a slider: nothing new, even if watching");
|
||||
assert_eq!(pace_for(true, false, 2.0, 0.2, false), "light", "the pointer is moving: one photo at a time");
|
||||
assert_eq!(pace_for(true, false, 5.0, 2.9, true), "light");
|
||||
assert_eq!(pace_for(true, false, 5.0, 4.0, false), "normal", "idle for a few seconds: half the machine");
|
||||
assert_eq!(pace_for(true, false, 60.0, 60.0, true), "full", "idle and looking at the progress: most of it");
|
||||
// minimized: gentle, whatever the user did last
|
||||
assert_eq!([pace_for(true, true, 0.0, 0.0, true), pace_for(false, true, 99.0, 99.0, true)], ["light", "light"]);
|
||||
// on screen behind another app: the user is elsewhere, so the scan gets on with it, however recent the last input
|
||||
assert_eq!([pace_for(false, false, 0.0, 0.0, true), pace_for(false, false, 99.0, 99.0, false)], ["normal", "normal"]);
|
||||
}
|
||||
|
||||
fn row(id: &str, state: &str) -> Value {
|
||||
json!({"id": id, "state": state})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_installed_download_is_announced_once_and_the_rest_are_followed() {
|
||||
let rows = vec![row("a", "running"), row("b", "done"), row("c", "installed"), row("d", "failed"), row("e", "cancelled")];
|
||||
let watched = ["a", "b", "c", "d", "e", "gone"].map(String::from).to_vec();
|
||||
let (installed, keep) = due(&rows, watched);
|
||||
assert_eq!(installed, ["c"]);
|
||||
// running and just-arrived ones are watched; failed, cancelled and vanished ones are dropped
|
||||
assert_eq!(keep, ["a", "b"]);
|
||||
assert_eq!(due(&[], vec!["x".into()]), (vec![], vec![]));
|
||||
}
|
||||
}
|
||||
|
|
@ -7,6 +7,7 @@ pub mod crop_overlay;
|
|||
pub mod detail;
|
||||
pub mod dialogs;
|
||||
pub mod edit;
|
||||
pub mod faces;
|
||||
pub mod filterbar;
|
||||
pub mod grid;
|
||||
pub mod keymap;
|
||||
|
|
@ -15,6 +16,7 @@ pub mod library_problem;
|
|||
pub mod masking;
|
||||
pub mod notices;
|
||||
pub mod people;
|
||||
pub mod person;
|
||||
pub mod presets;
|
||||
pub mod profiles;
|
||||
pub mod right;
|
||||
|
|
@ -23,6 +25,7 @@ pub mod second;
|
|||
pub mod settings;
|
||||
pub mod strip;
|
||||
pub mod topbar;
|
||||
pub mod unnamed;
|
||||
|
||||
use egui::{Align2, Rect, pos2, vec2};
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,13 @@
|
|||
//! People: a card per person named on faces (read from XMP): a close-up of their largest face, the
|
||||
//! name and how many photos they are in. A click shows that person's photos in the grid.
|
||||
//! People: a card per person named on faces (read from XMP): a close-up of their face and how many photos they are in
|
||||
//! (a small number on the picture). A click opens that person's page (`person.rs`): only their cropped faces, and the
|
||||
//! faces that look like them. Below the people, the faces nobody has named yet (`unnamed.rs`), to be named a group at a
|
||||
//! time.
|
||||
//!
|
||||
//! Only the rows on screen ask for a face render (the engine caches them, memory and disk).
|
||||
//! The thumbnail-size slider in the bottom bar sizes the faces here too (with its own smaller and larger limits), and
|
||||
//! every face is cut from the same kind of box (the detector's, once the scan has looked at it), so faces are shown
|
||||
//! equally close. Only the rows on screen ask for a face render (the engine caches them, memory and disk).
|
||||
|
||||
use egui::{Align2, Color32, Rect, Sense, Stroke, StrokeKind, pos2, vec2};
|
||||
use egui::{Align2, Color32, Rect, RichText, Sense, Stroke, StrokeKind, pos2, vec2};
|
||||
use lightcraft_catalog::Person;
|
||||
use serde_json::json;
|
||||
|
||||
|
|
@ -11,25 +15,81 @@ use crate::LightcraftApp;
|
|||
use crate::theme::Tokens;
|
||||
use crate::widgets::register;
|
||||
|
||||
/// A card's face picture (points); the name and count sit below it.
|
||||
const CARD: f32 = 150.0;
|
||||
const LABEL_H: f32 = 44.0;
|
||||
const GAP: f32 = 16.0;
|
||||
const PAD: f32 = 20.0;
|
||||
const GAP: f32 = 14.0;
|
||||
pub(super) const PAD: f32 = 20.0;
|
||||
const HEADER_H: f32 = 44.0;
|
||||
/// The name under a card's picture.
|
||||
const NAME_H: f32 = 26.0;
|
||||
/// The unnamed-faces heading and its hints.
|
||||
const SECTION_H: f32 = 70.0;
|
||||
|
||||
/// `thumb` (the thumbnail slider, 90 to 480) as a position between `low` and `high`.
|
||||
fn between(thumb: f32, low: f32, high: f32) -> f32 {
|
||||
let f = (thumb - 90.0) / (480.0 - 90.0);
|
||||
low + if f.is_finite() { f.clamp(0.0, 1.0) } else { 0.4 } * (high - low)
|
||||
}
|
||||
|
||||
/// A person card's face picture (points) for the thumbnail slider's value.
|
||||
pub fn card_edge(thumb: f32) -> f32 {
|
||||
between(thumb, 84.0, 260.0)
|
||||
}
|
||||
|
||||
/// A face tile's picture (points), for the unnamed faces and a person's page.
|
||||
pub fn tile_edge(thumb: f32) -> f32 {
|
||||
between(thumb, 56.0, 180.0)
|
||||
}
|
||||
|
||||
/// How many `edge`-wide cells (with the gap between them) fit across `width`.
|
||||
pub(super) fn columns(width: f32, edge: f32) -> usize {
|
||||
(((width - PAD * 2.0 + GAP) / (edge + GAP)).floor() as usize).max(1)
|
||||
}
|
||||
|
||||
pub(super) const TILE_GAP: f32 = GAP;
|
||||
|
||||
/// `text` shortened with "…" until it is no wider than `max_w` in `font` (measured, not guessed).
|
||||
pub(super) fn fit(painter: &egui::Painter, text: &str, font: egui::FontId, max_w: f32) -> String {
|
||||
let width = |s: &str| painter.layout_no_wrap(s.to_string(), font.clone(), Color32::WHITE).size().x;
|
||||
if width(text) <= max_w {
|
||||
return text.to_string();
|
||||
}
|
||||
let chars: Vec<char> = text.chars().collect();
|
||||
(1..chars.len()).rev().map(|n| chars[..n].iter().collect::<String>() + "…").find(|s| width(s) <= max_w).unwrap_or_else(|| "…".to_string())
|
||||
}
|
||||
|
||||
pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui) {
|
||||
if let Some(name) = app.ui.person_page.clone() {
|
||||
return super::person::show(app, ui, &name);
|
||||
}
|
||||
let t = Tokens::get(ui.ctx());
|
||||
let now = ui.input(|i| i.time);
|
||||
let thumb = app.ui.thumb_size;
|
||||
let people = app.caches.people(&app.session.catalog, &app.session.filter);
|
||||
let (head, _) = ui.allocate_exact_size(vec2(ui.available_width(), HEADER_H), Sense::hover());
|
||||
ui.painter().text(pos2(head.left() + PAD, head.center().y), Align2::LEFT_CENTER, crate::i18n::tr("Named People"), t.semibold(15.0), t.text);
|
||||
let title =
|
||||
ui.painter().text(pos2(head.left() + PAD, head.center().y), Align2::LEFT_CENTER, crate::i18n::tr("Named People"), t.semibold(15.0), t.text);
|
||||
// the person page last left (a face opened from it, or its back button): one click returns to it
|
||||
if let Some(last) = app.ui.last_person.clone().filter(|n| people.iter().any(|p| &p.name == n)) {
|
||||
let label = fit(ui.painter(), &format!("{last} ›"), t.font(13.0), 240.0);
|
||||
let width = ui.painter().layout_no_wrap(label.clone(), t.font(13.0), Color32::WHITE).size().x + 24.0;
|
||||
let r = Rect::from_min_size(pos2(title.right() + 14.0, head.center().y - 12.0), vec2(width, 24.0));
|
||||
let b = ui.put(r, egui::Button::new(RichText::new(label).font(t.font(13.0))).wrap_mode(egui::TextWrapMode::Extend));
|
||||
register(ui.ctx(), "people:last", b.rect);
|
||||
if b.on_hover_text(crate::i18n::tr("Back to this person")).clicked() {
|
||||
let _ = app.run("view.person", json!({"name": last}));
|
||||
}
|
||||
}
|
||||
ui.painter().text(pos2(head.right() - PAD, head.center().y), Align2::RIGHT_CENTER, people.len().to_string(), t.font(13.0), t.text_dim);
|
||||
// until face recognition is set up: what it takes, and a button that does the next step
|
||||
super::faces::setup_banner(app, ui);
|
||||
// the filters narrowing the list (a date, a keyword…), removable here
|
||||
let chips = lightcraft_engine::filter_chips(&app.session.filter, &app.session.catalog);
|
||||
super::chips::show(app, ui, &chips);
|
||||
if people.is_empty() {
|
||||
// what is selected below, and the box to name it with
|
||||
super::unnamed::naming_bar(app, ui);
|
||||
let unnamed = super::unnamed::faces(app, now);
|
||||
if people.is_empty() && unnamed.faces.is_empty() {
|
||||
let (title, body) = if chips.is_empty() {
|
||||
("No named people yet", "Face names written to XMP by Lightroom and other apps show up here.")
|
||||
("No people yet", "Face names written to XMP by Lightroom and other apps show up here, and faces found in your photos can be named.")
|
||||
} else {
|
||||
("No named people in these photos", "Remove a filter above, or choose Clear all")
|
||||
};
|
||||
|
|
@ -38,50 +98,108 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui) {
|
|||
}
|
||||
let ppp = ui.ctx().pixels_per_point();
|
||||
let active = app.session.filter.person.clone();
|
||||
let (edge, tile) = (card_edge(thumb), tile_edge(thumb));
|
||||
egui::ScrollArea::vertical().auto_shrink(false).show_viewport(ui, |ui, viewport| {
|
||||
let width = ui.available_width();
|
||||
let cols = (((width - PAD * 2.0 + GAP) / (CARD + GAP)).floor() as usize).max(1);
|
||||
let row_h = CARD + LABEL_H + GAP;
|
||||
let cols = columns(width, edge);
|
||||
let row_h = edge + NAME_H + GAP;
|
||||
let rows = people.len().div_ceil(cols);
|
||||
let (area, _) = ui.allocate_exact_size(vec2(width, PAD * 2.0 + rows as f32 * row_h), Sense::hover());
|
||||
let first = ((viewport.top() - PAD) / row_h).floor().max(0.0) as usize;
|
||||
let last = (((viewport.bottom() - PAD) / row_h).ceil().max(0.0) as usize).min(rows);
|
||||
for row in first..last {
|
||||
let named_h = if people.is_empty() { 40.0 } else { PAD + rows as f32 * row_h };
|
||||
let (tile_cols, tile_row) = (columns(width, tile), tile + GAP);
|
||||
let tile_rows = unnamed.faces.len().div_ceil(tile_cols);
|
||||
let section_top = named_h + 6.0;
|
||||
let (area, _) = ui.allocate_exact_size(vec2(width, section_top + SECTION_H + tile_rows as f32 * tile_row + PAD), Sense::hover());
|
||||
let visible = |top: f32, rows: usize, row_h: f32| {
|
||||
let first = ((viewport.top() - top) / row_h).floor().max(0.0) as usize;
|
||||
let last = (((viewport.bottom() - top) / row_h).ceil().max(0.0) as usize).min(rows);
|
||||
first..last
|
||||
};
|
||||
if people.is_empty() {
|
||||
ui.painter().text(
|
||||
pos2(area.left() + PAD, area.top() + 24.0),
|
||||
Align2::LEFT_CENTER,
|
||||
"No named people yet. Name the faces below to start.",
|
||||
t.font(13.0),
|
||||
t.text_dim,
|
||||
);
|
||||
}
|
||||
let mut shown = 0;
|
||||
for row in visible(PAD, rows, row_h) {
|
||||
for col in 0..cols {
|
||||
let Some(person) = people.get(row * cols + col) else { break };
|
||||
let min = area.min + vec2(PAD + col as f32 * (CARD + GAP), PAD + row as f32 * row_h);
|
||||
shown += 1;
|
||||
let min = area.min + vec2(PAD + col as f32 * (edge + GAP), PAD + row as f32 * row_h);
|
||||
let selected = active.as_deref().is_some_and(|a| a.eq_ignore_ascii_case(&person.name));
|
||||
card(app, ui, person, Rect::from_min_size(min, vec2(CARD, CARD + LABEL_H)), ppp, selected);
|
||||
card(app, ui, person, Rect::from_min_size(min, vec2(edge, edge + NAME_H)), edge, ppp, selected);
|
||||
}
|
||||
}
|
||||
// the faces nobody has named
|
||||
let top = area.top() + section_top;
|
||||
ui.painter().line_segment([pos2(area.left() + PAD, top), pos2(area.right() - PAD, top)], Stroke::new(1.0, t.button_border));
|
||||
super::unnamed::section_header(app, ui, &unnamed, pos2(area.left() + PAD, top + 14.0), area.right() - PAD);
|
||||
let grid_top = section_top + SECTION_H;
|
||||
for row in visible(grid_top, tile_rows, tile_row) {
|
||||
for col in 0..tile_cols {
|
||||
let i = row * tile_cols + col;
|
||||
let Some(face) = unnamed.faces.get(i) else { break };
|
||||
shown += 1;
|
||||
let min = area.min + vec2(PAD + col as f32 * (tile + GAP), grid_top + row as f32 * tile_row);
|
||||
super::unnamed::tile(app, ui, &unnamed, i, Rect::from_min_size(min, vec2(tile, tile)), face, ppp);
|
||||
}
|
||||
}
|
||||
// the picture cache keeps all of them (a fixed budget would evict and re-request the same few every frame)
|
||||
app.renderer.want_variants(shown);
|
||||
});
|
||||
}
|
||||
|
||||
fn card(app: &mut LightcraftApp, ui: &mut egui::Ui, person: &Person, r: Rect, ppp: f32, selected: bool) {
|
||||
fn card(app: &mut LightcraftApp, ui: &mut egui::Ui, person: &Person, r: Rect, edge: f32, ppp: f32, selected: bool) {
|
||||
let t = Tokens::get(ui.ctx());
|
||||
let face = Rect::from_min_size(r.min, vec2(CARD, CARD));
|
||||
let face = Rect::from_min_size(r.min, vec2(edge, edge));
|
||||
let resp = ui.interact(r, egui::Id::new(("person-card", &person.name)), Sense::click());
|
||||
register(ui.ctx(), format!("person:{}", person.name), r);
|
||||
let view = app.session.face_view(person.photo, person.face);
|
||||
let p = ui.painter();
|
||||
p.rect_filled(face, 3.0, t.canvas);
|
||||
if let Some(job) = app.session.face_job(person.photo, person.face, (CARD * ppp).ceil() as usize)
|
||||
if let Some(job) = app.session.face_job(person.photo, view, (edge * ppp).ceil() as usize)
|
||||
&& let Some(tex) = app.renderer.variant(job)
|
||||
{
|
||||
p.image(tex.tex.id(), face, Rect::from_min_max(pos2(0.0, 0.0), pos2(1.0, 1.0)), Color32::WHITE);
|
||||
}
|
||||
// how many photos they are in, on the picture
|
||||
let count = if person.count > 999 { "999+".to_string() } else { person.count.to_string() };
|
||||
let g = p.layout_no_wrap(count, t.font(11.5), Color32::WHITE);
|
||||
let pill = Rect::from_min_size(pos2(face.right() - g.size().x - 14.0, face.bottom() - g.size().y - 11.0), g.size() + vec2(10.0, 6.0));
|
||||
p.rect_filled(pill, 4.0, Color32::from_black_alpha(175));
|
||||
p.galley(pill.min + vec2(5.0, 3.0), g, Color32::WHITE);
|
||||
if selected {
|
||||
p.rect_stroke(face, 3.0, Stroke::new(2.0, Color32::WHITE), StrokeKind::Outside);
|
||||
} else if resp.hovered() {
|
||||
p.rect_stroke(face, 3.0, Stroke::new(1.0, t.text_dim), StrokeKind::Outside);
|
||||
}
|
||||
// a long name must not run past the card
|
||||
let name =
|
||||
if person.name.chars().count() > 19 { format!("{}…", person.name.chars().take(18).collect::<String>()) } else { person.name.clone() };
|
||||
let name = fit(p, &person.name, t.semibold(13.0), edge);
|
||||
p.text(pos2(r.left() + 2.0, face.bottom() + 14.0), Align2::LEFT_CENTER, name, t.semibold(13.0), t.text);
|
||||
let photos = crate::i18n::tr_format!("{n} photo{}", if person.count == 1 { "" } else { "s" }, n = person.count);
|
||||
p.text(pos2(r.left() + 2.0, face.bottom() + 32.0), Align2::LEFT_CENTER, photos, t.font(12.0), t.text_dim);
|
||||
if resp.on_hover_text(crate::i18n::tr_format!("{} — show their photos", person.name)).clicked() {
|
||||
let _ = app.run("library.filter", json!({"person": person.name}));
|
||||
let _ = app.run("view.photoGrid", json!({}));
|
||||
let _ = app.run("view.person", json!({"name": person.name}));
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn faces_follow_the_thumbnail_slider_within_their_own_limits() {
|
||||
assert_eq!((card_edge(90.0), card_edge(480.0)), (84.0, 260.0));
|
||||
assert_eq!((tile_edge(90.0), tile_edge(480.0)), (56.0, 180.0));
|
||||
// between, growing with the slider; beyond its ends, held at the limits; nonsense, a middle size
|
||||
assert!(card_edge(220.0) > card_edge(150.0) && tile_edge(300.0) > tile_edge(220.0));
|
||||
assert_eq!((card_edge(-5.0), card_edge(10_000.0)), (84.0, 260.0));
|
||||
assert!(card_edge(f32::NAN) > 84.0 && card_edge(f32::NAN) < 260.0);
|
||||
// there is always at least one column, whatever the width
|
||||
assert_eq!((columns(0.0, 260.0), columns(-50.0, 260.0), columns(f32::NAN, 100.0)), (1, 1, 1));
|
||||
assert!(columns(1400.0, 84.0) > columns(1400.0, 260.0));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,304 @@
|
|||
//! A person's page in the People view: only cropped faces, never whole photos. First the faces already named, then a
|
||||
//! "More" section with the unnamed faces that look like them (best first), to be confirmed with a click.
|
||||
//!
|
||||
//! The engine does the looking (`faces.person`); this draws it. Only the rows on screen ask for a face render.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use egui::{Align2, Color32, Pos2, Rect, RichText, Sense, Stroke, StrokeKind, pos2, vec2};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::LightcraftApp;
|
||||
use crate::theme::Tokens;
|
||||
use crate::widgets::register;
|
||||
|
||||
const PAD: f32 = 20.0;
|
||||
const SECTION_H: f32 = 56.0;
|
||||
|
||||
/// One face: which photo, which of its regions, the box to show it by, and (for the "More" faces) how alike.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct Face {
|
||||
pub photo: u64,
|
||||
pub index: usize,
|
||||
pub rect: lightcraft_geom::Rect,
|
||||
pub score: f32,
|
||||
}
|
||||
|
||||
/// What `faces.person` said about one person.
|
||||
#[derive(Clone, Debug, Default, PartialEq)]
|
||||
pub struct PersonPage {
|
||||
pub name: String,
|
||||
pub total: usize,
|
||||
pub confirmed: Vec<Face>,
|
||||
pub more: Vec<Face>,
|
||||
/// Recognition is on, so an empty "More" means "none yet" and not "not looking".
|
||||
pub ready: bool,
|
||||
pub pending: usize,
|
||||
}
|
||||
|
||||
fn face_of(v: &Value) -> Option<Face> {
|
||||
let rect = |r: &Value| Some(lightcraft_geom::Rect { x0: r["x0"].as_f64()?, y0: r["y0"].as_f64()?, x1: r["x1"].as_f64()?, y1: r["y1"].as_f64()? });
|
||||
Some(Face {
|
||||
photo: v["photo"].as_u64()?,
|
||||
index: usize::try_from(v["index"].as_u64()?).ok()?,
|
||||
// shown by the detector's box when the scan has one: every face equally close
|
||||
rect: rect(&v["view"]).or_else(|| rect(&v["rect"]))?,
|
||||
score: v["score"].as_f64().unwrap_or(0.0) as f32,
|
||||
})
|
||||
}
|
||||
|
||||
/// The engine's answer as a page (anything that does not make sense is left out).
|
||||
pub fn parse(v: &Value) -> PersonPage {
|
||||
let faces = |k: &str| v[k].as_array().map(|a| a.iter().filter_map(face_of).collect()).unwrap_or_default();
|
||||
PersonPage {
|
||||
name: v["name"].as_str().unwrap_or("").to_string(),
|
||||
total: v["total"].as_u64().map_or(0, |n| usize::try_from(n).unwrap_or(0)),
|
||||
confirmed: faces("confirmed"),
|
||||
more: faces("more"),
|
||||
ready: v["ready"] == true,
|
||||
pending: v["pendingPhotos"].as_u64().map_or(0, |n| usize::try_from(n).unwrap_or(0)),
|
||||
}
|
||||
}
|
||||
|
||||
/// The page for `name`, asked for again only when the catalog changed, or (while the scan runs) at most once a second
|
||||
/// as it learns more faces.
|
||||
fn page_for(app: &mut LightcraftApp, name: &str, now: f64) -> Arc<PersonPage> {
|
||||
let (rev, indexed) = (app.session.catalog.revision, app.caches.faces_indexed);
|
||||
if let Some((n, r, i, at, page)) = &app.caches.person_page
|
||||
&& n == name
|
||||
&& *r == rev
|
||||
&& (*i == indexed || now - *at < 1.0)
|
||||
{
|
||||
return page.clone();
|
||||
}
|
||||
let v = app.session.execute("faces.person", &json!({"name": name})).unwrap_or(Value::Null);
|
||||
let page = Arc::new(parse(&v));
|
||||
app.caches.person_page = Some((name.to_string(), rev, indexed, now, page.clone()));
|
||||
page
|
||||
}
|
||||
|
||||
#[derive(PartialEq)]
|
||||
enum Hit {
|
||||
Nothing,
|
||||
Open,
|
||||
Confirm,
|
||||
Dismiss,
|
||||
}
|
||||
|
||||
pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, name: &str) {
|
||||
let t = Tokens::get(ui.ctx());
|
||||
let now = ui.input(|i| i.time);
|
||||
let page = page_for(app, name, now);
|
||||
let shown_name = if page.name.is_empty() { name.to_string() } else { page.name.clone() };
|
||||
ui.add_space(8.0);
|
||||
ui.horizontal(|ui| {
|
||||
ui.add_space(PAD - 4.0);
|
||||
let back = ui.button("‹ People");
|
||||
register(ui.ctx(), "person:back", back.rect);
|
||||
if back.clicked() {
|
||||
let _ = app.run("view.people", json!({}));
|
||||
}
|
||||
ui.label(RichText::new(&shown_name).font(t.semibold(15.0)).color(t.text));
|
||||
let count = if page.total == 1 { "1 face".to_string() } else { format!("{} faces", page.total) };
|
||||
ui.label(RichText::new(count).color(t.text_dim));
|
||||
ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| {
|
||||
ui.add_space(PAD - 4.0);
|
||||
let photos = ui.button("Show photos").on_hover_text("Their photos in the grid");
|
||||
register(ui.ctx(), "person:photos", photos.rect);
|
||||
if photos.clicked() {
|
||||
let _ = app.run("library.filter", json!({"person": shown_name}));
|
||||
let _ = app.run("view.photoGrid", json!({}));
|
||||
}
|
||||
});
|
||||
});
|
||||
ui.add_space(6.0);
|
||||
if page.total == 0 {
|
||||
super::empty_message(
|
||||
ui,
|
||||
ui.available_rect_before_wrap(),
|
||||
"No faces by that name",
|
||||
"They may have been renamed or removed. Go back to People",
|
||||
);
|
||||
return;
|
||||
}
|
||||
let dismissed = &app.ui.dismissed_faces;
|
||||
let more: Vec<&Face> = page.more.iter().filter(|f| !dismissed.contains(&(f.photo, f.index))).collect();
|
||||
// what the "More" section says when it has no faces to show
|
||||
let recognising = app.caches.faces_active || page.ready;
|
||||
// with a model installed the button switches recognition on; without one it opens Settings ▸ Faces
|
||||
let step = if recognising { super::faces::Setup::Running } else { super::faces::setup(app, ui.ctx()) };
|
||||
let note: Option<String> = if step == super::faces::Setup::TurnOn {
|
||||
Some(format!("Face recognition is off. Turn it on to find more photos of {shown_name}."))
|
||||
} else if !recognising {
|
||||
Some(format!("Face recognition needs a model, in Settings, to find more photos of {shown_name}."))
|
||||
} else if more.is_empty() && page.pending > 0 {
|
||||
Some(format!("Looking through your photos… {} left", page.pending))
|
||||
} else if more.is_empty() {
|
||||
Some(format!("No other faces look like {shown_name} yet. Faces appear here as your photos are scanned."))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let ppp = ui.ctx().pixels_per_point();
|
||||
let edge = super::people::tile_edge(app.ui.thumb_size);
|
||||
let mut hit: Option<(Hit, Face)> = None;
|
||||
egui::ScrollArea::vertical().auto_shrink(false).show_viewport(ui, |ui, viewport| {
|
||||
let width = ui.available_width();
|
||||
let cols = super::people::columns(width, edge);
|
||||
let row_h = edge + super::people::TILE_GAP;
|
||||
let confirmed_rows = page.confirmed.len().div_ceil(cols);
|
||||
let more_rows = more.len().div_ceil(cols);
|
||||
let more_top = PAD + confirmed_rows as f32 * row_h + 12.0;
|
||||
let total_h = more_top + SECTION_H + more_rows as f32 * row_h + PAD;
|
||||
let (area, _) = ui.allocate_exact_size(vec2(width, total_h), Sense::hover());
|
||||
let visible = |top: f32, rows: usize| {
|
||||
let first = ((viewport.top() - top) / row_h).floor().max(0.0) as usize;
|
||||
let last = (((viewport.bottom() - top) / row_h).ceil().max(0.0) as usize).min(rows);
|
||||
first..last
|
||||
};
|
||||
let mut shown = 0;
|
||||
for row in visible(PAD, confirmed_rows) {
|
||||
for col in 0..cols {
|
||||
let Some(f) = page.confirmed.get(row * cols + col) else { break };
|
||||
shown += 1;
|
||||
let min = area.min + vec2(PAD + col as f32 * (edge + super::people::TILE_GAP), PAD + row as f32 * row_h);
|
||||
if let Hit::Open = tile(app, ui, Rect::from_min_size(min, vec2(edge, edge)), f, ppp, false, &shown_name) {
|
||||
hit = Some((Hit::Open, f.clone()));
|
||||
}
|
||||
}
|
||||
}
|
||||
// the "More" section
|
||||
let heading = pos2(area.left() + PAD, area.top() + more_top + 14.0);
|
||||
ui.painter().text(heading, Align2::LEFT_CENTER, "More", t.semibold(15.0), t.text);
|
||||
let sub = match ¬e {
|
||||
Some(n) => n.clone(),
|
||||
None => format!("Faces that look like {shown_name}. Click one to confirm it, or × to hide it."),
|
||||
};
|
||||
ui.painter().text(pos2(area.left() + PAD, area.top() + more_top + 36.0), Align2::LEFT_CENTER, sub, t.font(12.0), t.text_dim);
|
||||
if !recognising {
|
||||
let r = Rect::from_min_size(pos2(area.left() + PAD + 330.0, area.top() + more_top + 8.0), vec2(120.0, 24.0));
|
||||
let label = if step == super::faces::Setup::TurnOn { "Turn on" } else { "Open Settings" };
|
||||
let b = ui.put(r, egui::Button::new(label));
|
||||
register(ui.ctx(), "faces:setup", b.rect);
|
||||
if b.clicked() {
|
||||
if step == super::faces::Setup::TurnOn {
|
||||
super::faces::take_step(app, ui.ctx(), step);
|
||||
} else {
|
||||
let _ = app.run("app.settings", json!({"tab": "faces"}));
|
||||
}
|
||||
}
|
||||
}
|
||||
let grid_top = more_top + SECTION_H;
|
||||
for row in visible(grid_top, more_rows) {
|
||||
for col in 0..cols {
|
||||
let Some(f) = more.get(row * cols + col) else { break };
|
||||
shown += 1;
|
||||
let min = area.min + vec2(PAD + col as f32 * (edge + super::people::TILE_GAP), grid_top + row as f32 * row_h);
|
||||
match tile(app, ui, Rect::from_min_size(min, vec2(edge, edge)), f, ppp, true, &shown_name) {
|
||||
Hit::Confirm => hit = Some((Hit::Confirm, (*f).clone())),
|
||||
Hit::Dismiss => hit = Some((Hit::Dismiss, (*f).clone())),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
// the picture cache keeps every face on screen
|
||||
app.renderer.want_variants(shown);
|
||||
});
|
||||
if let Some((what, f)) = hit {
|
||||
match what {
|
||||
Hit::Open => {
|
||||
let _ = app.run("library.select", json!({"ids": [f.photo], "active": f.photo}));
|
||||
let _ = app.run("view.detail", json!({}));
|
||||
// Escape comes back to this page; the People button to everyone, with this person next to its title
|
||||
app.ui.person_from = Some((shown_name.clone(), f.photo));
|
||||
app.ui.last_person = Some(shown_name.clone());
|
||||
app.ui.person_page = None;
|
||||
}
|
||||
Hit::Confirm => {
|
||||
let _ = app.run("faces.setName", json!({"id": f.photo, "index": f.index, "name": shown_name}));
|
||||
}
|
||||
Hit::Dismiss => {
|
||||
app.ui.dismissed_faces.insert((f.photo, f.index));
|
||||
}
|
||||
Hit::Nothing => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// One face tile. A confirmed one opens its photo when clicked; a "More" one is confirmed when clicked, and has a × that
|
||||
/// hides it.
|
||||
fn tile(app: &mut LightcraftApp, ui: &mut egui::Ui, r: Rect, f: &Face, ppp: f32, more: bool, name: &str) -> Hit {
|
||||
let t = Tokens::get(ui.ctx());
|
||||
let salt = if more { "more-face" } else { "person-face" };
|
||||
let resp = ui.interact(r, egui::Id::new((salt, f.photo, f.index)), Sense::click());
|
||||
register(ui.ctx(), format!("{salt}:{}:{}", f.photo, f.index), r);
|
||||
// the × is added after the tile, so it wins where they overlap
|
||||
let x_rect = Rect::from_min_size(pos2(r.right() - 24.0, r.top() + 4.0), vec2(20.0, 20.0));
|
||||
let x_resp = more.then(|| ui.interact(x_rect, egui::Id::new(("more-face-x", f.photo, f.index)), Sense::click()));
|
||||
let hovered = resp.hovered() || x_resp.as_ref().is_some_and(|x| x.hovered());
|
||||
let p = ui.painter();
|
||||
p.rect_filled(r, 3.0, t.canvas);
|
||||
if let Some(job) = app.session.face_job(lightcraft_catalog::PhotoId(f.photo), f.rect, (r.width() * ppp).ceil() as usize)
|
||||
&& let Some(tex) = app.renderer.variant(job)
|
||||
{
|
||||
p.image(tex.tex.id(), r, Rect::from_min_max(pos2(0.0, 0.0), pos2(1.0, 1.0)), Color32::WHITE);
|
||||
}
|
||||
if hovered {
|
||||
p.rect_stroke(r, 3.0, Stroke::new(1.5, Color32::WHITE), StrokeKind::Outside);
|
||||
}
|
||||
let mut hit = Hit::Nothing;
|
||||
if more && hovered {
|
||||
// a check at the bottom right says a click confirms; the × at the top right hides the face
|
||||
let badge = Rect::from_min_size(pos2(r.right() - 28.0, r.bottom() - 28.0), vec2(24.0, 24.0));
|
||||
p.circle_filled(badge.center(), 12.0, Color32::from_rgba_unmultiplied(30, 30, 30, 220));
|
||||
let c = badge.center();
|
||||
let check = [pos2(c.x - 5.0, c.y), pos2(c.x - 1.5, c.y + 4.0), pos2(c.x + 5.5, c.y - 4.0)];
|
||||
p.line_segment([check[0], check[1]], Stroke::new(2.0, Color32::WHITE));
|
||||
p.line_segment([check[1], check[2]], Stroke::new(2.0, Color32::WHITE));
|
||||
p.circle_filled(x_rect.center(), 10.0, Color32::from_rgba_unmultiplied(30, 30, 30, 220));
|
||||
let xc: Pos2 = x_rect.center();
|
||||
p.line_segment([pos2(xc.x - 3.5, xc.y - 3.5), pos2(xc.x + 3.5, xc.y + 3.5)], Stroke::new(1.6, Color32::WHITE));
|
||||
p.line_segment([pos2(xc.x - 3.5, xc.y + 3.5), pos2(xc.x + 3.5, xc.y - 3.5)], Stroke::new(1.6, Color32::WHITE));
|
||||
}
|
||||
if more {
|
||||
let tip = format!("Looks like {name} (similarity {:.2}). Click to confirm.", f.score);
|
||||
let _ = resp.clone().on_hover_text(tip);
|
||||
if x_resp.as_ref().is_some_and(|x| x.clicked()) {
|
||||
hit = Hit::Dismiss;
|
||||
} else if resp.clicked() {
|
||||
hit = Hit::Confirm;
|
||||
}
|
||||
} else {
|
||||
let _ = resp.clone().on_hover_text(format!("{name} — click to open this photo"));
|
||||
if resp.clicked() {
|
||||
hit = Hit::Open;
|
||||
}
|
||||
}
|
||||
hit
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn the_engines_answer_becomes_a_page_and_odd_entries_are_left_out() {
|
||||
let v = json!({
|
||||
"name": "Jane Doe", "total": 2, "ready": true, "pendingPhotos": 7,
|
||||
"confirmed": [
|
||||
{"photo": 3, "index": 0, "rect": {"x0": 0.1, "y0": 0.2, "x1": 0.3, "y1": 0.5}},
|
||||
{"photo": 4, "index": 1},
|
||||
{"photo": "x", "index": 0, "rect": {"x0": 0.1, "y0": 0.2, "x1": 0.3, "y1": 0.5}},
|
||||
],
|
||||
"more": [{"photo": 9, "index": 2, "score": 0.5, "rect": {"x0": 0.0, "y0": 0.0, "x1": 1.0, "y1": 1.0}}],
|
||||
});
|
||||
let page = parse(&v);
|
||||
assert_eq!((page.name.as_str(), page.total, page.ready, page.pending), ("Jane Doe", 2, true, 7));
|
||||
assert_eq!(page.confirmed.len(), 1, "a face without a box or a photo is dropped");
|
||||
assert_eq!((page.confirmed[0].photo, page.confirmed[0].index), (3, 0));
|
||||
assert_eq!((page.more.len(), page.more[0].score), (1, 0.5));
|
||||
// nothing, or nonsense, is an empty page and not a crash
|
||||
for bad in [Value::Null, json!([]), json!({"confirmed": "x", "more": 5, "total": -1})] {
|
||||
assert_eq!(parse(&bad), PersonPage::default(), "{bad}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,4 +1,4 @@
|
|||
//! The Settings dialog (⌘,): General, Import, Performance, Interface.
|
||||
//! The Settings dialog (⌘,): General, Import, Performance, Interface, Faces.
|
||||
//!
|
||||
//! Changes apply immediately (no OK/Cancel). Where they are stored:
|
||||
//! - **app settings** ([`crate::state::AppSettings`]: startup view, delete confirmation, GPU,
|
||||
|
|
@ -17,7 +17,8 @@ use crate::theme::Tokens;
|
|||
use crate::widgets::register;
|
||||
|
||||
/// (id, label) of the tabs, in order.
|
||||
pub const TABS: &[(&str, &str)] = &[("general", "General"), ("import", "Import"), ("performance", "Performance"), ("interface", "Interface")];
|
||||
pub const TABS: &[(&str, &str)] =
|
||||
&[("general", "General"), ("import", "Import"), ("performance", "Performance"), ("interface", "Interface"), ("faces", "Faces")];
|
||||
|
||||
/// Thumbnail cache sizes offered (MB).
|
||||
const CACHE_SIZES: [u32; 5] = [512, 1024, 2048, 4096, 8192];
|
||||
|
|
@ -43,11 +44,12 @@ pub fn body(app: &mut LightcraftApp, ui: &mut egui::Ui, tab: &mut String) {
|
|||
"import" => import_tab(app, ui, &t),
|
||||
"performance" => performance_tab(app, ui, &t),
|
||||
"interface" => interface_tab(app, ui, &t),
|
||||
"faces" => super::faces::settings_tab(app, ui, &t),
|
||||
_ => general_tab(app, ui, &t),
|
||||
}
|
||||
}
|
||||
|
||||
fn heading(ui: &mut egui::Ui, t: &Tokens, text: &str) {
|
||||
pub(super) fn heading(ui: &mut egui::Ui, t: &Tokens, text: &str) {
|
||||
ui.add_space(4.0);
|
||||
ui.label(RichText::new(crate::i18n::tr(text)).font(t.semibold(12.5)).color(t.text));
|
||||
}
|
||||
|
|
@ -63,12 +65,12 @@ fn row<R>(ui: &mut egui::Ui, t: &Tokens, label: &str, add: impl FnOnce(&mut egui
|
|||
.inner
|
||||
}
|
||||
|
||||
fn hint(ui: &mut egui::Ui, t: &Tokens, text: &str) {
|
||||
pub(super) fn hint(ui: &mut egui::Ui, t: &Tokens, text: &str) {
|
||||
ui.label(RichText::new(crate::i18n::tr(text)).size(11.0).color(t.text_dim));
|
||||
}
|
||||
|
||||
/// A checkbox addressable as `check:{id}`; true when toggled.
|
||||
fn check(ui: &mut egui::Ui, id: &str, value: &mut bool, label: &str) -> bool {
|
||||
pub(super) fn check(ui: &mut egui::Ui, id: &str, value: &mut bool, label: &str) -> bool {
|
||||
let r = ui.checkbox(value, crate::i18n::tr(label));
|
||||
register(ui.ctx(), format!("check:{id}"), r.rect);
|
||||
r.changed()
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in New Issue