diff --git a/Cargo.lock b/Cargo.lock index 86d4bca6..397094ed 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -209,7 +209,7 @@ version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc890384c8602f339876ded803c97ad529f3842aba97f6392b3dba0dd171769b" dependencies = [ - "equator", + "equator 0.4.2", ] [[package]] @@ -967,6 +967,12 @@ dependencies = [ "unicode-width", ] +[[package]] +name = "coe-rs" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e8f1e641542c07631228b1e0dc04b69ae3c1d58ef65d5691a439711d805c698" + [[package]] name = "color" version = "0.3.3" @@ -1256,6 +1262,16 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d8b14ccef22fc6f5a8f4d7d768562a182c04ce9a3b3157b91390b52ddfdf1a76" +[[package]] +name = "dyn-stack" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56e53799688f5632f364f8fb387488dd05db9fe45db7011be066fc20e7027f8b" +dependencies = [ + "bytemuck", + "reborrow", +] + [[package]] name = "dyn-stack" version = "0.13.2" @@ -1537,13 +1553,33 @@ version = "0.36.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773fa9c96dd0dbef887e39d0ed6177f141cce4d68e6041df77570aa3702dfa13" +[[package]] +name = "equator" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3b0a88aa91d0ad2b9684e4479aed31a17d3f9051bdbbc634bd2c01bc5a5eee8" +dependencies = [ + "equator-macro 0.1.9", +] + [[package]] name = "equator" version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4711b213838dfee0117e3be6ac926007d7f433d7bbe33595975d4190cb07e6fc" dependencies = [ - "equator-macro", + "equator-macro 0.4.2", +] + +[[package]] +name = "equator-macro" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60d08acb9849f7fb4401564f251be5a526829183a3645a90197dea8e786cf3ae" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", ] [[package]] @@ -1608,6 +1644,42 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "faer-core" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32da0f6262445261106407bdfb548df4f8f1c89bcc3cff1bb6d08e7f525b575a" +dependencies = [ + "bytemuck", + "coe-rs", + "dyn-stack 0.10.0", + "equator 0.1.10", + "faer-entity", + "gemm 0.17.1", + "matrixcompare-core", + "num-complex", + "num-traits", + "paste", + "pulp 0.18.22", + "reborrow", + "seq-macro", +] + +[[package]] +name = "faer-entity" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fb12f36da3162b3fefc0c7343d905d5bd2c2bf444a7f63ceeaa2bc003311d3" +dependencies = [ + "bytemuck", + "coe-rs", + "libm", + "num-complex", + "num-traits", + "pulp 0.18.22", + "reborrow", +] + [[package]] name = "fastrand" version = "2.5.0" @@ -1783,13 +1855,33 @@ dependencies = [ "slab", ] +[[package]] +name = "gemm" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ab24cc62135b40090e31a76a9b2766a501979f3070fa27f689c27ec04377d32" +dependencies = [ + "dyn-stack 0.10.0", + "gemm-c32 0.17.1", + "gemm-c64 0.17.1", + "gemm-common 0.17.1", + "gemm-f16 0.17.1", + "gemm-f32 0.17.1", + "gemm-f64 0.17.1", + "num-complex", + "num-traits", + "paste", + "raw-cpuid 10.7.0", + "seq-macro", +] + [[package]] name = "gemm" version = "0.18.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ab96b703d31950f1aeddded248bc95543c9efc7ac9c4a21fda8703a83ee35451" dependencies = [ - "dyn-stack", + "dyn-stack 0.13.2", "gemm-c32 0.18.2", "gemm-c64 0.18.2", "gemm-common 0.18.2", @@ -1799,7 +1891,7 @@ dependencies = [ "num-complex", "num-traits", "paste", - "raw-cpuid", + "raw-cpuid 11.6.0", "seq-macro", ] @@ -1809,7 +1901,7 @@ version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "aa0673db364b12263d103b68337a68fbecc541d6f6b61ba72fe438654709eacb" dependencies = [ - "dyn-stack", + "dyn-stack 0.13.2", "gemm-c32 0.19.0", "gemm-c64 0.19.0", "gemm-common 0.19.0", @@ -1819,7 +1911,22 @@ dependencies = [ "num-complex", "num-traits", "paste", - "raw-cpuid", + "raw-cpuid 11.6.0", + "seq-macro", +] + +[[package]] +name = "gemm-c32" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9c030d0b983d1e34a546b86e08f600c11696fde16199f971cd46c12e67512c0" +dependencies = [ + "dyn-stack 0.10.0", + "gemm-common 0.17.1", + "num-complex", + "num-traits", + "paste", + "raw-cpuid 10.7.0", "seq-macro", ] @@ -1829,12 +1936,12 @@ version = "0.18.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f6db9fd9f40421d00eea9dd0770045a5603b8d684654816637732463f4073847" dependencies = [ - "dyn-stack", + "dyn-stack 0.13.2", "gemm-common 0.18.2", "num-complex", "num-traits", "paste", - "raw-cpuid", + "raw-cpuid 11.6.0", "seq-macro", ] @@ -1844,12 +1951,27 @@ version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "086936dbdcb99e37aad81d320f98f670e53c1e55a98bee70573e83f95beb128c" dependencies = [ - "dyn-stack", + "dyn-stack 0.13.2", "gemm-common 0.19.0", "num-complex", "num-traits", "paste", - "raw-cpuid", + "raw-cpuid 11.6.0", + "seq-macro", +] + +[[package]] +name = "gemm-c64" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fbb5f2e79fefb9693d18e1066a557b4546cd334b226beadc68b11a8f9431852a" +dependencies = [ + "dyn-stack 0.10.0", + "gemm-common 0.17.1", + "num-complex", + "num-traits", + "paste", + "raw-cpuid 10.7.0", "seq-macro", ] @@ -1859,12 +1981,12 @@ version = "0.18.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dfcad8a3d35a43758330b635d02edad980c1e143dc2f21e6fd25f9e4eada8edf" dependencies = [ - "dyn-stack", + "dyn-stack 0.13.2", "gemm-common 0.18.2", "num-complex", "num-traits", "paste", - "raw-cpuid", + "raw-cpuid 11.6.0", "seq-macro", ] @@ -1874,15 +1996,34 @@ version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20c8aeeeec425959bda4d9827664029ba1501a90a0d1e6228e48bef741db3a3f" dependencies = [ - "dyn-stack", + "dyn-stack 0.13.2", "gemm-common 0.19.0", "num-complex", "num-traits", "paste", - "raw-cpuid", + "raw-cpuid 11.6.0", "seq-macro", ] +[[package]] +name = "gemm-common" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2e7ea062c987abcd8db95db917b4ffb4ecdfd0668471d8dc54734fdff2354e8" +dependencies = [ + "bytemuck", + "dyn-stack 0.10.0", + "half", + "num-complex", + "num-traits", + "once_cell", + "paste", + "pulp 0.18.22", + "raw-cpuid 10.7.0", + "seq-macro", + "sysctl 0.5.5", +] + [[package]] name = "gemm-common" version = "0.18.2" @@ -1890,7 +2031,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a352d4a69cbe938b9e2a9cb7a3a63b7e72f9349174a2752a558a8a563510d0f3" dependencies = [ "bytemuck", - "dyn-stack", + "dyn-stack 0.13.2", "half", "libm", "num-complex", @@ -1898,10 +2039,10 @@ dependencies = [ "once_cell", "paste", "pulp 0.21.5", - "raw-cpuid", + "raw-cpuid 11.6.0", "rayon", "seq-macro", - "sysctl", + "sysctl 0.6.0", ] [[package]] @@ -1911,7 +2052,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "88027625910cc9b1085aaaa1c4bc46bb3a36aad323452b33c25b5e4e7c8e2a3e" dependencies = [ "bytemuck", - "dyn-stack", + "dyn-stack 0.13.2", "half", "libm", "num-complex", @@ -1919,10 +2060,27 @@ dependencies = [ "once_cell", "paste", "pulp 0.22.3", - "raw-cpuid", + "raw-cpuid 11.6.0", "rayon", "seq-macro", - "sysctl", + "sysctl 0.6.0", +] + +[[package]] +name = "gemm-f16" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ca4c06b9b11952071d317604acb332e924e817bd891bec8dfb494168c7cedd4" +dependencies = [ + "dyn-stack 0.10.0", + "gemm-common 0.17.1", + "gemm-f32 0.17.1", + "half", + "num-complex", + "num-traits", + "paste", + "raw-cpuid 10.7.0", + "seq-macro", ] [[package]] @@ -1931,14 +2089,14 @@ version = "0.18.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cff95ae3259432f3c3410eaa919033cd03791d81cebd18018393dc147952e109" dependencies = [ - "dyn-stack", + "dyn-stack 0.13.2", "gemm-common 0.18.2", "gemm-f32 0.18.2", "half", "num-complex", "num-traits", "paste", - "raw-cpuid", + "raw-cpuid 11.6.0", "rayon", "seq-macro", ] @@ -1949,30 +2107,45 @@ version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3df7a55202e6cd6739d82ae3399c8e0c7e1402859b30e4cb780e61525d9486e" dependencies = [ - "dyn-stack", + "dyn-stack 0.13.2", "gemm-common 0.19.0", "gemm-f32 0.19.0", "half", "num-complex", "num-traits", "paste", - "raw-cpuid", + "raw-cpuid 11.6.0", "rayon", "seq-macro", ] +[[package]] +name = "gemm-f32" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9a69f51aaefbd9cf12d18faf273d3e982d9d711f60775645ed5c8047b4ae113" +dependencies = [ + "dyn-stack 0.10.0", + "gemm-common 0.17.1", + "num-complex", + "num-traits", + "paste", + "raw-cpuid 10.7.0", + "seq-macro", +] + [[package]] name = "gemm-f32" version = "0.18.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bc8d3d4385393304f407392f754cd2dc4b315d05063f62cf09f47b58de276864" dependencies = [ - "dyn-stack", + "dyn-stack 0.13.2", "gemm-common 0.18.2", "num-complex", "num-traits", "paste", - "raw-cpuid", + "raw-cpuid 11.6.0", "seq-macro", ] @@ -1982,12 +2155,27 @@ version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "02e0b8c9da1fbec6e3e3ab2ce6bc259ef18eb5f6f0d3e4edf54b75f9fd41a81c" dependencies = [ - "dyn-stack", + "dyn-stack 0.13.2", "gemm-common 0.19.0", "num-complex", "num-traits", "paste", - "raw-cpuid", + "raw-cpuid 11.6.0", + "seq-macro", +] + +[[package]] +name = "gemm-f64" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa397a48544fadf0b81ec8741e5c0fba0043008113f71f2034def1935645d2b0" +dependencies = [ + "dyn-stack 0.10.0", + "gemm-common 0.17.1", + "num-complex", + "num-traits", + "paste", + "raw-cpuid 10.7.0", "seq-macro", ] @@ -1997,12 +2185,12 @@ version = "0.18.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "35b2a4f76ce4b8b16eadc11ccf2e083252d8237c1b589558a49b0183545015bd" dependencies = [ - "dyn-stack", + "dyn-stack 0.13.2", "gemm-common 0.18.2", "num-complex", "num-traits", "paste", - "raw-cpuid", + "raw-cpuid 11.6.0", "seq-macro", ] @@ -2012,12 +2200,12 @@ version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "056131e8f2a521bfab322f804ccd652520c79700d81209e9d9275bbdecaadc6a" dependencies = [ - "dyn-stack", + "dyn-stack 0.13.2", "gemm-common 0.19.0", "num-complex", "num-traits", "paste", - "raw-cpuid", + "raw-cpuid 11.6.0", "seq-macro", ] @@ -2296,6 +2484,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 +3034,7 @@ dependencies = [ "jxl-oxide", "lightcraft-color", "lightcraft-geom", + "lightcraft-heif", "lightcraft-raster", "moxcms 0.9.1", "png", @@ -2859,6 +3057,38 @@ dependencies = [ "serde", ] +[[package]] +name = "lightcraft-denoise" +version = "0.4.0" +dependencies = [ + "faer-core", + "lightcraft-denoise-core", + "lightcraft-raster", + "miniz_oxide 0.8.9", + "proptest", + "rayon", + "serde", + "serde_json", + "sha2", + "thiserror 2.0.21", + "web-time", +] + +[[package]] +name = "lightcraft-denoise-core" +version = "0.4.0" +dependencies = [ + "lightcraft-raster", + "miniz_oxide 0.8.9", + "proptest", + "rayon", + "serde", + "serde_json", + "sha2", + "thiserror 2.0.21", + "web-time", +] + [[package]] name = "lightcraft-develop" version = "0.4.0" @@ -2881,7 +3111,11 @@ dependencies = [ "lightcraft-catalog", "lightcraft-codecs", "lightcraft-color", + "lightcraft-denoise", + "lightcraft-denoise-core", "lightcraft-develop", + "lightcraft-faces", + "lightcraft-fetch", "lightcraft-geom", "lightcraft-gpu", "lightcraft-merge", @@ -2903,6 +3137,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" @@ -2928,6 +3176,7 @@ version = "0.4.0" dependencies = [ "bytemuck", "lightcraft-color", + "lightcraft-denoise", "lightcraft-develop", "lightcraft-geom", "lightcraft-pipeline", @@ -2938,6 +3187,13 @@ dependencies = [ "wgpu", ] +[[package]] +name = "lightcraft-heif" +version = "0.4.0" +dependencies = [ + "heic-rs", +] + [[package]] name = "lightcraft-mcp" version = "0.4.0" @@ -2986,6 +3242,7 @@ dependencies = [ "lightcraft-develop", "lightcraft-geom", "lightcraft-raster", + "lightcraft-raw", "lightcraft-scenes", "proptest", "rayon", @@ -3022,6 +3279,7 @@ name = "lightcraft-raw" version = "0.4.0" dependencies = [ "jpeg-encoder", + "jxl-oxide", "lightcraft-color", "lightcraft-geom", "lightcraft-meta", @@ -3035,6 +3293,7 @@ dependencies = [ "thiserror 2.0.21", "zune-core", "zune-jpeg", + "zune-jpegxl", ] [[package]] @@ -3082,8 +3341,10 @@ dependencies = [ "lightcraft-catalog", "lightcraft-codecs", "lightcraft-color", + "lightcraft-denoise", "lightcraft-develop", "lightcraft-engine", + "lightcraft-faces", "lightcraft-geom", "lightcraft-meta", "lightcraft-pipeline", @@ -3174,6 +3435,22 @@ dependencies = [ "libc", ] +[[package]] +name = "matrixcompare-core" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0bdabb30db18805d5290b3da7ceaccbddba795620b86c02145d688e04900a73" + +[[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 +3604,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" @@ -4291,6 +4583,18 @@ dependencies = [ "unarray", ] +[[package]] +name = "pulp" +version = "0.18.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a0a01a0dc67cf4558d279f0c25b0962bd08fc6dec0137699eae304103e882fe6" +dependencies = [ + "bytemuck", + "libm", + "num-complex", + "reborrow", +] + [[package]] name = "pulp" version = "0.21.5" @@ -4317,7 +4621,7 @@ dependencies = [ "num-complex", "paste", "pulp-wasm-simd-flag", - "raw-cpuid", + "raw-cpuid 11.6.0", "reborrow", "version_check", ] @@ -4518,6 +4822,15 @@ dependencies = [ "rgb", ] +[[package]] +name = "raw-cpuid" +version = "10.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c297679cb867470fa8c9f67dbba74a78d78e3e98d7cf2b08d6d71540f797332" +dependencies = [ + "bitflags 1.3.2", +] + [[package]] name = "raw-cpuid" version = "11.6.0" @@ -4545,6 +4858,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" @@ -5245,6 +5564,20 @@ dependencies = [ "syn 3.0.6", ] +[[package]] +name = "sysctl" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec7dddc5f0fee506baf8b9fdb989e242f17e4b11c61dfbb0635b705217199eea" +dependencies = [ + "bitflags 2.13.2", + "byteorder", + "enum-as-inner", + "libc", + "thiserror 1.0.69", + "walkdir", +] + [[package]] name = "sysctl" version = "0.6.0" @@ -6752,6 +7085,15 @@ dependencies = [ "zune-core", ] +[[package]] +name = "zune-jpegxl" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cba11ecd07cc23351500e840ae03841b7e4997923ec8e4832ee3ae199ea0b6b4" +dependencies = [ + "zune-core", +] + [[package]] name = "zvariant" version = "5.15.0" diff --git a/Cargo.toml b/Cargo.toml index 77b20666..5f7c85db 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,7 +11,10 @@ repository = "https://github.com/storytold/lightcraft" [workspace.dependencies] # internal +lightcraft-denoise-core = { path = "crates/denoise-core" } +lightcraft-denoise = { path = "crates/denoise", default-features = false } 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 +33,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 +66,7 @@ tiff = "0.11" image-webp = "0.2" moxcms = { version = "0.9", default-features = false, features = ["lut"] } jxl-oxide = { version = "0.12", default-features = false } +zune-jpegxl = { version = "0.5", default-features = false, features = ["std"] } # MIT/Apache-2.0/Zlib lossless JXL encoder (test fixtures only) ravif = { version = "0.13", default-features = false, features = ["threading"] } # no "asm": it needs nasm (non-Rust) # gpu (crates/gpu; native backends only, no GL: everything stays pure Rust) diff --git a/NOTICE b/NOTICE index f3f6499f..254e2628 100644 --- a/NOTICE +++ b/NOTICE @@ -27,6 +27,15 @@ the facebook/sam3 checkpoint only when they choose to (in the app, or by hand), License (https://github.com/facebookresearch/sam3/blob/main/LICENSE), a separate, non-OSI licence. See docs/ai-masks.md. +AI denoise weights are not part of LightCraft and are never distributed with it. +The default build offers user-supplied models only. The optional rawnind-model +feature offers a separately downloaded RawNIND UtNet2 Bayer checkpoint under +GPL-3.0 (https://www.gnu.org/licenses/gpl-3.0.html), after the user accepts its terms. +Commercial use is subject to those terms, including applicable source, licence and +notice obligations on redistribution. Enabling that offer in releases remains a +maintainer decision. The ONNX reader and CPU/GPU executor are our Rust implementation; +no darktable, RawSpeed, LibRaw or other GPL decoder/runtime code is included. + Builds made with craft-fonts (the optional CRAFT_FONTS_DIR build input; all official releases) also embed the fonts listed in https://github.com/storytold/craft-fonts/blob/main/ATTRIBUTION.md (the Japanese fonts BIZ UDPGothic, BIZ UDMincho and Shippori Mincho; the web build embeds only BIZ diff --git a/README.md b/README.md index 5aacbee0..e93683c6 100644 --- a/README.md +++ b/README.md @@ -260,7 +260,8 @@ lightcraft --control 7980 ~/Pictures/trip - **Pure Rust, no C.** Our own RAW decoders (DNG, Canon CR2/CR3, Sony ARW, Nikon NEF, Fujifilm RAF incl. X-Trans, Panasonic RW2 / Leica RWL, Pentax PEF, Olympus ORF), our own colour science, our own pipeline. JPEG, PNG, TIFF, WebP, - PSD composites and JPEG XL open today. + PSD composites and JPEG XL open today; HEIC/HEIF photos too with `--features heif` (HEVC is a build-time choice, + as in PhotoCraft). - **Scene-referred & wide-gamut.** Linear Rec.2020 float internally, Bradford-adapted white balance, gamut mapping instead of clipping, a filmic shoulder for raw and pixel-exact pass-through for JPEGs you haven't touched. - **Resolution-independent edits.** Radii and brush sizes are relative to the image, so a 400 px preview, your diff --git a/ROADMAP.md b/ROADMAP.md index 5e793e8b..4a6b101a 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -25,7 +25,7 @@ duplicate Local entries #22, black GPU exports on an Intel iGPU #78). | **Feature checklist** | 79.7% | P0 core and P1 nearly complete: import (Add / Copy / Move, templates, devices), library, grid/loupe/compare/survey, every Edit slider, curves, colour grading, masking tools, crop/Upright, heal/clone, presets/profiles, versions/history, sync, export, menus, shortcuts | P0: full-resolution zoom detail; P1: lens-profile database, content-aware fill (patch synthesis), video playback/trim | | **RAW coverage** (formats people shoot) | ~55% | DNG (all kinds), CR2, CR3 lossless Bayer and version 0x100/0x200 C-RAW (M50/R100/R8 pixel-exact corpus), ARW, NEF (uncompressed + Huffman lossless/lossy), RAF (uncompressed, lossless/lossy compressed Bayer and X-Trans), uncompressed ORF, RW2 / RWL / Panasonic RAW (every raw format, checked on 178 files from 118 Panasonic and Leica bodies), PEF; every container's embedded preview (incl. CR3) | **CR3** (every Canon since ~2018), compressed ORF, Nikon lossy-after-split, Canon sRAW, HEIC/AVIF. NEFs labelled compressed but stored uncompressed (Z 6 packed 14-bit, D850 12-bit uncompressed). Per-model verification is thin (~55 corpus files vs >1,000 models) | | **Colour & image quality** | ~55–65% | Pipeline is complete and fast; GPU path CPU-exact within 1/255 | **No measured camera calibration database**: ARW, NEF, RW2, RAF, CR3, CR2 and PEF have guarded camera-JPEG colour estimates, with bundled ILCE-7M4, X-H2S and X-T4 profiles (docs/camera-preview-colour.md); other non-DNG raws and rejected estimates use a neutral matrix. Colour fidelity remains incomplete. No lens-profile database. No measured fidelity against Lightroom (tone, highlights, texture/clarity, NR, sharpening are tuned by eye) | -| **AI & computational** | ~15–20% | Assisted culling (focus, bursts), auto tone, HDR/panorama merge; subject/sky/background masks as classical heuristics | Real segmentation masks (subject, sky, people, objects, landscape, depth), AI denoise, super resolution, lens blur, generative remove, faces/people, natural-language search. **Blocked on a model strategy** (licensable weights or our own training; pure-Rust inference is feasible) | +| **AI & computational** | ~15–20% | Assisted culling (focus, bursts), auto tone, HDR/panorama merge; subject/sky/background masks as classical heuristics, or the photo's own segmentation mattes when it ships them (DNG semantic masks, e.g. iPhone ProRAW's sky matte); non-destructive Bayer RAW denoise with opt-in weights on pure-Rust CPU/GPU; **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), denoise model policy / linear RGB / quality verification, 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 | @@ -40,7 +40,7 @@ duplicate Local entries #22, black GPU exports on an Intel iGPU #78). | Fujifilm raw shooter | ~65% | Uncompressed and lossless/lossy compressed RAF decoded; camera colour calibration remains missing. Verified on 13 bodies, including 14/16-bit GFX; see `docs/raf-compression.md` | | Canon CR3 / Olympus compressed-raw shooter | ~45% | CR3 lossless Bayer and version 0x100/0x200 C-RAW develop from sensor data (M50, R100, R8 verified pixel for pixel); other CRX variants and compressed ORF open as embedded JPEG previews. Camera colour and broader model verification remain | | Lightroom Classic power user | ~45% | Print, Book, Map, publish, tethering | -| Relies on AI (masks, denoise) | ~25% | Object / Describe masks via optional SAM 3 (download not yet hosted); no AI denoise | +| Relies on AI (masks, denoise) | ~25% | Object / Describe masks via optional SAM 3 (download not yet hosted); Bayer AI denoise needs separately installed weights; X-Trans/linear RGB, model policy and quality parity remain open | ## Where we're going @@ -59,9 +59,9 @@ 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. Bayer denoise inference is implemented in pure Rust with an opt-in GPL model; release policy and independently licensed weights remain open. 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. + to be set up, docs/ai-masks.md); Subject / Sky / People, Raw Details and Super Resolution are still open. 7. **Then:** HDR (Q), the Classic output modules (Print first, then Map view, Book, Slideshow), video (R), localisation and accessibility. @@ -86,7 +86,7 @@ yet guarantee original-pixel detail; visible-region rendering remains a gap (LR- | M4 | Crop, geometry, optics | crop tool + overlays, straighten, Upright (auto/level/vertical/full/guided), manual transforms, CA, defringe, manual lens corrections | 6–10 | ✅ | | M5 | Performance | source pyramids, wgpu compute pipeline (CPU oracle), draft/full renders, prefetch, budgets (16 ms slider updates on 24 MP) | 10–15 | 🚧 (stage cache, source pyramid, wgpu pipeline, prefetch, memory budget ✅; colour NR at half resolution, GPU histogram ⬜) | | M6 | Masking | brush, linear/radial gradients, colour/luminance/depth range, add/subtract/intersect/invert, all local adjustments, masks panel | 8–12 | 🚧 (brush/linear/radial/colour/luminance range, add/subtract/intersect, masks panel ✅; depth range, AI masks ⬜) | -| M7 | Detail | sharpening + masking preview, luminance/colour NR, Denoise, Raw Details, Super Resolution | 6–10 | 🚧 (sharpening, luminance/colour NR ✅; AI Denoise, Raw Details, Super Resolution ⬜) | +| M7 | Detail | sharpening + masking preview, luminance/colour NR, Denoise, Raw Details, Super Resolution | 6–10 | 🚧 (sharpening, luminance/colour NR ✅; AI Denoise 🟡 (Bayer, optional weights; policy/fidelity open); Raw Details, Super Resolution ⬜) | | M8 | Heal / Remove | content-aware remove (PatchMatch), heal, clone, brush spots, visualize spots, red/pet eye | 6–10 | 🚧 (heal, clone, auto source, visualize spots, red/pet eye ✅; PatchMatch remove ⬜) | | M9 | Presets, profiles, versions, sync | preset browser + amount, create/import presets, profile browser, versions, history, copy/paste/sync settings | 5–8 | ✅ | | M10 | Export & share | export dialog (JPEG/PNG/TIFF/DNG/AVIF/JXL/original), sizing, sharpening, metadata, watermark, naming, batch jobs, XMP sidecars, HDR export | 6–10 | 🚧 (all formats incl. DNG/original, sizing, presets, background jobs ✅; JXL encode, HDR export ⬜) | @@ -124,7 +124,7 @@ trip, slideshow, auto import…) in ≈ 12 h on 2026-10-02; four to six parallel | Work package | Tracker rows | Agent-hours | Risk | |---|---|---:|---| | Remaining P0/P1 UI and library features (folder rename/move, keyword painter, people view…) | ≈ 8 | 5–10 | low | -| Raw codecs: CR3 (CRX), compressed ORF, NEF lossy-after-split, HEIC/AVIF decode, JPEG XL DNG | LR-IMP-FORMATS | 40–80 | **high** — clean-room black-box analysis, no permissive specs | +| Raw codecs: CR3 (CRX), compressed ORF, NEF lossy-after-split, HEIC/AVIF decode | LR-IMP-FORMATS | 40–80 | **high** — clean-room black-box analysis, no permissive specs | | Lens-profile database of our own (calibration targets, fitting, data) | LR-EDIT-OPTICS-PROFILE | 15–30 | data collection | | Video: playback, trim, edits, export (pure-Rust decode, ideally shared with FilmCraft) | R. Video | 20–40 | medium | | AI: subject / sky / background / people / object masks, object-aware remove, AI denoise, super resolution, lens blur, people & faces, natural-language search, culling | ≈ 30 | 80–150 | **high** — permissively licensed weights, pure-Rust inference, maybe training | @@ -174,7 +174,8 @@ private 3:2 & 4:3 examples verified through native rendering & full-resolution 1 CR3 lossless Bayer and version 0x100/0x200 C-RAW sensor decoding is verified pixel-exactly on Canon M50, R100 and R8; any other CR3, or one that fails to decode, opens from its embedded JPEG as before. See [`docs/cr3.md`](docs/cr3.md) for coding limits, reference provenance and colour limitations. Decoded (CC0 corpus from raw.pixls.us, `cargo xtask corpus --download`, `crates/raw/tests/corpus.rs`): DNG (uncompressed, -LJ92, lossy JPEG / Smart Previews, Deflate, float, linear), CR2, ARW (uncompressed, ARW2, LJ92; as-shot white balance and black level of pre-2017 bodies from the enciphered +LJ92, lossy JPEG / Smart Previews, Deflate, float, linear, DNG 1.7 JPEG XL tiles: lossless ones bit-exact on synthetic files, +lossy ones decode but no real file has verified them), CR2, ARW (uncompressed, ARW2, LJ92; as-shot white balance and black level of pre-2017 bodies from the enciphered maker-note `Tag2010` and the encrypted `SR2SubIFD`, both recovered by black-box analysis, `crates/raw/src/vendor/arw.rs`), NEF/NRW uncompressed and Huffman-compressed (lossless, lossy type 1/2, 12/14-bit), RAF uncompressed and lossless/lossy compressed (Bayer and X-Trans, 14/16-bit real-file verification), RW2 / Leica RWL / Panasonic RAW in every raw format (compressed formats 4 and 6, the prefix-coded strips of format 8, packed formats 2, 5 and 7, the 16-bit words of the 2005–2007 bodies; all recovered by black-box analysis of 178 CC0 files from @@ -185,7 +186,7 @@ Not decoded yet — preview only (no permissively licensed description; black-bo files): rows above the split decode with the regular lossy table; from the split row on a different code is used that our black-box analysis has not recovered yet (none of the four regular tables fits, also not with byte alignment or reset predictors at the split row). The other Nikon Huffman variants are decoded (`crates/raw/src/vendor/nefc.rs` documents the analysis). -- **Olympus compressed ORF**, **Canon CR3 unverified CRX variants** (M11.1), **Canon sRAW/mRAW**, lossy DNG. +- **Olympus compressed ORF**, **Canon CR3 unverified CRX variants** (M11.1), **Canon sRAW/mRAW**. **Camera colour matrices:** ARW, NEF, RW2, RAF, CR3, CR2 and PEF files can use guarded, separate chromaticity and tone estimates from their own embedded JPEG (see `docs/camera-preview-colour.md`); this is a per-file camera-look estimate with relative WB, not measured calibration or absolute-Kelvin WB. Other non-DNG raws and rejected fits use the documented neutral fallback (camera RGB ≈ linear sRGB, flagged `matrix_is_fallback`) with the file's as-shot white-balance multipliers. Clean sources to evaluate next: manufacturer diff --git a/apps/lightcraft-cli/Cargo.toml b/apps/lightcraft-cli/Cargo.toml index 3f60fe32..6471ab86 100644 --- a/apps/lightcraft-cli/Cargo.toml +++ b/apps/lightcraft-cli/Cargo.toml @@ -25,7 +25,14 @@ serde_json = { workspace = true } dhat = { workspace = true, optional = true } [features] +default = ["denoise"] +denoise = ["lightcraft-engine/denoise"] +rawnind-model = ["denoise", "lightcraft-engine/rawnind-model"] # Heap profiling: `cargo run --release -p lightcraft-cli --features dhat-heap -- snapshot …` counts # every allocation (live/peak bytes in `library.memory` / `ui.inspect`) and writes dhat-heap.json # (open it in DHAT's viewer: allocation sites by bytes at the peak). Slows allocation-heavy work. dhat-heap = ["dep:dhat"] + +# Open HEIC/HEIF photos (the optional lightcraft-heif crate). Not default so distributors choose +# it (HEVC patents); build with `--features heif` to include it. +heif = ["lightcraft-codecs/heif"] diff --git a/apps/lightcraft-cli/src/main.rs b/apps/lightcraft-cli/src/main.rs index 6c5e6e44..f8ddf550 100644 --- a/apps/lightcraft-cli/src/main.rs +++ b/apps/lightcraft-cli/src/main.rs @@ -391,7 +391,7 @@ fn merge(args: &[String]) -> Result<(), String> { } i += 1; } - let mut s = Session::new().with_fs(); + let mut s = Session::new().with_fs().with_default_denoise_models(); let paths = expand_paths(&files); let r = s.execute("library.import", &json!({"paths": paths})).map_err(|e| e.to_string())?; let mut ids: Vec = r["imported"].as_array().map(|a| a.iter().filter_map(Value::as_u64).collect()).unwrap_or_default(); @@ -652,7 +652,7 @@ fn render(args: &[String]) -> Result<(), String> { } let input = input.ok_or("render: missing input file")?; let output = output.ok_or("render: missing -o OUTPUT")?; - let mut s = Session::new().with_fs(); + let mut s = Session::new().with_fs().with_default_denoise_models(); let abs = expand_paths(std::slice::from_ref(&input)); let r = s.execute("library.import", &json!({"paths": abs})).map_err(|e| e.to_string())?; let id = r["imported"][0].as_u64().ok_or_else(|| format!("{input}: not a readable photo"))?; @@ -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_denoise_models().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_denoise_models().with_default_face_models(), + None => Session::new().with_fs().with_default_denoise_models().with_default_face_models(), }; + // Apply compute policy before imports or the first scripted query can discover an adapter. + session.execute("app.gpu", &json!({"enabled": false})).map_err(|e| e.to_string())?; if !files.is_empty() { let ti = Instant::now(); let r = session.execute("library.import", &json!({"paths": expand_paths(&files)})).map_err(|e| e.to_string())?; @@ -755,7 +755,11 @@ fn snapshot(args: &[String]) -> Result<(), String> { })), ..Default::default() }; - let app = lightcraft_ui_egui::LightcraftApp::new(session, services); + let mut app = lightcraft_ui_egui::LightcraftApp::new(session, services); + // Snapshot sessions promise no GPU: disable photo compute as well as the compositor + // before the first frame can start background adapter discovery. + // This UI state is session-local: the CLI neither loads nor saves desktop ui.json. + app.ui.settings.gpu = false; let mut h = Headless::new(app, size, scale); let timeout = Duration::from_secs(60); let mut shots = 0usize; diff --git a/apps/lightcraft-cli/tests/cli.rs b/apps/lightcraft-cli/tests/cli.rs index d1b178e4..2009e818 100644 --- a/apps/lightcraft-cli/tests/cli.rs +++ b/apps/lightcraft-cli/tests/cli.rs @@ -233,6 +233,58 @@ fn snapshot_script_failure_exits_non_zero() { assert_eq!((d.width, d.height), (320, 200)); } +#[test] +fn snapshot_ui_zoom_keeps_requested_pixel_dimensions() { + let script = tmp("snap-ui-zoom.jsonl"); + std::fs::write( + &script, + format!("{}\n{}\n", json!({"method": "ui.key", "params": {"key": "Plus", "cmd": true}}), json!({"method": "ui.inspect"})), + ) + .unwrap(); + for (size, scale, expected) in [("400x240", "1", (400, 240)), ("345x200", "0.9", (311, 180))] { + let out = tmp(&format!("snap-ui-zoom-{scale}.png")); + let o = Command::new(BIN) + .args(["snapshot", "--script", script.to_str().unwrap(), "-o", out.to_str().unwrap(), "--size", size, "--scale", scale]) + .output() + .unwrap(); + assert!(o.status.success(), "{}", String::from_utf8_lossy(&o.stderr)); + let replies: Vec = String::from_utf8_lossy(&o.stdout).lines().map(|line| serde_json::from_str(line).unwrap()).collect(); + assert!(replies.iter().all(|reply| reply["ok"] == true), "{replies:?}"); + assert!(replies[1]["result"]["pixelsPerPoint"].as_f64().unwrap() > scale.parse::().unwrap()); + let decoded = lightcraft_codecs::decode(&std::fs::read(out).unwrap(), Default::default()).unwrap(); + assert_eq!((decoded.width, decoded.height), expected); + } +} + +/// Headless snapshots default to CPU photo rendering as well as CPU UI rasterization. +#[test] +fn snapshot_defaults_to_cpu_without_gpu_environment_overrides() { + let script = tmp("default-cpu.jsonl"); + std::fs::write( + &script, + format!("{}\n{}\n", json!({"method": "engine.execute", "params": {"command": "app.gpu"}}), json!({"method": "ui.screenshot"})), + ) + .unwrap(); + let image = tmp("default-cpu.png"); + let output = Command::new(BIN) + .env_remove("LIGHTCRAFT_GPU") + .env_remove("LIGHTCRAFT_GPU_BACKEND") + .env_remove("WGPU_BACKEND") + .args(["snapshot", "--demo", "--script", script.to_str().unwrap(), "-o", image.to_str().unwrap(), "--size", "480x320"]) + .output() + .unwrap(); + assert!(output.status.success(), "status {}: {}", output.status, String::from_utf8_lossy(&output.stderr)); + let replies: Vec = String::from_utf8_lossy(&output.stdout).lines().map(|line| serde_json::from_str(line).unwrap()).collect(); + assert!(replies.iter().all(|reply| reply["ok"] == true), "{replies:?}"); + let gpu = &replies[0]["result"]; + assert_eq!(gpu["enabled"], false, "{gpu}"); + assert_eq!(gpu["available"], false, "{gpu}"); + assert_eq!(gpu["adapter"], Value::Null, "no device is created: {gpu}"); + assert!(gpu["reason"].as_str().is_some_and(|reason| reason.contains("preference")), "{gpu}"); + let decoded = lightcraft_codecs::decode(&std::fs::read(image).unwrap(), Default::default()).unwrap(); + assert_eq!((decoded.width, decoded.height), (480, 320)); +} + /// Issue #136: with the GPU switched off from the environment the UI starts and renders on the CPU, /// without creating a GPU device (no driver is loaded), and says why. #[test] @@ -412,3 +464,79 @@ fn warnings_are_logged_on_stderr() { let off = run("off"); assert!(!off.contains("names no known backend"), "{off}"); } + +#[test] +fn snapshot_rejects_invalid_dimensions_before_creating_an_image() { + for (index, (size, scale)) in + [("NaNx100", "1"), ("100x100", "inf"), ("9000x9000", "1"), ("9000x100", "2"), ("100x100", "0.001")].into_iter().enumerate() + { + let image = tmp(&format!("invalid-viewport-{index}.png")); + let output = Command::new(BIN).args(["snapshot", "--size", size, "--scale", scale, "-o"]).arg(&image).output().unwrap(); + assert!(!output.status.success()); + assert!(String::from_utf8_lossy(&output.stderr).contains("snapshot: bad --size/--scale"), "{}", String::from_utf8_lossy(&output.stderr)); + assert!(!image.exists()); + } +} + +#[test] +fn snapshot_preserves_fractional_scale_rounding() { + let image = tmp("fractional-scale-rounding.png"); + let output = + Command::new(BIN).env("LIGHTCRAFT_GPU", "0").args(["snapshot", "--size", "345x200", "--scale", "0.9", "-o"]).arg(&image).output().unwrap(); + assert!(output.status.success(), "{}", String::from_utf8_lossy(&output.stderr)); + let decoded = lightcraft_codecs::decode(&std::fs::read(image).unwrap(), Default::default()).unwrap(); + assert_eq!((decoded.width, decoded.height), (311, 180)); +} + +/// Snapshot's CPU setting must never replace the desktop preference, including with --library. +#[cfg(target_os = "linux")] +#[test] +fn snapshot_leaves_desktop_gpu_preferences_untouched() { + let dir = tmp("snapshot-desktop-preferences"); + let library = dir.join("library"); + let mut session = lightcraft_engine::Session::new().with_fs(); + session.open_library(&library, true).unwrap(); + drop(session); + let config = dir.join("config"); + let app_config = config.join("lightcraft"); + std::fs::create_dir_all(&app_config).unwrap(); + let ui_path = app_config.join("ui.json"); + let mut ui = lightcraft_ui_egui::UiState::default(); + ui.settings.gpu = true; + let saved = serde_json::to_vec_pretty(&ui).unwrap(); + std::fs::write(&ui_path, &saved).unwrap(); + let script = dir.join("preferences.jsonl"); + std::fs::write( + &script, + format!("{}\n{}\n", json!({"method": "engine.execute", "params": {"command": "app.gpu"}}), json!({"method": "ui.inspect"})), + ) + .unwrap(); + for gpu_environment in [None, Some("1")] { + let image = dir.join(format!("{}.png", gpu_environment.unwrap_or("default"))); + let mut cmd = Command::new(BIN); + cmd.env("XDG_CONFIG_HOME", &config).env_remove("LIGHTCRAFT_GPU").env_remove("LIGHTCRAFT_GPU_BACKEND").env_remove("WGPU_BACKEND"); + if let Some(value) = gpu_environment { + cmd.env("LIGHTCRAFT_GPU", value); + } + let output = cmd + .args(["snapshot", "--library"]) + .arg(&library) + .args(["--script"]) + .arg(&script) + .args(["-o"]) + .arg(&image) + .args(["--size", "480x320"]) + .output() + .unwrap(); + assert!(output.status.success(), "{}", String::from_utf8_lossy(&output.stderr)); + let replies: Vec = String::from_utf8_lossy(&output.stdout).lines().map(|line| serde_json::from_str(line).unwrap()).collect(); + assert!(replies.iter().all(|reply| reply["ok"] == true), "{replies:?}"); + assert_eq!(replies[0]["result"]["enabled"], false); + assert_eq!(replies[0]["result"]["available"], false); + assert_eq!(replies[0]["result"]["adapter"], Value::Null); + assert_eq!(replies[1]["result"]["ui"]["settings"]["gpu"], false); + assert_eq!(std::fs::read(&ui_path).unwrap(), saved, "snapshot changed the desktop preferences"); + assert!(!app_config.join("ui.json.tmp").exists()); + assert!(image.exists()); + } +} diff --git a/apps/lightcraft-web/Cargo.toml b/apps/lightcraft-web/Cargo.toml index a3331a60..aaa14a64 100644 --- a/apps/lightcraft-web/Cargo.toml +++ b/apps/lightcraft-web/Cargo.toml @@ -71,3 +71,9 @@ web-sys = { workspace = true, features = [ "WorkerType", "WritableStream", ] } + +[features] +# Open HEIC/HEIF photos (the optional lightcraft-heif crate). Not default so distributors choose +# it (HEVC patents); build with `--features heif` to include it. On wasm32 panics abort, so +# heic-rs's `catch_unwind` guard does not apply there. +heif = ["lightcraft-codecs/heif"] diff --git a/apps/lightcraft-web/src/web.rs b/apps/lightcraft-web/src/web.rs index 9ffbf170..f78ac861 100644 --- a/apps/lightcraft-web/src/web.rs +++ b/apps/lightcraft-web/src/web.rs @@ -205,10 +205,13 @@ fn services(originals: Originals, backend: Option, files: Files, frozen open_picker(originals.clone(), backend.clone(), ctx.clone()); Vec::new() // files arrive asynchronously and are imported on a later frame })), + pick_denoise_model: None, // Preset files: browser pickers are asynchronous; not wired on the web yet. pick_preset_files: None, pick_tracklog: None, pick_lightroom_catalog: None, + // Face models need a folder to live in; the web has none. + pick_model_file: None, save_preset_file: None, pick_curve_preset_files: None, save_curve_preset_file: None, diff --git a/apps/lightcraft-web/src/wire.rs b/apps/lightcraft-web/src/wire.rs index 81a76e6c..b08ae150 100644 --- a/apps/lightcraft-web/src/wire.rs +++ b/apps/lightcraft-web/src/wire.rs @@ -368,7 +368,11 @@ mod tests { let a = core.render(&wire, None).unwrap(); let b = job.run().rendered.unwrap(); assert_eq!((a.image.width, a.image.height), (160, 120)); - assert_eq!(a.image.data, b.image.data); + assert_eq!((b.image.width, b.image.height), (160, 120)); + // the worker renders on the CPU; the engine job uses a GPU when the machine has one, which + // stays within docs/gpu-pipeline.md's bound (≤ 3 LSB per channel) rather than bit-exact + let max = a.image.data.iter().zip(&b.image.data).flat_map(|(p, q)| (0..3).map(move |c| p[c].abs_diff(q[c]))).max(); + assert!(max.is_some_and(|m| m <= 3), "max channel difference {max:?}"); } // a job from before windows existed has none diff --git a/apps/lightcraft/Cargo.toml b/apps/lightcraft/Cargo.toml index a29299da..e4d7f2b3 100644 --- a/apps/lightcraft/Cargo.toml +++ b/apps/lightcraft/Cargo.toml @@ -9,6 +9,14 @@ default-run = "lightcraft" [lints] workspace = true +[features] +default = ["denoise"] +denoise = ["lightcraft-engine/denoise"] +rawnind-model = ["denoise", "lightcraft-engine/rawnind-model"] +# Open HEIC/HEIF photos (the optional lightcraft-heif crate). Not default so distributors choose +# it (HEVC patents); build with `--features heif` to include it. +heif = ["lightcraft-codecs/heif"] + [dependencies] lightcraft-engine = { workspace = true, features = ["sam"] } lightcraft-sysmem = { workspace = true } diff --git a/apps/lightcraft/src/main.rs b/apps/lightcraft/src/main.rs index 4c7bb424..a0394da9 100644 --- a/apps/lightcraft/src/main.rs +++ b/apps/lightcraft/src/main.rs @@ -125,7 +125,7 @@ fn gpu_crash_notice(what: &str) -> String { } fn config_dir() -> Option { - lightcraft_engine::camera_profiles::config_dir() + lightcraft_engine::config::config_dir() } /// `/ui.json`, the saved UI state and app settings — `None` when nothing is read or @@ -387,6 +387,13 @@ fn services(ctx: egui::Context, log_file: Option<&std::path::Path>) -> Services .map(|p| p.to_string_lossy().to_string()) .collect() })), + pick_denoise_model: Some(Box::new(|| { + rfd::FileDialog::new() + .add_filter_nocase(lightcraft_ui_egui::i18n::tr("ONNX model"), &["onnx"]) + .pick_file() + .map(|p| vec![p.to_string_lossy().into_owned()]) + .unwrap_or_default() + })), pick_preset_files: Some(Box::new(|| { rfd::FileDialog::new() .set_title(lightcraft_ui_egui::i18n::tr("Import Presets")) @@ -400,6 +407,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,22 +461,29 @@ fn services(ctx: egui::Context, log_file: Option<&std::path::Path>) -> Services } } +/// The library session with the folder face models are kept in (`/models`). +fn open_session(in_memory: bool, dir: Option, seed_demo: bool) -> (Session, Option) { + let (mut s, problem) = open_library_session(in_memory, dir, seed_demo); + s.face_models_dir = lightcraft_engine::config::default_face_models_dir(); + (s, problem) +} + /// The persistent library session, or with `--memory` an in-memory one (demo photos). /// /// If the library can't be opened the session is empty and in memory — never seeded with the /// demo photos, never written anywhere — and the problem is returned: the window then says so /// and offers Try Again / Choose Another Library… / Continue Without Saving / Quit (issue #100). -fn open_session(in_memory: bool, dir: Option, seed_demo: bool) -> (Session, Option) { +fn open_library_session(in_memory: bool, dir: Option, seed_demo: bool) -> (Session, Option) { if in_memory { - return (if seed_demo { Session::with_demo() } else { Session::new() }.with_fs().with_system_clock(), None); + return (if seed_demo { Session::with_demo() } else { Session::new() }.with_fs().with_default_denoise_models().with_system_clock(), None); } - let unopened = || Session::new().with_fs().with_system_clock(); + let unopened = || Session::new().with_fs().with_default_denoise_models().with_system_clock(); let Some(dir) = dir else { log::warn!("no library location (set --library or LIGHTCRAFT_LIBRARY)"); return (unopened(), Some(LibraryProblem::new("", "There is no home folder to keep the library in. Choose a folder for it."))); }; let t0 = std::time::Instant::now(); - let mut s = Session::new().with_fs().with_system_clock(); + let mut s = Session::new().with_fs().with_default_denoise_models().with_system_clock(); match s.open_library(&dir, seed_demo) { Ok(r) => { let (replayed, torn) = (r.replayed, r.torn_bytes); diff --git a/contributors/people.toml b/contributors/people.toml index 89c6db20..fbda1f24 100644 --- a/contributors/people.toml +++ b/contributors/people.toml @@ -21,3 +21,6 @@ real_name = "Brandon Thomas" display_name = "Brandon Thomas" sync_github_name = false + +[people.sanyamsingh001] +display_name = "sanyamsingh001" diff --git a/crates/catalog/src/journal.rs b/crates/catalog/src/journal.rs index 96187cfa..5973975d 100644 --- a/crates/catalog/src/journal.rs +++ b/crates/catalog/src/journal.rs @@ -45,6 +45,7 @@ //! | (none) | early builds: a log without a snapshot | | //! | 1 | up to v0.2.0 | | //! | 2 | after v0.2.0 | `Op::SetBrowsed`, `Catalog.browsed`, `Photo.local_baseline` | +//! | 3 | albums by hand | `Op::SetAlbumOrder`, `Album.order` | //! //! Rules: //! - **Bump [`VERSION`]** (and add a row above) in the change that adds an [`Op`] variant or a @@ -69,7 +70,7 @@ pub const LOG: &str = "catalog.log"; const FORMAT: &str = "lightcraft-catalog"; /// The catalog format this build writes (and the newest it reads). See the module docs → /// *Format versions*; bump it whenever an [`Op`] variant or a serialized field is added. -pub const VERSION: u32 = 2; +pub const VERSION: u32 = 3; /// When [`Journal::wants_snapshot`] says it's time to compact the log. #[derive(Clone, Copy, Debug)] diff --git a/crates/catalog/src/lib.rs b/crates/catalog/src/lib.rs index 2934c7a6..884ec658 100644 --- a/crates/catalog/src/lib.rs +++ b/crates/catalog/src/lib.rs @@ -136,6 +136,11 @@ pub enum Op { id: AlbumId, parent: Option, }, + /// Where the album stands among its siblings (`None`: by name). Format version 3. + SetAlbumOrder { + id: AlbumId, + order: Option, + }, SetAlbumPhotos { id: AlbumId, photos: Vec, @@ -282,6 +287,28 @@ impl Catalog { pub fn albums(&self) -> impl Iterator { self.albums.values() } + /// The albums and folders directly inside `parent` (`None`: the top level), in the order the + /// sidebar lists them: folders first, then by their place if the user ordered them (those + /// without one after, by name), else by name. + pub fn album_children(&self, parent: Option) -> Vec<&Album> { + let mut kids: Vec<&Album> = self.albums.values().filter(|a| a.parent == parent).collect(); + kids.sort_by_cached_key(|a| album_order_key(a)); + kids + } + /// [`Self::album_children`] of every folder (and of the top level, under `None`) in one pass: + /// what a tree view that draws all of them wants each frame. + pub fn album_children_by_parent(&self) -> std::collections::HashMap, Vec<&Album>> { + let mut map: std::collections::HashMap, Vec<&Album>> = std::collections::HashMap::new(); + for a in self.albums.values() { + map.entry(a.parent).or_default().push(a); + } + map.values_mut().for_each(|kids| kids.sort_by_cached_key(|a| album_order_key(a))); + map + } + /// Whether anything inside `parent` has a place of its own (the folder is ordered by hand). + pub fn album_children_are_ordered(&self, parent: Option) -> bool { + self.albums.values().any(|a| a.parent == parent && a.order.is_some()) + } /// The Quick Collection, once something was added to it. pub fn quick_collection(&self) -> Option { self.albums.values().find(|a| a.quick).map(|a| a.id) @@ -479,7 +506,20 @@ impl Catalog { } } let a = self.album_mut(id)?; - Op::MoveAlbum { id, parent: std::mem::replace(&mut a.parent, parent) } + if a.parent == parent { + return Ok(Op::MoveAlbum { id, parent }); + } + // its place belonged to the old neighbours: the undo gives it back with the folder + let old_order = a.order.take(); + let old_parent = std::mem::replace(&mut a.parent, parent); + match old_order { + Some(_) => Op::Batch { ops: vec![Op::MoveAlbum { id, parent: old_parent }, Op::SetAlbumOrder { id, order: old_order }] }, + None => Op::MoveAlbum { id, parent: old_parent }, + } + } + Op::SetAlbumOrder { id, order } => { + let a = self.album_mut(id)?; + Op::SetAlbumOrder { id, order: std::mem::replace(&mut a.order, order) } } Op::SetAlbumPhotos { id, photos } => { let a = self.album_mut(id)?; @@ -668,9 +708,17 @@ impl Catalog { } } +/// How siblings are listed: folders first, then the ones with a place of their own by it, then by +/// name; the id settles any tie. +fn album_order_key(a: &Album) -> (bool, bool, u32, String, AlbumId) { + (!a.folder, a.order.is_none(), a.order.unwrap_or(0), a.name.to_lowercase(), a.id) +} + #[cfg(test)] mod tests; #[cfg(test)] +mod tests_album_order; +#[cfg(test)] mod tests_background; #[cfg(test)] mod tests_folders; diff --git a/crates/catalog/src/local.rs b/crates/catalog/src/local.rs index 79a4a42d..4291f45f 100644 --- a/crates/catalog/src/local.rs +++ b/crates/catalog/src/local.rs @@ -295,6 +295,7 @@ impl Op { Op::RemoveAlbum { .. } | Op::RenameAlbum { .. } | Op::MoveAlbum { .. } + | Op::SetAlbumOrder { .. } | Op::SetAlbumRules { .. } | Op::RemoveStack { .. } | Op::SetLabelName { .. } diff --git a/crates/catalog/src/model.rs b/crates/catalog/src/model.rs index ae53b486..957f5b68 100644 --- a/crates/catalog/src/model.rs +++ b/crates/catalog/src/model.rs @@ -405,12 +405,16 @@ pub struct Album { /// The Quick Collection: a regular album that B adds to (one per library). #[serde(default, skip_serializing_if = "std::ops::Not::not")] pub quick: bool, + /// Its place among the albums of the same kind in its folder, once the user ordered them by + /// hand (see [`crate::Catalog::album_children`]); `None`: listed by name. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub order: Option, } impl Album { /// A regular (manual) album. pub fn new(id: AlbumId, name: impl Into) -> Album { - Album { id, name: name.into(), parent: None, folder: false, photos: Vec::new(), cover: None, smart: None, quick: false } + Album { id, name: name.into(), parent: None, folder: false, photos: Vec::new(), cover: None, smart: None, quick: false, order: None } } pub fn is_smart(&self) -> bool { self.smart.is_some() diff --git a/crates/catalog/src/tests.rs b/crates/catalog/src/tests.rs index 49f0ac4d..8d34d5ef 100644 --- a/crates/catalog/src/tests.rs +++ b/crates/catalog/src/tests.rs @@ -58,12 +58,22 @@ fn albums_and_folders() { let a = photo(&mut c, "a.jpg", "2026-04-01"); let f = c.alloc_album_id(); c.apply(Op::AddAlbum { - album: Album { id: f, name: "Trips".into(), parent: None, folder: true, photos: vec![], cover: None, smart: None, quick: false }, + album: Album { id: f, name: "Trips".into(), parent: None, folder: true, photos: vec![], cover: None, smart: None, quick: false, order: None }, }) .unwrap(); let al = c.alloc_album_id(); c.apply(Op::AddAlbum { - album: Album { id: al, name: "Alps".into(), parent: Some(f), folder: false, photos: vec![a], cover: None, smart: None, quick: false }, + album: Album { + id: al, + name: "Alps".into(), + parent: Some(f), + folder: false, + photos: vec![a], + cover: None, + smart: None, + quick: false, + order: None, + }, }) .unwrap(); assert_eq!(c.albums_of(a), vec![al]); diff --git a/crates/catalog/src/tests_album_order.rs b/crates/catalog/src/tests_album_order.rs new file mode 100644 index 00000000..112633f4 --- /dev/null +++ b/crates/catalog/src/tests_album_order.rs @@ -0,0 +1,122 @@ +//! Album order: the children of a folder (or of the top level) are listed folders first, then +//! alphabetically, until some of them are given a place of their own (`Album.order`); moving an +//! album to another folder drops its place, and undoing the move brings it back. + +use crate::*; + +fn add(c: &mut Catalog, name: &str, parent: Option, folder: bool) -> AlbumId { + let id = c.alloc_album_id(); + c.apply(Op::AddAlbum { album: Album { parent, folder, ..Album::new(id, name) } }).unwrap(); + id +} + +fn names(c: &Catalog, parent: Option) -> Vec { + c.album_children(parent).iter().map(|a| a.name.clone()).collect() +} + +fn place(c: &mut Catalog, id: AlbumId, order: Option) { + c.apply(Op::SetAlbumOrder { id, order }).unwrap(); +} + +/// Given albums that were never ordered, when they are listed, then folders come first and each +/// group is alphabetical regardless of case. +#[test] +fn children_are_folders_first_then_alphabetical() { + let mut c = Catalog::new(); + add(&mut c, "beta", None, false); + add(&mut c, "Alpha", None, false); + add(&mut c, "zoo", None, true); + add(&mut c, "Archive", None, true); + assert_eq!(names(&c, None), ["Archive", "zoo", "Alpha", "beta"]); + assert!(!c.album_children_are_ordered(None)); +} + +/// Given some children with a place of their own, then they follow it; the others come after, by +/// name; folders stay ahead of albums. +#[test] +fn ordered_children_follow_their_order() { + let mut c = Catalog::new(); + let a = add(&mut c, "A", None, false); + let b = add(&mut c, "B", None, false); + let cc = add(&mut c, "C", None, false); + add(&mut c, "D", None, false); + let f = add(&mut c, "F", None, true); + place(&mut c, cc, Some(0)); + place(&mut c, a, Some(1)); + assert_eq!(names(&c, None), ["F", "C", "A", "B", "D"]); + assert!(c.album_children_are_ordered(None)); + // equal places fall back to the name, so a damaged catalog still lists deterministically + place(&mut c, b, Some(1)); + assert_eq!(names(&c, None), ["F", "C", "A", "B", "D"]); + // another container is unaffected + let inner = add(&mut c, "Z", Some(f), false); + add(&mut c, "Y", Some(f), false); + assert_eq!(names(&c, Some(f)), ["Y", "Z"]); + assert!(!c.album_children_are_ordered(Some(f))); + place(&mut c, inner, Some(0)); + assert_eq!(names(&c, Some(f)), ["Z", "Y"]); +} + +/// The one-pass map lists every folder's children exactly as `album_children` does. +#[test] +fn the_children_map_matches_the_children_of_each_folder() { + let mut c = Catalog::new(); + let f = add(&mut c, "F", None, true); + let g = add(&mut c, "G", Some(f), true); + for (name, parent) in [("b", None), ("a", None), ("y", Some(f)), ("x", Some(f)), ("q", Some(g))] { + add(&mut c, name, parent, false); + } + let x = c.albums().find(|a| a.name == "x").unwrap().id; + place(&mut c, x, Some(7)); + let map = c.album_children_by_parent(); + for parent in [None, Some(f), Some(g)] { + let from_map: Vec = map.get(&parent).map(|v| v.iter().map(|a| a.id).collect()).unwrap_or_default(); + let direct: Vec = c.album_children(parent).iter().map(|a| a.id).collect(); + assert_eq!(from_map, direct, "{parent:?}"); + } + assert!(!map.contains_key(&Some(AlbumId(999)))); +} + +/// A place is one undo step, and `None` puts the album back among the alphabetical ones. +#[test] +fn a_place_can_be_set_cleared_and_undone() { + let mut c = Catalog::new(); + let a = add(&mut c, "A", None, false); + add(&mut c, "B", None, false); + let inv = c.apply(Op::SetAlbumOrder { id: a, order: Some(5) }).unwrap(); + assert_eq!(c.album(a).unwrap().order, Some(5)); + c.apply(inv).unwrap(); + assert_eq!(c.album(a).unwrap().order, None); + assert!(c.apply(Op::SetAlbumOrder { id: AlbumId(999), order: Some(1) }).is_err(), "no such album"); +} + +/// Moving to another folder drops the album's place there (it belonged to the old neighbours), +/// and undoing the move gives it back. +#[test] +fn moving_drops_the_place_and_undo_restores_it() { + let mut c = Catalog::new(); + let from = add(&mut c, "From", None, true); + let to = add(&mut c, "To", None, true); + let a = add(&mut c, "A", Some(from), false); + place(&mut c, a, Some(3)); + let inv = c.apply(Op::MoveAlbum { id: a, parent: Some(to) }).unwrap(); + assert_eq!((c.album(a).unwrap().parent, c.album(a).unwrap().order), (Some(to), None)); + c.apply(inv).unwrap(); + assert_eq!((c.album(a).unwrap().parent, c.album(a).unwrap().order), (Some(from), Some(3))); + // "moving" to where it already is changes nothing + c.apply(Op::MoveAlbum { id: a, parent: Some(from) }).unwrap(); + assert_eq!(c.album(a).unwrap().order, Some(3)); +} + +/// Catalogs written before albums had an order load as they were; an unset order is not written. +#[test] +fn albums_without_an_order_load_and_save_unchanged() { + let old = r#"{"id":4,"name":"Old","parent":null,"folder":false,"photos":[],"cover":null}"#; + let a: Album = serde_json::from_str(old).unwrap(); + assert_eq!(a.order, None); + assert!(!serde_json::to_string(&a).unwrap().contains("order")); + let mut a = a; + a.order = Some(2); + let back: Album = serde_json::from_str(&serde_json::to_string(&a).unwrap()).unwrap(); + assert_eq!(back.order, Some(2)); +} diff --git a/crates/catalog/src/tests_format_version.rs b/crates/catalog/src/tests_format_version.rs index 3f2e8816..cb80ce67 100644 --- a/crates/catalog/src/tests_format_version.rs +++ b/crates/catalog/src/tests_format_version.rs @@ -47,6 +47,19 @@ fn v1_library_loads_and_is_upgraded() { assert_eq!((r.replayed, r.upgraded_from), (1, None)); } +/// A library saved before albums could be ordered by hand opens as it was and is rewritten as v3. +#[test] +fn v2_library_loads_and_is_upgraded() { + let (base, log, full) = legacy_parts(); + let m = MemStore::new(); + m.set(SNAPSHOT, format!("{{\"format\":\"lightcraft-catalog\",\"version\":2,\"seq\":2,\"catalog\":{}}}\n", base.to_snapshot()).into_bytes()); + m.set(LOG, log.into_bytes()); + let (_, c, r) = Journal::open(Box::new(m.clone())).unwrap(); + assert_eq!(c.to_snapshot(), full.to_snapshot()); + assert_eq!((r.replayed, r.upgraded_from), (1, Some(2))); + assert_eq!(snapshot_version(&m), 3); +} + #[test] fn versionless_log_only_library_is_upgraded() { let mut c = Catalog::new(); @@ -179,8 +192,9 @@ fn op_variants_are_versioned() { | Op::SetLabelName { .. } | Op::Batch { .. } => 1, Op::SetBrowsed { .. } => 2, + Op::SetAlbumOrder { .. } => 3, } } - let newest = since(&Op::SetBrowsed { folder: String::new(), at: None }); + let newest = since(&Op::SetAlbumOrder { id: AlbumId(0), order: None }); assert_eq!(newest, VERSION, "the newest op's version must be the current format version"); } diff --git a/crates/catalog/src/tests_journal.rs b/crates/catalog/src/tests_journal.rs index 0c9513e3..0594b6d8 100644 --- a/crates/catalog/src/tests_journal.rs +++ b/crates/catalog/src/tests_journal.rs @@ -59,6 +59,7 @@ fn op_for(c: &mut Catalog, kind: u8, a: u8, b: u8) -> Op { cover: None, smart: None, quick: false, + order: None, }, } } @@ -77,6 +78,7 @@ fn op_for(c: &mut Catalog, kind: u8, a: u8, b: u8) -> Op { 13 => Op::Batch { ops: c.remove_from_stacks_ops(&[pid]) }, 14 => c.set_top_ops(pid).unwrap_or(Op::Batch { ops: vec![] }), 9 if b.is_multiple_of(5) => Op::SetLabelName { label: ColorLabel::ALL[a as usize % 5], name: (a % 3 != 1).then(|| format!("Label {b}")) }, + 9 if b % 7 == 3 => Op::SetAlbumOrder { id: aid, order: (!b.is_multiple_of(3)).then_some(u32::from(b % 5)) }, 9 if b.is_multiple_of(2) => { Op::SetFile { id: pid, file_name: format!("Renamed_{b}.jpg"), source: Source::File { path: format!("/p/Renamed_{b}.jpg") } } } diff --git a/crates/codecs/Cargo.toml b/crates/codecs/Cargo.toml index 68d5e8cd..18d00ac3 100644 --- a/crates/codecs/Cargo.toml +++ b/crates/codecs/Cargo.toml @@ -12,6 +12,7 @@ workspace = true lightcraft-geom = { workspace = true } lightcraft-color = { workspace = true } lightcraft-raster = { workspace = true } +lightcraft-heif = { workspace = true, optional = true } serde = { workspace = true } bytemuck = { workspace = true } thiserror = { workspace = true } @@ -34,6 +35,10 @@ ravif = { workspace = true, optional = true } default = ["parallel", "jxl", "avif"] parallel = ["dep:rayon", "lightcraft-raster/parallel", "jpeg-decoder/rayon"] jxl = ["dep:jxl-oxide"] +# HEIC/HEIF decode via the optional lightcraft-heif crate (heic-rs, pure Rust). Off by default so +# distributors choose it at build time (HEVC patents are a distributor's call, same policy as +# PhotoCraft); `cargo xtask ci` clippies and tests it (`--features heif`), releases don't ship it yet. +heif = ["dep:lightcraft-heif"] ## AVIF encoding via ravif/rav1e (pure Rust, BSD). Native targets only; a no-op on wasm32. avif = ["dep:ravif"] diff --git a/crates/codecs/src/heif.rs b/crates/codecs/src/heif.rs new file mode 100644 index 00000000..fe91c0af --- /dev/null +++ b/crates/codecs/src/heif.rs @@ -0,0 +1,125 @@ +//! HEIF / HEIC (the iPhone and Mac photo format), read-only. +//! +//! Decoding lives in the optional `lightcraft-heif` crate (heic-rs, a pure-Rust HEVC +//! still-picture decoder), enabled by this crate's `heif` feature — off by default because HEVC +//! is patent-encumbered and whether a build carries an HEVC decoder is the distributor's call +//! (the same policy as the sibling PhotoCraft). Without the feature, HEIF files are still +//! detected and opening one is an [`Error::Unsupported`] error, never a panic. +//! +//! HEIF records orientation in the container (`irot`/`imir`, plus a `clap` crop), not in EXIF, +//! so the decoder applies them and `orientation` is reported as 1: the EXIF Orientation tag only +//! mirrors what the container says, and applying it as well would turn the photo twice. +//! The samples come back in the container's own primaries with its ICC profile (if any); like +//! every decoder here, the colour interpretation happens in [`crate::convert::finish`]. + +#[cfg(feature = "heif")] +use crate::convert::{Buf, Meta, Model, Raw, finish}; +use crate::{DecodeOptions, Decoded, Error, Format, Result}; + +const F: Format = Format::Heif; + +/// The reason given when this build has no HEIF decoder. +#[cfg(not(feature = "heif"))] +pub(crate) const NOT_IN_BUILD: &str = "HEIC/HEIF support isn't included in this build of LightCraft"; + +/// The reason a HEIF file that no build can open is refused (this crate decodes HEVC stills +/// only; the wrapper names the exact construct in its own errors). +#[cfg(feature = "heif")] +const NOT_DECODED: &str = "the image is not a decodable HEVC still (image sequence, overlay, alpha-less grid or AVIF-coded item)"; + +#[cfg(not(feature = "heif"))] +pub(crate) fn decode(_bytes: &[u8], _opts: &DecodeOptions) -> Result { + Err(Error::Unsupported(F, NOT_IN_BUILD)) +} + +#[cfg(feature = "heif")] +fn err(e: lightcraft_heif::Error) -> Error { + match e { + // This crate's Unsupported takes a fixed reason; the dynamic wording lives in the + // lightcraft-heif error itself. + lightcraft_heif::Error::Unsupported(_) => Error::Unsupported(F, NOT_DECODED), + lightcraft_heif::Error::Limit(m) | lightcraft_heif::Error::Malformed(m) => Error::Malformed(F, m), + } +} + +#[cfg(feature = "heif")] +pub(crate) fn decode(bytes: &[u8], opts: &DecodeOptions) -> Result { + // The container alone: the declared size is checked before any pixel is decoded. + let info = lightcraft_heif::probe(bytes).map_err(err)?; + if (info.width as u64).saturating_mul(info.height as u64) > opts.max_pixels { + return Err(Error::TooLarge(info.width as u64, info.height as u64)); + } + let options = lightcraft_heif::Options { max_pixels: opts.max_pixels, apply_transforms: true }; + let decoded = lightcraft_heif::decode(bytes, &options).map_err(err)?; + // The decoder re-reads the container; trust its output's shape, not the probe's. + let buf = if decoded.sixteen_bit { + let mut v = Vec::with_capacity(decoded.data.len() / 2); + for pair in decoded.data.as_chunks::<2>().0 { + v.push(u16::from_ne_bytes(*pair)); + } + Buf::U16(v) + } else { + Buf::U8(decoded.data) + }; + let raw = Raw { + width: decoded.width as usize, + height: decoded.height as usize, + model: Model::Rgb, + alpha: decoded.has_alpha, + premultiplied: false, + buf, + bit_depth: if decoded.sixteen_bit { 16 } else { 8 }, + }; + let meta = Meta { icc: decoded.icc, exif: decoded.exif, xmp: decoded.xmp, orientation: Some(1), ..Default::default() }; + finish(F, raw, meta, (decoded.width, decoded.height), opts) +} + +#[cfg(test)] +mod tests { + use crate::{Format, sniff}; + + /// The container header both builds recognise, whatever the feature says. + const HEIC_HEADER: &[u8] = b"\0\0\0\x18ftypheic\0\0\0\0mif1heic"; + + #[test] + fn heif_is_always_recognised() { + assert_eq!(sniff(HEIC_HEADER), Some(Format::Heif)); + } + + #[test] + fn can_decode_follows_the_feature() { + assert_eq!(Format::Heif.can_decode(), cfg!(feature = "heif")); + } + + #[cfg(not(feature = "heif"))] + #[test] + fn without_the_feature_heif_is_a_clear_unsupported_error() { + let r = crate::decode(HEIC_HEADER, crate::DecodeOptions::default()); + assert!(matches!(&r, Err(crate::Error::Unsupported(Format::Heif, why)) if why.contains("isn't included in this build")), "{r:?}"); + } + + /// Found by PhotoCraft's `decode_heif` fuzz target: a malformed box makes heic-rs 0.1.1 slice + /// out of range (`boxes.rs:130`). It must be an error, not a crash. + #[cfg(feature = "heif")] + const HEIC_RS_BOX_PANIC: [u8; 72] = [ + 0x00, 0x00, 0x00, 0x24, 0x66, 0x74, 0x79, 0x70, 0x68, 0x65, 0x69, 0x63, 0x00, 0x00, 0x00, 0x00, 0x6d, 0x69, 0x66, 0x31, 0x4d, 0x69, 0x50, + 0x72, 0x6d, 0x69, 0x61, 0x66, 0x4d, 0x69, 0x48, 0x42, 0x72, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x6d, 0x65, 0x74, 0x61, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x75, 0x75, 0x69, 0x64, 0x00, 0x00, 0x00, 0x00, 0x03, 0x08, 0x08, 0x08, 0x00, 0x03, 0x01, 0x03, 0x70, + 0x00, 0xa8, 0x00, + ]; + + #[cfg(feature = "heif")] + #[test] + fn a_heic_rs_panic_is_a_malformed_error() { + let r = crate::decode(&HEIC_RS_BOX_PANIC, crate::DecodeOptions::default()); + assert!(matches!(&r, Err(crate::Error::Malformed(Format::Heif, _))), "{r:?}"); + } + + #[cfg(feature = "heif")] + #[test] + fn an_image_sequence_is_unsupported() { + // An ftyp with no meta box: a video track, not a still. + let r = crate::decode(b"\0\0\0\x18ftypmsf1\0\0\0\0msf1hevc", crate::DecodeOptions::default()); + assert!(matches!(&r, Err(crate::Error::Unsupported(Format::Heif, _))), "{r:?}"); + } +} diff --git a/crates/codecs/src/lib.rs b/crates/codecs/src/lib.rs index 15990329..169f40c1 100644 --- a/crates/codecs/src/lib.rs +++ b/crates/codecs/src/lib.rs @@ -19,6 +19,7 @@ mod convert; pub mod encode; pub mod exif; +mod heif; pub mod icc; mod jpeg; pub mod jpeg_par; @@ -227,8 +228,8 @@ pub fn decode_unguarded(bytes: &[u8], format: Format, opts: &DecodeOptions) -> R Format::Jxl => jxl::decode(bytes, &opts), #[cfg(not(feature = "jxl"))] Format::Jxl => Err(Error::Unsupported(format, "built without the `jxl` feature")), + Format::Heif => heif::decode(bytes, &opts), Format::Avif => Err(Error::Unsupported(format, "no pure-Rust, permissively licensed AV1 decoder yet")), - Format::Heif => Err(Error::Unsupported(format, "no pure-Rust, permissively licensed HEVC decoder yet")), Format::RawTiffLike | Format::RawOther => Err(Error::Unsupported(format, "camera raw: decode with lightcraft-raw")), } } diff --git a/crates/codecs/src/sniff.rs b/crates/codecs/src/sniff.rs index a02d87d2..50fa3ccf 100644 --- a/crates/codecs/src/sniff.rs +++ b/crates/codecs/src/sniff.rs @@ -66,7 +66,8 @@ impl Format { pub fn can_decode(self) -> bool { match self { Format::Jxl => cfg!(feature = "jxl"), - Format::Avif | Format::Heif | Format::RawTiffLike | Format::RawOther => false, + Format::Heif => cfg!(feature = "heif"), + Format::Avif | Format::RawTiffLike | Format::RawOther => false, _ => true, } } diff --git a/crates/denoise-core/Cargo.toml b/crates/denoise-core/Cargo.toml new file mode 100644 index 00000000..b65f90ae --- /dev/null +++ b/crates/denoise-core/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "lightcraft-denoise-core" +version.workspace = true +edition.workspace = true +license.workspace = true +description = "LightCraft denoise model metadata, Bayer tiling and cache formats (L1; no inference backend)" + +[lints] +workspace = true + +[features] +default = [] +rawnind-model = [] + +[dependencies] +lightcraft-raster = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +sha2 = { workspace = true } +thiserror = { workspace = true } +miniz_oxide = { workspace = true } +rayon = { workspace = true } +web-time = { workspace = true } + +[dev-dependencies] +proptest = { workspace = true } diff --git a/crates/denoise-core/README.md b/crates/denoise-core/README.md new file mode 100644 index 00000000..74d1e324 --- /dev/null +++ b/crates/denoise-core/README.md @@ -0,0 +1,7 @@ +# lightcraft-denoise-core + +L1 model metadata, bounded archive/cache formats and Bayer packing/tiling helpers. +No weights or inference backend. See [the user workflow](../../docs/denoise.md). +`lightcraft-denoise` re-exports these types for CPU/GPU users; the engine uses this +crate directly when inference is disabled. The default download catalog is empty. +Run `cargo test -p lightcraft-denoise-core`; optional catalog checks use `--features rawnind-model`. diff --git a/crates/denoise-core/src/archive.rs b/crates/denoise-core/src/archive.rs new file mode 100644 index 00000000..ccfea471 --- /dev/null +++ b/crates/denoise-core/src/archive.rs @@ -0,0 +1,279 @@ +//! Taking one file out of a `.zip` (a `.dtmodel` is one), the little of the format a model download needs: stored and +//! deflated entries, found by name, checked against their CRC-32 and a size limit, written next to where they are +//! going and moved into place. It is a stranger's file, so every length and offset in it is checked, a name is only +//! ever compared (never used as a path), and zip64, encrypted and multi-disk archives are refused. + +use std::io::{Read, Seek, SeekFrom, Write}; +use std::path::Path; + +/// How far from the end the end-of-central-directory record can start (a 22-byte record and a 64 KiB comment). +const EOCD_SEARCH: u64 = 22 + 65535; +/// Largest central directory we read. +const MAX_DIRECTORY: u64 = 1 << 20; + +#[derive(Debug, PartialEq, thiserror::Error)] +pub enum ArchiveError { + #[error("the archive could not be read: {0}")] + Io(String), + #[error("not a usable archive: {0}")] + Format(String), + #[error("the archive has no file called `{0}`")] + Missing(String), +} + +impl From for ArchiveError { + fn from(e: std::io::Error) -> Self { + ArchiveError::Io(e.to_string()) + } +} + +fn format(why: &str) -> Result { + Err(ArchiveError::Format(why.into())) +} + +fn u16_at(b: &[u8], at: usize) -> Option { + Some(u16::from_le_bytes(b.get(at..at.checked_add(2)?)?.try_into().ok()?)) +} + +fn u32_at(b: &[u8], at: usize) -> Option { + Some(u32::from_le_bytes(b.get(at..at.checked_add(4)?)?.try_into().ok()?)) +} + +/// CRC-32 (IEEE) of `bytes`. +fn crc32(bytes: &[u8]) -> u32 { + let mut table = [0u32; 256]; + for (i, t) in table.iter_mut().enumerate() { + let mut c = i as u32; + for _ in 0..8 { + c = if c & 1 != 0 { 0xedb8_8320 ^ (c >> 1) } else { c >> 1 }; + } + *t = c; + } + !bytes.iter().fold(!0u32, |c, &b| table[((c ^ u32::from(b)) & 0xff) as usize] ^ (c >> 8)) +} + +struct Entry { + method: u16, + crc: u32, + compressed: u64, + size: u64, + local_header: u64, +} + +/// Find the entry called `name` (exactly, or as the last part of a longer name: `folder/name`) in the central directory. +fn find(f: &mut (impl Read + Seek), name: &str) -> Result { + let len = f.seek(SeekFrom::End(0))?; + let tail_len = len.min(EOCD_SEARCH); + f.seek(SeekFrom::Start(len - tail_len))?; + let mut tail = vec![0u8; tail_len as usize]; + f.read_exact(&mut tail)?; + let eocd = (0..tail.len().saturating_sub(21)).rev().find(|&i| u32_at(&tail, i) == Some(0x0605_4b50)); + let Some(e) = eocd else { return format("this is not a zip file") }; + let (entries, size, offset) = (u16_at(&tail, e + 10), u32_at(&tail, e + 12), u32_at(&tail, e + 16)); + let (Some(entries), Some(size), Some(offset)) = (entries, size, offset) else { return format("the end of the archive is cut short") }; + if entries == 0xffff || size == u32::MAX || offset == u32::MAX { + return format("zip64 archives are not supported"); + } + if u64::from(size) > MAX_DIRECTORY || u64::from(offset) + u64::from(size) > len { + return format("the list of files is not where it says"); + } + f.seek(SeekFrom::Start(u64::from(offset)))?; + let mut dir = vec![0u8; size as usize]; + f.read_exact(&mut dir)?; + let mut at = 0usize; + for _ in 0..entries { + if u32_at(&dir, at) != Some(0x0201_4b50) { + return format("the list of files is damaged"); + } + let get = |o: usize| at.checked_add(o); + let (flags, method, crc) = (u16_at(&dir, at + 8), u16_at(&dir, at + 10), u32_at(&dir, at + 16)); + let (compressed, uncompressed) = (u32_at(&dir, at + 20), u32_at(&dir, at + 24)); + let (n, x, c) = (u16_at(&dir, at + 28), u16_at(&dir, at + 30), u16_at(&dir, at + 32)); + let local = u32_at(&dir, at + 42); + let (Some(flags), Some(method), Some(crc), Some(compressed), Some(uncompressed), Some(n), Some(x), Some(c), Some(local), Some(start)) = + (flags, method, crc, compressed, uncompressed, n, x, c, local, get(46)) + else { + return format("the list of files is cut short"); + }; + let Some(entry_name) = dir.get(start..start + usize::from(n)) else { return format("a file name is cut short") }; + let matches = entry_name == name.as_bytes() || entry_name.strip_suffix(name.as_bytes()).is_some_and(|p| p.ends_with(b"/")); + if matches { + if flags & 1 != 0 { + return format("encrypted files are not supported"); + } + if compressed == u32::MAX || uncompressed == u32::MAX { + return format("zip64 files are not supported"); + } + return Ok(Entry { method, crc, compressed: u64::from(compressed), size: u64::from(uncompressed), local_header: u64::from(local) }); + } + at = start + usize::from(n) + usize::from(x) + usize::from(c); + } + Err(ArchiveError::Missing(name.into())) +} + +/// Take the file called `name` out of the archive at `archive` and write it to `dest` (through a temporary file next +/// to it); `max_bytes` bounds its size. Returns the number of bytes written. +pub fn extract(archive: &Path, name: &str, dest: &Path, max_bytes: u64) -> Result { + let mut f = std::io::BufReader::new(std::fs::File::open(archive)?); + let e = find(&mut f, name)?; + if e.size > max_bytes || e.compressed > max_bytes { + return format("the file in the archive is larger than allowed"); + } + f.seek(SeekFrom::Start(e.local_header))?; + let mut head = [0u8; 30]; + f.read_exact(&mut head)?; + if u32_at(&head, 0) != Some(0x0403_4b50) { + return format("the file's header is not where it says"); + } + let skip = u64::from(u16_at(&head, 26).unwrap_or(0)) + u64::from(u16_at(&head, 28).unwrap_or(0)); + f.seek_relative(i64::try_from(skip).map_err(|_| ArchiveError::Format("a header is too long".into()))?)?; + let len = usize::try_from(e.compressed) + .ok() + .filter(|&n| n <= 256 * 1024 * 1024) + .ok_or_else(|| ArchiveError::Format("the compressed model is too large".into()))?; + if e.size > 256 * 1024 * 1024 { + return format("the extracted model is too large"); + } + let mut packed = Vec::new(); + packed.try_reserve_exact(len).map_err(|_| ArchiveError::Format("not enough memory for the model archive".into()))?; + packed.resize(len, 0u8); + f.read_exact(&mut packed)?; + let data = match e.method { + 0 => packed, + 8 => miniz_oxide::inflate::decompress_to_vec_with_limit(&packed, max_bytes.min(256 * 1024 * 1024) as usize) + .map_err(|err| ArchiveError::Format(format!("{err:?}")))?, + m => return format(&format!("compression method {m} is not supported")), + }; + if data.len() as u64 != e.size || crc32(&data) != e.crc { + return format("the file in the archive is damaged"); + } + let tmp = dest.with_extension("part"); + let written = (|| -> std::io::Result<()> { + let mut out = std::fs::File::create(&tmp)?; + out.write_all(&data)?; + out.flush()?; + drop(out); + std::fs::rename(&tmp, dest) + })(); + if let Err(err) = written { + let _ = std::fs::remove_file(&tmp); + return Err(err.into()); + } + Ok(data.len() as u64) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A zip with the given (name, bytes, deflate) entries, written by hand. + pub(crate) fn make_zip(files: &[(&str, &[u8], bool)]) -> Vec { + let mut out = Vec::new(); + let mut dir = Vec::new(); + for (name, data, deflate) in files { + let packed = if *deflate { miniz_oxide::deflate::compress_to_vec(data, 6) } else { data.to_vec() }; + let offset = out.len() as u32; + let method: u16 = if *deflate { 8 } else { 0 }; + out.extend_from_slice(&0x0403_4b50u32.to_le_bytes()); + out.extend_from_slice(&[20, 0, 0, 0]); + out.extend_from_slice(&method.to_le_bytes()); + out.extend_from_slice(&[0, 0, 0, 0]); + out.extend_from_slice(&crc32(data).to_le_bytes()); + out.extend_from_slice(&(packed.len() as u32).to_le_bytes()); + out.extend_from_slice(&(data.len() as u32).to_le_bytes()); + out.extend_from_slice(&(name.len() as u16).to_le_bytes()); + out.extend_from_slice(&[0, 0]); + out.extend_from_slice(name.as_bytes()); + out.extend_from_slice(&packed); + dir.extend_from_slice(&0x0201_4b50u32.to_le_bytes()); + dir.extend_from_slice(&[20, 0, 20, 0, 0, 0]); + dir.extend_from_slice(&method.to_le_bytes()); + dir.extend_from_slice(&[0, 0, 0, 0]); + dir.extend_from_slice(&crc32(data).to_le_bytes()); + dir.extend_from_slice(&(packed.len() as u32).to_le_bytes()); + dir.extend_from_slice(&(data.len() as u32).to_le_bytes()); + dir.extend_from_slice(&(name.len() as u16).to_le_bytes()); + dir.extend_from_slice(&[0; 12]); + dir.extend_from_slice(&offset.to_le_bytes()); + dir.extend_from_slice(name.as_bytes()); + } + let cd_offset = out.len() as u32; + out.extend_from_slice(&dir); + out.extend_from_slice(&0x0605_4b50u32.to_le_bytes()); + out.extend_from_slice(&[0, 0, 0, 0]); + out.extend_from_slice(&(files.len() as u16).to_le_bytes()); + out.extend_from_slice(&(files.len() as u16).to_le_bytes()); + out.extend_from_slice(&(dir.len() as u32).to_le_bytes()); + out.extend_from_slice(&cd_offset.to_le_bytes()); + out.extend_from_slice(&[0, 0]); + out + } + + fn temp(name: &str) -> std::path::PathBuf { + std::env::temp_dir().join(format!("lc-denoise-zip-{name}-{}", std::process::id())) + } + + #[test] + fn a_stored_and_a_deflated_file_come_out_whole() { + let big: Vec = (0..200_000u32).map(|i| (i * 7 % 251) as u8).collect(); + let zip = make_zip(&[("m/config.json", b"{}", false), ("m/model.onnx", &big, true), ("m/other.onnx", b"zz", true)]); + let (a, out) = (temp("a.zip"), temp("a.onnx")); + std::fs::write(&a, &zip).unwrap(); + assert_eq!(extract(&a, "model.onnx", &out, 1 << 20).unwrap(), big.len() as u64); + assert_eq!(std::fs::read(&out).unwrap(), big); + assert_eq!(extract(&a, "m/config.json", &out, 1 << 20).unwrap(), 2); + assert_eq!(std::fs::read(&out).unwrap(), b"{}"); + // a name must match whole path parts: `del.onnx` is not `model.onnx` + assert!(matches!(extract(&a, "del.onnx", &out, 1 << 20), Err(ArchiveError::Missing(_)))); + assert!(matches!(extract(&a, "nothing", &out, 1 << 20), Err(ArchiveError::Missing(_)))); + // the size limit applies + assert!(matches!(extract(&a, "model.onnx", &out, 1000), Err(ArchiveError::Format(_)))); + for p in [a, out] { + let _ = std::fs::remove_file(p); + } + } + + /// Opt-in: `LC_DENOISE_DTMODEL= cargo test -p lightcraft-denoise -- --ignored archive` + #[test] + #[ignore = "needs the real download: set LC_DENOISE_DTMODEL"] + fn the_real_download_gives_the_model_it_should() { + let Some(path) = std::env::var_os("LC_DENOISE_DTMODEL") else { return }; + let known = crate::known::find("rawnind-bayer").unwrap(); + let (d, entry) = (known.download.unwrap(), known.entry.unwrap()); + assert_eq!(crate::hash::sha256_file(Path::new(&path)).unwrap(), d.sha256); + let out = temp("real.onnx"); + let n = extract(Path::new(&path), entry, &out, crate::manifest::MAX_MODEL_BYTES).unwrap(); + assert_eq!(Some(n), known.manifest.size_bytes); + assert_eq!(Some(crate::hash::sha256_file(&out).unwrap()), known.manifest.sha256); + let _ = std::fs::remove_file(out); + } + + #[test] + fn damaged_and_hostile_archives_are_errors_not_panics() { + let zip = make_zip(&[("m/model.onnx", &vec![9u8; 5000], true)]); + let (a, out) = (temp("b.zip"), temp("b.onnx")); + // cut anywhere + for n in (0..zip.len()).step_by(3) { + std::fs::write(&a, &zip[..n]).unwrap(); + assert!(extract(&a, "model.onnx", &out, 1 << 20).is_err(), "cut at {n}"); + } + // one byte changed anywhere: an error, or (when it falls in an unused field) the right file, never a panic + for at in 0..zip.len() { + let mut bad = zip.clone(); + bad[at] ^= 0x5a; + std::fs::write(&a, &bad).unwrap(); + if extract(&a, "model.onnx", &out, 1 << 20).is_ok() { + assert_eq!(std::fs::read(&out).unwrap(), vec![9u8; 5000], "a changed byte at {at} gave another file"); + } + } + // not a zip at all, and an empty file + for junk in [Vec::new(), b"PK".to_vec(), vec![0u8; 100_000]] { + std::fs::write(&a, junk).unwrap(); + assert!(extract(&a, "model.onnx", &out, 1 << 20).is_err()); + } + assert!(extract(&temp("missing.zip"), "model.onnx", &out, 1 << 20).is_err()); + for p in [a, out] { + let _ = std::fs::remove_file(p); + } + } +} diff --git a/crates/denoise-core/src/bayer.rs b/crates/denoise-core/src/bayer.rs new file mode 100644 index 00000000..a4a00d94 --- /dev/null +++ b/crates/denoise-core/src/bayer.rs @@ -0,0 +1,191 @@ +//! Bayer geometry: which way a mosaic is laid out, and how a patch of it is packed into the four planes +//! (red, green on the red row, green on the blue row, blue) a denoise model takes. +//! +//! Every Bayer layout is the same 2 × 2 cell shifted by one pixel at most, so a layout is just where its +//! first red sample is. The cell grid starts one pixel before the picture when that red sample is in +//! column or row 1, and the picture is read with its edge mirrored (without repeating the edge sample, so +//! the colour of a mirrored sample is the colour that belongs at that place). + +/// Where the first red sample of a Bayer mosaic sits, `(dx, dy)` in `{0, 1}²`: RGGB is `(0, 0)`, GRBG `(1, 0)`, +/// GBRG `(0, 1)`, BGGR `(1, 1)`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Layout { + pub dx: usize, + pub dy: usize, +} + +impl Layout { + /// The layout of a 2 × 2 colour pattern, row-major, colours 0 = red, 1 = green, 2 = blue; `None` when + /// it is not a Bayer cell (one red, one blue, the greens on the other diagonal). + pub fn from_cell(cell: [u8; 4]) -> Option { + let at = |c: u8| cell.iter().position(|&v| v == c); + let red = at(0)?; + let blue = at(2)?; + let greens = cell.iter().filter(|&&v| v == 1).count(); + // red and blue sit on one diagonal: indices 0+3 or 1+2 + if greens != 2 || red + blue != 3 { + return None; + } + Some(Layout { dx: red % 2, dy: red / 2 }) + } + + /// The cell grid's origin: `-1` along an axis whose first red sample is at 1, so that cell 0 starts on a red column or row. + fn origin(self) -> (isize, isize) { + (isize::try_from(self.dx).unwrap_or(isize::MAX).saturating_neg(), isize::try_from(self.dy).unwrap_or(isize::MAX).saturating_neg()) + } + + /// Cells across and down covering a `width × height` mosaic. + pub fn cells(self, width: usize, height: usize) -> (usize, usize) { + let cells = |n: usize, phase: usize| n / 2 + (n % 2 + phase.min(1)).div_ceil(2); + (cells(width, self.dx), cells(height, self.dy)) + } + + /// The mosaic position of cell `(cx, cy)`'s top-left sample (may lie outside the picture). + pub fn position(self, cx: isize, cy: isize) -> (isize, isize) { + let (ox, oy) = self.origin(); + (ox.saturating_add(cx.saturating_mul(2)), oy.saturating_add(cy.saturating_mul(2))) + } +} + +/// `i` mirrored into `0..n` (ping-pong, edge not repeated), so `-1` is `1` and `n` is `n - 2`. `n == 0` gives 0. +pub fn reflect(i: isize, n: usize) -> usize { + if n <= 1 { + return 0; + } + if let Some(period) = isize::try_from(n).ok().and_then(|n| (n - 1).checked_mul(2)) { + let m = i.rem_euclid(period); + return (if m < n as isize { m } else { period - m }) as usize; + } + let period = 2 * (n as i128 - 1); + let m = (i as i128).rem_euclid(period); + // The reflected value is in 0..n, so this conversion always fits usize. + (if m < n as i128 { m } else { period - m }) as usize +} + +/// Pack the `tile × tile` cells whose top-left cell is `(cx0, cy0)` into `out` as four planes `[R, G1, G2, B]` +/// (`4 · tile²` values, NCHW). G1 is the green beside the red sample, G2 the green under it. Samples outside +/// the picture are mirrored. Returns `false` (and writes nothing) when the sizes do not fit. +pub fn pack_tile(mosaic: &[f32], width: usize, height: usize, layout: Layout, cx0: isize, cy0: isize, tile: usize, out: &mut [f32]) -> bool { + let Some(plane) = tile.checked_mul(tile) else { return false }; + let Some(values) = plane.checked_mul(4) else { return false }; + if width == 0 + || height == 0 + || tile == 0 + || layout.dx > 1 + || layout.dy > 1 + || width.checked_mul(height) != Some(mosaic.len()) + || out.len() != values + { + return false; + } + let Ok(last) = isize::try_from(tile - 1) else { return false }; + let (ox, oy) = layout.origin(); + for (start, origin) in [(cx0, ox), (cy0, oy)] { + if start.checked_mul(2).and_then(|v| v.checked_add(origin)).is_none() + || start.checked_add(last).and_then(|v| v.checked_mul(2)).and_then(|v| v.checked_add(origin)).and_then(|v| v.checked_add(1)).is_none() + { + return false; + } + } + // the four samples of a cell, as (dx, dy) from its top-left sample + const SUB: [(isize, isize); 4] = [(0, 0), (1, 0), (0, 1), (1, 1)]; + for (p, &(sx, sy)) in SUB.iter().enumerate() { + let Some(dst) = out.get_mut(p * plane..(p + 1) * plane) else { return false }; + for (j, row) in dst.chunks_exact_mut(tile).enumerate() { + let (_, y) = layout.position(cx0, cy0 + j as isize); + let ry = reflect(y + sy, height); + let Some(src) = mosaic.get(ry * width..(ry + 1) * width) else { return false }; + for (i, v) in row.iter_mut().enumerate() { + let (x, _) = layout.position(cx0 + i as isize, cy0); + *v = src.get(reflect(x + sx, width)).copied().unwrap_or(0.0); + } + } + } + true +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn hostile_sizes_and_positions_never_panic_or_write_output() { + let layout = Layout::from_cell([0, 1, 1, 2]).unwrap(); + let mut out = [7.0; 4]; + for (w, h, tile, x) in [(usize::MAX, 2, 1, 0), (1, 1, usize::MAX, 0), (1, 1, 0, 0), (1, 1, 1, isize::MAX), (1, 1, 1, isize::MIN)] { + assert!(!pack_tile(&[0.5], w, h, layout, x, 0, tile, &mut out)); + assert_eq!(out, [7.0; 4]); + } + for n in [1, 2, isize::MAX as usize, usize::MAX] { + for i in [isize::MIN, -1, 0, isize::MAX] { + assert!(reflect(i, n) < n); + } + } + assert_eq!(layout.cells(usize::MAX, usize::MAX), (usize::MAX / 2 + 1, usize::MAX / 2 + 1)); + } + + #[test] + fn layouts_from_cells() { + assert_eq!(Layout::from_cell([0, 1, 1, 2]), Some(Layout { dx: 0, dy: 0 })); + assert_eq!(Layout::from_cell([1, 0, 2, 1]), Some(Layout { dx: 1, dy: 0 })); + assert_eq!(Layout::from_cell([1, 2, 0, 1]), Some(Layout { dx: 0, dy: 1 })); + assert_eq!(Layout::from_cell([2, 1, 1, 0]), Some(Layout { dx: 1, dy: 1 })); + // not Bayer: two reds, a missing blue, greens on the wrong diagonal + assert_eq!(Layout::from_cell([0, 0, 1, 2]), None); + assert_eq!(Layout::from_cell([0, 1, 2, 1]), None); + assert_eq!(Layout::from_cell([1, 1, 1, 1]), None); + } + + #[test] + fn reflection_keeps_the_colour_of_a_site() { + assert_eq!(reflect(-1, 10), 1); + assert_eq!(reflect(-2, 10), 2); + assert_eq!(reflect(10, 10), 8); + assert_eq!(reflect(11, 10), 7); + assert_eq!(reflect(25, 10), 7); + assert_eq!(reflect(5, 0), 0); + assert_eq!(reflect(5, 1), 0); + for i in -40..60 { + assert_eq!(reflect(i, 11) % 2, (i.rem_euclid(2)) as usize, "parity at {i}"); + } + } + + #[test] + fn packing_puts_each_colour_in_its_plane_for_every_layout() { + // a mosaic whose sample value says which colour it is: R = 10, G = 20, B = 30 + for cell in [[0u8, 1, 1, 2], [1, 0, 2, 1], [1, 2, 0, 1], [2, 1, 1, 0]] { + let layout = Layout::from_cell(cell).unwrap(); + let (w, h) = (10usize, 8usize); + let mosaic: Vec = (0..w * h).map(|i| [10.0, 20.0, 30.0][cell[(i / w % 2) * 2 + i % w % 2] as usize]).collect(); + let (nx, ny) = layout.cells(w, h); + let tile = nx.max(ny); + let mut out = vec![0.0; 4 * tile * tile]; + assert!(pack_tile(&mosaic, w, h, layout, 0, 0, tile, &mut out)); + let plane = tile * tile; + for (p, want) in [10.0, 20.0, 20.0, 30.0].into_iter().enumerate() { + assert!(out[p * plane..(p + 1) * plane].iter().all(|&v| v == want), "layout {cell:?} plane {p}"); + } + } + } + + #[test] + fn g1_is_beside_the_red_sample_and_g2_under_it() { + // RGGB with a distinct value at every site + let (w, h) = (4, 4); + let mosaic: Vec = (0..16).map(|i| i as f32).collect(); + let mut out = vec![0.0; 4 * 4]; + assert!(pack_tile(&mosaic, w, h, Layout { dx: 0, dy: 0 }, 0, 0, 2, &mut out)); + // cell (0,0): R = site (0,0)=0, G1 = site (1,0)=1, G2 = site (0,1)=4, B = site (1,1)=5 + assert_eq!([out[0], out[4], out[8], out[12]], [0.0, 1.0, 4.0, 5.0]); + // cell (1,0): sites (2,0)=2, (3,0)=3, (2,1)=6, (3,1)=7 + assert_eq!([out[1], out[5], out[9], out[13]], [2.0, 3.0, 6.0, 7.0]); + } + + #[test] + fn packing_rejects_wrong_sizes() { + let mut out = vec![0.0; 16]; + assert!(!pack_tile(&[0.0; 15], 4, 4, Layout { dx: 0, dy: 0 }, 0, 0, 2, &mut out)); + assert!(!pack_tile(&[0.0; 16], 4, 4, Layout { dx: 0, dy: 0 }, 0, 0, 3, &mut out)); + assert!(!pack_tile(&[], 0, 0, Layout { dx: 0, dy: 0 }, 0, 0, 2, &mut out)); + } +} diff --git a/crates/denoise-core/src/hash.rs b/crates/denoise-core/src/hash.rs new file mode 100644 index 00000000..50f83be3 --- /dev/null +++ b/crates/denoise-core/src/hash.rs @@ -0,0 +1,46 @@ +//! SHA-256 of a model file, streamed (a model can be hundreds of megabytes). + +use std::io::{self, Read}; +use std::path::Path; + +use sha2::{Digest, Sha256}; + +/// Lowercase hex SHA-256 of everything `r` yields. +pub fn sha256_reader(mut r: R) -> io::Result { + let mut h = Sha256::new(); + let mut buf = vec![0u8; 1 << 20]; + loop { + let n = r.read(&mut buf)?; + if n == 0 { + break; + } + h.update(buf.get(..n).unwrap_or(&[])); + } + Ok(hex(&h.finalize())) +} + +/// Lowercase hex SHA-256 of a file. +pub fn sha256_file(path: &Path) -> io::Result { + sha256_reader(io::BufReader::new(std::fs::File::open(path)?)) +} + +/// Lowercase hex SHA-256 of `bytes`. +pub fn sha256_hex(bytes: &[u8]) -> String { + hex(&Sha256::digest(bytes)) +} + +fn hex(bytes: &[u8]) -> String { + bytes.iter().map(|b| format!("{b:02x}")).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn known_digests() { + assert_eq!(sha256_hex(b""), "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"); + assert_eq!(sha256_hex(b"abc"), "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"); + assert_eq!(sha256_reader(&b"abc"[..]).unwrap(), sha256_hex(b"abc")); + } +} diff --git a/crates/denoise-core/src/known.rs b/crates/denoise-core/src/known.rs new file mode 100644 index 00000000..f9a51109 --- /dev/null +++ b/crates/denoise-core/src/known.rs @@ -0,0 +1,128 @@ +//! The denoise models LightCraft knows about: what each one is, who may use it, and where to fetch it. +//! +//! None is part of LightCraft. A model with an address here can be downloaded when the user asks and has accepted +//! its terms; anything else is brought by the user as an `.onnx` file with a `denoise-model.json` beside it. + +#[cfg(any(feature = "rawnind-model", test))] +use crate::licence::{Commercial, Licence}; + +use crate::manifest::DenoiserManifest; +#[cfg(any(feature = "rawnind-model", test))] +use crate::manifest::{Domain, Gain}; + +/// A known model, and how to get it. +pub struct Known { + pub manifest: DenoiserManifest, + /// What to fetch (an archive or the `.onnx` itself); `None` when the user must bring the file. + pub download: Option, + /// The `.onnx` inside the downloaded archive (`None`: the download is the `.onnx`). + pub entry: Option<&'static str>, +} + +/// RawNIND's UtNet2 for Bayer sensors, as published for darktable (`rawdenoise-nind.dtmodel`, a zip with a Bayer and a +/// linear model). Its weights are GPL-3.0; the training pictures are CC BY 4.0 and CC0. +#[cfg(any(feature = "rawnind-model", test))] +fn rawnind_bayer() -> Known { + let manifest = DenoiserManifest { + id: "rawnind-bayer".into(), + name: "RawNIND UtNet2 (Bayer)".into(), + version: "1.0".into(), + licence: Licence { + name: "GPL-3.0".into(), + commercial: Commercial::Yes, + url: Some("https://www.gnu.org/licenses/gpl-3.0.html".into()), + notice: "The weights are published under the GNU General Public License v3. LightCraft does not include them: they \ + are downloaded to your computer at your request and used by it only. If you pass the model file on, the \ + GPL applies to that copy. Commercial use is subject to the GPL terms, including applicable source, licence and notice obligations on redistribution." + .into(), + }, + source: Some("https://github.com/darktable-org/darktable-ai/tree/master/models/rawdenoise-nind".into()), + sha256: Some("da27509dab6a2915da67e988acd86cf71f9d5bbc8d1aa0ed32933578a887b901".into()), + size_bytes: Some(31_056_425), + provenance: "Trained on RawNIND, real photographs of the same scenes taken noisy and clean (images CC BY 4.0 and CC0, \ + published on Wikimedia Commons and by the University of Louvain). Paper: arXiv 2501.08924." + .into(), + domain: Domain::BayerToRgb, + tile: 512, + overlap: 64, + gain: Gain::MatchMean { nominal: 1.0e6, max_deviation: 0.05 }, + }; + let download = Download { + id: manifest.id.clone(), + url: "https://github.com/darktable-org/darktable-ai/releases/download/release-5.6.0/rawdenoise-nind.dtmodel".into(), + file_name: "rawdenoise-nind.dtmodel".into(), + size_bytes: 57_700_134, + sha256: "d71b5f1e727c85a359e6f74dca9e2016c9d8fc3e2f7ac3e9b347d80ceca969af".into(), + }; + Known { manifest, download: Some(download), entry: Some("model_bayer.onnx") } +} + +/// Every model LightCraft knows. +pub fn all() -> Vec { + #[cfg(feature = "rawnind-model")] + { + vec![rawnind_bayer()] + } + #[cfg(not(feature = "rawnind-model"))] + { + Vec::new() + } +} + +/// The known model with this id. +pub fn find(id: &str) -> Option { + all().into_iter().find(|k| k.manifest.id == id) +} + +/// The known model whose `.onnx` has this SHA-256. +pub fn by_sha256(sha256: &str) -> Option { + all().into_iter().map(|k| k.manifest).find(|m| m.sha256.as_deref() == Some(sha256)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::manifest::validate; + + #[test] + fn known_models_are_valid_unique_and_downloaded_over_https() { + let all = all(); + assert_eq!(all.is_empty(), !cfg!(feature = "rawnind-model")); + validate(&rawnind_bayer().manifest).unwrap(); + for (i, k) in all.iter().enumerate() { + validate(&k.manifest).unwrap(); + assert!(all.iter().skip(i + 1).all(|o| o.manifest.id != k.manifest.id), "unique ids"); + assert!(!k.manifest.licence.name.is_empty() && !k.manifest.licence.notice.is_empty(), "terms are shown before use"); + if let Some(d) = &k.download { + assert_eq!(d.id, k.manifest.id); + assert!(d.url.starts_with("https://") && d.sha256.len() == 64 && d.size_bytes > 0); + // a name that cannot leave the staging folder + assert!(!d.file_name.contains(['/', '\\']) && !d.file_name.starts_with('.')); + } + // a model that may not be used commercially is never offered for download + if k.manifest.licence.commercial == Commercial::No { + assert!(k.download.is_none(), "{}", k.manifest.id); + } + } + assert_eq!(find("rawnind-bayer").is_some(), cfg!(feature = "rawnind-model")); + assert!(find("nothing").is_none()); + assert_eq!( + by_sha256("da27509dab6a2915da67e988acd86cf71f9d5bbc8d1aa0ed32933578a887b901").map(|m| m.id), + if cfg!(feature = "rawnind-model") { Some("rawnind-bayer".into()) } else { None } + ); + } +} + +#[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, +} + +pub fn host(url: &str) -> &str { + url.strip_prefix("https://").and_then(|r| r.split('/').next()).unwrap_or("") +} diff --git a/crates/denoise-core/src/lib.rs b/crates/denoise-core/src/lib.rs new file mode 100644 index 00000000..568ed9f8 --- /dev/null +++ b/crates/denoise-core/src/lib.rs @@ -0,0 +1,16 @@ +//! Model metadata, bounded archive/cache formats and Bayer tiling, independent of inference. +//! Used by featureless engine and web builds; no neural-network or matrix backend is linked. +#![forbid(unsafe_code)] +#![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)] +pub mod archive; +pub mod bayer; +pub mod hash; +pub mod known; +pub mod licence; +pub mod manifest; +pub mod product; +pub mod run; +pub mod tiles; + +pub use manifest::{DenoiserManifest, Domain, Gain, ManifestError}; +pub use run::{Control, Error, Params, TileRunner, denoise_bayer}; diff --git a/crates/denoise-core/src/licence.rs b/crates/denoise-core/src/licence.rs new file mode 100644 index 00000000..1d89146b --- /dev/null +++ b/crates/denoise-core/src/licence.rs @@ -0,0 +1,34 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum Commercial { + Yes, + /// Research or personal use only. Never bundled or redistributed by LightCraft. + No, + /// The licence or the training data is unclear: the user is told so before installing. + Unknown, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub struct Licence { + /// SPDX id or a short name ("MIT", "Apache-2.0", "InsightFace non-commercial research"). + pub name: String, + pub commercial: Commercial, + pub url: Option, + /// Plain-language terms shown before the model is enabled. + pub notice: String, +} + +impl Default for Licence { + fn default() -> Self { + Self { name: String::new(), commercial: Commercial::Unknown, url: None, notice: String::new() } + } +} + +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'-')) +} diff --git a/crates/denoise-core/src/manifest.rs b/crates/denoise-core/src/manifest.rs new file mode 100644 index 00000000..9fbc9f49 --- /dev/null +++ b/crates/denoise-core/src/manifest.rs @@ -0,0 +1,217 @@ +//! A denoise model's manifest: what the model takes and gives, how big its tiles are and who may use it, as +//! plain data. +//! +//! It comes from a `denoise-model.json` next to the `.onnx` file or from the user's `catalog.json`. Both are +//! untrusted: [`validate`] must pass before a manifest is used, and [`parse`] applies it after reading JSON of +//! bounded size. The licence, id and address rules are the face models' (`crate::licence`). + +use crate::licence::{Licence, valid_id}; +use serde::{Deserialize, Serialize}; + +/// Largest `denoise-model.json` we read. +pub const MAX_MANIFEST_BYTES: usize = 64 * 1024; +/// Largest model file we accept (bytes). +pub const MAX_MODEL_BYTES: u64 = 2 * 1024 * 1024 * 1024; + +/// What the model takes and what it gives. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum Domain { + /// The sensor's Bayer mosaic (black subtracted, white at 1, not white balanced) packed as four planes + /// `[R, G1, G2, B]` of its 2 × 2 cells, forced to RGGB; gives camera RGB at the mosaic's resolution, so + /// the picture is denoised and demosaiced in one step and must not be demosaiced again. + BayerToRgb, +} + +/// How the model's output scale is brought back to the input's. +#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "camelCase")] +pub enum Gain { + /// The output already has the input's scale. + None, + /// The output has a scale of its own, `nominal` times the input's (a million, for one model) and a little + /// different from tile to tile: each tile is scaled so that its channel means equal the input's. A tile too + /// dark to measure, or whose measured scale is further than `maxDeviation` (a fraction, 0.05 = 5 %) from + /// `nominal`, is scaled as if it were within that limit. + MatchMean { nominal: f32, max_deviation: f32 }, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct DenoiserManifest { + /// Lowercase letters, digits, `.` `_` `-`. + pub id: String, + pub name: String, + pub version: String, + #[serde(default)] + pub licence: Licence, + /// Where the weights come from (an `https://` page or file), for the user's information. + #[serde(default)] + pub source: Option, + /// Lowercase hex SHA-256 of the `.onnx` file, when known. + #[serde(default)] + pub sha256: Option, + #[serde(default)] + pub size_bytes: Option, + /// What the model was trained on, in plain words; says "undisclosed" when it is. + #[serde(default)] + pub provenance: String, + pub domain: Domain, + /// Side of the square the model runs on, in packed cells (the file's fixed input size). + pub tile: u32, + /// Cells neighbouring tiles share; the picture is blended across them. + pub overlap: u32, + pub gain: Gain, +} + +#[derive(Debug, PartialEq, thiserror::Error)] +pub enum ManifestError { + #[error("`{0}` is not valid: {1}")] + Field(&'static str, String), + #[error("not a denoise-model manifest: {0}")] + Parse(String), + #[error("a manifest may be at most {MAX_MANIFEST_BYTES} bytes")] + TooLarge, +} + +fn bad(field: &'static str, why: impl Into) -> Result { + Err(ManifestError::Field(field, why.into())) +} + +fn text(field: &'static str, s: &str, min: usize, max: usize) -> Result<(), ManifestError> { + let n = s.chars().count(); + if n < min || n > max { + return bad(field, format!("must be {min} to {max} characters")); + } + if s.chars().any(|c| c.is_control() && c != '\n') { + return bad(field, "must not contain control characters"); + } + Ok(()) +} + +fn url(field: &'static str, u: &Option) -> Result<(), ManifestError> { + match u { + None => Ok(()), + Some(u) if u.len() <= 500 && u.starts_with("https://") && !u.chars().any(|c| c.is_control() || c.is_whitespace()) => Ok(()), + Some(_) => bad(field, "must be an https:// address of at most 500 characters"), + } +} + +/// Check a manifest from outside before trusting it. +pub fn validate(m: &DenoiserManifest) -> Result<(), ManifestError> { + if !valid_id(&m.id) { + return bad("id", "1 to 64 lowercase letters, digits, `.`, `_` or `-`, starting with a letter or digit"); + } + text("name", &m.name, 1, 80)?; + text("version", &m.version, 1, 40)?; + text("provenance", &m.provenance, 0, 1000)?; + text("licence.name", &m.licence.name, 0, 80)?; + text("licence.notice", &m.licence.notice, 0, 2000)?; + 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 hex digits"); + } + if let Some(n) = m.size_bytes + && (n == 0 || n > MAX_MODEL_BYTES) + { + return bad("sizeBytes", format!("must be between 1 and {MAX_MODEL_BYTES}")); + } + // a U-Net with four poolings wants a multiple of 16; nothing smaller than 64 or larger than 2048 is a tile + if !(64..=2048).contains(&m.tile) || !m.tile.is_multiple_of(16) { + return bad("tile", "must be a multiple of 16 from 64 to 2048"); + } + // tiles must overlap by less than half, or one tile's blend ramps would run into its other neighbour's + if !m.overlap.is_multiple_of(2) || m.overlap < 16 || m.overlap.saturating_mul(2) >= m.tile { + return bad("overlap", "must be even, at least 16 and less than half the tile"); + } + if let Gain::MatchMean { nominal, max_deviation } = m.gain { + if !(nominal.is_finite() && (1e-6..=1e12).contains(&nominal)) { + return bad("gain.nominal", "must be between 0.000001 and 1000000000000"); + } + if !(max_deviation.is_finite() && (0.0..=0.5).contains(&max_deviation)) { + return bad("gain.maxDeviation", "must be between 0 and 0.5"); + } + } + Ok(()) +} + +/// Read and validate a `denoise-model.json`. +pub fn parse(bytes: &[u8]) -> Result { + if bytes.len() > MAX_MANIFEST_BYTES { + return Err(ManifestError::TooLarge); + } + let bytes = bytes.strip_prefix(&[0xEF, 0xBB, 0xBF]).unwrap_or(bytes); + let m: DenoiserManifest = serde_json::from_slice(bytes).map_err(|e| ManifestError::Parse(e.to_string()))?; + validate(&m)?; + Ok(m) +} + +#[cfg(test)] +mod tests { + use super::*; + + pub(crate) fn sample() -> DenoiserManifest { + DenoiserManifest { + id: "nind-bayer".into(), + name: "Bayer denoiser".into(), + version: "1".into(), + licence: Licence::default(), + source: Some("https://example.org/model.onnx".into()), + sha256: None, + size_bytes: Some(31_000_000), + provenance: "trained on real raw noise pairs".into(), + domain: Domain::BayerToRgb, + tile: 512, + overlap: 64, + gain: Gain::MatchMean { nominal: 1.0e6, max_deviation: 0.05 }, + } + } + + #[test] + fn a_good_manifest_round_trips() { + let m = sample(); + let json = serde_json::to_vec(&m).unwrap(); + assert_eq!(parse(&json).unwrap(), m); + } + + #[test] + fn bad_fields_are_named() { + let cases: Vec<(&str, Box)> = vec![ + ("id", Box::new(|m| m.id = "../evil".into())), + ("id", Box::new(|m| m.id = String::new())), + ("name", Box::new(|m| m.name = "x\u{7}".into())), + ("source", Box::new(|m| m.source = Some("http://insecure".into()))), + ("sha256", Box::new(|m| m.sha256 = Some("XYZ".into()))), + ("sizeBytes", Box::new(|m| m.size_bytes = Some(0))), + ("tile", Box::new(|m| m.tile = 500)), + ("tile", Box::new(|m| m.tile = 4096)), + ("overlap", Box::new(|m| m.overlap = 300)), + ("overlap", Box::new(|m| m.overlap = 63)), + ("gain.maxDeviation", Box::new(|m| m.gain = Gain::MatchMean { nominal: 1.0, max_deviation: f32::NAN })), + ("gain.maxDeviation", Box::new(|m| m.gain = Gain::MatchMean { nominal: 1.0, max_deviation: 2.0 })), + ("gain.nominal", Box::new(|m| m.gain = Gain::MatchMean { nominal: 0.0, max_deviation: 0.1 })), + ("gain.nominal", Box::new(|m| m.gain = Gain::MatchMean { nominal: f32::INFINITY, max_deviation: 0.1 })), + ]; + for (field, edit) in cases { + let mut m = sample(); + edit(&mut m); + match validate(&m) { + Err(ManifestError::Field(f, _)) => assert_eq!(f, field), + other => panic!("{field}: {other:?}"), + } + } + } + + #[test] + fn parse_is_bounded_and_tolerant() { + assert_eq!(parse(&vec![b' '; MAX_MANIFEST_BYTES + 1]), Err(ManifestError::TooLarge)); + assert!(matches!(parse(b"{ nope"), Err(ManifestError::Parse(_)))); + assert!(matches!(parse(b""), Err(ManifestError::Parse(_)))); + let mut with_bom = vec![0xEF, 0xBB, 0xBF]; + with_bom.extend(serde_json::to_vec(&sample()).unwrap()); + assert!(parse(&with_bom).is_ok()); + } +} diff --git a/crates/denoise-core/src/product.rs b/crates/denoise-core/src/product.rs new file mode 100644 index 00000000..5a607c8f --- /dev/null +++ b/crates/denoise-core/src/product.rs @@ -0,0 +1,616 @@ +//! The denoised picture as cached data, and the blend the Amount slider controls. +//! +//! A product is camera RGB at the sensor's resolution (what the model gave, before the raw's own opcodes and +//! crop), kept as half floats in strips of rows, each strip byte-shuffled and deflated with its own checksum. +//! It is derived data: any problem reading it (missing, truncated, corrupt, made for another photo or model) is +//! just a cache miss, never a panic, and the picture can be made again. +//! +//! File: `LCDN`, version, header length, header JSON (`width`, `height`, `strip`, `key`), then for each strip its +//! compressed length and CRC-32, then the strips' bytes. `key` says what the product was made from (the photo's +//! content, the model, this code's version, the settings it depends on); a file with another key is stale. + +use std::io::{Read, Write}; +use std::path::Path; + +use lightcraft_raster::Rgb32f; +use rayon::prelude::*; +use serde::{Deserialize, Serialize}; + +const MAGIC: &[u8; 4] = b"LCDN"; +const VERSION: u8 = 1; +const MAX_HEADER: usize = 4096; +const MAX_SIDE: usize = 1 << 17; +/// Most pixels a product may have (the denoiser's own limit, [`crate::run::MAX_PIXELS`]). +const MAX_PIXELS: usize = crate::run::MAX_PIXELS; +/// Rows per strip: about a megapixel of three planes at 24 MP widths. +const STRIP_ROWS: usize = 128; +/// Bumped when the meaning of a product changes (included in keys by the caller). +pub const ALGORITHM: u32 = 1; + +#[derive(Debug, thiserror::Error)] +pub enum ProductError { + #[error("the denoised picture could not be read or written: {0}")] + Io(String), + #[error("not a usable denoised picture: {0}")] + Format(String), + /// A product made from something else (another photo state, model or version). + #[error("the denoised picture is out of date")] + Stale, +} + +impl From for ProductError { + fn from(e: std::io::Error) -> Self { + ProductError::Io(e.to_string()) + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct Header { + pub width: usize, + pub height: usize, + pub strip: usize, + pub key: String, +} + +fn allocate(n: usize, value: T) -> Result, ProductError> { + let mut v = Vec::new(); + v.try_reserve_exact(n).map_err(|_| ProductError::Format("not enough memory for the denoised picture".into()))?; + v.resize(n, value); + Ok(v) +} + +/// `f32` as IEEE half-float bits, rounded to nearest even; beyond the range it saturates, NaN becomes 0. +pub fn to_f16(v: f32) -> u16 { + if !v.is_finite() { + return if v.is_nan() { + 0 + } else if v > 0.0 { + 0x7bff + } else { + 0xfbff + }; + } + let b = v.to_bits(); + let sign = ((b >> 16) & 0x8000) as u16; + let exp = ((b >> 23) & 0xff) as i32 - 127 + 15; + let man = b & 0x7f_ffff; + if exp >= 31 { + return sign | 0x7bff; + } + if exp <= 0 { + if exp < -10 { + return sign; + } + let man = man | 0x80_0000; + let shift = (14 - exp) as u32; + let (half, rem, halfway) = (man >> shift, man & ((1u32 << shift) - 1), 1u32 << (shift - 1)); + let mut h = half as u16; + if rem > halfway || (rem == halfway && h & 1 == 1) { + h += 1; + } + return sign | h; + } + let mut h = (((exp as u32) << 10) | (man >> 13)) as u16; + let rem = man & 0x1fff; + if rem > 0x1000 || (rem == 0x1000 && h & 1 == 1) { + h += 1; + } + if h >= 0x7c00 { + h = 0x7bff; + } + sign | h +} + +/// Half-float bits as `f32`. +pub fn from_f16(h: u16) -> f32 { + let sign = if h & 0x8000 != 0 { -1.0f32 } else { 1.0 }; + let exp = i32::from((h >> 10) & 0x1f); + let man = f32::from(h & 0x3ff); + match exp { + 0 => sign * man * 2f32.powi(-24), + 31 => sign * 65504.0, + _ => sign * (1.0 + man / 1024.0) * 2f32.powi(exp - 15), + } +} + +fn crc(bytes: &[u8]) -> u32 { + // CRC-32 (IEEE), bitwise: strips are small and this runs once per strip + let mut table = [0u32; 256]; + for (i, t) in table.iter_mut().enumerate() { + let mut c = i as u32; + for _ in 0..8 { + c = if c & 1 != 0 { 0xedb8_8320 ^ (c >> 1) } else { c >> 1 }; + } + *t = c; + } + !bytes.iter().fold(!0u32, |c, &b| table[((c ^ u32::from(b)) & 0xff) as usize] ^ (c >> 8)) +} + +/// One strip as bytes: for each plane (R, G, B) the low bytes of its half floats, then the high bytes. +fn encode_strip(img: &Rgb32f, y0: usize, y1: usize) -> Result, ProductError> { + let n = (y1 - y0) * img.width; + let mut raw = allocate(6 * n, 0u8)?; + let rows = img.data.get(y0 * img.width..y1 * img.width).unwrap_or(&[]); + for (i, px) in rows.iter().enumerate() { + for (c, &v) in px.iter().enumerate() { + let h = to_f16(v); + if let Some(lo) = raw.get_mut(2 * c * n + i) { + *lo = h as u8; + } + if let Some(hi) = raw.get_mut((2 * c + 1) * n + i) { + *hi = (h >> 8) as u8; + } + } + } + Ok(miniz_oxide::deflate::compress_to_vec(&raw, 3)) +} + +/// Write `img` to `path` (through a temporary file next to it, so a reader never sees half a product). +pub fn write(path: &Path, img: &Rgb32f, key: &str) -> Result<(), ProductError> { + let pixels = img.width.checked_mul(img.height).filter(|&n| n <= MAX_PIXELS); + if img.width == 0 || img.height == 0 || img.width > MAX_SIDE || img.height > MAX_SIDE || pixels != Some(img.data.len()) { + return Err(ProductError::Format("the picture has no usable size".into())); + } + let strips: Vec<(usize, usize)> = (0..img.height).step_by(STRIP_ROWS).map(|y| (y, (y + STRIP_ROWS).min(img.height))).collect(); + let packed: Vec> = strips.par_iter().map(|&(a, b)| encode_strip(img, a, b)).collect::>()?; + let header = serde_json::to_vec(&Header { width: img.width, height: img.height, strip: STRIP_ROWS, key: key.to_string() }) + .map_err(|e| ProductError::Format(e.to_string()))?; + let tmp = path.with_extension("tmp"); + let result = (|| -> Result<(), ProductError> { + if let Some(dir) = path.parent() { + std::fs::create_dir_all(dir)?; + } + let mut f = std::io::BufWriter::new(std::fs::File::create(&tmp)?); + f.write_all(MAGIC)?; + f.write_all(&[VERSION, 0, 0, 0])?; + f.write_all(&(header.len() as u32).to_le_bytes())?; + f.write_all(&header)?; + for p in &packed { + f.write_all(&(p.len() as u32).to_le_bytes())?; + f.write_all(&crc(p).to_le_bytes())?; + } + for p in &packed { + f.write_all(p)?; + } + f.flush()?; + drop(f); + std::fs::rename(&tmp, path)?; + Ok(()) + })(); + if result.is_err() { + let _ = std::fs::remove_file(&tmp); + } + result +} + +/// The parsed start of a product file: header, each strip's (compressed length, CRC-32), and the file length the +/// start says the whole file has. +struct Start { + header: Header, + table: Vec<(usize, u32)>, + expected_len: u64, +} + +fn read_start(r: &mut impl Read) -> Result { + let mut fixed = [0u8; 12]; + r.read_exact(&mut fixed)?; + if fixed.get(0..4) != Some(&MAGIC[..]) || fixed.get(4) != Some(&VERSION) { + return Err(ProductError::Format("not a denoised picture of this version".into())); + } + let len = fixed.get(8..12).and_then(|b| <[u8; 4]>::try_from(b).ok()).map(u32::from_le_bytes).unwrap_or(u32::MAX) as usize; + if len > MAX_HEADER { + return Err(ProductError::Format("the header is too large".into())); + } + let mut hb = allocate(len, 0u8)?; + r.read_exact(&mut hb)?; + let h: Header = serde_json::from_slice(&hb).map_err(|e| ProductError::Format(e.to_string()))?; + let pixels = h.width.saturating_mul(h.height); + if h.width == 0 || h.height == 0 || h.width > MAX_SIDE || h.height > MAX_SIDE || pixels > MAX_PIXELS || h.strip != STRIP_ROWS { + return Err(ProductError::Format("the size is not usable".into())); + } + let n = h.height.div_ceil(h.strip); + let mut table = Vec::with_capacity(n); + let mut total = 12 + len as u64 + 8 * n as u64; + for _ in 0..n { + let mut e = [0u8; 8]; + r.read_exact(&mut e)?; + let l = u32::from_le_bytes([e[0], e[1], e[2], e[3]]) as usize; + let c = u32::from_le_bytes([e[4], e[5], e[6], e[7]]); + // a strip never compresses to much more than its raw size + if l > 12 * h.strip * h.width + 4096 { + return Err(ProductError::Format("a strip is larger than it can be".into())); + } + total += l as u64; + table.push((l, c)); + } + Ok(Start { header: h, table, expected_len: total }) +} + +/// The header of the product at `path`, when it is one. +pub fn read_header_at(path: &Path) -> Result { + let mut f = std::io::BufReader::new(std::fs::File::open(path)?); + Ok(read_start(&mut f)?.header) +} + +/// Whether `path` holds a complete-looking product made from `key` (the header and the file's length are checked; +/// strips are checked when read). +pub fn is_current(path: &Path, key: &str) -> bool { + let check = || -> Result { + let mut f = std::io::BufReader::new(std::fs::File::open(path)?); + let s = read_start(&mut f)?; + Ok(s.header.key == key && std::fs::metadata(path)?.len() == s.expected_len) + }; + check().unwrap_or(false) +} + +/// A rectangle of a product, in the product's pixels. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Window { + pub x: usize, + pub y: usize, + pub width: usize, + pub height: usize, +} + +/// Half floats as `f32`, for all 65536 bit patterns (a table is much faster than computing each). +fn f16_table() -> &'static [f32] { + static TABLE: std::sync::OnceLock> = std::sync::OnceLock::new(); + TABLE.get_or_init(|| (0..=u16::MAX).map(from_f16).collect()) +} + +/// Strips decoded together (they are inflated in parallel, then folded in order). +const GROUP: usize = 8; + +/// One strip, checked and inflated: for each plane the low bytes of its half floats, then the high bytes. +fn decode_strip(index: usize, packed: &[u8], expect_crc: u32, n: usize) -> Result, ProductError> { + if crc(packed) != expect_crc { + return Err(ProductError::Format(format!("strip {index} is damaged"))); + } + let raw = miniz_oxide::inflate::decompress_to_vec_with_limit(packed, 6 * n).map_err(|e| ProductError::Format(format!("strip {index}: {e:?}")))?; + if raw.len() != 6 * n { + return Err(ProductError::Format(format!("strip {index} has the wrong size"))); + } + Ok(raw) +} + +/// Read the product at `path`, made from `key`, with every `factor × factor` block of pixels averaged into one +/// (`factor` 1 = full size). The result is camera RGB. +/// +/// With a `window`, only that rectangle is read and the blocks start at its corner; only whole blocks are made, so +/// the result is `window.width / factor × window.height / factor` (the same blocks a binned development of the same +/// crop makes). Without one the whole product is read and blocks at the right and bottom edges may be smaller. +/// +/// With `clip` (white is 1), a block in which any sample of a colour is at or above it takes that colour's +/// maximum instead of the mean, so a clipped highlight stays clipped for highlight reconstruction. +pub fn read_window(path: &Path, key: &str, window: Option, factor: usize, clip: Option) -> Result { + let factor = factor.clamp(1, 1 << 12); + let mut f = std::io::BufReader::new(std::fs::File::open(path)?); + let Start { header: h, table, expected_len } = read_start(&mut f)?; + if h.key != key { + return Err(ProductError::Stale); + } + if std::fs::metadata(path)?.len() != expected_len { + return Err(ProductError::Format("the file is not the length its header says".into())); + } + let w = window.unwrap_or(Window { x: 0, y: 0, width: h.width, height: h.height }); + let fits = |a: usize, b: usize, limit: usize| a.checked_add(b).is_some_and(|e| e <= limit); + if w.width == 0 || w.height == 0 || !fits(w.x, w.width, h.width) || !fits(w.y, w.height, h.height) { + return Err(ProductError::Format("the window does not lie inside the picture".into())); + } + // whole blocks for a window; partial ones at the edges when the whole picture is read + let whole = window.is_none(); + let (ow, oh) = if whole { (w.width.div_ceil(factor), w.height.div_ceil(factor)) } else { (w.width / factor, w.height / factor) }; + if ow == 0 || oh == 0 { + return Err(ProductError::Format("the window is smaller than one block".into())); + } + // the picture rows and columns the blocks cover + let (x_end, y_end) = if whole { (w.x + w.width, w.y + w.height) } else { (w.x + ow * factor, w.y + oh * factor) }; + let mut out = Rgb32f { width: ow, height: oh, data: allocate(ow * oh, [0.0; 3])? }; + let lut = f16_table(); + let value = |raw: &[u8], c: usize, n: usize, i: usize| -> f32 { + let (lo, hi) = (raw.get(2 * c * n + i).copied().unwrap_or(0), raw.get((2 * c + 1) * n + i).copied().unwrap_or(0)); + lut.get(usize::from(lo) | usize::from(hi) << 8).copied().unwrap_or(0.0) + }; + // which strips are needed, and where each starts in the file (after the header and the strip table) + let wanted = |s: usize| s * h.strip < y_end && (s + 1) * h.strip > w.y; + let (mut sum, mut peak) = (allocate(ow, [0f32; 3])?, allocate(ow, [f32::MIN; 3])?); + let (mut rows_in, mut oy) = (0usize, 0usize); + let strips: Vec = (0..table.len()).filter(|&s| wanted(s)).collect(); + // skip the strips before the first one wanted + let skipped: usize = table.iter().take(strips.first().copied().unwrap_or(0)).map(|e| e.0).sum(); + f.seek_relative(i64::try_from(skipped).map_err(|_| ProductError::Format("the file is too large".into()))?)?; + for group in strips.chunks(GROUP) { + let mut packed = Vec::with_capacity(group.len()); + for &s in group { + let Some(&(len, _)) = table.get(s) else { return Err(ProductError::Format("a strip is missing".into())) }; + let mut p = allocate(len, 0u8)?; + f.read_exact(&mut p)?; + packed.push(p); + } + let decoded: Vec, ProductError>> = group + .par_iter() + .zip(packed.par_iter()) + .map(|(&s, p)| { + let rows = h.height.saturating_sub(s * h.strip).min(h.strip); + decode_strip(s, p, table.get(s).map_or(0, |e| e.1), rows * h.width) + }) + .collect(); + for (&s, raw) in group.iter().zip(decoded) { + let raw = raw?; + let y0 = s * h.strip; + let rows = h.height.saturating_sub(y0).min(h.strip); + let n = rows * h.width; + for r in 0..rows { + let y = y0 + r; + if y < w.y || y >= y_end { + continue; + } + let base = r * h.width; + for (ox, (sm, pk)) in sum.iter_mut().zip(peak.iter_mut()).enumerate() { + let x0 = w.x + ox * factor; + let x1 = (x0 + factor).min(x_end); + for x in x0..x1 { + for c in 0..3 { + let v = value(&raw, c, n, base + x); + if let (Some(s_), Some(p_)) = (sm.get_mut(c), pk.get_mut(c)) { + *s_ += v; + *p_ = p_.max(v); + } + } + } + } + rows_in += 1; + if rows_in == factor || y + 1 == y_end { + for (ox, (sm, pk)) in sum.iter_mut().zip(peak.iter_mut()).enumerate() { + // the block's size at the right and bottom edges may be smaller + let bw = (x_end - (w.x + ox * factor)).min(factor); + let k = 1.0 / (rows_in * bw).max(1) as f32; + if let Some(o) = out.data.get_mut(oy * ow + ox) { + for c in 0..3 { + let (s_, p_) = (sm.get(c).copied().unwrap_or(0.0), pk.get(c).copied().unwrap_or(0.0)); + if let Some(slot) = o.get_mut(c) { + *slot = if factor > 1 && clip.is_some_and(|t| p_ >= t) { p_ } else { s_ * k }; + } + } + } + *sm = [0.0; 3]; + *pk = [f32::MIN; 3]; + } + rows_in = 0; + oy += 1; + } + } + } + } + Ok(out) +} + +/// [`read_window`] of the whole product: every `factor × factor` block averaged into one. +pub fn read(path: &Path, key: &str, factor: usize) -> Result { + read_window(path, key, None, factor, None) +} + +/// `base + (denoised − base) · amount` for every pixel (`amount` 0 to 1), in place in `base`. The two pictures must +/// be the same size; if they are not, `base` is left as it was and `false` is returned. +pub fn blend(base: &mut Rgb32f, denoised: &Rgb32f, amount: f32) -> bool { + if base.width != denoised.width || base.height != denoised.height { + return false; + } + let a = if amount.is_finite() { amount.clamp(0.0, 1.0) } else { 0.0 }; + base.data.par_iter_mut().zip(denoised.data.par_iter()).for_each(|(b, d)| { + for c in 0..3 { + b[c] += (d[c] - b[c]) * a; + } + }); + true +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn oversized_product_header_is_rejected_before_allocating_pixels() { + let p = temp("oversized-header"); + let header = serde_json::to_vec(&Header { width: 10001, height: 10000, strip: STRIP_ROWS, key: "k".into() }).unwrap(); + let mut bytes = Vec::from(MAGIC.as_slice()); + bytes.extend_from_slice(&[VERSION, 0, 0, 0]); + bytes.extend_from_slice(&(header.len() as u32).to_le_bytes()); + bytes.extend_from_slice(&header); + std::fs::write(&p, bytes).unwrap(); + assert!(read_header_at(&p).is_err()); + let _ = std::fs::remove_file(p); + } + + fn temp(name: &str) -> std::path::PathBuf { + std::env::temp_dir().join(format!("lc-denoise-{name}-{}.lcdn", std::process::id())) + } + + fn picture(w: usize, h: usize) -> Rgb32f { + Rgb32f { + width: w, + height: h, + data: (0..w * h).map(|i| [(i % w) as f32 / w as f32, (i / w) as f32 / h as f32, 0.001 + 0.5 * ((i % 7) as f32 / 7.0)]).collect(), + } + } + + #[test] + fn half_floats_round_trip_closely() { + for v in [0.0f32, 1.0, -1.0, 0.5, 0.333_333, 1e-4, 6.1e-5, 3e-6, 123.456, 65000.0, -0.002] { + let back = from_f16(to_f16(v)); + assert!((back - v).abs() <= v.abs() * 0.001 + 1e-7, "{v} -> {back}"); + } + assert_eq!(from_f16(to_f16(1e10)), 65504.0); + assert_eq!(to_f16(f32::NAN), 0); + assert_eq!(from_f16(to_f16(f32::NEG_INFINITY)), -65504.0); + } + + #[test] + fn a_product_survives_a_round_trip_at_full_and_reduced_size() { + let (w, h) = (301, 517); + let img = picture(w, h); + let p = temp("round"); + write(&p, &img, "k1").unwrap(); + assert!(is_current(&p, "k1") && !is_current(&p, "k2")); + assert_eq!(read_header_at(&p).unwrap().width, w); + let full = read(&p, "k1", 1).unwrap(); + assert_eq!((full.width, full.height), (w, h)); + for (a, b) in full.data.iter().zip(&img.data) { + for c in 0..3 { + assert!((a[c] - b[c]).abs() < 0.002, "{a:?} vs {b:?}"); + } + } + // averaged into blocks of 4: the sizes round up, and the mean of the picture is kept + let small = read(&p, "k1", 4).unwrap(); + assert_eq!((small.width, small.height), (w.div_ceil(4), h.div_ceil(4))); + let mean = |i: &Rgb32f| i.data.iter().map(|p| f64::from(p[1])).sum::() / i.data.len() as f64; + assert!((mean(&small) - mean(&img)).abs() < 0.01); + // a factor that does not divide the strip height is the same average + let odd = read(&p, "k1", 5).unwrap(); + assert_eq!((odd.width, odd.height), (w.div_ceil(5), h.div_ceil(5))); + assert!((mean(&odd) - mean(&img)).abs() < 0.01); + assert!(matches!(read(&p, "other", 1), Err(ProductError::Stale))); + std::fs::remove_file(&p).unwrap(); + } + + /// What binning a window of `img` by `k` means, written the slow obvious way. + fn blocks(img: &Rgb32f, win: Window, k: usize, clip: Option) -> Rgb32f { + let (ow, oh) = (win.width / k, win.height / k); + let mut out = Rgb32f { width: ow, height: oh, data: vec![[0.0; 3]; ow * oh] }; + for by in 0..oh { + for bx in 0..ow { + for c in 0..3 { + let vals: Vec = (0..k * k).map(|i| img.data[(win.y + by * k + i / k) * img.width + win.x + bx * k + i % k][c]).collect(); + let max = vals.iter().copied().fold(f32::MIN, f32::max); + out.data[by * ow + bx][c] = + if k > 1 && clip.is_some_and(|t| max >= t) { max } else { vals.iter().sum::() / vals.len() as f32 }; + } + } + } + out + } + + #[test] + fn a_window_is_read_as_the_blocks_of_that_crop() { + let (w, h) = (301, 517); + let mut img = picture(w, h); + // a clipped sample in the middle of a block + img.data[(130 + 2) * w + 61 + 1] = [1.0, 0.4, 0.4]; + let p = temp("window"); + write(&p, &img, "k1").unwrap(); + let near = |a: &Rgb32f, b: &Rgb32f| { + assert_eq!((a.width, a.height), (b.width, b.height)); + for (x, y) in a.data.iter().zip(&b.data) { + for c in 0..3 { + assert!((x[c] - y[c]).abs() < 0.003, "{x:?} vs {y:?}"); + } + } + }; + // windows at odd places, across strip boundaries (128 rows), with factors that do and do not divide them + for (win, k) in [ + (Window { x: 3, y: 5, width: 200, height: 300 }, 1), + (Window { x: 3, y: 5, width: 200, height: 300 }, 4), + (Window { x: 61, y: 130, width: 121, height: 250 }, 5), + (Window { x: 0, y: 0, width: w, height: h }, 2), + (Window { x: 1, y: 380, width: 90, height: 137 }, 3), + ] { + let got = read_window(&p, "k1", Some(win), k, None).unwrap(); + near(&got, &blocks(&img, win, k, None)); + } + // with a clip level the block with the clipped sample keeps its highest value, the others stay means + let win = Window { x: 61, y: 130, width: 40, height: 20 }; + let clipped = read_window(&p, "k1", Some(win), 4, Some(0.99)).unwrap(); + near(&clipped, &blocks(&img, win, 4, Some(0.99))); + assert!(clipped.data[0][0] > 0.99, "the clipped block stays clipped: {:?}", clipped.data[0]); + let plain = read_window(&p, "k1", Some(win), 4, None).unwrap(); + assert!(plain.data[0][0] < 0.99, "an average falls below the clip level"); + // windows that do not fit, or are smaller than a block, are errors + for bad in [ + Window { x: 0, y: 0, width: 0, height: 5 }, + Window { x: 290, y: 0, width: 20, height: 5 }, + Window { x: 0, y: 510, width: 5, height: 10 }, + Window { x: usize::MAX, y: 0, width: 2, height: 2 }, + Window { x: 0, y: 0, width: 3, height: 3 }, + ] { + assert!(read_window(&p, "k1", Some(bad), 4, None).is_err(), "{bad:?}"); + } + std::fs::remove_file(&p).unwrap(); + } + + #[test] + fn a_smooth_picture_compresses_well() { + let img = Rgb32f { + width: 512, + height: 512, + data: (0..512 * 512) + .map(|i| { + let v = 0.2 + 0.0001 * (i % 512) as f32; + [v, v * 0.8, v * 0.6] + }) + .collect(), + }; + let p = temp("size"); + write(&p, &img, "k").unwrap(); + let bytes = std::fs::metadata(&p).unwrap().len() as usize; + assert!(bytes < 512 * 512 * 6 / 2, "{bytes} bytes for a smooth picture"); + std::fs::remove_file(&p).unwrap(); + } + + #[test] + fn damaged_files_are_errors_not_panics() { + let img = picture(100, 300); + let p = temp("damage"); + write(&p, &img, "k").unwrap(); + let good = std::fs::read(&p).unwrap(); + // truncated at every length + for cut in (0..good.len()).step_by(37) { + std::fs::write(&p, &good[..cut]).unwrap(); + assert!(read(&p, "k", 1).is_err(), "cut at {cut}"); + assert!(!is_current(&p, "k")); + } + // a flipped byte anywhere is caught by a checksum, the header parse, or the key + for at in (0..good.len()).step_by(53) { + let mut bad = good.clone(); + bad[at] ^= 0xff; + std::fs::write(&p, &bad).unwrap(); + let _ = read(&p, "k", 1); // must not panic; most are errors + } + // a header that claims an absurd size + let mut huge = b"LCDN\x01\0\0\0".to_vec(); + let hdr = br#"{"width":4000000000,"height":4000000000,"strip":128,"key":"k"}"#; + huge.extend((hdr.len() as u32).to_le_bytes()); + huge.extend(hdr); + std::fs::write(&p, &huge).unwrap(); + assert!(matches!(read(&p, "k", 1), Err(ProductError::Format(_)))); + std::fs::write(&p, b"").unwrap(); + assert!(read(&p, "k", 1).is_err()); + assert!(read(&temp("missing"), "k", 1).is_err()); + std::fs::remove_file(&p).unwrap(); + } + + #[test] + fn writing_rejects_unusable_pictures() { + let p = temp("bad"); + assert!(write(&p, &Rgb32f { width: 0, height: 0, data: vec![] }, "k").is_err()); + assert!(write(&p, &Rgb32f { width: 3, height: 3, data: vec![[0.0; 3]; 8] }, "k").is_err()); + assert!(!p.exists()); + } + + #[test] + fn blend_mixes_by_amount_and_rejects_other_sizes() { + let mut base = Rgb32f { width: 2, height: 1, data: vec![[0.0; 3], [1.0; 3]] }; + let dn = Rgb32f { width: 2, height: 1, data: vec![[1.0; 3], [0.0; 3]] }; + assert!(blend(&mut base, &dn, 0.25)); + assert_eq!(base.data, vec![[0.25; 3], [0.75; 3]]); + assert!(blend(&mut base, &dn, f32::NAN)); + assert_eq!(base.data, vec![[0.25; 3], [0.75; 3]]); + assert!(!blend(&mut base, &Rgb32f { width: 1, height: 1, data: vec![[0.0; 3]] }, 1.0)); + assert!(blend(&mut base, &dn, 1.0)); + assert_eq!(base.data, dn.data); + } + + #[test] + fn crc_matches_the_standard_check_value() { + assert_eq!(crc(b"123456789"), 0xcbf4_3926); + } +} diff --git a/crates/denoise-core/src/run.rs b/crates/denoise-core/src/run.rs new file mode 100644 index 00000000..1df96a9e --- /dev/null +++ b/crates/denoise-core/src/run.rs @@ -0,0 +1,589 @@ +//! Denoising a whole Bayer mosaic with a tile model: cut it into overlapping tiles, run the model on each, bring +//! each tile's scale back to the input's, keep clipped highlights as they were, and blend the tiles together. +//! +//! The model itself is a [`TileRunner`], so everything around it (packing, tiling, scale matching, highlight +//! protection, blending, cancelling, progress) is tested without one. + +use std::sync::atomic::{AtomicBool, Ordering}; + +use lightcraft_raster::Rgb32f; +use rayon::prelude::*; + +use crate::bayer::{Layout, pack_tile}; +use crate::manifest::Gain; +use crate::tiles::{starts, weight}; + +/// Most pixels one picture may have (guards the allocations that follow from its size). +pub const MAX_PIXELS: usize = 100_000_000; +/// A tile whose mean signal is below this (white = 1) is too dark to measure its scale from. +const MIN_SIGNAL: f32 = 0.003; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("{0}")] + Input(String), + #[error("the model gave {0}")] + Model(String), + #[error("the model could not run: {0}")] + Runtime(String), + #[error("the model could not be loaded: {0}")] + Load(String), + #[error("cancelled")] + Cancelled, +} + +fn allocate(n: usize, value: T) -> Result, Error> { + let mut v = Vec::new(); + v.try_reserve_exact(n).map_err(|_| Error::Input("not enough memory for denoise".into()))?; + v.resize(n, value); + Ok(v) +} + +/// Something that runs the model on one tile. +pub trait TileRunner: Sync { + /// `input` is the tile as four planes `[R, G1, G2, B]` of `tile × tile` cells (NCHW, batch 1); the answer is + /// camera RGB, three planes of `2·tile × 2·tile` pixels (CHW), in the model's own scale. + fn run(&self, input: &[f32]) -> Result, Error>; +} + +pub struct Params { + /// Side of the model's square, in cells. + pub tile: usize, + /// Cells neighbouring tiles share. + pub overlap: usize, + pub gain: Gain, + /// White is 1; samples at or above this (minus a little) are clipped and keep their value. + pub clip: Option, + /// Tiles run at once (each model call may itself be single-threaded). + pub parallel: usize, +} + +#[derive(Default)] +pub struct Control<'a> { + pub cancel: Option<&'a AtomicBool>, + /// `(tiles done, tiles in all)`. + pub progress: Option<&'a (dyn Fn(usize, usize) + Sync)>, +} + +/// The tile grid of a picture: `(starts across, starts down)`. +pub fn grid(width: usize, height: usize, layout: Layout, tile: usize, overlap: usize) -> (Vec, Vec) { + let (nx, ny) = layout.cells(width, height); + (starts(nx, tile, overlap), starts(ny, tile, overlap)) +} + +/// How many model calls [`denoise_bayer`] makes for a picture. +pub fn tile_count(width: usize, height: usize, layout: Layout, tile: usize, overlap: usize) -> usize { + let (sx, sy) = grid(width, height, layout, tile, overlap); + sx.len().saturating_mul(sy.len()) +} + +/// Denoise (and demosaic) the normalised mosaic `mosaic` of `width × height` samples, laid out as `layout`. +/// Gives camera RGB at the mosaic's resolution. +pub fn denoise_bayer( + mosaic: &[f32], + width: usize, + height: usize, + layout: Layout, + runner: &dyn TileRunner, + p: &Params, + ctl: &Control<'_>, +) -> Result { + let pixels = width + .checked_mul(height) + .filter(|&n| n > 0 && n <= MAX_PIXELS) + .ok_or_else(|| Error::Input("the picture has no size, or is too large".into()))?; + if mosaic.len() != pixels { + return Err(Error::Input("the mosaic does not have width × height samples".into())); + } + if !(16..=2048).contains(&p.tile) || p.overlap >= p.tile.div_ceil(2) || layout.dx > 1 || layout.dy > 1 { + return Err(Error::Input("the tile size and overlap do not fit".into())); + } + if let Gain::MatchMean { nominal, max_deviation } = p.gain + && !(nominal.is_finite() && (1e-6..=1e12).contains(&nominal) && max_deviation.is_finite() && (0.0..=0.5).contains(&max_deviation)) + { + return Err(Error::Input("the model gain is not finite or within its supported range".into())); + } + if p.clip.is_some_and(|v| !v.is_finite() || v <= 0.0) { + return Err(Error::Input("the highlight clip must be finite and positive".into())); + } + let (nx, ny) = layout.cells(width, height); + let (sx, sy) = (starts(nx, p.tile, p.overlap), starts(ny, p.tile, p.overlap)); + let (cols, rows) = (sx.len(), sy.len()); + if cols == 0 || rows == 0 { + return Err(Error::Input("the picture needs too many tiles".into())); + } + let tiles: Vec<(usize, usize, bool, bool, bool, bool)> = sy + .iter() + .enumerate() + .flat_map(|(j, &y0)| sx.iter().enumerate().map(move |(i, &x0)| (x0, y0, i > 0, i + 1 < cols, j > 0, j + 1 < rows))) + .collect(); + let total = tiles.len(); + + let mut acc = allocate(pixels, [0f32; 3])?; + let mut wsum = allocate(nx * ny, 0f32)?; + let group = p.parallel.clamp(1, 4); + let mut done = 0; + // where the time goes, printed under `LIGHTCRAFT_PROFILE` + let (mut t_wait, mut t_total) = (std::time::Duration::ZERO, std::time::Duration::ZERO); + let began = web_time::Instant::now(); + + // One group of tiles at a time: pack, run the model, check and bring each to the input's scale, in parallel. + // Blending (sequential: the tiles add into one picture, in tile order so the result does not depend on timing) of + // the previous group runs while the next group is on the model, so neither the card nor the cores wait for it. + type Tile = (usize, usize, bool, bool, bool, bool); + let compute = |chunk: &[Tile]| -> Vec, Error>> { + chunk + .par_iter() + .map(|&(x0, y0, bx, ax, by, ay)| { + let mut input = allocate(4 * p.tile * p.tile, 0f32)?; + if !pack_tile(mosaic, width, height, layout, x0 as isize, y0 as isize, p.tile, &mut input) { + return Err(Error::Input("a tile could not be packed".into())); + } + let mut out = runner.run(&input)?; + let side = 2 * p.tile; + if out.len() != 3 * side * side { + return Err(Error::Model(format!("{} numbers for a tile, not the {} expected", out.len(), 3 * side * side))); + } + if out.iter().any(|v| !v.is_finite()) { + return Err(Error::Model("numbers that are not finite".into())); + } + finish_tile(&input, &mut out, p, (bx, ax, by, ay)); + Ok(out) + }) + .collect() + }; + let mut blend = |chunk: &[Tile], outs: Vec, Error>>| -> Result<(), Error> { + for (&(x0, y0, bx, ax, by, ay), r) in chunk.iter().zip(outs) { + blend_tile(&r?, &mut acc, &mut wsum, (width, height), (nx, ny), layout, (x0, y0), p, (bx, ax, by, ay)); + done += 1; + if let Some(f) = ctl.progress { + f(done, total); + } + } + Ok(()) + }; + let mut previous: Option<(&[Tile], Vec, Error>>)> = None; + for chunk in tiles.chunks(group) { + if ctl.cancel.is_some_and(|c| c.load(Ordering::Relaxed)) { + return Err(Error::Cancelled); + } + let (outs, blended) = rayon::join( + || compute(chunk), + || match previous.take() { + Some((c, o)) => { + let started = web_time::Instant::now(); + let r = blend(c, o); + (r, started.elapsed()) + } + None => (Ok(()), std::time::Duration::ZERO), + }, + ); + blended.0?; + t_wait += blended.1; + previous = Some((chunk, outs)); + } + if let Some((c, o)) = previous.take() { + let started = web_time::Instant::now(); + blend(c, o)?; + t_wait += started.elapsed(); + } + t_total += began.elapsed(); + let started = web_time::Instant::now(); + // normalise by the weights + let (ox, oy) = (layout.position(0, 0).0, layout.position(0, 0).1); + acc.par_chunks_mut(width).enumerate().for_each(|(y, row)| { + let cy = ((y as isize - oy) / 2) as usize; + for (x, px) in row.iter_mut().enumerate() { + let cx = ((x as isize - ox) / 2) as usize; + let w = wsum.get(cy * nx + cx).copied().unwrap_or(0.0); + if w > 0.0 { + *px = px.map(|v| v / w); + } + } + }); + if std::env::var_os("LIGHTCRAFT_PROFILE").is_some() { + eprintln!( + "[denoise] {total} tiles, {} at once: tiles {:.0} ms (blending, {:.0} ms of it, overlaps the model), normalise {:.0} ms", + group, + t_total.as_secs_f64() * 1e3, + t_wait.as_secs_f64() * 1e3, + started.elapsed().as_secs_f64() * 1e3 + ); + } + Ok(Rgb32f { width, height, data: acc }) +} + +/// Bring a tile's output to the input's scale and keep clipped highlights as they were (in place). +fn finish_tile(input: &[f32], out: &mut [f32], p: &Params, edges: (bool, bool, bool, bool)) { + let t = p.tile; + let side = 2 * t; + if let Gain::MatchMean { nominal, max_deviation } = p.gain { + // the part of the tile whose scale is measured: away from the edges that blend into a neighbour + let m = p.overlap / 2; + let (i0, i1) = (if edges.0 { m } else { 0 }, t - if edges.1 { m } else { 0 }); + let (j0, j1) = (if edges.2 { m } else { 0 }, t - if edges.3 { m } else { 0 }); + let cells = ((i1 - i0) * (j1 - j0)) as f64; + let plane_sum = |pl: usize| -> f64 { + (j0..j1) + .map(|j| input.get(pl * t * t + j * t + i0..pl * t * t + j * t + i1).map_or(0.0, |r| r.iter().map(|&v| f64::from(v)).sum::())) + .sum() + }; + let in_mean = [plane_sum(0) / cells, (plane_sum(1) + plane_sum(2)) / (2.0 * cells), plane_sum(3) / cells]; + for (c, &mean_in) in in_mean.iter().enumerate() { + let sum_out: f64 = (2 * j0..2 * j1) + .map(|y| { + out.get(c * side * side + y * side + 2 * i0..c * side * side + y * side + 2 * i1) + .map_or(0.0, |r| r.iter().map(|&v| f64::from(v)).sum::()) + }) + .sum(); + let mean_out = sum_out / (4.0 * cells); + let measured = if mean_in as f32 >= MIN_SIGNAL && mean_out > 0.0 { (mean_out / mean_in) as f32 } else { nominal }; + let g = measured.clamp(nominal * (1.0 - max_deviation), nominal * (1.0 + max_deviation)); + if g.is_finite() && g > 0.0 { + let inv = 1.0 / g; + if let Some(plane) = out.get_mut(c * side * side..(c + 1) * side * side) { + plane.iter_mut().for_each(|v| *v *= inv); + } + } + } + } + if let Some(clip) = p.clip { + protect_clipped(input, out, t, clip); + } +} + +/// Where the mosaic is at (or within a cell of) white, the output becomes the plain cell values again, so a clipped +/// area stays exactly as clipped as the camera recorded it (highlight recovery later depends on that). +fn protect_clipped(input: &[f32], out: &mut [f32], t: usize, clip: f32) { + let side = 2 * t; + let plane = t * t; + let (Some(r), Some(g1), Some(g2), Some(b)) = + (input.get(0..plane), input.get(plane..2 * plane), input.get(2 * plane..3 * plane), input.get(3 * plane..4 * plane)) + else { + return; + }; + let mut m: Vec = (0..plane).map(|i| r[i].max(g1[i]).max(g2[i]).max(b[i])).collect(); + if !m.iter().any(|&v| v >= clip) { + return; + } + // 3×3 maximum over cells (separable) + let dilate = |src: &[f32], horizontal: bool| -> Vec { + (0..plane) + .map(|i| { + let (x, y) = (i % t, i / t); + let mut best = src[i]; + for d in [-1isize, 1] { + let (nx, ny) = if horizontal { (x as isize + d, y as isize) } else { (x as isize, y as isize + d) }; + if nx >= 0 && ny >= 0 && (nx as usize) < t && (ny as usize) < t { + best = best.max(src[ny as usize * t + nx as usize]); + } + } + best + }) + .collect() + }; + m = dilate(&dilate(&m, true), false); + let lo = clip - 0.03; + for j in 0..t { + for i in 0..t { + let k = j * t + i; + let s = ((m[k] - lo) / (clip - lo)).clamp(0.0, 1.0); + if s <= 0.0 { + continue; + } + let s = s * s * (3.0 - 2.0 * s); + let plain = [r[k], 0.5 * (g1[k] + g2[k]), b[k]]; + for (c, &pv) in plain.iter().enumerate() { + for (sy, sx) in [(0, 0), (0, 1), (1, 0), (1, 1)] { + if let Some(v) = out.get_mut(c * side * side + (2 * j + sy) * side + 2 * i + sx) { + *v += (pv - *v) * s; + } + } + } + } + } +} + +/// Add a finished tile to the picture with its blend weights. +#[allow(clippy::too_many_arguments)] +fn blend_tile( + out: &[f32], + acc: &mut [[f32; 3]], + wsum: &mut [f32], + size: (usize, usize), + cells: (usize, usize), + layout: Layout, + at: (usize, usize), + p: &Params, + edges: (bool, bool, bool, bool), +) { + let (width, height) = size; + let (nx, ny) = cells; + let (t, side) = (p.tile, 2 * p.tile); + let fade = p.overlap; + let wx: Vec = (0..t).map(|i| weight(i, t, fade, edges.0, edges.1)).collect(); + let wy: Vec = (0..t).map(|j| weight(j, t, fade, edges.2, edges.3)).collect(); + for j in 0..t { + let cy = at.1 + j; + for i in 0..t { + let cx = at.0 + i; + // cells past the picture's edge (a picture smaller than one tile is padded) belong to no pixel + if cx >= nx || cy >= ny { + continue; + } + let (Some(&a), Some(&b)) = (wx.get(i), wy.get(j)) else { continue }; + let w = a * b; + if let Some(s) = wsum.get_mut(cy * nx + cx) { + *s += w; + } + let (px0, py0) = layout.position(cx as isize, cy as isize); + for (sy, sx) in [(0isize, 0isize), (0, 1), (1, 0), (1, 1)] { + let (x, y) = (px0 + sx, py0 + sy); + if x < 0 || y < 0 || x as usize >= width || y as usize >= height { + continue; + } + let o = (2 * j + sy as usize) * side + 2 * i + sx as usize; + let (Some(&r), Some(&g), Some(&bl)) = (out.get(o), out.get(side * side + o), out.get(2 * side * side + o)) else { continue }; + if let Some(px) = acc.get_mut(y as usize * width + x as usize) { + px[0] += w * r; + px[1] += w * g; + px[2] += w * bl; + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A stand-in model: the cell's samples as an RGB picture at twice the resolution (a plain 2 × 2 "demosaic"), + /// times `scale`, so the output has a scale of its own like the real model's. + struct Mock { + tile: usize, + scale: f32, + } + impl TileRunner for Mock { + fn run(&self, input: &[f32]) -> Result, Error> { + let (t, side) = (self.tile, 2 * self.tile); + let plane = t * t; + let mut out = vec![0.0; 3 * side * side]; + for j in 0..t { + for i in 0..t { + let k = j * t + i; + let rgb = [input[k], 0.5 * (input[plane + k] + input[2 * plane + k]), input[3 * plane + k]]; + for (c, v) in rgb.iter().enumerate() { + for (sy, sx) in [(0, 0), (0, 1), (1, 0), (1, 1)] { + out[c * side * side + (2 * j + sy) * side + 2 * i + sx] = v * self.scale; + } + } + } + } + Ok(out) + } + } + + fn params(tile: usize, overlap: usize, gain: Gain) -> Params { + Params { tile, overlap, gain, clip: None, parallel: 2 } + } + + #[test] + fn overflowing_tiles_and_nan_gain_are_errors_before_the_model_runs() { + let layout = Layout::from_cell([0, 1, 1, 2]).unwrap(); + let model = Mock { tile: 16, scale: 1.0 }; + for (tile, overlap, gain) in [ + (usize::MAX, 0, Gain::None), + (16, usize::MAX, Gain::None), + (16, 0, Gain::MatchMean { nominal: 1.0, max_deviation: f32::NAN }), + (16, 0, Gain::MatchMean { nominal: f32::INFINITY, max_deviation: 0.1 }), + ] { + assert!(matches!(denoise_bayer(&[0.5], 1, 1, layout, &model, ¶ms(tile, overlap, gain), &Control::default()), Err(Error::Input(_)))); + } + } + + /// A smooth picture whose colour at a sample is `base[colour] · (1 + shading)`. + fn mosaic(w: usize, h: usize, cell: [u8; 4]) -> Vec { + (0..w * h) + .map(|i| { + let (x, y) = (i % w, i / w); + let c = cell[(y % 2) * 2 + x % 2] as usize; + [0.5, 0.3, 0.2][c] * (0.6 + 0.4 * ((x as f32 / 23.0).sin() * (y as f32 / 31.0).cos())) + }) + .collect() + } + + fn reference(m: &[f32], w: usize, h: usize, layout: Layout) -> Rgb32f { + // what the mock gives for the whole picture in one tile + let (nx, ny) = layout.cells(w, h); + let t = nx.max(ny).div_ceil(16) * 16; + denoise_bayer(m, w, h, layout, &Mock { tile: t, scale: 1.0 }, ¶ms(t, 16, Gain::None), &Control::default()).unwrap() + } + + #[test] + fn tiles_blend_to_the_same_picture_as_one_tile_for_every_layout() { + for cell in [[0u8, 1, 1, 2], [1, 0, 2, 1], [1, 2, 0, 1], [2, 1, 1, 0]] { + let layout = Layout::from_cell(cell).unwrap(); + let (w, h) = (150usize, 110usize); + let m = mosaic(w, h, cell); + let want = reference(&m, w, h, layout); + let got = denoise_bayer(&m, w, h, layout, &Mock { tile: 32, scale: 1.0 }, ¶ms(32, 8, Gain::None), &Control::default()).unwrap(); + assert_eq!((got.width, got.height), (w, h)); + for (a, b) in got.data.iter().zip(&want.data) { + for c in 0..3 { + assert!((a[c] - b[c]).abs() < 1e-5, "layout {cell:?}: {a:?} vs {b:?}"); + } + } + } + } + + #[test] + fn the_picture_does_not_depend_on_how_many_tiles_run_at_once_or_how_long_each_takes() { + // tiles that finish in a different order every time: blending still adds them in tile order + struct Jittery(Mock); + impl TileRunner for Jittery { + fn run(&self, input: &[f32]) -> Result, Error> { + let wait = (input.iter().take(64).sum::() * 1000.0) as u64 % 7; + std::thread::sleep(std::time::Duration::from_millis(wait)); + self.0.run(input) + } + } + let layout = Layout { dx: 0, dy: 1 }; + let (w, h) = (150usize, 110usize); + let m = mosaic(w, h, [1, 0, 2, 1]); + let run = |parallel: usize| { + let mut p = params(32, 8, Gain::MatchMean { nominal: 1.0e6, max_deviation: 0.05 }); + p.parallel = parallel; + denoise_bayer(&m, w, h, layout, &Jittery(Mock { tile: 32, scale: 1.01e6 }), &p, &Control::default()).unwrap() + }; + let one = run(1); + for parallel in [2, 3, 8] { + assert_eq!(run(parallel).data, one.data, "{parallel} at once"); + } + } + + #[test] + fn a_failing_tile_or_a_cancel_ends_the_picture_without_a_panic() { + struct Fails(Mock, std::sync::atomic::AtomicUsize); + impl TileRunner for Fails { + fn run(&self, input: &[f32]) -> Result, Error> { + if self.1.fetch_add(1, Ordering::Relaxed) == 4 { + return Err(Error::Runtime("boom".into())); + } + self.0.run(input) + } + } + let layout = Layout { dx: 0, dy: 0 }; + let (w, h) = (150usize, 110usize); + let m = mosaic(w, h, [0, 1, 1, 2]); + let r = denoise_bayer( + &m, + w, + h, + layout, + &Fails(Mock { tile: 32, scale: 1.0 }, Default::default()), + ¶ms(32, 8, Gain::None), + &Control::default(), + ); + assert!(matches!(r, Err(Error::Runtime(_))), "{r:?}"); + let stop = AtomicBool::new(true); + let ctl = Control { cancel: Some(&stop), progress: None }; + let r = denoise_bayer(&m, w, h, layout, &Mock { tile: 32, scale: 1.0 }, ¶ms(32, 8, Gain::None), &ctl); + assert!(matches!(r, Err(Error::Cancelled))); + } + + #[test] + fn a_scale_of_its_own_is_taken_out_tile_by_tile() { + let layout = Layout { dx: 0, dy: 0 }; + let (w, h) = (150usize, 110usize); + let m = mosaic(w, h, [0, 1, 1, 2]); + let want = reference(&m, w, h, layout); + let gain = Gain::MatchMean { nominal: 1.0e6, max_deviation: 0.05 }; + // the scale differs a little from nominal; each tile's own scale is measured and removed + let got = denoise_bayer(&m, w, h, layout, &Mock { tile: 32, scale: 1.02e6 }, ¶ms(32, 8, gain), &Control::default()).unwrap(); + let mean = |img: &Rgb32f, c: usize| img.data.iter().map(|p| f64::from(p[c])).sum::() / img.data.len() as f64; + for c in 0..3 { + let (a, b) = (mean(&got, c), mean(&want, c)); + assert!((a - b).abs() / b < 0.01, "channel {c}: {a} vs {b}"); + } + // a scale far from nominal is clamped rather than trusted + let wild = denoise_bayer(&m, w, h, layout, &Mock { tile: 32, scale: 3.0e6 }, ¶ms(32, 8, gain), &Control::default()).unwrap(); + assert!(mean(&wild, 1) > 2.0 * mean(&want, 1)); + } + + #[test] + fn a_picture_smaller_than_a_tile_is_padded_by_mirroring() { + let layout = Layout { dx: 1, dy: 1 }; + let (w, h) = (20usize, 14usize); + let m = mosaic(w, h, [2, 1, 1, 0]); + let got = denoise_bayer(&m, w, h, layout, &Mock { tile: 64, scale: 1.0 }, ¶ms(64, 16, Gain::None), &Control::default()).unwrap(); + assert_eq!((got.width, got.height), (w, h)); + assert!(got.data.iter().all(|p| p.iter().all(|v| v.is_finite() && *v > 0.0))); + } + + #[test] + fn clipped_highlights_stay_clipped() { + // a model that darkens everything by 10 % would pull a clipped plateau below white + struct Dim(Mock); + impl TileRunner for Dim { + fn run(&self, input: &[f32]) -> Result, Error> { + Ok(self.0.run(input)?.into_iter().map(|v| v * 0.9).collect()) + } + } + let layout = Layout { dx: 0, dy: 0 }; + let (w, h) = (96usize, 96usize); + let m: Vec = (0..w * h).map(|i| if i % w > 40 && i % w < 70 { 1.0 } else { 0.4 }).collect(); + let mut p = params(32, 8, Gain::None); + p.clip = Some(0.99); + let got = denoise_bayer(&m, w, h, layout, &Dim(Mock { tile: 32, scale: 1.0 }), &p, &Control::default()).unwrap(); + let mid = got.data[48 * w + 55]; + assert!(mid.iter().all(|&v| v > 0.995), "plateau {mid:?}"); + let dark = got.data[48 * w + 10]; + assert!(dark.iter().all(|&v| (v - 0.36).abs() < 0.01), "unclipped area still goes through the model: {dark:?}"); + } + + #[test] + fn cancelling_and_progress_work() { + let layout = Layout { dx: 0, dy: 0 }; + let (w, h) = (150usize, 110usize); + let m = mosaic(w, h, [0, 1, 1, 2]); + let cancel = AtomicBool::new(true); + let r = denoise_bayer( + &m, + w, + h, + layout, + &Mock { tile: 32, scale: 1.0 }, + ¶ms(32, 8, Gain::None), + &Control { cancel: Some(&cancel), progress: None }, + ); + assert!(matches!(r, Err(Error::Cancelled))); + let seen = std::sync::Mutex::new(Vec::new()); + let f = |d: usize, n: usize| seen.lock().unwrap().push((d, n)); + denoise_bayer(&m, w, h, layout, &Mock { tile: 32, scale: 1.0 }, ¶ms(32, 8, Gain::None), &Control { cancel: None, progress: Some(&f) }) + .unwrap(); + let seen = seen.into_inner().unwrap(); + let total = tile_count(w, h, layout, 32, 8); + assert!(total > 1); + assert_eq!(seen.last(), Some(&(total, total))); + assert!(seen.windows(2).all(|p| p[1].0 == p[0].0 + 1)); + } + + #[test] + fn hostile_inputs_are_errors() { + let layout = Layout { dx: 0, dy: 0 }; + let mock = Mock { tile: 32, scale: 1.0 }; + let p = params(32, 8, Gain::None); + assert!(matches!(denoise_bayer(&[], 0, 0, layout, &mock, &p, &Control::default()), Err(Error::Input(_)))); + assert!(matches!(denoise_bayer(&[0.0; 10], 4, 4, layout, &mock, &p, &Control::default()), Err(Error::Input(_)))); + assert!(matches!(denoise_bayer(&[0.0; 16], 4, 4, layout, &mock, ¶ms(32, 16, Gain::None), &Control::default()), Err(Error::Input(_)))); + struct Bad(Vec); + impl TileRunner for Bad { + fn run(&self, _: &[f32]) -> Result, Error> { + Ok(self.0.clone()) + } + } + assert!(matches!(denoise_bayer(&[0.5; 16], 4, 4, layout, &Bad(vec![0.0; 5]), &p, &Control::default()), Err(Error::Model(_)))); + assert!(matches!(denoise_bayer(&[0.5; 16], 4, 4, layout, &Bad(vec![f32::NAN; 3 * 64 * 64]), &p, &Control::default()), Err(Error::Model(_)))); + } +} diff --git a/crates/denoise-core/src/tiles.rs b/crates/denoise-core/src/tiles.rs new file mode 100644 index 00000000..59898c5a --- /dev/null +++ b/crates/denoise-core/src/tiles.rs @@ -0,0 +1,100 @@ +//! Cutting a picture into overlapping tiles and putting the results back together. +//! +//! A model runs on a fixed square, so a picture is covered by tiles spaced evenly along each axis, each sharing +//! at least `overlap` cells with its neighbours. Each tile's influence fades towards its own edges (except at the +//! picture's edge), and the result is the sum of tile × weight divided by the sum of weights, so any set of +//! overlapping tiles blends consistently and a seam between two tiles of slightly different scale is a smooth ramp. + +/// Where the tiles start along one axis of `n` cells: `[0]` when the axis fits one tile (the tile is then padded by +/// mirroring), else evenly spaced starts, the first at 0 and the last at `n − tile`, each tile sharing at least +/// `overlap` cells with the next. +pub fn starts(n: usize, tile: usize, overlap: usize) -> Vec { + if n <= tile || tile == 0 { + return vec![0]; + } + let step = tile.saturating_sub(overlap).max(1); + // tiles needed so that start k·spacing with spacing ≤ step reaches n − tile + let Some(k) = (n - tile).div_ceil(step).checked_add(1).filter(|&k| k <= 1_000_000) else { return Vec::new() }; + let last = n - tile; + (0..k).map(|i| if k == 1 { 0 } else { ((i as u128 * last as u128 + (k - 1) as u128 / 2) / (k - 1) as u128) as usize }).collect() +} + +/// Weight of the cell at offset `p` (0-based, inside a tile of `tile` cells) along one axis: 1 in the middle, easing +/// down to a sliver over `fade` cells at an edge that has a neighbour, 1 right up to an edge that is the picture's. +pub fn weight(p: usize, tile: usize, fade: usize, has_before: bool, has_after: bool) -> f32 { + let fade = fade.max(1) as f32; + let ramp = |d: usize| { + let t = ((d as f32 + 0.5) / fade).min(1.0); + // smoothstep, never exactly 0 so that a cell is always covered + (t * t * (3.0 - 2.0 * t)).max(1e-3) + }; + let before = if has_before { ramp(p) } else { 1.0 }; + let after = if has_after { ramp(tile.saturating_sub(1).saturating_sub(p)) } else { 1.0 }; + before.min(after) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn hostile_axes_have_bounded_tile_counts_and_checked_spacing() { + assert!(starts(usize::MAX, 1, 0).is_empty()); + assert!(starts(usize::MAX, 2, usize::MAX).is_empty()); + let s = starts(usize::MAX, usize::MAX / 2, 0); + assert_eq!(s.first(), Some(&0)); + assert_eq!(s.last(), Some(&(usize::MAX - usize::MAX / 2))); + assert!(s.windows(2).all(|s| s[0] < s[1])); + } + + #[test] + fn one_tile_when_it_fits() { + assert_eq!(starts(100, 512, 64), vec![0]); + assert_eq!(starts(512, 512, 64), vec![0]); + assert_eq!(starts(0, 512, 64), vec![0]); + assert_eq!(starts(10, 0, 0), vec![0]); + } + + #[test] + fn tiles_cover_the_axis_and_overlap_enough() { + for n in [513usize, 600, 900, 1000, 1024, 1500, 3024, 3600, 5000, 12345] { + let s = starts(n, 512, 64); + assert_eq!(s[0], 0, "n={n}"); + assert_eq!(*s.last().unwrap(), n - 512, "n={n}"); + for w in s.windows(2) { + assert!(w[1] > w[0], "starts increase, n={n}"); + assert!(w[0] + 512 >= w[1] + 64, "at least 64 shared cells between {} and {}, n={n}", w[0], w[1]); + } + } + } + + #[test] + fn weights_fade_only_towards_neighbours() { + let (t, f) = (512, 64); + // a lone tile: full weight everywhere + assert!((0..t).all(|p| weight(p, t, f, false, false) == 1.0)); + // the first tile fades on its far side only + assert_eq!(weight(0, t, f, false, true), 1.0); + assert!(weight(t - 1, t, f, false, true) < 0.01); + // the last tile fades on its near side only + assert!(weight(0, t, f, true, false) < 0.01); + assert_eq!(weight(t - 1, t, f, true, false), 1.0); + // a middle tile is 1 in the middle and everywhere positive + assert_eq!(weight(t / 2, t, f, true, true), 1.0); + assert!((0..t).all(|p| weight(p, t, f, true, true) > 0.0)); + } + + #[test] + fn two_overlapping_tiles_blend_to_a_partition_of_unity() { + // at every cell of the shared region the normalised weights are in [0, 1] and sum to 1 + let (t, f) = (512usize, 64usize); + let s = starts(900, t, 64); + assert_eq!(s.len(), 2); + for cell in s[1]..s[0] + t { + let a = weight(cell - s[0], t, f, false, true); + let b = weight(cell - s[1], t, f, true, false); + let (na, nb) = (a / (a + b), b / (a + b)); + assert!((na + nb - 1.0).abs() < 1e-6 && na >= 0.0 && nb >= 0.0); + } + } +} diff --git a/crates/denoise/Cargo.toml b/crates/denoise/Cargo.toml new file mode 100644 index 00000000..621fd3d6 --- /dev/null +++ b/crates/denoise/Cargo.toml @@ -0,0 +1,32 @@ +[package] +name = "lightcraft-denoise" +version.workspace = true +edition.workspace = true +license.workspace = true +description = "LightCraft AI denoise (L1): denoise-model manifests, Bayer packing, tiling and blending, the cache file, and pure-Rust CPU inference" + +[lints] +workspace = true + +[features] +default = [] +runtime = ["dep:faer"] +# Product-policy switch: absent by default; users may install their own model. +rawnind-model = ["lightcraft-denoise-core/rawnind-model"] + +[dependencies] +lightcraft-denoise-core = { workspace = true } +# Safe matrix views over the pure-Rust GEMM backend; no BLAS, cc/build scripts +# or inner thread pool. Its dependencies include architecture-specific assembly. +faer = { package = "faer-core", version = "=0.17.1", default-features = false, features = ["std"], optional = true } +lightcraft-raster = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +sha2 = { workspace = true } +thiserror = { workspace = true } +miniz_oxide = { workspace = true } +rayon = { workspace = true } +web-time = { workspace = true } + +[dev-dependencies] +proptest = { workspace = true } diff --git a/crates/denoise/README.md b/crates/denoise/README.md new file mode 100644 index 00000000..0397193f --- /dev/null +++ b/crates/denoise/README.md @@ -0,0 +1,18 @@ +# lightcraft-denoise + +L1 pure-Rust AI RAW denoise primitives and CPU inference. No weights are bundled. +See [user workflow and limitations](../../docs/denoise.md). + +`lightcraft-denoise-core` provides metadata, Bayer helpers and cache formats, +re-exported here. `manifest` and `known` validate opt-in models; `onnx` reads a bounded operator subset +into `net`; `cpu` executes it and `reference` supplies scalar checks. `bayer`, `tiles` +and `run` pack sensor data and blend overlapping tiles while preserving clipped +highlights. `product` reads/writes the disposable `.lcdn` cache. `runtime` supplies +the installation self-test. Synthetic networks are generated in Rust for tests. + +Run `cargo test -p lightcraft-denoise --features runtime`. The CPU backend is optional; +featureless and web builds keep the model metadata and cache format without faer/GEMM. +The `rawnind-model` feature enables the otherwise empty download catalog and requires +an explicit release-policy decision. Real-model tests are ignored by default and +require an explicitly supplied local model. The GPU twin is `lightcraft-gpu::nn`; +model management, background jobs and export integration live in `lightcraft-engine`. diff --git a/crates/denoise/src/cpu.rs b/crates/denoise/src/cpu.rs new file mode 100644 index 00000000..34e5269e --- /dev/null +++ b/crates/denoise/src/cpu.rs @@ -0,0 +1,386 @@ +//! Single-thread CPU execution of the checked denoise network. +//! +//! Convolutions lower at most 256 spatial positions at a time, then use faer's safe views over pure-Rust GEMM. +//! No inner thread pool: `run::denoise_bayer` already runs tiles on the application's Rayon pool. Tensor lifetimes +//! are planned once; buffers and im2col storage are recycled between calls, independently for concurrent callers. + +use crate::net::{Net, NetError, Op}; +use crate::run::{Error, TileRunner}; +use faer::{Parallelism, mat}; +use std::sync::{Arc, Mutex}; + +const BLOCK: usize = 256; +/// Cap all tensor buffers and convolution scratch for one tile together (512 MiB). +const MAX_VALUES: usize = 128 * 1024 * 1024; +const MAX_TILE: usize = 2048; +const MAX_IDLE: usize = 1; + +type Result = std::result::Result; +fn bad(why: &str) -> NetError { + NetError::Invalid(why.into()) +} +fn count(c: usize, side: usize) -> Result { + c.checked_mul(side) + .and_then(|n| n.checked_mul(side)) + .filter(|&n| n > 0 && n <= MAX_VALUES) + .ok_or_else(|| bad("a tensor has no size or is too large")) +} +fn zeros(n: usize) -> Result> { + let mut v = Vec::new(); + v.try_reserve_exact(n).map_err(|_| bad("not enough memory for a denoise tile"))?; + v.resize(n, 0.0); + Ok(v) +} +fn sources(op: &Op) -> (usize, Option) { + match op { + Op::Conv(c) => (c.src, c.src2), + Op::ConvT2(c) => (c.src, None), + Op::MaxPool2 { src } | Op::DepthToSpace2 { src } => (*src, None), + } +} +#[derive(Clone, Copy)] +struct Tensor { + slot: usize, + side: usize, + channels: usize, +} +struct Plan { + net: Net, + tensors: Vec, + sizes: Vec, + lower: usize, + product: usize, + up_weights: Vec>>, +} +struct Workspace { + buffers: Vec>, + lower: Vec, + product: Vec, +} + +/// A checked Net on square CHW tiles. Clones share immutable weights and the idle workspace cache. +#[derive(Clone)] +pub struct NetRunner { + plan: Arc, + idle: Arc>>, +} + +impl NetRunner { + pub fn new(net: &Net, tile: usize) -> Result { + if tile == 0 + || tile > MAX_TILE + || !net.fits_tile(tile) + || net.in_channels != 4 + || net.out_channels() != 3 + || net.shape(0).is_none_or(|s| s.channels != net.in_channels) + { + return Err(bad("the tile or Bayer-to-RGB contract does not fit the network")); + } + let mut last = vec![0; net.ops().len() + 1]; + for (i, op) in net.ops().iter().enumerate() { + let (a, b) = sources(op); + for id in std::iter::once(a).chain(b) { + *last.get_mut(id).ok_or_else(|| bad("missing tensor"))? = i + 1; + } + } + let mut tensors = vec![Tensor { slot: 0, side: tile, channels: net.in_channels }]; + let mut sizes = vec![count(net.in_channels, tile)?]; + let mut busy = vec![true]; + let (mut lower, mut product) = (0, 0); + let mut up_weights = Vec::new(); + for (i, op) in net.ops().iter().enumerate() { + let (src, src2) = sources(op); + let a = *tensors.get(src).ok_or_else(|| bad("missing tensor"))?; + let (channels, side) = match op { + Op::Conv(c) => { + if let Some(s) = src2 + && tensors.get(s).is_none_or(|b| b.side != a.side) + { + return Err(bad("a join has different sizes")); + } + (c.cout, a.side) + } + Op::ConvT2(c) => (c.cout, a.side.checked_mul(2).ok_or_else(|| bad("size overflow"))?), + Op::MaxPool2 { .. } => { + if a.side < 2 || !a.side.is_multiple_of(2) { + return Err(bad("pooling does not divide the tile")); + } + (a.channels, a.side / 2) + } + Op::DepthToSpace2 { .. } => (a.channels / 4, a.side.checked_mul(2).ok_or_else(|| bad("size overflow"))?), + }; + let n = count(channels, side)?; + let slot = busy + .iter() + .zip(&sizes) + .enumerate() + .filter(|(_, (used, _))| !**used) + .min_by_key(|(_, (_, size))| (n.saturating_sub(**size), **size)) + .map_or(busy.len(), |(i, _)| i); + if slot == sizes.len() { + sizes.push(n); + busy.push(true); + } else { + let sz = sizes.get_mut(slot).ok_or_else(|| bad("missing buffer"))?; + *sz = (*sz).max(n); + *busy.get_mut(slot).ok_or_else(|| bad("missing buffer"))? = true; + } + tensors.push(Tensor { slot, side, channels }); + // Release sources only after the output has its own buffer (also handles reading the same tensor twice). + for id in std::iter::once(src).chain(src2) { + if last.get(id) == Some(&(i + 1)) { + let t = tensors.get(id).ok_or_else(|| bad("missing tensor"))?; + *busy.get_mut(t.slot).ok_or_else(|| bad("missing buffer"))? = false; + } + } + let (k, m, positions) = match op { + Op::Conv(c) => (c.cin * c.k * c.k, c.cout, a.side * a.side), + Op::ConvT2(c) => (c.cin, c.cout * 4, a.side * a.side), + _ => (0, 0, 0), + }; + lower = lower.max(k.saturating_mul(BLOCK.min(positions))); + product = product.max(m.saturating_mul(BLOCK.min(positions))); + let packed = if let Op::ConvT2(c) = op { + let mut w = zeros(c.weight.len())?; + // [cin][cout][dy][dx] -> [cout*4][cin], so an up-convolution is one GEMM and a shuffle. + for (ci, row) in c.weight.chunks_exact(c.cout * 4).enumerate() { + for (co, &v) in row.iter().enumerate() { + *w.get_mut(co * c.cin + ci).ok_or_else(|| bad("invalid transposed weights"))? = v; + } + } + Some(w) + } else { + None + }; + up_weights.push(packed); + } + let total = sizes + .iter() + .try_fold(lower.saturating_add(product), |n, &s| n.checked_add(s)) + .filter(|&n| n <= MAX_VALUES) + .ok_or_else(|| bad("the tile needs too much working memory"))?; + let _ = total; + Ok(Self { plan: Arc::new(Plan { net: net.clone(), tensors, sizes, lower, product, up_weights }), idle: Arc::new(Mutex::new(Vec::new())) }) + } + /// Immutable data for the GPU runner, without re-reading the model. + pub fn net(&self) -> &Net { + &self.plan.net + } + pub fn bytes_per_tile(&self) -> usize { + (self.plan.sizes.iter().sum::() + self.plan.lower + self.plan.product) * 4 + } + fn workspace(&self) -> Result { + if let Some(w) = self.idle.lock().unwrap_or_else(std::sync::PoisonError::into_inner).pop() { + return Ok(w); + } + Ok(Workspace { + buffers: self.plan.sizes.iter().map(|&n| zeros(n)).collect::>()?, + lower: zeros(self.plan.lower)?, + product: zeros(self.plan.product)?, + }) + } + fn execute(&self, input: &[f32], w: &mut Workspace) -> Result> { + let first = self.plan.tensors.first().ok_or_else(|| bad("no input"))?; + let want = count(first.channels, first.side)?; + w.buffers.first_mut().and_then(|b| b.get_mut(..want)).ok_or_else(|| bad("missing input buffer"))?.copy_from_slice(input); + let profile = std::env::var_os("LIGHTCRAFT_PROFILE").is_some(); + for (i, op) in self.plan.net.ops().iter().enumerate() { + let started = web_time::Instant::now(); + let out = *self.plan.tensors.get(i + 1).ok_or_else(|| bad("missing output tensor"))?; + let (before, rest) = w.buffers.split_at_mut_checked(out.slot).ok_or_else(|| bad("missing output buffer"))?; + let (dest, after) = rest.split_first_mut().ok_or_else(|| bad("missing output buffer"))?; + let get = |id: usize| -> Result<(&[f32], Tensor)> { + let t = *self.plan.tensors.get(id).ok_or_else(|| bad("missing source tensor"))?; + let data = if t.slot < out.slot { + before.get(t.slot) + } else { + after.get(t.slot.checked_sub(out.slot + 1).ok_or_else(|| bad("aliased buffers"))?) + } + .and_then(|b| b.get(..count(t.channels, t.side).ok()?)) + .ok_or_else(|| bad("missing source buffer"))?; + Ok((data, t)) + }; + let (src, src2) = sources(op); + let (a, shape) = get(src)?; + let b = src2.map(get).transpose()?.map(|v| v.0); + let dest = dest.get_mut(..count(out.channels, out.side)?).ok_or_else(|| bad("short output buffer"))?; + match op { + Op::Conv(c) => conv(c, a, b, shape.side, dest, &mut w.lower, &mut w.product)?, + Op::ConvT2(c) => up( + c, + a, + shape.side, + dest, + self.plan.up_weights.get(i).and_then(Option::as_deref).ok_or_else(|| bad("missing up weights"))?, + &mut w.lower, + &mut w.product, + )?, + Op::MaxPool2 { .. } => pool(a, shape.side, dest)?, + Op::DepthToSpace2 { .. } => shuffle(a, shape.side, dest)?, + } + if dest.iter().any(|v| !v.is_finite()) { + return Err(bad("the model produced non-finite values")); + } + if profile { + eprintln!("[denoise-cpu] layer {i}: {}x{}x{} {:.2} ms", out.channels, out.side, out.side, started.elapsed().as_secs_f64() * 1e3); + } + } + let t = self.plan.tensors.last().ok_or_else(|| bad("no output"))?; + let data = w.buffers.get(t.slot).and_then(|b| b.get(..count(t.channels, t.side).ok()?)).ok_or_else(|| bad("missing result"))?; + let mut out = zeros(data.len())?; + out.copy_from_slice(data); + Ok(out) + } +} +impl TileRunner for NetRunner { + fn run(&self, input: &[f32]) -> std::result::Result, Error> { + let first = self.plan.tensors.first().ok_or_else(|| Error::Runtime("missing input shape".into()))?; + let want = count(first.channels, first.side).map_err(|e| Error::Runtime(e.to_string()))?; + if input.len() != want || input.iter().any(|v| !v.is_finite()) { + return Err(Error::Input("the tile has the wrong size or non-finite values".into())); + } + let mut w = self.workspace().map_err(|e| Error::Runtime(e.to_string()))?; + let result = self.execute(input, &mut w).map_err(|e| Error::Runtime(e.to_string())); + let mut idle = self.idle.lock().unwrap_or_else(std::sync::PoisonError::into_inner); + if idle.len() < MAX_IDLE { + idle.push(w); + } + result + } +} +fn activate(x: f32, bias: f32, alpha: Option) -> f32 { + let x = x + bias; + if x < 0.0 { x * alpha.unwrap_or(1.0) } else { x } +} + +/// All views are over exact, checked slices. The matrix library's assertions therefore never see file-derived sizes. +fn multiply(w: &[f32], x: &[f32], y: &mut [f32], m: usize, k: usize, n: usize) -> Result<()> { + if m.checked_mul(k) != Some(w.len()) || k.checked_mul(n) != Some(x.len()) || m.checked_mul(n) != Some(y.len()) || m == 0 || k == 0 || n == 0 { + return Err(bad("invalid matrix dimensions")); + } + faer::mul::matmul( + mat::from_row_major_slice_mut::(y, m, n), + mat::from_row_major_slice::(w, m, k), + mat::from_row_major_slice::(x, k, n), + None, + 1.0, + Parallelism::None, + ); + Ok(()) +} +fn lower(a: &[f32], b: Option<&[f32]>, side: usize, k: usize, start: usize, n: usize, dest: &mut [f32]) -> Result<()> { + let pad = k / 2; + let pixels = side * side; + // Whole row segments are copied, so bounds checks and edge branches stay out of the multiply loop. + for (ci, plane) in a.chunks_exact(pixels).chain(b.into_iter().flat_map(|b| b.chunks_exact(pixels))).enumerate() { + for ky in 0..k { + for kx in 0..k { + let row = (ci * k * k + ky * k + kx) * n; + let out = dest.get_mut(row..row + n).ok_or_else(|| bad("short im2col buffer"))?; + let mut at = 0; + while at < n { + let p = start + at; + let (y, x) = (p / side, p % side); + let len = (side - x).min(n - at); + if let Some(yy) = y.checked_add(ky).and_then(|v| v.checked_sub(pad)).filter(|&v| v < side) { + let lo = x.max(pad.saturating_sub(kx)); + let hi = (x + len).min((side + pad).saturating_sub(kx)); + if lo < hi { + let from = yy * side + lo + kx - pad; + let to = at + lo - x; + out.get_mut(at..to).ok_or_else(|| bad("short im2col padding"))?.fill(0.0); + out.get_mut(to..to + hi - lo) + .ok_or_else(|| bad("short im2col row"))? + .copy_from_slice(plane.get(from..from + hi - lo).ok_or_else(|| bad("short source row"))?); + out.get_mut(to + hi - lo..at + len).ok_or_else(|| bad("short im2col padding"))?.fill(0.0); + } else { + out.get_mut(at..at + len).ok_or_else(|| bad("short im2col row"))?.fill(0.0); + } + } else { + out.get_mut(at..at + len).ok_or_else(|| bad("short im2col row"))?.fill(0.0); + } + at += len; + } + } + } + } + Ok(()) +} +fn conv(c: &crate::net::Conv, a: &[f32], b: Option<&[f32]>, side: usize, dest: &mut [f32], scratch: &mut [f32], product: &mut [f32]) -> Result<()> { + let pixels = side * side; + let k = c.cin * c.k * c.k; + for start in (0..pixels).step_by(BLOCK) { + let n = BLOCK.min(pixels - start); + let x = scratch.get_mut(..k * n).ok_or_else(|| bad("short im2col storage"))?; + lower(a, b, side, c.k, start, n, x)?; + let y = product.get_mut(..c.cout * n).ok_or_else(|| bad("short product storage"))?; + multiply(&c.weight, x, y, c.cout, k, n)?; + for ((row, &bias), out) in y.chunks_exact(n).zip(&c.bias).zip(dest.chunks_exact_mut(pixels)) { + let to = out.get_mut(start..start + n).ok_or_else(|| bad("short convolution output"))?; + for (d, &v) in to.iter_mut().zip(row) { + *d = activate(v, bias, c.leaky); + } + } + } + Ok(()) +} +fn up(c: &crate::net::ConvT2, a: &[f32], side: usize, dest: &mut [f32], weights: &[f32], scratch: &mut [f32], product: &mut [f32]) -> Result<()> { + let pixels = side * side; + let out_side = 2 * side; + for start in (0..pixels).step_by(BLOCK) { + let n = BLOCK.min(pixels - start); + let x = scratch.get_mut(..c.cin * n).ok_or_else(|| bad("short up input storage"))?; + for (to, from) in x.chunks_exact_mut(n).zip(a.chunks_exact(pixels)) { + to.copy_from_slice(from.get(start..start + n).ok_or_else(|| bad("short up input"))?); + } + let y = product.get_mut(..c.cout * 4 * n).ok_or_else(|| bad("short up product storage"))?; + multiply(weights, x, y, c.cout * 4, c.cin, n)?; + for (co, (&bias, out)) in c.bias.iter().zip(dest.chunks_exact_mut(4 * pixels)).enumerate() { + for d in 0..4 { + let row = y.get((4 * co + d) * n..(4 * co + d + 1) * n).ok_or_else(|| bad("short up product"))?; + for (j, &v) in row.iter().enumerate() { + let p = start + j; + let at = (2 * (p / side) + d / 2) * out_side + 2 * (p % side) + d % 2; + *out.get_mut(at).ok_or_else(|| bad("short up output"))? = activate(v, bias, c.leaky); + } + } + } + } + Ok(()) +} +fn pool(a: &[f32], side: usize, dest: &mut [f32]) -> Result<()> { + let s = side / 2; + for (input, out) in a.chunks_exact(side * side).zip(dest.chunks_exact_mut(s * s)) { + for (y, row) in out.chunks_exact_mut(s).enumerate() { + let lo = input.get(2 * y * side..(2 * y + 1) * side).ok_or_else(|| bad("short pool input"))?; + let hi = input.get((2 * y + 1) * side..(2 * y + 2) * side).ok_or_else(|| bad("short pool input"))?; + for ((d, lo), hi) in row.iter_mut().zip(lo.as_chunks::<2>().0).zip(hi.as_chunks::<2>().0) { + let [a, b] = *lo; + let [c, e] = *hi; + *d = a.max(b).max(c).max(e); + } + } + } + Ok(()) +} +fn shuffle(a: &[f32], side: usize, dest: &mut [f32]) -> Result<()> { + let pixels = side * side; + for (input, out) in a.chunks_exact(4 * pixels).zip(dest.chunks_exact_mut(4 * pixels)) { + for d in 0..4 { + let plane = input.get(d * pixels..(d + 1) * pixels).ok_or_else(|| bad("short shuffle input"))?; + for (y, row) in plane.chunks_exact(side).enumerate() { + let at = (2 * y + d / 2) * 2 * side; + let to = out.get_mut(at..at + 2 * side).ok_or_else(|| bad("short shuffle output"))?; + for (to, &v) in to.as_chunks_mut::<2>().0.iter_mut().zip(row) { + if let Some(t) = to.get_mut(d % 2) { + *t = v; + } + } + } + } + } + Ok(()) +} + +#[cfg(test)] +mod tests; diff --git a/crates/denoise/src/cpu/tests.rs b/crates/denoise/src/cpu/tests.rs new file mode 100644 index 00000000..16a496b5 --- /dev/null +++ b/crates/denoise/src/cpu/tests.rs @@ -0,0 +1,151 @@ +use super::*; +use crate::{ + net::{Conv, ConvT2}, + reference, synthetic, +}; + +fn random(n: usize, seed: u32) -> Vec { + let mut s = seed | 1; + (0..n) + .map(|_| { + s ^= s << 13; + s ^= s >> 17; + s ^= s << 5; + (s >> 8) as f32 / (1u32 << 23) as f32 - 1.0 + }) + .collect() +} +fn net(tile: usize, ch: u64, depth: usize, seed: u64) -> Net { + let p = std::env::temp_dir().join(format!("lc-purecpu-{}-{tile}-{ch}-{depth}-{seed}.onnx", std::process::id())); + std::fs::write(&p, synthetic::unet_onnx(tile as u64, ch, depth, seed)).unwrap(); + let n = crate::onnx::read(&p).unwrap(); + std::fs::remove_file(p).unwrap(); + n +} +fn agree(a: &[f32], b: &[f32]) { + assert_eq!(a.len(), b.len()); + assert!(a.iter().chain(b).all(|v| v.is_finite())); + let scale = b.iter().fold(1e-6f32, |m, v| m.max(v.abs())); + let error = a.iter().zip(b).fold(0f32, |m, (a, b)| m.max((a - b).abs())) / scale; + assert!(error <= 1e-3, "relative maximum error {error:e}"); +} +#[test] +fn all_operators_match_scalar_on_random_nets_and_odd_and_tiny_shapes() { + for (t, ch, d) in [(1, 3, 0), (3, 5, 0), (5, 7, 0), (6, 3, 1), (20, 5, 2), (36, 9, 2), (64, 8, 3)] { + for seed in [1, 73] { + let n = net(t, ch, d, seed); + let r = NetRunner::new(&n, t).unwrap(); + let x = random(4 * t * t, seed as u32); + let want = reference::run(&n, t, &x).unwrap(); + agree(&r.run(&x).unwrap(), &want); + assert_eq!(r.run(&x).unwrap(), r.run(&x).unwrap()); + } + } +} + +#[test] +fn malformed_values_fail_before_execution_and_workspace_sizes_are_bounded() { + let mut n = net(4, 3, 1, 11); + for t in [0, 1, 3, usize::MAX, 4096] { + assert!(NetRunner::new(&n, t).is_err()); + } + for ch in [0, 3, usize::MAX] { + n.in_channels = ch; + assert!(NetRunner::new(&n, 4).is_err()); + } + let n = Net::new( + 4, + vec![ + Op::Conv(Conv { src: 0, src2: None, cin: 4, cout: 4096, k: 1, weight: vec![0.1; 4 * 4096], bias: vec![0.0; 4096], leaky: None }), + Op::Conv(Conv { src: 1, src2: None, cin: 4096, cout: 12, k: 1, weight: vec![0.1; 4096 * 12], bias: vec![0.0; 12], leaky: None }), + Op::DepthToSpace2 { src: 2 }, + ], + ) + .unwrap(); + assert!(NetRunner::new(&n, 2048).is_err()); + let ops = (0..100).map(|i| Op::MaxPool2 { src: i }).collect(); + let n = Net::new(4, ops).unwrap(); + assert!(NetRunner::new(&n, 16).is_err()); + let c = Conv { src: 0, src2: None, cin: 4, cout: 12, k: 1, weight: vec![0.25; 48], bias: vec![0.0; 12], leaky: None }; + for mode in 0..6 { + let mut bad = c.clone(); + match mode { + 0 => bad.src = usize::MAX, + 1 => bad.weight.pop().map(|_| ()).unwrap(), + 2 => bad.bias[0] = f32::NAN, + 3 => bad.leaky = Some(f32::INFINITY), + 4 => bad.cout = usize::MAX, + _ => bad.k = usize::MAX, + } + assert!(Net::new(4, vec![Op::Conv(bad), Op::DepthToSpace2 { src: 1 }]).is_err()); + } +} +#[test] +fn bad_tiles_and_arithmetic_overflow_are_errors_and_do_not_poison_reuse() { + let c = Conv { src: 0, src2: None, cin: 4, cout: 12, k: 1, weight: vec![f32::MAX; 48], bias: vec![0.0; 12], leaky: None }; + let n = Net::new(4, vec![Op::Conv(c), Op::DepthToSpace2 { src: 1 }]).unwrap(); + let r = NetRunner::new(&n, 1).unwrap(); + assert!(r.run(&[]).is_err()); + assert!(r.run(&[f32::NAN; 4]).is_err()); + assert!(r.run(&[f32::INFINITY; 4]).is_err()); + assert_eq!(r.idle.lock().unwrap().len(), 0, "invalid input never allocates a workspace"); + assert!(r.run(&[2.0; 4]).is_err()); + assert_eq!(r.run(&[0.0; 4]).unwrap(), vec![0.0; 12]); + // Matrix shape checks are the boundary in front of the dependency's asserting constructors. + assert!(multiply(&[], &[], &mut [], usize::MAX, 2, 2).is_err()); +} +#[test] +fn shared_runner_serves_concurrent_tiles_without_cross_contamination() { + let n = net(12, 5, 1, 97); + let r = NetRunner::new(&n, 12).unwrap(); + let inputs: Vec<_> = (0..8).map(|s| random(4 * 12 * 12, s)).collect(); + let answers: Vec<_> = inputs.iter().map(|x| reference::run(&n, 12, x).unwrap()).collect(); + std::thread::scope(|scope| { + let mut handles = Vec::new(); + for (x, want) in inputs.iter().zip(&answers) { + let r = &r; + handles.push(scope.spawn(move || { + for _ in 0..3 { + agree(&r.run(x).unwrap(), want); + } + })); + } + for h in handles { + h.join().unwrap(); + } + }); +} +#[test] +fn join_may_read_the_same_tensor_twice_and_leaky_up_convolution_is_supported() { + let c = ConvT2 { src: 0, cin: 4, cout: 3, weight: random(48, 2), bias: vec![0.1, -0.1, 0.0], leaky: Some(0.13) }; + let n = Net::new( + 4, + vec![ + Op::ConvT2(c), + Op::Conv(Conv { src: 1, src2: Some(1), cin: 6, cout: 3, k: 3, weight: random(162, 3), bias: vec![0.0; 3], leaky: Some(0.0) }), + ], + ) + .unwrap(); + let x = random(4 * 3 * 3, 7); + agree(&NetRunner::new(&n, 3).unwrap().run(&x).unwrap(), &reference::run(&n, 3, &x).unwrap()); +} + +#[test] +fn reused_im2col_clears_all_edge_padding_between_images() { + for t in [1, 3, 5, 36] { + let n = net(t, 5, 0, 121); + let r = NetRunner::new(&n, t).unwrap(); + for x in [vec![1.0; 4 * t * t], vec![0.0; 4 * t * t], random(4 * t * t, 9)] { + agree(&r.run(&x).unwrap(), &reference::run(&n, t, &x).unwrap()); + } + } +} +#[test] +#[ignore = "needs LC_DENOISE_MODEL; one tile with LIGHTCRAFT_PROFILE=1"] +fn profile_one_real_tile() { + let model = std::env::var_os("LC_DENOISE_MODEL").expect("set LC_DENOISE_MODEL"); + let m = crate::known::find("rawnind-bayer").unwrap().manifest; + let r = crate::runtime::CpuRunner::load(std::path::Path::new(&model), &m).unwrap(); + println!("workspace {:.1} MB", r.bytes_per_tile() as f64 / 1e6); + r.run(&crate::runtime::check_tile(512)).unwrap(); +} diff --git a/crates/denoise/src/lib.rs b/crates/denoise/src/lib.rs new file mode 100644 index 00000000..e72860ba --- /dev/null +++ b/crates/denoise/src/lib.rs @@ -0,0 +1,37 @@ +//! AI denoise for LightCraft, as a step between the sensor and the demosaic that the rest of the develop +//! pipeline never sees: the raw file stays untouched, the denoised picture is cached data that can always be +//! made again, and the amount is an ordinary develop setting. +//! +//! - [`manifest`]: what a denoise model takes and gives, and who may use it (untrusted JSON, validated). +//! - [`known`] and [`archive`]: the models LightCraft can fetch, and taking the model file out of the `.zip` it comes in. +//! - [`bayer`]: Bayer layouts and the packing of a mosaic into the four planes a model takes. +//! - [`tiles`] and [`run`]: cutting a picture into overlapping tiles for a fixed-size model, bringing each tile's +//! output scale back to the input's, keeping clipped highlights as they were, and blending the tiles together. +//! The model itself is a [`run::TileRunner`], so all of this is tested without one. +//! - [`product`]: the cache file the denoised picture is kept in (half floats, compressed in strips) and the +//! blend with the plain demosaic that the Amount slider controls. +//! - [`cpu`] and [`runtime`]: a single-thread pure-Rust CPU runner and its installation self-test. +//! - [`net`] and [`onnx`]: checked plain network data and bounded ONNX decoding, shared by CPU and GPU. +//! +//! Everything here treats model files, manifests and cache files as hostile input: sizes are capped, numbers must +//! be finite, and a malformed file is an error, never a panic. +#![forbid(unsafe_code)] +#![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)] + +#[cfg(feature = "runtime")] +pub mod cpu; + +pub mod net; +pub mod onnx; +mod onnx_proto; + +pub mod reference; + +#[cfg(feature = "runtime")] +pub mod runtime; +#[doc(hidden)] +pub mod synthetic; + +pub use lightcraft_denoise_core::*; + +pub mod synthetic_proto; diff --git a/crates/denoise/src/net.rs b/crates/denoise/src/net.rs new file mode 100644 index 00000000..42241833 --- /dev/null +++ b/crates/denoise/src/net.rs @@ -0,0 +1,345 @@ +//! A denoise network as plain data: the handful of layer kinds a U-Net denoiser is made of, with their weights. +//! +//! [`onnx`](crate::onnx) reads the file into this description, shared by the pure-Rust CPU and GPU runners. +//! Unsupported operators, groups, dilation or activations are reported as errors when loading the model. +//! +//! The description is checked ([`Net::new`]) because it comes from a stranger's file: channel counts must chain, +//! weights must have the sizes the shapes need and be finite, and the depth must stay small. + +/// Most layers a network may have. +pub const MAX_OPS: usize = 512; +/// Most channels a layer may have. +pub const MAX_CHANNELS: usize = 4096; +/// Most weights (numbers) a network may have in all. +pub const MAX_WEIGHTS: usize = 128 * 1024 * 1024; + +#[derive(Debug, PartialEq, thiserror::Error)] +pub enum NetError { + /// The model uses something this denoise network description cannot express. + #[error("not supported by the denoise runners: {0}")] + Unsupported(String), + /// The model is not a well-formed network. + #[error("not a valid network: {0}")] + Invalid(String), +} + +fn invalid(why: impl Into) -> Result { + Err(NetError::Invalid(why.into())) +} + +/// A tensor is named by its number: 0 is the network's input, and layer `i` makes tensor `i + 1`. +pub type TensorId = usize; + +/// A convolution with stride 1 and "same" padding (`k` is 1 or 3), over one tensor or two joined along the channels +/// (`src` first), then a bias and optionally a leaky ReLU. +#[derive(Clone, Debug, PartialEq)] +pub struct Conv { + pub src: TensorId, + /// A second source, whose channels follow `src`'s (a skip connection's concatenation). + pub src2: Option, + pub cin: usize, + pub cout: usize, + pub k: usize, + /// `[cout][cin][k][k]`, as ONNX has it. + pub weight: Vec, + pub bias: Vec, + /// `Some(alpha)`: a leaky ReLU with this slope follows. + pub leaky: Option, +} + +/// A transposed convolution with a 2 × 2 kernel and stride 2: every input pixel becomes a 2 × 2 block. +#[derive(Clone, Debug, PartialEq)] +pub struct ConvT2 { + pub src: TensorId, + pub cin: usize, + pub cout: usize, + /// `[cin][cout][2][2]`, as ONNX has it. + pub weight: Vec, + pub bias: Vec, + pub leaky: Option, +} + +#[derive(Clone, Debug, PartialEq)] +pub enum Op { + Conv(Conv), + ConvT2(ConvT2), + /// 2 × 2 maximum, stride 2. + MaxPool2 { + src: TensorId, + }, + /// Depth-to-space with block size 2 in "CRD" order: output channel `c` at `(2y + dy, 2x + dx)` is input channel + /// `4c + 2dy + dx` at `(y, x)`. + DepthToSpace2 { + src: TensorId, + }, +} + +/// What a tensor looks like: its channels, and its resolution as a power of two below the input's (a negative +/// level is above it). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Shape { + pub channels: usize, + pub level: i32, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Net { + pub in_channels: usize, + ops: Vec, + shapes: Vec, +} + +impl Net { + /// Check `ops` (layer `i` makes tensor `i + 1`) over an input of `in_channels` and describe every tensor. + pub fn new(in_channels: usize, ops: Vec) -> Result { + if in_channels == 0 || in_channels > MAX_CHANNELS { + return invalid("the input has no channels or too many"); + } + if ops.is_empty() || ops.len() > MAX_OPS { + return invalid("the network has no layers or too many"); + } + let mut shapes = vec![Shape { channels: in_channels, level: 0 }]; + let mut weights = 0usize; + for (i, op) in ops.iter().enumerate() { + let shape_of = |t: TensorId| { + shapes + .get(t) + .copied() + .filter(|_| t <= i) + .ok_or_else(|| NetError::Invalid(format!("layer {i} reads tensor {t}, which does not exist yet"))) + }; + let out = match op { + Op::Conv(c) => { + let a = shape_of(c.src)?; + let (cin, level) = match c.src2 { + Some(s2) => { + let b = shape_of(s2)?; + if b.level != a.level { + return invalid(format!("layer {i} joins tensors of different resolutions")); + } + (a.channels + b.channels, a.level) + } + None => (a.channels, a.level), + }; + if c.cin != cin || c.cout == 0 || c.cout > MAX_CHANNELS || !(c.k == 1 || c.k == 3) { + return invalid(format!( + "layer {i}: a convolution of {} to {} channels with a {}-wide kernel does not fit", + c.cin, c.cout, c.k + )); + } + check_weights( + i, + &c.weight, + c.cout.checked_mul(cin).and_then(|n| n.checked_mul(c.k * c.k)), + &c.bias, + c.cout, + c.leaky, + &mut weights, + )?; + Shape { channels: c.cout, level } + } + Op::ConvT2(c) => { + let a = shape_of(c.src)?; + if c.cin != a.channels || c.cout == 0 || c.cout > MAX_CHANNELS { + return invalid(format!("layer {i}: a transposed convolution of {} to {} channels does not fit", c.cin, c.cout)); + } + check_weights(i, &c.weight, c.cin.checked_mul(c.cout).and_then(|n| n.checked_mul(4)), &c.bias, c.cout, c.leaky, &mut weights)?; + Shape { channels: c.cout, level: a.level - 1 } + } + Op::MaxPool2 { src } => { + let a = shape_of(*src)?; + Shape { channels: a.channels, level: a.level + 1 } + } + Op::DepthToSpace2 { src } => { + let a = shape_of(*src)?; + if !a.channels.is_multiple_of(4) { + return invalid(format!("layer {i}: depth-to-space needs a multiple of 4 channels")); + } + Shape { channels: a.channels / 4, level: a.level - 1 } + } + }; + shapes.push(out); + } + Ok(Net { in_channels, ops, shapes }) + } + + /// The layers: layer `i` makes tensor `i + 1`. + pub fn ops(&self) -> &[Op] { + &self.ops + } + + /// The shape of tensor `t` (0 is the input). + pub fn shape(&self, t: TensorId) -> Option { + self.shapes.get(t).copied() + } + + /// The tensor the network gives (the last layer's). + pub fn output(&self) -> TensorId { + self.ops.len() + } + + pub fn out_channels(&self) -> usize { + self.shapes.last().map_or(0, |s| s.channels) + } + + /// The deepest resolution level: a tile must be a multiple of `2^this` cells wide. + pub fn depth(&self) -> u32 { + self.shapes.iter().map(|s| s.level.max(0) as u32).max().unwrap_or(0) + } + + /// Can the network run on a `tile × tile` input and give `[out_channels, 2·tile, 2·tile]`? (Pooling must divide + /// the tile, and the output must be one level above the input.) + pub fn fits_tile(&self, tile: usize) -> bool { + let Some(multiple) = 1usize.checked_shl(self.depth()) else { return false }; + tile > 0 && tile.is_multiple_of(multiple) && self.shapes.last().is_some_and(|s| s.level == -1) + } + + /// Multiply-adds for one tile, for estimating time. + pub fn macs(&self, tile: usize) -> u64 { + let pixels = |level: i32| -> u64 { + let side = if level >= 0 { + tile.checked_shr(level as u32).unwrap_or(0) + } else { + tile.saturating_mul(1usize.checked_shl(level.unsigned_abs()).unwrap_or(usize::MAX)) + }; + (side as u64).saturating_mul(side as u64) + }; + self.ops + .iter() + .enumerate() + .map(|(i, op)| { + let out = self.shapes.get(i + 1).map_or(0, |s| pixels(s.level)); + match op { + Op::Conv(c) => out.saturating_mul((c.cin * c.cout * c.k * c.k) as u64), + // every input pixel makes a 2 × 2 block: as many multiply-adds as the 2 × 2-wide kernel has weights + Op::ConvT2(c) => pixels(self.shapes.get(c.src).map_or(0, |s| s.level)).saturating_mul((c.cin * c.cout * 4) as u64), + _ => 0, + } + }) + .fold(0u64, u64::saturating_add) + } +} + +fn check_weights( + layer: usize, + weight: &[f32], + want: Option, + bias: &[f32], + cout: usize, + leaky: Option, + total: &mut usize, +) -> Result<(), NetError> { + if want != Some(weight.len()) || bias.len() != cout { + return invalid(format!("layer {layer}: the weights are not the size its shape needs")); + } + *total = total.saturating_add(weight.len()).saturating_add(bias.len()); + if *total > MAX_WEIGHTS { + return invalid("the network has too many weights"); + } + if weight.iter().chain(bias).any(|v| !v.is_finite()) || leaky.is_some_and(|a| !a.is_finite()) { + return invalid(format!("layer {layer}: a weight is not a finite number")); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn conv(src: TensorId, src2: Option, cin: usize, cout: usize, k: usize) -> Op { + Op::Conv(Conv { src, src2, cin, cout, k, weight: vec![0.1; cout * cin * k * k], bias: vec![0.0; cout], leaky: Some(0.2) }) + } + + fn convt(src: TensorId, cin: usize, cout: usize) -> Op { + Op::ConvT2(ConvT2 { src, cin, cout, weight: vec![0.1; cin * cout * 4], bias: vec![0.0; cout], leaky: None }) + } + + /// A two-level network of the shape the denoisers have: 4 → 8 channels, pool, 16 channels, up, join, 12, and the + /// 1 × 1 convolution to the colours. + fn unet() -> Vec { + vec![ + conv(0, None, 4, 8, 3), // 1: 8 ch, level 0 + Op::MaxPool2 { src: 1 }, // 2: level 1 + conv(2, None, 8, 16, 3), // 3: 16 ch, level 1 + convt(3, 16, 8), // 4: 8 ch, level 0 + conv(4, Some(1), 16, 8, 3), // 5: join [up, skip] + conv(5, None, 8, 12, 1), // 6: 12 ch + Op::DepthToSpace2 { src: 6 }, // 7: 3 ch, level -1 + ] + } + + #[test] + fn deep_networks_and_huge_cost_queries_never_overflow_or_truncate_depth() { + let n = Net::new(4, unet()).unwrap(); + assert_eq!(n.macs(usize::MAX), u64::MAX); + assert_eq!(n.macs(0), 0); + let mut ops = Vec::new(); + for i in 0..21 { + ops.push(Op::MaxPool2 { src: i }); + } + for i in 21..42 { + ops.push(convt(i, 4, 4)); + } + ops.push(conv(42, None, 4, 12, 1)); + ops.push(Op::DepthToSpace2 { src: 43 }); + let n = Net::new(4, ops).unwrap(); + assert!(!n.fits_tile(1 << 20)); + assert!(n.fits_tile(1 << 21)); + let _ = n.macs(16); + } + + #[test] + fn a_u_net_is_described_with_its_shapes_and_costs() { + let net = Net::new(4, unet()).unwrap(); + assert_eq!(net.out_channels(), 3); + assert_eq!(net.shape(3), Some(Shape { channels: 16, level: 1 })); + assert_eq!(net.shape(7), Some(Shape { channels: 3, level: -1 })); + assert_eq!(net.depth(), 1); + assert!(net.fits_tile(64) && !net.fits_tile(63) && !net.fits_tile(0)); + // 64 × 64 cells: layer 1 = 4096·4·8·9, and so on; just check it is the sum of what the layers cost + let t = 64u64; + let want = t * t * 4 * 8 * 9 + (t / 2) * (t / 2) * 8 * 16 * 9 + (t / 2) * (t / 2) * 16 * 8 * 4 + t * t * 16 * 8 * 9 + t * t * 8 * 12; + assert_eq!(net.macs(64), want); + } + + #[test] + fn a_network_that_does_not_chain_is_refused() { + // wrong input channel count for the first convolution + assert!(Net::new(3, unet()).is_err()); + // a layer reading a tensor that does not exist yet + let mut ops = unet(); + ops[2] = conv(9, None, 8, 16, 3); + assert!(Net::new(4, ops).is_err()); + // joining tensors of different resolutions + let mut ops = unet(); + ops[4] = conv(3, Some(1), 24, 8, 3); + assert!(Net::new(4, ops).is_err()); + // weights of the wrong size, and not finite + let mut ops = unet(); + if let Op::Conv(c) = &mut ops[0] { + c.weight.pop(); + } + assert!(Net::new(4, ops).is_err()); + let mut ops = unet(); + if let Op::Conv(c) = &mut ops[0] { + c.bias[0] = f32::NAN; + } + assert!(Net::new(4, ops).is_err()); + // a kernel that is neither 1 nor 3 wide, no layers, depth-to-space on a channel count it cannot split + let mut ops = unet(); + if let Op::Conv(c) = &mut ops[0] { + c.k = 5; + } + assert!(Net::new(4, ops).is_err()); + assert!(Net::new(4, vec![]).is_err()); + assert!(Net::new(4, vec![Op::DepthToSpace2 { src: 0 }]).is_ok()); + assert!(Net::new(4, vec![Op::DepthToSpace2 { src: 0 }, Op::DepthToSpace2 { src: 1 }]).is_err()); + assert!(Net::new(2, vec![Op::DepthToSpace2 { src: 0 }]).is_err()); + } + + #[test] + fn a_network_that_does_not_end_one_level_up_does_not_fit_a_tile() { + let net = Net::new(4, unet()[..6].to_vec()).unwrap(); + assert!(!net.fits_tile(64), "no depth-to-space: the output is not at twice the resolution"); + } +} diff --git a/crates/denoise/src/onnx.rs b/crates/denoise/src/onnx.rs new file mode 100644 index 00000000..fc84a36c --- /dev/null +++ b/crates/denoise/src/onnx.rs @@ -0,0 +1,369 @@ +//! Reading an `.onnx` file into a [`Net`] with bounded, pure-Rust protobuf decoding. +//! +//! Only what the denoise U-Nets are made of is understood: `Conv` (3 × 3 "same" or 1 × 1, stride 1, no groups, no +//! dilation) optionally followed by `LeakyRelu`, `ConvTranspose` (2 × 2, stride 2) likewise, `MaxPool` (2 × 2, stride 2), +//! a two-input `Concat` on the channels that feeds a convolution, and a final `DepthToSpace` (2, `CRD`). Any other +//! operator or attribute is [`NetError::Unsupported`], which is reported to the user by both runners. + +use std::collections::HashMap; +use std::path::Path; + +use crate::onnx_proto::{self, AttributeProto, NodeProto, TensorProto, ValueInfoProto}; +use std::io::Read; + +use crate::net::{Conv, ConvT2, MAX_WEIGHTS, Net, NetError, Op, TensorId}; + +fn unsupported(why: impl Into) -> Result { + Err(NetError::Unsupported(why.into())) +} + +fn invalid(why: impl Into) -> Result { + Err(NetError::Invalid(why.into())) +} + +/// Read the model at `path`. +pub fn read(path: &Path) -> Result { + let file = std::fs::File::open(path).map_err(|e| NetError::Invalid(e.to_string()))?; + let size = usize::try_from(file.metadata().map_err(|e| NetError::Invalid(e.to_string()))?.len()) + .ok() + .filter(|&n| n <= onnx_proto::MAX_BYTES) + .ok_or_else(|| NetError::Invalid("model file too large".into()))?; + let mut bytes = Vec::new(); + bytes.try_reserve_exact(size).map_err(|_| NetError::Invalid("not enough memory to load the model".into()))?; + file.take(onnx_proto::MAX_BYTES as u64 + 1).read_to_end(&mut bytes).map_err(|e| NetError::Invalid(e.to_string()))?; + read_bytes(&bytes) +} + +fn read_bytes(bytes: &[u8]) -> Result { + let graph = onnx_proto::decode(bytes)?; + + let initialisers: HashMap<&str, &TensorProto> = graph.initializer.iter().map(|t| (t.name.as_str(), t)).collect(); + if initialisers.len() != graph.initializer.len() { + return invalid("duplicate initializer names"); + } + let inputs: Vec<_> = graph.input.iter().filter(|i| !initialisers.contains_key(i.name.as_str())).collect(); + let ([input], [output]) = (inputs.as_slice(), graph.output.as_slice()) else { return unsupported("not one input and one output") }; + let in_channels = match channel_dim(input) { + Some(c) => c, + None => return unsupported("the input's channel count is not written in the file"), + }; + + // how many nodes read each name (a graph output counts as a reader) + let mut readers: HashMap<&str, usize> = HashMap::new(); + for n in &graph.node { + for i in &n.input { + *readers.entry(i.as_str()).or_default() += 1; + } + } + *readers.entry(output.name.as_str()).or_default() += 1; + + let mut ids: HashMap = HashMap::new(); + ids.insert(input.name.clone(), 0); + // a Concat is not a layer of its own: the convolution that reads it takes both sources + let mut joins: HashMap = HashMap::new(); + let mut ops: Vec = Vec::new(); + let mut weights_total = 0usize; + + let nodes = &graph.node; + let mut i = 0; + while i < nodes.len() { + let Some(node) = nodes.get(i) else { break }; + i += 1; + if node.output.len() != 1 || node.output.first().is_none_or(|n| n.is_empty() || ids.contains_key(n) || joins.contains_key(n)) { + return invalid("a layer has missing, repeated or multiple outputs"); + } + let allowed: &[&str] = match node.op_type.as_str() { + "Conv" => &["kernel_shape", "pads", "strides", "dilations", "group", "auto_pad"], + "ConvTranspose" => &["kernel_shape", "pads", "strides", "dilations", "group", "output_padding", "auto_pad"], + "MaxPool" => &["kernel_shape", "pads", "strides", "dilations", "ceil_mode", "storage_order", "auto_pad"], + "Concat" => &["axis"], + "DepthToSpace" => &["blocksize", "mode"], + _ => &[], + }; + if node.attribute.iter().any(|a| !allowed.contains(&a.name.as_str())) { + return unsupported("an unknown operator attribute"); + } + let this = ops.len() + 1; + let lookup = |name: &str| -> Result { + ids.get(name).copied().ok_or_else(|| NetError::Invalid(format!("`{}` reads `{name}`, which nothing makes", node.op_type))) + }; + // a LeakyRelu straight after this layer, which nothing else reads, is part of it + let mut fuse_leaky = |out: &str| -> Result<(Option, String), NetError> { + if let Some(next) = nodes.get(i) + && next.op_type == "LeakyRelu" + && next.input.first().map(String::as_str) == Some(out) + && readers.get(out).copied() == Some(1) + { + if next.input.len() != 1 || next.output.len() != 1 || next.attribute.iter().any(|a| a.name != "alpha") { + return invalid("malformed LeakyRelu"); + } + let alpha = attr(next, "alpha").map_or(0.01, |a| a.f); + i += 1; + let name = next.output.first().cloned().ok_or_else(|| NetError::Invalid("LeakyRelu without an output".into()))?; + return Ok((Some(alpha), name)); + } + Ok((None, out.to_string())) + }; + match node.op_type.as_str() { + "Concat" => { + if attr(node, "axis").map(|a| a.i) != Some(1) { + return unsupported("Concat on another axis than the channels"); + } + let [a, b] = node.input.as_slice() else { return unsupported("Concat of other than two tensors") }; + let out = node.output.first().ok_or_else(|| NetError::Invalid("Concat without an output".into()))?; + if readers.get(out.as_str()).copied() != Some(1) { + return unsupported("a Concat that is read by more than one layer"); + } + joins.insert(out.clone(), (lookup(a)?, lookup(b)?)); + } + "Conv" => { + let (w, b) = weights_of(node, &initialisers, &mut weights_total)?; + let (cout, cin, kh, kw) = four(&w.0)?; + let b = if b.is_empty() { vec![0.0; cout] } else { b }; + if kh != kw || !(kh == 1 || kh == 3) { + return unsupported(format!("a {kh} × {kw} convolution")); + } + if ints_or(node, "strides", &[1, 1]) != [1, 1] + || ints_or(node, "dilations", &[1, 1]) != [1, 1] + || attr(node, "group").is_some_and(|a| a.i != 1) + { + return unsupported("a convolution with strides, dilation or groups"); + } + if attr(node, "kernel_shape").is_some_and(|a| a.ints != [kh as i64, kw as i64]) { + return invalid("kernel shape does not match the weights"); + } + let pad = (kh as i64 - 1) / 2; + if ints_or(node, "pads", &[0; 4]) != [pad; 4] || attr(node, "auto_pad").is_some_and(|a| !a.s.is_empty() && a.s != b"NOTSET") { + return unsupported("a convolution that is not padded to keep its size"); + } + let src_name = node.input.first().ok_or_else(|| NetError::Invalid("Conv without an input".into()))?; + let (src, src2) = match joins.remove(src_name) { + Some((a, b)) => (a, Some(b)), + None => (lookup(src_name)?, None), + }; + let out = node.output.first().ok_or_else(|| NetError::Invalid("Conv without an output".into()))?; + let (leaky, out_name) = fuse_leaky(out)?; + ids.insert(out_name, this); + ops.push(Op::Conv(Conv { src, src2, cin, cout, k: kh, weight: w.1, bias: b, leaky })); + } + "ConvTranspose" => { + let (w, b) = weights_of(node, &initialisers, &mut weights_total)?; + let (cin, cout, kh, kw) = four(&w.0)?; + let b = if b.is_empty() { vec![0.0; cout] } else { b }; + if (kh, kw) != (2, 2) + || ints_or(node, "strides", &[1, 1]) != [2, 2] + || ints_or(node, "pads", &[0; 4]) != [0; 4] + || ints_or(node, "dilations", &[1, 1]) != [1, 1] + || attr(node, "group").is_some_and(|a| a.i != 1) + || attr(node, "output_padding").is_some_and(|a| a.ints.iter().any(|&v| v != 0)) + { + return unsupported("a transposed convolution other than 2 × 2 with stride 2"); + } + if attr(node, "auto_pad").is_some_and(|a| !a.s.is_empty() && a.s != b"NOTSET") { + return unsupported("transposed convolution auto padding"); + } + if attr(node, "kernel_shape").is_some_and(|a| a.ints != [2, 2]) { + return invalid("transposed kernel shape does not match the weights"); + } + let src = lookup(node.input.first().map_or("", String::as_str))?; + let out = node.output.first().ok_or_else(|| NetError::Invalid("ConvTranspose without an output".into()))?; + let (leaky, out_name) = fuse_leaky(out)?; + ids.insert(out_name, this); + ops.push(Op::ConvT2(ConvT2 { src, cin, cout, weight: w.1, bias: b, leaky })); + } + "MaxPool" => { + if ints_or(node, "kernel_shape", &[]) != [2, 2] + || ints_or(node, "strides", &[1, 1]) != [2, 2] + || ints_or(node, "pads", &[0; 4]) != [0; 4] + || ints_or(node, "dilations", &[1, 1]) != [1, 1] + || attr(node, "ceil_mode").is_some_and(|a| a.i != 0) + || attr(node, "storage_order").is_some_and(|a| a.i != 0) + || attr(node, "auto_pad").is_some_and(|a| !a.s.is_empty() && a.s != b"NOTSET") + { + return unsupported("a max-pool other than 2 × 2 with stride 2"); + } + let src = lookup(node.input.first().map_or("", String::as_str))?; + let out = node.output.first().ok_or_else(|| NetError::Invalid("MaxPool without an output".into()))?; + ids.insert(out.clone(), this); + ops.push(Op::MaxPool2 { src }); + } + "DepthToSpace" => { + if attr(node, "blocksize").map(|a| a.i) != Some(2) || attr(node, "mode").map(|a| a.s.as_slice()) != Some(&b"CRD"[..]) { + return unsupported("depth-to-space other than block size 2 in CRD order"); + } + let src = lookup(node.input.first().map_or("", String::as_str))?; + let out = node.output.first().ok_or_else(|| NetError::Invalid("DepthToSpace without an output".into()))?; + ids.insert(out.clone(), this); + ops.push(Op::DepthToSpace2 { src }); + } + other => return unsupported(format!("the `{other}` operator")), + } + } + if !joins.is_empty() { + return unsupported("a Concat whose result no convolution reads"); + } + if ids.get(output.name.as_str()) != Some(&ops.len()) { + return unsupported("the output is not the last layer"); + } + Net::new(in_channels, ops) +} + +fn attr<'a>(n: &'a NodeProto, name: &str) -> Option<&'a AttributeProto> { + n.attribute.iter().find(|a| a.name == name) +} + +/// The integer list attribute `name`, or `default` when the node does not have it. +fn ints_or(n: &NodeProto, name: &str, default: &[i64]) -> Vec { + attr(n, name).map_or_else(|| default.to_vec(), |a| a.ints.clone()) +} + +fn four(dims: &[i64]) -> Result<(usize, usize, usize, usize), NetError> { + match dims { + [a, b, c, d] if dims.iter().all(|&v| v > 0 && v <= 1 << 20) => Ok((*a as usize, *b as usize, *c as usize, *d as usize)), + _ => invalid("a weight tensor that is not four-dimensional"), + } +} + +type Weights = ((Vec, Vec), Vec); + +/// The weight and bias initialisers of a convolution node. +fn weights_of(node: &NodeProto, initialisers: &HashMap<&str, &TensorProto>, total: &mut usize) -> Result { + let mut get = |idx: usize| -> Result<(Vec, Vec), NetError> { + let name = node.input.get(idx).ok_or_else(|| NetError::Invalid(format!("`{}` without its weights", node.op_type)))?; + let t = initialisers.get(name.as_str()).ok_or_else(|| NetError::Unsupported(format!("weights that are computed (`{name}`)")))?; + let values = floats(t)?; + *total = total.saturating_add(values.len()); + if *total > MAX_WEIGHTS { + return invalid("the network has too many weights"); + } + Ok((t.dims.clone(), values)) + }; + let w = get(1)?; + let b = if node.input.len() > 2 { + let (dims, values) = get(2)?; + if dims.len() != 1 || values.is_empty() { + return invalid("a bias tensor must be a nonempty vector"); + } + values + } else { + Vec::new() + }; + Ok((w, b)) +} + +fn floats(t: &TensorProto) -> Result, NetError> { + if t.data_type != 1 { + return unsupported("weights that are not 32-bit floats"); + } + let want: usize = t + .dims + .iter() + .try_fold(1usize, |n, &d| usize::try_from(d).ok().and_then(|d| n.checked_mul(d))) + .ok_or_else(|| NetError::Invalid("a tensor with a negative or huge size".into()))?; + if want > MAX_WEIGHTS { + return invalid("a tensor with too many numbers"); + } + if !t.raw_data.is_empty() && (!t.float_data.is_empty() || !t.raw_data.len().is_multiple_of(4)) { + return invalid("ambiguous or truncated tensor data"); + } + let actual = if t.raw_data.is_empty() { t.float_data.len() } else { t.raw_data.len() / 4 }; + if actual != want { + return invalid("tensor data does not match its dimensions"); + } + let mut values = Vec::new(); + values.try_reserve_exact(want).map_err(|_| NetError::Invalid("not enough memory for weights".into()))?; + if !t.raw_data.is_empty() { + values.extend(t.raw_data.as_chunks::<4>().0.iter().map(|b| f32::from_le_bytes(*b))); + } else { + values.extend_from_slice(&t.float_data); + } + if values.len() != want { + return invalid(format!("tensor `{}` has {} numbers, not the {want} its shape says", t.name, values.len())); + } + Ok(values) +} + +/// The channel count in an input's declared shape `[batch, channels, …]`. +fn channel_dim(v: &ValueInfoProto) -> Option { + v.channels +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::synthetic::unet_onnx; + + fn temp_model(name: &str, bytes: &[u8]) -> std::path::PathBuf { + let p = std::env::temp_dir().join(format!("lc-denoise-onnx-{name}-{}.onnx", std::process::id())); + std::fs::write(&p, bytes).unwrap(); + p + } + + #[test] + fn a_u_net_is_read_with_its_layers_fused() { + let (tile, depth) = (32usize, 2usize); + let p = temp_model("unet", &unet_onnx(tile as u64, 8, depth, 7)); + let net = read(&p).unwrap(); + // per level two convolutions; a pool per level but the last; going up a transposed convolution and two convolutions + // per level (the join is part of the first); then the 1 × 1 convolution and depth-to-space + assert_eq!(net.ops().len(), 2 * (depth + 1) + depth + depth * 3 + 2); + assert!(matches!(net.ops().last(), Some(crate::net::Op::DepthToSpace2 { .. }))); + assert!(net.ops().iter().any(|o| matches!(o, crate::net::Op::Conv(c) if c.src2.is_some())), "the joins became two-source convolutions"); + assert!( + net.ops().iter().all(|o| !matches!(o, crate::net::Op::Conv(c) if c.leaky.is_none() && c.k == 3)), + "every 3 × 3 convolution has its activation" + ); + assert_eq!((net.in_channels, net.out_channels(), net.depth()), (4, 3, 2)); + assert!(net.fits_tile(tile) && !net.fits_tile(30)); + + let _ = std::fs::remove_file(p); + } + + #[test] + fn an_explicit_empty_or_nonvector_bias_is_not_treated_as_missing() { + let w = TensorProto { name: "w".into(), dims: vec![12, 4, 1, 1], data_type: 1, float_data: vec![0.1; 48], ..Default::default() }; + let n = NodeProto { input: vec!["data".into(), "w".into(), "b".into()], ..Default::default() }; + for dims in [vec![0], vec![3, 4]] { + let b = TensorProto { + name: "b".into(), + float_data: vec![0.0; dims.iter().product::() as usize], + dims, + data_type: 1, + ..Default::default() + }; + let inits = HashMap::from([("w", &w), ("b", &b)]); + assert!(weights_of(&n, &inits, &mut 0).is_err()); + } + } + + #[test] + fn truncated_and_mutated_protobuf_is_bounded_and_never_panics() { + let good = unet_onnx(4, 3, 1, 7); + for n in (0..good.len()).step_by(31) { + let _ = read_bytes(&good[..n]); + } + for at in (0..good.len()).step_by(37) { + let mut bad = good.clone(); + bad[at] ^= 0xff; + let _ = read_bytes(&bad); + } + assert!(read_bytes(&[0x3a, 0xff, 0xff, 0xff, 0xff, 0x7f]).is_err()); + assert!(read_bytes(&[0x3a, 0, 0x3a, 0]).is_err()); + } + proptest::proptest! { + #[test] + fn arbitrary_model_bytes_never_panic(bytes in proptest::collection::vec(proptest::prelude::any::(), 0..4096)) { + let _ = read_bytes(&bytes); + } + } + + #[test] + fn what_the_description_cannot_say_is_unsupported_and_garbage_is_an_error() { + // not a model + let p = temp_model("junk", b"this is not an onnx file"); + assert!(matches!(read(&p), Err(NetError::Invalid(_)) | Err(NetError::Unsupported(_)))); + let _ = std::fs::remove_file(p); + // a missing file + assert!(read(Path::new("/does/not/exist.onnx")).is_err()); + } +} diff --git a/crates/denoise/src/onnx_proto.rs b/crates/denoise/src/onnx_proto.rs new file mode 100644 index 00000000..d8bbba41 --- /dev/null +++ b/crates/denoise/src/onnx_proto.rs @@ -0,0 +1,295 @@ +//! Bounded, borrowed decoding of the ONNX fields used by denoise, from the public Apache-2.0 schema: +//! https://github.com/onnx/onnx/blob/main/onnx/onnx.proto3 +//! No generated code, protoc, recursive graph decoding, or external tensor files. + +use crate::net::{MAX_OPS, MAX_WEIGHTS, NetError}; + +const MAX_ITEMS: usize = MAX_OPS * 4; +const MAX_STRING: usize = 4096; +pub(super) const MAX_BYTES: usize = MAX_WEIGHTS * 4 + 16 * 1024 * 1024; + +type Result = std::result::Result; +fn bad() -> NetError { + NetError::Invalid("malformed or oversized ONNX protobuf".into()) +} + +struct Fields<'a> { + rest: &'a [u8], + count: usize, +} +enum Wire<'a> { + Int(u64), + Bytes(&'a [u8]), + F32(f32), + Other, +} + +fn varint(rest: &mut &[u8]) -> Result { + let mut value = 0u64; + for shift in (0..70).step_by(7) { + let (b, tail) = rest.split_first().ok_or_else(bad)?; + *rest = tail; + if shift == 63 && *b > 1 { + return Err(bad()); + } + value |= u64::from(*b & 127) << shift; + if b & 128 == 0 { + return Ok(value); + } + } + Err(bad()) +} + +impl<'a> Fields<'a> { + fn new(rest: &'a [u8]) -> Self { + Self { rest, count: 0 } + } + fn take(&mut self, n: usize) -> Result<&'a [u8]> { + let out = self.rest.get(..n).ok_or_else(bad)?; + self.rest = self.rest.get(n..).ok_or_else(bad)?; + Ok(out) + } + fn next(&mut self) -> Result)>> { + if self.rest.is_empty() { + return Ok(None); + } + self.count += 1; + if self.count > MAX_WEIGHTS { + return Err(bad()); + } + let tag = varint(&mut self.rest)?; + let number = u32::try_from(tag >> 3).map_err(|_| bad())?; + if number == 0 || number >= 1 << 29 { + return Err(bad()); + } + let wire = match tag & 7 { + 0 => Wire::Int(varint(&mut self.rest)?), + 1 => { + self.take(8)?; + Wire::Other + } + 2 => { + let n = usize::try_from(varint(&mut self.rest)?).map_err(|_| bad())?; + Wire::Bytes(self.take(n)?) + } + 5 => { + let b: [u8; 4] = self.take(4)?.try_into().map_err(|_| bad())?; + Wire::F32(f32::from_le_bytes(b)) + } + _ => return Err(bad()), + }; + Ok(Some((number, wire))) + } +} + +fn string(b: &[u8]) -> Result { + if b.len() > MAX_STRING { + return Err(bad()); + } + std::str::from_utf8(b).map(str::to_owned).map_err(|_| bad()) +} +fn push(v: &mut Vec, value: T, cap: usize) -> Result<()> { + if v.len() >= cap { + return Err(bad()); + } + v.try_reserve(1).map_err(|_| bad())?; + v.push(value); + Ok(()) +} +fn ints(w: Wire<'_>, out: &mut Vec, cap: usize) -> Result<()> { + match w { + Wire::Int(v) => push(out, v as i64, cap), + Wire::Bytes(mut b) => { + while !b.is_empty() { + let v = varint(&mut b)?; + push(out, v as i64, cap)?; + } + Ok(()) + } + _ => Err(bad()), + } +} + +#[derive(Default)] +pub(super) struct AttributeProto { + pub name: String, + pub f: f32, + pub i: i64, + pub s: Vec, + pub ints: Vec, +} +#[derive(Default)] +pub(super) struct NodeProto { + pub input: Vec, + pub output: Vec, + pub op_type: String, + pub attribute: Vec, +} +#[derive(Default)] +pub(super) struct TensorProto<'a> { + pub name: String, + pub dims: Vec, + pub data_type: i32, + pub raw_data: &'a [u8], + pub float_data: Vec, +} +pub(super) struct ValueInfoProto { + pub name: String, + pub channels: Option, +} +#[derive(Default)] +pub(super) struct Graph<'a> { + pub node: Vec, + pub initializer: Vec>, + pub input: Vec, + pub output: Vec, +} + +fn attribute(b: &[u8]) -> Result { + let mut a = AttributeProto::default(); + let mut fields = Fields::new(b); + while let Some((n, w)) = fields.next()? { + match (n, w) { + (1, Wire::Bytes(b)) => a.name = string(b)?, + (2, Wire::F32(v)) => a.f = v, + (3, Wire::Int(v)) => a.i = v as i64, + (4, Wire::Bytes(b)) => { + if b.len() > MAX_STRING { + return Err(bad()); + } + a.s = b.to_vec(); + } + (8, w) => ints(w, &mut a.ints, 16)?, + (20, Wire::Int(1 | 2 | 3 | 7)) => {} + _ => return Err(NetError::Unsupported("unknown attribute value".into())), + } + } + Ok(a) +} +fn node(b: &[u8]) -> Result { + let mut a = NodeProto::default(); + let mut fields = Fields::new(b); + while let Some((n, w)) = fields.next()? { + match (n, w) { + (1, Wire::Bytes(b)) => push(&mut a.input, string(b)?, 4)?, + (2, Wire::Bytes(b)) => push(&mut a.output, string(b)?, 2)?, + (4, Wire::Bytes(b)) => a.op_type = string(b)?, + (5, Wire::Bytes(b)) => push(&mut a.attribute, attribute(b)?, 16)?, + (7, Wire::Bytes(b)) if !b.is_empty() => return Err(NetError::Unsupported("a custom operator domain".into())), + _ => {} + } + } + Ok(a) +} +fn tensor(b: &[u8]) -> Result> { + let mut a = TensorProto::default(); + let mut fields = Fields::new(b); + while let Some((n, w)) = fields.next()? { + match (n, w) { + (1, w) => ints(w, &mut a.dims, 8)?, + (2, Wire::Int(v)) => a.data_type = i32::try_from(v).map_err(|_| bad())?, + (4, Wire::F32(v)) => push(&mut a.float_data, v, MAX_WEIGHTS)?, + (4, Wire::Bytes(b)) => { + if !b.len().is_multiple_of(4) { + return Err(bad()); + } + for chunk in b.as_chunks::<4>().0 { + push(&mut a.float_data, f32::from_le_bytes(*chunk), MAX_WEIGHTS)?; + } + } + (8, Wire::Bytes(b)) => a.name = string(b)?, + (9, Wire::Bytes(b)) => { + if b.len() > MAX_WEIGHTS * 4 { + return Err(bad()); + } + a.raw_data = b; + } + (13, _) | (14, Wire::Int(1..)) => return Err(NetError::Unsupported("external tensor data".into())), + _ => {} + } + } + Ok(a) +} +fn channel_shape(b: &[u8]) -> Result> { + let mut fields = Fields::new(b); + let mut dim = 0; + let mut channels = None; + while let Some((n, w)) = fields.next()? { + if let (1, Wire::Bytes(b)) = (n, w) { + if dim >= 8 { + return Err(bad()); + } + let mut df = Fields::new(b); + while let Some((n, w)) = df.next()? { + if let (1, Wire::Int(v)) = (n, w) + && dim == 1 + { + channels = usize::try_from(v).ok().filter(|&v| v > 0); + } + } + dim += 1; + } + } + Ok(channels) +} +fn channel_type(b: &[u8]) -> Result> { + let mut fields = Fields::new(b); + while let Some((n, w)) = fields.next()? { + if let (1, Wire::Bytes(b)) = (n, w) { + let mut tf = Fields::new(b); + let (mut elem, mut channels) = (0, None); + while let Some((n, w)) = tf.next()? { + match (n, w) { + (1, Wire::Int(v)) => elem = v, + (2, Wire::Bytes(b)) => channels = channel_shape(b)?, + _ => {} + } + } + return Ok(if elem == 1 { channels } else { None }); + } + } + Ok(None) +} +fn value_info(b: &[u8]) -> Result { + let mut out = ValueInfoProto { name: String::new(), channels: None }; + let mut fields = Fields::new(b); + while let Some((n, w)) = fields.next()? { + match (n, w) { + (1, Wire::Bytes(b)) => out.name = string(b)?, + (2, Wire::Bytes(b)) => out.channels = channel_type(b)?, + _ => {} + } + } + Ok(out) +} +fn graph(b: &[u8]) -> Result> { + let mut out = Graph::default(); + let mut fields = Fields::new(b); + while let Some((n, w)) = fields.next()? { + match (n, w) { + (1, Wire::Bytes(b)) => push(&mut out.node, node(b)?, MAX_ITEMS)?, + (5, Wire::Bytes(b)) => push(&mut out.initializer, tensor(b)?, MAX_ITEMS)?, + (11, Wire::Bytes(b)) => push(&mut out.input, value_info(b)?, MAX_ITEMS)?, + (12, Wire::Bytes(b)) => push(&mut out.output, value_info(b)?, 1)?, + (15, _) => return Err(NetError::Unsupported("sparse initializers".into())), + _ => {} + } + } + Ok(out) +} +pub(super) fn decode(b: &[u8]) -> Result> { + if b.len() > MAX_BYTES { + return Err(bad()); + } + let mut fields = Fields::new(b); + let mut out = None; + while let Some((n, w)) = fields.next()? { + if let (7, Wire::Bytes(b)) = (n, w) { + if out.is_some() { + return Err(bad()); + } + out = Some(graph(b)?); + } + } + out.ok_or_else(bad) +} diff --git a/crates/denoise/src/reference.rs b/crates/denoise/src/reference.rs new file mode 100644 index 00000000..ef58b5b9 --- /dev/null +++ b/crates/denoise/src/reference.rs @@ -0,0 +1,133 @@ +//! A plain-loop interpreter of a [`Net`], slow and obviously right: the reference the parser and the GPU runner are +//! tested against (tract's answer on the real model is the other). It is not used on photos. + +use crate::net::{Net, NetError, Op, TensorId}; +use crate::run::{Error, TileRunner}; + +/// A tensor in `[channels][side][side]` order. +struct T { + channels: usize, + side: usize, + data: Vec, +} + +/// Run `net` on `input` (`in_channels` planes of `tile × tile`). +pub fn run(net: &Net, tile: usize, input: &[f32]) -> Result, NetError> { + if !net.fits_tile(tile) || input.len() != net.in_channels * tile * tile { + return Err(NetError::Invalid("the tile does not fit the network".into())); + } + let mut tensors: Vec = vec![T { channels: net.in_channels, side: tile, data: input.to_vec() }]; + let get = |tensors: &[T], id: TensorId| -> Result { + if id < tensors.len() { Ok(id) } else { Err(NetError::Invalid("a layer reads a tensor that does not exist yet".into())) } + }; + for op in net.ops() { + let out = match op { + Op::Conv(c) => { + let a = get(&tensors, c.src)?; + let b = c.src2.map(|s| get(&tensors, s)).transpose()?; + let side = tensors[a].side; + let pad = (c.k - 1) / 2; + let mut data = vec![0f32; c.cout * side * side]; + // the input channels in order: the first source's, then the second's + let planes: Vec<&[f32]> = + std::iter::once(&tensors[a]).chain(b.map(|b| &tensors[b])).flat_map(|t| t.data.chunks(side * side)).collect(); + for co in 0..c.cout { + for y in 0..side { + for x in 0..side { + let mut sum = c.bias[co]; + for (ci, plane) in planes.iter().enumerate() { + for ky in 0..c.k { + for kx in 0..c.k { + let (yy, xx) = ((y + ky) as isize - pad as isize, (x + kx) as isize - pad as isize); + if yy < 0 || xx < 0 || yy >= side as isize || xx >= side as isize { + continue; + } + sum += plane[yy as usize * side + xx as usize] * c.weight[((co * c.cin + ci) * c.k + ky) * c.k + kx]; + } + } + } + data[(co * side + y) * side + x] = leaky(sum, c.leaky); + } + } + } + T { channels: c.cout, side, data } + } + Op::ConvT2(c) => { + let a = get(&tensors, c.src)?; + let side = tensors[a].side; + let out_side = 2 * side; + let mut data = vec![0f32; c.cout * out_side * out_side]; + for co in 0..c.cout { + for y in 0..side { + for x in 0..side { + for dy in 0..2 { + for dx in 0..2 { + let mut sum = c.bias[co]; + for ci in 0..c.cin { + sum += tensors[a].data[(ci * side + y) * side + x] * c.weight[((ci * c.cout + co) * 2 + dy) * 2 + dx]; + } + data[(co * out_side + 2 * y + dy) * out_side + 2 * x + dx] = leaky(sum, c.leaky); + } + } + } + } + } + T { channels: c.cout, side: out_side, data } + } + Op::MaxPool2 { src } => { + let a = get(&tensors, *src)?; + let (channels, side) = (tensors[a].channels, tensors[a].side / 2); + let mut data = vec![0f32; channels * side * side]; + for ch in 0..channels { + for y in 0..side { + for x in 0..side { + let at = |dy: usize, dx: usize| tensors[a].data[(ch * 2 * side + 2 * y + dy) * 2 * side + 2 * x + dx]; + data[(ch * side + y) * side + x] = at(0, 0).max(at(0, 1)).max(at(1, 0)).max(at(1, 1)); + } + } + } + T { channels, side, data } + } + Op::DepthToSpace2 { src } => { + let a = get(&tensors, *src)?; + let (channels, side) = (tensors[a].channels / 4, tensors[a].side); + let out_side = 2 * side; + let mut data = vec![0f32; channels * out_side * out_side]; + for ch in 0..channels { + for y in 0..side { + for x in 0..side { + for dy in 0..2 { + for dx in 0..2 { + data[(ch * out_side + 2 * y + dy) * out_side + 2 * x + dx] = + tensors[a].data[((4 * ch + 2 * dy + dx) * side + y) * side + x]; + } + } + } + } + } + T { channels, side: out_side, data } + } + }; + tensors.push(out); + } + tensors.pop().map(|t| t.data).ok_or_else(|| NetError::Invalid("no layers".into())) +} + +fn leaky(v: f32, alpha: Option) -> f32 { + match alpha { + Some(a) if v < 0.0 => v * a, + _ => v, + } +} + +/// The reference as a [`TileRunner`] (for tests of the code around the model). +pub struct ReferenceRunner { + pub net: Net, + pub tile: usize, +} + +impl TileRunner for ReferenceRunner { + fn run(&self, input: &[f32]) -> Result, Error> { + run(&self.net, self.tile, input).map_err(|e| Error::Runtime(e.to_string())) + } +} diff --git a/crates/denoise/src/runtime.rs b/crates/denoise/src/runtime.rs new file mode 100644 index 00000000..117f6775 --- /dev/null +++ b/crates/denoise/src/runtime.rs @@ -0,0 +1,433 @@ +//! Model loading, a shared installation self-test, and deterministic test tiles. +//! `CpuRunner` is the pure-Rust executor; scalar operator tests live in `reference`. + +use std::path::Path; + +use web_time::Instant; + +use serde::Serialize; + +use crate::manifest::{DenoiserManifest, Gain}; +use crate::run::{Error, TileRunner}; + +/// 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 run one tile. + pub tile_ms: f64, + /// Megapixels of picture one tile covers, for estimating a whole photo's time. + pub tile_megapixels: f64, + /// The model's output scale relative to its input on the test picture (what `Gain::MatchMean` takes out). + pub scale: f64, + /// Each check: what was checked and whether it held. + pub checks: Vec<(String, bool)>, +} + +/// The default pure-Rust CPU model, with the same manifest and installation self-test as the reference. +#[derive(Clone)] +pub struct CpuRunner { + runner: crate::cpu::NetRunner, + manifest: DenoiserManifest, + load_ms: f64, +} +impl CpuRunner { + pub fn load(path: &Path, manifest: &DenoiserManifest) -> Result { + crate::manifest::validate(manifest).map_err(|e| Error::Load(e.to_string()))?; + let started = Instant::now(); + let net = crate::onnx::read(path).map_err(|e| Error::Load(e.to_string()))?; + let runner = crate::cpu::NetRunner::new(&net, manifest.tile as usize).map_err(|e| Error::Load(e.to_string()))?; + Ok(Self { runner, manifest: manifest.clone(), load_ms: started.elapsed().as_secs_f64() * 1e3 }) + } + pub fn manifest(&self) -> &DenoiserManifest { + &self.manifest + } + pub fn net(&self) -> &crate::net::Net { + self.runner.net() + } + pub fn bytes_per_tile(&self) -> usize { + self.runner.bytes_per_tile() + } + pub fn self_test(&self) -> SelfTest { + self_test(self, &self.manifest, self.load_ms) + } +} +impl TileRunner for CpuRunner { + fn run(&self, input: &[f32]) -> Result, Error> { + self.runner.run(input) + } +} +fn self_test(runner: &dyn TileRunner, manifest: &DenoiserManifest, load_ms: f64) -> SelfTest { + let t = manifest.tile as usize; + let mut checks: Vec<(String, bool)> = Vec::new(); + let (input, base) = test_tile(t); + let (mut times, mut scale) = (Vec::new(), 0.0f64); + let mut run = || { + let started = Instant::now(); + let r = runner.run(&input); + times.push(started.elapsed().as_secs_f64() * 1000.0); + r + }; + let (a, b) = (run(), run()); + match (&a, &b) { + (Ok(a), Ok(b)) => { + checks.push((format!("gives {} finite numbers per tile", a.len()), true)); + checks.push(( + "gives the same answer for the same tile".into(), + a.iter().zip(b).all(|(x, y)| (x - y).abs() <= 1e-4 * x.abs().max(y.abs()).max(1e-9)), + )); + let side = 2 * t; + let mean_in = base.iter().map(|&v| f64::from(v)).sum::() / base.len().max(1) as f64; + let mean_out = a.iter().map(|&v| f64::from(v)).sum::() / a.len().max(1) as f64; + scale = if mean_in > 0.0 { mean_out / mean_in } else { 0.0 }; + if let Gain::MatchMean { nominal, max_deviation } = manifest.gain { + let ratio = scale / f64::from(nominal); + let tolerance = f64::from(max_deviation).max(0.05) * 4.0; + checks.push((format!("its scale is the {nominal:.3e} the description says (found {scale:.3e})"), (ratio - 1.0).abs() <= tolerance)); + } + // noise left after bringing the answer back to the input's scale, against the noise that went in + let g = if scale > 0.0 { 1.0 / scale } else { 1.0 }; + let (mut noise_in, mut noise_out, mut n_in, mut n_out) = (0.0f64, 0.0f64, 0usize, 0usize); + for (i, (&v, &s)) in input.iter().zip(base.iter().cycle()).enumerate() { + // centre of each plane only: edges are the model's guess + let (x, y) = ((i % (t * t)) % t, (i % (t * t)) / t); + if x >= t / 4 && x < 3 * t / 4 && y >= t / 4 && y < 3 * t / 4 { + noise_in += f64::from(v - s).powi(2); + n_in += 1; + } + } + for c in 0..3 { + for y in side / 4..3 * side / 4 { + for x in side / 4..3 * side / 4 { + if let (Some(&o), Some(&s)) = (a.get(c * side * side + y * side + x), base.get((y / 2) * t + x / 2)) { + noise_out += (f64::from(o) * g - f64::from(s)).powi(2); + n_out += 1; + } + } + } + } + let (rms_in, rms_out) = ((noise_in / n_in.max(1) as f64).sqrt(), (noise_out / n_out.max(1) as f64).sqrt()); + checks.push((format!("takes noise down (from {rms_in:.4} to {rms_out:.4})"), rms_out < 0.8 * rms_in)); + } + (Err(e), _) | (_, Err(e)) => checks.push((format!("runs: {e}"), false)), + } + times.sort_by(|x, y| x.total_cmp(y)); + SelfTest { + ok: !checks.is_empty() && checks.iter().all(|(_, ok)| *ok), + load_ms, + tile_ms: times.get(times.len() / 2).copied().unwrap_or(0.0), + tile_megapixels: (2 * t * 2 * t) as f64 / 1e6, + scale, + checks, + } +} + +/// A deterministic test tile: four planes of a smooth gradient with noise on top, and the gradient alone for one +/// plane (`tile × tile`). +pub fn test_tile(t: usize) -> (Vec, Vec) { + if t == 0 || t > 2048 { + return (Vec::new(), Vec::new()); + } + let mut seed = 0x2545_f491u32; + let mut next = move || { + seed ^= seed << 13; + seed ^= seed >> 17; + seed ^= seed << 5; + (seed >> 8) as f32 / (1u32 << 24) as f32 + }; + let base: Vec = (0..t * t).map(|i| 0.12 + 0.3 * (((i % t) + (i / t)) as f32 / (2 * t) as f32)).collect(); + let input = (0..4 * t * t) + .map(|i| { + // sum of four uniforms: close to Gaussian, deviation 0.0577 × 2 before scaling + let n = next() + next() + next() + next() - 2.0; + base.get(i % (t * t)).copied().unwrap_or(0.0) + 0.05 * n + }) + .collect(); + (input, base) +} + +/// A tile for checking that two runners of one model agree: the whole range a raw photo has rather than the test tile's +/// mid-greys — a noisy ramp from black to white across the tile, a flat black corner and a clipped white one (four +/// planes of `tile × tile`). +pub fn check_tile(t: usize) -> Vec { + if t == 0 || t > 2048 { + return Vec::new(); + } + let mut seed = 0x9e37_79b9u32; + let mut next = move || { + seed ^= seed << 13; + seed ^= seed >> 17; + seed ^= seed << 5; + (seed >> 8) as f32 / (1u32 << 24) as f32 + }; + let (edge, span) = ((t / 8).max(1), (2 * t).saturating_sub(2).max(1) as f32); + (0..4 * t * t) + .map(|i| { + let (x, y) = ((i % (t * t)) % t, (i % (t * t)) / t); + let n = next() + next() + next() + next() - 2.0; + if x < edge && y < edge { + 0.0 + } else if x + edge >= t && y + edge >= t { + 1.0 + } else { + ((x + y) as f32 / span + 0.05 * n).clamp(0.0, 1.0) + } + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::licence::Licence; + use crate::manifest::Domain; + use crate::synthetic_proto::{len_field, value_info_bytes}; + + fn varint(mut v: u64, out: &mut Vec) { + loop { + let b = (v & 0x7f) as u8; + v >>= 7; + if v == 0 { + out.push(b); + return; + } + out.push(b | 0x80); + } + } + + fn varint_field(n: u64, v: u64, out: &mut Vec) { + varint(n << 3, out); + varint(v, out); + } + + fn node(op: &str, inputs: &[&str], output: &str, ints: &[(&str, i64)]) -> Vec { + let mut n = Vec::new(); + for i in inputs { + len_field(1, i.as_bytes(), &mut n); + } + len_field(2, output.as_bytes(), &mut n); + len_field(4, op.as_bytes(), &mut n); + for (name, value) in ints { + let mut a = Vec::new(); + len_field(1, name.as_bytes(), &mut a); + varint_field(3, *value as u64, &mut a); + varint_field(20, 2, &mut a); + len_field(5, &a, &mut n); + } + n + } + + fn tensor(name: &str, dims: &[u64], values: &[f32]) -> Vec { + let mut t = Vec::new(); + for d in dims { + varint_field(1, *d, &mut t); + } + varint_field(2, 1, &mut t); + len_field(8, name.as_bytes(), &mut t); + let raw: Vec = values.iter().flat_map(|v| v.to_le_bytes()).collect(); + len_field(9, &raw, &mut t); + t + } + + /// A tiny working denoiser-shaped network: `[1, 4, t, t]` → a 1 × 1 convolution to 12 channels (each output + /// channel is `scale` times one input plane, so the scale is known) → depth-to-space → `[1, 3, 2t, 2t]`. + fn tiny_model(t: u64, scale: f32) -> Vec { + tiny_model_with(t, scale, 12) + } + + /// The same with another number of convolution channels; only 12 can be arranged into the three colours. + fn tiny_model_with(t: u64, scale: f32, channels: u64) -> Vec { + let mut graph = Vec::new(); + len_field(1, &node("Conv", &["data", "w"], if channels == 12 { "c" } else { "rgb" }, &[("group", 1)]), &mut graph); + if channels == 12 { + let mut depth_to_space = node("DepthToSpace", &["c"], "rgb", &[("blocksize", 2)]); + let mut mode = Vec::new(); + len_field(1, b"mode", &mut mode); + len_field(4, b"CRD", &mut mode); + varint_field(20, 3, &mut mode); + len_field(5, &mode, &mut depth_to_space); + len_field(1, &depth_to_space, &mut graph); + } + // CRD output channel colour·4 + dy·2 + dx reads input plane [R, G1, B] of the tile. + let mut w = vec![0f32; channels as usize * 4]; + for oc in 0..channels as usize { + let colour = (oc / 4).min(2); + let plane = [0, 1, 3][colour]; + w[oc * 4 + plane] = scale; + } + len_field(5, &tensor("w", &[channels, 4, 1, 1], &w), &mut graph); + len_field(11, &value_info_bytes("data", 1, &[Ok(1), Ok(4), Ok(t), Ok(t)]), &mut graph); + len_field(12, &value_info_bytes("rgb", 1, &[Ok(1), Ok(3), Ok(2 * t), Ok(2 * t)]), &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 + } + + fn manifest(tile: u32, nominal: f32) -> DenoiserManifest { + DenoiserManifest { + id: "tiny".into(), + name: "Tiny".into(), + version: "1".into(), + licence: Licence::default(), + source: None, + sha256: None, + size_bytes: None, + provenance: String::new(), + domain: Domain::BayerToRgb, + tile, + overlap: 16, + gain: Gain::MatchMean { nominal, max_deviation: 0.05 }, + } + } + + fn temp_model(name: &str, bytes: &[u8]) -> std::path::PathBuf { + let p = std::env::temp_dir().join(format!("lc-denoise-runtime-{name}-{}.onnx", std::process::id())); + std::fs::write(&p, bytes).unwrap(); + p + } + + #[test] + fn a_working_model_loads_runs_and_has_the_scale_it_says() { + let p = temp_model("ok", &tiny_model(64, 1000.0)); + let r = CpuRunner::load(&p, &manifest(64, 1000.0)).unwrap(); + let input: Vec = (0..4 * 64 * 64).map(|i| (i % 97) as f32 / 97.0).collect(); + let out = r.run(&input).unwrap(); + assert_eq!(out.len(), 3 * 128 * 128); + // output (colour c, pixel (2y+dy, 2x+dx)) = scale × input plane [R, G1, B][c] at (y, x) + for (c, plane) in [0usize, 1, 3].into_iter().enumerate() { + let (y, x) = (5usize, 9usize); + assert!((out[c * 128 * 128 + (2 * y + 1) * 128 + 2 * x] - 1000.0 * input[plane * 64 * 64 + y * 64 + x]).abs() < 1e-2); + } + let _ = std::fs::remove_file(p); + } + + #[test] + fn the_self_test_checks_scale_and_noise() { + // a model that only multiplies cannot take noise down: it must fail that check but pass the others + let p = temp_model("selftest", &tiny_model(64, 1000.0)); + let t = CpuRunner::load(&p, &manifest(64, 1000.0)).unwrap().self_test(); + assert!(!t.ok, "{t:?}"); + assert!(t.checks.iter().any(|(c, ok)| c.contains("same answer") && *ok), "{t:?}"); + assert!(t.checks.iter().any(|(c, ok)| c.contains("scale") && *ok), "{t:?}"); + assert!(t.checks.iter().any(|(c, ok)| c.contains("noise") && !*ok), "{t:?}"); + // and one that claims the wrong scale fails that check + let q = CpuRunner::load(&p, &manifest(64, 5.0)).unwrap().self_test(); + assert!(q.checks.iter().any(|(c, ok)| c.contains("scale") && !*ok), "{q:?}"); + let _ = std::fs::remove_file(p); + } + + #[test] + fn wrong_sizes_and_broken_files_are_errors() { + let p = temp_model("sizes", &tiny_model(64, 1.0)); + let r = CpuRunner::load(&p, &manifest(64, 1.0)).unwrap(); + assert!(matches!(r.run(&[0.0; 10]), Err(Error::Input(_)))); + // a fully convolutional model works at whatever tile size the manifest asks for + let wide = CpuRunner::load(&p, &manifest(128, 1.0)).unwrap(); + assert_eq!(wide.run(&vec![0.5; 4 * 128 * 128]).unwrap().len(), 3 * 256 * 256); + // a model that gives another number of values than the manifest promises is an error + let q = temp_model("channels", &tiny_model_with(64, 1.0, 8)); + assert!(CpuRunner::load(&q, &manifest(64, 1.0)).is_err(), "shape is refused before execution"); + let _ = std::fs::remove_file(q); + for (name, bytes) in [("junk", b"nonsense".to_vec()), ("empty", Vec::new())] { + let q = temp_model(name, &bytes); + assert!(CpuRunner::load(&q, &manifest(64, 1.0)).is_err(), "{name}"); + let _ = std::fs::remove_file(q); + } + assert!(CpuRunner::load(&std::env::temp_dir().join("does-not-exist.onnx"), &manifest(64, 1.0)).is_err()); + let _ = std::fs::remove_file(p); + } + + #[test] + fn the_check_tile_covers_black_to_white() { + for t in [1, 2, 7, 64] { + let c = check_tile(t); + assert_eq!(c.len(), 4 * t * t); + assert!(c.iter().all(|v| (0.0..=1.0).contains(v)), "{t}"); + } + let c = check_tile(64); + let near = |lo: f32, hi: f32| c.iter().filter(|v| (lo..=hi).contains(*v)).count(); + assert!(near(0.0, 0.0) >= 4 * 8 * 8 && near(1.0, 1.0) >= 4 * 8 * 8, "a flat black and a clipped white corner"); + for band in 0..10 { + let lo = band as f32 / 10.0; + assert!(near(lo, lo + 0.1) > 100, "values around {lo}"); + } + assert_eq!(check_tile(64), c, "the same tile every time"); + } + + #[test] + fn truncated_and_corrupted_models_never_panic() { + let good = tiny_model(64, 1.0); + for n in (0..good.len()).step_by(5) { + let q = temp_model("trunc", &good[..n]); + let _ = CpuRunner::load(&q, &manifest(64, 1.0)); + let _ = std::fs::remove_file(q); + } + for at in (0..good.len()).step_by(7) { + 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(r) = CpuRunner::load(&q, &manifest(64, 1.0)) { + let _ = r.run(&vec![0.5; 4 * 64 * 64]); + } + let _ = std::fs::remove_file(q); + } + } + + /// Opt-in: `LC_DENOISE_MODEL= cargo test -p lightcraft-denoise --release + /// -- --ignored --nocapture` self-tests a real model (RawNIND's Bayer model: tile 512, scale about 1e6). + #[test] + #[ignore = "needs a real model file: set LC_DENOISE_MODEL"] + fn a_real_model_passes_the_self_test() { + let Some(path) = std::env::var_os("LC_DENOISE_MODEL") else { return }; + let m = manifest(512, 1.0e6); + let r = CpuRunner::load(Path::new(&path), &m).unwrap(); + let t = r.self_test(); + println!("load {:.0} ms | {:.0} ms per tile | scale {:.4e} | {:?}", t.load_ms, t.tile_ms, t.scale, t.checks); + assert!(t.ok); + } +} + +#[cfg(test)] +mod cpu_tests { + use super::*; + #[test] + fn manifest_and_test_tile_sizes_are_checked_without_tract() { + for t in [0, 2049, usize::MAX] { + assert!(test_tile(t).0.is_empty()); + assert!(check_tile(t).is_empty()); + } + let path = std::env::temp_dir().join(format!("lc-cpu-load-{}.onnx", std::process::id())); + std::fs::write(&path, crate::synthetic::unet_onnx(64, 5, 1, 7)).unwrap(); + let mut m = DenoiserManifest { + id: "synthetic".into(), + name: "Synthetic".into(), + version: "1".into(), + licence: Default::default(), + source: None, + sha256: None, + size_bytes: None, + provenance: "synthetic test".into(), + domain: crate::manifest::Domain::BayerToRgb, + tile: 64, + overlap: 16, + gain: Gain::None, + }; + let r = CpuRunner::load(&path, &m).unwrap(); + let x = check_tile(64); + assert_eq!(r.run(&x).unwrap(), r.run(&x).unwrap()); + assert!(r.self_test().checks.iter().any(|(text, ok)| text.contains("same answer") && *ok)); + m.tile = u32::MAX; + assert!(CpuRunner::load(&path, &m).is_err()); + std::fs::remove_file(path).unwrap(); + } +} diff --git a/crates/denoise/src/synthetic.rs b/crates/denoise/src/synthetic.rs new file mode 100644 index 00000000..8ceafca7 --- /dev/null +++ b/crates/denoise/src/synthetic.rs @@ -0,0 +1,234 @@ +//! A small ONNX U-Net built in memory, for tests of the layers that read and run denoise networks (here, in the GPU +//! runner and in the engine): the same layer kinds and wiring as the real denoisers, with deterministic random weights. +//! Nothing here is a model anyone would use. + +use crate::synthetic_proto::{len_field, value_info_bytes}; + +fn varint(mut v: u64, out: &mut Vec) { + loop { + let b = (v & 0x7f) as u8; + v >>= 7; + if v == 0 { + out.push(b); + return; + } + out.push(b | 0x80); + } +} + +fn varint_field(n: u64, v: u64, out: &mut Vec) { + varint(n << 3, out); + varint(v, out); +} + +enum A<'a> { + Int(i64), + Ints(&'a [i64]), + Float(f32), + Str(&'a str), +} + +fn node(op: &str, inputs: &[&str], output: &str, attrs: &[(&str, A)]) -> Vec { + let mut n = Vec::new(); + for i in inputs { + len_field(1, i.as_bytes(), &mut n); + } + len_field(2, output.as_bytes(), &mut n); + len_field(4, op.as_bytes(), &mut n); + for (name, value) in attrs { + let mut a = Vec::new(); + len_field(1, name.as_bytes(), &mut a); + match value { + A::Int(v) => { + varint_field(3, *v as u64, &mut a); + varint_field(20, 2, &mut a); + } + A::Ints(vs) => { + for v in *vs { + varint_field(8, *v as u64, &mut a); + } + varint_field(20, 7, &mut a); + } + A::Float(v) => { + varint((2 << 3) | 5, &mut a); + a.extend_from_slice(&v.to_le_bytes()); + varint_field(20, 1, &mut a); + } + A::Str(s) => { + len_field(4, s.as_bytes(), &mut a); + varint_field(20, 3, &mut a); + } + } + len_field(5, &a, &mut n); + } + n +} + +fn tensor(name: &str, dims: &[u64], values: &[f32]) -> Vec { + let mut t = Vec::new(); + for d in dims { + varint_field(1, *d, &mut t); + } + varint_field(2, 1, &mut t); + len_field(8, name.as_bytes(), &mut t); + let raw: Vec = values.iter().flat_map(|v| v.to_le_bytes()).collect(); + len_field(9, &raw, &mut t); + t +} + +struct Rng(u64); + +impl Rng { + fn next(&mut self) -> f32 { + self.0 ^= self.0 >> 12; + self.0 ^= self.0 << 25; + self.0 ^= self.0 >> 27; + ((self.0.wrapping_mul(0x2545_f491_4f6c_dd1d) >> 40) as f32 / (1u64 << 24) as f32) * 2.0 - 1.0 + } +} + +struct Builder { + nodes: Vec>, + inits: Vec>, + rng: Rng, + n: usize, +} + +impl Builder { + fn name(&mut self, kind: &str) -> String { + self.n += 1; + format!("{kind}{}", self.n) + } + + /// Random weights scaled so activations keep a sane size through the layers. + fn weights(&mut self, name: &str, dims: &[u64], fan_in: usize) -> String { + let count: u64 = dims.iter().product(); + let scale = (3.0 / fan_in as f32).sqrt(); + let values: Vec = (0..count).map(|_| self.rng.next() * scale).collect(); + self.inits.push(tensor(name, dims, &values)); + name.to_string() + } + + fn bias(&mut self, name: &str, n: u64) -> String { + let values: Vec = (0..n).map(|_| self.rng.next() * 0.1).collect(); + self.inits.push(tensor(name, &[n], &values)); + name.to_string() + } + + /// `Conv` (+ `LeakyRelu` when `leaky`); gives the name of the tensor that results. + fn conv(&mut self, src: &str, cin: u64, cout: u64, k: u64, leaky: bool) -> String { + let (w, b, out) = (self.name("w"), self.name("b"), self.name("conv")); + let w = self.weights(&w, &[cout, cin, k, k], (cin * k * k) as usize); + let b = self.bias(&b, cout); + let pad = (k as i64 - 1) / 2; + self.nodes.push(node( + "Conv", + &[src, &w, &b], + &out, + &[("kernel_shape", A::Ints(&[k as i64, k as i64])), ("pads", A::Ints(&[pad; 4])), ("strides", A::Ints(&[1, 1])), ("group", A::Int(1))], + )); + if !leaky { + return out; + } + let act = self.name("act"); + self.nodes.push(node("LeakyRelu", &[&out], &act, &[("alpha", A::Float(0.2))])); + act + } + + fn pool(&mut self, src: &str) -> String { + let out = self.name("pool"); + self.nodes.push(node( + "MaxPool", + &[src], + &out, + &[("kernel_shape", A::Ints(&[2, 2])), ("strides", A::Ints(&[2, 2])), ("pads", A::Ints(&[0; 4]))], + )); + out + } + + fn up(&mut self, src: &str, cin: u64, cout: u64) -> String { + let (w, b, out) = (self.name("uw"), self.name("ub"), self.name("up")); + let w = self.weights(&w, &[cin, cout, 2, 2], (cin * 4) as usize); + let b = self.bias(&b, cout); + self.nodes.push(node( + "ConvTranspose", + &[src, &w, &b], + &out, + &[("kernel_shape", A::Ints(&[2, 2])), ("strides", A::Ints(&[2, 2])), ("pads", A::Ints(&[0; 4]))], + )); + out + } + + fn concat(&mut self, a: &str, b: &str) -> String { + let out = self.name("cat"); + self.nodes.push(node("Concat", &[a, b], &out, &[("axis", A::Int(1))])); + out + } +} + +/// A denoiser-shaped U-Net as ONNX bytes: `[1, 4, tile, tile]` → `[1, 3, 2·tile, 2·tile]`, with `depth` poolings +/// (so `tile` must be a multiple of `2^depth`) and `ch` channels on the first level, doubling on each (every count a +/// multiple of 4). Each level is two 3 × 3 convolutions with a leaky ReLU, as in the real networks; the output is a +/// 1 × 1 convolution to 12 channels and depth-to-space. +pub fn unet_onnx(tile: u64, ch: u64, depth: usize, seed: u64) -> Vec { + let mut b = Builder { nodes: Vec::new(), inits: Vec::new(), rng: Rng(seed | 1), n: 0 }; + let mut skips: Vec<(String, u64)> = Vec::new(); + let (mut cur, mut c_in, mut c) = ("data".to_string(), 4u64, ch); + for level in 0..=depth { + let x = b.conv(&cur, c_in, c, 3, true); + let x = b.conv(&x, c, c, 3, true); + if level < depth { + skips.push((x.clone(), c)); + cur = b.pool(&x); + c_in = c; + c *= 2; + } else { + cur = x; + } + } + for (skip, sc) in skips.into_iter().rev() { + let up = b.up(&cur, c, sc); + let joined = b.concat(&up, &skip); + let x = b.conv(&joined, 2 * sc, sc, 3, true); + cur = b.conv(&x, sc, sc, 3, true); + c = sc; + } + let rgb = b.conv(&cur, c, 12, 1, false); + b.nodes.push(node("DepthToSpace", &[&rgb], "rgb", &[("blocksize", A::Int(2)), ("mode", A::Str("CRD"))])); + + let mut graph = Vec::new(); + for n in &b.nodes { + len_field(1, n, &mut graph); + } + for t in &b.inits { + len_field(5, t, &mut graph); + } + len_field(11, &value_info_bytes("data", 1, &[Ok(1), Ok(4), Ok(tile), Ok(tile)]), &mut graph); + len_field(12, &value_info_bytes("rgb", 1, &[Ok(1), Ok(3), Ok(2 * tile), Ok(2 * tile)]), &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 tiny smoothing network for installing a real CPU model in setup-flow tests. +/// It averages a 3 × 3 neighborhood of the four Bayer planes and returns three identical channels. +pub fn smoothing_onnx() -> Vec { + let mut graph = Vec::new(); + len_field(1, &node("Conv", &["data", "weights"], "samples", &[("kernel_shape", A::Ints(&[3, 3])), ("pads", A::Ints(&[1, 1, 1, 1]))]), &mut graph); + len_field(1, &node("DepthToSpace", &["samples"], "rgb", &[("blocksize", A::Int(2)), ("mode", A::Str("CRD"))]), &mut graph); + len_field(5, &tensor("weights", &[12, 4, 3, 3], &[1.0 / 36.0; 12 * 4 * 3 * 3]), &mut graph); + len_field(11, &value_info_bytes("data", 1, &[Ok(1), Ok(4), Ok(64), Ok(64)]), &mut graph); + len_field(12, &value_info_bytes("rgb", 1, &[Ok(1), Ok(3), Ok(128), Ok(128)]), &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 +} diff --git a/crates/denoise/src/synthetic_proto.rs b/crates/denoise/src/synthetic_proto.rs new file mode 100644 index 00000000..a78d586b --- /dev/null +++ b/crates/denoise/src/synthetic_proto.rs @@ -0,0 +1,44 @@ +pub(crate) fn varint(mut v: u64, out: &mut Vec) { + loop { + let b = (v & 0x7f) as u8; + v >>= 7; + if v == 0 { + out.push(b); + return; + } + out.push(b | 0x80); + } +} + +pub fn len_field(n: u64, body: &[u8], out: &mut Vec) { + varint(n << 3 | 2, out); + varint(body.len() as u64, out); + out.extend_from_slice(body); +} + +pub(crate) fn varint_field(n: u64, v: u64, out: &mut Vec) { + varint(n << 3, out); + varint(v, out); +} + +/// A graph input or output: `Ok(n)` is a fixed dimension, `Err(name)` a named (dynamic) one. +pub fn value_info_bytes(name: &str, elem: u64, dims: &[std::result::Result]) -> Vec { + let mut shape = Vec::new(); + for d in dims { + let mut dim = Vec::new(); + match d { + Ok(v) => varint_field(1, *v, &mut dim), + Err(p) => len_field(2, p.as_bytes(), &mut dim), + } + len_field(1, &dim, &mut shape); + } + let mut tensor = Vec::new(); + varint_field(1, elem, &mut tensor); + len_field(2, &shape, &mut tensor); + let mut ty = Vec::new(); + len_field(1, &tensor, &mut ty); + let mut vi = Vec::new(); + len_field(1, name.as_bytes(), &mut vi); + len_field(2, &ty, &mut vi); + vi +} diff --git a/crates/denoise/tests/real_model.rs b/crates/denoise/tests/real_model.rs new file mode 100644 index 00000000..203e9937 --- /dev/null +++ b/crates/denoise/tests/real_model.rs @@ -0,0 +1,250 @@ +//! Opt-in end-to-end check of the real RawNIND Bayer model on real mosaics with noise added: +//! +//! ```text +//! LC_DENOISE_MODEL= LC_DENOISE_MOSAICS= LC_DENOISE_OUT= \ +//! cargo test --release -p lightcraft-denoise --features runtime --test real_model -- --ignored --nocapture +//! ``` +//! +//! `LC_DENOISE_MOSAICS` holds `.mosaic.f32` (little-endian f32, normalised: black 0, white 1, not white +//! balanced), `.rgb.f32` (a clean demosaic of it, camera RGB, interleaved) and `.txt` (`width`, +//! `height`, `cfa`, `wb` lines). Noise is Poisson-Gaussian, made here. Optional: `LC_DENOISE_PARALLEL` (tiles at +//! once, default 8), `LC_DENOISE_ONLY` (one name), `LC_DENOISE_ISO` (noise strength, default 1.0; 0 adds none), +//! `LC_DENOISE_FULL` (the whole frame instead of a 2048 × 1536 crop). + +#![cfg(feature = "runtime")] + +use std::path::{Path, PathBuf}; +use std::time::Instant; + +use lightcraft_denoise::bayer::Layout; +use lightcraft_denoise::manifest::{DenoiserManifest, Domain, Gain}; +use lightcraft_denoise::run::{Control, Params, denoise_bayer}; +use lightcraft_denoise::runtime::CpuRunner; + +struct Rng(u64); +impl Rng { + fn next(&mut self) -> f64 { + self.0 ^= self.0 >> 12; + self.0 ^= self.0 << 25; + self.0 ^= self.0 >> 27; + ((self.0.wrapping_mul(0x2545_f491_4f6c_dd1d) >> 11) as f64 + 0.5) / (1u64 << 53) as f64 + } + fn gauss(&mut self) -> f64 { + (-2.0 * self.next().ln()).sqrt() * (std::f64::consts::TAU * self.next()).cos() + } +} + +fn read_f32(path: &Path) -> Vec { + std::fs::read(path).unwrap().as_chunks::<4>().0.iter().map(|b| f32::from_le_bytes(*b)).collect() +} + +struct Dump { + width: usize, + height: usize, + layout: Layout, + wb: [f32; 3], + mosaic: Vec, + rgb: Vec, +} + +fn load(dir: &Path, name: &str) -> Dump { + let text = std::fs::read_to_string(dir.join(format!("{name}.txt"))).unwrap(); + let field = |k: &str| text.lines().find_map(|l| l.strip_prefix(k).map(|v| v.trim().to_string())).unwrap(); + let cell = match field("cfa").as_str() { + "RGGB" => [0, 1, 1, 2], + "BGGR" => [2, 1, 1, 0], + "GRBG" => [1, 0, 2, 1], + "GBRG" => [1, 2, 0, 1], + other => panic!("cfa {other}"), + }; + let wb: Vec = field("wb").split_whitespace().map(|v| v.parse().unwrap()).collect(); + Dump { + width: field("width").parse().unwrap(), + height: field("height").parse().unwrap(), + layout: Layout::from_cell(cell).unwrap(), + wb: [wb[0], wb[1], wb[2]], + mosaic: read_f32(&dir.join(format!("{name}.mosaic.f32"))), + rgb: read_f32(&dir.join(format!("{name}.rgb.f32"))), + } +} + +/// Crop with an even origin so the Bayer phase is kept. +fn crop(d: &Dump, x0: usize, y0: usize, w: usize, h: usize) -> (Vec, Vec) { + let mut m = Vec::with_capacity(w * h); + let mut c = Vec::with_capacity(w * h * 3); + for y in y0..y0 + h { + m.extend_from_slice(&d.mosaic[y * d.width + x0..y * d.width + x0 + w]); + c.extend_from_slice(&d.rgb[(y * d.width + x0) * 3..(y * d.width + x0 + w) * 3]); + } + (m, c) +} + +/// 0 = red, 1 = green, 2 = blue at a mosaic position. +fn site(layout: Layout, x: usize, y: usize) -> usize { + match ((x + layout.dx) % 2, (y + layout.dy) % 2) { + (0, 0) => 0, + (1, 1) => 2, + _ => 1, + } +} + +/// The simplest demosaic (each missing colour is the mean of the same-coloured samples in the 3 × 3 around), +/// as the baseline a user would get from the noisy picture without AI. +fn bilinear(m: &[f32], w: usize, h: usize, layout: Layout) -> Vec { + let mut out = vec![0f32; w * h * 3]; + for y in 0..h { + for x in 0..w { + let mut sum = [0f32; 3]; + let mut n = [0f32; 3]; + for yy in y.saturating_sub(1)..(y + 2).min(h) { + for xx in x.saturating_sub(1)..(x + 2).min(w) { + let c = site(layout, xx, yy); + sum[c] += m[yy * w + xx]; + n[c] += 1.0; + } + } + for c in 0..3 { + out[(y * w + x) * 3 + c] = if n[c] > 0.0 { sum[c] / n[c] } else { 0.0 }; + } + } + } + out +} + +/// Signal-to-error ratio in dB of `got` against `want` (both interleaved RGB), ignoring a border. +fn snr(got: &[f32], want: &[f32], w: usize, h: usize, border: usize) -> f64 { + let (mut sig, mut err) = (0f64, 0f64); + for y in border..h - border { + for x in border..w - border { + for c in 0..3 { + let (g, t) = (f64::from(got[(y * w + x) * 3 + c]), f64::from(want[(y * w + x) * 3 + c])); + sig += t * t; + err += (g - t).powi(2); + } + } + } + 10.0 * (sig / err.max(1e-30)).log10() +} + +fn crc32(data: &[u8]) -> u32 { + let mut c = !0u32; + for &b in data { + c ^= u32::from(b); + for _ in 0..8 { + c = if c & 1 != 0 { 0xedb8_8320 ^ (c >> 1) } else { c >> 1 }; + } + } + !c +} + +/// A viewable PNG of interleaved camera-RGB panels side by side: white balanced, one exposure for all, gamma. +fn write_png(path: &Path, panels: &[&[f32]], w: usize, h: usize, wb: [f32; 3], exposure: f32) { + let total_w = w * panels.len(); + let mut raw = Vec::with_capacity((total_w * 3 + 1) * h); + for y in 0..h { + raw.push(0u8); + for p in panels { + for x in 0..w { + for c in 0..3 { + let v = (p[(y * w + x) * 3 + c] * wb[c] * exposure).clamp(0.0, 1.0); + raw.push((v.powf(1.0 / 2.2) * 255.0 + 0.5) as u8); + } + } + } + } + let z = miniz_oxide::deflate::compress_to_vec_zlib(&raw, 6); + let mut png = b"\x89PNG\r\n\x1a\n".to_vec(); + let mut chunk = |kind: &[u8; 4], body: &[u8]| { + png.extend_from_slice(&(body.len() as u32).to_be_bytes()); + let mut k = kind.to_vec(); + k.extend_from_slice(body); + png.extend_from_slice(&k); + png.extend_from_slice(&crc32(&k).to_be_bytes()); + }; + let mut ihdr = Vec::new(); + ihdr.extend_from_slice(&(total_w as u32).to_be_bytes()); + ihdr.extend_from_slice(&(h as u32).to_be_bytes()); + ihdr.extend_from_slice(&[8, 2, 0, 0, 0]); + chunk(b"IHDR", &ihdr); + chunk(b"IDAT", &z); + chunk(b"IEND", &[]); + std::fs::write(path, png).unwrap(); +} + +#[test] +#[ignore = "needs a real model and mosaics: see the module docs"] +fn the_real_model_denoises_real_mosaics() { + let (Some(model), Some(dir)) = (std::env::var_os("LC_DENOISE_MODEL"), std::env::var_os("LC_DENOISE_MOSAICS")) else { return }; + let dir = PathBuf::from(dir); + let out_dir = std::env::var_os("LC_DENOISE_OUT").map(PathBuf::from); + let parallel: usize = std::env::var("LC_DENOISE_PARALLEL").ok().and_then(|v| v.parse().ok()).unwrap_or(8); + let iso: f64 = std::env::var("LC_DENOISE_ISO").ok().and_then(|v| v.parse().ok()).unwrap_or(1.0); + let names: Vec = match std::env::var("LC_DENOISE_ONLY") { + Ok(n) => vec![n], + Err(_) => ["arw-sony-a7m3-compressed", "cr2-canon-6d", "nef-nikon-d5100-lossless", "dng-google-pixel2xl", "pef-pentax-k3"] + .map(String::from) + .to_vec(), + }; + let manifest = DenoiserManifest { + id: "rawnind-bayer".into(), + name: "RawNIND Bayer".into(), + version: "1".into(), + licence: Default::default(), + source: None, + sha256: None, + size_bytes: None, + provenance: String::new(), + domain: Domain::BayerToRgb, + tile: 512, + overlap: 64, + gain: Gain::MatchMean { nominal: 1.0e6, max_deviation: 0.05 }, + }; + let runner = CpuRunner::load(Path::new(&model), &manifest).unwrap(); + let params = Params { tile: 512, overlap: 64, gain: manifest.gain, clip: Some(1.0), parallel }; + for name in names { + let d = load(&dir, &name); + // a crop of about 2 × 1.5 tiles, away from the edges, even origin + let full = std::env::var_os("LC_DENOISE_FULL").is_some(); + let (w, h) = if full { (d.width & !1, d.height & !1) } else { (2048usize, 1536usize) }; + let (x0, y0) = (((d.width - w) / 3) & !1, ((d.height - h) / 2) & !1); + let (clean, reference) = crop(&d, x0, y0, w, h); + // Poisson-Gaussian noise: variance a·signal + b, as a sensor at high ISO + let (a, b) = (1.5e-3 * iso, (4.0e-3 * iso.sqrt()).powi(2)); + let mut rng = Rng(0x9e37_79b9_7f4a_7c15 ^ name.len() as u64); + let noisy: Vec = clean + .iter() + .map(|&v| { + let v = f64::from(v); + let sigma = (a * v.max(0.0) + b).sqrt(); + ((v + sigma * rng.gauss()) as f32).min(1.0) + }) + .collect(); + let started = Instant::now(); + let done = std::sync::atomic::AtomicUsize::new(0); + let progress = |n: usize, total: usize| done.store(n * 1000 / total, std::sync::atomic::Ordering::Relaxed); + let denoised = denoise_bayer(&noisy, w, h, d.layout, &runner, ¶ms, &Control { cancel: None, progress: Some(&progress) }).unwrap(); + let secs = started.elapsed().as_secs_f64(); + let baseline = bilinear(&noisy, w, h, d.layout); + let clean_bilinear = bilinear(&clean, w, h, d.layout); + let den: Vec = denoised.data.iter().flat_map(|p| p.iter().copied()).collect(); + let (s_noisy, s_den) = (snr(&baseline, &reference, w, h, 64), snr(&den, &reference, w, h, 64)); + let s_floor = snr(&clean_bilinear, &reference, w, h, 64); + println!( + "{name}: {w}×{h} crop, {} tiles, parallel {parallel}, {secs:.1} s | SNR vs clean demosaic: noisy bilinear {s_noisy:.2} dB, denoised {s_den:.2} dB (clean bilinear {s_floor:.2} dB)", + lightcraft_denoise::run::tile_count(w, h, d.layout, 512, 64) + ); + if let Some(out) = &out_dir { + std::fs::create_dir_all(out).unwrap(); + let mean = reference.iter().map(|&v| f64::from(v)).sum::() / reference.len() as f64; + let exposure = (0.18 / mean.max(1e-6)) as f32; + // three panels, a 512 × 384 window each, at 1:1 + let (pw, ph, px, py) = (512usize, 384usize, 700usize, 500usize); + let window = + |buf: &[f32]| -> Vec { (0..ph).flat_map(|y| buf[((py + y) * w + px) * 3..((py + y) * w + px + pw) * 3].to_vec()).collect() }; + write_png(&out.join(format!("{name}.png")), &[&window(&baseline), &window(&den), &window(&reference)], pw, ph, d.wb, exposure); + } + if iso >= 4.0 { + assert!(s_den > s_noisy + 2.0, "{name}: denoised {s_den:.2} dB is not clearly better than noisy {s_noisy:.2} dB"); + } + } +} diff --git a/crates/develop/src/controls.rs b/crates/develop/src/controls.rs index 5696d670..c8a001b7 100644 --- a/crates/develop/src/controls.rs +++ b/crates/develop/src/controls.rs @@ -207,6 +207,7 @@ controls! { "grain.amount" => grain.amount, "Grain", Grain, 0, 100, 0, 1, 0, Plain; "grain.size" => grain.size, "Size", Grain, 0, 100, 25, 1, 0, Plain; "grain.roughness" => grain.roughness, "Roughness", Grain, 0, 100, 50, 1, 0, Plain; + "enhance.denoise" => enhance.denoise, "Amount", Detail, 0, 100, 0, 1, 0, Plain; "detail.sharpenAmount" => detail.sharpen_amount, "Sharpening", Detail, 0, 150, 0, 1, 0, Plain; "detail.sharpenRadius" => detail.sharpen_radius, "Radius", Detail, 0.5, 3, 1, 0.1, 1, Plain; "detail.sharpenDetail" => detail.sharpen_detail, "Detail", Detail, 0, 100, 25, 1, 0, Plain; diff --git a/crates/develop/src/lib.rs b/crates/develop/src/lib.rs index bfce1893..b7c95251 100644 --- a/crates/develop/src/lib.rs +++ b/crates/develop/src/lib.rs @@ -94,6 +94,13 @@ impl DevelopSettings { std::borrow::Cow::Owned(d) } + /// How much of the AI-denoised picture to mix in, 0 (none) to 1 (all of it): the Denoise amount, unless the Detail + /// section is switched off or the value is not a number. + pub fn denoise_amount(&self) -> f32 { + let a = self.enhance.denoise; + if a.is_finite() && self.enhance.denoise_enabled() && self.section_enabled("detail") { (a / 100.0).clamp(0.0, 1.0) as f32 } else { 0.0 } + } + pub fn set_section_enabled(&mut self, section: &str, on: bool) { self.disabled_sections.retain(|s| s != section); if !on { @@ -112,6 +119,7 @@ impl DevelopSettings { self.curve = d; } Section::Color => self.wb.mode = WbMode::AsShot, + Section::Detail => self.enhance.denoise_on = None, _ => {} } } @@ -193,4 +201,40 @@ mod tests { s.set_section_enabled("effects", true); assert!(s.section_enabled("effects")); } + + #[test] + fn the_denoise_amount_is_a_slider_and_follows_the_detail_section() { + let mut s = DevelopSettings::default(); + assert_eq!(s.denoise_amount(), 0.0); + assert!(controls::set(&mut s, "enhance.denoise", 60.0)); + assert_eq!(controls::get(&s, "enhance.denoise"), Some(60.0)); + assert!((s.denoise_amount() - 0.6).abs() < 1e-6); + // the slider's range is 0 to 100; values outside it are brought in + controls::set(&mut s, "enhance.denoise", 500.0); + assert_eq!(s.denoise_amount(), 1.0); + controls::set(&mut s, "enhance.denoise", -5.0); + assert_eq!(s.denoise_amount(), 0.0); + // switching the Detail section off switches the denoise off with it; a settings file with a bad number does nothing + controls::set(&mut s, "enhance.denoise", 80.0); + s.set_section_enabled("detail", false); + assert_eq!(s.denoise_amount(), 0.0); + s.set_section_enabled("detail", true); + s.enhance.denoise = f64::NAN; + assert_eq!(s.denoise_amount(), 0.0); + // it is part of Detail: a reset of the section clears it + controls::set(&mut s, "enhance.denoise", 40.0); + s.reset_section(Section::Detail); + assert_eq!(s.enhance.denoise, 0.0); + // The later independent switch preserves Amount when off, and can be on at zero. + controls::set(&mut s, "enhance.denoise", 75.0); + s.enhance.denoise_on = Some(false); + assert_eq!(s.denoise_amount(), 0.0); + assert_eq!(s.enhance.denoise, 75.0); + s.enhance.denoise_on = Some(true); + assert_eq!(s.denoise_amount(), 0.75); + controls::set(&mut s, "enhance.denoise", 0.0); + assert!(s.enhance.denoise_enabled()); + s.reset_section(Section::Detail); + assert_eq!(s.enhance.denoise_on, None); + } } diff --git a/crates/develop/src/settings.rs b/crates/develop/src/settings.rs index 910cd07d..5f0db8ab 100644 --- a/crates/develop/src/settings.rs +++ b/crates/develop/src/settings.rs @@ -915,6 +915,14 @@ pub struct LensBlur { #[serde(default)] pub struct Enhance { pub denoise: f64, + #[serde(skip_serializing_if = "Option::is_none")] + pub denoise_on: Option, pub raw_details: bool, pub super_resolution: bool, } + +impl Enhance { + pub fn denoise_enabled(&self) -> bool { + self.denoise_on.unwrap_or(self.denoise > 0.0) + } +} diff --git a/crates/engine/Cargo.toml b/crates/engine/Cargo.toml index a916d180..9bd9845f 100644 --- a/crates/engine/Cargo.toml +++ b/crates/engine/Cargo.toml @@ -21,7 +21,10 @@ web-time = { workspace = true } lightcraft-preview = { workspace = true } lightcraft-codecs = { workspace = true } lightcraft-raw = { workspace = true } +lightcraft-denoise-core = { workspace = true } +lightcraft-denoise = { workspace = true, optional = true } lightcraft-meta = { workspace = true } +lightcraft-faces = { workspace = true } miniz_oxide = { workspace = true } lightcraft-merge = { workspace = true } rayon = { workspace = true } @@ -35,9 +38,15 @@ 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"] +denoise = ["dep:lightcraft-denoise", "lightcraft-denoise/runtime", "lightcraft-gpu/denoise"] +rawnind-model = ["denoise", "lightcraft-denoise/rawnind-model", "lightcraft-denoise-core/rawnind-model"] [dev-dependencies] libc = { workspace = true } diff --git a/crates/engine/src/camera_profiles.rs b/crates/engine/src/camera_profiles.rs index 97fb8ef2..80192cf6 100644 --- a/crates/engine/src/camera_profiles.rs +++ b/crates/engine/src/camera_profiles.rs @@ -65,19 +65,8 @@ impl CameraProfile { } } -/// LightCraft's configuration folder (settings, GPU marker, camera profiles). -pub fn config_dir() -> Option { - if cfg!(target_os = "macos") { - std::env::var_os("HOME").map(|h| PathBuf::from(h).join("Library/Application Support/LightCraft")) - } else if cfg!(windows) { - std::env::var_os("APPDATA").map(|a| PathBuf::from(a).join("LightCraft")) - } else { - std::env::var_os("XDG_CONFIG_HOME") - .map(PathBuf::from) - .or_else(|| std::env::var_os("HOME").map(|h| PathBuf::from(h).join(".config"))) - .map(|c| c.join("lightcraft")) - } -} +/// LightCraft's configuration folder (settings, GPU marker, camera profiles); one helper for the app. +pub use crate::config::config_dir; /// Where camera profiles are read from and written to. pub fn dir() -> Option { diff --git a/crates/engine/src/cmd/denoise.rs b/crates/engine/src/cmd/denoise.rs new file mode 100644 index 00000000..e1cc0900 --- /dev/null +++ b/crates/engine/src/cmd/denoise.rs @@ -0,0 +1,659 @@ +//! Denoise commands (`denoise.*`): the models (list, install, download, test, remove, choose), the settings, and the +//! work (queue, cancel, status, pump, clear). The machinery they drive is [`crate::denoise`]. +//! +//! Models live in the folder the host provides (`Session::denoise.models_dir`): one subfolder per model with +//! `model.onnx`, `denoise-model.json` (its manifest) and `installed.json` (what the user accepted and the file's hash), +//! plus `settings.json` beside them. Model files, archives 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::path::{Path, PathBuf}; +use std::sync::{Arc, PoisonError, mpsc}; + +use lightcraft_catalog::PhotoId; +use lightcraft_denoise::archive; +use lightcraft_denoise::hash::sha256_file; +use lightcraft_denoise::known::{self, Known}; +use lightcraft_denoise::manifest::{self, DenoiserManifest, MAX_MANIFEST_BYTES, MAX_MODEL_BYTES}; +use serde_json::{Value, json}; + +use super::{CommandSpec, always, bad, bool_or, cmd, str_param}; +use crate::denoise::{Pace, RunOn, Settings, forget_gpu_set_ups, installed_models, read_capped, read_settings, write_atomic, write_settings}; +use crate::model_download::State; +use crate::{EngineError, Result, Session}; + +fn fail(what: &str, e: impl std::fmt::Display) -> EngineError { + EngineError::Other(format!("{what}: {e}")) +} + +fn models_dir(s: &Session, cmd: &str) -> Result { + s.denoise.models_dir.clone().ok_or_else(|| bad(cmd, "this build has nowhere to keep denoise models (the desktop app does)")) +} + +/// What a command says about a model. +fn row(m: &DenoiserManifest, download_host: Option, installed: bool, selected: bool, accepted: &Value) -> Value { + json!({ + "id": m.id, + "name": m.name, + "version": m.version, + "licence": m.licence, + "provenance": m.provenance, + "source": m.source, + "sizeBytes": m.size_bytes, + "sha256": m.sha256, + "tile": m.tile, + "known": known::find(&m.id).is_some(), + // a pinned address LightCraft can fetch it from (the user presses Download), and which site that is + "downloadHost": download_host, + "installed": installed, + "selected": selected, + "accepted": accepted, + }) +} + +fn host_of(k: &Known) -> Option { + k.download.as_ref().map(|d| lightcraft_denoise::known::host(&d.url).to_string()) +} + +fn list(s: &mut Session, _: &Value) -> Result { + s.denoise_refresh_active(true); + let dir = s.denoise.models_dir.clone(); + let st = dir.as_deref().map(read_settings).unwrap_or_default(); + let installed = dir.as_deref().map(installed_models).unwrap_or_default(); + let on_disk = |id: &str| installed.iter().find(|i| i.manifest.id == id); + let chosen = |id: &str| st.model.as_deref() == Some(id); + let mut models: Vec = Vec::new(); + for k in known::all() { + let found = on_disk(&k.manifest.id); + models.push(row(&k.manifest, host_of(&k), found.is_some(), chosen(&k.manifest.id), found.map_or(&Value::Null, |f| &f.accepted))); + } + for i in installed.iter().filter(|i| known::find(&i.manifest.id).is_none()) { + models.push(row(&i.manifest, None, true, chosen(&i.manifest.id), &i.accepted)); + } + Ok(json!({ + "dir": dir.as_ref().map(|d| d.display().to_string()), + "productsDir": s.denoise_products_dir().map(|d| d.display().to_string()), + "model": st.model, + // whether this build can run denoise models (the `denoise` feature: our pure-Rust CPU/GPU runners) + "runtime": cfg!(feature = "denoise"), + "auto": st.auto(), + "cacheGb": st.cache_bytes() >> 30, + "threads": st.threads, + "runOn": st.run_on().name(), + "models": models, + })) +} + +/// What a file the user pointed at is, as far as installing it goes. +struct Source { + manifest: DenoiserManifest, + /// The `.onnx` inside the archive; `None` when the file is the `.onnx`. + entry: Option<&'static str>, + file_name: String, +} + +fn identify(path: &Path, c: &str, sha_known: Option<&str>) -> Result { + let meta = std::fs::metadata(path).map_err(|e| bad(c, format!("cannot read `{}`: {e}", path.display())))?; + if !meta.is_file() { + return Err(bad(c, format!("`{}` is not a file", path.display()))); + } + if meta.len() == 0 || meta.len() > MAX_MODEL_BYTES { + return Err(bad(c, "a model file must be between 1 byte and 2 GiB")); + } + let sha256 = match sha_known { + Some(h) => h.to_string(), + None => sha256_file(path).map_err(|e| fail(&format!("could not read {}", path.display()), e))?, + }; + let file_name = path.file_name().map(|n| n.to_string_lossy().into_owned()).unwrap_or_default(); + // an archive LightCraft knows (the file darktable publishes the model in) + if let Some(k) = known::all().into_iter().find(|k| k.download.as_ref().is_some_and(|d| d.sha256 == sha256)) { + return Ok(Source { manifest: k.manifest, entry: k.entry, file_name }); + } + // the `.onnx` of a known model + if let Some(m) = known::by_sha256(&sha256) { + return Ok(Source { manifest: m, entry: None, file_name }); + } + // an `.onnx` with its manifest beside it + let sibling = path.with_file_name("denoise-model.json"); + if path.extension().is_some_and(|e| e.eq_ignore_ascii_case("onnx")) + && let Some(bytes) = read_capped(&sibling, MAX_MANIFEST_BYTES) + { + let m = manifest::parse(&bytes).map_err(|e| bad(c, format!("{}: {e}", sibling.display())))?; + return Ok(Source { manifest: m, entry: None, file_name }); + } + Err(bad( + c, + "LightCraft does not know this file. A denoise model is an .onnx file with a denoise-model.json beside it that says how to use it \ + (see docs/denoise.md), or an archive LightCraft can fetch itself (`denoise.models.download`)", + )) +} + +/// Install the model in `src`. A model that passes its self-test becomes the one in use (`activate`, the default): installing +/// a model is how someone says they want it. `verified` is for a download whose hash was already checked. +fn install_file( + dir: &Path, + loader: &crate::denoise::Loader, + run_on: RunOn, + accepted_at: String, + c: &str, + src: &Path, + acknowledged: bool, + verified: bool, +) -> Result<(Value, Arc)> { + let ident = identify(src, c, None)?; + let m = ident.manifest.clone(); + 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 (see `denoise.models.list`) and pass `acknowledged: true` once they agree", + )); + } + let home = dir.join(&m.id); + if home.join("model.onnx").exists() { + return Err(bad(c, "that model is already installed; remove it before installing a replacement")); + } + std::fs::create_dir_all(&home).map_err(|e| fail("could not create the model's folder", e))?; + let mut cleanup = InstallCleanup { home: home.clone(), armed: true }; + let (part, final_path) = (home.join("model.onnx.part"), home.join("model.onnx")); + let _ = std::fs::remove_file(&part); + let placed: Result<()> = match ident.entry { + Some(entry) => { + archive::extract(src, entry, &part, MAX_MODEL_BYTES).map(|_| ()).map_err(|e| fail("could not take the model out of the archive", e)) + } + None if verified && std::fs::rename(src, &part).is_ok() => Ok(()), + None => std::fs::copy(src, &part).map(|_| ()).map_err(|e| fail("could not copy the model (is the disk full?)", e)), + }; + let sha256 = placed.and_then(|_| sha256_file(&part).map_err(|e| fail("could not check the model", e))).and_then(|h| match &m.sha256 { + Some(want) if *want != h => Err(EngineError::Other("the model file is not the one its description names (its SHA-256 differs)".into())), + _ => Ok(h), + }); + let sha256 = match sha256 { + Ok(h) => h, + Err(e) => { + let _ = std::fs::remove_file(&part); + return Err(e); + } + }; + // Test the staged file first. A failed replacement preserves the installed model. + let model = loader(&part, &m).map_err(|e| bad(c, e))?; + let test = match model.self_test(run_on) { + Ok(t) => t, + Err(e) => { + let _ = std::fs::remove_file(&part); + return Err(bad(c, e)); + } + }; + // An installed id is immutable: installing again must never destroy a usable model. + if final_path.exists() { + let _ = std::fs::remove_file(&part); + return Err(bad(c, "that model is already installed; remove it before installing a replacement")); + } + std::fs::rename(&part, &final_path).map_err(|e| fail("could not finish the copy", e))?; + write_atomic(&home.join("denoise-model.json"), &serde_json::to_vec_pretty(&m).map_err(|e| fail("manifest", e))?).map_err(EngineError::Other)?; + let accepted = json!({ + "acceptedAt": accepted_at, + "licence": m.licence.name, + "commercial": m.licence.commercial, + "fileName": ident.file_name, + "sha256": sha256, + "selfTest": test, + }); + write_atomic(&home.join("installed.json"), &serde_json::to_vec_pretty(&accepted).map_err(|e| fail("record", e))?).map_err(EngineError::Other)?; + cleanup.armed = false; + Ok((json!({"installed": row(&m, known::find(&m.id).and_then(|k| host_of(&k)), true, false, &accepted), "model": m.id}), model)) +} + +/// Removes only files created for this new installation, including on a worker panic. +struct InstallCleanup { + home: PathBuf, + armed: bool, +} +impl Drop for InstallCleanup { + fn drop(&mut self) { + if self.armed { + for name in ["model.onnx.part", "model.onnx", "denoise-model.json", "installed.json"] { + let _ = std::fs::remove_file(self.home.join(name)); + } + let _ = std::fs::remove_dir(&self.home); + } + } +} + +pub(crate) struct InstallJob { + id: String, + activate: bool, + result: mpsc::Receiver), String>>, +} + +fn activate_install(s: &mut Session, result: &Value, model: Arc, activate: bool) -> Result { + let dir = models_dir(s, "denoise.models.install")?; + let id = result["installed"]["id"].as_str().ok_or_else(|| bad("denoise.models.install", "missing installed id"))?; + let mut st = read_settings(&dir); + if activate { + st.model = Some(id.to_string()); + write_settings(&dir, &st).map_err(EngineError::Other)?; + } + s.denoise.touch(); + s.denoise_refresh_active(true); + if let Some(a) = &s.denoise.active + && a.id == id + { + *a.model.lock().unwrap_or_else(PoisonError::into_inner) = Some(model); + } + let mut v = result.clone(); + v["installed"]["selected"] = json!(st.model.as_deref() == Some(id)); + v["model"] = json!(st.model); + Ok(v) +} + +fn start_install(s: &mut Session, src: PathBuf, id: String, activate: bool, verified: bool) -> Result { + const C: &str = "denoise.models.install"; + if s.denoise.install.is_some() { + return Err(bad(C, "another model is being installed")); + } + let dir = models_dir(s, C)?; + let loader = s.denoise.loader.clone(); + let run_on = s.denoise.settings.run_on(); + let at = (s.clock)(); + let (tx, rx) = mpsc::channel(); + std::thread::Builder::new() + .name("denoise-model-install".into()) + .stack_size(16 * 1024 * 1024) + .spawn(move || { + let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _held = crate::memory::work_gate().acquire(512 * 1024 * 1024); + install_file(&dir, &loader, run_on, at, C, &src, true, verified).map_err(|e| plain(&e)) + })) + .unwrap_or_else(|_| Err("the model installation stopped unexpectedly".into())); + if verified { + let _ = std::fs::remove_file(&src); + } + let _ = tx.send(outcome); + }) + .map_err(|e| bad(C, format!("could not start installation: {e}")))?; + s.denoise.install = Some(InstallJob { id: id.clone(), activate, result: rx }); + s.denoise.downloads.set_outcome(&id, State::Installing); + Ok(json!({"started": id})) +} + +fn install(s: &mut Session, p: &Value) -> Result { + const C: &str = "denoise.models.install"; + let path = str_param(p, "path").ok_or_else(|| bad(C, "missing `path`"))?; + let acknowledged = bool_or(p, "acknowledged", false); + let activate = bool_or(p, "activate", true); + if !acknowledged { + return Err(bad(C, "the licence has not been accepted: show its terms first")); + } + if bool_or(p, "background", false) { + return start_install(s, PathBuf::from(path), "local-install".into(), activate, false); + } + if s.denoise.install.is_some() { + return Err(bad(C, "another model is being installed")); + } + let (result, model) = + install_file(&models_dir(s, C)?, &s.denoise.loader, s.denoise.settings.run_on(), (s.clock)(), C, Path::new(path), acknowledged, false)?; + activate_install(s, &result, model, activate) +} + +/// Read only the small adjacent manifest; hashing and execution happen in the install worker. +fn inspect(s: &mut Session, p: &Value) -> Result { + const C: &str = "denoise.models.inspect"; + let path = str_param(p, "path").ok_or_else(|| bad(C, "missing `path`"))?; + let src = Path::new(path); + if !src.extension().is_some_and(|e| e.eq_ignore_ascii_case("onnx")) { + return Err(bad(C, "choose an .onnx file with denoise-model.json beside it")); + } + let bytes = read_capped(&src.with_file_name("denoise-model.json"), MAX_MANIFEST_BYTES) + .ok_or_else(|| bad(C, "denoise-model.json is missing or too large"))?; + let m = manifest::parse(&bytes).map_err(|e| bad(C, e.to_string()))?; + let dir = models_dir(s, C)?; + Ok(json!({"kind": "model", "domain": "denoise", "path": path, "fileName": src.file_name().map(|n| n.to_string_lossy()), "model": m, + "sizeBytes": std::fs::metadata(src).ok().map(|m| m.len()), "alreadyInstalled": dir.join(&m.id).join("model.onnx").is_file()})) +} + +/// 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 and chosen without +/// another question. Called by `denoise.models.downloads` and by every frame's `denoise.pump`. +pub(crate) fn finish_downloads(s: &mut Session) { + if let Some(job) = &s.denoise.install { + let outcome = match job.result.try_recv() { + Ok(v) => Some(v), + Err(mpsc::TryRecvError::Empty) => None, + Err(mpsc::TryRecvError::Disconnected) => Some(Err("the installation worker stopped".into())), + }; + if let Some(outcome) = outcome + && let Some(job) = s.denoise.install.take() + { + let state = match outcome { + Ok((v, model)) => match activate_install(s, &v, model, job.activate) { + Ok(_) => State::Installed, + Err(e) => State::Failed(plain(&e)), + }, + Err(e) => State::Failed(e), + }; + s.denoise.downloads.set_outcome(&job.id, state); + } + } + if s.denoise.install.is_none() + && let Some((id, path, _sha)) = s.denoise.downloads.finished().into_iter().next() + && let Err(e) = start_install(s, path, id.clone(), true, true) + { + s.denoise.downloads.set_outcome(&id, State::Failed(plain(&e))); + } +} + +fn download(s: &mut Session, p: &Value) -> Result { + const C: &str = "denoise.models.download"; + let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?; + let dir = models_dir(s, C)?; + let k = known::find(id).ok_or_else(|| bad(C, "LightCraft has no download for that model: get the file from its page, then add it"))?; + let spec = k.download.clone().ok_or_else(|| bad(C, "LightCraft has no download for that model: get the file from its page, then add it"))?; + if !cfg!(feature = "denoise") { + return Err(bad(C, "this build cannot run denoise models, so there is nothing to download for it")); + } + if !bool_or(p, "acknowledged", false) { + return Err(bad( + C, + "the licence has not been accepted: show the user the model's terms (see `denoise.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 = lightcraft_denoise::known::host(&spec.url).to_string(); + s.denoise.downloads.start(spec, &dir).map_err(|e| bad(C, e))?; + Ok(json!({"started": id, "from": host})) +} + +fn download_row(id: &str, state: &State) -> Value { + let from = known::find(id).and_then(|k| host_of(&k)); + 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::Installing => json!({"id": id, "state": "installing", "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}), + } +} + +fn downloads(s: &mut Session, _: &Value) -> Result { + finish_downloads(s); + let all: Vec = s.denoise.downloads.snapshot().iter().map(|(id, st)| download_row(id, st)).collect(); + Ok(json!({"running": s.denoise.downloads.running(), "downloads": all})) +} + +fn download_cancel(s: &mut Session, p: &Value) -> Result { + const C: &str = "denoise.models.downloadCancel"; + let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?; + Ok(json!({"discarded": s.denoise.downloads.discard(id)})) +} + +fn test(s: &mut Session, p: &Value) -> Result { + const C: &str = "denoise.models.test"; + let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?; + let dir = models_dir(s, C)?; + let found = installed_models(&dir).into_iter().find(|i| i.manifest.id == id).ok_or_else(|| bad(C, "that model is not installed"))?; + let run_on = s.denoise.settings.run_on(); + let result = (s.denoise.loader)(&found.onnx, &found.manifest).and_then(|m| m.self_test(run_on)); + 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(&found.onnx.with_file_name("installed.json"), &serde_json::to_vec_pretty(&record).map_err(|e| fail("record", e))?) + .map_err(EngineError::Other)?; + } + Ok(json!({"id": id, "ok": ok, "result": detail})) +} + +fn remove(s: &mut Session, p: &Value) -> Result { + const C: &str = "denoise.models.remove"; + let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?; + let dir = models_dir(s, C)?; + if !lightcraft_denoise::licence::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("denoise-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.model.as_deref() == Some(id) { + st.model = installed_models(&dir).into_iter().map(|i| i.manifest.id).find(|m| m != id); + write_settings(&dir, &st).map_err(EngineError::Other)?; + } + s.denoise.touch(); + s.denoise_refresh_active(true); + Ok(json!({"removed": id, "model": st.model})) +} + +fn select(s: &mut Session, p: &Value) -> Result { + const C: &str = "denoise.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) + { + return Err(bad(C, "that denoise model is not installed")); + } + let mut st = read_settings(&dir); + st.model = id; + write_settings(&dir, &st).map_err(EngineError::Other)?; + s.denoise.touch(); + s.denoise_refresh_active(true); + Ok(json!({"model": st.model})) +} + +fn settings(s: &mut Session, p: &Value) -> Result { + const C: &str = "denoise.settings"; + let dir = models_dir(s, C)?; + let mut st: Settings = read_settings(&dir); + if let Some(v) = p.get("auto") { + st.auto = Some(v.as_bool().ok_or_else(|| bad(C, "`auto` must be true or false"))?); + } + if let Some(v) = p.get("cacheGb") { + let n = + v.as_u64().filter(|n| (1..=10_000).contains(n)).ok_or_else(|| bad(C, "`cacheGb` must be a whole number of gigabytes from 1 to 10000"))?; + st.cache_gb = Some(n as u32); + } + if let Some(v) = p.get("threads") { + st.threads = match v { + Value::Null => None, + v => Some(v.as_u64().filter(|n| (1..=64).contains(n)).ok_or_else(|| bad(C, "`threads` must be from 1 to 64, or null"))? as u32), + }; + } + let run_on = match p.get("runOn") { + None => None, + Some(v) => Some(v.as_str().and_then(RunOn::parse).ok_or_else(|| bad(C, "`runOn` must be auto, gpu or cpu"))?), + }; + if let Some(r) = run_on { + st.run_on = Some(r.name().to_string()); + st.gpu = None; + } + write_settings(&dir, &st).map_err(EngineError::Other)?; + if run_on.is_some() { + // a choice of where denoise runs is also a go-ahead to try the card again: a set-up that crashed or gave up + // before is forgotten, and the loaded model sets the card up afresh for the next photo + forget_gpu_set_ups(&dir); + if let Some(a) = &s.denoise.active { + *a.model.lock().unwrap_or_else(PoisonError::into_inner) = None; + } + } + s.denoise.touch(); + s.denoise_refresh_active(true); + // a smaller limit takes effect now + s.denoise_evict(None); + Ok(json!({"auto": st.auto(), "cacheGb": st.cache_bytes() >> 30, "threads": st.threads, "runOn": st.run_on().name()})) +} + +fn ids_param(p: &Value) -> Option> { + p.get("ids").and_then(Value::as_array).map(|a| a.iter().filter_map(Value::as_u64).map(PhotoId).collect()) +} + +fn status_json(s: &Session, photo: Option) -> Value { + let (files, bytes) = s.denoise_cache_usage(); + let st = &s.denoise.settings; + let model = s.denoise.active.as_ref().map(|a| json!({"id": a.id, "name": a.manifest.name})); + let mut v = json!({ + "enabled": s.denoise.active.is_some(), + "model": model, + "runtime": cfg!(feature = "denoise"), + "productsDir": s.denoise_products_dir().map(|d| d.display().to_string()), + "queued": s.denoise.queued(), + "running": s.denoise_running_json(), + "ready": s.media.denoise.len(), + "made": s.denoise_made(), + "auto": st.auto(), + "runOn": st.run_on().name(), + "device": s.denoise_device(), + "cache": {"files": files, "bytes": bytes, "limitBytes": st.cache_bytes()}, + "generation": s.denoise.generation, + }); + if let (Some(id), Some(o)) = (photo, v.as_object_mut()) { + o.insert("photo".into(), serde_json::to_value(s.denoise_photo_state(id)).unwrap_or(Value::Null)); + } + v +} + +/// A per-photo switch, independent of the selection and Auto Sync. Turning on starts at 50%. +fn toggle(s: &mut Session, p: &Value) -> Result { + const C: &str = "denoise.toggle"; + let id = match p.get("id") { + Some(v) => PhotoId(v.as_u64().ok_or_else(|| bad(C, "`id` must be a photo id"))?), + None => s.active().ok_or_else(|| bad(C, "select a raw photo"))?, + }; + let photo = s.catalog.photo(id).ok_or_else(|| bad(C, "no such photo"))?; + if !crate::denoise::eligible(photo) { + return Err(bad(C, "AI Denoise applies to raw photos developed from their sensor data")); + } + let mut d = (*s.develop_of(id).ok_or_else(|| bad(C, "no such photo"))?).clone(); + let on = match p.get("enabled") { + Some(v) => v.as_bool().ok_or_else(|| bad(C, "`enabled` must be true or false"))?, + None => !d.enhance.denoise_enabled(), + }; + if on { + if !d.enhance.denoise.is_finite() { + return Err(bad(C, "the photo has an invalid Denoise amount")); + } + s.denoise_refresh_active(true); + if s.denoise.active.is_none() { + return Err(bad(C, "AI Denoise needs a model: install one in Settings > AI Denoise")); + } + } + // The switch is its own setting: the Amount stays where it is (0 included) when it is flipped, so flipping compares + // with and without. A photo that never had an Amount starts at 50 the first time. + let amount = if on && d.enhance.denoise_on.is_none() && d.enhance.denoise <= 0.0 { 50.0 } else { d.enhance.denoise }; + if d.enhance.denoise != amount || d.enhance.denoise_on != Some(on) { + s.end_interaction()?; + d.enhance.denoise = amount; + d.enhance.denoise_on = Some(on); + let op = s.develop_op(id, d, "AI Denoise").ok_or_else(|| bad(C, "no such photo"))?; + s.commit("AI Denoise", op)?; + } + if on { + s.denoise_enqueue(&[id], true, false); + } else { + s.denoise_cancel(Some(id), false); + } + Ok(json!({"id": id.0, "enabled": on, "amount": amount})) +} + +fn queue(s: &mut Session, p: &Value) -> Result { + const C: &str = "denoise.queue"; + s.denoise_refresh_active(true); + if s.denoise.active.is_none() { + return Err(bad(C, "no denoise model is chosen: install one first (`denoise.models.download` or `denoise.models.install`)")); + } + if s.denoise_products_dir().is_none() { + return Err(bad(C, "denoise keeps its pictures in the library's folder: open a library on disk first")); + } + let ids = match ids_param(p) { + Some(ids) => ids, + None => match str_param(p, "scope").unwrap_or("selected") { + "selected" => s.selection.active.iter().chain(s.selection.ids.iter()).copied().collect(), + "visible" => s.visible_cloned().to_vec(), + "withAmount" => s.catalog.photos().filter(|ph| ph.develop.denoise_amount() > 0.0).map(|ph| ph.id).collect(), + other => return Err(bad(C, format!("unknown scope `{other}`: use selected, visible or withAmount"))), + }, + }; + let queued = s.denoise_enqueue(&ids, false, bool_or(p, "retry", false)); + let mut v = status_json(s, None); + if let Some(o) = v.as_object_mut() { + o.insert("added".into(), json!(queued)); + } + Ok(v) +} + +fn cancel(s: &mut Session, p: &Value) -> Result { + let photo = p.get("id").and_then(Value::as_u64).map(PhotoId); + let n = s.denoise_cancel(photo, bool_or(p, "all", photo.is_none())); + Ok(json!({"cancelled": n})) +} + +fn status(s: &mut Session, p: &Value) -> Result { + s.denoise_refresh_active(false); + Ok(status_json(s, p.get("id").and_then(Value::as_u64).map(PhotoId))) +} + +/// `denoise.pump {pace?}`: cheap to call every frame. Takes in what the background job has finished and keeps one going. +fn pump(s: &mut Session, p: &Value) -> Result { + s.denoise_pump(Pace::parse(p.get("pace").and_then(Value::as_str))); + Ok(json!({ + "active": s.denoise.active.is_some(), + "running": s.denoise_running_json(), + "queued": s.denoise.queued(), + "ready": s.media.denoise.len(), + "generation": s.denoise.generation, + "downloading": s.denoise.downloads.running(), + })) +} + +fn clear(s: &mut Session, _: &Value) -> Result { + let (files, bytes) = s.denoise_clear(); + Ok(json!({"deleted": files, "bytes": bytes})) +} + +pub fn specs() -> Vec { + vec![ + cmd!(query "denoise.models.list", "Denoise Models", [], None, "{} → {dir, productsDir, model, runtime, auto, cacheGb, threads, runOn, models: [{id, name, version, licence{name, commercial, url, notice}, provenance, source, sizeBytes, sha256, tile, known, downloadHost, installed, selected, accepted}]}", always, list), + cmd!(query "denoise.models.inspect", "Inspect Denoise Model", [], None, "{path} → model terms from an adjacent manifest; never executes the file", always, inspect), + cmd!(query "denoise.models.install", "Install Denoise Model", [], None, "{path, acknowledged: true, activate?: true, background?: false} → {installed, model} or {started} — install a denoise model from a file: an .onnx with a denoise-model.json beside it, or an archive LightCraft knows (the darktable `.dtmodel`). `acknowledged` must be true: the user has been shown the model's licence and accepted it. Unless `activate` is false it becomes the model in use", always, install), + cmd!(query "denoise.models.download", "Download Denoise Model", [], None, "{id, acknowledged: true} → {started, from} — fetch a model LightCraft has a pinned address for (see `downloadHost` in the list) in the background over HTTPS with 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, installed and chosen by itself; `denoise.models.downloads` shows how far it is", always, download), + cmd!(query "denoise.models.downloads", "Denoise Model Downloads", [], None, "{} → {running, downloads: [{id, state: running | done | installed | failed | cancelled, bytes, total, error, from}]} — also installs any download that has arrived", always, downloads), + cmd!(query "denoise.models.downloadCancel", "Cancel Denoise Model Download", [], None, "{id} → {discarded} — stop a download, or delete a finished one that was not installed", always, download_cancel), + cmd!(query "denoise.models.test", "Test Denoise Model", [], None, "{id} → {ok, result} — load an installed model and check it gives sensible pictures; needs the denoise runtime", always, test), + cmd!(query "denoise.models.remove", "Remove Denoise Model", [], None, "{id} → {removed, model} — delete an installed model (its cached pictures stay until they are evicted)", always, remove), + cmd!(query "denoise.models.select", "Choose Denoise Model", [], None, "{id: installed model | null} → {model} — null switches denoise off", always, select), + cmd!(query "denoise.settings", "Denoise Settings", [], None, "{auto?: bool, cacheGb?: 1..10000, threads?: 1..64 | null, runOn?: auto | gpu | cpu} → the settings — `auto`: make the picture of the photos being looked at that have a Denoise amount; `cacheGb`: how much the cached pictures may take; `threads`: most tiles run at once; `runOn`: where the model runs — `auto` (the default) the graphics card where it is faster than the processor (both timed on this computer when the card is set up), `gpu` the card whenever it can run the model, `cpu` the processor. Setting `runOn` also lets the card be tried again after a set-up that failed or closed LightCraft", always, settings), + cmd!( + "denoise.toggle", + "AI Denoise", + [], + None, + "{id?: photoId, enabled?: bool} → {id, enabled, amount} — turn AI Denoise on (50% initially) or off for one raw photo, independent of selection and Auto Sync; needs an installed model to turn on", + always, + toggle + ), + cmd!(query "denoise.queue", "Denoise Photos", [], None, "{ids?: [photoId], scope?: selected | visible | withAmount, retry?: false} → {added, …status} — make the denoised pictures of raw photos that have none, in the background. `retry` tries photos that failed before", always, queue), + cmd!(query "denoise.cancel", "Cancel Denoise", [], None, "{id?, all?} → {cancelled} — stop the photo being made and/or forget the ones waiting (all of them when `id` is omitted)", always, cancel), + cmd!(query "denoise.status", "Denoise Status", [], None, "{id?} → {enabled, model, runtime, productsDir, queued, running: {photo, done, total, seconds} | null, ready, made, auto, runOn, device: {kind: gpu | cpu | pending | none, adapter?, reason?, cardMs?, cpuMs?, fellBack?}, cache: {files, bytes, limitBytes}, generation, photo?: {state: notApplicable | noModel | ready | queued | running | failed | idle, …}}", always, status), + cmd!(query "denoise.pump", "Denoise Pump", [], None, "{pace?: pause | light | normal | full} → {active, running, queued, ready, generation, downloading} — called every frame by the app: takes in finished work and keeps one photo going", always, pump), + cmd!(query "denoise.clear", "Clear Denoise Cache", [], None, "{} → {deleted, bytes} — delete every cached denoised picture of this library (they are made again when wanted)", always, clear), + ] +} diff --git a/crates/engine/src/cmd/face_detect.rs b/crates/engine/src/cmd/face_detect.rs new file mode 100644 index 00000000..ffdedbef --- /dev/null +++ b/crates/engine/src/cmd/face_detect.rs @@ -0,0 +1,184 @@ +//! `faces.detect`: find faces with the YuNet detector the user has downloaded (Settings > Faces). +//! +//! The photo is rendered upright and uncropped with default settings (so edits, crops and rotations you have +//! made do not matter), the detector runs on it, and the faces come back as fractions of the photo's width +//! and height, the same frame face regions from XMP use. Unless `apply` is false, they are added to the photo as +//! unnamed face regions in one undoable step; earlier detections are replaced, and regions that came from XMP (or +//! that you drew or named) are never touched. Like every region edit it is catalog-only: no sidecar is written. + +use std::path::PathBuf; +use std::sync::{Arc, Mutex, PoisonError}; +use std::time::Instant; + +use lightcraft_catalog::Op; +use lightcraft_develop::DevelopSettings; +use lightcraft_faces::known::YUNET_ID; +use lightcraft_faces::yunet::{Detector, Face, Options}; +use lightcraft_meta::{Region, RegionKind}; +use serde_json::{Value, json}; + +use super::{CommandSpec, bad, cmd, has_selection}; +use crate::{Result, Session}; + +const C: &str = "faces.detect"; +/// Regions made by this command say so in their description; only those are replaced on a new run. +pub(crate) const MARK: &str = "Detected by "; +pub(crate) const LABEL: &str = "YuNet 2023mar"; +/// Long edge of the picture the detector looks at (it shrinks it to its own 640 anyway). +const LOOK_EDGE: usize = 1280; +/// YuNet is 232 KB; a file this much larger in its folder is not it. +const MAX_MODEL_BYTES: u64 = 8 << 20; +/// How much a detection must overlap a region the photo already has to count as the same face. +const EXISTING_IOU: f64 = 0.3; + +/// What to tell someone who has not downloaded the detector yet. +const NOT_INSTALLED: &str = "finding faces needs the YuNet face detector (232 KB, MIT): download it in Settings > Faces"; + +/// The installed YuNet, loaded once and kept (a new one is loaded when the model folder changes). +/// (also used by the background scan, which aligns faces by the detector's landmarks and finds faces in photos that have none) +pub(crate) fn detector(s: &Session) -> std::result::Result, String> { + static LOADED: Mutex)>> = Mutex::new(None); + let home = s.face_models_dir.as_deref().ok_or("this build has nowhere to keep face models (the desktop app does)")?.join(YUNET_ID); + let path = home.join("model.onnx"); + // installing writes the manifest last, so a folder without one is not (yet) a model + if !path.is_file() || !home.join("face-model.json").is_file() { + return Err(NOT_INSTALLED.into()); + } + let mut loaded = LOADED.lock().unwrap_or_else(PoisonError::into_inner); + if let Some((p, d)) = loaded.as_ref() + && *p == path + { + return Ok(d.clone()); + } + let len = std::fs::metadata(&path).map_err(|e| format!("could not read the detector: {e}"))?.len(); + if len > MAX_MODEL_BYTES { + return Err("the detector's file is far bigger than YuNet: remove it in Settings > Faces and download it again".into()); + } + let bytes = std::fs::read(&path).map_err(|e| format!("could not read the detector: {e}"))?; + let d = Arc::new(Detector::new(&bytes).map_err(|e| format!("{e}: remove the detector in Settings > Faces and download it again"))?); + *loaded = Some((path, d.clone())); + Ok(d) +} + +fn face_json(f: &Face) -> Value { + json!({ + "rect": {"x0": f.x0, "y0": f.y0, "x1": f.x1, "y1": f.y1}, + "score": f.score, + "landmarks": f.landmarks.iter().map(|(x, y)| json!([x, y])).collect::>(), + }) +} + +/// A detected face that lies on a region the photo already has (from XMP, drawn or named) is that region, not a new face. +fn overlaps_existing(f: &Face, regions: &[Region]) -> bool { + regions.iter().any(|r| { + let (w, h) = ( + (r.rect.x1.min(f64::from(f.x1)) - r.rect.x0.max(f64::from(f.x0))).max(0.0), + (r.rect.y1.min(f64::from(f.y1)) - r.rect.y0.max(f64::from(f.y0))).max(0.0), + ); + let inter = w * h; + let union = (r.rect.x1 - r.rect.x0) * (r.rect.y1 - r.rect.y0) + f64::from((f.x1 - f.x0) * (f.y1 - f.y0)) - inter; + union > 0.0 && inter / union >= EXISTING_IOU + }) +} + +pub(crate) fn is_detected(r: &Region) -> bool { + r.description.as_deref().is_some_and(|d| d.starts_with(MARK)) +} + +/// The photo as the faces see it: upright, uncropped, default settings, as 8-bit RGB with its size, its long edge at +/// most `edge`. +pub(crate) fn render_rgb(s: &mut Session, id: lightcraft_catalog::PhotoId, edge: usize) -> std::result::Result<(Vec, usize, usize), String> { + let settings = DevelopSettings::default(); + let job = s.preview_job(id, edge, edge, false, &settings).ok_or("no such photo")?; + let image = job.run().rendered?.image; + let rgb: Vec = image.data.iter().flat_map(|px| [px[0], px[1], px[2]]).collect(); + Ok((rgb, image.width, image.height)) +} + +fn detect(s: &mut Session, p: &Value) -> Result { + let apply = p.get("apply").and_then(Value::as_bool).unwrap_or(true); + let defaults = Options::default(); + let opts = Options { + score: p.get("score").and_then(Value::as_f64).map_or(defaults.score, |v| v as f32), + nms_iou: p.get("nmsIou").and_then(Value::as_f64).map_or(defaults.nms_iou, |v| v as f32), + max_faces: p.get("maxFaces").and_then(Value::as_u64).map_or(defaults.max_faces, |v| usize::try_from(v).unwrap_or(defaults.max_faces)), + }; + let detector = detector(s).map_err(|e| bad(C, e))?; + let started = Instant::now(); + let (mut results, mut ops) = (Vec::new(), Vec::new()); + for id in s.targets(p) { + let (rgb, width, height) = match render_rgb(s, id, LOOK_EDGE) { + Ok(r) => r, + Err(e) => { + results.push(json!({"id": id.0, "error": e})); + continue; + } + }; + let faces = match detector.detect(&rgb, width, height, &opts) { + Ok(f) => f, + Err(e) => { + results.push(json!({"id": id.0, "error": e.to_string()})); + continue; + } + }; + if apply && let Some(photo) = s.catalog.photo(id) { + let mut meta = photo.meta.clone(); + meta.regions.retain(|r| !is_detected(r)); + let fresh: Vec<&Face> = faces.iter().filter(|f| !overlaps_existing(f, &meta.regions)).collect(); + for f in fresh { + meta.regions.push(Region { + rect: lightcraft_geom::Rect { x0: f64::from(f.x0), y0: f64::from(f.y0), x1: f64::from(f.x1), y1: f64::from(f.y1) }, + kind: RegionKind::Face, + name: None, + description: Some(format!("{MARK}{LABEL}")), + }); + } + ops.push(Op::SetMeta { id, meta: Box::new(meta) }); + } + results.push(json!({"id": id.0, "faces": faces.iter().map(face_json).collect::>()})); + } + let applied = !ops.is_empty(); + if applied { + s.commit("Detect Faces", Op::Batch { ops })?; + // regions are not written to XMP, so a sidecar must not be rewritten for this + s.skip_auto_write = true; + } + Ok(json!({"detector": LABEL, "applied": applied, "photos": results, "ms": started.elapsed().as_secs_f64() * 1000.0})) +} + +pub fn specs() -> Vec { + vec![cmd!( + "faces.detect", + "Detect Faces", + ["Photo"], + None, + "{ids?, apply?: bool (default true), score?: 0.6, nmsIou?: 0.3, maxFaces?} → {photos: [{id, faces: [{rect, score, landmarks}]}]}. Finds faces with the YuNet detector (download it first: `faces.models.download {id: \"yunet-2023mar\"}`) and, unless `apply` is false, adds them as unnamed face regions (one undo step), replacing earlier detections but never regions from XMP or ones you named; the sidecar is not touched", + has_selection, + detect + )] +} + +#[cfg(test)] +mod tests { + use super::*; + + fn face(x0: f32, y0: f32, x1: f32, y1: f32) -> Face { + Face { x0, y0, x1, y1, score: 0.9, landmarks: [(0.0, 0.0); 5] } + } + fn region(x0: f64, y0: f64, x1: f64, y1: f64) -> Region { + Region { rect: lightcraft_geom::Rect { x0, y0, x1, y1 }, kind: RegionKind::Face, name: Some("Ann".into()), description: None } + } + + #[test] + fn a_detection_on_an_existing_region_is_not_a_new_face() { + let existing = [region(0.40, 0.20, 0.60, 0.50), region(0.10, 0.10, 0.20, 0.20)]; + // the same face, a bit larger or shifted (a Lightroom box is not YuNet's box) + assert!(overlaps_existing(&face(0.38, 0.18, 0.64, 0.54), &existing)); + assert!(overlaps_existing(&face(0.42, 0.22, 0.58, 0.48), &existing)); + // another face, a sliver of overlap, and no regions at all + assert!(!overlaps_existing(&face(0.70, 0.20, 0.90, 0.50), &existing)); + assert!(!overlaps_existing(&face(0.59, 0.49, 0.80, 0.80), &existing)); + assert!(!overlaps_existing(&face(0.0, 0.0, 0.0, 0.0), &existing), "an empty box overlaps nothing"); + assert!(!overlaps_existing(&face(0.4, 0.2, 0.6, 0.5), &[])); + } +} diff --git a/crates/engine/src/cmd/face_models.rs b/crates/engine/src/cmd/face_models.rs new file mode 100644 index 00000000..44b359e4 --- /dev/null +++ b/crates/engine/src/cmd/face_models.rs @@ -0,0 +1,511 @@ +//! Face model commands (`faces.*`): which models exist, inspecting a file the user dropped in, installing +//! it once its licence is accepted, removing it, and the faces on/off setting. +//! +//! Models live in the folder the host provides (`Session::face_models_dir`): one subfolder per model with +//! `model.onnx`, `face-model.json` (its manifest) and `installed.json` (what the user accepted), plus +//! `settings.json` beside them. Model files and manifests are hostile input: sizes are capped, ids are +//! validated before they become folder names, and nothing outside a model's own folder is ever touched. + +use std::io::Read; +use std::path::{Path, PathBuf}; + +use lightcraft_faces::catalog::{self, Catalog}; +use lightcraft_faces::hash::sha256_file; +use lightcraft_faces::manifest::{self, MAX_MANIFEST_BYTES, MAX_MODEL_BYTES, ModelManifest, Role}; +use lightcraft_faces::suggest::{Suggestion, suggest_with}; +use lightcraft_faces::{known, onnx}; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; + +use super::{CommandSpec, always, bad, cmd, str_param}; +use crate::face_download::{self, State}; +use crate::{EngineError, Result, Session}; + +#[derive(Default, Serialize, Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub(super) struct FaceSettings { + pub(super) enabled: bool, + /// The installed recogniser in use. + pub(super) embedder: Option, +} + +fn fail(what: &str, e: impl std::fmt::Display) -> EngineError { + EngineError::Other(format!("{what}: {e}")) +} + +fn models_dir(s: &Session, cmd: &str) -> Result { + s.face_models_dir.clone().ok_or_else(|| bad(cmd, "this build has nowhere to keep face models (the desktop app does)")) +} + +/// Read at most `max` bytes of a small file. +fn read_capped(path: &Path, max: usize) -> Option> { + let mut buf = Vec::new(); + std::fs::File::open(path).ok()?.take(max as u64 + 1).read_to_end(&mut buf).ok()?; + (buf.len() <= max).then_some(buf) +} + +pub(super) fn read_settings(dir: &Path) -> FaceSettings { + read_capped(&dir.join("settings.json"), 4096).and_then(|b| serde_json::from_slice(&b).ok()).unwrap_or_default() +} + +fn write_atomic(path: &Path, bytes: &[u8]) -> Result<()> { + let part = path.with_extension("part"); + std::fs::write(&part, bytes).map_err(|e| fail(&format!("could not write {}", path.display()), e))?; + std::fs::rename(&part, path).map_err(|e| { + let _ = std::fs::remove_file(&part); + fail(&format!("could not write {}", path.display()), e) + }) +} + +fn write_settings(dir: &Path, st: &FaceSettings) -> Result<()> { + std::fs::create_dir_all(dir).map_err(|e| fail("could not create the face models folder", e))?; + write_atomic(&dir.join("settings.json"), &serde_json::to_vec_pretty(st).map_err(|e| fail("settings", e))?) +} + +pub(super) struct Installed { + pub(super) manifest: ModelManifest, + pub(super) accepted: Value, +} + +/// The models in the folder. A folder that is not a valid model (wrong name, no model file, a manifest that +/// does not validate or does not match its folder) is ignored, never trusted. +pub(super) fn installed_models(dir: &Path) -> Vec { + let Ok(entries) = std::fs::read_dir(dir) else { return Vec::new() }; + let mut out = Vec::new(); + for e in entries.flatten() { + let p = e.path(); + let Some(name) = p.file_name().and_then(|n| n.to_str()).map(str::to_string) else { continue }; + if !manifest::valid_id(&name) || !p.is_dir() || !p.join("model.onnx").is_file() { + continue; + } + let Some(m) = read_capped(&p.join("face-model.json"), MAX_MANIFEST_BYTES).and_then(|b| manifest::parse(&b).ok()) else { continue }; + if m.id != name { + continue; + } + let accepted = read_capped(&p.join("installed.json"), 4096).and_then(|b| serde_json::from_slice(&b).ok()).unwrap_or(Value::Null); + out.push(Installed { manifest: m, accepted }); + } + out.sort_by(|a, b| a.manifest.id.cmp(&b.manifest.id)); + out +} + +/// Read the user's catalog (`catalog.json` in the models folder) anew: a few kilobytes, so it is always current. +fn refresh_catalog(s: &mut Session) { + let Some(dir) = s.face_models_dir.clone() else { return }; + let file = dir.join("catalog.json"); + s.face_catalog = match read_capped(&file, catalog::MAX_CATALOG_BYTES) { + Some(bytes) => catalog::parse(&bytes), + None if file.is_file() => catalog::parse(&vec![b' '; catalog::MAX_CATALOG_BYTES + 1]), + None => Catalog::default(), + }; +} + +/// The manifests of the catalog's models (what a file's hash is matched against, besides the built-in ones). +fn catalog_manifests(s: &Session) -> Vec { + s.face_catalog.entries.iter().map(|e| e.manifest.clone()).collect() +} + +/// How to fetch a model: the built-in address, or the one in the user's catalog. +fn download_spec(s: &Session, id: &str) -> Option { + known::download(id).or_else(|| s.face_catalog.entries.iter().find(|e| e.manifest.id == id).and_then(catalog::Entry::download)) +} + +fn row(s: &Session, m: &ModelManifest, installed: bool, selected: bool, accepted: &Value) -> Value { + let from_catalog = s.face_catalog.entries.iter().any(|e| e.manifest.id == m.id); + json!({ + "id": m.id, + "name": m.name, + "role": m.role, + "licence": m.licence, + "provenance": m.provenance, + "source": m.source, + "sizeBytes": m.size_bytes, + "sha256": m.sha256, + // how fast it is as a multiple of a ResNet-100 model (computers differ, ratios between models hardly do), and + // that in words; null when nobody has measured it + "speed": m.speed, + "speedText": m.speed.map(manifest::describe_speed).filter(|t| !t.is_empty()), + "known": from_catalog || known::all().iter().any(|k| k.id == m.id), + // from the user's own catalog file + "fromCatalog": from_catalog, + // a pinned address LightCraft can fetch it from (the user presses Download), and which site that is + "downloadHost": download_spec(s, &m.id).map(|d| known::host(&d.url).to_string()), + "installed": installed, + "selected": selected, + "accepted": accepted, + }) +} + +struct Inspection { + sha256: String, + size: u64, + file_name: String, + suggestion: Suggestion, +} + +/// Look at a model file. `known_sha` is its SHA-256 when something already checked it (a finished download), so a +/// file of hundreds of megabytes is not read again just to be hashed. +fn inspect_file(path: &str, cmd: &str, known_sha: Option<&str>, extra: &[ModelManifest]) -> Result { + let p = Path::new(path); + let meta = std::fs::metadata(p).map_err(|e| bad(cmd, format!("cannot read `{path}`: {e}")))?; + if !meta.is_file() { + return Err(bad(cmd, format!("`{path}` is not a file"))); + } + if meta.len() == 0 || meta.len() > MAX_MODEL_BYTES { + return Err(bad(cmd, "a model file must be between 1 byte and 512 MiB")); + } + let sha256 = match known_sha { + Some(h) => h.to_string(), + None => sha256_file(p).map_err(|e| fail(&format!("could not read {path}"), e))?, + }; + let file_name = p.file_name().map(|n| n.to_string_lossy().into_owned()).unwrap_or_default(); + let suggestion = match onnx::probe_path(p) { + Ok(info) => suggest_with(&info, &sha256, meta.len(), &file_name, extra), + Err(e) => Suggestion::Unsupported(e.to_string()), + }; + Ok(Inspection { sha256, size: meta.len(), file_name, suggestion }) +} + +fn list(s: &mut Session, _: &Value) -> Result { + refresh_catalog(s); + let dir = s.face_models_dir.clone(); + let (settings, installed) = match &dir { + Some(d) => (read_settings(d), installed_models(d)), + None => (FaceSettings::default(), Vec::new()), + }; + let on_disk = |id: &str| installed.iter().find(|i| i.manifest.id == id); + let mut models: Vec = Vec::new(); + for k in known::all() { + let found = on_disk(&k.id); + models.push(row(s, &k, found.is_some(), settings.embedder.as_deref() == Some(k.id.as_str()), found.map_or(&Value::Null, |f| &f.accepted))); + } + // the user's own catalog + for e in &s.face_catalog.entries { + let found = on_disk(&e.manifest.id); + let chosen = settings.embedder.as_deref() == Some(e.manifest.id.as_str()); + models.push(row(s, &e.manifest, found.is_some(), chosen, found.map_or(&Value::Null, |f| &f.accepted))); + } + for i in &installed { + let listed = known::all().iter().any(|k| k.id == i.manifest.id) || s.face_catalog.entries.iter().any(|e| e.manifest.id == i.manifest.id); + if !listed { + models.push(row(s, &i.manifest, true, settings.embedder.as_deref() == Some(i.manifest.id.as_str()), &i.accepted)); + } + } + Ok(json!({ + "dir": dir.as_ref().map(|d| d.display().to_string()), + "enabled": settings.enabled, + "embedder": settings.embedder, + // Native builds always have the pure-Rust runner; model files/background scanning remain native-only. + "runtime": cfg!(not(target_arch = "wasm32")), + "models": models, + // the user's own list of models to download: where it goes, how many it holds, and what was wrong with the rest + "catalog": { + "path": dir.as_ref().map(|d| d.join("catalog.json").display().to_string()), + "models": s.face_catalog.entries.len(), + "errors": s.face_catalog.errors, + }, + })) +} + +fn inspect(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.models.inspect"; + let path = str_param(p, "path").ok_or_else(|| bad(C, "missing `path`"))?; + refresh_catalog(s); + let ins = inspect_file(path, C, None, &catalog_manifests(s))?; + let installed = s.face_models_dir.as_deref().map(installed_models).unwrap_or_default(); + let (kind, model, assumptions, reason) = match &ins.suggestion { + Suggestion::Known(m) => ("known", Some(m), Vec::new(), None), + Suggestion::Draft { manifest, assumptions } => ("draft", Some(manifest), assumptions.clone(), None), + Suggestion::Unsupported(why) => ("unsupported", None, Vec::new(), Some(why.clone())), + }; + let already = model.is_some_and(|m| installed.iter().any(|i| i.manifest.sha256 == m.sha256 && i.manifest.id == m.id)); + Ok(json!({ + "path": path, + "fileName": ins.file_name, + "sizeBytes": ins.size, + "sha256": ins.sha256, + "kind": kind, + "model": model.map(|m| row(s, m, already, false, &Value::Null)), + "assumptions": assumptions, + "reason": reason, + "alreadyInstalled": already, + })) +} + +fn install(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.models.install"; + let path = str_param(p, "path").ok_or_else(|| bad(C, "missing `path`"))?; + let acknowledged = p.get("acknowledged").and_then(Value::as_bool).unwrap_or(false); + let activate = p.get("activate").and_then(Value::as_bool).unwrap_or(true); + refresh_catalog(s); + let ins = inspect_file(path, C, None, &catalog_manifests(s))?; + install_file(s, C, Path::new(path), ins, acknowledged, activate, false) +} + +/// Install the inspected model file `src`. A recognition model that passes its self-test becomes the one in use and +/// recognition is switched on (`activate`, the default): installing a model is how someone says they want it. +/// `verified` is for a download that was already checked against its recorded hash: it is moved into place instead +/// of being copied and read again. +fn install_file(s: &mut Session, c: &str, src: &Path, ins: Inspection, acknowledged: bool, activate: bool, verified: bool) -> Result { + let dir = models_dir(s, c)?; + let m = match ins.suggestion { + Suggestion::Known(m) | Suggestion::Draft { manifest: m, .. } => m, + Suggestion::Unsupported(why) => return Err(bad(c, format!("this model cannot be used yet: {why}"))), + }; + if !matches!(m.role, Role::Embedder | Role::Detector) { + return Err(bad(c, "only face detection and recognition models can be installed")); + } + manifest::validate(&m).map_err(|e| bad(c, e.to_string()))?; + if !acknowledged { + return Err(bad(c, "the licence has not been accepted: show the user its terms and pass `acknowledged: true` once they agree")); + } + let home = dir.join(&m.id); + std::fs::create_dir_all(&home).map_err(|e| fail("could not create the model's folder", e))?; + let (part, final_path) = (home.join("model.onnx.part"), home.join("model.onnx")); + let moved = verified && std::fs::rename(src, &part).is_ok(); + let placed = if moved { + Ok(()) + } else { + std::fs::copy(src, &part).map_err(|e| fail("could not copy the model (is the disk full?)", e)).and_then(|_| match sha256_file(&part) { + Ok(h) if h == ins.sha256 => Ok(()), + Ok(_) => Err(EngineError::Other("the copy does not match the original (the file changed while copying?)".into())), + Err(e) => Err(fail("could not check the copy", e)), + }) + }; + if let Err(e) = placed { + let _ = std::fs::remove_file(&part); + return Err(e); + } + // Validate the candidate before replacing a working installation. + let test = match self_test(&part, &m) { + Ok(t) => t, + Err(e) => { + let _ = std::fs::remove_file(&part); + let _ = std::fs::remove_dir(&home); + return Err(bad(c, e)); + } + }; + if let Err(e) = std::fs::rename(&part, &final_path) { + let _ = std::fs::remove_file(&part); + return Err(fail("could not finish the copy", e)); + } + write_atomic(&home.join("face-model.json"), &serde_json::to_vec_pretty(&m).map_err(|e| fail("manifest", e))?)?; + let accepted = json!({"acceptedAt": (s.clock)(), "licence": m.licence.name, "commercial": m.licence.commercial, "fileName": ins.file_name, "selfTest": test}); + write_atomic(&home.join("installed.json"), &serde_json::to_vec_pretty(&accepted).map_err(|e| fail("record", e))?)?; + // a file that was downloaded and waited in the staging folder has served: the installed copy is the model now + if src.parent() == Some(dir.join(face_download::STAGING).as_path()) { + let _ = std::fs::remove_file(src); + s.face_downloads.staged_file_used(src); + } + // the newest model is the one in use (the earlier one stays installed, and `faces.models.select` switches back) + if activate && m.role == Role::Embedder { + let mut st = read_settings(&dir); + st.embedder = Some(m.id.clone()); + st.enabled = true; + write_settings(&dir, &st)?; + } + let now = read_settings(&dir); + let selected = now.embedder.as_deref() == Some(m.id.as_str()); + Ok(json!({"installed": row(s, &m, true, selected, &accepted), "active": {"embedder": now.embedder, "enabled": now.enabled}})) +} + +/// What an error says, without the command it came from (for a line shown to the user). +fn plain(e: &EngineError) -> String { + match e { + EngineError::BadParams { msg, .. } => msg.clone(), + other => other.to_string(), + } +} + +/// Install every download that has arrived: each was accepted when it was started, so it is installed, chosen and +/// switched on without another question. Called by `faces.models.downloads` and by every frame's `faces.pump`. +pub(super) fn finish_downloads(s: &mut Session) { + const C: &str = "faces.models.download"; + let extra = catalog_manifests(s); + for (id, path, sha) in s.face_downloads.finished() { + let outcome = inspect_file(&path.to_string_lossy(), C, Some(&sha), &extra).and_then(|ins| install_file(s, C, &path, ins, true, true, true)); + match outcome { + Ok(_) => s.face_downloads.set_outcome(&id, State::Installed), + Err(e) => { + // a file that cannot be installed is of no use to anyone: throw it away, say why + let _ = std::fs::remove_file(&path); + s.face_downloads.set_outcome(&id, State::Failed(plain(&e))); + } + } + } +} + +/// `faces.models.download {id, acknowledged: true}`: fetch a model LightCraft knows an address for, in the background; +/// when it has arrived and checked out it is installed. +fn download(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.models.download"; + let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?; + let dir = models_dir(s, C)?; + refresh_catalog(s); + let spec = download_spec(s, id).ok_or_else(|| bad(C, "LightCraft has no download for that model: get the file from its page, then add it"))?; + if cfg!(target_arch = "wasm32") { + return Err(bad(C, "the web build downloads no models")); + } + if p.get("acknowledged").and_then(Value::as_bool) != Some(true) { + return Err(bad( + C, + "the licence has not been accepted: show the user the model's terms (see `faces.models.list`) and pass `acknowledged: true` once they agree", + )); + } + if installed_models(&dir).iter().any(|i| i.manifest.id == spec.id) { + return Err(bad(C, "that model is already installed")); + } + let host = spec.host().to_string(); + s.face_downloads.start(spec, &dir).map_err(|e| bad(C, e))?; + Ok(json!({"started": id, "from": host})) +} + +fn download_row(s: &Session, id: &str, state: &State) -> Value { + let from = download_spec(s, id).map(|d| d.host().to_string()); + match state { + State::Running { bytes, total } => json!({"id": id, "state": "running", "bytes": bytes, "total": total, "from": from}), + State::Done { path, .. } => json!({"id": id, "state": "done", "path": path.display().to_string(), "from": from}), + State::Installed => json!({"id": id, "state": "installed", "from": from}), + State::Failed(why) => json!({"id": id, "state": "failed", "error": why, "from": from}), + State::Cancelled => json!({"id": id, "state": "cancelled", "from": from}), + } +} + +/// `faces.models.downloads`: where each download stands (a finished one is installed first). +fn downloads(s: &mut Session, _: &Value) -> Result { + finish_downloads(s); + let all: Vec = s.face_downloads.snapshot().iter().map(|(id, st)| download_row(s, id, st)).collect(); + Ok(json!({"running": s.face_downloads.running(), "downloads": all})) +} + +/// `faces.models.downloadCancel {id}`: stop a download, or clear a finished one (a staged file is deleted). +fn download_cancel(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.models.downloadCancel"; + let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?; + Ok(json!({"discarded": s.face_downloads.discard(id)})) +} + +/// Load the model and run its self-test: `Ok(null)` in a build without the recognition runtime (nothing can +/// be checked), `Ok(result)` when it passes, `Err(reason)` when it does not load or fails a check. +#[cfg(not(target_arch = "wasm32"))] +fn self_test(path: &Path, m: &ModelManifest) -> std::result::Result { + if m.role == Role::Detector { + let bytes = read_capped(path, 8 << 20).ok_or("the detector cannot be read within its 8 MiB limit")?; + let detector = lightcraft_faces::yunet::Detector::new(&bytes).map_err(|e| e.to_string())?; + if Some(detector.input_side()) != usize::try_from(m.input.width).ok() + || m.input.width != m.input.height + || !matches!(&m.output, manifest::OutputSpec::Detector { decoder } if decoder == "yunet-v2") + { + return Err("unsupported detector input or output decoder".into()); + } + detector.detect(&[0u8; 12], 2, 2, &lightcraft_faces::yunet::Options::default()).map_err(|e| e.to_string())?; + return Ok(json!({"ok": true})); + } + let embedder = lightcraft_faces::runtime::Embedder::load(path, m).map_err(|e| e.to_string())?; + let t = embedder.self_test(); + if !t.ok { + let failed: Vec<&str> = t.checks.iter().filter(|(_, ok)| !*ok).map(|(c, _)| c.as_str()).collect(); + return Err(format!("the model failed its self-test: it does not {}", failed.join(", and does not "))); + } + serde_json::to_value(&t).map_err(|e| e.to_string()) +} + +#[cfg(target_arch = "wasm32")] +fn self_test(_: &Path, _: &ModelManifest) -> std::result::Result { + Ok(Value::Null) +} + +/// `faces.models.test {id}`: run an installed model's self-test again. +fn test(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.models.test"; + let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?; + let dir = models_dir(s, C)?; + if !cfg!(not(target_arch = "wasm32")) { + return Err(bad(C, "this build cannot run face recognition models")); + } + let found = installed_models(&dir).into_iter().find(|i| i.manifest.id == id).ok_or_else(|| bad(C, "that model is not installed"))?; + if found.manifest.role != Role::Embedder { + return Err(bad(C, "only recognition models have a self-test")); + } + let result = self_test(&dir.join(id).join("model.onnx"), &found.manifest); + let (ok, detail) = match &result { + Ok(v) => (true, v.clone()), + Err(e) => (false, json!(e)), + }; + // keep the latest outcome with the model's record + let mut record = found.accepted; + if let Some(o) = record.as_object_mut() { + o.insert("selfTest".into(), if ok { detail.clone() } else { json!({"ok": false, "error": detail}) }); + write_atomic(&dir.join(id).join("installed.json"), &serde_json::to_vec_pretty(&record).map_err(|e| fail("record", e))?)?; + } + Ok(json!({"id": id, "ok": ok, "result": detail})) +} + +fn remove(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.models.remove"; + let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?; + let dir = models_dir(s, C)?; + if !manifest::valid_id(id) { + return Err(bad(C, "not a model id")); + } + let home = dir.join(id); + // only a folder that is a model (it has a manifest) is ever deleted + if !home.join("face-model.json").is_file() { + return Err(bad(C, "that model is not installed")); + } + std::fs::remove_dir_all(&home).map_err(|e| fail("could not remove the model", e))?; + let mut st = read_settings(&dir); + if st.embedder.as_deref() == Some(id) { + // another installed recogniser takes over, so recognition keeps working; with none left nothing is chosen + st.embedder = installed_models(&dir).into_iter().find(|i| i.manifest.id != id && i.manifest.role == Role::Embedder).map(|i| i.manifest.id); + write_settings(&dir, &st)?; + } + Ok(json!({"removed": id, "embedder": st.embedder})) +} + +fn select(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.models.select"; + let dir = models_dir(s, C)?; + let id = match p.get("id") { + None | Some(Value::Null) => None, + Some(v) => Some(v.as_str().ok_or_else(|| bad(C, "`id` must be a model id or null"))?.to_string()), + }; + if let Some(id) = &id + && !installed_models(&dir).iter().any(|i| &i.manifest.id == id && i.manifest.role == Role::Embedder) + { + return Err(bad(C, "that recognition model is not installed")); + } + let mut st = read_settings(&dir); + st.embedder = id; + write_settings(&dir, &st)?; + Ok(json!({"embedder": st.embedder})) +} + +fn enable(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.enable"; + let dir = models_dir(s, C)?; + let mut st = read_settings(&dir); + st.enabled = p.get("enabled").and_then(Value::as_bool).unwrap_or(!st.enabled); + write_settings(&dir, &st)?; + // the background scan notices at once, not within the second it trusts its last look + #[cfg(not(target_arch = "wasm32"))] + { + s.faces.enabled_seen = None; + } + Ok(json!({"enabled": st.enabled})) +} + +pub fn specs() -> Vec { + vec![ + cmd!(query "faces.models.list", "Face Models", [], None, "{} → {dir, enabled, embedder, runtime, models: [{id, name, role, licence{name, commercial, url, notice}, provenance, source, sizeBytes, known, installed, selected}]}", always, list), + cmd!(query "faces.models.inspect", "Inspect Face Model File", [], None, "{path} → what a .onnx file is: {kind: known | draft | unsupported, model, assumptions, reason, alreadyInstalled}; installs nothing", always, inspect), + cmd!(query "faces.models.install", "Install Face Model", [], None, "{path, acknowledged: true, activate?: true} → {installed, active: {embedder, enabled}} — copy a .onnx face recognition model into the models folder. `acknowledged` must be true: the user has been shown its licence (see inspect) and accepted it. Unless `activate` is false the model becomes the one in use and recognition is switched on; an earlier model stays installed", always, install), + cmd!(query "faces.models.download", "Download Face Model", [], None, "{id, acknowledged: true} → {started, from} — fetch a recognition model LightCraft has a pinned address for (see `downloadHost` in the list), in the background over pure-Rust https (lightcraft-fetch). `acknowledged` must be true: the user has been shown the model's terms and accepted them. It is checked against its size and SHA-256, then installed, chosen and switched on by itself; `faces.models.downloads` shows how far it is", always, download), + cmd!(query "faces.models.downloads", "Face Model Downloads", [], None, "{} → {running, downloads: [{id, state: running | done | installed | failed | cancelled, bytes, total, error, from}]} — also installs any download that has arrived; `installed` stays listed until `faces.models.downloadCancel` clears it", always, downloads), + cmd!(query "faces.models.downloadCancel", "Cancel Face Model Download", [], None, "{id} → {discarded} — stop a download, or delete a finished one that was not installed", always, download_cancel), + cmd!(query "faces.models.test", "Test Face Model", [], None, "{id} → {ok, result: {loadMs, embedMs, dimension, checks}} — load an installed recognition model and check it gives sensible faces; needs the recognition runtime", always, test), + cmd!(query "faces.models.remove", "Remove Face Model", [], None, "{id} — delete an installed model", always, remove), + cmd!(query "faces.models.select", "Choose Face Recognition Model", [], None, "{id: installed recogniser | null}", always, select), + cmd!(query "faces.enable", "Face Recognition On/Off", [], None, "{enabled?: bool} (toggles when omitted)", always, enable), + ] +} diff --git a/crates/engine/src/cmd/face_recognize.rs b/crates/engine/src/cmd/face_recognize.rs new file mode 100644 index 00000000..9efdb146 --- /dev/null +++ b/crates/engine/src/cmd/face_recognize.rs @@ -0,0 +1,1030 @@ +//! Recognition commands: `faces.index` (embed the faces in the library), `faces.suggest` (who might an unnamed +//! face be?), `faces.person` (one person's faces, and the unnamed faces that look like them), `faces.unnamed` (every +//! unnamed face, look-alikes together) and `faces.setName` / `faces.nameFaces` (name one face, or many at once, which is +//! what teaches the suggestions). +//! +//! Everything here is a suggestion for the user to confirm; nothing is ever named automatically. All of it is +//! catalog-only (no sidecar is written) and `faces.setName` is undoable. + +use lightcraft_catalog::Op; +use serde_json::{Value, json}; + +use super::{CommandSpec, bad, cmd}; +use crate::{Result, Session}; + +/// The `name` parameter of the naming commands: a name, or none (`null`, empty) to clear it. +fn name_param(p: &Value, c: &str) -> Result> { + let name = match p.get("name") { + None | Some(Value::Null) => None, + Some(v) => Some(v.as_str().ok_or_else(|| bad(c, "`name` must be text or null"))?.trim().to_string()).filter(|n| !n.is_empty()), + }; + if let Some(n) = &name + && (n.chars().count() > 200 || n.chars().any(char::is_control)) + { + return Err(bad(c, "a name is at most 200 characters, without control characters")); + } + Ok(name) +} + +/// Give a face region its name. Naming makes the face yours: a later "Detect Faces" run no longer replaces it. +fn apply_name(region: &mut lightcraft_meta::Region, name: &Option) { + region.name = name.clone(); + if name.is_some() && region.description.as_deref().is_some_and(|d| d.starts_with(super::face_detect::MARK)) { + region.description = None; + } +} + +/// `faces.nameFaces {faces: [{photo, index}], name}`: name (or clear) many face regions at once, as one undo step: what +/// naming a group of look-alikes does. +fn name_faces(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.nameFaces"; + let name = name_param(p, C)?; + let list = p.get("faces").and_then(Value::as_array).ok_or_else(|| bad(C, "missing `faces`"))?; + if list.len() > 5000 { + return Err(bad(C, "at most 5000 faces at a time")); + } + let mut by_photo: std::collections::BTreeMap> = std::collections::BTreeMap::new(); + for f in list { + let (Some(photo), Some(index)) = + (f.get("photo").and_then(Value::as_u64), f.get("index").and_then(Value::as_u64).and_then(|i| usize::try_from(i).ok())) + else { + return Err(bad(C, "each face needs a `photo` and an `index`")); + }; + by_photo.entry(photo).or_default().push(index); + } + let (mut ops, mut named) = (Vec::new(), 0usize); + for (photo, indexes) in by_photo { + let id = lightcraft_catalog::PhotoId(photo); + let Some(photo) = s.catalog.photo(id) else { continue }; + let mut meta = photo.meta.clone(); + let before = named; + for i in indexes { + if let Some(region) = meta.regions.get_mut(i).filter(|r| r.kind == lightcraft_meta::RegionKind::Face) { + apply_name(region, &name); + named += 1; + } + } + if named > before { + ops.push(Op::SetMeta { id, meta: Box::new(meta) }); + } + } + if ops.is_empty() { + return Err(bad(C, "none of those faces exist")); + } + let photos = ops.len(); + s.commit(if name.is_some() { "Name Faces" } else { "Clear Face Names" }, Op::Batch { ops })?; + s.skip_auto_write = true; + Ok(json!({"named": named, "photos": photos, "name": name})) +} + +/// `faces.setName {id?, index, name}`: name (or, with `null` or an empty name, un-name) one face region. Naming +/// makes the face yours: a later "Detect Faces" run no longer replaces it. +fn set_name(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.setName"; + let id = p.get("id").and_then(Value::as_u64).map(lightcraft_catalog::PhotoId).or_else(|| s.active()).ok_or_else(|| bad(C, "no photo"))?; + let index = p.get("index").and_then(Value::as_u64).and_then(|i| usize::try_from(i).ok()).ok_or_else(|| bad(C, "missing or invalid `index`"))?; + let name = name_param(p, C)?; + let mut meta = s.catalog.photo(id).ok_or_else(|| bad(C, "no such photo"))?.meta.clone(); + let region = meta.regions.get_mut(index).ok_or_else(|| bad(C, "no such region"))?; + apply_name(region, &name); + s.commit(if name.is_some() { "Name Face" } else { "Clear Face Name" }, Op::SetMeta { id, meta: Box::new(meta) })?; + s.skip_auto_write = true; + Ok(json!({"id": id.0, "index": index, "name": name})) +} + +/// Most faces `faces.unnamed` lists. +const MAX_UNNAMED: usize = 3000; + +/// How the unnamed faces were laid out. +struct Arranged { + /// Faces that look alike are next to each other (only the ones already embedded can be). + ordered: bool, + /// Recognition is on and a model is loaded. + ready: bool, + /// A name for some faces, when the faces already named say so with confidence: `(photo, index)` to `(name, score)`. + suggestions: std::collections::HashMap<(u64, usize), (String, f32)>, +} + +/// `faces.unnamed {limit?: 600}`: the faces nobody has named, to be named a group at a time. With recognition running, +/// faces that look alike are next to each other, and a face the named faces recognise carries the name they suggest. +fn unnamed(s: &mut Session, p: &Value) -> Result { + let limit = p.get("limit").and_then(Value::as_u64).map_or(600, |n| usize::try_from(n).unwrap_or(600)).clamp(1, MAX_UNNAMED); + let mut all: Vec<(u64, usize, lightcraft_geom::Rect)> = Vec::new(); + for ph in s.catalog.photos().filter(|ph| ph.in_library()) { + for (index, r) in ph.meta.regions.iter().enumerate() { + if r.kind == lightcraft_meta::RegionKind::Face && r.name.is_none() { + all.push((ph.id.0, index, r.rect)); + } + } + } + all.sort_by_key(|(photo, index, _)| (*photo, *index)); + let total = all.len(); + let arranged = imp::arrange_unnamed(s, &mut all, limit); + all.truncate(limit); + let faces: Vec = all + .iter() + .map(|(photo, index, rect)| { + let view = s.face_view(lightcraft_catalog::PhotoId(*photo), *rect); + let suggestion = arranged.suggestions.get(&(*photo, *index)).map(|(name, score)| json!({"name": name, "score": score})); + json!({"photo": photo, "index": index, "rect": rect_json(rect), "view": rect_json(&view), "suggestion": suggestion}) + }) + .collect(); + Ok(json!({"total": total, "faces": faces, "ordered": arranged.ordered, "ready": arranged.ready})) +} + +/// What looking for faces that resemble a person found. +struct Similar { + /// `{photo, index, rect, score}` of unnamed faces, the most alike first. + items: Vec, + /// Recognition is on and a model is loaded, so an empty list means "none yet" rather than "not looking". + ready: bool, + /// Photos the background scan has still to do. + pending: usize, +} + +fn rect_json(r: &lightcraft_geom::Rect) -> Value { + json!({"x0": r.x0, "y0": r.y0, "x1": r.x1, "y1": r.y1}) +} + +/// Most confirmed faces `faces.person` lists. +const MAX_CONFIRMED: usize = 5000; + +/// `faces.person {name, more?}`: the faces named `name` (case-insensitively) across the library, and, when recognition is +/// running, the unnamed faces that look most like them, to be confirmed one by one. +fn person(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.person"; + let name = p.get("name").and_then(Value::as_str).map(str::trim).filter(|n| !n.is_empty()).ok_or_else(|| bad(C, "missing `name`"))?; + let key = name.to_lowercase(); + let limit = p.get("more").and_then(Value::as_u64).map_or(120, |n| usize::try_from(n).unwrap_or(120)).min(1000); + let (mut shown, mut total, mut confirmed) = (name.to_string(), 0usize, Vec::new()); + for ph in s.catalog.photos().filter(|ph| ph.in_library()) { + for (index, r) in ph.meta.regions.iter().enumerate() { + let Some(n) = r.name.as_deref().filter(|n| r.kind == lightcraft_meta::RegionKind::Face && n.trim().to_lowercase() == key) else { + continue; + }; + if total == 0 { + shown = n.trim().to_string(); + } + total += 1; + if confirmed.len() < MAX_CONFIRMED { + confirmed.push(json!({"photo": ph.id.0, "index": index, "rect": rect_json(&r.rect), "view": rect_json(&s.face_view(ph.id, r.rect))})); + } + } + } + let similar = if total > 0 { imp::similar(s, &key, limit) } else { Similar { items: Vec::new(), ready: false, pending: 0 } }; + Ok( + json!({"name": shown, "total": total, "confirmed": confirmed, "more": similar.items, "ready": similar.ready, "pendingPhotos": similar.pending}), + ) +} + +#[cfg(not(target_arch = "wasm32"))] +mod imp { + use std::sync::Arc; + use std::time::{Duration, Instant}; + + use lightcraft_catalog::{Op, Photo, PhotoId}; + use lightcraft_develop::DevelopSettings; + use lightcraft_faces::matching; + use lightcraft_faces::runtime::Embedder; + use lightcraft_geom::Rect; + use lightcraft_meta::{Region, RegionKind}; + use serde_json::{Value, json}; + + use super::super::bad; + use super::super::face_detect::{LABEL, MARK}; + use super::super::face_models::{finish_downloads, installed_models, read_settings}; + use super::Similar; + use crate::faces_worker::{Done, EDGE, PREFETCH, Pace, Prepared, Worker, max_workers, process, target_workers}; + use crate::{Result, Session}; + + /// Defaults for suggestions, until a model's own threshold is known. + const DEFAULT_THRESHOLD: f32 = 0.45; + const DEFAULT_MARGIN: f32 = 0.06; + /// How long a look at the models folder is trusted by the per-frame `faces.pump`. + const RECHECK: Duration = Duration::from_secs(1); + /// How often the background pump writes new embeddings to the cache file (appending is cheap; this keeps it to a few + /// writes a minute, and a session that ends writes the rest). + const SAVE_EVERY: Duration = Duration::from_secs(3); + + /// The chosen recognition model, loaded, with the index reset to it when it changed. `fresh` looks at the models + /// folder again; otherwise a look made in the last second is trusted (the UI asks every frame). + fn current_with(s: &mut Session, cmd: &str, fresh: bool) -> Result> { + if !fresh + && let Some((_, e)) = &s.faces.embedder + && s.faces.checked.is_some_and(|t| t.elapsed() < RECHECK) + { + return Ok(e.clone()); + } + let dir = s.face_models_dir.clone().ok_or_else(|| bad(cmd, "this build has nowhere to keep face models"))?; + let id = read_settings(&dir).embedder.ok_or_else(|| bad(cmd, "no recognition model is chosen: add one in Settings > Faces and press Use"))?; + let installed = installed_models(&dir) + .into_iter() + .find(|i| i.manifest.id == id) + .ok_or_else(|| bad(cmd, "the chosen recognition model is not installed"))?; + let tag = format!("{id}@{}", installed.manifest.sha256.as_deref().unwrap_or("unknown")); + s.faces.checked = Some(Instant::now()); + // which photos the detector has searched is kept beside the library (or in memory when it is not on disk) + let scan_path = s.library.as_ref().filter(|l| l.on_disk).map(|l| l.dir.join("face-scanned.bin")); + if s.faces.scanned.path() != scan_path.as_deref() { + s.faces.scanned.reset(scan_path); + } + if let Some((have, e)) = &s.faces.embedder + && *have == tag + { + return Ok(e.clone()); + } + let embedder = Arc::new(Embedder::load(&dir.join(&id).join("model.onnx"), &installed.manifest).map_err(|e| bad(cmd, e.to_string()))?); + let dim = match installed.manifest.output { + lightcraft_faces::OutputSpec::Embedding { dim } => dim as usize, + _ => return Err(bad(cmd, "the chosen model is not a recognition model")), + }; + let cache = s.library.as_ref().filter(|l| l.on_disk).map(|l| l.dir.join(format!("face-embeddings-{id}.bin"))); + s.faces.index.reset(&tag, dim, cache); + s.faces.queue.clear(); + s.faces.queue_stamp = None; + s.faces.in_flight.clear(); + s.faces.embedder = Some((tag, embedder.clone())); + Ok(embedder) + } + + fn current(s: &mut Session, cmd: &str) -> Result> { + current_with(s, cmd, true) + } + + /// Whether the background scan should look for faces in this photo: it is in the library, has no face regions at all + /// (nothing from XMP, nothing drawn or named, no earlier detection) and has not been searched yet. + fn needs_search(s: &Session, photo: &Photo) -> bool { + photo.in_library() + && !photo.meta.regions.iter().any(|r| r.kind == RegionKind::Face) + && !s.faces.scanned.contains(photo.id.0) + && !s.faces.scan_failed.contains(&photo.id.0) + } + + /// The photo's faces that still need embedding (and, with `search`, the faces to find in it first), with the render job + /// for its picture. `None` when there is nothing to do. + fn prepare(s: &mut Session, id: PhotoId, search: bool) -> Option { + let tag = s.faces.index.tag.clone(); + let photo = s.catalog.photo(id)?.clone(); + let todo: Vec = + photo.meta.regions.iter().filter(|r| r.kind == RegionKind::Face && s.faces.index.needs(id.0, &r.rect)).map(|r| r.rect).collect(); + // finding faces needs the detector the user may not have downloaded: such a photo is left unsearched, not marked as searched + let detector = crate::cmd::face_detect::detector(s).ok(); + let detect = search && detector.is_some() && needs_search(s, &photo); + if todo.is_empty() && !detect { + return None; + } + let job = s.preview_job(id, EDGE, EDGE, false, &DevelopSettings::default())?; + let preview = s.scan_preview_of(&photo); + Some(Prepared { tag, epoch: s.faces.epoch, id, job, preview, todo, detect, detector }) + } + + /// The faces the detector found in a photo that had none become its face regions: unnamed, marked as detected (so a + /// manual Detect Faces replaces them), and not an undo step, since the user did not do this. + fn add_found(s: &mut Session, id: PhotoId, found: &[Rect]) { + if found.is_empty() { + return; + } + let Some(photo) = s.catalog.photo(id) else { return }; + // the user (or an import) got there first: leave their faces alone + if photo.meta.regions.iter().any(|r| r.kind == RegionKind::Face) { + return; + } + let mut meta = photo.meta.clone(); + for rect in found { + meta.regions.push(Region { rect: *rect, kind: RegionKind::Face, name: None, description: Some(format!("{MARK}{LABEL}")) }); + } + if s.apply_system(Op::SetMeta { id, meta: Box::new(meta) }).is_ok() { + // regions are not written to XMP, so no sidecar may be rewritten for this + s.skip_auto_write = true; + } + } + + /// Take a finished photo into the index (ignored when it was made for another model or another library). Returns + /// (embedded, failed). + fn ingest(s: &mut Session, d: Done) -> (usize, usize) { + s.faces.in_flight.remove(&d.id); + if d.epoch != s.faces.epoch || d.tag != s.faces.index.tag { + return (0, 0); + } + if d.detect_wanted { + match &d.found { + Some(found) => { + s.faces.scanned.insert(d.id.0); + add_found(s, d.id, found); + } + None => { + s.faces.scan_failed.insert(d.id.0); + } + } + } + let (mut embedded, mut failed) = (0, 0); + for (rect, v, view) in d.results { + match v { + Some(v) if v.len() == s.faces.index.dim => { + s.faces.index.insert_with_view(d.id.0, &rect, v, &view); + embedded += 1; + } + _ => { + s.faces.index.skip(d.id.0, &rect); + failed += 1; + } + } + } + (embedded, failed) + } + + /// Take in whatever the worker has finished. + fn collect(s: &mut Session) -> (usize, usize) { + let finished = s.faces.worker.as_ref().map(Worker::finished).unwrap_or_default(); + finished.into_iter().fold((0, 0), |(e, f), d| { + let (de, df) = ingest(s, d); + (e + de, f + df) + }) + } + + /// Embed every face region of `id` the index does not have yet, here and now. Returns (embedded, failed). + fn embed_photo(s: &mut Session, id: PhotoId, e: &Embedder) -> Result<(usize, usize)> { + let Some(job) = prepare(s, id, false) else { return Ok((0, 0)) }; + Ok(ingest(s, process(job, e))) + } + + /// Photos with a face that needs embedding (and, with `search`, photos to look for faces in): those the caller asked + /// for first, then those with named faces (they are the gallery everything else is compared with), then the rest, and + /// last the photos that only need searching. + fn pending(s: &Session, first: &[PhotoId], search: bool) -> Vec { + let mut out: Vec<(u8, PhotoId)> = Vec::new(); + for p in s.catalog.photos().filter(|p| p.in_library()) { + let faces: Vec<_> = p.meta.regions.iter().filter(|r| r.kind == RegionKind::Face).collect(); + let embed = faces.iter().any(|r| s.faces.index.needs(p.id.0, &r.rect)); + if embed || (search && needs_search(s, p)) { + let rank = if first.contains(&p.id) { + 0 + } else if faces.iter().any(|r| r.name.is_some()) { + 1 + } else if embed { + 2 + } else { + 3 + }; + out.push((rank, p.id)); + } + } + out.sort_by_key(|(rank, id)| (*rank, id.0)); + out.into_iter().map(|(_, id)| id).collect() + } + + pub fn index(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.index"; + let embedder = current(s, C)?; + collect(s); + let budget = Duration::from_millis(p.get("budgetMs").and_then(Value::as_u64).unwrap_or(200).min(600_000)); + let first: Vec = + p.get("ids").and_then(Value::as_array).map(|a| a.iter().filter_map(Value::as_u64).map(PhotoId).collect()).unwrap_or_default(); + let todo = pending(s, &first, false); + let started = Instant::now(); + let (mut embedded, mut failed, mut done) = (0, 0, 0); + for id in &todo { + // at least one photo per call, so a slow model still makes progress; none when only the status is asked for + if budget.is_zero() || (done > 0 && started.elapsed() >= budget) { + break; + } + let (e, f) = embed_photo(s, *id, &embedder)?; + (embedded, failed, done) = (embedded + e, failed + f, done + 1); + } + if let Err(e) = s.faces.index.save() { + log::warn!("could not save the face embeddings: {e}"); + } + Ok(json!({ + "model": s.faces.index.tag, + "embedded": embedded, + "failed": failed, + "photosDone": done, + "pendingPhotos": todo.len().saturating_sub(done), + "indexedFaces": s.faces.index.len(), + "ms": started.elapsed().as_secs_f64() * 1000.0, + })) + } + + fn str_scope(p: &Value) -> Option<&str> { + p.get("scope").and_then(Value::as_str) + } + + /// Seconds since 1970 of an ISO capture time (`2024-09-20T10:48:45…`, any zone suffix ignored): only for telling + /// shots taken close together, so the zone does not matter. + fn capture_secs(iso: &str) -> Option { + let n = |a: usize, b: usize| iso.get(a..b)?.parse::().ok(); + let (y, m, d, hh, mm, ss) = (n(0, 4)?, n(5, 7)?, n(8, 10)?, n(11, 13)?, n(14, 16)?, n(17, 19)?); + if !(1..=12).contains(&m) || !(1..=31).contains(&d) { + return None; + } + // days from civil (proleptic Gregorian) + let y = if m <= 2 { y - 1 } else { y }; + let era = y.div_euclid(400); + let yoe = y - era * 400; + let doy = (153 * (if m > 2 { m - 3 } else { m + 9 }) + 2) / 5 + d - 1; + let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy; + Some((era * 146_097 + doe - 719_468) * 86_400 + hh * 3600 + mm * 60 + ss) + } + + /// `faces.evaluate`: leave-one-out over the library's named faces, to see how well recognition works on *your* + /// photos and which threshold to trust. Each named face in turn is treated as unknown and ranked against the + /// other named faces, ignoring those from shots taken within `burstSecs` of it (near-duplicates would make + /// everything look easy). Reports counts only. + pub fn evaluate(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.evaluate"; + let embedder = current(s, C)?; + collect(s); + let burst = p.get("burstSecs").and_then(Value::as_i64).unwrap_or(5).clamp(0, 86_400 * 7); + let margin = p.get("margin").and_then(Value::as_f64).filter(|v| v.is_finite()).map_or(DEFAULT_MARGIN, |v| v as f32); + let thresholds: Vec = p + .get("thresholds") + .and_then(Value::as_array) + .map(|a| a.iter().filter_map(Value::as_f64).filter(|v| v.is_finite()).map(|v| v as f32).take(64).collect()) + .filter(|v: &Vec| !v.is_empty()) + .unwrap_or_else(|| (4..=16).map(|i| i as f32 * 0.05).collect()); + let budget = Duration::from_millis(p.get("budgetMs").and_then(Value::as_u64).unwrap_or(0).min(3_600_000)); + let started = Instant::now(); + // `scope: "visible"` looks only at the photos the current filter shows (say one year, or one album) + let scope: Option> = (str_scope(p) == Some("visible")).then(|| s.visible_cloned().into_iter().collect()); + let in_scope = |id: PhotoId| scope.as_ref().is_none_or(|v| v.contains(&id)); + let todo: Vec = pending(s, &[], false).into_iter().filter(|id| in_scope(*id)).collect(); + let mut done = 0; + for id in &todo { + if started.elapsed() >= budget { + break; + } + embed_photo(s, *id, &embedder)?; + done += 1; + } + // every named face with its embedding, the photo it is in, and when that was taken + struct Named { + name: String, + emb: Arc<[f32]>, + photo: PhotoId, + when: Option, + } + let named: Vec = s + .catalog + .photos() + .filter(|ph| ph.in_library() && in_scope(ph.id)) + .flat_map(|ph| { + let when = ph.captured.as_deref().and_then(capture_secs); + ph.meta + .regions + .iter() + .filter(|r| r.kind == RegionKind::Face) + .filter_map(|r| Some(Named { name: r.name.clone()?, emb: s.faces.index.get(ph.id.0, &r.rect)?.clone(), photo: ph.id, when })) + .collect::>() + }) + .collect(); + let near = |a: &Named, b: &Named| a.photo == b.photo || matches!((a.when, b.when), (Some(x), Some(y)) if (x - y).abs() <= burst); + let mut sweep: Vec<(f32, usize, usize)> = thresholds.iter().map(|t| (*t, 0, 0)).collect(); + let (mut queries, mut top1, mut unmatchable) = (0usize, 0usize, 0usize); + for q in &named { + let others: Vec<&Named> = named.iter().filter(|g| !near(q, g)).collect(); + if !others.iter().any(|g| g.name.eq_ignore_ascii_case(&q.name)) { + unmatchable += 1; + continue; + } + let refs: Vec<(&str, &[f32])> = others.iter().map(|g| (g.name.as_str(), &g.emb[..])).collect(); + let ranked = matching::rank(&q.emb, &refs); + queries += 1; + if ranked.first().is_some_and(|r| r.name.eq_ignore_ascii_case(&q.name)) { + top1 += 1; + } + for (t, suggested, correct) in sweep.iter_mut() { + if let Some(sg) = matching::suggest(&ranked, *t, margin) { + *suggested += 1; + if sg.name.eq_ignore_ascii_case(&q.name) { + *correct += 1; + } + } + } + } + let pct = |n: usize, d: usize| if d == 0 { Value::Null } else { json!((n as f64 / d as f64 * 1000.0).round() / 1000.0) }; + Ok(json!({ + "model": s.faces.index.tag, + "namedFaces": named.len(), + "queries": queries, + "unmatchable": unmatchable, + "pendingPhotos": todo.len().saturating_sub(done), + "top1": pct(top1, queries), + "margin": margin, + "burstSecs": burst, + "sweep": sweep.iter().map(|(t, sg, ok)| json!({"threshold": t, "suggested": sg, "correct": ok, "precision": pct(*ok, *sg), "recall": pct(*ok, queries)})).collect::>(), + })) + } + + /// `faces.pump`: cheap to call every frame. Takes in what the background worker has finished and, when it is idle, + /// hands it the next photo with faces to embed. Does nothing (and says so) until recognition is switched on in + /// Settings with a model chosen. + pub fn pump(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.pump"; + // how hard the scan may work now, from what the user is doing (the app says); nothing new is started while paused, + // and what is already running finishes + let pace = Pace::parse(p.get("pace").and_then(Value::as_str)); + let target = target_workers(pace); + // a model that has finished downloading is installed and switched on here, whether or not recognition was on + finish_downloads(s); + let enabled = s.face_models_dir.clone().is_some_and(|d| enabled_cached(s, &d)); + if !enabled { + return Ok(json!({"active": false})); + } + if s.faces.retry_at.is_some_and(|t| Instant::now() < t) { + return Ok(json!({"active": false})); + } + let Ok(embedder) = current_with(s, C, false) else { + // the chosen model is gone or broken: look again in a few seconds, not every frame + s.faces.retry_at = Some(Instant::now() + Duration::from_secs(3)); + return Ok(json!({"active": false})); + }; + s.faces.retry_at = None; + let (embedded, _) = collect(s); + if (s.faces.index.has_unsaved() || s.faces.scanned.has_unsaved()) && s.faces.saved_at.is_none_or(|t| t.elapsed() >= SAVE_EVERY) { + if let Err(e) = s.faces.index.save() { + log::warn!("could not save the face embeddings: {e}"); + } + if let Err(e) = s.faces.scanned.save() { + log::warn!("could not save the list of searched photos: {e}"); + } + s.faces.saved_at = Some(Instant::now()); + } + if s.faces.in_flight.is_empty() && s.faces.queue.is_empty() && s.faces.queue_stamp != Some(s.catalog.revision) { + // the queue is taken from the end, so it is made last first + let mut queue = pending(s, &[], true); + queue.reverse(); + s.faces.queue = queue; + s.faces.queue_stamp = Some(s.catalog.revision); + } + if target > 0 && s.faces.worker.is_none() && !s.faces.queue.is_empty() { + s.faces.worker = Worker::start(max_workers()); + } + // the workers are allowed `target` photos at once; a couple more wait behind them, so a worker that finishes has its + // next photo at once instead of waiting for the next frame to hand it one + if let Some(w) = &s.faces.worker { + w.allow(target, pace); + } + while target > 0 && s.faces.in_flight.len() < target + PREFETCH { + let Some(id) = s.faces.queue.pop() else { break }; + let Some(job) = prepare(s, id, true) else { continue }; + if s.faces.worker.as_ref().is_some_and(|w| w.submit(embedder.clone(), job)) { + s.faces.in_flight.insert(id); + } else { + break; + } + } + Ok(json!({ + "active": true, + "inFlight": s.faces.in_flight.len(), + "pendingPhotos": s.faces.queue.len() + s.faces.in_flight.len(), + "indexedFaces": s.faces.index.len(), + "searchedPhotos": s.faces.scanned.len(), + "workers": target, + "embedded": embedded, + })) + } + + /// The unnamed faces that look most like the person whose lower-cased name is `key`, best first, at most `limit`: each + /// is scored like a suggestion (the mean of its two closest faces of theirs) and kept above four fifths of the + /// model's suggestion bar, since a person looks at these before anything is named. Only faces already embedded can be + /// found; `ready` says whether recognition is running at all. + pub fn similar(s: &mut Session, key: &str, limit: usize) -> Similar { + let none = Similar { items: Vec::new(), ready: false, pending: 0 }; + let Some(dir) = s.face_models_dir.clone() else { return none }; + if !enabled_cached(s, &dir) { + return none; + } + let Ok(embedder) = current_with(s, "faces.person", false) else { return none }; + let pending = s.faces.queue.len() + s.faces.in_flight.len(); + let (mut gallery, mut candidates): (Vec>, Vec<(u64, usize, Rect, Arc<[f32]>)>) = (Vec::new(), Vec::new()); + for ph in s.catalog.photos().filter(|ph| ph.in_library()) { + for (index, r) in ph.meta.regions.iter().enumerate().filter(|(_, r)| r.kind == RegionKind::Face) { + let Some(e) = s.faces.index.get(ph.id.0, &r.rect) else { continue }; + match r.name.as_deref() { + Some(n) if n.trim().to_lowercase() == key => gallery.push(e.clone()), + Some(_) => {} + None => candidates.push((ph.id.0, index, r.rect, e.clone())), + } + } + } + let ready = Similar { items: Vec::new(), ready: true, pending }; + if gallery.is_empty() { + return ready; + } + let bar = embedder.manifest().thresholds.match_cosine.unwrap_or(DEFAULT_THRESHOLD) * 0.8; + let refs: Vec<(&str, &[f32])> = gallery.iter().map(|e| ("person", &e[..])).collect(); + let mut scored: Vec<(f32, u64, usize, Rect)> = candidates + .into_iter() + .filter_map(|(photo, index, rect, e)| { + let score = matching::rank(&e, &refs).first()?.score; + (score >= bar).then_some((score, photo, index, rect)) + }) + .collect(); + scored.sort_by(|a, b| b.0.total_cmp(&a.0).then(a.1.cmp(&b.1)).then(a.2.cmp(&b.2))); + scored.truncate(limit); + let items = scored + .into_iter() + .map(|(score, photo, index, rect)| { + let view = s.face_view(PhotoId(photo), rect); + json!({"photo": photo, "index": index, "rect": super::rect_json(&rect), "view": super::rect_json(&view), "score": score}) + }) + .collect(); + Similar { items, ..ready } + } + + /// Most faces put in order of likeness (the order is a pass over every pair of them). + const MAX_ARRANGED: usize = 1500; + + /// Faces in an order that puts look-alikes together: start anywhere, then always go to the closest face not yet placed. + fn nearest_chain(embeddings: &[&[f32]]) -> Vec { + let n = embeddings.len(); + let (mut placed, mut order) = (vec![false; n], Vec::with_capacity(n)); + let mut current = 0; + if n == 0 { + return order; + } + placed[0] = true; + order.push(0); + for _ in 1..n { + let mut best: Option<(usize, f32)> = None; + for (j, e) in embeddings.iter().enumerate() { + if placed[j] { + continue; + } + let c = matching::cosine(embeddings[current], e).unwrap_or(-1.0); + if best.is_none_or(|(_, b)| c > b) { + best = Some((j, c)); + } + } + let Some((j, _)) = best else { break }; + placed[j] = true; + order.push(j); + current = j; + } + order + } + + /// Put the unnamed faces `all` in order of likeness (the ones already embedded; the rest follow in photo order) and + /// find a name for those the named faces recognise, for the first `limit` of them. + pub fn arrange_unnamed(s: &mut Session, all: &mut Vec<(u64, usize, Rect)>, limit: usize) -> super::Arranged { + let none = super::Arranged { ordered: false, ready: false, suggestions: Default::default() }; + let Some(dir) = s.face_models_dir.clone() else { return none }; + if !enabled_cached(s, &dir) { + return none; + } + let Ok(embedder) = current_with(s, "faces.unnamed", false) else { return none }; + let (mut embedded, mut plain) = (Vec::new(), Vec::new()); + for f in all.drain(..) { + match s.faces.index.get(f.0, &f.2) { + Some(e) => embedded.push((f, e.clone())), + None => plain.push(f), + } + } + let beyond = embedded.split_off(embedded.len().min(MAX_ARRANGED)); + let order = nearest_chain(&embedded.iter().map(|(_, e)| &e[..]).collect::>()); + let ordered: Vec<((u64, usize, Rect), Arc<[f32]>)> = order.into_iter().filter_map(|i| embedded.get(i).cloned()).collect(); + // names the faces already named suggest, as `faces.suggest` would + let gallery: Vec<(String, Arc<[f32]>)> = s + .catalog + .photos() + .filter(|ph| ph.in_library()) + .flat_map(|ph| { + ph.meta + .regions + .iter() + .filter(|r| r.kind == RegionKind::Face) + .filter_map(|r| Some((r.name.clone()?, s.faces.index.get(ph.id.0, &r.rect)?.clone()))) + .collect::>() + }) + .collect(); + let refs: Vec<(&str, &[f32])> = gallery.iter().map(|(n, e)| (n.as_str(), &e[..])).collect(); + let threshold = embedder.manifest().thresholds.match_cosine.unwrap_or(DEFAULT_THRESHOLD); + let mut suggestions = std::collections::HashMap::new(); + if !refs.is_empty() { + for ((photo, index, _), e) in ordered.iter().take(limit) { + if let Some(x) = matching::suggest(&matching::rank(e, &refs), threshold, DEFAULT_MARGIN) { + suggestions.insert((*photo, *index), (x.name, x.score)); + } + } + } + let grouped = !ordered.is_empty(); + all.extend(ordered.into_iter().map(|(f, _)| f)); + all.extend(beyond.into_iter().map(|(f, _)| f)); + all.extend(plain); + super::Arranged { ordered: grouped, ready: true, suggestions } + } + + /// Whether face recognition is switched on, looked up at most once a second. + fn enabled_cached(s: &mut Session, dir: &std::path::Path) -> bool { + if let Some((t, on)) = s.faces.enabled_seen + && t.elapsed() < RECHECK + { + return on; + } + let on = read_settings(dir).enabled; + s.faces.enabled_seen = Some((Instant::now(), on)); + on + } + + pub fn suggest(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.suggest"; + let embedder = current(s, C)?; + collect(s); + let targets = s.targets(p); + let model_threshold = embedder.manifest().thresholds.match_cosine; + let num = |k: &str| p.get(k).and_then(Value::as_f64).filter(|v| v.is_finite()).map(|v| v as f32); + let threshold = num("threshold").or(model_threshold).unwrap_or(DEFAULT_THRESHOLD); + let margin = num("margin").unwrap_or(DEFAULT_MARGIN); + // the photos asked about are embedded first, then the named faces they are compared with, within the budget + let budget = Duration::from_millis(p.get("budgetMs").and_then(Value::as_u64).unwrap_or(1500).min(600_000)); + let started = Instant::now(); + let todo = pending(s, &targets, false); + let mut done = 0; + for id in &todo { + // a budget of zero only looks at what is already embedded + if budget.is_zero() || (done > 0 && started.elapsed() >= budget) { + break; + } + embed_photo(s, *id, &embedder)?; + done += 1; + } + let pending_photos: Vec = todo.iter().skip(done).map(|id| id.0).collect(); + let gallery: Vec<(String, Arc<[f32]>)> = s + .catalog + .photos() + .filter(|ph| ph.in_library()) + .flat_map(|ph| { + ph.meta + .regions + .iter() + .filter(|r| r.kind == RegionKind::Face) + .filter_map(|r| Some((r.name.clone()?, s.faces.index.get(ph.id.0, &r.rect)?.clone()))) + .collect::>() + }) + .collect(); + let refs: Vec<(&str, &[f32])> = gallery.iter().map(|(n, e)| (n.as_str(), &e[..])).collect(); + let mut photos = Vec::new(); + for id in &targets { + let Some(photo) = s.catalog.photo(*id) else { continue }; + let mut faces = Vec::new(); + for (index, r) in photo.meta.regions.iter().enumerate().filter(|(_, r)| r.kind == RegionKind::Face && r.name.is_none()) { + let Some(q) = s.faces.index.get(id.0, &r.rect) else { + faces.push(json!({"index": index, "pending": true})); + continue; + }; + let ranked = matching::rank(q, &refs); + let suggestion = matching::suggest(&ranked, threshold, margin); + faces.push(json!({ + "index": index, + "suggestion": suggestion.as_ref().map(|x| json!({"name": x.name, "score": x.score, "runnerUp": x.runner_up.as_ref().map(|(n, sc)| json!({"name": n, "score": sc}))})), + "candidates": ranked.iter().take(3).map(|c| json!({"name": c.name, "score": c.score, "faces": c.faces})).collect::>(), + })); + } + photos.push(json!({"id": id.0, "faces": faces})); + } + Ok( + json!({"model": s.faces.index.tag, "threshold": threshold, "margin": margin, "galleryFaces": gallery.len(), "pendingPhotos": pending_photos, "photos": photos}), + ) + } + #[cfg(test)] + mod tests { + use super::*; + use crate::tests_face_recognize::{install_yunet, region, set_regions, setup, temp, two_photos}; + + fn unit(parts: &[(usize, f32)]) -> Vec { + let mut v = vec![0.0; 64]; + for (i, x) in parts { + v[*i] = *x; + } + let n = v.iter().map(|x| x * x).sum::().sqrt(); + v.iter().map(|x| x / n).collect() + } + + /// A session with a tiny model loaded (so the index has its model's tag). + fn loaded(name: &str) -> (Session, std::path::PathBuf) { + let d = temp(name); + let (mut s, _) = setup(&d, 64); + s.execute("faces.index", &json!({"budgetMs": 0})).unwrap(); + (s, d) + } + + fn finished(s: &Session, id: PhotoId, found: Option>, results: Vec<(Rect, Option>, Rect)>) -> Done { + Done { tag: s.faces.index.tag.clone(), epoch: s.faces.epoch, id, results, detect_wanted: true, found } + } + + #[test] + fn faces_found_by_the_scan_become_unnamed_regions_without_an_undo_step() { + let (mut s, d) = loaded("found"); + let (a, b) = two_photos(&s); + set_regions(&mut s, b, vec![region(0.1, Some("Ann"))]); + let (undo, redo) = (s.undo.len(), s.redo.len()); + let found = region(0.3, None).rect; + let done = finished(&s, a, Some(vec![found]), vec![(found, Some(unit(&[(0, 1.0)])), found)]); + assert_eq!(ingest(&mut s, done), (1, 0)); + let regions = &s.catalog.photo(a).unwrap().meta.regions; + assert_eq!(regions.len(), 1); + assert!(regions[0].kind == RegionKind::Face && regions[0].name.is_none() && regions[0].rect == found); + assert!(regions[0].description.as_deref().unwrap().starts_with(MARK), "marked as detected"); + // not something the user did, and its embedding is already there + assert_eq!((s.undo.len(), s.redo.len()), (undo, redo)); + assert!(s.skip_auto_write, "no sidecar is rewritten for it"); + assert!(s.faces.scanned.contains(a.0) && s.faces.index.contains(a.0, &found)); + // a photo that has faces by now (the user, an import) keeps exactly those + let other = region(0.6, None).rect; + let done = finished(&s, b, Some(vec![other]), vec![(other, None, other)]); + ingest(&mut s, done); + assert_eq!(s.catalog.photo(b).unwrap().meta.regions.len(), 1); + assert!(s.faces.scanned.contains(b.0)); + let _ = std::fs::remove_dir_all(&d); + } + + #[test] + fn a_photo_that_could_not_be_searched_is_left_alone_for_the_session() { + let (mut s, d) = loaded("scanfail"); + let (a, _) = two_photos(&s); + assert!(needs_search(&s, s.catalog.photo(a).unwrap())); + let done = finished(&s, a, None, vec![]); + ingest(&mut s, done); + assert!(!s.faces.scanned.contains(a.0), "not remembered as searched: the file may come back"); + assert!(!needs_search(&s, s.catalog.photo(a).unwrap()), "but not retried on every pass"); + assert!(!pending(&s, &[], true).contains(&a)); + let _ = std::fs::remove_dir_all(&d); + } + + #[test] + fn only_the_background_scan_looks_for_faces_and_only_in_photos_without_any() { + let (mut s, d) = loaded("search"); + // finding faces needs the detector, which is not bundled; without it the scan only embeds + let detector = install_yunet(&mut s); + let (a, b) = two_photos(&s); + set_regions(&mut s, b, vec![region(0.1, Some("Ann")), region(0.5, None)]); + // `faces.index` embeds what is boxed; only the scan (search = true) also searches the rest, last + assert_eq!(pending(&s, &[], false), vec![b]); + let scan = pending(&s, &[], true); + assert_eq!(scan.first(), Some(&b), "named faces first"); + assert!(scan.contains(&a) && scan.len() > 2, "the unsearched photos follow"); + assert!(prepare(&mut s, a, false).is_none(), "nothing to embed in a photo without boxes"); + if detector { + assert!(prepare(&mut s, a, true).is_some_and(|p| p.detect && p.todo.is_empty())); + } else { + assert!(prepare(&mut s, a, true).is_none(), "no detector: nothing to do, and not marked as searched"); + } + assert!(prepare(&mut s, b, true).is_some_and(|p| !p.detect && p.todo.len() == 2), "a photo with boxes is not searched"); + // once searched, a photo is not searched again + s.faces.scanned.insert(a.0); + assert!(!pending(&s, &[], true).contains(&a)); + let _ = std::fs::remove_dir_all(&d); + } + + #[test] + fn opening_another_library_forgets_the_old_ones_faces() { + let (mut s, d) = loaded("newlib"); + let (a, _) = two_photos(&s); + let rect = region(0.2, None).rect; + s.faces.index.insert(a.0, &rect, unit(&[(0, 1.0)])); + s.faces.scanned.insert(a.0); + let epoch = s.faces.epoch; + let old = finished(&s, a, Some(vec![rect]), vec![(rect, Some(unit(&[(1, 1.0)])), rect)]); + s.open_library(d.join("second"), true).unwrap(); + assert_eq!((s.faces.index.len(), s.faces.scanned.len(), s.faces.queue.len(), s.faces.in_flight.len()), (0, 0, 0, 0)); + assert!(s.faces.epoch != epoch && s.faces.embedder.is_none()); + // a worker that was still busy with the old library cannot put its faces into the new one + assert_eq!(ingest(&mut s, old), (0, 0)); + assert!(s.faces.index.len() == 0 && !s.faces.scanned.contains(a.0)); + let _ = std::fs::remove_dir_all(&d); + } + + #[test] + fn the_unnamed_faces_that_look_like_a_person_come_best_first() { + let (mut s, d) = loaded("similar"); + let (a, b) = two_photos(&s); + let c = s.catalog.photos().map(|p| p.id).find(|id| *id != a && *id != b).unwrap(); + set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, Some("Bob"))]); + set_regions(&mut s, b, vec![region(0.1, None), region(0.5, None)]); + set_regions(&mut s, c, vec![region(0.3, None)]); + let put = |s: &mut Session, id: PhotoId, x: f64, v: Vec| s.faces.index.insert(id.0, ®ion(x, None).rect, v); + put(&mut s, a, 0.1, unit(&[(0, 1.0)])); // Ann + put(&mut s, a, 0.5, unit(&[(1, 1.0)])); // Bob + put(&mut s, b, 0.1, unit(&[(0, 0.95), (1, 0.05)])); // looks just like Ann + put(&mut s, b, 0.5, unit(&[(1, 1.0)])); // looks like Bob, not Ann + put(&mut s, c, 0.3, unit(&[(0, 0.8), (1, 0.6)])); // somewhat like Ann + let r = similar(&mut s, "ann", 10); + assert!(r.ready); + let found: Vec<(u64, u64)> = r.items.iter().map(|i| (i["photo"].as_u64().unwrap(), i["index"].as_u64().unwrap())).collect(); + assert_eq!(found, vec![(b.0, 0), (c.0, 0)], "best first; the face that looks like Bob is left out"); + assert!(r.items[0]["score"].as_f64().unwrap() > r.items[1]["score"].as_f64().unwrap()); + assert_eq!(similar(&mut s, "ann", 1).items.len(), 1, "capped"); + assert!(similar(&mut s, "nobody", 10).items.is_empty()); + let _ = std::fs::remove_dir_all(&d); + } + } +} + +#[cfg(target_arch = "wasm32")] +mod imp { + use serde_json::Value; + + use super::super::bad; + use crate::{Result, Session}; + + pub fn index(_: &mut Session, _: &Value) -> Result { + Err(bad("faces.index", "this build cannot run face recognition")) + } + + pub fn suggest(_: &mut Session, _: &Value) -> Result { + Err(bad("faces.suggest", "this build cannot run face recognition")) + } + + pub fn evaluate(_: &mut Session, _: &Value) -> Result { + Err(bad("faces.evaluate", "this build cannot run face recognition")) + } + + pub fn pump(_: &mut Session, _: &Value) -> Result { + Ok(serde_json::json!({"active": false})) + } + + pub fn similar(_: &mut Session, _: &str, _: usize) -> super::Similar { + super::Similar { items: Vec::new(), ready: false, pending: 0 } + } + + pub fn arrange_unnamed(_: &mut Session, _: &mut Vec<(u64, usize, lightcraft_geom::Rect)>, _: usize) -> super::Arranged { + super::Arranged { ordered: false, ready: false, suggestions: Default::default() } + } +} + +pub fn specs() -> Vec { + vec![ + cmd!( + "faces.index", + "Index Faces", + [], + None, + "{budgetMs?: 200, ids?: [photo ids to do first]} → {embedded, photosDone, pendingPhotos, indexedFaces, ms} — embed the faces in the library with the chosen recognition model, as many photos as fit in the time (at least one); call again until `pendingPhotos` is 0. `budgetMs: 0` only reports", + super::always, + imp::index + ), + cmd!( + "faces.suggest", + "Suggest Names for Faces", + [], + None, + "{ids?, threshold?, margin?, budgetMs?} → {photos: [{id, faces: [{index, suggestion: {name, score, runnerUp} | null, candidates: [{name, score}]}]}]} — who the unnamed faces in these photos might be, from the faces already named in the library. Only suggestions: nothing is named", + super::always, + imp::suggest + ), + cmd!( + query "faces.pump", + "Index Faces in the Background", + [], + None, + "{pace?: pause | light | normal | full} → {active, inFlight, pendingPhotos, indexedFaces, searchedPhotos, workers} — cheap to call every frame: takes in what the background workers have finished and gives them the next photos, as many at once as the pace allows (none when paused, one when light, half the machine's threads when normal, four fifths when full); inactive until face recognition is on in Settings with a model chosen", + super::always, + imp::pump + ), + cmd!( + "faces.evaluate", + "Evaluate Face Recognition", + [], + None, + "{budgetMs?: 0, burstSecs?: 5, margin?, thresholds?: [..], scope?: \"visible\"} → {queries, top1, sweep: [{threshold, suggested, correct, precision, recall}]} — leave-one-out over your named faces: each is ranked against the others (ignoring shots within burstSecs of it) to show how well the chosen model recognises people in *your* photos and which threshold to trust. Faces not yet embedded are done within budgetMs (`faces.index` does them all); counts only", + super::always, + imp::evaluate + ), + cmd!( + query "faces.unnamed", + "Unnamed Faces", + [], + None, + "{limit?: 600} → {total, faces: [{photo, index, rect, view, suggestion: {name, score} | null}], ordered, ready} — the faces nobody has named; with recognition running, look-alikes are next to each other and a face the named ones recognise carries the name they suggest. `view` is the box to show the face by. Name them with `faces.nameFaces`", + super::always, + unnamed + ), + cmd!( + "faces.nameFaces", + "Name Faces", + [], + None, + "{faces: [{photo, index}], name} → {named, photos} — name (or, with a null or empty name, clear) many face regions as one undoable step; the sidecar is not touched", + super::always, + name_faces + ), + cmd!( + query "faces.person", + "Faces of a Person", + [], + None, + "{name, more?: 120} → {name, total, confirmed: [{photo, index, rect}], more: [{photo, index, rect, score}], ready, pendingPhotos} — the faces named `name` across the library, and the unnamed faces that look most like them (best first, already embedded ones only, once recognition is on); confirm one with `faces.setName`", + super::always, + person + ), + cmd!( + "faces.setName", + "Name Face", + [], + None, + "{id?, index, name} — name a face region (null or empty clears it); undoable; the sidecar is not touched", + super::always, + set_name + ), + ] +} diff --git a/crates/engine/src/cmd/library.rs b/crates/engine/src/cmd/library.rs index 436ad123..d900eacf 100644 --- a/crates/engine/src/cmd/library.rs +++ b/crates/engine/src/cmd/library.rs @@ -119,6 +119,14 @@ fn album_param(p: &Value, key: &str, c: &str) -> Result { p.get(key).and_then(Value::as_u64).map(AlbumId).ok_or_else(|| bad(c, format!("missing album `{key}`"))) } +/// An optional album id: absent or `null` is `None`, anything but a number is an error. +fn opt_album_param(p: &Value, key: &str, c: &str) -> Result> { + match p.get(key) { + None | Some(Value::Null) => Ok(None), + Some(v) => v.as_u64().map(|n| Some(AlbumId(n))).ok_or_else(|| bad(c, format!("`{key}` must be an album id or null"))), + } +} + fn strs(p: &Value, key: &str) -> Vec { p.get(key).and_then(Value::as_array).map(|a| a.iter().filter_map(Value::as_str).map(str::to_string).collect()).unwrap_or_default() } @@ -784,7 +792,7 @@ pub fn specs() -> Vec { let cover = photos.first().copied(); s.commit( if folder { "New Folder" } else { "New Album" }, - Op::AddAlbum { album: Album { id, name, parent, folder, photos, cover, smart: None, quick: false } }, + Op::AddAlbum { album: Album { id, name, parent, folder, photos, cover, smart: None, quick: false, order: None } }, )?; Ok(json!({"id": id.0})) }), @@ -880,6 +888,80 @@ pub fn specs() -> Vec { s.commit("Move Album", Op::MoveAlbum { id, parent })?; ok() }), + cmd!( + "album.reorder", + "Reorder Album", + [], + None, + "{id, parent?: folderId|null, before?: albumId|null} — places the album (or folder) in `parent` (default: where it is; null: the top level) before the sibling `before` of the same kind, or last of its kind; the folder is then ordered by hand. One undo step. → {changed}: how many edits it took (0: already there)", + always, + |s, p| { + let id = album_param(p, "id", "album.reorder")?; + let me = s.catalog.album(id).ok_or_else(|| bad("album.reorder", "no such album"))?; + let parent = if p.get("parent").is_some() { opt_album_param(p, "parent", "album.reorder")? } else { me.parent }; + let before = opt_album_param(p, "before", "album.reorder")?; + if before == Some(id) { + // before itself: where it is + return Ok(json!({"changed": 0})); + } + let mut order: Vec = s.catalog.album_children(parent).into_iter().filter(|a| a.id != id).map(|a| a.id).collect(); + let at = match before { + Some(b) => { + let at = order.iter().position(|x| *x == b).ok_or_else(|| bad("album.reorder", "`before` is not inside that folder"))?; + if s.catalog.album(b).is_none_or(|x| x.folder != me.folder) { + return Err(bad("album.reorder", "folders are placed among folders and albums among albums")); + } + at + } + None if me.folder => order.iter().take_while(|x| s.catalog.album(**x).is_some_and(|a| a.folder)).count(), + None => order.len(), + }; + order.insert(at.min(order.len()), id); + let moved = me.parent != parent; + let mut ops = Vec::new(); + if moved { + ops.push(Op::MoveAlbum { id, parent }); + } + for (i, aid) in order.iter().enumerate() { + let want = u32::try_from(i).ok(); + // a moved album starts without a place in its new folder + let have = if moved && *aid == id { None } else { s.catalog.album(*aid).and_then(|a| a.order) }; + if have != want { + ops.push(Op::SetAlbumOrder { id: *aid, order: want }); + } + } + if ops.is_empty() { + return Ok(json!({"changed": 0})); + } + let changed = ops.len(); + s.commit("Reorder Album", Op::Batch { ops })?; + Ok(json!({"changed": changed})) + } + ), + cmd!( + "album.sort", + "Sort Albums by Name", + [], + None, + "{parent?: folderId|null} — drops the hand order of the albums inside the folder (default: the top level), so they are listed by name → {changed}: how many albums lost their place", + always, + |s, p| { + let parent = opt_album_param(p, "parent", "album.sort")?; + let ops: Vec = s + .catalog + .album_children(parent) + .into_iter() + .filter(|a| a.order.is_some()) + .map(|a| Op::SetAlbumOrder { id: a.id, order: None }) + .collect(); + if ops.is_empty() { + return Ok(json!({"changed": 0})); + } + let changed = ops.len(); + s.commit("Sort Albums by Name", Op::Batch { ops })?; + Ok(json!({"changed": changed})) + } + ), cmd!("album.addPhotos", "Add to Album", ["Photo"], None, "{id: albumId, ids?: [photoIds]} (default: selection)", has_selection, |s, p| { let id = album_param(p, "id", "album.addPhotos")?; let targets = ids_param(p).unwrap_or_else(|| s.targets(&Value::Null)); diff --git a/crates/engine/src/cmd/masks.rs b/crates/engine/src/cmd/masks.rs index 1615fb06..dd1fa77c 100644 --- a/crates/engine/src/cmd/masks.rs +++ b/crates/engine/src/cmd/masks.rs @@ -432,6 +432,7 @@ pub fn specs() -> Vec { "available": crate::segment::Segmenter::AVAILABLE, "installed": g.installed(), "dir": g.dir.as_ref().map(|d| d.display().to_string()), + "remote": g.remote_endpoint(), "loaded": g.loaded(), "busy": g.busy(), "analyzing": g.analyzing(), diff --git a/crates/engine/src/cmd/mod.rs b/crates/engine/src/cmd/mod.rs index 9fc83116..ef838fe0 100644 --- a/crates/engine/src/cmd/mod.rs +++ b/crates/engine/src/cmd/mod.rs @@ -10,9 +10,13 @@ mod color; pub(crate) mod convert; mod cull; pub mod curves; +pub(crate) mod denoise; 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 +144,10 @@ pub fn command_specs() -> &'static [CommandSpec] { v.extend(missing::specs()); v.extend(metadata::specs()); v.extend(filters::specs()); + v.extend(denoise::specs()); + v.extend(face_models::specs()); + v.extend(face_detect::specs()); + v.extend(face_recognize::specs()); v }) } diff --git a/crates/engine/src/config.rs b/crates/engine/src/config.rs new file mode 100644 index 00000000..7ab9b653 --- /dev/null +++ b/crates/engine/src/config.rs @@ -0,0 +1,58 @@ +//! Where LightCraft keeps per-user data on this machine: the config folder (app settings, camera +//! profiles, face and denoise 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- or denoise-models folder until a host asks for one ([`Session::with_default_face_models`], +//! [`Session::with_default_denoise_models`]), so tests stay hermetic. + +use std::path::PathBuf; + +use crate::Session; + +/// LightCraft's per-user config folder (`…/LightCraft`), if the platform tells us where. +pub fn config_dir() -> Option { + if cfg!(target_os = "macos") { + std::env::var_os("HOME").map(|h| PathBuf::from(h).join("Library/Application Support/LightCraft")) + } else if cfg!(windows) { + std::env::var_os("APPDATA").map(|a| PathBuf::from(a).join("LightCraft")) + } else { + std::env::var_os("XDG_CONFIG_HOME") + .map(PathBuf::from) + .or_else(|| std::env::var_os("HOME").map(|h| PathBuf::from(h).join(".config"))) + .map(|c| c.join("lightcraft")) + } +} + +/// Where face models are kept: `$LIGHTCRAFT_FACE_MODELS` if set, else `/models`. +pub fn default_face_models_dir() -> Option { + std::env::var_os("LIGHTCRAFT_FACE_MODELS").filter(|v| !v.is_empty()).map(PathBuf::from).or_else(|| config_dir().map(|d| d.join("models"))) +} + +/// Where opt-in denoise models are kept: `$LIGHTCRAFT_DENOISE_MODELS` if set, else `/denoise-models`. +pub fn default_denoise_models_dir() -> Option { + std::env::var_os("LIGHTCRAFT_DENOISE_MODELS") + .filter(|v| !v.is_empty()) + .map(PathBuf::from) + .or_else(|| config_dir().map(|d| d.join("denoise-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 + } + + /// Keep denoise models in the shared default folder ([`default_denoise_models_dir`]). + pub fn with_default_denoise_models(mut self) -> Self { + self.set_denoise_models_dir(default_denoise_models_dir()); + self + } + + /// Point AI denoise at `dir` (or at no folder), dropping what was learned about the old one. + pub fn set_denoise_models_dir(&mut self, dir: Option) { + self.denoise.models_dir = dir; + self.denoise.touch(); + } +} diff --git a/crates/engine/src/demo.rs b/crates/engine/src/demo.rs index dc5c3e5a..48da21e2 100644 --- a/crates/engine/src/demo.rs +++ b/crates/engine/src/demo.rs @@ -62,7 +62,17 @@ pub fn load(s: &mut Session) { } let folder = s.catalog.alloc_album_id(); ops.push(Op::AddAlbum { - album: Album { id: folder, name: "Travel 2026".into(), parent: None, folder: true, photos: vec![], cover: None, smart: None, quick: false }, + album: Album { + id: folder, + name: "Travel 2026".into(), + parent: None, + folder: true, + photos: vec![], + cover: None, + smart: None, + quick: false, + order: None, + }, }); let by_kw = |kw: &str| -> Vec<_> { scenes.iter().zip(&ids).filter(|(sc, _)| sc.meta.keywords.contains(&kw)).map(|(_, id)| *id).collect() }; for (name, parent, photos) in [ @@ -75,7 +85,9 @@ pub fn load(s: &mut Session) { ] { let id = s.catalog.alloc_album_id(); let cover = photos.first().copied(); - ops.push(Op::AddAlbum { album: Album { id, name: name.into(), parent, folder: false, photos, cover, smart: None, quick: false } }); + ops.push(Op::AddAlbum { + album: Album { id, name: name.into(), parent, folder: false, photos, cover, smart: None, quick: false, order: None }, + }); } for op in ops { let _ = s.catalog.apply(op); diff --git a/crates/engine/src/denoise.rs b/crates/engine/src/denoise.rs new file mode 100644 index 00000000..01bcdc78 --- /dev/null +++ b/crates/engine/src/denoise.rs @@ -0,0 +1,1644 @@ +//! AI denoise, the session's side (`docs/denoise.md`): which photos have a denoised picture, making the missing ones +//! in the background, and keeping the cache within its size. +//! +//! The denoised picture of a raw photo is *cached data*, never a file in the library: it is made from the untouched +//! original by the model the user chose, kept in `/denoise/.lcdn`, and thrown away whenever space is +//! needed (it is made again when it is next wanted). Its name is a hash of the model, the algorithm and the photo's +//! content, so a changed file or another model simply does not find it. The Denoise *amount* is an ordinary develop +//! setting (`enhance.denoise`): a render mixes the plain and the denoised picture by it ([`crate::media`]). +//! +//! What runs where: the session thread only decides (which photos, which model) and keeps an index of what is ready +//! ([`crate::media::DenoiseIndex`], memory only). One background thread at a time decodes a photo, runs the model on +//! its mosaic tile by tile and writes the product ([`make_product`]); an export that finds its photo's product missing +//! makes it itself, on its own thread, so it never comes out without the denoise it was asked for. +//! +//! Failure is a message, never a panic: a file that is not a Bayer raw, a model that does not load, a full disk, a +//! panic inside the runtime and a cancelled job all end as an [`Outcome`] the pump records. + +use std::collections::{HashMap, HashSet, VecDeque}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, Condvar, Mutex, OnceLock, PoisonError}; +use std::time::{Duration, SystemTime}; + +use lightcraft_catalog::{Photo, PhotoId, Source}; +use lightcraft_denoise::bayer::Layout; +use lightcraft_denoise::manifest::DenoiserManifest; +use lightcraft_denoise::run::{Control, Params, TileRunner, denoise_bayer}; +use lightcraft_denoise::{Error as RunError, product}; +use lightcraft_develop::DevelopSettings; +use lightcraft_preview::Hasher128; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; + +use crate::media::{DenoiseSpec, MakeProduct, PairLoader}; +use crate::merge::ByteReader; +use crate::model_download::Downloads; +use crate::{Selection, Session}; + +/// Extension of a product file. +pub const EXT: &str = "lcdn"; +/// Where the products of a library are kept, inside the library's folder (safe to delete, like `thumbs/`). +pub const DIR: &str = "denoise"; +/// What the cache may take up unless the user says otherwise (GB). A product is about as large as the raw file's +/// decoded samples, compressed: ~100 MB for 24 megapixels. +pub const DEFAULT_CACHE_GB: u32 = 20; +/// Samples at or above this (white = 1) are clipped: the model leaves them as the camera recorded them. +const CLIP: f32 = 0.99; +/// Most photos waiting in the queue. +const MAX_QUEUE: usize = 5000; +/// Most selected photos the automatic queue looks at in one pump. +const AUTO_LOOK: usize = 64; +/// How long a look at the model settings is trusted. +const SETTINGS_TTL: Duration = Duration::from_secs(2); +/// Stack of the threads that run the model (the runtime recurses through its graph). +const STACK: usize = 16 << 20; + +// --------------------------------------------------------------------------------------------------------------------- +// the model + +/// Where the model runs (Settings ▸ AI Denoise ▸ Speed; `denoise.settings {runOn}`). +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum RunOn { + /// The graphics card where it is faster than the processor at the pace the work goes at (both are timed on this + /// computer when the card is set up), else the processor. + #[default] + Auto, + /// The graphics card whenever it can run the model, even where the processor would be faster. + Gpu, + /// Always the processor. + Cpu, +} + +impl RunOn { + /// `auto`, `gpu` or `cpu`. + pub fn parse(name: &str) -> Option { + match name { + "auto" => Some(RunOn::Auto), + "gpu" => Some(RunOn::Gpu), + "cpu" => Some(RunOn::Cpu), + _ => None, + } + } + + pub fn name(self) -> &'static str { + match self { + RunOn::Auto => "auto", + RunOn::Gpu => "gpu", + RunOn::Cpu => "cpu", + } + } +} + +/// A loaded denoise model: what runs tiles, and a self-check. +pub(crate) trait Model: Send + Sync { + /// What runs a photo's tiles, and how many to run at once, when `threads` may be used and the setting is `run_on` + /// (a graphics card is kept busy by a few). The first time the card may be used it is set up and checked here. + fn runner(&self, run_on: RunOn, threads: usize) -> (&dyn TileRunner, usize); + /// Check it does something sensible: the test's result, or why not. Unless `run_on` is the processor the graphics + /// card is set up and checked too, and the result says where the work will run. + fn self_test(&self, run_on: RunOn) -> Result; + /// Conservative scratch bound for an unknown runner. + fn bytes_per_tile(&self) -> usize { + 512 * 1024 * 1024 + } + /// Where the tiles run: `{kind: "gpu", adapter}`, `{kind: "cpu", reason?}` or `{kind: "pending"}` (the card is not + /// set up yet), with the card's and the processor's time for the check tile (`cardMs`, `cpuMs`) once both were + /// timed. Never starts anything. + fn device(&self, _run_on: RunOn) -> Value { + json!({"kind": "cpu"}) + } +} + +/// Loads the model at a path as its manifest describes it. +pub(crate) type Loader = Arc Result, String> + Send + Sync>; + +/// Tiles run at once on a graphics card: enough to keep it fed while others are packed and blended. +#[cfg(feature = "denoise")] +const GPU_PARALLEL: usize = 4; +/// How far the GPU's answer on the check tile may be from the CPU's (of the answer's largest value) before it is not used. +#[cfg(feature = "denoise")] +const GPU_AGREE: f32 = 1e-3; +/// Longest the card may take to set up (build its kernels, run the check tile, time it) before the processor does the +/// work instead. +#[cfg(feature = "denoise")] +const GPU_SETUP_LIMIT: Duration = Duration::from_secs(30); +/// Tiles the card may get wrong (an error, numbers that are not finite) before it stops being used for the model. +#[cfg(feature = "denoise")] +const GPU_MAX_FALLBACKS: usize = 3; +/// The file beside a model while the graphics card is set up for it (see [`set_up_card`]). +pub(crate) const GPU_SETUP_MARKER: &str = "gpu-setup.marker"; +/// Why the card is not tried when a set-up marker is found. +#[cfg(feature = "denoise")] +const GPU_CRASHED: &str = "LightCraft closed while it was setting up the graphics card for AI Denoise last time"; +/// Why the card is not used when its set-up ran out of time (the start of it). +#[cfg(feature = "denoise")] +const GPU_SLOW: &str = "setting up the graphics card took longer than"; +/// Why the card is not used while a set-up that ran out of time is still going. +#[cfg(feature = "denoise")] +const GPU_STUCK: &str = "the graphics card is still setting up from before (it took too long)"; + +/// A set-up that may work another time (a crash, the time limit), unlike one the card or the model cannot do. +#[cfg(feature = "denoise")] +fn worth_retrying(why: &str) -> bool { + why == GPU_CRASHED || why == GPU_STUCK || why.starts_with(GPU_SLOW) +} + +/// The model on our pure-Rust CPU runner, and on the graphics card when there is one that can run it: the card is set up for the +/// first photo that may use it, checked against the CPU's answer on a check tile and timed against it, and used while +/// it keeps working; a tile it gets wrong is run on the CPU, so a photo is always finished. +#[cfg(feature = "denoise")] +struct Network { + cpu: lightcraft_denoise::runtime::CpuRunner, + path: PathBuf, + tile: usize, + gpu: OnceLock>, + /// Tiles the card got wrong (an error, numbers that are not finite) and the CPU ran. + fallbacks: AtomicUsize, + /// Tiles at once of the photo started last (0: none yet), to say where the work runs. + last_threads: AtomicUsize, +} + +#[cfg(feature = "denoise")] +struct GpuSide { + runner: Box, + adapter: String, + /// The check tile on the card, the fastest of a few runs (ms). + card_ms: f64, + /// The check tile on the CPU, on one thread (ms). + cpu_ms: f64, +} + +/// Whether the card (`card_ms` a tile) is faster than the CPU running `threads` tiles at once, from the CPU's time for +/// one tile alone (`cpu_ms`). Tiles at once share the memory and the cores' second threads, so each one more counts as +/// half a tile's speed: on a 16-core laptop one tile alone took 1.4 s and sixteen at once finished one every 0.2 s. +#[cfg(feature = "denoise")] +fn card_is_faster(card_ms: f64, cpu_ms: f64, threads: usize) -> bool { + let at_once = 1.0 + 0.5 * threads.saturating_sub(1) as f64; + card_ms.is_finite() && card_ms <= cpu_ms / at_once +} + +/// Card set-ups of this process that have not finished, by marker (their markers are not a crash's). +#[cfg(feature = "denoise")] +static SETTING_UP: Mutex> = Mutex::new(Vec::new()); + +/// Removes a set-up's marker when the set-up ends, however it ends — but not when the process dies inside the driver. +#[cfg(feature = "denoise")] +struct SetUpEnded(PathBuf); + +#[cfg(feature = "denoise")] +impl Drop for SetUpEnded { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.0); + SETTING_UP.lock().unwrap_or_else(PoisonError::into_inner).retain(|m| m != &self.0); + } +} + +/// Set the graphics card up with `card` on a thread of its own while `meanwhile` runs here, and give up on the card when +/// it takes longer than `limit` (a driver building kernels for minutes, or stuck). The `marker` file exists while `card` +/// runs: one found before it starts was left by a LightCraft that closed during a set-up (a driver crash takes the app +/// with it), so the card is not tried again until the user chooses where denoise runs (`denoise.settings {runOn}`). +#[cfg(feature = "denoise")] +fn set_up_card( + marker: &Path, + limit: Duration, + card: impl FnOnce() -> Result + Send + 'static, + meanwhile: impl FnOnce() -> R, +) -> Result<(T, R), String> { + let ended = { + let mut running = SETTING_UP.lock().unwrap_or_else(PoisonError::into_inner); + if running.iter().any(|m| m == marker) { + return Err(GPU_STUCK.into()); + } + if marker.exists() { + return Err(GPU_CRASHED.into()); + } + let secs = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).map_or(0, |d| d.as_secs()); + // a marker that cannot be written only loses the crash guard + let _ = std::fs::write(marker, format!("setting up the graphics card for AI Denoise (unix time {secs})\n")); + running.push(marker.to_path_buf()); + SetUpEnded(marker.to_path_buf()) + }; + let started = web_time::Instant::now(); + let (tx, rx) = std::sync::mpsc::channel(); + std::thread::Builder::new() + .name("denoise-gpu-setup".into()) + .stack_size(STACK) + .spawn(move || { + let r = std::panic::catch_unwind(std::panic::AssertUnwindSafe(card)).unwrap_or_else(|_| Err("the GPU runner gave up".into())); + drop(ended); + let _ = tx.send(r); + }) + .map_err(|e| format!("could not start setting up the graphics card: {e}"))?; + let here = meanwhile(); + match rx.recv_timeout(limit.saturating_sub(started.elapsed())) { + Ok(r) => r.map(|t| (t, here)), + Err(std::sync::mpsc::RecvTimeoutError::Timeout) => Err(format!("{GPU_SLOW} {} s, so it is not used", limit.as_secs_f64())), + Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => Err("the GPU runner gave up".into()), + } +} + +/// The card's time for one tile once it is warm: the fastest of up to three runs, stopping after about two seconds (a +/// slow card is slow enough to tell at once). +#[cfg(feature = "denoise")] +fn card_tile_ms(runner: &dyn TileRunner, input: &[f32]) -> f64 { + if runner.run(input).is_err() { + return f64::INFINITY; + } + let began = web_time::Instant::now(); + let mut best = f64::INFINITY; + for _ in 0..3 { + let one = web_time::Instant::now(); + if runner.run(input).is_err() { + break; + } + best = best.min(one.elapsed().as_secs_f64() * 1000.0); + if began.elapsed() > Duration::from_secs(2) { + break; + } + } + best +} + +/// Set the card up for the model at `path`: kernels built, the check tile run on both and compared, both timed. +#[cfg(feature = "denoise")] +fn make_gpu(cpu: &lightcraft_denoise::runtime::CpuRunner, path: &Path, tile: usize) -> Result { + let input = Arc::new(lightcraft_denoise::runtime::check_tile(tile)); + let (on_card, net) = (input.clone(), cpu.net().clone()); + let card = move || -> Result<_, String> { + let runner = lightcraft_gpu::nn::runner(&net, tile)?; + let got = runner.run(&on_card).map_err(|e| format!("the GPU could not run the check tile: {e}"))?; + let card_ms = card_tile_ms(&runner, &on_card); + Ok((runner, got, card_ms)) + }; + let on_cpu = || { + let started = web_time::Instant::now(); + (cpu.run(&input), started.elapsed().as_secs_f64() * 1000.0) + }; + let ((runner, got, card_ms), (want, cpu_ms)) = set_up_card(&path.with_file_name(GPU_SETUP_MARKER), GPU_SETUP_LIMIT, card, on_cpu)?; + let want = want.map_err(|e| format!("the CPU could not check the GPU: {e}"))?; + let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-12); + let worst = want.iter().zip(&got).fold(0f32, |m, (a, b)| m.max((a - b).abs())) / scale; + if got.len() != want.len() || !got.iter().all(|v| v.is_finite()) || worst > GPU_AGREE { + return Err(format!("its answer on the check tile is not the CPU's (off by {worst:.1e} of the largest value): not used")); + } + Ok(GpuSide { adapter: runner.adapter(), runner: Box::new(runner), card_ms, cpu_ms }) +} + +#[cfg(feature = "denoise")] +impl Network { + /// The card when it can be used now, else why not. `start` sets it up when that has not been tried yet; without it + /// nothing is started. + fn card(&self, start: bool) -> Result<&GpuSide, String> { + let side = if start { + Some(self.gpu.get_or_init(|| { + let made = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| make_gpu(&self.cpu, &self.path, self.tile))) + .unwrap_or_else(|_| Err("the GPU runner gave up".into())); + if let Err(why) = &made { + log::info!("denoise: running on the CPU ({why})"); + } + made + })) + } else { + self.gpu.get() + }; + let g = match side { + None => return Err("the graphics card is not set up yet".into()), + Some(Err(why)) => return Err(why.clone()), + Some(Ok(g)) => g, + }; + if let Some(why) = lightcraft_gpu::nn::broken_reason() { + return Err(format!("the graphics card stopped working ({why})")); + } + if self.fallbacks.load(Ordering::Relaxed) >= GPU_MAX_FALLBACKS { + return Err(format!("{} got {GPU_MAX_FALLBACKS} tiles wrong", g.adapter)); + } + Ok(g) + } + + /// The card, when a photo made `threads` tiles at a time under `run_on` runs on it; else why the CPU does. + fn choose(&self, run_on: RunOn, threads: usize, start: bool) -> Result<&GpuSide, String> { + if run_on == RunOn::Cpu { + return Err("the processor is chosen in Settings".into()); + } + let g = self.card(start)?; + if run_on == RunOn::Auto && !card_is_faster(g.card_ms, g.cpu_ms, threads) { + return Err(format!("it is faster than {} on this computer", g.adapter)); + } + Ok(g) + } + + fn fell_back(&self, why: &str) { + let n = self.fallbacks.fetch_add(1, Ordering::Relaxed) + 1; + if n == 1 { + log::warn!("denoise: a tile came back wrong from the GPU ({why}); the CPU runs it"); + } else if n == GPU_MAX_FALLBACKS { + log::warn!("denoise: the GPU got {n} tiles wrong; the CPU does the rest"); + } + } +} + +/// The card, with the CPU running any tile it gets wrong. Handed out only for photos that run on the card. +#[cfg(feature = "denoise")] +impl TileRunner for Network { + fn run(&self, input: &[f32]) -> Result, RunError> { + if let Ok(g) = self.card(false) { + match g.runner.run(input) { + Ok(out) if out.iter().all(|v| v.is_finite()) => return Ok(out), + Ok(_) => self.fell_back("numbers that are not finite"), + Err(e) => self.fell_back(&e.to_string()), + } + } + self.cpu.run(input) + } +} + +#[cfg(feature = "denoise")] +impl Model for Network { + fn bytes_per_tile(&self) -> usize { + self.cpu.bytes_per_tile() + } + fn runner(&self, run_on: RunOn, threads: usize) -> (&dyn TileRunner, usize) { + let threads = threads.max(1); + self.last_threads.store(threads, Ordering::Relaxed); + match self.choose(run_on, threads, true) { + Ok(_) => (self, threads.min(GPU_PARALLEL)), + Err(_) => (&self.cpu, threads), + } + } + + fn self_test(&self, run_on: RunOn) -> Result { + let t = self.cpu.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 "))); + } + let mut v = serde_json::to_value(&t).map_err(|e| e.to_string())?; + // the card is set up and checked here, so the test says where the work will run + if run_on != RunOn::Cpu { + let _ = self.card(true); + } + if let Some(o) = v.as_object_mut() { + o.insert("device".into(), self.device(run_on)); + } + Ok(v) + } + + fn device(&self, run_on: RunOn) -> Value { + if run_on != RunOn::Cpu && self.gpu.get().is_none() { + return json!({"kind": "pending"}); + } + let threads = match self.last_threads.load(Ordering::Relaxed) { + 0 => Pace::Full.parallel(std::thread::available_parallelism().map_or(1, |n| n.get()), None), + n => n, + }; + let mut v = match self.choose(run_on, threads, false) { + Ok(g) => json!({"kind": "gpu", "adapter": g.adapter}), + // a set-up that crashed or ran out of time can be tried again (`denoise.settings {runOn}`) + Err(why) => json!({"kind": "cpu", "retry": worth_retrying(&why), "reason": why}), + }; + if let (Some(Ok(g)), Some(o)) = (self.gpu.get(), v.as_object_mut()) { + // measured on this computer, on the check tile: the card's best run, and one tile alone on the CPU + o.insert("cardMs".into(), json!(g.card_ms)); + o.insert("cpuMs".into(), json!(g.cpu_ms)); + o.insert("fellBack".into(), json!(self.fallbacks.load(Ordering::Relaxed))); + } + v + } +} + +/// The loader of this build: our pure-Rust executor with the `denoise` feature, else one that says why nothing can run. +pub(crate) fn default_loader() -> Loader { + #[cfg(feature = "denoise")] + { + Arc::new(|path, manifest| { + let cpu = lightcraft_denoise::runtime::CpuRunner::load(path, manifest).map_err(|e| e.to_string())?; + Ok(Arc::new(Network { + cpu, + path: path.to_path_buf(), + tile: manifest.tile as usize, + gpu: OnceLock::new(), + fallbacks: AtomicUsize::new(0), + last_threads: AtomicUsize::new(0), + }) as Arc) + }) + } + #[cfg(not(feature = "denoise"))] + { + Arc::new(|_, _| Err("this build cannot run denoise models".to_string())) + } +} + +/// A model that is loaded on first use and kept, shared by the background job and the exports. +pub(crate) type ModelSlot = Arc>>>; + +// --------------------------------------------------------------------------------------------------------------------- +// settings and installed models + +/// `settings.json` in the models folder. +#[derive(Clone, Default, Serialize, Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub(crate) struct Settings { + /// The installed model in use; `None`: denoise is off. + pub model: Option, + /// Photos with a Denoise amount above 0 that are being looked at (active, selected) get their picture made without + /// being asked (`None`: yes). + pub auto: Option, + /// Cache size limit (GB; `None`: [`DEFAULT_CACHE_GB`]). + pub cache_gb: Option, + /// Most tiles run at once (`None`: by how hard the pump says to work). + pub threads: Option, + /// Where the model runs: `auto`, `gpu` or `cpu` (`None` or anything else: `auto`; see [`RunOn`]). + pub run_on: Option, + /// What came before `runOn`: `false` was the processor only. + pub gpu: Option, +} + +impl Settings { + pub fn auto(&self) -> bool { + self.auto.unwrap_or(true) + } + + pub fn run_on(&self) -> RunOn { + match self.run_on.as_deref().and_then(RunOn::parse) { + Some(r) => r, + None if self.gpu == Some(false) => RunOn::Cpu, + None => RunOn::Auto, + } + } + + pub fn cache_bytes(&self) -> u64 { + u64::from(self.cache_gb.unwrap_or(DEFAULT_CACHE_GB).clamp(1, 10_000)) << 30 + } +} + +pub(crate) fn read_capped(path: &Path, max: usize) -> Option> { + use std::io::Read; + 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(crate) fn read_settings(dir: &Path) -> Settings { + read_capped(&dir.join("settings.json"), 4096).and_then(|b| serde_json::from_slice(&b).ok()).unwrap_or_default() +} + +/// Write a file next to its final name and rename it into place. +pub(crate) fn write_atomic(path: &Path, bytes: &[u8]) -> Result<(), String> { + let part = path.with_extension("part"); + std::fs::write(&part, bytes).map_err(|e| format!("could not write {}: {e}", path.display()))?; + std::fs::rename(&part, path).map_err(|e| { + let _ = std::fs::remove_file(&part); + format!("could not write {}: {e}", path.display()) + }) +} + +pub(crate) fn write_settings(dir: &Path, st: &Settings) -> Result<(), String> { + std::fs::create_dir_all(dir).map_err(|e| format!("could not create the denoise models folder: {e}"))?; + write_atomic(&dir.join("settings.json"), &serde_json::to_vec_pretty(st).map_err(|e| e.to_string())?) +} + +/// An installed model folder: `/model.onnx`, `denoise-model.json`, `installed.json`. +pub(crate) struct Installed { + pub manifest: DenoiserManifest, + pub onnx: PathBuf, + /// What was recorded when it was installed (terms accepted, hash of the file, self-test). + pub accepted: Value, +} + +impl Installed { + /// Names the model file the cache is made with: a hash of the file as installed, so another file with the same id + /// never finds this one's products. + pub fn fingerprint(&self) -> String { + let sha = self.accepted.get("sha256").and_then(Value::as_str).or(self.manifest.sha256.as_deref()).unwrap_or(""); + let len = std::fs::metadata(&self.onnx).map(|m| m.len()).unwrap_or(0); + format!("{}@{}:{len}", self.manifest.id, sha) + } +} + +/// 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(crate) fn installed_models(dir: &Path) -> Vec { + let Ok(entries) = std::fs::read_dir(dir) else { return Vec::new() }; + let mut out = Vec::new(); + for e in entries.flatten() { + let p = e.path(); + let Some(name) = p.file_name().and_then(|n| n.to_str()).map(str::to_string) else { continue }; + let onnx = p.join("model.onnx"); + if !lightcraft_denoise::licence::valid_id(&name) || !p.is_dir() || !onnx.is_file() { + continue; + } + let Some(m) = read_capped(&p.join("denoise-model.json"), lightcraft_denoise::manifest::MAX_MANIFEST_BYTES) + .and_then(|b| lightcraft_denoise::manifest::parse(&b).ok()) + else { + continue; + }; + if m.id != name { + continue; + } + let accepted = read_capped(&p.join("installed.json"), 8192).and_then(|b| serde_json::from_slice(&b).ok()).unwrap_or(Value::Null); + out.push(Installed { manifest: m, onnx, accepted }); + } + out.sort_by(|a, b| a.manifest.id.cmp(&b.manifest.id)); + out +} + +/// Forget that setting the graphics card up for a model closed LightCraft before (the markers [`set_up_card`] leaves): +/// the user chose where denoise runs, so the card may be tried again. +pub(crate) fn forget_gpu_set_ups(dir: &Path) { + for i in installed_models(dir) { + let _ = std::fs::remove_file(i.onnx.with_file_name(GPU_SETUP_MARKER)); + } +} + +/// The model in use. +pub(crate) struct Active { + pub id: String, + pub manifest: DenoiserManifest, + pub onnx: PathBuf, + pub fingerprint: String, + pub model: ModelSlot, +} + +// --------------------------------------------------------------------------------------------------------------------- +// products + +/// The photos whose files can be denoised at all: raw files developed from their sensor data and read from disk. +pub(crate) fn eligible(p: &Photo) -> bool { + p.develops_raw() && matches!(p.source, Source::File { .. }) && !p.deleted +} + +/// A product's name: what it was made from and with. +pub(crate) fn product_key(fingerprint: &str, p: &Photo) -> String { + let mut h = Hasher128::new(); + h.str("lcdn").u64(u64::from(product::ALGORITHM)).str(fingerprint).str(&crate::media::content_key(p)); + h.finish().to_string() +} + +pub(crate) fn product_path(dir: &Path, key: &str) -> PathBuf { + dir.join(format!("{key}.{EXT}")) +} + +/// The `.lcdn` files in `dir`: path, size, last change. +fn products_in(dir: &Path) -> Vec<(PathBuf, u64, SystemTime)> { + let Ok(entries) = std::fs::read_dir(dir) else { return Vec::new() }; + entries + .flatten() + .filter(|e| e.path().extension().is_some_and(|x| x == EXT)) + .filter_map(|e| { + let m = e.metadata().ok()?; + m.is_file().then(|| (e.path(), m.len(), m.modified().unwrap_or(SystemTime::UNIX_EPOCH))) + }) + .collect() +} + +/// Make sure only one thread writes a product at a time: a second one waits, then finds it made. +fn making() -> &'static (Mutex>, Condvar) { + static MAKING: OnceLock<(Mutex>, Condvar)> = OnceLock::new(); + MAKING.get_or_init(|| (Mutex::new(HashSet::new()), Condvar::new())) +} + +struct Slot(PathBuf); + +impl Slot { + fn enter(path: &Path) -> Slot { + let (set, cv) = making(); + let mut g = set.lock().unwrap_or_else(PoisonError::into_inner); + while g.contains(path) { + g = cv.wait(g).unwrap_or_else(PoisonError::into_inner); + } + g.insert(path.to_path_buf()); + Slot(path.to_path_buf()) + } +} + +impl Drop for Slot { + fn drop(&mut self) { + let (set, cv) = making(); + set.lock().unwrap_or_else(PoisonError::into_inner).remove(&self.0); + cv.notify_all(); + } +} + +// --------------------------------------------------------------------------------------------------------------------- +// making a product + +/// Everything a thread needs to make one product, detached from the session. +#[derive(Clone)] +pub(crate) struct JobSpec { + /// The original file. + pub path: String, + pub read: Option, + pub product: PathBuf, + pub key: String, + pub manifest: DenoiserManifest, + pub onnx: PathBuf, + pub model: ModelSlot, + pub loader: Loader, + /// Tiles run at once. + pub parallel: usize, + /// Where the model runs (the setting). + pub run_on: RunOn, +} + +/// Why a product was not made. +#[derive(Debug, PartialEq)] +pub(crate) enum MakeError { + /// Denoise does not apply to this file (not a Bayer raw): nothing is wrong. + Unsupported(String), + Failed(String), + Cancelled, +} + +impl From for MakeError { + fn from(e: RunError) -> Self { + match e { + RunError::Cancelled => MakeError::Cancelled, + other => MakeError::Failed(other.to_string()), + } + } +} + +/// Progress and cancel of a running job. +#[derive(Default)] +pub(crate) struct Progress { + done: AtomicUsize, + total: AtomicUsize, + cancel: AtomicBool, +} + +impl Progress { + /// Ask the job to stop at the next tile. + pub(crate) fn cancel_now(&self) { + self.cancel.store(true, Ordering::Relaxed); + } +} + +/// The loaded model, loading it first when it is not (kept for the next photo). +fn loaded(spec: &JobSpec) -> Result, MakeError> { + if let Some(m) = spec.model.lock().unwrap_or_else(PoisonError::into_inner).clone() { + return Ok(m); + } + let m = (spec.loader)(&spec.onnx, &spec.manifest).map_err(|e| MakeError::Failed(format!("the denoise model could not be loaded: {e}")))?; + *spec.model.lock().unwrap_or_else(PoisonError::into_inner) = Some(m.clone()); + Ok(m) +} + +/// The cell `[(0,0), (1,0), (0,1), (1,1)]` colours of a Bayer mosaic and so its layout. +fn layout_of(cfa: &lightcraft_raw::Cfa) -> Option { + Layout::from_cell([cfa.color_at(0, 0), cfa.color_at(1, 0), cfa.color_at(0, 1), cfa.color_at(1, 1)]) +} + +/// Make the product of `spec` unless it is there: decode the photo, run the model over its mosaic, write the result. +/// `Ok(true)` when it was made now, `Ok(false)` when it already existed. +pub(crate) fn make_product(spec: &JobSpec, progress: Option<&Progress>) -> Result { + let _slot = Slot::enter(&spec.product); + if product::is_current(&spec.product, &spec.key) { + return Ok(false); + } + let began = web_time::Instant::now(); + let cancelled = || progress.is_some_and(|p| p.cancel.load(Ordering::Relaxed)); + if cancelled() { + return Err(MakeError::Cancelled); + } + let bytes = match &spec.read { + Some(r) => r(&spec.path).map_err(MakeError::Failed)?, + None => std::fs::read(&spec.path).map_err(|e| MakeError::Failed(format!("{}: {e}", spec.path)))?, + }; + // the decode and the normalised mosaic wait their turn at the process-wide memory gate, as every decode does + let (mosaic, layout) = { + let _held = crate::memory::work_gate().acquire(bytes.len().saturating_mul(4)); + if lightcraft_raw::probe(&bytes).is_none() { + return Err(MakeError::Unsupported("not a raw file".into())); + } + let raw = match lightcraft_raw::decode(&bytes) { + Ok(r) => r, + Err(lightcraft_raw::RawError::Unsupported(why)) => return Err(MakeError::Unsupported(why)), + Err(e) => return Err(MakeError::Failed(e.to_string())), + }; + drop(bytes); + if raw.cpp != 1 || !raw.cfa.as_ref().is_some_and(lightcraft_raw::Cfa::is_bayer) { + return Err(MakeError::Unsupported( + "only Bayer raw files can be denoised (this one has another sensor layout or is already demosaiced)".into(), + )); + } + let n = raw.normalized().map_err(|e| MakeError::Failed(e.to_string()))?; + drop(raw); + let layout = + n.cfa.as_ref().and_then(layout_of).ok_or_else(|| MakeError::Unsupported("the colour filter layout is not a Bayer pattern".into()))?; + (n, layout) + }; + if cancelled() { + return Err(MakeError::Cancelled); + } + let decoded = began.elapsed(); + let model = loaded(spec)?; + // Packed input, returned RGB and the preceding group's RGB overlap while tiles run. + let scratch = + model.bytes_per_tile().saturating_add(128usize.saturating_mul(spec.manifest.tile as usize).saturating_mul(spec.manifest.tile as usize)); + let (_, budget) = crate::memory::work_gate().usage(); + let parallel = spec.parallel.clamp(1, 4).min((budget / scratch.max(1)).max(1)); + let _held = crate::memory::work_gate().acquire(mosaic.data.len().saturating_mul(20).saturating_add(scratch.saturating_mul(parallel))); + let (runner, parallel) = model.runner(spec.run_on, parallel); + let params = Params { + tile: spec.manifest.tile as usize, + overlap: spec.manifest.overlap as usize, + gain: spec.manifest.gain, + clip: Some(CLIP), + parallel: parallel.max(1), + }; + let report = |done: usize, total: usize| { + if let Some(p) = progress { + p.total.store(total, Ordering::Relaxed); + p.done.store(done, Ordering::Relaxed); + } + }; + let ctl = Control { cancel: progress.map(|p| &p.cancel), progress: Some(&report) }; + // the tiles run on threads of their own, so a busy window's renders keep theirs + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(params.parallel) + .stack_size(STACK) + .thread_name(|i| format!("denoise-{i}")) + .build() + .map_err(|e| MakeError::Failed(format!("could not start the worker threads: {e}")))?; + let rgb = pool.install(|| denoise_bayer(&mosaic.data, mosaic.width, mosaic.height, layout, runner, ¶ms, &ctl))?; + drop(mosaic); + if cancelled() { + return Err(MakeError::Cancelled); + } + let ran = began.elapsed(); + product::write(&spec.product, &rgb, &spec.key).map_err(|e| MakeError::Failed(format!("could not save the denoised picture: {e}")))?; + if std::env::var_os("LIGHTCRAFT_PROFILE").is_some() { + eprintln!( + "[denoise] read + decode {:.0} ms, model (set-up, tiles, blend) {:.0} ms, write {:.0} ms", + decoded.as_secs_f64() * 1e3, + (ran - decoded).as_secs_f64() * 1e3, + (began.elapsed() - ran).as_secs_f64() * 1e3 + ); + } + Ok(true) +} + +// --------------------------------------------------------------------------------------------------------------------- +// the session's state + +/// How hard the background work may go (the app picks it from what the user is doing, like the face scan). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Pace { + /// Start nothing new. + Pause, + /// Two tiles at once. + Light, + /// Half the machine. + Normal, + /// Most of the machine: the user is watching the progress. + Full, +} + +impl Pace { + /// `pause`, `light`, `normal` or `full`; anything else 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, + } + } + + /// Tiles to run at once on a machine with `cores` threads (`cap`: the user's own limit). + pub fn parallel(self, cores: usize, cap: Option) -> usize { + let n = match self { + Pace::Pause | Pace::Light => 2, + Pace::Normal => cores / 2, + Pace::Full => cores * 4 / 5, + } + .clamp(1, 16); + cap.map_or(n, |c| n.min(c.max(1) as usize)) + } +} + +/// How a photo stands, for the interface. +#[derive(Clone, Debug, PartialEq, Serialize)] +#[serde(tag = "state", rename_all = "camelCase")] +pub enum PhotoState { + /// Not a photo denoise applies to (not a raw, offline, a demo). + NotApplicable, + /// No model is chosen: nothing can be made. + NoModel, + /// Its picture is made: the Amount slider works. + Ready, + Queued { + /// Photos ahead of it (0 = next). + ahead: usize, + }, + Running { + done: usize, + total: usize, + }, + /// Making it failed (`why`), or denoise does not apply to this file (`unsupported`). + Failed { + why: String, + unsupported: bool, + }, + /// Nothing has asked for it yet. + Idle, +} + +#[derive(Clone)] +struct Failure { + key: String, + why: String, + unsupported: bool, +} + +enum Outcome { + Done, + Failed(String, bool), + Cancelled, +} + +struct Running { + photo: PhotoId, + key: String, + product: PathBuf, + progress: Arc, + outcome: Arc>>, + started: web_time::Instant, +} + +pub(crate) struct State { + /// Where the models are kept (`None` where there is no file system). + pub models_dir: Option, + /// Models being downloaded at the user's request. + pub downloads: Downloads, + pub(crate) install: Option, + /// How a model is loaded (tests put another in). + pub loader: Loader, + pub(crate) active: Option, + active_seen: Option, + /// `settings.json` as last read (with the model choice). + pub(crate) settings: Settings, + /// The index must be made again (the library, the model or the cache changed). + dirty: bool, + /// The catalog revision and time of the last index, to look again when photos changed. + reindexed: Option<(u64, web_time::Instant)>, + queue: VecDeque, + running: Option, + failed: HashMap, + /// Counts every change of what is ready, so the interface knows to draw again. + pub generation: u64, + /// Photos made since the session started. + made: u64, +} + +impl Default for State { + fn default() -> Self { + State { + models_dir: None, + downloads: Downloads::default(), + install: None, + loader: default_loader(), + active: None, + active_seen: None, + settings: Settings::default(), + dirty: true, + reindexed: None, + queue: VecDeque::new(), + running: None, + failed: HashMap::new(), + generation: 0, + made: 0, + } + } +} + +impl Drop for State { + /// Quitting stops the job instead of leaving it to finish a photo nobody will see. + fn drop(&mut self) { + if let Some(r) = &self.running { + r.progress.cancel_now(); + } + } +} + +impl State { + /// The library was replaced: nothing carries over (photo ids are per library). + pub fn library_changed(&mut self) { + if let Some(r) = self.running.take() { + r.progress.cancel_now(); + } + self.queue.clear(); + self.failed.clear(); + self.dirty = true; + self.generation += 1; + } + + /// The model choice or the cache changed on disk: look again at the next pump. + pub fn touch(&mut self) { + self.active_seen = None; + self.dirty = true; + } + + pub fn queued(&self) -> usize { + self.queue.len() + } +} + +impl Session { + /// Where the active model's tiles run (see [`Model::device`]); `{kind: "none"}` while no model is loaded. + pub(crate) fn denoise_device(&self) -> Value { + let loaded = self.denoise.active.as_ref().and_then(|a| a.model.lock().unwrap_or_else(PoisonError::into_inner).clone()); + match loaded { + Some(m) => m.device(self.denoise.settings.run_on()), + None => json!({"kind": "none"}), + } + } + + /// The folder the products of the open library are kept in; `None` without a library on disk. + pub(crate) fn denoise_products_dir(&self) -> Option { + self.library.as_ref().filter(|l| l.on_disk).map(|l| l.dir.join(DIR)) + } + + /// Read which model is in use (at most every couple of seconds): the manifest and file of the model the settings + /// name, if it is installed. + pub(crate) fn denoise_refresh_active(&mut self, force: bool) { + let fresh = self.denoise.active_seen.is_some_and(|t| t.elapsed() < SETTINGS_TTL); + if fresh && !force { + return; + } + self.denoise.active_seen = Some(web_time::Instant::now()); + self.denoise.settings = self.denoise.models_dir.as_deref().map(read_settings).unwrap_or_default(); + let chosen = self.denoise.models_dir.as_deref().and_then(|dir| { + let id = self.denoise.settings.model.clone()?; + installed_models(dir).into_iter().find(|i| i.manifest.id == id) + }); + let found = chosen.map(|i| (i.fingerprint(), i)); + let same = match (&self.denoise.active, &found) { + (None, None) => true, + (Some(a), Some((fp, _))) => a.fingerprint == *fp, + _ => false, + }; + if same { + return; + } + // another model: what was made with the old one is not this one's, and what failed may work now + self.denoise.active = found.map(|(fingerprint, i)| Active { + id: i.manifest.id.clone(), + manifest: i.manifest, + onnx: i.onnx, + fingerprint, + model: Arc::new(Mutex::new(None)), + }); + if let Some(r) = self.denoise.running.take() { + r.progress.cancel_now(); + } + self.denoise.queue.clear(); + self.denoise.failed.clear(); + self.denoise.dirty = true; + } + + fn denoise_spec(&self, p: &Photo) -> Option<(DenoiseSpec, String)> { + let (dir, active) = (self.denoise_products_dir()?, self.denoise.active.as_ref()?); + let key = product_key(&active.fingerprint, p); + Some((DenoiseSpec { product: product_path(&dir, &key), key: key.clone(), make: None }, key)) + } + + /// Make the index of ready photos from the products on disk (a directory listing and a hash per photo: no file is + /// opened). Pictures that appeared or vanished drop the decoded sources they affect. + pub(crate) fn denoise_reindex(&mut self) { + self.denoise.dirty = false; + self.denoise.reindexed = Some((self.catalog.revision, web_time::Instant::now())); + let mut now: HashMap = HashMap::new(); + if let (Some(dir), Some(active)) = (self.denoise_products_dir(), self.denoise.active.as_ref()) { + let names: HashSet = + products_in(&dir).into_iter().filter_map(|(p, _, _)| p.file_stem().and_then(|s| s.to_str()).map(str::to_string)).collect(); + if !names.is_empty() { + for p in self.catalog.photos().filter(|p| eligible(p)) { + let key = product_key(&active.fingerprint, p); + if names.contains(&key) { + now.insert(p.id, DenoiseSpec { product: product_path(&dir, &key), key, make: None }); + } + } + } + } + let mut changed = false; + for id in self.media.denoise.ids() { + if !now.contains_key(&id) { + self.media.denoise.remove(id); + self.media.forget(id); + changed = true; + } + } + for (id, spec) in now { + if self.media.denoise.spec(id).is_none_or(|old| old.key != spec.key) { + self.media.denoise.set(id, spec); + self.media.forget(id); + changed = true; + } + } + if changed { + self.denoise.generation += 1; + } + } + + /// A photo whose file changed (its content key is not the one its index entry was made for) has no picture. + fn denoise_check_photo(&mut self, id: PhotoId) { + let Some(p) = self.catalog.photo(id).cloned() else { return }; + let Some((spec, _)) = self.denoise_spec(&p) else { return }; + match self.media.denoise.spec(id) { + Some(have) if have.key != spec.key => { + self.media.denoise.remove(id); + self.media.forget(id); + self.denoise.generation += 1; + } + None if spec.product.is_file() && product::is_current(&spec.product, &spec.key) => { + // made by an export, or by another session + self.media.denoise.set(id, spec); + self.media.forget(id); + self.denoise.generation += 1; + } + _ => {} + } + } + + /// How photo `id` stands. + pub fn denoise_photo_state(&self, id: PhotoId) -> PhotoState { + let Some(p) = self.catalog.photo(id) else { return PhotoState::NotApplicable }; + if !eligible(p) { + return PhotoState::NotApplicable; + } + if self.denoise.active.is_none() { + // the model folder has not been looked at yet (just started): not "no model" until it has been + return if self.denoise.active_seen.is_none() { PhotoState::Idle } else { PhotoState::NoModel }; + } + if self.denoise_products_dir().is_none() { + return PhotoState::Failed { + why: "AI Denoise keeps its pictures in the library folder, and this session has none.".into(), + unsupported: true, + }; + } + if self.media.denoise.spec(id).is_some() { + return PhotoState::Ready; + } + if let Some(r) = self.denoise.running.as_ref().filter(|r| r.photo == id) { + return PhotoState::Running { done: r.progress.done.load(Ordering::Relaxed), total: r.progress.total.load(Ordering::Relaxed) }; + } + if let Some(ahead) = self.denoise.queue.iter().position(|q| *q == id) { + return PhotoState::Queued { ahead: ahead + usize::from(self.denoise.running.is_some()) }; + } + let key = self.denoise_spec(p).map(|s| s.1); + match self.denoise.failed.get(&id) { + Some(f) if Some(&f.key) == key.as_ref() => PhotoState::Failed { why: f.why.clone(), unsupported: f.unsupported }, + _ => PhotoState::Idle, + } + } + + /// Ask for photos' pictures to be made (the ones that are not made, queued or known to be impossible). With `retry` + /// a photo that failed before is tried again. A photo that is not raw is skipped. Returns how many were queued. + pub(crate) fn denoise_enqueue(&mut self, ids: &[PhotoId], front: bool, retry: bool) -> usize { + self.denoise_refresh_active(false); + if self.denoise.active.is_none() || self.denoise_products_dir().is_none() { + return 0; + } + let mut added = 0; + for id in ids { + let Some(p) = self.catalog.photo(*id).cloned() else { continue }; + if !eligible(&p) || self.media.denoise.spec(*id).is_some() || self.denoise.queue.contains(id) { + continue; + } + if self.denoise.running.as_ref().is_some_and(|r| r.photo == *id) || self.denoise.queue.len() >= MAX_QUEUE { + continue; + } + let Some((_, key)) = self.denoise_spec(&p) else { continue }; + if retry { + self.denoise.failed.remove(id); + } else if self.denoise.failed.get(id).is_some_and(|f| f.key == key) { + continue; + } + if front { + self.denoise.queue.push_front(*id); + } else { + self.denoise.queue.push_back(*id); + } + added += 1; + } + added + } + + /// Stop the job that is running, and with `all` forget everything that waits. + pub(crate) fn denoise_cancel(&mut self, photo: Option, all: bool) -> usize { + let mut n = 0; + if let Some(r) = &self.denoise.running + && photo.is_none_or(|p| p == r.photo) + { + r.progress.cancel_now(); + n += 1; + } + let before = self.denoise.queue.len(); + match photo { + Some(p) => self.denoise.queue.retain(|q| *q != p), + None if all => self.denoise.queue.clear(), + None => {} + } + n + before - self.denoise.queue.len() + } + + /// The photos the automatic queue looks at: the one open and the ones selected, those with a Denoise amount. + fn denoise_wanted(&self) -> Vec { + let Selection { active, ids, .. } = &self.selection; + active + .iter() + .chain(ids.iter().take(AUTO_LOOK)) + .copied() + .filter(|id| self.catalog.photo(*id).is_some_and(|p| p.develop.denoise_amount() > 0.0)) + .collect() + } + + fn denoise_job_spec(&self, p: &Photo, parallel: usize) -> Option { + let Source::File { path } = &p.source else { return None }; + let (spec, _) = self.denoise_spec(p)?; + let active = self.denoise.active.as_ref()?; + Some(JobSpec { + path: path.clone(), + read: self.media.file_bytes.clone(), + product: spec.product, + key: spec.key, + manifest: active.manifest.clone(), + onnx: active.onnx.clone(), + model: active.model.clone(), + loader: self.denoise.loader.clone(), + parallel, + run_on: self.denoise.settings.run_on(), + }) + } + + /// Start the next queued photo on its own thread. + fn denoise_start_next(&mut self, parallel: usize) { + while let Some(id) = self.denoise.queue.pop_front() { + let Some(p) = self.catalog.photo(id).cloned() else { continue }; + if !eligible(&p) || self.media.denoise.spec(id).is_some() { + continue; + } + let Some(spec) = self.denoise_job_spec(&p, parallel) else { continue }; + // another session or an export may have made it since it was queued + if product::is_current(&spec.product, &spec.key) { + self.denoise_check_photo(id); + continue; + } + let progress = Arc::new(Progress::default()); + let outcome: Arc>> = Arc::new(Mutex::new(None)); + let (key, product) = (spec.key.clone(), spec.product.clone()); + let (worker_progress, worker_outcome) = (progress.clone(), outcome.clone()); + let started = std::thread::Builder::new().name("denoise".into()).stack_size(STACK).spawn(move || { + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| make_product(&spec, Some(&worker_progress)))); + let o = match result { + Ok(Ok(_)) => Outcome::Done, + Ok(Err(MakeError::Cancelled)) => Outcome::Cancelled, + Ok(Err(MakeError::Unsupported(why))) => Outcome::Failed(why, true), + Ok(Err(MakeError::Failed(why))) => Outcome::Failed(why, false), + Err(_) => Outcome::Failed("the denoiser stopped unexpectedly".into(), false), + }; + *worker_outcome.lock().unwrap_or_else(PoisonError::into_inner) = Some(o); + }); + match started { + Ok(_) => { + self.denoise.running = Some(Running { photo: id, key, product, progress, outcome, started: web_time::Instant::now() }); + } + Err(e) => { + self.denoise.failed.insert(id, Failure { key, why: format!("could not start the denoiser: {e}"), unsupported: false }); + continue; + } + } + return; + } + } + + /// Take in the running job's result. + fn denoise_poll(&mut self) { + let Some(r) = &self.denoise.running else { return }; + let Some(outcome) = r.outcome.lock().unwrap_or_else(PoisonError::into_inner).take() else { return }; + let Some(r) = self.denoise.running.take() else { return }; + match outcome { + Outcome::Done => { + self.denoise.made += 1; + self.denoise_check_photo(r.photo); + self.denoise_evict(Some(&r.product)); + } + Outcome::Failed(why, unsupported) => { + self.denoise.failed.insert(r.photo, Failure { key: r.key, why, unsupported }); + } + Outcome::Cancelled => {} + } + // the picture's state changed: the interface draws it again + self.denoise.generation += 1; + } + + /// Keep the cache within its limit: the pictures nobody asks for go first, the oldest first; `keep` (the one just + /// made) never goes. The index forgets what was deleted. + pub(crate) fn denoise_evict(&mut self, keep: Option<&Path>) { + let Some(models) = self.denoise.models_dir.clone() else { return }; + self.denoise_evict_to(read_settings(&models).cache_bytes(), keep); + } + + /// [`Self::denoise_evict`] for a limit of `limit` bytes. + pub(crate) fn denoise_evict_to(&mut self, limit: u64, keep: Option<&Path>) { + let Some(dir) = self.denoise_products_dir() else { return }; + let mut files = products_in(&dir); + let mut total: u64 = files.iter().map(|f| f.1).sum(); + if total <= limit { + return; + } + // the pictures photos use now (a Denoise amount above 0) are the last to go + let wanted: HashSet = self + .media + .denoise + .ids() + .into_iter() + .filter(|id| self.catalog.photo(*id).is_some_and(|p| p.develop.denoise_amount() > 0.0)) + .filter_map(|id| self.media.denoise.spec(id).map(|s| s.product.clone())) + .collect(); + files.sort_by_key(|(path, _, modified)| (wanted.contains(path), *modified)); + let mut gone: HashSet = HashSet::new(); + for (path, size, _) in files { + if total <= limit { + break; + } + if keep == Some(path.as_path()) { + continue; + } + if std::fs::remove_file(&path).is_ok() { + total = total.saturating_sub(size); + gone.insert(path); + } + } + if gone.is_empty() { + return; + } + for id in self.media.denoise.ids() { + if self.media.denoise.spec(id).is_some_and(|s| gone.contains(&s.product)) { + self.media.denoise.remove(id); + self.media.forget(id); + } + } + self.denoise.generation += 1; + } + + /// Delete every product of this library (they are made again when wanted). + pub(crate) fn denoise_clear(&mut self) -> (usize, u64) { + self.denoise_cancel(None, true); + let (mut files, mut bytes) = (0, 0); + if let Some(dir) = self.denoise_products_dir() { + for (path, size, _) in products_in(&dir) { + if std::fs::remove_file(&path).is_ok() { + files += 1; + bytes += size; + } + } + } + self.denoise.failed.clear(); + self.denoise.touch(); + (files, bytes) + } + + /// Per frame: take in what finished, look at what the user is looking at, and keep one job going. Cheap when + /// nothing is happening (no file is opened unless a photo with a Denoise amount lacks its picture). + pub(crate) fn denoise_pump(&mut self, pace: Pace) { + crate::cmd::denoise::finish_downloads(self); + self.denoise_refresh_active(false); + // photos that changed (a replaced file has another key) must not keep the picture of what they were + let stale = !self.media.denoise.is_empty() + && self.denoise.reindexed.is_none_or(|(rev, at)| rev != self.catalog.revision && at.elapsed() > SETTINGS_TTL); + if self.denoise.dirty || stale { + self.denoise_reindex(); + } + self.denoise_poll(); + if self.denoise.active.is_none() || self.denoise_products_dir().is_none() { + return; + } + let (cap, auto) = (self.denoise.settings.threads, self.denoise.settings.auto()); + let wanted = self.denoise_wanted(); + for id in &wanted { + self.denoise_check_photo(*id); + } + if auto { + // the photo being edited goes first + let (open, others): (Vec, Vec) = wanted.into_iter().partition(|id| self.selection.active == Some(*id)); + self.denoise_enqueue(&others, false, false); + self.denoise_enqueue(&open, true, false); + } + if self.denoise.running.is_none() && pace != Pace::Pause { + let cores = std::thread::available_parallelism().map_or(1, |n| n.get()); + self.denoise_start_next(pace.parallel(cores, cap)); + } + } + + /// The statistics `denoise.status` reports besides the queue (reads the cache folder). + pub(crate) fn denoise_cache_usage(&self) -> (usize, u64) { + match self.denoise_products_dir() { + Some(dir) => { + let files = products_in(&dir); + (files.len(), files.iter().map(|f| f.1).sum()) + } + None => (0, 0), + } + } + + /// Are pictures made without being asked for photos being looked at (Settings ▸ AI Denoise)? + pub fn denoise_auto(&self) -> bool { + self.denoise.settings.auto() + } + + /// Is a picture being made, or waiting for its turn with a model to make it? (Headless runs wait for this.) + pub fn denoise_busy(&self) -> bool { + self.denoise.running.is_some() || (!self.denoise.queue.is_empty() && self.denoise.active.is_some()) + } + + /// Seconds the running job has been going, and its progress. + pub(crate) fn denoise_running_json(&self) -> Value { + match &self.denoise.running { + Some(r) => json!({ + "photo": r.photo.0, + "done": r.progress.done.load(Ordering::Relaxed), + "total": r.progress.total.load(Ordering::Relaxed), + "seconds": r.started.elapsed().as_secs_f64(), + }), + None => Value::Null, + } + } + + pub(crate) fn denoise_made(&self) -> u64 { + self.denoise.made + } + + /// What an export of photo `id` with `settings` needs for its Denoise amount: its denoised picture, made now when the + /// queue has not got to it. `Ok(None)` when denoise does not apply (amount 0, not a Bayer raw), `Err` when it should + /// have and cannot (no model chosen, nowhere to keep the picture, the model failing). + pub(crate) fn denoise_for_export(&mut self, id: PhotoId, settings: &DevelopSettings) -> Result, String> { + if settings.denoise_amount() <= 0.0 { + return Ok(None); + } + let Some(p) = self.catalog.photo(id).cloned() else { return Ok(None) }; + let Some(loader) = self.media.denoise.loader.clone() else { return Ok(None) }; + if !eligible(&p) { + return Ok(None); + } + self.denoise_refresh_active(false); + let name = &p.file_name; + if self.denoise.active.is_none() { + return Err(format!("{name}: Denoise is on for this photo but no denoise model is chosen; choose one in Settings, or set Denoise to 0")); + } + if self.denoise_products_dir().is_none() { + return Err(format!("{name}: Denoise needs a library on disk to keep its pictures in")); + } + self.denoise_check_photo(id); + let Some(job) = self.denoise_job_spec(&p, Pace::Full.parallel(std::thread::available_parallelism().map_or(1, |n| n.get()), None)) else { + return Ok(None); + }; + let label = name.clone(); + let spec = job.clone(); + let make: MakeProduct = Arc::new(move || match make_product(&spec, None) { + Ok(_) => Ok(true), + Err(MakeError::Unsupported(_)) => Ok(false), + Err(MakeError::Failed(why)) => Err(format!("{label}: could not denoise: {why}")), + Err(MakeError::Cancelled) => Err(format!("{label}: denoise was cancelled")), + }); + Ok(Some((DenoiseSpec { product: job.product, key: job.key, make: Some(make) }, loader))) + } +} + +#[cfg(all(test, feature = "denoise"))] +mod gpu_tests { + use super::*; + use lightcraft_denoise::manifest::{Domain, Gain}; + + /// A folder of the test's own: the card's set-up leaves its marker beside the model. + fn folder(name: &str) -> PathBuf { + let dir = std::env::temp_dir().join(format!("lc-engine-gpu-{}-{name}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + dir + } + + /// The model file is read again when the card is set up (for the first photo), so it stays until the test ends. + fn model(name: &str, tile: u64) -> (Arc, lightcraft_denoise::runtime::CpuRunner, PathBuf) { + let path = folder(name).join("model.onnx"); + std::fs::write(&path, lightcraft_denoise::synthetic::unet_onnx(tile, 8, 2, 7)).unwrap(); + let manifest = DenoiserManifest { + id: "synthetic".into(), + name: "Synthetic".into(), + version: "1".into(), + licence: Default::default(), + source: None, + sha256: None, + size_bytes: None, + provenance: String::new(), + domain: Domain::BayerToRgb, + tile: tile as u32, + overlap: 16, + gain: Gain::MatchMean { nominal: 1.0, max_deviation: 0.05 }, + }; + let loaded = default_loader()(&path, &manifest).unwrap(); + let cpu = lightcraft_denoise::runtime::CpuRunner::load(&path, &manifest).unwrap(); + (loaded, cpu, path) + } + + fn worst(want: &[f32], got: &[f32]) -> f32 { + assert_eq!(want.len(), got.len()); + let scale = want.iter().fold(1e-9f32, |m, v| m.max(v.abs())); + want.iter().zip(got).fold(0f32, |m, (a, b)| m.max((a - b).abs())) / scale + } + + #[test] + fn a_tile_on_the_card_is_the_cpus_tile() { + let (model, cpu, path) = model("same", 64); + let (input, _) = lightcraft_denoise::runtime::test_tile(64); + let want = cpu.run(&input).unwrap(); + let (runner, at_once) = model.runner(RunOn::Gpu, 16); + assert!(worst(&want, &runner.run(&input).unwrap()) < 1e-3); + let device = model.device(RunOn::Gpu); + eprintln!("device: {device}"); + match device["kind"].as_str() { + Some("gpu") => { + assert!(device["adapter"].as_str().is_some_and(|a| !a.is_empty())); + assert!(at_once <= GPU_PARALLEL, "a card is fed by a few tiles at once"); + assert!(device["cardMs"].as_f64().is_some_and(|ms| ms > 0.0) && device["cpuMs"].as_f64().is_some_and(|ms| ms > 0.0), "{device}"); + } + Some("cpu") => { + assert!(device["reason"].as_str().is_some_and(|r| !r.is_empty()), "the CPU is only used for a reason: {device}"); + assert_eq!(at_once, 16); + } + other => panic!("the first photo decides where the work runs, not {other:?}"), + } + assert!(!path.with_file_name(GPU_SETUP_MARKER).exists(), "a set-up that ended leaves no marker"); + let _ = std::fs::remove_dir_all(path.parent().unwrap()); + } + + #[test] + fn with_the_processor_chosen_everything_runs_on_the_cpu() { + let (model, cpu, path) = model("cpu", 64); + let (input, _) = lightcraft_denoise::runtime::test_tile(64); + let (runner, at_once) = model.runner(RunOn::Cpu, 16); + assert_eq!(runner.run(&input).unwrap(), cpu.run(&input).unwrap(), "the CPU path is the CPU runner, bit for bit"); + assert_eq!(at_once, 16, "the threads asked for are the threads used"); + assert_eq!(model.device(RunOn::Cpu)["kind"], "cpu"); + assert_eq!(model.device(RunOn::Auto)["kind"], "pending", "the card was not even set up"); + let _ = std::fs::remove_dir_all(path.parent().unwrap()); + } + + #[test] + fn the_card_is_used_where_it_is_faster_than_the_processor() { + // measured: a tile alone on one CPU thread 1.4 s; an RTX 4090 Laptop 13 ms; a Radeon 610M 490 ms + for threads in [1, 2, 8, 16, 64] { + assert!(card_is_faster(13.0, 1400.0, threads), "a big card wins at {threads}"); + } + assert!(card_is_faster(490.0, 1400.0, 2), "a small card beats the two threads background work gets"); + assert!(!card_is_faster(490.0, 1400.0, 16), "and loses to sixteen"); + assert!(!card_is_faster(f64::INFINITY, 1400.0, 1) && !card_is_faster(f64::NAN, 1400.0, 1), "a card that never ran is never faster"); + assert!(!card_is_faster(10.0, 0.0, 1)); + } + + /// A card that gives numbers that are not finite. + struct Garbage(AtomicUsize); + + impl TileRunner for Garbage { + fn run(&self, input: &[f32]) -> Result, RunError> { + self.0.fetch_add(1, Ordering::Relaxed); + let n = input.len() / 4 * 12; + Ok(vec![f32::NAN; n]) + } + } + + #[test] + fn a_tile_the_card_gets_wrong_is_run_on_the_processor_and_three_stop_the_card() { + let (_, cpu, path) = model("garbage", 64); + let card = GpuSide { runner: Box::new(Garbage(AtomicUsize::new(0))), adapter: "Test card".into(), card_ms: 1.0, cpu_ms: 1000.0 }; + let model = Network { + cpu: cpu.clone(), + path: path.clone(), + tile: 32, + gpu: OnceLock::from(Ok(card)), + fallbacks: AtomicUsize::new(0), + last_threads: AtomicUsize::new(0), + }; + let (input, _) = lightcraft_denoise::runtime::test_tile(64); + let want = cpu.run(&input).unwrap(); + let (runner, at_once) = model.runner(RunOn::Auto, 8); + assert_eq!(at_once, GPU_PARALLEL, "the (fast) card is chosen"); + for _ in 0..GPU_MAX_FALLBACKS { + assert_eq!(runner.run(&input).unwrap(), want, "a wrong tile is the CPU's tile"); + } + let device = model.device(RunOn::Auto); + assert_eq!(device["kind"], "cpu", "{device}"); + assert!(device["reason"].as_str().unwrap().contains("tiles wrong"), "{device}"); + assert_eq!(device["fellBack"], GPU_MAX_FALLBACKS); + let (_, at_once) = model.runner(RunOn::Gpu, 8); + assert_eq!(at_once, 8, "the next photo runs on the CPU, even with the card chosen"); + let _ = std::fs::remove_dir_all(path.parent().unwrap()); + } + + #[test] + fn a_slow_card_is_left_out_unless_it_is_chosen() { + let (_, cpu, path) = model("slow", 64); + let slow = GpuSide { runner: Box::new(cpu.clone()), adapter: "Slow card".into(), card_ms: 490.0, cpu_ms: 1400.0 }; + let model = Network { + cpu, + path: path.clone(), + tile: 32, + gpu: OnceLock::from(Ok(slow)), + fallbacks: AtomicUsize::new(0), + last_threads: AtomicUsize::new(0), + }; + assert_eq!(model.runner(RunOn::Auto, 16).1, 16, "sixteen threads beat it"); + let device = model.device(RunOn::Auto); + assert_eq!((device["kind"].as_str(), device["reason"].as_str()), (Some("cpu"), Some("it is faster than Slow card on this computer"))); + assert_eq!(device["cardMs"], 490.0, "the times measured are reported"); + assert_eq!(device["retry"], false, "trying again would not make it faster"); + assert_eq!(model.runner(RunOn::Auto, 2).1, GPU_PARALLEL.min(2), "two threads do not"); + assert_eq!(model.device(RunOn::Auto)["kind"], "gpu"); + assert_eq!(model.runner(RunOn::Gpu, 16).1, GPU_PARALLEL, "chosen, it is used anyway"); + let _ = std::fs::remove_dir_all(path.parent().unwrap()); + } + + #[test] + fn a_set_up_marker_left_behind_keeps_the_card_off() { + let (model, _, path) = model("crashed", 64); + let marker = path.with_file_name(GPU_SETUP_MARKER); + std::fs::write(&marker, "setting up the graphics card for AI Denoise (unix time 0)").unwrap(); + let (_, at_once) = model.runner(RunOn::Gpu, 4); + assert_eq!(at_once, 4, "the CPU does the work"); + let device = model.device(RunOn::Gpu); + assert_eq!((device["kind"].as_str(), device["reason"].as_str()), (Some("cpu"), Some(GPU_CRASHED)), "{device}"); + assert_eq!(device["retry"], true, "the interface offers to try again"); + assert!(marker.exists(), "it stays until the user chooses where denoise runs"); + let _ = std::fs::remove_dir_all(path.parent().unwrap()); + } + + #[test] + fn a_set_up_runs_beside_the_cpu_and_removes_its_marker() { + let dir = folder("setup"); + let marker = dir.join(GPU_SETUP_MARKER); + let seen = Arc::new(AtomicBool::new(false)); + let (m, s) = (marker.clone(), seen.clone()); + let r = set_up_card( + &marker, + Duration::from_secs(10), + move || { + s.store(m.exists(), Ordering::Relaxed); + Ok(7) + }, + || "cpu", + ); + assert_eq!(r, Ok((7, "cpu"))); + assert!(seen.load(Ordering::Relaxed), "the marker is there while the card is set up"); + assert!(!marker.exists(), "and gone after"); + // a set-up that fails or panics removes it too + assert_eq!(set_up_card(&marker, Duration::from_secs(10), || Err::<(), _>("no card".into()), || ()), Err("no card".into())); + assert!(!marker.exists()); + let r = set_up_card(&marker, Duration::from_secs(10), || -> Result<(), String> { panic!("driver") }, || ()); + assert!(r.is_err() && !marker.exists(), "{r:?}"); + // a marker found before the set-up starts means the last one took LightCraft down: the card is not touched + std::fs::write(&marker, "left behind").unwrap(); + let touched = Arc::new(AtomicBool::new(false)); + let t = touched.clone(); + let r = set_up_card( + &marker, + Duration::from_secs(10), + move || { + t.store(true, Ordering::Relaxed); + Ok(()) + }, + || (), + ); + assert_eq!(r, Err(GPU_CRASHED.to_string())); + assert!(!touched.load(Ordering::Relaxed) && marker.exists()); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn a_set_up_that_takes_too_long_is_given_up_on() { + let dir = folder("slowsetup"); + let marker = dir.join(GPU_SETUP_MARKER); + let started = web_time::Instant::now(); + let slow = || -> Result<(), String> { + std::thread::sleep(Duration::from_millis(600)); + Ok(()) + }; + let r = set_up_card(&marker, Duration::from_millis(50), slow, || ()); + assert!(r.as_ref().is_err_and(|e| e.contains("took longer") && worth_retrying(e)), "{r:?}"); + assert!(started.elapsed() < Duration::from_millis(500), "the caller does not wait for the card"); + // while the first is still stuck, another set-up of the same model does not start a second one + let again = set_up_card(&marker, Duration::from_secs(5), || Ok(()), || ()); + assert!(again.is_err_and(|e| e.contains("still setting up") && worth_retrying(&e))); + // once it ends its marker goes, and the card can be set up again + while marker.exists() { + assert!(started.elapsed() < Duration::from_secs(10), "the stuck set-up ended"); + std::thread::sleep(Duration::from_millis(20)); + } + while SETTING_UP.lock().unwrap().iter().any(|m| m == &marker) { + std::thread::sleep(Duration::from_millis(5)); + } + assert_eq!(set_up_card(&marker, Duration::from_secs(5), || Ok(1), || 2), Ok((1, 2))); + let _ = std::fs::remove_dir_all(&dir); + } +} diff --git a/crates/engine/src/face_download.rs b/crates/engine/src/face_download.rs new file mode 100644 index 00000000..0c6ed4ee --- /dev/null +++ b/crates/engine/src/face_download.rs @@ -0,0 +1,372 @@ +//! Downloading a face model the user asked for (Settings ▸ Faces ▸ Download). +//! +//! The transfer is [`lightcraft_fetch`] (pure-Rust HTTPS with resume, size and SHA-256 checks), on a background thread, +//! so the UI never waits for the network. Only a model listed in [`lightcraft_faces::known::download`] can be fetched: +//! the address is a pinned constant, never something typed or passed in, and no token is ever sent with it. What arrives +//! matches the model's recorded size and SHA-256 or it is thrown away; it then waits in a staging folder until the model +//! is installed (`faces.models.install`, or by itself once the user has accepted the model's terms in the download). +//! +//! Failure is a message, never a panic: no network, a stalled or truncated transfer, a full disk, a file that is not the +//! one expected, and a panic inside the thread itself all end as [`State::Failed`]. + +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex, PoisonError}; +use std::time::Duration; + +use lightcraft_faces::known::Download; + +/// The folder inside the models folder where a downloaded file waits to be installed. Its name starts with a dot, so it +/// can never be a model's own folder (ids start with a letter or digit). +pub const STAGING: &str = ".downloads"; + +/// Where a download stands. +#[derive(Clone, Debug, PartialEq)] +pub enum State { + /// Fetching (`bytes` of `total` so far); the file is checked as the last step. + Running { + bytes: u64, + total: u64, + }, + /// Fetched and verified (its SHA-256 is `sha256`), waiting in the staging folder at `path` to be installed. + Done { + path: PathBuf, + sha256: String, + }, + /// Installed and in use: nothing more to do but tell the user. + Installed, + Failed(String), + Cancelled, +} + +struct Shared { + state: Mutex, + cancel: AtomicBool, +} + +impl Shared { + fn new(total: u64) -> Self { + Shared { state: Mutex::new(State::Running { bytes: 0, total }), cancel: AtomicBool::new(false) } + } + + fn set(&self, s: State) { + *self.state.lock().unwrap_or_else(PoisonError::into_inner) = s; + } + + fn get(&self) -> State { + self.state.lock().unwrap_or_else(PoisonError::into_inner).clone() + } +} + +/// The downloads of one session, by model id. +#[derive(Default)] +pub struct Downloads { + jobs: HashMap>, +} + +impl Downloads { + /// Start fetching `spec` into `/.downloads/`. Different models can download side by side; the same + /// one is started once. + pub fn start(&mut self, spec: Download, models_dir: &Path) -> Result<(), String> { + if self.jobs.get(spec.id.as_str()).is_some_and(|j| matches!(j.get(), State::Running { .. })) { + return Err("that model is already downloading".into()); + } + let staging = models_dir.join(STAGING); + std::fs::create_dir_all(&staging).map_err(|e| format!("could not create the download folder: {e}"))?; + let shared = Arc::new(Shared::new(spec.size_bytes)); + let worker = shared.clone(); + let id = spec.id.clone(); + std::thread::Builder::new() + .name("face-model-download".into()) + .spawn(move || { + let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| run(&spec, spec.base(), &staging, &worker))); + if outcome.is_err() { + worker.set(State::Failed("the download stopped unexpectedly".into())); + } + }) + .map_err(|e| format!("could not start the download: {e}"))?; + self.jobs.insert(id, shared); + Ok(()) + } + + /// A download that has arrived and been checked, as if the transfer had just ended (tests of installing). + #[cfg(test)] + pub(crate) fn arrived(&mut self, id: &str, path: PathBuf, sha256: &str) { + let j = Shared::new(0); + j.set(State::Done { path, sha256: sha256.to_string() }); + self.jobs.insert(id.to_string(), Arc::new(j)); + } + + /// Every download this session has started and not discarded, by model id. + pub fn snapshot(&self) -> Vec<(String, State)> { + let mut all: Vec<_> = self.jobs.iter().map(|(id, j)| (id.clone(), j.get())).filter(|(_, s)| *s != State::Cancelled).collect(); + all.sort_by(|a, b| a.0.cmp(&b.0)); + all + } + + pub fn running(&self) -> bool { + self.jobs.values().any(|j| matches!(j.get(), State::Running { .. })) + } + + /// Stop a running download, or discard a finished one (its staged file is deleted). `true` if there was one. + pub fn discard(&mut self, id: &str) -> bool { + let Some(job) = self.jobs.get(id) else { return false }; + job.cancel.store(true, Ordering::Relaxed); + match job.get() { + // the thread notices, deletes the partial file and marks itself cancelled + State::Running { .. } => {} + State::Done { path, .. } => { + let _ = std::fs::remove_file(path); + self.jobs.remove(id); + } + State::Installed | State::Failed(_) | State::Cancelled => { + self.jobs.remove(id); + } + } + true + } + + /// The staged file at `path` has been installed (its copy is the model now) or is no longer wanted. + pub fn staged_file_used(&mut self, path: &Path) { + self.jobs.retain(|_, j| !matches!(j.get(), State::Done { path: p, .. } if p == path)); + } + + /// Downloads that have arrived and been checked, ready to install: (model id, staged file, its SHA-256). + pub fn finished(&self) -> Vec<(String, PathBuf, String)> { + let mut out: Vec<_> = self + .jobs + .iter() + .filter_map(|(id, j)| match j.get() { + State::Done { path, sha256 } => Some((id.clone(), path, sha256)), + _ => None, + }) + .collect(); + out.sort_by(|a, b| a.0.cmp(&b.0)); + out + } + + /// Record how installing a finished download went (`Installed` or `Failed`). + pub fn set_outcome(&mut self, id: &str, state: State) { + match self.jobs.get(id) { + Some(j) => j.set(state), + None => { + let j = Shared::new(0); + j.set(state); + self.jobs.insert(id.to_string(), Arc::new(j)); + } + } + } +} + +impl Drop for Downloads { + /// Quitting while a download runs stops it instead of leaving it behind. + fn drop(&mut self) { + for job in self.jobs.values() { + job.cancel.store(true, Ordering::Relaxed); + } + let started = std::time::Instant::now(); + while self.running() && started.elapsed() < Duration::from_secs(1) { + std::thread::sleep(Duration::from_millis(10)); + } + } +} + +/// What went wrong, in words the user can act on. +#[cfg(not(target_arch = "wasm32"))] +fn explain(e: &lightcraft_fetch::DownloadError) -> String { + use lightcraft_fetch::DownloadError; + match e { + DownloadError::AllMirrorsFailed { errors, .. } => { + let detail: String = errors.join("; ").chars().take(300).collect(); + format!("could not download the model ({detail}): check the internet connection, or get the file from its page and add it yourself") + } + other => other.to_string(), + } +} + +/// Fetch `spec` from `mirror` (the address without the file name) into `staging`, check it and leave it there. +#[cfg(not(target_arch = "wasm32"))] +fn fetch(spec: &Download, mirror: &str, staging: &Path, shared: &Shared) -> Result<(), lightcraft_fetch::DownloadError> { + let file = lightcraft_fetch::FileSpec { name: &spec.file_name, size: Some(spec.size_bytes), sha256: Some(&spec.sha256), max: spec.size_bytes }; + lightcraft_fetch::download(&[file], &[mirror.to_string()], staging, &lightcraft_fetch::Options::default(), &shared.cancel, &mut |p| { + shared.set(State::Running { bytes: p.done.min(spec.size_bytes), total: spec.size_bytes }); + }) +} + +#[cfg(target_arch = "wasm32")] +fn fetch(_: &Download, _: &str, _: &Path, _: &Shared) -> Result<(), String> { + Err("the web build downloads no models".into()) +} + +/// The whole job: reuse a verified earlier copy, else fetch, check and stage. Always ends in a final [`State`]. +fn run(spec: &Download, mirror: &str, staging: &Path, shared: &Shared) { + let target = staging.join(&spec.file_name); + let part = staging.join(format!("{}.part", spec.file_name)); + match fetch(spec, mirror, staging, shared) { + Ok(()) => shared.set(State::Done { path: target, sha256: spec.sha256.to_string() }), + #[cfg(not(target_arch = "wasm32"))] + Err(lightcraft_fetch::DownloadError::Cancelled) => { + let _ = std::fs::remove_file(&part); + shared.set(State::Cancelled); + } + #[cfg(not(target_arch = "wasm32"))] + Err(e) => { + let _ = std::fs::remove_file(&part); + shared.set(State::Failed(explain(&e))); + } + #[cfg(target_arch = "wasm32")] + Err(e) => { + let _ = std::fs::remove_file(&part); + shared.set(State::Failed(e)); + } + } +} + +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::*; + use lightcraft_faces::hash::sha256_hex; + use std::io::{BufRead, BufReader, Write}; + use std::net::TcpListener; + + fn scratch(name: &str) -> PathBuf { + let d = std::env::temp_dir().join(format!("lc-face-dl-{name}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&d); + std::fs::create_dir_all(&d).unwrap(); + d + } + + fn spec_for(bytes: &[u8]) -> Download { + Download { + id: "test-model".into(), + url: "https://example.invalid/m/test-model.onnx".into(), + file_name: "test-model.onnx".into(), + size_bytes: bytes.len() as u64, + sha256: sha256_hex(bytes), + } + } + + /// A local server that answers every request with `status` and `body` (the whole file, no ranges). + fn serve(status: &'static str, body: Vec) -> String { + let l = TcpListener::bind("127.0.0.1:0").unwrap(); + let base = format!("http://{}/m", l.local_addr().unwrap()); + std::thread::spawn(move || { + for mut s in l.incoming().flatten() { + let mut r = BufReader::new(s.try_clone().unwrap()); + loop { + let mut line = String::new(); + if r.read_line(&mut line).unwrap_or(0) == 0 || line.trim().is_empty() { + break; + } + } + let _ = s.write_all(format!("HTTP/1.1 {status}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len()).as_bytes()); + let _ = s.write_all(&body); + } + }); + base + } + + #[test] + fn fetches_checks_and_stages_a_file() { + let dir = scratch("ok"); + let bytes: Vec = (0..200_000u32).map(|i| (i % 251) as u8).collect(); + let (staging, spec) = (dir.join(STAGING), spec_for(&bytes)); + std::fs::create_dir_all(&staging).unwrap(); + let mirror = serve("200 OK", bytes.clone()); + let shared = Shared::new(spec.size_bytes); + run(&spec, &mirror, &staging, &shared); + let State::Done { path, sha256 } = shared.get() else { panic!("{:?}", shared.get()) }; + assert_eq!((std::fs::read(&path).unwrap(), sha256), (bytes, spec.sha256.to_string())); + assert!(!staging.join("test-model.onnx.part").exists()); + // asked again with the copy still staged: no fetch needed (nothing listens there) + let again = Shared::new(spec.size_bytes); + run(&spec, "http://127.0.0.1:1/m", &staging, &again); + assert!(matches!(again.get(), State::Done { .. }), "{:?}", again.get()); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn a_file_that_does_not_match_is_thrown_away() { + let dir = scratch("bad"); + let staging = dir.join(STAGING); + std::fs::create_dir_all(&staging).unwrap(); + // the right size and the wrong content + let spec = spec_for(&[7u8; 5000]); + let shared = Shared::new(spec.size_bytes); + run(&spec, &serve("200 OK", vec![8u8; 5000]), &staging, &shared); + assert!(matches!(shared.get(), State::Failed(_)), "{:?}", shared.get()); + // nothing is left behind: not the file, not a partial one + assert_eq!(std::fs::read_dir(&staging).unwrap().count(), 0); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn an_unreachable_or_refusing_server_is_a_message() { + let dir = scratch("down"); + let staging = dir.join(STAGING); + std::fs::create_dir_all(&staging).unwrap(); + let spec = spec_for(b"abc"); + for mirror in ["http://127.0.0.1:1/m".to_string(), serve("404 Not Found", b"not found".to_vec())] { + let shared = Shared::new(3); + run(&spec, &mirror, &staging, &shared); + let State::Failed(why) = shared.get() else { panic!("{:?}", shared.get()) }; + assert!(why.contains("internet connection"), "{why}"); + } + assert_eq!(std::fs::read_dir(&staging).unwrap().count(), 0); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn a_cancelled_download_leaves_nothing() { + let dir = scratch("cancel"); + let staging = dir.join(STAGING); + std::fs::create_dir_all(&staging).unwrap(); + let spec = spec_for(b"abc"); + let shared = Shared::new(3); + shared.cancel.store(true, Ordering::Relaxed); + run(&spec, &serve("200 OK", b"abc".to_vec()), &staging, &shared); + assert_eq!(shared.get(), State::Cancelled); + assert_eq!(std::fs::read_dir(&staging).unwrap().count(), 0); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn dropping_the_session_stops_a_running_download() { + let shared = Arc::new(Shared::new(10)); + // a stand-in for the download thread: notices the cancel flag the way the downloader does + let worker = shared.clone(); + let handle = std::thread::spawn(move || { + while !worker.cancel.load(Ordering::Relaxed) { + std::thread::sleep(Duration::from_millis(5)); + } + worker.set(State::Cancelled); + }); + let mut d = Downloads::default(); + d.jobs.insert("m".into(), shared.clone()); + assert!(d.running()); + let started = std::time::Instant::now(); + drop(d); + assert!(started.elapsed() < Duration::from_millis(900), "quitting must not wait for a transfer to finish"); + assert_eq!(shared.get(), State::Cancelled); + handle.join().unwrap(); + } + + #[test] + fn discarding_a_finished_download_deletes_its_file() { + let dir = scratch("discard"); + let file = dir.join("staged.onnx"); + std::fs::write(&file, b"x").unwrap(); + let mut d = Downloads::default(); + let shared = Arc::new(Shared::new(1)); + shared.set(State::Done { path: file.clone(), sha256: String::new() }); + d.jobs.insert("m".into(), shared); + assert_eq!(d.snapshot().len(), 1); + assert!(!d.running()); + assert!(d.discard("m")); + assert!(!file.exists()); + assert!(d.snapshot().is_empty()); + assert!(!d.discard("m"), "nothing left to discard"); + let _ = std::fs::remove_dir_all(&dir); + } +} diff --git a/crates/engine/src/faces_index.rs b/crates/engine/src/faces_index.rs new file mode 100644 index 00000000..c250f868 --- /dev/null +++ b/crates/engine/src/faces_index.rs @@ -0,0 +1,455 @@ +//! The face index: one embedding per face region, kept in memory and cached beside the library. +//! +//! Embedding a face costs a model run (tens to hundreds of milliseconds), so each region is embedded once and +//! remembered, keyed by its photo and its box. The cache file (`face-embeddings-.bin` in the library folder, +//! one per model, so switching back to an earlier model does not start over) belongs to one model: its header names the +//! model and its file hash, and a file made by another model, an older format, or a damaged one is simply ignored and +//! rebuilt. Embeddings are kept out of the catalog on purpose: they are large, rebuildable, and not comparable between +//! models. +//! +//! Each face also keeps the box to cut its picture from (its "view"): the detector's box when the detector found the same +//! face, so every face is shown equally close however loosely or tightly its own region was drawn (a Lightroom box, an +//! older tool's box), else the region's own box. +//! +//! Layout (little-endian): `"LCFE"`, version `2`, tag length `u16`, the tag (`@`), dimension `u32`, +//! then records of `photo u64`, four `i32` (the region's box in ten-thousandths), four `i32` (the view box likewise), and +//! `dim` `f32`. + +use std::collections::HashMap; +use std::io::Write; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use lightcraft_geom::Rect; + +const MAGIC: &[u8; 4] = b"LCFE"; +const VERSION: u8 = 2; +/// Largest cache file we read. +const MAX_FILE: u64 = 1 << 30; + +pub(crate) type Key = (u64, [i32; 4]); + +/// One face's record: its embedding and its view box (quantized like a key). +type Entry = (Arc<[f32]>, [i32; 4]); + +/// A quantized box as a rectangle again. +fn rect_of(k: [i32; 4]) -> Rect { + let f = |v: i32| f64::from(v) / 10_000.0; + Rect { x0: f(k[0]), y0: f(k[1]), x1: f(k[2]), y1: f(k[3]) } +} + +/// A region's box as the integer key it is cached under. +pub(crate) fn region_key(r: &Rect) -> [i32; 4] { + let q = |v: f64| if v.is_finite() { (v.clamp(-1e3, 1e3) * 10_000.0).round() as i32 } else { 0 }; + [q(r.x0), q(r.y0), q(r.x1), q(r.y1)] +} + +#[derive(Default)] +pub(crate) struct Index { + /// `@` of the model the embeddings belong to. + pub tag: String, + pub dim: usize, + entries: HashMap, + unsaved: Vec<(Key, Entry)>, + /// Faces that could not be embedded this session (so they are not retried on every call). + skipped: std::collections::HashSet, + /// The cache file on disk belongs to this model, so new records can be appended to it. + file_matches: bool, + pub path: Option, +} + +impl Index { + pub fn len(&self) -> usize { + self.entries.len() + } + + pub fn get(&self, photo: u64, rect: &Rect) -> Option<&Arc<[f32]>> { + self.entries.get(&(photo, region_key(rect))).map(|e| &e.0) + } + + /// The box to cut this face's picture from, once the scan has looked at it. + pub fn view(&self, photo: u64, rect: &Rect) -> Option { + self.entries.get(&(photo, region_key(rect))).map(|e| rect_of(e.1)) + } + + #[cfg(test)] + pub fn contains(&self, photo: u64, rect: &Rect) -> bool { + self.entries.contains_key(&(photo, region_key(rect))) + } + + /// Whether this face still has to be embedded. + pub fn needs(&self, photo: u64, rect: &Rect) -> bool { + let key = (photo, region_key(rect)); + !self.entries.contains_key(&key) && !self.skipped.contains(&key) + } + + pub fn skip(&mut self, photo: u64, rect: &Rect) { + self.skipped.insert((photo, region_key(rect))); + } + + #[cfg(test)] + pub fn insert(&mut self, photo: u64, rect: &Rect, embedding: Vec) { + self.insert_with_view(photo, rect, embedding, rect); + } + + /// Remember a face's embedding and the box to show it by. + pub fn insert_with_view(&mut self, photo: u64, rect: &Rect, embedding: Vec, view: &Rect) { + let entry: Entry = (embedding.into(), region_key(view)); + let key = (photo, region_key(rect)); + self.entries.insert(key, entry.clone()); + self.unsaved.push((key, entry)); + } + + /// Whether there are embeddings not yet written to the cache file. + pub fn has_unsaved(&self) -> bool { + !self.unsaved.is_empty() + } + + /// Start over for `tag` (`dim` numbers per face), loading what the cache file at `path` holds for it. What the + /// previous model's index had not written yet is written first, so switching models loses nothing. + pub fn reset(&mut self, tag: &str, dim: usize, path: Option) { + let _ = self.save(); + self.tag = tag.to_string(); + self.dim = dim; + self.path = path; + self.entries.clear(); + self.unsaved.clear(); + self.skipped.clear(); + self.file_matches = false; + if let Some(p) = self.path.clone() { + self.load(&p); + } + } + + fn load(&mut self, path: &Path) { + let Ok(meta) = std::fs::metadata(path) else { return }; + if meta.len() > MAX_FILE { + return; + } + let Ok(bytes) = std::fs::read(path) else { return }; + let Some((tag_len, rest)) = bytes + .strip_prefix(MAGIC) + .and_then(|r| r.split_first()) + .filter(|(v, _)| **v == VERSION) + .and_then(|(_, r)| r.split_at_checked(2)) + .map(|(l, r)| (usize::from(u16::from_le_bytes([l[0], l[1]])), r)) + else { + return; + }; + let Some((tag, rest)) = rest.split_at_checked(tag_len) else { return }; + let Some((dim, mut rest)) = rest.split_at_checked(4).map(|(d, r)| (u32::from_le_bytes([d[0], d[1], d[2], d[3]]) as usize, r)) else { return }; + if tag != self.tag.as_bytes() || dim != self.dim || dim == 0 { + return; + } + self.file_matches = true; + let record = 8 + 16 + 16 + dim * 4; + while let Some((rec, tail)) = rest.split_at_checked(record) { + rest = tail; + let photo = u64::from_le_bytes([rec[0], rec[1], rec[2], rec[3], rec[4], rec[5], rec[6], rec[7]]); + let int = |at: usize| i32::from_le_bytes([rec[at], rec[at + 1], rec[at + 2], rec[at + 3]]); + let key = [int(8), int(12), int(16), int(20)]; + let view = [int(24), int(28), int(32), int(36)]; + let emb: Vec = rec[40..].as_chunks::<4>().0.iter().map(|c| f32::from_le_bytes(*c)).collect(); + if emb.iter().all(|v| v.is_finite()) { + self.entries.insert((photo, key), (emb.into(), view)); + } + } + } + + fn header(&self) -> Vec { + let tag = self.tag.as_bytes(); + let mut h = Vec::with_capacity(11 + tag.len()); + h.extend_from_slice(MAGIC); + h.push(VERSION); + h.extend_from_slice(&u16::try_from(tag.len()).unwrap_or(0).to_le_bytes()); + h.extend_from_slice(tag); + h.extend_from_slice(&(self.dim as u32).to_le_bytes()); + h + } + + fn record(&self, key: &Key, entry: &Entry) -> Vec { + let mut r = Vec::with_capacity(40 + entry.0.len() * 4); + r.extend_from_slice(&key.0.to_le_bytes()); + key.1.iter().chain(entry.1.iter()).for_each(|k| r.extend_from_slice(&k.to_le_bytes())); + entry.0.iter().for_each(|v| r.extend_from_slice(&v.to_le_bytes())); + r + } + + /// Write what was added since the last save: appended when the file already belongs to this model, else the + /// whole index is written fresh (atomically). A failure is returned, the records stay pending. + pub fn save(&mut self) -> std::io::Result<()> { + let Some(path) = self.path.clone() else { + self.unsaved.clear(); + return Ok(()); + }; + if self.unsaved.is_empty() { + return Ok(()); + } + if self.file_matches { + let mut f = std::fs::OpenOptions::new().append(true).open(&path)?; + let mut buf = Vec::new(); + for (k, e) in &self.unsaved { + buf.extend(self.record(k, e)); + } + f.write_all(&buf)?; + } else { + let mut buf = self.header(); + for (k, e) in &self.entries { + buf.extend(self.record(k, e)); + } + let part = path.with_extension("part"); + std::fs::write(&part, &buf)?; + std::fs::rename(&part, &path)?; + self.file_matches = true; + } + self.unsaved.clear(); + Ok(()) + } +} + +impl Drop for Index { + fn drop(&mut self) { + let _ = self.save(); + } +} + +const SCAN_MAGIC: &[u8; 4] = b"LCFS"; +const SCAN_VERSION: u8 = 1; +/// The detector a search was made with: a record made by another one is ignored (a better detector may find more). +pub(crate) const SCAN_TAG: &str = "yunet-2023mar"; +/// Largest record file we read. +const SCAN_MAX_FILE: u64 = 64 << 20; + +/// The photos the detector has already searched for faces, so each is searched once, however many times LightCraft +/// starts. Kept beside the library (`face-scanned.bin`: a header naming the detector, then photo ids, appended to). +#[derive(Default)] +pub(crate) struct Scanned { + set: std::collections::HashSet, + unsaved: Vec, + path: Option, + file_matches: bool, +} + +impl Scanned { + pub fn path(&self) -> Option<&Path> { + self.path.as_deref() + } + + pub fn contains(&self, photo: u64) -> bool { + self.set.contains(&photo) + } + + pub fn len(&self) -> usize { + self.set.len() + } + + pub fn insert(&mut self, photo: u64) { + if self.set.insert(photo) { + self.unsaved.push(photo); + } + } + + pub fn has_unsaved(&self) -> bool { + !self.unsaved.is_empty() + } + + /// Start over with the record at `path` (written first: what was not yet saved is not lost). + pub fn reset(&mut self, path: Option) { + let _ = self.save(); + self.set.clear(); + self.unsaved.clear(); + self.file_matches = false; + self.path = path; + if let Some(p) = self.path.clone() { + self.load(&p); + } + } + + fn header() -> Vec { + let mut h = Vec::new(); + h.extend_from_slice(SCAN_MAGIC); + h.push(SCAN_VERSION); + h.extend_from_slice(&(SCAN_TAG.len() as u16).to_le_bytes()); + h.extend_from_slice(SCAN_TAG.as_bytes()); + h + } + + fn load(&mut self, path: &Path) { + let Ok(meta) = std::fs::metadata(path) else { return }; + if meta.len() > SCAN_MAX_FILE { + return; + } + let Ok(bytes) = std::fs::read(path) else { return }; + let head = Self::header(); + let Some(rest) = bytes.strip_prefix(head.as_slice()) else { return }; + self.file_matches = true; + // a half-written last id (a crash while appending) is ignored + self.set.extend(rest.as_chunks::<8>().0.iter().map(|c| u64::from_le_bytes(*c))); + } + + /// Write what was added since the last save: appended when the file is this detector's, else the whole record + /// is written fresh (atomically). A failure is returned and the ids stay pending. + pub fn save(&mut self) -> std::io::Result<()> { + let Some(path) = self.path.clone() else { + self.unsaved.clear(); + return Ok(()); + }; + if self.unsaved.is_empty() { + return Ok(()); + } + if self.file_matches { + let mut f = std::fs::OpenOptions::new().append(true).open(&path)?; + let buf: Vec = self.unsaved.iter().flat_map(|id| id.to_le_bytes()).collect(); + f.write_all(&buf)?; + } else { + let mut buf = Self::header(); + buf.extend(self.set.iter().flat_map(|id| id.to_le_bytes())); + let part = path.with_extension("part"); + std::fs::write(&part, &buf)?; + std::fs::rename(&part, &path)?; + self.file_matches = true; + } + self.unsaved.clear(); + Ok(()) + } +} + +impl Drop for Scanned { + fn drop(&mut self) { + let _ = self.save(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn rect(x: f64) -> Rect { + Rect { x0: x, y0: 0.1, x1: x + 0.2, y1: 0.4 } + } + + fn temp(name: &str) -> PathBuf { + std::env::temp_dir().join(format!("lc-faceindex-{name}-{}.bin", std::process::id())) + } + + #[test] + fn embeddings_survive_a_save_and_load_and_append() { + let path = temp("roundtrip"); + let _ = std::fs::remove_file(&path); + let mut a = Index::default(); + a.reset("m@abc", 3, Some(path.clone())); + a.insert(7, &rect(0.1), vec![1.0, 0.0, 0.5]); + a.insert(8, &rect(0.3), vec![0.0, 1.0, 0.25]); + a.save().unwrap(); + a.insert(9, &rect(0.5), vec![0.5, 0.5, 0.5]); + a.save().unwrap(); // appended + let mut b = Index::default(); + b.reset("m@abc", 3, Some(path.clone())); + assert_eq!(b.len(), 3); + assert_eq!(b.get(8, &rect(0.3)).map(|e| e.to_vec()), Some(vec![0.0, 1.0, 0.25])); + assert!(b.contains(9, &rect(0.5)) && !b.contains(9, &rect(0.6)) && !b.contains(1, &rect(0.5))); + let _ = std::fs::remove_file(path); + } + + #[test] + fn another_model_a_damaged_file_or_a_torn_tail_are_handled() { + let path = temp("other"); + let _ = std::fs::remove_file(&path); + let mut a = Index::default(); + a.reset("m@abc", 3, Some(path.clone())); + a.insert(1, &rect(0.1), vec![1.0, 0.0, 0.0]); + a.insert(2, &rect(0.2), vec![0.0, 1.0, 0.0]); + a.save().unwrap(); + // another model: nothing is loaded, and saving replaces the file for the new model + let mut other = Index::default(); + other.reset("n@def", 4, Some(path.clone())); + assert_eq!(other.len(), 0); + other.insert(5, &rect(0.1), vec![1.0, 0.0, 0.0, 0.0]); + other.save().unwrap(); + let mut again = Index::default(); + again.reset("n@def", 4, Some(path.clone())); + assert_eq!(again.len(), 1); + let mut old = Index::default(); + old.reset("m@abc", 3, Some(path.clone())); + assert_eq!(old.len(), 0, "the old model's records are gone"); + // a torn last record loses only that record + let bytes = std::fs::read(&path).unwrap(); + std::fs::write(&path, &bytes[..bytes.len() - 5]).unwrap(); + let mut torn = Index::default(); + torn.reset("n@def", 4, Some(path.clone())); + assert_eq!(torn.len(), 0); + // garbage of every kind is ignored + for junk in [Vec::new(), b"LCFE".to_vec(), b"LCFE\x01\xff\xff".to_vec(), vec![0xff; 200], b"LCFE\x02\x00\x00".to_vec()] { + std::fs::write(&path, &junk).unwrap(); + let mut j = Index::default(); + j.reset("n@def", 4, Some(path.clone())); + assert_eq!(j.len(), 0); + } + // and a missing file is fine + let _ = std::fs::remove_file(&path); + let mut m = Index::default(); + m.reset("n@def", 4, Some(path.clone())); + assert_eq!(m.len(), 0); + } + + #[test] + fn non_finite_boxes_and_vectors_do_not_break_the_keys() { + assert_eq!(region_key(&Rect { x0: f64::NAN, y0: f64::INFINITY, x1: 1e300, y1: -1e300 }), [0, 0, 10_000_000, -10_000_000]); + let path = temp("nan"); + let _ = std::fs::remove_file(&path); + let mut a = Index::default(); + a.reset("m@abc", 2, Some(path.clone())); + a.insert(1, &rect(0.1), vec![f32::NAN, 1.0]); + a.save().unwrap(); + let mut b = Index::default(); + b.reset("m@abc", 2, Some(path.clone())); + assert_eq!(b.len(), 0, "a non-finite vector is not loaded back"); + let _ = std::fs::remove_file(path); + } +} + +/// What the recognition commands keep between calls: the loaded model and the embeddings made with it. +#[derive(Default)] +pub(crate) struct FacesState { + /// `@` and the loaded model. + pub embedder: Option<(String, Arc)>, + pub index: Index, + /// When the models folder was last looked at (the per-frame pump trusts a look for a second). + pub checked: Option, + /// Whether recognition was switched on when last looked up, and when. + pub enabled_seen: Option<(std::time::Instant, bool)>, + pub worker: Option, + /// After a failure to load the chosen model, the background pump leaves it alone until then. + pub retry_at: Option, + /// Photos the worker is busy with. + pub in_flight: std::collections::HashSet, + /// Photos still to hand to the worker, and the catalog revision this list was made at. + pub queue: Vec, + pub queue_stamp: Option, + /// When the index was last written to its cache file by the background pump. + pub saved_at: Option, + /// Bumped when another library is opened: what a worker was making for the old one is dropped. + pub epoch: u64, + /// Photos already searched for faces. + pub scanned: Scanned, + /// Photos that could not be searched this session (unreadable files): left alone until the next launch. + pub scan_failed: std::collections::HashSet, +} + +impl FacesState { + /// Another library is being opened: what was learned about the old one is saved, and nothing of it is kept (photo + /// ids are only unique within a library). The loaded model stays; the index is reloaded for the new library. + pub fn library_changed(&mut self) { + self.index.reset("", 0, None); + self.scanned.reset(None); + self.scan_failed.clear(); + self.embedder = None; + self.checked = None; + self.queue.clear(); + self.queue_stamp = None; + self.in_flight.clear(); + self.retry_at = None; + self.saved_at = None; + self.epoch = self.epoch.wrapping_add(1); + } +} diff --git a/crates/engine/src/faces_worker.rs b/crates/engine/src/faces_worker.rs new file mode 100644 index 00000000..81ff6ada --- /dev/null +++ b/crates/engine/src/faces_worker.rs @@ -0,0 +1,483 @@ +//! The heavy part of scanning faces, run off the session thread, on a few threads at once. +//! +//! For one photo: render it (upright, uncropped, default settings), find its faces with the YuNet detector (when installed) +//! so each region can be aligned by landmarks, cut each region out and run the recognition model. A photo that has +//! no face regions at all and has not been searched yet also gets its detections back, with their embeddings, so the +//! scan decodes each photo once. The session thread only prepares the job (which needs the catalog) and later ingests +//! the result, so the window never waits for a model. + +use std::panic::{AssertUnwindSafe, catch_unwind}; +use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering}; +use std::sync::mpsc::{Receiver, Sender, channel}; +use std::sync::{Arc, Condvar, Mutex, PoisonError}; +use std::time::Duration; + +use lightcraft_catalog::PhotoId; +use lightcraft_faces::align::{Rgb, align_to_template, crop_box}; +use lightcraft_faces::runtime::Embedder; +use lightcraft_faces::yunet::{Detector, Face, Options}; +use lightcraft_geom::Rect; + +use crate::media::{PreviewLoader, RenderJob}; + +/// Long edge of the picture faces are cut from. +pub(crate) const EDGE: usize = 2048; +/// A box from an XMP sidecar is cut with this much room around it (there are no landmarks to align by). +pub(crate) const BOX_PAD: f32 = 1.5; +/// How much a detected face must overlap a region to lend it its landmarks. +pub(crate) const MATCH_IOU: f32 = 0.3; +/// A detection must be this sure to become a face region (the bar `faces.detect` uses). +pub(crate) const REGION_SCORE: f32 = 0.6; +/// A raw's embedded preview is used for the scan only when its long edge is at least this (a thumbnail is too small to +/// find faces in; the raw itself is decoded then). +const MIN_PREVIEW_EDGE: usize = 1000; + +/// Most photos ever worked on at once. +const MAX_WORKERS: usize = 24; +/// What one photo in progress holds: the file, its decoded picture and the copies the detector and recogniser read. +/// Measured: peak memory grew by about 180 MB for each extra worker. +const WORKER_BYTES: usize = 180 << 20; +/// Photos queued behind the ones running, so a worker that finishes never waits for the next frame to be fed. +pub(crate) const PREFETCH: usize = 2; + +/// How hard the background scan may work right now, as the app judges from what the user is doing. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum Pace { + /// Start nothing new: the user is dragging, typing or scrolling. + Pause = 0, + /// One photo at a time: the user is around, or the window is minimized. + Light = 1, + /// Half the machine: the user is idle. + Normal = 2, + /// Most of the machine: the user is watching the progress. + Full = 3, +} + +impl Pace { + fn from_u8(n: u8) -> Pace { + match n { + 0 => Pace::Pause, + 1 => Pace::Light, + 3 => Pace::Full, + _ => Pace::Normal, + } + } + + /// `pause`, `light`, `normal` or `full`; anything else (or nothing) is `normal`. + pub fn parse(name: Option<&str>) -> Pace { + match name { + Some("pause") => Pace::Pause, + Some("light") => Pace::Light, + Some("full") => Pace::Full, + _ => Pace::Normal, + } + } +} + +/// How many photos are worked on at once at `pace` on a machine with `cores` threads (as the system reports them, which +/// respects container limits and affinity). `cap` is a limit the user set (`LIGHTCRAFT_FACE_THREADS`). Memory is not +/// decided here: decodes wait at the process-wide memory gate ([`crate::memory::work_gate`]), sized from the RAM. +pub(crate) fn workers_for(cores: usize, pace: Pace, cap: Option) -> usize { + let n = match pace { + Pace::Pause => return 0, + Pace::Light => 1, + Pace::Normal => cores / 2, + Pace::Full => cores * 4 / 5, + } + .clamp(1, MAX_WORKERS); + cap.map_or(n, |c| n.min(c.clamp(1, MAX_WORKERS))) +} + +fn cores() -> usize { + std::thread::available_parallelism().map_or(1, |n| n.get()) +} + +fn cap_from_env() -> Option { + std::env::var("LIGHTCRAFT_FACE_THREADS").ok().and_then(|v| v.trim().parse::().ok()) +} + +/// The most photos the memory budget can hold in progress at once: half of it (the rest is the window's, the caches', +/// the renders'), a worker at a time. Raising the budget (`LIGHTCRAFT_MEMORY_MB`) raises this. +fn memory_workers() -> usize { + (crate::memory::budget() / 2 / WORKER_BYTES).max(1) +} + +/// Photos to work on at once at `pace` on this machine: by the processor's threads and the pace, and never more than +/// the memory allows. A limit the user sets (`LIGHTCRAFT_FACE_THREADS`) replaces the memory one. +pub(crate) fn target_workers(pace: Pace) -> usize { + workers_for(cores(), pace, Some(cap_from_env().unwrap_or_else(memory_workers))) +} + +/// Threads the worker pool starts with: enough for the fullest pace (the idle ones cost nothing). +pub(crate) fn max_workers() -> usize { + target_workers(Pace::Full) +} + +/// Threads for the scan's own parallel work (decoding a picture, developing it) at `pace` on a machine with `cores` threads: +/// two when light, half the machine when normal, four fifths when watched. A photo's parallel parts are most of its cost, +/// so this, more than the number of photos at once, is how much of the processor the scan takes. Measured on a 32-thread +/// desktop (4 photos at once): 8 threads scanned 2.2 photos a second, 16 threads 5.8, 25 threads 6.4. +fn pool_threads_for(cores: usize, pace: Pace, cap: Option) -> usize { + let n = match pace { + Pace::Pause | Pace::Light => 2, + Pace::Normal => cores / 2, + Pace::Full => cores * 4 / 5, + } + .clamp(2, 32); + cap.map_or(n, |c| n.min(c.max(1))) +} + +/// One photo's faces to embed (and, if `detect`, to find first). +pub(crate) struct Prepared { + /// `@` the embeddings are for. + pub tag: String, + /// The session's face epoch when this was prepared: another library has been opened since if it differs. + pub epoch: u64, + pub id: PhotoId, + pub job: RenderJob, + /// A raw's embedded preview, tried before the render job (much faster than decoding the raw). + pub preview: Option<(String, PreviewLoader)>, + pub todo: Vec, + /// The photo has no face regions and has not been searched: look for faces too. + pub detect: bool, + /// The YuNet detector, when the user has installed it: faces are aligned by its landmarks, and it finds the faces of + /// photos that have none. Without it a scan only embeds the faces the photos already have, cut by their boxes. + pub detector: Option>, +} + +/// The embeddings made for a [`Prepared`]; `None` for a face that could not be embedded. +pub(crate) struct Done { + pub tag: String, + pub epoch: u64, + pub id: PhotoId, + /// Each face's region, its embedding (`None` when it could not be made) and the box to show it by. + pub results: Vec<(Rect, Option>, Rect)>, + /// Faces were to be looked for in this photo. + pub detect_wanted: bool, + /// What the detector found (their embeddings are in `results`); `None` when the photo could not be searched. + pub found: Option>, +} + +fn overlap(a: &Rect, f: &Face) -> f32 { + let (w, h) = ((a.x1.min(f64::from(f.x1)) - a.x0.max(f64::from(f.x0))).max(0.0), (a.y1.min(f64::from(f.y1)) - a.y0.max(f64::from(f.y0))).max(0.0)); + let inter = w * h; + let union = (a.x1 - a.x0) * (a.y1 - a.y0) + f64::from((f.x1 - f.x0) * (f.y1 - f.y0)) - inter; + if union > 0.0 { (inter / union) as f32 } else { 0.0 } +} + +/// A detection as a region box: inside the photo and not a sliver; `None` for anything else (it is only ever data). +fn face_rect(f: &Face) -> Option { + if ![f.x0, f.y0, f.x1, f.y1].iter().all(|v| v.is_finite()) { + return None; + } + let c = |v: f32| f64::from(v.clamp(0.0, 1.0)); + let r = Rect { x0: c(f.x0), y0: c(f.y0), x1: c(f.x1), y1: c(f.y1) }; + (r.x1 - r.x0 > 0.005 && r.y1 - r.y0 > 0.005).then_some(r) +} + +/// Do the work for one photo. Never panics: whatever goes wrong, its faces come back as `None`. +pub(crate) fn process(p: Prepared, embedder: &Embedder) -> Done { + let Prepared { tag, epoch, id, job, preview, todo, detect, detector } = p; + let fallback = |todo: Vec| todo.into_iter().map(|r| (r, None, r)).collect::>(); + let rects = todo.clone(); + let run = catch_unwind(AssertUnwindSafe(|| { + let from_preview = preview.and_then(|(path, load)| load(&path, EDGE)).filter(|img| img.width.max(img.height) >= MIN_PREVIEW_EDGE); + let image = match from_preview { + Some(img) => img, + None => match job.run().rendered { + Ok(rendered) => rendered.image, + Err(_) => return (fallback(todo), None), + }, + }; + let (w, h) = (image.width, image.height); + let rgb: Vec = image.data.iter().flat_map(|px| [px[0], px[1], px[2]]).collect(); + // the detector's landmarks, where it finds the same face: faces are aligned by them, not cut by a box + let detected = detector.as_deref().and_then(|d| d.detect(&rgb, w, h, &Options { score: 0.5, ..Options::default() }).ok()); + let found = detected.clone().unwrap_or_default(); + let img = Rgb { data: &rgb, width: w, height: h }; + let (iw, ih) = embedder.input_size(); + let scaled = |f: &Face| f.landmarks.map(|(x, y)| (x * w as f32, y * h as f32)); + let mut results: Vec<(Rect, Option>, Rect)> = todo + .into_iter() + .map(|rect| { + let best = found.iter().map(|f| (overlap(&rect, f), f)).filter(|(o, _)| *o >= MATCH_IOU).max_by(|a, b| a.0.total_cmp(&b.0)); + let crop = match best { + Some((_, f)) => align_to_template(&img, &scaled(f), iw, ih), + None => crop_box( + &img, + rect.x0 as f32 * w as f32, + rect.y0 as f32 * h as f32, + rect.x1 as f32 * w as f32, + rect.y1 as f32 * h as f32, + BOX_PAD, + iw, + ih, + ), + }; + let v = crop.and_then(|c| embedder.embed(&c).ok()); + // shown by the detector's box where it found this face: every face then has the same tightness + let view = best.and_then(|(_, f)| face_rect(f)).unwrap_or(rect); + (rect, v, view) + }) + .collect(); + // a photo with no face regions: what the detector is sure of becomes its faces, embedded in the same pass + let new_faces = if detect { + detected.map(|faces| { + let mut rects = Vec::new(); + for f in faces.iter().filter(|f| f.score >= REGION_SCORE) { + let Some(rect) = face_rect(f) else { continue }; + let v = align_to_template(&img, &scaled(f), iw, ih).and_then(|c| embedder.embed(&c).ok()); + results.push((rect, v, rect)); + rects.push(rect); + } + rects + }) + } else { + None + }; + (results, new_faces) + })); + let (results, found) = run.unwrap_or_else(|_| (fallback(rects), None)); + Done { tag, epoch, id, results, detect_wanted: detect, found } +} + +type Job = (Arc, Prepared); + +/// How many photos may be worked on at once, changed at any moment by the pace: a thread that has a job waits here until +/// fewer than that are running. A pace of zero holds every job where it is, with nothing decoded. +struct Limit { + allowed: AtomicUsize, + /// The pace the limit was last set for (a [`Pace`] as a number): the pool a photo's parallel work uses. + pace: AtomicU8, + running: Mutex, + wake: Condvar, +} + +impl Limit { + fn new() -> Limit { + Limit { allowed: AtomicUsize::new(0), pace: AtomicU8::new(Pace::Normal as u8), running: Mutex::new(0), wake: Condvar::new() } + } + + fn set(&self, n: usize, pace: Pace) { + self.pace.store(pace as u8, Ordering::Relaxed); + self.allowed.store(n, Ordering::Relaxed); + self.wake.notify_all(); + } + + fn pace(&self) -> Pace { + Pace::from_u8(self.pace.load(Ordering::Relaxed)) + } + + /// Wait for a place; the limit is looked at again whenever it changes and at least every 100 ms. + fn enter(&self) { + let mut running = self.running.lock().unwrap_or_else(PoisonError::into_inner); + while *running >= self.allowed.load(Ordering::Relaxed) { + running = self.wake.wait_timeout(running, Duration::from_millis(100)).unwrap_or_else(PoisonError::into_inner).0; + } + *running += 1; + } + + fn leave(&self) { + let mut running = self.running.lock().unwrap_or_else(PoisonError::into_inner); + *running = running.saturating_sub(1); + drop(running); + self.wake.notify_one(); + } +} + +/// The threads a photo's parallel work (decoding, developing a picture) runs on at `pace`: one pool for each of three sizes +/// ([`pool_threads_for`]), made when first needed, whose idle threads cost nothing. rayon's global pool serves the +/// interactive work (the loupe, exports) and has no priorities, so scan work queued there would make a slider drag wait +/// behind it; here a light scan can use only two threads however large the machine is. +fn scan_pool(pace: Pace) -> Option<&'static rayon::ThreadPool> { + static LIGHT: std::sync::OnceLock> = std::sync::OnceLock::new(); + static NORMAL: std::sync::OnceLock> = std::sync::OnceLock::new(); + static FULL: std::sync::OnceLock> = std::sync::OnceLock::new(); + let (cell, pace) = match pace { + Pace::Pause | Pace::Light => (&LIGHT, Pace::Light), + Pace::Normal => (&NORMAL, Pace::Normal), + Pace::Full => (&FULL, Pace::Full), + }; + cell.get_or_init(|| { + rayon::ThreadPoolBuilder::new() + .num_threads(pool_threads_for(cores(), pace, cap_from_env())) + .thread_name(move |i| format!("lightcraft-faces-{}-{i}", pace as u8)) + .build() + .ok() + }) + .as_ref() +} + +/// Run one photo on the pool for the pace, so its parallel parts stay there. +/// +/// It must NOT wait at the memory gate ([`crate::memory::work_gate`], which background decodes do): a job running on a +/// pool thread that waits inside a nested rayon call can pick up another job to run on top of itself, and if that one waits +/// for memory the first holds, neither can finish (seen as a few photos that never come back). The memory a scan may use is +/// instead limited by how many photos run at once ([`memory_workers`]); its decodes count as interactive ones at the gate, +/// which never wait. +fn process_in_background(prepared: Prepared, embedder: &Embedder, pace: Pace) -> Done { + let work = || process(prepared, embedder); + match scan_pool(pace) { + Some(pool) => pool.install(work), + None => work(), + } +} + +/// Threads that run [`process`] for the jobs they are given, each on one photo at a time. +pub(crate) struct Worker { + jobs: Sender, + done: Receiver, + limit: Arc, +} + +impl Worker { + pub fn start(threads: usize) -> Option { + let (jobs, job_rx) = channel::(); + let (done_tx, done) = channel::(); + let job_rx = Arc::new(Mutex::new(job_rx)); + let limit = Arc::new(Limit::new()); + let mut started = 0; + for i in 0..threads.clamp(1, MAX_WORKERS) { + let (rx, tx, limit) = (job_rx.clone(), done_tx.clone(), limit.clone()); + let spawned = std::thread::Builder::new().name(format!("lightcraft-faces-{i}")).spawn(move || { + loop { + // the lock is held only while waiting for a job, never while working on one; the threads end when + // the session (the only sender) is dropped + let next = rx.lock().unwrap_or_else(PoisonError::into_inner).recv(); + let Ok((embedder, prepared)) = next else { break }; + limit.enter(); + let done = process_in_background(prepared, &embedder, limit.pace()); + limit.leave(); + if tx.send(done).is_err() { + break; + } + } + }); + if spawned.is_ok() { + started += 1; + } + } + (started > 0).then_some(Worker { jobs, done, limit }) + } + + /// How many photos may be worked on at once from now on (what is running is not interrupted). + pub fn allow(&self, n: usize, pace: Pace) { + self.limit.set(n, pace); + } + + pub fn submit(&self, embedder: Arc, prepared: Prepared) -> bool { + self.jobs.send((embedder, prepared)).is_ok() + } + + /// Everything finished since the last call. + pub fn finished(&self) -> Vec { + self.done.try_iter().collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn the_pace_sets_how_much_of_the_machine_the_scan_may_use() { + use Pace::*; + // a 32-thread desktop: nothing while the user works, one photo around them, half when idle, most when watched + assert_eq!([Pause, Light, Normal, Full].map(|p| workers_for(32, p, None)), [0, 1, 16, 24]); + assert_eq!([Pause, Light, Normal, Full].map(|p| workers_for(8, p, None)), [0, 1, 4, 6]); + assert_eq!([Pause, Light, Normal, Full].map(|p| workers_for(4, p, None)), [0, 1, 2, 3]); + // small machines still get one worker (except while paused) + assert_eq!([Pause, Light, Normal, Full].map(|p| workers_for(1, p, None)), [0, 1, 1, 1]); + assert_eq!(workers_for(0, Full, None), 1); + assert_eq!(workers_for(10_000, Full, None), MAX_WORKERS); + // a limit set by the user caps every pace, and is itself kept sane + assert_eq!([Light, Normal, Full].map(|p| workers_for(32, p, Some(3))), [1, 3, 3]); + assert_eq!((workers_for(32, Full, Some(0)), workers_for(32, Pause, Some(5))), (1, 0)); + assert_eq!(workers_for(32, Full, Some(usize::MAX)), 24); + } + + #[test] + fn the_limit_holds_threads_to_the_pace_and_can_change_at_any_moment() { + let limit = Arc::new(Limit::new()); + let (now, peak) = (Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0))); + let started = Arc::new(AtomicUsize::new(0)); + let threads: Vec<_> = (0..6) + .map(|_| { + let (limit, now, peak, started) = (limit.clone(), now.clone(), peak.clone(), started.clone()); + std::thread::spawn(move || { + limit.enter(); + started.fetch_add(1, Ordering::SeqCst); + let n = now.fetch_add(1, Ordering::SeqCst) + 1; + peak.fetch_max(n, Ordering::SeqCst); + std::thread::sleep(Duration::from_millis(40)); + now.fetch_sub(1, Ordering::SeqCst); + limit.leave(); + }) + }) + .collect(); + // paused: nothing starts + std::thread::sleep(Duration::from_millis(150)); + assert_eq!(started.load(Ordering::SeqCst), 0); + // two at a time, never more + limit.set(2, Pace::Normal); + std::thread::sleep(Duration::from_millis(100)); + assert!(started.load(Ordering::SeqCst) >= 2); + // raised, the rest go; every thread finishes + limit.set(6, Pace::Full); + for t in threads { + t.join().unwrap(); + } + assert_eq!(started.load(Ordering::SeqCst), 6); + assert!(peak.load(Ordering::SeqCst) >= 2, "it did run in parallel"); + assert_eq!(*limit.running.lock().unwrap(), 0, "every place was given back"); + } + + #[test] + fn memory_limits_the_workers_and_a_users_limit_replaces_it() { + // the default budget (a quarter of the RAM, at most 1.5 GiB) holds a few photos in progress, not dozens + assert!((1..=8).contains(&memory_workers()), "{}", memory_workers()); + assert!(target_workers(Pace::Full) >= 1); + assert_eq!(target_workers(Pace::Pause), 0); + } + + #[test] + fn pace_names_are_forgiving() { + assert_eq!(["pause", "light", "normal", "full"].map(|n| Pace::parse(Some(n))), [Pace::Pause, Pace::Light, Pace::Normal, Pace::Full]); + assert_eq!( + (Pace::parse(None), Pace::parse(Some("")), Pace::parse(Some("FULL")), Pace::parse(Some("turbo"))), + (Pace::Normal, Pace::Normal, Pace::Normal, Pace::Normal) + ); + } + + #[test] + fn the_scans_parallel_work_gets_a_pool_the_size_of_the_pace() { + use Pace::*; + assert_eq!([Pause, Light, Normal, Full].map(|p| pool_threads_for(32, p, None)), [2, 2, 16, 25]); + assert_eq!([Light, Normal, Full].map(|p| pool_threads_for(8, p, None)), [2, 4, 6]); + assert_eq!([Light, Normal, Full].map(|p| pool_threads_for(2, p, None)), [2, 2, 2]); + assert_eq!(pool_threads_for(10_000, Full, None), 32); + assert_eq!([Light, Normal, Full].map(|p| pool_threads_for(32, p, Some(3))), [2, 3, 3]); + // the pools exist and are the sizes asked for + for pace in [Light, Normal, Full] { + let pool = scan_pool(pace).expect("a pool can be built"); + assert_eq!(pool.install(rayon::current_num_threads), pool_threads_for(cores(), pace, cap_from_env())); + } + // a paused scan has the light pool; the pace is remembered as a number and read back + assert_eq!(scan_pool(Pause).map(|p| p as *const _), scan_pool(Light).map(|p| p as *const _)); + assert_eq!([0, 1, 2, 3, 9].map(Pace::from_u8), [Pause, Light, Normal, Full, Normal]); + } + + #[test] + fn detections_become_boxes_only_when_they_make_sense() { + let f = |x0, y0, x1, y1| Face { x0, y0, x1, y1, score: 0.9, landmarks: [(0.0, 0.0); 5] }; + assert!(face_rect(&f(0.2, 0.2, 0.4, 0.5)).is_some()); + // a box that spills over the edge is brought back inside + let r = face_rect(&f(-0.1, 0.1, 0.3, 1.2)).unwrap(); + assert_eq!((r.x0, r.y1), (0.0, 1.0)); + // slivers, inverted boxes and non-numbers are dropped, never turned into regions + for bad in [f(0.2, 0.2, 0.2, 0.5), f(0.4, 0.2, 0.2, 0.5), f(f32::NAN, 0.2, 0.4, 0.5), f(0.2, 0.2, f32::INFINITY, 0.5), f(1.5, 0.2, 2.0, 0.5)] + { + assert!(face_rect(&bad).is_none(), "{bad:?}"); + } + } +} diff --git a/crates/engine/src/files.rs b/crates/engine/src/files.rs index c620be92..95481a4f 100644 --- a/crates/engine/src/files.rs +++ b/crates/engine/src/files.rs @@ -11,7 +11,7 @@ use lightcraft_pipeline::SourceInfo; use lightcraft_raster::Rgb32f; use lightcraft_raster::resample::{Filter, fit}; -use crate::media::{FileLoader, FileProbe, PreviewLoader, ProbeInfo}; +use crate::media::{DenoiseSpec, FileLoader, FileProbe, PairLoader, PreviewLoader, ProbeInfo}; fn meta_of(m: &lightcraft_meta::Metadata) -> (Meta, Option) { let shutter = m.exposure_time.map(|t| if t >= 1.0 { format!("{t:.0}") } else { format!("1/{:.0}", 1.0 / t) }).unwrap_or_default(); @@ -220,22 +220,56 @@ pub fn bin_factor(raw: &lightcraft_raw::RawImage, max_edge: usize) -> Option Result<(Rgb32f, SourceInfo), String> { - rayon::scope(|_| load_bytes_now(std::borrow::Cow::Borrowed(bytes), max_edge)) + rayon::scope(|_| load_bytes_now(std::borrow::Cow::Borrowed(bytes), max_edge, None).map(|(img, _, info)| (img, info))) } /// [`load_bytes`] taking the file's bytes: a raw file's bytes are freed as soon as it is decoded /// (less memory held while it is developed). pub fn load_vec(bytes: Vec, max_edge: usize) -> Result<(Rgb32f, SourceInfo), String> { - rayon::scope(move |_| load_bytes_now(std::borrow::Cow::Owned(bytes), max_edge)) + rayon::scope(move |_| load_bytes_now(std::borrow::Cow::Owned(bytes), max_edge, None).map(|(img, _, info)| (img, info))) } -fn load_bytes_now(bytes: std::borrow::Cow<'_, [u8]>, max_edge: usize) -> Result<(Rgb32f, SourceInfo), String> { +/// [`load_vec`] that also develops the photo's cached denoised picture (`spec`) from the same decode: the plain +/// picture, the denoised one (`None` when it cannot be read or does not fit this raw: the plain one is then all +/// there is) and the source facts. Both pictures are the same size. +pub fn load_vec_pair(bytes: Vec, max_edge: usize, spec: &DenoiseSpec) -> Result<(Rgb32f, Option, SourceInfo), String> { + rayon::scope(move |_| load_bytes_now(std::borrow::Cow::Owned(bytes), max_edge, Some(spec))) +} + +/// The denoised camera RGB of `raw` from its cached product, as the plain development would have it: the default +/// crop, and `factor` × `factor` blocks averaged into one pixel when the plain one was binned. `None` when the product +/// is missing, damaged, made from something else or not the size of this raw's active area. +fn denoised_camera_rgb(raw: &lightcraft_raw::RawImage, spec: &DenoiseSpec, factor: usize) -> Option { + use lightcraft_denoise::product::{self, Window}; + let a = raw.active_area; + let header = product::read_header_at(&spec.product).ok()?; + if header.key != spec.key || (header.width, header.height) != (a.width, a.height) { + return None; + } + if raw.opcodes.list3.is_empty() { + let c = raw.develop_crop(a.width, a.height); + let window = Window { x: c.x, y: c.y, width: c.width, height: c.height }; + product::read_window(&spec.product, &spec.key, Some(window), factor, Some(HIGHLIGHT_CLIP)).ok() + } else { + // (the plain development is not binned either when the file has opcodes of this kind: `factor` is 1) + let whole = product::read(&spec.product, &spec.key, 1).ok()?; + Some(raw.finish_demosaiced(whole)) + } +} + +fn load_bytes_now( + bytes: std::borrow::Cow<'_, [u8]>, + max_edge: usize, + denoise: Option<&DenoiseSpec>, +) -> Result<(Rgb32f, Option, SourceInfo), String> { if lightcraft_raw::probe(&bytes).is_some() { let mut raw = match lightcraft_raw::decode(&bytes).map_err(|e| preview_reason(&bytes, e)) { Ok(r) => r, Err(Ok(why)) => { // show the camera's embedded JPEG (rendered, not raw) until the variant is supported - return load_embedded_preview(&bytes, max_edge).ok_or(format!("unsupported raw ({why}) without an embedded preview")); + return load_embedded_preview(&bytes, max_edge) + .map(|(img, info)| (img, None, info)) + .ok_or(format!("unsupported raw ({why}) without an embedded preview")); } Err(Err(e)) => return Err(e), }; @@ -245,53 +279,63 @@ 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. let t0 = web_time::Instant::now(); - let binned = match bin_factor(&raw, max_edge) { + let bin = bin_factor(&raw, max_edge); + let binned = match bin { Some(k) => raw.develop_binned(k, HIGHLIGHT_CLIP).map_err(|e| e.to_string())?, None => None, }; - let mut img = match binned { + let factor = if binned.is_some() { bin.unwrap_or(1) } else { 1 }; + let img = match binned { Some(img) => img, None => { let method = if max_edge <= 600 { lightcraft_raw::Method::Bilinear } else { lightcraft_raw::Method::Ahd }; raw.develop(method).map_err(|e| e.to_string())? } }; + // the same development from the denoised mosaic, when the photo has been denoised + let twin = denoise.and_then(|spec| denoised_camera_rgb(&raw, spec, factor)).filter(|d| d.width == img.width && d.height == img.height); // the samples aren't needed any more (the colour model below reads only the tags) raw.data = lightcraft_raw::RawData::U16(Vec::new()); - let mut stages = vec![("develop", t0.elapsed())]; - stages.push(("transform", t0.elapsed())); - lightcraft_raw::highlight::reconstruct(&mut img, t.wb, HIGHLIGHT_CLIP); - stages.push(("highlights", t0.elapsed())); + let wb = t.wb; let m = camera_look.as_ref().map(|p| p.matrix.mul(&t.matrix)).unwrap_or(t.matrix).to_f32(); let hue_sat = camera_look.as_ref().and_then(|p| p.hue_sat.as_ref()).and_then(crate::camera_preview::HueSat::new); let gain = 2f32.powf(t.baseline_exposure as f32); - let wb = t.wb; // A DNG's own profile look (hue/saturation map, look table), DNG spec chapter 6. let tables = lightcraft_raw::profile::ProfileTables::new(&raw.color.profile, lightcraft_raw::color::illuminant_weight(&raw.color, xy)); - img.map_in_place(|p| { - let c = [p[0] * wb[0], p[1] * wb[1], p[2] * wb[2]]; - let rgb = [ - m[0][0] * c[0] + m[0][1] * c[1] + m[0][2] * c[2], - m[1][0] * c[0] + m[1][1] * c[1] + m[1][2] * c[2], - m[2][0] * c[0] + m[2][1] * c[1] + m[2][2] * c[2], - ]; - let rgb = match &tables { - Some(tables) => tables.apply(rgb, gain), - None => rgb.map(|v| v * gain), - }; - // after the baseline exposure, as when it was fitted - hue_sat.as_ref().map_or(rgb, |h| h.apply(rgb)).map(|v| v.max(0.0)) - }); - stages.push(("colour", t0.elapsed())); - let img = fit(&img, max_edge, max_edge, Filter::Box); - stages.push(("fit", t0.elapsed())); - let img = img.into_oriented(raw.orientation); - stages.push(("orient", t0.elapsed())); + // camera RGB → what the pipeline takes: clipped highlights rebuilt, white balance and the colour model applied, + // fitted to `max_edge` and upright (the same for the plain and the denoised picture) + let finish = |mut img: Rgb32f| { + let mut stages = vec![("develop", t0.elapsed())]; + stages.push(("transform", t0.elapsed())); + lightcraft_raw::highlight::reconstruct(&mut img, wb, HIGHLIGHT_CLIP); + stages.push(("highlights", t0.elapsed())); + img.map_in_place(|p| { + let c = [p[0] * wb[0], p[1] * wb[1], p[2] * wb[2]]; + let rgb = [ + m[0][0] * c[0] + m[0][1] * c[1] + m[0][2] * c[2], + m[1][0] * c[0] + m[1][1] * c[1] + m[1][2] * c[2], + m[2][0] * c[0] + m[2][1] * c[1] + m[2][2] * c[2], + ]; + let rgb = match &tables { + Some(tables) => tables.apply(rgb, gain), + None => rgb.map(|v| v * gain), + }; + hue_sat.as_ref().map_or(rgb, |h| h.apply(rgb)).map(|v| v.max(0.0)) + }); + stages.push(("colour", t0.elapsed())); + let img = fit(&img, max_edge, max_edge, Filter::Box); + stages.push(("fit", t0.elapsed())); + let img = img.into_oriented(raw.orientation); + stages.push(("orient", t0.elapsed())); + (img, stages) + }; + let (img, stages) = finish(img); if lightcraft_pipeline::profiling() { let mut prev = std::time::Duration::ZERO; let parts: Vec = stages @@ -304,17 +348,64 @@ fn load_bytes_now(bytes: std::borrow::Cow<'_, [u8]>, max_edge: usize) -> Result< .collect(); eprintln!("[profile] raw source {}×{} (max {max_edge}, ms after decode): {}", img.width, img.height, parts.join(", ")); } + let twin = twin.map(|picture| finish(picture).0); 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, + twin, + 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); let img = d.to_working(); let img = if img.width.max(img.height) > max_edge { fit(&img, max_edge, max_edge, Filter::Mitchell) } else { img }; - Ok((img.into_oriented(Orientation::from_exif(d.orientation)), SourceInfo::default())) + Ok((img.into_oriented(Orientation::from_exif(d.orientation)), None, SourceInfo::default())) +} + +/// The semantic masks of a DNG that AI masks understand, over the developed image (default crop, +/// oriented like it). Person mattes without a confidently selected pixel are left out (the +/// Subject mask then falls back to its heuristic); a sky matte counts even when empty (no sky). +fn dng_mattes(bytes: &[u8], raw: &lightcraft_raw::RawImage) -> Option> { + use lightcraft_pipeline::masks::{MatteKind, Mattes}; + if raw.format != lightcraft_raw::RawFormat::Dng { + return None; + } + let mut mattes = Mattes::default(); + for m in lightcraft_raw::semantic_masks(bytes) { + let Some(kind) = matte_kind(&m.name) else { continue }; + let Some(img) = m.developed(raw.active_area, raw.crop) else { continue }; + if kind != MatteKind::Sky && !img.data.iter().any(|&v| v >= 128) { + continue; + } + mattes.push(kind, img.into_oriented(raw.orientation)); + } + (!mattes.is_empty()).then(|| Arc::new(mattes)) +} + +/// What a semantic mask selects, by its `SemanticName`: Apple's (iPhone ProRAW) are named after +/// their auxiliary image types, `urn:com:apple:photo::aux:`. Other names: `None`. +fn matte_kind(name: &str) -> Option { + use lightcraft_pipeline::masks::MatteKind::*; + let kind = name.strip_prefix("urn:com:apple:photo:")?.rsplit(':').next()?; + Some(match kind { + "semanticskymatte" => Sky, + "semanticskinmatte" => Skin, + "semantichairmatte" => Hair, + "semanticteethmatte" => Teeth, + "semanticglassesmatte" => Glasses, + "portraiteffectsmatte" => Person, + _ => return None, + }) } /// A DNG `ProfileToneCurve` (linear in, linear out, 1.0 = white after exposure compensation) as the @@ -404,8 +495,31 @@ 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) + }) +} + +/// The filesystem-backed [`PairLoader`]: the file is read and decoded once, and developed as usual and from the photo's +/// denoised picture. It holds two pictures at once, so it counts for more at the memory gate than a plain load. +pub fn fs_pair_loader() -> PairLoader { + Arc::new(|path: &str, max_edge: usize, spec: &DenoiseSpec| { + let len = std::fs::metadata(path).map(|m| m.len() as usize).unwrap_or(0); + let weight = len * if max_edge <= crate::media::SourceLevel::Thumb.max_edge() { 4 } else { 10 }; + let gate = crate::memory::work_gate(); + let _permit = if crate::memory::is_background() { gate.acquire(weight) } else { gate.acquire_urgent(weight) }; + let bytes = std::fs::read(path).map_err(|e| format!("{path}: {e}"))?; + let r = load_vec_pair(bytes, max_edge, spec); + crate::memory::release(); + r + }) } /// Filesystem-backed hooks (native). On the web the host installs bytes-based hooks instead. @@ -417,7 +531,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 +542,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) @@ -443,6 +557,7 @@ impl crate::Session { self.media.file_loader = Some(l); self.media.file_probe = Some(p); self.media.preview_loader = Some(fs_preview_loader()); + self.media.denoise.loader = Some(fs_pair_loader()); self } } @@ -483,6 +598,37 @@ mod tests { assert!(probe_bytes("x.jpg", b"not an image").is_err()); } + #[test] + fn cached_denoise_uses_the_same_white_balance_profile_and_orientation_as_plain_raw() { + use lightcraft_denoise::product; + use lightcraft_raw::{DngWriteOptions, Method, Orientation, profile::HsvTable, write_dng}; + let dir = crate::tests_xmp::temp_dir("denoise-color-finish"); + let fixture = crate::tests_xmp::synthetic_dng_with(None, Default::default()); + let mut raw = lightcraft_raw::decode(&fixture).unwrap(); + raw.orientation = Orientation::Rotate90; + raw.color.baseline_exposure = 0.5; + raw.color.profile.look_table = + Some(HsvTable { hue_divisions: 1, sat_divisions: 2, val_divisions: 1, data: vec![[23.0, 0.65, 0.8]; 2], srgb_value: false }); + let bytes = write_dng(&raw, &DngWriteOptions::default()).unwrap(); + let decoded = lightcraft_raw::decode(&bytes).unwrap(); + assert!(decoded.color.profile.look_table.is_some()); + // A cache containing the ordinary demosaic must undergo the identical downstream finish. + let camera_rgb = decoded.develop(Method::Ahd).unwrap(); + let key = "synthetic-color-finish"; + let path = dir.join("same.lcdn"); + product::write(&path, &camera_rgb, key).unwrap(); + let spec = DenoiseSpec { product: path, key: key.into(), make: None }; + let (ordinary, _) = load_bytes(&bytes, 2048).unwrap(); + let (plain, cached, _) = load_vec_pair(bytes, 2048, &spec).unwrap(); + assert_eq!(plain.data, ordinary.data, "requesting a twin must not change the ordinary path"); + let cached = cached.unwrap(); + assert_eq!((plain.width, plain.height), (24, 32)); + assert_eq!((cached.width, cached.height), (plain.width, plain.height)); + let worst = plain.data.iter().zip(&cached.data).flat_map(|(a, b)| a.iter().zip(b).map(|(a, b)| (a - b).abs())).fold(0.0f32, f32::max); + assert!(worst < 0.002, "color finish differs beyond cache half-float rounding: {worst}"); + let _ = std::fs::remove_dir_all(dir); + } + /// Luminance (from linear RGB) minus its 9 × 9 local mean: tone differences between a render and a camera JPEG /// mostly cancel, edges remain. fn detail(w: usize, h: usize, lum: impl Fn(usize, usize) -> f32) -> Vec { @@ -676,8 +822,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 +845,119 @@ mod tests { assert!((mean(&dim) / mean(&before) - 0.5).abs() < 0.02, "{} vs {}", mean(&dim), mean(&before)); } + /// iPhone ProRAW-style semantic masks drive the Sky mask: Apple's sky matte (here the right + /// half of the sensor, i.e. the bottom of the photo once turned 90° clockwise — where the sky + /// heuristic would never look) is used; a matte with an unknown name is ignored. + #[test] + fn dng_sky_matte_drives_the_sky_mask() { + use lightcraft_develop::{DevelopSettings, LocalAdjustments, Mask, MaskComponent, MaskOp, MaskShape}; + use lightcraft_pipeline::{RenderRequest, render}; + use lightcraft_tiff::tags::{self as t, photometric}; + use lightcraft_tiff::{ByteOrder, IfdBuilder, ImageData, TiffWriter, Value}; + let (w, h) = (16u32, 8u32); + let image = |ifd: &mut IfdBuilder, w: u32, h: u32, cpp: u16, bits: u16, data: Vec| { + ifd.set(t::IMAGE_WIDTH, Value::Long(vec![w])); + ifd.set(t::IMAGE_LENGTH, Value::Long(vec![h])); + ifd.set(t::BITS_PER_SAMPLE, Value::Short(vec![bits; cpp as usize])); + ifd.set(t::SAMPLES_PER_PIXEL, Value::Short(vec![cpp])); + ifd.set(t::COMPRESSION, Value::Short(vec![1])); + ifd.set_image(ImageData::Strips { rows_per_strip: h, strips: vec![data] }); + }; + let mut raw = IfdBuilder::new(); + raw.set(t::NEW_SUBFILE_TYPE, Value::Long(vec![0])); + raw.set(t::PHOTOMETRIC, Value::Short(vec![photometric::LINEAR_RAW])); + image(&mut raw, w, h, 3, 16, 12000u16.to_le_bytes().repeat((w * h * 3) as usize)); + let matte = |name: &str, right: bool| { + let mut m = IfdBuilder::new(); + m.set(t::NEW_SUBFILE_TYPE, Value::Long(vec![t::SUBFILE_SEMANTIC_MASK])); + m.set(t::PHOTOMETRIC, Value::Short(vec![photometric::MASK])); + m.set(t::SEMANTIC_NAME, Value::Ascii(name.into())); + image(&mut m, w / 2, h / 2, 1, 8, (0..w / 2 * h / 2).map(|i| if (i % (w / 2) >= w / 4) == right { 255 } else { 0 }).collect()); + m + }; + let mut ifd0 = IfdBuilder::new(); + ifd0.set(t::NEW_SUBFILE_TYPE, Value::Long(vec![1])); + ifd0.set(t::DNG_VERSION, Value::Byte(vec![1, 6, 0, 0])); + ifd0.set(t::MAKE, Value::Ascii("Apple".into())); + ifd0.set(t::ORIENTATION, Value::Short(vec![6])); + ifd0.set(t::COLOR_MATRIX_1, Value::SRational(vec![(1, 1), (0, 1), (0, 1), (0, 1), (1, 1), (0, 1), (0, 1), (0, 1), (1, 1)])); + ifd0.set(t::CALIBRATION_ILLUMINANT_1, Value::Short(vec![21])); + ifd0.set(t::PHOTOMETRIC, Value::Short(vec![photometric::RGB])); + image(&mut ifd0, 4, 2, 3, 8, vec![128; 24]); + ifd0.add_sub_ifd(raw); + ifd0.add_sub_ifd(matte("urn:com:apple:photo:2020:aux:semanticskymatte", true)); + ifd0.add_sub_ifd(matte("urn:com:apple:photo:2020:aux:semanticsomethingelse", false)); + let bytes = TiffWriter::new(ByteOrder::Little, false).write(&[ifd0]).unwrap(); + + let (img, info) = load_bytes(&bytes, 64).unwrap(); + assert_eq!((img.width, img.height), (8, 16)); + assert!(info.mattes.is_some()); + let sky = Mask { + components: vec![MaskComponent { name: None, op: MaskOp::Add, invert: false, shape: MaskShape::Sky }], + adjust: LocalAdjustments { exposure: -3.0, ..Default::default() }, + ..Default::default() + }; + let s = DevelopSettings { masks: vec![sky], ..Default::default() }; + let req = RenderRequest::fit(8, 16); + let green = |s: &DevelopSettings, info: &SourceInfo, y: usize| render(&img, info, s, &req).image.get(4, y)[1] as i32; + let plain = DevelopSettings::default(); + assert_eq!(green(&s, &info, 2), green(&plain, &info, 2), "no sky at the top"); + assert!(green(&s, &info, 13) + 40 < green(&plain, &info, 13), "the sky at the bottom is darkened"); + // without the matte the heuristic looks at the top of the frame instead + let heuristic = SourceInfo { mattes: None, ..info.clone() }; + assert!(green(&s, &heuristic, 2) < green(&plain, &heuristic, 2)); + assert_eq!(green(&s, &heuristic, 13), green(&plain, &heuristic, 13)); + } + + /// A DNG's white balance is re-evaluated through its own colour model: the decoder hands the + /// pipeline the file's colour tags (without the profile look, applied at load) and the white + /// the pixels were developed for, and a custom white balance renders a neutral differently + /// from the Bradford adaptation it replaces. + #[test] + fn dng_white_balance_uses_the_camera_colour_model() { + use lightcraft_pipeline::{RenderRequest, render}; + let dng = crate::tests_xmp::synthetic_dng_with(None, Default::default()); + let raw = lightcraft_raw::decode(&dng).unwrap(); + let (img, info) = load_bytes(&dng, 64).unwrap(); + let cc = info.camera_color.as_ref().expect("a DNG with a colour matrix"); + assert_eq!(cc.developed_for, lightcraft_raw::color::as_shot_white_xy(&raw)); + assert!(cc.tags.profile.is_empty()); + assert_eq!(cc.tags.color_matrix, raw.color.color_matrix); + let mut s = lightcraft_develop::DevelopSettings::default(); + (s.wb.mode, s.wb.temp, s.wb.tint) = (lightcraft_develop::WbMode::Custom, 3000.0, 20.0); + let req = RenderRequest::fit(32, 32); + let camera = render(&img, &info, &s, &req).image; + let adapted = render(&img, &SourceInfo { camera_color: None, ..info.clone() }, &s, &req).image; + assert_ne!(camera.data, adapted.data); + } + + /// Apple ProRAW: like Lightroom Classic, the default render ignores the file's + /// `ProfileGainTableMap` (Apple's local tone mapping); a map that would double every pixel + /// changes nothing. + #[test] + fn dng_gain_table_map_is_not_rendered() { + use lightcraft_raw::gaintable::GainTableMap; + let plain = crate::tests_xmp::synthetic_dng_with(None, Default::default()); + let mut raw = lightcraft_raw::decode(&plain).unwrap(); + let (before, _) = load_bytes(&lightcraft_raw::write_dng(&raw, &Default::default()).unwrap(), 64).unwrap(); + raw.color.profile.gain_table_map = Some(GainTableMap { + points_v: 1, + points_h: 1, + points_n: 1, + spacing_v: 1.0, + spacing_h: 1.0, + origin_v: 0.0, + origin_h: 0.0, + weights: [0.2, 0.2, 0.2, 0.2, 0.2], + gamma: 1.0, + gains: vec![2.0], + }); + let with_map = lightcraft_raw::write_dng(&raw, &Default::default()).unwrap(); + assert!(lightcraft_raw::decode(&with_map).unwrap().color.profile.gain_table_map.is_some(), "the map is kept"); + let (after, _) = load_bytes(&with_map, 64).unwrap(); + assert_eq!(before.data, after.data); + } + #[test] fn unsupported_raw_falls_back_to_embedded_preview() { let b = cr3_with_preview(48, 32); diff --git a/crates/engine/src/lib.rs b/crates/engine/src/lib.rs index f26b4ca7..a6c5f642 100644 --- a/crates/engine/src/lib.rs +++ b/crates/engine/src/lib.rs @@ -10,15 +10,26 @@ //! run them off the UI thread. #![forbid(unsafe_code)] +// Model/cache types remain available without linking the optional inference crate. +#[cfg(not(feature = "denoise"))] +extern crate lightcraft_denoise_core as lightcraft_denoise; + 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 denoise; 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; @@ -33,6 +44,7 @@ pub mod logging; pub mod media; pub mod memory; pub mod merge; +mod model_download; pub mod originals; pub mod preset_import; pub mod preset_luminar; @@ -132,8 +144,10 @@ pub struct Interaction { /// Source of [`Session::visible_shared`] generations (process-wide, so two sessions never share one). static VISIBLE_GEN: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1); +static LIBRARY_GEN: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1); pub struct Session { + library_generation: u64, /// Auto Sync: edits to the active photo also change the other selected photos (the settings /// that changed, nothing else). pub auto_sync: bool, @@ -157,6 +171,17 @@ 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, + /// AI denoise: the model in use, the photos that have their picture and the work in progress. + pub(crate) denoise: denoise::State, + /// Where the host keeps face models (one folder each); `None` where there is no file system (the web). + pub face_models_dir: Option, + /// Face model downloads started this session (the staged files wait in `/.downloads`). + pub face_downloads: face_download::Downloads, + /// The user's own list of models to download (`catalog.json` in the models folder), as last read. + pub(crate) face_catalog: lightcraft_faces::catalog::Catalog, + /// The loaded recognition model and the face embeddings made with it. + #[cfg(not(target_arch = "wasm32"))] + pub(crate) faces: faces_index::FacesState, /// Copied develop settings (partial JSON) for Paste. pub clipboard: Option, /// The folder on disk the [`LibrarySource::Folder`] view browses. @@ -247,8 +272,13 @@ impl Default for Session { } impl Session { + pub fn library_generation(&self) -> u64 { + self.library_generation + } + pub fn new() -> Session { Session { + library_generation: LIBRARY_GEN.fetch_add(1, std::sync::atomic::Ordering::Relaxed), auto_sync: false, catalog: Catalog::new(), source: LibrarySource::All, @@ -263,6 +293,12 @@ impl Session { redo: Vec::new(), interaction: None, skip_auto_write: false, + denoise: Default::default(), + 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 +439,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. @@ -781,10 +838,19 @@ pub fn json_delta(old: &Value, new: &Value) -> Option { #[cfg(test)] mod tests; #[cfg(test)] +mod tests_album_order; +#[cfg(test)] mod tests_color; #[cfg(test)] +mod tests_denoise; +#[cfg(test)] mod tests_export; #[cfg(test)] +mod tests_face_models; +#[cfg(test)] +#[cfg(not(target_arch = "wasm32"))] +mod tests_face_recognize; +#[cfg(test)] mod tests_folders; #[cfg(test)] mod tests_forget_local; diff --git a/crates/engine/src/library.rs b/crates/engine/src/library.rs index c1b91c4c..a83dc4a8 100644 --- a/crates/engine/src/library.rs +++ b/crates/engine/src/library.rs @@ -297,6 +297,12 @@ impl Session { { log::error!("library: {e}"); } + // Photo ids belong to one library; denoise work must not carry into the next. + self.denoise.library_changed(); + self.media.denoise.clear(); + // 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 @@ -305,6 +311,7 @@ impl Session { static LOADS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); catalog.revision = LOADS.fetch_add(1, std::sync::atomic::Ordering::Relaxed).wrapping_add(1) << 32; self.catalog = catalog; + self.library_generation = crate::LIBRARY_GEN.fetch_add(1, std::sync::atomic::Ordering::Relaxed); self.undo.clear(); self.redo.clear(); self.interaction = None; diff --git a/crates/engine/src/media.rs b/crates/engine/src/media.rs index c3b3c468..b0439390 100644 --- a/crates/engine/src/media.rs +++ b/crates/engine/src/media.rs @@ -62,7 +62,7 @@ impl SettingsHashes { } /// Bump when the pipeline's output changes, to invalidate cached thumbnails. -pub const RENDER_CACHE_VERSION: u64 = 15; +pub const RENDER_CACHE_VERSION: u64 = 16; /// Thumbnails render at one of these long edges (so window/cell size changes reuse the cache). pub const THUMB_SIZES: [usize; 4] = [128, 256, 384, 512]; @@ -125,6 +125,43 @@ pub type FileLoader = Arc Result<(Rgb32f, SourceInfo), St /// `max_edge` (set by the app). `None`: no usable preview. pub type PreviewLoader = Arc Option + Send + Sync>; +/// Makes a photo's denoised picture when it is not cached yet (an export does not wait for the background queue): +/// `Ok(true)` once it is there, `Ok(false)` when denoise does not apply to this photo (render it plain), `Err` when it +/// should have worked and did not. +pub type MakeProduct = Arc Result + Send + Sync>; + +/// A raw photo's cached AI-denoised picture (see [`crate::denoise`]): where the file is and the key it must have been +/// made with. A file that is missing, damaged or made for something else is a cache miss, never an error. +#[derive(Clone)] +pub struct DenoiseSpec { + pub product: std::path::PathBuf, + pub key: String, + /// Run before the picture is read, by a job that must not render without it (an export). + pub make: Option, +} + +/// Decodes a raw file once and develops it twice: as usual, and from the denoised picture in `spec` when that can be +/// read (`None` when it cannot). Both are at most `max_edge` long and the same size. +pub type PairLoader = Arc Result<(Rgb32f, Option, SourceInfo), String> + Send + Sync>; + +/// A photo's denoised picture next to its plain one, and the last mix of the two that was asked for. +pub struct Twin { + pub image: Arc, + /// The Amount (in thousandths) and the picture mixed for it, so dragging other sliders keeps one source picture. + mix: std::sync::Mutex)>>, +} + +impl Twin { + pub fn new(image: Arc) -> Twin { + Twin { image, mix: std::sync::Mutex::new(None) } + } + + /// Bytes reserved for the twin and its lazy mix, before either cache cost can grow. + fn bytes(&self) -> usize { + self.image.data.len().saturating_mul(24) + } +} + /// Pixels and the source interpretation learned while decoding; they must be cached together. #[derive(Clone)] pub struct DecodedSource { @@ -133,17 +170,52 @@ pub struct DecodedSource { /// A smart preview's stored camera tone curve: the one decoder fact its pixels need that the /// catalog's header facts lack (used when `info` is `None`). pub camera_tone: Option, + /// The same picture developed from the AI-denoised mosaic, when the photo has one cached. + pub denoised: Option>, } impl DecodedSource { pub fn new(image: Arc, info: Option) -> Self { - DecodedSource { image, info, camera_tone: None } + DecodedSource { image, info, camera_tone: None, denoised: None } + } + + /// The picture to develop for a Denoise `amount` (0 to 1, see [`DevelopSettings::denoise_amount`]): the plain one, + /// the denoised one, or a mix of the two. A mix is made once per amount and kept, so the stages after it are + /// reused until the amount changes. + pub fn image_for(&self, amount: f32) -> Arc { + let Some(twin) = &self.denoised else { return self.image.clone() }; + if amount.is_nan() || amount <= 0.0 || twin.image.width != self.image.width || twin.image.height != self.image.height { + return self.image.clone(); + } + if amount >= 1.0 { + return twin.image.clone(); + } + let thousandths = (amount * 1000.0).round().clamp(1.0, 999.0) as u32; + let mut memo = twin.mix.lock().unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some((t, img)) = memo.as_ref() + && *t == thousandths + { + return img.clone(); + } + let mut mixed = (*self.image).clone(); + if !lightcraft_denoise::product::blend(&mut mixed, &twin.image, thousandths as f32 / 1000.0) { + return self.image.clone(); + } + let mixed = Arc::new(mixed); + *memo = Some((thousandths, mixed.clone())); + mixed + } + + /// Cache cost: the picture, its segmentation mattes, its denoised twin and room for the lazy mix. + pub fn bytes(&self) -> usize { + let mattes = self.info.as_ref().and_then(|i| i.mattes.as_ref()).map_or(0, |m| m.bytes()); + self.image.data.len() * 12 + 64 + std::mem::size_of::() + mattes + self.denoised.as_ref().map_or(0, |t| t.bytes()) } /// 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 }) } } @@ -161,6 +233,8 @@ pub enum SourceRef { /// The photo's smart preview, used when the original can't be read (offline drive). The /// UI thread never checks whether the original is there; the render worker finds out. fallback: Option, + /// The photo has a cached denoised picture: load it with the plain one (see [`PairLoader`]). + denoise: Option<(DenoiseSpec, PairLoader)>, }, /// A smart preview standing in for a missing original. Smart { @@ -177,15 +251,45 @@ impl SourceRef { let image = match self { SourceRef::Loaded(a) => return Ok((**a).clone()), SourceRef::Demo { scene, max_edge } => crate::demo::scene_pixels(scene, *max_edge), - SourceRef::File { path, max_edge, loader, fallback } => { - let r = match loader { + SourceRef::File { path, max_edge, loader, fallback, denoise } => { + let plain = || match loader { Some(l) => l(path, *max_edge).map(|(image, info)| DecodedSource::new(Arc::new(image), Some(info))), None => Err(format!("no decoder available for {path}")), }; + // the denoised picture comes with the plain one from a single decode; when it cannot be read the + // pair loader gives the plain picture alone (the cached file is never a reason to fail a render) + let r = match denoise { + Some((spec, pair)) => { + // a job that needs the denoised picture has it made first; if that fails the job fails (no smart + // preview stands in: an export must not quietly come out without the denoise it was asked for) + let made = match &spec.make { + Some(make) => make()?, + None => true, + }; + if made { + pair(path, *max_edge, spec).and_then(|(image, twin, info)| { + if spec.make.is_some() && twin.is_none() { + return Err("the cached AI Denoise picture could not be read; clear the denoise cache and retry".into()); + } + Ok(DecodedSource { + image: Arc::new(image), + info: Some(info), + camera_tone: None, + denoised: twin.map(|t| Arc::new(Twin::new(Arc::new(t)))), + }) + }) + } else { + plain() + } + } + None => plain(), + }; // An offline original renders from its smart preview (no decoder facts: header ones). return match (r, fallback) { #[cfg(not(target_arch = "wasm32"))] - (Err(e), Some(sp)) if crate::smart::is_valid(sp) => crate::smart::load(sp).map_err(|e2| format!("{e}; smart preview: {e2}")), + (Err(e), Some(sp)) if denoise.as_ref().is_none_or(|(spec, _)| spec.make.is_none()) && crate::smart::is_valid(sp) => { + crate::smart::load(sp).map_err(|e2| format!("{e}; smart preview: {e2}")) + } (r, _) => r, }; } @@ -224,6 +328,54 @@ pub struct MediaCache { scenes: Vec, /// Rendered thumbnails (memory, plus disk once a library is attached). pub rendered: Arc, + /// Which photos have a cached denoised picture, and how to read it. + pub denoise: DenoiseIndex, +} + +/// The photos that have a cached denoised picture for the denoise model in use. Only memory: nothing here touches the +/// file system, so it can be asked on the UI thread. +#[derive(Default)] +pub struct DenoiseIndex { + ready: std::collections::HashMap, + /// Reads the pictures (`None`: this build cannot, and nothing is ever ready). + pub loader: Option, +} + +impl DenoiseIndex { + /// What to load for `id` besides its plain picture, and how. + pub fn get(&self, id: PhotoId) -> Option<(&DenoiseSpec, &PairLoader)> { + self.ready.get(&id).zip(self.loader.as_ref()) + } + + pub fn spec(&self, id: PhotoId) -> Option<&DenoiseSpec> { + self.ready.get(&id) + } + + pub fn set(&mut self, id: PhotoId, spec: DenoiseSpec) { + self.ready.insert(id, spec); + } + + pub fn remove(&mut self, id: PhotoId) -> bool { + self.ready.remove(&id).is_some() + } + + pub fn clear(&mut self) { + self.ready.clear(); + } + + pub fn len(&self) -> usize { + self.ready.len() + } + + pub fn is_empty(&self) -> bool { + self.ready.is_empty() + } + + pub fn ids(&self) -> Vec { + let mut v: Vec = self.ready.keys().copied().collect(); + v.sort(); + v + } } impl Default for MediaCache { @@ -244,6 +396,7 @@ impl Default for MediaCache { availability: Default::default(), scenes: Vec::new(), rendered: Arc::new(PreviewCache::memory(rendered_budget(budget))), + denoise: DenoiseIndex::default(), } } } @@ -253,10 +406,6 @@ 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::() -} - impl MediaCache { /// Keep rendered thumbnails on disk in `dir` as well. The cache this replaces is retired /// (its files stay): render jobs still holding it can't write through it any more, even @@ -319,8 +468,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| self.previews.remove(i).1.bytes()), + 2 => self.full.take().map(|e| e.1.bytes()), _ => self.rendered.evict_oldest(), }; match freed { @@ -372,7 +521,7 @@ impl MediaCache { let tick = lightcraft_preview::next_tick(); match level { SourceLevel::Thumb => { - let cost = source_bytes(&img.image); + let cost = img.bytes(); self.thumbs.insert(id, img, cost); } SourceLevel::Preview => { @@ -415,13 +564,42 @@ 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| e.1.bytes()).sum()); + let full = self.full.as_ref().map(|e| Usage::new(1, e.1.bytes())).unwrap_or_default(); (Usage::new(self.thumbs.len(), self.thumbs.cost()), previews, full) } + /// What a render of photo `id` with settings `s` adds to its identity for the denoised picture it will use: 0 unless + /// it uses one (the Denoise amount is above 0 and the photo has a cached denoised picture). + pub fn denoise_salt(&self, id: PhotoId, s: &DevelopSettings) -> u64 { + match self.denoise.spec(id) { + Some(spec) if s.denoise_amount() > 0.0 => Hasher128::new().str(&spec.key).finish().0 as u64 | 1, + _ => 0, + } + } + + /// How to load photo `p` at `level`: from the cache when its source is there, else from its file (with its denoised + /// picture beside it when it has one), its smart preview or the demo scene. pub fn source_ref(&mut self, p: &Photo, level: SourceLevel) -> SourceRef { - if let Some(a) = self.get_source(p.id, level) { + self.source_ref_with(p, level, true) + } + + /// [`Self::source_ref`] that never takes a decoded source from memory: for a job that needs something the cached + /// one lacks (an export, whose denoised picture is made when it runs). + pub fn source_ref_uncached(&mut self, p: &Photo, level: SourceLevel) -> SourceRef { + self.source_ref_with(p, level, false) + } + + fn source_ref_with(&mut self, p: &Photo, level: SourceLevel, cached: bool) -> SourceRef { + let mut r = self.source_ref_plain(p, level, cached); + if let (SourceRef::File { denoise, .. }, Some((spec, loader))) = (&mut r, self.denoise.get(p.id)) { + *denoise = Some((spec.clone(), loader.clone())); + } + r + } + + fn source_ref_plain(&mut self, p: &Photo, level: SourceLevel, cached: bool) -> SourceRef { + if cached && let Some(a) = self.get_source(p.id, level) { return SourceRef::Loaded(Box::new(a)); } // Procedural scenes have a nominal size: "full" is that size, not unbounded. @@ -459,10 +637,12 @@ impl MediaCache { } match self.scenes.iter().find(|s| s.id == *scene) { Some(s) => SourceRef::Demo { scene: Box::new(s.clone()), max_edge }, - None => SourceRef::File { path: format!("demo:{scene}"), max_edge, loader: None, fallback: None }, + None => SourceRef::File { path: format!("demo:{scene}"), max_edge, loader: None, fallback: None, denoise: None }, } } - Source::File { path } => SourceRef::File { path: path.clone(), max_edge, loader: self.file_loader.clone(), fallback: None }, + Source::File { path } => { + SourceRef::File { path: path.clone(), max_edge, loader: self.file_loader.clone(), fallback: None, denoise: None } + } } } } @@ -573,7 +753,8 @@ impl RenderJob { let was_loaded = matches!(self.source, SourceRef::Loaded(_)); match self.source.load_source() { Ok(source) => { - let src = &source.image; + // the plain picture, or (Denoise amount above 0 and a denoised picture cached) the mix the amount asks for + let src = &source.image_for(self.settings.denoise_amount()); let info = source.info_or(self.info); // Thumbnails (many small jobs side by side) stay on the CPU; views and exports use // the GPU when there is one. @@ -757,22 +938,23 @@ impl crate::Session { let content = source_key.0 as u64; // Temporary "before" settings must not replace the photo's cached develop hash. let settings_hash = if before { settings.hash64() } else { self.media.settings_hashes.get(id, &settings) }; + let salt = self.media.denoise_salt(id, &settings); let key = settings_hash ^ ((max_w as u64) << 40) ^ ((max_h as u64) << 20) ^ (apply_crop as u64) ^ (level as u64) << 60 ^ id.0.wrapping_mul(0x9e37_79b9_7f4a_7c15) - ^ content.rotate_left(17); + ^ content.rotate_left(17) + ^ salt; let cache = thumb_bucket.map(|b| { - let k = Hasher128::new() - .str(&content_key) - .u64(settings_hash) - .u64(b as u64) - .u64(RENDER_CACHE_VERSION) - .u64(crate::camera_profiles::cache_key()) - .finish(); - (self.media.rendered.clone(), k) + let mut k = Hasher128::new(); + k.str(&content_key).u64(settings_hash).u64(b as u64).u64(RENDER_CACHE_VERSION).u64(crate::camera_profiles::cache_key()); + // (only a render that uses a denoised picture has it in its key, so every other thumbnail stays valid) + if salt != 0 { + k.u64(salt); + } + (self.media.rendered.clone(), k.finish()) }); Some(RenderJob { request_id: 0, @@ -979,6 +1161,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 { @@ -1052,6 +1245,15 @@ impl crate::Session { depth: lightcraft_pipeline::OutputDepth, ) -> Result { let mut job = self.render_job(id, max_w, max_h, false, true).ok_or("no such photo")?; + // An export is made from the denoised picture its Denoise amount asks for, made when the job runs if the + // background queue has not got to it. (A source kept in memory may lack the picture, so the job reads the file.) + if let Some(wanted) = self.denoise_for_export(id, &job.settings)? { + let p = self.catalog.photo(id).ok_or("no such photo")?.clone(); + job.source = self.media.source_ref_uncached(&p, job.level); + if let SourceRef::File { denoise, .. } = &mut job.source { + *denoise = Some(wanted); + } + } job.request.space = space; job.request.depth = depth; job.key ^= (space as u64 + 1).wrapping_mul(0xa076_1d64_78bd_642f) ^ (depth as u64 + 1).wrapping_mul(0xe703_7ed1_a0b4_28db); @@ -1199,6 +1401,20 @@ mod thumbnail_hash_tests { mod tests { use super::*; + #[test] + fn denoise_mix_is_reserved_before_the_source_enters_the_cache() { + let mut plain = Rgb32f::new(8, 8); + plain.data.fill([1.0; 3]); + let mut source = DecodedSource::new(Arc::new(plain), None); + source.denoised = Some(Arc::new(Twin::new(Arc::new(Rgb32f::new(8, 8))))); + let reserved = source.bytes(); + assert!(reserved >= 3 * 8 * 8 * 12); + assert_eq!(source.image_for(0.5).data[0], [0.5; 3]); + assert_eq!(source.bytes(), reserved, "creating a mix cannot evade the LRU's recorded cost"); + let _ = source.image_for(0.25); + assert_eq!(source.bytes(), reserved, "changing Amount replaces the reserved mix"); + } + #[test] fn decoder_info_survives_render_jobs_cache_and_eviction() { let tone = lightcraft_pipeline::tone::CameraTone::new(std::array::from_fn(|i| { @@ -1210,20 +1426,22 @@ 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, + denoise: None, }; job.info = SourceInfo::default(); // Header facts cannot override decoder facts. job.settings = Arc::new(DevelopSettings::default()); let r = job.clone().run(); - assert_eq!(r.loaded.as_ref().unwrap().info, Some(info)); + assert_eq!(r.loaded.as_ref().unwrap().info.as_ref(), Some(&info)); let expected = lightcraft_pipeline::render(&r.loaded.as_ref().unwrap().image, &info, &job.settings, &job.request); assert_eq!(r.rendered.as_ref().unwrap().image.data, expected.image.data); s.accept(&r); diff --git a/crates/engine/src/model_download.rs b/crates/engine/src/model_download.rs new file mode 100644 index 00000000..98707f78 --- /dev/null +++ b/crates/engine/src/model_download.rs @@ -0,0 +1,359 @@ +//! Downloading a denoise model explicitly requested after showing its terms. +//! +//! 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_denoise::known::Download`] can be fetched: +//! the address and hash come from our pinned catalog or the user's model catalog, and no token is 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 (`denoise.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_denoise::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. + Installing, + Installed, + Failed(String), + Cancelled, +} + +struct Shared { + state: Mutex, + cancel: AtomicBool, +} + +impl Shared { + fn new(total: u64) -> Self { + Shared { state: Mutex::new(State::Running { bytes: 0, total }), cancel: AtomicBool::new(false) } + } + + fn set(&self, s: State) { + *self.state.lock().unwrap_or_else(PoisonError::into_inner) = s; + } + + fn get(&self) -> State { + self.state.lock().unwrap_or_else(PoisonError::into_inner).clone() + } +} + +/// The downloads of one session, by model id. +#[derive(Default)] +pub struct Downloads { + jobs: HashMap>, +} + +impl Downloads { + /// Start fetching `spec` into `/.downloads/`. Different models can download side by side; the same + /// one is started once. + pub fn start(&mut self, spec: Download, models_dir: &Path) -> Result<(), String> { + if self.jobs.get(&spec.id).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("denoise-model-download".into()) + .spawn(move || { + let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| run(&spec, &spec.url, &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(()) + } + + /// 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::Installing => {} + State::Done { path, .. } => { + let _ = std::fs::remove_file(path); + self.jobs.remove(id); + } + State::Installed | State::Failed(_) | State::Cancelled => { + self.jobs.remove(id); + } + } + true + } + + /// 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 = web_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 its exact URL 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_url(&file, mirror, 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) => { + 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_denoise::hash::sha256_hex; + use std::io::{BufRead, BufReader, Write}; + use std::net::TcpListener; + + fn scratch(name: &str) -> PathBuf { + let d = std::env::temp_dir().join(format!("lc-face-dl-{name}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&d); + std::fs::create_dir_all(&d).unwrap(); + d + } + + fn spec_for(bytes: &[u8]) -> Download { + Download { + id: "test-model".into(), + url: "https://example.invalid/m/test-model.onnx".into(), + file_name: "test-model.onnx".into(), + size_bytes: bytes.len() as u64, + sha256: sha256_hex(bytes), + } + } + + /// A local server that answers every request with `status` and `body` (the whole file, no ranges). + fn serve(status: &'static str, body: Vec) -> String { + let l = TcpListener::bind("127.0.0.1:0").unwrap(); + let base = format!("http://{}/m", l.local_addr().unwrap()); + std::thread::spawn(move || { + for mut s in l.incoming().flatten() { + let mut r = BufReader::new(s.try_clone().unwrap()); + loop { + let mut line = String::new(); + if r.read_line(&mut line).unwrap_or(0) == 0 || line.trim().is_empty() { + break; + } + } + let _ = s.write_all(format!("HTTP/1.1 {status}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len()).as_bytes()); + let _ = s.write_all(&body); + } + }); + base + } + + #[test] + fn fetches_checks_and_stages_a_file() { + let dir = scratch("ok"); + let bytes: Vec = (0..200_000u32).map(|i| (i % 251) as u8).collect(); + let (staging, spec) = (dir.join(STAGING), spec_for(&bytes)); + std::fs::create_dir_all(&staging).unwrap(); + let mirror = serve("200 OK", bytes.clone()); + let shared = Shared::new(spec.size_bytes); + run(&spec, &mirror, &staging, &shared); + let State::Done { path, sha256 } = shared.get() else { panic!("{:?}", shared.get()) }; + assert_eq!((std::fs::read(&path).unwrap(), sha256), (bytes, spec.sha256.to_string())); + assert!(!staging.join("test-model.onnx.part").exists()); + // asked again with the copy still staged: no fetch needed (nothing listens there) + let again = Shared::new(spec.size_bytes); + run(&spec, "http://127.0.0.1:1/m", &staging, &again); + assert!(matches!(again.get(), State::Done { .. }), "{:?}", again.get()); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn a_file_that_does_not_match_is_thrown_away() { + let dir = scratch("bad"); + let staging = dir.join(STAGING); + std::fs::create_dir_all(&staging).unwrap(); + // the right size and the wrong content + let spec = spec_for(&[7u8; 5000]); + let shared = Shared::new(spec.size_bytes); + run(&spec, &serve("200 OK", vec![8u8; 5000]), &staging, &shared); + assert!(matches!(shared.get(), State::Failed(_)), "{:?}", shared.get()); + // nothing is left behind: not the file, not a partial one + assert_eq!(std::fs::read_dir(&staging).unwrap().count(), 0); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn an_unreachable_or_refusing_server_is_a_message() { + let dir = scratch("down"); + let staging = dir.join(STAGING); + std::fs::create_dir_all(&staging).unwrap(); + let spec = spec_for(b"abc"); + for mirror in ["http://127.0.0.1:1/m".to_string(), serve("404 Not Found", b"not found".to_vec())] { + let shared = Shared::new(3); + run(&spec, &mirror, &staging, &shared); + let State::Failed(why) = shared.get() else { panic!("{:?}", shared.get()) }; + assert!(why.contains("internet connection"), "{why}"); + } + assert_eq!(std::fs::read_dir(&staging).unwrap().count(), 0); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn a_cancelled_download_leaves_nothing() { + let dir = scratch("cancel"); + let staging = dir.join(STAGING); + std::fs::create_dir_all(&staging).unwrap(); + let spec = spec_for(b"abc"); + let shared = Shared::new(3); + shared.cancel.store(true, Ordering::Relaxed); + run(&spec, &serve("200 OK", b"abc".to_vec()), &staging, &shared); + assert_eq!(shared.get(), State::Cancelled); + assert_eq!(std::fs::read_dir(&staging).unwrap().count(), 0); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn dropping_the_session_stops_a_running_download() { + let shared = Arc::new(Shared::new(10)); + // a stand-in for the download thread: notices the cancel flag the way the downloader does + let worker = shared.clone(); + let handle = std::thread::spawn(move || { + while !worker.cancel.load(Ordering::Relaxed) { + std::thread::sleep(Duration::from_millis(5)); + } + worker.set(State::Cancelled); + }); + let mut d = Downloads::default(); + d.jobs.insert("m".into(), shared.clone()); + assert!(d.running()); + let started = web_time::Instant::now(); + drop(d); + assert!(started.elapsed() < Duration::from_millis(900), "quitting must not wait for a transfer to finish"); + assert_eq!(shared.get(), State::Cancelled); + handle.join().unwrap(); + } + + #[test] + fn discarding_a_finished_download_deletes_its_file() { + let dir = scratch("discard"); + let file = dir.join("staged.onnx"); + std::fs::write(&file, b"x").unwrap(); + let mut d = Downloads::default(); + let shared = Arc::new(Shared::new(1)); + shared.set(State::Done { path: file.clone(), sha256: String::new() }); + d.jobs.insert("m".into(), shared); + assert_eq!(d.snapshot().len(), 1); + assert!(!d.running()); + assert!(d.discard("m")); + assert!(!file.exists()); + assert!(d.snapshot().is_empty()); + assert!(!d.discard("m"), "nothing left to discard"); + let _ = std::fs::remove_dir_all(&dir); + } +} diff --git a/crates/engine/src/segment/mod.rs b/crates/engine/src/segment/mod.rs index 75e3d11b..c21873f5 100644 --- a/crates/engine/src/segment/mod.rs +++ b/crates/engine/src/segment/mod.rs @@ -157,7 +157,14 @@ impl Segmenter { /// Whether this build can compute AI masks at all. pub const AVAILABLE: bool = cfg!(feature = "sam"); - /// Whether the model's files are in place. + pub fn remote_endpoint(&self) -> Option { + #[cfg(feature = "sam")] + return lightcraft_segment::remote::configured(); + #[cfg(not(feature = "sam"))] + None + } + + /// Whether the local model files are in place. pub fn installed(&self) -> bool { #[cfg(feature = "sam")] if let Some(dir) = &self.dir { @@ -230,6 +237,10 @@ impl Segmenter { if !Self::AVAILABLE { return Err("AI masks are not available in this build".into()); } + #[cfg(feature = "sam")] + if lightcraft_segment::remote::configured().is_some() { + return Ok(PathBuf::new()); + } let dir = self.dir.clone().ok_or("no folder is set for the SAM 3 model")?; if self.installed() { return Ok(dir); diff --git a/crates/engine/src/segment/worker.rs b/crates/engine/src/segment/worker.rs index b4a1eedf..5769fd2a 100644 --- a/crates/engine/src/segment/worker.rs +++ b/crates/engine/src/segment/worker.rs @@ -17,7 +17,10 @@ use std::sync::mpsc::{self, RecvTimeoutError}; use std::time::Duration; use lightcraft_develop::SegMask; -use lightcraft_segment::{Encoded, Sam3}; +use lightcraft_segment::{ + Encoded, Sam3, + remote::{self, Operation, Request, text_logits}, +}; use super::{Input, Prompt, Tag}; use crate::media::RenderJob; @@ -175,6 +178,7 @@ impl Worker { struct State { model: Option<(PathBuf, Sam3)>, cache: Option<(u64, Encoded)>, + remote: Option, } /// Whether `later` makes `job` pointless (it asks for the same selection again). @@ -195,7 +199,7 @@ fn run(rx: &mpsc::Receiver, shared: &Shared, idle: Duration) { match rx.recv_timeout(idle) { Ok(j) => queue.push_back(j), Err(RecvTimeoutError::Timeout) => { - if state.model.is_some() { + if state.model.is_some() || state.remote.is_some() { state = State::default(); shared.loaded.store(false, Ordering::SeqCst); log::info!("SAM 3 unloaded after {} minutes without use", idle.as_secs() / 60); @@ -231,7 +235,7 @@ fn run(rx: &mpsc::Receiver, shared: &Shared, idle: Duration) { if result.is_err() && !is_detail { state.cache = None; } - shared.loaded.store(state.model.is_some(), Ordering::SeqCst); + shared.loaded.store(state.model.is_some() || state.remote.is_some(), Ordering::SeqCst); drop(done); let _ = job.reply.send(Outcome { tag: job.tag, result, superseded: false }); } @@ -284,29 +288,68 @@ fn encoded<'a>(state: &'a mut State, shared: &Shared, job: &mut Job) -> Result<( log::info!("SAM 3 encoded {}×{} in {:?}", input.w, input.h, t.elapsed()); state.cache = Some((job.key, enc)); } - let State { model, cache } = state; + let State { model, cache, .. } = state; match (model.as_mut(), cache.as_mut()) { (Some((_, m)), Some((_, e))) => Ok((m, e)), _ => Err("the model is not ready".into()), } } -/// Every phrase of `text` ("car, road"), merged (the per-pixel maximum); `None` when nothing -/// matches. -fn text_logits(model: &mut Sam3, enc: &mut Encoded, text: &str) -> Result>, String> { - let mut merged: Option> = None; - for phrase in text.split([',', ';']).map(str::trim).filter(|p| !p.is_empty()) { - if let Some(p) = model.segment_text(enc, phrase, 0.5).map_err(|e| e.to_string())? { - merged = Some(match merged { - None => p.logits, - Some(m) => m.iter().zip(&p.logits).map(|(a, b)| a.max(*b)).collect(), - }); - } +fn remote_job(state: &mut State, shared: &Shared, job: &mut Job, endpoint: &str) -> Result, String> { + let _busy = Flag::raise(&shared.encoding); + if state.remote.is_none() { + state.remote = Some(remote::Client::connect(endpoint)?); } - Ok(merged) + let client = state.remote.as_mut().ok_or("remote SAM unavailable")?; + let mut region = None; + let operation = match std::mem::replace(&mut job.kind, Kind::Prepare) { + Kind::Prepare => Operation::Prepare, + Kind::Text(text) => Operation::Text(text), + Kind::Clicks(clicks) => Operation::Clicks(clicks.iter().map(|c| [c.x, c.y, if c.positive { 1.0 } else { 0.0 }]).collect()), + Kind::Detail { prompt, region: r } => { + region = Some(r); + let (rw, rh) = (r[2] - r[0], r[3] - r[1]); + if !(rw > 0.0 && rh > 0.0) { + return Ok(None); + } + match prompt { + Prompt::Text(text) => Operation::Text(text), + Prompt::Clicks(clicks) => { + let inside: Vec<_> = clicks + .iter() + .map(|c| [((c.at.x - r[0]) / rw) as f32, ((c.at.y - r[1]) / rh) as f32, if c.include { 1.0 } else { 0.0 }]) + .filter(|p| (0.0..=1.0).contains(&p[0]) && (0.0..=1.0).contains(&p[1])) + .collect(); + if !inside.iter().any(|p| p[2] == 1.0) { + return Ok(None); + } + Operation::Clicks(inside) + } + } + } + #[cfg(test)] + Kind::Panic => return Err("test panic request".into()), + }; + let input = if region.is_some() || client.key != Some(job.key) { render(job)? } else { Input { rgb: Vec::new(), w: 0, h: 0 } }; + let request = Request { width: input.w, height: input.h, temporary: region.is_some(), operation }; + let logits = client.request(&request, &input.rgb)?; + if region.is_none() { + client.key = Some(job.key); + } + Ok(logits.map(|l| match region { + Some(r) => SegMask::from_logits_in(lightcraft_segment::MASK_SIDE, &l, r), + None => SegMask::from_logits(lightcraft_segment::MASK_SIDE, &l), + })) } fn run_job(state: &mut State, shared: &Shared, job: &mut Job) -> Result, String> { + if let Some(endpoint) = remote::configured() { + let result = remote_job(state, shared, job, &endpoint); + if result.is_err() { + state.remote = None; + } + return result.map_err(|e| format!("Remote SAM 3 failed: {e}. Check the worker and SSH tunnel, then retry.")); + } let side = lightcraft_segment::MASK_SIDE; match std::mem::replace(&mut job.kind, Kind::Prepare) { Kind::Prepare => { diff --git a/crates/engine/src/smart.rs b/crates/engine/src/smart.rs index 8deff7bd..a48fcd04 100644 --- a/crates/engine/src/smart.rs +++ b/crates/engine/src/smart.rs @@ -190,7 +190,7 @@ pub fn is_valid(path: &Path) -> bool { /// Load the proxy at `path`. pub fn load(path: &Path) -> Result { let b = std::fs::read(path).map_err(|e| format!("{}: {e}", path.display()))?; - decode(&b).map(|(image, camera_tone)| crate::media::DecodedSource { image: Arc::new(image), info: None, camera_tone }) + decode(&b).map(|(image, camera_tone)| crate::media::DecodedSource { image: Arc::new(image), info: None, camera_tone, denoised: None }) } #[cfg(test)] diff --git a/crates/engine/src/tests.rs b/crates/engine/src/tests.rs index 978e8067..9806a81e 100644 --- a/crates/engine/src/tests.rs +++ b/crates/engine/src/tests.rs @@ -118,6 +118,29 @@ fn albums_crud() { s.execute("album.delete", &json!({"id": f})).unwrap(); } +/// Given an album and a folder, when something is created inside each, then only the folder +/// accepts it: the sidebar only lists children of folders, so a child of a plain album would +/// be invisible. +#[test] +fn albums_and_smart_albums_are_created_only_inside_folders() { + let mut s = demo(); + let folder = s.execute("album.create", &json!({"name": "Trips", "folder": true})).unwrap()["id"].as_u64().unwrap(); + let album = s.execute("album.create", &json!({"name": "Best", "parent": folder})).unwrap()["id"].as_u64().unwrap(); + for (cmd, params) in [ + ("album.create", json!({"name": "Nested", "parent": album})), + ("album.create", json!({"name": "Nested", "folder": true, "parent": album})), + ("album.createSmart", json!({"name": "Nested", "parent": album})), + ("album.create", json!({"name": "Nested", "parent": 987_654})), + ] { + assert!(s.execute(cmd, ¶ms).is_err(), "{cmd} {params} must be refused"); + } + let n = s.catalog.albums().count(); + let sub = s.execute("album.create", &json!({"name": "Sub", "folder": true, "parent": folder})).unwrap()["id"].as_u64().unwrap(); + let smart = s.execute("album.createSmart", &json!({"name": "Rated", "parent": sub})).unwrap()["id"].as_u64().unwrap(); + assert_eq!(s.catalog.albums().count(), n + 2); + assert_eq!(s.catalog.album(lightcraft_catalog::AlbumId(smart)).unwrap().parent, Some(lightcraft_catalog::AlbumId(sub))); +} + #[test] fn filters_and_delete_restore() { let mut s = demo(); diff --git a/crates/engine/src/tests_album_order.rs b/crates/engine/src/tests_album_order.rs new file mode 100644 index 00000000..cff524cb --- /dev/null +++ b/crates/engine/src/tests_album_order.rs @@ -0,0 +1,159 @@ +//! Ordering albums by hand: `album.reorder` puts an album (or folder) before a sibling, or last of +//! its kind, in one undo step, also when it changes folder; `album.sort` goes back to by name. + +use lightcraft_catalog::AlbumId; +use serde_json::{Value, json}; + +use crate::Session; + +fn make(s: &mut Session, name: &str, parent: Option, folder: bool) -> u64 { + let mut p = json!({"name": name, "folder": folder}); + if let Some(parent) = parent { + p["parent"] = json!(parent); + } + s.execute("album.create", &p).unwrap()["id"].as_u64().unwrap() +} + +fn names(s: &Session, parent: Option) -> Vec { + s.catalog.album_children(parent.map(AlbumId)).iter().map(|a| a.name.clone()).collect() +} + +fn reorder(s: &mut Session, id: u64, params: Value) -> Result { + let mut p = params; + p["id"] = json!(id); + s.execute("album.reorder", &p).map_err(|e| e.to_string()) +} + +/// Given three albums listed by name, when one is placed before another, then the list follows; +/// "no sibling" puts it last (among albums). +#[test] +fn an_album_is_placed_before_a_sibling_or_last() { + let mut s = Session::new(); + let (a, b, c) = (make(&mut s, "A", None, false), make(&mut s, "B", None, false), make(&mut s, "C", None, false)); + reorder(&mut s, c, json!({"before": a})).unwrap(); + assert_eq!(names(&s, None), ["C", "A", "B"]); + assert!(s.catalog.album_children_are_ordered(None)); + reorder(&mut s, c, json!({"before": null})).unwrap(); + assert_eq!(names(&s, None), ["A", "B", "C"]); + reorder(&mut s, a, json!({})).unwrap(); + assert_eq!(names(&s, None), ["B", "C", "A"], "no `before`: last"); + assert_eq!(reorder(&mut s, b, json!({"before": b})).unwrap()["changed"], 0, "before itself is where it is"); + assert_eq!(names(&s, None), ["B", "C", "A"]); +} + +/// Folders stay ahead of albums: a folder is placed among folders, an album among albums. +#[test] +fn folders_and_albums_keep_to_their_own_group() { + let mut s = Session::new(); + let (f1, f2) = (make(&mut s, "F1", None, true), make(&mut s, "F2", None, true)); + let (a1, a2) = (make(&mut s, "A1", None, false), make(&mut s, "A2", None, false)); + reorder(&mut s, f2, json!({"before": f1})).unwrap(); + reorder(&mut s, a2, json!({"before": a1})).unwrap(); + assert_eq!(names(&s, None), ["F2", "F1", "A2", "A1"]); + reorder(&mut s, f1, json!({"before": a1})).expect_err("a folder is not placed among albums"); + reorder(&mut s, a1, json!({"before": f1})).expect_err("an album is not placed among folders"); + reorder(&mut s, f2, json!({"before": null})).unwrap(); + assert_eq!(names(&s, None), ["F1", "F2", "A2", "A1"], "last of the folders"); +} + +/// Placing an album in another folder moves it there, and one undo undoes both. +#[test] +fn placing_across_folders_moves_it_and_undoes_in_one_step() { + let mut s = Session::new(); + let folder = make(&mut s, "Trips", None, true); + let (x, y) = (make(&mut s, "X", Some(folder), false), make(&mut s, "Y", Some(folder), false)); + let loose = make(&mut s, "Loose", None, false); + reorder(&mut s, y, json!({"before": x})).unwrap(); + assert_eq!(names(&s, Some(folder)), ["Y", "X"]); + reorder(&mut s, loose, json!({"parent": folder, "before": x})).unwrap(); + assert_eq!(names(&s, Some(folder)), ["Y", "Loose", "X"]); + assert_eq!(s.catalog.album(AlbumId(loose)).unwrap().parent, Some(AlbumId(folder))); + s.execute("edit.undo", &json!({})).unwrap(); + assert_eq!(names(&s, Some(folder)), ["Y", "X"]); + assert_eq!(s.catalog.album(AlbumId(loose)).unwrap().parent, None); + assert_eq!(s.catalog.album(AlbumId(loose)).unwrap().order, None); +} + +/// Undo and redo of placing across folders, of sorting, and of moving a folder that has children. +#[test] +fn placing_and_sorting_undo_and_redo() { + let mut s = Session::new(); + let top = make(&mut s, "Top", None, true); + let child = make(&mut s, "Child", Some(top), false); + let (a, b) = (make(&mut s, "A", None, false), make(&mut s, "B", None, false)); + reorder(&mut s, b, json!({"before": a})).unwrap(); + // a folder with children, placed inside another folder + let other = make(&mut s, "Other", None, true); + reorder(&mut s, top, json!({"parent": other})).unwrap(); + assert_eq!(names(&s, Some(other)), ["Top"]); + assert_eq!(names(&s, Some(top)), ["Child"], "its children came along"); + s.execute("edit.undo", &json!({})).unwrap(); + assert_eq!(s.catalog.album(AlbumId(top)).unwrap().parent, None); + s.execute("edit.redo", &json!({})).unwrap(); + assert_eq!(s.catalog.album(AlbumId(top)).unwrap().parent, Some(AlbumId(other))); + // sort, undo, redo + s.execute("album.sort", &json!({})).unwrap(); + assert_eq!(names(&s, None), ["Other", "A", "B"]); + s.execute("edit.undo", &json!({})).unwrap(); + assert_eq!(names(&s, None), ["Other", "B", "A"], "the hand order is back"); + s.execute("edit.redo", &json!({})).unwrap(); + assert_eq!(names(&s, None), ["Other", "A", "B"]); + assert_eq!(s.catalog.album(AlbumId(child)).unwrap().parent, Some(AlbumId(top))); +} + +/// Bad requests change nothing: a `before` that is not in the target folder, a target that is not +/// a folder or lies inside the album, an unknown album. +#[test] +fn bad_placements_are_refused_and_change_nothing() { + let mut s = Session::new(); + let f = make(&mut s, "F", None, true); + let inner = make(&mut s, "Inner", Some(f), true); + let (a, b) = (make(&mut s, "A", None, false), make(&mut s, "B", Some(f), false)); + let before = s.catalog.to_snapshot(); + reorder(&mut s, a, json!({"before": b})).expect_err("B is in F, not at the top level"); + reorder(&mut s, a, json!({"parent": b})).expect_err("a plain album is no folder"); + reorder(&mut s, f, json!({"parent": inner})).expect_err("a folder into itself"); + reorder(&mut s, 987_654, json!({})).expect_err("no such album"); + reorder(&mut s, a, json!({"before": 987_654})).expect_err("no such sibling"); + assert_eq!(s.catalog.to_snapshot(), before, "nothing changed"); +} + +/// `album.sort` drops the hand order of one folder (and only that one): by name again. +#[test] +fn sorting_goes_back_to_by_name_for_one_folder() { + let mut s = Session::new(); + let f = make(&mut s, "F", None, true); + let (x, y) = (make(&mut s, "X", Some(f), false), make(&mut s, "Y", Some(f), false)); + let (a, b) = (make(&mut s, "A", None, false), make(&mut s, "B", None, false)); + reorder(&mut s, y, json!({"before": x})).unwrap(); + reorder(&mut s, b, json!({"before": a})).unwrap(); + assert_eq!(names(&s, Some(f)), ["Y", "X"]); + assert_eq!(names(&s, None), ["F", "B", "A"]); + let r = s.execute("album.sort", &json!({"parent": f})).unwrap(); + assert_eq!(r["changed"], 2); + assert_eq!(names(&s, Some(f)), ["X", "Y"]); + assert_eq!(names(&s, None), ["F", "B", "A"], "the top level keeps its order"); + assert!(!s.catalog.album_children_are_ordered(Some(AlbumId(f)))); + // nothing to do: no undo step is made + let r = s.execute("album.sort", &json!({"parent": f})).unwrap(); + assert_eq!(r["changed"], 0); + s.execute("edit.undo", &json!({})).unwrap(); + assert_eq!(names(&s, Some(f)), ["Y", "X"], "the one undo step was the sort"); + let r = s.execute("album.sort", &json!({})).unwrap(); + assert_eq!(r["changed"], 3, "every child of the top level had a place: F, B and A"); + assert_eq!(names(&s, None), ["F", "A", "B"]); +} + +/// Moving an album with `album.move` out of an ordered folder leaves the old neighbours' order. +#[test] +fn moving_out_of_an_ordered_folder_keeps_the_others_in_place() { + let mut s = Session::new(); + let f = make(&mut s, "F", None, true); + let (x, y, z) = (make(&mut s, "X", Some(f), false), make(&mut s, "Y", Some(f), false), make(&mut s, "Z", Some(f), false)); + reorder(&mut s, z, json!({"before": x})).unwrap(); + assert_eq!(names(&s, Some(f)), ["Z", "X", "Y"]); + s.execute("album.move", &json!({"id": x, "parent": null})).unwrap(); + assert_eq!(names(&s, Some(f)), ["Z", "Y"]); + assert_eq!(names(&s, None), ["F", "X"]); + let _ = y; +} diff --git a/crates/engine/src/tests_denoise.rs b/crates/engine/src/tests_denoise.rs new file mode 100644 index 00000000..ef87747f --- /dev/null +++ b/crates/engine/src/tests_denoise.rs @@ -0,0 +1,759 @@ +//! AI denoise through the engine: installing a model, making a photo's denoised picture in the background and for an +//! export, the Amount slider mixing it in, and the cache. A tiny stand-in model (it halves every sample) takes the +//! place of a real one, which is enough to see what the machinery does with its answer; how well real models denoise is +//! measured separately (`crates/denoise/tests/real_model.rs`). + +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use lightcraft_denoise::manifest::{DenoiserManifest, Domain, Gain}; +use lightcraft_denoise::run::TileRunner; +use serde_json::{Value, json}; + +use crate::Session; +use crate::denoise::{Loader, Model, PhotoState, RunOn}; +use crate::export::{ExportOptions, export_photo}; +use crate::tests_xmp::{synthetic_dng_of, temp_dir}; + +/// A model that makes every output pixel half of its cell's colour: a picture that is plainly not the plain one. +struct Dim { + runs: Arc, +} + +impl TileRunner for Dim { + fn run(&self, input: &[f32]) -> Result, lightcraft_denoise::Error> { + self.runs.fetch_add(1, Ordering::Relaxed); + let t = ((input.len() / 4) as f64).sqrt() as usize; + let side = 2 * t; + let mut out = vec![0f32; 3 * side * side]; + for y in 0..side { + for x in 0..side { + let i = (y / 2) * t + x / 2; + let (r, g, b) = (input[i], (input[t * t + i] + input[2 * t * t + i]) / 2.0, input[3 * t * t + i]); + out[y * side + x] = r * 0.5; + out[side * side + y * side + x] = g * 0.5; + out[2 * side * side + y * side + x] = b * 0.5; + } + } + Ok(out) + } +} + +impl Model for Dim { + fn runner(&self, _: RunOn, threads: usize) -> (&dyn TileRunner, usize) { + (self, threads) + } + + fn self_test(&self, _: RunOn) -> Result { + Ok(json!({"ok": true})) + } +} + +fn dim_loader(runs: &Arc) -> Loader { + let runs = runs.clone(); + Arc::new(move |_, _| Ok(Arc::new(Dim { runs: runs.clone() }) as Arc)) +} + +fn manifest(id: &str) -> DenoiserManifest { + DenoiserManifest { + id: id.into(), + name: "Stand-in".into(), + version: "1".into(), + licence: Default::default(), + source: None, + sha256: None, + size_bytes: None, + provenance: "made up for a test".into(), + domain: Domain::BayerToRgb, + tile: 64, + overlap: 16, + gain: Gain::None, + } +} + +/// The model's `.onnx` (not a real one: the stand-in loader never reads it) with its manifest beside it. +fn model_files(dir: &Path, id: &str) -> PathBuf { + let src = dir.join(format!("src-{id}")); + std::fs::create_dir_all(&src).unwrap(); + std::fs::write(src.join("denoise-model.json"), serde_json::to_vec(&manifest(id)).unwrap()).unwrap(); + std::fs::write(src.join("model.onnx"), format!("not an onnx file, only a stand-in for {id}")).unwrap(); + src.join("model.onnx") +} + +#[test] +fn model_installation_keeps_the_frame_pump_responsive_and_reuses_the_loaded_model() { + use std::sync::{Mutex, mpsc}; + use std::time::Duration; + let mut x = setup("background-install", false); + let path = model_files(&x.dir, "stand-in"); + let (entered_tx, entered_rx) = mpsc::channel(); + let (release_tx, release_rx) = mpsc::channel(); + let release_rx = Mutex::new(release_rx); + let runs = x.runs.clone(); + let loads = Arc::new(AtomicUsize::new(0)); + let worker_loads = loads.clone(); + x.s.denoise.loader = Arc::new(move |_, _| { + worker_loads.fetch_add(1, Ordering::Relaxed); + entered_tx.send(()).map_err(|e| e.to_string())?; + release_rx.lock().unwrap().recv_timeout(Duration::from_secs(3)).map_err(|e| e.to_string())?; + Ok(Arc::new(Dim { runs: runs.clone() })) + }); + let info = x.s.execute("denoise.models.inspect", &json!({"path": path})).unwrap(); + assert_eq!(info["model"]["id"], "stand-in"); + assert_eq!(loads.load(Ordering::Relaxed), 0, "showing terms never executes the model"); + let r = x.s.execute("denoise.models.install", &json!({"path": path, "acknowledged": true, "background": true})).unwrap(); + assert_eq!(r["started"], "local-install"); + entered_rx.recv_timeout(Duration::from_secs(3)).unwrap(); + // The loader is blocked: the frame pump must still return, reporting installation in progress. + let state = x.s.execute("denoise.models.downloads", &json!({})).unwrap(); + assert_eq!(state["downloads"][0]["state"], "installing"); + assert_eq!(x.s.execute("denoise.pump", &json!({"pace": "pause"})).unwrap()["active"], false); + assert!(x.s.execute("denoise.models.install", &json!({"path": path, "acknowledged": true, "background": true})).is_err()); + release_tx.send(()).unwrap(); + for _ in 0..600 { + let v = x.s.execute("denoise.models.downloads", &json!({})).unwrap(); + if v["downloads"][0]["state"] == "installed" { + break; + } + std::thread::sleep(Duration::from_millis(5)); + } + assert_eq!(x.s.execute("denoise.models.downloads", &json!({})).unwrap()["downloads"][0]["state"], "installed"); + let a = x.s.denoise.active.as_ref().unwrap(); + assert!(a.model.lock().unwrap().is_some(), "the worker's tested model is reused"); + assert_eq!(loads.load(Ordering::Relaxed), 1); + drop(x.s); + let _ = std::fs::remove_dir_all(x.dir); +} + +#[test] +fn failed_model_replacement_preserves_the_installed_model() { + let mut x = setup("replacement", true); + let installed = x.dir.join("denoise-models/stand-in/model.onnx"); + let original = std::fs::read(&installed).unwrap(); + let source = model_files(&x.dir, "stand-in"); + std::fs::write(&source, "a different broken model").unwrap(); + assert!(x.s.execute("denoise.models.install", &json!({"path": source, "acknowledged": true})).is_err()); + assert_eq!(std::fs::read(installed).unwrap(), original); + assert_eq!(x.s.execute("denoise.models.list", &json!({})).unwrap()["model"], "stand-in"); + drop(x.s); + let _ = std::fs::remove_dir_all(x.dir); +} + +/// A Bayer DNG that is a smooth ramp, big enough for a 64-cell tile and a second one. +fn ramp_dng(w: usize, h: usize, cfa: &str) -> Vec { + let data: Vec = (0..w * h).map(|i| 600 + ((i % w) * 30 + (i / w) * 35) as u16).collect(); + synthetic_dng_of(w, h, cfa, data, None, Default::default()) +} + +struct Setup { + dir: PathBuf, + s: Session, + runs: Arc, +} + +/// A session with a library on disk holding one raw photo, and the stand-in model installed and chosen. +fn setup(tag: &str, with_model: bool) -> Setup { + let dir = temp_dir(tag); + let originals = dir.join("originals"); + std::fs::create_dir_all(&originals).unwrap(); + std::fs::write(originals.join("ramp.dng"), ramp_dng(192, 160, "RGGB")).unwrap(); + let runs = Arc::new(AtomicUsize::new(0)); + let mut s = Session::new().with_fs(); + s.open_library(dir.join("library"), false).unwrap(); + s.execute("library.import", &json!({"paths": [originals.join("ramp.dng").to_string_lossy()]})).unwrap(); + let id = s.catalog.photos().next().unwrap().id; + s.execute("library.select", &json!({"ids": [id.0]})).unwrap(); + s.set_denoise_models_dir(Some(dir.join("denoise-models"))); + s.denoise.loader = dim_loader(&runs); + if with_model { + let onnx = model_files(&dir, "stand-in"); + s.execute("denoise.models.install", &json!({"path": onnx.to_string_lossy(), "acknowledged": true})).unwrap(); + } + Setup { dir, s, runs } +} + +fn photo(s: &Session) -> lightcraft_catalog::PhotoId { + s.catalog.photos().next().unwrap().id +} + +/// Pump until `n` photos are ready (the job runs on its own thread). +fn pump_until_ready(s: &mut Session, n: u64) -> Value { + for _ in 0..600 { + let r = s.execute("denoise.pump", &json!({"pace": "full"})).unwrap(); + if r["ready"].as_u64() >= Some(n) { + return r; + } + std::thread::sleep(std::time::Duration::from_millis(50)); + } + panic!("not ready: {}", s.execute("denoise.status", &json!({})).unwrap()); +} + +fn mean(img: &lightcraft_raster::Rgba8) -> f64 { + let sum: u64 = img.data.iter().map(|p| u64::from(p[0]) + u64::from(p[1]) + u64::from(p[2])).sum(); + sum as f64 / (img.data.len() * 3) as f64 +} + +fn rendered_mean(s: &mut Session, id: lightcraft_catalog::PhotoId) -> f64 { + mean(&s.render_now(id, 160, 160).unwrap().image) +} + +fn set_amount(s: &mut Session, amount: f64) { + s.execute("develop.set", &json!({"control": "enhance.denoise", "value": amount})).unwrap(); +} + +#[test] +fn a_model_is_installed_chosen_and_removed_with_the_terms_accepted() { + let mut x = setup("model", false); + let onnx = model_files(&x.dir, "stand-in"); + // the terms come first + let e = x.s.execute("denoise.models.install", &json!({"path": onnx.to_string_lossy()})).unwrap_err().to_string(); + assert!(e.contains("licence has not been accepted"), "{e}"); + let r = x.s.execute("denoise.models.list", &json!({})).unwrap(); + assert_eq!(r["model"], Value::Null); + assert_eq!(r["runtime"], cfg!(feature = "denoise")); + // a file with no manifest is not guessed at + let stray = x.dir.join("stray.onnx"); + std::fs::write(&stray, b"who knows").unwrap(); + let e = x.s.execute("denoise.models.install", &json!({"path": stray.to_string_lossy(), "acknowledged": true})).unwrap_err().to_string(); + assert!(e.contains("does not know this file"), "{e}"); + // installed, it is the model in use + let r = x.s.execute("denoise.models.install", &json!({"path": onnx.to_string_lossy(), "acknowledged": true})).unwrap(); + assert_eq!((r["installed"]["id"].clone(), r["model"].clone()), (json!("stand-in"), json!("stand-in")), "{r}"); + assert!(r["installed"]["accepted"]["sha256"].as_str().is_some_and(|h| h.len() == 64), "the file's hash is recorded: {r}"); + assert!(x.s.denoise.active.is_some()); + assert_eq!(x.s.execute("denoise.status", &json!({})).unwrap()["model"]["id"], "stand-in"); + // it can be checked again, switched off and on, and removed + assert_eq!(x.s.execute("denoise.models.test", &json!({"id": "stand-in"})).unwrap()["ok"], true); + assert_eq!(x.s.execute("denoise.models.select", &json!({"id": null})).unwrap()["model"], Value::Null); + assert!(x.s.denoise.active.is_none()); + assert!(x.s.execute("denoise.models.select", &json!({"id": "nothing"})).is_err()); + x.s.execute("denoise.models.select", &json!({"id": "stand-in"})).unwrap(); + assert!(x.s.denoise.active.is_some()); + let r = x.s.execute("denoise.models.remove", &json!({"id": "stand-in"})).unwrap(); + assert_eq!(r["model"], Value::Null); + assert!(x.s.denoise.active.is_none()); + assert!(!x.dir.join("denoise-models").join("stand-in").exists()); + // ids that are not folder names, and folders that are not models, are never touched + for bad in ["../library", "", "a/b", "settings.json"] { + assert!(x.s.execute("denoise.models.remove", &json!({"id": bad})).is_err(), "{bad}"); + } + assert!(x.dir.join("library").is_dir()); + // a model that fails its self-test is not left installed + x.s.denoise.loader = Arc::new(|_, _| Err("it does not load".into())); + let e = x.s.execute("denoise.models.install", &json!({"path": onnx.to_string_lossy(), "acknowledged": true})).unwrap_err().to_string(); + assert!(e.contains("it does not load"), "{e}"); + assert!(!x.dir.join("denoise-models").join("stand-in").exists()); + let _ = std::fs::remove_dir_all(&x.dir); +} + +#[test] +fn a_manifest_that_does_not_validate_or_name_the_file_is_refused() { + let mut x = setup("manifest", false); + let onnx = model_files(&x.dir, "stand-in"); + let json_path = onnx.with_file_name("denoise-model.json"); + // a tile the model cannot run on + let mut m = manifest("stand-in"); + m.tile = 50; + std::fs::write(&json_path, serde_json::to_vec(&m).unwrap()).unwrap(); + let e = x.s.execute("denoise.models.install", &json!({"path": onnx.to_string_lossy(), "acknowledged": true})).unwrap_err().to_string(); + assert!(e.contains("tile"), "{e}"); + // a hash that is not the file's + let mut m = manifest("stand-in"); + m.sha256 = Some("0".repeat(64)); + std::fs::write(&json_path, serde_json::to_vec(&m).unwrap()).unwrap(); + let e = x.s.execute("denoise.models.install", &json!({"path": onnx.to_string_lossy(), "acknowledged": true})).unwrap_err().to_string(); + assert!(e.contains("SHA-256"), "{e}"); + assert!(!x.dir.join("denoise-models").join("stand-in").join("model.onnx").exists(), "nothing half-installed is left"); + // an id that would leave the folder + let mut m = manifest("../escape"); + m.id = "../escape".into(); + std::fs::write(&json_path, serde_json::to_vec(&m).unwrap()).unwrap(); + assert!(x.s.execute("denoise.models.install", &json!({"path": onnx.to_string_lossy(), "acknowledged": true})).is_err()); + assert!(!x.dir.join("escape").exists()); + // garbage instead of a manifest + std::fs::write(&json_path, b"{ not json").unwrap(); + assert!(x.s.execute("denoise.models.install", &json!({"path": onnx.to_string_lossy(), "acknowledged": true})).is_err()); + let _ = std::fs::remove_dir_all(&x.dir); +} + +#[test] +fn a_photo_is_denoised_in_the_background_and_the_amount_mixes_it_in() { + let mut x = setup("background", true); + let id = photo(&x.s); + // nothing is made until a photo has an amount + let r = x.s.execute("denoise.pump", &json!({"pace": "full"})).unwrap(); + assert_eq!((r["queued"].clone(), r["running"].clone(), r["ready"].clone()), (json!(0), Value::Null, json!(0)), "{r}"); + assert_eq!(x.s.denoise_photo_state(id), PhotoState::Idle); + let plain = rendered_mean(&mut x.s, id); + let key_before = x.s.render_job(id, 160, 160, false, true).unwrap().key; + + set_amount(&mut x.s, 100.0); + // the photo being edited is picked up by itself, and made on another thread + let r = pump_until_ready(&mut x.s, 1); + assert_eq!(r["queued"], 0, "{r}"); + assert_eq!(x.s.denoise_photo_state(id), PhotoState::Ready); + assert!(x.runs.load(Ordering::Relaxed) >= 4, "a 192 × 160 mosaic is four 64-cell tiles: {}", x.runs.load(Ordering::Relaxed)); + let products: Vec<_> = std::fs::read_dir(x.dir.join("library").join("denoise")).unwrap().flatten().collect(); + assert_eq!(products.len(), 1, "one cached picture, no DNG or anything else in the library"); + assert!(products[0].path().extension().is_some_and(|e| e == "lcdn")); + let originals: Vec<_> = std::fs::read_dir(x.dir.join("originals")).unwrap().flatten().collect(); + assert_eq!(originals.len(), 1, "the original is untouched and alone"); + + // the render now comes from the denoised picture: darker at 100, in between at 50, the plain one at 0 + let full = rendered_mean(&mut x.s, id); + assert!(full < plain - 3.0, "denoised {full} vs plain {plain}"); + set_amount(&mut x.s, 50.0); + let half = rendered_mean(&mut x.s, id); + assert!(half < plain - 1.0 && half > full + 1.0, "{full} < {half} < {plain}"); + set_amount(&mut x.s, 0.0); + assert!((rendered_mean(&mut x.s, id) - plain).abs() < 0.01, "amount 0 is the plain picture again"); + // …and a render that uses it has another identity than one that does not + assert_ne!(key_before, { + set_amount(&mut x.s, 100.0); + x.s.render_job(id, 160, 160, false, true).unwrap().key + }); + // switching the Detail section off switches the denoise off with it + set_amount(&mut x.s, 100.0); + x.s.execute("develop.sectionEnabled", &json!({"section": "detail", "enabled": false})).unwrap(); + assert!((rendered_mean(&mut x.s, id) - plain).abs() < 0.01); + assert_eq!(x.s.media.denoise_salt(id, &x.s.catalog.photo(id).unwrap().develop), 0); + let _ = std::fs::remove_dir_all(&x.dir); +} + +#[test] +fn what_was_made_is_found_again_by_a_new_session_and_a_new_model_starts_over() { + let mut x = setup("again", true); + let id = photo(&x.s); + set_amount(&mut x.s, 100.0); + pump_until_ready(&mut x.s, 1); + let made = x.runs.load(Ordering::Relaxed); + x.s.close_library().unwrap(); + drop(x.s); + + // reopened: ready without running the model again + let mut s2 = Session::new().with_fs(); + s2.set_denoise_models_dir(Some(x.dir.join("denoise-models"))); + s2.denoise.loader = dim_loader(&x.runs); + s2.open_library(x.dir.join("library"), false).unwrap(); + s2.denoise_pump(crate::denoise::Pace::Full); + assert_eq!(s2.denoise_photo_state(id), PhotoState::Ready); + assert_eq!(x.runs.load(Ordering::Relaxed), made, "nothing was made again"); + + // another model's pictures are different pictures + let onnx = model_files(&x.dir, "other"); + s2.execute("denoise.models.install", &json!({"path": onnx.to_string_lossy(), "acknowledged": true})).unwrap(); + s2.denoise_pump(crate::denoise::Pace::Pause); + assert!( + matches!(s2.denoise_photo_state(id), PhotoState::Queued { .. }), + "the first model's picture is not this model's: it is to be made: {:?}", + s2.denoise_photo_state(id) + ); + // and choosing the first again finds its picture where it was + s2.execute("denoise.models.select", &json!({"id": "stand-in"})).unwrap(); + s2.denoise_pump(crate::denoise::Pace::Pause); + assert_eq!(s2.denoise_photo_state(id), PhotoState::Ready); + let _ = std::fs::remove_dir_all(&x.dir); +} + +#[test] +fn an_export_makes_its_picture_when_the_queue_has_not_and_never_comes_out_without_it() { + let mut x = setup("export", true); + let id = photo(&x.s); + let o = ExportOptions::from_json(&json!({"format": "png", "width": 160, "height": 160})); + let decode = |bytes: &[u8]| { + let d = lightcraft_codecs::decode(bytes, lightcraft_codecs::DecodeOptions::fit(160, 160)).unwrap(); + let w = d.to_working(); + w.data.iter().map(|p| f64::from(p[0] + p[1] + p[2])).sum::() / (w.data.len() * 3) as f64 + }; + let plain = decode(&export_photo(&mut x.s, id, &o, 1).unwrap().bytes); + assert_eq!(x.runs.load(Ordering::Relaxed), 0, "no amount, no denoise"); + set_amount(&mut x.s, 100.0); + // no pump: the export makes the picture itself + let denoised = decode(&export_photo(&mut x.s, id, &o, 1).unwrap().bytes); + assert!(x.runs.load(Ordering::Relaxed) >= 4); + assert!(denoised < plain * 0.9, "denoised {denoised} vs plain {plain}"); + let products = std::fs::read_dir(x.dir.join("library").join("denoise")).unwrap().flatten().count(); + assert_eq!(products, 1); + // a second export finds it + let runs = x.runs.load(Ordering::Relaxed); + let again = decode(&export_photo(&mut x.s, id, &o, 1).unwrap().bytes); + assert_eq!(x.runs.load(Ordering::Relaxed), runs); + assert!((again - denoised).abs() < 1e-6); + // the pump takes in what the export made + x.s.denoise_pump(crate::denoise::Pace::Pause); + assert_eq!(x.s.denoise_photo_state(id), PhotoState::Ready); + + // no model: the export says so instead of leaving the denoise out + x.s.execute("denoise.models.select", &json!({"id": null})).unwrap(); + let e = export_photo(&mut x.s, id, &o, 1).err().unwrap(); + assert!(e.contains("no denoise model is chosen") && e.contains("ramp.dng"), "{e}"); + set_amount(&mut x.s, 0.0); + assert!(export_photo(&mut x.s, id, &o, 1).is_ok(), "without an amount nothing is asked of the model"); + let _ = std::fs::remove_dir_all(&x.dir); +} + +#[test] +fn a_model_that_fails_fails_the_export_and_is_not_retried_by_the_queue() { + let mut x = setup("fails", true); + let id = photo(&x.s); + x.s.denoise.loader = Arc::new(|_, _| Err("the model went away".into())); + // Simulate loss of the in-memory model; successful installation now retains its tested instance. + *x.s.denoise.active.as_ref().unwrap().model.lock().unwrap() = None; + set_amount(&mut x.s, 80.0); + let o = ExportOptions::from_json(&json!({"format": "png", "width": 100, "height": 100})); + let e = export_photo(&mut x.s, id, &o, 1).err().unwrap(); + assert!(e.contains("ramp.dng") && e.contains("the model went away"), "{e}"); + // the background job fails too, says why once and does not try again every frame + for _ in 0..200 { + x.s.execute("denoise.pump", &json!({"pace": "full"})).unwrap(); + if matches!(x.s.denoise_photo_state(id), PhotoState::Failed { .. }) { + break; + } + std::thread::sleep(std::time::Duration::from_millis(20)); + } + match x.s.denoise_photo_state(id) { + PhotoState::Failed { why, unsupported } => assert!(why.contains("the model went away") && !unsupported, "{why}"), + other => panic!("{other:?}"), + } + let r = x.s.execute("denoise.pump", &json!({"pace": "full"})).unwrap(); + assert_eq!((r["queued"].clone(), r["running"].clone()), (json!(0), Value::Null), "{r}"); + // asking again by hand does retry + x.s.denoise.loader = dim_loader(&x.runs); + let r = x.s.execute("denoise.queue", &json!({"retry": true})).unwrap(); + assert_eq!(r["added"], 1, "{r}"); + pump_until_ready(&mut x.s, 1); + let _ = std::fs::remove_dir_all(&x.dir); +} + +#[test] +fn a_corrupt_cache_with_a_current_header_never_silently_exports_the_plain_photo() { + let mut x = setup("corrupt-export", true); + let id = photo(&x.s); + set_amount(&mut x.s, 100.0); + let options = ExportOptions::from_json(&json!({"format": "png", "width": 160, "height": 160})); + export_photo(&mut x.s, id, &options, 1).unwrap(); + x.s.denoise_pump(crate::denoise::Pace::Pause); + let spec = x.s.media.denoise.spec(id).unwrap().clone(); + let mut bytes = std::fs::read(&spec.product).unwrap(); + for byte in bytes.iter_mut().rev().take(64) { + *byte ^= 0xff; + } + std::fs::write(&spec.product, bytes).unwrap(); + assert!(lightcraft_denoise::product::is_current(&spec.product, &spec.key), "header/length still look current"); + x.s.media.forget(id); + let error = export_photo(&mut x.s, id, &options, 1).err().unwrap(); + assert!(error.contains("AI Denoise picture could not be read"), "{error}"); + let _ = std::fs::remove_dir_all(x.dir); +} + +#[test] +fn only_raw_photos_on_disk_are_denoised_and_nothing_is_queued_without_a_model() { + let dir = temp_dir("eligible"); + let mut demo = Session::with_demo(); + demo.set_denoise_models_dir(Some(dir.join("models"))); + let id = demo.active().unwrap(); + assert_eq!(demo.denoise_photo_state(id), PhotoState::NotApplicable, "a generated photo has no sensor data"); + let e = demo.execute("denoise.queue", &json!({})).unwrap_err().to_string(); + assert!(e.contains("no denoise model is chosen"), "{e}"); + // no models folder at all (the web build): the commands say so, nothing panics + let mut web = Session::with_demo(); + for cmd in ["denoise.models.install", "denoise.models.select", "denoise.models.remove", "denoise.settings"] { + assert!(web.execute(cmd, &json!({"path": "x", "acknowledged": true, "id": "x"})).is_err(), "{cmd}"); + } + assert_eq!(web.execute("denoise.pump", &json!({})).unwrap()["active"], false); + assert_eq!(web.execute("denoise.status", &json!({})).unwrap()["enabled"], false); + // no library open: nowhere to keep pictures, and the queue says so + let x = setup("eligible2", true); + let mut bare = Session::new().with_fs(); + bare.set_denoise_models_dir(Some(x.dir.join("denoise-models"))); + assert!(bare.denoise_products_dir().is_none()); + let e = bare.execute("denoise.queue", &json!({})).unwrap_err().to_string(); + assert!(e.contains("open a library on disk first"), "{e}"); + assert_eq!(bare.execute("denoise.pump", &json!({})).unwrap()["queued"], 0); + let _ = std::fs::remove_dir_all(&dir); + let _ = std::fs::remove_dir_all(&x.dir); +} + +#[test] +fn files_that_cannot_be_denoised_fail_with_a_reason_and_never_a_panic() { + let dir = temp_dir("makefails"); + let spec = |path: &Path| crate::denoise::JobSpec { + path: path.to_string_lossy().into_owned(), + read: None, + product: dir.join("p.lcdn"), + key: "k".into(), + manifest: manifest("stand-in"), + onnx: dir.join("model.onnx"), + model: Default::default(), + loader: dim_loader(&Arc::new(AtomicUsize::new(0))), + parallel: 2, + run_on: RunOn::Auto, + }; + use crate::denoise::{MakeError, make_product}; + // not a file + assert!(matches!(make_product(&spec(&dir.join("missing.dng")), None), Err(MakeError::Failed(_)))); + // not a raw + let junk = dir.join("junk.dng"); + std::fs::write(&junk, b"this is not a raw file at all").unwrap(); + assert!(matches!(make_product(&spec(&junk), None), Err(MakeError::Unsupported(_)))); + // a raw cut short + let whole = ramp_dng(192, 160, "RGGB"); + let cut = dir.join("cut.dng"); + std::fs::write(&cut, &whole[..whole.len() / 3]).unwrap(); + assert!(make_product(&spec(&cut), None).is_err()); + // a model that gives the wrong number of values is an error from the model, not a crash + struct Wrong; + impl TileRunner for Wrong { + fn run(&self, _: &[f32]) -> Result, lightcraft_denoise::Error> { + Ok(vec![0.0; 7]) + } + } + impl Model for Wrong { + fn runner(&self, _: RunOn, threads: usize) -> (&dyn TileRunner, usize) { + (self, threads) + } + fn self_test(&self, _: RunOn) -> Result { + Err("no".into()) + } + } + let good = dir.join("good.dng"); + std::fs::write(&good, &whole).unwrap(); + let mut s = spec(&good); + s.loader = Arc::new(|_, _| Ok(Arc::new(Wrong) as Arc)); + match make_product(&s, None) { + Err(MakeError::Failed(why)) => assert!(why.contains("7 numbers"), "{why}"), + other => panic!("{other:?}"), + } + assert!(!dir.join("p.lcdn").exists(), "a failed job leaves no product"); + // a job told to stop stops before doing anything + let progress = crate::denoise::Progress::default(); + progress.cancel_now(); + assert_eq!(make_product(&spec(&good), Some(&progress)), Err(MakeError::Cancelled)); + // a Bayer layout of the other kinds is denoised too, and a product that exists is not made again + for cfa in ["BGGR", "GRBG", "GBRG"] { + let f = dir.join(format!("{cfa}.dng")); + std::fs::write(&f, ramp_dng(192, 160, cfa)).unwrap(); + let mut s = spec(&f); + s.product = dir.join(format!("{cfa}.lcdn")); + assert_eq!(make_product(&s, None), Ok(true), "{cfa}"); + assert_eq!(make_product(&s, None), Ok(false), "{cfa} is already there"); + } + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn two_threads_making_the_same_picture_make_it_once() { + let dir = temp_dir("twice"); + let f = dir.join("a.dng"); + std::fs::write(&f, ramp_dng(192, 160, "RGGB")).unwrap(); + let runs = Arc::new(AtomicUsize::new(0)); + let spec = crate::denoise::JobSpec { + path: f.to_string_lossy().into_owned(), + read: None, + product: dir.join("p.lcdn"), + key: "k".into(), + manifest: manifest("stand-in"), + onnx: dir.join("model.onnx"), + model: Default::default(), + loader: dim_loader(&runs), + parallel: 2, + run_on: RunOn::Auto, + }; + let made: Vec = std::thread::scope(|sc| { + let hs: Vec<_> = (0..4).map(|_| sc.spawn(|| crate::denoise::make_product(&spec, None).unwrap())).collect(); + hs.into_iter().map(|h| h.join().unwrap()).collect() + }); + assert_eq!(made.iter().filter(|m| **m).count(), 1, "{made:?}"); + assert_eq!(runs.load(Ordering::Relaxed), 4, "one photo's four tiles, once"); + assert!(lightcraft_denoise::product::is_current(&spec.product, "k")); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn the_cache_is_kept_within_its_limit_and_what_photos_use_goes_last() { + let mut x = setup("evict", true); + let id = photo(&x.s); + set_amount(&mut x.s, 100.0); + pump_until_ready(&mut x.s, 1); + let used = x.dir.join("library").join("denoise"); + // two older pictures nobody uses + for (i, name) in ["aaaa", "bbbb"].iter().enumerate() { + let f = used.join(format!("{name}.lcdn")); + std::fs::write(&f, vec![7u8; 1000]).unwrap(); + let t = std::time::SystemTime::now() - std::time::Duration::from_secs(1000 - i as u64); + std::fs::File::options().write(true).open(&f).unwrap().set_modified(t).unwrap(); + } + let size_of = |p: &Path| std::fs::metadata(p).unwrap().len(); + let mine: PathBuf = std::fs::read_dir(&used).unwrap().flatten().map(|e| e.path()).find(|p| p.file_stem().is_some_and(|s| s.len() == 32)).unwrap(); + // room for the photo's own picture and one stray: the oldest stray goes, then the other + x.s.denoise_evict_to(size_of(&mine) + 1500, None); + assert!(!used.join("aaaa.lcdn").exists() && used.join("bbbb.lcdn").exists() && mine.exists()); + x.s.denoise_evict_to(size_of(&mine) + 10, None); + assert!(!used.join("bbbb.lcdn").exists() && mine.exists()); + assert_eq!(x.s.denoise_photo_state(id), PhotoState::Ready); + // too small for even that: the photo's own picture is the last thing to go, and the photo knows + x.s.denoise_evict_to(10, None); + assert!(!mine.exists()); + assert_ne!(x.s.denoise_photo_state(id), PhotoState::Ready); + // …and it is made again, because the photo still has its amount + pump_until_ready(&mut x.s, 1); + assert!(mine.exists()); + // clearing deletes it all + let r = x.s.execute("denoise.clear", &json!({})).unwrap(); + assert_eq!(r["deleted"], 1, "{r}"); + assert_eq!(std::fs::read_dir(&used).unwrap().flatten().count(), 0); + let _ = std::fs::remove_dir_all(&x.dir); +} + +#[test] +fn settings_are_checked_saved_and_a_paused_pump_starts_nothing() { + let mut x = setup("settings", true); + let r = x.s.execute("denoise.settings", &json!({"auto": false, "cacheGb": 5, "threads": 3})).unwrap(); + assert_eq!((r["auto"].clone(), r["cacheGb"].clone(), r["threads"].clone()), (json!(false), json!(5), json!(3)), "{r}"); + for bad in [ + json!({"cacheGb": 0}), + json!({"cacheGb": "big"}), + json!({"threads": 0}), + json!({"threads": 1000}), + json!({"auto": "yes"}), + json!({"runOn": "fastest"}), + json!({"runOn": true}), + ] { + assert!(x.s.execute("denoise.settings", &bad).is_err(), "{bad}"); + } + assert_eq!(x.s.execute("denoise.models.list", &json!({})).unwrap()["cacheGb"], 5, "a refused change changes nothing"); + // the faster of the card and the processor unless the user says otherwise, and the interface can read the choice back + assert_eq!(x.s.execute("denoise.models.list", &json!({})).unwrap()["runOn"], "auto"); + assert_eq!(x.s.execute("denoise.settings", &json!({"runOn": "cpu"})).unwrap()["runOn"], "cpu"); + assert_eq!(x.s.execute("denoise.models.list", &json!({})).unwrap()["runOn"], "cpu"); + let status = x.s.execute("denoise.status", &json!({})).unwrap(); + assert_eq!((status["runOn"].clone(), status["device"]["kind"].clone()), (json!("cpu"), json!("none")), "no model is loaded yet: {status}"); + // the setting before `runOn` still reads as what it meant + let models = x.s.denoise.models_dir.clone().unwrap(); + let mut st = crate::denoise::read_settings(&models); + (st.run_on, st.gpu) = (None, Some(false)); + crate::denoise::write_settings(&models, &st).unwrap(); + assert_eq!(x.s.execute("denoise.models.list", &json!({})).unwrap()["runOn"], "cpu"); + // choosing where it runs lets the card be tried again after a set-up that closed LightCraft + let onnx = crate::denoise::installed_models(&models).first().unwrap().onnx.clone(); + let marker = onnx.with_file_name(crate::denoise::GPU_SETUP_MARKER); + std::fs::write(&marker, "setting up").unwrap(); + assert_eq!(x.s.execute("denoise.settings", &json!({"runOn": "auto"})).unwrap()["runOn"], "auto"); + assert!(!marker.exists(), "the crash is forgotten"); + assert_eq!(crate::denoise::read_settings(&models).gpu, None, "the old setting is not kept beside the new one"); + // with auto off nothing is started by itself, and the photo is made when asked + let id = photo(&x.s); + set_amount(&mut x.s, 100.0); + for _ in 0..3 { + x.s.execute("denoise.pump", &json!({"pace": "full"})).unwrap(); + } + assert_eq!(x.s.denoise_photo_state(id), PhotoState::Idle); + x.s.execute("denoise.settings", &json!({"auto": true})).unwrap(); + // paused: it is queued but nothing starts + for _ in 0..3 { + let r = x.s.execute("denoise.pump", &json!({"pace": "pause"})).unwrap(); + assert_eq!(r["running"], Value::Null); + } + assert!(matches!(x.s.denoise_photo_state(id), PhotoState::Queued { .. }), "{:?}", x.s.denoise_photo_state(id)); + // …and cancelling forgets it + assert_eq!(x.s.execute("denoise.cancel", &json!({})).unwrap()["cancelled"], 1); + assert_eq!(x.s.denoise_photo_state(id), PhotoState::Idle); + // when the pace allows, it is made + x.s.execute("denoise.queue", &json!({"scope": "withAmount"})).unwrap(); + pump_until_ready(&mut x.s, 1); + let _ = std::fs::remove_dir_all(&x.dir); +} + +#[test] +fn a_photo_whose_file_changed_has_no_picture_until_it_is_made_again() { + let mut x = setup("changed", true); + let id = photo(&x.s); + set_amount(&mut x.s, 100.0); + pump_until_ready(&mut x.s, 1); + // the photo's content is something else now (a replaced file): its old picture is not its picture + let p = x.s.catalog.photo(id).unwrap().clone(); + x.s.commit( + "test", + lightcraft_catalog::Op::SetContent { + id, + width: p.width, + height: p.height, + file_size: p.file_size + 1, + content_hash: Some("ffffffffffffffffffffffffffffffff".into()), + preview_only: None, + }, + ) + .unwrap(); + x.s.denoise_pump(crate::denoise::Pace::Pause); + assert_ne!(x.s.denoise_photo_state(id), PhotoState::Ready); + // the old picture is not used by a render either + assert_eq!(x.s.media.denoise_salt(id, &x.s.catalog.photo(id).unwrap().develop), 0); + let _ = std::fs::remove_dir_all(&x.dir); +} + +#[test] +fn a_photo_waits_for_its_turn_unless_pictures_are_not_made_on_their_own_or_have_nowhere_to_go() { + let mut x = setup("waits", true); + let id = photo(&x.s); + set_amount(&mut x.s, 80.0); + // automatic off: the open photo is left alone until it is asked for + x.s.execute("denoise.settings", &json!({"auto": false})).unwrap(); + assert!(!x.s.denoise_auto()); + x.s.execute("denoise.pump", &json!({"pace": "full"})).unwrap(); + assert_eq!(x.s.denoise_photo_state(id), PhotoState::Idle); + assert!(!x.s.denoise_busy()); + x.s.execute("denoise.queue", &json!({"ids": [id.0]})).unwrap(); + assert!(x.s.denoise_busy(), "queued with a model: the headless runs wait for it"); + pump_until_ready(&mut x.s, 1); + assert_eq!(x.s.denoise_photo_state(id), PhotoState::Ready); + assert!(!x.s.denoise_busy()); + // a session whose library is not on disk has nowhere to keep it: the photo says so instead of waiting for ever + x.s.denoise.settings.auto = Some(true); + if let Some(l) = x.s.library.as_mut() { + l.on_disk = false; + } + match x.s.denoise_photo_state(id) { + PhotoState::Failed { why, unsupported } => assert!(unsupported && why.contains("library folder"), "{why}"), + other => panic!("{other:?}"), + } + let _ = std::fs::remove_dir_all(&x.dir); +} + +#[test] +fn the_denoise_switch_validates_arguments_and_accepts_explicit_targets_without_selection() { + let mut x = setup("toggle", false); + let id = photo(&x.s); + for params in [json!({"id": "bad"}), json!({"id": u64::MAX}), json!({"id": id.0, "enabled": "yes"})] { + assert!(x.s.execute("denoise.toggle", ¶ms).is_err()); + } + assert!(x.s.execute("denoise.toggle", &json!({"id": id.0, "enabled": true})).is_err()); + assert_eq!(x.s.develop_of(id).unwrap().enhance.denoise, 0.0); + let model = model_files(&x.dir, "toggle-model"); + x.s.execute("denoise.models.install", &json!({"path": model.to_string_lossy(), "acknowledged": true})).unwrap(); + x.s.selection = Default::default(); + let undo = x.s.undo.len(); + let r = x.s.execute("denoise.toggle", &json!({"id": id.0, "enabled": true})).unwrap(); + assert_eq!(r["amount"], 50.0); + assert_eq!(x.s.undo.len(), undo + 1); + x.s.execute("denoise.toggle", &json!({"id": id.0, "enabled": false})).unwrap(); + assert_eq!(x.s.develop_of(id).unwrap().denoise_amount(), 0.0); + // the switch is its own setting: flipping it keeps the Amount (so it compares with and without), and an Amount of 0 + // does not switch it off + x.s.execute("library.select", &json!({"ids": [id.0]})).unwrap(); + x.s.execute("denoise.toggle", &json!({"id": id.0, "enabled": true})).unwrap(); + x.s.execute("develop.set", &json!({"control": "enhance.denoise", "value": 80.0})).unwrap(); + let off = x.s.execute("denoise.toggle", &json!({"id": id.0, "enabled": false})).unwrap(); + assert_eq!(off["amount"], 80.0); + let d = x.s.develop_of(id).unwrap(); + assert!(!d.enhance.denoise_enabled() && d.denoise_amount() == 0.0); + assert_eq!(x.s.execute("denoise.toggle", &json!({"id": id.0, "enabled": true})).unwrap()["amount"], 80.0); + assert!((x.s.develop_of(id).unwrap().denoise_amount() - 0.8).abs() < 1e-6); + x.s.execute("develop.set", &json!({"control": "enhance.denoise", "value": 0.0})).unwrap(); + assert!(x.s.develop_of(id).unwrap().enhance.denoise_enabled(), "sliding to 0 leaves it on"); + assert_eq!(x.s.execute("denoise.toggle", &json!({"id": id.0, "enabled": true})).unwrap()["amount"], 0.0); + // with no `enabled`, the toggle flips what the switch says + assert_eq!(x.s.execute("denoise.toggle", &json!({"id": id.0})).unwrap()["enabled"], false); + drop(x.s); + let _ = std::fs::remove_dir_all(x.dir); +} diff --git a/crates/engine/src/tests_face_models.rs b/crates/engine/src/tests_face_models.rs new file mode 100644 index 00000000..c71f7afe --- /dev/null +++ b/crates/engine/src/tests_face_models.rs @@ -0,0 +1,586 @@ +//! The `faces.*` model commands: list, inspect, install behind the licence, select, remove, the on/off +//! setting, and the refusals (no folder, licence not accepted, unsupported or hostile files and ids). + +use std::path::PathBuf; + +use serde_json::{Value, json}; + +use crate::Session; + +fn temp(name: &str) -> PathBuf { + let d = std::env::temp_dir().join(format!("lc-facemodels-{name}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&d); + std::fs::create_dir_all(&d).unwrap(); + d +} + +fn session(dir: &std::path::Path) -> Session { + let mut s = Session::new(); + s.face_models_dir = Some(dir.join("models")); + s +} + +fn model_file(dir: &std::path::Path, name: &str, dim: u64) -> String { + let p = dir.join(name); + std::fs::write(&p, lightcraft_faces::synthetic::tiny_embedder_model(dim)).unwrap(); + p.to_string_lossy().into_owned() +} + +fn find<'a>(list: &'a Value, id: &str) -> &'a Value { + list["models"].as_array().unwrap().iter().find(|m| m["id"] == id).unwrap_or(&Value::Null) +} + +#[test] +fn the_list_knows_the_known_models_and_is_honest_about_the_runtime() { + let d = temp("list"); + let mut s = session(&d); + let l = s.execute("faces.models.list", &json!({})).unwrap(); + assert_eq!(l["enabled"], false); + assert_eq!(l["runtime"], cfg!(not(target_arch = "wasm32"))); + assert_eq!(find(&l, "yunet-2023mar")["installed"], false, "nothing ships with LightCraft"); + assert_eq!(find(&l, "auraface-v1")["installed"], false); + assert_eq!(find(&l, "auraface-v1")["licence"]["commercial"], "yes"); + assert_eq!(find(&l, "sface-2021dec")["licence"]["commercial"], "unknown"); + // speed is a ratio to a ResNet-100 model, in words, for models that are not installed yet too; never milliseconds + assert_eq!(find(&l, "sface-2021dec")["speedText"], "6.9× faster than a ResNet-100 model"); + assert_eq!(find(&l, "auraface-v1")["speedText"], "Same speed as a ResNet-100 model"); + assert_eq!(find(&l, "yunet-2023mar")["speedText"], Value::Null, "nobody timed the detector against it"); + // a build with no folder lists too, and cannot install + let mut web = Session::new(); + assert!(web.execute("faces.models.list", &json!({})).is_ok()); + let f = model_file(&d, "m.onnx", 512); + assert!(web.execute("faces.models.install", &json!({"path": f, "acknowledged": true})).is_err()); + assert!(web.execute("faces.enable", &json!({"enabled": true})).is_err()); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn inspect_describes_a_file_and_installs_nothing() { + let d = temp("inspect"); + let mut s = session(&d); + let f = model_file(&d, "My Model.onnx", 512); + let r = s.execute("faces.models.inspect", &json!({"path": f})).unwrap(); + assert_eq!(r["kind"], "draft"); + assert_eq!(r["model"]["licence"]["commercial"], "unknown"); + assert_eq!(r["model"]["known"], false); + assert!(r["assumptions"].as_array().unwrap().iter().any(|a| a.as_str().unwrap().contains("127.5"))); + assert_eq!(r["alreadyInstalled"], false); + assert!(!d.join("models").exists(), "inspecting writes nothing"); + // not a model: a plain answer, not an error + let junk = d.join("junk.onnx"); + std::fs::write(&junk, b"this is not a model").unwrap(); + let r = s.execute("faces.models.inspect", &json!({"path": junk.to_string_lossy()})).unwrap(); + assert_eq!(r["kind"], "unsupported"); + assert!(r["reason"].as_str().unwrap().len() > 5); + // missing and empty files and folders are errors + assert!(s.execute("faces.models.inspect", &json!({"path": d.join("nope.onnx").to_string_lossy()})).is_err()); + assert!(s.execute("faces.models.inspect", &json!({"path": d.to_string_lossy()})).is_err()); + let empty = d.join("empty.onnx"); + std::fs::write(&empty, b"").unwrap(); + assert!(s.execute("faces.models.inspect", &json!({"path": empty.to_string_lossy()})).is_err()); + assert!(s.execute("faces.models.inspect", &json!({})).is_err()); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn install_needs_the_licence_accepted_then_select_enable_and_remove_work() { + let d = temp("flow"); + let mut s = session(&d); + let f = model_file(&d, "Mine.onnx", 512); + assert!(s.execute("faces.models.install", &json!({"path": f})).is_err(), "no acknowledgement"); + assert!(s.execute("faces.models.install", &json!({"path": f, "acknowledged": false})).is_err()); + assert!(s.execute("faces.models.install", &json!({"path": f, "acknowledged": "yes"})).is_err(), "must be the boolean true"); + assert!(!d.join("models").exists(), "refused installs leave nothing behind"); + + let r = s.execute("faces.models.install", &json!({"path": f, "acknowledged": true})).unwrap(); + let id = r["installed"]["id"].as_str().unwrap().to_string(); + assert!(id.starts_with("custom-mine-")); + let home = d.join("models").join(&id); + assert!(home.join("model.onnx").is_file() && home.join("face-model.json").is_file() && home.join("installed.json").is_file()); + assert!(!home.join("model.onnx.part").exists()); + assert_eq!(std::fs::read(home.join("model.onnx")).unwrap(), std::fs::read(&f).unwrap()); + let rec: Value = serde_json::from_slice(&std::fs::read(home.join("installed.json")).unwrap()).unwrap(); + assert_eq!(rec["commercial"], "unknown"); + assert_eq!(rec["fileName"], "Mine.onnx"); + + // listed as installed and custom; installing again is harmless + let l = s.execute("faces.models.list", &json!({})).unwrap(); + assert_eq!(find(&l, &id)["installed"], true); + assert_eq!(find(&l, &id)["known"], false); + s.execute("faces.models.install", &json!({"path": f, "acknowledged": true})).unwrap(); + assert_eq!(s.execute("faces.models.inspect", &json!({"path": f})).unwrap()["alreadyInstalled"], true); + + // choose it; the choice and the on/off setting survive a new session on the same folder + assert!(s.execute("faces.models.select", &json!({"id": "auraface-v1"})).is_err(), "not installed"); + assert!(s.execute("faces.models.select", &json!({"id": "yunet-2023mar"})).is_err(), "a detector is not a recogniser"); + s.execute("faces.models.select", &json!({"id": id})).unwrap(); + assert_eq!(s.execute("faces.enable", &json!({"enabled": true})).unwrap()["enabled"], true); + let mut again = session(&d); + let l = again.execute("faces.models.list", &json!({})).unwrap(); + assert_eq!((l["enabled"].clone(), l["embedder"].clone()), (json!(true), json!(id))); + assert_eq!(again.execute("faces.enable", &json!({})).unwrap()["enabled"], false, "toggles when omitted"); + + // removing it clears the choice + assert_eq!(s.execute("faces.models.remove", &json!({"id": id})).unwrap()["removed"], id); + assert!(!home.exists()); + assert_eq!(s.execute("faces.models.list", &json!({})).unwrap()["embedder"], Value::Null); + assert!(s.execute("faces.models.remove", &json!({"id": id})).is_err(), "already gone"); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn hostile_ids_unsupported_models_and_odd_folders_are_handled() { + let d = temp("hostile"); + let mut s = session(&d); + let models = d.join("models"); + std::fs::create_dir_all(&models).unwrap(); + // something that is not a model sits in the folder: it is never listed and never deleted + let precious = models.join("precious"); + std::fs::create_dir_all(&precious).unwrap(); + std::fs::write(precious.join("keep.txt"), b"mine").unwrap(); + for id in ["precious", "..", ".", "", "../x", "a/b", r"a\b", ".hidden", "yunet-2023mar", "UPPER"] { + assert!(s.execute("faces.models.remove", &json!({"id": id})).is_err(), "{id:?}"); + } + assert!(precious.join("keep.txt").is_file()); + assert!(d.join("models").is_dir()); + // a folder whose manifest lies about its id is ignored + let liar = models.join("liar"); + std::fs::create_dir_all(&liar).unwrap(); + std::fs::write(liar.join("model.onnx"), lightcraft_faces::synthetic::embedder_model(64)).unwrap(); + let mut m = lightcraft_faces::known::auraface(); + m.id = "someone-else".into(); + std::fs::write(liar.join("face-model.json"), serde_json::to_vec(&m).unwrap()).unwrap(); + std::fs::write(models.join("settings.json"), b"{ not json").unwrap(); + let l = s.execute("faces.models.list", &json!({})).unwrap(); + assert!(find(&l, "someone-else").is_null() && find(&l, "liar").is_null()); + assert_eq!(l["enabled"], false, "a corrupt settings file is ignored"); + // not-a-model and wrong-shaped files cannot be installed + let junk = d.join("junk.onnx"); + std::fs::write(&junk, b"nonsense").unwrap(); + let e = s.execute("faces.models.install", &json!({"path": junk.to_string_lossy(), "acknowledged": true})).unwrap_err().to_string(); + assert!(e.contains("cannot be used yet"), "{e}"); + let tokens = d.join("tokens.onnx"); + { + use lightcraft_faces::synthetic::{len_field, value_info_bytes}; + let mut graph = Vec::new(); + len_field(11, &value_info_bytes("pixels", 1, &[Err("N"), Ok(3), Ok(224), Ok(224)]), &mut graph); + len_field(12, &value_info_bytes("tokens", 1, &[Err("N"), Ok(257), Ok(384)]), &mut graph); + let mut model = Vec::new(); + len_field(7, &graph, &mut model); + std::fs::write(&tokens, model).unwrap(); + } + assert!(s.execute("faces.models.install", &json!({"path": tokens.to_string_lossy(), "acknowledged": true})).is_err()); + assert!(s.execute("faces.models.select", &json!({"id": 5})).is_err()); + assert!(s.execute("faces.models.select", &json!({"id": ""})).is_err()); + let _ = std::fs::remove_dir_all(&d); +} + +/// A demo session with the real YuNet installed, for the tests that need it: `LC_YUNET_MODEL=` (what +/// Settings > Faces downloads; nothing is committed). Without it those tests say so and pass. +fn demo_with_yunet() -> Option { + let Some(file) = std::env::var_os("LC_YUNET_MODEL") else { + eprintln!("LC_YUNET_MODEL is not set: skipping"); + return None; + }; + let dir = temp("yunet"); + let mut s = Session::with_demo(); + s.face_models_dir = Some(dir.join("models")); + let r = s.execute("faces.models.install", &json!({"path": file.to_string_lossy(), "acknowledged": true})); + assert!(r.is_ok(), "{r:?}"); + Some(s) +} + +#[test] +fn installing_the_detector_does_not_choose_it_as_the_recogniser_or_switch_recognition_on() { + let Some(mut s) = demo_with_yunet() else { return }; + let l = s.execute("faces.models.list", &json!({})).unwrap(); + assert_eq!(find(&l, "yunet-2023mar")["installed"], true); + assert_eq!(l["enabled"], false, "{l}"); + assert!(l["embedder"].is_null(), "{l}"); + assert!(s.execute("faces.models.test", &json!({"id": "yunet-2023mar"})).is_err(), "a detector has no recognition self-test"); +} + +#[test] +fn detect_without_the_model_says_where_to_get_it() { + let d = temp("nodetector"); + let mut s = Session::with_demo(); + // no folder for models at all, then a folder that has none yet + assert!(s.execute("faces.detect", &json!({})).is_err()); + s.face_models_dir = Some(d.join("models")); + let e = s.execute("faces.detect", &json!({})).unwrap_err().to_string(); + assert!(e.contains("Settings > Faces"), "{e}"); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn the_detector_and_the_recognisers_can_be_downloaded_and_the_download_is_refused_otherwise() { + let d = temp("dl-refusals"); + let mut s = session(&d); + let l = s.execute("faces.models.list", &json!({})).unwrap(); + for (id, host) in [("yunet-2023mar", "github.com"), ("sface-2021dec", "github.com"), ("auraface-v1", "huggingface.co")] { + assert_eq!(find(&l, id)["downloadHost"], host, "{id}"); + } + // without the user's acceptance, for an id nobody knows, and without a folder: nothing starts + assert!(s.execute("faces.models.download", &json!({"id": "yunet-2023mar"})).is_err()); + assert!(s.execute("faces.models.download", &json!({"id": "yunet-2023mar", "acknowledged": false})).is_err()); + for id in ["nope", "", "../yunet-2023mar"] { + assert!(s.execute("faces.models.download", &json!({"id": id, "acknowledged": true})).is_err(), "{id}"); + } + assert!(Session::new().execute("faces.models.download", &json!({"id": "yunet-2023mar", "acknowledged": true})).is_err()); + assert_eq!(s.execute("faces.models.downloads", &json!({})).unwrap()["downloads"], json!([])); + assert!(!d.join("models").join(".downloads").exists(), "a refused download touches nothing"); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn a_download_that_has_arrived_is_installed_by_itself_and_its_staged_file_goes() { + let d = temp("dl-arrived"); + let mut s = session(&d); + let staging = d.join("models").join(".downloads"); + std::fs::create_dir_all(&staging).unwrap(); + let staged = staging.join("model.onnx"); + std::fs::write(&staged, lightcraft_faces::synthetic::tiny_embedder_model(512)).unwrap(); + let sha = lightcraft_faces::hash::sha256_file(&staged).unwrap(); + s.face_downloads.arrived("my-model", staged.clone(), &sha); + + let r = s.execute("faces.models.downloads", &json!({})).unwrap(); + let row = &r["downloads"][0]; + assert_eq!((row["id"].as_str(), row["state"].as_str()), (Some("my-model"), Some("installed")), "{r}"); + assert!(!staged.exists(), "the staged file was moved into the model's folder"); + let l = s.execute("faces.models.list", &json!({})).unwrap(); + assert_eq!(l["models"].as_array().unwrap().iter().filter(|m| m["installed"] == true).count(), 1); + // announced once, cleared on request + assert_eq!(s.execute("faces.models.downloadCancel", &json!({"id": "my-model"})).unwrap()["discarded"], true); + assert_eq!(s.execute("faces.models.downloads", &json!({})).unwrap()["downloads"], json!([])); + assert_eq!(s.execute("faces.models.downloadCancel", &json!({"id": "my-model"})).unwrap()["discarded"], false); + + // a file that cannot be installed is thrown away with a reason + std::fs::write(&staged, b"this is not a model").unwrap(); + let sha = lightcraft_faces::hash::sha256_file(&staged).unwrap(); + s.face_downloads.arrived("bad", staged.clone(), &sha); + let r = s.execute("faces.models.downloads", &json!({})).unwrap(); + assert_eq!(r["downloads"][0]["state"], "failed", "{r}"); + assert!(!staged.exists()); + let _ = std::fs::remove_dir_all(&d); +} + +/// `faces.detect` on photos without faces: it reports nothing, and a new run replaces earlier detections but +/// never regions that came from XMP or were named; one undo step restores everything. +#[test] +fn detect_replaces_only_earlier_detections_and_is_one_undo_step() { + use lightcraft_catalog::Op; + use lightcraft_meta::{Region, RegionKind}; + let region = |name: Option<&str>, description: Option<&str>, x: f64| Region { + rect: lightcraft_geom::Rect { x0: x, y0: 0.2, x1: x + 0.2, y1: 0.5 }, + kind: RegionKind::Face, + name: name.map(str::to_string), + description: description.map(str::to_string), + }; + let Some(mut s) = demo_with_yunet() else { return }; + let id = s.active().unwrap(); + let mut meta = s.catalog.photo(id).unwrap().meta.clone(); + meta.regions = + vec![region(Some("Jane Doe"), None, 0.1), region(None, Some("Detected by YuNet 2023mar"), 0.5), region(None, Some("Drawn by hand"), 0.7)]; + s.commit("setup", Op::SetMeta { id, meta: Box::new(meta) }).unwrap(); + let undo_before = s.undo.len(); + + // a dry run reports and changes nothing + let r = s.execute("faces.detect", &json!({"apply": false})).unwrap(); + assert_eq!(r["applied"], false); + assert_eq!(s.catalog.photo(id).unwrap().meta.regions.len(), 3); + assert_eq!(s.undo.len(), undo_before); + + // a real run (the default) drops the earlier detection (the demo has no faces) and keeps the others + let r = s.execute("faces.detect", &json!({})).unwrap(); + assert_eq!(r["applied"], true); + assert_eq!(r["detector"], "YuNet 2023mar"); + let photo = r["photos"].as_array().unwrap().iter().find(|p| p["id"] == id.0).unwrap(); + assert_eq!(photo["faces"].as_array().unwrap().len(), 0, "the procedural demo photos have no faces"); + let names: Vec<_> = s.catalog.photo(id).unwrap().meta.regions.iter().map(|r| (r.name.clone(), r.description.clone())).collect(); + assert_eq!(names, vec![(Some("Jane Doe".to_string()), None), (None, Some("Drawn by hand".to_string()))]); + assert_eq!(s.undo.len(), undo_before + 1, "one undo step for the whole run"); + s.execute("edit.undo", &json!({})).unwrap(); + assert_eq!(s.catalog.photo(id).unwrap().meta.regions.len(), 3, "undo brings the detection back"); + + // odd parameters are tolerated, and nothing is detected without a selection + for p in [json!({"score": 5}), json!({"score": -1}), json!({"nmsIou": 99}), json!({"maxFaces": 0}), json!({"score": "x"})] { + assert!(s.execute("faces.detect", &p).is_ok(), "{p}"); + } + let mut empty = Session::new(); + assert!(empty.execute("faces.detect", &json!({})).is_err()); +} + +/// Opt-in, needs the internet: `cargo test -p lightcraft-engine real_download -- --ignored --nocapture`. Downloads YuNet +/// from its pinned address the way the app does, waits for it to be installed, then finds faces with it. +#[test] +#[ignore = "downloads from github.com"] +fn real_download_of_yunet_installs_it_and_detects() { + let d = temp("real-download"); + let mut s = Session::with_demo(); + s.face_models_dir = Some(d.join("models")); + let started = s.execute("faces.models.download", &json!({"id": "yunet-2023mar", "acknowledged": true})).unwrap(); + assert_eq!(started["from"], "github.com"); + let t0 = std::time::Instant::now(); + loop { + let r = s.execute("faces.models.downloads", &json!({})).unwrap(); + let row = r["downloads"][0].clone(); + match row["state"].as_str() { + Some("installed") => break, + Some("running") | Some("done") => {} + _ => panic!("{r}"), + } + assert!(t0.elapsed() < std::time::Duration::from_secs(120), "{r}"); + std::thread::sleep(std::time::Duration::from_millis(100)); + } + println!("downloaded and installed in {:?}", t0.elapsed()); + let l = s.execute("faces.models.list", &json!({})).unwrap(); + assert_eq!(find(&l, "yunet-2023mar")["installed"], true); + let r = s.execute("faces.detect", &json!({"apply": false})).unwrap(); + assert_eq!(r["detector"], "YuNet 2023mar"); + // asking again once installed is refused + assert!(s.execute("faces.models.download", &json!({"id": "yunet-2023mar", "acknowledged": true})).is_err()); + let _ = std::fs::remove_dir_all(&d); +} + +/// Opt-in, needs the internet: the recognisers are big, so this only checks that bytes start to arrive (through the +/// hosts' redirects) and that cancelling leaves nothing behind. +#[test] +#[ignore = "downloads from github.com and huggingface.co"] +fn real_downloads_of_the_recognisers_start_and_can_be_cancelled() { + let d = temp("real-start"); + let mut s = Session::with_demo(); + s.face_models_dir = Some(d.join("models")); + for id in ["sface-2021dec", "auraface-v1"] { + s.execute("faces.models.download", &json!({"id": id, "acknowledged": true})).unwrap(); + let t0 = std::time::Instant::now(); + loop { + let r = s.execute("faces.models.downloads", &json!({})).unwrap(); + let row = r["downloads"].as_array().unwrap().iter().find(|x| x["id"] == id).unwrap().clone(); + if row["state"] == "running" && row["bytes"].as_u64().unwrap_or(0) > 100_000 { + println!("{id}: {} of {} bytes after {:?}", row["bytes"], row["total"], t0.elapsed()); + break; + } + assert!(row["state"] == "running", "{r}"); + assert!(t0.elapsed() < std::time::Duration::from_secs(60), "{r}"); + std::thread::sleep(std::time::Duration::from_millis(50)); + } + s.execute("faces.models.downloadCancel", &json!({"id": id})).unwrap(); + } + drop(s); + let staging = d.join("models").join(".downloads"); + let left: Vec<_> = std::fs::read_dir(&staging).map(|r| r.flatten().map(|e| e.file_name()).collect()).unwrap_or_default(); + println!("left in staging: {left:?}"); + let _ = std::fs::remove_dir_all(&d); +} + +/// With the recognition runtime: installing runs the model's self-test, a model that does not work is refused +/// and leaves nothing behind, and `faces.models.test` runs it again. +#[cfg(not(target_arch = "wasm32"))] +#[test] +fn installing_runs_the_self_test_and_refuses_models_that_do_not_work() { + let d = temp("selftest"); + let mut s = session(&d); + // a working model passes; its record keeps the result + let f = model_file(&d, "Works.onnx", 64); + let r = s.execute("faces.models.install", &json!({"path": f, "acknowledged": true})).unwrap(); + let id = r["installed"]["id"].as_str().unwrap().to_string(); + assert_eq!(r["installed"]["accepted"]["selfTest"]["ok"], true, "{r}"); + assert_eq!(r["installed"]["accepted"]["selfTest"]["dimension"], 64); + let t = s.execute("faces.models.test", &json!({"id": id})).unwrap(); + assert_eq!(t["ok"], true, "{t}"); + assert!(t["result"]["embedMs"].as_f64().unwrap() >= 0.0); + assert!(s.execute("faces.models.test", &json!({"id": "nope"})).is_err()); + assert!(s.execute("faces.models.test", &json!({})).is_err()); + // a graph with no layers passes the shape check but cannot run: refused, and nothing is left on disk + let broken = d.join("Broken.onnx"); + std::fs::write(&broken, lightcraft_faces::synthetic::embedder_model(64)).unwrap(); + let e = s.execute("faces.models.install", &json!({"path": broken.to_string_lossy(), "acknowledged": true})).unwrap_err().to_string(); + assert!(e.contains("could not be loaded") || e.contains("self-test"), "{e}"); + let leftovers: Vec<_> = std::fs::read_dir(d.join("models")) + .unwrap() + .flatten() + .map(|e| e.file_name().to_string_lossy().into_owned()) + .filter(|n| n != "settings.json") + .collect(); + assert!(leftovers.len() == 1 && leftovers[0].starts_with("custom-works-"), "only the working model remains: {leftovers:?}"); + let _ = std::fs::remove_dir_all(&d); +} + +#[cfg(not(target_arch = "wasm32"))] +#[test] +fn failed_replacement_keeps_the_installed_model_and_records() { + let d = temp("replacement"); + let models = d.join("models"); + std::fs::create_dir_all(&models).unwrap(); + let good = model_file(&d, "Good.onnx", 64); + let write_catalog = |path: &std::path::Path| { + let catalog = json!({"models": [{ + "id": "replacement", "name": "Replacement", "version": "1", "role": "embedder", + "url": "https://models.example.org/weights.onnx", + "sha256": lightcraft_faces::hash::sha256_file(path).unwrap(), + "sizeBytes": std::fs::metadata(path).unwrap().len(), + "licence": {"name": "MIT", "commercial": "yes"}, + "output": {"kind": "embedding", "dim": 64} + }]}); + std::fs::write(models.join("catalog.json"), catalog.to_string()).unwrap(); + }; + write_catalog(std::path::Path::new(&good)); + let mut s = session(&d); + s.execute("faces.models.install", &json!({"path": good, "acknowledged": true})).unwrap(); + let home = models.join("replacement"); + let files = ["model.onnx", "face-model.json", "installed.json"]; + let before: Vec<_> = files.iter().map(|name| std::fs::read(home.join(name)).unwrap()).collect(); + let settings = std::fs::read(models.join("settings.json")).unwrap(); + let bad = d.join("Bad.onnx"); + std::fs::write(&bad, lightcraft_faces::synthetic::embedder_model(64)).unwrap(); + write_catalog(&bad); + let mut s = session(&d); + assert!(s.execute("faces.models.install", &json!({"path": bad, "acknowledged": true})).is_err()); + for (name, bytes) in files.iter().zip(before) { + assert_eq!(std::fs::read(home.join(name)).unwrap(), bytes, "{name} was replaced"); + } + assert_eq!(std::fs::read(models.join("settings.json")).unwrap(), settings); + assert!(!home.join("model.onnx.part").exists()); + let _ = std::fs::remove_dir_all(d); +} + +fn settings_of(d: &std::path::Path) -> Value { + serde_json::from_slice(&std::fs::read(d.join("models").join("settings.json")).unwrap()).unwrap() +} + +#[test] +fn installing_a_model_chooses_it_and_switches_recognition_on() { + let d = temp("activate"); + let mut s = session(&d); + let a = s.execute("faces.models.install", &json!({"path": model_file(&d, "First.onnx", 64), "acknowledged": true})).unwrap(); + let a_id = a["installed"]["id"].as_str().unwrap().to_string(); + assert_eq!(a["active"], json!({"embedder": a_id, "enabled": true})); + assert_eq!(a["installed"]["selected"], true); + assert_eq!(settings_of(&d), json!({"enabled": true, "embedder": a_id})); + // a second model takes over (the newest is the one the user just asked for); the first stays installed + let b = s.execute("faces.models.install", &json!({"path": model_file(&d, "Second.onnx", 96), "acknowledged": true})).unwrap(); + let b_id = b["installed"]["id"].as_str().unwrap().to_string(); + assert_ne!(a_id, b_id); + assert_eq!(b["active"]["embedder"], json!(b_id)); + let l = s.execute("faces.models.list", &json!({})).unwrap(); + assert_eq!((find(&l, &a_id)["installed"].clone(), find(&l, &a_id)["selected"].clone()), (json!(true), json!(false))); + assert_eq!(find(&l, &b_id)["selected"], true); + // switching back is a choice, and a model added with `activate: false` leaves the choice alone + s.execute("faces.models.select", &json!({"id": a_id})).unwrap(); + let c = s.execute("faces.models.install", &json!({"path": model_file(&d, "Third.onnx", 128), "acknowledged": true, "activate": false})).unwrap(); + let c_id = c["installed"]["id"].as_str().unwrap().to_string(); + assert_eq!(c["active"]["embedder"], json!(a_id)); + assert_eq!(c["installed"]["selected"], false); + // a switched-off recognition is switched back on by the next install, not by a refused one + s.execute("faces.enable", &json!({"enabled": false})).unwrap(); + assert!(s.execute("faces.models.install", &json!({"path": model_file(&d, "Nope.onnx", 64)})).is_err()); + assert_eq!(settings_of(&d)["enabled"], false); + s.execute("faces.models.install", &json!({"path": model_file(&d, "Fourth.onnx", 64), "acknowledged": true})).unwrap(); + assert_eq!(settings_of(&d)["enabled"], true); + // removing the model in use hands over to another installed one instead of leaving nothing chosen + let in_use = settings_of(&d)["embedder"].as_str().unwrap().to_string(); + let r = s.execute("faces.models.remove", &json!({"id": in_use})).unwrap(); + let next = r["embedder"].as_str().unwrap().to_string(); + assert_ne!(next, in_use); + assert!([&a_id, &b_id, &c_id].contains(&&next) || next.starts_with("custom-"), "{next}"); + assert_eq!(settings_of(&d)["embedder"], json!(next)); + // removing a model that is not in use changes nothing + let other = [&a_id, &b_id, &c_id].into_iter().find(|i| **i != next && **i != in_use).unwrap().clone(); + s.execute("faces.models.remove", &json!({"id": other})).unwrap(); + assert_eq!(settings_of(&d)["embedder"], json!(next)); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn a_download_that_arrives_is_installed_and_switched_on_by_itself() { + use crate::face_download::{STAGING, State}; + let d = temp("finish"); + let mut s = session(&d); + let staging = d.join("models").join(STAGING); + std::fs::create_dir_all(&staging).unwrap(); + // what the download thread leaves when it succeeds: a verified file in the staging folder and a finished job + let staged = model_file(&staging, "arrived.onnx", 64); + let sha = lightcraft_faces::hash::sha256_file(std::path::Path::new(&staged)).unwrap(); + s.face_downloads.set_outcome("fake-model", State::Done { path: staged.clone().into(), sha256: sha }); + let r = s.execute("faces.models.downloads", &json!({})).unwrap(); + assert_eq!(r["downloads"][0]["id"], "fake-model"); + assert_eq!(r["downloads"][0]["state"], "installed", "{r}"); + let st = settings_of(&d); + assert_eq!(st["enabled"], true); + assert!(st["embedder"].as_str().unwrap().starts_with("custom-arrived-"), "{st}"); + assert!(!std::path::Path::new(&staged).exists(), "the staged file was moved into place"); + // the record stays until it is cleared, and is cleared once + assert_eq!(s.execute("faces.models.downloads", &json!({})).unwrap()["downloads"][0]["state"], "installed"); + assert_eq!(s.execute("faces.models.downloadCancel", &json!({"id": "fake-model"})).unwrap()["discarded"], true); + assert_eq!(s.execute("faces.models.downloads", &json!({})).unwrap()["downloads"], json!([])); + + // a download that turns out not to be a usable model is thrown away, and the reason is shown + let junk = staging.join("junk.onnx"); + std::fs::write(&junk, b"definitely not a model").unwrap(); + let sha = lightcraft_faces::hash::sha256_file(&junk).unwrap(); + s.face_downloads.set_outcome("fake-junk", State::Done { path: junk.clone(), sha256: sha }); + let r = s.execute("faces.models.downloads", &json!({})).unwrap(); + assert_eq!(r["downloads"][0]["state"], "failed", "{r}"); + assert!(r["downloads"][0]["error"].as_str().unwrap().contains("cannot be used"), "{r}"); + assert!(!junk.exists()); + assert_eq!(settings_of(&d)["embedder"], st["embedder"], "a failed download leaves the choice alone"); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn the_users_own_catalog_adds_models_with_the_same_download_and_install() { + use crate::face_download::State; + let d = temp("catalog"); + let mut s = session(&d); + let models = d.join("models"); + std::fs::create_dir_all(&models).unwrap(); + // a model file, and a catalog that lists it by its hash (plus an entry that cannot be used) + let file = model_file(&d, "weights.onnx", 64); + let (sha, size) = (lightcraft_faces::hash::sha256_file(std::path::Path::new(&file)).unwrap(), std::fs::metadata(&file).unwrap().len()); + let catalog = json!({"models": [ + {"id": "my-research-model", "name": "My research model", "version": "1", "role": "embedder", "url": "https://models.example.org/w/weights.onnx", + "sha256": sha, "sizeBytes": size, "licence": {"name": "Research only", "commercial": "no", "notice": "Not for commercial use."}, + "provenance": "Some dataset", "output": {"kind": "embedding", "dim": 64}, "thresholds": {"matchCosine": 0.42}}, + {"id": "broken", "name": "No hash", "version": "1", "role": "embedder", "url": "https://models.example.org/x.onnx", "output": {"kind": "embedding", "dim": 64}}, + ]}); + std::fs::write(models.join("catalog.json"), catalog.to_string()).unwrap(); + // it is listed like a built-in model, with its own terms, a download address, and the problem with the other entry + let l = s.execute("faces.models.list", &json!({})).unwrap(); + let m = find(&l, "my-research-model"); + assert_eq!( + (m["fromCatalog"].clone(), m["downloadHost"].clone(), m["licence"]["commercial"].clone()), + (json!(true), json!("models.example.org"), json!("no")) + ); + assert_eq!(m["installed"], false); + assert_eq!(find(&l, "broken"), &Value::Null); + assert_eq!(l["catalog"]["models"], 1); + assert!(l["catalog"]["errors"][0].as_str().unwrap().contains("broken"), "{}", l["catalog"]); + // nothing is fetched without the terms being accepted + assert!(s.execute("faces.models.download", &json!({"id": "my-research-model"})).is_err()); + // a file that matches is recognised as that model (not as a draft), so its terms and settings are used + let seen = s.execute("faces.models.inspect", &json!({"path": file})).unwrap(); + assert_eq!((seen["kind"].clone(), seen["model"]["id"].clone()), (json!("known"), json!("my-research-model"))); + // a finished download of it installs under the catalog's id, with the catalog's licence record + let staging = models.join(crate::face_download::STAGING); + std::fs::create_dir_all(&staging).unwrap(); + let staged = staging.join("weights.onnx"); + std::fs::copy(&file, &staged).unwrap(); + s.face_downloads.set_outcome("my-research-model", State::Done { path: staged, sha256: sha }); + let r = s.execute("faces.models.downloads", &json!({})).unwrap(); + assert_eq!(r["downloads"][0]["state"], "installed", "{r}"); + let st = settings_of(&d); + assert_eq!((st["embedder"].clone(), st["enabled"].clone()), (json!("my-research-model"), json!(true))); + let l = s.execute("faces.models.list", &json!({})).unwrap(); + assert_eq!( + (find(&l, "my-research-model")["installed"].clone(), find(&l, "my-research-model")["accepted"]["commercial"].clone()), + (json!(true), json!("no")) + ); + // a broken or hostile catalog never stops the list + for bytes in [&b"{"[..], b"[]", &vec![b'x'; 400_000]] { + std::fs::write(models.join("catalog.json"), bytes).unwrap(); + let l = s.execute("faces.models.list", &json!({})).unwrap(); + assert!(l["catalog"]["errors"].as_array().is_some_and(|e| !e.is_empty()), "{}", l["catalog"]); + assert!(find(&l, "sface-2021dec")["id"] == "sface-2021dec"); + } + let _ = std::fs::remove_dir_all(&d); +} diff --git a/crates/engine/src/tests_face_recognize.rs b/crates/engine/src/tests_face_recognize.rs new file mode 100644 index 00000000..2fd53c67 --- /dev/null +++ b/crates/engine/src/tests_face_recognize.rs @@ -0,0 +1,489 @@ +//! Recognition: indexing faces, suggesting names, naming, with a tiny working network standing in for a real +//! model (its embedding depends on the picture's colours, which is enough to test the machinery; how well real +//! models recognise people is measured separately, on real photos). + +use std::path::PathBuf; + +use lightcraft_catalog::{Op, PhotoId}; +use lightcraft_geom::Rect; +use lightcraft_meta::{Region, RegionKind}; +use serde_json::{Value, json}; + +use crate::Session; + +pub(crate) fn temp(name: &str) -> PathBuf { + let d = std::env::temp_dir().join(format!("lc-facerec-{name}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&d); + std::fs::create_dir_all(&d).unwrap(); + d +} + +/// A demo session with a tiny model installed and chosen. +pub(crate) fn setup(d: &std::path::Path, dim: u64) -> (Session, String) { + let mut s = Session::with_demo(); + s.face_models_dir = Some(d.join("models")); + let model = d.join(format!("tiny{dim}.onnx")); + std::fs::write(&model, lightcraft_faces::synthetic::tiny_embedder_model(dim)).unwrap(); + let r = s.execute("faces.models.install", &json!({"path": model.to_string_lossy(), "acknowledged": true})).unwrap(); + let id = r["installed"]["id"].as_str().unwrap().to_string(); + s.execute("faces.models.select", &json!({"id": id})).unwrap(); + (s, id) +} + +/// Installs the real YuNet from `LC_YUNET_MODEL=` (what Settings > Faces downloads; nothing is +/// committed) and says whether it did: the parts of a test that need to find faces run only then. +pub(crate) fn install_yunet(s: &mut Session) -> bool { + let Some(file) = std::env::var_os("LC_YUNET_MODEL") else { + eprintln!("LC_YUNET_MODEL is not set: skipping the face-finding part"); + return false; + }; + let r = s.execute("faces.models.install", &json!({"path": file.to_string_lossy(), "acknowledged": true})); + assert!(r.is_ok(), "{r:?}"); + true +} + +pub(crate) fn region(x: f64, name: Option<&str>) -> Region { + Region { rect: Rect { x0: x, y0: 0.2, x1: x + 0.25, y1: 0.6 }, kind: RegionKind::Face, name: name.map(str::to_string), description: None } +} + +pub(crate) fn set_regions(s: &mut Session, id: PhotoId, regions: Vec) { + let mut meta = s.catalog.photo(id).unwrap().meta.clone(); + meta.regions = regions; + s.commit("setup", Op::SetMeta { id, meta: Box::new(meta) }).unwrap(); +} + +pub(crate) fn two_photos(s: &Session) -> (PhotoId, PhotoId) { + let mut ids = s.catalog.photos().map(|p| p.id).collect::>(); + ids.sort(); + (ids[0], ids[1]) +} + +#[test] +fn indexing_embeds_each_face_once_and_reports_progress() { + let d = temp("index"); + let (mut s, _) = setup(&d, 64); + let (a, b) = two_photos(&s); + set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, None)]); + set_regions(&mut s, b, vec![region(0.3, None)]); + // only a report + let r = s.execute("faces.index", &json!({"budgetMs": 0})).unwrap(); + assert_eq!((r["embedded"].clone(), r["photosDone"].clone(), r["pendingPhotos"].clone()), (json!(0), json!(0), json!(2)), "{r}"); + // a generous budget does it all, the named photo first + let r = s.execute("faces.index", &json!({"budgetMs": 60_000})).unwrap(); + assert_eq!((r["embedded"].clone(), r["pendingPhotos"].clone(), r["indexedFaces"].clone()), (json!(3), json!(0), json!(3)), "{r}"); + // nothing is embedded twice + let r = s.execute("faces.index", &json!({"budgetMs": 60_000})).unwrap(); + assert_eq!((r["embedded"].clone(), r["photosDone"].clone()), (json!(0), json!(0))); + // a small budget still makes progress (at least one photo per call), and `ids` go first + let c = s.catalog.photos().map(|p| p.id).find(|id| *id != a && *id != b).unwrap(); + set_regions(&mut s, c, vec![region(0.2, None)]); + set_regions(&mut s, b, vec![region(0.3, None), region(0.6, None)]); + let r = s.execute("faces.index", &json!({"budgetMs": 1, "ids": [c.0]})).unwrap(); + assert_eq!(r["photosDone"], 1, "{r}"); + assert_eq!(r["pendingPhotos"], 1, "{r}"); + assert!(s.faces.index.contains(c.0, ®ion(0.2, None).rect), "the photo asked for first was done first"); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn suggestions_come_from_the_named_faces_and_are_never_applied() { + let d = temp("suggest"); + let (mut s, _) = setup(&d, 64); + let (a, b) = two_photos(&s); + set_regions(&mut s, a, vec![region(0.1, Some("Ann"))]); + set_regions(&mut s, b, vec![region(0.1, None), region(0.5, Some("Pet?")), region(0.7, None)]); + // thresholds that accept anything: Ann is the only person with a named face except "Pet?" + let r = s.execute("faces.suggest", &json!({"ids": [b.0], "threshold": -1.0, "margin": -2.0, "budgetMs": 60_000})).unwrap(); + assert_eq!(r["galleryFaces"], 2, "{r}"); + let faces = r["photos"][0]["faces"].as_array().unwrap(); + assert_eq!(faces.len(), 2, "only the unnamed faces get suggestions: {r}"); + assert!(faces.iter().all(|f| f["suggestion"]["name"].is_string()), "{r}"); + assert_eq!(faces[0]["index"], 0); + assert_eq!(faces[1]["index"], 2); + assert!(faces[0]["candidates"].as_array().unwrap().len() == 2); + // an impossible threshold suggests nobody but still lists the candidates + let r = s.execute("faces.suggest", &json!({"ids": [b.0], "threshold": 2.0, "budgetMs": 60_000})).unwrap(); + assert!(r["photos"][0]["faces"][0]["suggestion"].is_null()); + assert!(!r["photos"][0]["faces"][0]["candidates"].as_array().unwrap().is_empty()); + // nothing was named by suggesting + assert!(s.catalog.photo(b).unwrap().meta.regions.iter().filter(|r| r.name.is_none()).count() == 2); + // pets and other kinds are not people + let mut meta = s.catalog.photo(b).unwrap().meta.clone(); + meta.regions[1].kind = RegionKind::Pet; + s.commit("pet", Op::SetMeta { id: b, meta: Box::new(meta) }).unwrap(); + let r = s.execute("faces.suggest", &json!({"ids": [b.0], "threshold": -1.0, "margin": -2.0, "budgetMs": 60_000})).unwrap(); + assert_eq!(r["galleryFaces"], 1, "{r}"); + // odd parameters are tolerated + for p in [json!({"threshold": "x"}), json!({"margin": null}), json!({"budgetMs": 0}), json!({"ids": []}), json!({"ids": [999999]})] { + assert!(s.execute("faces.suggest", &p).is_ok(), "{p}"); + } + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn naming_a_face_is_undoable_and_makes_it_yours() { + let d = temp("name"); + let (mut s, _) = setup(&d, 64); + let (a, _) = two_photos(&s); + let mut detected = region(0.1, None); + detected.description = Some("Detected by YuNet 2023mar".into()); + set_regions(&mut s, a, vec![detected, region(0.5, None)]); + let undo = s.undo.len(); + let r = s.execute("faces.setName", &json!({"id": a.0, "index": 0, "name": " Ann "})).unwrap(); + assert_eq!(r["name"], "Ann"); + let regions = s.catalog.photo(a).unwrap().meta.regions.clone(); + assert_eq!((regions[0].name.as_deref(), regions[0].description.as_deref()), (Some("Ann"), None), "a named detection loses its detected marker"); + assert_eq!(s.undo.len(), undo + 1); + // a new detection run now leaves it alone + if install_yunet(&mut s) { + s.execute("faces.detect", &json!({"id": a.0})).unwrap(); + assert_eq!(s.catalog.photo(a).unwrap().meta.regions.iter().filter(|r| r.name.as_deref() == Some("Ann")).count(), 1); + } + // clearing, undo + s.execute("faces.setName", &json!({"id": a.0, "index": 0, "name": null})).unwrap(); + assert_eq!(s.catalog.photo(a).unwrap().meta.regions[0].name, None); + s.execute("edit.undo", &json!({})).unwrap(); + assert_eq!(s.catalog.photo(a).unwrap().meta.regions[0].name.as_deref(), Some("Ann")); + // bad input + for p in [ + json!({"id": a.0}), + json!({"id": a.0, "index": 9, "name": "x"}), + json!({"id": a.0, "index": 0, "name": 5}), + json!({"id": a.0, "index": 0, "name": "a\u{7}b"}), + json!({"id": a.0, "index": 0, "name": "x".repeat(201)}), + json!({"id": 999999, "index": 0, "name": "x"}), + ] { + assert!(s.execute("faces.setName", &p).is_err(), "{p}"); + } + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn embeddings_are_cached_beside_the_library_and_reset_for_another_model() { + let d = temp("persist"); + let lib = d.join("lib"); + let models = d.join("models"); + let install = |s: &mut Session, dim: u64| -> String { + let model = d.join(format!("tiny{dim}.onnx")); + std::fs::write(&model, lightcraft_faces::synthetic::tiny_embedder_model(dim)).unwrap(); + let r = s.execute("faces.models.install", &json!({"path": model.to_string_lossy(), "acknowledged": true})).unwrap(); + let id = r["installed"]["id"].as_str().unwrap().to_string(); + s.execute("faces.models.select", &json!({"id": id})).unwrap(); + id + }; + let mut s = Session::new().with_fs(); + s.face_models_dir = Some(models.clone()); + s.open_library(&lib, true).unwrap(); + let (a, _) = two_photos(&s); + set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, None)]); + install(&mut s, 64); + s.execute("faces.index", &json!({"budgetMs": 60_000})).unwrap(); + assert!(lib.join(format!("face-embeddings-{}.bin", s.faces.index.tag.split('@').next().unwrap())).is_file()); + s.close_library().unwrap(); + drop(s); + + // a new session finds the embeddings on disk and has nothing left to do + let mut s2 = Session::new().with_fs(); + s2.face_models_dir = Some(models); + s2.open_library(&lib, true).unwrap(); + let r = s2.execute("faces.index", &json!({"budgetMs": 60_000})).unwrap(); + assert_eq!((r["embedded"].clone(), r["indexedFaces"].clone(), r["pendingPhotos"].clone()), (json!(0), json!(2), json!(0)), "{r}"); + // choosing another model starts over: its faces are different numbers, the old ones are not used + install(&mut s2, 32); + let r = s2.execute("faces.index", &json!({"budgetMs": 0})).unwrap(); + assert_eq!((r["indexedFaces"].clone(), r["pendingPhotos"].clone()), (json!(0), json!(1)), "{r}"); + let r = s2.execute("faces.index", &json!({"budgetMs": 60_000})).unwrap(); + assert_eq!(r["embedded"], 2); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn without_a_chosen_model_recognition_says_so() { + let d = temp("nomodel"); + let mut s = Session::with_demo(); + assert!(s.execute("faces.index", &json!({})).is_err(), "no models folder"); + s.face_models_dir = Some(d.join("models")); + for cmd in ["faces.index", "faces.suggest"] { + let e = s.execute(cmd, &json!({})).unwrap_err().to_string(); + assert!(e.contains("no recognition model is chosen"), "{cmd}: {e}"); + } + let _: Value = json!(null); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn the_background_pump_indexes_by_itself_once_recognition_is_on() { + let d = temp("pump"); + let (mut s, _) = setup(&d, 64); + let (a, b) = two_photos(&s); + set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, None)]); + set_regions(&mut s, b, vec![region(0.3, None)]); + // installing the model switched recognition on; switched off, the pump does nothing, and says so + assert_eq!(s.execute("faces.enable", &json!({"enabled": false})).unwrap()["enabled"], false); + assert_eq!(s.execute("faces.pump", &json!({})).unwrap()["active"], false); + assert_eq!(s.faces.index.len(), 0); + s.execute("faces.enable", &json!({"enabled": true})).unwrap(); + // `enabled` is looked up at most once a second, so the first pump may still see "off" + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(60); + let mut last = json!(null); + while std::time::Instant::now() < deadline { + last = s.execute("faces.pump", &json!({})).unwrap(); + if last["active"] == true && last["pendingPhotos"] == 0 && last["indexedFaces"] == 3 { + break; + } + std::thread::sleep(std::time::Duration::from_millis(50)); + } + assert_eq!((last["active"].clone(), last["pendingPhotos"].clone(), last["indexedFaces"].clone()), (json!(true), json!(0), json!(3)), "{last}"); + // a new face is picked up (the catalog changed), and suggestions use what the pump embedded + set_regions(&mut s, b, vec![region(0.3, None), region(0.6, None)]); + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(60); + while std::time::Instant::now() < deadline && s.faces.index.len() < 4 { + s.execute("faces.pump", &json!({})).unwrap(); + std::thread::sleep(std::time::Duration::from_millis(50)); + } + assert_eq!(s.faces.index.len(), 4); + let r = s.execute("faces.suggest", &json!({"ids": [b.0], "threshold": -1.0, "margin": -2.0, "budgetMs": 0})).unwrap(); + assert_eq!(r["photos"][0]["faces"].as_array().unwrap().len(), 2); + assert!(r["photos"][0]["faces"][0]["suggestion"]["name"].is_string(), "{r}"); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn the_background_pump_saves_what_it_embeds_for_the_next_launch() { + let d = temp("pumpsave"); + let (lib, models) = (d.join("lib"), d.join("models")); + let mut s = Session::new().with_fs(); + s.face_models_dir = Some(models); + s.open_library(&lib, true).unwrap(); + let model = d.join("tiny64.onnx"); + std::fs::write(&model, lightcraft_faces::synthetic::tiny_embedder_model(64)).unwrap(); + // installing switches recognition on: nothing else is asked for + s.execute("faces.models.install", &json!({"path": model.to_string_lossy(), "acknowledged": true})).unwrap(); + // without the detector (not bundled) photos without faces are left unsearched + let searching = install_yunet(&mut s); + let (a, b) = two_photos(&s); + set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, None)]); + set_regions(&mut s, b, vec![region(0.3, None)]); + let photos = s.catalog.photos().count(); + let searched_photos = if searching { photos - 2 } else { 0 }; + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(120); + let mut last = json!(null); + while std::time::Instant::now() < deadline { + last = s.execute("faces.pump", &json!({})).unwrap(); + if last["pendingPhotos"] == 0 && s.faces.index.len() == 3 && s.faces.scanned.len() == searched_photos { + break; + } + std::thread::sleep(std::time::Duration::from_millis(50)); + } + assert_eq!((s.faces.index.len(), s.faces.scanned.len()), (3, searched_photos), "{last}"); + let (tag, path) = (s.faces.index.tag.clone(), s.faces.index.path.clone().unwrap()); + assert!(path.file_name().unwrap().to_string_lossy().starts_with("face-embeddings-custom-"), "one cache file per model: {}", path.display()); + // the first batch is written at once; whatever is left is written when the session ends, so a launch never redoes + // work: neither the embeddings nor the list of photos already searched for faces + s.close_library().unwrap(); + drop(s); + let mut again = crate::faces_index::Index::default(); + again.reset(&tag, 64, Some(path)); + assert_eq!(again.len(), 3); + let mut searched = crate::faces_index::Scanned::default(); + searched.reset(Some(lib.join("face-scanned.bin"))); + assert_eq!(searched.len(), searched_photos, "every photo without faces is remembered as searched"); + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn faces_person_lists_a_persons_faces_and_looks_for_more_only_while_recognition_runs() { + let d = temp("person"); + let (mut s, _) = setup(&d, 64); + s.execute("faces.index", &json!({"budgetMs": 0})).unwrap(); + let (a, b) = two_photos(&s); + set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, Some("Bob"))]); + set_regions(&mut s, b, vec![region(0.1, None), region(0.5, Some("ann"))]); + let unit = |parts: &[(usize, f32)]| { + let mut v = vec![0.0f32; 64]; + for (i, x) in parts { + v[*i] = *x; + } + let n = v.iter().map(|x| x * x).sum::().sqrt(); + v.iter().map(|x| x / n).collect::>() + }; + for (id, x, v) in [(a, 0.1, unit(&[(0, 1.0)])), (a, 0.5, unit(&[(1, 1.0)])), (b, 0.1, unit(&[(0, 0.9), (1, 0.1)])), (b, 0.5, unit(&[(0, 1.0)]))] { + s.faces.index.insert(id.0, ®ion(x, None).rect, v); + } + // the name is matched without regard to case or surrounding spaces; every face is listed, wherever it is + let r = s.execute("faces.person", &json!({"name": " ANN "})).unwrap(); + assert_eq!((r["name"].clone(), r["total"].clone()), (json!("Ann"), json!(2))); + let mut at: Vec<(u64, u64)> = + r["confirmed"].as_array().unwrap().iter().map(|f| (f["photo"].as_u64().unwrap(), f["index"].as_u64().unwrap())).collect(); + at.sort(); + let mut want = vec![(a.0, 0), (b.0, 1)]; + want.sort(); + assert_eq!(at, want); + assert!(r["confirmed"].as_array().unwrap().iter().any(|f| f["rect"]["x0"] == 0.1 && f["rect"]["y1"] == 0.6)); + // recognition is on (installing the model switched it on): the unnamed face that looks like Ann is offered + assert_eq!(r["ready"], true); + let more = r["more"].as_array().unwrap(); + assert_eq!((more.len(), more[0]["photo"].clone(), more[0]["index"].clone()), (1, json!(b.0), json!(0)), "{r}"); + assert!(more[0]["score"].as_f64().unwrap() > 0.9); + // confirming it moves it from one list to the other (and is one undo step) + s.execute("faces.setName", &json!({"id": b.0, "index": 0, "name": "Ann"})).unwrap(); + let r = s.execute("faces.person", &json!({"name": "Ann"})).unwrap(); + assert_eq!((r["total"].clone(), r["more"].as_array().unwrap().len()), (json!(3), 0)); + // switched off, the faces are still listed but nothing is searched for + s.execute("faces.enable", &json!({"enabled": false})).unwrap(); + let r = s.execute("faces.person", &json!({"name": "Ann"})).unwrap(); + assert_eq!((r["total"].clone(), r["ready"].clone(), r["more"].as_array().unwrap().len()), (json!(3), json!(false), 0)); + // nobody by that name, and bad input, are plain answers and errors + let r = s.execute("faces.person", &json!({"name": "Nobody"})).unwrap(); + assert_eq!((r["total"].clone(), r["confirmed"].as_array().unwrap().len(), r["ready"].clone()), (json!(0), 0, json!(false))); + for bad in [json!({}), json!({"name": ""}), json!({"name": " "}), json!({"name": 7})] { + assert!(s.execute("faces.person", &bad).is_err(), "{bad}"); + } + assert!(s.execute("faces.person", &json!({"name": "Ann", "more": u64::MAX})).is_ok()); + let _ = std::fs::remove_dir_all(&d); +} + +fn unit64(parts: &[(usize, f32)]) -> Vec { + let mut v = vec![0.0f32; 64]; + for (i, x) in parts { + v[*i] = *x; + } + let n = v.iter().map(|x| x * x).sum::().sqrt(); + v.iter().map(|x| x / n).collect() +} + +#[test] +fn many_faces_are_named_in_one_undo_step() { + let d = temp("namemany"); + let (mut s, _) = setup(&d, 64); + let (a, b) = two_photos(&s); + let detected = |x: f64| Region { description: Some("Detected by YuNet 2023mar".into()), ..region(x, None) }; + set_regions(&mut s, a, vec![detected(0.1), region(0.5, None)]); + set_regions(&mut s, b, vec![detected(0.3)]); + let undo = s.undo.len(); + let r = s + .execute( + "faces.nameFaces", + &json!({"faces": [{"photo": a.0, "index": 0}, {"photo": a.0, "index": 1}, {"photo": b.0, "index": 0}], "name": " Jane Doe "}), + ) + .unwrap(); + assert_eq!((r["named"].clone(), r["photos"].clone()), (json!(3), json!(2))); + assert_eq!(s.undo.len(), undo + 1, "one step for the whole group"); + for (id, i) in [(a, 0), (a, 1), (b, 0)] { + let reg = &s.catalog.photo(id).unwrap().meta.regions[i]; + assert_eq!(reg.name.as_deref(), Some("Jane Doe")); + assert!(reg.description.is_none() || !reg.description.as_deref().unwrap().starts_with("Detected by"), "a named face is the user's now"); + } + // undo restores every one of them together + s.execute("edit.undo", &json!({})).unwrap(); + assert!(s.catalog.photo(a).unwrap().meta.regions.iter().chain(s.catalog.photo(b).unwrap().meta.regions.iter()).all(|r| r.name.is_none())); + // an empty name clears; faces that do not exist are skipped, and with none left it is an error + s.execute("faces.nameFaces", &json!({"faces": [{"photo": a.0, "index": 0}], "name": "Ann"})).unwrap(); + let r = s + .execute( + "faces.nameFaces", + &json!({"faces": [{"photo": a.0, "index": 0}, {"photo": a.0, "index": 99}, {"photo": 987_654, "index": 0}], "name": null}), + ) + .unwrap(); + assert_eq!(r["named"], 1); + assert!(s.catalog.photo(a).unwrap().meta.regions[0].name.is_none()); + for bad in [ + json!({"name": "x"}), + json!({"faces": "x", "name": "x"}), + json!({"faces": [{"photo": a.0}], "name": "x"}), + json!({"faces": [{"photo": 987_654, "index": 0}], "name": "x"}), + json!({"faces": [{"photo": a.0, "index": 0}], "name": 7}), + json!({"faces": [{"photo": a.0, "index": 0}], "name": "bad\u{7}name"}), + json!({"faces": vec![json!({"photo": a.0, "index": 0}); 5001], "name": "x"}), + ] { + assert!(s.execute("faces.nameFaces", &bad).is_err(), "{bad}"); + } + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn unnamed_faces_come_with_look_alikes_together_and_the_name_the_named_ones_suggest() { + let d = temp("unnamed"); + let (mut s, _) = setup(&d, 64); + s.execute("faces.index", &json!({"budgetMs": 0})).unwrap(); + let (a, b) = two_photos(&s); + let c = s.catalog.photos().map(|p| p.id).find(|id| *id != a && *id != b).unwrap(); + set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, None), region(0.7, None)]); + set_regions(&mut s, b, vec![region(0.1, None), region(0.5, None)]); + set_regions(&mut s, c, vec![region(0.3, None)]); + // three faces like Ann, two like Bob (unnamed), one unembedded + let put = |s: &mut Session, id: lightcraft_catalog::PhotoId, x: f64, v: Vec| s.faces.index.insert(id.0, ®ion(x, None).rect, v); + put(&mut s, a, 0.1, unit64(&[(0, 1.0)])); // Ann, named + put(&mut s, a, 0.5, unit64(&[(1, 1.0)])); // like Bob + put(&mut s, b, 0.1, unit64(&[(0, 0.97), (2, 0.2)])); // like Ann + put(&mut s, b, 0.5, unit64(&[(1, 0.98), (2, 0.2)])); // like Bob + put(&mut s, c, 0.3, unit64(&[(0, 0.95), (3, 0.3)])); // like Ann + // (a's last face, at 0.7, is not embedded yet) + let r = s.execute("faces.unnamed", &json!({})).unwrap(); + assert_eq!((r["total"].clone(), r["ordered"].clone(), r["ready"].clone()), (json!(5), json!(true), json!(true)), "{r}"); + let faces = r["faces"].as_array().unwrap(); + let at = |f: &Value| (f["photo"].as_u64().unwrap(), f["index"].as_u64().unwrap()); + let order: Vec<(u64, u64)> = faces.iter().map(at).collect(); + // look-alikes are neighbours: the two like Bob are adjacent, and the three like Ann are adjacent; the unembedded one is last + let pos = |x: (u64, u64)| order.iter().position(|o| *o == x).unwrap(); + assert_eq!(pos((a.0, 1)).abs_diff(pos((b.0, 1))), 1, "{order:?}"); + assert_eq!(order.last(), Some(&(a.0, 2)), "{order:?}"); + // a face the named ones recognise carries the name they suggest + let like_ann = faces.iter().find(|f| at(f) == (b.0, 0)).unwrap(); + assert_eq!(like_ann["suggestion"]["name"], "Ann", "{like_ann}"); + assert!(faces.iter().find(|f| at(f) == (b.0, 1)).unwrap()["suggestion"].is_null(), "Bob is nobody the named faces know"); + // each face says which box to show it by; with no scan view it is its own box + for k in ["x0", "y0", "x1", "y1"] { + assert!((like_ann["view"][k].as_f64().unwrap() - like_ann["rect"][k].as_f64().unwrap()).abs() < 1e-3, "{like_ann}"); + } + // a limit cuts the list, not the count + let r = s.execute("faces.unnamed", &json!({"limit": 2})).unwrap(); + assert_eq!((r["total"].clone(), r["faces"].as_array().unwrap().len()), (json!(5), 2)); + // switched off, the faces are still listed, in photo order, without suggestions + s.execute("faces.enable", &json!({"enabled": false})).unwrap(); + let r = s.execute("faces.unnamed", &json!({})).unwrap(); + assert_eq!((r["total"].clone(), r["ordered"].clone(), r["ready"].clone()), (json!(5), json!(false), json!(false))); + assert!(r["faces"].as_array().unwrap().iter().all(|f| f["suggestion"].is_null())); + for odd in [json!({"limit": 0}), json!({"limit": "x"}), json!({"limit": u64::MAX})] { + assert!(s.execute("faces.unnamed", &odd).is_ok(), "{odd}"); + } + let _ = std::fs::remove_dir_all(&d); +} + +#[test] +fn every_scanned_face_is_shown_by_the_same_kind_of_box() { + let d = temp("view"); + let (mut s, _) = setup(&d, 64); + s.execute("faces.index", &json!({"budgetMs": 0})).unwrap(); + let (a, _) = two_photos(&s); + // a loosely drawn box (from some other tool) and the detector's tighter box of the same face + let loose = lightcraft_geom::Rect { x0: 0.2, y0: 0.1, x1: 0.8, y1: 0.9 }; + let tight = lightcraft_geom::Rect { x0: 0.35, y0: 0.25, x1: 0.65, y1: 0.65 }; + let mut meta = s.catalog.photo(a).unwrap().meta.clone(); + meta.regions = vec![Region { rect: loose, kind: RegionKind::Face, name: Some("Ann".into()), description: None }]; + s.commit("setup", Op::SetMeta { id: a, meta: Box::new(meta) }).unwrap(); + // before the scan has looked at it, the face is shown by its own box + assert_eq!(s.face_view(a, loose), loose); + s.faces.index.insert_with_view(a.0, &loose, unit64(&[(0, 1.0)]), &tight); + let shown = s.face_view(a, loose); + assert!((shown.x0 - tight.x0).abs() < 1e-3 && (shown.y1 - tight.y1).abs() < 1e-3, "shown by the detector's box once known: {shown:?}"); + // a face nobody has looked at, or in another photo, is shown by its own box + let other = lightcraft_geom::Rect { x0: 0.1, y0: 0.1, x1: 0.2, y1: 0.2 }; + assert_eq!(s.face_view(a, other), other); + // the page of the person uses it too + let r = s.execute("faces.person", &json!({"name": "Ann"})).unwrap(); + let view = &r["confirmed"][0]["view"]; + assert!((view["x0"].as_f64().unwrap() - tight.x0).abs() < 1e-3, "{r}"); + // and it survives a restart: the index file keeps it beside the embedding + let (tag, path) = (s.faces.index.tag.clone(), d.join("view-cache.bin")); + let mut written = crate::faces_index::Index::default(); + written.reset(&tag, 64, Some(path.clone())); + written.insert_with_view(a.0, &loose, unit64(&[(0, 1.0)]), &tight); + written.save().unwrap(); + drop(written); + let mut again = crate::faces_index::Index::default(); + again.reset(&tag, 64, Some(path)); + let v = again.view(a.0, &loose).expect("the face is remembered"); + assert!((v.x0 - tight.x0).abs() < 1e-3 && (v.y0 - tight.y0).abs() < 1e-3, "{v:?}"); + let _ = std::fs::remove_dir_all(&d); +} diff --git a/crates/engine/src/tests_xmp.rs b/crates/engine/src/tests_xmp.rs index fede1fec..ab801e23 100644 --- a/crates/engine/src/tests_xmp.rs +++ b/crates/engine/src/tests_xmp.rs @@ -335,10 +335,15 @@ fn synthetic_dng(xmp: &str) -> Vec { /// A tiny Bayer DNG written by our own DNG writer (optional embedded XMP, camera metadata). pub(crate) fn synthetic_dng_with(xmp: Option<&str>, metadata: lightcraft_meta::Metadata) -> Vec { - use lightcraft_raw::*; let (w, h) = (32usize, 24usize); - let cfa = Cfa::bayer("RGGB").unwrap(); let data: Vec = (0..w * h).map(|i| 256 + ((i % w) * 300 + (i / w) * 200) as u16).collect(); + synthetic_dng_of(w, h, "RGGB", data, xmp, metadata) +} + +/// A synthetic Bayer DNG for denoise and cache tests. +pub(crate) fn synthetic_dng_of(w: usize, h: usize, cfa: &str, data: Vec, xmp: Option<&str>, metadata: lightcraft_meta::Metadata) -> Vec { + use lightcraft_raw::*; + let cfa = Cfa::bayer(cfa).unwrap(); let raw = RawImage { format: RawFormat::Dng, width: w, diff --git a/crates/engine/tests/denoise_cpu.rs b/crates/engine/tests/denoise_cpu.rs new file mode 100644 index 00000000..d01b7c5a --- /dev/null +++ b/crates/engine/tests/denoise_cpu.rs @@ -0,0 +1,101 @@ +//! Opt-in same-machine denoise model-stage benchmark on one public Bayer raw. +//! LC_DENOISE_MODEL= LC_DENOISE_RAW= cargo test --release +//! -p lightcraft-engine --features denoise --test denoise_cpu -- --ignored --nocapture +//! LC_DENOISE_ONLY selects a file (default arw-sony-a7m3-compressed.arw). +//! The whole model stage includes packing, scale/clip protection, blending and normalization; +//! file decode, model loading and GPU setup are outside the timer. It uses the application's tile pool convention. +#![cfg(feature = "denoise")] + +use lightcraft_denoise::{ + TileRunner, + bayer::Layout, + run::{Control, Params, denoise_bayer}, + runtime::{CpuRunner, check_tile}, +}; +use lightcraft_raster::Rgb32f; +use std::{path::Path, time::Instant}; + +fn agree(a: &[f32], b: &[f32]) -> f32 { + assert_eq!(a.len(), b.len()); + assert!(a.iter().chain(b).all(|v| v.is_finite())); + let scale = b.iter().fold(1e-6f32, |m, v| m.max(v.abs())); + let error = a.iter().zip(b).fold(0f32, |m, (a, b)| m.max((a - b).abs())) / scale; + assert!(error <= 1e-3, "maximum relative error {error:e}"); + error +} +fn tile_time(name: &str, r: &dyn TileRunner, x: &[f32]) { + r.run(x).unwrap(); + let mut times: Vec<_> = (0..5) + .map(|_| { + let t = Instant::now(); + r.run(x).unwrap(); + t.elapsed().as_secs_f64() * 1000.0 + }) + .collect(); + times.sort_by(f64::total_cmp); + println!("{name}, one 512-cell tile, one thread: {:.1} ms median ({:.1}..{:.1})", times[2], times[0], times[4]); +} +fn photo_time(name: &str, r: &dyn TileRunner, raw: &lightcraft_raw::Normalized, layout: Layout, p: &Params) -> Rgb32f { + let pool = rayon::ThreadPoolBuilder::new().num_threads(p.parallel).build().unwrap(); + let run = || pool.install(|| denoise_bayer(&raw.data, raw.width, raw.height, layout, r, p, &Control::default())).unwrap(); + drop(run()); // warm the workers and their workspace cache + let mut times = Vec::new(); + let mut answer = Rgb32f { width: 0, height: 0, data: Vec::new() }; + for i in 0..3 { + let t = Instant::now(); + answer = run(); + let secs = t.elapsed().as_secs_f64(); + times.push(secs); + println!("{name}, whole model stage, {} workers, run {i}: {secs:.3} s", p.parallel); + } + times.sort_by(f64::total_cmp); + println!("{name}, whole model stage median {:.3} s", times[1]); + answer +} +#[test] +#[ignore = "needs LC_DENOISE_MODEL and LC_DENOISE_RAW; release timing and real-model accuracy"] +fn denoise_stage_compares_cpu_and_gpu() { + let model = std::env::var_os("LC_DENOISE_MODEL").expect("set LC_DENOISE_MODEL"); + let folder = std::env::var_os("LC_DENOISE_RAW").expect("set LC_DENOISE_RAW to public CC0 corpus/raw"); + let name = std::env::var("LC_DENOISE_ONLY").unwrap_or_else(|_| "arw-sony-a7m3-compressed.arw".into()); + let bytes = std::fs::read(Path::new(&folder).join(&name)).unwrap(); + let raw = lightcraft_raw::decode(&bytes).unwrap().normalized().unwrap(); + let cfa = raw.cfa.as_ref().unwrap(); + assert!(cfa.is_bayer()); + let layout = Layout::from_cell([cfa.color_at(0, 0), cfa.color_at(1, 0), cfa.color_at(0, 1), cfa.color_at(1, 1)]).unwrap(); + let m = lightcraft_denoise::known::find("rawnind-bayer").unwrap().manifest; + let cpu = CpuRunner::load(Path::new(&model), &m).unwrap(); + let threads = std::env::var("LC_DENOISE_PARALLEL") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or_else(|| std::thread::available_parallelism().map_or(1, |n| n.get()).clamp(1, 16)); + let p = Params { tile: 512, overlap: 64, gain: m.gain, clip: Some(1.0), parallel: threads }; + println!( + "{name}: {}x{}, {:.2} MP, {} tiles, {} CPU workers, {:.1} MB/workspace, {:.1} GMAC/tile", + raw.width, + raw.height, + raw.width as f64 * raw.height as f64 / 1e6, + lightcraft_denoise::run::tile_count(raw.width, raw.height, layout, 512, 64), + threads, + cpu.bytes_per_tile() as f64 / 1e6, + cpu.net().macs(512) as f64 / 1e9 + ); + let x = check_tile(512); + let want = cpu.run(&x).unwrap(); + tile_time("pure Rust", &cpu, &x); + let done = photo_time("pure Rust", &cpu, &raw, layout, &p); + + match lightcraft_gpu::nn::runner(cpu.net(), 512) { + Ok(gpu) => { + println!("GPU: {}", gpu.adapter()); + println!("CPU/check tile vs GPU: {:e}", agree(&want, &gpu.run(&x).unwrap())); + tile_time("GPU", &gpu, &x); + let gp = Params { parallel: 4, ..p }; + let answer = photo_time("GPU", &gpu, &raw, layout, &gp); + let a: Vec<_> = done.data.iter().flatten().copied().collect(); + let b: Vec<_> = answer.data.iter().flatten().copied().collect(); + println!("CPU/whole photo vs GPU: {:e}", agree(&a, &b)); + } + Err(e) => println!("GPU unavailable: {e}"), + } +} diff --git a/crates/engine/tests/denoise_real.rs b/crates/engine/tests/denoise_real.rs new file mode 100644 index 00000000..8d858175 --- /dev/null +++ b/crates/engine/tests/denoise_real.rs @@ -0,0 +1,117 @@ +//! Opt-in end-to-end check of AI denoise with the real RawNIND Bayer model, through the engine: install it, make a real +//! raw photo's picture in the background, render and export it with the Amount slider, and see what a file that is not +//! Bayer says. +//! +//! ```text +//! LC_DENOISE_MODEL= LC_DENOISE_RAW= \ +//! cargo test --release -p lightcraft-engine --features denoise --test denoise_real -- --ignored --nocapture +//! ``` +//! +//! `LC_DENOISE_RAW` is a folder of CC0 raw files (`cargo xtask corpus --download` → `corpus/raw`); every file in it is +//! tried. Optional: `LC_DENOISE_ONLY` (a file name to try alone), `LC_DENOISE_RUN_ON=cpu|gpu|auto` (where the model runs; `cpu` to compare with the card). +#![cfg(feature = "denoise")] + +use std::path::{Path, PathBuf}; +use std::time::{Duration, Instant}; + +use lightcraft_engine::Session; +use serde_json::{Value, json}; + +fn temp(tag: &str) -> PathBuf { + let d = std::env::temp_dir().join(format!("lc-denoise-real-{tag}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&d); + std::fs::create_dir_all(&d).unwrap(); + d +} + +/// Mean absolute second difference of the luma: the fine detail and the noise of a picture. +fn roughness(img: &lightcraft_raster::Rgba8) -> f64 { + let (w, h) = (img.width, img.height); + let luma = |x: usize, y: usize| { + let p = img.data[y * w + x]; + 0.299 * f64::from(p[0]) + 0.587 * f64::from(p[1]) + 0.114 * f64::from(p[2]) + }; + let (mut sum, mut n) = (0.0, 0usize); + for y in 1..h - 1 { + for x in 1..w - 1 { + sum += (4.0 * luma(x, y) - luma(x - 1, y) - luma(x + 1, y) - luma(x, y - 1) - luma(x, y + 1)).abs(); + n += 1; + } + } + sum / n.max(1) as f64 +} + +fn files(dir: &Path) -> Vec { + let only = std::env::var("LC_DENOISE_ONLY").ok(); + let mut v: Vec = std::fs::read_dir(dir) + .unwrap() + .flatten() + .map(|e| e.path()) + .filter(|p| p.is_file() && only.as_deref().is_none_or(|o| p.file_name().is_some_and(|n| n.to_string_lossy() == o))) + .collect(); + v.sort(); + v +} + +#[test] +#[ignore = "needs the real model and raw files (see the file's header)"] +fn the_real_model_denoises_real_raws_through_the_engine() { + let (Some(model), Some(raws)) = (std::env::var_os("LC_DENOISE_MODEL"), std::env::var_os("LC_DENOISE_RAW")) else { + panic!("set LC_DENOISE_MODEL and LC_DENOISE_RAW"); + }; + let dir = temp("e2e"); + let mut s = Session::new().with_fs(); + s.open_library(dir.join("library"), false).unwrap(); + s.set_denoise_models_dir(Some(dir.join("models"))); + let t = Instant::now(); + let r = s.execute("denoise.models.install", &json!({"path": model.to_string_lossy(), "acknowledged": true})).unwrap(); + println!("installed {} in {:.1} s: self-test {}", r["installed"]["id"], t.elapsed().as_secs_f64(), r["installed"]["accepted"]["selfTest"]); + assert_eq!(r["installed"]["id"], "rawnind-bayer"); + if let Ok(run_on) = std::env::var("LC_DENOISE_RUN_ON") { + s.execute("denoise.settings", &json!({"runOn": run_on})).unwrap(); + } + + for file in files(Path::new(&raws)) { + let name = file.file_name().unwrap().to_string_lossy().into_owned(); + let imported = s.execute("library.import", &json!({"paths": [file.to_string_lossy()]})).unwrap(); + let Some(id) = imported["imported"].as_array().and_then(|a| a.first()).and_then(Value::as_u64) else { + println!("{name}: not imported ({imported})"); + continue; + }; + s.execute("library.select", &json!({"ids": [id]})).unwrap(); + s.execute("develop.set", &json!({"control": "enhance.denoise", "value": 100})).unwrap(); + let photo = lightcraft_engine::catalog::PhotoId(id); + let started = Instant::now(); + let mut state = Value::Null; + while started.elapsed() < Duration::from_secs(600) { + s.execute("denoise.pump", &json!({"pace": "full"})).unwrap(); + state = s.execute("denoise.status", &json!({"id": id})).unwrap()["photo"].clone(); + if matches!(state["state"].as_str(), Some("ready" | "failed")) { + break; + } + std::thread::sleep(Duration::from_millis(250)); + } + println!("{name}: {} in {:.1} s", state, started.elapsed().as_secs_f64()); + if state["state"] == "failed" { + // a file that is not a Bayer raw says why, and its export is the plain one + assert_eq!(state["unsupported"], true, "{name}: {state}"); + continue; + } + assert_eq!(state["state"], "ready", "{name}: {state}"); + let denoised = s.render_now(photo, 1600, 1600).unwrap().image; + s.execute("develop.set", &json!({"control": "enhance.denoise", "value": 0})).unwrap(); + let plain = s.render_now(photo, 1600, 1600).unwrap().image; + let (a, b) = (roughness(&plain), roughness(&denoised)); + println!("{name}: roughness at 1600 px: plain {a:.3}, denoised {b:.3} ({:+.1} %)", 100.0 * (b - a) / a); + assert!(b < a, "{name}: the denoised picture is not smoother"); + // the export is made from the same picture, without asking the queue + s.execute("develop.set", &json!({"control": "enhance.denoise", "value": 100})).unwrap(); + let o = lightcraft_engine::export::ExportOptions::from_json(&json!({"format": "jpeg", "width": 1200, "height": 1200})); + let e = lightcraft_engine::export::export_photo(&mut s, photo, &o, 1).unwrap(); + assert!(e.bytes.len() > 1000); + } + let status = s.execute("denoise.status", &json!({})).unwrap(); + println!("device: {}", status["device"]); + println!("cache: {}", status["cache"]); + let _ = std::fs::remove_dir_all(&dir); +} diff --git a/crates/faces/Cargo.toml b/crates/faces/Cargo.toml new file mode 100644 index 00000000..112967ce --- /dev/null +++ b/crates/faces/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "lightcraft-faces" +version.workspace = true +edition.workspace = true +license.workspace = true +description = "LightCraft faces (L1): bounded ONNX CPU inference, face-model manifests, validation and detection" + +[lints] +workspace = true + +[features] +default = [] + +[dependencies] +ndarray = { version = "=0.17.2", default-features = false, features = ["std"] } +# Rust intrinsics only, CPU-dispatched; no BLAS, cc/assembly or nested threading. +matrixmultiply = { version = "=0.3.11", default-features = false, features = ["std", "avx512"] } +web-time = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +sha2 = { workspace = true } +thiserror = { workspace = true } + +[dev-dependencies] +proptest = { workspace = true } diff --git a/crates/faces/src/align.rs b/crates/faces/src/align.rs new file mode 100644 index 00000000..4982325a --- /dev/null +++ b/crates/faces/src/align.rs @@ -0,0 +1,273 @@ +//! Cutting a face out of a photo the way recognition models expect it. +//! +//! With the five landmarks the detector finds (eyes, nose tip, mouth corners), the face is turned and scaled +//! (a similarity transform, found by least squares) so those points land on the standard 112 × 112 +//! ArcFace template. Without landmarks (a box from an XMP sidecar) a padded square around the box is used +//! instead. Both warp with supersampling, so shrinking a large face does not alias. +//! +//! Pure geometry on 8-bit RGB: no image types, no panics (every access is checked). + +/// The five template points of a 112 × 112 aligned face, image-left eye first: eye, eye, nose tip, mouth +/// corner, mouth corner. (The standard ArcFace template.) +pub const TEMPLATE_112: [(f32, f32); 5] = [(38.2946, 51.6963), (73.5318, 51.5014), (56.0252, 71.7366), (41.5493, 92.3655), (70.7299, 92.2041)]; + +/// A borrowed 8-bit RGB picture. +#[derive(Clone, Copy)] +pub struct Rgb<'a> { + pub data: &'a [u8], + pub width: usize, + pub height: usize, +} + +impl Rgb<'_> { + fn valid(&self) -> bool { + self.width > 0 && self.height > 0 && self.width.checked_mul(self.height).and_then(|n| n.checked_mul(3)) == Some(self.data.len()) + } + + /// The pixel at integer coordinates, clamped to the picture. + fn at(&self, x: isize, y: isize) -> [f32; 3] { + let (x, y) = (x.clamp(0, self.width as isize - 1) as usize, y.clamp(0, self.height as isize - 1) as usize); + let i = (y * self.width + x) * 3; + match self.data.get(i..i + 3) { + Some(p) => [f32::from(p[0]), f32::from(p[1]), f32::from(p[2])], + None => [0.0; 3], + } + } + + /// Bilinear sample at a fractional position (pixel centres at integers). + fn sample(&self, x: f32, y: f32) -> [f32; 3] { + let (x0, y0) = (x.floor(), y.floor()); + let (fx, fy) = (x - x0, y - y0); + let (x0, y0) = (x0 as isize, y0 as isize); + let (a, b, c, d) = (self.at(x0, y0), self.at(x0 + 1, y0), self.at(x0, y0 + 1), self.at(x0 + 1, y0 + 1)); + let mut out = [0.0; 3]; + for k in 0..3 { + let top = a[k] * (1.0 - fx) + b[k] * fx; + let bottom = c[k] * (1.0 - fx) + d[k] * fx; + out[k] = top * (1.0 - fy) + bottom * fy; + } + out + } +} + +/// `q = a · p + t` with complex `a` (scale and rotation) and `t`, mapping source positions to output positions. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Similarity { + pub a: (f32, f32), + pub t: (f32, f32), +} + +impl Similarity { + pub fn scale(&self) -> f32 { + self.a.0.hypot(self.a.1) + } + + /// Where an output position `q` comes from in the source. + fn invert(&self, q: (f32, f32)) -> Option<(f32, f32)> { + let n = self.a.0 * self.a.0 + self.a.1 * self.a.1; + if !(n.is_finite() && n > 1e-12) { + return None; + } + let (dx, dy) = (q.0 - self.t.0, q.1 - self.t.1); + Some(((dx * self.a.0 + dy * self.a.1) / n, (dy * self.a.0 - dx * self.a.1) / n)) + } +} + +/// The least-squares similarity transform taking `from[i]` to `to[i]` (no reflection). `None` when the points +/// do not determine one (all the same point, or non-finite). +pub fn fit_similarity(from: &[(f32, f32)], to: &[(f32, f32)]) -> Option { + if from.len() != to.len() || from.len() < 2 { + return None; + } + let n = from.len() as f32; + let mean = |pts: &[(f32, f32)]| (pts.iter().map(|p| p.0).sum::() / n, pts.iter().map(|p| p.1).sum::() / n); + let (mp, mq) = (mean(from), mean(to)); + // a = Σ conj(p') q' / Σ |p'|² (complex numbers) + let (mut re, mut im, mut norm) = (0.0f32, 0.0f32, 0.0f32); + for (p, q) in from.iter().zip(to) { + let (px, py, qx, qy) = (p.0 - mp.0, p.1 - mp.1, q.0 - mq.0, q.1 - mq.1); + re += px * qx + py * qy; + im += px * qy - py * qx; + norm += px * px + py * py; + } + if !(norm.is_finite() && norm > 1e-9) { + return None; + } + let a = (re / norm, im / norm); + let t = (mq.0 - (a.0 * mp.0 - a.1 * mp.1), mq.1 - (a.0 * mp.1 + a.1 * mp.0)); + (a.0.is_finite() && a.1.is_finite() && t.0.is_finite() && t.1.is_finite()).then_some(Similarity { a, t }) +} + +/// Fill an `out_w` × `out_h` picture by taking, for each output pixel, the source under it, averaging +/// enough sub-samples that shrinking does not alias. +fn warp(img: &Rgb, m: &Similarity, out_w: usize, out_h: usize) -> Option> { + if !img.valid() || out_w == 0 || out_h == 0 || out_w > 4096 || out_h > 4096 { + return None; + } + if ![m.a.0, m.a.1, m.t.0, m.t.1].iter().all(|v| v.is_finite()) { + return None; + } + // how many source pixels one output pixel spans, per side + let span = 1.0 / m.scale(); + let n = if span.is_finite() { (span.ceil() as usize).clamp(1, 4) } else { 1 }; + let mut out = vec![0u8; out_w * out_h * 3]; + for (i, px) in out.as_chunks_mut::<3>().0.iter_mut().enumerate() { + let (ox, oy) = ((i % out_w) as f32, (i / out_w) as f32); + let mut acc = [0.0f32; 3]; + for sy in 0..n { + for sx in 0..n { + let q = (ox + (sx as f32 + 0.5) / n as f32 - 0.5, oy + (sy as f32 + 0.5) / n as f32 - 0.5); + let (px_, py_) = m.invert(q)?; + let s = img.sample(px_, py_); + acc.iter_mut().zip(s).for_each(|(a, v)| *a += v); + } + } + let k = (n * n) as f32; + for (o, a) in px.iter_mut().zip(acc) { + *o = (a / k).round().clamp(0.0, 255.0) as u8; + } + } + Some(out) +} + +/// The face aligned to the ArcFace template, `out_w` × `out_h` (the template is for 112 × 112 and is scaled for +/// other sizes). `landmarks` are in pixels of `img`. `None` when the landmarks cannot define a transform. +pub fn align_to_template(img: &Rgb, landmarks: &[(f32, f32); 5], out_w: usize, out_h: usize) -> Option> { + let (sx, sy) = (out_w as f32 / 112.0, out_h as f32 / 112.0); + let template: Vec<(f32, f32)> = TEMPLATE_112.iter().map(|(x, y)| (x * sx, y * sy)).collect(); + let m = fit_similarity(landmarks, &template)?; + // a face smaller than a few pixels, or absurdly large, is not a face we can use + let scale = m.scale(); + if !(scale.is_finite() && (0.01..=100.0).contains(&scale)) { + return None; + } + warp(img, &m, out_w, out_h) +} + +/// A square around a face box (pixels of `img`), `pad` times its larger side, as `out_w` × `out_h`. +pub fn crop_box(img: &Rgb, x0: f32, y0: f32, x1: f32, y1: f32, pad: f32, out_w: usize, out_h: usize) -> Option> { + if ![x0, y0, x1, y1, pad].iter().all(|v| v.is_finite()) { + return None; + } + let side = (x1 - x0).abs().max((y1 - y0).abs()) * pad; + if !(side.is_finite() && side >= 2.0) { + return None; + } + let (cx, cy) = ((x0 + x1) / 2.0, (y0 + y1) / 2.0); + let s = out_w.min(out_h) as f32 / side; + let m = Similarity { a: (s, 0.0), t: (out_w as f32 / 2.0 - s * cx, out_h as f32 / 2.0 - s * cy) }; + warp(img, &m, out_w, out_h) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn close(a: (f32, f32), b: (f32, f32), tol: f32) -> bool { + (a.0 - b.0).abs() <= tol && (a.1 - b.1).abs() <= tol + } + + fn apply(m: &Similarity, p: (f32, f32)) -> (f32, f32) { + (m.a.0 * p.0 - m.a.1 * p.1 + m.t.0, m.a.0 * p.1 + m.a.1 * p.0 + m.t.1) + } + + #[test] + fn a_known_similarity_is_recovered() { + // rotate 20 degrees, scale 0.7, translate + let (s, th) = (0.7f32, 20f32.to_radians()); + let truth = Similarity { a: (s * th.cos(), s * th.sin()), t: (13.0, -5.0) }; + let from: Vec<(f32, f32)> = TEMPLATE_112.to_vec(); + let to: Vec<(f32, f32)> = from.iter().map(|p| apply(&truth, *p)).collect(); + let m = fit_similarity(&from, &to).unwrap(); + assert!(close(m.a, truth.a, 1e-4) && close(m.t, truth.t, 1e-3), "{m:?} vs {truth:?}"); + // noisy points still give the best fit (the residual is small) + let noisy: Vec<(f32, f32)> = to.iter().enumerate().map(|(i, p)| (p.0 + (i as f32 - 2.0) * 0.2, p.1 - (i as f32 - 2.0) * 0.1)).collect(); + let m2 = fit_similarity(&from, &noisy).unwrap(); + assert!(from.iter().zip(&noisy).all(|(p, q)| close(apply(&m2, *p), *q, 1.0))); + } + + #[test] + fn degenerate_points_give_nothing() { + let same = vec![(5.0, 5.0); 5]; + assert!(fit_similarity(&same, &TEMPLATE_112).is_none()); + assert!(fit_similarity(&TEMPLATE_112, &TEMPLATE_112[..3]).is_none()); + assert!(fit_similarity(&[(0.0, 0.0)], &[(1.0, 1.0)]).is_none()); + let nan = vec![(f32::NAN, 0.0); 5]; + assert!(fit_similarity(&nan, &TEMPLATE_112).is_none()); + let inf = vec![(f32::INFINITY, 0.0), (1.0, 1.0), (2.0, 2.0), (3.0, 3.0), (4.0, 4.0)]; + assert!(fit_similarity(&inf, &TEMPLATE_112).is_none()); + } + + /// A picture with a bright dot at each of five landmarks, the landmarks being the template moved by a known + /// similarity: aligning must put the dots back on the template. + #[test] + fn alignment_puts_the_landmarks_on_the_template() { + let (w, h) = (400usize, 300usize); + let (s, th) = (1.7f32, -15f32.to_radians()); + let place = Similarity { a: (s * th.cos(), s * th.sin()), t: (150.0, 60.0) }; + let landmarks: Vec<(f32, f32)> = TEMPLATE_112.iter().map(|p| apply(&place, *p)).collect(); + let mut data = vec![0u8; w * h * 3]; + for l in &landmarks { + for dy in -3i32..=3 { + for dx in -3i32..=3 { + let (x, y) = ((l.0 as i32 + dx) as usize, (l.1 as i32 + dy) as usize); + if x < w && y < h { + data[(y * w + x) * 3] = 255; + } + } + } + } + let img = Rgb { data: &data, width: w, height: h }; + let lm: [(f32, f32); 5] = [landmarks[0], landmarks[1], landmarks[2], landmarks[3], landmarks[4]]; + let out = align_to_template(&img, &lm, 112, 112).unwrap(); + assert_eq!(out.len(), 112 * 112 * 3); + // the red channel's centre of mass near each template point is on the point + for t in TEMPLATE_112 { + let (mut sx, mut sy, mut sw) = (0.0f32, 0.0f32, 0.0f32); + for y in (t.1 as usize).saturating_sub(8)..(t.1 as usize + 9).min(112) { + for x in (t.0 as usize).saturating_sub(8)..(t.0 as usize + 9).min(112) { + let v = f32::from(out[(y * 112 + x) * 3]); + sx += v * x as f32; + sy += v * y as f32; + sw += v; + } + } + assert!(sw > 100.0, "no dot near {t:?}"); + assert!(close((sx / sw, sy / sw), t, 1.5), "dot at {:?}, template {t:?}", (sx / sw, sy / sw)); + } + } + + #[test] + fn box_crops_are_centred_and_bad_inputs_are_refused() { + let (w, h) = (200usize, 100usize); + let mut data = vec![0u8; w * h * 3]; + for y in 40..60 { + for x in 90..110 { + data[(y * w + x) * 3 + 1] = 255; + } + } + let img = Rgb { data: &data, width: w, height: h }; + let out = crop_box(&img, 90.0, 40.0, 110.0, 60.0, 2.0, 64, 64).unwrap(); + // the 20 px block fills the middle half of the 64 px output + let at = |x: usize, y: usize| out[(y * 64 + x) * 3 + 1]; + assert!(at(32, 32) > 250 && at(20, 20) > 250 && at(44, 44) > 250 && at(5, 5) < 5 && at(60, 32) < 5); + assert!(crop_box(&img, 5.0, 5.0, 5.0, 5.0, 2.0, 64, 64).is_none(), "an empty box"); + assert!(crop_box(&img, f32::NAN, 0.0, 10.0, 10.0, 2.0, 64, 64).is_none()); + assert!(crop_box(&img, 0.0, 0.0, 10.0, 10.0, 2.0, 0, 64).is_none()); + assert!(crop_box(&img, 0.0, 0.0, 10.0, 10.0, 2.0, 100_000, 64).is_none()); + let bad = Rgb { data: &data[..10], width: w, height: h }; + assert!(crop_box(&bad, 0.0, 0.0, 10.0, 10.0, 2.0, 64, 64).is_none(), "wrong data length"); + } + + #[test] + fn faces_off_the_edge_and_tiny_or_huge_landmarks_do_not_panic() { + let data = vec![80u8; 50 * 50 * 3]; + let img = Rgb { data: &data, width: 50, height: 50 }; + // a box mostly outside the picture is padded by clamping, not an error + assert!(crop_box(&img, -30.0, -30.0, 20.0, 20.0, 1.5, 32, 32).is_some()); + let far: [(f32, f32); 5] = [(1e6, 1e6), (1e6 + 5.0, 1e6), (1e6 + 2.0, 1e6 + 3.0), (1e6, 1e6 + 6.0), (1e6 + 5.0, 1e6 + 6.0)]; + let _ = align_to_template(&img, &far, 112, 112); + let tiny: [(f32, f32); 5] = [(1.0, 1.0), (1.01, 1.0), (1.0, 1.01), (1.01, 1.01), (1.005, 1.005)]; + assert!(align_to_template(&img, &tiny, 112, 112).is_none(), "landmarks a hundredth of a pixel apart"); + } +} diff --git a/crates/faces/src/catalog.rs b/crates/faces/src/catalog.rs new file mode 100644 index 00000000..07a20622 --- /dev/null +++ b/crates/faces/src/catalog.rs @@ -0,0 +1,228 @@ +//! Your own list of models to download: `catalog.json` in the models folder. +//! +//! LightCraft's built-in list only holds models whose terms let the project point at them. Anything else (a stronger +//! recogniser whose weights are for research use, a model you trained, a mirror you trust) can be added by you: each entry +//! is a model manifest (the same fields as a `face-model.json`) plus the address to fetch it from, and it then gets the same +//! Download button as a built-in one, with its licence notice shown before anything is fetched. LightCraft ships and links to +//! none of them; what the file says, and whether you may use the weights, is yours to check. +//! +//! ```json +//! { "models": [ { +//! "id": "my-recogniser", "name": "My recogniser", "version": "1", "role": "embedder", +//! "url": "https://example.org/weights/recogniser.onnx", +//! "sha256": "…64 lowercase hex digits…", "sizeBytes": 166000000, +//! "licence": { "name": "Research use only", "commercial": "no", "notice": "Not for commercial use." }, +//! "provenance": "Trained on …", +//! "output": { "kind": "embedding", "dim": 512 }, +//! "thresholds": { "matchCosine": 0.4 }, +//! "speed": 1.7 +//! } ] } +//! ``` +//! +//! `speed` (optional) is how many times faster the model is than a ResNet-100, shown in Settings instead of timings. +//! `input` may be left out (112 × 112 RGB, `(x − 127.5) / 127.5`, the ArcFace convention). The file is untrusted: it is read +//! with a size limit, every entry is validated, and a bad entry is reported and skipped, never a reason to refuse the rest. +//! The `sha256` and `sizeBytes` are required, since they are what a downloaded file is checked against. + +use serde::Deserialize; +use serde_json::Value; + +use crate::known::{self, Download}; +use crate::manifest::{self, ModelManifest, Role}; + +/// Largest `catalog.json` we read. +pub const MAX_CATALOG_BYTES: usize = 256 * 1024; +/// Most entries we take from it. +pub const MAX_ENTRIES: usize = 50; + +/// One model from the catalog: its manifest, where to fetch it and what to call the file meanwhile. +#[derive(Clone, Debug, PartialEq)] +pub struct Entry { + pub manifest: ModelManifest, + pub url: String, + pub file_name: String, +} + +impl Entry { + /// What the download machinery needs. (`sha256` and `size_bytes` are present: [`parse`] requires them.) + pub fn download(&self) -> Option { + Some(Download { + id: self.manifest.id.clone(), + url: self.url.clone(), + file_name: self.file_name.clone(), + size_bytes: self.manifest.size_bytes?, + sha256: self.manifest.sha256.clone()?, + }) + } +} + +/// What a catalog file holds. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct Catalog { + pub entries: Vec, + /// One plain sentence for each entry (or the file) that could not be used. + pub errors: Vec, +} + +#[derive(Deserialize)] +struct Raw { + url: String, + #[serde(flatten)] + manifest: ModelManifest, +} + +/// A file name for the download, from the address: letters, digits, `.`, `_`, `-`, ending in `.onnx`. +fn file_name_of(url: &str) -> String { + let last = url.split(['?', '#']).next().unwrap_or("").rsplit('/').next().unwrap_or(""); + let clean: String = last.chars().filter(|c| c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '-')).take(80).collect(); + if clean.ends_with(".onnx") && clean.len() > 5 && !clean.starts_with('.') { clean } else { "model.onnx".to_string() } +} + +fn entry(v: &Value) -> Result { + let raw: Raw = serde_json::from_value(v.clone()).map_err(|e| e.to_string())?; + manifest::validate(&raw.manifest).map_err(|e| e.to_string())?; + let m = &raw.manifest; + if m.role != Role::Embedder { + return Err("only face recognition models (role \"embedder\") can be downloaded".into()); + } + if m.sha256.is_none() || m.size_bytes.is_none() { + return Err("needs `sha256` and `sizeBytes`: a download is checked against them".into()); + } + if !(raw.url.len() <= 500 && raw.url.starts_with("https://") && !raw.url.chars().any(|c| c.is_control() || c.is_whitespace())) { + return Err("`url` must be an https:// address of at most 500 characters, without spaces".into()); + } + if known::all().iter().any(|k| k.id == m.id) { + return Err("that id belongs to a model LightCraft already knows".into()); + } + if m.sha256.as_deref().is_some_and(|h| known::lookup(h).is_some()) { + return Err("LightCraft already knows a model with that file".into()); + } + Ok(Entry { file_name: file_name_of(&raw.url), url: raw.url, manifest: raw.manifest }) +} + +/// Read a catalog file's bytes: the usable entries, and why the others were left out. +pub fn parse(bytes: &[u8]) -> Catalog { + let mut out = Catalog::default(); + if bytes.len() > MAX_CATALOG_BYTES { + out.errors.push(format!("catalog.json is larger than {} KB: not read", MAX_CATALOG_BYTES / 1024)); + return out; + } + // Notepad and older PowerShell begin a UTF-8 file with a byte-order mark + let bytes = bytes.strip_prefix(&[0xEF, 0xBB, 0xBF]).unwrap_or(bytes); + let value: Value = match serde_json::from_slice(bytes) { + Ok(v) => v, + Err(e) => { + out.errors.push(format!("catalog.json is not valid JSON: {e}")); + return out; + } + }; + let Some(models) = value.get("models").and_then(Value::as_array) else { + out.errors.push("catalog.json needs a \"models\" list".into()); + return out; + }; + if models.len() > MAX_ENTRIES { + out.errors.push(format!("only the first {MAX_ENTRIES} models of the catalog are read")); + } + for (i, v) in models.iter().take(MAX_ENTRIES).enumerate() { + let label = v + .get("id") + .and_then(Value::as_str) + .map_or_else(|| format!("model {}", i + 1), |id| format!("`{}`", id.chars().take(40).collect::())); + match entry(v) { + Ok(e) if out.entries.iter().any(|x| x.manifest.id == e.manifest.id) => { + out.errors.push(format!("{label}: listed twice, the first is used")) + } + Ok(e) => out.entries.push(e), + Err(why) => out.errors.push(format!("{label}: {why}")), + } + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn good(id: &str) -> Value { + json!({ + "id": id, "name": "Test model", "version": "1", "role": "embedder", + "url": "https://example.org/w/recogniser-r50.onnx?download=1", + "sha256": "ab".repeat(32), "sizeBytes": 1234, + "licence": {"name": "Research only", "commercial": "no", "notice": "Not for commercial use."}, + "output": {"kind": "embedding", "dim": 512}, "thresholds": {"matchCosine": 0.4}, "speed": 1.7, + }) + } + + fn cat(models: Vec) -> Catalog { + parse(json!({"models": models}).to_string().as_bytes()) + } + + #[test] + fn a_manifest_with_an_address_is_a_downloadable_model() { + let c = cat(vec![good("my-model")]); + assert_eq!(c.errors, Vec::::new()); + let e = &c.entries[0]; + assert_eq!((e.manifest.id.as_str(), e.file_name.as_str()), ("my-model", "recogniser-r50.onnx"), "the file name comes from the address"); + assert_eq!(e.manifest.licence.commercial, crate::Commercial::No, "a non-commercial model is the user's to add"); + let d = e.download().unwrap(); + assert_eq!((d.size_bytes, d.sha256.len(), d.url.as_str()), (1234, 64, "https://example.org/w/recogniser-r50.onnx?download=1")); + assert_eq!(e.manifest.input, crate::InputSpec::default(), "the input defaults to the ArcFace convention"); + assert_eq!(e.manifest.speed, Some(1.7), "the speed ratio is the user's to state"); + let mut nonsense = good("nonsense-speed"); + nonsense["speed"] = json!(-3); + let c = cat(vec![nonsense]); + assert!(c.entries.is_empty() && c.errors.iter().any(|e| e.contains("speed")), "{:?}", c.errors); + } + + #[test] + fn a_file_saved_with_a_byte_order_mark_still_reads() { + let mut bytes = vec![0xEF, 0xBB, 0xBF]; + bytes.extend_from_slice(json!({"models": [good("with-bom")]}).to_string().as_bytes()); + let c = parse(&bytes); + assert_eq!((c.entries.len(), c.errors.len()), (1, 0), "{:?}", c.errors); + } + + #[test] + fn bad_entries_are_reported_and_the_rest_still_load() { + let mut no_hash = good("no-hash"); + no_hash.as_object_mut().unwrap().remove("sha256"); + let mut http = good("plain-http"); + http["url"] = json!("http://example.org/m.onnx"); + let mut spaces = good("spaces"); + spaces["url"] = json!("https://example.org/a b.onnx"); + let mut detector = good("a-detector"); + detector["role"] = json!("detector"); + detector["output"] = json!({"kind": "detector", "decoder": "yunet-v2"}); + let mut bad_id = good("Bad/Id"); + bad_id["id"] = json!("../escape"); + let builtin = good("sface-2021dec"); + let mut known_file = good("known-file"); + known_file["sha256"] = json!(known::SFACE_SHA256); + let c = cat(vec![good("fine"), no_hash, http, spaces, detector, bad_id, builtin, known_file, good("fine"), json!(7), json!({"id": "x"})]); + assert_eq!(c.entries.iter().map(|e| e.manifest.id.as_str()).collect::>(), ["fine"]); + assert_eq!(c.errors.len(), 10, "{:#?}", c.errors); + assert!(c.errors.iter().any(|e| e.contains("`no-hash`") && e.contains("sha256"))); + assert!(c.errors.iter().any(|e| e.contains("https"))); + assert!(c.errors.iter().any(|e| e.contains("already knows"))); + assert!(c.errors.iter().any(|e| e.contains("listed twice"))); + } + + #[test] + fn a_hostile_or_broken_file_is_an_error_message_not_a_crash() { + for bytes in [&b""[..], b"not json", b"[]", b"{}", b"{\"models\": 5}", b"\xff\xfe\x00", &vec![b' '; MAX_CATALOG_BYTES + 1]] { + let c = parse(bytes); + assert!(c.entries.is_empty() && !c.errors.is_empty(), "{:?}", String::from_utf8_lossy(&bytes[..bytes.len().min(20)])); + } + // more entries than allowed: the extra are ignored, with a note + let many: Vec = (0..MAX_ENTRIES + 20).map(|i| good(&format!("m{i}"))).collect(); + let c = cat(many); + assert_eq!(c.entries.len(), MAX_ENTRIES); + assert!(c.errors.iter().any(|e| e.contains("first"))); + // file names from odd addresses are safe + for url in ["https://x.org/a/b/c.onnx", "https://x.org/", "https://x.org/..%2f..%2fevil.onnx", "https://x.org/q.bin", "https://x.org/.onnx"] { + let got = file_name_of(url); + assert!(!got.contains(['/', '\\']) && got.ends_with(".onnx") && !got.starts_with('.'), "{url} -> {got}"); + } + } +} diff --git a/crates/faces/src/graph.rs b/crates/faces/src/graph.rs new file mode 100644 index 00000000..a5a22d26 --- /dev/null +++ b/crates/faces/src/graph.rs @@ -0,0 +1,442 @@ +//! A full, bounded read of a small ONNX model: its nodes, attributes and weights. +//! +//! The detector and installed recognisers share this reader and [`crate::net`]. Files are capped at +//! 512 MiB, with at most 64 million weight elements. The bytes are still treated as hostile: the +//! number of nodes, weight tensors and elements is capped, every length is checked against its parent, +//! and a weight blob whose size disagrees with its shape is an error. + +use std::collections::HashMap; +use std::io::Cursor; + +use crate::onnx::{ProbeError, Walker, Wire}; + +type Result = std::result::Result; + +const MAX_NODES: usize = 20_000; +const MAX_TENSORS: usize = 20_000; +/// Elements across all weight tensors (4 bytes each as float32). +const MAX_ELEMENTS: usize = 64 * 1024 * 1024; +const MAX_NAME: usize = 4096; +const MAX_LIST: usize = 1_000_000; + +#[derive(Clone, Debug, PartialEq)] +pub enum Attr { + Int(i64), + Float(f32), + Str(String), + Ints(Vec), + Floats(Vec), + /// A kind we do not read (a tensor, a sub-graph…). + Other, +} + +#[derive(Clone, Debug)] +pub struct Node { + pub op: String, + pub inputs: Vec, + pub outputs: Vec, + pub attrs: HashMap, +} + +impl Node { + pub fn int(&self, name: &str) -> Option { + match self.attrs.get(name) { + Some(Attr::Int(v)) => Some(*v), + _ => None, + } + } + pub fn ints(&self, name: &str) -> Option<&[i64]> { + match self.attrs.get(name) { + Some(Attr::Ints(v)) => Some(v), + _ => None, + } + } + pub fn string(&self, name: &str) -> Option<&str> { + match self.attrs.get(name) { + Some(Attr::Str(v)) => Some(v), + _ => None, + } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub enum Data { + F32(Vec), + I64(Vec), +} + +/// A weight tensor. +#[derive(Clone, Debug, PartialEq)] +pub struct Weight { + pub dims: Vec, + pub data: Data, +} + +#[derive(Clone, Debug)] +pub struct Graph { + pub nodes: Vec, + pub weights: HashMap, + /// The inputs the caller feeds (weights some exporters also list as inputs are left out), and their shapes + /// (0 for a dynamic size). + pub inputs: Vec<(String, Vec)>, + pub outputs: Vec, +} + +fn malformed(why: &'static str) -> Result { + Err(ProbeError::Malformed(why)) +} + +fn f32s(raw: &[u8]) -> Vec { + raw.as_chunks::<4>().0.iter().map(|b| f32::from_le_bytes(*b)).collect() +} + +fn i64s(raw: &[u8]) -> Vec { + raw.as_chunks::<8>().0.iter().map(|b| i64::from_le_bytes(*b)).collect() +} + +/// A list of varints in a packed field (`start..end`), appended to `out`. +fn packed_varints(w: &mut Walker>, end: u64, out: &mut Vec) -> Result<()> { + while w.pos()? < end { + if out.len() >= MAX_LIST { + return malformed("a list is too long"); + } + out.push(w.varint()?); + if w.pos()? > end { + return malformed("a packed number runs past its field"); + } + } + Ok(()) +} + +fn attribute(w: &mut Walker>, end: u64) -> Result<(String, Attr)> { + let (mut name, mut kind) = (String::new(), 0u64); + let (mut i, mut f, mut s) = (0i64, 0f32, String::new()); + let (mut ints, mut floats): (Vec, Vec) = (Vec::new(), Vec::new()); + while let Some((n, f_)) = w.next(end)? { + match (n, f_) { + (1, Wire::Len { start, end: e }) => { + name = w.string(start, e)?; + w.seek_to(e)?; + } + (2, Wire::Fixed32(v)) => f = f32::from_bits(v), + (3, Wire::Varint(v)) => i = v as i64, + (4, Wire::Len { start, end: e }) => { + s = w.string(start, e)?; + w.seek_to(e)?; + } + (7, Wire::Len { end: e, .. }) => { + // packed floats + let at = w.pos()?; + let raw = w.bytes(at, e, (MAX_LIST * 4) as u64)?; + if !raw.len().is_multiple_of(4) || floats.len().saturating_add(raw.len() / 4) > MAX_LIST { + return malformed("an attribute float list is invalid or too long"); + } + floats.extend(f32s(&raw)); + } + (7, Wire::Fixed32(v)) => { + if floats.len() >= MAX_LIST { + return malformed("an attribute float list is too long"); + } + floats.push(f32::from_bits(v)); + } + (8, Wire::Len { end: e, .. }) => { + packed_varints(w, e, &mut ints)?; + w.seek_to(e)?; + } + (8, Wire::Varint(v)) => { + if ints.len() >= MAX_LIST { + return malformed("a list is too long"); + } + ints.push(v); + } + (20, Wire::Varint(v)) => kind = v, + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + let ints = ints.into_iter().map(|v| v as i64).collect(); + let value = match kind { + 1 => Attr::Float(f), + 2 => Attr::Int(i), + 3 => Attr::Str(s), + 6 => Attr::Floats(floats), + 7 => Attr::Ints(ints), + _ => Attr::Other, + }; + Ok((name, value)) +} + +fn node(w: &mut Walker>, end: u64) -> Result { + let mut n = Node { op: String::new(), inputs: Vec::new(), outputs: Vec::new(), attrs: HashMap::new() }; + while let Some((num, f)) = w.next(end)? { + match (num, f) { + (1, Wire::Len { start, end: e }) => { + n.inputs.push(w.string(start, e)?); + w.seek_to(e)?; + } + (2, Wire::Len { start, end: e }) => { + n.outputs.push(w.string(start, e)?); + w.seek_to(e)?; + } + (4, Wire::Len { start, end: e }) => { + n.op = w.string(start, e)?; + w.seek_to(e)?; + } + (5, Wire::Len { end: e, .. }) => { + let (k, v) = attribute(w, e)?; + if k.is_empty() || n.attrs.insert(k, v).is_some() { + return malformed("an attribute name is empty or repeated"); + } + w.seek_to(e)?; + } + (7, Wire::Len { start, end: e }) => { + let domain = w.string(start, e)?; + if !domain.is_empty() && domain != "ai.onnx" { + return malformed("custom operator domains are unsupported"); + } + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + if n.inputs.len() > 64 || n.outputs.len() > 64 || n.attrs.len() > 64 { + return malformed("a node is too large"); + } + } + Ok(n) +} + +/// A weight tensor (`TensorProto`): its name, shape and float32 or int64 values. +fn tensor(w: &mut Walker>, end: u64, budget: &mut usize) -> Result<(String, Option)> { + let (mut name, mut dtype) = (String::new(), 0u64); + let mut dims: Vec = Vec::new(); + let mut raw: Option<&[u8]> = None; + let (mut float_data, mut int_data): (Vec, Vec) = (Vec::new(), Vec::new()); + let max_bytes = (*budget as u64).saturating_mul(8); + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Len { end: e, .. }) => { + packed_varints(w, e, &mut dims)?; + w.seek_to(e)?; + } + (1, Wire::Varint(v)) => dims.push(v), + (2, Wire::Varint(v)) => dtype = v, + (4, Wire::Len { end: e, .. }) => { + let at = w.pos()?; + let bytes = w.bytes(at, e, max_bytes)?; + if !bytes.len().is_multiple_of(4) || float_data.len().saturating_add(bytes.len() / 4) > *budget { + return malformed("invalid or excessive packed float weights"); + } + float_data.extend(f32s(&bytes)); + } + (4, Wire::Fixed32(v)) => { + if float_data.len() >= *budget { + return malformed("too many float weights"); + } + float_data.push(f32::from_bits(v)); + } + (7, Wire::Varint(v)) => { + if int_data.len() >= *budget { + return malformed("too many integer weights"); + } + int_data.push(v); + } + (7, Wire::Len { end: e, .. }) => { + packed_varints(w, e, &mut int_data)?; + w.seek_to(e)?; + } + (8, Wire::Len { start, end: e }) => { + name = w.string(start, e)?; + w.seek_to(e)?; + } + (9, Wire::Len { start, end: e }) => { + if raw.is_some() || e.saturating_sub(start) > max_bytes { + return malformed("repeated or excessive raw weights"); + } + let from = usize::try_from(start).map_err(|_| ProbeError::Malformed("weight offset overflow"))?; + let to = usize::try_from(e).map_err(|_| ProbeError::Malformed("weight offset overflow"))?; + raw = Some((*w.r.get_ref()).get(from..to).ok_or(ProbeError::Malformed("short raw weights"))?); + w.seek_to(e)?; + } + (13, Wire::Len { .. }) | (14, Wire::Varint(1..)) => return malformed("external tensor data is unsupported"), + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + if dims.len() > 8 { + return malformed("a tensor has too many dimensions"); + } + } + let mut elements: usize = 1; + let mut shape = Vec::new(); + for d in &dims { + let d = usize::try_from(*d).map_err(|_| ProbeError::Malformed("a tensor dimension is too large"))?; + elements = elements.checked_mul(d).ok_or(ProbeError::Malformed("a tensor is too large"))?; + shape.push(d); + } + if elements > *budget { + return malformed("the weights are too large"); + } + *budget -= elements; + if raw.is_some() && (!float_data.is_empty() || !int_data.is_empty()) { + return malformed("ambiguous raw and typed weights"); + } + let data = match dtype { + 1 => { + let v = if let Some(raw) = raw { + if elements.checked_mul(4) != Some(raw.len()) { + return malformed("float weight bytes disagree with shape"); + } + f32s(raw) + } else { + float_data + }; + if v.len() != elements { + return malformed("a weight tensor's size disagrees with its shape"); + } + Some(Data::F32(v)) + } + 7 => { + let v = if let Some(raw) = raw { + if elements.checked_mul(8) != Some(raw.len()) { + return malformed("integer weight bytes disagree with shape"); + } + i64s(raw) + } else { + int_data.into_iter().map(|v| v as i64).collect() + }; + if v.len() != elements { + return malformed("a weight tensor's size disagrees with its shape"); + } + Some(Data::I64(v)) + } + _ => return malformed("only float32 and int64 initializers are supported"), + }; + if name.is_empty() { + return malformed("a weight name is empty"); + } + Ok((name, data.map(|data| Weight { dims: shape, data }))) +} + +fn value_info(w: &mut Walker>, end: u64) -> Result<(String, Vec)> { + let (mut name, mut shape) = (String::new(), Vec::new()); + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Len { start, end: e }) => { + name = w.string(start, e)?; + w.seek_to(e)?; + } + (2, Wire::Len { end: e, .. }) => { + // TypeProto > tensor_type (1) > shape (2) > dim (1) > dim_value (1) + let mut info = crate::onnx::TensorInfo { name: String::new(), elem_type: 0, shape: Vec::new() }; + crate::onnx::type_proto(w, e, &mut info)?; + shape = info + .shape + .iter() + .map(|d| match d { + crate::onnx::Dim::Fixed(v) => { + usize::try_from(*v).map_err(|_| ProbeError::Malformed("a declared dimension exceeds this platform")) + } + crate::onnx::Dim::Dynamic(_) => Ok(0), + }) + .collect::>()?; + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + if name.len() > MAX_NAME { + return malformed("a name is too long"); + } + Ok((name, shape)) +} + +/// Read a bounded float32 ONNX inference model. +pub fn load(bytes: &[u8]) -> Result { + let total = bytes.len() as u64; + if total > 512 * 1024 * 1024 { + return malformed("the model exceeds 512 MiB"); + } + let info = crate::onnx::probe_reader(&mut Cursor::new(bytes))?; + if info.opset.is_none_or(|v| !(10..=17).contains(&v)) { + return malformed("only ONNX default-domain opsets 10 through 17 are supported"); + } + if info.inputs.iter().chain(&info.outputs).any(|v| v.shape.iter().any(|d| matches!(d, crate::onnx::Dim::Fixed(0)))) { + return malformed("a declared input/output has a zero dimension"); + } + if info.inputs.iter().chain(&info.outputs).any(|v| v.elem_type != 1) { + return malformed("model inputs and outputs must be float32"); + } + if total == 0 { + return malformed("the file is empty"); + } + let mut cursor = Cursor::new(bytes); + let mut w = Walker { r: &mut cursor, total, fields: 0 }; + let (mut nodes, mut weights, mut inputs, mut outputs) = (Vec::new(), HashMap::new(), Vec::new(), Vec::new()); + let mut budget = MAX_ELEMENTS; + let mut saw_graph = false; + while let Some((n, f)) = w.next(total)? { + match (n, f) { + (7, Wire::Len { end, .. }) => { + if saw_graph { + return malformed("the model contains multiple graphs"); + } + saw_graph = true; + while let Some((gn, gf)) = w.next(end)? { + match (gn, gf) { + (1, Wire::Len { end: e, .. }) => { + if nodes.len() >= MAX_NODES { + return malformed("too many nodes"); + } + nodes.push(node(&mut w, e)?); + w.seek_to(e)?; + } + (5, Wire::Len { end: e, .. }) => { + if weights.len() >= MAX_TENSORS { + return malformed("too many weight tensors"); + } + let (name, wt) = tensor(&mut w, e, &mut budget)?; + if let Some(wt) = wt + && weights.insert(name, wt).is_some() + { + return malformed("a weight is defined twice"); + } + w.seek_to(e)?; + } + (11, Wire::Len { end: e, .. }) => { + if inputs.len() >= MAX_TENSORS { + return malformed("too many graph inputs"); + } + inputs.push(value_info(&mut w, e)?); + w.seek_to(e)?; + } + (12, Wire::Len { end: e, .. }) => { + if outputs.len() >= 64 { + return malformed("too many graph outputs"); + } + outputs.push(value_info(&mut w, e)?.0); + w.seek_to(e)?; + } + (15, Wire::Len { .. }) => return malformed("sparse initializers are unsupported"), + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + w.seek_to(end)?; + } + (_, Wire::Len { end, .. }) => w.seek_to(end)?, + _ => {} + } + } + if !saw_graph || nodes.is_empty() { + return malformed("no graph"); + } + inputs.retain(|(name, _)| !weights.contains_key(name)); + if inputs.is_empty() || outputs.is_empty() { + return malformed("the graph has no inputs or no outputs"); + } + Ok(Graph { nodes, weights, inputs, outputs }) +} + +#[cfg(test)] +mod tests; diff --git a/crates/faces/src/graph/tests.rs b/crates/faces/src/graph/tests.rs new file mode 100644 index 00000000..38635459 --- /dev/null +++ b/crates/faces/src/graph/tests.rs @@ -0,0 +1,117 @@ +use super::*; +use crate::synthetic::{len_field, value_info_bytes, varint, varint_field}; +fn node() -> Vec { + let mut n = Vec::new(); + for (id, v) in [(1, "x"), (2, "y"), (4, "Identity")] { + len_field(id, v.as_bytes(), &mut n); + } + n +} +fn weight(raw: &[u8]) -> Vec { + let mut t = Vec::new(); + varint_field(1, 1, &mut t); + varint_field(2, 1, &mut t); + len_field(8, b"w", &mut t); + len_field(9, raw, &mut t); + t +} +fn model(node: &[u8], weights: &[Vec]) -> Vec { + let mut g = Vec::new(); + len_field(1, node, &mut g); + for t in weights { + len_field(5, t, &mut g); + } + for (id, name) in [(11, "x"), (12, "y")] { + len_field(id, &value_info_bytes(name, 1, &[Ok(1), Ok(3), Ok(112), Ok(112)]), &mut g); + } + let mut ops = Vec::new(); + varint_field(2, 11, &mut ops); + let mut out = Vec::new(); + len_field(7, &g, &mut out); + len_field(8, &ops, &mut out); + out +} +#[test] +fn raw_and_typed_weights_are_exact_not_truncated_or_ambiguous() { + assert!(load(&model(&node(), &[weight(&1f32.to_le_bytes())])).is_ok()); + for raw in [vec![0; 3], vec![0; 5], vec![0; 8]] { + assert!(load(&model(&node(), &[weight(&raw)])).is_err()); + } + let mut t = weight(&1f32.to_le_bytes()); + varint(4 << 3 | 5, &mut t); + t.extend_from_slice(&1f32.to_le_bytes()); + assert!(load(&model(&node(), &[t])).is_err()); + let mut t = weight(&1f32.to_le_bytes()); + len_field(9, &1f32.to_le_bytes(), &mut t); + assert!(load(&model(&node(), &[t])).is_err()); + // The protobuf float_data form must work as well as raw_data. + let mut t = Vec::new(); + varint_field(1, 1, &mut t); + varint_field(2, 1, &mut t); + len_field(8, b"w", &mut t); + varint(4 << 3 | 5, &mut t); + t.extend_from_slice(&1f32.to_le_bytes()); + assert!(load(&model(&node(), &[t])).is_ok()); +} +#[test] +fn external_sparse_duplicate_custom_domain_and_wrong_dtype_are_errors() { + let t = weight(&1f32.to_le_bytes()); + assert!(load(&model(&node(), &[t.clone(), t.clone()])).is_err()); + for field in [13, 14] { + let mut bad = t.clone(); + if field == 13 { + len_field(field, b"external", &mut bad); + } else { + varint_field(field, 1, &mut bad); + } + assert!(load(&model(&node(), &[bad])).is_err()); + } + let mut bad = t.clone(); + varint_field(2, 10, &mut bad); + assert!(load(&model(&node(), &[bad])).is_err()); + let mut n = node(); + len_field(7, b"vendor.example", &mut n); + assert!(load(&model(&n, &[])).is_err()); + let mut attr = Vec::new(); + len_field(1, b"ratio", &mut attr); + varint_field(20, 1, &mut attr); + let mut n = node(); + len_field(5, &attr, &mut n); + len_field(5, &attr, &mut n); + assert!(load(&model(&n, &[])).is_err()); + let mut g = Vec::new(); + len_field(15, b"sparse", &mut g); + let mut bad = model(&node(), &[]); + len_field(7, &g, &mut bad); + assert!(load(&bad).is_err()); +} +#[test] +fn overflowing_and_cross_boundary_varints_long_strings_and_bad_shapes_are_errors() { + assert!(load(&[0x08, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x02]).is_err()); + let mut t = Vec::new(); + len_field(1, &[0x80], &mut t); + varint_field(2, 1, &mut t); + len_field(8, b"w", &mut t); + assert!(load(&model(&node(), &[t])).is_err()); + let mut n = node(); + len_field(1, &vec![b'x'; 4097], &mut n); + assert!(load(&model(&n, &[])).is_err()); + let mut n = node(); + len_field(1, &[0xff], &mut n); + assert!(load(&model(&n, &[])).is_err()); + let mut t = Vec::new(); + varint_field(1, u64::MAX, &mut t); + varint_field(1, 8, &mut t); + varint_field(2, 1, &mut t); + len_field(8, b"w", &mut t); + assert!(load(&model(&node(), &[t])).is_err()); + let mut bad = model(&node(), &[]); + let mut ops = Vec::new(); + varint_field(2, 999, &mut ops); + len_field(8, &ops, &mut bad); + assert!(load(&bad).is_err()); +} +proptest::proptest! { + #[test] + fn arbitrary_graph_bytes_do_not_panic(bytes in proptest::collection::vec(proptest::num::u8::ANY,0..4096)) { let _=load(&bytes); } +} diff --git a/crates/faces/src/hash.rs b/crates/faces/src/hash.rs new file mode 100644 index 00000000..50f83be3 --- /dev/null +++ b/crates/faces/src/hash.rs @@ -0,0 +1,46 @@ +//! SHA-256 of a model file, streamed (a model can be hundreds of megabytes). + +use std::io::{self, Read}; +use std::path::Path; + +use sha2::{Digest, Sha256}; + +/// Lowercase hex SHA-256 of everything `r` yields. +pub fn sha256_reader(mut r: R) -> io::Result { + let mut h = Sha256::new(); + let mut buf = vec![0u8; 1 << 20]; + loop { + let n = r.read(&mut buf)?; + if n == 0 { + break; + } + h.update(buf.get(..n).unwrap_or(&[])); + } + Ok(hex(&h.finalize())) +} + +/// Lowercase hex SHA-256 of a file. +pub fn sha256_file(path: &Path) -> io::Result { + sha256_reader(io::BufReader::new(std::fs::File::open(path)?)) +} + +/// Lowercase hex SHA-256 of `bytes`. +pub fn sha256_hex(bytes: &[u8]) -> String { + hex(&Sha256::digest(bytes)) +} + +fn hex(bytes: &[u8]) -> String { + bytes.iter().map(|b| format!("{b:02x}")).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn known_digests() { + assert_eq!(sha256_hex(b""), "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"); + assert_eq!(sha256_hex(b"abc"), "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"); + assert_eq!(sha256_reader(&b"abc"[..]).unwrap(), sha256_hex(b"abc")); + } +} diff --git a/crates/faces/src/known.rs b/crates/faces/src/known.rs new file mode 100644 index 00000000..f1ff541c --- /dev/null +++ b/crates/faces/src/known.rs @@ -0,0 +1,223 @@ +//! The models LightCraft recognises by their SHA-256, with what is honestly known about each. +//! +//! The licence and provenance texts here are what the user reads before enabling a model, so they say what +//! is *not* known too. Weights are never part of LightCraft: every model is an opt-in the user installs. + +use crate::manifest::{Colour, Commercial, InputSpec, Licence, ModelManifest, OutputSpec, Resize, Role, Thresholds}; + +/// Detector output decoders built into LightCraft. +pub const DECODERS: &[&str] = &["yunet-v2"]; + +/// The id of the face detector LightCraft runs (see [`yunet`]). +pub const YUNET_ID: &str = "yunet-2023mar"; +pub const YUNET_SHA256: &str = "8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4"; +pub const SFACE_SHA256: &str = "0ba9fbfa01b5270c96627c4ef784da859931e02f04419c829e83484087c34e79"; +pub const AURAFACE_SHA256: &str = "a7933ea5330113b01c9b60351d8f4c33003f145d8470ac5f0e52ee2effe25c60"; + +fn licence(name: &str, commercial: Commercial, url: &str, notice: &str) -> Licence { + Licence { name: name.into(), commercial, url: Some(url.into()), notice: notice.into() } +} + +/// YuNet 2023mar (OpenCV Zoo): a small face detector. +pub fn yunet() -> ModelManifest { + ModelManifest { + id: YUNET_ID.into(), + name: "YuNet (face detector)".into(), + version: "2023mar".into(), + role: Role::Detector, + licence: licence( + "MIT", + Commercial::Yes, + "https://github.com/opencv/opencv_zoo/blob/main/models/face_detection_yunet/LICENSE", + "MIT licence, copyright Shiqi Yu. A separate opt-in download.", + ), + source: Some("https://github.com/opencv/opencv_zoo/tree/main/models/face_detection_yunet".into()), + sha256: Some(YUNET_SHA256.into()), + size_bytes: Some(232_589), + provenance: + "Trained on WIDER FACE, whose terms are not settled by the weights' MIT licence. A detector: it finds faces, it does not identify anyone." + .into(), + input: InputSpec { width: 640, height: 640, colour: Colour::Bgr, mean: [0.0; 3], std: [1.0; 3], resize: Resize::Letterbox }, + output: OutputSpec::Detector { decoder: "yunet-v2".into() }, + thresholds: Thresholds { match_cosine: None, score: Some(0.6), nms_iou: Some(0.3) }, + speed: None, + } +} + +/// SFace 2021dec (OpenCV Zoo). Its match threshold is a starting point: on 755 named faces (busts, matched across +/// shots more than five seconds apart) 0.55 gave 97.5% right suggestions for 88% of faces; `faces.evaluate` checks it on +/// your own photos. +pub fn sface() -> ModelManifest { + ModelManifest { + id: "sface-2021dec".into(), + name: "SFace (face recogniser)".into(), + version: "2021dec".into(), + role: Role::Embedder, + licence: licence( + "Apache-2.0 (as labelled by OpenCV Zoo)", + Commercial::Unknown, + "https://github.com/opencv/opencv_zoo/tree/main/models/face_recognition_sface", + "Labelled Apache-2.0, but what it was trained on is not documented, and two questions about commercial use (opencv_zoo issues 313 and 318) are unanswered. Fine to try for yourself; do not redistribute.", + ), + source: Some("https://github.com/opencv/opencv_zoo/tree/main/models/face_recognition_sface".into()), + sha256: Some(SFACE_SHA256.into()), + size_bytes: Some(38_696_353), + provenance: "Undocumented. The original SFace repository mentions CASIA-WebFace, VGGFace2 and MS1MV2.".into(), + input: InputSpec { width: 112, height: 112, colour: Colour::Rgb, mean: [0.0; 3], std: [1.0; 3], resize: Resize::Stretch }, + output: OutputSpec::Embedding { dim: 128 }, + thresholds: Thresholds { match_cosine: Some(0.55), ..Thresholds::default() }, + speed: Some(6.9), + } +} + +/// AuraFace v1 `glintr100` (fal.ai). Starting-point threshold: on the same 755 faces 0.40 gave 96.6% right suggestions +/// for 56% of faces (0.30: 97.3% for 85%). +pub fn auraface() -> ModelManifest { + ModelManifest { + id: "auraface-v1".into(), + name: "AuraFace v1 (face recogniser)".into(), + version: "1".into(), + role: Role::Embedder, + licence: licence( + "Apache-2.0", + Commercial::Yes, + "https://huggingface.co/fal/AuraFace-v1", + "Apache-2.0. Its training data is described only as a commercial dataset, with no dataset named and no consent statement, and it covers some ethnicities less well. About 261 MB, and slow on a CPU.", + ), + source: Some("https://huggingface.co/fal/AuraFace-v1".into()), + sha256: Some(AURAFACE_SHA256.into()), + size_bytes: Some(260_694_151), + provenance: "Undisclosed: \"a commercial dataset comprising face images from various sources\".".into(), + input: InputSpec::default(), + output: OutputSpec::Embedding { dim: 512 }, + thresholds: Thresholds { match_cosine: Some(0.40), ..Thresholds::default() }, + speed: Some(1.0), + } +} + +/// Every model we know by hash. +pub fn all() -> Vec { + vec![yunet(), sface(), auraface()] +} + +/// The known model whose file has this SHA-256 (lowercase hex). +pub fn lookup(sha256: &str) -> Option { + all().into_iter().find(|m| m.sha256.as_deref() == Some(sha256)) +} + +/// Where the user's "Download" button fetches a model from. The URL is pinned to a commit of the model's own +/// repository, and what arrives is checked against the manifest's size and SHA-256 before anything uses it. +#[derive(Clone, Debug, PartialEq)] +pub struct Download { + pub id: String, + pub url: String, + /// The name the file is kept under while it waits for the user to accept its terms. + pub file_name: String, + pub size_bytes: u64, + pub sha256: String, +} + +/// Pinned download addresses by model id. Only models whose terms allow us to point at them are here: a model +/// that is marked non-commercial is never offered for download (the user brings the file themselves). +const SOURCES: &[(&str, &str, &str)] = &[ + ( + YUNET_ID, + "https://github.com/opencv/opencv_zoo/raw/25f423d0e04c31a17254620e58febd7386da523b/models/face_detection_yunet/face_detection_yunet_2023mar.onnx", + "face_detection_yunet_2023mar.onnx", + ), + ( + "sface-2021dec", + "https://github.com/opencv/opencv_zoo/raw/25f423d0e04c31a17254620e58febd7386da523b/models/face_recognition_sface/face_recognition_sface_2021dec.onnx", + "face_recognition_sface_2021dec.onnx", + ), + ("auraface-v1", "https://huggingface.co/fal/AuraFace-v1/resolve/af6d057c9b0ec4071d4c49c80e3539258798b609/glintr100.onnx", "glintr100.onnx"), +]; + +/// How to download the model with this id, if LightCraft offers to. +pub fn download(id: &str) -> Option { + let (_, url, file_name) = SOURCES.iter().find(|(i, _, _)| *i == id)?; + let m = all().into_iter().find(|m| m.id == id)?; + if m.licence.commercial == Commercial::No { + return None; + } + Some(Download { id: m.id, url: (*url).into(), file_name: (*file_name).into(), size_bytes: m.size_bytes?, sha256: m.sha256? }) +} + +/// "github.com" for a download address: what the user is told the file comes from. +pub fn host(url: &str) -> &str { + url.strip_prefix("https://").and_then(|r| r.split('/').next()).unwrap_or("") +} + +impl Download { + /// The address without the file name (the downloader's "mirror": it fetches `/`). + pub fn base(&self) -> &str { + self.url.strip_suffix(self.file_name.as_str()).map_or(self.url.as_str(), |b| b.trim_end_matches('/')) + } + + /// The site the file comes from, as the user is told. + pub fn host(&self) -> &str { + host(&self.url) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::manifest::validate; + + #[test] + fn known_models_are_valid_and_unique() { + let all = all(); + for m in &all { + assert_eq!(validate(m), Ok(()), "{}", m.id); + } + let mut ids: Vec<_> = all.iter().map(|m| &m.id).collect(); + ids.sort(); + ids.dedup(); + assert_eq!(ids.len(), all.len(), "unique ids"); + let mut hashes: Vec<_> = all.iter().filter_map(|m| m.sha256.as_deref()).collect(); + hashes.sort(); + hashes.dedup(); + assert_eq!(hashes.len(), all.len(), "unique hashes"); + } + + #[test] + fn lookup_finds_by_hash() { + assert_eq!(lookup(AURAFACE_SHA256).map(|m| m.id), Some("auraface-v1".to_string())); + assert_eq!(lookup(YUNET_SHA256).map(|m| m.role), Some(Role::Detector)); + assert!(lookup(&"0".repeat(64)).is_none()); + assert!(lookup("").is_none()); + } + + #[test] + fn downloads_are_pinned_checked_and_only_for_models_we_may_point_at() { + for m in all() { + let Some(d) = download(&m.id) else { continue }; + assert!(d.url.starts_with("https://"), "{}: https only", d.id); + assert!(!d.url.contains("/main/") && !d.url.contains("/master/"), "{}: pinned to a commit, not a branch", d.id); + assert!(d.file_name.ends_with(".onnx") && !d.file_name.contains(['/', '\\']), "{}", d.file_name); + assert!(d.url.ends_with(&format!("/{}", d.file_name)), "{}: the address ends in the file's name", d.id); + assert!(d.base().starts_with("https://") && !d.base().ends_with('/') && !d.base().ends_with(".onnx"), "{}", d.base()); + assert_eq!((Some(d.size_bytes), Some(d.sha256.as_str())), (m.size_bytes, m.sha256.as_deref()), "{}", d.id); + assert!(m.licence.commercial != Commercial::No, "{} is non-commercial and must not be offered", d.id); + } + assert_eq!(download("yunet-2023mar").map(|d| d.host().to_string()).as_deref(), Some("github.com")); + assert_eq!(download("sface-2021dec").map(|d| host(&d.url).to_string()).as_deref(), Some("github.com")); + assert_eq!(download("auraface-v1").map(|d| host(&d.url).to_string()).as_deref(), Some("huggingface.co")); + // unknown ids and path tricks have nothing to download + for id in ["nope", "", "../sface-2021dec", "SFACE-2021DEC"] { + assert!(download(id).is_none(), "{id}"); + } + // every address in the table belongs to a known model + for (id, _, _) in SOURCES { + assert!(all().iter().any(|m| m.id == *id), "{id}"); + } + } + + #[test] + fn unresolved_models_say_so() { + assert_eq!(sface().licence.commercial, Commercial::Unknown); + assert!(sface().provenance.to_lowercase().contains("undocumented")); + assert!(auraface().provenance.to_lowercase().contains("undisclosed")); + } +} diff --git a/crates/faces/src/lib.rs b/crates/faces/src/lib.rs new file mode 100644 index 00000000..b4b7066d --- /dev/null +++ b/crates/faces/src/lib.rs @@ -0,0 +1,26 @@ +//! Face models for LightCraft: what a model is ([`ModelManifest`]), whether to trust a description of one +//! ([`manifest::validate`]), the models we know by hash ([`known`]), and a probe that reads an unknown +//! `.onnx` file's input and output shapes without running it ([`onnx`]), so installing a model of your +//! own can be one drop and one licence prompt. +//! +//! Everything here treats model files and manifests as hostile input: sizes are capped, numbers must be +//! finite and sane, and a malformed file is an error, never a panic. Detection and recognition run +//! on the same checked, single-thread pure-Rust CPU interpreter. +#![forbid(unsafe_code)] +#![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)] + +pub mod align; +pub mod catalog; +pub mod graph; +pub mod hash; +pub mod known; +pub mod manifest; +pub mod matching; +pub mod net; +pub mod onnx; +pub mod runtime; +pub mod suggest; +pub mod synthetic; +pub mod yunet; + +pub use manifest::{Colour, Commercial, InputSpec, Licence, ManifestError, ModelManifest, OutputSpec, Resize, Role, Thresholds}; diff --git a/crates/faces/src/manifest.rs b/crates/faces/src/manifest.rs new file mode 100644 index 00000000..d3d99501 --- /dev/null +++ b/crates/faces/src/manifest.rs @@ -0,0 +1,386 @@ +//! A face model's manifest: what the model is, what it needs and who may use it, as plain data. +//! +//! A manifest comes from the known-model table, from a suggestion made for a file the user dropped in, or +//! from a `face-model.json` next to a model. All of those are untrusted: [`validate`] must pass before a +//! manifest is used, and [`parse`] applies it after reading JSON of bounded size. + +use serde::{Deserialize, Serialize}; + +/// Largest `face-model.json` we read. +pub const MAX_MANIFEST_BYTES: usize = 64 * 1024; +/// Largest model file we accept (bytes). +pub const MAX_MODEL_BYTES: u64 = 512 * 1024 * 1024; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum Role { + /// Finds faces (boxes, and landmarks when the model has them). + Detector, + /// Turns an aligned face into a vector whose distance says whether two faces are the same person. + Embedder, +} + +/// Whether the weights may be used commercially. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum Commercial { + Yes, + /// Research or personal use only. Never bundled or redistributed by LightCraft. + No, + /// The licence or the training data is unclear: the user is told so before installing. + Unknown, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub struct Licence { + /// SPDX id or a short name ("MIT", "Apache-2.0", "InsightFace non-commercial research"). + pub name: String, + pub commercial: Commercial, + pub url: Option, + /// Plain-language terms shown before the model is enabled. + pub notice: String, +} + +impl Default for Licence { + fn default() -> Self { + Self { name: String::new(), commercial: Commercial::Unknown, url: None, notice: String::new() } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum Colour { + Rgb, + Bgr, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum Resize { + /// Stretch the crop to the input size (aligned face crops). + Stretch, + /// Keep the aspect ratio and pad (whole photos for a detector). + Letterbox, +} + +/// How an image becomes the model's input: NCHW float, `(pixel − mean) / std` per channel, pixel in 0..255. +#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub struct InputSpec { + pub width: u32, + pub height: u32, + pub colour: Colour, + pub mean: [f32; 3], + pub std: [f32; 3], + pub resize: Resize, +} + +impl Default for InputSpec { + /// ArcFace-style: 112 × 112, RGB, `(x − 127.5) / 127.5`. + fn default() -> Self { + Self { width: 112, height: 112, colour: Colour::Rgb, mean: [127.5; 3], std: [127.5; 3], resize: Resize::Stretch } + } +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "camelCase")] +pub enum OutputSpec { + /// One vector of `dim` numbers per face. + Embedding { dim: u32 }, + /// A detector whose outputs a built-in decoder understands (`"yunet-v2"`). + Detector { decoder: String }, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub struct Thresholds { + /// Embedders: cosine similarity at or above which two faces are suggested as a match. + pub match_cosine: Option, + /// Detectors: lowest detection score kept. + pub score: Option, + /// Detectors: overlap above which two boxes are one face. + pub nms_iou: Option, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ModelManifest { + /// Lowercase letters, digits, `.` `_` `-`. + pub id: String, + pub name: String, + pub version: String, + pub role: Role, + #[serde(default)] + pub licence: Licence, + /// Where the weights come from (an `https://` page or file), for the user's information. + #[serde(default)] + pub source: Option, + /// Lowercase hex SHA-256 of the `.onnx` file, when known. + #[serde(default)] + pub sha256: Option, + #[serde(default)] + pub size_bytes: Option, + /// What the model was trained on, in plain words; says "undisclosed" when it is. + #[serde(default)] + pub provenance: String, + #[serde(default)] + pub input: InputSpec, + pub output: OutputSpec, + #[serde(default)] + pub thresholds: Thresholds, + /// How fast it is as a multiple of a ResNet-100 model's speed (`1.0`): `2.0` is twice as fast. Shown instead of + /// timings, which depend on the computer; leave it out when you have not measured it. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub speed: Option, +} + +/// The reference model the `speed` ratios are measured against. +pub const SPEED_REFERENCE: &str = "a ResNet-100 model"; + +/// A speed ratio in plain words, as times faster than the reference ("4.6× faster than a ResNet-100 model"). +pub fn describe_speed(speed: f32) -> String { + if !(speed.is_finite() && speed > 0.0) { + return String::new(); + } + let times = |x: f32| { + let s = format!("{x:.1}"); + s.strip_suffix(".0").map_or(s.clone(), str::to_string) + }; + if speed >= 1.15 { + format!("{}× faster than {SPEED_REFERENCE}", times(speed)) + } else if speed > 0.87 { + format!("Same speed as {SPEED_REFERENCE}") + } else { + format!("{}× slower than {SPEED_REFERENCE}", times(1.0 / speed)) + } +} + +#[derive(Debug, PartialEq, thiserror::Error)] +pub enum ManifestError { + #[error("`{0}` is not valid: {1}")] + Field(&'static str, String), + #[error("not a face-model manifest: {0}")] + Parse(String), + #[error("a manifest may be at most {MAX_MANIFEST_BYTES} bytes")] + TooLarge, +} + +fn bad(field: &'static str, why: impl Into) -> Result { + Err(ManifestError::Field(field, why.into())) +} + +/// A model id is a folder name: 1 to 64 characters, starting with a lowercase letter or digit, then +/// lowercase letters, digits, `.`, `_` or `-`. So it is never `.`, `..`, hidden, or a path. +pub fn valid_id(id: &str) -> bool { + let mut bytes = id.bytes(); + let first_ok = bytes.next().is_some_and(|b| b.is_ascii_lowercase() || b.is_ascii_digit()); + first_ok && id.len() <= 64 && bytes.all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || matches!(b, b'.' | b'_' | b'-')) +} + +fn text(field: &'static str, s: &str, min: usize, max: usize) -> Result<(), ManifestError> { + let n = s.chars().count(); + if n < min || n > max { + return bad(field, format!("must be {min} to {max} characters")); + } + if s.chars().any(|c| c.is_control() && c != '\n') { + return bad(field, "must not contain control characters"); + } + Ok(()) +} + +fn url(field: &'static str, u: &Option) -> Result<(), ManifestError> { + match u { + None => Ok(()), + Some(u) if u.len() <= 500 && u.starts_with("https://") && !u.chars().any(|c| c.is_control() || c.is_whitespace()) => Ok(()), + Some(_) => bad(field, "must be an https:// address of at most 500 characters"), + } +} + +fn ratio(field: &'static str, v: Option, lo: f32, hi: f32) -> Result<(), ManifestError> { + match v { + Some(v) if !(v.is_finite() && (lo..=hi).contains(&v)) => bad(field, format!("must be a number from {lo} to {hi}")), + _ => Ok(()), + } +} + +/// Check a manifest from any source. Everything the rest of LightCraft relies on (sizes it will allocate, +/// numbers it will divide by, text it will show) is bounded here. +pub fn validate(m: &ModelManifest) -> Result<(), ManifestError> { + if !valid_id(&m.id) { + return bad("id", "must be 1 to 64 characters: lowercase letters and digits, then also '.', '_' or '-'"); + } + text("name", &m.name, 1, 120)?; + text("version", &m.version, 1, 40)?; + text("licence.name", &m.licence.name, 0, 80)?; + text("licence.notice", &m.licence.notice, 0, 800)?; + text("provenance", &m.provenance, 0, 800)?; + url("licence.url", &m.licence.url)?; + url("source", &m.source)?; + if let Some(h) = &m.sha256 + && !(h.len() == 64 && h.bytes().all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b))) + { + return bad("sha256", "must be 64 lowercase hexadecimal digits"); + } + if m.size_bytes.is_some_and(|s| s == 0 || s > MAX_MODEL_BYTES) { + return bad("sizeBytes", format!("must be 1 to {MAX_MODEL_BYTES}")); + } + let i = &m.input; + if !(16..=4096).contains(&i.width) || !(16..=4096).contains(&i.height) { + return bad("input.width/height", "must be 16 to 4096"); + } + if i.mean.iter().any(|v| !v.is_finite() || v.abs() > 1000.0) { + return bad("input.mean", "must be finite numbers within ±1000"); + } + if i.std.iter().any(|v| !v.is_finite() || v.abs() < 1e-3 || v.abs() > 1000.0) { + return bad("input.std", "must be finite, not zero, within ±1000"); + } + match &m.output { + OutputSpec::Embedding { dim } => { + if m.role != Role::Embedder { + return bad("output", "an embedding output needs the embedder role"); + } + if !(1..=4096).contains(dim) { + return bad("output.dim", "must be 1 to 4096"); + } + } + OutputSpec::Detector { decoder } => { + if m.role != Role::Detector { + return bad("output", "a detector output needs the detector role"); + } + if !crate::known::DECODERS.contains(&decoder.as_str()) { + return bad("output.decoder", format!("unknown decoder `{}`", decoder.chars().take(40).collect::())); + } + } + } + ratio("thresholds.matchCosine", m.thresholds.match_cosine, -1.0, 1.0)?; + ratio("thresholds.score", m.thresholds.score, 0.0, 1.0)?; + ratio("thresholds.nmsIou", m.thresholds.nms_iou, 0.0, 1.0)?; + ratio("speed", m.speed, 0.01, 1000.0)?; + Ok(()) +} + +/// Read a `face-model.json` (at most [`MAX_MANIFEST_BYTES`]) and validate it. +pub fn parse(bytes: &[u8]) -> Result { + if bytes.len() > MAX_MANIFEST_BYTES { + return Err(ManifestError::TooLarge); + } + let m: ModelManifest = serde_json::from_slice(bytes).map_err(|e| ManifestError::Parse(e.to_string()))?; + validate(&m)?; + Ok(m) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn good() -> ModelManifest { + ModelManifest { + id: "example-r50".into(), + name: "Example R50".into(), + version: "1".into(), + role: Role::Embedder, + licence: Licence { + name: "MIT".into(), + commercial: Commercial::Yes, + url: Some("https://example.org/licence".into()), + notice: String::new(), + }, + source: Some("https://example.org/model.onnx".into()), + sha256: Some("a".repeat(64)), + size_bytes: Some(1234), + provenance: "synthetic".into(), + input: InputSpec::default(), + output: OutputSpec::Embedding { dim: 512 }, + thresholds: Thresholds { match_cosine: Some(0.4), ..Default::default() }, + speed: Some(1.7), + } + } + + #[test] + fn speed_reads_as_times_faster_than_the_reference() { + assert_eq!(describe_speed(7.0), "7× faster than a ResNet-100 model"); + assert_eq!(describe_speed(4.6), "4.6× faster than a ResNet-100 model"); + assert_eq!(describe_speed(1.7), "1.7× faster than a ResNet-100 model"); + assert_eq!(describe_speed(1.0), "Same speed as a ResNet-100 model"); + assert_eq!(describe_speed(0.5), "2× slower than a ResNet-100 model"); + for odd in [0.0, -1.0, f32::NAN, f32::INFINITY] { + assert_eq!(describe_speed(odd), "", "{odd}"); + } + } + + #[test] + fn a_good_manifest_passes_and_round_trips() { + let m = good(); + assert_eq!(validate(&m), Ok(())); + let json = serde_json::to_vec(&m).unwrap(); + assert_eq!(parse(&json).unwrap(), m); + } + + #[test] + fn every_field_is_bounded() { + let cases: Vec<(&str, Box)> = vec![ + ("empty id", Box::new(|m| m.id = String::new())), + ("uppercase id", Box::new(|m| m.id = "Bad".into())), + ("path in id", Box::new(|m| m.id = "../etc".into())), + ("dot id", Box::new(|m| m.id = ".".into())), + ("dotdot id", Box::new(|m| m.id = "..".into())), + ("hidden id", Box::new(|m| m.id = ".hidden".into())), + ("slash id", Box::new(|m| m.id = "a/b".into())), + ("backslash id", Box::new(|m| m.id = r"a\b".into())), + ("long id", Box::new(|m| m.id = "a".repeat(65))), + ("empty name", Box::new(|m| m.name = String::new())), + ("control in name", Box::new(|m| m.name = "a\u{7}b".into())), + ("huge notice", Box::new(|m| m.licence.notice = "x".repeat(801))), + ("http source", Box::new(|m| m.source = Some("http://example.org/m.onnx".into()))), + ("file source", Box::new(|m| m.source = Some("file:///etc/passwd".into()))), + ("space in url", Box::new(|m| m.source = Some("https://a b".into()))), + ("short sha", Box::new(|m| m.sha256 = Some("abc".into()))), + ("uppercase sha", Box::new(|m| m.sha256 = Some("A".repeat(64)))), + ("zero size", Box::new(|m| m.size_bytes = Some(0))), + ("huge size", Box::new(|m| m.size_bytes = Some(u64::MAX))), + ("tiny input", Box::new(|m| m.input.width = 4)), + ("huge input", Box::new(|m| m.input.height = 1 << 20)), + ("nan mean", Box::new(|m| m.input.mean[1] = f32::NAN)), + ("inf mean", Box::new(|m| m.input.mean[0] = f32::INFINITY)), + ("zero std", Box::new(|m| m.input.std[2] = 0.0)), + ("nan std", Box::new(|m| m.input.std[0] = f32::NAN)), + ("zero dim", Box::new(|m| m.output = OutputSpec::Embedding { dim: 0 })), + ("huge dim", Box::new(|m| m.output = OutputSpec::Embedding { dim: 10_000_000 })), + ("role mismatch", Box::new(|m| m.role = Role::Detector)), + ( + "unknown decoder", + Box::new(|m| { + m.role = Role::Detector; + m.output = OutputSpec::Detector { decoder: "made-up".into() }; + }), + ), + ("nan threshold", Box::new(|m| m.thresholds.match_cosine = Some(f32::NAN))), + ("threshold range", Box::new(|m| m.thresholds.score = Some(1.5))), + ("nan speed", Box::new(|m| m.speed = Some(f32::NAN))), + ("zero speed", Box::new(|m| m.speed = Some(0.0))), + ("huge speed", Box::new(|m| m.speed = Some(1e9))), + ]; + for (what, f) in cases { + let mut m = good(); + f(&mut m); + assert!(validate(&m).is_err(), "{what} must be rejected"); + } + } + + #[test] + fn hostile_json_is_an_error() { + assert!(matches!(parse(&vec![b' '; MAX_MANIFEST_BYTES + 1]), Err(ManifestError::TooLarge))); + for bytes in [ + &b""[..], + b"null", + b"[]", + b"{", + b"{\"id\": 1}", + b"\xff\xfe\x00", + b"{\"id\":\"x\",\"name\":\"n\",\"version\":\"1\",\"role\":\"embedder\"}", + ] { + assert!(parse(bytes).is_err(), "{:?}", String::from_utf8_lossy(bytes)); + } + } +} diff --git a/crates/faces/src/matching.rs b/crates/faces/src/matching.rs new file mode 100644 index 00000000..de50ad3e --- /dev/null +++ b/crates/faces/src/matching.rs @@ -0,0 +1,128 @@ +//! Who does this face look like? Rank the people in a gallery of named faces by how close their faces are, +//! and suggest one only when the match is strong *and* clearly ahead of the runner-up. +//! +//! Embeddings are unit vectors, so closeness is the dot product (cosine similarity). A person's score is the +//! mean of their best few faces (one stray, wrongly named face cannot make a stranger look like them). The +//! result is only ever a suggestion for the user to confirm: it is better to leave a face unnamed than to +//! name it wrongly. + +use std::cmp::Ordering; + +/// How many of a person's best matching faces are averaged into their score. +pub const TOP_FACES: usize = 2; + +#[derive(Clone, Debug, PartialEq)] +pub struct Ranked { + pub name: String, + pub score: f32, + /// How many of their faces the score is based on. + pub faces: usize, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Suggestion { + pub name: String, + pub score: f32, + /// The next best person, and how far behind. + pub runner_up: Option<(String, f32)>, +} + +/// Cosine similarity of two unit vectors. `None` when their lengths differ or either has a non-finite number. +pub fn cosine(a: &[f32], b: &[f32]) -> Option { + if a.len() != b.len() || a.is_empty() { + return None; + } + let d: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum(); + d.is_finite().then_some(d) +} + +/// Rank the people in `gallery` (a name and one face's embedding per entry; a person appears once per face), +/// best first. Names are compared case-insensitively and shown as first seen. +pub fn rank(query: &[f32], gallery: &[(&str, &[f32])]) -> Vec { + let mut people: Vec<(String, String, Vec)> = Vec::new(); // (key, shown name, similarities) + for (name, emb) in gallery { + let Some(c) = cosine(query, emb) else { continue }; + let key = name.to_lowercase(); + match people.iter_mut().find(|p| p.0 == key) { + Some(p) => p.2.push(c), + None => people.push((key, (*name).to_string(), vec![c])), + } + } + let mut out: Vec = people + .into_iter() + .map(|(_, name, mut sims)| { + sims.sort_by(|a, b| b.partial_cmp(a).unwrap_or(Ordering::Equal)); + let top = sims.get(..sims.len().min(TOP_FACES)).unwrap_or(&[]); + Ranked { name, score: top.iter().sum::() / top.len().max(1) as f32, faces: sims.len() } + }) + .collect(); + out.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(Ordering::Equal).then_with(|| a.name.cmp(&b.name))); + out +} + +/// The best person if they pass both bars: a score of at least `threshold`, and at least `margin` ahead of the +/// next person (when there is one). +pub fn suggest(ranked: &[Ranked], threshold: f32, margin: f32) -> Option { + let best = ranked.first()?; + let next = ranked.get(1); + if best.score.is_nan() || best.score < threshold || next.is_some_and(|n| best.score - n.score < margin) { + return None; + } + Some(Suggestion { name: best.name.clone(), score: best.score, runner_up: next.map(|n| (n.name.clone(), n.score)) }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn unit(v: &[f32]) -> Vec { + let n = v.iter().map(|x| x * x).sum::().sqrt(); + v.iter().map(|x| x / n).collect() + } + + #[test] + fn people_are_ranked_by_their_closest_faces() { + let (a1, a2, b1) = (unit(&[1.0, 0.1, 0.0]), unit(&[0.9, 0.2, 0.0]), unit(&[0.0, 1.0, 0.2])); + let gallery: Vec<(&str, &[f32])> = vec![("Ann", &a1), ("ann", &a2), ("Bob", &b1)]; + let ranked = rank(&unit(&[1.0, 0.0, 0.0]), &gallery); + assert_eq!(ranked.iter().map(|r| r.name.as_str()).collect::>(), vec!["Ann", "Bob"], "case-insensitive, shown as first seen"); + assert_eq!(ranked[0].faces, 2); + assert!(ranked[0].score > 0.95 && ranked[1].score < 0.2); + // a stray, wrongly named face cannot carry a person: Cy has one far face and one good one + let (c_good, c_bad) = (unit(&[1.0, 0.0, 0.0]), unit(&[0.0, 0.0, 1.0])); + let g2: Vec<(&str, &[f32])> = vec![("Cy", &c_good), ("Cy", &c_bad), ("Di", &a1), ("Di", &a2)]; + let r2 = rank(&unit(&[1.0, 0.0, 0.0]), &g2); + assert_eq!(r2[0].name, "Di", "{r2:?}"); + } + + #[test] + fn a_suggestion_needs_a_strong_match_clearly_ahead() { + let r = |name: &str, score: f32| Ranked { name: name.into(), score, faces: 1 }; + // strong and clear + let s = suggest(&[r("Ann", 0.62), r("Bob", 0.20)], 0.45, 0.08).unwrap(); + assert_eq!((s.name.as_str(), s.runner_up.as_ref().map(|x| x.0.as_str())), ("Ann", Some("Bob"))); + // strong but too close to the runner-up: nobody + assert!(suggest(&[r("Ann", 0.62), r("Bob", 0.58)], 0.45, 0.08).is_none()); + // clear but not strong enough + assert!(suggest(&[r("Ann", 0.40), r("Bob", 0.10)], 0.45, 0.08).is_none()); + // a single person has no runner-up to be ahead of + assert!(suggest(&[r("Ann", 0.5)], 0.45, 0.08).is_some()); + assert!(suggest(&[], 0.45, 0.08).is_none()); + assert!(suggest(&[r("Ann", f32::NAN)], 0.45, 0.08).is_none()); + } + + #[test] + fn mismatched_or_non_finite_embeddings_are_skipped_not_panics() { + assert!(cosine(&[1.0], &[1.0, 0.0]).is_none()); + assert!(cosine(&[], &[]).is_none()); + assert!(cosine(&[f32::NAN], &[1.0]).is_none()); + let good = unit(&[1.0, 0.0]); + let bad = [f32::NAN, 0.0]; + let short = [1.0f32]; + let gallery: Vec<(&str, &[f32])> = vec![("Ann", &good), ("Nan", &bad), ("Short", &short)]; + let ranked = rank(&good, &gallery); + assert_eq!(ranked.len(), 1); + assert!(rank(&[], &gallery).is_empty()); + assert!(rank(&good, &[]).is_empty()); + } +} diff --git a/crates/faces/src/net.rs b/crates/faces/src/net.rs new file mode 100644 index 00000000..6222314f --- /dev/null +++ b/crates/faces/src/net.rs @@ -0,0 +1,570 @@ +//! A small CPU interpreter for the convolutional networks LightCraft bundles (the YuNet face detector). +//! +//! It runs checked convolution, affine/activation, broadcast, pooling, matrix and shape operators on +//! float32 tensors. Unsupported semantics fail at load time. Dense convolutions use bounded im2col and +//! single-thread GEMM; depthwise convolutions use row-wise Rust loops. Every file-derived shape is checked. +//! +//! The same checked float32 engine runs YuNet and installed SFace/AuraFace recognisers. +//! Unknown operators, attributes and shapes are rejected before recognition starts. + +use std::collections::HashMap; +use std::rc::Rc; + +mod kernels; +mod shape; + +use crate::graph::{Data, Graph, Node}; + +/// Elements of any one tensor the interpreter will allocate. +const MAX_TENSOR: usize = 64 * 1024 * 1024; + +#[derive(Debug, thiserror::Error, PartialEq)] +pub enum NetError { + #[error("this model uses something the built-in runtime does not support: {0}")] + Unsupported(String), + #[error("the model does not fit together: {0}")] + Shape(String), +} + +type Result = std::result::Result; + +fn shape_err(why: impl Into) -> Result { + Err(NetError::Shape(why.into())) +} + +fn unsupported(why: impl Into) -> Result { + Err(NetError::Unsupported(why.into())) +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Tensor { + pub shape: Vec, + pub data: Vec, +} + +impl Tensor { + pub fn new(shape: Vec, data: Vec) -> Result { + let n = shape.iter().try_fold(1usize, |a, d| a.checked_mul(*d)); + if shape.len() > 8 || shape.contains(&0) || n.is_none_or(|n| n > MAX_TENSOR) || n != Some(data.len()) || data.iter().any(|v| !v.is_finite()) { + return shape_err(format!("{} values do not fill the shape {shape:?}", data.len())); + } + Ok(Tensor { shape, data }) + } + + fn zeros(shape: Vec) -> Result { + match shape.iter().try_fold(1usize, |a, d| a.checked_mul(*d)) { + Some(n) if n > 0 && n <= MAX_TENSOR && shape.len() <= 8 && !shape.contains(&0) => { + let mut data = Vec::new(); + data.try_reserve_exact(n).map_err(|_| NetError::Shape("not enough memory for a model tensor".into()))?; + data.resize(n, 0.0); + Ok(Tensor { shape, data }) + } + _ => shape_err("a tensor would be too large"), + } + } +} + +/// A loaded network, ready to run. +pub struct Net { + graph: Graph, +} + +/// Validate semantics rather than silently ignoring unknown or mistyped attributes. +fn check(node: &Node, g: &Graph) -> Result<()> { + shape::check(node, g) +} + +impl Net { + pub fn new(graph: Graph) -> Result { + shape::check_graph(&graph)?; + for n in &graph.nodes { + check(n, &graph)?; + } + Ok(Net { graph }) + } + + /// The shape the model declares for its input (0 for a dynamic size). + pub fn input_shape(&self) -> Option<&[usize]> { + self.graph.inputs.first().map(|(_, s)| s.as_slice()) + } + + /// Infer and check every intermediate without allocating its pixel data. + /// Recognition calls this at load time with the manifest's fixed input. + pub fn validate_input(&self, shape: &[usize]) -> Result>> { + shape::infer(&self.graph, shape) + } + + /// Run the network on `input` (its single input); returns each output by name. + pub fn run(&self, input: Tensor) -> Result> { + shape::tensor(&input.shape, input.data.len(), &input.data)?; + self.validate_input(&input.shape)?; + let mut scratch = kernels::Scratch::default(); + let profile = std::env::var_os("LIGHTCRAFT_PROFILE").is_some(); + let g = &self.graph; + let Some((input_name, _)) = g.inputs.first() else { return shape_err("the model has no input") }; + // when each value is last read, so big intermediates are dropped as soon as nothing needs them + let mut last_use: HashMap<&str, usize> = HashMap::new(); + for (i, n) in g.nodes.iter().enumerate() { + for name in &n.inputs { + last_use.insert(name.as_str(), i); + } + } + for o in &g.outputs { + last_use.insert(o.as_str(), usize::MAX); + } + let mut values: HashMap> = HashMap::new(); + values.insert(input_name.clone(), Rc::new(input)); + for (i, node) in g.nodes.iter().enumerate() { + let get = |idx: usize| -> Result> { + let name = node.inputs.get(idx).ok_or_else(|| NetError::Shape(format!("{} is missing an input", node.op)))?; + if let Some(value) = values.get(name) { + return Ok(value.clone()); + } + if let Some(crate::graph::Weight { dims, data: Data::F32(data) }) = g.weights.get(name) { + return Ok(Rc::new(Tensor::new(dims.clone(), data.clone())?)); + } + values.get(name).cloned().ok_or_else(|| NetError::Shape(format!("`{name}` is read before it is computed"))) + }; + let started = web_time::Instant::now(); + let out = match node.op.as_str() { + "Conv" => kernels::conv(node, &*get(0)?, g, &mut scratch)?, + "Relu" => map(&*get(0)?, |v| v.max(0.0)), + "Sigmoid" => map(&*get(0)?, |v| 1.0 / (1.0 + (-v).exp())), + "Add" | "Sub" | "Mul" => kernels::binary(node.op.as_str(), &*get(0)?, &*get(1)?)?, + "PRelu" => kernels::prelu(node, &*get(0)?, g)?, + "BatchNormalization" => kernels::batch_norm(node, &*get(0)?, g)?, + "Gemm" => kernels::gemm(node, &*get(0)?, g)?, + "Flatten" => kernels::flatten(node, &*get(0)?)?, + "Dropout" | "Identity" => (*get(0)?).clone(), + "GlobalAveragePool" => kernels::global_average(&*get(0)?)?, + "MaxPool" => max_pool(node, &*get(0)?)?, + "Resize" => resize(node, &*get(0)?, g)?, + "Transpose" => transpose(node, &*get(0)?)?, + "Reshape" => reshape(node, &*get(0)?, g)?, + other => return unsupported(format!("the operator `{other}`")), + }; + if out.data.iter().any(|v| !v.is_finite()) { + return shape_err(format!("{} produced non-finite values", node.op)); + } + if profile { + eprintln!("[faces-cpu] layer {i} {} {:?}: {:.3} ms", node.op, out.shape, started.elapsed().as_secs_f64() * 1000.0); + } + let Some(name) = node.outputs.first() else { return shape_err("a node has no output") }; + values.insert(name.clone(), Rc::new(out)); + for name in &node.inputs { + if last_use.get(name.as_str()) == Some(&i) { + values.remove(name); + } + } + } + g.outputs + .iter() + .map(|name| { + let t = values.remove(name).ok_or_else(|| NetError::Shape(format!("output `{name}` was never computed")))?; + let t = Rc::try_unwrap(t).map_err(|_| NetError::Shape(format!("output `{name}` is still shared")))?; + Ok((name.clone(), t)) + }) + .collect() + } +} + +fn map(x: &Tensor, f: impl Fn(f32) -> f32) -> Tensor { + Tensor { shape: x.shape.clone(), data: x.data.iter().map(|v| f(*v)).collect() } +} + +/// `[1, c, h, w]` → `(c, h, w)`. +fn chw(x: &Tensor) -> Result<(usize, usize, usize)> { + match x.shape.as_slice() { + [1, c, h, w] => Ok((*c, *h, *w)), + s => unsupported(format!("a tensor of shape {s:?} (only [1, channels, height, width])")), + } +} + +/// `pads` of an ONNX node as (top, left, bottom, right). +fn pads(node: &Node) -> Result<(usize, usize, usize, usize)> { + match node.ints("pads") { + None => Ok((0, 0, 0, 0)), + Some([t, l, b, r]) => { + let u = |v: &i64| usize::try_from(*v).map_err(|_| NetError::Shape("a negative pad".into())); + Ok((u(t)?, u(l)?, u(b)?, u(r)?)) + } + Some(_) => unsupported("pads that are not [top, left, bottom, right]"), + } +} + +fn pair(node: &Node, name: &str, default: usize) -> Result<(usize, usize)> { + match node.ints(name) { + None => Ok((default, default)), + Some([a, b]) => match (usize::try_from(*a), usize::try_from(*b)) { + (Ok(a), Ok(b)) if a >= 1 && b >= 1 => Ok((a, b)), + _ => shape_err(format!("`{name}` must be positive")), + }, + Some(_) => unsupported(format!("`{name}` that is not two numbers")), + } +} + +fn conv_direct(node: &Node, x: &Tensor, g: &Graph) -> Result { + let (cin, h, w) = chw(x)?; + let weight = node.inputs.get(1).and_then(|n| g.weights.get(n)).ok_or_else(|| NetError::Shape("a Conv lost its weights".into()))?; + let Data::F32(wd) = &weight.data else { return unsupported("non-float Conv weights") }; + let [cout, cin_g, kh, kw] = weight.dims.as_slice() else { return unsupported("a Conv that is not 2-D") }; + let (cout, cin_g, kh, kw) = (*cout, *cin_g, *kh, *kw); + let groups = + usize::try_from(node.int("group").unwrap_or(1)).ok().filter(|g| *g >= 1).ok_or_else(|| NetError::Shape("a bad Conv group".into()))?; + if groups.checked_mul(cin_g) != Some(cin) || cout % groups != 0 || kh == 0 || kw == 0 { + return shape_err("the Conv weights do not match its input"); + } + let bias: Option<&[f32]> = match node.inputs.get(2).filter(|b| !b.is_empty()).and_then(|b| g.weights.get(b)) { + Some(b) => match &b.data { + Data::F32(v) if v.len() == cout => Some(v), + _ => return shape_err("the Conv bias does not match"), + }, + None => None, + }; + let (sh, sw) = pair(node, "strides", 1)?; + let (pt, pl, pb, pr) = pads(node)?; + let ph = h.checked_add(pt).and_then(|v| v.checked_add(pb)).ok_or_else(|| NetError::Shape("Conv padding overflow".into()))?; + let pw = w.checked_add(pl).and_then(|v| v.checked_add(pr)).ok_or_else(|| NetError::Shape("Conv padding overflow".into()))?; + if ph < kh || pw < kw { + return shape_err("the Conv kernel is larger than its input"); + } + let (oh, ow) = ((ph - kh) / sh + 1, (pw - kw) / sw + 1); + let mut out = Tensor::zeros(vec![1, cout, oh, ow])?; + let plane = oh * ow; + let in_planes: Vec<&[f32]> = x.data.chunks_exact((h * w).max(1)).collect(); + let cout_g = cout / groups; + for (co, out_plane) in out.data.chunks_exact_mut(plane.max(1)).enumerate() { + if let Some(b) = bias.and_then(|b| b.get(co)) { + out_plane.iter_mut().for_each(|o| *o = *b); + } + let group = co / cout_g; + for ci_l in 0..cin_g { + let Some(in_plane) = in_planes.get(group * cin_g + ci_l) else { return shape_err("a Conv input channel is missing") }; + let base = (co * cin_g + ci_l) * kh * kw; + let Some(kernel) = wd.get(base..base + kh * kw) else { return shape_err("the Conv weights are too short") }; + for (k, wv) in kernel.iter().enumerate() { + if *wv == 0.0 { + continue; + } + let (ky, kx) = (k / kw, k % kw); + for oy in 0..oh { + // input row for this output row and kernel row (in padded coordinates minus the pad) + let Some(iy) = (oy * sh + ky).checked_sub(pt).filter(|iy| *iy < h) else { continue }; + let in_row = in_plane.get(iy * w..(iy + 1) * w).unwrap_or(&[]); + let out_row = out_plane.get_mut(oy * ow..(oy + 1) * ow).unwrap_or(&mut []); + if sw == 1 { + // ox runs over outputs whose input column ox + kx - pl lies inside the row + let lo = pl.saturating_sub(kx); + let hi = ow.min((w + pl).saturating_sub(kx)); + if lo < hi { + let src = in_row.get(lo + kx - pl..hi + kx - pl).unwrap_or(&[]); + let dst = out_row.get_mut(lo..hi).unwrap_or(&mut []); + dst.iter_mut().zip(src).for_each(|(o, i)| *o += wv * i); + } + } else { + for (ox, o) in out_row.iter_mut().enumerate() { + if let Some(ix) = (ox * sw + kx).checked_sub(pl).filter(|ix| *ix < w) + && let Some(i) = in_row.get(ix) + { + *o += wv * i; + } + } + } + } + } + } + } + Ok(out) +} + +fn max_pool(node: &Node, x: &Tensor) -> Result { + let (c, h, w) = chw(x)?; + let (kh, kw) = pair(node, "kernel_shape", 1)?; + let (sh, sw) = pair(node, "strides", 1)?; + let (pt, pl, pb, pr) = pads(node)?; + if (pt, pl, pb, pr) != (0, 0, 0, 0) { + return unsupported("a padded MaxPool"); + } + if h < kh || w < kw { + return shape_err("the MaxPool window is larger than its input"); + } + let (oh, ow) = ((h - kh) / sh + 1, (w - kw) / sw + 1); + let mut out = Tensor::zeros(vec![1, c, oh, ow])?; + for (plane_in, plane_out) in x.data.chunks_exact((h * w).max(1)).zip(out.data.chunks_exact_mut((oh * ow).max(1))) { + for oy in 0..oh { + for ox in 0..ow { + let mut m = f32::NEG_INFINITY; + for ky in 0..kh { + let row = plane_in.get((oy * sh + ky) * w..(oy * sh + ky + 1) * w).unwrap_or(&[]); + for kx in 0..kw { + if let Some(v) = row.get(ox * sw + kx) { + m = m.max(*v); + } + } + } + if let Some(o) = plane_out.get_mut(oy * ow + ox) { + *o = m; + } + } + } + } + Ok(out) +} + +/// Nearest-neighbour upscale by whole numbers (YuNet's feature pyramid doubles each map). +fn resize(node: &Node, x: &Tensor, g: &Graph) -> Result { + let (c, h, w) = chw(x)?; + let scales = match node.inputs.get(2).and_then(|n| g.weights.get(n)).map(|w| &w.data) { + Some(Data::F32(s)) if s.len() == 4 => s, + _ => return unsupported("a Resize whose scales are not four floats"), + }; + let whole = |v: f32| ((1.0..=16.0).contains(&v) && v.fract() == 0.0).then_some(v as usize); + let (Some(sh), Some(sw)) = (whole(scales.get(2).copied().unwrap_or(0.0)), whole(scales.get(3).copied().unwrap_or(0.0))) else { + return unsupported("a Resize by a scale that is not a whole number from 1 to 16"); + }; + if scales.first() != Some(&1.0) || scales.get(1) != Some(&1.0) { + return unsupported("a Resize that scales channels"); + } + let (oh, ow) = (h * sh, w * sw); + let mut out = Tensor::zeros(vec![1, c, oh, ow])?; + for (plane_in, plane_out) in x.data.chunks_exact((h * w).max(1)).zip(out.data.chunks_exact_mut((oh * ow).max(1))) { + for (oy, row_out) in plane_out.chunks_exact_mut(ow.max(1)).enumerate() { + let row_in = plane_in.get((oy / sh) * w..(oy / sh + 1) * w).unwrap_or(&[]); + for (ox, o) in row_out.iter_mut().enumerate() { + *o = row_in.get(ox / sw).copied().unwrap_or(0.0); + } + } + } + Ok(out) +} + +fn transpose(node: &Node, x: &Tensor) -> Result { + let default_perm: Vec = (0..x.shape.len() as i64).rev().collect(); + let perm: Vec = node + .ints("perm") + .unwrap_or(&default_perm) + .iter() + .map(|p| usize::try_from(*p).map_err(|_| NetError::Shape("a negative permutation".into()))) + .collect::>()?; + let rank = x.shape.len(); + let mut seen = vec![false; rank]; + if perm.len() != rank || perm.iter().any(|p| *p >= rank || std::mem::replace(seen.get_mut(*p).unwrap_or(&mut true), true)) { + return shape_err("the Transpose permutation does not match the tensor"); + } + let out_shape: Vec = perm.iter().filter_map(|p| x.shape.get(*p).copied()).collect(); + // strides of the input, in elements + let mut in_strides = vec![1usize; rank]; + for i in (0..rank.saturating_sub(1)).rev() { + let v = in_strides.get(i + 1).copied().unwrap_or(1) * x.shape.get(i + 1).copied().unwrap_or(1); + *in_strides.get_mut(i).ok_or_else(|| NetError::Shape("missing transpose stride".into()))? = v; + } + let strides_for_out: Vec = perm.iter().filter_map(|p| in_strides.get(*p).copied()).collect(); + let mut out = Tensor::zeros(out_shape.clone())?; + let mut counter = vec![0usize; rank]; + let mut src = 0usize; + for o in out.data.iter_mut() { + *o = x.data.get(src).copied().unwrap_or(0.0); + // advance the multi-index of the output and the matching input offset + for d in (0..rank).rev() { + let (Some(c), Some(len), Some(stride)) = (counter.get_mut(d), out_shape.get(d), strides_for_out.get(d)) else { break }; + *c += 1; + src += stride; + if *c < *len { + break; + } + src -= stride * *len; + *c = 0; + } + } + Ok(out) +} + +fn reshape(node: &Node, x: &Tensor, g: &Graph) -> Result { + let Some(Data::I64(dims)) = node.inputs.get(1).and_then(|n| g.weights.get(n)).map(|w| &w.data) else { + return unsupported("a Reshape with a computed shape"); + }; + let total = x.data.len(); + let mut shape: Vec = Vec::with_capacity(dims.len()); + let mut infer = None; + for (i, d) in dims.iter().enumerate() { + match *d { + -1 if infer.is_none() => { + infer = Some(i); + shape.push(1); + } + 0 => shape.push(x.shape.get(i).copied().ok_or_else(|| NetError::Shape("a Reshape copies a missing dimension".into()))?), + d if d > 0 => shape.push(usize::try_from(d).map_err(|_| NetError::Shape("a Reshape dimension is too large".into()))?), + _ => return shape_err("a bad Reshape dimension"), + } + } + let known = shape.iter().try_fold(1usize, |a, d| a.checked_mul(*d)).filter(|k| *k > 0); + if let (Some(i), Some(known)) = (infer, known) { + if !total.is_multiple_of(known) { + return shape_err("the Reshape does not divide evenly"); + } + if let Some(s) = shape.get_mut(i) { + *s = total / known; + } + } + Tensor::new(shape, x.data.clone()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::graph::{Attr, Weight}; + + fn node(op: &str, inputs: &[&str], attrs: &[(&str, Attr)]) -> Node { + Node { + op: op.into(), + inputs: inputs.iter().map(|s| s.to_string()).collect(), + outputs: vec!["y".into()], + attrs: attrs.iter().map(|(k, v)| (k.to_string(), v.clone())).collect(), + } + } + fn graph(nodes: Vec, weights: Vec<(&str, Weight)>) -> Graph { + Graph { + nodes, + weights: weights.into_iter().map(|(k, v)| (k.to_string(), v)).collect(), + inputs: vec![("x".into(), vec![1, 1, 3, 3])], + outputs: vec!["y".into()], + } + } + fn w(dims: &[usize], v: Vec) -> Weight { + Weight { dims: dims.to_vec(), data: Data::F32(v) } + } + fn input(c: usize, h: usize, wd: usize, f: impl Fn(usize) -> f32) -> Tensor { + Tensor::new(vec![1, c, h, wd], (0..c * h * wd).map(f).collect()).unwrap() + } + fn run(g: Graph, x: Tensor) -> Result { + Ok(Net::new(g)?.run(x)?.remove(0).1) + } + + /// A direct, slow reference convolution to check the fast one against. + fn reference_conv(x: &Tensor, wd: &[f32], dims: [usize; 4], bias: &[f32], s: usize, p: usize, groups: usize) -> Tensor { + let (c, h, wi) = (x.shape[1], x.shape[2], x.shape[3]); + let [co_n, cin_g, kh, kw] = dims; + let (oh, ow) = ((h + 2 * p - kh) / s + 1, (wi + 2 * p - kw) / s + 1); + let mut out = vec![0.0f32; co_n * oh * ow]; + for co in 0..co_n { + let g = co / (co_n / groups); + for oy in 0..oh { + for ox in 0..ow { + let mut acc = bias[co]; + for cl in 0..cin_g { + for ky in 0..kh { + for kx in 0..kw { + let (iy, ix) = ((oy * s + ky) as isize - p as isize, (ox * s + kx) as isize - p as isize); + if iy >= 0 && ix >= 0 && (iy as usize) < h && (ix as usize) < wi { + acc += + x.data[((g * cin_g + cl) * h + iy as usize) * wi + ix as usize] * wd[((co * cin_g + cl) * kh + ky) * kw + kx]; + } + } + } + } + out[(co * oh + oy) * ow + ox] = acc; + } + } + } + let _ = c; + Tensor::new(vec![1, co_n, oh, ow], out).unwrap() + } + + #[test] + fn convolutions_match_a_slow_reference() { + // (cin, cout, groups, kernel, stride, pad, h, w) + for (cin, cout, groups, k, s, p, h, wi) in [ + (3, 4, 1, 3, 1, 1, 7, 9), + (3, 4, 1, 3, 2, 1, 8, 8), + (4, 4, 4, 3, 1, 1, 6, 5), + (2, 6, 1, 1, 1, 0, 5, 5), + (4, 8, 2, 3, 1, 1, 5, 6), + (1, 1, 1, 3, 1, 0, 3, 3), + ] { + let cin_g = cin / groups; + let wd: Vec = (0..cout * cin_g * k * k).map(|i| ((i * 37 % 11) as f32 - 5.0) * 0.1).collect(); + let bias: Vec = (0..cout).map(|i| i as f32 * 0.25 - 0.5).collect(); + let x = input(cin, h, wi, |i| ((i * 13 % 17) as f32 - 8.0) * 0.2); + let attrs = [("group", Attr::Int(groups as i64)), ("strides", Attr::Ints(vec![s as i64; 2])), ("pads", Attr::Ints(vec![p as i64; 4]))]; + let g = graph( + vec![node("Conv", &["x", "w", "b"], &attrs)], + vec![("w", w(&[cout, cin_g, k, k], wd.clone())), ("b", w(&[cout], bias.clone()))], + ); + let got = run(g, x.clone()).unwrap(); + let want = reference_conv(&x, &wd, [cout, cin_g, k, k], &bias, s, p, groups); + assert_eq!(got.shape, want.shape, "{cin} {cout} {groups} {k} {s} {p}"); + for (a, b) in got.data.iter().zip(&want.data) { + assert!((a - b).abs() < 1e-4, "{cin} {cout} {groups} {k} {s} {p}: {a} vs {b}"); + } + } + } + + #[test] + fn the_small_operators() { + let x = input(1, 4, 4, |i| i as f32 - 6.0); + // Relu, Sigmoid + let r = run(graph(vec![node("Relu", &["x"], &[])], vec![]), x.clone()).unwrap(); + assert!(r.data.iter().all(|v| *v >= 0.0) && r.data.get(7) == Some(&1.0)); + let s = run(graph(vec![node("Sigmoid", &["x"], &[])], vec![]), x.clone()).unwrap(); + assert!((s.data[6] - 0.5).abs() < 1e-6 && s.data.iter().all(|v| (0.0..=1.0).contains(v))); + // MaxPool 2x2/2 + let m = run( + graph(vec![node("MaxPool", &["x"], &[("kernel_shape", Attr::Ints(vec![2, 2])), ("strides", Attr::Ints(vec![2, 2]))])], vec![]), + x.clone(), + ) + .unwrap(); + assert_eq!((m.shape.clone(), m.data.clone()), (vec![1, 1, 2, 2], vec![-1.0, 1.0, 7.0, 9.0])); + // Resize x2 nearest + let z = run(graph(vec![node("Resize", &["x", "", "s"], &[])], vec![("s", w(&[4], vec![1.0, 1.0, 2.0, 2.0]))]), input(1, 2, 2, |i| i as f32)) + .unwrap(); + assert_eq!(z.shape, vec![1, 1, 4, 4]); + assert_eq!(&z.data[..8], &[0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0]); + // Add + let a = run(graph(vec![node("Add", &["x", "x"], &[])], vec![]), x.clone()).unwrap(); + assert_eq!(a.data[5], 2.0 * x.data[5]); + // Transpose [0,2,3,1] then Reshape [1,-1,2] + let t = run(graph(vec![node("Transpose", &["x"], &[("perm", Attr::Ints(vec![0, 2, 3, 1]))])], vec![]), input(2, 2, 3, |i| i as f32)).unwrap(); + assert_eq!(t.shape, vec![1, 2, 3, 2]); + assert_eq!(t.data, vec![0.0, 6.0, 1.0, 7.0, 2.0, 8.0, 3.0, 9.0, 4.0, 10.0, 5.0, 11.0]); + let rs = run( + graph(vec![node("Reshape", &["x", "s"], &[])], vec![("s", Weight { dims: vec![3], data: Data::I64(vec![1, -1, 2]) })]), + input(2, 2, 3, |i| i as f32), + ) + .unwrap(); + assert_eq!(rs.shape, vec![1, 6, 2]); + } + + #[test] + fn unsupported_and_inconsistent_models_are_errors() { + let x = input(1, 3, 3, |i| i as f32); + for (what, g) in [ + ("unknown op", graph(vec![node("Softmax", &["x"], &[])], vec![])), + ( + "dilated conv", + graph(vec![node("Conv", &["x", "w"], &[("dilations", Attr::Ints(vec![2, 2]))])], vec![("w", w(&[1, 1, 1, 1], vec![1.0]))]), + ), + ("conv without weights", graph(vec![node("Conv", &["x", "w"], &[])], vec![])), + ("channels do not match", graph(vec![node("Conv", &["x", "w"], &[])], vec![("w", w(&[1, 5, 1, 1], vec![1.0; 5]))])), + ("kernel too big", graph(vec![node("Conv", &["x", "w"], &[])], vec![("w", w(&[1, 1, 5, 5], vec![1.0; 25]))])), + ("zero group", graph(vec![node("Conv", &["x", "w"], &[("group", Attr::Int(0))])], vec![("w", w(&[1, 1, 1, 1], vec![1.0]))])), + ("zero stride", graph(vec![node("Conv", &["x", "w"], &[("strides", Attr::Ints(vec![0, 0]))])], vec![("w", w(&[1, 1, 1, 1], vec![1.0]))])), + ( + "negative pad", + graph(vec![node("Conv", &["x", "w"], &[("pads", Attr::Ints(vec![-1, 0, 0, 0]))])], vec![("w", w(&[1, 1, 1, 1], vec![1.0]))]), + ), + ("bad perm", graph(vec![node("Transpose", &["x"], &[("perm", Attr::Ints(vec![0, 0, 1, 2]))])], vec![])), + ( + "huge reshape", + graph(vec![node("Reshape", &["x", "s"], &[])], vec![("s", Weight { dims: vec![2], data: Data::I64(vec![i64::MAX, 1]) })]), + ), + ("reshape mismatch", graph(vec![node("Reshape", &["x", "s"], &[])], vec![("s", Weight { dims: vec![2], data: Data::I64(vec![4, 4]) })])), + ("non-integer resize", graph(vec![node("Resize", &["x", "", "s"], &[])], vec![("s", w(&[4], vec![1.0, 1.0, 1.5, 1.5]))])), + ("add that broadcasts", graph(vec![node("Add", &["x", "c"], &[])], vec![])), + ] { + assert!(run(g, x.clone()).is_err(), "{what}"); + } + } +} + +#[cfg(test)] +mod recognition_tests; diff --git a/crates/faces/src/net/kernels.rs b/crates/faces/src/net/kernels.rs new file mode 100644 index 00000000..73cc72b9 --- /dev/null +++ b/crates/faces/src/net/kernels.rs @@ -0,0 +1,306 @@ +//! Safe CPU kernels shared by detection and recognition. No inner thread pool. +use super::*; +use ndarray::{ArrayView2, ArrayViewMut2, linalg::general_mat_mul}; +const BLOCK: usize = 256; +#[derive(Default)] +pub(super) struct Scratch { + lower: Vec, + product: Vec, +} +fn storage(v: &mut Vec, n: usize) -> Result<&mut [f32]> { + if n > MAX_TENSOR { + return shape_err("convolution scratch exceeds the limit"); + } + if v.len() < n { + v.try_reserve_exact(n - v.len()).map_err(|_| NetError::Shape("not enough memory for convolution scratch".into()))?; + v.resize(n, 0.0); + } + v.get_mut(..n).ok_or_else(|| NetError::Shape("missing convolution scratch".into())) +} +/// The safe matrix views assert lengths, so validate exact slices before creating any view. +fn multiply(w: &[f32], x: &[f32], y: &mut [f32], m: usize, k: usize, n: usize) -> Result<()> { + if m == 0 || k == 0 || n == 0 || m.checked_mul(k) != Some(w.len()) || k.checked_mul(n) != Some(x.len()) || m.checked_mul(n) != Some(y.len()) { + return shape_err("invalid GEMM dimensions"); + } + let a = ArrayView2::from_shape((m, k), w).map_err(|e| NetError::Shape(e.to_string()))?; + let b = ArrayView2::from_shape((k, n), x).map_err(|e| NetError::Shape(e.to_string()))?; + let mut c = ArrayViewMut2::from_shape((m, n), y).map_err(|e| NetError::Shape(e.to_string()))?; + general_mat_mul(1.0, &a, &b, 0.0, &mut c); + Ok(()) +} +/// For equal-width, unit-stride maps a shifted plane is one contiguous copy. Only horizontal +/// border stripes need masking; clipping the flat source range also removes vertical padding. +fn shifted_plane(plane: &[f32], width: usize, ky: usize, kx: usize, pt: usize, pl: usize, start: usize, row: &mut [f32]) -> Result<()> { + let end = start + row.len(); + let delta = (ky as i128 - pt as i128) * width as i128 + kx as i128 - pl as i128; + let lo = (-delta).clamp(start as i128, end as i128); + let hi = (plane.len() as i128 - delta).clamp(start as i128, end as i128); + row.fill(0.0); + if lo < hi { + let from = usize::try_from(lo + delta).map_err(|_| NetError::Shape("shifted source overflow".into()))?; + let at = usize::try_from(lo - start as i128).map_err(|_| NetError::Shape("shifted destination overflow".into()))?; + let len = usize::try_from(hi - lo).map_err(|_| NetError::Shape("shifted length overflow".into()))?; + row.get_mut(at..at + len) + .ok_or_else(|| NetError::Shape("short shifted destination".into()))? + .copy_from_slice(plane.get(from..from + len).ok_or_else(|| NetError::Shape("short shifted source".into()))?); + } + let left = pl.saturating_sub(kx).min(width); + let right = (width + pl).saturating_sub(kx).min(width); + if left > 0 || right < width { + for y in start / width..=(end - 1) / width { + for (a, b) in [(y * width, y * width + left), (y * width + right, (y + 1) * width)] { + let lo = a.max(start); + let hi = b.min(end); + if lo < hi { + row.get_mut(lo - start..hi - start).ok_or_else(|| NetError::Shape("short border stripe".into()))?.fill(0.0); + } + } + } + } + Ok(()) +} +pub(super) fn conv(n: &Node, x: &Tensor, g: &Graph, scratch: &mut Scratch) -> Result { + let (cin, h, w) = chw(x)?; + let wt = shape::weight(n, g, 1)?; + let wd = shape::floats(n, g, 1)?; + let [cout, cg, kh, kw] = wt.dims.as_slice() else { + return shape_err("invalid convolution weights"); + }; + let (cout, cg, kh, kw) = (*cout, *cg, *kh, *kw); + let groups = usize::try_from(n.int("group").unwrap_or(1)).map_err(|_| NetError::Shape("invalid group".into()))?; + if cg == 1 { + return conv_direct(n, x, g); + } + let (sh, sw) = pair(n, "strides", 1)?; + let (pt, pl, pb, pr) = pads(n)?; + let ph = h.checked_add(pt).and_then(|v| v.checked_add(pb)).ok_or_else(|| NetError::Shape("padding overflow".into()))?; + let pw = w.checked_add(pl).and_then(|v| v.checked_add(pr)).ok_or_else(|| NetError::Shape("padding overflow".into()))?; + if ph < kh || pw < kw || groups == 0 || cin != groups * cg { + return shape_err("invalid convolution shape"); + } + let (oh, ow) = ((ph - kh) / sh + 1, (pw - kw) / sw + 1); + let mut out = Tensor::zeros(vec![1, cout, oh, ow])?; + let pixels = oh * ow; + let k = cg * kh * kw; + let m = cout / groups; + let bias = if n.inputs.get(2).is_some_and(|s| !s.is_empty()) { Some(shape::floats(n, g, 2)?) } else { None }; + for group in 0..groups { + let weights = wd.get(group * m * k..(group + 1) * m * k).ok_or_else(|| NetError::Shape("short convolution weights".into()))?; + // A pointwise convolution is already a matrix: no lowering, copy or spatial blocking needed. + if kh == 1 && kw == 1 && sh == 1 && sw == 1 && (pt, pl, pb, pr) == (0, 0, 0, 0) { + let input = x.data.get(group * cg * pixels..(group + 1) * cg * pixels).ok_or_else(|| NetError::Shape("short pointwise input".into()))?; + let output = + out.data.get_mut(group * m * pixels..(group + 1) * m * pixels).ok_or_else(|| NetError::Shape("short pointwise output".into()))?; + multiply(weights, input, output, m, cg, pixels)?; + for (co, plane) in output.chunks_exact_mut(pixels).enumerate() { + let b = bias.and_then(|v| v.get(group * m + co)).copied().unwrap_or(0.0); + if b != 0.0 { + for v in plane { + *v += b; + } + } + } + continue; + } + for start in (0..pixels).step_by(BLOCK) { + let len = BLOCK.min(pixels - start); + let lower = storage(&mut scratch.lower, k * len)?; + for cl in 0..cg { + let plane = x + .data + .get((group * cg + cl) * h * w..(group * cg + cl + 1) * h * w) + .ok_or_else(|| NetError::Shape("short convolution input".into()))?; + for ky in 0..kh { + for kx in 0..kw { + let row = lower + .get_mut((cl * kh * kw + ky * kw + kx) * len..(cl * kh * kw + ky * kw + kx + 1) * len) + .ok_or_else(|| NetError::Shape("short im2col row".into()))?; + if sh == 1 && sw == 1 && ow == w { + shifted_plane( + plane, + w, + ky, + kx, + pt, + pl, + start, + row.get_mut(..len).ok_or_else(|| NetError::Shape("short lower panel".into()))?, + )?; + continue; + } + let mut at = 0; + while at < len { + let p = start + at; + let (oy, ox) = (p / ow, p % ow); + let span = (ow - ox).min(len - at); + if let Some(yy) = (oy * sh + ky).checked_sub(pt).filter(|&v| v < h) { + let lo = ox.max(pl.saturating_sub(kx).div_ceil(sw)); + let hi = (ox + span).min((w + pl).saturating_sub(kx).div_ceil(sw)); + if lo < hi { + let from = yy * w + lo * sw + kx - pl; + let to = at + lo - ox; + row.get_mut(at..to).ok_or_else(|| NetError::Shape("short im2col padding".into()))?.fill(0.0); + let src = plane + .get(from..from + (hi - lo - 1) * sw + 1) + .ok_or_else(|| NetError::Shape("short im2col source row".into()))?; + let dst = row.get_mut(to..to + hi - lo).ok_or_else(|| NetError::Shape("short im2col row".into()))?; + if sw == 1 { + dst.copy_from_slice(src); + } else { + for (o, &v) in dst.iter_mut().zip(src.iter().step_by(sw)) { + *o = v; + } + } + row.get_mut(to + hi - lo..at + span).ok_or_else(|| NetError::Shape("short im2col padding".into()))?.fill(0.0); + } else { + row.get_mut(at..at + span).ok_or_else(|| NetError::Shape("short im2col row".into()))?.fill(0.0); + } + } else { + row.get_mut(at..at + span).ok_or_else(|| NetError::Shape("short im2col row".into()))?.fill(0.0); + } + at += span; + } + } + } + } + let product = storage(&mut scratch.product, m * len)?; + multiply(weights, lower, product, m, k, len)?; + for (co, row) in product.chunks_exact(len).enumerate() { + let channel = group * m + co; + let b = bias.and_then(|v| v.get(channel)).copied().unwrap_or(0.0); + let dest = out + .data + .get_mut(channel * pixels + start..channel * pixels + start + len) + .ok_or_else(|| NetError::Shape("short convolution output".into()))?; + for (to, &value) in dest.iter_mut().zip(row) { + *to = value + b; + } + } + } + } + Ok(out) +} +/// Flat offset under ONNX's trailing-dimension broadcasting; shapes are checked before execution. +fn broadcast_offset(mut at: usize, out: &[usize], input: &[usize]) -> usize { + let mut offset = 0; + let mut stride = 1; + for (axis, &size) in out.iter().enumerate().rev() { + let coordinate = at % size; + at /= size; + if let Some(&d) = axis.checked_sub(out.len() - input.len()).and_then(|i| input.get(i)) { + if d != 1 { + offset += coordinate * stride; + } + stride *= d; + } + } + offset +} +pub(super) fn binary(op: &str, a: &Tensor, b: &Tensor) -> Result { + let dims = shape::broadcast(&a.shape, &b.shape)?; + let mut out = Tensor::zeros(dims)?; + let f = |x: f32, y: f32| match op { + "Add" => x + y, + "Sub" => x - y, + _ => x * y, + }; + if a.shape == b.shape { + for ((o, &x), &y) in out.data.iter_mut().zip(&a.data).zip(&b.data) { + *o = f(x, y); + } + } else if b.data.len() == 1 && a.shape == out.shape { + let y = *b.data.first().ok_or_else(|| NetError::Shape("empty scalar".into()))?; + for (o, &x) in out.data.iter_mut().zip(&a.data) { + *o = f(x, y); + } + } else { + for (i, o) in out.data.iter_mut().enumerate() { + let x = a.data.get(broadcast_offset(i, &out.shape, &a.shape)).ok_or_else(|| NetError::Shape("bad broadcast A".into()))?; + let y = b.data.get(broadcast_offset(i, &out.shape, &b.shape)).ok_or_else(|| NetError::Shape("bad broadcast B".into()))?; + *o = f(*x, *y); + } + } + Ok(out) +} +pub(super) fn prelu(n: &Node, x: &Tensor, g: &Graph) -> Result { + let wt = shape::weight(n, g, 1)?; + let slope = shape::floats(n, g, 1)?; + let mut out = Tensor::zeros(x.shape.clone())?; + let channels = x.shape.get(1).copied().unwrap_or(1); + let mut expanded = vec![1; x.shape.len().saturating_sub(wt.dims.len())]; + expanded.extend_from_slice(&wt.dims); + let channel_slope = expanded.iter().enumerate().all(|(i, &d)| d == 1 || i == 1) && slope.len() == channels; + if channel_slope || slope.len() == 1 { + let plane = shape::elements(x.shape.get(2..).unwrap_or(&[]))?; + for (c, (src, dst)) in x.data.chunks_exact(plane).zip(out.data.chunks_exact_mut(plane)).enumerate() { + let alpha = *slope.get(if slope.len() == 1 { 0 } else { c % channels }).ok_or_else(|| NetError::Shape("short PRelu slope".into()))?; + for (to, &v) in dst.iter_mut().zip(src) { + *to = if v < 0.0 { v * alpha } else { v }; + } + } + } else { + for (i, (o, &v)) in out.data.iter_mut().zip(&x.data).enumerate() { + let alpha = *slope.get(broadcast_offset(i, &x.shape, &wt.dims)).ok_or_else(|| NetError::Shape("bad PRelu broadcast".into()))?; + *o = if v < 0.0 { v * alpha } else { v }; + } + } + Ok(out) +} +pub(super) fn batch_norm(n: &Node, x: &Tensor, g: &Graph) -> Result { + let channels = *x.shape.get(1).ok_or_else(|| NetError::Shape("BatchNormalization rank".into()))?; + let (scale, bias, mean, var) = (shape::floats(n, g, 1)?, shape::floats(n, g, 2)?, shape::floats(n, g, 3)?, shape::floats(n, g, 4)?); + let eps = shape::float(n, "epsilon", 1e-5); + let plane = shape::elements(x.shape.get(2..).unwrap_or(&[]))?; + let mut out = Tensor::zeros(x.shape.clone())?; + for (j, (src, dst)) in x.data.chunks_exact(plane).zip(out.data.chunks_exact_mut(plane)).enumerate() { + let c = j % channels; + let coefficient = *scale.get(c).ok_or_else(|| NetError::Shape("short BN scale".into()))? + / (var.get(c).copied().ok_or_else(|| NetError::Shape("short BN variance".into()))? + eps).sqrt(); + let offset = bias.get(c).copied().ok_or_else(|| NetError::Shape("short BN bias".into()))? + - mean.get(c).copied().ok_or_else(|| NetError::Shape("short BN mean".into()))? * coefficient; + for (to, &value) in dst.iter_mut().zip(src) { + *to = value * coefficient + offset; + } + } + Ok(out) +} +pub(super) fn gemm(n: &Node, a: &Tensor, g: &Graph) -> Result { + let wt = shape::weight(n, g, 1)?; + let b = shape::floats(n, g, 1)?; + let [ar, ac] = a.shape.as_slice() else { + return shape_err("Gemm A is not a matrix"); + }; + let [br, bc] = wt.dims.as_slice() else { + return shape_err("Gemm B is not a matrix"); + }; + if ar.checked_mul(*ac) != Some(a.data.len()) || br.checked_mul(*bc) != Some(b.len()) { + return shape_err("invalid matrix buffers"); + } + let av = ArrayView2::from_shape((*ar, *ac), a.data.as_slice()).map_err(|e| NetError::Shape(e.to_string()))?; + let bv = ArrayView2::from_shape((*br, *bc), b).map_err(|e| NetError::Shape(e.to_string()))?; + let av = if n.int("transA").unwrap_or(0) != 0 { av.reversed_axes() } else { av }; + let bv = if n.int("transB").unwrap_or(0) != 0 { bv.reversed_axes() } else { bv }; + if av.ncols() != bv.nrows() { + return shape_err("incompatible matrices"); + } + let mut out = Tensor::zeros(vec![av.nrows(), bv.ncols()])?; + let mut view = ArrayViewMut2::from_shape((av.nrows(), bv.ncols()), out.data.as_mut_slice()).map_err(|e| NetError::Shape(e.to_string()))?; + general_mat_mul(shape::float(n, "alpha", 1.0), &av, &bv, 0.0, &mut view); + if n.inputs.get(2).is_some_and(|s| !s.is_empty()) { + let wt = shape::weight(n, g, 2)?; + let bias = shape::floats(n, g, 2)?; + let beta = shape::float(n, "beta", 1.0); + for (i, to) in out.data.iter_mut().enumerate() { + *to += beta * bias.get(broadcast_offset(i, &out.shape, &wt.dims)).copied().ok_or_else(|| NetError::Shape("short matrix bias".into()))?; + } + } + Ok(out) +} +pub(super) fn flatten(n: &Node, x: &Tensor) -> Result { + let axis = shape::axis(n, x.shape.len())?; + Tensor::new(vec![shape::elements(x.shape.get(..axis).unwrap_or(&[]))?, shape::elements(x.shape.get(axis..).unwrap_or(&[]))?], x.data.clone()) +} +pub(super) fn global_average(x: &Tensor) -> Result { + let (c, h, w) = chw(x)?; + Tensor::new(vec![1, c, 1, 1], x.data.chunks_exact(h * w).map(|v| v.iter().sum::() / (h * w) as f32).collect()) +} diff --git a/crates/faces/src/net/recognition_tests.rs b/crates/faces/src/net/recognition_tests.rs new file mode 100644 index 00000000..e245bb33 --- /dev/null +++ b/crates/faces/src/net/recognition_tests.rs @@ -0,0 +1,337 @@ +use super::*; +use crate::graph::{Attr, Weight}; +fn weight(shape: &[usize], values: Vec) -> Weight { + Weight { dims: shape.to_vec(), data: Data::F32(values) } +} +fn graph(op: &str, extra: &[&str], weights: Vec<(&str, Weight)>, attrs: Vec<(&str, Attr)>) -> Graph { + Graph { + nodes: vec![Node { + op: op.into(), + inputs: std::iter::once("x").chain(extra.iter().copied()).map(str::to_string).collect(), + outputs: vec!["y".into()], + attrs: attrs.into_iter().map(|(s, a)| (s.into(), a)).collect(), + }], + weights: weights.into_iter().map(|(s, w)| (s.into(), w)).collect(), + inputs: vec![("x".into(), vec![1, 3, 2, 2])], + outputs: vec!["y".into()], + } +} +fn execute(g: Graph, shape: &[usize], values: Vec) -> Tensor { + Net::new(g).unwrap().run(Tensor::new(shape.to_vec(), values).unwrap()).unwrap().remove(0).1 +} +#[test] +fn batch_norm_matches_the_inference_formula_for_images_and_matrices() { + for shape in [vec![1, 3, 2, 5], vec![2, 3]] { + let values: Vec<_> = (0..shape.iter().product()).map(|i| i as f32 * 0.07 - 0.8).collect(); + let g = graph( + "BatchNormalization", + &["scale", "bias", "mean", "var"], + vec![ + ("scale", weight(&[3], vec![0.7, -1.1, 2.3])), + ("bias", weight(&[3], vec![-0.4, 0.6, 0.2])), + ("mean", weight(&[3], vec![0.5, -0.3, 0.9])), + ("var", weight(&[3], vec![0.2, 0.0, 1.4])), + ], + vec![("epsilon", Attr::Float(1e-3))], + ); + let out = execute(g.clone(), &shape, values.clone()); + let spatial: usize = shape.get(2..).unwrap_or(&[]).iter().product(); + for (i, (&value, &got)) in values.iter().zip(&out.data).enumerate() { + let c = (i / spatial) % 3; + let [s, b, m, v] = ["scale", "bias", "mean", "var"].map(|name| { + let Data::F32(v) = &g.weights[name].data else { panic!() }; + v[c] + }); + let want = (value - m) * s / (v + 1e-3).sqrt() + b; + assert!((got - want).abs() < 1e-5); + } + } +} +#[test] +fn prelu_obeys_trailing_broadcasting_including_spatial_slopes() { + for (shape, slopes) in + [(vec![3, 1, 1], vec![0.1, 0.2, 0.3]), (vec![1, 3, 1, 1], vec![0.1, 0.2, 0.3]), (vec![1], vec![0.3]), (vec![2, 2], vec![0.1, 0.2, 0.3, 0.4])] + { + let values: Vec<_> = (0..12).map(|i| i as f32 - 6.0).collect(); + let out = execute(graph("PRelu", &["s"], vec![("s", weight(&shape, slopes.clone()))], vec![]), &[1, 3, 2, 2], values.clone()); + for (i, (&v, &got)) in values.iter().zip(&out.data).enumerate() { + let alpha = if slopes.len() == 1 { + slopes[0] + } else if slopes.len() == 3 { + slopes[i / 4] + } else { + slopes[i % 4] + }; + assert_eq!(got, if v < 0.0 { v * alpha } else { v }); + } + } +} +#[test] +fn gemm_transpositions_scalars_vectors_matrices_and_optional_bias() { + for ta in [0, 1] { + for tb in [0, 1] { + for bias_shape in [vec![], vec![2], vec![1, 2], vec![2, 1], vec![2, 2]] { + let (ash, bsh) = (if ta == 0 { vec![2, 3] } else { vec![3, 2] }, if tb == 0 { vec![3, 2] } else { vec![2, 3] }); + let a: Vec<_> = (0..6).map(|i| i as f32 * 0.4 - 0.8).collect(); + let b: Vec<_> = (0..6).map(|i| i as f32 * 0.1 - 0.3).collect(); + let bias: Vec<_> = (0..bias_shape.iter().product()).map(|i| i as f32 * 0.25 + 0.3).collect(); + let g = graph( + "Gemm", + &["w", "b"], + vec![("w", weight(&bsh, b.clone())), ("b", weight(&bias_shape, bias.clone()))], + vec![("transA", Attr::Int(ta)), ("transB", Attr::Int(tb)), ("alpha", Attr::Float(0.7)), ("beta", Attr::Float(-0.4))], + ); + let out = execute(g, &ash, a.clone()); + assert_eq!(out.shape, vec![2, 2]); + for i in 0..2 { + for j in 0..2 { + let sum = (0..3) + .map(|k| a[if ta == 0 { i * 3 + k } else { k * 2 + i }] * b[if tb == 0 { k * 2 + j } else { j * 3 + k }]) + .sum::(); + let off = match bias_shape.as_slice() { + [] => 0, + [2] | [1, 2] => j, + [2, 1] => i, + _ => i * 2 + j, + }; + assert!((out.data[i * 2 + j] - (0.7 * sum - 0.4 * bias[off])).abs() < 1e-6); + } + } + } + } + } + let out = execute(graph("Gemm", &["w"], vec![("w", weight(&[2, 1], vec![2.0, 3.0]))], vec![]), &[1, 2], vec![4.0, 5.0]); + assert_eq!(out.data, vec![23.0]); +} +#[test] +fn arithmetic_broadcasts_and_keeps_operand_order() { + for op in ["Add", "Sub", "Mul"] { + let out = + execute(graph(op, &["w"], vec![("w", weight(&[3, 1], vec![0.5, 1.5, 2.5]))], vec![]), &[1, 3, 2], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + for (i, &got) in out.data.iter().enumerate() { + let x = (i + 1) as f32; + let y = (i / 2) as f32 + 0.5; + let want = match op { + "Add" => x + y, + "Sub" => x - y, + _ => x * y, + }; + assert_eq!(got, want); + } + } +} +#[test] +fn flatten_negative_axis_and_default_transpose_are_correct() { + let out = execute(graph("Flatten", &[], vec![], vec![("axis", Attr::Int(-1))]), &[1, 3, 2, 5], (0..30).map(|i| i as f32).collect()); + assert_eq!(out.shape, vec![6, 5]); + let out = execute(graph("Transpose", &[], vec![], vec![]), &[2, 3], (0..6).map(|i| i as f32).collect()); + assert_eq!(out.shape, vec![3, 2]); + assert_eq!(out.data, vec![0.0, 3.0, 1.0, 4.0, 2.0, 5.0]); +} +#[test] +fn all_recognition_operators_run_on_seeded_synthetic_networks() { + let input = Tensor::new(vec![1, 3, 112, 112], (0..3 * 112 * 112).map(|i| (i * 37 % 256) as f32).collect()).unwrap(); + for seed in 0..20 { + let net = Net::new(crate::graph::load(&crate::synthetic::recognition_model(seed)).unwrap()).unwrap(); + let a = net.run(input.clone()).unwrap(); + let b = net.run(input.clone()).unwrap(); + assert_eq!(a, b); + assert_eq!(a[0].1.shape, vec![1, 8]); + } +} +#[test] +fn hostile_attributes_weights_shapes_and_public_tensors_return_errors() { + for (name, attr) in [ + ("strides", Attr::Ints(vec![0, 1])), + ("pads", Attr::Ints(vec![i64::MAX; 4])), + ("group", Attr::Int(0)), + ("kernel_shape", Attr::Ints(vec![7, 7])), + ("strides", Attr::Float(1.0)), + ("new_attribute", Attr::Int(1)), + ] { + let g = graph("Conv", &["w"], vec![("w", weight(&[3, 3, 1, 1], vec![0.1; 9]))], vec![(name, attr)]); + let result = Net::new(g).and_then(|n| n.validate_input(&[1, 3, 3, 3])); + assert!(result.is_err(), "{name}"); + } + for op in ["Loop", "MatMul", "If", "Concat", "Unknown"] { + assert!(Net::new(graph(op, &[], vec![], vec![])).is_err()); + } + for value in [f32::NAN, f32::INFINITY] { + assert!(Net::new(graph("PRelu", &["w"], vec![("w", weight(&[1], vec![value]))], vec![])).is_err()); + } + let net = Net::new(graph("Relu", &[], vec![], vec![])).unwrap(); + for input in [ + Tensor { shape: vec![usize::MAX, 2], data: vec![] }, + Tensor { shape: vec![0, 3], data: vec![] }, + Tensor { shape: vec![1], data: vec![f32::NAN] }, + Tensor { shape: vec![1; 9], data: vec![1.0] }, + Tensor { shape: vec![2], data: vec![1.0] }, + ] { + assert!(net.run(input).is_err()); + } + let mut g = graph("Relu", &[], vec![], vec![]); + g.nodes.push(g.nodes[0].clone()); + assert!(Net::new(g).is_err()); + let mut g = graph("Relu", &[], vec![], vec![]); + g.nodes[0].inputs[0] = "future".into(); + assert!(Net::new(g).is_err()); +} +#[test] +fn malformed_normalization_prelu_matrix_dropout_and_axes_are_errors() { + let bn = || { + graph( + "BatchNormalization", + &["s", "b", "m", "v"], + vec![ + ("s", weight(&[3], vec![1.0; 3])), + ("b", weight(&[3], vec![0.0; 3])), + ("m", weight(&[3], vec![0.0; 3])), + ("v", weight(&[3], vec![1.0; 3])), + ], + vec![], + ) + }; + for mutation in 0..4 { + let mut g = bn(); + match mutation { + 0 => { + g.nodes[0].attrs.insert("training_mode".into(), Attr::Int(1)); + } + 1 => { + g.nodes[0].attrs.insert("epsilon".into(), Attr::Float(0.0)); + } + 2 => { + g.weights.insert("v".into(), weight(&[3], vec![-1.0; 3])); + } + _ => { + g.weights.insert("b".into(), weight(&[1, 3], vec![0.0; 3])); + } + }; + assert!(Net::new(g).and_then(|n| n.validate_input(&[1, 3, 2, 2])).is_err()); + } + for g in [ + graph("PRelu", &["s"], vec![("s", weight(&[3], vec![0.1; 3]))], vec![]), + graph("Gemm", &["w"], vec![("w", weight(&[4, 2], vec![0.1; 8]))], vec![]), + graph("Flatten", &[], vec![], vec![("axis", Attr::Int(i64::MIN))]), + graph("Dropout", &[], vec![], vec![("ratio", Attr::Float(1.0))]), + graph("Transpose", &[], vec![], vec![("perm", Attr::Ints(vec![0, 1, 1, 3]))]), + ] { + assert!(Net::new(g).and_then(|n| n.validate_input(&[1, 3, 2, 2])).is_err()); + } +} +#[test] +fn arithmetic_overflow_is_reported_and_next_run_succeeds() { + let net = Net::new(graph("Mul", &["w"], vec![("w", weight(&[1], vec![f32::MAX]))], vec![])).unwrap(); + assert!(net.run(Tensor::new(vec![1], vec![2.0]).unwrap()).is_err()); + assert!(net.run(Tensor::new(vec![1], vec![0.0]).unwrap()).is_ok()); +} + +#[test] +fn multiple_batches_prelu_keeps_channels_independent() { + let out = execute(graph("PRelu", &["s"], vec![("s", weight(&[3, 1, 1], vec![0.1, 0.2, 0.3]))], vec![]), &[2, 3, 2, 2], vec![-1.0; 24]); + for (i, &v) in out.data.iter().enumerate() { + assert_eq!(v, -[0.1, 0.2, 0.3][i / 4 % 3]); + } +} +#[test] +fn integer_activations_and_constant_outputs_fail_at_load() { + let integer = Weight { dims: vec![1], data: Data::I64(vec![1]) }; + assert!(Net::new(graph("Add", &["w"], vec![("w", integer.clone())], vec![])).is_err()); + let mut g = graph("Relu", &[], vec![("w", integer)], vec![]); + g.outputs = vec!["w".into()]; + assert!(Net::new(g).is_err()); +} +#[test] +fn normalisation_coefficient_and_working_memory_overflows_fail_before_running() { + let g = graph( + "BatchNormalization", + &["s", "b", "m", "v"], + vec![ + ("s", weight(&[3], vec![f32::MAX; 3])), + ("b", weight(&[3], vec![0.0; 3])), + ("m", weight(&[3], vec![f32::MAX; 3])), + ("v", weight(&[3], vec![0.0; 3])), + ], + vec![], + ); + assert!(Net::new(g).unwrap().validate_input(&[1, 3, 2, 2]).is_err()); + let g = graph("Conv", &["w"], vec![("w", weight(&[4096, 3, 1, 1], vec![0.1; 4096 * 3]))], vec![]); + assert!(Net::new(g).unwrap().validate_input(&[1, 3, 4096, 4096]).is_err()); +} + +#[test] +fn empty_broadcast_operands_are_rejected_before_allocation() { + for op in ["Add", "Sub", "Mul", "PRelu"] { + let g = graph(op, &["w"], vec![("w", weight(&[0, 3, 1, 1], vec![]))], vec![]); + assert!(Net::new(g).and_then(|n| n.validate_input(&[1, 3, 2, 2])).is_err()); + } +} + +#[test] +fn odd_maps_cross_panel_boundaries_and_reused_padding_is_zero() { + for (c, co, groups) in [(3, 11, 1), (6, 10, 2)] { + let weights: Vec<_> = (0..co * (c / groups) * 9).map(|i| (i as f32 % 19.0 - 9.0) * 0.013).collect(); + let g = graph( + "Conv", + &["w", "b"], + vec![("w", weight(&[co, c / groups, 3, 3], weights)), ("b", weight(&[co], vec![0.1; co]))], + vec![("group", Attr::Int(groups as i64)), ("strides", Attr::Ints(vec![1, 1])), ("pads", Attr::Ints(vec![2, 1, 0, 1]))], + ); + let mut scratch = kernels::Scratch::default(); + for zero in [false, true, false] { + let input = Tensor::new(vec![1, c, 23, 29], (0..c * 23 * 29).map(|i| if zero { 0.0 } else { (i as f32 % 29.0 - 14.0) * 0.12 }).collect()) + .unwrap(); + let got = kernels::conv(&g.nodes[0], &input, &g, &mut scratch).unwrap(); + let want = conv_direct(&g.nodes[0], &input, &g).unwrap(); + assert_eq!(got.shape, want.shape); + for (&a, &b) in got.data.iter().zip(&want.data) { + assert!((a - b).abs() < 1e-5, "{a} vs {b}"); + } + } + } +} + +#[test] +fn retained_convolution_scratch_counts_towards_the_working_limit() { + let network = |output_channels: usize| { + let pads = vec![("pads", Attr::Ints(vec![1, 0, 0, 0]))]; + let mut g = graph("Conv", &["w"], vec![("w", weight(&[2, 130_000, 2, 1], vec![0.0; 2 * 130_000 * 2]))], pads.clone()); + g.inputs[0].1 = vec![1, 130_000, 2, 256]; + g.weights.insert("v".into(), weight(&[output_channels, 2, 2, 1], vec![0.0; output_channels * 4])); + g.nodes.push(Node { + op: "Conv".into(), + inputs: vec!["y".into(), "v".into()], + outputs: vec!["z".into()], + attrs: pads.into_iter().map(|(s, a)| (s.into(), a)).collect(), + }); + g.outputs = vec!["z".into()]; + Net::new(g).unwrap() + }; + // The first layer retains a large im2col buffer; the second grows the product buffer. + // Both layers individually fit, but keeping both scratch peaks would exceed 512 MiB. + let shape = [1, 130_000, 2, 256]; + assert!(network(131_072).validate_input(&shape).is_err()); + assert!(network(65_536).validate_input(&shape).is_ok()); +} + +#[test] +fn constant_input_copies_count_towards_the_working_limit() { + let network = |retained_outputs: usize| { + let dims = [1, 1, 2048, 2048]; + let mut g = graph("Identity", &[], vec![("w", weight(&dims, vec![0.0; 2048 * 2048]))], vec![]); + g.inputs[0].1 = dims.to_vec(); + g.nodes.clear(); + g.outputs.clear(); + for i in 0..retained_outputs { + let name = format!("retained-{i}"); + g.nodes.push(Node { op: "Identity".into(), inputs: vec!["x".into()], outputs: vec![name.clone()], attrs: HashMap::new() }); + g.outputs.push(name); + } + g.nodes.push(Node { op: "Add".into(), inputs: vec!["x".into(), "w".into()], outputs: vec!["z".into()], attrs: HashMap::new() }); + g.outputs.push("z".into()); + Net::new(g).unwrap() + }; + // Infer only: no pixel tensors are allocated. The constant's temporary copy must be counted. + assert!(network(30).validate_input(&[1, 1, 2048, 2048]).is_err()); + assert!(network(29).validate_input(&[1, 1, 2048, 2048]).is_ok()); +} diff --git a/crates/faces/src/net/shape.rs b/crates/faces/src/net/shape.rs new file mode 100644 index 00000000..d76271bc --- /dev/null +++ b/crates/faces/src/net/shape.rs @@ -0,0 +1,496 @@ +//! Bounded, allocation-free (apart from small shape lists) graph checking. +use super::*; +use crate::graph::{Attr, Weight}; + +pub(super) fn elements(s: &[usize]) -> Result { + if s.len() > 8 || s.contains(&0) { + return shape_err("a tensor has a zero dimension or excessive rank"); + } + s.iter() + .try_fold(1usize, |n, &d| n.checked_mul(d)) + .filter(|&n| n <= MAX_TENSOR) + .ok_or_else(|| NetError::Shape("a tensor exceeds the element limit".into())) +} +pub(super) fn tensor(s: &[usize], n: usize, data: &[f32]) -> Result<()> { + if elements(s)? != n || data.iter().any(|v| !v.is_finite()) { + return shape_err("invalid tensor size or non-finite values"); + } + Ok(()) +} +pub(super) fn float(n: &Node, name: &str, default: f32) -> f32 { + match n.attrs.get(name) { + Some(Attr::Float(v)) => *v, + _ => default, + } +} +pub(super) fn weight<'a>(n: &Node, g: &'a Graph, i: usize) -> Result<&'a Weight> { + n.inputs.get(i).and_then(|s| g.weights.get(s)).ok_or_else(|| NetError::Unsupported(format!("{} needs constant input {i}", n.op))) +} +pub(super) fn floats<'a>(n: &Node, g: &'a Graph, i: usize) -> Result<&'a [f32]> { + match &weight(n, g, i)?.data { + Data::F32(v) => Ok(v), + _ => unsupported(format!("{} needs float32 weights", n.op)), + } +} +pub(super) fn check_graph(g: &Graph) -> Result<()> { + if g.inputs.len() != 1 || g.outputs.is_empty() || g.outputs.len() > 64 || g.nodes.len() > 20_000 || g.weights.len() > 20_000 { + return unsupported("one input and at most 64 outputs are required"); + } + let mut known = std::collections::HashSet::new(); + let mut total_weights = 0usize; + for (name, w) in &g.weights { + let count = w.dims.iter().try_fold(1usize, |n, &d| n.checked_mul(d)); + let len = match &w.data { + Data::F32(v) => { + if v.iter().any(|v| !v.is_finite()) { + return shape_err("non-finite weights"); + } + v.len() + } + Data::I64(v) => v.len(), + }; + if name.is_empty() + || name.len() > 4096 + || w.dims.len() > 8 + || w.dims.iter().any(|&d| d > MAX_TENSOR) + || count != Some(len) + || len > MAX_TENSOR + { + return shape_err("invalid weight name, shape or size"); + } + total_weights = total_weights + .checked_add(len) + .filter(|&v| v <= MAX_TENSOR) + .ok_or_else(|| NetError::Shape("combined weights exceed 64 Mi elements".into()))?; + known.insert(name.as_str()); + } + let input = g.inputs.first().ok_or_else(|| NetError::Shape("missing input".into()))?; + if input.0.is_empty() || !known.insert(&input.0) { + return shape_err("invalid or repeated input name"); + } + for n in &g.nodes { + if n.inputs.len() > 64 || n.outputs.len() != 1 || n.outputs.first().is_none_or(|s| s.is_empty() || s.len() > 4096) { + return unsupported(format!("{} requires one named output", n.op)); + } + if n.inputs.iter().any(|s| !s.is_empty() && !known.contains(s.as_str())) { + return shape_err(format!("{} reads an unknown or forward value", n.op)); + } + if !known.insert(n.outputs.first().ok_or_else(|| NetError::Shape("missing output".into()))?) { + return shape_err("a graph value is defined twice"); + } + } + if g.outputs.iter().any(|s| g.weights.contains_key(s)) { + return unsupported("constant graph outputs"); + } + let mut output_names = std::collections::HashSet::new(); + if g.outputs.iter().any(|s| !known.contains(s.as_str()) || !output_names.insert(s)) { + return shape_err("unknown or repeated graph output"); + } + Ok(()) +} +pub(super) fn check(n: &Node, g: &Graph) -> Result<()> { + let (min, max, attrs): (usize, usize, &[(&str, u8)]) = match n.op.as_str() { + "Conv" => (2, 3, &[("group", 0), ("kernel_shape", 2), ("strides", 2), ("pads", 2), ("dilations", 2), ("auto_pad", 3)]), + "BatchNormalization" => (5, 5, &[("epsilon", 1), ("momentum", 1), ("training_mode", 0)]), + "PRelu" | "Add" | "Sub" | "Mul" => (2, 2, &[]), + "Relu" | "Sigmoid" | "GlobalAveragePool" | "Identity" => (1, 1, &[]), + "Dropout" => (1, 1, &[("ratio", 1), ("seed", 0)]), + "Flatten" => (1, 1, &[("axis", 0)]), + "Gemm" => (2, 3, &[("alpha", 1), ("beta", 1), ("transA", 0), ("transB", 0)]), + "MaxPool" => { + (1, 1, &[("kernel_shape", 2), ("strides", 2), ("pads", 2), ("dilations", 2), ("ceil_mode", 0), ("auto_pad", 3), ("storage_order", 0)]) + } + "Resize" => ( + 3, + 4, + &[ + ("mode", 3), + ("coordinate_transformation_mode", 3), + ("nearest_mode", 3), + ("cubic_coeff_a", 1), + ("exclude_outside", 0), + ("extrapolation_value", 1), + ], + ), + "Transpose" => (1, 1, &[("perm", 2)]), + "Reshape" => (2, 2, &[("allowzero", 0)]), + op => return unsupported(format!("the operator `{op}`")), + }; + if !(min..=max).contains(&n.inputs.len()) || n.inputs.first().is_none_or(String::is_empty) { + return shape_err(format!("{} has invalid inputs", n.op)); + } + for (name, value) in &n.attrs { + let actual = match value { + Attr::Int(_) => 0, + Attr::Float(v) if v.is_finite() => 1, + Attr::Ints(_) => 2, + Attr::Str(_) => 3, + _ => 255, + }; + if attrs.iter().find(|(key, _)| *key == name).is_none_or(|(_, kind)| *kind != actual) { + return unsupported(format!("{} attribute `{name}` (or its type)", n.op)); + } + } + for (i, name) in n.inputs.iter().enumerate() { + if let Some(w) = g.weights.get(name) + && matches!(w.data, Data::I64(_)) + && !(n.op == "Reshape" && i == 1) + { + return unsupported(format!("{} needs float32 input {i}", n.op)); + } + } + match n.op.as_str() { + "Conv" => { + let w = weight(n, g, 1)?; + floats(n, g, 1)?; + if w.dims.len() != 4 || w.dims.contains(&0) { + return shape_err("Conv needs nonempty 2-D weights"); + } + if n.ints("kernel_shape").is_some_and(|v| { + v.iter().map(|&d| usize::try_from(d).ok()).collect::>>().as_deref() != Some(w.dims.get(2..).unwrap_or(&[])) + }) { + return shape_err("Conv kernel_shape differs from its weights"); + } + if n.ints("dilations").is_some_and(|d| d != [1, 1]) || n.string("auto_pad").is_some_and(|s| s != "NOTSET") { + return unsupported("Conv dilation or automatic padding"); + } + let (sh, sw) = pair(n, "strides", 1)?; + let (pt, pl, pb, pr) = pads(n)?; + if [sh, sw, pt, pl, pb, pr].iter().any(|&v| v > 65536) { + return shape_err("Conv stride or padding exceeds 65536"); + } + if n.int("group").unwrap_or(1) < 1 { + return shape_err("Conv group must be positive"); + } + if n.inputs.get(2).is_some_and(|s| !s.is_empty()) { + floats(n, g, 2)?; + } + } + "BatchNormalization" => { + if n.int("training_mode").unwrap_or(0) != 0 { + return unsupported("training BatchNormalization"); + } + if float(n, "epsilon", 1e-5) <= 0.0 { + return shape_err("BatchNormalization epsilon must be positive"); + } + for i in 1..5 { + floats(n, g, i)?; + } + } + "PRelu" => { + floats(n, g, 1)?; + } + "Gemm" => { + floats(n, g, 1)?; + if n.inputs.get(2).is_some_and(|s| !s.is_empty()) { + floats(n, g, 2)?; + } + } + "Dropout" => { + if !(0.0..1.0).contains(&float(n, "ratio", 0.5)) { + return shape_err("Dropout ratio is invalid"); + } + } + "MaxPool" => { + pair(n, "kernel_shape", 1)?; + pair(n, "strides", 1)?; + if n.ints("kernel_shape").is_none() + || n.int("ceil_mode").unwrap_or(0) != 0 + || n.int("storage_order").unwrap_or(0) != 0 + || n.ints("dilations").is_some_and(|v| v != [1, 1]) + || pads(n)? != (0, 0, 0, 0) + || n.string("auto_pad").is_some_and(|s| s != "NOTSET") + { + return unsupported("this MaxPool padding, dilation or ceil mode"); + } + } + "Resize" => { + floats(n, g, 2)?; + if n.string("mode").unwrap_or("nearest") != "nearest" || n.inputs.get(3).is_some_and(|s| !s.is_empty()) { + return unsupported("Resize requires nearest constant scales"); + } + if n.int("exclude_outside").unwrap_or(0) != 0 || float(n, "extrapolation_value", 0.0) != 0.0 { + return unsupported("Resize excludes or extrapolates outside values"); + } + let coord = n.string("coordinate_transformation_mode").unwrap_or("half_pixel"); + let nearest = n.string("nearest_mode").unwrap_or("round_prefer_floor"); + if !matches!((coord, nearest), ("asymmetric", "floor") | ("half_pixel", "round_prefer_floor")) { + return unsupported("this Resize coordinate/nearest mode"); + } + if n.inputs.get(1).is_some_and(|s| !s.is_empty()) && !floats(n, g, 1)?.is_empty() { + return unsupported("Resize with an ROI"); + } + } + "Reshape" => { + if n.int("allowzero").unwrap_or(0) != 0 { + return unsupported("Reshape allowzero"); + } + if !matches!(weight(n, g, 1)?.data, Data::I64(_)) { + return unsupported("Reshape needs a constant int64 shape"); + } + } + _ => {} + } + Ok(()) +} +pub(super) fn broadcast(a: &[usize], b: &[usize]) -> Result> { + elements(a)?; + elements(b)?; + let rank = a.len().max(b.len()); + let mut out = vec![1; rank]; + for (i, o) in out.iter_mut().enumerate() { + let aa = i.checked_sub(rank - a.len()).and_then(|j| a.get(j)).copied().unwrap_or(1); + let bb = i.checked_sub(rank - b.len()).and_then(|j| b.get(j)).copied().unwrap_or(1); + if aa != bb && aa != 1 && bb != 1 { + return shape_err("incompatible broadcast shapes"); + } + *o = aa.max(bb); + } + elements(&out)?; + Ok(out) +} +pub(super) fn axis(n: &Node, rank: usize) -> Result { + let raw = n.int("axis").unwrap_or(1); + let raw = if raw < 0 { raw.checked_add(rank as i64).ok_or_else(|| NetError::Shape("axis overflow".into()))? } else { raw }; + usize::try_from(raw).ok().filter(|&a| a <= rank).ok_or_else(|| NetError::Shape("Flatten axis is outside its rank".into())) +} +fn image(s: &[usize]) -> Result<(usize, usize, usize)> { + match s { + [1, c, h, w] => Ok((*c, *h, *w)), + _ => unsupported("a spatial operator needs [1,C,H,W]"), + } +} +pub(super) fn infer(g: &Graph, input: &[usize]) -> Result>> { + elements(input)?; + let mut shapes: HashMap<&str, Vec> = g.weights.iter().map(|(n, w)| (n.as_str(), w.dims.clone())).collect(); + shapes.insert(g.inputs.first().ok_or_else(|| NetError::Shape("no input".into()))?.0.as_str(), input.to_vec()); + let mut live = elements(input)?; + // Scratch grows independently and remains allocated across all subsequent layers. + let mut scratch_lower = 0; + let mut scratch_product = 0; + let mut last = HashMap::new(); + for (i, n) in g.nodes.iter().enumerate() { + for s in &n.inputs { + last.insert(s.as_str(), i); + } + } + for s in &g.outputs { + last.insert(s.as_str(), usize::MAX); + } + for (i, n) in g.nodes.iter().enumerate() { + let a = n.inputs.first().and_then(|s| shapes.get(s.as_str())).ok_or_else(|| NetError::Shape("missing input shape".into()))?; + elements(a)?; + let out = match n.op.as_str() { + "Conv" => { + let (cin, h, w) = image(a)?; + let dims = &weight(n, g, 1)?.dims; + let [co, cg, kh, kw] = dims.as_slice() else { + return shape_err("Conv dimensions"); + }; + let groups = usize::try_from(n.int("group").unwrap_or(1)).map_err(|_| NetError::Shape("Conv group".into()))?; + if groups.checked_mul(*cg) != Some(cin) || !co.is_multiple_of(groups) { + return shape_err("Conv channels do not match weights/groups"); + } + if n.inputs.get(2).is_some_and(|s| !s.is_empty()) && weight(n, g, 2)?.dims != [*co] { + return shape_err("Conv bias is not a channel vector"); + } + let (sh, sw) = pair(n, "strides", 1)?; + let (pt, pl, pb, pr) = pads(n)?; + let padded = |size: usize, p: usize, q: usize, k: usize, s: usize| { + size.checked_add(p) + .and_then(|v| v.checked_add(q)) + .and_then(|v| v.checked_sub(k)) + .and_then(|v| (v / s).checked_add(1)) + .ok_or_else(|| NetError::Shape("Conv padding overflow or oversized kernel".into())) + }; + vec![1, *co, padded(h, pt, pb, *kh, sh)?, padded(w, pl, pr, *kw, sw)?] + } + "BatchNormalization" => { + let channels = *a.get(1).ok_or_else(|| NetError::Shape("BatchNormalization rank must be at least two".into()))?; + for j in 1..5 { + if weight(n, g, j)?.dims != [channels] { + return shape_err("BatchNormalization parameters must be channel vectors"); + } + } + if floats(n, g, 4)?.iter().any(|v| *v < 0.0) { + return shape_err("BatchNormalization has negative variance"); + } + let eps = float(n, "epsilon", 1e-5); + let (s, b, m, v) = (floats(n, g, 1)?, floats(n, g, 2)?, floats(n, g, 3)?, floats(n, g, 4)?); + for (((&s, &b), &m), &v) in s.iter().zip(b).zip(m).zip(v) { + let variance = v + eps; + let coefficient = s / variance.sqrt(); + let offset = b - m * coefficient; + if !variance.is_finite() || !coefficient.is_finite() || !offset.is_finite() { + return shape_err("BatchNormalization coefficients overflow"); + } + } + + a.clone() + } + "PRelu" => { + if broadcast(a, &weight(n, g, 1)?.dims)? != *a { + return shape_err("PRelu slope would enlarge its input"); + } + a.clone() + } + "Add" | "Sub" | "Mul" => { + let b = n.inputs.get(1).and_then(|s| shapes.get(s.as_str())).ok_or_else(|| NetError::Shape("missing second input".into()))?; + broadcast(a, b)? + } + "Gemm" => { + let b = &weight(n, g, 1)?.dims; + let [ar, ac] = a.as_slice() else { + return shape_err("Gemm needs matrix A"); + }; + let [br, bc] = b.as_slice() else { + return shape_err("Gemm needs matrix B"); + }; + let (m, k) = if n.int("transA").unwrap_or(0) != 0 { (*ac, *ar) } else { (*ar, *ac) }; + let (bk, cols) = if n.int("transB").unwrap_or(0) != 0 { (*bc, *br) } else { (*br, *bc) }; + if k != bk { + return shape_err("Gemm inner dimensions differ"); + } + let s = vec![m, cols]; + if n.inputs.get(2).is_some_and(|s| !s.is_empty()) && broadcast(&s, &weight(n, g, 2)?.dims)? != s { + return shape_err("Gemm bias would enlarge the output"); + } + s + } + "Flatten" => { + let at = axis(n, a.len())?; + vec![ + elements(a.get(..at).ok_or_else(|| NetError::Shape("bad axis".into()))?)?, + elements(a.get(at..).ok_or_else(|| NetError::Shape("bad axis".into()))?)?, + ] + } + "GlobalAveragePool" => { + let (c, _, _) = image(a)?; + vec![1, c, 1, 1] + } + "MaxPool" => { + let (c, h, w) = image(a)?; + let (kh, kw) = pair(n, "kernel_shape", 1)?; + let (sh, sw) = pair(n, "strides", 1)?; + if h < kh || w < kw { + return shape_err("MaxPool kernel exceeds input"); + } + vec![1, c, (h - kh) / sh + 1, (w - kw) / sw + 1] + } + "Resize" => { + let (c, h, w) = image(a)?; + let s = floats(n, g, 2)?; + if s.len() != 4 + || s.first() != Some(&1.0) + || s.get(1) != Some(&1.0) + || s.iter().any(|&v| !(1.0..=16.0).contains(&v) || v.fract() != 0.0) + { + return unsupported("Resize supports whole spatial scales from 1 to 16"); + } + vec![ + 1, + c, + h.checked_mul(s.get(2).copied().unwrap_or(0.0) as usize).ok_or_else(|| NetError::Shape("Resize overflow".into()))?, + w.checked_mul(s.get(3).copied().unwrap_or(0.0) as usize).ok_or_else(|| NetError::Shape("Resize overflow".into()))?, + ] + } + "Transpose" => { + let perm: Vec = n.ints("perm").map_or_else(|| (0..a.len() as i64).rev().collect(), Vec::from); + let mut seen = std::collections::HashSet::new(); + let mut s = Vec::new(); + if perm.len() != a.len() { + return shape_err("Transpose rank mismatch"); + } + for d in perm { + let d = usize::try_from(d).map_err(|_| NetError::Shape("negative permutation".into()))?; + if !seen.insert(d) { + return shape_err("repeated permutation axis"); + } + s.push(*a.get(d).ok_or_else(|| NetError::Shape("bad permutation axis".into()))?); + } + s + } + "Reshape" => { + let Data::I64(dims) = &weight(n, g, 1)?.data else { + return shape_err("Reshape datatype"); + }; + if dims.len() > 8 { + return shape_err("Reshape rank exceeds eight"); + } + let total = elements(a)?; + let mut s = Vec::new(); + let mut infer = None; + for (j, &v) in dims.iter().enumerate() { + s.push(match v { + -1 if infer.is_none() => { + infer = Some(j); + 1 + } + 0 => *a.get(j).ok_or_else(|| NetError::Shape("Reshape copies missing dimension".into()))?, + v if v > 0 => usize::try_from(v).map_err(|_| NetError::Shape("Reshape dimension overflow".into()))?, + _ => return shape_err("invalid Reshape dimension"), + }); + } + let known = elements(&s)?; + if let Some(j) = infer { + if !total.is_multiple_of(known) { + return shape_err("Reshape does not divide evenly"); + } + *s.get_mut(j).ok_or_else(|| NetError::Shape("missing inferred dimension".into()))? = total / known; + } + if elements(&s)? != total { + return shape_err("Reshape changes element count"); + } + s + } + _ => a.clone(), + }; + let out_count = elements(&out)?; + let mut copied_constants = 0usize; + for (j, name) in n.inputs.iter().enumerate() { + // Execution borrows operator weights, but get(0) and binary get(1) materialize constants. + if (j == 0 || (j == 1 && matches!(n.op.as_str(), "Add" | "Sub" | "Mul"))) + && let Some(w) = g.weights.get(name) + { + copied_constants = + copied_constants.checked_add(elements(&w.dims)?).ok_or_else(|| NetError::Shape("constant copies overflow".into()))?; + } + } + if n.op == "Conv" { + let wd = &weight(n, g, 1)?.dims; + let [co, cg, kh, kw] = wd.as_slice() else { + return shape_err("Conv rank"); + }; + let block = 256.min(out.get(2).copied().unwrap_or(1) * out.get(3).copied().unwrap_or(1)); + let lower = cg + .checked_mul(*kh) + .and_then(|v| v.checked_mul(*kw)) + .and_then(|v| v.checked_mul(block)) + .filter(|&v| v <= MAX_TENSOR) + .ok_or_else(|| NetError::Shape("convolution scratch exceeds the limit".into()))?; + scratch_lower = scratch_lower.max(lower); + scratch_product = scratch_product.max(co * block); + } + let scratch = scratch_lower + scratch_product; + if live + .checked_add(out_count) + .and_then(|v| v.checked_add(scratch)) + .and_then(|v| v.checked_add(copied_constants)) + .is_none_or(|v| v > 128 * 1024 * 1024) + { + return shape_err("tensors and scratch exceed the 512 MiB working limit"); + } + live = live + .checked_add(out_count) + .filter(|&n| n <= 128 * 1024 * 1024) + .ok_or_else(|| NetError::Shape("live tensors exceed the 512 MiB working limit".into()))?; + shapes.insert(n.outputs.first().ok_or_else(|| NetError::Shape("no output".into()))?, out); + let mut removed = std::collections::HashSet::new(); + for name in &n.inputs { + if last.get(name.as_str()) == Some(&i) && !g.weights.contains_key(name) && removed.insert(name) { + live = live.saturating_sub(shapes.get(name.as_str()).map_or(0, |s| elements(s).unwrap_or(0))); + } + } + } + g.outputs.iter().map(|s| shapes.get(s.as_str()).cloned().ok_or_else(|| NetError::Shape("missing output".into()))).collect() +} diff --git a/crates/faces/src/onnx.rs b/crates/faces/src/onnx.rs new file mode 100644 index 00000000..453be241 --- /dev/null +++ b/crates/faces/src/onnx.rs @@ -0,0 +1,514 @@ +//! A streaming probe of an ONNX file: the graph's inputs and outputs (name, element type, shape), without +//! running or even fully loading the model. +//! +//! ONNX is protobuf. We read only the few fields we need (the producer, the opset, and the graph's input, +//! output and initializer names) and seek past everything else, so a 260 MB model costs a few kilobytes and +//! a few milliseconds. The file is hostile input: every length is checked against its parent message and +//! the file size, the number of fields read is capped, strings are bounded, and nothing indexes or unwraps. + +use std::collections::HashSet; +use std::io::{self, Read, Seek, SeekFrom}; +use std::path::Path; + +const MAX_FIELDS: u32 = 5_000_000; +const MAX_STRING: u64 = 4096; +const MAX_DIMS: usize = 16; +const MAX_TENSORS: usize = 64; +/// Older exporters list every weight as a graph input too; that many are read before the weights are filtered out. +const MAX_LISTED_INPUTS: usize = 200_000; +const MAX_INITIALIZERS: usize = 1_000_000; + +#[derive(Clone, Debug, PartialEq)] +pub enum Dim { + Fixed(u64), + /// A named or unknown size (typically the batch). + Dynamic(String), +} + +#[derive(Clone, Debug, PartialEq)] +pub struct TensorInfo { + pub name: String, + /// ONNX element type (1 = float32). + pub elem_type: u32, + pub shape: Vec, +} + +#[derive(Clone, Debug, Default, PartialEq)] +pub struct OnnxInfo { + pub producer: String, + /// The default-domain opset version. + pub opset: Option, + /// Real inputs only (weights that older exporters also list as inputs are left out). + pub inputs: Vec, + pub outputs: Vec, +} + +#[derive(Debug, thiserror::Error)] +pub enum ProbeError { + #[error("not an ONNX model: {0}")] + Malformed(&'static str), + #[error("could not read the model: {0}")] + Io(#[from] io::Error), +} + +pub(crate) type Result = std::result::Result; + +pub(crate) enum Wire { + Varint(u64), + /// A length-delimited field: its bytes are `start..end` and the reader sits at `start`. + Len { + start: u64, + end: u64, + }, + Fixed32(u32), + #[allow(dead_code)] + Fixed64(u64), +} + +pub(crate) struct Walker<'a, R> { + pub(crate) r: &'a mut R, + pub(crate) total: u64, + pub(crate) fields: u32, +} + +impl Walker<'_, R> { + pub(crate) fn pos(&mut self) -> Result { + Ok(self.r.stream_position()?) + } + + pub(crate) fn seek_to(&mut self, at: u64) -> Result<()> { + self.r.seek(SeekFrom::Start(at))?; + Ok(()) + } + + pub(crate) fn byte(&mut self) -> Result { + let mut b = [0u8; 1]; + self.r + .read_exact(&mut b) + .map_err(|e| if e.kind() == io::ErrorKind::UnexpectedEof { ProbeError::Malformed("file ends inside a field") } else { e.into() })?; + let [x] = b; + Ok(x) + } + + pub(crate) fn varint(&mut self) -> Result { + let mut v = 0u64; + for shift in (0..70u32).step_by(7) { + let b = self.byte()?; + if shift == 63 && b > 1 { + return Err(ProbeError::Malformed("number overflows 64 bits")); + } + v |= u64::from(b & 0x7f) << shift; + if b & 0x80 == 0 { + return Ok(v); + } + } + Err(ProbeError::Malformed("number too long")) + } + + /// The next field of the message that ends at `end`, or `None` at its end. + pub(crate) fn next(&mut self, end: u64) -> Result> { + let here = self.pos()?; + if here >= end { + return Ok(None); + } + self.fields += 1; + if self.fields > MAX_FIELDS { + return Err(ProbeError::Malformed("too many fields")); + } + let tag = self.varint()?; + let number = u32::try_from(tag >> 3).map_err(|_| ProbeError::Malformed("field number too large"))?; + if number == 0 || number > 0x1fff_ffff { + return Err(ProbeError::Malformed("invalid field number")); + } + let wire = match tag & 7 { + 0 => Wire::Varint(self.varint()?), + 1 => { + let at = self.pos()?.checked_add(8).ok_or(ProbeError::Malformed("offset overflow"))?; + if at > end || at > self.total { + return Err(ProbeError::Malformed("field runs past its message")); + } + let mut b = [0u8; 8]; + self.r.read_exact(&mut b).map_err(|_| ProbeError::Malformed("file ends inside a field"))?; + Wire::Fixed64(u64::from_le_bytes(b)) + } + 5 => { + let at = self.pos()?.checked_add(4).ok_or(ProbeError::Malformed("offset overflow"))?; + if at > end || at > self.total { + return Err(ProbeError::Malformed("field runs past its message")); + } + let mut b = [0u8; 4]; + self.r.read_exact(&mut b).map_err(|_| ProbeError::Malformed("file ends inside a field"))?; + Wire::Fixed32(u32::from_le_bytes(b)) + } + 2 => { + let len = self.varint()?; + let start = self.pos()?; + let stop = start.checked_add(len).ok_or(ProbeError::Malformed("offset overflow"))?; + if stop > end || stop > self.total { + return Err(ProbeError::Malformed("field runs past its message")); + } + Wire::Len { start, end: stop } + } + _ => return Err(ProbeError::Malformed("unsupported field type")), + }; + if self.pos()? > end { + return Err(ProbeError::Malformed("field runs past its message")); + } + Ok(Some((number, wire))) + } + + /// A bounded UTF-8 string occupying `start..end`; never silently truncate graph names. + pub(crate) fn string(&mut self, start: u64, end: u64) -> Result { + let len = end.saturating_sub(start); + if len > MAX_STRING { + return Err(ProbeError::Malformed("a string exceeds 4096 bytes")); + } + let mut buf = Vec::new(); + self.r.by_ref().take(len).read_to_end(&mut buf)?; + String::from_utf8(buf).map_err(|_| ProbeError::Malformed("a string is not UTF-8")) + } + + /// The bytes of `start..end`, which must be at most `max` long. + pub(crate) fn bytes(&mut self, start: u64, end: u64, max: u64) -> Result> { + let len = end.saturating_sub(start); + if len > max { + return Err(ProbeError::Malformed("a field is larger than allowed")); + } + let mut buf = Vec::new(); + self.r.by_ref().take(len).read_to_end(&mut buf)?; + if buf.len() as u64 != len { + return Err(ProbeError::Malformed("file ends inside a field")); + } + Ok(buf) + } +} + +fn dim(w: &mut Walker, end: u64) -> Result> { + let mut out = None; + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Varint(v)) => out = Some(Dim::Fixed(v)), + (2, Wire::Len { start, end: e }) => { + out = Some(Dim::Dynamic(w.string(start, e)?)); + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(out) +} + +fn shape(w: &mut Walker, end: u64) -> Result> { + let mut dims = Vec::new(); + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Len { end: e, .. }) => { + if dims.len() >= MAX_DIMS { + return Err(ProbeError::Malformed("too many dimensions")); + } + dims.push(dim(w, e)?.unwrap_or(Dim::Dynamic(String::new()))); + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(dims) +} + +fn tensor_type(w: &mut Walker, end: u64, info: &mut TensorInfo) -> Result<()> { + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Varint(v)) => info.elem_type = u32::try_from(v).unwrap_or(0), + (2, Wire::Len { end: e, .. }) => { + info.shape = shape(w, e)?; + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(()) +} + +pub(crate) fn type_proto(w: &mut Walker, end: u64, info: &mut TensorInfo) -> Result<()> { + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Len { end: e, .. }) => { + tensor_type(w, e, info)?; + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(()) +} + +fn value_info(w: &mut Walker, end: u64) -> Result { + let mut info = TensorInfo { name: String::new(), elem_type: 0, shape: Vec::new() }; + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Len { start, end: e }) => { + info.name = w.string(start, e)?; + w.seek_to(e)?; + } + (2, Wire::Len { end: e, .. }) => { + type_proto(w, e, &mut info)?; + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(info) +} + +/// The `name` (field 8) of a weight tensor; its bulk data is skipped, never read. +fn tensor_name(w: &mut Walker, end: u64) -> Result> { + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (8, Wire::Len { start, end: e }) => return Ok(Some(w.string(start, e)?)), + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(None) +} + +fn graph(w: &mut Walker, end: u64, info: &mut OnnxInfo, weights: &mut HashSet) -> Result<()> { + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (5, Wire::Len { end: e, .. }) => { + if weights.len() < MAX_INITIALIZERS + && let Some(name) = tensor_name(w, e)? + { + weights.insert(name); + } + w.seek_to(e)?; + } + (11, Wire::Len { end: e, .. }) => { + if info.inputs.len() >= MAX_LISTED_INPUTS { + return Err(ProbeError::Malformed("too many inputs")); + } + info.inputs.push(value_info(w, e)?); + w.seek_to(e)?; + } + (12, Wire::Len { end: e, .. }) => { + if info.outputs.len() >= MAX_TENSORS { + return Err(ProbeError::Malformed("too many outputs")); + } + info.outputs.push(value_info(w, e)?); + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(()) +} + +fn opset(w: &mut Walker, end: u64) -> Result> { + let (mut domain, mut version) = (String::new(), None); + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Len { start, end: e }) => { + domain = w.string(start, e)?; + w.seek_to(e)?; + } + (2, Wire::Varint(v)) => version = Some(v), + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(if domain.is_empty() || domain == "ai.onnx" { version } else { None }) +} + +/// Read an ONNX model's graph inputs and outputs from any seekable source. +pub fn probe_reader(r: &mut R) -> Result { + let total = r.seek(SeekFrom::End(0))?; + if total > crate::manifest::MAX_MODEL_BYTES { + return Err(ProbeError::Malformed("the model exceeds 512 MiB")); + } + r.seek(SeekFrom::Start(0))?; + if total == 0 { + return Err(ProbeError::Malformed("the file is empty")); + } + let mut w = Walker { r, total, fields: 0 }; + let (mut info, mut weights, mut saw_graph) = (OnnxInfo::default(), HashSet::new(), false); + while let Some((n, f)) = w.next(total)? { + match (n, f) { + (2, Wire::Len { start, end }) => { + info.producer = w.string(start, end)?; + w.seek_to(end)?; + } + (7, Wire::Len { end, .. }) => { + saw_graph = true; + graph(&mut w, end, &mut info, &mut weights)?; + w.seek_to(end)?; + } + (8, Wire::Len { end, .. }) => { + if let Some(v) = opset(&mut w, end)? { + info.opset = Some(v); + } + w.seek_to(end)?; + } + (_, Wire::Len { end, .. }) => w.seek_to(end)?, + _ => {} + } + } + if !saw_graph { + return Err(ProbeError::Malformed("no graph")); + } + info.inputs.retain(|t| !weights.contains(&t.name)); + if info.inputs.len() > MAX_TENSORS { + return Err(ProbeError::Malformed("too many inputs")); + } + if info.inputs.is_empty() || info.outputs.is_empty() { + return Err(ProbeError::Malformed("the graph has no inputs or no outputs")); + } + Ok(info) +} + +/// Read an ONNX file's graph inputs and outputs. +pub fn probe_path(path: &Path) -> Result { + probe_reader(&mut io::BufReader::new(std::fs::File::open(path)?)) +} + +#[cfg(test)] +pub(crate) mod tests { + use super::*; + use std::io::Cursor; + + pub(crate) use crate::synthetic::embedder_model as model_bytes; + use crate::synthetic::{len_field, value_info_bytes, varint}; + + #[test] + fn reads_inputs_outputs_and_skips_weights() { + let info = probe_reader(&mut Cursor::new(model_bytes(512))).unwrap(); + assert_eq!(info.producer, "pytorch"); + assert_eq!(info.opset, Some(17)); + assert_eq!(info.inputs.len(), 1, "the weight listed as an input is left out"); + assert_eq!(info.inputs[0].name, "data"); + assert_eq!(info.inputs[0].elem_type, 1); + assert_eq!(info.inputs[0].shape, vec![Dim::Dynamic("N".into()), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]); + assert_eq!(info.outputs[0].shape, vec![Dim::Dynamic("N".into()), Dim::Fixed(512)]); + } + + #[test] + fn old_style_files_listing_hundreds_of_weights_as_inputs_still_probe() { + let mut graph = Vec::new(); + len_field(11, &value_info_bytes("data", 1, &[Err("N"), Ok(3), Ok(112), Ok(112)]), &mut graph); + for i in 0..500 { + let name = format!("weight{i}"); + len_field(11, &value_info_bytes(&name, 1, &[Ok(4)]), &mut graph); + let mut tensor = Vec::new(); + len_field(8, name.as_bytes(), &mut tensor); + len_field(5, &tensor, &mut graph); + } + len_field(12, &value_info_bytes("emb", 1, &[Err("N"), Ok(128)]), &mut graph); + let mut m = Vec::new(); + len_field(7, &graph, &mut m); + let info = probe_reader(&mut Cursor::new(m)).unwrap(); + assert_eq!(info.inputs.len(), 1); + assert_eq!(info.inputs[0].name, "data"); + } + + #[test] + fn many_real_inputs_are_refused() { + let mut graph = Vec::new(); + for i in 0..70 { + len_field(11, &value_info_bytes(&format!("in{i}"), 1, &[Ok(4)]), &mut graph); + } + len_field(12, &value_info_bytes("emb", 1, &[Ok(4)]), &mut graph); + let mut m = Vec::new(); + len_field(7, &graph, &mut m); + assert!(probe_reader(&mut Cursor::new(m)).is_err()); + } + + #[test] + fn truncated_at_every_length_never_panics() { + let bytes = model_bytes(128); + for n in 0..bytes.len() { + let _ = probe_reader(&mut Cursor::new(bytes[..n].to_vec())); + } + assert!(probe_reader(&mut Cursor::new(bytes)).is_ok()); + } + + #[test] + fn hostile_files_are_errors() { + let mut lying = Vec::new(); + varint(7 << 3 | 2, &mut lying); + varint(1 << 40, &mut lying); // a graph 1 TB long in a 10-byte file + let mut overflow = Vec::new(); + varint(7 << 3 | 2, &mut overflow); + varint(u64::MAX, &mut overflow); + let mut long_number = vec![0x08]; + long_number.extend_from_slice(&[0xff; 11]); + for (what, bytes) in [ + ("empty", Vec::new()), + ("zeros", vec![0u8; 64]), + ("text", b"this is not an onnx model at all".to_vec()), + ("lying length", lying), + ("overflowing length", overflow), + ("endless number", long_number), + ("no graph", { + let mut m = Vec::new(); + len_field(2, b"x", &mut m); + m + }), + ("graph without io", { + let mut m = Vec::new(); + len_field(7, &[], &mut m); + m + }), + ] { + assert!(probe_reader(&mut Cursor::new(bytes)).is_err(), "{what}"); + } + } + + #[test] + fn too_many_dimensions_or_tensors_are_errors() { + let dims: Vec> = (0..40).map(Ok).collect(); + let mut graph = Vec::new(); + len_field(11, &value_info_bytes("data", 1, &dims), &mut graph); + len_field(12, &value_info_bytes("out", 1, &[Ok(1)]), &mut graph); + let mut m = Vec::new(); + len_field(7, &graph, &mut m); + assert!(probe_reader(&mut Cursor::new(m)).is_err()); + } + + proptest::proptest! { + #[test] + fn arbitrary_bytes_never_panic(bytes in proptest::collection::vec(proptest::prelude::any::(), 0..400)) { + let _ = probe_reader(&mut Cursor::new(bytes)); + } + + #[test] + fn a_valid_model_with_flipped_bytes_never_panics(flips in proptest::collection::vec((0usize..400, proptest::prelude::any::()), 1..6)) { + let mut bytes = model_bytes(512); + for (at, v) in flips { + if let Some(b) = bytes.get_mut(at) { + *b = v; + } + } + let _ = probe_reader(&mut Cursor::new(bytes)); + } + } + #[test] + fn oversized_models_are_rejected_without_reading_or_allocating() { + struct Oversized; + impl std::io::Read for Oversized { + fn read(&mut self, _: &mut [u8]) -> std::io::Result { + panic!("an oversized model must not be read") + } + } + impl std::io::Seek for Oversized { + fn seek(&mut self, _: std::io::SeekFrom) -> std::io::Result { + Ok(crate::manifest::MAX_MODEL_BYTES + 1) + } + } + assert!(probe_reader(&mut Oversized).is_err()); + } +} diff --git a/crates/faces/src/runtime.rs b/crates/faces/src/runtime.rs new file mode 100644 index 00000000..bbb1457e --- /dev/null +++ b/crates/faces/src/runtime.rs @@ -0,0 +1,352 @@ +//! Installed recognition models on LightCraft's own checked, single-thread float32 CPU engine. +//! The detector and recognisers share graph validation and convolution kernels. + +use std::panic::{AssertUnwindSafe, catch_unwind}; +use std::path::Path; +use std::sync::Arc; +use web_time::Instant; + +use crate::{ + graph, + net::{Net, Tensor}, +}; +use serde::Serialize; + +use crate::manifest::{Colour, ModelManifest, OutputSpec}; + +#[derive(Debug, thiserror::Error)] +pub enum RuntimeError { + #[error("the model could not be loaded: {0}")] + Load(String), + #[error("the model failed while running: {0}")] + Run(String), + #[error("{0}")] + Input(String), +} + +/// What `self_test` found. +#[derive(Clone, Debug, PartialEq, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct SelfTest { + pub ok: bool, + /// Milliseconds to load and optimise the model. + pub load_ms: f64, + /// Median milliseconds to embed one face. + pub embed_ms: f64, + pub dimension: usize, + /// Each check: what was checked and whether it held. + pub checks: Vec<(String, bool)>, +} + +type Plan = Net; + +/// A loaded recognition model. +#[derive(Clone)] +pub struct Embedder { + plan: Arc, + manifest: ModelManifest, + width: usize, + height: usize, + dim: usize, + load_ms: f64, +} + +fn guarded(what: &str, f: impl FnOnce() -> Result) -> Result { + match catch_unwind(AssertUnwindSafe(f)) { + Ok(r) => r, + Err(_) => Err(RuntimeError::Run(format!("the runtime gave up on this model while {what}"))), + } +} + +impl Embedder { + /// Load the `.onnx` at `path` as described by `manifest` (an embedder). + pub fn load(path: &Path, manifest: &ModelManifest) -> Result { + let OutputSpec::Embedding { dim } = manifest.output else { + return Err(RuntimeError::Load("this is not a recognition model".into())); + }; + let (w, h) = (manifest.input.width as usize, manifest.input.height as usize); + let started = Instant::now(); + crate::manifest::validate(manifest).map_err(|e| RuntimeError::Load(e.to_string()))?; + let mut file = std::fs::File::open(path).map_err(|e| RuntimeError::Load(e.to_string()))?; + let length = file.metadata().map_err(|e| RuntimeError::Load(e.to_string()))?.len(); + if length == 0 || length > 512 * 1024 * 1024 { + return Err(RuntimeError::Load("the model file is empty or exceeds 512 MiB".into())); + } + let mut bytes = Vec::new(); + bytes + .try_reserve_exact(usize::try_from(length).map_err(|_| RuntimeError::Load("model size overflow".into()))?) + .map_err(|_| RuntimeError::Load("not enough memory for the model".into()))?; + use std::io::Read; + (&mut file).take(512 * 1024 * 1024 + 1).read_to_end(&mut bytes).map_err(|e| RuntimeError::Load(e.to_string()))?; + let graph = graph::load(&bytes).map_err(|e| RuntimeError::Load(e.to_string()))?; + let input = graph.inputs.first().ok_or_else(|| RuntimeError::Load("no input".into()))?; + let want = [1, 3, h, w]; + if input.1.len() != 4 || input.1.iter().zip(want).any(|(&d, w)| d != 0 && d != w) { + return Err(RuntimeError::Load("the graph input differs from the manifest's [1,3,H,W]".into())); + } + let plan = Net::new(graph).map_err(|e| RuntimeError::Load(e.to_string()))?; + let shapes = plan.validate_input(&want).map_err(|e| RuntimeError::Load(e.to_string()))?; + if shapes.len() != 1 || shapes.first().is_none_or(|s| s.first() != Some(&1) || s.iter().product::() != dim as usize) { + return Err(RuntimeError::Load("the graph output differs from the manifest's embedding dimension".into())); + } + Ok(Embedder { + plan: Arc::new(plan), + manifest: manifest.clone(), + width: w, + height: h, + dim: dim as usize, + load_ms: started.elapsed().as_secs_f64() * 1000.0, + }) + } + + /// The size of the aligned face picture the model wants (width, height). + pub fn input_size(&self) -> (usize, usize) { + (self.width, self.height) + } + + pub fn manifest(&self) -> &ModelManifest { + &self.manifest + } + + /// The face as a unit-length vector. `rgb` is an aligned face at exactly [`Self::input_size`], 8-bit RGB. + pub fn embed(&self, rgb: &[u8]) -> Result, RuntimeError> { + if rgb.len() != self.width * self.height * 3 { + return Err(RuntimeError::Input("the face picture is not the size the model wants".into())); + } + let (mean, std, bgr) = (self.manifest.input.mean, self.manifest.input.std, self.manifest.input.colour == Colour::Bgr); + let plane = self.width * self.height; + let mut data = vec![0.0f32; 3 * plane]; + for (i, px) in rgb.as_chunks::<3>().0.iter().enumerate() { + for c in 0..3 { + let src = if bgr { 2 - c } else { c }; + if let (Some(v), Some(m), Some(s), Some(o)) = (px.get(src), mean.get(c), std.get(c), data.get_mut(c * plane + i)) { + *o = (f32::from(*v) - m) / s; + } + } + } + let out = guarded("running it", || { + let input = Tensor::new(vec![1, 3, self.height, self.width], data).map_err(|e| RuntimeError::Input(e.to_string()))?; + let mut result = self.plan.run(input).map_err(|e| RuntimeError::Run(e.to_string()))?; + if result.len() != 1 { + return Err(RuntimeError::Run("the model must give one embedding".into())); + } + Ok(result.pop().ok_or_else(|| RuntimeError::Run("no output".into()))?.1.data) + })?; + if out.len() != self.dim { + return Err(RuntimeError::Run(format!("the model gave {} numbers per face, not the {} its description says", out.len(), self.dim))); + } + if out.iter().any(|v| !v.is_finite()) { + return Err(RuntimeError::Run("the model gave numbers that are not finite".into())); + } + let norm = out.iter().map(|&v| f64::from(v).powi(2)).sum::().sqrt(); + if !norm.is_finite() || norm < 1e-9 { + return Err(RuntimeError::Run("the model gave an all-zero face".into())); + } + Ok(out.into_iter().map(|v| (f64::from(v) / norm) as f32).collect()) + } + + /// Check the model does something sensible before it is ever used on a photo: the right number of finite + /// numbers, the same picture twice gives the same face, and two different pictures give different faces. + pub fn self_test(&self) -> SelfTest { + let mut checks: Vec<(String, bool)> = Vec::new(); + let n = self.width * self.height * 3; + // two deterministic, different test pictures: noise, and a diagonal gradient with a colour cast + let mut seed = 0x2545_f491u32; + let noise: Vec = (0..n) + .map(|_| { + seed ^= seed << 13; + seed ^= seed >> 17; + seed ^= seed << 5; + (seed >> 8) as u8 + }) + .collect(); + let gradient: Vec = (0..n) + .map(|i| (((i / 3) % self.width + (i / 3) / self.width) * 255 / (self.width + self.height) + (i % 3) * 40).min(255) as u8) + .collect(); + let mut times = Vec::new(); + let mut run = |img: &[u8]| { + let t = Instant::now(); + let r = self.embed(img); + times.push(t.elapsed().as_secs_f64() * 1000.0); + r + }; + let (a, a2, b) = (run(&noise), run(&noise), run(&gradient)); + let dot = |x: &[f32], y: &[f32]| x.iter().zip(y).map(|(p, q)| p * q).sum::(); + match (&a, &a2, &b) { + (Ok(a), Ok(a2), Ok(b)) => { + checks.push((format!("gives {} finite numbers per face", a.len()), true)); + checks.push(("gives the same face for the same picture".into(), dot(a, a2) > 0.9999)); + checks.push(("gives different faces for different pictures".into(), dot(a, b) < 0.9995)); + } + (Err(e), _, _) | (_, Err(e), _) | (_, _, Err(e)) => checks.push((format!("runs: {e}"), false)), + } + times.sort_by(|x, y| x.total_cmp(y)); + SelfTest { + ok: checks.iter().all(|(_, ok)| *ok), + load_ms: self.load_ms, + embed_ms: times.get(times.len() / 2).copied().unwrap_or(0.0), + dimension: self.dim, + checks, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::manifest::{InputSpec, Licence, OutputSpec, Role, Thresholds}; + + fn manifest(dim: u32) -> ModelManifest { + ModelManifest { + id: "tiny".into(), + name: "Tiny".into(), + version: "1".into(), + role: Role::Embedder, + licence: Licence::default(), + source: None, + sha256: None, + size_bytes: None, + provenance: String::new(), + input: InputSpec::default(), + output: OutputSpec::Embedding { dim }, + thresholds: Thresholds::default(), + speed: None, + } + } + + fn temp_model(name: &str, bytes: &[u8]) -> std::path::PathBuf { + let p = std::env::temp_dir().join(format!("lc-faces-runtime-{name}-{}.onnx", std::process::id())); + std::fs::write(&p, bytes).unwrap(); + p + } + + #[test] + fn a_working_model_loads_embeds_and_passes_the_self_test() { + let p = temp_model("ok", &crate::synthetic::tiny_embedder_model(16)); + let e = Embedder::load(&p, &manifest(16)).unwrap(); + assert_eq!(e.input_size(), (112, 112)); + let v = e.embed(&vec![100u8; 112 * 112 * 3]).unwrap(); + assert_eq!(v.len(), 16); + assert!((v.iter().map(|x| x * x).sum::() - 1.0).abs() < 1e-4, "unit length"); + let t = e.self_test(); + assert!(t.ok, "{t:?}"); + assert_eq!(t.dimension, 16); + let _ = std::fs::remove_file(p); + } + + #[test] + fn wrong_sizes_wrong_dimensions_and_broken_files_are_errors() { + let p = temp_model("dim", &crate::synthetic::tiny_embedder_model(16)); + let e = Embedder::load(&p, &manifest(16)).unwrap(); + assert!(e.embed(&[0u8; 10]).is_err(), "wrong picture size"); + // the manifest promises 32 numbers, the model gives 16 + assert!(Embedder::load(&p, &manifest(32)).is_err()); + // not a model, an empty file, a graph with no layers, a detector-shaped manifest + for (name, bytes) in [("junk", b"nonsense".to_vec()), ("empty", Vec::new()), ("nolayers", crate::synthetic::embedder_model(16))] { + let q = temp_model(name, &bytes); + assert!(Embedder::load(&q, &manifest(16)).is_err(), "{name}"); + let _ = std::fs::remove_file(q); + } + let mut detector = manifest(16); + detector.output = OutputSpec::Detector { decoder: "yunet-v2".into() }; + assert!(Embedder::load(&p, &detector).is_err()); + assert!(Embedder::load(&std::env::temp_dir().join("does-not-exist.onnx"), &manifest(16)).is_err()); + let _ = std::fs::remove_file(p); + } + + #[test] + fn truncated_and_corrupted_models_never_panic() { + let good = crate::synthetic::tiny_embedder_model(16); + for n in (0..good.len()).step_by(7) { + let q = temp_model("trunc", &good[..n]); + let _ = Embedder::load(&q, &manifest(16)); + let _ = std::fs::remove_file(q); + } + for at in (0..good.len()).step_by(11) { + let mut bad = good.clone(); + if let Some(b) = bad.get_mut(at) { + *b = b.wrapping_add(0x55); + } + let q = temp_model("flip", &bad); + if let Ok(e) = Embedder::load(&q, &manifest(16)) { + let _ = e.embed(&vec![7u8; 112 * 112 * 3]); + } + let _ = std::fs::remove_file(q); + } + } + + /// Opt-in: `LC_FACE_MODELS= cargo test -p lightcraft-faces -- --ignored --nocapture` + /// self-tests every known recognition model found there. + #[test] + #[ignore = "needs real model files: set LC_FACE_MODELS"] + fn real_models_pass_the_self_test() { + let Some(dir) = std::env::var_os("LC_FACE_MODELS") else { return }; + for entry in std::fs::read_dir(dir).unwrap().flatten() { + let path = entry.path(); + let Ok(sha) = crate::hash::sha256_file(&path) else { continue }; + let Some(m) = crate::known::lookup(&sha).filter(|m| m.role == Role::Embedder) else { continue }; + let e = Embedder::load(&path, &m).unwrap(); + let t = e.self_test(); + println!("{}: ok {} | load {:.0} ms | {:.0} ms per face | {} numbers | {:?}", m.id, t.ok, t.load_ms, t.embed_ms, t.dimension, t.checks); + assert!(t.ok); + } + } + + #[test] + fn bgr_changes_the_input() { + let p = temp_model("bgr", &crate::synthetic::tiny_embedder_model(16)); + let rgb = Embedder::load(&p, &manifest(16)).unwrap(); + let mut m = manifest(16); + m.input.colour = crate::manifest::Colour::Bgr; + let bgr = Embedder::load(&p, &m).unwrap(); + // a picture with different channel levels: swapping the order must change the embedding + let img: Vec = (0..112 * 112).flat_map(|_| [200u8, 100, 20]).collect(); + let (a, b) = (rgb.embed(&img).unwrap(), bgr.embed(&img).unwrap()); + let cos: f32 = a.iter().zip(&b).map(|(x, y)| x * y).sum(); + assert!(cos < 0.9999, "{cos}"); + let _ = std::fs::remove_file(p); + } + #[test] + fn large_finite_embeddings_have_a_finite_unit_norm() { + let p = temp_model("large-norm", &crate::synthetic::norm_overflow_model()); + let e = Embedder::load(&p, &manifest(3)).unwrap(); + let out = e.embed(&vec![255; 112 * 112 * 3]).unwrap(); + assert!(out.iter().all(|v| v.is_finite())); + assert!((out.iter().map(|v| v * v).sum::() - 1.0).abs() < 1e-6); + let _ = std::fs::remove_file(p); + } + #[test] + fn invalid_manifest_numbers_are_rejected_before_preprocessing() { + let p = temp_model("invalid-manifest", &crate::synthetic::tiny_embedder_model(16)); + for case in 0..5 { + let mut m = manifest(16); + match case { + 0 => m.input.std = [0.0; 3], + 1 => m.input.mean = [f32::NAN; 3], + 2 => m.input.width = 0, + 3 => m.input.width = u32::MAX, + _ => m.output = OutputSpec::Embedding { dim: 0 }, + }; + assert!(Embedder::load(&p, &m).is_err()); + } + let _ = std::fs::remove_file(p); + } + #[test] + fn concurrent_faces_share_weights_and_give_deterministic_embeddings() { + let p = temp_model("concurrent", &crate::synthetic::tiny_embedder_model(16)); + let e = Embedder::load(&p, &manifest(16)).unwrap(); + let input = vec![123; 112 * 112 * 3]; + let want = e.embed(&input).unwrap(); + let threads: Vec<_> = (0..8) + .map(|_| { + let e = e.clone(); + let input = input.clone(); + std::thread::spawn(move || e.embed(&input)) + }) + .collect(); + for thread in threads { + assert_eq!(thread.join().unwrap().unwrap(), want); + } + let _ = std::fs::remove_file(p); + } +} diff --git a/crates/faces/src/suggest.rs b/crates/faces/src/suggest.rs new file mode 100644 index 00000000..377452e4 --- /dev/null +++ b/crates/faces/src/suggest.rs @@ -0,0 +1,237 @@ +//! What to do with a model file the user drops in: recognise it by hash, or look at its shape and propose a +//! manifest the user confirms, or say plainly why it cannot be used yet. + +use crate::manifest::{Commercial, InputSpec, Licence, ModelManifest, OutputSpec, Role, Thresholds}; +use crate::onnx::{Dim, OnnxInfo, TensorInfo}; + +/// ONNX element type for float32. +const FLOAT32: u32 = 1; + +#[derive(Clone, Debug, PartialEq)] +pub enum Suggestion { + /// A model we know by its SHA-256: its manifest as it is. + Known(ModelManifest), + /// An unknown model shaped like a face embedder. `assumptions` is what we guessed, in plain words. + Draft { manifest: ModelManifest, assumptions: Vec }, + /// Not usable yet, and why. + Unsupported(String), +} + +fn fixed(d: &Dim) -> Option { + match d { + Dim::Fixed(v) => Some(*v), + Dim::Dynamic(_) => None, + } +} + +/// An id made of the file name and the start of the hash: lowercase letters, digits, `.` `_` `-`. +fn make_id(file_name: &str, sha256: &str) -> String { + let stem = file_name.rsplit_once('.').map_or(file_name, |(s, _)| s); + let clean: String = stem.chars().map(|c| if c.is_ascii_alphanumeric() { c.to_ascii_lowercase() } else { '-' }).collect(); + let clean: String = clean.trim_matches('-').chars().take(40).collect(); + let short: String = sha256.chars().take(8).collect(); + match (clean.is_empty(), short.is_empty()) { + (true, true) => "custom".into(), + (true, false) => format!("custom-{short}"), + (false, true) => format!("custom-{clean}"), + (false, false) => format!("custom-{clean}-{short}"), + } +} + +fn display_name(file_name: &str) -> String { + let stem = file_name.rsplit_once('.').map_or(file_name, |(s, _)| s); + let name: String = stem.chars().filter(|c| !c.is_control()).take(60).collect(); + if name.trim().is_empty() { "Custom face model".into() } else { format!("{} (your model)", name.trim()) } +} + +/// The input of a face embedder: float32, `[batch, 3, height, width]`. +fn face_input(t: &TensorInfo) -> Option<(u32, u32)> { + if t.elem_type != FLOAT32 || t.shape.len() != 4 || t.shape.get(1).and_then(fixed) != Some(3) { + return None; + } + let size = |d: Option<&Dim>| match d.and_then(fixed) { + Some(v) if (16..=1024).contains(&v) => u32::try_from(v).ok(), + Some(_) => None, + None => Some(112), // a dynamic size: the common face-crop size + }; + Some((size(t.shape.get(3))?, size(t.shape.get(2))?)) +} + +/// Propose what to do with a probed model. `sha256` and `size` describe the file, `file_name` names it. +pub fn suggest(info: &OnnxInfo, sha256: &str, size: u64, file_name: &str) -> Suggestion { + suggest_with(info, sha256, size, file_name, &[]) +} + +/// [`suggest`], also recognising the models in `extra` (the user's own catalog) by their SHA-256. +pub fn suggest_with(info: &OnnxInfo, sha256: &str, size: u64, file_name: &str, extra: &[ModelManifest]) -> Suggestion { + if let Some(m) = crate::known::lookup(sha256).or_else(|| extra.iter().find(|m| m.sha256.as_deref() == Some(sha256)).cloned()) { + return Suggestion::Known(m); + } + let [input] = info.inputs.as_slice() else { + return Suggestion::Unsupported(format!("the model has {} inputs; face models here take one image input", info.inputs.len())); + }; + let Some((width, height)) = face_input(input) else { + return Suggestion::Unsupported("its input is not a float image of shape [batch, 3, height, width]".into()); + }; + let [output] = info.outputs.as_slice() else { + return Suggestion::Unsupported( + "it has several outputs, like a face detector; only models that give one vector per face are supported for now".into(), + ); + }; + if output.elem_type != FLOAT32 || output.shape.len() != 2 { + return Suggestion::Unsupported("its output is not one vector per face (shape [batch, dimension])".into()); + } + let dim = match output.shape.get(1).and_then(fixed).map(u32::try_from) { + Some(Ok(d)) if (8..=4096).contains(&d) => d, + _ => return Suggestion::Unsupported("its output size is not a fixed number between 8 and 4096".into()), + }; + let manifest = ModelManifest { + id: make_id(file_name, sha256), + name: display_name(file_name), + version: "1".into(), + role: Role::Embedder, + licence: Licence { + name: "Unknown".into(), + commercial: Commercial::Unknown, + url: None, + notice: "LightCraft does not know where this model came from or what its licence allows. Use it only if its licence allows what you are doing, and do not redistribute it.".into(), + }, + source: None, + sha256: Some(sha256.to_string()).filter(|h| h.len() == 64), + size_bytes: Some(size).filter(|s| *s > 0), + provenance: "Unknown.".into(), + input: InputSpec { width, height, ..InputSpec::default() }, + output: OutputSpec::Embedding { dim }, + thresholds: Thresholds::default(), + speed: None, + }; + let assumptions = vec![ + format!("Faces are aligned crops of {width} × {height} pixels."), + "Pixels are RGB, normalised as (x − 127.5) / 127.5: the usual ArcFace convention, which InsightFace models use.".into(), + format!("Each face becomes a vector of {dim} numbers; two faces match when their cosine similarity is high."), + ]; + Suggestion::Draft { manifest, assumptions } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::manifest::validate; + use crate::onnx::probe_reader; + use std::io::Cursor; + + fn tensor(elem: u32, shape: &[Dim]) -> TensorInfo { + TensorInfo { name: "t".into(), elem_type: elem, shape: shape.to_vec() } + } + fn dynamic() -> Dim { + Dim::Dynamic("N".into()) + } + fn info(input: TensorInfo, output: TensorInfo) -> OnnxInfo { + OnnxInfo { producer: "x".into(), opset: Some(17), inputs: vec![input], outputs: vec![output] } + } + const SHA: &str = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"; + + #[test] + fn a_known_hash_wins() { + let i = info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(10), Dim::Fixed(10)]), tensor(1, &[dynamic(), Dim::Fixed(5)])); + assert!(matches!(suggest(&i, crate::known::AURAFACE_SHA256, 1, "whatever.onnx"), Suggestion::Known(m) if m.id == "auraface-v1")); + } + + #[test] + fn an_unknown_embedder_gets_a_draft_that_validates() { + let i = info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(512)])); + let Suggestion::Draft { manifest, assumptions } = suggest(&i, SHA, 5000, "My Model (r50).onnx") else { unreachable_in_test() }; + assert_eq!(validate(&manifest), Ok(())); + assert_eq!(manifest.id, "custom-my-model--r50-01234567"); + assert_eq!(manifest.output, OutputSpec::Embedding { dim: 512 }); + assert_eq!((manifest.input.width, manifest.input.height), (112, 112)); + assert_eq!(manifest.licence.commercial, Commercial::Unknown); + assert!(assumptions.iter().any(|a| a.contains("127.5"))); + } + + #[allow(clippy::panic)] + fn unreachable_in_test() -> ! { + panic!("expected a draft"); + } + + #[test] + fn dynamic_sizes_default_and_odd_shapes_are_refused() { + let ok = info(tensor(1, &[dynamic(), Dim::Fixed(3), dynamic(), dynamic()]), tensor(1, &[dynamic(), Dim::Fixed(128)])); + assert!(matches!(suggest(&ok, SHA, 1, "a.onnx"), Suggestion::Draft { manifest, .. } if manifest.input.width == 112)); + let refused = [ + info(tensor(1, &[dynamic(), Dim::Fixed(1), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(128)])), // grey + info(tensor(7, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(128)])), // int64 input + info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(128)])), // rank 3 + info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(1_000_000), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(128)])), + info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(257), Dim::Fixed(384)])), // tokens + info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), dynamic()])), + info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(2)])), + info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(1_000_000)])), + ]; + for r in refused { + assert!(matches!(suggest(&r, SHA, 1, "a.onnx"), Suggestion::Unsupported(_)), "{r:?}"); + } + let mut two = ok.clone(); + two.outputs.push(two.outputs[0].clone()); + assert!(matches!(suggest(&two, SHA, 1, "a.onnx"), Suggestion::Unsupported(_)), "several outputs"); + } + + #[test] + fn file_names_cannot_break_the_manifest() { + let i = info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(512)])); + for name in ["", ".onnx", "../../etc/passwd", "\u{7}\u{0}x.onnx", &"x".repeat(5000), "ÅÄÖ 🙂.onnx", "UPPER lower.ONNX"] { + let Suggestion::Draft { manifest, .. } = suggest(&i, SHA, 1, name) else { unreachable_in_test() }; + assert_eq!(validate(&manifest), Ok(()), "{name:?}"); + } + // and with no usable hash at all + let Suggestion::Draft { manifest, .. } = suggest(&i, "", 0, "x.onnx") else { unreachable_in_test() }; + assert_eq!(validate(&manifest), Ok(())); + assert!(manifest.sha256.is_none() && manifest.size_bytes.is_none()); + } + + /// Opt-in: `LC_FACE_MODELS=/folder/with/onnx/files cargo test -p lightcraft-faces -- --ignored --nocapture` + /// probes real models (never committed) and prints what would be suggested, with timings. + #[test] + #[ignore = "needs real model files: set LC_FACE_MODELS"] + fn probing_real_models() { + let Some(dir) = std::env::var_os("LC_FACE_MODELS") else { return }; + for entry in std::fs::read_dir(dir).unwrap().flatten() { + let path = entry.path(); + if path.extension().is_none_or(|e| e != "onnx") { + continue; + } + let t = std::time::Instant::now(); + let sha = crate::hash::sha256_file(&path).unwrap(); + let hashed = t.elapsed(); + let t = std::time::Instant::now(); + let info = crate::onnx::probe_path(&path); + let probed = t.elapsed(); + let size = std::fs::metadata(&path).unwrap().len(); + let name = path.file_name().unwrap().to_string_lossy().into_owned(); + let verdict = match &info { + Ok(i) => match suggest(i, &sha, size, &name) { + Suggestion::Known(m) => format!("KNOWN {} ({:?})", m.id, m.licence.commercial), + Suggestion::Draft { manifest, .. } => format!("DRAFT {} {:?}", manifest.id, manifest.output), + Suggestion::Unsupported(why) => format!("UNSUPPORTED: {why}"), + }, + Err(e) => format!("PROBE ERROR: {e}"), + }; + println!("{name}: {} MB | hash {hashed:?} | probe {probed:?} | {verdict}", size / 1_048_576); + if let Ok(i) = info { + println!( + " in {:?} + out {:?}", + i.inputs, i.outputs + ); + } + } + } + + #[test] + fn probing_then_suggesting_end_to_end() { + let probed = probe_reader(&mut Cursor::new(crate::onnx::tests::model_bytes(512))).unwrap(); + assert!( + matches!(suggest(&probed, SHA, 9, "m.onnx"), Suggestion::Draft { manifest, .. } if manifest.output == OutputSpec::Embedding { dim: 512 }) + ); + } +} diff --git a/crates/faces/src/synthetic.rs b/crates/faces/src/synthetic.rs new file mode 100644 index 00000000..61010a24 --- /dev/null +++ b/crates/faces/src/synthetic.rs @@ -0,0 +1,250 @@ +//! A minimal, valid ONNX model made in memory: for tests, and later for the install self-check. It is not a +//! working network (it has no layers), only a graph with the right inputs, outputs and a weight blob. + +pub(crate) fn varint(mut v: u64, out: &mut Vec) { + loop { + let b = (v & 0x7f) as u8; + v >>= 7; + if v == 0 { + out.push(b); + return; + } + out.push(b | 0x80); + } +} + +pub fn len_field(n: u64, body: &[u8], out: &mut Vec) { + varint(n << 3 | 2, out); + varint(body.len() as u64, out); + out.extend_from_slice(body); +} + +pub(crate) fn varint_field(n: u64, v: u64, out: &mut Vec) { + varint(n << 3, out); + varint(v, out); +} + +/// A graph input or output: `Ok(n)` is a fixed dimension, `Err(name)` a named (dynamic) one. +pub fn value_info_bytes(name: &str, elem: u64, dims: &[std::result::Result]) -> Vec { + let mut shape = Vec::new(); + for d in dims { + let mut dim = Vec::new(); + match d { + Ok(v) => varint_field(1, *v, &mut dim), + Err(p) => len_field(2, p.as_bytes(), &mut dim), + } + len_field(1, &dim, &mut shape); + } + let mut tensor = Vec::new(); + varint_field(1, elem, &mut tensor); + len_field(2, &shape, &mut tensor); + let mut ty = Vec::new(); + len_field(1, &tensor, &mut ty); + let mut vi = Vec::new(); + len_field(1, name.as_bytes(), &mut vi); + len_field(2, &ty, &mut vi); + vi +} + +/// A small valid model shaped like a face embedder: float input `data` `[N, 3, 112, 112]`, an old-style +/// weight also listed as an input, a weight with bulk data, float output `emb` `[N, output_dim]`, opset 17. +pub fn embedder_model(output_dim: u64) -> Vec { + let mut graph = Vec::new(); + len_field(11, &value_info_bytes("data", 1, &[Err("N"), Ok(3), Ok(112), Ok(112)]), &mut graph); + len_field(11, &value_info_bytes("w", 1, &[Ok(4)]), &mut graph); + let mut tensor = Vec::new(); + varint_field(1, 4, &mut tensor); + len_field(8, b"w", &mut tensor); + len_field(9, &[7u8; 5000], &mut tensor); + len_field(5, &tensor, &mut graph); + len_field(12, &value_info_bytes("emb", 1, &[Err("N"), Ok(output_dim)]), &mut graph); + let mut ops = Vec::new(); + len_field(1, b"", &mut ops); + varint_field(2, 17, &mut ops); + let mut model = Vec::new(); + varint_field(1, 8, &mut model); + len_field(2, b"pytorch", &mut model); + len_field(7, &graph, &mut model); + len_field(8, &ops, &mut model); + model +} + +/// A node attribute for [`node_bytes`]. +enum Attribute<'a> { + Int(i64), + Ints(&'a [i64]), + Float(f32), +} + +fn node_bytes(op: &str, inputs: &[&str], output: &str, attrs: &[(&str, Attribute)]) -> Vec { + let mut n = Vec::new(); + for i in inputs { + len_field(1, i.as_bytes(), &mut n); + } + len_field(2, output.as_bytes(), &mut n); + len_field(4, op.as_bytes(), &mut n); + for (name, value) in attrs { + let mut a = Vec::new(); + len_field(1, name.as_bytes(), &mut a); + match value { + Attribute::Float(v) => { + varint(2 << 3 | 5, &mut a); + a.extend_from_slice(&v.to_le_bytes()); + varint_field(20, 1, &mut a); + } + Attribute::Int(v) => { + varint_field(3, *v as u64, &mut a); + varint_field(20, 2, &mut a); // INT + } + Attribute::Ints(values) => { + for v in *values { + varint_field(8, *v as u64, &mut a); + } + varint_field(20, 7, &mut a); // INTS + } + } + len_field(5, &a, &mut n); + } + n +} + +fn float_tensor_bytes(name: &str, dims: &[u64], values: &[f32]) -> Vec { + let mut t = Vec::new(); + for d in dims { + varint_field(1, *d, &mut t); + } + varint_field(2, 1, &mut t); // FLOAT + len_field(8, name.as_bytes(), &mut t); + let raw: Vec = values.iter().flat_map(|v| v.to_le_bytes()).collect(); + len_field(9, &raw, &mut t); + t +} + +/// A tiny working network shaped like a face embedder, for tests of the runtime: input `data` +/// `[1, 3, 112, 112]`, a 1 × 1 convolution to `dim` channels, global average pooling and flattening to +/// `[1, dim]`. Its output depends on the picture's per-channel brightness, which is all the tests need. +pub fn tiny_embedder_model(dim: u64) -> Vec { + let mut graph = Vec::new(); + len_field(1, &node_bytes("Conv", &["data", "w", "b"], "c", &[("kernel_shape", Attribute::Ints(&[1, 1]))]), &mut graph); + len_field(1, &node_bytes("GlobalAveragePool", &["c"], "p", &[]), &mut graph); + len_field(1, &node_bytes("Flatten", &["p"], "emb", &[("axis", Attribute::Int(1))]), &mut graph); + let w: Vec = (0..dim * 3).map(|i| ((i * 37 % 19) as f32 - 9.0) / 20.0).collect(); + let b: Vec = (0..dim).map(|i| (i % 5) as f32 / 10.0 - 0.2).collect(); + len_field(5, &float_tensor_bytes("w", &[dim, 3, 1, 1], &w), &mut graph); + len_field(5, &float_tensor_bytes("b", &[dim], &b), &mut graph); + len_field(11, &value_info_bytes("data", 1, &[Ok(1), Ok(3), Ok(112), Ok(112)]), &mut graph); + len_field(12, &value_info_bytes("emb", 1, &[Ok(1), Ok(dim)]), &mut graph); + let mut ops = Vec::new(); + len_field(1, b"", &mut ops); + varint_field(2, 13, &mut ops); + let mut model = Vec::new(); + varint_field(1, 8, &mut model); + len_field(7, &graph, &mut model); + len_field(8, &ops, &mut model); + model +} + +/// A deterministic small recognition graph using all SFace/AuraFace operators (no downloaded weights). +/// Intended for runtime regression and independent reference comparisons. +pub fn recognition_model(mut seed: u32) -> Vec { + let mut random = || { + seed = seed.wrapping_mul(1664525).wrapping_add(1013904223); + ((seed >> 16) as f32 / 65535.0 - 0.5) * 0.2 + }; + let mut g = Vec::new(); + for (op, inputs, out, attrs) in [ + ("Sub", vec!["data", "center"], "centered", vec![]), + ("Mul", vec!["centered", "scale"], "scaled", vec![]), + ("Conv", vec!["scaled", "conv.w", "conv.b"], "conv", vec![("pads", Attribute::Ints(&[1, 1, 1, 1])), ("strides", Attribute::Ints(&[2, 2]))]), + ("BatchNormalization", vec!["conv", "bn.s", "bn.b", "bn.m", "bn.v"], "bn", vec![("epsilon", Attribute::Float(1e-5))]), + ("PRelu", vec!["bn", "slope"], "act", vec![]), + ("Add", vec!["act", "bn"], "skip", vec![]), + ("Dropout", vec!["skip"], "drop", vec![("ratio", Attribute::Float(0.4))]), + ("GlobalAveragePool", vec!["drop"], "pool", vec![]), + ("Flatten", vec!["pool"], "flat", vec![]), + ("Gemm", vec!["flat", "fc.w", "fc.b"], "emb", vec![("transB", Attribute::Int(1))]), + ] { + len_field(1, &node_bytes(op, &inputs, out, &attrs), &mut g); + } + for (name, dims, data) in [ + ("center", vec![1], vec![127.5]), + ("scale", vec![1], vec![1.0 / 127.5]), + ("conv.w", vec![5, 3, 3, 3], (0..135).map(|_| random()).collect()), + ("conv.b", vec![5], (0..5).map(|_| random()).collect()), + ("bn.s", vec![5], (0..5).map(|_| 1.0 + random()).collect()), + ("bn.b", vec![5], (0..5).map(|_| random()).collect()), + ("bn.m", vec![5], (0..5).map(|_| random()).collect()), + ("bn.v", vec![5], (0..5).map(|_| 1.0 + random()).collect()), + ("slope", vec![5, 1, 1], (0..5).map(|_| 0.2 + random()).collect()), + ("fc.w", vec![8, 5], (0..40).map(|_| random()).collect()), + ("fc.b", vec![8], (0..8).map(|_| random()).collect()), + ] { + len_field(5, &float_tensor_bytes(name, &dims, &data), &mut g); + } + len_field(11, &value_info_bytes("data", 1, &[Ok(1), Ok(3), Ok(112), Ok(112)]), &mut g); + len_field(12, &value_info_bytes("emb", 1, &[Ok(1), Ok(8)]), &mut g); + let mut ops = Vec::new(); + varint_field(2, 11, &mut ops); + let mut model = Vec::new(); + varint_field(1, 8, &mut model); + len_field(7, &g, &mut model); + len_field(8, &ops, &mut model); + model +} + +#[cfg(test)] +pub(crate) fn norm_overflow_model() -> Vec { + let mut graph = Vec::new(); + len_field(1, &node_bytes("Flatten", &["data"], "flat", &[]), &mut graph); + len_field(1, &node_bytes("Gemm", &["flat", "w"], "emb", &[("transB", Attribute::Int(1))]), &mut graph); + let plane = 112 * 112; + let mut weights = vec![0.0; 3 * 3 * plane]; + for c in 0..3 { + weights[c * 3 * plane + c * plane] = 5e37; + } + len_field(5, &float_tensor_bytes("w", &[3, (3 * plane) as u64], &weights), &mut graph); + len_field(11, &value_info_bytes("data", 1, &[Ok(1), Ok(3), Ok(112), Ok(112)]), &mut graph); + len_field(12, &value_info_bytes("emb", 1, &[Ok(1), Ok(3)]), &mut graph); + let mut ops = Vec::new(); + varint_field(2, 11, &mut ops); + let mut model = Vec::new(); + len_field(7, &graph, &mut model); + len_field(8, &ops, &mut model); + model +} + +/// A small YuNet-shaped stand-in that reports no faces, independent of the real downloaded model. +pub fn tiny_detector_model() -> Vec { + let mut graph = Vec::new(); + len_field(11, &value_info_bytes("input", 1, &[Ok(1), Ok(3), Ok(64), Ok(64)]), &mut graph); + for stride in [8u64, 16, 32] { + let cells = (64 / stride) * (64 / stride); + for (kind, components) in [("cls", 1u64), ("obj", 1), ("bbox", 4), ("kps", 10)] { + let name = format!("{kind}_{stride}"); + let weight = format!("weight_{name}"); + let values = vec![0.0; (cells * components) as usize]; + len_field(5, &float_tensor_bytes(&weight, &[1, cells, components], &values), &mut graph); + len_field(1, &node_bytes("Identity", &[&weight], &name, &[]), &mut graph); + len_field(12, &value_info_bytes(&name, 1, &[Ok(1), Ok(cells), Ok(components)]), &mut graph); + } + } + let mut model = Vec::new(); + varint_field(1, 8, &mut model); + len_field(7, &graph, &mut model); + let mut ops = Vec::new(); + varint_field(2, 13, &mut ops); + len_field(8, &ops, &mut model); + model +} + +/// Manifest for the stand-in, for model-install and background-scan tests only. +pub fn tiny_detector_manifest() -> crate::manifest::ModelManifest { + let bytes = tiny_detector_model(); + let mut m = crate::known::yunet(); + m.name = "Synthetic detector (tests only)".into(); + m.input.width = 64; + m.input.height = 64; + m.sha256 = Some(crate::hash::sha256_hex(&bytes)); + m.size_bytes = Some(bytes.len() as u64); + m +} diff --git a/crates/faces/src/yunet.rs b/crates/faces/src/yunet.rs new file mode 100644 index 00000000..b77f318b --- /dev/null +++ b/crates/faces/src/yunet.rs @@ -0,0 +1,370 @@ +//! The YuNet face detector (232 KB, MIT, downloaded by the user in Settings ▸ Faces), run by LightCraft's own small +//! interpreter ([`crate::net`]): no extra runtime. +//! +//! A photo is letterboxed (aspect ratio kept, padded at the right and bottom) into the network's square +//! input, the network gives a score, a box and five landmarks for every cell of three grids (strides 8, 16 and +//! 32), and the cells above the score threshold are decoded, de-duplicated by non-maximum suppression and +//! mapped back to the photo as fractions of its width and height. +//! +//! The pixels come from the caller as plain 8-bit RGB; this module has no image types of its own. + +use crate::graph; +use crate::net::{Net, NetError, Tensor}; + +const STRIDES: [usize; 3] = [8, 16, 32]; +/// Largest photo side we accept. +const MAX_SIDE: usize = 65_536; +/// Candidate cells kept before suppression, so a hostile or noisy output cannot make it slow. +const MAX_CANDIDATES: usize = 4000; + +#[derive(Debug, thiserror::Error)] +pub enum DetectError { + #[error("the detector model is not usable: {0}")] + Model(String), + #[error("{0}")] + Input(String), +} + +impl From for DetectError { + fn from(e: NetError) -> Self { + DetectError::Model(e.to_string()) + } +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Options { + /// Lowest detection score kept (0 to 1). + pub score: f32, + /// Boxes overlapping a better one by more than this (intersection over union) are dropped. + pub nms_iou: f32, + /// Most faces returned. + pub max_faces: usize, +} + +impl Default for Options { + fn default() -> Self { + Options { score: 0.6, nms_iou: 0.3, max_faces: 200 } + } +} + +/// A detected face, in fractions of the photo's width and height (0 to 1, y down). +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Face { + pub x0: f32, + pub y0: f32, + pub x1: f32, + pub y1: f32, + pub score: f32, + /// Right eye, left eye, nose tip, right and left mouth corner, as the detector names them (from the + /// subject's point of view, so the "right eye" is on the left of the picture). May lie slightly outside 0 to 1. + pub landmarks: [(f32, f32); 5], +} + +pub struct Detector { + net: Net, + /// Side of the square network input. + side: usize, +} + +impl Detector { + /// A detector from YuNet-architecture ONNX bytes (a retrained model of the same shape works too). + pub fn new(onnx: &[u8]) -> Result { + let g = graph::load(onnx).map_err(|e| DetectError::Model(e.to_string()))?; + let outputs = g.outputs.clone(); + let net = Net::new(g)?; + let side = match net.input_shape() { + Some(&[1, 3, h, w]) if h == w && (64..=4096).contains(&h) && h % 32 == 0 => h, + _ => return Err(DetectError::Model("the model's input is not a fixed square image [1, 3, N, N] with N a multiple of 32".into())), + }; + let shapes = net.validate_input(&[1, 3, side, side])?; + for stride in STRIDES { + for (kind, components) in [("cls", 1), ("obj", 1), ("bbox", 4), ("kps", 10)] { + let name = format!("{kind}_{stride}"); + let index = outputs.iter().position(|n| n == &name).ok_or_else(|| DetectError::Model(format!("the model has no output `{name}`")))?; + let count = shapes.get(index).and_then(|s| s.iter().try_fold(1usize, |n, d| n.checked_mul(*d))); + if count != Some((side / stride) * (side / stride) * components) { + return Err(DetectError::Model(format!("the output `{name}` has the wrong shape"))); + } + } + } + Ok(Detector { net, side }) + } + + /// The square input side, used to check an installed manifest. + pub fn input_side(&self) -> usize { + self.side + } + + /// Find faces in an upright `width` × `height` photo of 8-bit RGB pixels (`rgb` is `width * height * 3` bytes). + pub fn detect(&self, rgb: &[u8], width: usize, height: usize, opts: &Options) -> Result, DetectError> { + if width == 0 || height == 0 || width > MAX_SIDE || height > MAX_SIDE { + return Err(DetectError::Input(format!("a {width} × {height} picture cannot be searched"))); + } + if width.checked_mul(height).and_then(|n| n.checked_mul(3)) != Some(rgb.len()) { + return Err(DetectError::Input("the pixel data does not match the picture's size".into())); + } + let opts = Options { + score: if opts.score.is_finite() { opts.score.clamp(0.0, 1.0) } else { 0.6 }, + nms_iou: if opts.nms_iou.is_finite() { opts.nms_iou.clamp(0.0, 1.0) } else { 0.3 }, + max_faces: opts.max_faces.min(10_000), + }; + let scale = (self.side as f32 / width as f32).min(self.side as f32 / height as f32); + let input = self.letterbox(rgb, width, height, scale)?; + let outputs = self.net.run(input)?; + let find = |name: String| { + outputs.iter().find(|(n, _)| *n == name).map(|(_, t)| t).ok_or_else(|| DetectError::Model(format!("the model has no output `{name}`"))) + }; + let mut candidates: Vec = Vec::new(); + for stride in STRIDES { + let (cls, obj) = (find(format!("cls_{stride}"))?, find(format!("obj_{stride}"))?); + let (bbox, kps) = (find(format!("bbox_{stride}"))?, find(format!("kps_{stride}"))?); + let cols = self.side / stride; + let cells = cols * cols; + if cls.data.len() != cells || obj.data.len() != cells || bbox.data.len() != cells * 4 || kps.data.len() != cells * 10 { + return Err(DetectError::Model(format!("the outputs for stride {stride} have the wrong size"))); + } + for (i, (c, o)) in cls.data.iter().zip(&obj.data).enumerate() { + let score = (c.clamp(0.0, 1.0) * o.clamp(0.0, 1.0)).sqrt(); + if score.is_nan() || score < opts.score { + continue; + } + let (row, col) = ((i / cols) as f32, (i % cols) as f32); + let (Some(b), Some(k)) = (bbox.data.get(i * 4..i * 4 + 4), kps.data.get(i * 10..i * 10 + 10)) else { continue }; + let s = stride as f32; + let (cx, cy) = ((col + b[0]) * s, (row + b[1]) * s); + let (w, h) = (b[2].clamp(-10.0, 10.0).exp() * s, b[3].clamp(-10.0, 10.0).exp() * s); + let mut landmarks = [(0.0, 0.0); 5]; + for (n, l) in landmarks.iter_mut().enumerate() { + let (lx, ly) = (k.get(n * 2).copied().unwrap_or(0.0), k.get(n * 2 + 1).copied().unwrap_or(0.0)); + *l = (((lx + col) * s) / scale / width as f32, ((ly + row) * s) / scale / height as f32); + } + candidates.push(Face { + x0: (cx - w / 2.0) / scale / width as f32, + y0: (cy - h / 2.0) / scale / height as f32, + x1: (cx + w / 2.0) / scale / width as f32, + y1: (cy + h / 2.0) / scale / height as f32, + score, + landmarks, + }); + } + } + candidates.sort_by(|a, b| b.score.total_cmp(&a.score)); + candidates.truncate(MAX_CANDIDATES); + let mut kept: Vec = Vec::new(); + for c in candidates { + if kept.len() >= opts.max_faces { + break; + } + // the box as it lies inside the photo + let clipped = Face { x0: c.x0.clamp(0.0, 1.0), y0: c.y0.clamp(0.0, 1.0), x1: c.x1.clamp(0.0, 1.0), y1: c.y1.clamp(0.0, 1.0), ..c }; + if !(clipped.x1 - clipped.x0 > 1e-4 && clipped.y1 - clipped.y0 > 1e-4) || kept.iter().any(|k| iou(k, &clipped) > opts.nms_iou) { + continue; + } + kept.push(clipped); + } + Ok(kept) + } + + /// The photo scaled by `scale` into the top left of a zero-filled square, as BGR 0..255 (the model's own input). + fn letterbox(&self, rgb: &[u8], width: usize, height: usize, scale: f32) -> Result { + let side = self.side; + let nw = ((width as f32 * scale).round() as usize).clamp(1, side); + let nh = ((height as f32 * scale).round() as usize).clamp(1, side); + let (xt, yt) = (taps(width, nw), taps(height, nh)); + // horizontal pass: every source row, nw columns, three channels + let mut mid = vec![0.0f32; height * nw * 3]; + for (src_row, mid_row) in rgb.chunks_exact(width * 3).zip(mid.chunks_exact_mut(nw * 3)) { + for (x, tap) in xt.iter().enumerate() { + let mut acc = [0.0f32; 3]; + for (j, wt) in tap.weights.iter().enumerate() { + if let Some(px) = src_row.get((tap.start + j) * 3..(tap.start + j) * 3 + 3) { + acc[0] += wt * f32::from(px[0]); + acc[1] += wt * f32::from(px[1]); + acc[2] += wt * f32::from(px[2]); + } + } + if let Some(o) = mid_row.get_mut(x * 3..x * 3 + 3) { + o.copy_from_slice(&acc); + } + } + } + // vertical pass into the planes B, G, R + let plane = side * side; + let mut data = vec![0.0f32; 3 * plane]; + for (y, tap) in yt.iter().enumerate() { + for x in 0..nw { + let mut acc = [0.0f32; 3]; + for (j, wt) in tap.weights.iter().enumerate() { + if let Some(px) = mid.get(((tap.start + j) * nw + x) * 3..((tap.start + j) * nw + x) * 3 + 3) { + acc[0] += wt * px[0]; + acc[1] += wt * px[1]; + acc[2] += wt * px[2]; + } + } + let at = y * side + x; + for (c, v) in [acc[2], acc[1], acc[0]].into_iter().enumerate() { + if let Some(o) = data.get_mut(c * plane + at) { + *o = v; + } + } + } + } + Ok(Tensor::new(vec![1, 3, side, side], data)?) + } +} + +struct Tap { + start: usize, + weights: Vec, +} + +/// For each of `dst` output positions, which `src` positions to average and with what weights: exact area +/// averaging when shrinking, linear interpolation when enlarging. +fn taps(src: usize, dst: usize) -> Vec { + let r = src as f32 / dst as f32; + (0..dst) + .map(|i| { + if r >= 1.0 { + let (a, b) = (i as f32 * r, (i as f32 + 1.0) * r); + let first = a.floor() as usize; + let last = ((b.ceil() as usize).min(src)).max(first + 1); + let weights: Vec = (first..last).map(|j| ((j as f32 + 1.0).min(b) - (j as f32).max(a)).max(0.0) / r).collect(); + Tap { start: first.min(src - 1), weights } + } else { + let c = ((i as f32 + 0.5) * r - 0.5).max(0.0); + let lo = (c.floor() as usize).min(src - 1); + let hi = (lo + 1).min(src - 1); + let f = c - lo as f32; + if hi == lo { Tap { start: lo, weights: vec![1.0] } } else { Tap { start: lo, weights: vec![1.0 - f, f] } } + } + }) + .collect() +} + +fn iou(a: &Face, b: &Face) -> f32 { + let (w, h) = ((a.x1.min(b.x1) - a.x0.max(b.x0)).max(0.0), (a.y1.min(b.y1) - a.y0.max(b.y0)).max(0.0)); + let inter = w * h; + let union = (a.x1 - a.x0) * (a.y1 - a.y0) + (b.x1 - b.x0) * (b.y1 - b.y0) - inter; + if union > 0.0 { inter / union } else { 0.0 } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The real model, for the tests that need it: `LC_YUNET_MODEL=` (the file + /// Settings ▸ Faces downloads; nothing is committed). Without it those tests say so and pass. + fn model() -> Option> { + let path = std::env::var_os("LC_YUNET_MODEL")?; + let bytes = std::fs::read(path).ok()?; + assert_eq!(crate::hash::sha256_hex(&bytes), crate::known::YUNET_SHA256, "LC_YUNET_MODEL is not YuNet 2023mar"); + Some(bytes) + } + + fn detector() -> Option { + let Some(bytes) = model() else { + eprintln!("LC_YUNET_MODEL is not set: skipping"); + return None; + }; + Some(Detector::new(&bytes).unwrap()) + } + + #[test] + fn the_model_loads_and_has_a_square_input() { + let Some(d) = detector() else { return }; + assert_eq!(d.side, 640); + } + + #[test] + fn a_blank_picture_has_no_faces_and_odd_sizes_are_handled() { + let Some(d) = detector() else { return }; + for (w, h) in [(640, 480), (100, 700), (1, 1), (37, 41), (1000, 20)] { + let rgb = vec![128u8; w * h * 3]; + let faces = d.detect(&rgb, w, h, &Options::default()).unwrap(); + assert!(faces.is_empty(), "{w}×{h}: {faces:?}"); + } + } + + #[test] + fn bad_input_is_an_error_not_a_panic() { + let d = Detector::new(&crate::synthetic::tiny_detector_model()).unwrap(); + let o = Options::default(); + assert!(d.detect(&[], 0, 0, &o).is_err()); + assert!(d.detect(&[0; 12], 2, 3, &o).is_err(), "wrong byte count"); + assert!(d.detect(&[0; 3], usize::MAX, 2, &o).is_err(), "overflowing size"); + assert!(d.detect(&[0; 3], MAX_SIDE + 1, 1, &o).is_err()); + let nan = Options { score: f32::NAN, nms_iou: f32::INFINITY, max_faces: usize::MAX }; + assert!(d.detect(&vec![0u8; 12 * 12 * 3], 12, 12, &nan).is_ok(), "non-finite options fall back to defaults"); + } + + #[test] + fn missing_decoder_outputs_are_refused_at_load_time() { + let mut bytes = crate::synthetic::tiny_detector_model(); + let mut i = 0; + while i + 5 <= bytes.len() { + if &bytes[i..i + 5] == b"cls_8" { + bytes[i..i + 5].copy_from_slice(b"bad_8"); + } + i += 1; + } + let error = Detector::new(&bytes).err().unwrap().to_string(); + assert!(error.contains("cls_8"), "{error}"); + } + + #[test] + fn a_model_that_is_not_yunet_shaped_is_refused() { + let d = Detector::new(&crate::synthetic::tiny_detector_model()).unwrap(); + assert!(d.detect(&[128; 3], 1, 1, &Options::default()).unwrap().is_empty()); + assert!(Detector::new(b"nonsense").is_err()); + assert!(Detector::new(&crate::synthetic::embedder_model(512)).is_err()); + if let Some(bytes) = model() { + for n in [0, 1, 100, bytes.len() / 2, bytes.len() - 1] { + assert!(Detector::new(&bytes[..n]).is_err(), "truncated at {n}"); + } + } + } + + /// Opt-in: `LC_YUNET_REF= cargo test -p lightcraft-faces -- --ignored --nocapture` compares every output + /// of the network with ONNX Runtime's on the same input (`input.bin`, `out_.bin`, raw little-endian f32). + #[test] + #[ignore = "needs reference files: set LC_YUNET_REF"] + fn the_network_matches_onnx_runtime() { + let Some(dir) = std::env::var_os("LC_YUNET_REF") else { return }; + let dir = std::path::PathBuf::from(dir); + let floats = + |name: &str| -> Vec { std::fs::read(dir.join(name)).unwrap().as_chunks::<4>().0.iter().map(|b| f32::from_le_bytes(*b)).collect() }; + let Some(d) = detector() else { return }; + let started = std::time::Instant::now(); + let outs = d.net.run(Tensor::new(vec![1, 3, 640, 640], floats("input.bin")).unwrap()).unwrap(); + println!("network run: {:?}", started.elapsed()); + assert_eq!(outs.len(), 12); + for (name, t) in outs { + let want = floats(&format!("out_{name}.bin")); + assert_eq!(t.data.len(), want.len(), "{name}"); + let worst = t.data.iter().zip(&want).map(|(a, b)| (a - b).abs()).fold(0.0f32, f32::max); + println!("{name}: {} values, worst difference {worst:e}", want.len()); + assert!(worst < 2e-3, "{name} differs by {worst}"); + } + } + + #[test] + fn taps_cover_the_source_and_sum_to_one() { + for (src, dst) in [(6000, 640), (640, 640), (37, 640), (3, 5), (1, 1), (1, 640), (641, 640)] { + let t = taps(src, dst); + assert_eq!(t.len(), dst); + for tap in &t { + assert!(tap.start + tap.weights.len() <= src, "{src}→{dst}"); + let sum: f32 = tap.weights.iter().sum(); + assert!((sum - 1.0).abs() < 1e-3, "{src}→{dst}: {sum}"); + } + } + } + + #[test] + fn overlap_is_symmetric() { + let f = |x0, y0, x1, y1| Face { x0, y0, x1, y1, score: 1.0, landmarks: [(0.0, 0.0); 5] }; + let (a, b) = (f(0.0, 0.0, 0.5, 0.5), f(0.25, 0.25, 0.75, 0.75)); + assert!((iou(&a, &b) - iou(&b, &a)).abs() < 1e-6 && (iou(&a, &b) - 0.0625 / 0.4375).abs() < 1e-5); + assert_eq!(iou(&a, &f(0.6, 0.6, 0.9, 0.9)), 0.0); + } +} diff --git a/crates/gpu/Cargo.toml b/crates/gpu/Cargo.toml index 765e4c6d..b19acf54 100644 --- a/crates/gpu/Cargo.toml +++ b/crates/gpu/Cargo.toml @@ -9,6 +9,7 @@ description = "LightCraft GPU develop pipeline: wgpu compute (WGSL) twin of the workspace = true [dependencies] +lightcraft-denoise = { workspace = true, optional = true } lightcraft-geom = { workspace = true } lightcraft-color = { workspace = true } lightcraft-raster = { workspace = true } @@ -23,5 +24,9 @@ wgpu = { workspace = true } pollster = { workspace = true } bytemuck = { workspace = true } +[features] +default = [] +denoise = ["dep:lightcraft-denoise", "lightcraft-denoise/runtime"] + [dev-dependencies] lightcraft-scenes = { workspace = true } diff --git a/crates/gpu/src/lib.rs b/crates/gpu/src/lib.rs index 312d4263..9c234964 100644 --- a/crates/gpu/src/lib.rs +++ b/crates/gpu/src/lib.rs @@ -29,6 +29,8 @@ use lightcraft_raster::Rgb32f; pub mod backend; #[cfg(not(target_arch = "wasm32"))] mod ctx; +#[cfg(all(feature = "denoise", not(target_arch = "wasm32")))] +pub mod nn; #[cfg(not(target_arch = "wasm32"))] mod params; #[cfg(not(target_arch = "wasm32"))] @@ -145,6 +147,16 @@ pub fn enabled() -> bool { ENABLED.load(Ordering::Relaxed) && !BROKEN.load(Ordering::Relaxed) && !env_disabled() } +/// Why GPU compute that has a device of its own (denoise) is switched off by the user's settings, or `None`: the +/// environment switches and the rendering preference, but not a render-device failure (the devices are separate). +#[cfg(all(feature = "denoise", not(target_arch = "wasm32")))] +pub(crate) fn switched_off() -> Option { + if env_disabled() { + return Some(if backend::env_off() { "disabled by LIGHTCRAFT_GPU_BACKEND=off" } else { "disabled by LIGHTCRAFT_GPU=0" }.into()); + } + (!ENABLED.load(Ordering::Relaxed)).then(|| "disabled by the GPU rendering preference (app.gpu)".into()) +} + /// `LIGHTCRAFT_GPU=0` (or `LIGHTCRAFT_GPU_BACKEND=off`): no GPU for the whole process. fn env_disabled() -> bool { static OFF: std::sync::OnceLock = std::sync::OnceLock::new(); diff --git a/crates/gpu/src/nn.rs b/crates/gpu/src/nn.rs new file mode 100644 index 00000000..4231dae8 --- /dev/null +++ b/crates/gpu/src/nn.rs @@ -0,0 +1,1048 @@ +//! Denoise networks on the GPU: a [`NetRunner`] executes a [`Net`] (the plain description of a U-Net's layers) with +//! WGSL compute kernels, one tile at a time, behind the same [`TileRunner`] interface the CPU runner has. +//! +//! The runner has a device of its own, so a denoise job and the interactive renders are separate contexts that the +//! driver time-slices, and a failure of one never takes the other down. It respects the same switches as GPU rendering +//! (`LIGHTCRAFT_GPU`, `LIGHTCRAFT_GPU_BACKEND`, the GPU preference) and the same crash sentinel around device creation. +//! `LIGHTCRAFT_GPU_ADAPTER` picks an adapter by (part of) its name, e.g. to try another card. +//! +//! Every failure is an `Err` (never a panic), and a device that errs or is lost stops being used: callers then run the +//! CPU runner. The kernels (`wgsl/nn_conv.wgsl`, `wgsl/nn_pool.wgsl`) are tested against the plain-loop reference +//! interpreter in `lightcraft_denoise::reference` and, with the real model, against the pure-Rust CPU runner. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Condvar, Mutex, OnceLock}; +use std::time::Duration; + +use bytemuck::{Pod, Zeroable}; +use lightcraft_denoise::net::{Net, Op}; +use lightcraft_denoise::run::{Error, TileRunner}; +use wgpu::util::DeviceExt; + +/// Longest wait for one tile's GPU work before the device is given up on. +const TIMEOUT: Duration = Duration::from_secs(60); +/// The largest tile side (in input cells) the dispatch sizes and 32-bit indexing are checked for. +const MAX_TILE: usize = 1024; +/// The largest single buffer used: indices in the kernels are 32-bit element offsets. +const MAX_BUFFER: u64 = 2 << 30; +/// Tiles that can be on the GPU at once (each has buffers of its own). +const WORKSPACES: usize = 2; + +// --------------------------------------------------------------------------------------------------------------------- +// the device + +struct Dev { + device: wgpu::Device, + queue: wgpu::Queue, + info: wgpu::AdapterInfo, + limits: wgpu::Limits, +} + +static DEV: OnceLock> = OnceLock::new(); +static BROKEN: AtomicBool = AtomicBool::new(false); +static BROKEN_WHY: Mutex> = Mutex::new(None); + +fn mark_broken(why: &str) { + let mut w = BROKEN_WHY.lock().unwrap_or_else(|e| e.into_inner()); + if w.is_none() { + *w = Some(why.to_string()); + } + BROKEN.store(true, Ordering::Relaxed); +} + +/// Why the GPU stopped being used for denoise (a device error), if it did. +pub fn broken_reason() -> Option { + BROKEN.load(Ordering::Relaxed).then(|| BROKEN_WHY.lock().unwrap_or_else(|e| e.into_inner()).clone().unwrap_or_else(|| "device error".into())) +} + +fn dev() -> Result<&'static Dev, String> { + if let Some(why) = broken_reason() { + return Err(format!("stopped after a GPU failure: {why}")); + } + if let Some(why) = crate::switched_off() { + return Err(why); + } + DEV.get_or_init(create_device).as_ref().map_err(|e| e.clone()) +} + +fn create_device() -> Result { + let Some(backends) = crate::backend::compute_backends() else { return Err("disabled by LIGHTCRAFT_GPU_BACKEND=off".into()) }; + crate::backend::with_init_marker(backends, || { + std::panic::catch_unwind(|| make_device(backends)).unwrap_or_else(|_| Err("device creation panicked".into())) + }) +} + +fn make_device(backends: wgpu::Backends) -> Result { + let mut desc = wgpu::InstanceDescriptor::new_without_display_handle(); + desc.backends = backends; + let instance = wgpu::Instance::new(desc); + let wanted = std::env::var("LIGHTCRAFT_GPU_ADAPTER").ok().map(|s| s.trim().to_lowercase()).filter(|s| !s.is_empty()); + let adapter = match &wanted { + Some(w) => { + let all = pollster::block_on(instance.enumerate_adapters(backends)); + let names: Vec = all.iter().map(|a| a.get_info().name).collect(); + all.into_iter() + .find(|a| { + let i = a.get_info(); + i.device_type != wgpu::DeviceType::Cpu && i.name.to_lowercase().contains(w.as_str()) + }) + .ok_or_else(|| format!("no GPU adapter named like `{w}` (found: {})", names.join(", ")))? + } + None => pollster::block_on( + instance.request_adapter(&wgpu::RequestAdapterOptions { power_preference: wgpu::PowerPreference::HighPerformance, ..Default::default() }), + ) + .map_err(|e| format!("no GPU adapter among {backends:?} ({e})"))?, + }; + let info = adapter.get_info(); + if info.device_type == wgpu::DeviceType::Cpu { + return Err(format!("software adapter ({}) skipped: the CPU is faster", info.name)); + } + let limits = adapter.limits(); + if limits.max_storage_buffers_per_shader_stage < 6 + || limits.max_compute_invocations_per_workgroup < 256 + || limits.max_compute_workgroup_size_x < 256 + { + return Err(format!("{} is too limited for the denoise kernels", info.name)); + } + if limits.max_compute_workgroup_storage_size < 16 * 1024 { + return Err(format!("{} has too little workgroup memory for the denoise kernels", info.name)); + } + let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor { + label: Some("lightcraft-denoise"), + required_limits: limits.clone(), + ..Default::default() + })) + .map_err(|e| format!("{}: device creation failed ({e})", info.name))?; + device.on_uncaptured_error(std::sync::Arc::new(|e| { + log::error!("gpu (denoise): {e}"); + mark_broken(&format!("device error: {e}")); + })); + device.set_device_lost_callback(|reason, msg| { + log::error!("gpu (denoise): device lost ({reason:?}): {msg}"); + mark_broken(&format!("device lost ({reason:?}): {msg}")); + }); + log::info!("gpu (denoise): {} ({:?})", info.name, info.backend); + Ok(Dev { device, queue, info, limits }) +} + +/// The adapter denoise would use ("NVIDIA GeForce RTX 4090 Laptop GPU (Dx12)"), or why it cannot use one. Creates the +/// device on first use. +pub fn adapter() -> Result { + dev().map(|d| format!("{} ({:?})", d.info.name, d.info.backend)) +} + +// --------------------------------------------------------------------------------------------------------------------- +// the plan: what to run, in what buffers + +#[repr(C)] +#[derive(Clone, Copy, Pod, Zeroable)] +struct ConvParams { + h: u32, + w: u32, + c0: u32, + c1: u32, + n: u32, + npad: u32, + k: u32, + mode: u32, + cout: u32, + leaky: u32, + alpha: f32, + pad: u32, +} + +#[repr(C)] +#[derive(Clone, Copy, Pod, Zeroable)] +struct PoolParams { + h: u32, + w: u32, + c4: u32, + total: u32, +} + +#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)] +enum Kernel { + Conv { bm: u32, bn: u32, bk: u32 }, + Pool, +} + +/// Where a step's tensor lives. +#[derive(Clone, Copy, Debug)] +enum Place { + Input, + Slot(usize), + Output, +} + +struct Step { + kernel: Kernel, + groups: [u32; 3], + src: Place, + src2: Option, + dst: Place, + /// Index into the per-layer GPU data (weights, bias), for the convolutions. + layer: Option, + params: Vec, +} + +struct Layout { + steps: Vec, + /// Bytes of each activation buffer. + slots: Vec, + in_floats: usize, + out_floats: usize, +} + +/// Weights arranged for the kernel: `[k][npad]` rows with `k = tap · cin + ci` (columns are output channels, zero past +/// `n`), and the bias padded the same. +struct Layer { + weight: Vec, + bias: Vec, +} + +fn unsupported(why: impl Into) -> Result { + Err(why.into()) +} + +/// The kernel shape for `n` output columns: 64 × 64 blocks when `n` is a multiple of 64, else 128 × 32. +fn blocks(n: usize, cin: usize) -> (u32, u32, u32) { + let bn = if n.is_multiple_of(64) { 64 } else { 32 }; + let bk = if cin.is_multiple_of(16) { + 16 + } else if cin.is_multiple_of(8) { + 8 + } else { + 4 + }; + (4096 / bn, bn, bk) +} + +fn round_up(v: usize, to: usize) -> usize { + v.div_ceil(to) * to +} + +/// Work out the steps, the buffers and the arranged weights for `net` on `tile × tile` input cells. +fn plan(net: &Net, tile: usize, max_binding: u64) -> Result<(Layout, Vec), String> { + if !net.fits_tile(tile) { + return unsupported("the tile does not fit the network"); + } + if tile > MAX_TILE { + return unsupported("the tile is larger than the kernels' dispatches reach"); + } + let side = |t: usize| -> usize { + let level = net.shape(t).map_or(0, |s| s.level); + if level >= 0 { tile >> level } else { tile << (-level) } + }; + let chan = |t: usize| net.shape(t).map_or(0, |s| s.channels); + if !net.in_channels.is_multiple_of(4) { + return unsupported("the input channel count is not a multiple of 4"); + } + let ops = net.ops(); + // the network must end in a 1 × 1 convolution and a depth-to-space (they run as one kernel) + let n_ops = ops.len(); + let ends_right = matches!( + (n_ops.checked_sub(2).and_then(|i| ops.get(i)), ops.last()), + (Some(Op::Conv(c)), Some(Op::DepthToSpace2 { src })) if c.k == 1 && c.src2.is_none() && *src == n_ops - 1 && c.cout == 12 && c.leaky.is_none() + ); + if !ends_right { + return unsupported("the network does not end in a 1 × 1 convolution to 12 channels and depth-to-space"); + } + // the last layer that reads each tensor + let mut last_use = vec![0usize; n_ops + 1]; + for (i, op) in ops.iter().enumerate() { + let reads: Vec = match op { + Op::Conv(c) => std::iter::once(c.src).chain(c.src2).collect(), + Op::ConvT2(c) => vec![c.src], + Op::MaxPool2 { src } | Op::DepthToSpace2 { src } => vec![*src], + }; + for t in reads { + if let Some(l) = last_use.get_mut(t) { + *l = i; + } + } + } + + use Place::{Input, Output, Slot}; + let mut slots: Vec = Vec::new(); + let mut free: Vec = Vec::new(); + let mut place: Vec> = vec![None; n_ops + 1]; + if let Some(p) = place.first_mut() { + *p = Some(Input); + } + let mut steps: Vec = Vec::new(); + let mut layers: Vec = Vec::new(); + /// A buffer of at least `bytes`: the smallest free one that is big enough, else the biggest free one grown, else new. + fn alloc(bytes: u64, slots: &mut Vec, free: &mut Vec) -> usize { + let size = |s: usize| slots.get(s).copied().unwrap_or(0); + let fits = free.iter().enumerate().filter(|&(_, &s)| size(s) >= bytes).min_by_key(|&(_, &s)| size(s)).map(|(i, _)| i); + let pick = fits.or_else(|| free.iter().enumerate().max_by_key(|&(_, &s)| size(s)).map(|(i, _)| i)); + match pick { + Some(i) => { + let s = free.swap_remove(i); + if let Some(b) = slots.get_mut(s) { + *b = (*b).max(bytes); + } + s + } + None => { + slots.push(bytes); + slots.len() - 1 + } + } + } + let get = |place: &[Option], t: usize| place.get(t).copied().flatten().ok_or_else(|| format!("tensor {t} has no buffer")); + + for (i, op) in ops.iter().enumerate() { + let out_t = i + 1; + match op { + Op::Conv(c) => { + if !c.cin.is_multiple_of(4) || !c.cout.is_multiple_of(4) { + return unsupported("a convolution whose channel counts are not multiples of 4"); + } + let s = side(c.src); + let last = i + 2 == n_ops; + let (n, mode) = (c.cout, if last { 2u32 } else { 0u32 }); + let (bm, bn, bk) = blocks(n, c.cin); + let npad = round_up(n, bn as usize); + // arranged weights: row (tap · cin + ci), column co + let taps = c.k * c.k; + let mut weight = vec![0f32; taps * c.cin * npad]; + for co in 0..c.cout { + for ci in 0..c.cin { + for ky in 0..c.k { + for kx in 0..c.k { + let v = c.weight.get(((co * c.cin + ci) * c.k + ky) * c.k + kx).copied().unwrap_or(0.0); + if let Some(w) = weight.get_mut(((ky * c.k + kx) * c.cin + ci) * npad + co) { + *w = v; + } + } + } + } + } + let mut bias = vec![0f32; npad]; + if let Some(b) = bias.get_mut(..c.cout) { + b.copy_from_slice(&c.bias); + } + let m = s * s; + let params = ConvParams { + h: s as u32, + w: s as u32, + c0: chan(c.src) as u32, + c1: c.src2.map_or(0, chan) as u32, + n: n as u32, + npad: npad as u32, + k: c.k as u32, + mode, + cout: c.cout as u32, + leaky: u32::from(c.leaky.is_some()), + alpha: c.leaky.unwrap_or(0.0), + pad: 0, + }; + let src = get(&place, c.src)?; + let src2 = c.src2.map(|t| get(&place, t)).transpose()?; + let dst = if last { + Output + } else { + let slot = alloc((m * c.cout * 4) as u64, &mut slots, &mut free); + Slot(slot) + }; + if let Some(p) = place.get_mut(out_t) { + *p = Some(dst); + } + layers.push(Layer { weight, bias }); + steps.push(Step { + kernel: Kernel::Conv { bm, bn, bk }, + groups: [m.div_ceil(bm as usize) as u32, (npad as u32) / bn, 1], + src, + src2, + dst, + layer: Some(layers.len() - 1), + params: bytemuck::bytes_of(¶ms).to_vec(), + }); + } + Op::ConvT2(c) => { + if !c.cin.is_multiple_of(4) || !c.cout.is_multiple_of(4) { + return unsupported("a transposed convolution whose channel counts are not multiples of 4"); + } + let s = side(c.src); + let n = 4 * c.cout; + let (bm, bn, bk) = blocks(n, c.cin); + let npad = round_up(n, bn as usize); + // arranged weights: row ci, column (ky · 2 + kx) · cout + co + let mut weight = vec![0f32; c.cin * npad]; + for ci in 0..c.cin { + for co in 0..c.cout { + for ky in 0..2 { + for kx in 0..2 { + let v = c.weight.get(((ci * c.cout + co) * 2 + ky) * 2 + kx).copied().unwrap_or(0.0); + if let Some(w) = weight.get_mut(ci * npad + (ky * 2 + kx) * c.cout + co) { + *w = v; + } + } + } + } + } + let mut bias = vec![0f32; npad]; + for q in 0..4 { + for co in 0..c.cout { + if let (Some(b), Some(v)) = (bias.get_mut(q * c.cout + co), c.bias.get(co)) { + *b = *v; + } + } + } + let m = s * s; + let params = ConvParams { + h: s as u32, + w: s as u32, + c0: c.cin as u32, + c1: 0, + n: n as u32, + npad: npad as u32, + k: 1, + mode: 1, + cout: c.cout as u32, + leaky: u32::from(c.leaky.is_some()), + alpha: c.leaky.unwrap_or(0.0), + pad: 0, + }; + let src = get(&place, c.src)?; + let slot = alloc((4 * m * c.cout * 4) as u64, &mut slots, &mut free); + if let Some(p) = place.get_mut(out_t) { + *p = Some(Slot(slot)); + } + layers.push(Layer { weight, bias }); + steps.push(Step { + kernel: Kernel::Conv { bm, bn, bk }, + groups: [m.div_ceil(bm as usize) as u32, (npad as u32) / bn, 1], + src, + src2: None, + dst: Slot(slot), + layer: Some(layers.len() - 1), + params: bytemuck::bytes_of(¶ms).to_vec(), + }); + } + Op::MaxPool2 { src } => { + let (s, c) = (side(*src), chan(*src)); + if !c.is_multiple_of(4) { + return unsupported("a max-pool over a channel count that is not a multiple of 4"); + } + let total = (s / 2) * (s / 2) * (c / 4); + let params = PoolParams { h: s as u32, w: s as u32, c4: (c / 4) as u32, total: total as u32 }; + let from = get(&place, *src)?; + let slot = alloc((total * 16) as u64, &mut slots, &mut free); + if let Some(p) = place.get_mut(out_t) { + *p = Some(Slot(slot)); + } + let groups = total.div_ceil(64); + steps.push(Step { + kernel: Kernel::Pool, + groups: [groups.min(8192) as u32, groups.div_ceil(8192) as u32, 1], + src: from, + src2: None, + dst: Slot(slot), + layer: None, + params: bytemuck::bytes_of(¶ms).to_vec(), + }); + } + // run as part of the 1 × 1 convolution before it + Op::DepthToSpace2 { .. } => {} + } + // buffers whose last reader was this layer are free for the next layer's result + for t in 0..=i { + if last_use.get(t).copied() == Some(i) + && let Some(Some(Slot(s))) = place.get(t) + && !free.contains(s) + { + free.push(*s); + } + } + } + if slots.iter().any(|&b| b > max_binding) { + return unsupported("an activation buffer is larger than this GPU allows"); + } + if layers.iter().any(|l| (l.weight.len() as u64).saturating_mul(4) > max_binding) { + return unsupported("a weight buffer is larger than this GPU allows"); + } + Ok((Layout { steps, slots, in_floats: net.in_channels * tile * tile, out_floats: 3 * (2 * tile) * (2 * tile) }, layers)) +} + +// --------------------------------------------------------------------------------------------------------------------- +// the GPU objects + +struct Pipelines { + conv: HashMap<(u32, u32, u32), wgpu::ComputePipeline>, + conv_layout: wgpu::BindGroupLayout, + pool: wgpu::ComputePipeline, + pool_layout: wgpu::BindGroupLayout, + /// How long each kernel took to build (the driver compiles them), for the log and the tests. + built: Vec<(String, f64)>, +} + +fn entry(binding: u32, ty: wgpu::BufferBindingType) -> wgpu::BindGroupLayoutEntry { + wgpu::BindGroupLayoutEntry { + binding, + visibility: wgpu::ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { ty, has_dynamic_offset: false, min_binding_size: None }, + count: None, + } +} + +fn pipeline(d: &Dev, label: &str, src: &str, layout: &wgpu::BindGroupLayout) -> wgpu::ComputePipeline { + let module = d.device.create_shader_module(wgpu::ShaderModuleDescriptor { label: Some(label), source: wgpu::ShaderSource::Wgsl(src.into()) }); + let pl = d.device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { label: None, bind_group_layouts: &[Some(layout)], immediate_size: 0 }); + d.device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some(label), + layout: Some(&pl), + module: &module, + entry_point: Some("main"), + // The kernels write all of their workgroup memory before reading any of it, so the zeroing wgpu adds by default + // is not needed; on DX12 it is written out as one assignment per element, which makes the shader compiler take + // a minute over a 2k-element array. + compilation_options: wgpu::PipelineCompilationOptions { zero_initialize_workgroup_memory: false, ..Default::default() }, + cache: None, + }) +} + +fn conv_source(bm: u32, bn: u32, bk: u32) -> String { + include_str!("wgsl/nn_conv.wgsl") + .replace("{{BM}}", &bm.to_string()) + .replace("{{BN}}", &bn.to_string()) + .replace("{{BK}}", &bk.to_string()) + .replace("{{A_VECS}}", &(bk * bm / 4).to_string()) + .replace("{{B_VECS}}", &(bk * bn / 4).to_string()) +} + +fn conv_bgl(d: &Dev) -> wgpu::BindGroupLayout { + use wgpu::BufferBindingType::{Storage, Uniform}; + d.device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: None, + entries: &[ + entry(0, Uniform), + entry(1, Storage { read_only: true }), + entry(2, Storage { read_only: true }), + entry(3, Storage { read_only: true }), + entry(4, Storage { read_only: true }), + entry(5, Storage { read_only: false }), + ], + }) +} + +fn pool_bgl(d: &Dev) -> wgpu::BindGroupLayout { + use wgpu::BufferBindingType::{Storage, Uniform}; + d.device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: None, + entries: &[entry(0, Uniform), entry(1, Storage { read_only: true }), entry(2, Storage { read_only: false })], + }) +} + +fn build_pipelines(d: &Dev, layout: &Layout) -> Pipelines { + let (conv_layout, pool_layout) = (conv_bgl(d), pool_bgl(d)); + let mut conv = HashMap::new(); + let mut built = Vec::new(); + for s in &layout.steps { + if let Kernel::Conv { bm, bn, bk } = s.kernel { + conv.entry((bm, bn, bk)).or_insert_with(|| { + let started = std::time::Instant::now(); + let p = pipeline(d, "nn_conv", &conv_source(bm, bn, bk), &conv_layout); + built.push((format!("conv {bm}x{bn}x{bk}"), started.elapsed().as_secs_f64() * 1000.0)); + p + }); + } + } + let started = std::time::Instant::now(); + let pool = pipeline(d, "nn_pool", include_str!("wgsl/nn_pool.wgsl"), &pool_layout); + built.push(("pool".into(), started.elapsed().as_secs_f64() * 1000.0)); + Pipelines { conv, conv_layout, pool, pool_layout, built } +} + +/// One tile's buffers and the bound steps. +struct Workspace { + input: wgpu::Buffer, + output: wgpu::Buffer, + staging: wgpu::Buffer, + bound: Vec<(Kernel, [u32; 3], wgpu::BindGroup)>, + _slots: Vec, +} + +struct Shared { + dev: &'static Dev, + layout: Layout, + pipes: Pipelines, + weights: Vec<(wgpu::Buffer, wgpu::Buffer)>, + uniforms: Vec, + dummy: wgpu::Buffer, +} + +fn storage(d: &Dev, label: &str, bytes: u64, extra: wgpu::BufferUsages) -> wgpu::Buffer { + d.device.create_buffer(&wgpu::BufferDescriptor { + label: Some(label), + size: bytes.max(16), + usage: wgpu::BufferUsages::STORAGE | extra, + mapped_at_creation: false, + }) +} + +impl Shared { + fn workspace(&self) -> Workspace { + let d = self.dev; + let slots: Vec = self.layout.slots.iter().map(|&b| storage(d, "activations", b, wgpu::BufferUsages::empty())).collect(); + let input = storage(d, "input", (self.layout.in_floats * 4) as u64, wgpu::BufferUsages::COPY_DST); + let output = storage(d, "output", (self.layout.out_floats * 4) as u64, wgpu::BufferUsages::COPY_SRC); + let staging = d.device.create_buffer(&wgpu::BufferDescriptor { + label: Some("readback"), + size: (self.layout.out_floats * 4) as u64, + usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + let buf = |p: Place| -> &wgpu::Buffer { + match p { + Place::Input => &input, + Place::Output => &output, + Place::Slot(i) => slots.get(i).unwrap_or(&self.dummy), + } + }; + let mut bound = Vec::new(); + for (i, s) in self.layout.steps.iter().enumerate() { + let uniform = self.uniforms.get(i).unwrap_or(&self.dummy); + let bg = match s.kernel { + Kernel::Conv { .. } => { + let (w, b) = s.layer.and_then(|l| self.weights.get(l)).map_or((&self.dummy, &self.dummy), |(w, b)| (w, b)); + d.device.create_bind_group(&wgpu::BindGroupDescriptor { + label: None, + layout: &self.pipes.conv_layout, + entries: &[ + wgpu::BindGroupEntry { binding: 0, resource: uniform.as_entire_binding() }, + wgpu::BindGroupEntry { binding: 1, resource: buf(s.src).as_entire_binding() }, + wgpu::BindGroupEntry { binding: 2, resource: s.src2.map_or(&self.dummy, buf).as_entire_binding() }, + wgpu::BindGroupEntry { binding: 3, resource: w.as_entire_binding() }, + wgpu::BindGroupEntry { binding: 4, resource: b.as_entire_binding() }, + wgpu::BindGroupEntry { binding: 5, resource: buf(s.dst).as_entire_binding() }, + ], + }) + } + Kernel::Pool => d.device.create_bind_group(&wgpu::BindGroupDescriptor { + label: None, + layout: &self.pipes.pool_layout, + entries: &[ + wgpu::BindGroupEntry { binding: 0, resource: uniform.as_entire_binding() }, + wgpu::BindGroupEntry { binding: 1, resource: buf(s.src).as_entire_binding() }, + wgpu::BindGroupEntry { binding: 2, resource: buf(s.dst).as_entire_binding() }, + ], + }), + }; + bound.push((s.kernel, s.groups, bg)); + } + Workspace { input, output, staging, bound, _slots: slots } + } +} + +/// A network on the GPU for one tile size. Cheap to share between threads: each call takes a free workspace. +pub struct NetRunner { + shared: Shared, + tile: usize, + in_channels: usize, + free: Mutex>, + /// Workspaces made so far (they are made when a call finds none free, up to [`WORKSPACES`]). + made: Mutex, + wake: Condvar, + bytes: u64, +} + +/// Put `net` on the GPU for `tile × tile` input cells, or say why it cannot be (no GPU, a layer or a size the kernels +/// do not do, the device refusing the buffers). Creates the device on first use and compiles the kernels. +pub fn runner(net: &Net, tile: usize) -> Result { + let d = dev()?; + let max_binding = d.limits.max_storage_buffer_binding_size.min(d.limits.max_buffer_size).min(MAX_BUFFER); + let (layout, layers) = plan(net, tile, max_binding)?; + let scope = [ + d.device.push_error_scope(wgpu::ErrorFilter::OutOfMemory), + d.device.push_error_scope(wgpu::ErrorFilter::Validation), + d.device.push_error_scope(wgpu::ErrorFilter::Internal), + ]; + let pipes = build_pipelines(d, &layout); + let weights: Vec<(wgpu::Buffer, wgpu::Buffer)> = layers + .iter() + .map(|l| { + let make = |v: &[f32]| { + d.device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("weights"), + contents: bytemuck::cast_slice(v), + usage: wgpu::BufferUsages::STORAGE, + }) + }; + (make(&l.weight), make(&l.bias)) + }) + .collect(); + let uniforms: Vec = layout + .steps + .iter() + .map(|s| { + d.device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("params"), + contents: &s.params, + usage: wgpu::BufferUsages::UNIFORM, + }) + }) + .collect(); + let dummy = storage(d, "dummy", 16, wgpu::BufferUsages::empty()); + let bytes: u64 = layout.slots.iter().sum::() + ((layout.in_floats + layout.out_floats) * 8) as u64; + let shared = Shared { dev: d, layout, pipes, weights, uniforms, dummy }; + let first = shared.workspace(); + let [oom, validation, internal] = scope; + let errors = [pollster::block_on(internal.pop()), pollster::block_on(validation.pop()), pollster::block_on(oom.pop())]; + if let Some(e) = errors.into_iter().flatten().next() { + return Err(format!("{}: the denoise kernels could not be set up ({e})", d.info.name)); + } + Ok(NetRunner { shared, tile, in_channels: net.in_channels, free: Mutex::new(vec![first]), made: Mutex::new(1), wake: Condvar::new(), bytes }) +} + +impl NetRunner { + /// The adapter this runner uses. + pub fn adapter(&self) -> String { + format!("{} ({:?})", self.shared.dev.info.name, self.shared.dev.info.backend) + } + + /// How long each kernel took to build, e.g. `conv 64x64x16: 812 ms`. + pub fn build_report(&self) -> String { + self.shared.pipes.built.iter().map(|(k, ms)| format!("{k}: {ms:.0} ms")).collect::>().join(", ") + } + + /// Device memory one tile in flight uses (bytes), without the weights. + pub fn bytes_per_tile(&self) -> u64 { + self.bytes + } + + /// A free workspace, a new one when fewer than [`WORKSPACES`] exist, else wait for one to come back. + fn take(&self) -> Result { + let mut free = self.free.lock().unwrap_or_else(|e| e.into_inner()); + loop { + if let Some(w) = free.pop() { + return Ok(w); + } + { + let mut made = self.made.lock().unwrap_or_else(|e| e.into_inner()); + if *made < WORKSPACES { + *made += 1; + drop(made); + drop(free); + let d = self.shared.dev; + let scopes = + [d.device.push_error_scope(wgpu::ErrorFilter::OutOfMemory), d.device.push_error_scope(wgpu::ErrorFilter::Validation)]; + let w = self.shared.workspace(); + let [oom, validation] = scopes; + let (validation, oom) = (pollster::block_on(validation.pop()), pollster::block_on(oom.pop())); + if let Some(e) = oom.or(validation) { + self.forget(); + return Err(Error::Runtime(format!("the GPU could not make room for another tile: {e}"))); + } + return Ok(w); + } + } + free = self.wake.wait(free).unwrap_or_else(|e| e.into_inner()); + } + } + + fn give(&self, w: Workspace) { + self.free.lock().unwrap_or_else(|e| e.into_inner()).push(w); + self.wake.notify_one(); + } + + /// A workspace is gone (dropped after a failure): room for a new one, and a waiting caller may make it. + fn forget(&self) { + let mut made = self.made.lock().unwrap_or_else(|e| e.into_inner()); + *made = made.saturating_sub(1); + drop(made); + self.wake.notify_one(); + } + + fn execute(&self, w: &Workspace, input: &[f32]) -> Result, Error> { + let d = self.shared.dev; + let (t, c) = (self.tile, self.in_channels); + // planar → channels innermost + let mut nhwc = vec![0f32; input.len()]; + for (ch, plane) in input.chunks_exact(t * t).enumerate() { + for (i, v) in plane.iter().enumerate() { + if let Some(o) = nhwc.get_mut(i * c + ch) { + *o = *v; + } + } + } + let mut calls = Vec::with_capacity(w.bound.len()); + for (kernel, groups, bg) in &w.bound { + let pipe = match kernel { + Kernel::Conv { bm, bn, bk } => self.shared.pipes.conv.get(&(*bm, *bn, *bk)), + Kernel::Pool => Some(&self.shared.pipes.pool), + }; + let Some(pipe) = pipe else { return Err(Error::Runtime("a kernel is missing".into())) }; + calls.push((pipe, groups, bg)); + } + let scopes = [ + d.device.push_error_scope(wgpu::ErrorFilter::OutOfMemory), + d.device.push_error_scope(wgpu::ErrorFilter::Validation), + d.device.push_error_scope(wgpu::ErrorFilter::Internal), + ]; + d.queue.write_buffer(&w.input, 0, bytemuck::cast_slice(&nhwc)); + let mut enc = d.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None }); + { + let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("denoise tile"), timestamp_writes: None }); + for (pipe, groups, bg) in calls { + pass.set_pipeline(pipe); + pass.set_bind_group(0, bg, &[]); + pass.dispatch_workgroups(groups[0], groups[1], groups[2]); + } + } + let bytes = (self.shared.layout.out_floats * 4) as u64; + enc.copy_buffer_to_buffer(&w.output, 0, &w.staging, 0, bytes); + d.queue.submit([enc.finish()]); + let slice = w.staging.slice(..bytes); + let (tx, rx) = std::sync::mpsc::channel(); + slice.map_async(wgpu::MapMode::Read, move |r| { + let _ = tx.send(r); + }); + let waited = d.device.poll(wgpu::PollType::Wait { submission_index: None, timeout: Some(TIMEOUT) }); + let [oom, validation, internal] = scopes; + let errors = [pollster::block_on(internal.pop()), pollster::block_on(validation.pop()), pollster::block_on(oom.pop())]; + let [internal, validation, oom] = errors; + if let Some(e) = oom { + // not a broken device: this tile runs on the CPU, and the next may fit + return Err(Error::Runtime(format!("the GPU ran out of memory: {e}"))); + } + if let Some(e) = internal.or(validation) { + mark_broken(&format!("device error: {e}")); + return Err(Error::Runtime(format!("the GPU reported an error: {e}"))); + } + if let Err(e) = waited { + mark_broken(&format!("the GPU did not finish a tile: {e}")); + return Err(Error::Runtime(format!("the GPU did not finish: {e}"))); + } + match rx.recv_timeout(Duration::from_secs(10)) { + Ok(Ok(())) => {} + Ok(Err(e)) => return Err(Error::Runtime(format!("reading the result back failed: {e}"))), + Err(e) => return Err(Error::Runtime(format!("reading the result back failed: {e}"))), + } + let out = { + let data = slice.get_mapped_range().map_err(|e| Error::Runtime(format!("reading the result back failed: {e:?}")))?; + data.as_chunks::<4>().0.iter().map(|b| f32::from_le_bytes(*b)).collect::>() + }; + w.staging.unmap(); + Ok(out) + } +} + +impl TileRunner for NetRunner { + fn run(&self, input: &[f32]) -> Result, Error> { + if let Some(why) = broken_reason() { + return Err(Error::Runtime(format!("the GPU is not used any more: {why}"))); + } + if input.len() != self.shared.layout.in_floats { + return Err(Error::Input("the tile is not the size the model wants".into())); + } + let w = self.take()?; + let r = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| self.execute(&w, input))); + match r { + Ok(Ok(out)) => { + self.give(w); + Ok(out) + } + Ok(Err(e)) => { + // its readback buffer may be left half-mapped: a fresh workspace is made when one is next needed + drop(w); + self.forget(); + Err(e) + } + Err(_) => { + // a panic inside wgpu: this workspace is not given back, and the GPU is not used again + mark_broken("a denoise tile panicked"); + drop(w); + self.forget(); + Err(Error::Runtime("the GPU runner gave up".into())) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use lightcraft_denoise::{onnx, reference, synthetic}; + + fn net(tile: u64, ch: u64, depth: usize, seed: u64) -> Net { + let path = std::env::temp_dir().join(format!("lc-nn-{}-{tile}-{ch}-{depth}-{seed}.onnx", std::process::id())); + std::fs::write(&path, synthetic::unet_onnx(tile, ch, depth, seed)).unwrap(); + let n = onnx::read(&path).unwrap(); + let _ = std::fs::remove_file(&path); + n + } + + fn picture(len: usize, seed: u64) -> Vec { + let mut s = seed | 1; + (0..len) + .map(|_| { + s ^= s >> 12; + s ^= s << 25; + s ^= s >> 27; + (s.wrapping_mul(0x2545_f491_4f6c_dd1d) >> 40) as f32 / (1u64 << 24) as f32 + }) + .collect() + } + + /// The runner, or `None` (with a note) on a machine with no usable GPU. + fn gpu(net: &Net, tile: usize) -> Option { + match runner(net, tile) { + Ok(r) => Some(r), + Err(e) => { + eprintln!("skipped: no GPU runner here ({e})"); + None + } + } + } + + fn worst(a: &[f32], b: &[f32]) -> f32 { + assert_eq!(a.len(), b.len()); + let scale = b.iter().fold(1e-3f32, |m, v| m.max(v.abs())); + a.iter().zip(b).fold(0f32, |m, (x, y)| m.max((x - y).abs())) / scale + } + + #[test] + fn the_gpu_gives_what_the_reference_gives() { + // (tile, channels, depth): a tile that is not a multiple of the block size, odd channel counts, a deeper net + for (tile, ch, depth, seed) in + [(1u64, 4u64, 0usize, 1u64), (3, 4, 0, 8), (6, 4, 1, 9), (20, 12, 2, 2), (32, 8, 2, 3), (48, 12, 2, 4), (64, 16, 3, 5)] + { + let n = net(tile, ch, depth, seed); + let Some(g) = gpu(&n, tile as usize) else { return }; + let x = picture(4 * (tile * tile) as usize, seed); + let want = reference::run(&n, tile as usize, &x).unwrap(); + let got = g.run(&x).unwrap(); + let cpu = lightcraft_denoise::cpu::NetRunner::new(&n, tile as usize).unwrap().run(&x).unwrap(); + assert!(worst(&got, &cpu) < 1e-3, "GPU and pure-Rust CPU differ"); + let e = worst(&got, &want); + eprintln!("{}: tile {tile}, {ch} channels, depth {depth}: off by {e:e} of the output's size", g.adapter()); + assert!(e < 2e-4, "tile {tile}, {ch} channels, depth {depth}: the GPU is off by {e:e} of the output's size"); + } + } + + #[test] + fn a_runner_serves_many_tiles_from_many_threads() { + let n = net(32, 8, 2, 7); + let Some(g) = gpu(&n, 32) else { return }; + let inputs: Vec> = (0..12).map(|i| picture(4 * 32 * 32, 100 + i)).collect(); + let wants: Vec> = inputs.iter().map(|x| reference::run(&n, 32, x).unwrap()).collect(); + std::thread::scope(|s| { + for chunk in inputs.chunks(3).zip(wants.chunks(3)) { + let g = &g; + s.spawn(move || { + for (x, want) in chunk.0.iter().zip(chunk.1) { + let e = worst(&g.run(x).unwrap(), want); + assert!(e < 2e-4, "off by {e:e}"); + } + }); + } + }); + } + + /// The real model on the GPU against the pure-Rust CPU runner, with timings (the numbers in docs/denoise.md come from here): + /// `LC_DENOISE_MODEL= cargo test --release -p lightcraft-gpu --lib real_model -- --ignored --nocapture` + #[test] + #[ignore = "needs the real model: see the doc comment"] + fn real_model_on_the_gpu_matches_cpu() { + use lightcraft_denoise::manifest::{DenoiserManifest, Domain, Gain}; + use lightcraft_denoise::runtime::CpuRunner; + use std::time::Instant; + let Some(model) = std::env::var_os("LC_DENOISE_MODEL") else { return }; + let net = onnx::read(std::path::Path::new(&model)).unwrap(); + let tile = 512usize; + let started = Instant::now(); + let g = runner(&net, tile).unwrap(); + println!("{}: set up in {:.2} s, {:.0} MB per tile in flight", g.adapter(), started.elapsed().as_secs_f64(), g.bytes_per_tile() as f64 / 1e6); + println!("kernels: {}", g.build_report()); + let manifest = DenoiserManifest { + id: "rawnind-bayer".into(), + name: "RawNIND Bayer".into(), + version: "1".into(), + licence: Default::default(), + source: None, + sha256: None, + size_bytes: None, + provenance: String::new(), + domain: Domain::BayerToRgb, + tile: 512, + overlap: 64, + gain: Gain::MatchMean { nominal: 1.0e6, max_deviation: 0.05 }, + }; + let cpu = CpuRunner::load(std::path::Path::new(&model), &manifest).unwrap(); + // a gradient with noise, in the model's input range + let noise = picture(4 * tile * tile, 42); + let input: Vec = + (0..4 * tile * tile).map(|i| (0.2 + 0.3 * ((i % tile) as f32 / tile as f32) + 0.1 * (noise[i] - 0.5)).clamp(0.0, 1.0)).collect(); + let t = Instant::now(); + let want = cpu.run(&input).unwrap(); + println!("CPU: {:.0} ms for one tile (one call runs on one core)", t.elapsed().as_secs_f64() * 1000.0); + let t = Instant::now(); + let got = g.run(&input).unwrap(); + println!("gpu, first tile: {:.0} ms", t.elapsed().as_secs_f64() * 1000.0); + let e = worst(&got, &want); + println!("worst difference {e:e} of the output's size"); + assert!(e < 1e-3, "the GPU does not give the CPU answer ({e:e})"); + let mut times: Vec = (0..8) + .map(|_| { + let t = Instant::now(); + g.run(&input).unwrap(); + t.elapsed().as_secs_f64() * 1000.0 + }) + .collect(); + times.sort_by(|a, b| a.total_cmp(b)); + println!("gpu, one tile at a time: median {:.0} ms (best {:.0}, worst {:.0})", times[times.len() / 2], times[0], times[times.len() - 1]); + // two threads, as the denoiser feeds it + let t = Instant::now(); + std::thread::scope(|s| { + for _ in 0..2 { + s.spawn(|| { + for _ in 0..8 { + g.run(&input).unwrap(); + } + }); + } + }); + println!("gpu, 16 tiles from 2 threads: {:.0} ms per tile", t.elapsed().as_secs_f64() * 1000.0 / 16.0); + } + + /// How long the driver takes to build one conv kernel variant: `LC_SHADER= LC_BLOCKS=128,32,16`. + #[test] + #[ignore = "a probe for shader compile times"] + fn time_a_shader_build() { + let Some(path) = std::env::var_os("LC_SHADER") else { return }; + let blocks: Vec = std::env::var("LC_BLOCKS").unwrap_or_else(|_| "128,32,16".into()).split(',').filter_map(|v| v.parse().ok()).collect(); + let (bm, bn, bk) = (blocks[0], blocks[1], blocks[2]); + let src = std::fs::read_to_string(path) + .unwrap() + .replace("{{BM}}", &bm.to_string()) + .replace("{{BN}}", &bn.to_string()) + .replace("{{BK}}", &bk.to_string()) + .replace("{{A_VECS}}", &(bk * bm / 4).to_string()) + .replace("{{B_VECS}}", &(bk * bn / 4).to_string()); + let d = dev().unwrap(); + let layout = conv_bgl(d); + let started = std::time::Instant::now(); + let _ = pipeline(d, "probe", &src, &layout); + println!("{} {bm}x{bn}x{bk}: built in {:.0} ms", d.info.backend.to_str(), started.elapsed().as_secs_f64() * 1000.0); + } + + #[test] + fn a_tile_of_the_wrong_size_is_an_error_not_a_crash() { + let n = net(16, 4, 1, 9); + let Some(g) = gpu(&n, 16) else { return }; + assert!(g.run(&[0.0; 10]).is_err()); + assert!(g.run(&[]).is_err()); + } + + #[test] + fn networks_the_kernels_cannot_run_are_refused() { + let n = net(16, 4, 1, 11); + // a tile the network does not fit + assert!(plan(&n, 15, u64::MAX).is_err()); + // a buffer limit smaller than the activations + assert!(plan(&n, 16, 64).is_err()); + } +} diff --git a/crates/gpu/src/render.rs b/crates/gpu/src/render.rs index 60ffd806..7e7ac8a3 100644 --- a/crates/gpu/src/render.rs +++ b/crates/gpu/src/render.rs @@ -717,7 +717,7 @@ pub fn render( Err(m) => { let img = host.lin.get_or_insert_with(|| Arc::new(cx.read_rgb(&lin, w, h))).clone(); let l = host.log_l.get_or_insert_with(|| Arc::new(cx.read_plane(&prep.log_l, w, h))).clone(); - lightcraft_pipeline::masks::evaluate_one(m, &plan.frame, w, h, &img, &l, s.light.exposure as f32) + lightcraft_pipeline::masks::evaluate_one(m, &plan.frame, w, h, &img, &l, s.light.exposure as f32, plan.mattes.as_deref()) } }); lightcraft_pipeline::visualize::apply(&mut image, req.overlay, &plan, overlay_mask.as_ref()); @@ -1018,7 +1018,7 @@ fn masks(cx: &mut Cx<'_>, lin: &Buf, prep: &Prep, plan: &Plan<'_>, host: &mut Ho // no kernel: evaluate on the CPU let img = host.lin.get_or_insert_with(|| Arc::new(cx.read_rgb(lin, w, h))).clone(); let l = host.log_l.get_or_insert_with(|| Arc::new(cx.read_plane(&prep.log_l, w, h))).clone(); - let mut v = lightcraft_pipeline::masks::shape_alpha(&comp.shape, frame, w, h, &img, &l, ev); + let mut v = lightcraft_pipeline::masks::shape_alpha(&comp.shape, frame, w, h, &img, &l, ev, plan.mattes.as_deref()); if comp.invert { v.data.iter_mut().for_each(|x| *x = 1.0 - *x); } diff --git a/crates/gpu/src/wgsl/nn_conv.wgsl b/crates/gpu/src/wgsl/nn_conv.wgsl new file mode 100644 index 00000000..d991e9e7 --- /dev/null +++ b/crates/gpu/src/wgsl/nn_conv.wgsl @@ -0,0 +1,162 @@ +// Convolution for the denoise networks, as an implicit matrix product: every output pixel is a row, every input +// channel of every kernel tap is a column of K, and the output channels are the columns of N. +// +// out[m][n] = bias[n] + sum_k A[m][k] * W[k][n] A[m][k] = input pixel (m + tap offset), channel ci +// +// Tensors are NHWC, channels innermost, read four at a time (every channel count is a multiple of 4). A workgroup of +// 256 threads makes a BM x BN block of the output, 4 x 4 values per thread, and walks K in steps of BK through +// workgroup memory. The source can be two tensors whose channels follow each other (a skip connection's join, so it +// is never made), and a leaky ReLU is applied on the way out. Specialised by the host: BM, BN, BK (BM * BN = 4096). +// +// Output modes: +// 0 NHWC at the input's size. +// 1 a 2x2 stride-2 transposed convolution: the columns are (dy, dx, co) and each lands in the output pixel +// (2y + dy, 2x + dx) of an output twice the size. +// 2 the last 1x1 convolution followed by depth-to-space (CRD): columns 4c + 2dy + dx land in plane c at +// (2y + dy, 2x + dx) of a planar (CHW) output twice the size. + +struct P { + h: u32, + w: u32, + c0: u32, + c1: u32, + n: u32, + npad: u32, + k: u32, + mode: u32, + cout: u32, + leaky: u32, + alpha: f32, + pad: u32, +} + +@group(0) @binding(0) var p: P; +@group(0) @binding(1) var src0: array>; +@group(0) @binding(2) var src1: array>; +@group(0) @binding(3) var wgt: array>; +@group(0) @binding(4) var bias: array>; +@group(0) @binding(5) var dst: array; + +const BM: u32 = {{BM}}u; +const BN: u32 = {{BN}}u; +const BK: u32 = {{BK}}u; +const TX: u32 = BN / 4u; +const A_ITEMS: u32 = BM * (BK / 4u); +const B_ITEMS: u32 = BK * (BN / 4u); + +// A[k][m] four pixels to a vec4, B[k][n] four channels to a vec4 +var As: array, {{A_VECS}}>; +var Bs: array, {{B_VECS}}>; + +@compute @workgroup_size(256) +fn main(@builtin(workgroup_id) wg: vec3, @builtin(local_invocation_index) tid: u32) { + let m0 = wg.x * BM; + let n0 = wg.y * BN; + let tx = tid % TX; + let ty = tid / TX; + let hw = p.h * p.w; + let cin = p.c0 + p.c1; + let taps = p.k * p.k; + let pad = (p.k - 1u) / 2u; + let chunks = cin / BK; + + var acc0 = vec4(0.0); + var acc1 = vec4(0.0); + var acc2 = vec4(0.0); + var acc3 = vec4(0.0); + + for (var tap = 0u; tap < taps; tap++) { + let ty_off = i32(tap / p.k) - i32(pad); + let tx_off = i32(tap % p.k) - i32(pad); + for (var chunk = 0u; chunk < chunks; chunk++) { + let ci0 = chunk * BK; + for (var item = tid; item < A_ITEMS; item += 256u) { + let ml = item / (BK / 4u); + let q = item % (BK / 4u); + let m = m0 + ml; + var v = vec4(0.0); + if (m < hw) { + let y = i32(m / p.w) + ty_off; + let x = i32(m % p.w) + tx_off; + if (y >= 0 && y < i32(p.h) && x >= 0 && x < i32(p.w)) { + let pix = u32(y) * p.w + u32(x); + let ci = ci0 + q * 4u; + if (ci < p.c0) { + v = src0[pix * (p.c0 / 4u) + ci / 4u]; + } else { + v = src1[pix * (p.c1 / 4u) + (ci - p.c0) / 4u]; + } + } + } + let base = (q * 4u) * (BM / 4u) + ml / 4u; + let lane = ml % 4u; + As[base][lane] = v.x; + As[base + BM / 4u][lane] = v.y; + As[base + 2u * (BM / 4u)][lane] = v.z; + As[base + 3u * (BM / 4u)][lane] = v.w; + } + for (var item = tid; item < B_ITEMS; item += 256u) { + let r = item / (BN / 4u); + let c4 = item % (BN / 4u); + let krow = tap * cin + ci0 + r; + Bs[item] = wgt[krow * (p.npad / 4u) + n0 / 4u + c4]; + } + workgroupBarrier(); + for (var kk = 0u; kk < BK; kk++) { + let a = As[kk * (BM / 4u) + ty]; + let b = Bs[kk * TX + tx]; + acc0 += a.x * b; + acc1 += a.y * b; + acc2 += a.z * b; + acc3 += a.w * b; + } + workgroupBarrier(); + } + } + + let n = n0 + tx * 4u; + if (n >= p.n) { + return; + } + let bv = bias[n / 4u]; + for (var i = 0u; i < 4u; i++) { + let m = m0 + ty * 4u + i; + if (m >= hw) { + continue; + } + var v = bv; + if (i == 0u) { v += acc0; } else if (i == 1u) { v += acc1; } else if (i == 2u) { v += acc2; } else { v += acc3; } + if (p.leaky != 0u) { + v = select(v * p.alpha, v, v >= vec4(0.0)); + } + if (p.mode == 0u) { + let o = m * p.n + n; + dst[o] = v.x; + dst[o + 1u] = v.y; + dst[o + 2u] = v.z; + dst[o + 3u] = v.w; + } else if (p.mode == 1u) { + let q = n / p.cout; + let co = n % p.cout; + let y = m / p.w; + let x = m % p.w; + let opix = (2u * y + q / 2u) * (2u * p.w) + 2u * x + q % 2u; + let o = opix * p.cout + co; + dst[o] = v.x; + dst[o + 1u] = v.y; + dst[o + 2u] = v.z; + dst[o + 3u] = v.w; + } else { + let c = n / 4u; + let y = m / p.w; + let x = m % p.w; + let ow = 2u * p.w; + let plane = c * (2u * p.h) * ow; + let o = plane + (2u * y) * ow + 2u * x; + dst[o] = v.x; + dst[o + 1u] = v.y; + dst[o + ow] = v.z; + dst[o + ow + 1u] = v.w; + } + } +} diff --git a/crates/gpu/src/wgsl/nn_pool.wgsl b/crates/gpu/src/wgsl/nn_pool.wgsl new file mode 100644 index 00000000..6117eb06 --- /dev/null +++ b/crates/gpu/src/wgsl/nn_pool.wgsl @@ -0,0 +1,28 @@ +// 2x2 maximum with stride 2 over an NHWC tensor read and written four channels at a time. + +struct P { + h: u32, + w: u32, + c4: u32, + total: u32, +} + +@group(0) @binding(0) var p: P; +@group(0) @binding(1) var src: array>; +@group(0) @binding(2) var dst: array>; + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) nwg: vec3) { + let o = gid.x + gid.y * nwg.x * 64u; + if (o >= p.total) { + return; + } + let c = o % p.c4; + let pix = o / p.c4; + let ow = p.w / 2u; + let oy = pix / ow; + let ox = pix % ow; + let i00 = ((2u * oy) * p.w + 2u * ox) * p.c4 + c; + let i10 = i00 + p.w * p.c4; + dst[o] = max(max(src[i00], src[i00 + p.c4]), max(src[i10], src[i10 + p.c4])); +} diff --git a/crates/heif/Cargo.toml b/crates/heif/Cargo.toml new file mode 100644 index 00000000..9a12bb2e --- /dev/null +++ b/crates/heif/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "lightcraft-heif" +description = "Optional HEIF/HEIC still-image decoder for LightCraft (a thin, panic-guarded wrapper around heic-rs)" +version.workspace = true +edition.workspace = true +license.workspace = true + +[dependencies] +# Pure Rust, no unsafe. Pinned exactly: a young, single-owner crate that parses untrusted files; +# moving the pin is a reviewed change (re-run the heif tests and the codecs corpus tests). +# Mirrors PhotoCraft's photocraft-heif (same team, same decision): heic-rs was chosen over the +# other permissive candidates (heif-oxide lacks 4:2:2; heifer was published the same week). +heic-rs = { version = "=0.1.1", default-features = false, features = ["std"] } + +# The tiles of a HEIC photo decode in parallel on native; wasm stays single-threaded. +[target.'cfg(not(target_arch = "wasm32"))'.dependencies] +heic-rs = { version = "=0.1.1", default-features = false, features = ["std", "parallel"] } + +[lints] +workspace = true diff --git a/crates/heif/README.md b/crates/heif/README.md new file mode 100644 index 00000000..e008a01c --- /dev/null +++ b/crates/heif/README.md @@ -0,0 +1,33 @@ +# lightcraft-heif + +The optional HEIF / HEIC decoder of LightCraft: iPhone and Mac photos. A thin, panic-guarded wrapper +around [`heic-rs`](https://github.com/tbraun96/heic-rs), a pure-Rust HEVC still-picture decoder +(no `unsafe`, MIT OR Apache-2.0): single pictures and grid-tiled photos, 8- and 10-bit (10-bit +decodes to 16-bit), alpha auxiliary images, the container's rotation/mirror/crop, ICC, EXIF and XMP. +Read-only: writing would need an HEVC encoder. + +```rust +let info = lightcraft_heif::probe(&bytes)?; // size, depth, alpha: check limits first +let img = lightcraft_heif::decode(&bytes, &lightcraft_heif::Options::default())?; +// img.width, img.height, img.has_alpha, img.sixteen_bit, img.data (RGB/RGBA), img.icc, img.exif, img.xmp +``` + +Errors are `Error::Unsupported` (image sequences, overlays, …), `Error::Limit` or +`Error::Malformed`. It never panics: every call into heic-rs runs under `catch_unwind`, and a panic +inside it (PhotoCraft's fuzzing found two in 0.1.1) becomes `Error::Malformed`. The crate is a port +of PhotoCraft's `photocraft-heif` — the sibling apps deliberately share one decoder, one pin and +one set of regression fixtures. + +## Why it is a separate, optional crate + +- **A young decoder.** heic-rs is new and has a single maintainer. It is pinned exactly + (`=0.1.1`); moving the pin is a reviewed change (re-run the crate tests and + `cargo test -p lightcraft-codecs --features heif`). +- **HEVC patents are a distributor's call.** HEVC is patent-encumbered in some jurisdictions, so + whether a build includes an HEVC decoder is a build-time choice. `lightcraft-codecs` uses this + crate only behind its non-default `heif` feature; the apps forward it + (`cargo build -p lightcraft --features heif`). Without it, HEIC files are still recognised and + opening one is a clear "HEIC/HEIF support isn't included in this build" error — the same policy + as PhotoCraft. + +The crate depends on no other LightCraft crate; `codecs` → `heif` is the only edge between them. diff --git a/crates/heif/src/lib.rs b/crates/heif/src/lib.rs new file mode 100644 index 00000000..7b19b9b9 --- /dev/null +++ b/crates/heif/src/lib.rs @@ -0,0 +1,229 @@ +//! `lightcraft-heif`: the optional HEIF / HEIC (iPhone and Mac photo) decoder. +//! +//! A thin wrapper around `heic-rs`, a pure-Rust HEVC still-picture decoder: single pictures and +//! grid-tiled photos, 8- and 10-bit (10-bit decodes to 16-bit), alpha auxiliary images, ICC, +//! EXIF and XMP. Its API is plain data ([`Info`], [`Decoded`], [`Error`]) so the crate knows +//! nothing about the rest of LightCraft; `lightcraft-codecs` adapts it behind its `heif` feature. +//! A port of PhotoCraft's `photocraft-heif` (same team, same licence) so the sibling apps share +//! one decision, one dependency pin and one set of regression fixtures. +//! +//! HEIF records orientation in the container (`irot`/`imir`, plus a `clap` crop), not in EXIF. +//! [`Options::apply_transforms`] applies them; with it off, the pixels come back as coded. +//! +//! Never panics: heic-rs is young, so every call into it runs under `catch_unwind` and a panic +//! inside it becomes [`Error::Malformed`]. + +#![forbid(unsafe_code)] +#![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)] + +use std::borrow::Cow; + +use heic_rs::PixelLayout; + +/// What a HEIF file declares, read from the container alone (no pixel is decoded). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Info { + /// Size after the container transforms (rotation, mirror, crop): what a viewer shows. + pub width: u32, + pub height: u32, + /// Size as coded, before the transforms. + pub coded_width: u32, + pub coded_height: u32, + /// The primary image has an alpha auxiliary image. + pub has_alpha: bool, + /// More than 8 bits per sample: it decodes to 16-bit. + pub sixteen_bit: bool, +} + +/// Decode settings. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Options { + /// Refuse images with more pixels than this (checked before decoding). + pub max_pixels: u64, + /// Apply the container's rotation, mirror and crop (`irot`/`imir`/`clap`). + pub apply_transforms: bool, +} + +impl Default for Options { + fn default() -> Self { + Options { max_pixels: 1 << 28, apply_transforms: true } + } +} + +/// A decoded image: interleaved RGB or RGBA samples, row-major, no padding. 16-bit samples are +/// native-endian `u16`s. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Decoded { + pub width: u32, + pub height: u32, + /// Four channels (RGBA) instead of three (RGB). + pub has_alpha: bool, + /// Two bytes per sample instead of one. + pub sixteen_bit: bool, + pub data: Vec, + /// The primary image's ICC profile. + pub icc: Option>, + /// EXIF as a TIFF structure (without HEIF's offset header). Its Orientation tag only mirrors + /// what the container says; it is returned verbatim. + pub exif: Option>, + /// The XMP packet. + pub xmp: Option, +} + +/// Why a file could not be decoded. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Error { + /// A valid file using something this decoder does not handle (image sequences, overlays…). + Unsupported(String), + /// The image is larger than [`Options::max_pixels`] or holds an oversized box. + Limit(String), + /// Broken or truncated data (or a decoder bug, reported the same way). + Malformed(String), +} + +impl std::fmt::Display for Error { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Error::Unsupported(m) => write!(f, "unsupported HEIF: {m}"), + Error::Limit(m) => write!(f, "limit exceeded: {m}"), + Error::Malformed(m) => write!(f, "malformed HEIF: {m}"), + } + } +} + +impl std::error::Error for Error {} + +const SEQUENCE: &str = "this file holds an image sequence; only HEIF still images can be opened"; + +fn err(e: heic_rs::Error) -> Error { + match e { + heic_rs::Error::Unsupported(what) => Error::Unsupported(what.to_string()), + // The primary item is neither an HEVC picture nor a grid of them. + heic_rs::Error::MissingBox("hvcC") => Error::Unsupported("the image is an overlay, an identity derivation or not HEVC-coded".into()), + heic_rs::Error::PixelLimit { .. } | heic_rs::Error::BoxTooLarge { .. } => Error::Limit(e.to_string()), + e => Error::Malformed(e.to_string()), + } +} + +/// Runs `f`, turning a heic-rs panic into an error. Fuzzing in the sibling app found two +/// out-of-range slices in heic-rs 0.1.1 on malformed files: in its box parser (`boxes.rs:130`) +/// and during reconstruction (`hevc/decode/recon.rs:70`). +fn guarded(f: impl FnOnce() -> Result) -> Result { + std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)).unwrap_or_else(|_| Err(Error::Malformed("the HEIF decoder failed on this file".into()))) +} + +/// Reads the container's declared size, depth and alpha without decoding pixels, so callers can +/// check their limits first. +pub fn probe(bytes: &[u8]) -> Result { + guarded(|| probe_unguarded(bytes)) +} + +fn probe_unguarded(bytes: &[u8]) -> Result { + let info = heic_rs::probe(bytes).map_err(|e| match e { + // No `meta` box: no still image, only a `moov` image sequence (a video track). + heic_rs::Error::MissingBox("meta") => Error::Unsupported(SEQUENCE.into()), + e => err(e), + })?; + Ok(Info { + width: info.width, + height: info.height, + coded_width: info.coded_width, + coded_height: info.coded_height, + has_alpha: info.has_alpha, + sixteen_bit: info.bit_depth > 8, + }) +} + +/// Decodes the primary image, with its metadata. +pub fn decode(bytes: &[u8], options: &Options) -> Result { + guarded(|| decode_unguarded(bytes, options)) +} + +fn decode_unguarded(bytes: &[u8], options: &Options) -> Result { + let info = probe_unguarded(bytes)?; + let pixel_layout = match (info.has_alpha, info.sixteen_bit) { + (true, true) => PixelLayout::Rgba16, + (true, false) => PixelLayout::Rgba8, + (false, true) => PixelLayout::Rgb16, + (false, false) => PixelLayout::Rgb8, + }; + let raw = heic_rs::DecodeOptions::default() + .with_layout(pixel_layout) + .with_max_pixels(Some(options.max_pixels)) + .with_transforms(options.apply_transforms) + .with_alpha(info.has_alpha); + let decoded = heic_rs::decode(bytes, &raw).map_err(err)?; + let (icc, exif, xmp) = metadata(bytes); + Ok(Decoded { + width: decoded.width, + height: decoded.height, + has_alpha: info.has_alpha, + sixteen_bit: info.sixteen_bit, + data: decoded.data, + icc, + exif, + xmp, + }) +} + +/// The primary image's ICC profile, EXIF (TIFF structure) and XMP packet. Metadata is a +/// courtesy: a broken metadata item never fails a decode whose pixels came out fine. +fn metadata(bytes: &[u8]) -> (Option>, Option>, Option) { + let Ok(ctx) = heic_rs::context::Context::open(bytes) else { + return (None, None, None); + }; + let id = ctx.meta.primary; + let icc = ctx.props(id).ok().and_then(|p| ctx.icc(&p).map(<[u8]>::to_vec)); + let exif = ctx.exif(id).ok().flatten().map(<[u8]>::to_vec); + let xmp = ctx.xmp_item(id).and_then(|x| ctx.item_data(x).ok()).and_then(|d| match d { + Cow::Borrowed(b) => String::from_utf8(b.to_vec()).ok(), + Cow::Owned(b) => String::from_utf8(b).ok(), + }); + (icc, exif, xmp.map(|x| x.trim_end_matches('\0').to_string())) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Found by PhotoCraft's `decode_heif` fuzz target: a malformed box makes heic-rs 0.1.1 slice + /// out of range (`boxes.rs:130`). It must be an error, not a crash. + const HEIC_RS_BOX_PANIC: [u8; 72] = [ + 0x00, 0x00, 0x00, 0x24, 0x66, 0x74, 0x79, 0x70, 0x68, 0x65, 0x69, 0x63, 0x00, 0x00, 0x00, 0x00, 0x6d, 0x69, 0x66, 0x31, 0x4d, 0x69, 0x50, + 0x72, 0x6d, 0x69, 0x61, 0x66, 0x4d, 0x69, 0x48, 0x42, 0x72, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x6d, 0x65, 0x74, 0x61, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x75, 0x75, 0x69, 0x64, 0x00, 0x00, 0x00, 0x00, 0x03, 0x08, 0x08, 0x08, 0x00, 0x03, 0x01, 0x03, 0x70, + 0x00, 0xa8, 0x00, + ]; + + #[test] + fn a_heic_rs_panic_is_a_malformed_error() { + for apply_transforms in [true, false] { + let r = decode(&HEIC_RS_BOX_PANIC, &Options { apply_transforms, ..Default::default() }); + assert!(matches!(r, Err(Error::Malformed(_))), "{r:?}"); + } + } + + #[test] + fn image_sequence_is_unsupported_with_a_clear_message() { + // An ftyp and nothing else: a HEIF image sequence would have a moov here instead of a meta. + let bytes = b"\0\0\0\x18ftypmsf1\0\0\0\0msf1hevc"; + for r in [probe(bytes).map(|_| ()), decode(bytes, &Options::default()).map(|_| ())] { + assert!(matches!(&r, Err(Error::Unsupported(m)) if m.contains("sequence")), "{r:?}"); + } + } + + #[test] + fn garbage_is_an_error() { + for bytes in [&b""[..], b"\0\0\0\x18ftypheic\0\0\0\0mif1heic", &[0xFF; 64]] { + assert!(probe(bytes).is_err()); + assert!(decode(bytes, &Options::default()).is_err()); + } + } + + #[test] + fn errors_display_their_kind() { + assert!(Error::Unsupported("x".into()).to_string().contains("unsupported")); + assert!(Error::Limit("x".into()).to_string().contains("limit")); + assert!(Error::Malformed("x".into()).to_string().contains("malformed")); + } +} diff --git a/crates/mcp/src/headless.rs b/crates/mcp/src/headless.rs index bdebec60..837302d6 100644 --- a/crates/mcp/src/headless.rs +++ b/crates/mcp/src/headless.rs @@ -31,7 +31,7 @@ impl Drop for Headless { impl Default for Headless { fn default() -> Self { - Self::new(Session::new().with_fs()) + Self::new(Session::new().with_fs().with_default_denoise_models().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_denoise_models().with_default_face_models()) } fn photo_or_active(&self, p: &Value) -> Result { diff --git a/crates/merge/src/output.rs b/crates/merge/src/output.rs index 442fa2ee..e020ac45 100644 --- a/crates/merge/src/output.rs +++ b/crates/merge/src/output.rs @@ -29,10 +29,16 @@ pub enum DngSamples { U16, } -/// The DNG colour tags for frames of `color`. +/// The DNG colour tags for frames of `color`. A source's gain table map is dropped: it is +/// positioned on that one capture's active area, which a merged image (aligned, oriented, +/// stitched, cropped) no longer shares. pub fn color_data(color: &FrameColor) -> ColorData { match color { - FrameColor::Camera(c) => (**c).clone(), + FrameColor::Camera(c) => { + let mut c = (**c).clone(); + c.profile.gain_table_map = None; + c + } FrameColor::Linear { to_xyz_d50 } => { let to_xyz_d65 = bradford(D50, D65).mul(to_xyz_d50); let cm = to_xyz_d65.inverse().unwrap_or(Mat3::IDENTITY); @@ -45,6 +51,7 @@ pub fn color_data(color: &FrameColor) -> ColorData { as_shot_neutral: Some([1.0, 1.0, 1.0]), as_shot_white_xy: None, baseline_exposure: 0.0, + baseline_sharpness: None, profile: Default::default(), } } diff --git a/crates/pipeline/Cargo.toml b/crates/pipeline/Cargo.toml index 27d77816..c8219f02 100644 --- a/crates/pipeline/Cargo.toml +++ b/crates/pipeline/Cargo.toml @@ -13,6 +13,7 @@ lightcraft-geom = { workspace = true } lightcraft-color = { workspace = true } lightcraft-raster = { workspace = true } lightcraft-develop = { workspace = true } +lightcraft-raw = { workspace = true } rayon = { workspace = true } serde = { workspace = true } diff --git a/crates/pipeline/src/auto.rs b/crates/pipeline/src/auto.rs index 45100a1e..3455fff2 100644 --- a/crates/pipeline/src/auto.rs +++ b/crates/pipeline/src/auto.rs @@ -120,10 +120,17 @@ pub fn auto_wb(src: &Rgb32f, info: &SourceInfo) -> (f64, f64) { return (info.as_shot_temp, info.as_shot_tint); } let avg = acc.map(|v| v / wsum); - let xyz = REC2020.to_xyz().apply(avg); - let shot = lightcraft_color::cct::temp_tint_to_xy(info.as_shot_temp, info.as_shot_tint); - let seen = bradford(REC2020.white, shot).apply(xyz); - let (t, tint) = xy_to_temp_tint(Xy::from_xyz(seen)); + // the white that makes `avg` neutral, through the same model the white balance renders with + // (`local::wb_matrix_for`): the camera's own colour model, else a Bradford adaptation + let camera = info.camera_color.as_deref().filter(|_| info.raw && !info.relative_wb); + let seen = match camera.and_then(|cc| lightcraft_raw::color::neutral_white(&cc.tags, cc.developed_for, avg)) { + Some(white) => white, + None => { + let shot = lightcraft_color::cct::temp_tint_to_xy(info.as_shot_temp, info.as_shot_tint); + Xy::from_xyz(bradford(REC2020.white, shot).apply(REC2020.to_xyz().apply(avg))) + } + }; + let (t, tint) = xy_to_temp_tint(seen); (t.clamp(2000.0, 50000.0).round(), tint.clamp(-150.0, 150.0).round()) } diff --git a/crates/pipeline/src/lib.rs b/crates/pipeline/src/lib.rs index 057fb42f..84f5bd98 100644 --- a/crates/pipeline/src/lib.rs +++ b/crates/pipeline/src/lib.rs @@ -54,8 +54,19 @@ use lightcraft_raster::{Histogram, Plane, Rgb32f, Rgba8, par_rows}; pub use tone::ToneMap; +/// A raw source's own colour model, for white balance: like Lightroom, a white balance +/// re-evaluates the camera's colour matrices at the chosen white (camera-space white balance, DNG +/// spec ch. 6) instead of adapting the as-shot rendering ([`local::wb_matrix_for`]). +#[derive(Debug, PartialEq)] +pub struct CameraColor { + /// The file's colour tags (its profile look left out: it is applied at load). + pub tags: lightcraft_raw::ColorData, + /// The white the source pixels were developed for. + pub developed_for: lightcraft_color::Xy, +} + /// Facts about the source the settings are interpreted against. -#[derive(Clone, Copy, Debug, PartialEq)] +#[derive(Clone, Debug, PartialEq)] pub struct SourceInfo { /// Lens corrections embedded in the file (DNG opcodes), relative to the EXIF-oriented source. pub lens: Option, @@ -66,12 +77,26 @@ pub struct SourceInfo { pub as_shot_tint: f64, /// No measured camera illuminant: WB adjustments are relative to the camera's rendered look. pub relative_wb: bool, + /// The camera's colour model (raw files with a colour matrix); `None`: white balance adapts + /// the developed pixels (Bradford, linear Rec.2020). + pub camera_color: Option>, pub camera_tone: Option, + /// Segmentation mattes stored in the file (DNG semantic masks): AI masks use them. + pub mattes: Option>, } impl Default for SourceInfo { fn default() -> Self { - Self { raw: false, as_shot_temp: 6500.0, as_shot_tint: 0.0, lens: None, relative_wb: false, camera_tone: None } + Self { + raw: false, + as_shot_temp: 6500.0, + as_shot_tint: 0.0, + lens: None, + relative_wb: false, + camera_color: None, + camera_tone: None, + mattes: None, + } } } @@ -332,6 +357,8 @@ pub struct Plan<'a> { /// [`spots::window_for_reads`]); this is the requested window inside the rendered one, which /// the result is cut to. pub keep: Option, + /// The source's segmentation mattes ([`SourceInfo::mattes`]). + pub mattes: Option>, } /// Long edge of the small render a windowed render estimates the whole frame's airlight from. @@ -392,7 +419,7 @@ pub fn plan<'a>(src: &Rgb32f, info: &SourceInfo, s: &'a DevelopSettings, req: &R [d.nr_luminance, d.nr_detail, d.nr_color, d.nr_color_detail, d.nr_color_smoothness].map(f64::to_bits), src_long, )); - Plan { settings, frame, w, h, px_per_long, src_long, geo, lin_key, eyes, fixed_air, keep } + Plan { settings, frame, w, h, px_per_long, src_long, geo, lin_key, eyes, fixed_air, keep, mattes: info.mattes.clone() } } /// Whether the scene-linear stage needs work only the CPU does (defringe, spot removal). @@ -474,7 +501,7 @@ fn render_impl(src: Src<'_>, info: &SourceInfo, s: &DevelopSettings, req: &Rende Some(p) if p.key == lin_key => p, _ => local::Planes { key: lin_key, ..Default::default() }, }; - let mut prep = local::prepare(lin.clone(), s, frame, px_per_long, req.quality, &mut planes); + let mut prep = local::prepare(lin.clone(), s, frame, px_per_long, req.quality, &mut planes, plan.mattes.as_deref()); if let Some(a) = plan.fixed_air { prep.air = a; } @@ -543,7 +570,7 @@ fn overlay_alpha(o: Overlay, plan: &Plan<'_>, prep: &Prepared) -> Option return Some(e.alpha.clone()); } let ev = plan.settings.light.exposure as f32; - Some(masks::evaluate_one(m, &plan.frame, plan.w, plan.h, &prep.img, &prep.log_l, ev)) + Some(masks::evaluate_one(m, &plan.frame, plan.w, plan.h, &prep.img, &prep.log_l, ev, plan.mattes.as_deref())) } /// Convenience: render a before/after pair side by side is up to the UI; this renders "before" diff --git a/crates/pipeline/src/local.rs b/crates/pipeline/src/local.rs index 61ec8482..fd7fefd3 100644 --- a/crates/pipeline/src/local.rs +++ b/crates/pipeline/src/local.rs @@ -22,13 +22,21 @@ pub fn white_balance(img: &mut Rgb32f, info: &SourceInfo, s: &DevelopSettings) { wb_gain(img, info, s, 1.0); } -/// The white-balance matrix (linear Rec.2020, luminance-preserving) for the settings, or `None` -/// when the as-shot white is kept. +/// The white-balance matrix (linear Rec.2020) for the settings, or `None` when the as-shot white +/// is kept. A raw source with its camera colour model ([`crate::CameraColor`]) is re-developed for +/// the new white in camera space, as Lightroom does (so a neutral under the chosen white renders +/// neutral and saturated colours move as the camera's matrices say); other sources are adapted +/// with Bradford, luminance-preserving. pub fn wb_matrix_for(info: &SourceInfo, s: &DevelopSettings) -> Option<[[f32; 3]; 3]> { let (t, tint) = effective_wb(info, s); if (t - info.as_shot_temp).abs() < 1e-6 && (tint - info.as_shot_tint).abs() < 1e-6 { return None; } + if let Some(cc) = info.camera_color.as_deref().filter(|_| info.raw && !info.relative_wb) + && let Some(m) = lightcraft_raw::color::rebalance(&cc.tags, cc.developed_for, temp_tint_to_xy(t, tint)) + { + return Some(m.to_f32()); + } let set = wb_matrix(&REC2020, temp_tint_to_xy(t, tint)); let shot = wb_matrix(&REC2020, temp_tint_to_xy(info.as_shot_temp, info.as_shot_tint)); let m = set.mul(&shot.inverse().unwrap_or(lightcraft_color::Mat3::IDENTITY)); @@ -279,7 +287,16 @@ pub fn plane_sigmas(s: &DevelopSettings, px_per_long: f64, q: Quality) -> PlaneS /// The spatial planes the per-pixel stage needs for `img` (white-balanced, before exposure), /// reusing whatever `planes` already holds for it. Missing planes are computed side by side (each /// one alone scales poorly: the guided filters work on small subsampled grids). -pub(crate) fn prepare(img: Arc, s: &DevelopSettings, frame: &Frame, px_per_long: f64, q: Quality, planes: &mut Planes) -> Prepared { +#[allow(clippy::too_many_arguments)] +pub(crate) fn prepare( + img: Arc, + s: &DevelopSettings, + frame: &Frame, + px_per_long: f64, + q: Quality, + planes: &mut Planes, + mattes: Option<&masks::Mattes>, +) -> Prepared { let log_l = planes.log_l.get_or_insert_with(|| Arc::new(timed("log_l", || img.map(log_lum)))).clone(); let PlaneSigmas { base: base_sigma, clarity: clarity_sigma, texture: texture_sigma, dark: dark_sigma, chroma: chroma_sigma } = plane_sigmas(s, px_per_long, q); @@ -324,7 +341,7 @@ pub(crate) fn prepare(img: Arc, s: &DevelopSettings, frame: &Frame, px_p None => (None, 1.0), }; let ev = s.light.exposure as f32; - let masks = timed("masks", || masks::evaluate(&s.masks, frame, img.width, img.height, &img, &log_l, ev)); + let masks = timed("masks", || masks::evaluate(&s.masks, frame, img.width, img.height, &img, &log_l, ev, mattes)); Prepared { img, log_l, base, clarity_blur, texture_blur, dark, chroma_blur, air, masks, px_per_long } } diff --git a/crates/pipeline/src/masks.rs b/crates/pipeline/src/masks.rs index aae261c5..4777c878 100644 --- a/crates/pipeline/src/masks.rs +++ b/crates/pipeline/src/masks.rs @@ -2,14 +2,96 @@ //! //! Shapes are defined in normalized oriented-image coordinates; sizes are in "long-edge units". //! Components combine with Add (max), Subtract (a·(1−c)) and Intersect (a·c). +//! +//! AI shapes (Sky, Subject, Background, People) use the segmentation [`Mattes`] the source carries +//! (DNG semantic masks, e.g. iPhone ProRAW's) when it has a matching one, else a heuristic. use lightcraft_develop::{BrushStroke, LocalAdjustments, Mask, MaskOp, MaskShape, SegMask}; -use lightcraft_geom::Point; -use lightcraft_raster::{Plane, Rgb32f}; +use lightcraft_geom::{Orientation, Point}; +use lightcraft_raster::{Image, Plane, Rgb32f}; use crate::for_rows; use crate::geometry::Frame; +/// What an embedded segmentation matte selects. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum MatteKind { + Sky, + /// A whole person (all people in the photo). + Person, + Skin, + Hair, + Teeth, + Glasses, +} + +/// Segmentation mattes that came with the source: 0..255 alpha planes, each covering the whole +/// (EXIF-oriented) source image at its own resolution. Several mattes of one kind (e.g. one per +/// person) are combined with max. +#[derive(Clone, Default, PartialEq)] +pub struct Mattes { + mattes: Vec<(MatteKind, Image)>, +} + +impl std::fmt::Debug for Mattes { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_list().entries(self.mattes.iter().map(|(k, m)| (k, m.width, m.height))).finish() + } +} + +impl Mattes { + /// Add a matte (empty ones are ignored). + pub fn push(&mut self, kind: MatteKind, matte: Image) { + if matte.width > 0 && matte.height > 0 && matte.data.len() == matte.width * matte.height { + self.mattes.push((kind, matte)); + } + } + + pub fn is_empty(&self) -> bool { + self.mattes.is_empty() + } + + /// Memory held (bytes). + pub fn bytes(&self) -> usize { + self.mattes.iter().map(|(_, m)| m.data.len()).sum() + } + + fn of(&self, kinds: &[MatteKind]) -> Vec<&Image> { + self.mattes.iter().filter(|(k, _)| kinds.contains(k)).map(|(_, m)| m).collect() + } + + /// The mattes a shape uses instead of its heuristic (empty: none fits). + fn for_shape(&self, shape: &MaskShape) -> Vec<&Image> { + use MatteKind::*; + // the whole person: a person matte, else the union of the parts + let person = || { + let p = self.of(&[Person]); + if p.is_empty() { self.of(&[Skin, Hair, Teeth, Glasses]) } else { p } + }; + match shape { + MaskShape::Sky => self.of(&[Sky]), + MaskShape::Subject => person(), + MaskShape::People { parts, .. } => { + let wanted: Vec = parts + .iter() + .filter_map(|p| { + let p = p.to_ascii_lowercase(); + [("skin", Skin), ("hair", Hair), ("teeth", Teeth), ("glasses", Glasses)] + .into_iter() + .find(|(n, _)| p.contains(n)) + .map(|(_, k)| k) + }) + .collect(); + let whole = + parts.is_empty() || parts.iter().any(|p| matches!(p.to_ascii_lowercase().as_str(), "person" | "entire person" | "entireperson")); + // parts without a matte (lips, clothes, …) keep the heuristic + if whole { person() } else { self.of(&wanted) } + } + _ => Vec::new(), + } + } +} + pub struct Evaluated { /// The mask's id ([`Mask::id`]). pub id: u32, @@ -23,22 +105,24 @@ fn smooth(e0: f32, e1: f32, x: f32) -> f32 { t * t * (3.0 - 2.0 * t) } -/// The visible masks with components, evaluated in order. -pub fn evaluate(masks: &[Mask], frame: &Frame, w: usize, h: usize, img: &Rgb32f, log_l: &Plane, ev: f32) -> Vec { +/// The visible masks with components, evaluated in order (`mattes`: the source's, see [`Mattes`]). +#[allow(clippy::too_many_arguments)] +pub fn evaluate(masks: &[Mask], frame: &Frame, w: usize, h: usize, img: &Rgb32f, log_l: &Plane, ev: f32, mattes: Option<&Mattes>) -> Vec { masks .iter() .filter(|m| m.visible && !m.components.is_empty()) - .map(|m| Evaluated { id: m.id, alpha: evaluate_one(m, frame, w, h, img, log_l, ev), adjust: m.adjust }) + .map(|m| Evaluated { id: m.id, alpha: evaluate_one(m, frame, w, h, img, log_l, ev, mattes), adjust: m.adjust }) .collect() } /// The alpha plane of mask `m` (whether visible or not): its components combined, inverted and /// scaled by its amount. -pub fn evaluate_one(m: &Mask, frame: &Frame, w: usize, h: usize, img: &Rgb32f, log_l: &Plane, ev: f32) -> Plane { +#[allow(clippy::too_many_arguments)] +pub fn evaluate_one(m: &Mask, frame: &Frame, w: usize, h: usize, img: &Rgb32f, log_l: &Plane, ev: f32, mattes: Option<&Mattes>) -> Plane { let mut alpha = Plane::new(w, h); let mut first = true; for comp in &m.components { - let mut c = shape_alpha(&comp.shape, frame, w, h, img, log_l, ev); + let mut c = shape_alpha(&comp.shape, frame, w, h, img, log_l, ev, mattes); if comp.invert { c.data.iter_mut().for_each(|v| *v = 1.0 - *v); } @@ -88,7 +172,13 @@ fn for_each_pos(frame: &Frame, w: usize, h: usize, out: &mut Plane, f: impl Fn(P }); } -pub fn shape_alpha(shape: &MaskShape, frame: &Frame, w: usize, h: usize, img: &Rgb32f, log_l: &Plane, ev: f32) -> Plane { +/// The alpha plane of one component shape. AI shapes use the source's matching `mattes`, if any. +#[allow(clippy::too_many_arguments)] +pub fn shape_alpha(shape: &MaskShape, frame: &Frame, w: usize, h: usize, img: &Rgb32f, log_l: &Plane, ev: f32, mattes: Option<&Mattes>) -> Plane { + let fitting = mattes.map(|m| m.for_shape(shape)).unwrap_or_default(); + if !fitting.is_empty() { + return matte_alpha(&fitting, frame, w, h); + } let mut out = Plane::new(w, h); let to_long = |p: Point| frame.norm_to_long(p); // `img`/`log_l` are before exposure: range masks select on the exposed values. @@ -180,7 +270,7 @@ pub fn shape_alpha(shape: &MaskShape, frame: &Frame, w: usize, h: usize, img: &R smooth_plane(&mut out, 0.015 * frame_px(frame, w)); } MaskShape::Background => { - let mut s = shape_alpha(&MaskShape::Subject, frame, w, h, img, log_l, ev); + let mut s = shape_alpha(&MaskShape::Subject, frame, w, h, img, log_l, ev, mattes); s.data.iter_mut().for_each(|v| *v = 1.0 - *v); out = s; } @@ -252,6 +342,71 @@ fn smooth_plane(p: &mut Plane, sigma: f32) { *p = lightcraft_raster::blur::gaussian(p, sigma.max(0.5)); } +/// The union (max) of `mattes` (0..255, over the EXIF-oriented source) at output resolution, +/// sampled like the image: user orientation, lens / perspective warp, crop and flips; area-averaged +/// first when the output is much smaller than the matte, bilinear otherwise. +fn matte_alpha(mattes: &[&Image], frame: &Frame, w: usize, h: usize) -> Plane { + let mut out = Plane::new(w, h); + let o2t = frame.out_to_oriented(w, h); + let ppl = frame.px_per_long(w).max(1e-9); + for m in mattes { + let oriented = + if frame.orient == Orientation::Normal { std::borrow::Cow::Borrowed(*m) } else { std::borrow::Cow::Owned(m.oriented(frame.orient)) }; + // matte px per output px + let k = oriented.width as f64 / frame.ow * frame.ow.max(frame.oh) / ppl; + let base = box_down(&oriented, if k >= 2.0 { (k.floor() as usize).min(64) } else { 1 }); + let (sx, sy) = ((base.width as f64 / frame.ow) as f32, (base.height as f64 / frame.oh) as f32); + for_rows(&mut out.data, w, |y, row| { + for (x, v) in row.iter_mut().enumerate() { + let s = match &frame.warp { + Some(wp) => { + let Some((_, s)) = wp.frame(&o2t, x, y) else { + continue; + }; + s + } + None => o2t.apply(Point::new(x as f64 + 0.5, y as f64 + 0.5)), + }; + *v = v.max(bilinear(&base, s.x as f32 * sx, s.y as f32 * sy) / 255.0); + } + }); + } + out +} + +/// `m` area-averaged by `f × f` (0..255 values kept); `f == 1` converts only. +fn box_down(m: &Image, f: usize) -> Plane { + let f = f.max(1); + let (bw, bh) = (m.width.div_ceil(f).max(1), m.height.div_ceil(f).max(1)); + let mut out = Plane::new(bw, bh); + for_rows(&mut out.data, bw, |by, row| { + let rows = by * f..((by + 1) * f).min(m.height); + for (bx, v) in row.iter_mut().enumerate() { + let cols = bx * f..((bx + 1) * f).min(m.width); + let (mut sum, mut n) = (0u32, 0u32); + for y in rows.clone() { + for x in cols.clone() { + sum += m.data.get(y * m.width + x).copied().unwrap_or(0) as u32; + n += 1; + } + } + *v = sum as f32 / n.max(1) as f32; + } + }); + out +} + +/// Bilinear sample at continuous pixel coordinates (pixel centres at +0.5), edges extended. +fn bilinear(p: &Plane, x: f32, y: f32) -> f32 { + let (fx, fy) = (x - 0.5, y - 0.5); + let (x0, y0) = (fx.floor(), fy.floor()); + let (tx, ty) = (fx - x0, fy - y0); + let (x0, y0) = (x0 as isize, y0 as isize); + let a = p.get_clamped(x0, y0) + (p.get_clamped(x0 + 1, y0) - p.get_clamped(x0, y0)) * tx; + let b = p.get_clamped(x0, y0 + 1) + (p.get_clamped(x0 + 1, y0 + 1) - p.get_clamped(x0, y0 + 1)) * tx; + a + (b - a) * ty +} + /// A rough display mapping for colour picking (so samples taken on screen match). pub(crate) fn tonemap_for_select(c: [f32; 3]) -> [f32; 3] { c.map(|v| v / (1.0 + v)) @@ -436,7 +591,7 @@ mod tests { let f = frame(100, 50); let img = Rgb32f::new(100, 50); let l = Plane::new(100, 50); - let a = shape_alpha(&MaskShape::Linear { start: Point::new(0.0, 0.0), end: Point::new(1.0, 0.0) }, &f, 100, 50, &img, &l, 0.0); + let a = shape_alpha(&MaskShape::Linear { start: Point::new(0.0, 0.0), end: Point::new(1.0, 0.0) }, &f, 100, 50, &img, &l, 0.0, None); assert!(a.get(1, 25) > 0.99 && a.get(98, 25) < 0.01); assert!((a.get(50, 25) - 0.5).abs() < 0.05); } @@ -452,7 +607,7 @@ mod tests { MaskShape::Prompt { text: "left".into(), seg: Some(seg.clone()), detail: vec![], edge: 0.0 }, ] { for (w, h) in [(40, 20), (400, 200)] { - let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0); + let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0, None); assert!(a.get(1, h / 2) > 0.99 && a.get(w - 2, h / 2) < 0.01, "{w}×{h}"); if w >= 400 { let soft = (0..w).filter(|x| (0.05..0.95).contains(&a.get(*x, h / 2))).count(); @@ -465,7 +620,7 @@ mod tests { d.crop.geometry.rect = lightcraft_geom::Rect::from_center(Point::new(0.75, 0.5), 0.5, 1.0); let f = Frame::new(40, 20, &d, true); let shape = MaskShape::Prompt { text: "left".into(), seg: Some(seg), detail: vec![], edge: 0.0 }; - let a = shape_alpha(&shape, &f, 20, 20, &Rgb32f::new(20, 20), &Plane::new(20, 20), 0.0); + let a = shape_alpha(&shape, &f, 20, 20, &Rgb32f::new(20, 20), &Plane::new(20, 20), 0.0, None); assert!(a.data.iter().all(|v| *v < 0.05)); } @@ -478,7 +633,7 @@ mod tests { let patch = SegMask::from_logits_in(side, &l, [0.5, 0.0, 1.0, 1.0]); let shape = MaskShape::Object { hint: vec![Point::new(0.6, 0.5)], exclude: vec![], seg: Some(coarse), detail: vec![patch], edge: 0.0 }; let (w, h) = (200, 100); - let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0); + let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0, None); assert!(a.get(20, 50) < 0.01, "outside the patch: the coarse mask"); assert!(a.get(120, 50) > 0.99, "inside the patch, its selected half"); assert!(a.get(190, 50) < 0.01, "inside the patch, its unselected half"); @@ -494,7 +649,7 @@ mod tests { let shape = MaskShape::Object { hint: vec![Point::new(0.2, 0.5)], exclude: vec![], seg: None, detail: vec![patch; 20_000], edge: 0.0 }; let (w, h) = (64, 32); let t = std::time::Instant::now(); - let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0); + let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0, None); assert!(a.get(5, 16) > 0.99 && a.get(60, 16) < 0.01); assert!(t.elapsed() < std::time::Duration::from_secs(20), "{:?}", t.elapsed()); } @@ -507,7 +662,7 @@ mod tests { let (w, h) = (400, 40); let soft_px = |edge: f64| { let shape = MaskShape::Prompt { text: "x".into(), seg: Some(seg.clone()), detail: vec![], edge }; - let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0); + let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0, None); assert!(a.get(2, h / 2) > 0.9 && a.get(w - 3, h / 2) < 0.1, "still the same selection at {edge}"); (0..w).filter(|x| (0.1..0.9).contains(&a.get(*x, h / 2))).count() }; @@ -524,7 +679,7 @@ mod tests { MaskShape::Prompt { text: "sky".into(), seg: None, detail: vec![], edge: 0.0 }, MaskShape::Prompt { text: "sky".into(), seg: Some(SegMask { side: 4, data: "damaged!".into(), rect: None }), detail: vec![], edge: 0.0 }, ] { - assert!(shape_alpha(&shape, &f, 30, 30, &img, &l, 0.0).data.iter().all(|v| *v == 0.0)); + assert!(shape_alpha(&shape, &f, 30, 30, &img, &l, 0.0, None).data.iter().all(|v| *v == 0.0)); } } @@ -534,10 +689,10 @@ mod tests { let img = Rgb32f::new(100, 100); let l = Plane::new(100, 100); let shape = MaskShape::Radial { center: Point::new(0.5, 0.5), rx: 0.2, ry: 0.2, angle: 0.0, feather: 20.0, invert: false }; - let a = shape_alpha(&shape, &f, 100, 100, &img, &l, 0.0); + let a = shape_alpha(&shape, &f, 100, 100, &img, &l, 0.0, None); assert!(a.get(50, 50) > 0.99 && a.get(5, 5) < 0.01); let m = Mask { components: vec![MaskComponent { name: None, op: MaskOp::Add, invert: true, shape }], ..Default::default() }; - let e = evaluate(&[m], &f, 100, 100, &img, &l, 0.0); + let e = evaluate(&[m], &f, 100, 100, &img, &l, 0.0, None); assert!(e[0].alpha.get(50, 50) < 0.01); } @@ -552,7 +707,7 @@ mod tests { components: vec![MaskComponent { name: None, op: MaskOp::Add, invert: false, shape: MaskShape::Brush { strokes: vec![stroke, erase] } }], ..Default::default() }; - let e = evaluate(&[m], &f, 200, 100, &img, &l, 0.0); + let e = evaluate(&[m], &f, 200, 100, &img, &l, 0.0, None); let a = &e[0].alpha; assert!(a.get(40, 50) > 0.9, "{}", a.get(40, 50)); assert!(a.get(100, 50) < 0.05, "erased centre {}", a.get(100, 50)); @@ -574,8 +729,8 @@ mod tests { ..Default::default() }; let comp = |auto| MaskShape::Brush { strokes: vec![stroke(auto)] }; - let plain = shape_alpha(&comp(false), &f, w, h, &img, &l, 0.0); - let auto = shape_alpha(&comp(true), &f, w, h, &img, &l, 0.0); + let plain = shape_alpha(&comp(false), &f, w, h, &img, &l, 0.0, None); + let auto = shape_alpha(&comp(true), &f, w, h, &img, &l, 0.0, None); // without Auto Mask the brush spills over the edge; with it, it stays on the dark side assert!(plain.get(105, 50) > 0.9, "{}", plain.get(105, 50)); assert!(auto.get(105, 50) < 0.05, "spill {}", auto.get(105, 50)); @@ -585,8 +740,8 @@ mod tests { // on a flat area Auto Mask paints like the plain brush let flat = Rgb32f::from_fn(w, h, |_, _| [0.2; 3]); let fl = flat.map(crate::local::log_lum); - let a = shape_alpha(&comp(true), &f, w, h, &flat, &fl, 0.0); - let b = shape_alpha(&comp(false), &f, w, h, &flat, &fl, 0.0); + let a = shape_alpha(&comp(true), &f, w, h, &flat, &fl, 0.0, None); + let b = shape_alpha(&comp(false), &f, w, h, &flat, &fl, 0.0, None); assert!((a.get(94, 50) - b.get(94, 50)).abs() < 0.02, "{} vs {}", a.get(94, 50), b.get(94, 50)); } @@ -609,7 +764,7 @@ mod tests { }; let refined = Mask { refine: 100.0, ..soft.clone() }; let step = |m: &Mask| { - let a = evaluate_one(m, &f, w, h, &img, &log_l, 0.0); + let a = evaluate_one(m, &f, w, h, &img, &log_l, 0.0, None); a.data[40 * w + 63] - a.data[40 * w + 56] }; let (plain, sharp) = (step(&soft), step(&refined)); @@ -640,7 +795,7 @@ mod tests { } let alpha = |refine| { let shape = MaskShape::ColorRange { samples: vec![sample], refine }; - let a = shape_alpha(&shape, &f, w, h, &img, &log_l, 0.0); + let a = shape_alpha(&shape, &f, w, h, &img, &log_l, 0.0, None); (a.get(30, 20), a.get(90, 20)) }; for refine in [0.0, 50.0, 85.0] { @@ -653,4 +808,63 @@ mod tests { let (l, r) = alpha(100.0); assert!(l > 0.95 && r < 0.3, "refine 100: {l} vs {r}"); } + + /// A half-resolution matte selecting the left `frac` of a `w × h` source. + fn left_matte(w: usize, h: usize, frac: f64) -> Image { + Image::from_fn(w / 2, h / 2, |x, _| if (x as f64 + 0.5) < frac * (w / 2) as f64 { 255 } else { 0 }) + } + + /// The Sky mask follows the source's sky matte where it has one (here: the bottom-left of a + /// uniformly bright blue photo, where the heuristic sees no sky, and not its top right, where + /// it does), through user orientation, at any output size; without one it keeps the heuristic. + #[test] + fn sky_uses_the_sources_matte() { + let (w, h) = (200, 100); + let img = Rgb32f::filled(w, h, [0.5, 0.7, 1.4]); + let l = img.map(crate::local::log_lum); + let mut mattes = Mattes::default(); + mattes.push(MatteKind::Sky, left_matte(w, h, 0.5)); + let f = frame(w, h); + let heuristic = shape_alpha(&MaskShape::Sky, &f, w, h, &img, &l, 0.0, None); + let matte = shape_alpha(&MaskShape::Sky, &f, w, h, &img, &l, 0.0, Some(&mattes)); + assert!(heuristic.get(150, 5) > 0.9 && heuristic.get(25, 90) < 0.1, "{} {}", heuristic.get(150, 5), heuristic.get(25, 90)); + assert!(matte.get(150, 5) < 0.01 && matte.get(25, 90) > 0.99, "{} {}", matte.get(150, 5), matte.get(25, 90)); + // mattes of other kinds leave the heuristic alone + let mut hair = Mattes::default(); + hair.push(MatteKind::Hair, left_matte(w, h, 0.5)); + assert_eq!(shape_alpha(&MaskShape::Sky, &f, w, h, &img, &l, 0.0, Some(&hair)), heuristic); + // rotated 90° clockwise, the source's left half is the top half; a preview and a larger + // render agree + let s = lightcraft_develop::DevelopSettings { orientation: Orientation::Rotate90, ..Default::default() }; + let f = Frame::new(w, h, &s, true); + for (ow, oh) in [(100, 200), (25, 50)] { + let small = Rgb32f::filled(ow, oh, [0.5, 0.7, 1.4]); + let a = shape_alpha(&MaskShape::Sky, &f, ow, oh, &small, &small.map(crate::local::log_lum), 0.0, Some(&mattes)); + let at = |x: f64, y: f64| a.get((x * ow as f64) as usize, (y * oh as f64) as usize); + assert!(at(0.5, 0.2) > 0.99 && at(0.5, 0.8) < 0.01, "{ow}×{oh}: {} {}", at(0.5, 0.2), at(0.5, 0.8)); + } + } + + /// Subject, Background and People use person mattes: the whole person is the union of the + /// parts; a People part without a matte keeps the heuristic. + #[test] + fn people_and_subject_use_person_mattes() { + let (w, h) = (200, 100); + let img = Rgb32f::filled(w, h, [0.2; 3]); + let l = img.map(crate::local::log_lum); + let mut mattes = Mattes::default(); + mattes.push(MatteKind::Skin, left_matte(w, h, 0.25)); + let right: Image = Image::from_fn(w / 2, h / 2, |x, _| if x >= 75 { 255 } else { 0 }); + mattes.push(MatteKind::Hair, right); + let f = frame(w, h); + let eval = |shape: &MaskShape| shape_alpha(shape, &f, w, h, &img, &l, 0.0, Some(&mattes)); + let subject = eval(&MaskShape::Subject); + assert!(subject.get(10, 50) > 0.99 && subject.get(190, 50) > 0.99 && subject.get(100, 50) < 0.01); + let background = eval(&MaskShape::Background); + assert!(background.get(100, 50) > 0.99 && background.get(10, 50) < 0.01); + let hair = eval(&MaskShape::People { person: 0, parts: vec!["Hair".into()] }); + assert!(hair.get(190, 50) > 0.99 && hair.get(10, 50) < 0.01); + let lips = MaskShape::People { person: 0, parts: vec!["Lips".into()] }; + assert_eq!(eval(&lips), shape_alpha(&lips, &f, w, h, &img, &l, 0.0, None)); + } } diff --git a/crates/pipeline/src/tests.rs b/crates/pipeline/src/tests.rs index 697a0f97..f572b21d 100644 --- a/crates/pipeline/src/tests.rs +++ b/crates/pipeline/src/tests.rs @@ -369,6 +369,8 @@ fn tint_negative_is_green_and_positive_is_magenta() { SourceInfo::default(), SourceInfo { raw: true, relative_wb: true, ..Default::default() }, SourceInfo { raw: true, as_shot_temp: 4200.0, as_shot_tint: 15.0, ..Default::default() }, + // camera-space white balance: the same sign through the camera's own colour model + SourceInfo { raw: true, as_shot_temp: 4200.0, as_shot_tint: 15.0, camera_color: Some(camera_color(4200.0, 15.0)), ..Default::default() }, ] { let mut s = DevelopSettings::default(); let neutral = render(&src, &info, &s, &RenderRequest::fit(16, 16)).image.get(8, 8); @@ -384,3 +386,52 @@ fn tint_negative_is_green_and_positive_is_magenta() { } } } + +/// A dual-illuminant (A / D65) camera colour model, its pixels developed for `temp` / `tint`. +fn camera_color(temp: f64, tint: f64) -> std::sync::Arc { + use lightcraft_color::Mat3; + let tags = lightcraft_raw::ColorData { + illuminant: [17, 21], + color_matrix: [ + Some(Mat3([[0.9, 0.2, -0.15], [-0.3, 1.25, 0.08], [0.02, -0.12, 0.85]])), + Some(Mat3([[0.7, 0.3, -0.1], [-0.35, 1.3, 0.1], [0.05, -0.2, 1.0]])), + ], + ..Default::default() + }; + std::sync::Arc::new(crate::CameraColor { tags, developed_for: lightcraft_color::cct::temp_tint_to_xy(temp, tint) }) +} + +/// Camera-space white balance: a raw source with its colour model is re-developed for the chosen +/// white, so a grey card lit by that white renders neutral. Without the model, or with a relative +/// white balance, the developed pixels are adapted (Bradford, neutral luminance kept). +#[test] +fn custom_white_balance_redevelops_in_camera_space() { + use crate::local::wb_matrix_for; + use lightcraft_color::cct::temp_tint_to_xy; + use lightcraft_develop::WbMode; + let cc = camera_color(5500.0, 0.0); + let info = SourceInfo { raw: true, as_shot_temp: 5500.0, as_shot_tint: 0.0, camera_color: Some(cc.clone()), ..Default::default() }; + let mut s = DevelopSettings::default(); + (s.wb.mode, s.wb.temp, s.wb.tint) = (WbMode::Custom, 3200.0, 10.0); + // a grey card lit by 3200 K / +10, as the camera records it, developed for the as-shot white + let a = lightcraft_raw::color::camera_transform_of(&cc.tags, cc.developed_for); + let n = lightcraft_raw::color::camera_neutral(&cc.tags, temp_tint_to_xy(3200.0, 10.0)); + let card = a.matrix.apply(std::array::from_fn(|i| n[i] * a.wb[i] as f64)).map(|v| v as f32); + let apply = |m: [[f32; 3]; 3], c: [f32; 3]| -> [f32; 3] { std::array::from_fn(|r| m[r][0] * c[0] + m[r][1] * c[1] + m[r][2] * c[2]) }; + let m = wb_matrix_for(&info, &s).unwrap(); + let out = apply(m, card); + assert!(out.iter().all(|v| (v / out[1] - 1.0).abs() < 1e-4), "{card:?} -> {out:?}"); + // the adaptation path for sources without a colour model, and for relative white balance + let luma = |c: [f32; 3]| lightcraft_color::luminance_2020(c); + for other in [SourceInfo { camera_color: None, ..info.clone() }, SourceInfo { relative_wb: true, ..info.clone() }] { + let b = wb_matrix_for(&other, &s).unwrap(); + assert_ne!(b, m); + assert!((luma(apply(b, [1.0; 3])) - 1.0).abs() < 1e-4); + } + // the white-balance picker / auto white balance invert the same model: the card names its white + let picked = crate::auto::auto_wb(&Rgb32f::filled(4, 4, card), &info); + assert_eq!(picked, (3200.0, 10.0)); + // the as-shot white needs nothing + s.wb.mode = WbMode::AsShot; + assert!(wb_matrix_for(&info, &s).is_none()); +} diff --git a/crates/raw/Cargo.toml b/crates/raw/Cargo.toml index a40fcc43..2ca9ef24 100644 --- a/crates/raw/Cargo.toml +++ b/crates/raw/Cargo.toml @@ -20,9 +20,16 @@ rayon = { workspace = true } miniz_oxide = { workspace = true } zune-jpeg = { workspace = true } zune-core = { workspace = true } +jxl-oxide = { workspace = true, optional = true } + +[features] +default = ["jxl"] +## JPEG XL compressed DNGs (DNG 1.7, compression 52546) via jxl-oxide (pure Rust, MIT OR Apache-2.0). +jxl = ["dep:jxl-oxide"] [dev-dependencies] jpeg-encoder = { workspace = true } proptest = { workspace = true } +zune-jpegxl = { workspace = true } # Pinned CC0 corpus identities; not a product dependency. sha2 = { version = "0.10", default-features = false, features = ["std"] } diff --git a/crates/raw/src/color.rs b/crates/raw/src/color.rs index 050d21ee..d738bb1f 100644 --- a/crates/raw/src/color.rs +++ b/crates/raw/src/color.rs @@ -181,7 +181,11 @@ pub struct CameraTransform { /// Camera transform for white `wb_xy`: `rec2020 = matrix · (wb ⊙ camera)`. pub fn camera_transform(raw: &RawImage, wb_xy: Xy) -> CameraTransform { - let color = &raw.color; + camera_transform_of(&raw.color, wb_xy) +} + +/// [`camera_transform`] from the colour tags alone. +pub fn camera_transform_of(color: &ColorData, wb_xy: Xy) -> CameraTransform { let fallback = !has_matrix(color); let wb = wb_multipliers(color, wb_xy); let to_d50 = camera_to_xyz_d50(color, wb_xy); @@ -192,6 +196,34 @@ pub fn camera_transform(raw: &RawImage, wb_xy: Xy) -> CameraTransform { CameraTransform { matrix, wb: wb.map(|v| v as f32), white_xy: wb_xy, matrix_is_fallback: fallback, baseline_exposure: color.baseline_exposure } } +/// Linear Rec.2020 → linear Rec.2020 matrix that re-develops pixels developed with the camera +/// transform for white `from` as if developed for white `to`: the camera's own colour model is +/// evaluated at the new white (camera-space white balance, DNG spec ch. 6), as Lightroom does, so +/// a neutral lit by `to` comes out neutral. `None` when the tags have no colour matrix or the +/// transform is singular. +pub fn rebalance(color: &ColorData, from: Xy, to: Xy) -> Option { + if !has_matrix(color) { + return None; + } + let (a, b) = (camera_transform_of(color, from), camera_transform_of(color, to)); + let ratio = Mat3::diag(f64::from(b.wb[0] / a.wb[0]), f64::from(b.wb[1] / a.wb[1]), f64::from(b.wb[2] / a.wb[2])); + Some(b.matrix.mul(&ratio).mul(&a.matrix.inverse()?)) +} + +/// The white under which `rgb` (linear Rec.2020, developed with the camera transform for white +/// `developed_for`) is a neutral: the white [`rebalance`] must re-develop for to make it neutral +/// (white-balance picker, auto white balance). `None` without a colour matrix, or when `rgb` maps +/// to a non-positive camera value. +pub fn neutral_white(color: &ColorData, developed_for: Xy, rgb: [f64; 3]) -> Option { + if !has_matrix(color) { + return None; + } + let a = camera_transform_of(color, developed_for); + let balanced = a.matrix.inverse()?.apply(rgb); + let camera: [f64; 3] = std::array::from_fn(|i| balanced[i] / f64::from(a.wb[i])); + camera.iter().all(|v| v.is_finite() && *v > 0.0).then(|| neutral_to_xy(color, camera)) +} + /// `(matrix, multipliers)` with `rec2020 = matrix · (multipliers ⊙ camera)`, linear Rec.2020 D65. pub fn camera_to_rec2020(raw: &RawImage, wb_xy: Xy) -> (Mat3, [f32; 3]) { let t = camera_transform(raw, wb_xy); @@ -373,6 +405,30 @@ mod tests { assert_eq!(illuminant_temperature(255), None); } + /// Camera-space white balance: re-developing for another white turns a neutral lit by that + /// white neutral (the camera's own matrices, interpolated for the new white, say how); the same + /// white is the identity; without a colour matrix there is nothing to re-evaluate. + #[test] + fn rebalance_redevelops_for_the_new_white() { + let color = ColorData { illuminant: [17, 21], color_matrix: [Some(cam_a()), Some(cam_b())], ..Default::default() }; + let (from, to) = (cct::temp_tint_to_xy(5500.0, 0.0), cct::temp_tint_to_xy(3200.0, 10.0)); + let a = camera_transform_of(&color, from); + // a grey card lit by `to`, as the camera records it, developed for `from`: warm + let n = camera_neutral(&color, to); + let developed = a.matrix.apply(std::array::from_fn(|i| n[i] * a.wb[i] as f64)); + assert!(developed[0] > 1.2 * developed[2], "{developed:?}"); + let out = rebalance(&color, from, to).unwrap().apply(developed); + assert!(close(out.map(|v| v / out[1]), [1.0; 3], 1e-6), "{out:?}"); + let same = rebalance(&color, from, from).unwrap(); + assert!((0..3).all(|r| close(same.0[r], Mat3::IDENTITY.0[r], 1e-9)), "{same:?}"); + assert!(rebalance(&ColorData::default(), from, to).is_none()); + // and back: the developed grey card names the white it was lit by + let white = neutral_white(&color, from, developed).unwrap(); + assert!((white.x - to.x).abs() < 1e-6 && (white.y - to.y).abs() < 1e-6, "{white:?} vs {to:?}"); + assert!(neutral_white(&color, from, [-1.0, 0.5, 0.5]).is_none()); + assert!(neutral_white(&ColorData::default(), from, developed).is_none()); + } + #[test] fn grey_world_balances() { let img = lightcraft_raster::Rgb32f::from_fn(20, 20, |x, _| { diff --git a/crates/raw/src/dng.rs b/crates/raw/src/dng.rs index 1e80c0b3..63ddf91a 100644 --- a/crates/raw/src/dng.rs +++ b/crates/raw/src/dng.rs @@ -1,13 +1,14 @@ //! DNG (Adobe Digital Negative Specification 1.7): raw IFD selection, pixel data (via [`crate::tiffraw`]), //! linearization, black/white levels, active area, default crop, CFA description, colour tags, opcode lists. +use crate::gaintable::GainTableMap; use crate::profile::{HsvTable, ProfileLook, ToneCurve}; use crate::tiffraw::{Packing, read_image_in}; use crate::{BlackLevel, Cfa, ColorData, Mat3, Mode, RawData, RawError, RawFormat, RawImage, Rect, Result, opcodes}; use lightcraft_color::Xy; use lightcraft_geom::Orientation; use lightcraft_tiff::tags::{self as t, photometric}; -use lightcraft_tiff::{Ifd, Tiff}; +use lightcraft_tiff::{ByteOrder, Ifd, Tiff}; /// The main raw IFD: full-resolution (NewSubfileType 0) CFA or LinearRaw image with the most pixels. pub(crate) fn raw_ifd(tiff: &Tiff) -> Option<&Ifd> { @@ -33,7 +34,8 @@ fn vec3(v: Option>) -> Option<[f64; 3]> { } /// Colour tags: DNG puts them in IFD0, but some writers use the raw IFD; prefer the raw IFD. -pub(crate) fn color_data(ifd0: &Ifd, raw: &Ifd) -> ColorData { +/// `order` is the file's byte order (for the binary gain-table-map tags). +pub(crate) fn color_data(ifd0: &Ifd, raw: &Ifd, order: ByteOrder) -> ColorData { let get = |tag: u16| raw.f64s(tag).or_else(|| ifd0.f64s(tag)); let geti = |tag: u16| raw.u16(tag).or_else(|| ifd0.u16(tag)); ColorData { @@ -46,13 +48,16 @@ pub(crate) fn color_data(ifd0: &Ifd, raw: &Ifd) -> ColorData { as_shot_white_xy: get(t::AS_SHOT_WHITE_XY).filter(|v| v.len() == 2 && v[0] > 0.0 && v[1] > 0.0).map(|v| Xy::new(v[0], v[1])), baseline_exposure: get(t::BASELINE_EXPOSURE).and_then(|v| v.first().copied()).filter(|v| v.is_finite()).unwrap_or(0.0) + get(t::BASELINE_EXPOSURE_OFFSET).and_then(|v| v.first().copied()).filter(|v| v.is_finite()).unwrap_or(0.0), - profile: profile_look(ifd0, raw), + // DNG 1.7 BaselineSharpness (IFD 0 or the enhanced IFD); implausible values are ignored + baseline_sharpness: get(t::BASELINE_SHARPNESS).and_then(|v| v.first().copied()).filter(|v| v.is_finite() && *v > 0.0 && *v <= 16.0), + profile: profile_look(ifd0, raw, order), } } -/// The profile look tags (`ProfileHueSatMap*`, `ProfileLookTable*`, `ProfileToneCurve`); malformed -/// ones are ignored. Like the colour tags, read from the raw IFD first, else IFD 0. -pub(crate) fn profile_look(ifd0: &Ifd, raw: &Ifd) -> ProfileLook { +/// The profile look tags (`ProfileHueSatMap*`, `ProfileLookTable*`, `ProfileToneCurve`, +/// `ProfileGainTableMap*`); malformed ones are ignored. Like the colour tags, read from the raw IFD +/// first, else IFD 0. +pub(crate) fn profile_look(ifd0: &Ifd, raw: &Ifd, order: ByteOrder) -> ProfileLook { let pick = |tag: u16| if raw.contains(tag) { raw } else { ifd0 }; let table = |dims: u16, data: u16, enc: u16| { let ifd = pick(dims); @@ -67,9 +72,22 @@ pub(crate) fn profile_look(ifd0: &Ifd, raw: &Ifd) -> ProfileLook { ], look_table: table(t::PROFILE_LOOK_TABLE_DIMS, t::PROFILE_LOOK_TABLE_DATA, t::PROFILE_LOOK_TABLE_ENCODING), tone_curve: pick(t::PROFILE_TONE_CURVE).f64s(t::PROFILE_TONE_CURVE).and_then(|v| ToneCurve::from_tag(&v)), + gain_table_map: gain_table_map(ifd0, raw, order), } } +/// The gain table map (DNG 1.7 precedence): `ProfileGainTableMap2` from IFD 0 (the embedded +/// camera profile), else from the raw IFD; else `ProfileGainTableMap` from the raw IFD (where DNG +/// 1.6 put it, and where Apple writes it), else from IFD 0 (where DNG 1.7 says it was meant to be). +/// A malformed tag falls through to the next one. +fn gain_table_map(ifd0: &Ifd, raw: &Ifd, order: ByteOrder) -> Option { + let parse = |ifd: &Ifd, tag: u16| ifd.bytes(tag).and_then(|b| GainTableMap::parse(b, order, tag == t::PROFILE_GAIN_TABLE_MAP_2)); + parse(ifd0, t::PROFILE_GAIN_TABLE_MAP_2) + .or_else(|| parse(raw, t::PROFILE_GAIN_TABLE_MAP_2)) + .or_else(|| parse(raw, t::PROFILE_GAIN_TABLE_MAP)) + .or_else(|| parse(ifd0, t::PROFILE_GAIN_TABLE_MAP)) +} + pub(crate) fn decode(bytes: &[u8], mode: Mode) -> Result { let tiff = Tiff::parse(bytes)?; let ifd0 = &tiff.ifds[0]; @@ -79,9 +97,6 @@ pub(crate) fn decode(bytes: &[u8], mode: Mode) -> Result { if !(1..=4).contains(&cpp) { return Err(RawError::Unsupported(format!("{cpp} samples per pixel"))); } - if info.compression == t::compression::JPEG_XL { - return Err(RawError::Unsupported("JPEG XL DNG".into())); - } let mut data = read_image_in(mode, bytes, &info, tiff.order, Packing::Msb)?; let bits = info.bits() as u32; @@ -175,7 +190,7 @@ pub(crate) fn decode(bytes: &[u8], mode: Mode) -> Result { active_area, crop, orientation: Orientation::from_exif(ifd0.u16(t::ORIENTATION).unwrap_or(1)), - color: color_data(ifd0, raw), + color: color_data(ifd0, raw, tiff.order), wb_multipliers: None, linearized, opcodes, diff --git a/crates/raw/src/dngwrite.rs b/crates/raw/src/dngwrite.rs index bd46a2c5..2f04f016 100644 --- a/crates/raw/src/dngwrite.rs +++ b/crates/raw/src/dngwrite.rs @@ -182,6 +182,9 @@ pub fn write_dng(raw: &RawImage, opts: &DngWriteOptions) -> Result> { ifd.set(t::AS_SHOT_NEUTRAL, rat_vec(&m.map(|v| 1.0 / v as f64))); } ifd.set(t::BASELINE_EXPOSURE, Value::SRational(vec![srational(c.baseline_exposure)])); + if let Some(v) = c.baseline_sharpness { + ifd.set(t::BASELINE_SHARPNESS, rat_vec(&[v])); + } // the source's own profile look travels with its data (conversions, smart previews, merges) let p = &c.profile; let table_tags = |ifd: &mut IfdBuilder, table: &crate::profile::HsvTable, dims: u16, enc: u16| { @@ -214,6 +217,11 @@ pub fn write_dng(raw: &RawImage, opts: &DngWriteOptions) -> Result> { if let Some(curve) = &p.tone_curve { ifd.set(t::PROFILE_TONE_CURVE, Value::Float(curve.points.iter().flatten().copied().collect())); } + if let Some(map) = &p.gain_table_map { + // the raw IFD is IFD 0 here, where both tag versions are valid + let (version2, bytes) = map.to_bytes(opts.order); + ifd.set(if version2 { t::PROFILE_GAIN_TABLE_MAP_2 } else { t::PROFILE_GAIN_TABLE_MAP }, Value::Undefined(bytes)); + } for (list, tag) in [(&raw.opcodes.list1, t::OPCODE_LIST_1), (&raw.opcodes.list2, t::OPCODE_LIST_2), (&raw.opcodes.list3, t::OPCODE_LIST_3)] { if !list.is_empty() { ifd.set(tag, Value::Undefined(opcodes::write_list(list))); diff --git a/crates/raw/src/gaintable.rs b/crates/raw/src/gaintable.rs new file mode 100644 index 00000000..750a1c70 --- /dev/null +++ b/crates/raw/src/gaintable.rs @@ -0,0 +1,325 @@ +//! DNG `ProfileGainTableMap` (tag 52525, DNG 1.6) and `ProfileGainTableMap2` (tag 52544, DNG 1.7): +//! a coarse grid of 1D gain tables carrying the maker's local tone mapping. Apple ProRAW has one. +//! +//! We read and write the tag (a DNG export keeps it) but don't render it: Lightroom Classic +//! renders Apple ProRAW without the map (see [`crate::profile`]), and so do we. +//! +//! - The grid has `MapPointsV × MapPointsH` tables of `MapPointsN` gains. Its origin and spacing +//! are relative to the active area (1.0 = the active area's height or width). +//! - The table input is the dot product of (R, G, B, min, max) with the tag's five weights; +//! version 2 then raises it to the tag's `Gamma`. +//! - Version 2 stores gains as u8 or u16 (mapped linearly onto `GainMin..=GainMax`), f16 or f32. +//! It takes precedence over version 1 when both are present. +//! +//! Malformed tags are ignored: a wrong size, zero or huge dimensions, non-finite or negative gains, +//! non-positive spacing, or gamma outside 0.25..=4. +//! +//! Byte order: the tag's fields are TIFF types (LONG, DOUBLE, FLOAT, …), and neither DNG 1.6 nor +//! 1.7.1 gives this tag a byte order of its own, as they do for opcode lists and +//! `OriginalRawFileData` (always big-endian). What applies is the general rule of the DNG 1.6 +//! specification ("DNG Format Overview" → "Byte Order", worded the same in 1.7.1): +//! +//! > DNG readers are required to support either byte order, even for files from a particular +//! > camera model. +//! +//! (, "DNG Specification".) So we read +//! the fields in the file's byte order, and fall back to big-endian for writers that followed the +//! opcode-list convention. A read in the wrong order can't pass for a valid map: one of +//! `MapPointsV` / `MapPointsH` is below 256 (their product is at most [`MAX_TABLES`]), so swapped +//! it is at least 2^24 and the size checks reject it. We write in the file's byte order. + +use lightcraft_tiff::ByteOrder; +use serde::{Deserialize, Serialize}; + +/// Upper bound on `MapPointsV × MapPointsH` (Apple uses 6 × 8). +const MAX_TABLES: usize = 1 << 12; +/// Upper bound on the total number of gains (Apple: 6 × 8 × 257 = 12 336). +const MAX_GAINS: usize = 1 << 22; + +/// A parsed `ProfileGainTableMap` / `ProfileGainTableMap2`. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct GainTableMap { + /// Tables vertically (`MapPointsV`) and horizontally (`MapPointsH`), and points per table (`MapPointsN`). + pub points_v: usize, + pub points_h: usize, + pub points_n: usize, + /// Grid spacing and origin, relative to the active area's height (`_v`) and width (`_h`). + pub spacing_v: f64, + pub spacing_h: f64, + pub origin_v: f64, + pub origin_h: f64, + /// Weights of (R, G, B, min, max) for the table input. + pub weights: [f32; 5], + /// Exponent for the table input (1.0 for the version-1 tag). + pub gamma: f32, + /// `points_v × points_h × points_n` gains, row-major with the table points innermost. Integer + /// storage has already been mapped onto `GainMin..=GainMax`. Mapping before interpolating is + /// the same as mapping after, because the mapping is linear. + pub gains: Vec, +} + +impl GainTableMap { + /// Parse the tag's bytes from a file in byte order `order` (`version2`: `ProfileGainTableMap2`): + /// in that order, else big-endian (see the module docs). `None` when malformed. + pub fn parse(b: &[u8], order: ByteOrder, version2: bool) -> Option { + Self::parse_in(b, order, version2).or_else(|| if order == ByteOrder::Big { None } else { Self::parse_in(b, ByteOrder::Big, version2) }) + } + + /// Parse the tag's bytes with every multi-byte field in `order`. + fn parse_in(b: &[u8], order: ByteOrder, version2: bool) -> Option { + let take = |o: usize, n: usize| b.get(o..o.checked_add(n)?); + let u32_at = |o: usize| take(o, 4).and_then(|s| s.try_into().ok()).map(|a: [u8; 4]| order.u32(a)); + let f32_at = |o: usize| u32_at(o).map(f32::from_bits); + let f64_at = |o: usize| take(o, 8).and_then(|s| s.try_into().ok()).map(|a: [u8; 8]| f64::from_bits(order.u64(a))); + let (points_v, points_h) = (u32_at(0)? as usize, u32_at(4)? as usize); + let (spacing_v, spacing_h, origin_v, origin_h) = (f64_at(8)?, f64_at(16)?, f64_at(24)?, f64_at(32)?); + let points_n = u32_at(40)? as usize; + let mut weights = [0f32; 5]; + for (i, w) in weights.iter_mut().enumerate() { + *w = f32_at(44 + 4 * i)?; + } + let tables = points_v.checked_mul(points_h).filter(|&t| t > 0 && t <= MAX_TABLES)?; + let count = tables.checked_mul(points_n).filter(|&c| c > 0 && c <= MAX_GAINS)?; + let spacing_ok = |s: f64, points: usize| s.is_finite() && (s > 0.0 || points == 1); + if !spacing_ok(spacing_v, points_v) + || !spacing_ok(spacing_h, points_h) + || !origin_v.is_finite() + || !origin_h.is_finite() + || !weights.iter().all(|w| w.is_finite()) + { + return None; + } + let (header, data_type, gamma, gain_min, gain_max) = + if version2 { (80, u32_at(64)?, f32_at(68)?, f32_at(72)?, f32_at(76)?) } else { (64, 3, 1.0, 0.0, 1.0) }; + if !(0.25..=4.0).contains(&gamma) { + return None; + } + let width = match data_type { + 0 => 1, + 1 | 2 => 2, + 3 => 4, + _ => return None, + }; + if b.len() != count.checked_mul(width)?.checked_add(header)? { + return None; + } + let data = b.get(header..)?; + let scaled = |q: f32, max: f32| gain_min + q / max * (gain_max - gain_min); + let gains: Vec = match data_type { + 0 => data.iter().map(|&q| scaled(q as f32, 255.0)).collect(), + 1 => data.as_chunks::<2>().0.iter().map(|c| scaled(order.u16(*c) as f32, 65535.0)).collect(), + 2 => data.as_chunks::<2>().0.iter().map(|c| crate::unpack::f16_to_f32(order.u16(*c))).collect(), + _ => data.as_chunks::<4>().0.iter().map(|c| f32::from_bits(order.u32(*c))).collect(), + }; + if gains.len() != count || !gains.iter().all(|g| g.is_finite() && *g >= 0.0) { + return None; + } + Some(GainTableMap { points_v, points_h, points_n, spacing_v, spacing_h, origin_v, origin_h, weights, gamma, gains }) + } + + /// The tag's bytes in `order`: version 1 (DNG 1.6, f32 gains) when the gamma is 1, else + /// version 2 with f32 gains. Returns `(version2, bytes)`. + pub fn to_bytes(&self, order: ByteOrder) -> (bool, Vec) { + let version2 = self.gamma != 1.0; + let mut out = Vec::with_capacity(80 + 4 * self.gains.len()); + let u32b = |v: u32| match order { + ByteOrder::Big => v.to_be_bytes(), + ByteOrder::Little => v.to_le_bytes(), + }; + let f64b = |v: f64| match order { + ByteOrder::Big => v.to_be_bytes(), + ByteOrder::Little => v.to_le_bytes(), + }; + let dim = |v: usize| u32b(u32::try_from(v).unwrap_or(u32::MAX)); + out.extend_from_slice(&dim(self.points_v)); + out.extend_from_slice(&dim(self.points_h)); + for v in [self.spacing_v, self.spacing_h, self.origin_v, self.origin_h] { + out.extend_from_slice(&f64b(v)); + } + out.extend_from_slice(&dim(self.points_n)); + for w in self.weights { + out.extend_from_slice(&u32b(w.to_bits())); + } + if version2 { + // DataType 3 (f32); GainMin/GainMax are ignored for float data + for v in [3, self.gamma.to_bits(), 0f32.to_bits(), 1f32.to_bits()] { + out.extend_from_slice(&u32b(v)); + } + } + for g in &self.gains { + out.extend_from_slice(&u32b(g.to_bits())); + } + (version2, out) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Tag bytes: the header fields (`v2`: data type, gamma, gain min, gain max), then `data`. + fn tag( + order: ByteOrder, + dims: [u32; 3], + spacing: [f64; 2], + origin: [f64; 2], + weights: [f32; 5], + v2: Option<(u32, f32, f32, f32)>, + data: &[u8], + ) -> Vec { + let be = order == ByteOrder::Big; + let u = |v: u32| if be { v.to_be_bytes() } else { v.to_le_bytes() }; + let d = |v: f64| if be { v.to_be_bytes() } else { v.to_le_bytes() }; + let mut b = Vec::new(); + b.extend_from_slice(&u(dims[0])); + b.extend_from_slice(&u(dims[1])); + for v in [spacing[0], spacing[1], origin[0], origin[1]] { + b.extend_from_slice(&d(v)); + } + b.extend_from_slice(&u(dims[2])); + for w in weights { + b.extend_from_slice(&u(w.to_bits())); + } + if let Some((dt, gamma, min, max)) = v2 { + for v in [dt, gamma.to_bits(), min.to_bits(), max.to_bits()] { + b.extend_from_slice(&u(v)); + } + } + b.extend_from_slice(data); + b + } + + fn floats(order: ByteOrder, g: &[f32]) -> Vec { + g.iter().flat_map(|v| if order == ByteOrder::Big { v.to_be_bytes() } else { v.to_le_bytes() }).collect() + } + + const GREEN: [f32; 5] = [0.0, 1.0, 0.0, 0.0, 0.0]; + + /// 2 × 2 tables of N = 4 points: table (r, c) = base · [1, 2, 3, 4] with base 1, 2, 3, 4. + fn grid() -> Vec { + (0..4).flat_map(|t| (1..=4).map(move |i| (t + 1) as f32 * i as f32)).collect() + } + + #[test] + fn version_2_storage_types_and_gamma() { + let order = ByteOrder::Little; + let v2 = |dt: u32, gamma: f32, data: &[u8]| { + GainTableMap::parse(&tag(order, [1, 1, 4], [1.0, 1.0], [0.0, 0.0], GREEN, Some((dt, gamma, 0.5, 2.5)), data), order, true) + }; + // u8: 0 → GainMin, 255 → GainMax + let m = v2(0, 1.0, &[0, 51, 255, 255]).unwrap(); + let want = [0.5, 0.9, 2.5, 2.5]; + assert!(m.gains.iter().zip(want).all(|(g, w)| (g - w).abs() < 1e-6), "{:?}", m.gains); + // u16 + let m = v2(1, 1.0, &[0, 0, 0xff, 0xff, 0, 0, 0xff, 0xff]).unwrap(); + assert_eq!(m.gains, vec![0.5, 2.5, 0.5, 2.5]); + // f16 ignores GainMin/GainMax: 1.0, 2.0, 0.5, 0 + let m = v2(2, 1.0, &[0x00, 0x3c, 0x00, 0x40, 0x00, 0x38, 0, 0]).unwrap(); + assert_eq!(m.gains, vec![1.0, 2.0, 0.5, 0.0]); + // f32 with gamma 2 + let m = v2(3, 2.0, &floats(order, &[1.0, 2.0, 3.0, 4.0])).unwrap(); + assert_eq!((m.gains.clone(), m.gamma), (vec![1.0, 2.0, 3.0, 4.0], 2.0)); + // unknown data type, gamma out of range, size for another type: ignored + assert!(v2(4, 1.0, &[0; 4]).is_none()); + assert!(v2(0, 0.1, &[0; 4]).is_none()); + assert!(v2(0, 5.0, &[0; 4]).is_none()); + assert!(v2(1, 1.0, &[0; 4]).is_none()); + // negative integer range + let neg = tag(order, [1, 1, 4], [1.0, 1.0], [0.0, 0.0], GREEN, Some((0, 1.0, -1.0, 1.0)), &[0; 4]); + assert!(GainTableMap::parse(&neg, order, true).is_none()); + } + + #[test] + fn malformed_tags_are_rejected() { + let order = ByteOrder::Big; + let good = tag(order, [2, 2, 4], [0.5, 0.5], [0.25, 0.25], GREEN, None, &floats(order, &grid())); + assert!(GainTableMap::parse(&good, order, false).is_some()); + // every truncation (and an extra byte) + for n in 0..good.len() { + assert!(GainTableMap::parse(&good[..n], order, false).is_none(), "{n} bytes"); + } + let mut long = good.clone(); + long.push(0); + assert!(GainTableMap::parse(&long, order, false).is_none()); + // the wrong byte order reads absurd dimensions + assert!(GainTableMap::parse_in(&good, ByteOrder::Little, false).is_none()); + // read as version 2 the header is too short for the data + assert!(GainTableMap::parse(&good, order, true).is_none()); + let bad = |dims: [u32; 3], spacing: [f64; 2], origin: [f64; 2], gains: &[f32]| { + GainTableMap::parse(&tag(order, dims, spacing, origin, GREEN, None, &floats(order, gains)), order, false) + }; + let mut g = grid(); + g[5] = f32::NAN; + assert!(bad([2, 2, 4], [0.5, 0.5], [0.0, 0.0], &g).is_none(), "NaN gain"); + g[5] = -1.0; + assert!(bad([2, 2, 4], [0.5, 0.5], [0.0, 0.0], &g).is_none(), "negative gain"); + g[5] = f32::INFINITY; + assert!(bad([2, 2, 4], [0.5, 0.5], [0.0, 0.0], &g).is_none(), "infinite gain"); + assert!(bad([2, 2, 4], [0.0, 0.5], [0.0, 0.0], &grid()).is_none(), "zero spacing"); + assert!(bad([2, 2, 4], [0.5, f64::NAN], [0.0, 0.0], &grid()).is_none(), "NaN spacing"); + assert!(bad([2, 2, 4], [0.5, 0.5], [f64::INFINITY, 0.0], &grid()).is_none(), "infinite origin"); + assert!(bad([0, 2, 4], [0.5, 0.5], [0.0, 0.0], &[]).is_none(), "no rows"); + assert!(bad([1, 1, 0], [0.5, 0.5], [0.0, 0.0], &[]).is_none(), "no points"); + // huge dimensions (whose product overflows or exceeds the cap) are rejected before allocating + assert!(bad([u32::MAX, u32::MAX, u32::MAX], [0.5, 0.5], [0.0, 0.0], &grid()).is_none()); + assert!(bad([4096, 4096, 1], [0.5, 0.5], [0.0, 0.0], &grid()).is_none()); + assert!(bad([1, 1, u32::MAX], [0.5, 0.5], [0.0, 0.0], &grid()).is_none()); + // a single table needs no spacing + assert!(bad([1, 1, 2], [0.0, 0.0], [0.0, 0.0], &[1.0, 2.0]).is_some()); + let mut w = GREEN; + w[3] = f32::NAN; + assert!(GainTableMap::parse(&tag(order, [1, 1, 1], [1.0, 1.0], [0.0, 0.0], w, None, &floats(order, &[1.0])), order, false).is_none()); + } + + /// The payload is read in the file's byte order, else big-endian; a little-endian payload in + /// a big-endian file is malformed. + #[test] + fn byte_order_is_the_files_with_a_big_endian_fallback() { + use ByteOrder::{Big, Little}; + let v1 = |o: ByteOrder| tag(o, [2, 2, 4], [0.5, 0.5], [0.25, 0.25], GREEN, None, &floats(o, &grid())); + let v2 = |o: ByteOrder| { + let mut gains = Vec::new(); + for q in [0u16, 0x1234, 0xffff, 0x8000] { + gains.extend_from_slice(&if o == Big { q.to_be_bytes() } else { q.to_le_bytes() }); + } + tag(o, [1, 1, 4], [1.0, 1.0], [0.0, 0.0], GREEN, Some((1, 2.0, 0.5, 2.5)), &gains) + }; + let (want1, want2) = (GainTableMap::parse_in(&v1(Big), Big, false).unwrap(), GainTableMap::parse_in(&v2(Big), Big, true).unwrap()); + assert_eq!(want1.gains, grid()); + assert_eq!((want2.gains[0], want2.gains[2], want2.gamma), (0.5, 2.5, 2.0)); + // in the file's order + assert_eq!(GainTableMap::parse(&v1(Little), Little, false).as_ref(), Some(&want1)); + assert_eq!(GainTableMap::parse(&v2(Little), Little, true).as_ref(), Some(&want2)); + assert_eq!(GainTableMap::parse(&v1(Big), Big, false).as_ref(), Some(&want1)); + // a big-endian payload in a little-endian file: the fallback + assert_eq!(GainTableMap::parse(&v1(Big), Little, false).as_ref(), Some(&want1)); + assert_eq!(GainTableMap::parse(&v2(Big), Little, true).as_ref(), Some(&want2)); + // no little-endian fallback in a big-endian file + assert!(GainTableMap::parse(&v1(Little), Big, false).is_none()); + assert!(GainTableMap::parse(&v2(Little), Big, true).is_none()); + // whatever the grid, a read in the wrong order fails the size checks + for dims in [[1, 1, 1], [6, 8, 257], [48, 64, 2], [255, 16, 1], [256, 16, 1], [1, 4096, 1]] { + for (o, wrong) in [(Big, Little), (Little, Big)] { + let g = vec![1.0; dims.iter().product::() as usize]; + let b = tag(o, dims, [0.5, 0.5], [0.0, 0.0], GREEN, None, &floats(o, &g)); + assert!(GainTableMap::parse_in(&b, o, false).is_some(), "{dims:?}"); + assert!(GainTableMap::parse_in(&b, wrong, false).is_none(), "{dims:?} read {wrong:?}"); + } + } + } + + #[test] + fn writes_what_it_reads() { + for order in [ByteOrder::Big, ByteOrder::Little] { + let m = + GainTableMap::parse(&tag(order, [2, 2, 4], [0.5, 0.5], [0.25, 0.25], GREEN, None, &floats(order, &grid())), order, false).unwrap(); + let (v2, b) = m.to_bytes(order); + assert!(!v2); + assert_eq!(GainTableMap::parse(&b, order, false).unwrap(), m); + let g = GainTableMap { gamma: 2.0, ..m }; + let (v2, b) = g.to_bytes(order); + assert!(v2); + assert_eq!(GainTableMap::parse(&b, order, true).unwrap(), g); + } + } +} diff --git a/crates/raw/src/jxl.rs b/crates/raw/src/jxl.rs new file mode 100644 index 00000000..9e5bfc5f --- /dev/null +++ b/crates/raw/src/jxl.rs @@ -0,0 +1,121 @@ +//! JPEG XL tiles and strips (DNG 1.7 compression 52546), decoded with `jxl-oxide` (pure Rust). +//! +//! Per the DNG 1.7 specification a JPEG XL chunk holds N-bit unsigned integers (8 ≤ N ≤ 16) or 16-bit +//! floats, in 1 or 3 planes, as either a bare codestream or an ISO-BMFF container; jxl-oxide reads both. +//! +//! jxl-oxide renders every channel as `f32`: integer samples come out as `value / (2^bits − 1)` with +//! `bits` the codestream's own declared bit depth, so scaling back by that maximum and rounding gives the +//! stored integers exactly (an `f32` holds `v / 65535` to far better than half a code value) — lossless +//! tiles decode bit-exactly, whatever bit depth the encoder declared relative to the IFD's +//! `BitsPerSample`. Float samples are returned as decoded. +//! +//! We never request an output colour encoding, so jxl-oxide renders in the codestream's *original* +//! colour encoding: for lossless (non-XYB) data that is the identity, keeping the values in the raw IFD's +//! sample space. Lossy (XYB) tiles are converted from XYB back to that signalled encoding by jxl-oxide; +//! [INFERENCE] DNG writers signal a linear encoding matching the raw data there, so the result is the raw +//! sample space up to the codec's loss — not verified against a real lossy JXL DNG. + +use crate::tiffraw::ChunkPx; +use crate::{RawError, Result}; +use jxl_oxide::frame::Encoding; +use jxl_oxide::image::BitDepth; +use jxl_oxide::{AllocTracker, JxlImage, JxlThreadPool}; + +/// Decoder memory (bytes) allowed per sample of a tile. jxl-oxide 0.12 peaks at about three 32-bit +/// buffers per sample: measured ≤ 10.3 bytes for lossless (modular) tiles and ≤ 13.5 bytes per +/// colour sample for lossy (VarDCT) ones, photon noise included. +const BYTES_PER_SAMPLE: usize = 16; +/// Fixed decoder overhead on top (headers, entropy-coding tables, group metadata): under 10 KiB +/// on a 16 × 16 tile. +const MARGIN_BYTES: usize = 64 << 10; + +fn corrupt(e: impl std::fmt::Display) -> RawError { + RawError::Corrupt(format!("JPEG XL tile: {e}")) +} + +/// Parse a chunk's headers (no pixel decoding) and check them against the chunk it must fill: +/// nominal `cw × ch`, or the part of it inside the image `vw × vh` (edge tiles), `cpp` colour planes, +/// a sample type matching the IFD. Returns the decoder and the codestream's integer maximum (0 for float). +#[allow(clippy::too_many_arguments)] +fn open(src: &[u8], cw: usize, ch: usize, vw: usize, vh: usize, cpp: usize, bits: u32, float: bool) -> Result<(JxlImage, f32)> { + if !matches!(cpp, 1 | 3) { + return Err(RawError::Unsupported(format!("JPEG XL with {cpp} samples per pixel"))); + } + if float && bits != 16 || !float && !(8..=16).contains(&bits) { + return Err(RawError::Unsupported(format!("JPEG XL with {bits}-bit {} samples", if float { "float" } else { "integer" }))); + } + // The decoder's memory is bounded by the tile itself: a codestream that needs more (extra + // channels, layers, …) fails to decode instead of allocating what its header asks for. + let pixels = cw.saturating_mul(ch); + let budget = |planes: usize| pixels.saturating_mul(planes).saturating_mul(BYTES_PER_SAMPLE); + let tracker = AllocTracker::with_limit(budget(cpp).saturating_add(MARGIN_BYTES)); + let img = JxlImage::builder() + .pool(JxlThreadPool::none()) // chunks are already decoded in parallel + .alloc_tracker(tracker.clone()) + .read(std::io::Cursor::new(src)) + .map_err(corrupt)?; + // lossy (VarDCT) decoding always works on three colour channels, even for a grey tile + if cpp == 1 && img.frame_header(0).is_some_and(|f| f.encoding == Encoding::VarDct) { + tracker.expand_limit(budget(2)); + } + let header = img.image_header(); + let (jw, jh) = (img.width() as usize, img.height() as usize); + if (jw, jh) != (cw, ch) && (jw, jh) != (vw, vh) { + return Err(RawError::Corrupt(format!("JPEG XL tile is {jw}×{jh}, expected {cw}×{ch} (or {vw}×{vh} at the image edge)"))); + } + if header.metadata.orientation != 1 { + return Err(RawError::Corrupt("JPEG XL tile with a non-identity orientation".into())); + } + let planes = if header.metadata.grayscale() { 1 } else { 3 }; + if planes != cpp { + return Err(RawError::Corrupt(format!("JPEG XL tile has {planes} colour planes, the IFD {cpp} samples per pixel"))); + } + let max = match header.metadata.bit_depth { + BitDepth::IntegerSample { bits_per_sample } if !float && (1..=16).contains(&bits_per_sample) => ((1u32 << bits_per_sample) - 1) as f32, + BitDepth::FloatSample { .. } if float => 0.0, + d => return Err(RawError::Corrupt(format!("JPEG XL tile sample type {d:?} does not match the IFD"))), + }; + Ok((img, max)) +} + +/// Check the first chunk's headers ([`open`]) without decoding pixels (header-only probing). +#[allow(clippy::too_many_arguments)] +pub(crate) fn check(src: &[u8], cw: usize, ch: usize, vw: usize, vh: usize, cpp: usize, bits: u32, float: bool) -> Result<()> { + open(src, cw, ch, vw, vh, cpp, bits, float).map(|_| ()) +} + +/// Decode one chunk into `cw × ch × cpp` samples (row-major, interleaved); a chunk coded at its +/// in-image size `vw × vh` fills the top-left of that, the rest stays 0. +#[allow(clippy::too_many_arguments)] +pub(crate) fn decode(src: &[u8], cw: usize, ch: usize, vw: usize, vh: usize, cpp: usize, bits: u32, float: bool) -> Result { + let (img, max) = open(src, cw, ch, vw, vh, cpp, bits, float)?; + if img.num_loaded_keyframes() == 0 { + return Err(corrupt("no complete frame")); + } + let render = img.render_frame(0).map_err(corrupt)?; + let fb = render.image_all_channels(); + let (jw, jh, fc) = (fb.width(), fb.height(), fb.channels()); + if jw == 0 || jh == 0 || jw > cw || jh > ch || fc < cpp { + return Err(corrupt(format!("rendered {jw}×{jh}×{fc}, expected {cw}×{ch}×{cpp}"))); + } + let n = cw.saturating_mul(ch).saturating_mul(cpp); + let rows = fb.buf().chunks_exact(jw * fc); + Ok(if float { + let mut out = vec![0f32; n]; + for (dst, src) in out.chunks_exact_mut(cw * cpp).zip(rows) { + for (d, s) in dst.chunks_exact_mut(cpp).zip(src.chunks_exact(fc)) { + d.iter_mut().zip(s).for_each(|(d, s)| *d = *s); + } + } + ChunkPx::F32(out) + } else { + let mut out = vec![0u16; n]; + for (dst, src) in out.chunks_exact_mut(cw * cpp).zip(rows) { + for (d, s) in dst.chunks_exact_mut(cpp).zip(src.chunks_exact(fc)) { + // `as` saturates (NaN → 0): lossy overshoot clips to the integer range + d.iter_mut().zip(s).for_each(|(d, s)| *d = (s * max).round() as u16); + } + } + ChunkPx::U16(out) + }) +} diff --git a/crates/raw/src/lib.rs b/crates/raw/src/lib.rs index d63e5254..d38b2fc6 100644 --- a/crates/raw/src/lib.rs +++ b/crates/raw/src/lib.rs @@ -2,16 +2,18 @@ //! //! - [`probe`] recognises raw containers; [`decode`] turns a file into a [`RawImage`] (sensor data + everything //! needed to render it: CFA, black/white levels, active area, default crop, orientation, DNG colour tags, -//! opcode lists, [`Metadata`]); [`embedded_preview`] returns the largest embedded JPEG. +//! opcode lists, [`Metadata`]); [`embedded_preview`] returns the largest embedded JPEG (or DNG 1.7 JPEG XL) preview; +//! [`semantic_masks`] reads a DNG's semantic masks (segmentation mattes, e.g. iPhone ProRAW's sky matte). //! - [`RawImage::normalized`] subtracts black, scales white to 1.0 and crops to the active area (applying DNG //! `OpcodeList1`/`OpcodeList2`); [`demosaic`] turns CFA data into camera-RGB [`Rgb32f`]; //! [`RawImage::develop`] does all of it plus `OpcodeList3` and the default crop; [`RawImage::develop_binned`] //! produces the same at 1/k of the size straight from the mosaic (previews, thumbnails). //! - [`color`] implements the DNG colour model (dual-illuminant interpolation, forward matrices, white balance) //! and produces camera → linear Rec.2020 D65 matrices; [`profile`] reads and applies a DNG's own profile -//! look tables and tone curve. +//! look tables and tone curve, and [`gaintable`] its gain table map (Apple ProRAW's local tone mapping). //! -//! Formats: DNG (uncompressed, lossless JPEG, lossy JPEG (Smart Previews), Deflate incl. floating point, tiled/stripped, CFA and LinearRaw), +//! Formats: DNG (uncompressed, lossless JPEG, lossy JPEG (Smart Previews), Deflate incl. floating point, JPEG XL (DNG 1.7, +//! `jxl` feature, on by default), tiled/stripped, CFA and LinearRaw), //! Canon CR2 / CR3 (lossless CRX Bayer and version 0x100/0x200 C-RAW), Nikon NEF/NRW (uncompressed, Huffman lossless / lossy compressed), Sony ARW (uncompressed, ARW2, lossless), Fujifilm RAF (uncompressed Bayer //! and X-Trans, lossless and lossy compressed), Panasonic RW2 / Leica RWL / Panasonic RAW (every raw format: compressed 4 and 6, the prefix-coded strips of 8, //! packed 2/5/7, the 16-bit words of the oldest bodies), Pentax PEF (uncompressed, Huffman), Olympus ORF (uncompressed). @@ -27,11 +29,15 @@ pub mod color; pub mod demosaic; mod dng; pub mod dngwrite; +pub mod gaintable; pub mod highlight; +#[cfg(feature = "jxl")] +mod jxl; pub mod ljpeg; pub mod opcodes; mod preview; pub mod profile; +pub mod semantic; mod tiffraw; mod unpack; mod vendor; @@ -44,6 +50,7 @@ pub use lightcraft_meta::Metadata; pub use lightcraft_raster::Rgb32f; pub use opcodes::{Opcode, OpcodeLists}; pub use preview::{PreviewColorSpace, embedded_preview, embedded_preview_color_space}; +pub use semantic::{SemanticMask, semantic_masks}; use lightcraft_color::Xy; use lightcraft_tiff::{Tiff, TiffError}; @@ -509,8 +516,13 @@ pub struct ColorData { pub as_shot_white_xy: Option, /// EV to add for a "normal" rendering (`BaselineExposure` + `BaselineExposureOffset`). pub baseline_exposure: f64, + /// `BaselineSharpness`: sharpening relative to a reference camera (`None`: the DNG default, + /// 1). Read and written back by DNG export; rendering doesn't use it yet. + #[serde(default)] + pub baseline_sharpness: Option, /// The file's own camera-profile look (`ProfileHueSatMap*`, `ProfileLookTable*`, - /// `ProfileToneCurve`), applied by [`color`]'s users at render time. + /// `ProfileToneCurve`), applied by [`color`]'s users at render time, and its + /// `ProfileGainTableMap*`, kept (a DNG export writes it back) but not rendered. #[serde(default)] pub profile: profile::ProfileLook, } @@ -677,14 +689,27 @@ impl RawImage { /// The result is camera RGB (not white balanced), white level = 1.0, not oriented. pub fn develop(&self, method: Method) -> Result { let n = self.normalized()?; - let mut rgb = demosaic(&n, method); + let rgb = demosaic(&n, method); drop(n); + Ok(self.finish_demosaiced(rgb)) + } + + /// The last steps of [`Self::develop`] for camera RGB made some other way (a denoiser that demosaics too) from + /// [`Self::normalized`]'s mosaic: apply `OpcodeList3`, then the default crop. + pub fn finish_demosaiced(&self, mut rgb: Rgb32f) -> Rgb32f { opcodes::apply_list3(&self.opcodes.list3, &mut rgb); - let c = self.crop.clipped(rgb.width, rgb.height); - if c.width == 0 || c.height == 0 || (c.x == 0 && c.y == 0 && c.width == rgb.width && c.height == rgb.height) { - return Ok(rgb); + let c = self.develop_crop(rgb.width, rgb.height); + if c.x == 0 && c.y == 0 && c.width == rgb.width && c.height == rgb.height { + return rgb; } - Ok(rgb.into_crop(c.x, c.y, c.width, c.height)) + rgb.into_crop(c.x, c.y, c.width, c.height) + } + + /// The part of a `width × height` demosaiced picture (the active area) that [`Self::develop`] keeps: the default + /// crop, or all of it when there is none. + pub fn develop_crop(&self, width: usize, height: usize) -> Rect { + let c = self.crop.clipped(width, height); + if c.width == 0 || c.height == 0 { Rect::new(0, 0, width, height) } else { c } } } @@ -717,6 +742,43 @@ mod tests { } } + #[test] + fn a_picture_made_elsewhere_is_finished_like_develop_does() { + let (w, h) = (12usize, 10usize); + let raw = RawImage { + format: RawFormat::Dng, + width: w, + height: h, + cpp: 1, + data: RawData::U16((0..w * h).map(|i| (100 + (i * 37) % 900) as u16).collect()), + cfa: Cfa::bayer("RGGB"), + bits: 12, + black: BlackLevel::uniform(64.0), + white: vec![1023.0], + active_area: Rect::new(2, 2, 8, 6), + crop: Rect::new(1, 1, 5, 4), + orientation: Orientation::Normal, + color: ColorData::default(), + wb_multipliers: None, + linearized: false, + opcodes: OpcodeLists::default(), + metadata: Metadata::default(), + }; + let direct = raw.develop(Method::Bilinear).unwrap(); + let by_hand = raw.finish_demosaiced(demosaic(&raw.normalized().unwrap(), Method::Bilinear)); + assert_eq!((direct.width, direct.height), (5, 4)); + assert_eq!(direct, by_hand); + // the crop is what develop keeps; no usable crop keeps everything + assert_eq!(raw.develop_crop(8, 6), Rect::new(1, 1, 5, 4)); + let mut whole = raw.clone(); + whole.crop = Rect::default(); + assert_eq!(whole.develop_crop(8, 6), Rect::new(0, 0, 8, 6)); + assert_eq!(whole.develop(Method::Bilinear).unwrap().width, 8); + // a crop that pokes out is clipped to the picture + whole.crop = Rect::new(6, 4, 10, 10); + assert_eq!(whole.develop_crop(8, 6), Rect::new(6, 4, 2, 2)); + } + #[test] fn black_levels() { let b = BlackLevel { repeat_rows: 2, repeat_cols: 2, values: vec![1.0, 2.0, 3.0, 4.0], delta_h: vec![0.5, 0.0], delta_v: vec![0.0, 10.0] }; diff --git a/crates/raw/src/preview.rs b/crates/raw/src/preview.rs index 6ec326ea..371b5e4a 100644 --- a/crates/raw/src/preview.rs +++ b/crates/raw/src/preview.rs @@ -1,5 +1,7 @@ //! Embedded preview extraction: the largest baseline/progressive JPEG stored in a TIFF-based raw (IFD strips -//! with JPEG compression, `JPEGInterchangeFormat` pointers in any IFD, Nikon/others' maker-note preview IFDs). +//! with JPEG compression, `JPEGInterchangeFormat` pointers in any IFD, Nikon/others' maker-note preview IFDs), +//! or a DNG 1.7 JPEG XL preview IFD stored as a single tile/strip (a standalone `.jxl` file, as the spec +//! recommends for previews). use lightcraft_tiff::image::chunk_bytes; use lightcraft_tiff::tags as t; @@ -56,6 +58,11 @@ fn is_dct_jpeg(b: &[u8]) -> bool { false } +/// Whether `b` is a JPEG XL file: a bare codestream or the ISO-BMFF container's signature box. +fn is_jxl(b: &[u8]) -> bool { + b.starts_with(&[0xff, 0x0a]) || b.starts_with(&[0, 0, 0, 0x0c, b'J', b'X', b'L', b' ', 0x0d, 0x0a, 0x87, 0x0a]) +} + fn candidates<'a>(data: &'a [u8], ifd: &Ifd, base: u64, out: &mut Vec<&'a [u8]>) { // whole JPEG files stored as an undefined-type tag value (e.g. Panasonic `JpgFromRaw` 0x002e) for e in &ifd.entries { @@ -83,9 +90,20 @@ fn candidates<'a>(data: &'a [u8], ifd: &Ifd, base: u64, out: &mut Vec<&'a [u8]>) out.push(s); } } + // a rendered (RGB or grey, never CFA/LinearRaw/mask) JPEG XL preview in one chunk + if ifd.u16(t::COMPRESSION) == Some(t::compression::JPEG_XL) + && matches!(ifd.u16(t::PHOTOMETRIC), Some(t::photometric::BLACK_IS_ZERO | t::photometric::RGB)) + && let Ok(info) = ifd.image() + && let [chunk] = info.chunks(data.len() as u64).as_slice() + && let Some(s) = chunk_bytes(data, chunk) + && is_jxl(s) + { + out.push(s); + } } -/// The largest embedded JPEG preview, if any. +/// The largest embedded preview, if any: a JPEG, or (DNG 1.7) a JPEG XL file — both decode with +/// `lightcraft_codecs::decode`. pub fn embedded_preview(bytes: &[u8]) -> Option> { if bytes.starts_with(b"FUJIFILMCCD-RAW") { let j = crate::vendor::raf::header(bytes).ok()?.jpeg?; @@ -121,7 +139,11 @@ pub fn embedded_preview(bytes: &[u8]) -> Option> { if let Some(p) = crate::vendor::orf::preview(bytes) { found.push(p); } - let best = found.into_iter().filter(|s| is_dct_jpeg(s)).max_by_key(|s| s.len()).map(|s| trim_eoi(s).to_vec()); + let best = found + .into_iter() + .filter(|s| is_dct_jpeg(s) || is_jxl(s)) + .max_by_key(|s| s.len()) + .map(|s| if is_jxl(s) { s.to_vec() } else { trim_eoi(s).to_vec() }); // a TIFF raw we can't decode whose preview no tag points to (Leaf MOS, Epson ERF) best.or_else(|| format.filter(|f| !f.is_supported()).and_then(|_| scan_for_jpeg(bytes))) } diff --git a/crates/raw/src/profile.rs b/crates/raw/src/profile.rs index 8b9761ab..a21444c4 100644 --- a/crates/raw/src/profile.rs +++ b/crates/raw/src/profile.rs @@ -1,5 +1,6 @@ //! DNG camera-profile look data carried by the file itself (Adobe DNG Specification 1.6/1.7, -//! chapter 6 and the `ProfileHueSatMap*`, `ProfileLookTable*` and `ProfileToneCurve` tags). +//! chapter 6 and the `ProfileHueSatMap*`, `ProfileLookTable*`, `ProfileToneCurve` and +//! `ProfileGainTableMap*` tags). //! //! - A hue/saturation/value table holds, for a grid of (value, hue, saturation) inputs, a hue shift //! in degrees, a saturation scale and a value scale (value-major, hue-middle, saturation-minor). @@ -10,11 +11,17 @@ //! - `ProfileHueSatMapData1/2` follow the colour matrices (interpolated by the same illuminant //! weight); `ProfileLookTableData` comes after exposure compensation and before any tone curve. //! - `ProfileToneCurve`: (input, output) pairs in linear gamma from (0, 0) to (1, 1). +//! - `ProfileGainTableMap*` ([`crate::gaintable`]) is read and kept (a DNG export writes it back) +//! but not rendered: Lightroom Classic renders Apple ProRAW without it (measured on the iPhone +//! 12 Pro corpus file against Lightroom Classic 15.6: with the map the default render is far +//! brighter than Lightroom's, mean ΔE00 19; without it Lightroom's output is a tight function of +//! the hue/saturation map, exposure compensation and tone curve alone). //! //! The data is read from the user's own file at run time; nothing here ships profile data. //! Values above 1.0 (highlight headroom) keep their headroom: the table's value scale only ever //! lowers them, so a lookup can't clip what the scene-referred pipeline still needs. +use crate::gaintable::GainTableMap; use lightcraft_color::transfer::{linear_to_srgb, srgb_to_linear}; use lightcraft_color::{D50, D65, Mat3, PROPHOTO, REC2020, bradford}; use serde::{Deserialize, Serialize}; @@ -207,11 +214,14 @@ pub struct ProfileLook { pub look_table: Option, /// `ProfileToneCurve`. pub tone_curve: Option, + /// `ProfileGainTableMap2`, else `ProfileGainTableMap`. + #[serde(default)] + pub gain_table_map: Option, } impl ProfileLook { pub fn is_empty(&self) -> bool { - self.hue_sat_map.iter().all(Option::is_none) && self.look_table.is_none() && self.tone_curve.is_none() + self.hue_sat_map.iter().all(Option::is_none) && self.look_table.is_none() && self.tone_curve.is_none() && self.gain_table_map.is_none() } /// The hue/saturation map for calibration-1 weight `g` (one table: used for any white). @@ -360,7 +370,7 @@ mod tests { let m = HsvTable::blend(&a, &b, 0.25).unwrap(); assert_eq!(m.data[1], [7.5, 1.75, 1.0]); assert!(HsvTable::blend(&a, &identity(3, 2, 1), 0.5).is_none()); - let look = ProfileLook { hue_sat_map: [Some(a.clone()), None], look_table: Some(a), tone_curve: None }; + let look = ProfileLook { hue_sat_map: [Some(a.clone()), None], look_table: Some(a), tone_curve: None, gain_table_map: None }; let t = ProfileTables::new(&look, 0.5).unwrap(); let c = [0.3, 0.2, 0.1]; assert!(close(t.apply(c, 2.0), [0.6, 0.4, 0.2], 1e-4), "{:?}", t.apply(c, 2.0)); diff --git a/crates/raw/src/semantic.rs b/crates/raw/src/semantic.rs new file mode 100644 index 00000000..e9af87c0 --- /dev/null +++ b/crates/raw/src/semantic.rs @@ -0,0 +1,136 @@ +//! DNG semantic masks (DNG 1.6+, chapter "Semantic Masks"): single-channel segmentation mattes +//! stored in their own IFDs (NewSubFileType 10004.H, PhotometricInterpretation 52527) with a +//! `SemanticName`, an optional `SemanticInstanceID` and an optional `MaskSubArea`. iPhone ProRAW +//! stores e.g. a sky matte this way. +//! +//! [`semantic_masks`] reads them separately from [`crate::decode`] (only callers that use them +//! pay for them); malformed or unsupported mask IFDs are skipped. + +use crate::Rect; +use crate::tiffraw::{Packing, read_image}; +use crate::{RawData, Result}; +use lightcraft_raster::Image; +use lightcraft_tiff::Tiff; +use lightcraft_tiff::tags::{self as t, photometric}; + +/// At most this many mask IFDs are read from one file. +pub const MAX_MASKS: usize = 16; +/// Masks with more samples than this are skipped (a mask is a low-resolution matte; this is +/// 8192 × 4096). +pub const MAX_MASK_SAMPLES: usize = 1 << 25; + +/// Where a cropped mask sits in the uncropped mask (DNG `MaskSubArea`). The uncropped mask +/// (`full_width × full_height`) spans the main image's active area. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct MaskSubArea { + pub top: u32, + pub left: u32, + pub full_width: u32, + pub full_height: u32, +} + +/// One semantic mask: `width × height` values, 0 = not part of the attribute, 65535 = fully part. +#[derive(Clone, Debug, PartialEq)] +pub struct SemanticMask { + /// `SemanticName`, e.g. `urn:com:apple:photo:2020:aux:semanticskymatte`. + pub name: String, + pub instance_id: Option, + pub width: usize, + pub height: usize, + /// `None`: the mask covers the whole active area. + pub sub_area: Option, + pub data: Vec, +} + +/// The semantic masks of a DNG, in file order (none for other files). Only the first +/// [`MAX_MASKS`] mask IFDs are looked at, malformed ones included, so a file full of mask IFDs +/// costs at most that many decodes. +pub fn semantic_masks(bytes: &[u8]) -> Vec { + let Ok(tiff) = Tiff::parse(bytes) else { return Vec::new() }; + tiff.all_ifds() + .into_iter() + .filter(|i| i.u16(t::PHOTOMETRIC) == Some(photometric::MASK) || i.u32(t::NEW_SUBFILE_TYPE) == Some(t::SUBFILE_SEMANTIC_MASK)) + .take(MAX_MASKS) + .filter_map(|ifd| read_mask(bytes, &tiff, ifd).ok().flatten()) + .collect() +} + +fn read_mask(bytes: &[u8], tiff: &Tiff, ifd: &lightcraft_tiff::Ifd) -> Result> { + let Some(name) = ifd.string(t::SEMANTIC_NAME).filter(|n| !n.is_empty()) else { return Ok(None) }; + let info = ifd.image()?; + let (w, h) = (info.width as usize, info.height as usize); + if info.samples_per_pixel != 1 || w.checked_mul(h).is_none_or(|n| n > MAX_MASK_SAMPLES) { + return Ok(None); + } + let bits = info.bits() as u32; + let data: Vec = match read_image(bytes, &info, tiff.order, Packing::Msb)? { + RawData::U16(v) if (1..=16).contains(&bits) => { + let max = ((1u32 << bits) - 1) as f32; + if bits == 16 { v } else { v.into_iter().map(|s| (s as f32 / max * 65535.0).round().min(65535.0) as u16).collect() } + } + RawData::F32(v) => v.into_iter().map(|s| if s.is_finite() { (s.clamp(0.0, 1.0) * 65535.0).round() as u16 } else { 0 }).collect(), + RawData::U16(_) => return Ok(None), + }; + if data.len() != w * h { + return Ok(None); + } + // the tag is ignored when the stored mask does not fit inside the uncropped one + let sub_area = match ifd.u64s(t::MASK_SUB_AREA).as_deref() { + Some(&[top, left, fw, fh]) => { + let fits = fw > 0 && fh > 0 && left.saturating_add(w as u64) <= fw && top.saturating_add(h as u64) <= fh && fw.max(fh) <= u32::MAX as u64; + fits.then_some(MaskSubArea { top: top as u32, left: left as u32, full_width: fw as u32, full_height: fh as u32 }) + } + _ => None, + }; + Ok(Some(SemanticMask { name, instance_id: ifd.string(t::SEMANTIC_INSTANCE_ID).filter(|s| !s.is_empty()), width: w, height: h, sub_area, data })) +} + +impl SemanticMask { + /// The mask value (0..1) at a continuous position of the uncropped mask (pixel centres at + /// +0.5); 0 outside the stored (cropped) mask. + pub fn value_at(&self, x: f64, y: f64) -> f32 { + let (left, top) = self.sub_area.map_or((0.0, 0.0), |a| (a.left as f64, a.top as f64)); + let (x, y) = (x - left, y - top); + if self.width == 0 || self.height == 0 || !(x >= 0.0 && y >= 0.0 && x <= self.width as f64 && y <= self.height as f64) { + return 0.0; + } + let (fx, fy) = ((x - 0.5).max(0.0), (y - 0.5).max(0.0)); + let (x0, y0) = (fx.floor() as usize, fy.floor() as usize); + let (tx, ty) = ((fx - x0 as f64) as f32, (fy - y0 as f64) as f32); + let at = |x: usize, y: usize| { + let (x, y) = (x.min(self.width - 1), y.min(self.height - 1)); + self.data.get(y * self.width + x).copied().unwrap_or(0) as f32 / 65535.0 + }; + let top = at(x0, y0) + (at(x0 + 1, y0) - at(x0, y0)) * tx; + let bottom = at(x0, y0 + 1) + (at(x0 + 1, y0 + 1) - at(x0, y0 + 1)) * tx; + top + (bottom - top) * ty + } + + /// The mask over the developed image ([`crate::RawImage::develop`]: the default `crop` of the + /// `active_area`, before orientation), at the mask's own resolution, as 0..255 alpha. `None` + /// for an empty crop. + pub fn developed(&self, active_area: Rect, crop: Rect) -> Option> { + let c = crop.clipped(active_area.width, active_area.height); + if c.width == 0 || c.height == 0 || self.width == 0 || self.height == 0 { + return None; + } + let (fw, fh) = self.sub_area.map_or((self.width, self.height), |a| (a.full_width as usize, a.full_height as usize)); + // uncropped-mask px per active-area px + let (sx, sy) = (fw as f64 / active_area.width as f64, fh as f64 / active_area.height as f64); + let ow = ((c.width as f64 * sx).round() as usize).clamp(1, self.width.max(fw)); + let oh = ((c.height as f64 * sy).round() as usize).clamp(1, self.height.max(fh)); + if ow.checked_mul(oh).is_none_or(|n| n > MAX_MASK_SAMPLES) { + return None; + } + let (kx, ky) = (c.width as f64 / ow as f64, c.height as f64 / oh as f64); + let mut out = Image::new(ow, oh); + lightcraft_raster::par_rows(&mut out.data, ow, |y, row| { + let ay = c.y as f64 + (y as f64 + 0.5) * ky; + for (x, v) in row.iter_mut().enumerate() { + let ax = c.x as f64 + (x as f64 + 0.5) * kx; + *v = (self.value_at(ax * sx, ay * sy) * 255.0).round() as u8; + } + }); + Some(out) + } +} diff --git a/crates/raw/src/tiffraw.rs b/crates/raw/src/tiffraw.rs index 200dd391..c28b8b64 100644 --- a/crates/raw/src/tiffraw.rs +++ b/crates/raw/src/tiffraw.rs @@ -1,4 +1,4 @@ -//! Reading the pixel data of a TIFF IFD (strips or tiles; uncompressed, lossless JPEG, Deflate), in parallel. +//! Reading the pixel data of a TIFF IFD (strips or tiles; uncompressed, lossless/lossy JPEG, Deflate, JPEG XL), in parallel. use crate::unpack::*; use crate::{MAX_SAMPLES, RawData, RawError, Result, ljpeg}; @@ -16,9 +16,12 @@ pub enum Packing { Lsb, /// Each sample stored in 16 bits in file byte order, regardless of `bits`. Word16, + /// MSB-first bit stream inside 32-bit words stored in file byte order ([`unpack_words32_msb`]); rows start on a + /// word boundary. + Words32Msb, } -enum ChunkPx { +pub(crate) enum ChunkPx { U16(Vec), F32(Vec), } @@ -61,11 +64,24 @@ pub(crate) fn read_image_in(mode: crate::Mode, data: &[u8], info: &ImageInfo, or crate::Mode::Header => { check_image(data, info)?; // a lossless-JPEG layout the decoder rejects (subsampled components, e.g. Sony's - // lossless M/S sizes) must fail here too, as the full decode will - if info.compression == 7 - && let Some(src) = info.chunks(data.len() as u64).first().and_then(|c| chunk_bytes(data, c)) + // lossless M/S sizes) must fail here too, as the full decode will; likewise a JPEG XL + // chunk whose headers don't match the IFD + let first = info.chunks(data.len() as u64).first().copied(); + if let Some(c) = first + && let Some(src) = chunk_bytes(data, &c) { - ljpeg::frame_info(src)?; + match info.compression { + 7 => { + ljpeg::frame_info(src)?; + } + comp::JPEG_XL => { + let planar = info.planar == 2 && info.samples_per_pixel > 1; + let cpp = if planar { 1 } else { info.samples_per_pixel as usize }; + let (cw, ch, vw, vh) = chunk_dims(info, &c); + jxl_chunk_check(src, cw, ch, vw, vh, cpp, info.bits() as u32, info.sample_format == 3)?; + } + _ => {} + } } Ok(if info.sample_format == 3 { RawData::F32(Vec::new()) } else { RawData::U16(Vec::new()) }) } @@ -138,7 +154,7 @@ pub fn read_image(data: &[u8], info: &ImageInfo, order: ByteOrder, packing: Pack #[allow(clippy::too_many_arguments)] fn decode_chunk(data: &[u8], info: &ImageInfo, c: &Chunk, order: ByteOrder, packing: Packing, cpp: usize, bits: u32, float: bool) -> Result { let src = chunk_bytes(data, c).ok_or_else(|| RawError::Corrupt("chunk offset past end of file".into()))?; - let (cw, ch) = (c.width as usize, c.height as usize); + let (cw, ch, vw, vh) = chunk_dims(info, c); let n = cw.checked_mul(ch).and_then(|v| v.checked_mul(cpp)).ok_or(RawError::Limit("chunk too large"))?; if n > MAX_SAMPLES { return Err(RawError::Limit("chunk too large")); @@ -176,10 +192,34 @@ fn decode_chunk(data: &[u8], info: &ImageInfo, c: &Chunk, order: ByteOrder, pack let raw = inflate(src, row_bytes * ch)?; unpack_chunk(&raw, order, packing, cw, ch, cpp, bits, float, info.predictor) } + comp::JPEG_XL => jxl_chunk(src, cw, ch, vw, vh, cpp, bits, float), other => Err(RawError::Unsupported(format!("TIFF compression {other}"))), } } +/// A chunk's nominal size and the part of it inside the image (smaller for edge tiles). +fn chunk_dims(info: &ImageInfo, c: &Chunk) -> (usize, usize, usize, usize) { + let (cw, ch) = (c.width as usize, c.height as usize); + let vw = cw.min(info.width.saturating_sub(c.x) as usize); + let vh = ch.min(info.height.saturating_sub(c.y) as usize); + (cw, ch, vw, vh) +} + +#[cfg(feature = "jxl")] +use crate::jxl::{check as jxl_chunk_check, decode as jxl_chunk}; + +#[cfg(not(feature = "jxl"))] +#[allow(clippy::too_many_arguments)] +fn jxl_chunk(_: &[u8], _: usize, _: usize, _: usize, _: usize, _: usize, _: u32, _: bool) -> Result { + Err(RawError::Unsupported("JPEG XL DNG (built without the `jxl` feature)".into())) +} + +#[cfg(not(feature = "jxl"))] +#[allow(clippy::too_many_arguments)] +fn jxl_chunk_check(src: &[u8], cw: usize, ch: usize, vw: usize, vh: usize, cpp: usize, bits: u32, float: bool) -> Result<()> { + jxl_chunk(src, cw, ch, vw, vh, cpp, bits, float).map(|_| ()) +} + /// Decode one baseline JPEG chunk to 8-bit samples: (samples, width, height, channels). fn lossy_jpeg(src: &[u8], cpp: usize) -> Result<(Vec, usize, usize, usize)> { use zune_core::bytestream::ZCursor; @@ -263,6 +303,7 @@ fn unpack_chunk( (16, _) | (_, Packing::Word16) => (row_n * 2, 1), (_, Packing::Msb) => ((row_n * bits as usize).div_ceil(8), 2), (_, Packing::Lsb) => ((row_n * bits as usize).div_ceil(8), 3), + (_, Packing::Words32Msb) => ((row_n * bits as usize).div_ceil(32) * 4, 4), }; for y in 0..ch { let s = src.get(y * row_bytes..).unwrap_or(&[]); @@ -272,6 +313,7 @@ fn unpack_chunk( 0 => row.iter_mut().zip(s).for_each(|(o, &b)| *o = b as u16), 1 => read_u16s(s, order, row), 2 => unpack_msb(s, bits, row), + 4 => unpack_words32_msb(s, order, bits, row), _ => unpack_lsb(s, bits, row), } if factor > 0 { diff --git a/crates/raw/src/unpack.rs b/crates/raw/src/unpack.rs index 15451548..77c69c00 100644 --- a/crates/raw/src/unpack.rs +++ b/crates/raw/src/unpack.rs @@ -40,6 +40,19 @@ pub fn unpack_lsb(src: &[u8], bits: u32, out: &mut [u16]) { } } +/// Unpack MSB-first packed samples of `bits` from a stream of 32-bit words stored in `order` (a word is read as a +/// number first, then its bits are taken from the top): the layout of some cameras' uncompressed 12- and 14-bit +/// data, which is not the byte-wise MSB stream of [`unpack_msb`] unless `order` is big-endian. A short last word +/// is padded with zeros. +pub fn unpack_words32_msb(src: &[u8], order: ByteOrder, bits: u32, out: &mut [u16]) { + let words = src.chunks(4).flat_map(|c| { + let mut w = [0u8; 4]; + w[..c.len()].copy_from_slice(c); + order.u32(w).to_be_bytes() + }); + unpack_msb(&words.collect::>(), bits, out); +} + pub fn read_u16s(src: &[u8], order: ByteOrder, out: &mut [u16]) { for (o, c) in out.iter_mut().zip(src.as_chunks::<2>().0) { *o = order.u16([c[0], c[1]]); @@ -128,6 +141,30 @@ pub fn undo_float_predictor(row: &mut [u8], n: usize, bytes_per: usize, stride: mod tests { use super::*; + #[test] + fn words32_are_read_as_numbers_then_msb_first() { + // eight 12-bit samples occupy 96 bits = three 32-bit words, filled from the top + let samples: [u16; 8] = [0x123, 0x456, 0x789, 0xabc, 0xdef, 0x012, 0x345, 0x678]; + let stream = samples.iter().fold(0u128, |acc, &v| (acc << 12) | v as u128); + let words = [(stream >> 64) as u32, (stream >> 32) as u32, stream as u32]; + let le: Vec = words.iter().flat_map(|w| w.to_le_bytes()).collect(); + let be: Vec = words.iter().flat_map(|w| w.to_be_bytes()).collect(); + let mut o = [0u16; 8]; + unpack_words32_msb(&le, ByteOrder::Little, 12, &mut o); + assert_eq!(o, samples); + o = [0; 8]; + unpack_words32_msb(&be, ByteOrder::Big, 12, &mut o); + assert_eq!(o, samples); + // read as a plain byte stream the little-endian words give different samples + o = [0; 8]; + unpack_msb(&le, 12, &mut o); + assert_ne!(o, samples); + // a short last word is zero padded + let mut o = [9u16; 2]; + unpack_words32_msb(&[0x34, 0x12], ByteOrder::Little, 12, &mut o); + assert_eq!(o, [0x000, 0x012]); // the word 0x00001234, top bits first + } + #[test] fn packing() { // 12-bit MSB: 0xABC, 0xDEF -> AB CD EF diff --git a/crates/raw/src/vendor/nef.rs b/crates/raw/src/vendor/nef.rs index 8d1898e0..5ca66af4 100644 --- a/crates/raw/src/vendor/nef.rs +++ b/crates/raw/src/vendor/nef.rs @@ -18,6 +18,25 @@ const WB_RB_LEVELS: u16 = 0x000c; const BLACK_LEVEL: u16 = 0x003d; const LINEARIZATION_TABLE: u16 = 0x0096; const CROP_AREA: u16 = 0x0045; +/// Maker note `0x0014`, the "NRW" data block of Coolpix raws: `"NRW "` + a 4-character version, then fields. In +/// version `0104` the `u32` at byte `0x20` is the black level in sample units (200 in the four 12-bit files here +/// that have 200 as their darkest samples, 0 in the one whose samples reach 0). +const NRW_DATA: u16 = 0x0014; +const NRW_BLACK_AT: usize = 0x20; + +/// The NRW data block, when the maker note has one. +fn nrw_data(mn: Option<&makernote::MakerNote>) -> Option<(&[u8], ByteOrder)> { + let m = mn?; + m.ifd.bytes(NRW_DATA).filter(|b| b.starts_with(b"NRW ")).map(|b| (b, m.order)) +} + +/// Black level from the NRW data block, bounded to a quarter of the sample range so an unknown version's field +/// can't turn into a large offset. +fn nrw_black(mn: Option<&makernote::MakerNote>, bits: u32) -> Option { + let (b, order) = nrw_data(mn)?; + let v = order.u32(b.get(NRW_BLACK_AT..NRW_BLACK_AT + 4)?.try_into().ok()?); + (u64::from(v) < (1u64 << bits.min(16)) / 4).then_some(v as f32) +} /// Nikon CropArea is `[left, top, width, height]` in sensor pixels (maker-note tag documentation). /// Keep the active area/CFA origin unchanged: the default crop is applied after demosaicing. @@ -65,8 +84,11 @@ pub(crate) fn decode(bytes: &[u8]) -> Result { let make = ifd0.string(t::MAKE).unwrap_or_default(); let mn = tiff.exif().and_then(|e| e.get(t::MAKER_NOTE)).and_then(|e| makernote::parse_makernote(bytes, e.offset, e.count() as u64, tiff.order, &make)); - let data = - if info.compression == t::compression::NIKON { compressed(bytes, &info, mn.as_ref())? } else { uncompressed(bytes, &info, tiff.order)? }; + let data = if info.compression == t::compression::NIKON { + compressed(bytes, &info, mn.as_ref())? + } else { + uncompressed(bytes, &info, tiff.order, mn.as_ref())? + }; let RawData::U16(ref samples) = data else { return Err(RawError::Unsupported("float NEF".into())) }; // Maker note 0x003d is in 14-bit units whatever the sample depth: 12-bit files from the D750, D780, D850, @@ -80,7 +102,7 @@ pub(crate) fn decode(bytes: &[u8]) -> Result { values: [a, b, c, d].iter().map(|v| (**v * black_scale) as f32).collect(), ..Default::default() }, - _ => BlackLevel::uniform(0.0), + _ => BlackLevel::uniform(nrw_black(mn.as_ref(), bits).unwrap_or(0.0)), }; let wb = mn .as_ref() @@ -120,8 +142,10 @@ pub(crate) fn decode(bytes: &[u8]) -> Result { Ok(img) } -/// Uncompressed strips: 16-bit words or 12-bit MSB-packed, told apart by the strip size. -fn uncompressed(bytes: &[u8], info: &ImageInfo, order: ByteOrder) -> Result { +/// Uncompressed strips: 16-bit words or 12-bit packed, told apart by the strip size. Packed 12-bit data is an MSB-first +/// byte stream, except in files with an NRW data block (Coolpix), where it is MSB-first inside little-endian 32-bit +/// words (the rows are whole words). +fn uncompressed(bytes: &[u8], info: &ImageInfo, order: ByteOrder, mn: Option<&makernote::MakerNote>) -> Result { let (w, h) = (info.width as usize, info.height as usize); let bits = info.bits() as u32; let row_samples = w * info.samples_per_pixel as usize; @@ -131,7 +155,7 @@ fn uncompressed(bytes: &[u8], info: &ImageInfo, order: ByteOrder) -> Result= row_samples * 2 { Packing::Word16 } else if bits == 12 && bytes_per_row * 8 >= row_samples * 12 && bytes_per_row * 8 < row_samples * 13 { - Packing::Msb + if nrw_data(mn).is_some() && bytes_per_row.is_multiple_of(4) { Packing::Words32Msb } else { Packing::Msb } } else if info.compression == 1 && bits != 8 && bits != 16 { return Err(RawError::Unsupported(format!("NEF uncompressed packing ({bytes_per_row} bytes per {row_samples}-sample row)"))); } else { @@ -171,6 +195,12 @@ mod tests { /// [`nef`] with an optional maker note (`Nikon\0` v2 header + embedded big-endian TIFF holding `note`). fn nef_with_note(compression: u16, bits: u16, strips: Vec>, w: u32, h: u32, rps: u32, note: Option) -> Vec { + nef_in(ByteOrder::Big, compression, bits, strips, w, h, rps, note) + } + + /// [`nef_with_note`] in either byte order (file and maker note). + #[allow(clippy::too_many_arguments)] + fn nef_in(order: ByteOrder, compression: u16, bits: u16, strips: Vec>, w: u32, h: u32, rps: u32, note: Option) -> Vec { let mut raw = IfdBuilder::new(); raw.set(t::NEW_SUBFILE_TYPE, Value::Long(vec![0])); raw.set(t::IMAGE_WIDTH, Value::Long(vec![w])); @@ -186,13 +216,13 @@ mod tests { ifd0.set(t::MODEL, Value::Ascii("NIKON TEST".into())); if let Some(mn) = note { let mut bytes = b"Nikon\0\x02\x10\0\0".to_vec(); - bytes.extend(TiffWriter::new(ByteOrder::Big, false).write(&[mn]).unwrap()); + bytes.extend(TiffWriter::new(order, false).write(&[mn]).unwrap()); let mut exif = IfdBuilder::new(); exif.set(t::MAKER_NOTE, Value::Undefined(bytes)); ifd0.set_child(t::EXIF_IFD, exif); } ifd0.add_sub_ifd(raw); - TiffWriter::new(ByteOrder::Big, false).write(&[ifd0]).unwrap() + TiffWriter::new(order, false).write(&[ifd0]).unwrap() } #[test] @@ -266,6 +296,61 @@ mod tests { assert_eq!(crate::decode(&bytes).unwrap().data, RawData::U16(px)); } + /// An NRW data block (`"NRW 0104"`, fields from byte 8) with `black` at byte `0x20`, as a maker note. + fn nrw_note(black: u32) -> IfdBuilder { + let mut block = b"NRW 0104".to_vec(); + block.resize(0x20, 0); + block.extend_from_slice(&black.to_le_bytes()); + block.resize(0x40, 0); + let mut mn = IfdBuilder::new(); + mn.set(NRW_DATA, Value::Undefined(block)); + mn + } + + /// 12-bit samples as an MSB-first bit stream inside 32-bit words in `order` (rows are whole words). + fn words32(px: &[u16], w: usize, order: ByteOrder) -> Vec { + let mut out = Vec::new(); + for row in px.chunks(w) { + let stream = row.iter().fold(Vec::new(), |mut bits: Vec, v| { + bits.extend((0..12).rev().map(|b| v >> b & 1 == 1)); + bits + }); + for word in stream.chunks(32) { + let v = word.iter().enumerate().fold(0u32, |a, (i, b)| a | (*b as u32) << (31 - i)); + out.extend_from_slice(&match order { + ByteOrder::Little => v.to_le_bytes(), + ByteOrder::Big => v.to_be_bytes(), + }); + } + } + out + } + + #[test] + fn nrw_block_selects_32_bit_word_packing_and_supplies_the_black_level() { + let (w, h) = (16usize, 6usize); + let px: Vec = (0..w * h).map(|i| (200 + i * 131 % 3800) as u16).collect(); + let packed = words32(&px, w, ByteOrder::Little); + assert_eq!(packed.len(), w * h * 3 / 2); + let decode = |note: Option| { + crate::decode(&nef_in(ByteOrder::Little, 1, 12, vec![packed.clone()], w as u32, h as u32, h as u32, note)).unwrap() + }; + let r = decode(Some(nrw_note(200))); + assert_eq!(r.data, RawData::U16(px.clone())); + assert_eq!(r.black.values, vec![200.0]); + assert_eq!(decode(Some(nrw_note(0))).black.values, vec![0.0]); + // an implausible field is not a black level + assert_eq!(decode(Some(nrw_note(60000))).black.values, vec![0.0]); + // without the block the same bytes are an MSB-first byte stream + let plain = decode(None); + assert_ne!(plain.data, RawData::U16(px)); + assert_eq!(plain.black.values, vec![0.0]); + // the maker-note black level (14-bit units) wins when present + let mut both = nrw_note(200); + both.set(BLACK_LEVEL, Value::Short(vec![400, 404, 408, 412])); + assert_eq!(decode(Some(both)).black.values, vec![100.0, 101.0, 102.0, 103.0]); + } + #[test] fn compressed_without_table_is_unsupported() { let bytes = nef(34713, 14, vec![vec![0; 64]], 8, 8, 8); diff --git a/crates/raw/src/vendor/raf.rs b/crates/raw/src/vendor/raf.rs index 9d950121..61205a3f 100644 --- a/crates/raw/src/vendor/raf.rs +++ b/crates/raw/src/vendor/raf.rs @@ -12,6 +12,8 @@ //! IFD; offsets in it are relative to the raw block. //! - Sample packing, found by testing candidate bit orders for the smoothest image: 16-bit little-endian words; //! 12-bit LSB-first bit stream; 14-bit (and other depths) stored as little-endian 32-bit words read MSB-first. +//! Raw-IFD flag `0xf006` = 1 selects the 32-bit-word packing at any depth: 12-bit files with the flag set +//! (FinePix S1, X-A10, SL1000) are word-packed like the 14-bit ones, the 12-bit files without it are LSB-first. //! `0xf005`, when non-zero, is the row stride in 32-bit words. //! - CFA: the X-Trans layout record (36 bytes, 0 = R, 1 = G, 2 = B) is stored in reverse order and, reversed, //! is anchored at raw pixel (0, 0) — verified on X-Trans I and III samples by minimising the difference @@ -20,7 +22,7 @@ //! - Striped predictive compression is decoded by `rafc`; see `docs/raf-compression.md`. use super::white_from_data; -use crate::unpack::{unpack_lsb, unpack_msb}; +use crate::unpack::{unpack_lsb, unpack_words32_msb}; use crate::{BlackLevel, Cfa, ColorData, Mode, OpcodeLists, RawData, RawError, RawFormat, RawImage, Rect, Result}; use lightcraft_geom::Orientation; use lightcraft_tiff::{ByteOrder, Ifd, Tiff}; @@ -29,6 +31,9 @@ use rayon::prelude::*; const CROP_TOP_LEFT: u16 = 0x0110; const CROPPED_SIZE: u16 = 0x0111; const XTRANS_LAYOUT: u16 = 0x0131; +/// Raw-IFD flag `0xf006`: 1 in every 12- and 14-bit file here whose samples are packed into 32-bit words (the 12-bit +/// ones are otherwise an LSB-first stream), 0 where they are an LSB-first stream or 16-bit words. +const PACKED_WORDS: u16 = 0xf006; /// The fixed-position header pointers. pub(crate) struct Header<'a> { @@ -95,22 +100,14 @@ fn cfa(h: &Header) -> Cfa { } } -/// Unpack one row of `bits`-bit samples in Fujifilm's packing. -pub(crate) fn unpack_row(src: &[u8], bits: u32, out: &mut [u16]) { +/// Unpack one row of `bits`-bit samples in Fujifilm's packing. `words32`: the raw IFD's `0xf006` flag is set, the +/// samples are MSB-first in little-endian 32-bit words whatever the depth; otherwise 16-bit words, a 12-bit +/// LSB-first bit stream, or (other depths) the same 32-bit words. +pub(crate) fn unpack_row(src: &[u8], bits: u32, words32: bool, out: &mut [u16]) { match bits { 16 => out.iter_mut().zip(src.as_chunks::<2>().0).for_each(|(o, c)| *o = u16::from_le_bytes([c[0], c[1]])), - 12 => unpack_lsb(src, 12, out), - _ => { - let swapped: Vec = src - .chunks(4) - .flat_map(|c| { - let mut w = [0u8; 4]; - w[..c.len()].copy_from_slice(c); - [w[3], w[2], w[1], w[0]] - }) - .collect(); - unpack_msb(&swapped, bits, out) - } + 12 if !words32 => unpack_lsb(src, 12, out), + _ => unpack_words32_msb(src, ByteOrder::Little, bits, out), } } @@ -147,6 +144,7 @@ pub(crate) fn decode(bytes: &[u8], mode: Mode) -> Result { } data = super::rafc::decode(src, w, hgt, bits, cfa.width == 6, mode)?; } else { + let words32 = ifd.u64(PACKED_WORDS).is_some_and(|v| v != 0); // Samples may use 16-bit words (e.g. 14-bit X-T20); otherwise they are packed at `bits`. let store = if len as u64 >= (n as u64) * 2 { 16 } else { bits }; let packed_row = w.checked_mul(store as usize).map(|n| n.div_ceil(8)).ok_or(RawError::Limit("RAF row too large"))?; @@ -163,7 +161,7 @@ pub(crate) fn decode(bytes: &[u8], mode: Mode) -> Result { data.par_chunks_mut(w).enumerate().for_each(|(y, row)| { // `required <= src.len()` proves these row offsets fit, including the final row. let s = src.get(y * stride..y * stride + packed_row).unwrap_or(&[]); - unpack_row(s, store, row); + unpack_row(s, store, words32, row); }); } } @@ -220,19 +218,25 @@ pub(super) mod tests { /// Build a minimal RAF: header, records, raw block (LE TIFF whose IFD0 points to a 0xf000 raw IFD). pub(crate) fn raf(w: u32, h: u32, bits: u32, strip: Vec, layout: Option<[u8; 36]>, jpeg: &[u8]) -> Vec { + raf_flagged(w, h, bits, strip, layout, jpeg, 0) + } + + /// [`raf`] with the raw IFD's `0xf006` flag set to `words32`. + pub(crate) fn raf_flagged(w: u32, h: u32, bits: u32, strip: Vec, layout: Option<[u8; 36]>, jpeg: &[u8], words32: u32) -> Vec { let mut block = b"II*\0\x08\0\0\0".to_vec(); let sub_off: u32 = 8 + 2 + 12 + 4; block.extend_from_slice(&1u16.to_le_bytes()); block.extend_from_slice(&[0x00, 0xf0, 13, 0, 1, 0, 0, 0]); block.extend_from_slice(&sub_off.to_le_bytes()); block.extend_from_slice(&0u32.to_le_bytes()); - let n = 7u32; + let n = 8u32; let wb_off = sub_off + 2 + 12 * n + 4; let strip_off = wb_off + 12; - let entries: [(u16, u32, u32); 7] = [ + let entries: [(u16, u32, u32); 8] = [ (0xf001, 1, w), (0xf002, 1, h), (0xf003, 1, bits), + (0xf006, 1, words32), (0xf007, 1, strip_off), (0xf008, 1, strip.len() as u32), (0xf00a, 1, 64), @@ -345,14 +349,39 @@ pub(super) mod tests { assert!(matches!(crate::decode(&bytes), Err(RawError::Unsupported(_)))); } + /// Raw IFD flag `0xf006` selects the 32-bit-word packing at 12 bits; without it 12 bits are an LSB-first stream. + #[test] + fn flag_selects_word_packed_12_bit() { + let (w, h) = (16u32, 6u32); + let px: Vec = (0..w * h).map(|i| ((i * 977) % 4096) as u16).collect(); + // the same 12-bit samples as little-endian 32-bit words (what the 14-bit branch of `pack` writes) + let mut stream: Vec = Vec::new(); + let (mut acc, mut n) = (0u64, 0u32); + for &v in &px { + acc = (acc << 12) | v as u64; + n += 12; + while n >= 8 { + stream.push((acc >> (n - 8)) as u8); + n -= 8; + } + } + let le: Vec = stream.chunks(4).flat_map(|c| [c[3], c[2], c[1], c[0]]).collect(); + let decode = |bytes: Vec| crate::decode(&bytes).unwrap().data; + assert_eq!(decode(raf_flagged(w, h, 12, le.clone(), None, b"", 1)), RawData::U16(px.clone())); + // flag clear: the same bytes are read as the LSB-first stream, which gives other samples + assert_ne!(decode(raf_flagged(w, h, 12, le, None, b"", 0)), RawData::U16(px.clone())); + // flag clear: the LSB-first stream still decodes + assert_eq!(decode(raf_flagged(w, h, 12, pack(&px, 12), None, b"", 0)), RawData::U16(px)); + } + #[test] fn overflowing_row_stride_is_an_error() { let mut bytes = raf(16, 6, 16, vec![0; 192], None, &[]); let off = be32(&bytes, 100).unwrap(); // Replace WB with a LONG8 stride, pointing to the existing out-of-line WB storage. - bytes[off + 100..off + 104].copy_from_slice(&[0x05, 0xf0, 16, 0]); - bytes[off + 104..off + 108].copy_from_slice(&1u32.to_le_bytes()); - bytes[off + 116..off + 124].copy_from_slice(&u64::MAX.to_le_bytes()); + bytes[off + 112..off + 116].copy_from_slice(&[0x05, 0xf0, 16, 0]); + bytes[off + 116..off + 120].copy_from_slice(&1u32.to_le_bytes()); + bytes[off + 128..off + 136].copy_from_slice(&u64::MAX.to_le_bytes()); assert!(matches!(crate::decode(&bytes), Err(RawError::Limit(_)))); } } diff --git a/crates/raw/tests/corpus.rs b/crates/raw/tests/corpus.rs index b3c0ade9..53193f1f 100644 --- a/crates/raw/tests/corpus.rs +++ b/crates/raw/tests/corpus.rs @@ -43,7 +43,9 @@ fn corpus_raw_decodes() { // DNG previews are optional (and some carry only an uncompressed RGB thumbnail); vendor raws embed a JPEG, // except the Panasonic `.RAW` files of 2005–2007 and Canon's explicit HEVC preview tracks. if let Some(p) = &preview { - assert!(p.starts_with(&[0xff, 0xd8]) && p.ends_with(&[0xff, 0xd9]), "{name}: preview is not a JPEG"); + let jpeg = p.starts_with(&[0xff, 0xd8]) && p.ends_with(&[0xff, 0xd9]); + let jxl = p.starts_with(&[0xff, 0x0a]) || p.starts_with(b"\0\0\0\x0cJXL "); + assert!(jpeg || jxl, "{name}: preview is neither a JPEG nor a JPEG XL file"); } else { let hevc_preview = fmt == RawFormat::Cr3 && lightcraft_meta::cr3::parse_cr3(&bytes) @@ -250,9 +252,11 @@ fn corpus_nef_12_bit_black_level_matches_14_bit() { } /// Issue #138: DNGs converted by Adobe software carry their camera profile's hue/saturation map and -/// look table; we read them (and render with them). Camera-written DNGs here carry none. +/// look table; we read them (and render with them). Apple ProRAW carries a tone curve and a gain +/// table map (its local tone mapping, read and kept but not rendered, as in Lightroom Classic). +/// Other camera-written DNGs here carry none. #[test] -fn corpus_adobe_dngs_carry_profile_looks() { +fn corpus_dngs_carry_profile_looks() { let dir = corpus_root().join("raw"); let Ok(rd) = std::fs::read_dir(&dir) else { eprintln!("skip: {} absent", dir.display()); @@ -276,14 +280,51 @@ fn corpus_adobe_dngs_carry_profile_looks() { let g = t.apply([0.18; 3], 1.0); assert!(g.iter().all(|v| (v - g[0]).abs() < 0.01 * g[0].max(0.01)), "{name}: grey → {g:?}"); } + if name.starts_with("dng-apple-") { + seen += 1; + let map = look.gain_table_map.as_ref().unwrap_or_else(|| panic!("{name}: no gain table map")); + assert!(map.points_v > 1 && map.points_h > 1 && map.points_n > 1, "{name}"); + assert!(look.tone_curve.is_some(), "{name}: no tone curve"); + // the map lifts dark tones (Apple's local tone mapping), most at the darkest input + let first = &map.gains[..map.points_n]; + assert!(first[0] > 1.5 && first[0] >= first[map.points_n - 1], "{name}: first table {first:?}"); + } eprintln!( - "{name:44} profile look: hsm {} look {} tone {}", + "{name:44} profile look: hsm {} look {} tone {} gain map {}", look.hue_sat_map[0].is_some(), look.look_table.is_some(), - look.tone_curve.is_some() + look.tone_curve.is_some(), + look.gain_table_map.is_some() ); } - eprintln!("{seen} Adobe-converted DNGs checked"); + eprintln!("{seen} DNGs with profile looks checked"); +} + +/// iPhone ProRAW (`dng-apple-iphone12pro-proraw.dng`, iPhone 12 Pro, CC0 from raw.pixls.us, see +/// `cargo xtask corpus`) carries a lossy-JPEG sky matte at half resolution as a DNG semantic mask: +/// it decodes, and it is a sky. The matte is stored like the raw (landscape; the photo is shown +/// rotated 90° clockwise), so the sky is on its left and the lake and shore on its right. +#[test] +fn corpus_proraw_sky_matte() { + let path = corpus_root().join("raw/dng-apple-iphone12pro-proraw.dng"); + let Ok(bytes) = std::fs::read(&path) else { + eprintln!("skip: {} absent", path.display()); + return; + }; + let t0 = Instant::now(); + let masks = lightcraft_raw::semantic_masks(&bytes); + eprintln!("semantic masks read in {:.1} ms", t0.elapsed().as_secs_f64() * 1e3); + let names: Vec<&str> = masks.iter().map(|m| m.name.as_str()).collect(); + assert_eq!(names, ["urn:com:apple:photo:2020:aux:semanticskymatte"]); + let m = &masks[0]; + assert_eq!((m.width, m.height, m.sub_area), (2016, 1512, None)); + let cols = |x0: usize, x1: usize| { + let sum: f64 = (0..m.height).flat_map(|y| m.data[y * m.width + x0..y * m.width + x1].iter()).map(|&v| v as f64 / 65535.0).sum(); + sum / ((x1 - x0) * m.height) as f64 + }; + let (left, right) = (cols(0, 100), cols(1916, 2016)); + eprintln!("sky matte: left columns {left:.3}, right columns {right:.3}"); + assert!(left > 0.5 && right < 0.02, "left {left}, right {right}"); } /// Issue #148: Sony ARWs from before ~2017 carry no plain white-balance, black-level or crop tags in the raw IFD. diff --git a/crates/raw/tests/dng.rs b/crates/raw/tests/dng.rs index 0d38ebb4..bba5ac81 100644 --- a/crates/raw/tests/dng.rs +++ b/crates/raw/tests/dng.rs @@ -39,6 +39,7 @@ fn synthetic(w: usize, h: usize, cfa: Option, cpp: usize) -> RawImage { ], as_shot_neutral: Some([0.5, 1.0, 0.7]), baseline_exposure: 0.35, + baseline_sharpness: Some(1.5), ..Default::default() }; let metadata = lightcraft_meta::Metadata { make: Some("Synth".into()), model: Some("Cam 1".into()), rating: Some(3), ..Default::default() }; @@ -84,6 +85,7 @@ fn assert_same(a: &RawImage, b: &RawImage) { let (n1, n2) = (a.color.as_shot_neutral.unwrap(), b.color.as_shot_neutral.unwrap()); assert!((0..3).all(|i| (n1[i] - n2[i]).abs() < 1e-6)); assert!((a.color.baseline_exposure - b.color.baseline_exposure).abs() < 1e-6); + assert!(b.color.baseline_sharpness.is_some_and(|v| (v - 1.5).abs() < 1e-6), "{:?}", b.color.baseline_sharpness); assert_eq!(b.metadata.make.as_deref(), Some("Synth")); assert_eq!(b.metadata.rating, Some(3)); } @@ -548,3 +550,422 @@ fn malformed_profile_look_tags_are_ignored() { let plain = decode(&profile_dng(ByteOrder::Little, |_| {})).unwrap(); assert!(plain.color.profile.is_empty()); } + +/// Lossless JPEG XL codestream of `px` (`w × h × cpp` interleaved), 16-bit unless `eight`. +fn jxl_encode(px: &[u16], w: usize, h: usize, cpp: usize, eight: bool) -> Vec { + use zune_core::bit_depth::BitDepth; + use zune_core::colorspace::ColorSpace; + use zune_core::options::EncoderOptions; + let cs = if cpp == 1 { ColorSpace::Luma } else { ColorSpace::RGB }; + let (bytes, depth): (Vec, _) = if eight { + (px.iter().map(|&v| v as u8).collect(), BitDepth::Eight) + } else { + (px.iter().flat_map(|v| v.to_ne_bytes()).collect(), BitDepth::Sixteen) + }; + let mut out = Vec::new(); + zune_jpegxl::JxlSimpleEncoder::new(&bytes, EncoderOptions::new(w, h, cs, depth)).encode(&mut out).unwrap(); + out +} + +/// A bare codestream wrapped in the ISO-BMFF container (signature, `ftyp`, `jxlc`). +fn jxl_container(cs: &[u8]) -> Vec { + let mut b = vec![0, 0, 0, 0x0c, b'J', b'X', b'L', b' ', 0x0d, 0x0a, 0x87, 0x0a, 0, 0, 0, 0x14]; + b.extend_from_slice(b"ftypjxl \0\0\0\0jxl "); + b.extend_from_slice(&((cs.len() + 8) as u32).to_be_bytes()); + b.extend_from_slice(b"jxlc"); + b.extend_from_slice(cs); + b +} + +/// A tiled JPEG XL raw IFD of `px`. Edge tiles alternate between being coded at their in-image size +/// and padded to the nominal tile size (with junk the decoder must drop); tiles alternate between bare +/// codestreams and containers. +fn jxl_tiled(px: &[u16], w: usize, h: usize, cpp: usize, bits: u16, cfa: bool, tile: (usize, usize)) -> IfdBuilder { + let (tw, th) = tile; + let mut tiles = Vec::new(); + for ty in 0..h.div_ceil(th) { + for tx in 0..w.div_ceil(tw) { + let i = tiles.len(); + let (vw, vh) = (tw.min(w - tx * tw), th.min(h - ty * th)); + let (cw, ch) = if i % 2 == 0 { (tw, th) } else { (vw, vh) }; + let mut t = vec![bits.min(8) * 7; cw * ch * cpp]; + for y in 0..vh { + for x in 0..vw { + for s in 0..cpp { + t[(y * cw + x) * cpp + s] = px[((ty * th + y) * w + tx * tw + x) * cpp + s]; + } + } + } + let cs = jxl_encode(&t, cw, ch, cpp, bits == 8); + tiles.push(if i % 3 == 1 { jxl_container(&cs) } else { cs }); + } + } + let mut raw = base_ifd(w, h, bits, cpp, cfa); + raw.set(t::COMPRESSION, Value::Short(vec![compression::JPEG_XL])); + raw.set_image(ImageData::Tiles { tile_width: tw as u32, tile_height: th as u32, tiles }); + raw +} + +#[test] +fn jpeg_xl_tiles_decode_bit_exactly() { + let (w, h) = (37, 29); + // LinearRaw 3 × 16-bit, the full code range + let px = pattern(w, h, 3, 16); + let bytes = dng(jxl_tiled(&px, w, h, 3, 16, false, (16, 16)), ByteOrder::Big); + let img = decode(&bytes).unwrap(); + assert_eq!((img.width, img.height, img.cpp), (w, h, 3)); + assert_eq!(img.data, RawData::U16(px)); + assert_eq!(lightcraft_raw::probe_info(&bytes).unwrap(), img.info()); + // CFA 1 × 8/12/14-bit + for bits in [8u32, 12, 14] { + let px = pattern(w, h, 1, bits); + let img = decode(&dng(jxl_tiled(&px, w, h, 1, bits as u16, true, (16, 8)), ByteOrder::Little)).unwrap(); + assert_eq!(img.data, RawData::U16(px), "{bits} bits"); + assert_eq!(img.white, vec![((1u32 << bits) - 1) as f32]); + assert_eq!(img.cfa.unwrap().name(), "GRBG"); + } + // strips: a single JPEG XL strip per 8 rows (the last one shorter) + let px = pattern(w, h, 1, 12); + let strips = px.chunks(w * 8).map(|c| jxl_encode(c, w, c.len() / w, 1, false)).collect(); + let mut raw = base_ifd(w, h, 12, 1, true); + raw.set(t::COMPRESSION, Value::Short(vec![compression::JPEG_XL])); + raw.set_image(ImageData::Strips { rows_per_strip: 8, strips }); + assert_eq!(decode(&dng(raw, ByteOrder::Big)).unwrap().data, RawData::U16(px)); +} + +#[test] +fn malformed_jpeg_xl_tiles_are_errors() { + use lightcraft_raw::RawError; + let (w, h) = (16, 16); + let px = pattern(w, h, 1, 12); + let good = jxl_encode(&px, w, h, 1, false); + let one = |tile: Vec, cpp: usize| { + let mut raw = base_ifd(w, h, 12, cpp, cpp == 1); + raw.set(t::COMPRESSION, Value::Short(vec![compression::JPEG_XL])); + raw.set_image(ImageData::Tiles { tile_width: w as u32, tile_height: h as u32, tiles: vec![tile] }); + dng(raw, ByteOrder::Little) + }; + // garbage, every truncation, and flipped bytes: errors (or, for a truncation that only drops + // trailing padding, the exact samples), never panics + let mut garbage = vec![0xff, 0x0a]; + garbage.extend((0..400u32).map(|i| (i.wrapping_mul(2654435761) >> 13) as u8)); + assert!(matches!(decode(&one(garbage, 1)), Err(RawError::Corrupt(_)))); + for n in 0..good.len() { + let bytes = one(good[..n].to_vec(), 1); + if let Ok(img) = decode(&bytes) { + assert_eq!(img.data, RawData::U16(px.clone()), "truncated to {n}"); + } + let _ = lightcraft_raw::probe_info(&bytes); + } + for i in 0..good.len() { + let mut b = good.clone(); + b[i] ^= 0x5a; + let _ = decode(&one(b, 1)); + } + // a tile of the wrong size, or the wrong plane count for the IFD + let small = jxl_encode(&px[..w * 8], w, 8, 1, false); + assert!(matches!(decode(&one(small, 1)), Err(RawError::Corrupt(_)))); + assert!(matches!(decode(&one(good.clone(), 3)), Err(RawError::Corrupt(_)))); + assert!(matches!(lightcraft_raw::probe_info(&one(good, 3)), Err(RawError::Corrupt(_)))); +} + +/// The decoder's memory is bounded by the tile: a grey tile that also codes an alpha channel needs +/// about twice what one plane may use, and is an error rather than an allocation the header asks +/// for. The same samples without the alpha channel decode. +#[test] +fn jpeg_xl_tiles_beyond_their_decode_budget_are_errors() { + use lightcraft_raw::RawError; + use zune_core::bit_depth::BitDepth; + use zune_core::colorspace::ColorSpace; + use zune_core::options::EncoderOptions; + let (w, h) = (512, 512); + let grey: Vec = (0..w * h).map(|i| ((i % w) * 97 + (i / w) * 31) as u16 & 0x0fff).collect(); + let with_alpha: Vec = grey.iter().flat_map(|&v| [v, 4095 - v]).flat_map(u16::to_ne_bytes).collect(); + let mut tile = Vec::new(); + zune_jpegxl::JxlSimpleEncoder::new(&with_alpha, EncoderOptions::new(w, h, ColorSpace::LumaA, BitDepth::Sixteen)).encode(&mut tile).unwrap(); + let one = |tile: Vec| { + let mut raw = base_ifd(w, h, 12, 1, true); + raw.set(t::COMPRESSION, Value::Short(vec![compression::JPEG_XL])); + raw.set_image(ImageData::Tiles { tile_width: w as u32, tile_height: h as u32, tiles: vec![tile] }); + dng(raw, ByteOrder::Little) + }; + assert!(matches!(decode(&one(tile)), Err(RawError::Corrupt(_)))); + assert_eq!(decode(&one(jxl_encode(&grey, w, h, 1, false))).unwrap().data, RawData::U16(grey)); +} + +#[test] +fn jpeg_xl_preview_ifd_is_offered() { + let raw = synthetic(16, 16, Some(Cfa::bayer("RGGB").unwrap()), 1); + let RawData::U16(px) = &raw.data else { unreachable!() }; + let mut sub = base_ifd(16, 16, 14, 1, true); + sub.set(t::COMPRESSION, Value::Short(vec![compression::JPEG_XL])); + sub.set_image(ImageData::Tiles { tile_width: 16, tile_height: 16, tiles: vec![jxl_encode(px, 16, 16, 1, false)] }); + let preview = jxl_container(&jxl_encode(&[200; 4 * 2 * 3], 4, 2, 3, true)); + let bytes = dng_with(sub, ByteOrder::Big, |ifd0| { + ifd0.set(t::COMPRESSION, Value::Short(vec![compression::JPEG_XL])); + ifd0.set_image(ImageData::Strips { rows_per_strip: 2, strips: vec![preview.clone()] }); + }); + // the rendered RGB preview is offered whole; the (single-tile) CFA raw never is + assert_eq!(embedded_preview(&bytes).unwrap(), preview); + assert_eq!(decode(&bytes).unwrap().data, raw.data); +} + +/// `ProfileGainTableMap` / `ProfileGainTableMap2` bytes in `order`: a `dims` = (V, H, N) grid with +/// spacing 0.5 and origin 0.25, green-only input weights, `v2` = (data type, gamma, gain min, +/// gain max), then `data` (the stored gains, already in `order`). +fn gain_map_tag(order: ByteOrder, dims: [u32; 3], v2: Option<(u32, f32, f32, f32)>, data: &[u8]) -> Vec { + let be = order == ByteOrder::Big; + let u = |v: u32| if be { v.to_be_bytes() } else { v.to_le_bytes() }; + let mut b = Vec::new(); + b.extend_from_slice(&u(dims[0])); + b.extend_from_slice(&u(dims[1])); + for v in [0.5f64, 0.5, 0.25, 0.25] { + b.extend_from_slice(&if be { v.to_be_bytes() } else { v.to_le_bytes() }); + } + b.extend_from_slice(&u(dims[2])); + for w in [0.0f32, 1.0, 0.0, 0.0, 0.0] { + b.extend_from_slice(&u(w.to_bits())); + } + if let Some((dt, gamma, min, max)) = v2 { + for v in [dt, gamma.to_bits(), min.to_bits(), max.to_bits()] { + b.extend_from_slice(&u(v)); + } + } + b.extend_from_slice(data); + b +} + +fn f32_bytes(v: &[f32], order: ByteOrder) -> Vec { + v.iter().flat_map(|f| if order == ByteOrder::Big { f.to_be_bytes() } else { f.to_le_bytes() }).collect() +} + +/// [`profile_dng`] with binary tags in the raw IFD (where Apple writes `ProfileGainTableMap`) and +/// in IFD 0. +fn gain_map_dng(order: ByteOrder, raw_tags: &[(u16, &[u8])], ifd0_tags: &[(u16, &[u8])]) -> Vec { + let (w, h) = (8, 6); + let px = pattern(w, h, 3, 16); + let mut raw = base_ifd(w, h, 16, 3, false); + raw.set(t::COMPRESSION, Value::Short(vec![1])); + raw.set_image(ImageData::Strips { rows_per_strip: h as u32, strips: vec![u16_bytes(&px, order)] }); + for (tag, b) in raw_tags { + raw.set(*tag, Value::Undefined(b.to_vec())); + } + dng_with(raw, order, |ifd0| { + for (tag, b) in ifd0_tags { + ifd0.set(*tag, Value::Undefined(b.to_vec())); + } + }) +} + +#[test] +fn gain_table_maps_follow_dng_precedence() { + const V1: u16 = t::PROFILE_GAIN_TABLE_MAP; + const V2: u16 = t::PROFILE_GAIN_TABLE_MAP_2; + for order in [ByteOrder::Big, ByteOrder::Little] { + let v1 = gain_map_tag(order, [2, 2, 4], None, &f32_bytes(&(1..=16).map(|i| i as f32 * 0.25).collect::>(), order)); + // Apple's layout: version 1 in the raw IFD, in the file's byte order + let apple = gain_map_dng(order, &[(V1, &v1[..])], &[]); + let img = decode(&apple).unwrap(); + let map = img.color.profile.gain_table_map.clone().expect("gain table map"); + assert_eq!((map.points_v, map.points_h, map.points_n, map.gamma), (2, 2, 4, 1.0)); + assert_eq!((map.spacing_v, map.origin_h, map.weights[1]), (0.5, 0.25, 1.0)); + assert_eq!(map.gains[5], 1.5); + assert_eq!(lightcraft_raw::probe_info(&apple).unwrap().color.profile, img.color.profile); + // our DNG writer keeps it + let again = decode(&write_dng(&img, &DngWriteOptions::default()).unwrap()).unwrap(); + assert_eq!(again.color.profile.gain_table_map.as_ref(), Some(&map)); + // version 1 in IFD 0 (where DNG 1.7 says it was meant to go) is read too + let in_ifd0 = decode(&gain_map_dng(order, &[], &[(V1, &v1[..])])).unwrap(); + assert_eq!(in_ifd0.color.profile.gain_table_map.as_ref(), Some(&map)); + // version 2 in IFD 0 supersedes version 1: 8-bit gains over 0.5..=1.5, gamma 2 + let v2 = gain_map_tag(order, [1, 1, 2], Some((0, 2.0, 0.5, 1.5)), &[0, 255]); + let img = decode(&gain_map_dng(order, &[(V1, &v1[..])], &[(V2, &v2[..])])).unwrap(); + let map2 = img.color.profile.gain_table_map.clone().unwrap(); + assert_eq!((map2.points_n, map2.gamma, map2.gains.clone()), (2, 2.0, vec![0.5, 1.5])); + // a gamma ≠ 1 survives our writer (as version 2) + let again = decode(&write_dng(&img, &DngWriteOptions::default()).unwrap()).unwrap(); + assert_eq!(again.color.profile.gain_table_map, Some(map2)); + // a malformed version 2 (gamma outside 0.25..=4) falls back to version 1 + let bad_v2 = gain_map_tag(order, [1, 1, 2], Some((0, 9.0, 0.5, 1.5)), &[0, 255]); + assert_eq!(decode(&gain_map_dng(order, &[(V1, &v1[..])], &[(V2, &bad_v2[..])])).unwrap().color.profile.gain_table_map, Some(map)); + } +} + +#[test] +fn malformed_gain_table_maps_are_ignored() { + let order = ByteOrder::Little; + let ok = gain_map_tag(order, [1, 1, 2], None, &f32_bytes(&[1.0, 2.0], order)); + let cases: Vec<(&str, Vec, bool)> = vec![ + ("valid", ok.clone(), false), + ("empty", vec![], false), + ("header only", ok[..64].to_vec(), false), + ("one byte short", ok[..ok.len() - 1].to_vec(), false), + ("NaN gain", gain_map_tag(order, [1, 1, 2], None, &f32_bytes(&[1.0, f32::NAN], order)), false), + ("negative gain", gain_map_tag(order, [1, 1, 2], None, &f32_bytes(&[1.0, -2.0], order)), false), + ("huge dimensions", gain_map_tag(order, [u32::MAX, u32::MAX, u32::MAX], None, &f32_bytes(&[1.0, 2.0], order)), false), + ("zero points", gain_map_tag(order, [1, 0, 2], None, &[]), false), + ("unknown data type", gain_map_tag(order, [1, 1, 2], Some((9, 1.0, 0.0, 1.0)), &[0, 0]), true), + ("NaN gain range", gain_map_tag(order, [1, 1, 2], Some((0, 1.0, f32::NAN, 1.0)), &[0, 255]), true), + ]; + for (what, bytes, v2) in cases { + let tag = if v2 { t::PROFILE_GAIN_TABLE_MAP_2 } else { t::PROFILE_GAIN_TABLE_MAP }; + let img = decode(&gain_map_dng(order, &[(tag, &bytes[..])], &[])).unwrap_or_else(|e| panic!("{what}: {e}")); + assert_eq!(img.color.profile.gain_table_map.is_some(), what == "valid", "{what}"); + } +} + +/// The map is read in the file's byte order, else big-endian (`lightcraft_raw::gaintable`): a +/// big-endian payload in a little-endian file is read, a little-endian one in a big-endian file +/// is not. +#[test] +fn gain_table_map_byte_orders() { + use ByteOrder::{Big, Little}; + let payload = |o: ByteOrder| gain_map_tag(o, [1, 1, 2], None, &f32_bytes(&[1.0, 2.0], o)); + for (file, data, read) in [(Little, Little, true), (Little, Big, true), (Big, Big, true), (Big, Little, false)] { + let img = decode(&gain_map_dng(file, &[(t::PROFILE_GAIN_TABLE_MAP, &payload(data)[..])], &[])).unwrap(); + assert_eq!(img.color.profile.gain_table_map.map(|m| m.gains), read.then(|| vec![1.0, 2.0]), "{data:?} payload, {file:?} file"); + } +} + +/// A semantic mask IFD (8-bit, uncompressed) named `name`. +fn mask_ifd(w: usize, h: usize, name: &str, px: Vec) -> IfdBuilder { + let mut m = IfdBuilder::new(); + m.set(t::NEW_SUBFILE_TYPE, Value::Long(vec![t::SUBFILE_SEMANTIC_MASK])); + m.set(t::IMAGE_WIDTH, Value::Long(vec![w as u32])); + m.set(t::IMAGE_LENGTH, Value::Long(vec![h as u32])); + m.set(t::BITS_PER_SAMPLE, Value::Short(vec![8])); + m.set(t::SAMPLES_PER_PIXEL, Value::Short(vec![1])); + m.set(t::PHOTOMETRIC, Value::Short(vec![photometric::MASK])); + m.set(t::COMPRESSION, Value::Short(vec![1])); + if !name.is_empty() { + m.set(t::SEMANTIC_NAME, Value::Ascii(name.into())); + } + m.set_image(ImageData::Strips { rows_per_strip: h as u32, strips: vec![px] }); + m +} + +/// Semantic masks: a cropped mask is placed by its `MaskSubArea` (zero outside it), mapped onto the +/// developed image (default crop of the active area), and doesn't disturb the raw decode. +#[test] +fn semantic_masks_are_read_and_placed() { + use lightcraft_raw::semantic::MaskSubArea; + for order in [ByteOrder::Little, ByteOrder::Big] { + let bytes = profile_dng(order, |ifd0| { + // a 4 × 2 crop at (top 1, left 2) of an 8 × 4 mask over the 8 × 6 image + let mut sky = mask_ifd(4, 2, "sky", vec![10, 20, 30, 40, 50, 60, 70, 255]); + sky.set(t::MASK_SUB_AREA, Value::Long(vec![1, 2, 8, 4])); + ifd0.add_sub_ifd(sky); + let mut skin = mask_ifd(2, 2, "skin", vec![0, 255, 255, 0]); + skin.set(t::SEMANTIC_INSTANCE_ID, Value::Ascii("person_a".into())); + ifd0.add_sub_ifd(skin); + }); + let raw = decode(&bytes).unwrap(); + assert_eq!((raw.width, raw.height, raw.cpp), (8, 6, 3), "the masks are not the raw image"); + let masks = lightcraft_raw::semantic_masks(&bytes); + assert_eq!(masks.len(), 2); + let (sky, skin) = (&masks[0], &masks[1]); + assert_eq!((sky.name.as_str(), sky.instance_id.as_deref(), sky.width, sky.height), ("sky", None, 4, 2)); + assert_eq!(sky.sub_area, Some(MaskSubArea { top: 1, left: 2, full_width: 8, full_height: 4 })); + assert_eq!((sky.data[0], sky.data[7]), (10 * 257, 65535)); + assert_eq!((skin.name.as_str(), skin.instance_id.as_deref(), skin.sub_area), ("skin", Some("person_a"), None)); + // positions are in uncropped-mask pixels: zero outside the stored crop + assert!((sky.value_at(2.5, 1.5) - 10.0 / 255.0).abs() < 1e-4); + assert!((sky.value_at(5.5, 2.5) - 1.0).abs() < 1e-4); + assert_eq!((sky.value_at(1.5, 1.5), sky.value_at(2.5, 0.5), sky.value_at(6.5, 1.5)), (0.0, 0.0, 0.0)); + // over the whole image: the uncropped mask's 8 × 4 grid + let full = sky.developed(Rect::new(0, 0, 8, 6), Rect::new(0, 0, 8, 6)).unwrap(); + assert_eq!((full.width, full.height), (8, 4)); + assert_eq!((full.get(2, 1), full.get(5, 2), full.get(1, 1), full.get(6, 2), full.get(2, 0)), (10, 255, 0, 0, 0)); + // a default crop keeping columns 2..6 of the image: the stored crop fills it + let cropped = sky.developed(Rect::new(0, 0, 8, 6), Rect::new(2, 0, 4, 6)).unwrap(); + assert_eq!((cropped.width, cropped.height), (4, 4)); + assert_eq!((cropped.get(0, 1), cropped.get(3, 2), cropped.get(0, 0)), (10, 255, 0)); + } +} + +/// Malformed semantic masks are skipped (never a panic), and so is the `MaskSubArea` of a mask it +/// doesn't fit; the raw image still decodes. +#[test] +fn malformed_semantic_masks_are_skipped() { + let bytes = profile_dng(ByteOrder::Little, |ifd0| { + ifd0.add_sub_ifd(mask_ifd(2, 2, "", vec![1, 2, 3, 4])); // no SemanticName + let mut rgb = mask_ifd(2, 2, "rgb", vec![0; 12]); + rgb.set(t::SAMPLES_PER_PIXEL, Value::Short(vec![3])); + rgb.set(t::BITS_PER_SAMPLE, Value::Short(vec![8, 8, 8])); + ifd0.add_sub_ifd(rgb); + let mut huge = mask_ifd(2, 2, "huge", vec![0; 4]); + huge.set(t::IMAGE_WIDTH, Value::Long(vec![100_000])); + huge.set(t::IMAGE_LENGTH, Value::Long(vec![100_000])); + ifd0.add_sub_ifd(huge); + let mut tiny = mask_ifd(2, 2, "tiny data", vec![0; 4]); + tiny.set(t::IMAGE_WIDTH, Value::Long(vec![4000])); + tiny.set(t::IMAGE_LENGTH, Value::Long(vec![3000])); + ifd0.add_sub_ifd(tiny); + let mut jxl = mask_ifd(2, 2, "unknown coding", vec![0; 4]); + jxl.set(t::COMPRESSION, Value::Short(vec![60_000])); + ifd0.add_sub_ifd(jxl); + let mut zero = mask_ifd(2, 2, "zero size", vec![0; 4]); + zero.set(t::IMAGE_WIDTH, Value::Long(vec![0])); + ifd0.add_sub_ifd(zero); + let mut short = mask_ifd(4, 4, "short strip", vec![200; 3]); + short.set(t::MASK_SUB_AREA, Value::Long(vec![u32::MAX, u32::MAX, u32::MAX, u32::MAX])); + ifd0.add_sub_ifd(short); + let mut outside = mask_ifd(2, 2, "sub area too small", vec![255; 4]); + outside.set(t::MASK_SUB_AREA, Value::Long(vec![3, 0, 2, 4])); + ifd0.add_sub_ifd(outside); + }); + assert_eq!(decode(&bytes).unwrap().cpp, 3); + let masks = lightcraft_raw::semantic_masks(&bytes); + let names: Vec<&str> = masks.iter().map(|m| m.name.as_str()).collect(); + assert_eq!(names, ["short strip", "sub area too small"]); + assert!(masks.iter().all(|m| m.sub_area.is_none()), "sub areas the mask doesn't fit are ignored"); + for m in &masks { + let d = m.developed(Rect::new(0, 0, 8, 6), Rect::new(0, 0, 8, 6)).unwrap(); + assert_eq!((d.width, d.height), (m.width, m.height)); + assert!(m.developed(Rect::new(0, 0, 8, 6), Rect::new(8, 6, 0, 0)).is_none()); + } + // not a TIFF at all + assert!(lightcraft_raw::semantic_masks(b"not a dng").is_empty()); +} + +/// Only the first `MAX_MASKS` mask IFDs are decoded, malformed ones included, so a file full of +/// mask IFDs costs a bounded number of decodes: here the first `MAX_MASKS` can't be decoded, and +/// the valid masks after them are never looked at. +#[test] +fn semantic_mask_ifds_past_the_cap_are_not_decoded() { + use lightcraft_raw::semantic::MAX_MASKS; + let with = |undecodable: usize, valid: usize| { + profile_dng(ByteOrder::Little, |ifd0| { + for i in 0..undecodable { + let mut m = mask_ifd(2, 2, &format!("undecodable {i}"), vec![0; 4]); + m.set(t::COMPRESSION, Value::Short(vec![60_000])); + ifd0.add_sub_ifd(m); + } + for i in 0..valid { + ifd0.add_sub_ifd(mask_ifd(2, 2, &format!("valid {i}"), vec![255; 4])); + } + }) + }; + let bytes = with(MAX_MASKS, 3 * MAX_MASKS); + assert_eq!(decode(&bytes).unwrap().cpp, 3); + assert!(lightcraft_raw::semantic_masks(&bytes).is_empty()); + // one fewer undecodable mask leaves room for exactly one valid one + let names: Vec = lightcraft_raw::semantic_masks(&with(MAX_MASKS - 1, 3 * MAX_MASKS)).into_iter().map(|m| m.name).collect(); + assert_eq!(names, ["valid 0"]); +} + +/// `BaselineSharpness` is read from IFD 0; a zero, negative or absurd value is ignored. +#[test] +fn baseline_sharpness_is_read_when_plausible() { + for (value, want) in [((3, 2), Some(1.5)), ((0, 1), None), ((1, 0), None), ((100, 1), None)] { + let bytes = profile_dng(ByteOrder::Big, |ifd0| { + ifd0.set(t::BASELINE_SHARPNESS, Value::Rational(vec![value])); + }); + assert_eq!(decode(&bytes).unwrap().color.baseline_sharpness, want, "{value:?}"); + } + let bytes = profile_dng(ByteOrder::Big, |ifd0| { + ifd0.set(t::BASELINE_SHARPNESS, Value::SRational(vec![(-1, 1)])); + }); + assert_eq!(decode(&bytes).unwrap().color.baseline_sharpness, None); + assert_eq!(decode(&profile_dng(ByteOrder::Big, |_| {})).unwrap().color.baseline_sharpness, None); +} diff --git a/crates/segment/src/bin/lightcraft-sam3-worker.rs b/crates/segment/src/bin/lightcraft-sam3-worker.rs new file mode 100644 index 00000000..f8e52541 --- /dev/null +++ b/crates/segment/src/bin/lightcraft-sam3-worker.rs @@ -0,0 +1,18 @@ +//! Serves SAM 3 on this Mac's Metal GPU to one LightCraft editor over a loopback socket (an SSH +//! tunnel); see docs/ai-masks.md → Remote Metal inference. +#![forbid(unsafe_code)] +#![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)] + +#[cfg(not(target_arch = "wasm32"))] +fn main() -> Result<(), Box> { + let mut args = std::env::args().skip(1); + let dir = args.next().ok_or("usage: lightcraft-sam3-worker MODEL_DIR [127.0.0.1:8793]")?; + let address = args.next().unwrap_or_else(|| "127.0.0.1:8793".into()).parse()?; + let device = candle_core::Device::new_metal(0)?; + lightcraft_segment::remote::serve(address, std::path::Path::new(&dir), device)?; + Ok(()) +} + +/// The worker serves a local GPU over a socket; there is nothing to serve in a browser. +#[cfg(target_arch = "wasm32")] +fn main() {} diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index 42aed538..6d488a91 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -37,6 +37,8 @@ pub mod mask; mod neck; #[cfg(not(target_arch = "wasm32"))] mod nn; +#[cfg(not(target_arch = "wasm32"))] +pub mod remote; pub mod tokenizer; #[cfg(not(target_arch = "wasm32"))] mod tracker; diff --git a/crates/segment/src/remote.rs b/crates/segment/src/remote.rs new file mode 100644 index 00000000..cdee096f --- /dev/null +++ b/crates/segment/src/remote.rs @@ -0,0 +1,207 @@ +//! Private SAM 3 transport. Bind the worker to loopback and forward it over SSH. +use std::io::{self, Read, Write}; +use std::net::{SocketAddr, TcpListener, TcpStream}; +use std::path::Path; +use std::time::Duration; + +use crate::{Click, Encoded, INPUT, MASK_SIDE, Sam3}; +use serde::{Deserialize, Serialize}; + +const CONTROL_LIMIT: usize = 16 * 1024; +const RESPONSE_LIMIT: usize = 4 * 1024 * 1024; +const TIMEOUT: Duration = Duration::from_secs(9 * 60); +pub const ENV: &str = "LIGHTCRAFT_SAM3_REMOTE"; + +#[derive(Serialize, Deserialize)] +pub enum Operation { + Prepare, + Clicks(Vec<[f32; 3]>), + Text(String), +} + +#[derive(Serialize, Deserialize)] +pub struct Request { + pub width: usize, + pub height: usize, + pub temporary: bool, + pub operation: Operation, +} + +impl Request { + pub fn validate(&self) -> Result<(), String> { + if self.width > INPUT || self.height > INPUT || (self.width == 0) != (self.height == 0) { + return Err("invalid SAM image dimensions".into()); + } + if self.temporary && self.width == 0 { + return Err("detail pass requires an image".into()); + } + match &self.operation { + Operation::Prepare => {} + Operation::Text(text) if text.len() <= 4096 && text.split([',', ';']).any(|p| !p.trim().is_empty()) => {} + Operation::Clicks(clicks) + if !clicks.is_empty() + && clicks.len() <= 64 + && clicks.iter().all(|c| { + c.iter().all(|v| v.is_finite()) && (0.0..=1.0).contains(&c[0]) && (0.0..=1.0).contains(&c[1]) && (c[2] == 0.0 || c[2] == 1.0) + }) + && clicks.iter().any(|c| c[2] == 1.0) => {} + _ => return Err("invalid SAM prompt".into()), + } + Ok(()) + } +} + +pub fn configured() -> Option { + std::env::var(ENV).ok().filter(|v| !v.trim().is_empty()) +} + +pub fn read_frame(reader: &mut impl Read, limit: usize) -> io::Result> { + let mut size = [0; 4]; + reader.read_exact(&mut size)?; + let size = u32::from_be_bytes(size) as usize; + if size > limit { + return Err(io::Error::new(io::ErrorKind::InvalidData, "SAM frame exceeds limit")); + } + let mut data = vec![0; size]; + reader.read_exact(&mut data)?; + Ok(data) +} + +fn write_frame(writer: &mut impl Write, data: &[u8]) -> io::Result<()> { + let size = u32::try_from(data.len()).map_err(io::Error::other)?; + writer.write_all(&size.to_be_bytes())?; + writer.write_all(data) +} + +pub struct Client { + stream: TcpStream, + pub key: Option, +} + +impl Client { + pub fn connect(endpoint: &str) -> Result { + let address: SocketAddr = endpoint.parse().map_err(|e| format!("SAM endpoint must be an IP:port: {e}"))?; + if !address.ip().is_loopback() { + return Err("SAM endpoint must be loopback; use an SSH tunnel for remote access".into()); + } + let stream = TcpStream::connect_timeout(&address, Duration::from_secs(5)).map_err(|e| e.to_string())?; + stream.set_read_timeout(Some(TIMEOUT)).map_err(|e| e.to_string())?; + stream.set_write_timeout(Some(Duration::from_secs(30))).map_err(|e| e.to_string())?; + Ok(Self { stream, key: None }) + } + + pub fn request(&mut self, request: &Request, rgb: &[u8]) -> Result>, String> { + request.validate()?; + if rgb.len() != request.width * request.height * 3 { + return Err("SAM RGB buffer size does not match dimensions".into()); + } + let bytes = serde_json::to_vec(request).map_err(|e| e.to_string())?; + write_frame(&mut self.stream, &bytes).map_err(|e| e.to_string())?; + write_frame(&mut self.stream, rgb).map_err(|e| e.to_string())?; + let response = read_frame(&mut self.stream, RESPONSE_LIMIT).map_err(|e| e.to_string())?; + let result: Result>, String> = serde_json::from_slice(&response).map_err(|e| e.to_string())?; + let logits = result?; + if logits.as_ref().is_some_and(|v| v.len() != MASK_SIDE * MASK_SIDE || v.iter().any(|n| !n.is_finite())) { + return Err("invalid SAM mask response".into()); + } + Ok(logits) + } +} + +pub fn text_logits(model: &mut Sam3, enc: &mut Encoded, text: &str) -> Result>, String> { + let mut merged: Option> = None; + for phrase in text.split([',', ';']).map(str::trim).filter(|p| !p.is_empty()) { + if let Some(p) = model.segment_text(enc, phrase, 0.5).map_err(|e| e.to_string())? { + if let Some(m) = &mut merged { + for (a, b) in m.iter_mut().zip(p.logits) { + *a = a.max(b); + } + } else { + merged = Some(p.logits); + } + } + } + Ok(merged) +} + +fn predict(model: &mut Sam3, enc: &mut Encoded, operation: Operation) -> Result>, String> { + match operation { + Operation::Prepare => { + model.prepare_clicks(enc).map_err(|e| e.to_string())?; + Ok(None) + } + Operation::Text(text) => text_logits(model, enc, &text), + Operation::Clicks(points) => { + let clicks: Vec<_> = points.into_iter().map(|p| Click { x: p[0], y: p[1], positive: p[2] == 1.0 }).collect(); + Ok(Some(model.segment_clicks(enc, &clicks).map_err(|e| e.to_string())?.logits)) + } + } +} + +fn serve_client(stream: &mut TcpStream, dir: &Path, device: candle_core::Device) -> Result<(), String> { + stream.set_read_timeout(Some(Duration::from_secs(10 * 60))).map_err(|e| e.to_string())?; + stream.set_write_timeout(Some(Duration::from_secs(30))).map_err(|e| e.to_string())?; + let mut model = None; + let mut cache = None; + loop { + let frame = read_frame(stream, CONTROL_LIMIT).map_err(|e| e.to_string())?; + let request: Request = serde_json::from_slice(&frame).map_err(|e| e.to_string())?; + request.validate()?; + let rgb = read_frame(stream, request.width * request.height * 3).map_err(|e| e.to_string())?; + if rgb.len() != request.width * request.height * 3 { + return Err("invalid RGB payload".into()); + } + let started = std::time::Instant::now(); + let result = (|| { + if model.is_none() { + model = Some(Sam3::load_on(dir, device.clone()).map_err(|e| e.to_string())?); + } + let m = model.as_mut().ok_or("SAM model unavailable")?; + let mut temporary = None; + if request.width > 0 { + let encoded = m.encode(&rgb, request.width, request.height).map_err(|e| e.to_string())?; + if request.temporary { + temporary = Some(encoded); + } else { + cache = Some(encoded); + } + } + let enc = if request.temporary { temporary.as_mut() } else { cache.as_mut() }.ok_or("SAM image cache empty")?; + predict(m, enc, request.operation) + })(); + eprintln!( + "SAM3 Metal request elapsed={:.3}s image={}x{} detail={} success={}", + started.elapsed().as_secs_f64(), + request.width, + request.height, + request.temporary, + result.is_ok() + ); + let response = serde_json::to_vec(&result).map_err(|e| e.to_string())?; + write_frame(stream, &response).map_err(|e| e.to_string())?; + if result.is_err() { + return Err("SAM inference failed; session released".into()); + } + } +} + +pub fn serve(address: SocketAddr, dir: &Path, device: candle_core::Device) -> Result<(), String> { + if !address.ip().is_loopback() || !device.is_metal() { + return Err("SAM worker requires a loopback address and a Metal device".into()); + } + if !crate::is_model_dir(dir) { + return Err("SAM model files are missing".into()); + } + let listener = TcpListener::bind(address).map_err(|e| e.to_string())?; + eprintln!("SAM3 worker listening on {address}; device=Metal; idle unload=600s"); + for stream in listener.incoming() { + let mut stream = stream.map_err(|e| e.to_string())?; + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| serve_client(&mut stream, dir, device.clone()))); + match result { + Ok(Err(error)) => eprintln!("SAM3 session closed: {error}"), + Err(_) => eprintln!("SAM3 session panicked; model released"), + Ok(Ok(())) => {} + } + } + Ok(()) +} diff --git a/crates/segment/tests/remote.rs b/crates/segment/tests/remote.rs new file mode 100644 index 00000000..ee20ea20 --- /dev/null +++ b/crates/segment/tests/remote.rs @@ -0,0 +1,21 @@ +use lightcraft_segment::remote::{Operation, Request, read_frame}; +use std::io::Cursor; + +#[test] +fn rejects_unbounded_and_truncated_frames() { + assert!(read_frame(&mut Cursor::new(u32::MAX.to_be_bytes()), 1024).is_err()); + assert!(read_frame(&mut Cursor::new([0, 0, 0, 3, 1]), 1024).is_err()); +} + +#[test] +fn validates_images_and_prompt_boundaries() { + let request = |width, height, operation| Request { width, height, temporary: false, operation }; + assert!(request(1009, 1, Operation::Prepare).validate().is_err()); + assert!(request(0, 1, Operation::Prepare).validate().is_err()); + assert!(request(0, 0, Operation::Text("dog".into())).validate().is_ok()); + assert!(request(4, 3, Operation::Clicks(vec![[0.5, 0.5, 1.0]])).validate().is_ok()); + assert!(request(4, 3, Operation::Clicks(vec![[1.5, 0.5, 1.0]])).validate().is_err()); + assert!(request(4, 3, Operation::Clicks(vec![[0.5, f32::NAN, 1.0]])).validate().is_err()); + assert!(request(4, 3, Operation::Clicks(vec![[0.5, 0.5, 0.0]])).validate().is_err()); + assert!(request(4, 3, Operation::Text(" ".into())).validate().is_err()); +} diff --git a/crates/tiff/src/tags.rs b/crates/tiff/src/tags.rs index aa24ec87..5771316a 100644 --- a/crates/tiff/src/tags.rs +++ b/crates/tiff/src/tags.rs @@ -172,6 +172,7 @@ pub const OPCODE_LIST_2: u16 = 51009; pub const OPCODE_LIST_3: u16 = 51022; pub const NOISE_PROFILE: u16 = 51041; pub const DEFAULT_USER_CROP: u16 = 51125; +pub const PROFILE_GAIN_TABLE_MAP: u16 = 52525; pub const CALIBRATION_ILLUMINANT_3: u16 = 52529; pub const CAMERA_CALIBRATION_3: u16 = 52530; pub const COLOR_MATRIX_3: u16 = 52531; @@ -179,6 +180,13 @@ pub const FORWARD_MATRIX_3: u16 = 52532; pub const ILLUMINANT_DATA_1: u16 = 52533; pub const ILLUMINANT_DATA_2: u16 = 52534; pub const ILLUMINANT_DATA_3: u16 = 52535; +// DNG 1.6 semantic masks (spec chapter 4, "Semantic Masks") +pub const SEMANTIC_NAME: u16 = 52526; +pub const SEMANTIC_INSTANCE_ID: u16 = 52528; +pub const MASK_SUB_AREA: u16 = 52536; +/// NewSubFileType of a semantic mask IFD. +pub const SUBFILE_SEMANTIC_MASK: u32 = 0x10004; +pub const PROFILE_GAIN_TABLE_MAP_2: u16 = 52544; // --- Compression codes (TIFF 6.0, DNG, and widely used registrations) pub mod compression { @@ -206,4 +214,6 @@ pub mod photometric { pub const YCBCR: u16 = 6; pub const CFA: u16 = 32803; pub const LINEAR_RAW: u16 = 34892; + /// DNG 1.6 semantic mask. + pub const MASK: u16 = 52527; } diff --git a/crates/ui-egui/Cargo.toml b/crates/ui-egui/Cargo.toml index 49bdfc90..d67494b7 100644 --- a/crates/ui-egui/Cargo.toml +++ b/crates/ui-egui/Cargo.toml @@ -25,7 +25,9 @@ serde_json = { workspace = true } log = { workspace = true } [dev-dependencies] +lightcraft-denoise = { workspace = true, features = ["runtime"] } lightcraft-codecs = { workspace = true } +lightcraft-faces = { workspace = true } [build-dependencies] serde_json = { workspace = true } diff --git a/crates/ui-egui/locales/de-formats.json b/crates/ui-egui/locales/de-formats.json index af10202f..12672394 100644 --- a/crates/ui-egui/locales/de-formats.json +++ b/crates/ui-egui/locales/de-formats.json @@ -187,5 +187,24 @@ "Moved; {} photo(s) relinked": "Verschoben; {} Fotos neu verknüpft", "{} label (click more labels to show any of them)": "Markierung {} (weitere Markierungen auswählen, um alle passenden Fotos zu zeigen)", ", {} star{}": ", Sterne: {}{:.0}", - ", {} label": ", Markierung: {}" + ", {} label": ", Markierung: {}", + "Downloads from {host}. Once it has arrived and checked out, LightCraft installs it, {after}.": "Downloads from {host}. Once it has arrived and checked out, LightCraft installs it, {after}.", + "{:.1} MB": "{:.1} MB", + "Running on the processor: {}.": "Running on the processor: {}.", + "{} · {}. Kept in the library's “denoise” folder; the oldest ones that no photo uses go first when the limit is reached.": "{} · {}. Kept in the library's “denoise” folder; the oldest ones that no photo uses go first when the limit is reached.", + "Denoising {name}: {done} of {total} tiles{}": "Denoising {name}: {done} of {total} tiles{}", + "{size} GB": "{size} GB", + "{} of {} from {host}": "{} of {} from {host}", + "{} KB": "{} KB", + "{} MB": "{} MB", + "Making the denoised picture… {}%": "Making the denoised picture… {}%", + "{name} · {terms}": "{name} · {terms}", + "{n} {noun}": "{n} {noun}", + "Trained on: {provenance}": "Trained on: {provenance}", + " · {queued} waiting": " · {queued} waiting", + "Shows the model's terms, then downloads {} from {host}, installs it and starts using it. Nothing else is sent.": "Shows the model's terms, then downloads {} from {host}, installs it and starts using it. Nothing else is sent.", + "In line for AI Denoise ({ahead} ahead)…": "In line for AI Denoise ({ahead} ahead)…", + "Running on {}.": "Running on {}.", + "{file} · {size}": "{file} · {size}", + "• {assumption}": "• {assumption}" } diff --git a/crates/ui-egui/locales/de.json b/crates/ui-egui/locales/de.json index 2a472186..64a485ed 100644 --- a/crates/ui-egui/locales/de.json +++ b/crates/ui-egui/locales/de.json @@ -1797,5 +1797,95 @@ "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", + "AI Denoise": "AI Denoise", + "AI Denoise Model": "AI Denoise Model", + "AI Denoise is off: press Use on a model.": "AI Denoise is off: press Use on a model.", + "AI Denoise needs a model. Install one in Settings.": "AI Denoise needs a model. Install one in Settings.", + "Accept & Download": "Accept & Download", + "Already installed. Remove it before installing a replacement.": "Already installed. Remove it before installing a replacement.", + "Automatic uses the graphics card where it is faster than the processor: both are timed on this computer the first time a photo is made.": "Automatic uses the graphics card where it is faster than the processor: both are timed on this computer the first time a photo is made.", + "Cache": "Cache", + "Checking and installing the model…": "Checking and installing the model…", + "Checking…": "Checking…", + "Clear cache": "Clear cache", + "Denoise Models": "Denoise Models", + "Denoise Settings": "Denoise Settings", + "Denoise all photos with an Amount": "Denoise all photos with an Amount", + "Downloaded and checked. Installing…": "Downloaded and checked. Installing…", + "Failed its last test": "Failed its last test", + "Graphics card": "Graphics card", + "I have read these terms and accept them for my own use": "I have read these terms and accept them for my own use", + "Inspect Denoise Model": "Inspect Denoise Model", + "Install Denoise Model": "Install Denoise Model", + "Install from file…": "Install from file…", + "Install your own ONNX model with denoise-model.json beside it. This build offers no model downloads.": "Install your own ONNX model with denoise-model.json beside it. This build offers no model downloads.", + "Its denoised picture is not made yet, and pictures are not made on their own (Settings, AI Denoise).": "Its denoised picture is not made yet, and pictures are not made on their own (Settings, AI Denoise).", + "LightCraft does not know this model, so it assumed:": "LightCraft does not know this model, so it assumed:", + "Limit": "Limit", + "Make it now": "Make it now", + "Make the denoised picture of the photos I am looking at": "Make the denoised picture of the photos I am looking at", + "Next up for AI Denoise…": "Next up for AI Denoise…", + "Nothing is fetched until you accept. The model is kept on this computer only; LightCraft never uploads or shares it.": "Nothing is fetched until you accept. The model is kept on this computer only; LightCraft never uploads or shares it.", + "Only photos whose AI Denoise amount is above 0 are made. Exports always make theirs if it is missing.": "Only photos whose AI Denoise amount is above 0 are made. Exports always make theirs if it is missing.", + "Open page": "Open page", + "Open the model's page": "Open the model's page", + "Opens the model's own page in your browser, to read about it or get the file yourself.": "Opens the model's own page in your browser, to read about it or get the file yourself.", + "Processor": "Processor", + "Set up AI Denoise…": "Set up AI Denoise…", + "Speed": "Speed", + "Starting…": "Starting…", + "Tested and works on this computer": "Tested and works on this computer", + "The denoise model is installed and in use: the AI Denoise Amount slider under Detail now works": "The denoise model is installed and in use: the AI Denoise Amount slider under Detail now works", + "The graphics card is set up when the next photo is made.": "The graphics card is set up when the next photo is made.", + "The model is kept on this computer only. LightCraft never uploads or shares it.": "The model is kept on this computer only. LightCraft never uploads or shares it.", + "The processor does the work.": "The processor does the work.", + "There is no model file to install": "There is no model file to install", + "This build cannot run denoise models.": "This build cannot run denoise models.", + "This build has nowhere to keep denoise models, so AI Denoise is not available. (The desktop app has.)": "This build has nowhere to keep denoise models, so AI Denoise is not available. (The desktop app has.)", + "This file cannot be used yet": "This file cannot be used yet", + "Tick the box to accept the model's terms first": "Tick the box to accept the model's terms first", + "Try again": "Try again", + "Try the graphics card again": "Try the graphics card again", + "Turn off": "Turn off", + "Unknown licence": "Unknown licence", + "Use": "Use", + "Waiting to start…": "Waiting to start…", + "Works on raw files from cameras with a standard (Bayer) sensor. Other files keep their normal noise reduction.": "Works on raw files from cameras with a standard (Bayer) sensor. Other files keep their normal noise reduction.", + "cached picture": "cached picture", + "cached pictures": "cached pictures", + "commercial use allowed": "commercial use allowed", + "denoiseModel.accept": "denoiseModel.accept", + "in use": "in use", + "installed": "installed", + "makes AI Denoise available": "makes AI Denoise available", + "non-commercial use only": "non-commercial use only", + "photo has its denoised picture.": "photo has its denoised picture.", + "photo waiting": "photo waiting", + "photos have their denoised picture.": "photos have their denoised picture.", + "photos waiting": "photos waiting", + "starts using it for AI Denoise": "starts using it for AI Denoise", + "terms unclear": "terms unclear", + "the graphics card": "the graphics card", + "the graphics card is not used": "the graphics card is not used", + "Cancel Denoise": "Cancel Denoise", + "Cancel Denoise Model Download": "Cancel Denoise Model Download", + "Choose Denoise Model": "Choose Denoise Model", + "Clear Denoise Cache": "Clear Denoise Cache", + "Denoise Model Downloads": "Denoise Model Downloads", + "Denoise Photos": "Denoise Photos", + "Denoise Pump": "Denoise Pump", + "Denoise Status": "Denoise Status", + "Download Denoise Model": "Download Denoise Model", + "Remove Denoise Model": "Remove Denoise Model", + "Test Denoise Model": "Test Denoise Model", + "AI Denoise is enabled on the requesting photo": "AI Denoise is enabled on the requesting photo", + "After installation, AI Denoise will turn on for the requesting photo.": "After installation, AI Denoise will turn on for the requesting photo.", + "Cancel pending action": "Cancel pending action", + "Install a denoise model to enable AI Denoise on this photo": "Install a denoise model to enable AI Denoise on this photo", + "Pending photo action cancelled": "Pending photo action cancelled", + "Pending photo action cancelled because its library or photo changed": "Pending photo action cancelled because its library or photo changed", + "ONNX model": "ONNX model", + "commercial use subject to GPL conditions": "commercial use subject to GPL conditions" } diff --git a/crates/ui-egui/locales/es-formats.json b/crates/ui-egui/locales/es-formats.json index b67d6720..54418c76 100644 --- a/crates/ui-egui/locales/es-formats.json +++ b/crates/ui-egui/locales/es-formats.json @@ -187,5 +187,24 @@ "Moved; {} photo(s) relinked": "Movida; {} foto(s) vinculada(s) de nuevo", "{} label (click more labels to show any of them)": "Etiqueta {} (hacer clic en más etiquetas para mostrar cualquiera de ellas)", ", {} star{}": ", {} estrella{}", - ", {} label": ", etiqueta {}" + ", {} label": ", etiqueta {}", + "Downloads from {host}. Once it has arrived and checked out, LightCraft installs it, {after}.": "Downloads from {host}. Once it has arrived and checked out, LightCraft installs it, {after}.", + "{:.1} MB": "{:.1} MB", + "Running on the processor: {}.": "Running on the processor: {}.", + "{} · {}. Kept in the library's “denoise” folder; the oldest ones that no photo uses go first when the limit is reached.": "{} · {}. Kept in the library's “denoise” folder; the oldest ones that no photo uses go first when the limit is reached.", + "Denoising {name}: {done} of {total} tiles{}": "Denoising {name}: {done} of {total} tiles{}", + "{size} GB": "{size} GB", + "{} of {} from {host}": "{} of {} from {host}", + "{} KB": "{} KB", + "{} MB": "{} MB", + "Making the denoised picture… {}%": "Making the denoised picture… {}%", + "{name} · {terms}": "{name} · {terms}", + "{n} {noun}": "{n} {noun}", + "Trained on: {provenance}": "Trained on: {provenance}", + " · {queued} waiting": " · {queued} waiting", + "Shows the model's terms, then downloads {} from {host}, installs it and starts using it. Nothing else is sent.": "Shows the model's terms, then downloads {} from {host}, installs it and starts using it. Nothing else is sent.", + "In line for AI Denoise ({ahead} ahead)…": "In line for AI Denoise ({ahead} ahead)…", + "Running on {}.": "Running on {}.", + "{file} · {size}": "{file} · {size}", + "• {assumption}": "• {assumption}" } diff --git a/crates/ui-egui/locales/es.json b/crates/ui-egui/locales/es.json index c0986c19..461fa2f6 100644 --- a/crates/ui-egui/locales/es.json +++ b/crates/ui-egui/locales/es.json @@ -1642,5 +1642,95 @@ "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", + "AI Denoise": "AI Denoise", + "AI Denoise Model": "AI Denoise Model", + "AI Denoise is off: press Use on a model.": "AI Denoise is off: press Use on a model.", + "AI Denoise needs a model. Install one in Settings.": "AI Denoise needs a model. Install one in Settings.", + "Accept & Download": "Accept & Download", + "Already installed. Remove it before installing a replacement.": "Already installed. Remove it before installing a replacement.", + "Automatic uses the graphics card where it is faster than the processor: both are timed on this computer the first time a photo is made.": "Automatic uses the graphics card where it is faster than the processor: both are timed on this computer the first time a photo is made.", + "Cache": "Cache", + "Checking and installing the model…": "Checking and installing the model…", + "Checking…": "Checking…", + "Clear cache": "Clear cache", + "Denoise Models": "Denoise Models", + "Denoise Settings": "Denoise Settings", + "Denoise all photos with an Amount": "Denoise all photos with an Amount", + "Downloaded and checked. Installing…": "Downloaded and checked. Installing…", + "Failed its last test": "Failed its last test", + "Graphics card": "Graphics card", + "I have read these terms and accept them for my own use": "I have read these terms and accept them for my own use", + "Inspect Denoise Model": "Inspect Denoise Model", + "Install Denoise Model": "Install Denoise Model", + "Install from file…": "Install from file…", + "Install your own ONNX model with denoise-model.json beside it. This build offers no model downloads.": "Install your own ONNX model with denoise-model.json beside it. This build offers no model downloads.", + "Its denoised picture is not made yet, and pictures are not made on their own (Settings, AI Denoise).": "Its denoised picture is not made yet, and pictures are not made on their own (Settings, AI Denoise).", + "LightCraft does not know this model, so it assumed:": "LightCraft does not know this model, so it assumed:", + "Limit": "Limit", + "Make it now": "Make it now", + "Make the denoised picture of the photos I am looking at": "Make the denoised picture of the photos I am looking at", + "Next up for AI Denoise…": "Next up for AI Denoise…", + "Nothing is fetched until you accept. The model is kept on this computer only; LightCraft never uploads or shares it.": "Nothing is fetched until you accept. The model is kept on this computer only; LightCraft never uploads or shares it.", + "Only photos whose AI Denoise amount is above 0 are made. Exports always make theirs if it is missing.": "Only photos whose AI Denoise amount is above 0 are made. Exports always make theirs if it is missing.", + "Open page": "Open page", + "Open the model's page": "Open the model's page", + "Opens the model's own page in your browser, to read about it or get the file yourself.": "Opens the model's own page in your browser, to read about it or get the file yourself.", + "Processor": "Processor", + "Set up AI Denoise…": "Set up AI Denoise…", + "Speed": "Speed", + "Starting…": "Starting…", + "Tested and works on this computer": "Tested and works on this computer", + "The denoise model is installed and in use: the AI Denoise Amount slider under Detail now works": "The denoise model is installed and in use: the AI Denoise Amount slider under Detail now works", + "The graphics card is set up when the next photo is made.": "The graphics card is set up when the next photo is made.", + "The model is kept on this computer only. LightCraft never uploads or shares it.": "The model is kept on this computer only. LightCraft never uploads or shares it.", + "The processor does the work.": "The processor does the work.", + "There is no model file to install": "There is no model file to install", + "This build cannot run denoise models.": "This build cannot run denoise models.", + "This build has nowhere to keep denoise models, so AI Denoise is not available. (The desktop app has.)": "This build has nowhere to keep denoise models, so AI Denoise is not available. (The desktop app has.)", + "This file cannot be used yet": "This file cannot be used yet", + "Tick the box to accept the model's terms first": "Tick the box to accept the model's terms first", + "Try again": "Try again", + "Try the graphics card again": "Try the graphics card again", + "Turn off": "Turn off", + "Unknown licence": "Unknown licence", + "Use": "Use", + "Waiting to start…": "Waiting to start…", + "Works on raw files from cameras with a standard (Bayer) sensor. Other files keep their normal noise reduction.": "Works on raw files from cameras with a standard (Bayer) sensor. Other files keep their normal noise reduction.", + "cached picture": "cached picture", + "cached pictures": "cached pictures", + "commercial use allowed": "commercial use allowed", + "denoiseModel.accept": "denoiseModel.accept", + "in use": "in use", + "installed": "installed", + "makes AI Denoise available": "makes AI Denoise available", + "non-commercial use only": "non-commercial use only", + "photo has its denoised picture.": "photo has its denoised picture.", + "photo waiting": "photo waiting", + "photos have their denoised picture.": "photos have their denoised picture.", + "photos waiting": "photos waiting", + "starts using it for AI Denoise": "starts using it for AI Denoise", + "terms unclear": "terms unclear", + "the graphics card": "the graphics card", + "the graphics card is not used": "the graphics card is not used", + "Cancel Denoise": "Cancel Denoise", + "Cancel Denoise Model Download": "Cancel Denoise Model Download", + "Choose Denoise Model": "Choose Denoise Model", + "Clear Denoise Cache": "Clear Denoise Cache", + "Denoise Model Downloads": "Denoise Model Downloads", + "Denoise Photos": "Denoise Photos", + "Denoise Pump": "Denoise Pump", + "Denoise Status": "Denoise Status", + "Download Denoise Model": "Download Denoise Model", + "Remove Denoise Model": "Remove Denoise Model", + "Test Denoise Model": "Test Denoise Model", + "AI Denoise is enabled on the requesting photo": "AI Denoise is enabled on the requesting photo", + "After installation, AI Denoise will turn on for the requesting photo.": "After installation, AI Denoise will turn on for the requesting photo.", + "Cancel pending action": "Cancel pending action", + "Install a denoise model to enable AI Denoise on this photo": "Install a denoise model to enable AI Denoise on this photo", + "Pending photo action cancelled": "Pending photo action cancelled", + "Pending photo action cancelled because its library or photo changed": "Pending photo action cancelled because its library or photo changed", + "ONNX model": "ONNX model", + "commercial use subject to GPL conditions": "commercial use subject to GPL conditions" } diff --git a/crates/ui-egui/locales/ja-formats.json b/crates/ui-egui/locales/ja-formats.json index e462871c..b0e0f2f9 100644 --- a/crates/ui-egui/locales/ja-formats.json +++ b/crates/ui-egui/locales/ja-formats.json @@ -187,5 +187,24 @@ "Moved; {} photo(s) relinked": "移動しました。{}枚の写真の参照先を更新しました", "{} label (click more labels to show any of them)": "{}ラベル（ほかのラベルもクリックすると、いずれかに一致する写真を表示）", ", {} star{}": "、星{}つ{:.0}", - ", {} label": "、{}ラベル" + ", {} label": "、{}ラベル", + "Downloads from {host}. Once it has arrived and checked out, LightCraft installs it, {after}.": "Downloads from {host}. Once it has arrived and checked out, LightCraft installs it, {after}.", + "{:.1} MB": "{:.1} MB", + "Running on the processor: {}.": "Running on the processor: {}.", + "{} · {}. Kept in the library's “denoise” folder; the oldest ones that no photo uses go first when the limit is reached.": "{} · {}. Kept in the library's “denoise” folder; the oldest ones that no photo uses go first when the limit is reached.", + "Denoising {name}: {done} of {total} tiles{}": "Denoising {name}: {done} of {total} tiles{}", + "{size} GB": "{size} GB", + "{} of {} from {host}": "{} of {} from {host}", + "{} KB": "{} KB", + "{} MB": "{} MB", + "Making the denoised picture… {}%": "Making the denoised picture… {}%", + "{name} · {terms}": "{name} · {terms}", + "{n} {noun}": "{n} {noun}", + "Trained on: {provenance}": "Trained on: {provenance}", + " · {queued} waiting": " · {queued} waiting", + "Shows the model's terms, then downloads {} from {host}, installs it and starts using it. Nothing else is sent.": "Shows the model's terms, then downloads {} from {host}, installs it and starts using it. Nothing else is sent.", + "In line for AI Denoise ({ahead} ahead)…": "In line for AI Denoise ({ahead} ahead)…", + "Running on {}.": "Running on {}.", + "{file} · {size}": "{file} · {size}", + "• {assumption}": "• {assumption}" } diff --git a/crates/ui-egui/locales/ja.json b/crates/ui-egui/locales/ja.json index 31c61a7a..49f2ef80 100644 --- a/crates/ui-egui/locales/ja.json +++ b/crates/ui-egui/locales/ja.json @@ -1642,5 +1642,95 @@ "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", + "AI Denoise": "AI Denoise", + "AI Denoise Model": "AI Denoise Model", + "AI Denoise is off: press Use on a model.": "AI Denoise is off: press Use on a model.", + "AI Denoise needs a model. Install one in Settings.": "AI Denoise needs a model. Install one in Settings.", + "Accept & Download": "Accept & Download", + "Already installed. Remove it before installing a replacement.": "Already installed. Remove it before installing a replacement.", + "Automatic uses the graphics card where it is faster than the processor: both are timed on this computer the first time a photo is made.": "Automatic uses the graphics card where it is faster than the processor: both are timed on this computer the first time a photo is made.", + "Cache": "Cache", + "Checking and installing the model…": "Checking and installing the model…", + "Checking…": "Checking…", + "Clear cache": "Clear cache", + "Denoise Models": "Denoise Models", + "Denoise Settings": "Denoise Settings", + "Denoise all photos with an Amount": "Denoise all photos with an Amount", + "Downloaded and checked. Installing…": "Downloaded and checked. Installing…", + "Failed its last test": "Failed its last test", + "Graphics card": "Graphics card", + "I have read these terms and accept them for my own use": "I have read these terms and accept them for my own use", + "Inspect Denoise Model": "Inspect Denoise Model", + "Install Denoise Model": "Install Denoise Model", + "Install from file…": "Install from file…", + "Install your own ONNX model with denoise-model.json beside it. This build offers no model downloads.": "Install your own ONNX model with denoise-model.json beside it. This build offers no model downloads.", + "Its denoised picture is not made yet, and pictures are not made on their own (Settings, AI Denoise).": "Its denoised picture is not made yet, and pictures are not made on their own (Settings, AI Denoise).", + "LightCraft does not know this model, so it assumed:": "LightCraft does not know this model, so it assumed:", + "Limit": "Limit", + "Make it now": "Make it now", + "Make the denoised picture of the photos I am looking at": "Make the denoised picture of the photos I am looking at", + "Next up for AI Denoise…": "Next up for AI Denoise…", + "Nothing is fetched until you accept. The model is kept on this computer only; LightCraft never uploads or shares it.": "Nothing is fetched until you accept. The model is kept on this computer only; LightCraft never uploads or shares it.", + "Only photos whose AI Denoise amount is above 0 are made. Exports always make theirs if it is missing.": "Only photos whose AI Denoise amount is above 0 are made. Exports always make theirs if it is missing.", + "Open page": "Open page", + "Open the model's page": "Open the model's page", + "Opens the model's own page in your browser, to read about it or get the file yourself.": "Opens the model's own page in your browser, to read about it or get the file yourself.", + "Processor": "Processor", + "Set up AI Denoise…": "Set up AI Denoise…", + "Speed": "Speed", + "Starting…": "Starting…", + "Tested and works on this computer": "Tested and works on this computer", + "The denoise model is installed and in use: the AI Denoise Amount slider under Detail now works": "The denoise model is installed and in use: the AI Denoise Amount slider under Detail now works", + "The graphics card is set up when the next photo is made.": "The graphics card is set up when the next photo is made.", + "The model is kept on this computer only. LightCraft never uploads or shares it.": "The model is kept on this computer only. LightCraft never uploads or shares it.", + "The processor does the work.": "The processor does the work.", + "There is no model file to install": "There is no model file to install", + "This build cannot run denoise models.": "This build cannot run denoise models.", + "This build has nowhere to keep denoise models, so AI Denoise is not available. (The desktop app has.)": "This build has nowhere to keep denoise models, so AI Denoise is not available. (The desktop app has.)", + "This file cannot be used yet": "This file cannot be used yet", + "Tick the box to accept the model's terms first": "Tick the box to accept the model's terms first", + "Try again": "Try again", + "Try the graphics card again": "Try the graphics card again", + "Turn off": "Turn off", + "Unknown licence": "Unknown licence", + "Use": "Use", + "Waiting to start…": "Waiting to start…", + "Works on raw files from cameras with a standard (Bayer) sensor. Other files keep their normal noise reduction.": "Works on raw files from cameras with a standard (Bayer) sensor. Other files keep their normal noise reduction.", + "cached picture": "cached picture", + "cached pictures": "cached pictures", + "commercial use allowed": "commercial use allowed", + "denoiseModel.accept": "denoiseModel.accept", + "in use": "in use", + "installed": "installed", + "makes AI Denoise available": "makes AI Denoise available", + "non-commercial use only": "non-commercial use only", + "photo has its denoised picture.": "photo has its denoised picture.", + "photo waiting": "photo waiting", + "photos have their denoised picture.": "photos have their denoised picture.", + "photos waiting": "photos waiting", + "starts using it for AI Denoise": "starts using it for AI Denoise", + "terms unclear": "terms unclear", + "the graphics card": "the graphics card", + "the graphics card is not used": "the graphics card is not used", + "Cancel Denoise": "Cancel Denoise", + "Cancel Denoise Model Download": "Cancel Denoise Model Download", + "Choose Denoise Model": "Choose Denoise Model", + "Clear Denoise Cache": "Clear Denoise Cache", + "Denoise Model Downloads": "Denoise Model Downloads", + "Denoise Photos": "Denoise Photos", + "Denoise Pump": "Denoise Pump", + "Denoise Status": "Denoise Status", + "Download Denoise Model": "Download Denoise Model", + "Remove Denoise Model": "Remove Denoise Model", + "Test Denoise Model": "Test Denoise Model", + "AI Denoise is enabled on the requesting photo": "AI Denoise is enabled on the requesting photo", + "After installation, AI Denoise will turn on for the requesting photo.": "After installation, AI Denoise will turn on for the requesting photo.", + "Cancel pending action": "Cancel pending action", + "Install a denoise model to enable AI Denoise on this photo": "Install a denoise model to enable AI Denoise on this photo", + "Pending photo action cancelled": "Pending photo action cancelled", + "Pending photo action cancelled because its library or photo changed": "Pending photo action cancelled because its library or photo changed", + "ONNX model": "ONNX model", + "commercial use subject to GPL conditions": "commercial use subject to GPL conditions" } diff --git a/crates/ui-egui/locales/pt-br-formats.json b/crates/ui-egui/locales/pt-br-formats.json index 59bec59d..ab64e082 100644 --- a/crates/ui-egui/locales/pt-br-formats.json +++ b/crates/ui-egui/locales/pt-br-formats.json @@ -187,5 +187,24 @@ "Moved; {} photo(s) relinked": "Movida; {} foto(s) vinculada(s) novamente", "{} label (click more labels to show any of them)": "Rótulo {} (clique em mais rótulos para mostrar qualquer um deles)", ", {} star{}": ", {} estrela{}", - ", {} label": ", rótulo {}" + ", {} label": ", rótulo {}", + "Downloads from {host}. Once it has arrived and checked out, LightCraft installs it, {after}.": "Downloads from {host}. Once it has arrived and checked out, LightCraft installs it, {after}.", + "{:.1} MB": "{:.1} MB", + "Running on the processor: {}.": "Running on the processor: {}.", + "{} · {}. Kept in the library's “denoise” folder; the oldest ones that no photo uses go first when the limit is reached.": "{} · {}. Kept in the library's “denoise” folder; the oldest ones that no photo uses go first when the limit is reached.", + "Denoising {name}: {done} of {total} tiles{}": "Denoising {name}: {done} of {total} tiles{}", + "{size} GB": "{size} GB", + "{} of {} from {host}": "{} of {} from {host}", + "{} KB": "{} KB", + "{} MB": "{} MB", + "Making the denoised picture… {}%": "Making the denoised picture… {}%", + "{name} · {terms}": "{name} · {terms}", + "{n} {noun}": "{n} {noun}", + "Trained on: {provenance}": "Trained on: {provenance}", + " · {queued} waiting": " · {queued} waiting", + "Shows the model's terms, then downloads {} from {host}, installs it and starts using it. Nothing else is sent.": "Shows the model's terms, then downloads {} from {host}, installs it and starts using it. Nothing else is sent.", + "In line for AI Denoise ({ahead} ahead)…": "In line for AI Denoise ({ahead} ahead)…", + "Running on {}.": "Running on {}.", + "{file} · {size}": "{file} · {size}", + "• {assumption}": "• {assumption}" } diff --git a/crates/ui-egui/locales/pt-br.json b/crates/ui-egui/locales/pt-br.json index 03ae0773..81e331c5 100644 --- a/crates/ui-egui/locales/pt-br.json +++ b/crates/ui-egui/locales/pt-br.json @@ -1642,5 +1642,95 @@ "≤ at most": "≤ no máximo", "≥ at least": "≥ no mínimo", "日本語": "日本語", - "繁體中文（台灣）": "繁體中文（台灣）" + "繁體中文（台灣）": "繁體中文（台灣）", + "Back to this person": "Back to this person", + "AI Denoise": "AI Denoise", + "AI Denoise Model": "AI Denoise Model", + "AI Denoise is off: press Use on a model.": "AI Denoise is off: press Use on a model.", + "AI Denoise needs a model. Install one in Settings.": "AI Denoise needs a model. Install one in Settings.", + "Accept & Download": "Accept & Download", + "Already installed. Remove it before installing a replacement.": "Already installed. Remove it before installing a replacement.", + "Automatic uses the graphics card where it is faster than the processor: both are timed on this computer the first time a photo is made.": "Automatic uses the graphics card where it is faster than the processor: both are timed on this computer the first time a photo is made.", + "Cache": "Cache", + "Checking and installing the model…": "Checking and installing the model…", + "Checking…": "Checking…", + "Clear cache": "Clear cache", + "Denoise Models": "Denoise Models", + "Denoise Settings": "Denoise Settings", + "Denoise all photos with an Amount": "Denoise all photos with an Amount", + "Downloaded and checked. Installing…": "Downloaded and checked. Installing…", + "Failed its last test": "Failed its last test", + "Graphics card": "Graphics card", + "I have read these terms and accept them for my own use": "I have read these terms and accept them for my own use", + "Inspect Denoise Model": "Inspect Denoise Model", + "Install Denoise Model": "Install Denoise Model", + "Install from file…": "Install from file…", + "Install your own ONNX model with denoise-model.json beside it. This build offers no model downloads.": "Install your own ONNX model with denoise-model.json beside it. This build offers no model downloads.", + "Its denoised picture is not made yet, and pictures are not made on their own (Settings, AI Denoise).": "Its denoised picture is not made yet, and pictures are not made on their own (Settings, AI Denoise).", + "LightCraft does not know this model, so it assumed:": "LightCraft does not know this model, so it assumed:", + "Limit": "Limit", + "Make it now": "Make it now", + "Make the denoised picture of the photos I am looking at": "Make the denoised picture of the photos I am looking at", + "Next up for AI Denoise…": "Next up for AI Denoise…", + "Nothing is fetched until you accept. The model is kept on this computer only; LightCraft never uploads or shares it.": "Nothing is fetched until you accept. The model is kept on this computer only; LightCraft never uploads or shares it.", + "Only photos whose AI Denoise amount is above 0 are made. Exports always make theirs if it is missing.": "Only photos whose AI Denoise amount is above 0 are made. Exports always make theirs if it is missing.", + "Open page": "Open page", + "Open the model's page": "Open the model's page", + "Opens the model's own page in your browser, to read about it or get the file yourself.": "Opens the model's own page in your browser, to read about it or get the file yourself.", + "Processor": "Processor", + "Set up AI Denoise…": "Set up AI Denoise…", + "Speed": "Speed", + "Starting…": "Starting…", + "Tested and works on this computer": "Tested and works on this computer", + "The denoise model is installed and in use: the AI Denoise Amount slider under Detail now works": "The denoise model is installed and in use: the AI Denoise Amount slider under Detail now works", + "The graphics card is set up when the next photo is made.": "The graphics card is set up when the next photo is made.", + "The model is kept on this computer only. LightCraft never uploads or shares it.": "The model is kept on this computer only. LightCraft never uploads or shares it.", + "The processor does the work.": "The processor does the work.", + "There is no model file to install": "There is no model file to install", + "This build cannot run denoise models.": "This build cannot run denoise models.", + "This build has nowhere to keep denoise models, so AI Denoise is not available. (The desktop app has.)": "This build has nowhere to keep denoise models, so AI Denoise is not available. (The desktop app has.)", + "This file cannot be used yet": "This file cannot be used yet", + "Tick the box to accept the model's terms first": "Tick the box to accept the model's terms first", + "Try again": "Try again", + "Try the graphics card again": "Try the graphics card again", + "Turn off": "Turn off", + "Unknown licence": "Unknown licence", + "Use": "Use", + "Waiting to start…": "Waiting to start…", + "Works on raw files from cameras with a standard (Bayer) sensor. Other files keep their normal noise reduction.": "Works on raw files from cameras with a standard (Bayer) sensor. Other files keep their normal noise reduction.", + "cached picture": "cached picture", + "cached pictures": "cached pictures", + "commercial use allowed": "commercial use allowed", + "denoiseModel.accept": "denoiseModel.accept", + "in use": "in use", + "installed": "installed", + "makes AI Denoise available": "makes AI Denoise available", + "non-commercial use only": "non-commercial use only", + "photo has its denoised picture.": "photo has its denoised picture.", + "photo waiting": "photo waiting", + "photos have their denoised picture.": "photos have their denoised picture.", + "photos waiting": "photos waiting", + "starts using it for AI Denoise": "starts using it for AI Denoise", + "terms unclear": "terms unclear", + "the graphics card": "the graphics card", + "the graphics card is not used": "the graphics card is not used", + "Cancel Denoise": "Cancel Denoise", + "Cancel Denoise Model Download": "Cancel Denoise Model Download", + "Choose Denoise Model": "Choose Denoise Model", + "Clear Denoise Cache": "Clear Denoise Cache", + "Denoise Model Downloads": "Denoise Model Downloads", + "Denoise Photos": "Denoise Photos", + "Denoise Pump": "Denoise Pump", + "Denoise Status": "Denoise Status", + "Download Denoise Model": "Download Denoise Model", + "Remove Denoise Model": "Remove Denoise Model", + "Test Denoise Model": "Test Denoise Model", + "AI Denoise is enabled on the requesting photo": "AI Denoise is enabled on the requesting photo", + "After installation, AI Denoise will turn on for the requesting photo.": "After installation, AI Denoise will turn on for the requesting photo.", + "Cancel pending action": "Cancel pending action", + "Install a denoise model to enable AI Denoise on this photo": "Install a denoise model to enable AI Denoise on this photo", + "Pending photo action cancelled": "Pending photo action cancelled", + "Pending photo action cancelled because its library or photo changed": "Pending photo action cancelled because its library or photo changed", + "ONNX model": "ONNX model", + "commercial use subject to GPL conditions": "commercial use subject to GPL conditions" } diff --git a/crates/ui-egui/locales/ru-formats.json b/crates/ui-egui/locales/ru-formats.json index 2beb4a5b..f3a0f647 100644 --- a/crates/ui-egui/locales/ru-formats.json +++ b/crates/ui-egui/locales/ru-formats.json @@ -187,5 +187,24 @@ "{} — show their photos": "{} — показать фото этого человека", "· {dups} already in the library": "· уже в библиотеке: {dups}", "· {n} photo{}": "· фото: {n}{:.0}", - "… and {more} more": "… и ещё {more}" + "… and {more} more": "… и ещё {more}", + "Downloads from {host}. Once it has arrived and checked out, LightCraft installs it, {after}.": "Downloads from {host}. Once it has arrived and checked out, LightCraft installs it, {after}.", + "{:.1} MB": "{:.1} MB", + "Running on the processor: {}.": "Running on the processor: {}.", + "{} · {}. Kept in the library's “denoise” folder; the oldest ones that no photo uses go first when the limit is reached.": "{} · {}. Kept in the library's “denoise” folder; the oldest ones that no photo uses go first when the limit is reached.", + "Denoising {name}: {done} of {total} tiles{}": "Denoising {name}: {done} of {total} tiles{}", + "{size} GB": "{size} GB", + "{} of {} from {host}": "{} of {} from {host}", + "{} KB": "{} KB", + "{} MB": "{} MB", + "Making the denoised picture… {}%": "Making the denoised picture… {}%", + "{name} · {terms}": "{name} · {terms}", + "{n} {noun}": "{n} {noun}", + "Trained on: {provenance}": "Trained on: {provenance}", + " · {queued} waiting": " · {queued} waiting", + "Shows the model's terms, then downloads {} from {host}, installs it and starts using it. Nothing else is sent.": "Shows the model's terms, then downloads {} from {host}, installs it and starts using it. Nothing else is sent.", + "In line for AI Denoise ({ahead} ahead)…": "In line for AI Denoise ({ahead} ahead)…", + "Running on {}.": "Running on {}.", + "{file} · {size}": "{file} · {size}", + "• {assumption}": "• {assumption}" } diff --git a/crates/ui-egui/locales/ru.json b/crates/ui-egui/locales/ru.json index 530cdfcb..382bfc05 100644 --- a/crates/ui-egui/locales/ru.json +++ b/crates/ui-egui/locales/ru.json @@ -1312,5 +1312,96 @@ "of the following rules:": "из следующих условий:", "saved in": "сохранено в", "this photo": "это фото", - "− harder border · + softer border": "− жёстче край · + мягче край" + "− harder border · + softer border": "− жёстче край · + мягче край", + "Back to this person": "Back to this person", + "AI Denoise": "AI Denoise", + "AI Denoise Model": "AI Denoise Model", + "AI Denoise is off: press Use on a model.": "AI Denoise is off: press Use on a model.", + "AI Denoise needs a model. Install one in Settings.": "AI Denoise needs a model. Install one in Settings.", + "Accept & Download": "Accept & Download", + "Already installed. Remove it before installing a replacement.": "Already installed. Remove it before installing a replacement.", + "Automatic uses the graphics card where it is faster than the processor: both are timed on this computer the first time a photo is made.": "Automatic uses the graphics card where it is faster than the processor: both are timed on this computer the first time a photo is made.", + "Cache": "Cache", + "Checking and installing the model…": "Checking and installing the model…", + "Checking…": "Checking…", + "Clear cache": "Clear cache", + "Denoise Models": "Denoise Models", + "Denoise Settings": "Denoise Settings", + "Denoise all photos with an Amount": "Denoise all photos with an Amount", + "Downloaded and checked. Installing…": "Downloaded and checked. Installing…", + "Failed its last test": "Failed its last test", + "Graphics card": "Graphics card", + "I have read these terms and accept them for my own use": "I have read these terms and accept them for my own use", + "Inspect Denoise Model": "Inspect Denoise Model", + "Install Denoise Model": "Install Denoise Model", + "Install from file…": "Install from file…", + "Install your own ONNX model with denoise-model.json beside it. This build offers no model downloads.": "Install your own ONNX model with denoise-model.json beside it. This build offers no model downloads.", + "Its denoised picture is not made yet, and pictures are not made on their own (Settings, AI Denoise).": "Its denoised picture is not made yet, and pictures are not made on their own (Settings, AI Denoise).", + "LightCraft does not know this model, so it assumed:": "LightCraft does not know this model, so it assumed:", + "Limit": "Limit", + "Make it now": "Make it now", + "Make the denoised picture of the photos I am looking at": "Make the denoised picture of the photos I am looking at", + "Models": "Models", + "Next up for AI Denoise…": "Next up for AI Denoise…", + "Nothing is fetched until you accept. The model is kept on this computer only; LightCraft never uploads or shares it.": "Nothing is fetched until you accept. The model is kept on this computer only; LightCraft never uploads or shares it.", + "Only photos whose AI Denoise amount is above 0 are made. Exports always make theirs if it is missing.": "Only photos whose AI Denoise amount is above 0 are made. Exports always make theirs if it is missing.", + "Open page": "Open page", + "Open the model's page": "Open the model's page", + "Opens the model's own page in your browser, to read about it or get the file yourself.": "Opens the model's own page in your browser, to read about it or get the file yourself.", + "Processor": "Processor", + "Set up AI Denoise…": "Set up AI Denoise…", + "Speed": "Speed", + "Starting…": "Starting…", + "Tested and works on this computer": "Tested and works on this computer", + "The denoise model is installed and in use: the AI Denoise Amount slider under Detail now works": "The denoise model is installed and in use: the AI Denoise Amount slider under Detail now works", + "The graphics card is set up when the next photo is made.": "The graphics card is set up when the next photo is made.", + "The model is kept on this computer only. LightCraft never uploads or shares it.": "The model is kept on this computer only. LightCraft never uploads or shares it.", + "The processor does the work.": "The processor does the work.", + "There is no model file to install": "There is no model file to install", + "This build cannot run denoise models.": "This build cannot run denoise models.", + "This build has nowhere to keep denoise models, so AI Denoise is not available. (The desktop app has.)": "This build has nowhere to keep denoise models, so AI Denoise is not available. (The desktop app has.)", + "This file cannot be used yet": "This file cannot be used yet", + "Tick the box to accept the model's terms first": "Tick the box to accept the model's terms first", + "Try again": "Try again", + "Try the graphics card again": "Try the graphics card again", + "Turn off": "Turn off", + "Unknown licence": "Unknown licence", + "Use": "Use", + "Waiting to start…": "Waiting to start…", + "Works on raw files from cameras with a standard (Bayer) sensor. Other files keep their normal noise reduction.": "Works on raw files from cameras with a standard (Bayer) sensor. Other files keep their normal noise reduction.", + "cached picture": "cached picture", + "cached pictures": "cached pictures", + "commercial use allowed": "commercial use allowed", + "denoiseModel.accept": "denoiseModel.accept", + "in use": "in use", + "installed": "installed", + "makes AI Denoise available": "makes AI Denoise available", + "non-commercial use only": "non-commercial use only", + "photo has its denoised picture.": "photo has its denoised picture.", + "photo waiting": "photo waiting", + "photos have their denoised picture.": "photos have their denoised picture.", + "photos waiting": "photos waiting", + "starts using it for AI Denoise": "starts using it for AI Denoise", + "terms unclear": "terms unclear", + "the graphics card": "the graphics card", + "the graphics card is not used": "the graphics card is not used", + "Cancel Denoise": "Cancel Denoise", + "Cancel Denoise Model Download": "Cancel Denoise Model Download", + "Choose Denoise Model": "Choose Denoise Model", + "Clear Denoise Cache": "Clear Denoise Cache", + "Denoise Model Downloads": "Denoise Model Downloads", + "Denoise Photos": "Denoise Photos", + "Denoise Pump": "Denoise Pump", + "Denoise Status": "Denoise Status", + "Download Denoise Model": "Download Denoise Model", + "Remove Denoise Model": "Remove Denoise Model", + "Test Denoise Model": "Test Denoise Model", + "AI Denoise is enabled on the requesting photo": "AI Denoise is enabled on the requesting photo", + "After installation, AI Denoise will turn on for the requesting photo.": "After installation, AI Denoise will turn on for the requesting photo.", + "Cancel pending action": "Cancel pending action", + "Install a denoise model to enable AI Denoise on this photo": "Install a denoise model to enable AI Denoise on this photo", + "Pending photo action cancelled": "Pending photo action cancelled", + "Pending photo action cancelled because its library or photo changed": "Pending photo action cancelled because its library or photo changed", + "ONNX model": "ONNX model", + "commercial use subject to GPL conditions": "commercial use subject to GPL conditions" } diff --git a/crates/ui-egui/locales/zh-hans-formats.json b/crates/ui-egui/locales/zh-hans-formats.json index eca40f8e..e32c7f66 100644 --- a/crates/ui-egui/locales/zh-hans-formats.json +++ b/crates/ui-egui/locales/zh-hans-formats.json @@ -187,5 +187,24 @@ "Moved; {} photo(s) relinked": "已移动；已重新链接 {0} 张照片", "{} label (click more labels to show any of them)": "{0}标签（单击更多标签可显示符合其中任一标签的照片）", ", {} star{}": "，{0} 星{1:.0}", - ", {} label": "，{0}标签" + ", {} label": "，{0}标签", + "Downloads from {host}. Once it has arrived and checked out, LightCraft installs it, {after}.": "Downloads from {host}. Once it has arrived and checked out, LightCraft installs it, {after}.", + "{:.1} MB": "{:.1} MB", + "Running on the processor: {}.": "Running on the processor: {}.", + "{} · {}. Kept in the library's “denoise” folder; the oldest ones that no photo uses go first when the limit is reached.": "{} · {}. Kept in the library's “denoise” folder; the oldest ones that no photo uses go first when the limit is reached.", + "Denoising {name}: {done} of {total} tiles{}": "Denoising {name}: {done} of {total} tiles{}", + "{size} GB": "{size} GB", + "{} of {} from {host}": "{} of {} from {host}", + "{} KB": "{} KB", + "{} MB": "{} MB", + "Making the denoised picture… {}%": "Making the denoised picture… {}%", + "{name} · {terms}": "{name} · {terms}", + "{n} {noun}": "{n} {noun}", + "Trained on: {provenance}": "Trained on: {provenance}", + " · {queued} waiting": " · {queued} waiting", + "Shows the model's terms, then downloads {} from {host}, installs it and starts using it. Nothing else is sent.": "Shows the model's terms, then downloads {} from {host}, installs it and starts using it. Nothing else is sent.", + "In line for AI Denoise ({ahead} ahead)…": "In line for AI Denoise ({ahead} ahead)…", + "Running on {}.": "Running on {}.", + "{file} · {size}": "{file} · {size}", + "• {assumption}": "• {assumption}" } diff --git a/crates/ui-egui/locales/zh-hans.json b/crates/ui-egui/locales/zh-hans.json index cea313d0..4074dfbd 100644 --- a/crates/ui-egui/locales/zh-hans.json +++ b/crates/ui-egui/locales/zh-hans.json @@ -1642,5 +1642,95 @@ "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", + "AI Denoise": "AI Denoise", + "AI Denoise Model": "AI Denoise Model", + "AI Denoise is off: press Use on a model.": "AI Denoise is off: press Use on a model.", + "AI Denoise needs a model. Install one in Settings.": "AI Denoise needs a model. Install one in Settings.", + "Accept & Download": "Accept & Download", + "Already installed. Remove it before installing a replacement.": "Already installed. Remove it before installing a replacement.", + "Automatic uses the graphics card where it is faster than the processor: both are timed on this computer the first time a photo is made.": "Automatic uses the graphics card where it is faster than the processor: both are timed on this computer the first time a photo is made.", + "Cache": "Cache", + "Checking and installing the model…": "Checking and installing the model…", + "Checking…": "Checking…", + "Clear cache": "Clear cache", + "Denoise Models": "Denoise Models", + "Denoise Settings": "Denoise Settings", + "Denoise all photos with an Amount": "Denoise all photos with an Amount", + "Downloaded and checked. Installing…": "Downloaded and checked. Installing…", + "Failed its last test": "Failed its last test", + "Graphics card": "Graphics card", + "I have read these terms and accept them for my own use": "I have read these terms and accept them for my own use", + "Inspect Denoise Model": "Inspect Denoise Model", + "Install Denoise Model": "Install Denoise Model", + "Install from file…": "Install from file…", + "Install your own ONNX model with denoise-model.json beside it. This build offers no model downloads.": "Install your own ONNX model with denoise-model.json beside it. This build offers no model downloads.", + "Its denoised picture is not made yet, and pictures are not made on their own (Settings, AI Denoise).": "Its denoised picture is not made yet, and pictures are not made on their own (Settings, AI Denoise).", + "LightCraft does not know this model, so it assumed:": "LightCraft does not know this model, so it assumed:", + "Limit": "Limit", + "Make it now": "Make it now", + "Make the denoised picture of the photos I am looking at": "Make the denoised picture of the photos I am looking at", + "Next up for AI Denoise…": "Next up for AI Denoise…", + "Nothing is fetched until you accept. The model is kept on this computer only; LightCraft never uploads or shares it.": "Nothing is fetched until you accept. The model is kept on this computer only; LightCraft never uploads or shares it.", + "Only photos whose AI Denoise amount is above 0 are made. Exports always make theirs if it is missing.": "Only photos whose AI Denoise amount is above 0 are made. Exports always make theirs if it is missing.", + "Open page": "Open page", + "Open the model's page": "Open the model's page", + "Opens the model's own page in your browser, to read about it or get the file yourself.": "Opens the model's own page in your browser, to read about it or get the file yourself.", + "Processor": "Processor", + "Set up AI Denoise…": "Set up AI Denoise…", + "Speed": "Speed", + "Starting…": "Starting…", + "Tested and works on this computer": "Tested and works on this computer", + "The denoise model is installed and in use: the AI Denoise Amount slider under Detail now works": "The denoise model is installed and in use: the AI Denoise Amount slider under Detail now works", + "The graphics card is set up when the next photo is made.": "The graphics card is set up when the next photo is made.", + "The model is kept on this computer only. LightCraft never uploads or shares it.": "The model is kept on this computer only. LightCraft never uploads or shares it.", + "The processor does the work.": "The processor does the work.", + "There is no model file to install": "There is no model file to install", + "This build cannot run denoise models.": "This build cannot run denoise models.", + "This build has nowhere to keep denoise models, so AI Denoise is not available. (The desktop app has.)": "This build has nowhere to keep denoise models, so AI Denoise is not available. (The desktop app has.)", + "This file cannot be used yet": "This file cannot be used yet", + "Tick the box to accept the model's terms first": "Tick the box to accept the model's terms first", + "Try again": "Try again", + "Try the graphics card again": "Try the graphics card again", + "Turn off": "Turn off", + "Unknown licence": "Unknown licence", + "Use": "Use", + "Waiting to start…": "Waiting to start…", + "Works on raw files from cameras with a standard (Bayer) sensor. Other files keep their normal noise reduction.": "Works on raw files from cameras with a standard (Bayer) sensor. Other files keep their normal noise reduction.", + "cached picture": "cached picture", + "cached pictures": "cached pictures", + "commercial use allowed": "commercial use allowed", + "denoiseModel.accept": "denoiseModel.accept", + "in use": "in use", + "installed": "installed", + "makes AI Denoise available": "makes AI Denoise available", + "non-commercial use only": "non-commercial use only", + "photo has its denoised picture.": "photo has its denoised picture.", + "photo waiting": "photo waiting", + "photos have their denoised picture.": "photos have their denoised picture.", + "photos waiting": "photos waiting", + "starts using it for AI Denoise": "starts using it for AI Denoise", + "terms unclear": "terms unclear", + "the graphics card": "the graphics card", + "the graphics card is not used": "the graphics card is not used", + "Cancel Denoise": "Cancel Denoise", + "Cancel Denoise Model Download": "Cancel Denoise Model Download", + "Choose Denoise Model": "Choose Denoise Model", + "Clear Denoise Cache": "Clear Denoise Cache", + "Denoise Model Downloads": "Denoise Model Downloads", + "Denoise Photos": "Denoise Photos", + "Denoise Pump": "Denoise Pump", + "Denoise Status": "Denoise Status", + "Download Denoise Model": "Download Denoise Model", + "Remove Denoise Model": "Remove Denoise Model", + "Test Denoise Model": "Test Denoise Model", + "AI Denoise is enabled on the requesting photo": "AI Denoise is enabled on the requesting photo", + "After installation, AI Denoise will turn on for the requesting photo.": "After installation, AI Denoise will turn on for the requesting photo.", + "Cancel pending action": "Cancel pending action", + "Install a denoise model to enable AI Denoise on this photo": "Install a denoise model to enable AI Denoise on this photo", + "Pending photo action cancelled": "Pending photo action cancelled", + "Pending photo action cancelled because its library or photo changed": "Pending photo action cancelled because its library or photo changed", + "ONNX model": "ONNX model", + "commercial use subject to GPL conditions": "commercial use subject to GPL conditions" } diff --git a/crates/ui-egui/locales/zh-hant-formats.json b/crates/ui-egui/locales/zh-hant-formats.json index 640e3eb7..182a7dc6 100644 --- a/crates/ui-egui/locales/zh-hant-formats.json +++ b/crates/ui-egui/locales/zh-hant-formats.json @@ -187,5 +187,24 @@ "Moved; {} photo(s) relinked": "已移動；已重新連結 {} 張照片", "{} label (click more labels to show any of them)": "{}標籤（點擊更多標籤可顯示符合其中任一標籤的照片）", ", {} star{}": "，{} 顆星{:.0}", - ", {} label": "，{}標籤" + ", {} label": "，{}標籤", + "Downloads from {host}. Once it has arrived and checked out, LightCraft installs it, {after}.": "Downloads from {host}. Once it has arrived and checked out, LightCraft installs it, {after}.", + "{:.1} MB": "{:.1} MB", + "Running on the processor: {}.": "Running on the processor: {}.", + "{} · {}. Kept in the library's “denoise” folder; the oldest ones that no photo uses go first when the limit is reached.": "{} · {}. Kept in the library's “denoise” folder; the oldest ones that no photo uses go first when the limit is reached.", + "Denoising {name}: {done} of {total} tiles{}": "Denoising {name}: {done} of {total} tiles{}", + "{size} GB": "{size} GB", + "{} of {} from {host}": "{} of {} from {host}", + "{} KB": "{} KB", + "{} MB": "{} MB", + "Making the denoised picture… {}%": "Making the denoised picture… {}%", + "{name} · {terms}": "{name} · {terms}", + "{n} {noun}": "{n} {noun}", + "Trained on: {provenance}": "Trained on: {provenance}", + " · {queued} waiting": " · {queued} waiting", + "Shows the model's terms, then downloads {} from {host}, installs it and starts using it. Nothing else is sent.": "Shows the model's terms, then downloads {} from {host}, installs it and starts using it. Nothing else is sent.", + "In line for AI Denoise ({ahead} ahead)…": "In line for AI Denoise ({ahead} ahead)…", + "Running on {}.": "Running on {}.", + "{file} · {size}": "{file} · {size}", + "• {assumption}": "• {assumption}" } diff --git a/crates/ui-egui/locales/zh-hant.json b/crates/ui-egui/locales/zh-hant.json index 597ecbf9..96fcbb18 100644 --- a/crates/ui-egui/locales/zh-hant.json +++ b/crates/ui-egui/locales/zh-hant.json @@ -1642,5 +1642,95 @@ "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", + "AI Denoise": "AI Denoise", + "AI Denoise Model": "AI Denoise Model", + "AI Denoise is off: press Use on a model.": "AI Denoise is off: press Use on a model.", + "AI Denoise needs a model. Install one in Settings.": "AI Denoise needs a model. Install one in Settings.", + "Accept & Download": "Accept & Download", + "Already installed. Remove it before installing a replacement.": "Already installed. Remove it before installing a replacement.", + "Automatic uses the graphics card where it is faster than the processor: both are timed on this computer the first time a photo is made.": "Automatic uses the graphics card where it is faster than the processor: both are timed on this computer the first time a photo is made.", + "Cache": "Cache", + "Checking and installing the model…": "Checking and installing the model…", + "Checking…": "Checking…", + "Clear cache": "Clear cache", + "Denoise Models": "Denoise Models", + "Denoise Settings": "Denoise Settings", + "Denoise all photos with an Amount": "Denoise all photos with an Amount", + "Downloaded and checked. Installing…": "Downloaded and checked. Installing…", + "Failed its last test": "Failed its last test", + "Graphics card": "Graphics card", + "I have read these terms and accept them for my own use": "I have read these terms and accept them for my own use", + "Inspect Denoise Model": "Inspect Denoise Model", + "Install Denoise Model": "Install Denoise Model", + "Install from file…": "Install from file…", + "Install your own ONNX model with denoise-model.json beside it. This build offers no model downloads.": "Install your own ONNX model with denoise-model.json beside it. This build offers no model downloads.", + "Its denoised picture is not made yet, and pictures are not made on their own (Settings, AI Denoise).": "Its denoised picture is not made yet, and pictures are not made on their own (Settings, AI Denoise).", + "LightCraft does not know this model, so it assumed:": "LightCraft does not know this model, so it assumed:", + "Limit": "Limit", + "Make it now": "Make it now", + "Make the denoised picture of the photos I am looking at": "Make the denoised picture of the photos I am looking at", + "Next up for AI Denoise…": "Next up for AI Denoise…", + "Nothing is fetched until you accept. The model is kept on this computer only; LightCraft never uploads or shares it.": "Nothing is fetched until you accept. The model is kept on this computer only; LightCraft never uploads or shares it.", + "Only photos whose AI Denoise amount is above 0 are made. Exports always make theirs if it is missing.": "Only photos whose AI Denoise amount is above 0 are made. Exports always make theirs if it is missing.", + "Open page": "Open page", + "Open the model's page": "Open the model's page", + "Opens the model's own page in your browser, to read about it or get the file yourself.": "Opens the model's own page in your browser, to read about it or get the file yourself.", + "Processor": "Processor", + "Set up AI Denoise…": "Set up AI Denoise…", + "Speed": "Speed", + "Starting…": "Starting…", + "Tested and works on this computer": "Tested and works on this computer", + "The denoise model is installed and in use: the AI Denoise Amount slider under Detail now works": "The denoise model is installed and in use: the AI Denoise Amount slider under Detail now works", + "The graphics card is set up when the next photo is made.": "The graphics card is set up when the next photo is made.", + "The model is kept on this computer only. LightCraft never uploads or shares it.": "The model is kept on this computer only. LightCraft never uploads or shares it.", + "The processor does the work.": "The processor does the work.", + "There is no model file to install": "There is no model file to install", + "This build cannot run denoise models.": "This build cannot run denoise models.", + "This build has nowhere to keep denoise models, so AI Denoise is not available. (The desktop app has.)": "This build has nowhere to keep denoise models, so AI Denoise is not available. (The desktop app has.)", + "This file cannot be used yet": "This file cannot be used yet", + "Tick the box to accept the model's terms first": "Tick the box to accept the model's terms first", + "Try again": "Try again", + "Try the graphics card again": "Try the graphics card again", + "Turn off": "Turn off", + "Unknown licence": "Unknown licence", + "Use": "Use", + "Waiting to start…": "Waiting to start…", + "Works on raw files from cameras with a standard (Bayer) sensor. Other files keep their normal noise reduction.": "Works on raw files from cameras with a standard (Bayer) sensor. Other files keep their normal noise reduction.", + "cached picture": "cached picture", + "cached pictures": "cached pictures", + "commercial use allowed": "commercial use allowed", + "denoiseModel.accept": "denoiseModel.accept", + "in use": "in use", + "installed": "installed", + "makes AI Denoise available": "makes AI Denoise available", + "non-commercial use only": "non-commercial use only", + "photo has its denoised picture.": "photo has its denoised picture.", + "photo waiting": "photo waiting", + "photos have their denoised picture.": "photos have their denoised picture.", + "photos waiting": "photos waiting", + "starts using it for AI Denoise": "starts using it for AI Denoise", + "terms unclear": "terms unclear", + "the graphics card": "the graphics card", + "the graphics card is not used": "the graphics card is not used", + "Cancel Denoise": "Cancel Denoise", + "Cancel Denoise Model Download": "Cancel Denoise Model Download", + "Choose Denoise Model": "Choose Denoise Model", + "Clear Denoise Cache": "Clear Denoise Cache", + "Denoise Model Downloads": "Denoise Model Downloads", + "Denoise Photos": "Denoise Photos", + "Denoise Pump": "Denoise Pump", + "Denoise Status": "Denoise Status", + "Download Denoise Model": "Download Denoise Model", + "Remove Denoise Model": "Remove Denoise Model", + "Test Denoise Model": "Test Denoise Model", + "AI Denoise is enabled on the requesting photo": "AI Denoise is enabled on the requesting photo", + "After installation, AI Denoise will turn on for the requesting photo.": "After installation, AI Denoise will turn on for the requesting photo.", + "Cancel pending action": "Cancel pending action", + "Install a denoise model to enable AI Denoise on this photo": "Install a denoise model to enable AI Denoise on this photo", + "Pending photo action cancelled": "Pending photo action cancelled", + "Pending photo action cancelled because its library or photo changed": "Pending photo action cancelled because its library or photo changed", + "ONNX model": "ONNX model", + "commercial use subject to GPL conditions": "commercial use subject to GPL conditions" } diff --git a/crates/ui-egui/src/control.rs b/crates/ui-egui/src/control.rs index cfcb287f..dcd3d1f8 100644 --- a/crates/ui-egui/src/control.rs +++ b/crates/ui-egui/src/control.rs @@ -105,6 +105,8 @@ pub fn inspect(app: &LightcraftApp, ctx: &egui::Context) -> Value { "status": app.ui.status, "notices": app.notices, "quitPrompt": app.quit_prompt, + // how hard the background face scan may work right now, and why it is judged so (see `panels::faces::scan_pace`) + "faceScan": {"pace": app.caches.faces_pace, "focused": app.caches.faces_in_front, "pending": app.caches.faces_pending, "indexed": app.caches.faces_indexed}, "unsaved": app.session.unsaved().map(|(n, e)| json!({"ops": n, "error": e})), "libraryProblem": app.library_problem.as_ref().map(crate::panels::library_problem::LibraryProblem::to_json), "scan": app.scan.as_ref().map(crate::import::ScanTask::status), @@ -323,8 +325,18 @@ pub fn handle(app: &mut LightcraftApp, ctx: &egui::Context, req: &ControlRequest ok(Value::Null) } "ui.resize" => { + for name in ["width", "height"] { + if p.get(name).is_some() && f(name).is_none() { + return err(format!("ui.resize: {name} must be a number")); + } + } let w = f("width").unwrap_or(1600.0) as f32; let h = f("height").unwrap_or(1000.0) as f32; + let zoom = ctx.zoom_factor(); + let native_scale = ctx.input(|i| i.viewport().native_pixels_per_point).unwrap_or_else(|| ctx.pixels_per_point() / zoom); + if let Err(e) = crate::headless::resized_viewport([w, h], native_scale, zoom) { + return err(e); + } ctx.send_viewport_cmd(egui::ViewportCommand::InnerSize(egui::vec2(w, h))); ok(Value::Null) } diff --git a/crates/ui-egui/src/headless.rs b/crates/ui-egui/src/headless.rs index 206d75b1..d1fae1e1 100644 --- a/crates/ui-egui/src/headless.rs +++ b/crates/ui-egui/src/headless.rs @@ -28,13 +28,56 @@ const CLEAR: Color32 = Color32::from_rgb(12, 12, 12); /// Simulated frame interval (deterministic animation time). const FRAME_DT: f64 = 1.0 / 60.0; +/// Validate the scaled screenshot dimensions before layout or pixel allocation. +pub fn viewport_pixels(size: [f32; 2], pixels_per_point: f32) -> Result<[usize; 2], String> { + if size.iter().any(|edge| *edge < 1.0) { + return Err("viewport dimensions must be at least one logical point".into()); + } + native_viewport_pixels(size, pixels_per_point) +} + +fn native_viewport_pixels(size: [f32; 2], pixels_per_point: f32) -> Result<[usize; 2], String> { + let [width, height] = size; + if !width.is_finite() || !height.is_finite() || width <= 0.0 || height <= 0.0 || !pixels_per_point.is_finite() || pixels_per_point <= 0.0 { + return Err("viewport dimensions and scale must be finite and positive".into()); + } + let axis = |value: f32| { + // Match egui/scissor rounding and the existing screenshot dimension contract. + let pixels = (value * pixels_per_point).round(); + if !(1.0..=softpaint::MAX_IMAGE_EDGE as f32).contains(&pixels) { + return Err(format!("scaled viewport edges must be between 1 and {} pixels", softpaint::MAX_IMAGE_EDGE)); + } + Ok(pixels as usize) + }; + let (w, h) = (axis(width)?, axis(height)?); + if w.checked_mul(h).is_none_or(|pixels| pixels > softpaint::MAX_IMAGE_PIXELS) { + return Err(format!("scaled viewport exceeds {} pixels", softpaint::MAX_IMAGE_PIXELS)); + } + Ok([w, h]) +} + +/// Convert an effective-point resize to the native viewport before validating pixel bounds. +/// Sub-point native edges after zooming out are valid when they round to at least one pixel. +pub fn resized_viewport(size: [f32; 2], native_pixels_per_point: f32, zoom: f32) -> Result { + if size.iter().any(|edge| !edge.is_finite() || *edge < 1.0) || !zoom.is_finite() || zoom <= 0.0 { + return Err("resize dimensions must be at least one finite logical point with a positive zoom".into()); + } + let native_size = egui::Vec2::from(size) * zoom; + native_viewport_pixels([native_size.x, native_size.y], native_pixels_per_point)?; + Ok(native_size) +} + +fn bounded_viewport(size: egui::Vec2, scale: f32) -> (egui::Vec2, f32) { + if native_viewport_pixels([size.x, size.y], scale).is_ok() { (size, scale) } else { (egui::vec2(1600.0, 1000.0), 1.0) } +} + /// An offscreen egui context with our fonts and theme, and a CPU mirror of its textures. pub struct HeadlessView { pub ctx: egui::Context, pub textures: TextureStore, pub(crate) shapes: Vec, pixels_per_point: f32, - size: egui::Vec2, + size: [usize; 2], frames: u64, } @@ -49,11 +92,12 @@ impl HeadlessView { let ctx = egui::Context::default(); crate::theme::install_fonts(&ctx); crate::theme::apply(&ctx); - HeadlessView { ctx, textures: TextureStore::default(), shapes: vec![], pixels_per_point: 1.0, size: egui::vec2(1600.0, 1000.0), frames: 0 } + HeadlessView { ctx, textures: TextureStore::default(), shapes: vec![], pixels_per_point: 1.0, size: [1600, 1000], frames: 0 } } /// Input for one frame of a `size` (points) viewport at `pixels_per_point`. pub fn raw_input(size: egui::Vec2, pixels_per_point: f32, time: f64, events: Vec) -> RawInput { + let (size, pixels_per_point) = bounded_viewport(size, pixels_per_point); let rect = egui::Rect::from_min_size(egui::Pos2::ZERO, size); let mut raw = RawInput { screen_rect: Some(rect), @@ -73,9 +117,20 @@ impl HeadlessView { } /// Run one frame; keeps the shapes for [`Self::paint`]. Returns the root viewport's commands. - pub fn run(&mut self, raw: RawInput, run_ui: impl FnMut(&mut egui::Ui)) -> Vec { - if let Some(r) = raw.screen_rect { - self.size = r.size(); + pub fn run(&mut self, mut raw: RawInput, run_ui: impl FnMut(&mut egui::Ui)) -> Vec { + if let Some(rect) = raw.screen_rect { + let native_scale = raw.viewports.get(&ViewportId::ROOT).and_then(|info| info.native_pixels_per_point).unwrap_or(1.0); + self.size = native_viewport_pixels([rect.width(), rect.height()], native_scale).unwrap_or([0, 0]); + // The host viewport is measured before UI zoom. Egui lays out in zoomed points; + // retain the host's physical size while allowing text and controls to grow. + let zoom = self.ctx.zoom_factor(); + if zoom.is_finite() && zoom > 0.0 { + raw.screen_rect = Some(rect / zoom); + if let Some(info) = raw.viewports.get_mut(&ViewportId::ROOT) { + info.inner_rect = info.inner_rect.map(|rect| rect / zoom); + info.outer_rect = info.outer_rect.map(|rect| rect / zoom); + } + } } let mut out = self.ctx.run_ui(raw, run_ui); self.frames += 1; @@ -92,11 +147,17 @@ impl HeadlessView { /// Size in pixels of the last frame. pub fn size_px(&self) -> [usize; 2] { - [(self.size.x * self.pixels_per_point).round() as usize, (self.size.y * self.pixels_per_point).round() as usize] + self.size } /// Rasterize the last frame. `extra` textures (by id) take precedence over the context's own. pub fn paint(&self, extra: &HashMap) -> ColorImage { + if viewport_pixels([self.size[0] as f32, self.size[1] as f32], 1.0).is_err() + || !self.pixels_per_point.is_finite() + || self.pixels_per_point <= 0.0 + { + return ColorImage::new([0, 0], Vec::new()); + } let prims = self.ctx.tessellate(self.shapes.clone(), self.pixels_per_point); softpaint::paint(&prims, &Layered { over: extra, base: &self.textures }, self.size_px(), self.pixels_per_point, CLEAR) } @@ -125,7 +186,14 @@ pub struct Headless { impl Headless { /// Wrap `app` (its control channel is replaced by the driver's). + /// Invalid dimensions fall back to 1600×1000 at scale 1; external callers can use + /// [`viewport_pixels`] to reject them before constructing the session. pub fn new(app: LightcraftApp, size: [f32; 2], pixels_per_point: f32) -> Self { + let (size, pixels_per_point) = if viewport_pixels(size, pixels_per_point).is_ok() { + (egui::Vec2::from(size), pixels_per_point) + } else { + (egui::vec2(1600.0, 1000.0), 1.0) + }; let (tx, rx) = std::sync::mpsc::channel(); let mut app = app.with_control(rx); app.headless_host = true; @@ -133,7 +201,7 @@ impl Headless { app, view: HeadlessView::new(), max_texture_side: 16384, - size: egui::vec2(size[0], size[1]), + size, pixels_per_point, time: 0.0, frames: 0, @@ -156,6 +224,7 @@ impl Headless { /// Run one frame. pub fn step(&mut self) { + (self.size, self.pixels_per_point) = bounded_viewport(self.size, self.pixels_per_point); let mut raw = HeadlessView::raw_input(self.size, self.pixels_per_point, self.time, std::mem::take(&mut self.events)); raw.viewports.entry(ViewportId::ROOT).or_default().maximized = Some(self.window_maximized); raw.max_texture_side = Some(self.max_texture_side); @@ -173,7 +242,12 @@ impl Headless { let image = Arc::new(self.paint()); self.events.push(egui::Event::Screenshot { viewport_id: ViewportId::ROOT, user_data, image }); } - ViewportCommand::InnerSize(s) if s.x >= 1.0 && s.y >= 1.0 => self.size = s, + ViewportCommand::InnerSize(size) => { + // Viewport commands use current egui points, which include UI zoom. + if let Ok(native_size) = resized_viewport([size.x, size.y], self.pixels_per_point, self.view.ctx.zoom_factor()) { + self.size = native_size; + } + } ViewportCommand::Close => self.quit = true, ViewportCommand::Maximized(on) => { self.window_maximized = on; @@ -194,6 +268,7 @@ impl Headless { || self.app.scan.is_some() || self.app.import.is_some() || self.app.export.is_some() + || self.app.session.denoise_busy() || !self.app.tasks.is_empty() || !self.app.synthetic.is_empty() || !self.events.is_empty() @@ -291,8 +366,209 @@ impl Headless { #[cfg(test)] mod tests { + #[test] + fn detail_offers_ai_denoise_for_raw_photos_and_leads_to_its_settings() { + use crate::state::Dialog; + use lightcraft_catalog::{MediaKind, Op, Photo, PhotoId, Source}; + let mut h = demo([1400.0, 900.0]); + let t = Duration::from_secs(10); + let dir = std::env::temp_dir().join(format!("lc-ui-denoise-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + h.app.session.set_denoise_models_dir(Some(dir.join("denoise-models"))); + // a raw photo (its file need not exist: nothing here reads it) + let id = PhotoId(h.app.session.catalog.photos().map(|p| p.id.0).max().unwrap_or(0) + 1); + let mut raw = Photo::new( + id, + Source::File { path: dir.join("IMG_1.dng").to_string_lossy().into_owned() }, + "IMG_1.dng", + "DNG", + 6000, + 4000, + "2026-10-01T00:00:00", + ); + raw.kind = MediaKind::Raw; + h.app.session.commit("setup", Op::AddPhoto { photo: Box::new(raw) }).unwrap(); + let widgets = |h: &mut Headless| h.request("ui.widgets", json!({}), t).to_string(); + h.request("engine.execute", json!({"command": "library.select", "params": {"ids": [id.0]}}), t); + h.request("ui.set", json!({"view": "detail"}), t); + h.app.ui.open_sections = vec!["detail".to_string()]; + h.settle(SETTLE); + h.step(); + h.step(); + let w = widgets(&mut h); + assert!(w.contains("\"denoise:setup\""), "a raw photo with no model is offered the setup: {w}"); + // One click lands in Settings, with local installation available even without a download offer. + let r = h.request("ui.clickWidget", json!({"id": "denoise:setup"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + h.step(); + assert_eq!(h.app.ui.dialog, Some(Dialog::Settings { tab: "denoise".into() })); + h.settle(SETTLE); + let listed = h.app.session.execute("denoise.models.list", &json!({})).unwrap(); + assert_eq!(listed["models"].as_array().unwrap().len(), lightcraft_denoise::known::all().len(), "{listed}"); + assert!(widgets(&mut h).contains("\"denoise:installFile\"")); + h.app.ui.dialog = None; + // a photo that is not raw: no offer + h.request( + "engine.execute", + json!({"command": "library.select", "params": {"ids": [h.app.session.catalog.photos().find(|p| p.id != id).unwrap().id.0]}}), + t, + ); + h.settle(SETTLE); + h.step(); + h.step(); + assert!(!widgets(&mut h).contains("\"denoise:setup\""), "nothing is offered for a photo denoise does not apply to"); + let _ = std::fs::remove_dir_all(&dir); + } + 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 +677,494 @@ mod tests { h.settle(SETTLE); } + /// Naming a face in the loupe: point at an unnamed face, click its "Add name" label, type, Enter. Needs the + /// recognition runtime only for the background indexer, so the naming itself is tested whatever the build. + #[test] + fn naming_a_face_in_the_loupe() { + use crate::state::NameEdit; + use lightcraft_catalog::Op; + let mut h = demo([1400.0, 900.0]); + let t = Duration::from_secs(10); + let id = h.app.session.active().unwrap(); + let mut meta = h.app.session.catalog.photo(id).unwrap().meta.clone(); + meta.regions = vec![lightcraft_meta::Region { + rect: lightcraft_geom::Rect { x0: 0.3, y0: 0.2, x1: 0.55, y1: 0.6 }, + kind: lightcraft_meta::RegionKind::Face, + name: None, + description: None, + }]; + h.app.session.commit("setup", Op::SetMeta { id, meta: Box::new(meta) }).unwrap(); + h.request("ui.set", json!({"view": "detail", "right": "none"}), t); + h.settle(SETTLE); + // pointing at the box shows the invitation; clicking it opens the name box + h.request("ui.pointer", json!({"events": [{"kind": "move", "x": 0.42, "y": 0.4}]}), t); + h.step(); + h.step(); + let r = h.request("ui.clickWidget", json!({"id": "regionLabel:0"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + assert!(matches!(&h.app.ui.name_edit, Some(NameEdit { index: 0, .. })), "{:?}", h.app.ui.name_edit); + // Escape closes it without naming anything + let r = h.request("ui.key", json!({"key": "escape"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + h.step(); + assert!(h.app.ui.name_edit.is_none(), "{:?}", h.app.ui.name_edit); + assert_eq!(h.app.session.catalog.photo(id).unwrap().meta.regions[0].name, None); + // open it again, type a name, Enter: the face is named, in one undo step + h.request("ui.pointer", json!({"events": [{"kind": "move", "x": 0.42, "y": 0.4}]}), t); + h.step(); + h.step(); + h.request("ui.clickWidget", json!({"id": "regionLabel:0"}), t); + h.step(); + h.step(); + let undo = h.app.session.undo.len(); + h.request("ui.text", json!({"text": "Ann"}), t); + h.step(); + h.request("ui.key", json!({"key": "enter"}), t); + h.step(); + h.step(); + assert_eq!(h.app.session.catalog.photo(id).unwrap().meta.regions[0].name.as_deref(), Some("Ann")); + assert_eq!(h.app.session.undo.len(), undo + 1); + assert!(h.app.ui.name_edit.is_none()); + } + + /// Adding a face model: the dialog shows the file's terms, the model is installed only once they are + /// accepted, and a file LightCraft cannot use only says why. + #[test] + fn adding_a_face_model_shows_its_terms_and_installs_only_once_accepted() { + use crate::state::Dialog; + let mut h = demo([1400.0, 900.0]); + let t = Duration::from_secs(10); + let dir = std::env::temp_dir().join(format!("lc-ui-facemodel-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + h.app.session.face_models_dir = Some(dir.join("models")); + let model = dir.join("Mine.onnx"); + std::fs::write(&model, lightcraft_faces::synthetic::tiny_embedder_model(512)).unwrap(); + let open = |h: &mut Headless, path: &std::path::Path| { + h.request("engine.execute", json!({"command": "dialog.faceModel", "params": {"path": path.to_string_lossy()}}), t) + }; + + // the dialog opens for the chosen file, with nothing accepted and nothing installed + assert_eq!(open(&mut h, &model)["ok"], true); + h.settle(SETTLE); + h.step(); + let Some(dlg) = h.app.ui.dialog.clone() else { panic!("no dialog") }; + let Dialog::FaceModel { info, accepted, .. } = &dlg else { panic!("{dlg:?}") }; + assert!(!accepted); + assert_eq!(info["kind"], "draft"); + assert!(!dir.join("models").exists()); + // OK does nothing yet + assert!(crate::panels::dialogs::confirm_dialog(&mut h.app, &dlg).is_err()); + assert!(!dir.join("models").exists()); + // ticking the box and confirming installs it + let r = h.request("ui.clickWidget", json!({"id": "check:faceModel.accept"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + assert!(matches!(&h.app.ui.dialog, Some(Dialog::FaceModel { accepted: true, .. }))); + let r = h.request("ui.dialog.confirm", json!({}), t); + assert_eq!(r["ok"], true, "{r}"); + let listed = h.app.run("faces.models.list", json!({})).unwrap(); + assert!(listed["models"].as_array().unwrap().iter().any(|m| m["installed"] == true && m["known"] == false), "{listed}"); + // installed, in use and recognition on: back in Settings ▸ Faces, where the scan can be watched + assert_eq!(h.app.ui.dialog, Some(Dialog::Settings { tab: "faces".into() })); + assert_eq!(listed["enabled"], true); + h.app.ui.dialog = None; + + // a file that is not a usable model says why and offers no install + let junk = dir.join("junk.onnx"); + std::fs::write(&junk, b"not a model").unwrap(); + assert_eq!(open(&mut h, &junk)["ok"], true); + h.step(); + let Some(Dialog::FaceModel { info, .. }) = h.app.ui.dialog.clone() else { panic!("no dialog") }; + assert_eq!(info["kind"], "unsupported"); + // a path that does not exist is an error, not a dialog + h.app.ui.dialog = None; + assert_eq!(open(&mut h, &dir.join("nope.onnx"))["ok"], false); + assert!(h.app.ui.dialog.is_none()); + let _ = std::fs::remove_dir_all(&dir); + } + + /// Photo > Detect Faces without the detector offers the download: the terms first, nothing fetched before they are + /// accepted, and the model's row in Settings has its Download button. + #[test] + fn detect_faces_without_the_detector_offers_the_download_and_fetches_nothing_before_acceptance() { + use crate::state::Dialog; + let mut h = demo([1400.0, 900.0]); + let dir = std::env::temp_dir().join(format!("lc-ui-detector-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + h.app.session.face_models_dir = Some(dir.join("models")); + + let r = h.app.run("faces.detect", json!({})); + assert!(r.is_err(), "no detector yet"); + let Some(dlg) = h.app.ui.dialog.clone() else { panic!("the download was not offered") }; + let Dialog::FaceModel { info, accepted, .. } = &dlg else { panic!("{dlg:?}") }; + assert!(!accepted); + assert_eq!((info["download"].as_str(), info["host"].as_str()), (Some("yunet-2023mar"), Some("github.com"))); + assert_eq!(info["model"]["licence"]["name"], "MIT"); + // OK without accepting the terms starts nothing + assert!(crate::panels::dialogs::confirm_dialog(&mut h.app, &dlg).is_err()); + assert!(h.app.session.face_downloads.snapshot().is_empty()); + assert!(!dir.join("models").join(".downloads").exists()); + assert!(h.app.caches.faces_dl_watch.is_empty()); + + // Settings > Faces shows the detector with a Download button + h.app.ui.dialog = Some(Dialog::Settings { tab: "faces".into() }); + h.step(); + // (the first request switches the widget list on, the next frame fills it) + h.request("ui.widgets", json!({}), Duration::from_secs(10)); + h.step(); + let w = h.request("ui.widgets", json!({"filter": "faces:"}), Duration::from_secs(10)); + assert!(w.to_string().contains("faces:download:yunet-2023mar"), "{w}"); + let _ = std::fs::remove_dir_all(&dir); + } + + /// Pressing Download in Settings shows the model's terms first and fetches nothing until they are accepted; a build + /// that cannot run recognition models offers no download. The licence dialog replaces Settings, which it was opened from. + #[test] + fn download_shows_the_terms_first_and_fetches_nothing_until_accepted() { + use crate::state::Dialog; + let mut h = demo([1400.0, 900.0]); + let t = Duration::from_secs(10); + let dir = std::env::temp_dir().join(format!("lc-ui-facedl-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + h.app.session.face_models_dir = Some(dir.join("models")); + h.request("engine.execute", json!({"command": "app.settings", "params": {"tab": "faces"}}), t); + h.settle(SETTLE); + h.step(); + let runtime = h.app.session.execute("faces.models.list", &json!({})).unwrap()["runtime"] == true; + let r = h.request("ui.clickWidget", json!({"id": "faces:download:sface-2021dec"}), t); + if !runtime { + assert_eq!(r["ok"], false, "a build that cannot run the model offers no download: {r}"); + let _ = std::fs::remove_dir_all(&dir); + return; + } + assert_eq!(r["ok"], true, "{r}"); + h.step(); + h.step(); + let Some(dlg) = h.app.ui.dialog.clone() else { panic!("no dialog") }; + let Dialog::FaceModel { path, info, accepted } = &dlg else { panic!("the terms did not replace Settings: {dlg:?}") }; + assert_eq!((info["download"].as_str(), *accepted, path.as_str()), (Some("sface-2021dec"), false, "")); + assert_eq!(info["model"]["licence"]["commercial"], "unknown"); + // OK does nothing until the terms are accepted, and nothing has been fetched + assert!(crate::panels::dialogs::confirm_dialog(&mut h.app, &dlg).is_err()); + assert_eq!(h.app.session.execute("faces.models.downloads", &json!({})).unwrap()["downloads"], json!([])); + // cancelling leaves it that way + h.request("ui.key", json!({"key": "escape"}), t); + h.step(); + assert!(h.app.ui.dialog.is_none()); + assert_eq!(h.app.session.execute("faces.models.downloads", &json!({})).unwrap()["downloads"], json!([])); + let _ = std::fs::remove_dir_all(&dir); + } + + /// Until face recognition is set up, the People view and the loupe's name box say so and offer the next step: with no + /// model, a button that opens Settings ▸ Faces; with a model installed but recognition off, one that switches it on. + #[test] + fn people_and_the_name_box_offer_to_set_face_recognition_up() { + use crate::state::{Dialog, ViewMode}; + use lightcraft_catalog::Op; + let mut h = demo([1400.0, 900.0]); + let t = Duration::from_secs(10); + let dir = std::env::temp_dir().join(format!("lc-ui-facesetup-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + h.app.session.face_models_dir = Some(dir.join("models")); + let runtime = h.app.session.execute("faces.models.list", &json!({})).unwrap()["runtime"] == true; + let id = h.app.session.active().unwrap(); + let mut meta = h.app.session.catalog.photo(id).unwrap().meta.clone(); + meta.regions = vec![lightcraft_meta::Region { + rect: lightcraft_geom::Rect { x0: 0.3, y0: 0.2, x1: 0.55, y1: 0.6 }, + kind: lightcraft_meta::RegionKind::Face, + name: None, + description: None, + }]; + h.app.session.commit("setup", Op::SetMeta { id, meta: Box::new(meta) }).unwrap(); + let offers = |h: &mut Headless| h.request("ui.widgets", json!({}), t).to_string().contains("\"faces:setup\""); + + h.request("engine.execute", json!({"command": "view.people"}), t); + h.settle(SETTLE); + h.step(); + assert_eq!(h.app.ui.view, ViewMode::People); + if !runtime { + assert!(!offers(&mut h), "a build that cannot run recognition offers nothing"); + let _ = std::fs::remove_dir_all(&dir); + return; + } + // no model: the offer is there, and one click lands in Settings ▸ Faces + assert!(offers(&mut h)); + let r = h.request("ui.clickWidget", json!({"id": "faces:setup"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + h.step(); + assert_eq!(h.app.ui.dialog, Some(Dialog::Settings { tab: "faces".into() })); + h.app.ui.dialog = None; + + // the loupe's name box offers it too, and its button leaves for Settings with the box closed + h.request("ui.set", json!({"view": "detail", "right": "none"}), t); + h.settle(SETTLE); + h.request("ui.pointer", json!({"events": [{"kind": "move", "x": 0.42, "y": 0.4}]}), t); + h.step(); + h.step(); + h.request("ui.clickWidget", json!({"id": "regionLabel:0"}), t); + h.step(); + h.step(); + assert!(h.app.ui.name_edit.is_some() && offers(&mut h), "the name box offers the setup"); + let r = h.request("ui.clickWidget", json!({"id": "faces:setup"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + h.step(); + assert!(h.app.ui.name_edit.is_none()); + assert_eq!(h.app.ui.dialog, Some(Dialog::Settings { tab: "faces".into() })); + h.app.ui.dialog = None; + + // a model installed but recognition off: "Turn on" does it on the spot + let model = dir.join("Mine.onnx"); + std::fs::write(&model, lightcraft_faces::synthetic::tiny_embedder_model(512)).unwrap(); + let installed = h.app.run("faces.models.install", json!({"path": model.to_string_lossy(), "acknowledged": true, "activate": false})).unwrap(); + h.app.run("faces.models.select", json!({"id": installed["installed"]["id"]})).unwrap(); + h.app.caches.faces_epoch += 1; + h.request("engine.execute", json!({"command": "view.people"}), t); + h.settle(SETTLE); + h.step(); + assert!(offers(&mut h)); + assert_eq!(h.app.session.execute("faces.models.list", &json!({})).unwrap()["enabled"], false); + let r = h.request("ui.clickWidget", json!({"id": "faces:setup"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + h.step(); + assert_eq!(h.app.session.execute("faces.models.list", &json!({})).unwrap()["enabled"], true); + assert_eq!(h.app.ui.dialog, None, "nothing to go to Settings for"); + h.step(); + assert!(!offers(&mut h), "once it is on the offer is gone"); + h.settle(SETTLE); + let _ = std::fs::remove_dir_all(&dir); + } + + /// A click on a person's card opens their page, which shows only cropped faces (one per face, not per photo); a click + /// on one opens its photo; Back, the People button and Escape return to everyone. + #[test] + fn a_persons_page_shows_their_faces_and_back_returns_to_everyone() { + use crate::state::ViewMode; + use lightcraft_catalog::Op; + let mut h = demo([1400.0, 900.0]); + let t = Duration::from_secs(10); + let ids: Vec<_> = h.app.session.catalog.photos().map(|p| p.id).take(3).collect(); + let face = |x: f64, name: &str| lightcraft_meta::Region { + rect: lightcraft_geom::Rect { x0: x, y0: 0.2, x1: x + 0.25, y1: 0.6 }, + kind: lightcraft_meta::RegionKind::Face, + name: Some(name.to_string()), + description: None, + }; + // Jane Doe twice in the first photo and once in the second; John Roe in the third + for (id, regions) in [ + (ids[0], vec![face(0.1, "Jane Doe"), face(0.5, "jane doe")]), + (ids[1], vec![face(0.3, "Jane Doe")]), + (ids[2], vec![face(0.3, "John Roe")]), + ] { + let mut meta = h.app.session.catalog.photo(id).unwrap().meta.clone(); + meta.regions = regions; + h.app.session.commit("setup", Op::SetMeta { id, meta: Box::new(meta) }).unwrap(); + } + let open_people = |h: &mut Headless| { + h.request("engine.execute", json!({"command": "view.people"}), t); + h.settle(SETTLE); + h.step(); + }; + open_people(&mut h); + assert_eq!((h.app.ui.view, h.app.ui.person_page.clone()), (ViewMode::People, None)); + // the card opens the page: a face tile for each of the three faces, not a tile per photo + let r = h.request("ui.clickWidget", json!({"id": "person:Jane Doe"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + h.step(); + assert_eq!((h.app.ui.view, h.app.ui.person_page.as_deref()), (ViewMode::People, Some("Jane Doe"))); + let widgets = h.request("ui.widgets", json!({}), t).to_string(); + for (id, index) in [(ids[0], 0), (ids[0], 1), (ids[1], 0)] { + assert!(widgets.contains(&format!("\"person-face:{}:{index}\"", id.0)), "a tile for each face: {widgets}"); + } + assert!(!widgets.contains(&format!("\"person-face:{}:0\"", ids[2].0)), "someone else's face is not on the page"); + // a face opens its photo in the detail view + h.request("ui.clickWidget", json!({"id": format!("person-face:{}:0", ids[1].0)}), t); + h.step(); + h.step(); + assert_eq!((h.app.ui.view, h.app.session.active()), (ViewMode::Detail, Some(ids[1]))); + // Escape from that photo goes back to the page, not to the grid + h.request("ui.key", json!({"key": "escape"}), t); + h.step(); + assert_eq!((h.app.ui.view, h.app.ui.person_page.as_deref()), (ViewMode::People, Some("Jane Doe"))); + // the People button from a photo shows everyone, with the person last left one click away next to the title + h.request("ui.clickWidget", json!({"id": format!("person-face:{}:0", ids[1].0)}), t); + h.step(); + h.step(); + assert_eq!(h.app.ui.view, ViewMode::Detail); + open_people(&mut h); + assert_eq!((h.app.ui.view, h.app.ui.person_page.clone()), (ViewMode::People, None)); + h.step(); + let r = h.request("ui.clickWidget", json!({"id": "people:last"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + h.step(); + assert_eq!((h.app.ui.view, h.app.ui.person_page.as_deref()), (ViewMode::People, Some("Jane Doe"))); + // a photo reached any other way goes back to the grid on Escape + h.request("engine.execute", json!({"command": "view.detail"}), t); + h.request("ui.key", json!({"key": "escape"}), t); + h.step(); + assert_eq!(h.app.ui.view, ViewMode::PhotoGrid); + h.request("engine.execute", json!({"command": "view.person", "params": {"name": "Jane Doe"}}), t); + h.step(); + let r = h.request("ui.clickWidget", json!({"id": "person:back"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + h.step(); + assert_eq!(h.app.ui.person_page, None); + h.request("engine.execute", json!({"command": "view.person", "params": {"name": "Jane Doe"}}), t); + h.request("ui.key", json!({"key": "escape"}), t); + h.step(); + assert_eq!((h.app.ui.view, h.app.ui.person_page.clone()), (ViewMode::People, None), "Escape goes back to everyone"); + h.request("engine.execute", json!({"command": "view.person", "params": {"name": "John Roe"}}), t); + open_people(&mut h); + assert_eq!(h.app.ui.person_page, None, "the People button shows everyone"); + // a missing or empty name is refused + assert_eq!(h.request("engine.execute", json!({"command": "view.person", "params": {}}), t)["ok"], false); + assert_eq!(h.request("engine.execute", json!({"command": "view.person", "params": {"name": " "}}), t)["ok"], false); + } + + /// The People view lists the unnamed faces below the named people: select some (a click each, or all), type a name, + /// press Enter, and they are all named at once, in one undo step; the new person appears among the named. + #[test] + fn unnamed_faces_are_selected_and_named_together() { + use lightcraft_catalog::Op; + let mut h = demo([1400.0, 900.0]); + let t = Duration::from_secs(10); + let ids: Vec<_> = h.app.session.catalog.photos().map(|p| p.id).take(3).collect(); + let face = |x: f64, name: Option<&str>| lightcraft_meta::Region { + rect: lightcraft_geom::Rect { x0: x, y0: 0.2, x1: x + 0.2, y1: 0.55 }, + kind: lightcraft_meta::RegionKind::Face, + name: name.map(str::to_string), + description: None, + }; + for (id, regions) in [ + (ids[0], vec![face(0.1, Some("Jane Doe")), face(0.5, None)]), + (ids[1], vec![face(0.3, None), face(0.6, None)]), + (ids[2], vec![face(0.3, None)]), + ] { + let mut meta = h.app.session.catalog.photo(id).unwrap().meta.clone(); + meta.regions = regions; + h.app.session.commit("setup", Op::SetMeta { id, meta: Box::new(meta) }).unwrap(); + } + h.request("engine.execute", json!({"command": "view.people"}), t); + h.settle(SETTLE); + h.step(); + h.step(); + let tiles = |h: &mut Headless| h.request("ui.widgets", json!({"filter": "unnamed-face:"}), t)["result"].as_array().map_or(0, Vec::len); + // the named person is a card, and the four unnamed faces are tiles below + let r = h.request("ui.clickWidget", json!({"id": "person:Jane Doe"}), t); + assert_eq!(r["ok"], true, "{r}"); + h.request("engine.execute", json!({"command": "view.people"}), t); + h.step(); + h.step(); + assert_eq!(tiles(&mut h), 4); + // select two (nothing is named by selecting), and the naming bar appears + for (id, i) in [(ids[1], 0), (ids[2], 0)] { + let r = h.request("ui.clickWidget", json!({"id": format!("unnamed-face:{}:{i}", id.0)}), t); + assert_eq!(r["ok"], true, "{r}"); + h.step(); + h.step(); + } + assert_eq!(h.app.ui.unnamed_selected.len(), 2); + let bar = h.request("ui.widgets", json!({"filter": "unnamed:"}), t).to_string(); + assert!(bar.contains("unnamed:name") && bar.contains("unnamed:clear"), "{bar}"); + assert!(h.request("ui.widgets", json!({"filter": "field:unnamedName"}), t).to_string().contains("field:unnamedName")); + // type a name and press Enter: both faces are named, in one undo step, and the selection is gone + let undo = h.app.session.undo.len(); + h.request("ui.text", json!({"text": "Ann Example"}), t); + h.step(); + h.request("ui.key", json!({"key": "enter"}), t); + h.step(); + h.step(); + let named = |h: &Headless, id: lightcraft_catalog::PhotoId, i: usize| h.app.session.catalog.photo(id).unwrap().meta.regions[i].name.clone(); + assert_eq!((named(&h, ids[1], 0), named(&h, ids[2], 0)), (Some("Ann Example".to_string()), Some("Ann Example".to_string()))); + assert_eq!(named(&h, ids[1], 1), None, "a face that was not selected stays unnamed"); + assert_eq!(h.app.session.undo.len(), undo + 1, "one step for both"); + assert!(h.app.ui.unnamed_selected.is_empty() && h.app.ui.unnamed_name.is_empty()); + h.settle(SETTLE); + h.step(); + assert_eq!(tiles(&mut h), 2, "the named faces left the unnamed list"); + assert_eq!(h.app.session.catalog.people().len(), 2, "and the new person is among the named"); + // select all, then clear: nothing is named + h.request("ui.clickWidget", json!({"id": "unnamed:selectAll"}), t); + h.step(); + h.step(); + assert_eq!(h.app.ui.unnamed_selected.len(), 2); + h.request("ui.clickWidget", json!({"id": "unnamed:clear"}), t); + h.step(); + assert!(h.app.ui.unnamed_selected.is_empty()); + // undo gives the two faces back + h.request("engine.execute", json!({"command": "edit.undo"}), t); + h.step(); + assert_eq!(named(&h, ids[1], 0), None); + } + + /// A screenful of faces larger than the picture cache's usual budget (96) is all kept: with a fixed budget the same few + /// tiles were evicted and re-requested every frame and stayed blank. + #[test] + fn a_screenful_of_small_faces_is_not_evicted_and_left_blank() { + use lightcraft_catalog::Op; + let mut h = demo([1500.0, 1000.0]); + let t = Duration::from_secs(20); + let ids: Vec<_> = h.app.session.catalog.photos().map(|p| p.id).collect(); + // seven unnamed faces in every photo of the demo library + for id in &ids { + let mut meta = h.app.session.catalog.photo(*id).unwrap().meta.clone(); + meta.regions = (0..7) + .map(|k| lightcraft_meta::Region { + rect: lightcraft_geom::Rect { x0: 0.05 + 0.12 * k as f64, y0: 0.2, x1: 0.15 + 0.12 * k as f64, y1: 0.45 }, + kind: lightcraft_meta::RegionKind::Face, + name: None, + description: None, + }) + .collect(); + h.app.session.commit("setup", Op::SetMeta { id: *id, meta: Box::new(meta) }).unwrap(); + } + // the smallest faces, so a lot of them fit on the screen + h.request("ui.set", json!({"thumbSize": 90.0, "view": "people"}), t); + h.settle(SETTLE); + h.step(); + let drawn = h.request("ui.widgets", json!({"filter": "unnamed-face:"}), t)["result"].as_array().map_or(0, Vec::len); + assert!(drawn > 96, "the test needs more faces on screen than the usual budget, got {drawn}"); + // the pictures are made and kept, well beyond the old budget (a few demo photos share a scene, so two faces can share a + // picture: the count is a little under the number of tiles) + let started = Instant::now(); + while h.app.renderer.variant_textures() < drawn - 20 && started.elapsed() < Duration::from_secs(20) { + h.step(); + std::thread::sleep(Duration::from_millis(20)); + } + let kept = h.app.renderer.variant_textures(); + assert!(kept > 96 + 30, "{kept} pictures kept for {drawn} faces on screen"); + // and they stay: many more frames, and none is evicted to be asked for again + for _ in 0..30 { + h.step(); + } + assert!(h.app.renderer.variant_textures() >= kept, "{} pictures kept after more frames, {kept} before", h.app.renderer.variant_textures()); + } + + /// `ui.inspect` says how hard the face scan is allowed to work and whether the window counts as in front, so a slow + /// scan can be told from a stuck one. + #[test] + fn inspect_reports_the_face_scan_pace() { + let mut h = demo([1000.0, 700.0]); + let t = Duration::from_secs(10); + h.settle(SETTLE); + h.step(); + let r = h.request("ui.inspect", json!({}), t); + let scan = &r["result"]["faceScan"]; + assert!(["pause", "light", "normal", "full"].contains(&scan["pace"].as_str().unwrap_or("")), "{scan}"); + assert!(scan["focused"].is_boolean() && scan["pending"].is_u64() && scan["indexed"].is_u64(), "{scan}"); + } + /// Profile browser: live variant thumbnails, hover previews in the loupe without touching the /// photo or its history, click applies, the star toggles the favourite. #[test] diff --git a/crates/ui-egui/src/lib.rs b/crates/ui-egui/src/lib.rs index aeab2794..7684f00d 100644 --- a/crates/ui-egui/src/lib.rs +++ b/crates/ui-egui/src/lib.rs @@ -19,6 +19,7 @@ pub mod menu_level; pub mod menubar; pub mod menus; pub mod merge; +mod model_setup; pub mod panels; pub mod pick; pub mod region; @@ -106,6 +107,10 @@ pub struct Services { pub picker: Option, /// Show an open dialog for photos; returns paths. pub pick_files: Option, + /// Local ONNX model; a manifest must accompany it. + pub pick_denoise_model: Option, + /// Open dialog for a face model (`.onnx`; Settings ▸ Faces ▸ Add a model file…). + pub pick_model_file: Option, /// Open dialog for preset files (`.lcpreset`, `.xmp`, `.lrtemplate`, `.zip`, `.dng`, Luminar `.lmp` / `.mplumpack`). pub pick_preset_files: Option, /// Open dialog for a GPS track log (`.gpx`; Photo ▸ Auto-Tag from Tracklog…). @@ -157,6 +162,7 @@ pub struct Perf { } pub struct LightcraftApp { + pub(crate) model_setup: model_setup::Pending, /// Per-catalog-revision caches of library-wide results the panels show every frame /// (expensive on big libraries). pub caches: Caches, @@ -312,6 +318,7 @@ impl LightcraftApp { gpu_applied: None, memory_applied: None, library_problem: None, + model_setup: Default::default(), } } @@ -324,6 +331,9 @@ impl LightcraftApp { /// Run a UI or engine command by id. The single entry point for every frontend path. pub fn run(&mut self, id: &str, params: Value) -> Result { + if let Some(result) = model_setup::intercept(self, id, ¶ms) { + return result; + } if let Some(r) = menus::run_ui_command(self, id, ¶ms) { return r; } @@ -342,9 +352,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 +624,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 +729,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 +773,9 @@ impl LightcraftApp { ctx.request_repaint_after(std::time::Duration::from_secs(3)); } self.session.persist_if_dirty(); + panels::denoise::pump(self, ctx); + model_setup::pump(self, ctx); + panels::faces::pump(self, ctx); self.collect_screenshots(ctx); self.issue_screenshots(ctx); if self.fonts_ready { @@ -761,6 +786,13 @@ impl LightcraftApp { let dropped: Vec = ctx.input(|i| i.raw.dropped_files.iter().map(|f| f.path().to_string_lossy().to_string()).filter(|p| !p.is_empty()).collect()); // preset files import as presets, everything else as photos + // a dropped face model opens its licence dialog; presets and photos as before + let (models, dropped): (Vec, Vec) = dropped.into_iter().partition(|p| is_model_file(p)); + if let Some(model) = models.first() + && let Err(e) = panels::faces::open_dialog(self, model) + { + self.toast(ctx, e); + } let (presets, photos): (Vec, Vec) = dropped.into_iter().partition(|p| is_preset_file(p)); if !presets.is_empty() { let _ = self.run("file.importPresets", serde_json::json!({"paths": presets})); @@ -912,6 +944,7 @@ impl LightcraftApp { export_task::poll(self, &ctx); pick::poll(self, &ctx); panels::grid::drag_feedback(self, &ctx); + panels::left::album_drag_feedback(self, &ctx); panels::toast(self, &ctx); self.widgets = widgets::take_registry(&ctx); self.end_frame(t0); @@ -956,6 +989,25 @@ pub fn is_bw(d: &lightcraft_develop::DevelopSettings) -> bool { d.treatment == lightcraft_develop::Treatment::Bw || d.profile.id == "lc.mono" || d.profile.id.starts_with("lc.bw.") } +/// "Found 3 faces in 2 photos" for a `faces.detect` result. +fn detect_summary(v: &Value) -> String { + let photos = v["photos"].as_array().map(Vec::as_slice).unwrap_or_default(); + let counts: Vec = photos.iter().map(|p| p["faces"].as_array().map_or(0, Vec::len)).collect(); + let (faces, with) = (counts.iter().sum::(), counts.iter().filter(|n| **n > 0).count()); + let plural = |n: usize, w: &str| format!("{n} {w}{}", if n == 1 { "" } else { "s" }); + match (faces, photos.len()) { + (0, 1) => "No faces found".into(), + (0, n) => format!("No faces found in {}", plural(n, "photo")), + (f, 1) => format!("Found {}", plural(f, "face")), + (f, _) => format!("Found {} in {}", plural(f, "face"), plural(with, "photo")), + } +} + +/// A dropped face model file (`.onnx`). +pub fn is_model_file(path: &str) -> bool { + std::path::Path::new(path).extension().is_some_and(|e| e.eq_ignore_ascii_case("onnx")) +} + /// Files dropped on the window that are presets rather than photos. pub fn is_preset_file(path: &str) -> bool { let ext = std::path::Path::new(path).extension().map(|e| e.to_string_lossy().to_ascii_lowercase()).unwrap_or_default(); @@ -964,6 +1016,28 @@ pub fn is_preset_file(path: &str) -> bool { #[cfg(test)] mod drop_tests { + #[test] + fn detection_results_are_summarised_in_words() { + let r = + |counts: &[usize]| serde_json::json!({"photos": counts.iter().map(|n| serde_json::json!({"faces": vec![0; *n]})).collect::>()}); + assert_eq!(super::detect_summary(&r(&[0])), "No faces found"); + assert_eq!(super::detect_summary(&r(&[1])), "Found 1 face"); + assert_eq!(super::detect_summary(&r(&[3])), "Found 3 faces"); + assert_eq!(super::detect_summary(&r(&[0, 0])), "No faces found in 2 photos"); + assert_eq!(super::detect_summary(&r(&[2, 0, 1])), "Found 3 faces in 2 photos"); + assert_eq!(super::detect_summary(&serde_json::json!({})), "No faces found in 0 photos"); + } + + #[test] + fn dropped_face_models_are_told_apart() { + for p in ["/a/model.onnx", "/a/dir/M.ONNX"] { + assert!(super::is_model_file(p), "{p}"); + } + for p in ["/a/model.onnx.jpg", "/a/onnx", "/a/b.png"] { + assert!(!super::is_model_file(p), "{p}"); + } + } + #[test] fn dropped_presets_are_told_apart_from_photos() { for p in ["/a/Look.lrtemplate", "/a/b.XMP", "/a/pack.zip", "/a/x.lcpreset", "/a/Magic Hour.mplumpack", "/a/Pop.lmp", "/a/Bundle.LMP"] { @@ -988,6 +1062,33 @@ pub struct Caches { album_counts: Option<(u64, std::sync::Arc>)>, /// How often the album counts were recomputed (tests check that unchanged frames don't). pub album_count_scans: usize, + /// AI denoise: what the pump last saw, the model list and the downloads being watched. + pub denoise: panels::denoise::Ui, + /// Bumped when a face model is installed, removed or chosen, so Settings re-reads the list at once. + pub faces_epoch: u64, + /// Face model downloads the user started, followed until installed (see `panels::faces::pump`). + pub faces_dl_watch: Vec, + /// Whether the background face indexer is running, how many faces it has embedded, how many photos are left, and + /// when to ask it again. + pub faces_active: bool, + pub faces_indexed: u64, + pub faces_pending: u64, + pub faces_next_pump: f64, + /// When the user last dragged, typed or scrolled, and when they last moved the pointer (egui time). + pub last_input: f64, + /// The pace the face scan was last given, and whether the window was in front then (for `ui.inspect`). + pub faces_pace: &'static str, + pub faces_in_front: bool, + pub last_move: f64, + /// Name suggestions for the photo in the loupe: (photo, catalog revision, faces indexed, suggestions by region). + pub face_hints: Option<(u64, u64, u64, std::sync::Arc)>, + /// The open person page: (name, catalog revision, faces indexed, when it was asked for, the page). + pub person_page: Option<(String, u64, u64, f64, std::sync::Arc)>, + /// The unnamed faces: (catalog revision, faces indexed, when it was asked for, the list). + pub unnamed: Option<(u64, u64, f64, std::sync::Arc)>, + /// The most photos the face scan has had left at once since it last finished (the progress bar's whole). + pub faces_peak: u64, + person_names: Option<(u64, std::sync::Arc>)>, /// The grid's date runs, layout and indexes (by the visible list's generation). pub grid: panels::grid::GridCache, /// What the grid did on its frames (benchmarks and tests check unchanged frames stay cheap). @@ -1076,6 +1177,17 @@ impl Caches { self.counts = Some((cat.revision, c)); c } + /// Everyone named on a face in the library (for completing a name as it is typed). + pub fn person_names(&mut self, cat: &lightcraft_catalog::Catalog) -> std::sync::Arc> { + match &self.person_names { + Some((r, v)) if *r == cat.revision => v.clone(), + _ => { + let v = std::sync::Arc::new(cat.people().into_iter().map(|p| p.name).collect::>()); + self.person_names = Some((cat.revision, v.clone())); + v + } + } + } /// The By Date tree. pub fn date_groups(&mut self, cat: &lightcraft_catalog::Catalog) -> std::sync::Arc> { match &self.date_groups { @@ -1213,3 +1325,6 @@ mod cache_tests { assert_eq!(c.album_count_scans, 1); } } + +#[cfg(test)] +mod tests_model_setup; diff --git a/crates/ui-egui/src/menubar.rs b/crates/ui-egui/src/menubar.rs index 6ecdb169..7d964856 100644 --- a/crates/ui-egui/src/menubar.rs +++ b/crates/ui-egui/src/menubar.rs @@ -753,10 +753,29 @@ pub fn show_in_window(app: &mut LightcraftApp, ui: &mut egui::Ui, max_width: f32 if total <= max_width { let saved = ui.spacing().item_spacing.x; ui.spacing_mut().item_spacing.x = TITLE_GAP; - for (title, items) in &bar { - let r = ui.add(egui::Button::new(egui::RichText::new(crate::i18n::tr(title)).font(font.clone()).color(t.text_label)).frame(false)); - crate::widgets::register(ui.ctx(), format!("menu:{title}"), r.rect); - egui::Popup::menu(&r).show(|ui| crate::menu_level::level(ui, 1, bar_bottom, |ui| nodes_ui(ui, items, mac, &mut clicked, 1, bar_bottom))); + // All titles first, then their popups, so a title the pointer moves onto can take over + // from the open menu before either is drawn (no frame with two menus or none). + let titles: Vec = bar + .iter() + .map(|(title, _)| { + let r = ui.add(egui::Button::new(egui::RichText::new(crate::i18n::tr(title)).font(font.clone()).color(t.text_label)).frame(false)); + crate::widgets::register(ui.ctx(), format!("menu:{title}"), r.rect); + r + }) + .collect(); + let confirmed = switch_open_title(ui.ctx(), &titles); + for ((_, items), r) in bar.iter().zip(&titles) { + // (opening and closing by click is done in `switch_open_title`) + // (the click confirming a menu a hover opened must not close it again) + let close = if confirmed == Some(egui::Popup::default_response_id(r)) { + egui::PopupCloseBehavior::IgnoreClicks + } else { + egui::PopupCloseBehavior::CloseOnClick + }; + egui::Popup::menu(r) + .close_behavior(close) + .open_memory(None) + .show(|ui| crate::menu_level::level(ui, 1, bar_bottom, |ui| nodes_ui(ui, items, mac, &mut clicked, 1, bar_bottom))); } ui.spacing_mut().item_spacing.x = saved; } else { @@ -782,6 +801,47 @@ pub fn show_in_window(app: &mut LightcraftApp, ui: &mut egui::Ui, max_width: f32 ui.cursor().left() - start } +/// Native menu-bar behaviour for the top-level titles: while one menu is open, resting the pointer +/// on another title closes it and opens that one, and a click on another title opens it in that +/// same click (instead of only closing the open menu as a click outside). Without an open menu +/// hovering does nothing; clicking a title toggles its menu. egui keeps one popup open at a time. +/// Returns the menu whose title was just clicked right after a hover opened it. +fn switch_open_title(ctx: &egui::Context, titles: &[egui::Response]) -> Option { + let popup = egui::Popup::default_response_id; + let key = egui::Id::new("lc-menubar-hover-opened"); + let open = titles.iter().position(|r| egui::Popup::is_id_open(ctx, popup(r))); + // the menu a hover opened, until it closes or is clicked + let mut confirmed = None; + let mut hover_opened: Option = ctx.data(|d| d.get_temp(key)).filter(|id| open.is_some_and(|o| popup(&titles[o]) == *id)); + for (i, r) in titles.iter().enumerate() { + if r.clicked() { + // the click that follows a hover switch confirms the menu the hover opened + if hover_opened == Some(popup(r)) { + confirmed = Some(popup(r)); + } else { + egui::Popup::toggle_id(ctx, popup(r)); + } + hover_opened = None; + } else if r.hovered() && open.is_some_and(|o| o != i) { + egui::Popup::open_id(ctx, popup(r)); + hover_opened = Some(popup(r)); + // a pointer move alone doesn't schedule the frame that draws the new menu + ctx.request_repaint(); + } + } + // (leaving the titles ends it: a later click on the open title closes the menu) + let hover_opened = hover_opened.filter(|_| titles.iter().any(|r| r.hovered())); + ctx.data_mut(|d| match hover_opened { + Some(id) => { + d.insert_temp(key, id); + } + None => { + d.remove::(key); + } + }); + confirmed +} + /// A submenu row showing `text` at `depth`, whose rows `children` draws one level deeper, /// bounded like every level. The row is registered as `menusub: