From 7767d21ccfa54b1f3c131a885d5489e13c2285a8 Mon Sep 17 00:00:00 2001 From: kieran Date: Wed, 29 May 2024 14:53:11 +0100 Subject: [PATCH] fix: probe files for no-compress --- migrations/20240529112336_upload_labels.sql | 19 ++++++ src/filesystem.rs | 30 +++++++++- src/processing/mod.rs | 8 ++- src/processing/probe.rs | 64 +++++++++++++++++++++ 4 files changed, 119 insertions(+), 2 deletions(-) create mode 100644 migrations/20240529112336_upload_labels.sql create mode 100644 src/processing/probe.rs diff --git a/migrations/20240529112336_upload_labels.sql b/migrations/20240529112336_upload_labels.sql new file mode 100644 index 0000000..8a3c2be --- /dev/null +++ b/migrations/20240529112336_upload_labels.sql @@ -0,0 +1,19 @@ +-- Add migration script here +alter table uploads + add column blur_hash varchar(512), + add column width integer unsigned, + add column height integer unsigned; + +create table upload_labels +( + file binary(32) not null, + label varchar(255) not null, + created timestamp default current_timestamp, + model varchar(255) not null, + + constraint fk_upload_labels_file_id + foreign key (file) references uploads (id) + on delete cascade + on update restrict +); +create unique index ix_upload_labels_file_label_model on upload_labels (file, label, model); \ No newline at end of file diff --git a/src/filesystem.rs b/src/filesystem.rs index 9380092..6bbc159 100644 --- a/src/filesystem.rs +++ b/src/filesystem.rs @@ -13,7 +13,7 @@ use tokio::fs::File; use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeekExt}; use crate::db::{FileLabel, FileUpload}; -use crate::processing::{compress_file, FileProcessorResult}; +use crate::processing::{compress_file, FileProcessorResult, probe_file, ProbeResult, ProbeStream}; use crate::processing::labeling::label_frame; use crate::settings::Settings; @@ -148,7 +148,34 @@ impl FileStore { }, }); } + } else if let FileProcessorResult::Probe(p) = probe_file(tmp_path.clone())? { + let video_stream_size = p.streams.iter().find_map(|s| match s { + ProbeStream::Video { width, height, .. } => Some((width, height)), + _ => None + }); + let n = file.metadata().await?.len(); + let hash = FileStore::hash_file(&mut file).await?; + return Ok(FileSystemResult { + path: tmp_path, + upload: FileUpload { + id: hash, + name: "".to_string(), + size: n, + created: Utc::now(), + mime_type: mime_type.to_string(), + width: match video_stream_size { + Some((w, _h)) => Some(*w), + _ => None + }, + height: match video_stream_size { + Some((_w, h)) => Some(*h), + _ => None + }, + ..Default::default() + }, + }); } + let n = file.metadata().await?.len(); let hash = FileStore::hash_file(&mut file).await?; Ok(FileSystemResult { @@ -158,6 +185,7 @@ impl FileStore { name: "".to_string(), size: n, created: Utc::now(), + mime_type: mime_type.to_string(), ..Default::default() }, }) diff --git a/src/processing/mod.rs b/src/processing/mod.rs index b462192..34ad94f 100644 --- a/src/processing/mod.rs +++ b/src/processing/mod.rs @@ -4,6 +4,7 @@ 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::probe::FFProbe; use crate::processing::webp::WebpProcessor; @@ -56,6 +57,11 @@ pub fn compress_file(in_file: PathBuf, mime_type: &str) -> Result Result { + let proc = FFProbe::new(); + proc.process_file(in_file, "") +} + 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, @@ -76,4 +82,4 @@ unsafe fn resize_image(frame: *const AVFrame, width: usize, height: usize, pix_f sws_freeContext(sws_ctx); Ok(dst_frame) -} \ No newline at end of file +} diff --git a/src/processing/probe.rs b/src/processing/probe.rs new file mode 100644 index 0000000..0f92bbd --- /dev/null +++ b/src/processing/probe.rs @@ -0,0 +1,64 @@ +use std::ffi::CStr; +use std::path::PathBuf; +use std::ptr; + +use anyhow::Error; +use ffmpeg_sys_the_third::{avcodec_get_name, avformat_close_input, avformat_find_stream_info, avformat_free_context, avformat_open_input, AVFormatContext}; +use ffmpeg_sys_the_third::AVMediaType::{AVMEDIA_TYPE_AUDIO, AVMEDIA_TYPE_VIDEO}; + +use crate::processing::{FileProcessorResult, ProbeResult, ProbeStream}; + +/// Image converter to WEBP +pub struct FFProbe {} + +impl FFProbe { + pub fn new() -> Self { + Self {} + } + + pub fn process_file(mut self, in_file: PathBuf, mime_type: &str) -> Result { + unsafe { + let mut dec_fmt: *mut AVFormatContext = ptr::null_mut(); + let ret = avformat_open_input(&mut dec_fmt, + format!("{}\0", in_file.into_os_string().into_string().unwrap()).as_ptr() as *const libc::c_char, + ptr::null_mut(), + ptr::null_mut()); + if ret < 0 { + return Err(Error::msg("Failed to create input context")); + } + + let ret = avformat_find_stream_info(dec_fmt, ptr::null_mut()); + if ret < 0 { + return Err(Error::msg("Failed to probe input")); + } + + let mut stream_info = vec![]; + let mut ptr_x = 0; + while ptr_x < (*dec_fmt).nb_streams { + let ptr = *(*dec_fmt).streams.add(ptr_x as usize); + let codec_par = (*ptr).codecpar; + let codec = CStr::from_ptr(avcodec_get_name((*codec_par).codec_id)).to_str()?.to_string(); + if (*codec_par).codec_type == AVMEDIA_TYPE_VIDEO { + stream_info.push(ProbeStream::Video { + width: (*codec_par).width as u32, + height: (*codec_par).height as u32, + codec, + }); + } else if (*codec_par).codec_type == AVMEDIA_TYPE_AUDIO { + stream_info.push(ProbeStream::Audio { + sample_rate: (*codec_par).sample_rate as u32, + codec, + }); + } + ptr_x += 1; + } + + avformat_close_input(&mut dec_fmt); + avformat_free_context(dec_fmt); + + Ok(FileProcessorResult::Probe(ProbeResult { + streams: stream_info + })) + } + } +} \ No newline at end of file