Scale chroma distortion.

src/encoder.rs

@@ -501,6 +501,7 @@ pub struct FrameInvariants<T: Pixel> {
   pub ac_delta_q: [i8; 3],
   pub lambda: f64,
   pub me_lambda: f64,
+  pub dist_scale: [f64; 3],
   pub me_range_scale: u8,
   pub use_tx_domain_distortion: bool,
   pub use_tx_domain_rate: bool,
@@ -677,6 +678,7 @@ impl<T: Pixel> FrameInvariants<T> {
       dc_delta_q: [0; 3],
       ac_delta_q: [0; 3],
       lambda: 0.0,
+      dist_scale: [1.0; 3],
       me_lambda: 0.0,
       me_range_scale: 1,
       use_tx_domain_distortion,
@@ -859,6 +861,7 @@ impl<T: Pixel> FrameInvariants<T> {
     self.lambda =
       qps.lambda * ((1 << (2 * (self.sequence.bit_depth - 8))) as f64);
     self.me_lambda = self.lambda.sqrt();
+    self.dist_scale = qps.dist_scale;
 
     let q = bexp64(qps.log_target_q as i64 + q57(QSCALE)) as f32;
     /* These coefficients were trained on libaom. */
@@ -1795,6 +1798,7 @@ pub fn write_tx_blocks<T: Pixel>(
     }
   }
 
+  tx_dist = ((tx_dist as f64) * fi.dist_scale[0]) as i64;
   if luma_only {
     return tx_dist;
   };
@@ -1879,7 +1883,7 @@ pub fn write_tx_blocks<T: Pixel>(
               || skip
               || dist >= 0
           );
-          tx_dist += dist;
+          tx_dist += ((dist as f64) * fi.dist_scale[p]) as i64;
         }
       }
     }
@@ -1939,7 +1943,7 @@ pub fn write_tx_tree<T: Pixel>(
   assert!(
     !fi.use_tx_domain_distortion || need_recon_pixel || skip || dist >= 0
   );
-  tx_dist += dist;
+  tx_dist += ((dist as f64) * fi.dist_scale[0]) as i64;
 
   if luma_only {
     return tx_dist;
@@ -2017,7 +2021,7 @@ pub fn write_tx_tree<T: Pixel>(
               || skip
               || dist >= 0
           );
-          tx_dist += dist;
+          tx_dist += ((dist as f64) * fi.dist_scale[p]) as i64;
         }
       }
     }

src/rate.rs

@@ -552,6 +552,7 @@ pub struct QuantizerParameters {
   pub dc_qi: [u8; 3],
   pub ac_qi: [u8; 3],
   pub lambda: f64,
+  pub dist_scale: [f64; 3],
 }
 
 const Q57_SQUARE_EXP_SCALE: f64 =
@@ -572,9 +573,17 @@ impl QuantizerParameters {
   ) -> QuantizerParameters {
     let scale = q57(QSCALE + bit_depth as i32 - 8);
     let quantizer = bexp64(log_target_q + scale);
-    let (offset_u, offset_v) = chroma_offset(log_target_q);
-    let quantizer_u = bexp64(log_target_q + offset_u + scale);
-    let quantizer_v = bexp64(log_target_q + offset_v + scale);
+    let (offset_u, offset_v) = chroma_offset(log_target_q, uv_m_q8);
+    let log_target_q_u = log_target_q + offset_u;
+    let log_target_q_v = log_target_q + offset_v;
+    let quantizer_u = bexp64(log_target_q_u + scale);
+    let quantizer_v = bexp64(log_target_q_v + scale);
+    let lambda = (::std::f64::consts::LN_2 / 6.0)
+      * ((log_target_q as f64) * Q57_SQUARE_EXP_SCALE).exp();
+    let lambda_u = (::std::f64::consts::LN_2 / 6.0)
+      * ((log_target_q_u as f64) * Q57_SQUARE_EXP_SCALE).exp();
+    let lambda_v = (::std::f64::consts::LN_2 / 6.0)
+      * ((log_target_q_v as f64) * Q57_SQUARE_EXP_SCALE).exp();
     QuantizerParameters {
       log_base_q,
       log_target_q,
@@ -589,8 +598,8 @@ impl QuantizerParameters {
         select_ac_qi(quantizer_u, bit_depth).max(1),
         select_ac_qi(quantizer_v, bit_depth).max(1),
       ],
-      lambda: (::std::f64::consts::LN_2 / 6.0)
-        * ((log_target_q as f64) * Q57_SQUARE_EXP_SCALE).exp(),
+      lambda,
+      dist_scale: [1.0, lambda / lambda_u, lambda / lambda_v],
     }
   }
 }

src/rdo.rs

@@ -355,6 +355,8 @@ fn compute_distortion<T: Pixel>(
     ),
   };
 
+  distortion = (fi.dist_scale[0] * distortion as f64) as u64;
+
   if !luma_only {
     let PlaneConfig { xdec, ydec, .. } = ts.input.planes[1].cfg;
 
@@ -372,7 +374,7 @@ fn compute_distortion<T: Pixel>(
       for p in 1..3 {
         let input_region = ts.input_tile.planes[p].subregion(area);
         let rec_region = ts.rec.planes[p].subregion(area);
-        distortion += sse_wxh(
+        distortion += (sse_wxh(
           &input_region,
           &rec_region,
           w_uv,
@@ -384,7 +386,8 @@ fn compute_distortion<T: Pixel>(
               bsize,
             )
           },
-        );
+        ) as f64
+          * fi.dist_scale[p]) as u64;
       }
     };
   }
@@ -422,6 +425,8 @@ fn compute_tx_distortion<T: Pixel>(
     (tx_dist as f64 * bias) as u64
   };
 
+  distortion = (fi.dist_scale[0] * distortion as f64) as u64;
+
   if !luma_only && skip {
     let PlaneConfig { xdec, ydec, .. } = ts.input.planes[1].cfg;
 
@@ -439,7 +444,7 @@ fn compute_tx_distortion<T: Pixel>(
       for p in 1..3 {
         let input_region = ts.input_tile.planes[p].subregion(area);
         let rec_region = ts.rec.planes[p].subregion(area);
-        distortion += sse_wxh(
+        distortion += (sse_wxh(
           &input_region,
           &rec_region,
           w_uv,
@@ -451,7 +456,8 @@ fn compute_tx_distortion<T: Pixel>(
               bsize,
             )
           },
-        );
+        ) as f64
+          * fi.dist_scale[p]) as u64;
       }
     }
   }
@@ -1620,7 +1626,7 @@ fn rdo_loop_plane_error<T: Pixel>(
           fi,
           ts.to_frame_block_offset(bo),
           BlockSize::BLOCK_8X8,
-        );
+        ) * fi.dist_scale[pli];
         err += (value as f64 * bias) as u64;
       }
     }