Skip to content

Commit

Permalink
Enable photon noise arg in stable
Browse files Browse the repository at this point in the history
Add pre-processing denoising

Align denoising arrays

Autovectorization pass
  • Loading branch information
shssoichiro committed Sep 5, 2022
1 parent 8c260d2 commit eb5b8f8
Show file tree
Hide file tree
Showing 13 changed files with 674 additions and 10 deletions.
10 changes: 8 additions & 2 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,6 @@ binaries = [
"fern",
"console",
"av-metrics",
"nom",
]
default = ["binaries", "asm", "threading", "signal_support"]
asm = ["nasm-rs", "cc", "regex"]
Expand Down Expand Up @@ -100,11 +99,18 @@ wasm-bindgen = { version = "0.2.63", optional = true }
rust_hawktracer = "0.7.0"
arrayref = "0.3.6"
const_fn_assert = "0.1.2"
nom = { version = "7.0.0", optional = true }
# `unreachable!` macro which panics in debug mode
# and optimizes away in release mode
new_debug_unreachable = "1.0.4"
once_cell = "1.13.0"
av1-grain = { version = "0.2.0", features = ["serialize"] }
serde-big-array = { version = "0.4.1", optional = true }
# Used for parsing film grain table files
nom = "7.0.0"
# Used as a data holder during denoising
ndarray = "0.15.4"
# Used for running FFTs during denoising
ndrustfft = "0.3.0"

[dependencies.image]
version = "0.24.3"
Expand Down
1 change: 1 addition & 0 deletions clippy.toml
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
too-many-arguments-threshold = 16
cognitive-complexity-threshold = 40
trivial-copy-size-limit = 16 # 128-bits = 2 64-bit registers
doc-valid-idents = ["DFTTest"]
msrv = "1.59"
3 changes: 3 additions & 0 deletions src/api/config/encoder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,8 @@ pub struct EncoderConfig {
pub tune: Tune,
/// Parameters for grain synthesis.
pub film_grain_params: Option<Vec<GrainTableSegment>>,
/// Strength of denoising, 0 = disabled
pub denoise_strength: u8,
/// Number of tiles horizontally. Must be a power of two.
///
/// Overridden by [`tiles`], if present.
Expand Down Expand Up @@ -159,6 +161,7 @@ impl EncoderConfig {
bitrate: 0,
tune: Tune::default(),
film_grain_params: None,
denoise_strength: 0,
tile_cols: 0,
tile_rows: 0,
tiles: 0,
Expand Down
34 changes: 33 additions & 1 deletion src/api/internal.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ use crate::api::{
};
use crate::color::ChromaSampling::Cs400;
use crate::cpu_features::CpuFeatureLevel;
use crate::denoise::{DftDenoiser, TB_MIDPOINT};
use crate::dist::get_satd;
use crate::encoder::*;
use crate::frame::*;
Expand Down Expand Up @@ -220,7 +221,7 @@ impl<T: Pixel> FrameData<T> {
}
}

type FrameQueue<T> = BTreeMap<u64, Option<Arc<Frame<T>>>>;
pub(crate) type FrameQueue<T> = BTreeMap<u64, Option<Arc<Frame<T>>>>;
type FrameDataQueue<T> = BTreeMap<u64, Option<FrameData<T>>>;

// the fields pub(super) are accessed only by the tests
Expand Down Expand Up @@ -248,6 +249,7 @@ pub(crate) struct ContextInner<T: Pixel> {
/// Maps `output_frameno` to `gop_input_frameno_start`.
pub(crate) gop_input_frameno_start: BTreeMap<u64, u64>,
keyframe_detector: SceneChangeDetector<T>,
denoiser: Option<DftDenoiser<T>>,
pub(crate) config: Arc<EncoderConfig>,
seq: Arc<Sequence>,
pub(crate) rc_state: RCState,
Expand Down Expand Up @@ -295,6 +297,17 @@ impl<T: Pixel> ContextInner<T> {
lookahead_distance,
seq.clone(),
),
denoiser: if enc.denoise_strength > 0 {
Some(DftDenoiser::new(
enc.denoise_strength as f32 / 10.0,
enc.width,
enc.height,
enc.bit_depth as u8,
enc.chroma_sampling,
))
} else {
None
},
config: Arc::new(enc.clone()),
seq,
rc_state: RCState::new(
Expand Down Expand Up @@ -359,6 +372,25 @@ impl<T: Pixel> ContextInner<T> {
self.t35_q.insert(input_frameno, params.t35_metadata);
}

// If denoising is enabled, run it now because we want the entire
// encoding process, including lookahead, to see the denoised frame.
if let Some(ref mut denoiser) = self.denoiser {
loop {
let denoiser_frame = denoiser.cur_frameno;
if (!is_flushing
&& input_frameno >= denoiser_frame + TB_MIDPOINT as u64)
|| (is_flushing && Some(denoiser_frame) < self.limit)
{
self.frame_q.insert(
denoiser_frame,
Some(Arc::new(denoiser.filter_frame(&self.frame_q).unwrap())),
);
} else {
break;
}
}
}

if !self.needs_more_frame_q_lookahead(self.next_lookahead_frame) {
let lookahead_frames = self
.frame_q
Expand Down
2 changes: 2 additions & 0 deletions src/api/test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2131,6 +2131,7 @@ fn log_q_exp_overflow() {
tile_cols: 0,
tile_rows: 0,
tiles: 0,
denoise_strength: 0,
speed_settings: SpeedSettings {
multiref: false,
fast_deblock: true,
Expand Down Expand Up @@ -2207,6 +2208,7 @@ fn guess_frame_subtypes_assert() {
tile_cols: 0,
tile_rows: 0,
tiles: 0,
denoise_strength: 0,
speed_settings: SpeedSettings {
multiref: false,
fast_deblock: true,
Expand Down
26 changes: 23 additions & 3 deletions src/bin/common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,6 @@ pub struct CliOptions {
pub still_picture: bool,
/// Uses grain synthesis to add photon noise to the resulting encode.
/// Takes a strength value 0-64.
#[cfg(feature = "unstable")]
#[clap(
long,
conflicts_with = "film-grain-table",
Expand All @@ -185,6 +184,17 @@ pub struct CliOptions {
help_heading = "ENCODE SETTINGS"
)]
pub photon_noise: u8,
/// Enable spatio-temporal denoising, intended to be used with grain synthesis.
/// Takes a strength value 0-50.
///
/// Default strength is 1/2 of photon noise strength,
/// or 4 if a photon noise table is specified.
#[clap(
long,
value_parser = clap::value_parser!(u8).range(0..=50),
help_heading = "ENCODE SETTINGS"
)]
pub denoise: Option<u8>,
/// Uses a film grain table file to apply grain synthesis to the encode.
/// Uses the same table file format as aomenc and svt-av1.
#[clap(
Expand Down Expand Up @@ -334,7 +344,6 @@ pub struct ParsedCliOptions {
pub save_config: Option<PathBuf>,
#[cfg(feature = "unstable")]
pub slots: usize,
#[cfg(feature = "unstable")]
pub generate_grain_strength: u8,
}

Expand Down Expand Up @@ -482,7 +491,6 @@ pub fn parse_cli() -> Result<ParsedCliOptions, CliError> {
save_config: save_config_path,
#[cfg(feature = "unstable")]
slots,
#[cfg(feature = "unstable")]
generate_grain_strength: matches.photon_noise,
})
}
Expand Down Expand Up @@ -676,7 +684,19 @@ fn parse_config(matches: &CliOptions) -> Result<EncoderConfig, CliError> {
.expect("Failed to parse film grain table");
if !table.is_empty() {
cfg.film_grain_params = Some(table);
cfg.denoise_strength = 4;
}
} else if matches.photon_noise > 0 {
cfg.denoise_strength = matches.photon_noise / 2;
// We have to know the video resolution before we can generate a table,
// so we must handle that elsewhere.
}
// A user set denoise strength overrides the defaults above
if let Some(denoise_str) = matches.denoise {
if denoise_str > 50 {
panic!("Denoising strength must be between 0-50");
}
cfg.denoise_strength = denoise_str;
}

if let Some(frame_rate) = matches.frame_rate {
Expand Down
1 change: 0 additions & 1 deletion src/bin/rav1e-ch.rs
Original file line number Diff line number Diff line change
Expand Up @@ -473,7 +473,6 @@ fn run() -> Result<(), error::CliError> {
cli.enc.time_base = video_info.time_base;
}

#[cfg(feature = "unstable")]
if cli.generate_grain_strength > 0 && cli.enc.film_grain_params.is_none() {
cli.enc.film_grain_params = Some(vec![generate_photon_noise_params(
0,
Expand Down
1 change: 0 additions & 1 deletion src/bin/rav1e.rs
Original file line number Diff line number Diff line change
Expand Up @@ -456,7 +456,6 @@ fn run() -> Result<(), error::CliError> {
cli.enc.time_base = video_info.time_base;
}

#[cfg(feature = "unstable")]
if cli.generate_grain_strength > 0 && cli.enc.film_grain_params.is_none() {
cli.enc.film_grain_params = Some(vec![generate_photon_noise_params(
0,
Expand Down
Loading

0 comments on commit eb5b8f8

Please sign in to comment.