Skip to content

Commit

Permalink
Add additional flags, options (master-of-zen#671)
Browse files Browse the repository at this point in the history
* Add never overwrite -n flag

Acts as the opposite of the overwrite -y flag

* Add --scaler flag

This allows the user to control which scaling algorithm is used during scene detection downscaling and VMAF calculations, and also allows the width of the lanczos scaler to be selected from 1 to 9

* Add inputres option for --vmaf-res

This allows the user to have VMAF calculations use the input resolution automatically, without having to type it in manually per file

* Change progress bar characters for Windows

This looks smoother on the default Windows command prompt

* Add --extra-splits-sec option

Allows the user to specify extra splits in seconds. If both frames and seconds are specified, frames will take priority

* Add --ignore-frame-mismatch flag

Allows the user to ignore any reported frame mismatches between the encoder and chunk frame counts, which is useful if an ffmpeg filter changes the frame count, or the input video is encoded badly (ex. "Missing key frame while searching for timestamp" ffmpeg warning)

* formatting

* Fix compilation

oops

---------

Co-authored-by: Josh Holmer <jholmer.in@gmail.com>
  • Loading branch information
woot000 and shssoichiro authored Jun 25, 2023
1 parent ac687ab commit e10880d
Show file tree
Hide file tree
Showing 10 changed files with 162 additions and 33 deletions.
26 changes: 20 additions & 6 deletions av1an-core/src/broker.rs
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,12 @@ impl Display for EncoderCrash {

impl<'a> Broker<'a> {
/// Main encoding loop. set_thread_affinity may be ignored if the value is invalid.
pub fn encoding_loop(self, tx: Sender<()>, mut set_thread_affinity: Option<usize>) {
pub fn encoding_loop(
self,
tx: Sender<()>,
mut set_thread_affinity: Option<usize>,
ignore_frame_mismatch: bool,
) {
assert!(!self.chunk_queue.is_empty());

if !self.chunk_queue.is_empty() {
Expand Down Expand Up @@ -149,7 +154,7 @@ impl<'a> Broker<'a> {
}

while let Ok(mut chunk) = rx.recv() {
if let Err(e) = queue.encode_chunk(&mut chunk, worker_id) {
if let Err(e) = queue.encode_chunk(&mut chunk, worker_id, ignore_frame_mismatch) {
error!("[chunk {}] {}", chunk.index, e);

tx.send(()).unwrap();
Expand All @@ -170,7 +175,12 @@ impl<'a> Broker<'a> {
}
}

fn encode_chunk(&self, chunk: &mut Chunk, worker_id: usize) -> Result<(), Box<EncoderCrash>> {
fn encode_chunk(
&self,
chunk: &mut Chunk,
worker_id: usize,
ignore_frame_mismatch: bool,
) -> Result<(), Box<EncoderCrash>> {
let st_time = Instant::now();

if let Some(ref tq) = self.project.args.target_quality {
Expand All @@ -190,9 +200,13 @@ impl<'a> Broker<'a> {
let passes = chunk.passes;
for current_pass in 1..=passes {
for r#try in 1..=self.project.args.max_tries {
let res = self
.project
.create_pipes(chunk, current_pass, worker_id, padding);
let res = self.project.create_pipes(
chunk,
current_pass,
worker_id,
padding,
ignore_frame_mismatch,
);
if let Err((e, frames)) = res {
dec_bar(frames);

Expand Down
36 changes: 29 additions & 7 deletions av1an-core/src/context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -308,7 +308,11 @@ impl Av1anContext {

let (tx, rx) = mpsc::channel();
let handle = s.spawn(|_| {
broker.encoding_loop(tx, self.args.set_thread_affinity);
broker.encoding_loop(
tx,
self.args.set_thread_affinity,
self.args.ignore_frame_mismatch,
);
});

// Queue::encoding_loop only sends a message if there was an error (meaning a chunk crashed)
Expand Down Expand Up @@ -348,11 +352,25 @@ impl Av1anContext {
}

if let Some(ref tq) = self.args.target_quality {
let mut temp_res = tq.vmaf_res.to_string();
if tq.vmaf_res == "inputres" {
let inputres = self.args.input.resolution()?;
temp_res.push_str(&format!(
"{}x{}",
&inputres.0.to_string(),
&inputres.1.to_string()
));
temp_res.to_string();
} else {
temp_res = tq.vmaf_res.to_string();
}

if let Err(e) = vmaf::plot(
self.args.output_file.as_ref(),
&self.args.input,
tq.model.as_deref(),
tq.vmaf_res.as_str(),
temp_res.as_str(),
tq.vmaf_scaler.as_str(),
1,
tq.vmaf_filter.as_deref(),
tq.vmaf_threads,
Expand Down Expand Up @@ -400,6 +418,7 @@ impl Av1anContext {
current_pass: u8,
worker_id: usize,
padding: usize,
ignore_frame_mismatch: bool,
) -> Result<(), (Box<EncoderCrash>, u64)> {
update_mp_chunk(worker_id, chunk.index, padding);

Expand Down Expand Up @@ -618,11 +637,13 @@ impl Av1anContext {
let encoded_frames = num_frames(chunk.output().as_ref());

let err_str = match encoded_frames {
Ok(encoded_frames) if encoded_frames != chunk.frames() => Some(format!(
"FRAME MISMATCH: chunk {}: {encoded_frames}/{} (actual/expected frames)",
chunk.index,
chunk.frames()
)),
Ok(encoded_frames) if !ignore_frame_mismatch && encoded_frames != chunk.frames() => {
Some(format!(
"FRAME MISMATCH: chunk {}: {encoded_frames}/{} (actual/expected frames)",
chunk.index,
chunk.frames()
))
}
Err(error) => Some(format!(
"FAILED TO COUNT FRAMES: chunk {}: {error}",
chunk.index
Expand Down Expand Up @@ -689,6 +710,7 @@ impl Av1anContext {
self.frames,
self.args.min_scene_len,
self.args.verbosity,
self.args.scaler.as_str(),
self.args.sc_pix_format,
self.args.sc_method,
self.args.sc_downscale_height,
Expand Down
8 changes: 6 additions & 2 deletions av1an-core/src/progress_bar.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,16 @@ use once_cell::sync::OnceCell;
use crate::util::printable_base10_digits;
use crate::{get_done, Verbosity};

const PROGRESS_CHARS: &str = "█▉▊▋▌▍▎▏ ";
const PROGRESS_CHARS: &str = if cfg!(windows) {
"█▓▒░ "
} else {
"█▉▊▋▌▍▎▏ "
};

const INDICATIF_PROGRESS_TEMPLATE: &str = if cfg!(windows) {
// Do not use a spinner on Windows since the default console cannot display
// the characters used for the spinner
"{elapsed_precise:.bold} {wide_bar:.blue/white.dim} {percent:.bold} {pos} ({fps:.bold}, eta {fixed_eta}{msg})"
"{elapsed_precise:.bold} {wide_bar:.blue/white.dim} {percent:.bold} {pos} ({fps:.bold}, eta {fixed_eta}{msg})"
} else {
"{spinner:.green.bold} {elapsed_precise:.bold} ▕{wide_bar:.blue/white.dim}▏ {percent:.bold} {pos} ({fps:.bold}, eta {fixed_eta}{msg})"
};
Expand Down
21 changes: 17 additions & 4 deletions av1an-core/src/scene_detect.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ pub fn av_scenechange_detect(
total_frames: usize,
min_scene_len: usize,
verbosity: Verbosity,
sc_scaler: &str,
sc_pix_format: Option<Pixel>,
sc_method: ScenecutMethod,
sc_downscale_height: Option<usize>,
Expand Down Expand Up @@ -54,6 +55,7 @@ pub fn av_scenechange_detect(
})
},
min_scene_len,
sc_scaler,
sc_pix_format,
sc_method,
sc_downscale_height,
Expand All @@ -75,12 +77,19 @@ pub fn scene_detect(
total_frames: usize,
callback: Option<&dyn Fn(usize)>,
min_scene_len: usize,
sc_scaler: &str,
sc_pix_format: Option<Pixel>,
sc_method: ScenecutMethod,
sc_downscale_height: Option<usize>,
zones: &[Scene],
) -> anyhow::Result<Vec<Scene>> {
let (mut decoder, bit_depth) = build_decoder(input, encoder, sc_pix_format, sc_downscale_height)?;
let (mut decoder, bit_depth) = build_decoder(
input,
encoder,
sc_scaler,
sc_pix_format,
sc_downscale_height,
)?;

let mut scenes = Vec::new();
let mut cur_zone = zones.first().filter(|frame| frame.start_frame == 0);
Expand Down Expand Up @@ -194,6 +203,7 @@ pub fn scene_detect(
fn build_decoder(
input: &Input,
encoder: Encoder,
sc_scaler: &str,
sc_pix_format: Option<Pixel>,
sc_downscale_height: Option<usize>,
) -> anyhow::Result<(y4m::Decoder<impl Read>, usize)> {
Expand All @@ -202,12 +212,15 @@ fn build_decoder(
(Some(sdh), Some(spf)) => into_smallvec![
"-vf",
format!(
"format={},scale=-2:'min({},ih)'",
"format={},scale=-2:'min({},ih)':flags={}",
spf.descriptor().unwrap().name(),
sdh
sdh,
sc_scaler
)
],
(Some(sdh), None) => into_smallvec!["-vf", format!("scale=-2:'min({sdh},ih)'")],
(Some(sdh), None) => {
into_smallvec!["-vf", format!("scale=-2:'min({sdh},ih)':flags={sc_scaler}")]
}
(None, Some(spf)) => into_smallvec!["-pix_fmt", spf.descriptor().unwrap().name()],
(None, None) => smallvec![],
};
Expand Down
2 changes: 2 additions & 0 deletions av1an-core/src/scenes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -295,6 +295,8 @@ fn get_test_args() -> Av1anContext {
workers: 1,
set_thread_affinity: None,
zones: None,
scaler: String::new(),
ignore_frame_mismatch: false,
};
Av1anContext {
vs_script: None,
Expand Down
6 changes: 6 additions & 0 deletions av1an-core/src/settings.rs
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ pub struct EncodeArgs {

pub chunk_method: ChunkMethod,
pub chunk_order: ChunkOrdering,
pub scaler: String,
pub scenes: Option<PathBuf>,
pub split_method: SplitMethod,
pub sc_pix_format: Option<Pixel>,
Expand All @@ -46,6 +47,7 @@ pub struct EncodeArgs {
pub extra_splits_len: Option<usize>,
pub min_scene_len: usize,
pub force_keyframes: Vec<usize>,
pub ignore_frame_mismatch: bool,

pub max_tries: usize,

Expand Down Expand Up @@ -127,6 +129,10 @@ properly into a mkv file. Specify mkvmerge as the concatenation method by settin
warn!("It is not recommended to use the \"select\" chunk method, as it is very slow");
}

if self.ignore_frame_mismatch {
warn!("The output video's frame count may differ, and VMAF calculations may be incorrect");
}

if let Some(vmaf_path) = &self
.target_quality
.as_ref()
Expand Down
2 changes: 2 additions & 0 deletions av1an-core/src/target_quality.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ const VMAF_PERCENTILE: f64 = 0.01;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TargetQuality {
pub vmaf_res: String,
pub vmaf_scaler: String,
pub vmaf_filter: Option<String>,
pub vmaf_threads: usize,
pub model: Option<PathBuf>,
Expand Down Expand Up @@ -255,6 +256,7 @@ impl TargetQuality {
&fl_path,
self.model.as_ref(),
&self.vmaf_res,
&self.vmaf_scaler,
self.probing_rate,
self.vmaf_filter.as_deref(),
self.vmaf_threads,
Expand Down
9 changes: 6 additions & 3 deletions av1an-core/src/vmaf.rs
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,7 @@ pub fn plot(
reference: &Input,
model: Option<impl AsRef<Path>>,
res: &str,
scaler: &str,
sample_rate: usize,
filter: Option<&str>,
threads: usize,
Expand Down Expand Up @@ -153,6 +154,7 @@ pub fn plot(
&json_file,
model,
res,
scaler,
sample_rate,
filter,
threads,
Expand All @@ -168,6 +170,7 @@ pub fn run_vmaf(
stat_file: impl AsRef<Path>,
model: Option<impl AsRef<Path>>,
res: &str,
scaler: &str,
sample_rate: usize,
vmaf_filter: Option<&str>,
threads: usize,
Expand Down Expand Up @@ -229,10 +232,10 @@ pub fn run_vmaf(
cmd.arg(encoded);
cmd.args(["-r", "60", "-i", "-", "-filter_complex"]);

let distorted = format!("[0:v]scale={}:flags=bicubic:force_original_aspect_ratio=decrease,setpts=PTS-STARTPTS[distorted];", &res);
let distorted = format!("[0:v]scale={}:flags={}:force_original_aspect_ratio=decrease,setpts=PTS-STARTPTS,setsar=1[distorted];", &res, &scaler);
let reference = format!(
"[1:v]{}scale={}:flags=bicubic:force_original_aspect_ratio=decrease,setpts=PTS-STARTPTS[ref];",
filter, &res
"[1:v]{}scale={}:flags={}:force_original_aspect_ratio=decrease,setpts=PTS-STARTPTS,setsar=1[ref];",
filter, &res, &scaler
);

cmd.arg(format!("{distorted}{reference}{vmaf}"));
Expand Down
Loading

0 comments on commit e10880d

Please sign in to comment.