diff --git a/Cargo.toml b/Cargo.toml index c315a5b..74f6016 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,9 +9,14 @@ edition = "2021" anyhow = "1.0.70" axum = { version = "0.6.12", features = ["tokio", "multipart"] } axum-macros = "0.3.7" +ffmpeg-cli = "0.1.0" +ffprobe = "0.3.3" +filepath = "0.1.2" +futures-util = "0.3.28" log = "0.4.17" pretty_env_logger = "0.4.0" reqwest = { version = "0.11.16", features = ["json", "multipart"] } serde = { version = "1.0.159", features = ["derive"] } serde_json = "1.0.95" +tempfile = "3.5.0" tokio = { version = "1.27.0", features = ["full"] } diff --git a/src/main.rs b/src/main.rs index 4560d9c..a498dc1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,7 +5,13 @@ use axum::response::{IntoResponse, Response}; use axum::routing::post; use axum::{Json, Router}; use axum_macros::debug_handler; +use filepath::FilePath; +use reqwest::header::WWW_AUTHENTICATE; +use serde::__private::de::FlatInternallyTaggedAccess; use serde::{Deserialize, Serialize}; +use std::collections::HashSet; +use std::io::Read; +use std::io::Write; #[tokio::main] async fn main() { @@ -97,6 +103,56 @@ async fn send_image_to_dd( Ok(json_response) } +use std::process::Stdio; + +use ffmpeg_cli::Parameter; + +use futures_util::{future::ready, StreamExt}; + +async fn fetch_frame_as_image( + input_path: &str, + output_path: &str, + frame_index: usize, +) -> anyhow::Result<()> { + let frame_index_param = format!("select=eq(n\\,{})", frame_index); + let builder = ffmpeg_cli::FfmpegBuilder::new() + .stderr(Stdio::piped()) + .option(Parameter::Single("nostdin")) + .option(Parameter::Single("y")) + .input(ffmpeg_cli::File::new(input_path)) + .output( + ffmpeg_cli::File::new(output_path) + .option(Parameter::KeyValue("vf", &frame_index_param)) + .option(Parameter::KeyValue("vframes", "1")), + ); + + log::debug!("running {:?}", builder); + let ffmpeg = builder.run().await.unwrap(); + + log::debug!("run"); + ffmpeg + .progress + .for_each(|x| { + dbg!(&x); + // lmao + // x.unwrap(); + ready(()) + }) + .await; + + log::debug!("run"); + let output = ffmpeg.process.wait_with_output().unwrap(); + + log::debug!("out"); + println!( + "{}\nstderr:\n{}", + output.status, + std::str::from_utf8(&output.stderr).unwrap() + ); + + Ok(()) +} + #[debug_handler] async fn upload_file( options: Query, @@ -121,14 +177,68 @@ async fn upload_file( } if let Some(file_contents) = maybe_file_contents { - let json_response = send_image_to_dd( - file_contents.to_vec(), - maybe_file_name.unwrap(), - &maybe_file_type.unwrap(), - &options, - ) - .await?; - Ok((StatusCode::OK, Json(json_response))) + let file_type = maybe_file_type.unwrap(); + let file_name = maybe_file_name.unwrap(); + let is_video = file_type.starts_with("video/") || file_name.ends_with(".mp4"); + if is_video { + let mut final_tag_set = HashSet::new(); + + let mut temp_file = tempfile::NamedTempFile::new()?; + temp_file.write_all(&file_contents.to_vec())?; + + log::debug!("tmp path: {:?}", temp_file.path()); + + let info = ffprobe::ffprobe(temp_file.path())?; + let stream = info.streams.get(0).unwrap(); + + let all_frames = stream.nb_frames.clone().unwrap().parse::()?; + let frame_rate_str = stream.r_frame_rate.clone(); + + let parts = frame_rate_str.split("/").into_iter().collect::>(); + let frame_rate = + parts.get(0).unwrap().parse::()? / parts.get(1).unwrap().parse::()?; + + let wanted_frame_skip = frame_rate.try_into().unwrap(); // every second + + let temporary_frame_dir = tempfile::tempdir()?; + + let temporary_frame_path = + format!("{}/frame.png", temporary_frame_dir.path().to_string_lossy()); + log::info!("path:{}", &temporary_frame_path); + + for frame_number in (0..all_frames).step_by(wanted_frame_skip) { + log::info!("extracting frame {:?}", frame_number); + fetch_frame_as_image( + temp_file.path().to_str().unwrap(), + &temporary_frame_path, + frame_number.try_into().unwrap(), + ) + .await?; + + let mut actual_frame_file = std::fs::File::open(&temporary_frame_path)?; + let mut frame_data = vec![]; + actual_frame_file.read_to_end(&mut frame_data)?; + let tags_from_frame = if let WrappedResponse::Tags(tags_from_frame) = + send_image_to_dd(frame_data, "amongus.png".to_string(), "image/png", &options) + .await? + { + tags_from_frame + } else { + todo!() + }; + + for tag in tags_from_frame { + final_tag_set.insert(tag); + } + } + + let response = WrappedResponse::Tags(final_tag_set.into_iter().collect::>()); + Ok((StatusCode::OK, Json(response))) + } else { + let json_response = + send_image_to_dd(file_contents.to_vec(), file_name, &file_type, &options).await?; + Ok((StatusCode::OK, Json(json_response))) + } } else { Ok(( StatusCode::INTERNAL_SERVER_ERROR,