diff --git a/Cargo.lock b/Cargo.lock index 7134498..d33c3f3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -131,9 +131,9 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.86" +version = "1.0.93" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3d1d046238990b9cf5bcde22a3fb3584ee5cf65fb2765f454ed428c7a0063da" +checksum = "4c95c10ba0b00a02636238b814946408b1322d5ac4760326e6fb8ec956d85775" [[package]] name = "arbitrary" @@ -413,12 +413,6 @@ dependencies = [ "generic-array", ] -[[package]] -name = "blurhash" -version = "0.2.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e79769241dcd44edf79a732545e8b5cec84c247ac060f5252cd51885d093a8fc" - [[package]] name = "bumpalo" version = "3.16.0" @@ -592,6 +586,16 @@ dependencies = [ "zeroize", ] +[[package]] +name = "clang" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "84c044c781163c001b913cd018fc95a628c50d0d2dfea8bca77dad71edb16e37" +dependencies = [ + "clang-sys", + "libc", +] + [[package]] name = "clang-sys" version = "1.8.1" @@ -1070,14 +1074,26 @@ version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e8c02a5121d4ea3eb16a80748c74f5549a5665e4c21333c6098f283870fbdea6" +[[package]] +name = "ffmpeg-rs-raw" +version = "0.1.0" +source = "git+https://git.v0l.io/Kieran/ffmpeg-rs-raw.git?rev=363ad4f55ce3b7047d13197bdff6ed9bd059c099#363ad4f55ce3b7047d13197bdff6ed9bd059c099" +dependencies = [ + "anyhow", + "ffmpeg-sys-the-third", + "libc", + "log", + "slimbox", +] + [[package]] name = "ffmpeg-sys-the-third" -version = "2.0.0+ffmpeg-7.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a82bfdb0a7925996707f0a7dc37b2f3251ff5a15d26e78c586adb60c240dedc5" +version = "2.1.0+ffmpeg-7.1" +source = "git+https://github.com/shssoichiro/ffmpeg-the-third.git?branch=master#1b5d24091a4345dc65f0520974732cd58661c3e2" dependencies = [ "bindgen", "cc", + "clang", "libc", "pkg-config", "vcpkg", @@ -1885,9 +1901,9 @@ checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" [[package]] name = "libc" -version = "0.2.158" +version = "0.2.162" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d8adc4bb1803a324070e64a98ae98f38934d91957a99cfb3a43dcbc01bc56439" +checksum = "18d287de67fe55fd7e1581fe933d965a5a9477b38e949cfa9f8574ef01506398" [[package]] name = "libloading" @@ -2087,11 +2103,12 @@ dependencies = [ [[package]] name = "nostr" -version = "0.35.0" +version = "0.36.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "56db234b2e07901e372f34e9463f91590579cd8e6dbd34ed2ccc7e461e4ba639" +checksum = "14ad56c1d9a59f4edc46b17bc64a217b38b99baefddc0080f85ad98a0855336d" dependencies = [ "aes", + "async-trait", "base64 0.22.1", "bech32", "bip39", @@ -2231,9 +2248,9 @@ dependencies = [ [[package]] name = "once_cell" -version = "1.19.0" +version = "1.20.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" +checksum = "1261fe7e33c73b354eab43b1273a57c8f967d0391e80353e51f764ac02cf6775" [[package]] name = "opaque-debug" @@ -2962,7 +2979,6 @@ version = "0.2.0" dependencies = [ "anyhow", "base64 0.22.1", - "blurhash", "candle-core", "candle-nn", "candle-transformers", @@ -2970,7 +2986,7 @@ dependencies = [ "clap", "clap_derive", "config", - "ffmpeg-sys-the-third", + "ffmpeg-rs-raw", "hex", "libc", "log", @@ -3371,6 +3387,12 @@ dependencies = [ "autocfg", ] +[[package]] +name = "slimbox" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26dfcf7e4fe830e4b9245b9e0def30d3df9ea194aca707e9a78b079d2b646b1a" + [[package]] name = "smallvec" version = "1.13.2" diff --git a/Cargo.toml b/Cargo.toml index ce32263..8dc32e6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,7 +16,7 @@ name = "route96" [features] default = ["nip96", "blossom", "analytics"] -media-compression = ["dep:ffmpeg-sys-the-third", "dep:blurhash", "dep:libc"] +media-compression = ["dep:ffmpeg-rs-raw", "dep:libc"] labels = ["nip96", "dep:candle-core", "dep:candle-nn", "dep:candle-transformers"] nip96 = ["media-compression"] blossom = [] @@ -27,7 +27,7 @@ void-cat-redirects = ["dep:sqlx-postgres"] [dependencies] log = "0.4.21" -nostr = "0.35.0" +nostr = "0.36.0" pretty_env_logger = "0.5.0" rocket = { version = "0.5.0", features = ["json"] } tokio = { version = "1.37.0", features = ["rt", "rt-multi-thread", "macros"] } @@ -35,7 +35,7 @@ base64 = "0.22.1" hex = { version = "0.4.3", features = ["serde"] } serde = { version = "1.0.198", features = ["derive"] } uuid = { version = "1.8.0", features = ["v4"] } -anyhow = "1.0.82" +anyhow = "^1.0.82" sha2 = "0.10.8" sqlx = { version = "0.8.1", features = ["mysql", "runtime-tokio", "chrono", "uuid"] } config = { version = "0.14.0", features = ["toml"] } @@ -45,8 +45,7 @@ serde_with = { version = "3.8.1", features = ["hex"] } reqwest = "0.12.8" libc = { version = "0.2.153", optional = true } -blurhash = { version = "0.2.1", optional = true } -ffmpeg-sys-the-third = { version = "2.0.0", features = ["default"], optional = true } +ffmpeg-rs-raw = { git = "https://git.v0l.io/Kieran/ffmpeg-rs-raw.git", rev = "363ad4f55ce3b7047d13197bdff6ed9bd059c099", optional = true} candle-core = { git = "https://git.v0l.io/Kieran/candle.git", version = "^0.7.2", optional = true } candle-nn = { git = "https://git.v0l.io/Kieran/candle.git", version = "^0.7.2", optional = true } candle-transformers = { git = "https://git.v0l.io/Kieran/candle.git", version = "^0.7.2", optional = true } diff --git a/src/filesystem.rs b/src/filesystem.rs index 16c4213..99401d7 100644 --- a/src/filesystem.rs +++ b/src/filesystem.rs @@ -18,7 +18,7 @@ use crate::db::FileUpload; #[cfg(feature = "labels")] use crate::processing::labeling::label_frame; #[cfg(feature = "media-compression")] -use crate::processing::{compress_file, probe_file, FileProcessorResult, ProbeStream}; +use crate::processing::{compress_file, probe_file, FileProcessorResult}; use crate::settings::Settings; #[derive(Clone, Default, Serialize)] @@ -101,32 +101,18 @@ 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)?; 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_blur_hash = SystemTime::now().duration_since(start)?; - 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(), - new_temp.width, - new_temp.height, - mp.clone(), - )? - .iter() - .map(|l| FileLabel::new(l.clone(), "vit224".to_string())) - .collect() + label_frame(&new_temp.result, mp.clone())? + .iter() + .map(|l| FileLabel::new(l.clone(), "vit224".to_string())) + .collect() } else { vec![] }; @@ -145,12 +131,11 @@ impl FileStore { 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_blur_hash={:.2}ms, duration_labels={:.2}ms", + info!("Processed media: ratio={:.2}x, old_size={:.3}kb, new_size={:.3}kb, duration_compress={:.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_blur_hash.as_micros() as f64 / 1000.0, time_labels.as_micros() as f64 / 1000.0 ); @@ -162,7 +147,7 @@ impl FileStore { size: n, width: Some(new_temp.width as u32), height: Some(new_temp.height as u32), - blur_hash: Some(blur_hash), + blur_hash: None, mime_type: new_temp.mime_type, #[cfg(feature = "labels")] labels, @@ -171,11 +156,7 @@ 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, - }); + } else if let Ok(p) = probe_file(tmp_path.clone()) { let n = file.metadata().await?.len(); let hash = FileStore::hash_file(&mut file).await?; return Ok(FileSystemResult { @@ -186,14 +167,8 @@ impl FileStore { 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, - }, + width: p.map(|v| v.0 as u32), + height: p.map(|v| v.1 as u32), ..Default::default() }, }); diff --git a/src/processing/labeling.rs b/src/processing/labeling.rs index 14f589e..b3a7300 100644 --- a/src/processing/labeling.rs +++ b/src/processing/labeling.rs @@ -1,27 +1,18 @@ -use std::mem::transmute; -use std::path::PathBuf; -use std::{fs, ptr, slice}; +use std::path::{Path, PathBuf}; +use std::slice; -use anyhow::Error; +use anyhow::{Error, Result}; use candle_core::{DType, Device, IndexOp, Tensor, D}; use candle_nn::VarBuilder; use candle_transformers::models::vit; -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 ffmpeg_sys_the_third::{av_frame_alloc, av_frame_free}; +use ffmpeg_rs_raw::ffmpeg_sys_the_third::AVPixelFormat::AV_PIX_FMT_RGB24; +use ffmpeg_rs_raw::ffmpeg_sys_the_third::{av_frame_free, av_packet_free}; +use ffmpeg_rs_raw::{Decoder, Demuxer, Scaler}; -use crate::processing::resize_image; - -pub fn label_frame( - frame: &mut [u8], - width: usize, - height: usize, - model: PathBuf, -) -> Result, Error> { +pub fn label_frame(frame: &Path, model: PathBuf) -> Result> { unsafe { let device = Device::Cpu; - let image = load_frame_224(frame, width, height)?.to_device(&device)?; + let image = load_frame_224(frame)?.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)?; @@ -41,41 +32,50 @@ pub fn label_frame( } } -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; +/// Load an image from disk into RGB pixel buffer +unsafe fn load_image(path_buf: &Path, width: usize, height: usize) -> Result> { + let mut demux = Demuxer::new(path_buf.to_str().unwrap())?; + let info = demux.probe_input()?; + let image_stream = info + .best_video() + .ok_or(Error::msg("No image stream found"))?; - 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, - ); + let mut decoder = Decoder::new(); + decoder.setup_decoder(image_stream, None)?; - 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 mut scaler = Scaler::new(); + while let Ok((mut pkt, _)) = demux.get_packet() { + if let Some(mut frame) = decoder.decode_pkt(pkt)?.into_iter().next() { + let mut new_frame = + scaler.process_frame(frame, width as u16, height as u16, AV_PIX_FMT_RGB24)?; + let mut dst_vec = Vec::with_capacity(3 * width * height); + + for row in 0..height { + let line_size = (*new_frame).linesize[0] as usize; + let row_offset = line_size * row; + let row_slice = + slice::from_raw_parts((*new_frame).data[0].add(row_offset), line_size); + dst_vec.extend_from_slice(row_slice); + } + av_frame_free(&mut frame); + av_frame_free(&mut new_frame); + av_packet_free(&mut pkt); + return Ok(dst_vec); + } + } + Err(Error::msg("No image data found")) +} + +unsafe fn load_frame_224(path: &Path) -> Result { + let pic = load_image(path, 224, 224)?; + + //fs::write("frame_224.raw", &pic_slice)?; + let data = Tensor::from_vec(pic, (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) } diff --git a/src/processing/mod.rs b/src/processing/mod.rs index a390e7d..7cd5117 100644 --- a/src/processing/mod.rs +++ b/src/processing/mod.rs @@ -1,19 +1,68 @@ -use std::intrinsics::transmute; use std::path::PathBuf; -use std::ptr; - -use anyhow::Error; -use ffmpeg_sys_the_third::{ - av_frame_alloc, sws_freeContext, sws_getContext, sws_scale_frame, AVFrame, AVPixelFormat, -}; use crate::processing::probe::FFProbe; -use crate::processing::webp::WebpProcessor; +use anyhow::{bail, Error, Result}; +use ffmpeg_rs_raw::ffmpeg_sys_the_third::AVPixelFormat::AV_PIX_FMT_YUV420P; +use ffmpeg_rs_raw::{Encoder, StreamType, Transcoder}; #[cfg(feature = "labels")] pub mod labeling; mod probe; -mod webp; + +pub struct WebpProcessor; + +impl Default for WebpProcessor { + fn default() -> Self { + Self::new() + } +} + +impl WebpProcessor { + pub fn new() -> Self { + Self + } + + pub fn process_file(&mut self, input: PathBuf, mime_type: &str) -> Result { + use ffmpeg_rs_raw::ffmpeg_sys_the_third::AVCodecID::AV_CODEC_ID_WEBP; + + if !mime_type.starts_with("image/") { + bail!("MIME type not supported"); + } + + if mime_type == "image/webp" { + return Ok(FileProcessorResult::Skip); + } + + let mut out_path = input.clone(); + out_path.set_extension("compressed.webp"); + unsafe { + let mut trans = Transcoder::new(input.to_str().unwrap(), out_path.to_str().unwrap())?; + + let probe = trans.prepare()?; + let image_stream = probe + .streams + .iter() + .find(|c| c.stream_type == StreamType::Video) + .ok_or(Error::msg("No image found, cant compress"))?; + + let enc = Encoder::new(AV_CODEC_ID_WEBP)? + .with_height(image_stream.height as i32) + .with_width(image_stream.width as i32) + .with_pix_fmt(AV_PIX_FMT_YUV420P) + .open(None)?; + + trans.transcode_stream(image_stream, enc)?; + trans.run()?; + + Ok(FileProcessorResult::NewFile(NewFileProcessorResult { + result: out_path, + mime_type: "image/webp".to_string(), + width: image_stream.width, + height: image_stream.height, + })) + } + } +} pub struct ProbeResult { pub streams: Vec, @@ -33,7 +82,6 @@ pub enum ProbeStream { pub enum FileProcessorResult { NewFile(NewFileProcessorResult), - Probe(ProbeResult), Skip, } @@ -42,9 +90,6 @@ pub struct NewFileProcessorResult { pub mime_type: String, pub width: usize, pub height: usize, - - /// The image as RBGA - pub image: Vec, } pub fn compress_file(in_file: PathBuf, mime_type: &str) -> Result { @@ -53,46 +98,15 @@ pub fn compress_file(in_file: PathBuf, mime_type: &str) -> Result Result { +pub fn probe_file(in_file: PathBuf) -> 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, - 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) + let info = proc.process_file(in_file)?; + Ok(info.best_video().map(|v| (v.width, v.height))) } diff --git a/src/processing/probe.rs b/src/processing/probe.rs index a35bf43..86153e6 100644 --- a/src/processing/probe.rs +++ b/src/processing/probe.rs @@ -1,15 +1,6 @@ -use std::ffi::CStr; +use anyhow::Result; +use ffmpeg_rs_raw::{Demuxer, DemuxerInfo}; use std::path::PathBuf; -use std::ptr; - -use anyhow::Error; -use ffmpeg_sys_the_third::AVMediaType::{AVMEDIA_TYPE_AUDIO, AVMEDIA_TYPE_VIDEO}; -use ffmpeg_sys_the_third::{ - avcodec_get_name, avformat_close_input, avformat_find_stream_info, avformat_free_context, - avformat_open_input, AVFormatContext, -}; - -use crate::processing::{FileProcessorResult, ProbeResult, ProbeStream}; /// Image converter to WEBP pub struct FFProbe {} @@ -19,55 +10,10 @@ impl FFProbe { Self {} } - pub fn process_file(self, in_file: PathBuf) -> Result { + pub fn process_file(self, in_file: PathBuf) -> 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 { - // input might not be media - return Ok(FileProcessorResult::Skip); - } - - 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, - })) + let mut demuxer = Demuxer::new(in_file.to_str().unwrap())?; + demuxer.probe_input() } } } diff --git a/src/processing/webp.rs b/src/processing/webp.rs deleted file mode 100644 index 775e610..0000000 --- a/src/processing/webp.rs +++ /dev/null @@ -1,422 +0,0 @@ -use std::collections::HashMap; -use std::mem::transmute; -use std::path::PathBuf; -use std::{ptr, slice}; - -use anyhow::Error; -use ffmpeg_sys_the_third::AVMediaType::{AVMEDIA_TYPE_AUDIO, AVMEDIA_TYPE_VIDEO}; -use ffmpeg_sys_the_third::AVPixelFormat::{AV_PIX_FMT_RGBA, AV_PIX_FMT_YUV420P}; -use ffmpeg_sys_the_third::{ - 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_read_frame, av_write_trailer, 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, 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, avio_open, sws_freeContext, sws_getContext, - sws_scale_frame, AVCodec, AVCodecContext, AVCodecID, AVFormatContext, AVMediaType, AVPacket, - SwsContext, AVERROR, AVERROR_EOF, AVERROR_STREAM_NOT_FOUND, AVFMT_GLOBALHEADER, - AVIO_FLAG_WRITE, AV_CODEC_FLAG_GLOBAL_HEADER, AV_PROFILE_H264_HIGH, -}; -use libc::EAGAIN; - -use crate::processing::{resize_image, FileProcessorResult, NewFileProcessorResult}; - -/// Image converter to WEBP -pub struct WebpProcessor { - encoders: HashMap, - decoders: HashMap, - scalers: HashMap, - stream_map: HashMap, - width: Option, - height: Option, - image: Option>, -} - -unsafe impl Sync for WebpProcessor {} - -unsafe impl Send for WebpProcessor {} - -impl WebpProcessor { - pub fn new() -> Self { - Self { - encoders: HashMap::new(), - decoders: HashMap::new(), - scalers: HashMap::new(), - stream_map: HashMap::new(), - width: None, - height: None, - image: None, - } - } - - unsafe fn transcode_pkt( - &mut self, - pkt: *mut AVPacket, - in_fmt: *mut AVFormatContext, - out_fmt: *mut AVFormatContext, - ) -> Result<(), Error> { - let idx = (*pkt).stream_index as usize; - let out_idx = match self.stream_map.get(&idx) { - Some(i) => i, - None => return Ok(()), - }; - let in_stream = *(*in_fmt).streams.add(idx); - let out_stream = *(*out_fmt).streams.add(*out_idx); - av_packet_rescale_ts(pkt, (*in_stream).time_base, (*out_stream).time_base); - - let dec_ctx = self.decoders.get_mut(&idx).expect("Missing decoder config"); - let enc_ctx = self - .encoders - .get_mut(out_idx) - .expect("Missing encoder config"); - - let ret = avcodec_send_packet(*dec_ctx, pkt); - if ret < 0 { - return Err(Error::msg("Failed to decode packet")); - } - - let mut frame = av_frame_alloc(); - let mut frame_out = av_frame_alloc(); - loop { - let ret = avcodec_receive_frame(*dec_ctx, frame); - if ret == AVERROR_EOF || ret == AVERROR(EAGAIN) { - break; - } else if ret < 0 { - return Err(Error::msg("Frame read error")); - } - - let frame_out = match self.scalers.get_mut(out_idx) { - Some(sws) => { - av_frame_copy_props(frame_out, frame); - let ret = sws_scale_frame(*sws, frame_out, frame); - if ret < 0 { - return Err(Error::msg("Failed to scale frame")); - } - frame_out - } - None => frame, - }; - - // 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); - if ret < 0 { - return Err(Error::msg("Failed to encode frame")); - } - av_packet_unref(pkt); - loop { - let ret = avcodec_receive_packet(*enc_ctx, pkt); - if ret == AVERROR_EOF || ret == AVERROR(EAGAIN) { - break; - } else if ret < 0 { - return Err(Error::msg("Frame read error")); - } - - av_packet_rescale_ts(pkt, (*in_stream).time_base, (*out_stream).time_base); - let ret = av_interleaved_write_frame(out_fmt, pkt); - if ret < 0 { - return Err(Error::msg("Failed to encode frame")); - } - } - } - av_frame_free(&mut frame_out); - av_frame_free(&mut frame); - Ok(()) - } - - unsafe fn setup_decoder( - &mut self, - in_fmt: *mut AVFormatContext, - av_type: AVMediaType, - ) -> Result { - let mut decoder: *const AVCodec = ptr::null_mut(); - let stream_idx = av_find_best_stream(in_fmt, av_type, -1, -1, &mut decoder, 0); - if stream_idx == AVERROR_STREAM_NOT_FOUND { - return Ok(stream_idx); - } - let decoder_ctx = avcodec_alloc_context3(decoder); - if decoder_ctx.is_null() { - return Err(Error::msg("Failed to open video decoder")); - } - - let in_stream = *(*in_fmt).streams.add(stream_idx as usize); - let ret = avcodec_parameters_to_context(decoder_ctx, (*in_stream).codecpar); - if ret < 0 { - return Err(Error::msg("Failed to copy codec params to decoder")); - } - - let ret = avcodec_open2(decoder_ctx, decoder, ptr::null_mut()); - if ret < 0 { - return Err(Error::msg("Failed to open decoder")); - } - - self.decoders.insert(stream_idx as usize, decoder_ctx); - Ok(stream_idx) - } - - unsafe fn setup_encoder( - &mut self, - in_fmt: *mut AVFormatContext, - out_fmt: *mut AVFormatContext, - in_idx: i32, - ) -> Result<(), Error> { - let in_stream = *(*in_fmt).streams.add(in_idx as usize); - let stream_type = (*(*in_stream).codecpar).codec_type; - let out_codec = match stream_type { - AVMEDIA_TYPE_VIDEO => avcodec_find_encoder((*(*out_fmt).oformat).video_codec), - AVMEDIA_TYPE_AUDIO => avcodec_find_encoder((*(*out_fmt).oformat).audio_codec), - _ => ptr::null_mut(), - }; - // not mapped ignore - if out_codec.is_null() { - return Ok(()); - } - let stream = avformat_new_stream(out_fmt, out_codec); - - let encoder_ctx = avcodec_alloc_context3(out_codec); - if encoder_ctx.is_null() { - return Err(Error::msg("Failed to create encoder context")); - } - - match stream_type { - AVMEDIA_TYPE_VIDEO => { - (*encoder_ctx).width = (*(*in_stream).codecpar).width; - (*encoder_ctx).height = (*(*in_stream).codecpar).height; - (*encoder_ctx).pix_fmt = AV_PIX_FMT_YUV420P; - (*encoder_ctx).time_base = (*in_stream).time_base; - (*encoder_ctx).framerate = (*in_stream).avg_frame_rate; - if (*out_codec).id == AVCodecID::AV_CODEC_ID_H264 { - (*encoder_ctx).profile = AV_PROFILE_H264_HIGH; - (*encoder_ctx).level = 50; - (*encoder_ctx).qmin = 20; - (*encoder_ctx).qmax = 30; - } - (*stream).time_base = (*encoder_ctx).time_base; - (*stream).avg_frame_rate = (*encoder_ctx).framerate; - (*stream).r_frame_rate = (*encoder_ctx).framerate; - } - AVMEDIA_TYPE_AUDIO => { - (*encoder_ctx).sample_rate = (*(*in_stream).codecpar).sample_rate; - (*encoder_ctx).sample_fmt = transmute((*(*in_stream).codecpar).format); - (*encoder_ctx).ch_layout = (*(*in_stream).codecpar).ch_layout.clone(); - (*encoder_ctx).time_base = (*in_stream).time_base; - (*stream).time_base = (*encoder_ctx).time_base; - } - _ => {} - } - - if (*(*out_fmt).oformat).flags & AVFMT_GLOBALHEADER == AVFMT_GLOBALHEADER { - (*encoder_ctx).flags |= AV_CODEC_FLAG_GLOBAL_HEADER as libc::c_int; - } - - let ret = avcodec_open2(encoder_ctx, out_codec, ptr::null_mut()); - if ret < 0 { - return Err(Error::msg("Failed to open encoder")); - } - - let ret = avcodec_parameters_from_context((*stream).codecpar, encoder_ctx); - if ret < 0 { - return Err(Error::msg("Failed to open encoder")); - } - - let out_idx = (*stream).index as usize; - // setup scaler if pix_fmt doesnt match - if stream_type == AVMEDIA_TYPE_VIDEO - && (*(*in_stream).codecpar).format != (*(*stream).codecpar).format - { - let sws_ctx = sws_getContext( - (*(*in_stream).codecpar).width, - (*(*in_stream).codecpar).height, - transmute((*(*in_stream).codecpar).format), - (*(*stream).codecpar).width, - (*(*stream).codecpar).height, - transmute((*(*stream).codecpar).format), - 0, - ptr::null_mut(), - ptr::null_mut(), - ptr::null_mut(), - ); - if sws_ctx.is_null() { - return Err(Error::msg("Failed to create sws context")); - } - self.scalers.insert(out_idx, sws_ctx); - } - - self.encoders.insert(out_idx, encoder_ctx); - self.stream_map.insert(in_idx as usize, out_idx); - Ok(()) - } - - unsafe fn flush_output(&mut self, out_fmt: *mut AVFormatContext) -> Result<(), Error> { - let mut pkt = av_packet_alloc(); - for encoder in self.encoders.values() { - avcodec_send_frame(*encoder, ptr::null_mut()); - loop { - let ret = avcodec_receive_packet(*encoder, pkt); - if ret == AVERROR_EOF || ret == AVERROR(EAGAIN) { - break; - } else if ret < 0 { - return Err(Error::msg("Frame read error")); - } - - let ret = av_interleaved_write_frame(out_fmt, pkt); - if ret < 0 { - return Err(Error::msg("Failed to encode frame")); - } - } - } - av_packet_free(&mut pkt); - Ok(()) - } - - unsafe fn free(&mut self) -> Result<(), Error> { - for decoders in self.decoders.values_mut() { - avcodec_free_context(&mut *decoders); - } - self.decoders.clear(); - for encoders in self.encoders.values_mut() { - avcodec_free_context(&mut *encoders); - } - self.encoders.clear(); - for scaler in self.scalers.values_mut() { - sws_freeContext(*scaler); - } - self.scalers.clear(); - - Ok(()) - } - - 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"); - - 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 in_video_stream = self.setup_decoder(dec_fmt, AVMEDIA_TYPE_VIDEO)?; - let in_audio_stream = self.setup_decoder(dec_fmt, AVMEDIA_TYPE_AUDIO)?; - - let out_format = if mime_type.starts_with("image/") { - av_guess_format( - "webp\0".as_ptr() as *const libc::c_char, - ptr::null_mut(), - ptr::null_mut(), - ) - } else if mime_type.starts_with("video/") { - av_guess_format( - "matroska\0".as_ptr() as *const libc::c_char, - ptr::null_mut(), - ptr::null_mut(), - ) - } else { - return Err(Error::msg("Mime type not supported")); - }; - - let out_filename = format!( - "{}\0", - out_path.clone().into_os_string().into_string().unwrap() - ); - let mut out_fmt: *mut AVFormatContext = ptr::null_mut(); - let ret = avformat_alloc_output_context2( - &mut out_fmt, - out_format, - ptr::null_mut(), - out_filename.as_ptr() as *const libc::c_char, - ); - if ret < 0 { - return Err(Error::msg("Failed to create output context")); - } - - let ret = avio_open(&mut (*out_fmt).pb, (*out_fmt).url, AVIO_FLAG_WRITE); - if ret < 0 { - return Err(Error::msg("Failed to open output IO")); - } - - if in_video_stream != AVERROR_STREAM_NOT_FOUND { - self.setup_encoder(dec_fmt, out_fmt, in_video_stream)?; - let video_stream = *(*dec_fmt).streams.add(in_video_stream as usize); - self.width = Some((*(*video_stream).codecpar).width as usize); - self.height = Some((*(*video_stream).codecpar).height as usize); - } - if in_audio_stream != AVERROR_STREAM_NOT_FOUND { - self.setup_encoder(dec_fmt, out_fmt, in_audio_stream)?; - } - - let ret = avformat_init_output(out_fmt, ptr::null_mut()); - if ret < 0 { - return Err(Error::msg("Failed to write output")); - } - - av_dump_format(dec_fmt, 0, ptr::null_mut(), 0); - av_dump_format(out_fmt, 0, ptr::null_mut(), 1); - - let ret = avformat_write_header(out_fmt, ptr::null_mut()); - if ret < 0 { - return Err(Error::msg("Failed to write header to output")); - } - - let mut pkt = av_packet_alloc(); - loop { - let ret = av_read_frame(dec_fmt, pkt); - if ret < 0 { - break; - } - self.transcode_pkt(pkt, dec_fmt, out_fmt)?; - } - - // flush encoder - self.flush_output(out_fmt)?; - - av_write_trailer(out_fmt); - av_packet_free(&mut pkt); - - self.free()?; - - avformat_close_input(&mut dec_fmt); - avformat_free_context(dec_fmt); - avformat_free_context(out_fmt); - - Ok(FileProcessorResult::NewFile(NewFileProcessorResult { - result: out_path, - mime_type: "image/webp".to_string(), - width: self.width.unwrap_or(0), - height: self.height.unwrap_or(0), - image: self.image.unwrap_or_default(), - })) - } - } -}