more robust logic to fetch video frame count

This commit is contained in:
Luna 2023-06-12 17:08:28 -03:00
parent ad268357bd
commit ae9a6f19ba

View file

@ -7,6 +7,7 @@ use axum_macros::debug_handler;
use base64::{engine::general_purpose, Engine as _}; use base64::{engine::general_purpose, Engine as _};
use core::panic; use core::panic;
use ffmpeg_cli::Parameter; use ffmpeg_cli::Parameter;
use ffprobe::{Format, Stream};
use futures_util::{future::ready, StreamExt}; use futures_util::{future::ready, StreamExt};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet}; use std::collections::{HashMap, HashSet};
@ -209,6 +210,41 @@ async fn fetch_frame_as_image(
Ok(()) Ok(())
} }
fn fetch_frame_count_full_decode(path: &std::path::Path) -> anyhow::Result<u64> {
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::<u64>()?;
Ok(res)
}
fn calculate_frame_count(
path: &std::path::Path,
stream: &Stream,
format: &Format,
frame_rate: f64,
) -> anyhow::Result<u64> {
Ok(if let Some(parseable_data) = stream.nb_frames.clone() {
// if we can get it from the stream metadata, use it
parseable_data.parse::<u64>()?
} 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] #[debug_handler]
async fn upload_file( async fn upload_file(
options: Query<Options>, options: Query<Options>,
@ -250,13 +286,17 @@ async fn upload_file(
let info = ffprobe::ffprobe(temp_file.path())?; let info = ffprobe::ffprobe(temp_file.path())?;
let stream = info.streams.get(0).unwrap(); let stream = info.streams.get(0).unwrap();
let total_frame_count = stream.nb_frames.clone().unwrap().parse::<u32>()?; log::debug!("stream = {:?}", stream);
let frame_rate_str = stream.r_frame_rate.clone(); log::debug!("format = {:?}", info.format);
let frame_rate_str = stream.r_frame_rate.clone();
let parts = frame_rate_str.split("/").into_iter().collect::<Vec<_>>(); let parts = frame_rate_str.split("/").into_iter().collect::<Vec<_>>();
let frame_rate: f64 = let frame_rate: f64 =
parts.get(0).unwrap().parse::<f64>()? / parts.get(1).unwrap().parse::<f64>()?; parts.get(0).unwrap().parse::<f64>()? / parts.get(1).unwrap().parse::<f64>()?;
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 total_length_in_seconds = total_frame_count as f64 / frame_rate;
let wanted_frame_skip_seconds = match total_length_in_seconds as usize { let wanted_frame_skip_seconds = match total_length_in_seconds as usize {
0..=10 => 2, 0..=10 => 2,