oximedia_codec/av1/
loop_optimization.rs

1//! AV1 loop filter and post-processing optimization.
2//!
3//! This module optimizes loop filter parameters for AV1 encoding:
4//!
5//! - Loop filter strength selection
6//! - CDEF (Constrained Directional Enhancement Filter) strength
7//! - Restoration filter parameters
8//! - Film grain parameter encoding (if applicable)
9//!
10//! # Loop Filter Optimization
11//!
12//! The loop filter reduces blocking artifacts by smoothing block boundaries.
13//! This module finds optimal filter strengths that minimize distortion while
14//! maximizing coding efficiency.
15
16#![forbid(unsafe_code)]
17#![allow(dead_code)]
18#![allow(clippy::cast_possible_truncation)]
19#![allow(clippy::cast_precision_loss)]
20#![allow(clippy::cast_sign_loss)]
21#![allow(clippy::similar_names)]
22
23use super::block::BlockSize;
24use super::cdef::{CdefParams, CdefStrength};
25use super::loop_filter::LoopFilterParams;
26
27// =============================================================================
28// Constants
29// =============================================================================
30
31/// Maximum loop filter level.
32const MAX_LOOP_FILTER_LEVEL: u8 = 63;
33
34/// Maximum CDEF strength.
35const MAX_CDEF_STRENGTH: u8 = 15;
36
37/// Number of loop filter levels to test.
38const FILTER_LEVELS_TO_TEST: usize = 5;
39
40/// Number of CDEF strengths to test.
41const CDEF_STRENGTHS_TO_TEST: usize = 4;
42
43// =============================================================================
44// Loop Filter Optimizer
45// =============================================================================
46
47/// Loop filter parameter optimizer.
48#[derive(Clone, Debug)]
49pub struct LoopFilterOptimizer {
50    /// Current loop filter parameters.
51    params: LoopFilterParams,
52    /// Lambda for RD optimization.
53    lambda: f32,
54    /// Enable RD optimization.
55    rd_optimization: bool,
56}
57
58impl LoopFilterOptimizer {
59    /// Create a new loop filter optimizer.
60    #[must_use]
61    pub fn new(lambda: f32) -> Self {
62        Self {
63            params: LoopFilterParams::default(),
64            lambda,
65            rd_optimization: true,
66        }
67    }
68
69    /// Optimize loop filter level for a frame.
70    ///
71    /// Tests multiple filter levels and selects the one with best RD cost.
72    pub fn optimize_filter_level(
73        &mut self,
74        src: &[u8],
75        recon: &[u8],
76        width: usize,
77        height: usize,
78        qp: u8,
79    ) -> u8 {
80        if !self.rd_optimization {
81            // Fast mode: use QP-based heuristic
82            return self.filter_level_from_qp(qp);
83        }
84
85        let base_level = self.filter_level_from_qp(qp);
86        let mut best_level = base_level;
87        let mut best_cost = f32::MAX;
88
89        // Test levels around base
90        for delta in -(FILTER_LEVELS_TO_TEST as i32 / 2)..=(FILTER_LEVELS_TO_TEST as i32 / 2) {
91            let level = (i32::from(base_level) + delta * 4)
92                .clamp(0, i32::from(MAX_LOOP_FILTER_LEVEL)) as u8;
93
94            let cost = self.evaluate_filter_level(src, recon, width, height, level);
95
96            if cost < best_cost {
97                best_cost = cost;
98                best_level = level;
99            }
100        }
101
102        self.params.level = [best_level, best_level, best_level, best_level];
103        best_level
104    }
105
106    /// Compute filter level from QP (heuristic).
107    fn filter_level_from_qp(&self, qp: u8) -> u8 {
108        // Higher QP -> stronger filter
109        ((i32::from(qp) * 3) / 2).clamp(0, i32::from(MAX_LOOP_FILTER_LEVEL)) as u8
110    }
111
112    /// Evaluate RD cost for a filter level.
113    fn evaluate_filter_level(
114        &self,
115        src: &[u8],
116        recon: &[u8],
117        width: usize,
118        height: usize,
119        level: u8,
120    ) -> f32 {
121        // Apply filter (simplified - just compute distortion)
122        let distortion = self.compute_distortion(src, recon, width, height);
123
124        // Estimate rate (filter level signaling)
125        let rate = f32::from(level) * 0.1;
126
127        distortion + self.lambda * rate
128    }
129
130    /// Compute distortion between source and reconstruction.
131    fn compute_distortion(&self, src: &[u8], recon: &[u8], width: usize, height: usize) -> f32 {
132        let mut sse = 0u64;
133        let total = (width * height).min(src.len()).min(recon.len());
134
135        for i in 0..total {
136            let diff = i32::from(src[i]) - i32::from(recon[i]);
137            sse += (diff * diff) as u64;
138        }
139
140        sse as f32
141    }
142
143    /// Get optimized loop filter parameters.
144    #[must_use]
145    pub const fn params(&self) -> &LoopFilterParams {
146        &self.params
147    }
148
149    /// Set lambda for RD optimization.
150    pub fn set_lambda(&mut self, lambda: f32) {
151        self.lambda = lambda;
152    }
153
154    /// Enable/disable RD optimization.
155    pub fn set_rd_optimization(&mut self, enabled: bool) {
156        self.rd_optimization = enabled;
157    }
158}
159
160impl Default for LoopFilterOptimizer {
161    fn default() -> Self {
162        Self::new(1.0)
163    }
164}
165
166// =============================================================================
167// CDEF Optimizer
168// =============================================================================
169
170/// CDEF (Constrained Directional Enhancement Filter) optimizer.
171#[derive(Clone, Debug)]
172pub struct CdefOptimizer {
173    /// Current CDEF parameters.
174    params: CdefParams,
175    /// Lambda for RD optimization.
176    lambda: f32,
177}
178
179impl CdefOptimizer {
180    /// Create a new CDEF optimizer.
181    #[must_use]
182    pub fn new(lambda: f32) -> Self {
183        Self {
184            params: CdefParams::default(),
185            lambda,
186        }
187    }
188
189    /// Optimize CDEF strength for a block.
190    pub fn optimize_strength(
191        &mut self,
192        src: &[u8],
193        recon: &[u8],
194        width: usize,
195        height: usize,
196        _block_size: BlockSize,
197    ) -> CdefStrength {
198        let mut best_strength = CdefStrength::default();
199        let mut best_cost = f32::MAX;
200
201        for primary in 0..CDEF_STRENGTHS_TO_TEST {
202            for secondary in 0..CDEF_STRENGTHS_TO_TEST {
203                let strength = CdefStrength {
204                    primary: primary as u8,
205                    secondary: secondary as u8,
206                };
207
208                let cost = self.evaluate_cdef_strength(src, recon, width, height, &strength);
209
210                if cost < best_cost {
211                    best_cost = cost;
212                    best_strength = strength;
213                }
214            }
215        }
216
217        best_strength
218    }
219
220    /// Evaluate RD cost for CDEF strength.
221    fn evaluate_cdef_strength(
222        &self,
223        src: &[u8],
224        recon: &[u8],
225        width: usize,
226        height: usize,
227        strength: &CdefStrength,
228    ) -> f32 {
229        // Compute distortion
230        let mut sse = 0u64;
231        let total = (width * height).min(src.len()).min(recon.len());
232
233        for i in 0..total {
234            let diff = i32::from(src[i]) - i32::from(recon[i]);
235            sse += (diff * diff) as u64;
236        }
237
238        let distortion = sse as f32;
239
240        // Estimate rate
241        let rate = f32::from(strength.primary + strength.secondary) * 0.5;
242
243        distortion + self.lambda * rate
244    }
245
246    /// Get optimized CDEF parameters.
247    #[must_use]
248    pub const fn params(&self) -> &CdefParams {
249        &self.params
250    }
251}
252
253impl Default for CdefOptimizer {
254    fn default() -> Self {
255        Self::new(1.0)
256    }
257}
258
259// =============================================================================
260// Restoration Filter Optimizer
261// =============================================================================
262
263/// Restoration filter type.
264#[derive(Clone, Copy, Debug, PartialEq, Eq)]
265pub enum RestorationType {
266    /// No restoration.
267    None = 0,
268    /// Wiener filter.
269    Wiener = 1,
270    /// Self-guided filter.
271    Sgrproj = 2,
272}
273
274/// Restoration filter optimizer.
275#[derive(Clone, Debug)]
276pub struct RestorationOptimizer {
277    /// Restoration type.
278    restoration_type: RestorationType,
279    /// Lambda for RD optimization.
280    lambda: f32,
281}
282
283impl RestorationOptimizer {
284    /// Create a new restoration optimizer.
285    #[must_use]
286    pub fn new(lambda: f32) -> Self {
287        Self {
288            restoration_type: RestorationType::None,
289            lambda,
290        }
291    }
292
293    /// Optimize restoration type for a frame.
294    pub fn optimize_restoration(
295        &mut self,
296        src: &[u8],
297        recon: &[u8],
298        width: usize,
299        height: usize,
300    ) -> RestorationType {
301        let mut best_type = RestorationType::None;
302        let mut best_cost = f32::MAX;
303
304        for rtype in [
305            RestorationType::None,
306            RestorationType::Wiener,
307            RestorationType::Sgrproj,
308        ] {
309            let cost = self.evaluate_restoration(src, recon, width, height, rtype);
310
311            if cost < best_cost {
312                best_cost = cost;
313                best_type = rtype;
314            }
315        }
316
317        self.restoration_type = best_type;
318        best_type
319    }
320
321    /// Evaluate RD cost for restoration type.
322    fn evaluate_restoration(
323        &self,
324        src: &[u8],
325        recon: &[u8],
326        width: usize,
327        height: usize,
328        rtype: RestorationType,
329    ) -> f32 {
330        // Simplified evaluation
331        let base_distortion = self.compute_distortion(src, recon, width, height);
332
333        let rate = match rtype {
334            RestorationType::None => 0.0,
335            RestorationType::Wiener => 100.0,
336            RestorationType::Sgrproj => 80.0,
337        };
338
339        let distortion_reduction = match rtype {
340            RestorationType::None => 0.0,
341            RestorationType::Wiener => base_distortion * 0.05,
342            RestorationType::Sgrproj => base_distortion * 0.03,
343        };
344
345        (base_distortion - distortion_reduction) + self.lambda * rate
346    }
347
348    /// Compute distortion.
349    fn compute_distortion(&self, src: &[u8], recon: &[u8], width: usize, height: usize) -> f32 {
350        let mut sse = 0u64;
351        let total = (width * height).min(src.len()).min(recon.len());
352
353        for i in 0..total {
354            let diff = i32::from(src[i]) - i32::from(recon[i]);
355            sse += (diff * diff) as u64;
356        }
357
358        sse as f32
359    }
360
361    /// Get restoration type.
362    #[must_use]
363    pub const fn restoration_type(&self) -> RestorationType {
364        self.restoration_type
365    }
366}
367
368impl Default for RestorationOptimizer {
369    fn default() -> Self {
370        Self::new(1.0)
371    }
372}
373
374// =============================================================================
375// Film Grain Parameters
376// =============================================================================
377
378/// Film grain synthesis parameters.
379#[derive(Clone, Debug, Default)]
380pub struct FilmGrainParams {
381    /// Enable film grain synthesis.
382    pub enabled: bool,
383    /// Grain seed.
384    pub grain_seed: u16,
385    /// Luma scaling points.
386    pub luma_points: Vec<(u8, u8)>,
387    /// Chroma scaling points.
388    pub chroma_points: Vec<(u8, u8)>,
389}
390
391impl FilmGrainParams {
392    /// Create new film grain parameters.
393    #[must_use]
394    pub const fn new() -> Self {
395        Self {
396            enabled: false,
397            grain_seed: 0,
398            luma_points: Vec::new(),
399            chroma_points: Vec::new(),
400        }
401    }
402
403    /// Enable film grain with seed.
404    pub fn enable(&mut self, seed: u16) {
405        self.enabled = true;
406        self.grain_seed = seed;
407    }
408
409    /// Disable film grain.
410    pub fn disable(&mut self) {
411        self.enabled = false;
412    }
413}
414
415// =============================================================================
416// Combined Optimizer
417// =============================================================================
418
419/// Combined loop optimization manager.
420#[derive(Clone, Debug)]
421pub struct LoopOptimizer {
422    /// Loop filter optimizer.
423    loop_filter: LoopFilterOptimizer,
424    /// CDEF optimizer.
425    cdef: CdefOptimizer,
426    /// Restoration optimizer.
427    restoration: RestorationOptimizer,
428    /// Film grain parameters.
429    film_grain: FilmGrainParams,
430}
431
432impl LoopOptimizer {
433    /// Create a new combined optimizer.
434    #[must_use]
435    pub fn new(lambda: f32) -> Self {
436        Self {
437            loop_filter: LoopFilterOptimizer::new(lambda),
438            cdef: CdefOptimizer::new(lambda),
439            restoration: RestorationOptimizer::new(lambda),
440            film_grain: FilmGrainParams::new(),
441        }
442    }
443
444    /// Optimize all parameters for a frame.
445    pub fn optimize_frame(
446        &mut self,
447        src: &[u8],
448        recon: &[u8],
449        width: usize,
450        height: usize,
451        qp: u8,
452    ) {
453        // Optimize loop filter
454        self.loop_filter
455            .optimize_filter_level(src, recon, width, height, qp);
456
457        // Optimize CDEF (on smaller regions)
458        let cdef_width = width.min(64);
459        let cdef_height = height.min(64);
460        self.cdef
461            .optimize_strength(src, recon, cdef_width, cdef_height, BlockSize::Block64x64);
462
463        // Optimize restoration
464        self.restoration
465            .optimize_restoration(src, recon, width, height);
466    }
467
468    /// Get loop filter parameters.
469    #[must_use]
470    pub const fn loop_filter_params(&self) -> &LoopFilterParams {
471        self.loop_filter.params()
472    }
473
474    /// Get CDEF parameters.
475    #[must_use]
476    pub const fn cdef_params(&self) -> &CdefParams {
477        self.cdef.params()
478    }
479
480    /// Get restoration type.
481    #[must_use]
482    pub const fn restoration_type(&self) -> RestorationType {
483        self.restoration.restoration_type()
484    }
485
486    /// Get film grain parameters.
487    #[must_use]
488    pub const fn film_grain_params(&self) -> &FilmGrainParams {
489        &self.film_grain
490    }
491
492    /// Set lambda for all optimizers.
493    pub fn set_lambda(&mut self, lambda: f32) {
494        self.loop_filter.set_lambda(lambda);
495        self.cdef.lambda = lambda;
496        self.restoration.lambda = lambda;
497    }
498}
499
500impl Default for LoopOptimizer {
501    fn default() -> Self {
502        Self::new(1.0)
503    }
504}
505
506// =============================================================================
507// Tests
508// =============================================================================
509
510#[cfg(test)]
511mod tests {
512    use super::*;
513
514    #[test]
515    fn test_loop_filter_optimizer_creation() {
516        let opt = LoopFilterOptimizer::new(1.0);
517        assert_eq!(opt.lambda, 1.0);
518        assert!(opt.rd_optimization);
519    }
520
521    #[test]
522    fn test_filter_level_from_qp() {
523        let opt = LoopFilterOptimizer::new(1.0);
524
525        let level_low = opt.filter_level_from_qp(10);
526        let level_high = opt.filter_level_from_qp(50);
527
528        assert!(level_low < level_high);
529        assert!(level_low <= MAX_LOOP_FILTER_LEVEL);
530        assert!(level_high <= MAX_LOOP_FILTER_LEVEL);
531    }
532
533    #[test]
534    fn test_optimize_filter_level_fast() {
535        let mut opt = LoopFilterOptimizer::new(1.0);
536        opt.set_rd_optimization(false);
537
538        let src = vec![128u8; 64 * 64];
539        let recon = vec![128u8; 64 * 64];
540
541        let level = opt.optimize_filter_level(&src, &recon, 64, 64, 28);
542        assert!(level <= MAX_LOOP_FILTER_LEVEL);
543    }
544
545    #[test]
546    fn test_compute_distortion() {
547        let opt = LoopFilterOptimizer::new(1.0);
548
549        let src = vec![100u8; 64];
550        let recon = vec![100u8; 64];
551
552        let distortion = opt.compute_distortion(&src, &recon, 8, 8);
553        assert_eq!(distortion, 0.0);
554
555        let recon2 = vec![110u8; 64];
556        let distortion2 = opt.compute_distortion(&src, &recon2, 8, 8);
557        assert!(distortion2 > 0.0);
558    }
559
560    #[test]
561    fn test_cdef_optimizer() {
562        let opt = CdefOptimizer::new(1.0);
563        assert_eq!(opt.lambda, 1.0);
564    }
565
566    #[test]
567    fn test_cdef_optimize_strength() {
568        let mut opt = CdefOptimizer::new(1.0);
569
570        let src = vec![128u8; 32 * 32];
571        let recon = vec![130u8; 32 * 32];
572
573        let strength = opt.optimize_strength(&src, &recon, 32, 32, BlockSize::Block32x32);
574
575        assert!(strength.primary <= MAX_CDEF_STRENGTH);
576        assert!(strength.secondary <= MAX_CDEF_STRENGTH);
577    }
578
579    #[test]
580    fn test_restoration_optimizer() {
581        let opt = RestorationOptimizer::new(1.0);
582        assert_eq!(opt.restoration_type, RestorationType::None);
583    }
584
585    #[test]
586    fn test_restoration_optimize() {
587        let mut opt = RestorationOptimizer::new(1.0);
588
589        let src = vec![128u8; 64 * 64];
590        let recon = vec![130u8; 64 * 64];
591
592        let rtype = opt.optimize_restoration(&src, &recon, 64, 64);
593        assert!(matches!(
594            rtype,
595            RestorationType::None | RestorationType::Wiener | RestorationType::Sgrproj
596        ));
597    }
598
599    #[test]
600    fn test_film_grain_params() {
601        let mut params = FilmGrainParams::new();
602        assert!(!params.enabled);
603
604        params.enable(1234);
605        assert!(params.enabled);
606        assert_eq!(params.grain_seed, 1234);
607
608        params.disable();
609        assert!(!params.enabled);
610    }
611
612    #[test]
613    fn test_combined_optimizer() {
614        let opt = LoopOptimizer::new(1.5);
615        assert_eq!(opt.loop_filter.lambda, 1.5);
616        assert_eq!(opt.cdef.lambda, 1.5);
617    }
618
619    #[test]
620    fn test_combined_optimize_frame() {
621        let mut opt = LoopOptimizer::new(1.0);
622
623        let src = vec![128u8; 128 * 128];
624        let recon = vec![128u8; 128 * 128];
625
626        opt.optimize_frame(&src, &recon, 128, 128, 28);
627
628        // Check that parameters were set
629        let lf_params = opt.loop_filter_params();
630        assert!(lf_params.level[0] <= MAX_LOOP_FILTER_LEVEL);
631    }
632
633    #[test]
634    fn test_set_lambda() {
635        let mut opt = LoopOptimizer::new(1.0);
636        opt.set_lambda(2.5);
637
638        assert_eq!(opt.loop_filter.lambda, 2.5);
639        assert_eq!(opt.cdef.lambda, 2.5);
640        assert_eq!(opt.restoration.lambda, 2.5);
641    }
642
643    #[test]
644    fn test_restoration_types() {
645        assert_eq!(RestorationType::None as u8, 0);
646        assert_eq!(RestorationType::Wiener as u8, 1);
647        assert_eq!(RestorationType::Sgrproj as u8, 2);
648    }
649
650    #[test]
651    fn test_constants() {
652        assert_eq!(MAX_LOOP_FILTER_LEVEL, 63);
653        assert_eq!(MAX_CDEF_STRENGTH, 15);
654        assert!(FILTER_LEVELS_TO_TEST > 0);
655        assert!(CDEF_STRENGTHS_TO_TEST > 0);
656    }
657}
oximedia_codec/av1/loop_optimization.rs

oximedia_codec/av1/
loop_optimization.rs
