Add pre-processing denoising

Cargo.toml

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

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<GrainTableParams>>,
+  /// 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

@@ -13,6 +13,7 @@ use crate::api::lookahead::*;
 use crate::api::{EncoderConfig, EncoderStatus, FrameType, Opaque, Packet};
 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::*;
@@ -218,7 +219,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
@@ -246,6 +247,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,
@@ -291,6 +293,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(
@@ -353,6 +366,25 @@ impl<T: Pixel> ContextInner<T> {
       }
     }
 
+    // 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

@@ -2081,6 +2081,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,
@@ -2157,6 +2158,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

@@ -289,6 +289,14 @@ pub fn parse_cli() -> Result<CliOptions, CliError> {
         .takes_value(true)
         .conflicts_with("PHOTON_NOISE")
     )
+    .arg(
+      Arg::new("DENOISING")
+        .help("Enable spatio-temporal denoising, intended to be used with grain synthesis.\n\
+        Takes a strength value 0-50.\nDefault strength is 1/2 of photon noise strength, \
+        or 4 if a photon noise table is specified.")
+        .long("denoise")
+        .takes_value(true)
+    )
     // MASTERING
     .arg(
       Arg::new("PIXEL_RANGE")
@@ -782,6 +790,7 @@ fn parse_config(matches: &ArgMatches) -> Result<EncoderConfig, CliError> {
     if grain_str > 64 {
       panic!("Film grain strength must be between 0-64");
     }
+    cfg.denoise_strength = grain_str / 2;
     // We have to know the video resolution before we can generate a table,
     // so we must handle that elsewhere.
   } else if let Some(table_file) = matches.value_of("PHOTON_NOISE_TABLE") {
@@ -791,7 +800,16 @@ fn parse_config(matches: &ArgMatches) -> Result<EncoderConfig, CliError> {
       parse_grain_table(&contents).expect("Failed to parse film grain table");
     if !table.is_empty() {
       cfg.film_grain_params = Some(table);
+      cfg.denoise_strength = 4;
+    }
+  }
+  if let Some(denoise_str) =
+    matches.value_of("DENOISING").map(|s| s.parse::<u8>().unwrap())
+  {
+    if denoise_str > 50 {
+      panic!("Denoising strength must be between 0-50");
     }
+    cfg.denoise_strength = denoise_str;
   }
 
   if let Some(frame_rate) = matches.value_of("FRAME_RATE") {