From 324d43c908a285391c25e9457c909e090aa9218e Mon Sep 17 00:00:00 2001 From: Alex Butler Date: Thu, 27 Jun 2024 18:26:37 +0100 Subject: [PATCH] Add encoding/vmaf + progress logs --- src/command/crf_search.rs | 13 +++++++-- src/command/encode.rs | 14 ++++++++-- src/command/sample_encode.rs | 17 ++++++++---- src/command/vmaf.rs | 18 +++++++++--- src/log.rs | 54 ++++++++++++++++++++++++++++++++++++ src/main.rs | 1 + src/vmaf.rs | 8 +++++- 7 files changed, 110 insertions(+), 15 deletions(-) create mode 100644 src/log.rs diff --git a/src/command/crf_search.rs b/src/command/crf_search.rs index e6379e7..c340875 100644 --- a/src/command/crf_search.rs +++ b/src/command/crf_search.rs @@ -104,7 +104,7 @@ pub async fn crf_search(mut args: Args) -> anyhow::Result<()> { let best = best?; if std::io::stderr().is_terminal() { - // encode how-to hint + predictions + // encode how-to hint eprintln!( "\n{} {}\n", style("Encode with:").dim(), @@ -112,13 +112,22 @@ pub async fn crf_search(mut args: Args) -> anyhow::Result<()> { ); } - info!("crf {} successful", best.crf()); StdoutFormat::Human.print_result(&best, input_is_image); Ok(()) } pub async fn run( + args: &Args, + input_probe: Arc, + bar: ProgressBar, +) -> Result { + _run(args, input_probe, bar) + .await + .inspect(|s| info!("crf {} successful", s.crf())) +} + +async fn _run( Args { args, min_vmaf, diff --git a/src/command/encode.rs b/src/command/encode.rs index b792dbc..7e998cf 100644 --- a/src/command/encode.rs +++ b/src/command/encode.rs @@ -6,16 +6,18 @@ use crate::{ console_ext::style, ffmpeg, ffprobe::{self, Ffprobe}, + log::ProgressLogger, process::FfmpegOut, temporary::{self, TempKind}, }; use clap::Parser; use console::style; use indicatif::{HumanBytes, ProgressBar, ProgressStyle}; +use log::info; use std::{ path::{Path, PathBuf}, sync::Arc, - time::Duration, + time::{Duration, Instant}, }; use tokio::fs; use tokio_stream::StreamExt; @@ -89,8 +91,13 @@ pub async fn run( anyhow::bail!("--stereo-downmix cannot be used with --acodec copy"); } - let mut enc = ffmpeg::encode(enc_args, &output, has_audio, audio_codec, stereo_downmix)?; + info!( + "encoding {}", + output.file_name().and_then(|n| n.to_str()).unwrap_or("") + ); + let mut enc = ffmpeg::encode(enc_args, &output, has_audio, audio_codec, stereo_downmix)?; + let mut logger = ProgressLogger::new(module_path!(), Instant::now()); let mut stream_sizes = None; while let Some(progress) = enc.next().await { match progress? { @@ -98,8 +105,9 @@ pub async fn run( if fps > 0.0 { bar.set_message(format!("{fps} fps, ")); } - if probe.duration.is_ok() { + if let Ok(d) = &probe.duration { bar.set_position(time.as_micros_u64()); + logger.update(*d, time, fps); } } FfmpegOut::StreamSizes { diff --git a/src/command/sample_encode.rs b/src/command/sample_encode.rs index 8c259d8..496c483 100644 --- a/src/command/sample_encode.rs +++ b/src/command/sample_encode.rs @@ -8,9 +8,10 @@ use crate::{ console_ext::style, ffmpeg::{self, FfmpegEncodeArgs}, ffprobe::{self, Ffprobe}, + log::ProgressLogger, process::FfmpegOut, - sample, temporary, vmaf, - vmaf::VmafOut, + sample, temporary, + vmaf::{self, VmafOut}, SAMPLE_SIZE, SAMPLE_SIZE_S, }; use anyhow::ensure; @@ -158,6 +159,8 @@ pub async fn run( let (sample, sample_size) = sample?; + info!("encoding sample {sample_n}/{samples} crf {crf}",); + // encode sample let result = match cache::cached_encode( cache, @@ -184,7 +187,7 @@ pub async fn run( ); if samples > 1 { info!( - "Sample {sample_n}/{samples} crf {crf} VMAF {:.2} ({:.0}%) (cache)", + "sample {sample_n}/{samples} crf {crf} VMAF {:.2} ({:.0}%) (cache)", result.vmaf_score, 100.0 * result.encoded_size as f32 / sample_size as f32, ); @@ -194,6 +197,7 @@ pub async fn run( (None, key) => { bar.set_message("encoding,"); let b = Instant::now(); + let mut logger = ProgressLogger::new(module_path!(), b); let (encoded_sample, mut output) = ffmpeg::encode_sample( FfmpegEncodeArgs { input: &sample, @@ -210,6 +214,7 @@ pub async fn run( if fps > 0.0 { bar.set_message(format!("enc {fps} fps,")); } + logger.update(SAMPLE_SIZE, time, fps); } } let encode_time = b.elapsed(); @@ -229,6 +234,7 @@ pub async fn run( args.vfilter.as_deref(), ), )?; + let mut logger = ProgressLogger::new("ab_av1::vmaf", Instant::now()); let mut vmaf_score = -1.0; while let Some(vmaf) = vmaf.next().await { match vmaf { @@ -246,6 +252,7 @@ pub async fn run( if fps > 0.0 { bar.set_message(format!("vmaf {fps} fps,")); } + logger.update(SAMPLE_SIZE, time, fps); } VmafOut::Progress(_) => {} VmafOut::Err(e) => return Err(e), @@ -262,7 +269,7 @@ pub async fn run( ); if samples > 1 { info!( - "Sample {sample_n}/{samples} crf {crf} VMAF {vmaf_score:.2} ({:.0}%)", + "sample {sample_n}/{samples} crf {crf} VMAF {vmaf_score:.2} ({:.0}%)", 100.0 * encoded_size as f32 / sample_size as f32, ); } @@ -320,7 +327,7 @@ pub async fn run( if print_output { if std::io::stderr().is_terminal() { - // encode how-to hint + predictions + // encode how-to hint eprintln!( "\n{} {}\n", style("Encode with:").dim(), diff --git a/src/command/vmaf.rs b/src/command/vmaf.rs index 3eed6c8..acee14c 100644 --- a/src/command/vmaf.rs +++ b/src/command/vmaf.rs @@ -4,13 +4,16 @@ use crate::{ PROGRESS_CHARS, }, ffprobe, + log::ProgressLogger, process::FfmpegOut, - vmaf, - vmaf::VmafOut, + vmaf::{self, VmafOut}, }; use clap::Parser; use indicatif::{ProgressBar, ProgressStyle}; -use std::{path::PathBuf, time::Duration}; +use std::{ + path::PathBuf, + time::{Duration, Instant}, +}; use tokio_stream::StreamExt; /// Full VMAF score calculation, distorted file vs reference file. @@ -62,6 +65,7 @@ pub async fn vmaf( let rprobe = ffprobe::probe(&reference); let rpix_fmt = rprobe.pixel_format().unwrap_or(PixelFormat::Yuv444p10le); let nframes = dprobe.nframes().or_else(|_| rprobe.nframes()); + let duration = dprobe.duration.as_ref().or(rprobe.duration.as_ref()); if let Ok(nframes) = nframes { bar.set_length(nframes); } @@ -75,6 +79,7 @@ pub async fn vmaf( reference_vfilter.as_deref(), ), )?; + let mut logger = ProgressLogger::new(module_path!(), Instant::now()); let mut vmaf_score = -1.0; while let Some(vmaf) = vmaf.next().await { match vmaf { @@ -82,13 +87,18 @@ pub async fn vmaf( vmaf_score = score; break; } - VmafOut::Progress(FfmpegOut::Progress { frame, fps, .. }) => { + VmafOut::Progress(FfmpegOut::Progress { + frame, fps, time, .. + }) => { if fps > 0.0 { bar.set_message(format!("vmaf {fps} fps, ")); } if nframes.is_ok() { bar.set_position(frame); } + if let Ok(total) = duration { + logger.update(*total, time, fps); + } } VmafOut::Progress(FfmpegOut::StreamSizes { .. }) => {} VmafOut::Err(e) => return Err(e), diff --git a/src/log.rs b/src/log.rs new file mode 100644 index 0000000..2c2dea5 --- /dev/null +++ b/src/log.rs @@ -0,0 +1,54 @@ +use indicatif::HumanDuration; +use log::{info, log_enabled, Level}; +use std::time::{Duration, Instant}; + +/// Struct that info logs progress messages on a stream action like encoding. +#[derive(Debug)] +pub struct ProgressLogger { + target: &'static str, + start: Instant, + log_count: u32, +} + +impl ProgressLogger { + pub fn new(target: &'static str, start: Instant) -> Self { + Self { + target, + start, + log_count: 0, + } + } + + /// Update and potentially log progress on a stream action. + /// * `total` total duration of the stream + /// * `complete` the duration that has been completed at this time + /// * `fps` frames per second + pub fn update(&mut self, total: Duration, completed: Duration, fps: f32) { + if log_enabled!(Level::Info) && completed > Duration::ZERO { + let done = completed.as_secs_f64() / total.as_secs_f64(); + + let elapsed = self.start.elapsed(); + + let before_count = self.log_count; + while elapsed > self.next_log() { + self.log_count += 1; + } + if before_count == self.log_count { + return; + } + + let eta = Duration::from_secs_f64(elapsed.as_secs_f64() / done) - elapsed; + info!( + target: self.target, + "{:.0}%, {fps} fps, eta {}", + done * 100.0, + HumanDuration(eta) + ); + } + } + + /// First log after >=16s, then >=32s etc + fn next_log(&self) -> Duration { + Duration::from_secs(2_u64.pow(self.log_count + 4)) + } +} diff --git a/src/main.rs b/src/main.rs index dc6e858..df19aa2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,6 +3,7 @@ mod console_ext; mod ffmpeg; mod ffprobe; mod float; +mod log; mod process; mod sample; mod temporary; diff --git a/src/vmaf.rs b/src/vmaf.rs index 1e9174a..91cf09f 100644 --- a/src/vmaf.rs +++ b/src/vmaf.rs @@ -1,7 +1,7 @@ //! vmaf logic use crate::process::{exit_ok_stderr, Chunks, CommandExt, FfmpegOut}; use anyhow::Context; -use log::debug; +use log::{debug, info}; use std::path::Path; use tokio::process::Command; use tokio_process_stream::{Item, ProcessChunkStream}; @@ -14,6 +14,12 @@ pub fn run( distorted: &Path, filter_complex: &str, ) -> anyhow::Result> { + info!( + "vmaf {} vs reference {}", + distorted.file_name().and_then(|n| n.to_str()).unwrap_or(""), + reference.file_name().and_then(|n| n.to_str()).unwrap_or(""), + ); + let mut cmd = Command::new("ffmpeg"); cmd.kill_on_drop(true) .arg2("-r", "24")