diff --git a/src/encoder.rs b/src/encoder.rs index c9f12cb156..fc5f93dfaf 100644 --- a/src/encoder.rs +++ b/src/encoder.rs @@ -501,6 +501,7 @@ pub struct FrameInvariants { 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 FrameInvariants { 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 FrameInvariants { 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( } } + 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( || skip || dist >= 0 ); - tx_dist += dist; + tx_dist += ((dist as f64) * fi.dist_scale[p]) as i64; } } } diff --git a/src/rate.rs b/src/rate.rs index ca1ad4ef11..73e5a84bb9 100644 --- a/src/rate.rs +++ b/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 = @@ -573,8 +574,16 @@ impl 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 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], } } } diff --git a/src/rdo.rs b/src/rdo.rs index ff5ea0740a..2c8d345955 100644 --- a/src/rdo.rs +++ b/src/rdo.rs @@ -355,6 +355,8 @@ fn compute_distortion( ), }; + 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( 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( bsize, ) }, - ); + ) as f64 + * fi.dist_scale[p]) as u64; } }; } @@ -422,6 +425,8 @@ fn compute_tx_distortion( (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( 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( bsize, ) }, - ); + ) as f64 + * fi.dist_scale[p]) as u64; } } }