From cf2950bf996c312eddda6008603d78729d6b515c Mon Sep 17 00:00:00 2001 From: kieran Date: Tue, 14 May 2024 17:06:35 +0100 Subject: [PATCH] Automatic image labeling --- Cargo.lock | 521 ++++++++++- Cargo.toml | 3 + build.rs | 5 + config.toml | 5 +- migrations/20240514150647_user_files.sql | 26 + src/db.rs | 54 +- src/filesystem.rs | 58 +- src/processing/blurhash.rs | 42 - src/processing/labeling.rs | 1067 ++++++++++++++++++++++ src/processing/mod.rs | 59 +- src/processing/webp.rs | 38 +- src/routes/blossom.rs | 3 +- src/routes/mod.rs | 29 +- src/routes/nip96.rs | 15 +- src/settings.rs | 6 +- 15 files changed, 1791 insertions(+), 140 deletions(-) create mode 100644 build.rs create mode 100644 migrations/20240514150647_user_files.sql delete mode 100644 src/processing/blurhash.rs create mode 100644 src/processing/labeling.rs diff --git a/Cargo.lock b/Cargo.lock index 13ad357..1c11b05 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -87,6 +87,15 @@ version = "1.0.82" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f538837af36e6f6a9be0faa67f9a314f8119e4e4b5867c6ab40ed60360142519" +[[package]] +name = "arbitrary" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d5a26814d8dcb93b0e5a0ff3c6d80a8843bafb21b39e8e18a6f05471870e110" +dependencies = [ + "derive_arbitrary", +] + [[package]] name = "async-stream" version = "0.3.5" @@ -226,6 +235,21 @@ dependencies = [ "unicode-normalization", ] +[[package]] +name = "bit-set" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0700ddab506f33b20a03b13996eccd309a48e5ff77d0d95926aa0210fb4e95f1" +dependencies = [ + "bit-vec", +] + +[[package]] +name = "bit-vec" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "349f9b6a179ed607305526ca489b34ad0a41aed5f7980fa90eb03160b69598fb" + [[package]] name = "bitcoin" version = "0.31.2" @@ -317,6 +341,20 @@ name = "bytemuck" version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d6d68c57235a3a081186990eca2867354726650f42f7516ca50c28d6281fd15" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4da9a32f3fed317401fa3c862968128267c3106685286e15d5aaa3d7389c2f60" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.60", +] [[package]] name = "byteorder" @@ -330,6 +368,58 @@ version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "514de17de45fdb8dc022b1a7975556c53c86f9f0aa5f534b98977b171857c2c9" +[[package]] +name = "candle-core" +version = "0.5.1" +source = "git+https://github.com/huggingface/candle.git#13c64f6828360a9cb9b58b4f817e4f3b8316388c" +dependencies = [ + "byteorder", + "gemm", + "half", + "memmap2", + "num-traits", + "num_cpus", + "rand", + "rand_distr", + "rayon", + "safetensors", + "thiserror", + "yoke", + "zip", +] + +[[package]] +name = "candle-nn" +version = "0.5.1" +source = "git+https://github.com/huggingface/candle.git#13c64f6828360a9cb9b58b4f817e4f3b8316388c" +dependencies = [ + "candle-core", + "half", + "num-traits", + "rayon", + "safetensors", + "serde", + "thiserror", +] + +[[package]] +name = "candle-transformers" +version = "0.5.1" +source = "git+https://github.com/huggingface/candle.git#13c64f6828360a9cb9b58b4f817e4f3b8316388c" +dependencies = [ + "byteorder", + "candle-core", + "candle-nn", + "fancy-regex", + "num-traits", + "rand", + "rayon", + "serde", + "serde_json", + "serde_plain", + "tracing", +] + [[package]] name = "cbc" version = "0.1.2" @@ -517,6 +607,34 @@ version = "2.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19d374276b40fb8bbdee95aef7c7fa6b5316ec764510eb64b8dd0e2ed0d7e7f5" +[[package]] +name = "crc32fast" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3855a8a784b474f333699ef2bbca9db2c4a1f6d9088a90a2d25b1eb53111eaa" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613f8cc01fe9cf1a3eb3d7f488fd2fa8388403e97039e2f73692932e291a770d" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-queue" version = "0.3.11" @@ -569,6 +687,17 @@ dependencies = [ "powerfmt", ] +[[package]] +name = "derive_arbitrary" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67e77553c4162a157adbf834ebae5b415acbecbeafc7a74b0e886657506a7611" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.60", +] + [[package]] name = "devise" version = "0.4.1" @@ -614,6 +743,17 @@ dependencies = [ "subtle", ] +[[package]] +name = "displaydoc" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "487585f4d0c6655fe74905e2504d8ad6908e4db67f744eb140876906c2f3175d" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.60", +] + [[package]] name = "dlv-list" version = "0.5.2" @@ -629,6 +769,16 @@ version = "0.15.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" +[[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 = "either" version = "1.11.0" @@ -647,6 +797,18 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "enum-as-inner" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ffccbb6966c05b32ef8fbac435df276c4ae4d3dc55a8cd0eb9745e6c12f546a" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.60", +] + [[package]] name = "env_logger" version = "0.10.2" @@ -693,6 +855,17 @@ version = "2.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0206175f82b8d6bf6652ff7d71a1e27fd2e4efde587fd368662814d6ec1d9ce0" +[[package]] +name = "fancy-regex" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "531e46835a22af56d1e3b66f04844bed63158bc094a628bec1d321d9b4c44bf2" +dependencies = [ + "bit-set", + "regex-automata 0.4.6", + "regex-syntax 0.8.3", +] + [[package]] name = "fastrand" version = "2.0.2" @@ -846,6 +1019,124 @@ dependencies = [ "slab", ] +[[package]] +name = "gemm" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ab24cc62135b40090e31a76a9b2766a501979f3070fa27f689c27ec04377d32" +dependencies = [ + "dyn-stack", + "gemm-c32", + "gemm-c64", + "gemm-common", + "gemm-f16", + "gemm-f32", + "gemm-f64", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "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", + "gemm-common", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "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", + "gemm-common", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "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", + "half", + "num-complex", + "num-traits", + "once_cell", + "paste", + "pulp", + "raw-cpuid", + "rayon", + "seq-macro", + "sysctl", +] + +[[package]] +name = "gemm-f16" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ca4c06b9b11952071d317604acb332e924e817bd891bec8dfb494168c7cedd4" +dependencies = [ + "dyn-stack", + "gemm-common", + "gemm-f32", + "half", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "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", + "gemm-common", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "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", + "gemm-common", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + [[package]] name = "generator" version = "0.7.5" @@ -913,6 +1204,20 @@ dependencies = [ "tracing", ] +[[package]] +name = "half" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dd08c532ae367adf81c312a4580bc67f1d0fe8bc9c460520283f4c0ff277888" +dependencies = [ + "bytemuck", + "cfg-if", + "crunchy", + "num-traits", + "rand", + "rand_distr", +] + [[package]] name = "hashbrown" version = "0.13.2" @@ -1304,7 +1609,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c2a198fb6b0eada2a8df47933734e6d35d350665a33a3593d7164fa52c75c19" dependencies = [ "cfg-if", - "windows-targets 0.52.5", + "windows-targets 0.48.5", ] [[package]] @@ -1392,6 +1697,16 @@ version = "2.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6c8640c5d730cb13ebd907d8d04b52f55ac9a2eec55b440c8892f40d56c76c1d" +[[package]] +name = "memmap2" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe751422e4a8caa417e13c3ea66452215d7d63e19e604f4980461212f3ae1322" +dependencies = [ + "libc", + "stable_deref_trait", +] + [[package]] name = "mime" version = "0.3.17" @@ -1517,6 +1832,16 @@ dependencies = [ "zeroize", ] +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "bytemuck", + "num-traits", +] + [[package]] name = "num-conv" version = "0.1.0" @@ -1859,6 +2184,18 @@ dependencies = [ "yansi", ] +[[package]] +name = "pulp" +version = "0.18.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e14989307e408d9f4245d4fda09a7b144a08114ba124e26cab60ab83dc98db10" +dependencies = [ + "bytemuck", + "libm", + "num-complex", + "reborrow", +] + [[package]] name = "quote" version = "1.0.36" @@ -1898,6 +2235,51 @@ dependencies = [ "getrandom", ] +[[package]] +name = "rand_distr" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32cb0b9bc82b0a0876c2dd994a7e7a2683d3e7390ca40e6886785ef0c7e3ee31" +dependencies = [ + "num-traits", + "rand", +] + +[[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 = "rayon" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b418a60154510ca1a002a752ca9714984e21e4241e804d32555251faf8b78ffa" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1465873a3dfdaa8ae7cb14b4383657caab0b3e8a0aa9ae8e04b044854c8dfce2" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "reborrow" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03251193000f4bd3b042892be858ee50e8b3719f2b08e5833ac4353724632430" + [[package]] name = "redox_syscall" version = "0.4.1" @@ -2239,6 +2621,16 @@ version = "1.0.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e86697c916019a8588c99b5fac3cead74ec0b4b819707a682fd4d23fa0ce1ba1" +[[package]] +name = "safetensors" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ced76b22c7fba1162f11a5a75d9d8405264b467a07ae0c9c29be119b9297db9" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "salsa20" version = "0.10.2" @@ -2248,6 +2640,15 @@ dependencies = [ "cipher", ] +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + [[package]] name = "scoped-tls" version = "1.0.1" @@ -2293,6 +2694,12 @@ dependencies = [ "cc", ] +[[package]] +name = "seq-macro" +version = "0.3.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3f0bf26fd526d2a95683cd0f87bf103b8539e2ca1ef48ce002d67aad59aa0b4" + [[package]] name = "serde" version = "1.0.198" @@ -2325,6 +2732,15 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_plain" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ce1fc6db65a611022b23a0dec6975d63fb80a302cb3388835ff02c097258d50" +dependencies = [ + "serde", +] + [[package]] name = "serde_spanned" version = "0.6.5" @@ -2670,6 +3086,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "stable_deref_trait" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" + [[package]] name = "state" version = "0.6.0" @@ -2724,6 +3146,31 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2047c6ded9c721764247e62cd3b03c09ffc529b2ba5b10ec482ae507a4a70160" +[[package]] +name = "synstructure" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8af7666ab7b6390ab78131fb5b0fce11d6b7a6951602017c35fa82800708971" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.60", +] + +[[package]] +name = "sysctl" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec7dddc5f0fee506baf8b9fdb989e242f17e4b11c61dfbb0635b705217199eea" +dependencies = [ + "bitflags 2.5.0", + "byteorder", + "enum-as-inner", + "libc", + "thiserror", + "walkdir", +] + [[package]] name = "tempfile" version = "3.10.1" @@ -3175,6 +3622,9 @@ dependencies = [ "anyhow", "base64 0.21.7", "blurhash", + "candle-core", + "candle-nn", + "candle-transformers", "chrono", "config", "ffmpeg-sys-the-third", @@ -3191,6 +3641,16 @@ dependencies = [ "uuid", ] +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + [[package]] name = "want" version = "0.3.1" @@ -3532,6 +3992,30 @@ dependencies = [ "is-terminal", ] +[[package]] +name = "yoke" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65e71b2e4f287f467794c671e2b8f8a5f3716b3c829079a1c44740148eff07e4" +dependencies = [ + "serde", + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e6936f0cce458098a201c245a11bef556c6a0181129c7034d10d76d1ec3a2b8" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.60", + "synstructure", +] + [[package]] name = "zerocopy" version = "0.7.32" @@ -3552,8 +4036,43 @@ dependencies = [ "syn 2.0.60", ] +[[package]] +name = "zerofrom" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "655b0814c5c0b19ade497851070c640773304939a6c0fd5f5fb43da0696d05b7" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6a647510471d372f2e6c2e6b7219e44d8c574d24fdc11c610a61455782f18c3" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.60", + "synstructure", +] + [[package]] name = "zeroize" version = "1.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "525b4ec142c6b68a2d10f01f7bbf6755599ca3f81ea53b8431b7dd348f5fdb2d" + +[[package]] +name = "zip" +version = "1.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c700ea425e148de30c29c580c1f9508b93ca57ad31c9f4e96b83c194c37a7a8f" +dependencies = [ + "arbitrary", + "crc32fast", + "crossbeam-utils", + "displaydoc", + "indexmap", + "thiserror", +] diff --git a/Cargo.toml b/Cargo.toml index 179f27a..f946aba 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,3 +23,6 @@ chrono = { version = "0.4.38", features = ["serde"] } ffmpeg-sys-the-third = { version = "1.1.1",features = ["default"] } libc = "0.2.153" blurhash = "0.2.1" +candle-core = { git = "https://github.com/huggingface/candle.git", version = "0.5.1" } +candle-nn = { git = "https://github.com/huggingface/candle.git", version = "0.5.1" } +candle-transformers = { git = "https://github.com/huggingface/candle.git", version = "0.5.1" } diff --git a/build.rs b/build.rs new file mode 100644 index 0000000..7609593 --- /dev/null +++ b/build.rs @@ -0,0 +1,5 @@ +// generated by `sqlx migrate build-script` +fn main() { + // trigger recompilation when a new migration is added + println!("cargo:rerun-if-changed=migrations"); +} \ No newline at end of file diff --git a/config.toml b/config.toml index c5ded69..40cf28d 100644 --- a/config.toml +++ b/config.toml @@ -14,4 +14,7 @@ max_upload_bytes = 104857600 public_url = "http://localhost:8000" # Whitelisted pubkeys, leave out to disable -# whitelist = ["63fe6318dc58583cfe16810f86dd09e18bfd76aabc24a0081ce2856f330504ed"] \ No newline at end of file +# whitelist = ["63fe6318dc58583cfe16810f86dd09e18bfd76aabc24a0081ce2856f330504ed"] + +# Path for ViT image model (https://huggingface.co/google/vit-base-patch16-224) +# vit_model_path = "model.safetennsors" \ No newline at end of file diff --git a/migrations/20240514150647_user_files.sql b/migrations/20240514150647_user_files.sql new file mode 100644 index 0000000..914f827 --- /dev/null +++ b/migrations/20240514150647_user_files.sql @@ -0,0 +1,26 @@ +-- Add migration script here +alter table uploads + drop constraint fk_uploads_user; +create table user_uploads +( + file binary(32) not null, + user_id integer unsigned not null, + created timestamp default current_timestamp, + + constraint fk_user_uploads_file_id + foreign key (file) references uploads (id) + on delete cascade + on update restrict, + constraint fk_user_uploads_user_id + foreign key (user_id) references users (id) + on delete cascade + on update restrict +); +create unique index ix_user_uploads_file_pubkey on user_uploads (file, user_id); + +insert into user_uploads(file, user_id, created) +select uploads.id, uploads.user_id, uploads.created +from uploads; + +alter table uploads + drop column user_id; \ No newline at end of file diff --git a/src/db.rs b/src/db.rs index b10e9e6..d2ba53f 100644 --- a/src/db.rs +++ b/src/db.rs @@ -1,17 +1,23 @@ use chrono::{DateTime, Utc}; +use sqlx::{Error, Executor, FromRow, Row}; use sqlx::migrate::MigrateError; -use sqlx::{Error, FromRow, Row}; #[derive(Clone, FromRow)] pub struct FileUpload { pub id: Vec, - pub user_id: u64, pub name: String, pub size: u64, pub mime_type: String, pub created: DateTime, } +#[derive(Clone, FromRow)] +pub struct User { + pub id: u64, + pub pubkey: Vec, + pub created: DateTime, +} + #[derive(Clone)] pub struct Database { pool: sqlx::pool::Pool, @@ -50,28 +56,48 @@ impl Database { .try_get(0) } - pub async fn add_file(&self, file: &FileUpload) -> Result<(), Error> { - sqlx::query("insert into uploads(id,user_id,name,size,mime_type) values(?,?,?,?,?)") + pub async fn add_file(&self, file: &FileUpload, user_id: u64) -> Result<(), Error> { + let mut tx = self.pool.begin().await?; + let q = sqlx::query("insert ignore into uploads(id,name,size,mime_type) values(?,?,?,?)") .bind(&file.id) - .bind(&file.user_id) .bind(&file.name) - .bind(&file.size) - .bind(&file.mime_type) - .execute(&self.pool) - .await?; + .bind(file.size) + .bind(&file.mime_type); + let q2 = sqlx::query("insert into user_uploads(file,user_id) values(?,?)") + .bind(&file.id) + .bind(user_id); + tx.execute(q).await?; + tx.execute(q2).await?; + tx.commit().await?; Ok(()) } pub async fn get_file(&self, file: &Vec) -> Result, Error> { sqlx::query_as("select * from uploads where id = ?") - .bind(&file) + .bind(file) .fetch_optional(&self.pool) .await } + pub async fn get_file_owners(&self, file: &Vec) -> Result, Error> { + sqlx::query_as("select users.* from users, user_uploads where user.id = user_uploads.user_id and user_uploads.file = ?") + .bind(file) + .fetch_all(&self.pool) + .await + } + + pub async fn delete_file_owner(&self, file: &Vec, owner: u64) -> Result<(), Error> { + sqlx::query("delete from user_uploads where file = ? and user_id = ?") + .bind(file) + .bind(owner) + .execute(&self.pool) + .await?; + Ok(()) + } + pub async fn delete_file(&self, file: &Vec) -> Result<(), Error> { sqlx::query("delete from uploads where id = ?") - .bind(&file) + .bind(file) .execute(&self.pool) .await?; Ok(()) @@ -81,9 +107,9 @@ impl Database { let results: Vec = sqlx::query_as( "select * from uploads where user_id = (select id from users where pubkey = ?)", ) - .bind(&pubkey) - .fetch_all(&self.pool) - .await?; + .bind(pubkey) + .fetch_all(&self.pool) + .await?; Ok(results) } } diff --git a/src/filesystem.rs b/src/filesystem.rs index 005f788..cc69e75 100644 --- a/src/filesystem.rs +++ b/src/filesystem.rs @@ -11,10 +11,11 @@ use sha2::{Digest, Sha256}; use tokio::fs::File; use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeekExt}; -use crate::processing::{FileProcessor, FileProcessorResult, MediaProcessor}; +use crate::processing::{compress_file, FileProcessorResult}; +use crate::processing::labeling::label_frame; use crate::settings::Settings; -#[derive(Clone)] +#[derive(Clone, Default)] pub struct FileSystemResult { pub path: PathBuf, pub sha256: Vec, @@ -23,18 +24,17 @@ pub struct FileSystemResult { pub width: Option, pub height: Option, pub blur_hash: Option, + pub labels: Option>, } pub struct FileStore { settings: Settings, - processor: Arc>, } impl FileStore { pub fn new(settings: Settings) -> Self { Self { settings, - processor: Arc::new(Mutex::new(MediaProcessor::new())), } } @@ -86,19 +86,30 @@ impl FileStore { if compress { let start = SystemTime::now(); - let proc_result = { - let mut p_lock = self.processor.lock().expect("asd"); - p_lock.process_file(tmp_path.clone(), mime_type)? - }; - if let FileProcessorResult::NewFile(new_temp) = proc_result { + let proc_result = compress_file(tmp_path.clone(), mime_type)?; + if let FileProcessorResult::NewFile(mut new_temp) = proc_result { let old_size = tmp_path.metadata()?.len(); let new_size = new_temp.result.metadata()?.len(); - info!("Compressed media: ratio={:.2}x, old_size={:.3}kb, new_size={:.3}kb, duration={:.2}ms", - old_size as f32 / new_size as f32, - old_size as f32 / 1024.0, - new_size as f32 / 1024.0, - SystemTime::now().duration_since(start).unwrap().as_micros() as f64 / 1000.0 - ); + let time_compress = SystemTime::now().duration_since(start).unwrap(); + let start = SystemTime::now(); + let blur_hash = blurhash::encode( + 9, 9, + new_temp.width as u32, + new_temp.height as u32, + new_temp.image.as_slice(), + )?; + let time_blurhash = SystemTime::now().duration_since(start).unwrap(); + let start = SystemTime::now(); + let labels = if let Some(mp) = &self.settings.vit_model_path { + label_frame( + new_temp.image.as_mut_slice(), + new_temp.width, + new_temp.height, + mp.clone())? + } else { + vec![] + }; + let time_labels = SystemTime::now().duration_since(start).unwrap(); // delete old temp fs::remove_file(tmp_path)?; @@ -111,14 +122,25 @@ impl FileStore { .await?; let n = file.metadata().await?.len(); let hash = FileStore::hash_file(&mut file).await?; + + info!("Processed media: ratio={:.2}x, old_size={:.3}kb, new_size={:.3}kb, duration_compress={:.2}ms, duration_blurhash={:.2}ms, duration_labels={:.2}ms", + old_size as f32 / new_size as f32, + old_size as f32 / 1024.0, + new_size as f32 / 1024.0, + time_compress.as_micros() as f64 / 1000.0, + time_blurhash.as_micros() as f64 / 1000.0, + time_labels.as_micros() as f64 / 1000.0 + ); + return Ok(FileSystemResult { size: n, sha256: hash, path: new_temp.result, width: Some(new_temp.width), height: Some(new_temp.height), - blur_hash: Some(new_temp.blur_hash), + blur_hash: Some(blur_hash), mime_type: new_temp.mime_type, + labels: Some(labels), }); } } @@ -129,9 +151,7 @@ impl FileStore { sha256: hash, size: n, mime_type: mime_type.to_string(), - width: None, - height: None, - blur_hash: None, + ..Default::default() }) } diff --git a/src/processing/blurhash.rs b/src/processing/blurhash.rs deleted file mode 100644 index 224aa6e..0000000 --- a/src/processing/blurhash.rs +++ /dev/null @@ -1,42 +0,0 @@ -use std::{ptr, slice}; -use std::intrinsics::transmute; -use std::time::SystemTime; - -use anyhow::Error; -use blurhash::encode; -use ffmpeg_sys_the_third::{av_frame_alloc, av_frame_free, AVFrame, sws_freeContext, sws_getContext, sws_scale_frame}; -use ffmpeg_sys_the_third::AVPixelFormat::AV_PIX_FMT_RGBA; -use log::info; - -pub unsafe fn make_blur_hash(frame: *mut AVFrame, detail: u32) -> Result { - let start = SystemTime::now(); - let sws_ctx = sws_getContext((*frame).width, - (*frame).height, - transmute((*frame).format), - (*frame).width, - (*frame).height, - AV_PIX_FMT_RGBA, - 0, ptr::null_mut(), ptr::null_mut(), ptr::null_mut()); - if sws_ctx.is_null() { - return Err(Error::msg("Failed to create sws context")); - } - - let mut dst_frame = av_frame_alloc(); - let ret = sws_scale_frame(sws_ctx, dst_frame, frame); - if ret < 0 { - return Err(Error::msg("Failed to scale frame (blurhash)")); - } - - let pic_slice = slice::from_raw_parts_mut((*dst_frame).data[0], ((*frame).width * (*frame).height * 4) as usize); - let bh = encode(detail, detail, - (*frame).width as u32, - (*frame).height as u32, - pic_slice, - )?; - - av_frame_free(&mut dst_frame); - sws_freeContext(sws_ctx); - - info!("Generated blurhash in {}ms", SystemTime::now().duration_since(start).unwrap().as_millis()); - Ok(bh) -} \ No newline at end of file diff --git a/src/processing/labeling.rs b/src/processing/labeling.rs new file mode 100644 index 0000000..501651b --- /dev/null +++ b/src/processing/labeling.rs @@ -0,0 +1,1067 @@ +use std::{fs, ptr, slice}; +use std::mem::transmute; +use std::path::PathBuf; + +use anyhow::Error; +use candle_core::{D, Device, DType, IndexOp, Tensor}; +use candle_nn::VarBuilder; +use candle_transformers::models::vit; +use ffmpeg_sys_the_third::{av_frame_alloc, av_frame_free}; +use ffmpeg_sys_the_third::AVColorRange::AVCOL_RANGE_JPEG; +use ffmpeg_sys_the_third::AVColorSpace::AVCOL_SPC_RGB; +use ffmpeg_sys_the_third::AVPixelFormat::{AV_PIX_FMT_RGB24, AV_PIX_FMT_RGBA}; + +use crate::processing::resize_image; + +pub fn label_frame(frame: &mut [u8], width: usize, height: usize, model: PathBuf) -> Result, Error> { + unsafe { + let device = Device::Cpu; + let image = load_frame_224(frame, width, height)?.to_device(&device)?; + + let vb = VarBuilder::from_mmaped_safetensors(&[model], DType::F32, &device)?; + let model = vit::Model::new(&vit::Config::vit_base_patch16_224(), 1000, vb)?; + let logits = model.forward(&image.unsqueeze(0)?)?; + let prs = candle_nn::ops::softmax(&logits, D::Minus1)? + .i(0)? + .to_vec1::()?; + let mut prs = prs.iter().enumerate().collect::>(); + prs.sort_by(|(_, p1), (_, p2)| p2.total_cmp(p1)); + let res = prs.iter() + //.filter(|&(_c, q)| **q >= 0.50f32) + .take(5) + .map(|&(c, _q)| CLASSES[c].to_string()).collect(); + Ok(res) + } +} + +unsafe fn load_frame_224(data: &mut [u8], width: usize, height: usize) -> Result { + let frame = av_frame_alloc(); + (*frame).extended_data = &mut data.as_mut_ptr(); + (*frame).data = [*(*frame).extended_data, ptr::null_mut(), ptr::null_mut(), ptr::null_mut(), ptr::null_mut(), ptr::null_mut(), ptr::null_mut(), ptr::null_mut()]; + (*frame).linesize = [(width * 4) as libc::c_int, 0, 0, 0, 0, 0, 0, 0]; + (*frame).format = transmute(AV_PIX_FMT_RGBA); + (*frame).width = width as libc::c_int; + (*frame).height = height as libc::c_int; + (*frame).color_range = AVCOL_RANGE_JPEG; + (*frame).colorspace = AVCOL_SPC_RGB; + + let mut dst_frame = resize_image(frame, + 224, + 224, + AV_PIX_FMT_RGB24)?; + let pic_slice = slice::from_raw_parts_mut((*dst_frame).data[0], ((*dst_frame).width * (*dst_frame).height * 3) as usize); + + fs::write("frame_224.raw", &pic_slice)?; + let data = Tensor::from_vec(pic_slice.to_vec(), (224, 224, 3), &Device::Cpu)?.permute((2, 0, 1))?; + let mean = Tensor::new(&[0.485f32, 0.456, 0.406], &Device::Cpu)?.reshape((3, 1, 1))?; + let std = Tensor::new(&[0.229f32, 0.224, 0.225], &Device::Cpu)?.reshape((3, 1, 1))?; + let res = (data.to_dtype(DType::F32)? / 255.)? + .broadcast_sub(&mean)? + .broadcast_div(&std)?; + av_frame_free(&mut dst_frame); + Ok(res) +} + + +pub const CLASSES: [&str; 1000] = [ + "tench, Tinca tinca", + "goldfish, Carassius auratus", + "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias", + "tiger shark, Galeocerdo cuvieri", + "hammerhead, hammerhead shark", + "electric ray, crampfish, numbfish, torpedo", + "stingray", + "cock", + "hen", + "ostrich, Struthio camelus", + "brambling, Fringilla montifringilla", + "goldfinch, Carduelis carduelis", + "house finch, linnet, Carpodacus mexicanus", + "junco, snowbird", + "indigo bunting, indigo finch, indigo bird, Passerina cyanea", + "robin, American robin, Turdus migratorius", + "bulbul", + "jay", + "magpie", + "chickadee", + "water ouzel, dipper", + "kite", + "bald eagle, American eagle, Haliaeetus leucocephalus", + "vulture", + "great grey owl, great gray owl, Strix nebulosa", + "European fire salamander, Salamandra salamandra", + "common newt, Triturus vulgaris", + "eft", + "spotted salamander, Ambystoma maculatum", + "axolotl, mud puppy, Ambystoma mexicanum", + "bullfrog, Rana catesbeiana", + "tree frog, tree-frog", + "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui", + "loggerhead, loggerhead turtle, Caretta caretta", + "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea", + "mud turtle", + "terrapin", + "box turtle, box tortoise", + "banded gecko", + "common iguana, iguana, Iguana iguana", + "American chameleon, anole, Anolis carolinensis", + "whiptail, whiptail lizard", + "agama", + "frilled lizard, Chlamydosaurus kingi", + "alligator lizard", + "Gila monster, Heloderma suspectum", + "green lizard, Lacerta viridis", + "African chameleon, Chamaeleo chamaeleon", + "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis", + "African crocodile, Nile crocodile, Crocodylus niloticus", + "American alligator, Alligator mississipiensis", + "triceratops", + "thunder snake, worm snake, Carphophis amoenus", + "ringneck snake, ring-necked snake, ring snake", + "hognose snake, puff adder, sand viper", + "green snake, grass snake", + "king snake, kingsnake", + "garter snake, grass snake", + "water snake", + "vine snake", + "night snake, Hypsiglena torquata", + "boa constrictor, Constrictor constrictor", + "rock python, rock snake, Python sebae", + "Indian cobra, Naja naja", + "green mamba", + "sea snake", + "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus", + "diamondback, diamondback rattlesnake, Crotalus adamanteus", + "sidewinder, horned rattlesnake, Crotalus cerastes", + "trilobite", + "harvestman, daddy longlegs, Phalangium opilio", + "scorpion", + "black and gold garden spider, Argiope aurantia", + "barn spider, Araneus cavaticus", + "garden spider, Aranea diademata", + "black widow, Latrodectus mactans", + "tarantula", + "wolf spider, hunting spider", + "tick", + "centipede", + "black grouse", + "ptarmigan", + "ruffed grouse, partridge, Bonasa umbellus", + "prairie chicken, prairie grouse, prairie fowl", + "peacock", + "quail", + "partridge", + "African grey, African gray, Psittacus erithacus", + "macaw", + "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita", + "lorikeet", + "coucal", + "bee eater", + "hornbill", + "hummingbird", + "jacamar", + "toucan", + "drake", + "red-breasted merganser, Mergus serrator", + "goose", + "black swan, Cygnus atratus", + "tusker", + "echidna, spiny anteater, anteater", + "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus", + "wallaby, brush kangaroo", + "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus", + "wombat", + "jellyfish", + "sea anemone, anemone", + "brain coral", + "flatworm, platyhelminth", + "nematode, nematode worm, roundworm", + "conch", + "snail", + "slug", + "sea slug, nudibranch", + "chiton, coat-of-mail shell, sea cradle, polyplacophore", + "chambered nautilus, pearly nautilus, nautilus", + "Dungeness crab, Cancer magister", + "rock crab, Cancer irroratus", + "fiddler crab", + "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica", + "American lobster, Northern lobster, Maine lobster, Homarus americanus", + "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish", + "crayfish, crawfish, crawdad, crawdaddy", + "hermit crab", + "isopod", + "white stork, Ciconia ciconia", + "black stork, Ciconia nigra", + "spoonbill", + "flamingo", + "little blue heron, Egretta caerulea", + "American egret, great white heron, Egretta albus", + "bittern", + "crane", + "limpkin, Aramus pictus", + "European gallinule, Porphyrio porphyrio", + "American coot, marsh hen, mud hen, water hen, Fulica americana", + "bustard", + "ruddy turnstone, Arenaria interpres", + "red-backed sandpiper, dunlin, Erolia alpina", + "redshank, Tringa totanus", + "dowitcher", + "oystercatcher, oyster catcher", + "pelican", + "king penguin, Aptenodytes patagonica", + "albatross, mollymawk", + "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus", + "killer whale, killer, orca, grampus, sea wolf, Orcinus orca", + "dugong, Dugong dugon", + "sea lion", + "Chihuahua", + "Japanese spaniel", + "Maltese dog, Maltese terrier, Maltese", + "Pekinese, Pekingese, Peke", + "Shih-Tzu", + "Blenheim spaniel", + "papillon", + "toy terrier", + "Rhodesian ridgeback", + "Afghan hound, Afghan", + "basset, basset hound", + "beagle", + "bloodhound, sleuthhound", + "bluetick", + "black-and-tan coonhound", + "Walker hound, Walker foxhound", + "English foxhound", + "redbone", + "borzoi, Russian wolfhound", + "Irish wolfhound", + "Italian greyhound", + "whippet", + "Ibizan hound, Ibizan Podenco", + "Norwegian elkhound, elkhound", + "otterhound, otter hound", + "Saluki, gazelle hound", + "Scottish deerhound, deerhound", + "Weimaraner", + "Staffordshire bullterrier, Staffordshire bull terrier", + "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier", + "Bedlington terrier", + "Border terrier", + "Kerry blue terrier", + "Irish terrier", + "Norfolk terrier", + "Norwich terrier", + "Yorkshire terrier", + "wire-haired fox terrier", + "Lakeland terrier", + "Sealyham terrier, Sealyham", + "Airedale, Airedale terrier", + "cairn, cairn terrier", + "Australian terrier", + "Dandie Dinmont, Dandie Dinmont terrier", + "Boston bull, Boston terrier", + "miniature schnauzer", + "giant schnauzer", + "standard schnauzer", + "Scotch terrier, Scottish terrier, Scottie", + "Tibetan terrier, chrysanthemum dog", + "silky terrier, Sydney silky", + "soft-coated wheaten terrier", + "West Highland white terrier", + "Lhasa, Lhasa apso", + "flat-coated retriever", + "curly-coated retriever", + "golden retriever", + "Labrador retriever", + "Chesapeake Bay retriever", + "German short-haired pointer", + "vizsla, Hungarian pointer", + "English setter", + "Irish setter, red setter", + "Gordon setter", + "Brittany spaniel", + "clumber, clumber spaniel", + "English springer, English springer spaniel", + "Welsh springer spaniel", + "cocker spaniel, English cocker spaniel, cocker", + "Sussex spaniel", + "Irish water spaniel", + "kuvasz", + "schipperke", + "groenendael", + "malinois", + "briard", + "kelpie", + "komondor", + "Old English sheepdog, bobtail", + "Shetland sheepdog, Shetland sheep dog, Shetland", + "collie", + "Border collie", + "Bouvier des Flandres, Bouviers des Flandres", + "Rottweiler", + "German shepherd, German shepherd dog, German police dog, alsatian", + "Doberman, Doberman pinscher", + "miniature pinscher", + "Greater Swiss Mountain dog", + "Bernese mountain dog", + "Appenzeller", + "EntleBucher", + "boxer", + "bull mastiff", + "Tibetan mastiff", + "French bulldog", + "Great Dane", + "Saint Bernard, St Bernard", + "Eskimo dog, husky", + "malamute, malemute, Alaskan malamute", + "Siberian husky", + "dalmatian, coach dog, carriage dog", + "affenpinscher, monkey pinscher, monkey dog", + "basenji", + "pug, pug-dog", + "Leonberg", + "Newfoundland, Newfoundland dog", + "Great Pyrenees", + "Samoyed, Samoyede", + "Pomeranian", + "chow, chow chow", + "keeshond", + "Brabancon griffon", + "Pembroke, Pembroke Welsh corgi", + "Cardigan, Cardigan Welsh corgi", + "toy poodle", + "miniature poodle", + "standard poodle", + "Mexican hairless", + "timber wolf, grey wolf, gray wolf, Canis lupus", + "white wolf, Arctic wolf, Canis lupus tundrarum", + "red wolf, maned wolf, Canis rufus, Canis niger", + "coyote, prairie wolf, brush wolf, Canis latrans", + "dingo, warrigal, warragal, Canis dingo", + "dhole, Cuon alpinus", + "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus", + "hyena, hyaena", + "red fox, Vulpes vulpes", + "kit fox, Vulpes macrotis", + "Arctic fox, white fox, Alopex lagopus", + "grey fox, gray fox, Urocyon cinereoargenteus", + "tabby, tabby cat", + "tiger cat", + "Persian cat", + "Siamese cat, Siamese", + "Egyptian cat", + "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor", + "lynx, catamount", + "leopard, Panthera pardus", + "snow leopard, ounce, Panthera uncia", + "jaguar, panther, Panthera onca, Felis onca", + "lion, king of beasts, Panthera leo", + "tiger, Panthera tigris", + "cheetah, chetah, Acinonyx jubatus", + "brown bear, bruin, Ursus arctos", + "American black bear, black bear, Ursus americanus, Euarctos americanus", + "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus", + "sloth bear, Melursus ursinus, Ursus ursinus", + "mongoose", + "meerkat, mierkat", + "tiger beetle", + "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle", + "ground beetle, carabid beetle", + "long-horned beetle, longicorn, longicorn beetle", + "leaf beetle, chrysomelid", + "dung beetle", + "rhinoceros beetle", + "weevil", + "fly", + "bee", + "ant, emmet, pismire", + "grasshopper, hopper", + "cricket", + "walking stick, walkingstick, stick insect", + "cockroach, roach", + "mantis, mantid", + "cicada, cicala", + "leafhopper", + "lacewing, lacewing fly", + "dragonfly, darning needle, devil's darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk", + "damselfly", + "admiral", + "ringlet, ringlet butterfly", + "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus", + "cabbage butterfly", + "sulphur butterfly, sulfur butterfly", + "lycaenid, lycaenid butterfly", + "starfish, sea star", + "sea urchin", + "sea cucumber, holothurian", + "wood rabbit, cottontail, cottontail rabbit", + "hare", + "Angora, Angora rabbit", + "hamster", + "porcupine, hedgehog", + "fox squirrel, eastern fox squirrel, Sciurus niger", + "marmot", + "beaver", + "guinea pig, Cavia cobaya", + "sorrel", + "zebra", + "hog, pig, grunter, squealer, Sus scrofa", + "wild boar, boar, Sus scrofa", + "warthog", + "hippopotamus, hippo, river horse, Hippopotamus amphibius", + "ox", + "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis", + "bison", + "ram, tup", + "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis", + "ibex, Capra ibex", + "hartebeest", + "impala, Aepyceros melampus", + "gazelle", + "Arabian camel, dromedary, Camelus dromedarius", + "llama", + "weasel", + "mink", + "polecat, fitch, foulmart, foumart, Mustela putorius", + "black-footed ferret, ferret, Mustela nigripes", + "otter", + "skunk, polecat, wood pussy", + "badger", + "armadillo", + "three-toed sloth, ai, Bradypus tridactylus", + "orangutan, orang, orangutang, Pongo pygmaeus", + "gorilla, Gorilla gorilla", + "chimpanzee, chimp, Pan troglodytes", + "gibbon, Hylobates lar", + "siamang, Hylobates syndactylus, Symphalangus syndactylus", + "guenon, guenon monkey", + "patas, hussar monkey, Erythrocebus patas", + "baboon", + "macaque", + "langur", + "colobus, colobus monkey", + "proboscis monkey, Nasalis larvatus", + "marmoset", + "capuchin, ringtail, Cebus capucinus", + "howler monkey, howler", + "titi, titi monkey", + "spider monkey, Ateles geoffroyi", + "squirrel monkey, Saimiri sciureus", + "Madagascar cat, ring-tailed lemur, Lemur catta", + "indri, indris, Indri indri, Indri brevicaudatus", + "Indian elephant, Elephas maximus", + "African elephant, Loxodonta africana", + "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens", + "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca", + "barracouta, snoek", + "eel", + "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch", + "rock beauty, Holocanthus tricolor", + "anemone fish", + "sturgeon", + "gar, garfish, garpike, billfish, Lepisosteus osseus", + "lionfish", + "puffer, pufferfish, blowfish, globefish", + "abacus", + "abaya", + "academic gown, academic robe, judge's robe", + "accordion, piano accordion, squeeze box", + "acoustic guitar", + "aircraft carrier, carrier, flattop, attack aircraft carrier", + "airliner", + "airship, dirigible", + "altar", + "ambulance", + "amphibian, amphibious vehicle", + "analog clock", + "apiary, bee house", + "apron", + "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin", + "assault rifle, assault gun", + "backpack, back pack, knapsack, packsack, rucksack, haversack", + "bakery, bakeshop, bakehouse", + "balance beam, beam", + "balloon", + "ballpoint, ballpoint pen, ballpen, Biro", + "Band Aid", + "banjo", + "bannister, banister, balustrade, balusters, handrail", + "barbell", + "barber chair", + "barbershop", + "barn", + "barometer", + "barrel, cask", + "barrow, garden cart, lawn cart, wheelbarrow", + "baseball", + "basketball", + "bassinet", + "bassoon", + "bathing cap, swimming cap", + "bath towel", + "bathtub, bathing tub, bath, tub", + "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon", + "beacon, lighthouse, beacon light, pharos", + "beaker", + "bearskin, busby, shako", + "beer bottle", + "beer glass", + "bell cote, bell cot", + "bib", + "bicycle-built-for-two, tandem bicycle, tandem", + "bikini, two-piece", + "binder, ring-binder", + "binoculars, field glasses, opera glasses", + "birdhouse", + "boathouse", + "bobsled, bobsleigh, bob", + "bolo tie, bolo, bola tie, bola", + "bonnet, poke bonnet", + "bookcase", + "bookshop, bookstore, bookstall", + "bottlecap", + "bow", + "bow tie, bow-tie, bowtie", + "brass, memorial tablet, plaque", + "brassiere, bra, bandeau", + "breakwater, groin, groyne, mole, bulwark, seawall, jetty", + "breastplate, aegis, egis", + "broom", + "bucket, pail", + "buckle", + "bulletproof vest", + "bullet train, bullet", + "butcher shop, meat market", + "cab, hack, taxi, taxicab", + "caldron, cauldron", + "candle, taper, wax light", + "cannon", + "canoe", + "can opener, tin opener", + "cardigan", + "car mirror", + "carousel, carrousel, merry-go-round, roundabout, whirligig", + "carpenter's kit, tool kit", + "carton", + "car wheel", + "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM", + "cassette", + "cassette player", + "castle", + "catamaran", + "CD player", + "cello, violoncello", + "cellular telephone, cellular phone, cellphone, cell, mobile phone", + "chain", + "chainlink fence", + "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour", + "chain saw, chainsaw", + "chest", + "chiffonier, commode", + "chime, bell, gong", + "china cabinet, china closet", + "Christmas stocking", + "church, church building", + "cinema, movie theater, movie theatre, movie house, picture palace", + "cleaver, meat cleaver, chopper", + "cliff dwelling", + "cloak", + "clog, geta, patten, sabot", + "cocktail shaker", + "coffee mug", + "coffeepot", + "coil, spiral, volute, whorl, helix", + "combination lock", + "computer keyboard, keypad", + "confectionery, confectionary, candy store", + "container ship, containership, container vessel", + "convertible", + "corkscrew, bottle screw", + "cornet, horn, trumpet, trump", + "cowboy boot", + "cowboy hat, ten-gallon hat", + "cradle", + "crane", + "crash helmet", + "crate", + "crib, cot", + "Crock Pot", + "croquet ball", + "crutch", + "cuirass", + "dam, dike, dyke", + "desk", + "desktop computer", + "dial telephone, dial phone", + "diaper, nappy, napkin", + "digital clock", + "digital watch", + "dining table, board", + "dishrag, dishcloth", + "dishwasher, dish washer, dishwashing machine", + "disk brake, disc brake", + "dock, dockage, docking facility", + "dogsled, dog sled, dog sleigh", + "dome", + "doormat, welcome mat", + "drilling platform, offshore rig", + "drum, membranophone, tympan", + "drumstick", + "dumbbell", + "Dutch oven", + "electric fan, blower", + "electric guitar", + "electric locomotive", + "entertainment center", + "envelope", + "espresso maker", + "face powder", + "feather boa, boa", + "file, file cabinet, filing cabinet", + "fireboat", + "fire engine, fire truck", + "fire screen, fireguard", + "flagpole, flagstaff", + "flute, transverse flute", + "folding chair", + "football helmet", + "forklift", + "fountain", + "fountain pen", + "four-poster", + "freight car", + "French horn, horn", + "frying pan, frypan, skillet", + "fur coat", + "garbage truck, dustcart", + "gasmask, respirator, gas helmet", + "gas pump, gasoline pump, petrol pump, island dispenser", + "goblet", + "go-kart", + "golf ball", + "golfcart, golf cart", + "gondola", + "gong, tam-tam", + "gown", + "grand piano, grand", + "greenhouse, nursery, glasshouse", + "grille, radiator grille", + "grocery store, grocery, food market, market", + "guillotine", + "hair slide", + "hair spray", + "half track", + "hammer", + "hamper", + "hand blower, blow dryer, blow drier, hair dryer, hair drier", + "hand-held computer, hand-held microcomputer", + "handkerchief, hankie, hanky, hankey", + "hard disc, hard disk, fixed disk", + "harmonica, mouth organ, harp, mouth harp", + "harp", + "harvester, reaper", + "hatchet", + "holster", + "home theater, home theatre", + "honeycomb", + "hook, claw", + "hoopskirt, crinoline", + "horizontal bar, high bar", + "horse cart, horse-cart", + "hourglass", + "iPod", + "iron, smoothing iron", + "jack-o'-lantern", + "jean, blue jean, denim", + "jeep, landrover", + "jersey, T-shirt, tee shirt", + "jigsaw puzzle", + "jinrikisha, ricksha, rickshaw", + "joystick", + "kimono", + "knee pad", + "knot", + "lab coat, laboratory coat", + "ladle", + "lampshade, lamp shade", + "laptop, laptop computer", + "lawn mower, mower", + "lens cap, lens cover", + "letter opener, paper knife, paperknife", + "library", + "lifeboat", + "lighter, light, igniter, ignitor", + "limousine, limo", + "liner, ocean liner", + "lipstick, lip rouge", + "Loafer", + "lotion", + "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system", + "loupe, jeweler's loupe", + "lumbermill, sawmill", + "magnetic compass", + "mailbag, postbag", + "mailbox, letter box", + "maillot", + "maillot, tank suit", + "manhole cover", + "maraca", + "marimba, xylophone", + "mask", + "matchstick", + "maypole", + "maze, labyrinth", + "measuring cup", + "medicine chest, medicine cabinet", + "megalith, megalithic structure", + "microphone, mike", + "microwave, microwave oven", + "military uniform", + "milk can", + "minibus", + "miniskirt, mini", + "minivan", + "missile", + "mitten", + "mixing bowl", + "mobile home, manufactured home", + "Model T", + "modem", + "monastery", + "monitor", + "moped", + "mortar", + "mortarboard", + "mosque", + "mosquito net", + "motor scooter, scooter", + "mountain bike, all-terrain bike, off-roader", + "mountain tent", + "mouse, computer mouse", + "mousetrap", + "moving van", + "muzzle", + "nail", + "neck brace", + "necklace", + "nipple", + "notebook, notebook computer", + "obelisk", + "oboe, hautboy, hautbois", + "ocarina, sweet potato", + "odometer, hodometer, mileometer, milometer", + "oil filter", + "organ, pipe organ", + "oscilloscope, scope, cathode-ray oscilloscope, CRO", + "overskirt", + "oxcart", + "oxygen mask", + "packet", + "paddle, boat paddle", + "paddlewheel, paddle wheel", + "padlock", + "paintbrush", + "pajama, pyjama, pj's, jammies", + "palace", + "panpipe, pandean pipe, syrinx", + "paper towel", + "parachute, chute", + "parallel bars, bars", + "park bench", + "parking meter", + "passenger car, coach, carriage", + "patio, terrace", + "pay-phone, pay-station", + "pedestal, plinth, footstall", + "pencil box, pencil case", + "pencil sharpener", + "perfume, essence", + "Petri dish", + "photocopier", + "pick, plectrum, plectron", + "pickelhaube", + "picket fence, paling", + "pickup, pickup truck", + "pier", + "piggy bank, penny bank", + "pill bottle", + "pillow", + "ping-pong ball", + "pinwheel", + "pirate, pirate ship", + "pitcher, ewer", + "plane, carpenter's plane, woodworking plane", + "planetarium", + "plastic bag", + "plate rack", + "plow, plough", + "plunger, plumber's helper", + "Polaroid camera, Polaroid Land camera", + "pole", + "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria", + "poncho", + "pool table, billiard table, snooker table", + "pop bottle, soda bottle", + "pot, flowerpot", + "potter's wheel", + "power drill", + "prayer rug, prayer mat", + "printer", + "prison, prison house", + "projectile, missile", + "projector", + "puck, hockey puck", + "punching bag, punch bag, punching ball, punchball", + "purse", + "quill, quill pen", + "quilt, comforter, comfort, puff", + "racer, race car, racing car", + "racket, racquet", + "radiator", + "radio, wireless", + "radio telescope, radio reflector", + "rain barrel", + "recreational vehicle, RV, R.V.", + "reel", + "reflex camera", + "refrigerator, icebox", + "remote control, remote", + "restaurant, eating house, eating place, eatery", + "revolver, six-gun, six-shooter", + "rifle", + "rocking chair, rocker", + "rotisserie", + "rubber eraser, rubber, pencil eraser", + "rugby ball", + "rule, ruler", + "running shoe", + "safe", + "safety pin", + "saltshaker, salt shaker", + "sandal", + "sarong", + "sax, saxophone", + "scabbard", + "scale, weighing machine", + "school bus", + "schooner", + "scoreboard", + "screen, CRT screen", + "screw", + "screwdriver", + "seat belt, seatbelt", + "sewing machine", + "shield, buckler", + "shoe shop, shoe-shop, shoe store", + "shoji", + "shopping basket", + "shopping cart", + "shovel", + "shower cap", + "shower curtain", + "ski", + "ski mask", + "sleeping bag", + "slide rule, slipstick", + "sliding door", + "slot, one-armed bandit", + "snorkel", + "snowmobile", + "snowplow, snowplough", + "soap dispenser", + "soccer ball", + "sock", + "solar dish, solar collector, solar furnace", + "sombrero", + "soup bowl", + "space bar", + "space heater", + "space shuttle", + "spatula", + "speedboat", + "spider web, spider's web", + "spindle", + "sports car, sport car", + "spotlight, spot", + "stage", + "steam locomotive", + "steel arch bridge", + "steel drum", + "stethoscope", + "stole", + "stone wall", + "stopwatch, stop watch", + "stove", + "strainer", + "streetcar, tram, tramcar, trolley, trolley car", + "stretcher", + "studio couch, day bed", + "stupa, tope", + "submarine, pigboat, sub, U-boat", + "suit, suit of clothes", + "sundial", + "sunglass", + "sunglasses, dark glasses, shades", + "sunscreen, sunblock, sun blocker", + "suspension bridge", + "swab, swob, mop", + "sweatshirt", + "swimming trunks, bathing trunks", + "swing", + "switch, electric switch, electrical switch", + "syringe", + "table lamp", + "tank, army tank, armored combat vehicle, armoured combat vehicle", + "tape player", + "teapot", + "teddy, teddy bear", + "television, television system", + "tennis ball", + "thatch, thatched roof", + "theater curtain, theatre curtain", + "thimble", + "thresher, thrasher, threshing machine", + "throne", + "tile roof", + "toaster", + "tobacco shop, tobacconist shop, tobacconist", + "toilet seat", + "torch", + "totem pole", + "tow truck, tow car, wrecker", + "toyshop", + "tractor", + "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi", + "tray", + "trench coat", + "tricycle, trike, velocipede", + "trimaran", + "tripod", + "triumphal arch", + "trolleybus, trolley coach, trackless trolley", + "trombone", + "tub, vat", + "turnstile", + "typewriter keyboard", + "umbrella", + "unicycle, monocycle", + "upright, upright piano", + "vacuum, vacuum cleaner", + "vase", + "vault", + "velvet", + "vending machine", + "vestment", + "viaduct", + "violin, fiddle", + "volleyball", + "waffle iron", + "wall clock", + "wallet, billfold, notecase, pocketbook", + "wardrobe, closet, press", + "warplane, military plane", + "washbasin, handbasin, washbowl, lavabo, wash-hand basin", + "washer, automatic washer, washing machine", + "water bottle", + "water jug", + "water tower", + "whiskey jug", + "whistle", + "wig", + "window screen", + "window shade", + "Windsor tie", + "wine bottle", + "wing", + "wok", + "wooden spoon", + "wool, woolen, woollen", + "worm fence, snake fence, snake-rail fence, Virginia fence", + "wreck", + "yawl", + "yurt", + "web site, website, internet site, site", + "comic book", + "crossword puzzle, crossword", + "street sign", + "traffic light, traffic signal, stoplight", + "book jacket, dust cover, dust jacket, dust wrapper", + "menu", + "plate", + "guacamole", + "consomme", + "hot pot, hotpot", + "trifle", + "ice cream, icecream", + "ice lolly, lolly, lollipop, popsicle", + "French loaf", + "bagel, beigel", + "pretzel", + "cheeseburger", + "hotdog, hot dog, red hot", + "mashed potato", + "head cabbage", + "broccoli", + "cauliflower", + "zucchini, courgette", + "spaghetti squash", + "acorn squash", + "butternut squash", + "cucumber, cuke", + "artichoke, globe artichoke", + "bell pepper", + "cardoon", + "mushroom", + "Granny Smith", + "strawberry", + "orange", + "lemon", + "fig", + "pineapple, ananas", + "banana", + "jackfruit, jak, jack", + "custard apple", + "pomegranate", + "hay", + "carbonara", + "chocolate sauce, chocolate syrup", + "dough", + "meat loaf, meatloaf", + "pizza, pizza pie", + "potpie", + "burrito", + "red wine", + "espresso", + "cup", + "eggnog", + "alp", + "bubble", + "cliff, drop, drop-off", + "coral reef", + "geyser", + "lakeside, lakeshore", + "promontory, headland, head, foreland", + "sandbar, sand bar", + "seashore, coast, seacoast, sea-coast", + "valley, vale", + "volcano", + "ballplayer, baseball player", + "groom, bridegroom", + "scuba diver", + "rapeseed", + "daisy", + "yellow lady's slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum", + "corn", + "acorn", + "hip, rose hip, rosehip", + "buckeye, horse chestnut, conker", + "coral fungus", + "agaric", + "gyromitra", + "stinkhorn, carrion fungus", + "earthstar", + "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa", + "bolete", + "ear, spike, capitulum", + "toilet tissue, toilet paper, bathroom tissue", +]; \ No newline at end of file diff --git a/src/processing/mod.rs b/src/processing/mod.rs index 69267b9..230d47f 100644 --- a/src/processing/mod.rs +++ b/src/processing/mod.rs @@ -1,11 +1,14 @@ +use std::intrinsics::transmute; use std::path::PathBuf; +use std::ptr; use anyhow::Error; +use ffmpeg_sys_the_third::{av_frame_alloc, AVFrame, AVPixelFormat, sws_freeContext, sws_getContext, sws_scale_frame}; use crate::processing::webp::WebpProcessor; mod webp; -mod blurhash; +pub mod labeling; pub(crate) enum FileProcessorResult { NewFile(NewFileProcessorResult), @@ -17,32 +20,42 @@ pub(crate) struct NewFileProcessorResult { pub mime_type: String, pub width: usize, pub height: usize, - pub blur_hash: String, + + /// The image as RBGA + pub image: Vec, } -pub(crate) trait FileProcessor { - fn process_file(&mut self, in_file: PathBuf, mime_type: &str) -> Result; -} - -pub(crate) struct MediaProcessor {} - -impl MediaProcessor { - pub fn new() -> Self { - Self {} +pub fn compress_file(in_file: PathBuf, mime_type: &str) -> Result { + let proc = if mime_type.starts_with("image/") { + Some(WebpProcessor::new()) + } else { + None + }; + if let Some(mut proc) = proc { + proc.process_file(in_file, mime_type) + } else { + Ok(FileProcessorResult::Skip) } } -impl FileProcessor for MediaProcessor { - fn process_file(&mut self, in_file: PathBuf, mime_type: &str) -> Result { - let proc = if mime_type.starts_with("image/") { - Some(WebpProcessor::new()) - } else { - None - }; - if let Some(mut proc) = proc { - proc.process_file(in_file, mime_type) - } else { - Ok(FileProcessorResult::Skip) - } +unsafe fn resize_image(frame: *const AVFrame, width: usize, height: usize, pix_fmt: AVPixelFormat) -> Result<*mut AVFrame, Error> { + let sws_ctx = sws_getContext((*frame).width, + (*frame).height, + transmute((*frame).format), + width as libc::c_int, + height as libc::c_int, + pix_fmt, + 0, ptr::null_mut(), ptr::null_mut(), ptr::null_mut()); + if sws_ctx.is_null() { + return Err(Error::msg("Failed to create sws context")); } + + let dst_frame = av_frame_alloc(); + let ret = sws_scale_frame(sws_ctx, dst_frame, frame); + if ret < 0 { + return Err(Error::msg("Failed to scale frame")); + } + + sws_freeContext(sws_ctx); + Ok(dst_frame) } \ No newline at end of file diff --git a/src/processing/webp.rs b/src/processing/webp.rs index 12e0247..5c46051 100644 --- a/src/processing/webp.rs +++ b/src/processing/webp.rs @@ -1,16 +1,15 @@ +use std::{ptr, slice}; use std::collections::HashMap; use std::mem::transmute; use std::path::PathBuf; -use std::ptr; use anyhow::Error; use ffmpeg_sys_the_third::{AV_CODEC_FLAG_GLOBAL_HEADER, av_dump_format, av_find_best_stream, av_frame_alloc, av_frame_copy_props, av_frame_free, av_guess_format, av_interleaved_write_frame, av_packet_alloc, av_packet_free, av_packet_rescale_ts, av_packet_unref, AV_PROFILE_H264_HIGH, av_read_frame, av_write_trailer, AVCodec, avcodec_alloc_context3, avcodec_find_encoder, avcodec_free_context, avcodec_open2, avcodec_parameters_from_context, avcodec_parameters_to_context, avcodec_receive_frame, avcodec_receive_packet, avcodec_send_frame, avcodec_send_packet, AVCodecContext, AVCodecID, AVERROR, AVERROR_EOF, AVERROR_STREAM_NOT_FOUND, AVFMT_GLOBALHEADER, avformat_alloc_output_context2, avformat_close_input, avformat_find_stream_info, avformat_free_context, avformat_init_output, avformat_new_stream, avformat_open_input, avformat_write_header, AVFormatContext, AVIO_FLAG_WRITE, avio_open, AVMediaType, AVPacket, sws_freeContext, sws_getContext, sws_scale_frame, SwsContext}; use ffmpeg_sys_the_third::AVMediaType::{AVMEDIA_TYPE_AUDIO, AVMEDIA_TYPE_VIDEO}; -use ffmpeg_sys_the_third::AVPixelFormat::AV_PIX_FMT_YUV420P; +use ffmpeg_sys_the_third::AVPixelFormat::{AV_PIX_FMT_RGBA, AV_PIX_FMT_YUV420P}; use libc::EAGAIN; -use crate::processing::{FileProcessor, FileProcessorResult, NewFileProcessorResult}; -use crate::processing::blurhash::make_blur_hash; +use crate::processing::{FileProcessorResult, NewFileProcessorResult, resize_image}; /// Image converter to WEBP pub struct WebpProcessor { @@ -20,7 +19,7 @@ pub struct WebpProcessor { stream_map: HashMap, width: Option, height: Option, - blur_hash: Option, + image: Option>, } unsafe impl Sync for WebpProcessor {} @@ -36,7 +35,7 @@ impl WebpProcessor { stream_map: HashMap::new(), width: None, height: None, - blur_hash: None, + image: None, } } @@ -80,9 +79,15 @@ impl WebpProcessor { None => frame }; - // take blur_hash from first video frame - if (*(*out_stream).codecpar).codec_type == AVMEDIA_TYPE_VIDEO && self.blur_hash.is_none() { - self.blur_hash = Some(make_blur_hash(frame_out, 9)?); + // take the first frame as "image" + if (*(*out_stream).codecpar).codec_type == AVMEDIA_TYPE_VIDEO && self.image.is_none() { + let mut dst_frame = resize_image(frame_out, + (*frame_out).width as usize, + (*frame_out).height as usize, + AV_PIX_FMT_RGBA)?; + let pic_slice = slice::from_raw_parts_mut((*dst_frame).data[0], ((*dst_frame).width * (*dst_frame).height * 4) as usize); + self.image = Some(pic_slice.to_vec()); + av_frame_free(&mut dst_frame); } let ret = avcodec_send_frame(*enc_ctx, frame_out); @@ -256,16 +261,8 @@ impl WebpProcessor { Ok(()) } -} -impl Drop for WebpProcessor { - fn drop(&mut self) { - unsafe { self.free().unwrap(); } - } -} - -impl FileProcessor for WebpProcessor { - fn process_file(&mut self, in_file: PathBuf, mime_type: &str) -> Result { + pub fn process_file(mut self, in_file: PathBuf, mime_type: &str) -> Result { unsafe { let mut out_path = in_file.clone(); out_path.set_extension("_compressed"); @@ -363,10 +360,7 @@ impl FileProcessor for WebpProcessor { mime_type: "image/webp".to_string(), width: self.width.unwrap_or(0), height: self.height.unwrap_or(0), - blur_hash: match &self.blur_hash { - Some(s) => s.clone(), - None => "".to_string() - }, + image: self.image.unwrap_or_default(), })) } } diff --git a/src/routes/blossom.rs b/src/routes/blossom.rs index 13d8b79..d8a6373 100644 --- a/src/routes/blossom.rs +++ b/src/routes/blossom.rs @@ -130,13 +130,12 @@ async fn upload( }; let f = FileUpload { id: blob.sha256, - user_id, name: name.unwrap_or("".to_string()), size: blob.size, mime_type: blob.mime_type, created: Utc::now(), }; - if let Err(e) = db.add_file(&f).await { + if let Err(e) = db.add_file(&f, user_id).await { error!("{}", e.to_string()); let _ = fs::remove_file(blob.path); if let Some(dbe) = e.as_database_error() { diff --git a/src/routes/mod.rs b/src/routes/mod.rs index 25f743f..26d7114 100644 --- a/src/routes/mod.rs +++ b/src/routes/mod.rs @@ -4,10 +4,10 @@ use std::str::FromStr; use anyhow::Error; use nostr::Event; +use rocket::{Request, State}; use rocket::fs::NamedFile; use rocket::http::{ContentType, Header, Status}; use rocket::response::Responder; -use rocket::{Request, State}; use crate::db::{Database, FileUpload}; use crate::filesystem::FileStore; @@ -56,20 +56,25 @@ async fn delete_file( if id.len() != 32 { return Err(Error::msg("Invalid file id")); } - if let Ok(Some(info)) = db.get_file(&id).await { + if let Ok(Some(_info)) = db.get_file(&id).await { let pubkey_vec = auth.pubkey.to_bytes().to_vec(); - let user = match db.get_user_id(&pubkey_vec).await { - Ok(u) => u, - Err(_e) => return Err(Error::msg("User not found")), + let owners = db.get_file_owners(&id).await?; + + let this_owner = match owners.iter().find(|o| o.pubkey.eq(&pubkey_vec)) { + Some(o) => o, + None => return Err(Error::msg("You dont own this file, you cannot delete it")) }; - if user != info.user_id { - return Err(Error::msg("You dont own this file, you cannot delete it")); - } - if let Err(e) = db.delete_file(&id).await { + if let Err(e) = db.delete_file_owner(&id, this_owner.id).await { return Err(Error::msg(format!("Failed to delete (db): {}", e))); } - if let Err(e) = fs::remove_file(fs.get(&id)) { - return Err(Error::msg(format!("Failed to delete (fs): {}", e))); + // only 1 owner was left, delete file completely + if owners.len() == 1 { + if let Err(e) = db.delete_file(&id).await { + return Err(Error::msg(format!("Failed to delete (fs): {}", e))); + } + if let Err(e) = fs::remove_file(fs.get(&id)) { + return Err(Error::msg(format!("Failed to delete (fs): {}", e))); + } } Ok(()) } else { @@ -111,7 +116,7 @@ pub async fn get_blob( return Ok(FilePayload { file: f, info }); } } - return Err(Status::NotFound); + Err(Status::NotFound) } #[rocket::head("/")] diff --git a/src/routes/nip96.rs b/src/routes/nip96.rs index 81a1078..fa20561 100644 --- a/src/routes/nip96.rs +++ b/src/routes/nip96.rs @@ -184,7 +184,6 @@ async fn upload( }; let file_upload = FileUpload { id: blob.sha256, - user_id, name: match &form.caption { Some(c) => c.to_string(), None => "".to_string(), @@ -193,7 +192,7 @@ async fn upload( mime_type: blob.mime_type, created: Utc::now(), }; - if let Err(e) = db.add_file(&file_upload).await { + if let Err(e) = db.add_file(&file_upload, user_id).await { error!("{}", e.to_string()); let _ = fs::remove_file(blob.path); if let Some(dbe) = e.as_database_error() { @@ -221,7 +220,17 @@ async fn upload( if let (Some(w), Some(h)) = (blob.width, blob.height) { tags.push(vec!["dim".to_string(), format!("{}x{}", w, h)]) } - + if let Some(lbls) = blob.labels { + for l in lbls { + let val = if l.contains(',') { + let split_val: Vec<&str> = l.split(',').collect(); + split_val[0].to_string() + } else { + l + }; + tags.push(vec!["t".to_string(), val]) + } + } Nip96Response::UploadResult(Json(Nip96UploadResult { status: "success".to_string(), nip94_event: Some(Nip94Event { diff --git a/src/settings.rs b/src/settings.rs index e91240f..21af902 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -1,3 +1,4 @@ +use std::path::PathBuf; use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Serialize, Deserialize)] @@ -18,5 +19,8 @@ pub struct Settings { pub public_url: String, /// Whitelisted pubkeys - pub whitelist: Option> + pub whitelist: Option>, + + /// Path for ViT image model + pub vit_model_path: Option }