Enable photon noise arg in stable

Add pre-processing denoising

Align denoising arrays

Autovectorization pass

Author: shssoichiro

Cargo.toml

@@ -36,7 +36,6 @@ binaries = [
     "fern",
     "console",
     "av-metrics",
-    "nom",
 ]
 default = ["binaries", "asm", "threading", "signal_support"]
 asm = ["nasm-rs", "cc", "regex"]
@@ -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"

clippy.toml

@@ -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"

src/api/config/encoder.rs

@@ -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.
@@ -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,

src/api/internal.rs

@@ -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::*;
@@ -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
@@ -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,
@@ -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(
@@ -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

src/api/test.rs

@@ -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,
@@ -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,

src/bin/common.rs

@@ -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",
@@ -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(
@@ -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,
 }
 
@@ -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,
   })
 }
@@ -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 {

src/bin/rav1e-ch.rs

@@ -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,

src/bin/rav1e.rs

@@ -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,