noest 0.1.0

Tool to estimate noise in a video and output the results, optionally as ISO for use with photon noise in AV1 encoding.
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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

#[cfg(not(target_env = "msvc"))]
use tikv_jemallocator::Jemalloc;

#[cfg(not(target_env = "msvc"))]
#[global_allocator]
static GLOBAL: Jemalloc = Jemalloc;

use std::{path::{PathBuf, Path}, time::Duration, io::{stderr, IsTerminal}};

use clap::{Parser, ValueEnum};
use indicatif::{ProgressBar, ProgressDrawTarget};

mod noise_estimate;
mod progress_stuff;

#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
enum Format {
    /// Output the median for each plane
    Medians,
    /// Output just the median
    Median,

    /// Output the iso value (LUMA:ISO, CHROMA:ISO)
    ISO,

    /// Returns Av1an arguments to use
    Av1an
}

#[derive(Parser)]
#[command(author, version, about, long_about = None)]
struct Cli {
    /// The video to run noise estimation on
    input: PathBuf,

    /// Formatting to use when outputting
    #[arg(long, short, value_enum, default_value_t = Format::Medians)]
    format: Format,

    /// Whether to output ISO as strength or not
    /// A Strength value is ISO/100 (4 = 400 ISO)
    /// 
    /// Only applies when using --format iso
    #[arg(long, short, default_value_t = false, verbatim_doc_comment)]
    strength: bool,

    /// Amount of threads to use for decoding and noise estimation
    #[arg(long, short, default_value_t = 1)]
    threads: usize,
}

fn main() {
    let cli = Cli::parse();
    let (y, u, v) = run(&cli.input, cli.threads);

    let avg_y: f64 = y.iter().sum::<f64>() / y.len() as f64;
    let avg_u: f64 = u.iter().sum::<f64>() / u.len() as f64;
    let avg_v: f64 = v.iter().sum::<f64>() / v.len() as f64;
    //println!("Y: {avg_y}\nU: {avg_u}\nV: {avg_v}");

    match cli.format {
        Format::Medians => {
            println!("Y:{avg_y}");
            println!("U:{avg_u}");
            println!("V:{avg_v}");
        },
        Format::ISO => {
            if avg_y > avg_u+avg_v {
                // It's likely luma noise
                let iso = iso_from_noise(avg_y, cli.strength);
                println!("LUMA:{iso}");
            } else {
                // It's likely chroma noise

                let noise_floor = (avg_y+avg_u+avg_v)/3.0;
                let iso = iso_from_noise(noise_floor, cli.strength);
                println!("CHROMA:{iso}");
            }
        },
        Format::Av1an => {
            if avg_y > avg_u+avg_v {
                // It's likely luma noise
                let iso = iso_from_noise(avg_y, true);
                println!("--photon noise {iso}");
            } else {
                // It's likely chroma noise

                let noise_floor = (avg_y+avg_u+avg_v)/3.0;
                let iso = iso_from_noise(noise_floor, true);
                println!("--photon-noise {iso} --chroma-noise");
            }
        },
        Format::Median => {
            if avg_y > avg_u+avg_v {
                // It's likely luma noise
                println!("{avg_y}");
            } else {
                let noise = (avg_y+avg_u+avg_v)/3.0;
                println!("{noise}");
            }
        },
    }


}

// Converts the estimated noise value to a photon-noise iso
// Outputs in iso*k
// 
// Todo: More precise coefficients
fn iso_from_noise(noise: f64, stength: bool) -> usize {

    let (a, b, c) = (0.745974214, 1570.33842, -5.16323056);

    let res = libm::exp((noise-c)/a)-b;
    let iso = libm::round(res/100.0).clamp(0.0, 64.0);

    if stength { iso as usize } else { iso as usize * 100 }
}

fn run(path: &Path, threads: usize) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
    let (frame_count_opt, receiver) = noise_estimate::run(path, threads);

    let progress = if stderr().is_terminal() {
        let pb = if let Some(frame_count) = frame_count_opt {
            ProgressBar::new(frame_count as u64)
                .with_style(progress_stuff::pretty_progress_style())
        } else {
            ProgressBar::new_spinner().with_style(progress_stuff::pretty_spinner_style())
        };
        pb.set_draw_target(ProgressDrawTarget::stderr());
        pb.enable_steady_tick(Duration::from_millis(100));
        pb.reset();
        pb.reset_eta();
        pb.reset_elapsed();
        pb.set_position(0);
        pb
    } else {
        ProgressBar::hidden()
    };

    let mut results = Vec::new();
    for score in receiver {
        results.push(score);
        progress.inc(1);
    }

    progress.finish();

    let results_y: Vec<f64> = results.iter().map(|r| { r[0] }).collect();
    let results_u: Vec<f64> = results.iter().map(|r| { r[1] }).collect();
    let results_v: Vec<f64> = results.iter().map(|r| { r[2] }).collect();

    (results_y, results_u, results_v)
}