Kieran/void-cat-rs
2
0
Fork 0
Code Issues 5 Pull Requests Packages Projects Releases Wiki Activity

feat: feature labels
All checks were successful
continuous-integration/drone/push Build is passing
Details

Browse Source
This commit is contained in:
kieran 2024-05-29 15:23:14 +01:00
parent 243b334948
commit 19e995ac16
Signed by: Kieran
GPG Key ID: DE71CEB3925BE941
8 changed files with 100 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

93
Cargo.lock generated
View File

@ -347,9 +347,9 @@ dependencies = [
[[package]]
name = "bytemuck_derive"
version = "1.6.0"
version = "1.7.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4da9a32f3fed317401fa3c862968128267c3106685286e15d5aaa3d7389c2f60"
checksum = "1ee891b04274a59bd38b412188e24b849617b2e45a0fd8d057deb63e7403761b"
dependencies = [
"proc-macro2",
"quote",
@ -371,7 +371,7 @@ checksum = "514de17de45fdb8dc022b1a7975556c53c86f9f0aa5f534b98977b171857c2c9"
[[package]]
name = "candle-core"
version = "0.5.1"
source = "git+https://github.com/huggingface/candle.git#13c64f6828360a9cb9b58b4f817e4f3b8316388c"
source = "git+https://github.com/huggingface/candle.git#ea260aeffdf97d0231cdb5cf6d31d33adfc4d4b0"
dependencies = [
"byteorder",
"gemm",
@ -391,7 +391,7 @@ dependencies = [
[[package]]
name = "candle-nn"
version = "0.5.1"
source = "git+https://github.com/huggingface/candle.git#13c64f6828360a9cb9b58b4f817e4f3b8316388c"
source = "git+https://github.com/huggingface/candle.git#ea260aeffdf97d0231cdb5cf6d31d33adfc4d4b0"
dependencies = [
"candle-core",
"half",
@ -405,7 +405,7 @@ dependencies = [
[[package]]
name = "candle-transformers"
version = "0.5.1"
source = "git+https://github.com/huggingface/candle.git#13c64f6828360a9cb9b58b4f817e4f3b8316388c"
source = "git+https://github.com/huggingface/candle.git#ea260aeffdf97d0231cdb5cf6d31d33adfc4d4b0"
dependencies = [
"byteorder",
"candle-core",
@ -609,9 +609,9 @@ checksum = "19d374276b40fb8bbdee95aef7c7fa6b5316ec764510eb64b8dd0e2ed0d7e7f5"
[[package]]
name = "crc32fast"
version = "1.4.0"
version = "1.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b3855a8a784b474f333699ef2bbca9db2c4a1f6d9088a90a2d25b1eb53111eaa"
checksum = "a97769d94ddab943e4510d138150169a2758b5ef3eb191a9ee688de3e23ef7b3"
dependencies = [
"cfg-if",
]
@ -1957,6 +1957,27 @@ dependencies = [
"libc",
]
[[package]]
name = "num_enum"
version = "0.7.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "02339744ee7253741199f897151b38e72257d13802d4ee837285cc2990a90845"
dependencies = [
"num_enum_derive",
]
[[package]]
name = "num_enum_derive"
version = "0.7.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "681030a937600a36906c185595136d26abfebb4aa9c65701cefcaf8578bb982b"
dependencies = [
"proc-macro-crate",
"proc-macro2",
"quote",
"syn 2.0.60",
]
[[package]]
name = "object"
version = "0.32.2"
@ -2231,6 +2252,15 @@ dependencies = [
"log",
]
[[package]]
name = "proc-macro-crate"
version = "3.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6d37c51ca738a55da99dc0c4a34860fd675453b8b36209178c2249bb13651284"
dependencies = [
"toml_edit 0.21.1",
]
[[package]]
name = "proc-macro2"
version = "1.0.81"
@ -2255,9 +2285,9 @@ dependencies = [
[[package]]
name = "pulp"
version = "0.18.10"
version = "0.18.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e14989307e408d9f4245d4fda09a7b144a08114ba124e26cab60ab83dc98db10"
checksum = "0ec8d02258294f59e4e223b41ad7e81c874aa6b15bc4ced9ba3965826da0eed5"
dependencies = [
"bytemuck",
"libm",
@ -3468,7 +3498,7 @@ dependencies = [
"serde",
"serde_spanned",
"toml_datetime",
"toml_edit",
"toml_edit 0.22.12",
]
[[package]]
@ -3480,6 +3510,17 @@ dependencies = [
"serde",
]
[[package]]
name = "toml_edit"
version = "0.21.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6a8534fd7f78b5405e860340ad6575217ce99f38d4d5c8f2442cb5ecb50090e1"
dependencies = [
"indexmap 2.2.6",
"toml_datetime",
"winnow 0.5.40",
]
[[package]]
name = "toml_edit"
version = "0.22.12"
@ -3490,7 +3531,7 @@ dependencies = [
"serde",
"serde_spanned",
"toml_datetime",
"winnow",
"winnow 0.6.7",
]
[[package]]
@ -4082,6 +4123,15 @@ version = "0.52.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bec47e5bfd1bff0eeaf6d8b485cc1074891a197ab4225d504cb7a1ab88b02bf0"
[[package]]
name = "winnow"
version = "0.5.40"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f593a95398737aeed53e489c785df13f3618e41dbcd6718c6addbf1395aa6876"
dependencies = [
"memchr",
]
[[package]]
name = "winnow"
version = "0.6.7"
@ -4121,9 +4171,9 @@ dependencies = [
[[package]]
name = "yoke"
version = "0.7.3"
version = "0.7.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "65e71b2e4f287f467794c671e2b8f8a5f3716b3c829079a1c44740148eff07e4"
checksum = "6c5b1314b079b0930c31e3af543d8ee1757b1951ae1e1565ec704403a7240ca5"
dependencies = [
"serde",
"stable_deref_trait",
@ -4133,9 +4183,9 @@ dependencies = [
[[package]]
name = "yoke-derive"
version = "0.7.3"
version = "0.7.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9e6936f0cce458098a201c245a11bef556c6a0181129c7034d10d76d1ec3a2b8"
checksum = "28cc31741b18cb6f1d5ff12f5b7523e3d6eb0852bbbad19d73905511d9849b95"
dependencies = [
"proc-macro2",
"quote",
@ -4165,18 +4215,18 @@ dependencies = [
[[package]]
name = "zerofrom"
version = "0.1.3"
version = "0.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "655b0814c5c0b19ade497851070c640773304939a6c0fd5f5fb43da0696d05b7"
checksum = "91ec111ce797d0e0784a1116d0ddcdbea84322cd79e5d5ad173daeba4f93ab55"
dependencies = [
"zerofrom-derive",
]
[[package]]
name = "zerofrom-derive"
version = "0.1.3"
version = "0.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e6a647510471d372f2e6c2e6b7219e44d8c574d24fdc11c610a61455782f18c3"
checksum = "0ea7b4a3637ea8669cedf0f1fd5c286a17f3de97b8dd5a70a6c167a1730e63a5"
dependencies = [
"proc-macro2",
"quote",
@ -4192,14 +4242,15 @@ checksum = "525b4ec142c6b68a2d10f01f7bbf6755599ca3f81ea53b8431b7dd348f5fdb2d"
[[package]]
name = "zip"
version = "1.2.3"
version = "1.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c700ea425e148de30c29c580c1f9508b93ca57ad31c9f4e96b83c194c37a7a8f"
checksum = "9cc23c04387f4da0374be4533ad1208cbb091d5c11d070dfef13676ad6497164"
dependencies = [
"arbitrary",
"crc32fast",
"crossbeam-utils",
"displaydoc",
"indexmap 2.2.6",
"num_enum",
"thiserror",
]

10
Cargo.toml
View File

@ -5,6 +5,9 @@ edition = "2021"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[features]
labels = ["dep:candle-core", "dep:candle-nn", "dep:candle-transformers"]
[dependencies]
log = "0.4.21"
nostr = "0.30.0"
@ -23,9 +26,10 @@ 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" }
ureq = { version = "2.9.7", features = ["json"] }
url = "2.5.0"
serde_with = { version = "3.8.1", features = ["hex"] }
candle-core = { git = "https://github.com/huggingface/candle.git", version = "0.5.1", optional = true }
candle-nn = { git = "https://github.com/huggingface/candle.git", version = "0.5.1", optional = true }
candle-transformers = { git = "https://github.com/huggingface/candle.git", version = "0.5.1", optional = true }

5
src/db.rs
View File

@ -15,6 +15,7 @@ pub struct FileUpload {
pub blur_hash: Option<String>,
#[sqlx(skip)]
#[cfg(feature = "labels")]
pub labels: Vec<FileLabel>,
}
@ -25,6 +26,7 @@ pub struct User {
pub created: DateTime<Utc>,
}
#[cfg(feature = "labels")]
#[derive(Clone, FromRow, Serialize)]
pub struct FileLabel {
pub file: Vec<u8>,
@ -33,6 +35,7 @@ pub struct FileLabel {
pub model: String,
}
#[cfg(feature = "labels")]
impl FileLabel {
pub fn new(label: String, model: String) -> Self {
Self {
@ -99,6 +102,7 @@ impl Database {
.bind(user_id);
tx.execute(q2).await?;
#[cfg(feature = "labels")]
for lbl in &file.labels {
let q3 = sqlx::query("insert ignore into upload_labels(file,label,model) values(?,?,?)")
.bind(&file.id)
@ -124,6 +128,7 @@ impl Database {
.await
}
#[cfg(feature = "labels")]
pub async fn get_file_labels(&self, file: &Vec<u8>) -> Result<Vec<FileLabel>, Error> {
sqlx::query_as("select upload_labels.* from uploads, upload_labels where uploads.id = ? and uploads.id = upload_labels.file")
.bind(file)

13
src/filesystem.rs
View File

@ -12,8 +12,11 @@ use sha2::{Digest, Sha256};
use tokio::fs::File;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeekExt};
use crate::db::{FileLabel, FileUpload};
use crate::processing::{compress_file, FileProcessorResult, probe_file, ProbeResult, ProbeStream};
#[cfg(feature = "labels")]
use crate::db::FileLabel;
use crate::db::FileUpload;
use crate::processing::{compress_file, FileProcessorResult, probe_file, ProbeStream};
#[cfg(feature = "labels")]
use crate::processing::labeling::label_frame;
use crate::settings::Settings;
@ -86,7 +89,7 @@ impl FileStore {
if compress {
let start = SystemTime::now();
let proc_result = compress_file(tmp_path.clone(), mime_type)?;
if let FileProcessorResult::NewFile(mut new_temp) = proc_result {
if let FileProcessorResult::NewFile(new_temp) = proc_result {
let old_size = tmp_path.metadata()?.len();
let new_size = new_temp.result.metadata()?.len();
let time_compress = SystemTime::now().duration_since(start).unwrap();
@ -99,6 +102,8 @@ impl FileStore {
)?;
let time_blur_hash = SystemTime::now().duration_since(start).unwrap();
let start = SystemTime::now();
#[cfg(feature = "labels")]
let labels = if let Some(mp) = &self.settings.vit_model_path {
label_frame(
new_temp.image.as_mut_slice(),
@ -110,6 +115,7 @@ impl FileStore {
} else {
vec![]
};
let time_labels = SystemTime::now().duration_since(start).unwrap();
// delete old temp
@ -143,6 +149,7 @@ impl FileStore {
height: Some(new_temp.height as u32),
blur_hash: Some(blur_hash),
mime_type: new_temp.mime_type,
#[cfg(feature = "labels")]
labels,
created: Utc::now(),
},

5
src/processing/mod.rs
View File

@ -9,6 +9,7 @@ use crate::processing::probe::FFProbe;
use crate::processing::webp::WebpProcessor;
mod webp;
#[cfg(feature = "labels")]
pub mod labeling;
mod probe;
@ -50,7 +51,7 @@ pub fn compress_file(in_file: PathBuf, mime_type: &str) -> Result<FileProcessorR
} else {
None
};
if let Some(mut proc) = proc {
if let Some(proc) = proc {
proc.process_file(in_file, mime_type)
} else {
Ok(FileProcessorResult::Skip)
@ -59,7 +60,7 @@ pub fn compress_file(in_file: PathBuf, mime_type: &str) -> Result<FileProcessorR
pub fn probe_file(in_file: PathBuf) -> Result<FileProcessorResult, Error> {
let proc = FFProbe::new();
proc.process_file(in_file, "")
proc.process_file(in_file)
}
unsafe fn resize_image(frame: *const AVFrame, width: usize, height: usize, pix_fmt: AVPixelFormat) -> Result<*mut AVFrame, Error> {

2
src/processing/probe.rs
View File

@ -16,7 +16,7 @@ impl FFProbe {
Self {}
}
pub fn process_file(mut self, in_file: PathBuf, mime_type: &str) -> Result<FileProcessorResult, Error> {
pub fn process_file(self, in_file: PathBuf) -> Result<FileProcessorResult, Error> {
unsafe {
let mut dec_fmt: *mut AVFormatContext = ptr::null_mut();
let ret = avformat_open_input(&mut dec_fmt,

3
src/routes/blossom.rs
View File

@ -1,6 +1,5 @@
use std::fs;
use chrono::Utc;
use log::error;
use nostr::prelude::hex;
use nostr::Tag;
@ -13,7 +12,7 @@ use serde::{Deserialize, Serialize};
use crate::auth::blossom::BlossomAuth;
use crate::blob::BlobDescriptor;
use crate::db::{Database, FileUpload};
use crate::db::{Database};
use crate::filesystem::FileStore;
use crate::routes::delete_file;
use crate::settings::Settings;

1
src/routes/nip96.rs
View File

@ -130,6 +130,7 @@ impl Nip96UploadResult {
if let (Some(w), Some(h)) = (upload.width, upload.height) {
tags.push(vec!["dim".to_string(), format!("{}x{}", w, h)])
}
#[cfg(feature = "labels")]
for l in &upload.labels {
let val = if l.label.contains(',') {
let split_val: Vec<&str> = l.label.split(',').collect();