ab-av1 0.3.3

AV1 encoding with fast VMAF sampling
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79

mod command;
mod console_ext;
mod ffprobe;
mod process;
mod sample;
mod svtav1;
mod temporary;
mod vmaf;
mod yuv;

use anyhow::anyhow;
use clap::Parser;
use futures::FutureExt;
use std::{io, time::Duration};
use tokio::signal;

const SAMPLE_SIZE_S: u64 = 20;
const SAMPLE_SIZE: Duration = Duration::from_secs(SAMPLE_SIZE_S);

#[derive(Parser)]
#[clap(version, about)]
enum Command {
    SampleEncode(command::sample_encode::Args),
    Vmaf(command::vmaf::Args),
    Encode(command::encode::Args),
    CrfSearch(command::crf_search::Args),
    AutoEncode(command::auto_encode::Args),
    PrintCompletions(command::print_completions::Args),
}

#[tokio::main(flavor = "current_thread")]
async fn main() -> anyhow::Result<()> {
    let action = Command::parse();

    action.ensure_temp_dir_exists().await?;
    let keep = action.keep_temp_files();

    let local = tokio::task::LocalSet::new();

    let command = local.run_until(match action {
        Command::SampleEncode(args) => command::sample_encode(args).boxed_local(),
        Command::Vmaf(args) => command::vmaf(args).boxed_local(),
        Command::Encode(args) => command::encode(args).boxed_local(),
        Command::CrfSearch(args) => command::crf_search(args).boxed_local(),
        Command::AutoEncode(args) => command::auto_encode(args).boxed_local(),
        Command::PrintCompletions(args) => return command::print_completions(args),
    });

    let out = tokio::select! {
        r = command => r,
        _ = signal::ctrl_c() => Err(anyhow!("ctrl_c")),
    };

    temporary::clean(keep).await;

    out
}

impl Command {
    fn keep_temp_files(&self) -> bool {
        match self {
            Self::SampleEncode(args) => args.keep,
            _ => false,
        }
    }

    async fn ensure_temp_dir_exists(&self) -> io::Result<()> {
        let temp_dir = match self {
            Self::SampleEncode(args) => &args.sample.temp_dir,
            Self::CrfSearch(args) => &args.sample.temp_dir,
            Self::AutoEncode(args) => &args.search.sample.temp_dir,
            _ => &None,
        };
        if let Some(dir) = temp_dir {
            tokio::fs::create_dir_all(dir).await?;
        }
        Ok(())
    }
}