diff --git a/src/main.rs b/src/main.rs index 6e5365d..c4998f0 100644 --- a/src/main.rs +++ b/src/main.rs @@ -7,6 +7,7 @@ use axum_macros::debug_handler; use base64::{engine::general_purpose, Engine as _}; use core::panic; use ffmpeg_cli::Parameter; +use ffprobe::{Format, Stream}; use futures_util::{future::ready, StreamExt}; use serde::{Deserialize, Serialize}; use std::collections::{HashMap, HashSet}; @@ -209,6 +210,41 @@ async fn fetch_frame_as_image( Ok(()) } +fn fetch_frame_count_full_decode(path: &std::path::Path) -> anyhow::Result { + let config = ffprobe::ConfigBuilder::new().count_frames(true).build(); + let new_info = ffprobe::ffprobe_config(config, path)?; + let res = new_info + .streams + .get(0) + .unwrap() + .nb_read_frames + .clone() + .unwrap() + .parse::()?; + Ok(res) +} + +fn calculate_frame_count( + path: &std::path::Path, + stream: &Stream, + format: &Format, + frame_rate: f64, +) -> anyhow::Result { + Ok(if let Some(parseable_data) = stream.nb_frames.clone() { + // if we can get it from the stream metadata, use it + parseable_data.parse::()? + } else if let Some(parseable_data) = format.try_get_duration() { + // this is a std::time::duration + // multiply that by frame rate and we get total frame count (approximate) + log::warn!("fetching duration from format metadata..."); + let seconds = parseable_data?.as_secs_f64(); + (seconds * frame_rate) as u64 + } else { + log::warn!("file didn't provide frame metadata, calculating it ourselves..."); + fetch_frame_count_full_decode(path)? + }) +} + #[debug_handler] async fn upload_file( options: Query, @@ -250,13 +286,17 @@ async fn upload_file( let info = ffprobe::ffprobe(temp_file.path())?; let stream = info.streams.get(0).unwrap(); - let total_frame_count = stream.nb_frames.clone().unwrap().parse::()?; - let frame_rate_str = stream.r_frame_rate.clone(); + log::debug!("stream = {:?}", stream); + log::debug!("format = {:?}", info.format); + let frame_rate_str = stream.r_frame_rate.clone(); let parts = frame_rate_str.split("/").into_iter().collect::>(); let frame_rate: f64 = parts.get(0).unwrap().parse::()? / parts.get(1).unwrap().parse::()?; + let total_frame_count = + calculate_frame_count(temp_file.path(), &stream, &info.format, frame_rate)?; + let total_length_in_seconds = total_frame_count as f64 / frame_rate; let wanted_frame_skip_seconds = match total_length_in_seconds as usize { 0..=10 => 2,