Compute segmentation offsets dynamically

src/activity.rs

@@ -52,6 +52,7 @@ impl ActivityMask {
         variances.push(variance);
       }
     }
+
     ActivityMask { variances: variances.into_boxed_slice() }
   }
 

src/api/config/encoder.rs

@@ -108,6 +108,9 @@ pub struct EncoderConfig {
   /// [`tile_rows`]: #structfield.tile_rows
   pub tiles: usize,
 
+  /// Strength of variance AQ. Defaults to 1.0. Disabled automatically by Tune=Psnr.
+  pub aq_strength: f64,
+
   /// Settings which affect the encoding speed vs. quality trade-off.
   pub speed_settings: SpeedSettings,
 }
@@ -165,6 +168,7 @@ impl EncoderConfig {
       tile_cols: 0,
       tile_rows: 0,
       tiles: 0,
+      aq_strength: 1.0,
       speed_settings: SpeedSettings::from_preset(speed),
     }
   }

src/api/internal.rs

@@ -1314,6 +1314,7 @@ impl<T: Pixel> ContextInner<T> {
             frame_data.fi.sequence.bit_depth,
             &mut coded_data.activity_scales,
           );
+          frame_data.fi.compute_spatiotemporal_scores();
         } else {
           coded_data.activity_mask = ActivityMask::default();
         }

src/api/test.rs

@@ -2082,6 +2082,7 @@ fn log_q_exp_overflow() {
     tile_cols: 0,
     tile_rows: 0,
     tiles: 0,
+    aq_strength: 1.0,
     speed_settings: SpeedSettings {
       multiref: false,
       fast_deblock: true,
@@ -2159,6 +2160,7 @@ fn guess_frame_subtypes_assert() {
     tile_cols: 0,
     tile_rows: 0,
     tiles: 0,
+    aq_strength: 1.0,
     speed_settings: SpeedSettings {
       multiref: false,
       fast_deblock: true,

src/bin/common.rs

@@ -162,6 +162,15 @@ pub struct CliOptions {
   /// Still picture mode
   #[clap(long, help_heading = "ENCODE SETTINGS")]
   pub still_picture: bool,
+  /// Adjusts the strength of adaptive quantization
+  #[clap(
+    long,
+    value_parser,
+    hide = true,
+    default_value_t = 1.0,
+    help_heading = "ENCODE SETTINGS"
+  )]
+  pub aq_strength: f64,
   /// Uses grain synthesis to add photon noise to the resulting encode.
   /// Takes a strength value 0-64.
   #[cfg(feature = "unstable")]
@@ -601,6 +610,11 @@ fn parse_config(matches: &CliOptions) -> Result<EncoderConfig, CliError> {
   if cfg.tune == Tune::Psychovisual {
     cfg.speed_settings.transform.tx_domain_distortion = false;
   }
+  if cfg.tune == Tune::Psnr {
+    cfg.aq_strength = 0.0;
+  } else {
+    cfg.aq_strength = matches.aq_strength;
+  }
 
   cfg.tile_cols = matches.tile_cols;
   cfg.tile_rows = matches.tile_rows;

src/context/partition_unit.rs

@@ -249,7 +249,9 @@ impl<'a> ContextWriter<'a> {
     symbol_with_update!(self, w, skip as u32, cdf, 2);
   }
 
-  pub fn get_segment_pred(&self, bo: TileBlockOffset) -> (u8, u8) {
+  pub fn get_segment_pred(
+    &self, bo: TileBlockOffset, last_active_segid: u8,
+  ) -> (u8, u8) {
     let mut prev_ul = -1;
     let mut prev_u = -1;
     let mut prev_l = -1;
@@ -288,7 +290,8 @@ impl<'a> ContextWriter<'a> {
     } else {
       r = if prev_ul == prev_u { prev_u } else { prev_l };
     }
-    (r as u8, cdf_index)
+
+    ((r as u8).min(last_active_segid), cdf_index)
   }
 
   pub fn write_cfl_alphas<W: Writer>(&mut self, w: &mut W, cfl: CFLParams) {
@@ -434,7 +437,7 @@ impl<'a> ContextWriter<'a> {
     &mut self, w: &mut W, bo: TileBlockOffset, bsize: BlockSize, skip: bool,
     last_active_segid: u8,
   ) {
-    let (pred, cdf_index) = self.get_segment_pred(bo);
+    let (pred, cdf_index) = self.get_segment_pred(bo, last_active_segid);
     if skip {
       self.bc.blocks.set_segmentation_idx(bo, bsize, pred);
       return;

src/encoder.rs

@@ -572,6 +572,7 @@ pub struct SegmentationState {
   pub last_active_segid: u8,
   pub features: [[bool; SegLvl::SEG_LVL_MAX as usize]; 8],
   pub data: [[i16; SegLvl::SEG_LVL_MAX as usize]; 8],
+  pub segment_map: [usize; 8],
 }
 
 // Frame Invariants are invariant inside a frame
@@ -681,6 +682,11 @@ pub struct CodedFrameData<T: Pixel> {
   /// Pre-computed activity_scale.
   pub activity_scales: Box<[DistortionScale]>,
   pub activity_mask: ActivityMask,
+  /// Combined metric of activity and distortion
+  pub spatiotemporal_scores: Box<[DistortionScale]>,
+  /// Preliminary segmentation decisions, 0-7, at the importance block level.
+  /// These are used for determining how many and how strong the segmentation should be.
+  pub imp_b_segments: Box<[u8]>,
 }
 
 impl<T: Pixel> CodedFrameData<T> {
@@ -708,6 +714,8 @@ impl<T: Pixel> CodedFrameData<T> {
       ]
       .into_boxed_slice(),
       activity_mask: Default::default(),
+      spatiotemporal_scores: Default::default(),
+      imp_b_segments: Default::default(),
     }
   }
 }
@@ -1156,6 +1164,38 @@ impl<T: Pixel> FrameInvariants<T> {
     }
   }
 
+  // Assumes that we have already computed activity scales and distortion scales
+  pub fn compute_spatiotemporal_scores(&mut self) {
+    let coded_data = self.coded_frame_data.as_ref().unwrap();
+    let mut scores = vec![
+      DistortionScale::default();
+      coded_data.w_in_imp_b * coded_data.h_in_imp_b
+    ]
+    .into_boxed_slice();
+    let bsize = BlockSize::from_width_and_height(
+      IMPORTANCE_BLOCK_SIZE,
+      IMPORTANCE_BLOCK_SIZE,
+    );
+    for y_in_imp_b in 0..coded_data.h_in_imp_b {
+      for x_in_imp_b in 0..coded_data.w_in_imp_b {
+        let block_offset = PlaneBlockOffset(BlockOffset {
+          x: x_in_imp_b << IMPORTANCE_BLOCK_TO_BLOCK_SHIFT,
+          y: y_in_imp_b << IMPORTANCE_BLOCK_TO_BLOCK_SHIFT,
+        });
+        let scale = spatiotemporal_scale(self, block_offset, bsize);
+        let imp_b_idx = y_in_imp_b * coded_data.w_in_imp_b + x_in_imp_b;
+        scores[imp_b_idx] = scale;
+      }
+    }
+
+    let segments =
+      scores.iter().map(|&s| segment_idx_from_distortion(s)).collect();
+
+    let coded_data = self.coded_frame_data.as_mut().unwrap();
+    coded_data.spatiotemporal_scores = scores;
+    coded_data.imp_b_segments = segments;
+  }
+
   #[inline(always)]
   pub fn sb_size_log2(&self) -> usize {
     self.sequence.tiling.sb_size_log2

src/fuzzing.rs

@@ -232,6 +232,8 @@ impl Arbitrary for ArbitraryEncoder {
       tile_cols: u.int_in_range(0..=2)?,
       tile_rows: u.int_in_range(0..=2)?,
       tiles: u.int_in_range(0..=16)?,
+      aq_strength: *u
+        .choose(&[-1.0, -0.5, 0.0, 0.1, 0.5, 1.0, 1.5, 2.5, 69.0])?,
 
       chroma_sampling: *u.choose(&[
         ChromaSampling::Cs420,