Automatic image labeling
Some checks reported errors
continuous-integration/drone/push Build was killed
continuous-integration/drone Build was killed

This commit is contained in:
kieran 2024-05-14 17:06:35 +01:00
parent 30e4e8ed1c
commit cf2950bf99
Signed by: Kieran
GPG Key ID: DE71CEB3925BE941
15 changed files with 1791 additions and 140 deletions

521
Cargo.lock generated
View File

@ -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",
]

View File

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

5
build.rs Normal file
View File

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

View File

@ -14,4 +14,7 @@ max_upload_bytes = 104857600
public_url = "http://localhost:8000"
# Whitelisted pubkeys, leave out to disable
# whitelist = ["63fe6318dc58583cfe16810f86dd09e18bfd76aabc24a0081ce2856f330504ed"]
# whitelist = ["63fe6318dc58583cfe16810f86dd09e18bfd76aabc24a0081ce2856f330504ed"]
# Path for ViT image model (https://huggingface.co/google/vit-base-patch16-224)
# vit_model_path = "model.safetennsors"

View File

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

View File

@ -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<u8>,
pub user_id: u64,
pub name: String,
pub size: u64,
pub mime_type: String,
pub created: DateTime<Utc>,
}
#[derive(Clone, FromRow)]
pub struct User {
pub id: u64,
pub pubkey: Vec<u8>,
pub created: DateTime<Utc>,
}
#[derive(Clone)]
pub struct Database {
pool: sqlx::pool::Pool<sqlx::mysql::MySql>,
@ -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<u8>) -> Result<Option<FileUpload>, 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<u8>) -> Result<Vec<User>, 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<u8>, 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<u8>) -> 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<FileUpload> = 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)
}
}

View File

@ -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<u8>,
@ -23,18 +24,17 @@ pub struct FileSystemResult {
pub width: Option<usize>,
pub height: Option<usize>,
pub blur_hash: Option<String>,
pub labels: Option<Vec<String>>,
}
pub struct FileStore {
settings: Settings,
processor: Arc<Mutex<MediaProcessor>>,
}
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()
})
}

View File

@ -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<String, Error> {
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)
}

1067
src/processing/labeling.rs Normal file

File diff suppressed because it is too large Load Diff

View File

@ -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<u8>,
}
pub(crate) trait FileProcessor {
fn process_file(&mut self, in_file: PathBuf, mime_type: &str) -> Result<FileProcessorResult, Error>;
}
pub(crate) struct MediaProcessor {}
impl MediaProcessor {
pub fn new() -> Self {
Self {}
pub fn compress_file(in_file: PathBuf, mime_type: &str) -> Result<FileProcessorResult, Error> {
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<FileProcessorResult, Error> {
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)
}

View File

@ -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<usize, usize>,
width: Option<usize>,
height: Option<usize>,
blur_hash: Option<String>,
image: Option<Vec<u8>>,
}
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<FileProcessorResult, Error> {
pub fn process_file(mut self, in_file: PathBuf, mime_type: &str) -> Result<FileProcessorResult, Error> {
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(),
}))
}
}

View File

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

View File

@ -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("/<sha256>")]

View File

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

View File

@ -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<Vec<String>>
pub whitelist: Option<Vec<String>>,
/// Path for ViT image model
pub vit_model_path: Option<PathBuf>
}