pito_codecs/vp9/enc/rc/
aq.rs

1use super::q::{q_index_for, qstep};
2
3pub const SEGMENTS: usize = 8;
4const SCALE: [f64; SEGMENTS] = [0.76, 0.82, 0.88, 0.94, 1.0, 1.06, 1.13, 1.2];
5const STRENGTH: f64 = 0.12;
6
7#[derive(Clone, Debug, Default)]
8pub struct Segmentation {
9    pub enabled: bool,
10    pub update_map: bool,
11    pub update_data: bool,
12    pub abs_delta: bool,
13    pub tree_probs: [u8; 7],
14    pub q: [i16; SEGMENTS],
15    pub lf: [i8; SEGMENTS],
16    pub map: Vec<u8>,
17    pub mi_cols: usize,
18    pub mi_rows: usize,
19    block_seg: Vec<u8>,
20}
21
22impl Segmentation {
23    pub fn segment(&self, mi_row: usize, mi_col: usize) -> u8 {
24        if !self.enabled {
25            return 0;
26        }
27        self.map[mi_row * self.mi_cols + mi_col]
28    }
29
30    pub fn q_index(&self, base: u8, seg: u8) -> u8 {
31        if !self.enabled {
32            return base;
33        }
34        let q = self.q[seg as usize] as i32;
35        if self.abs_delta {
36            q.clamp(0, 255) as u8
37        } else {
38            (base as i32 + q).clamp(0, 255) as u8
39        }
40    }
41
42    pub fn disable(&mut self) {
43        self.enabled = false;
44        self.update_map = false;
45        self.update_data = false;
46    }
47
48    pub fn variance_aq(
49        &mut self,
50        base_q: u8,
51        mi_cols: usize,
52        mi_rows: usize,
53        activity: &[u32],
54        blocks: (usize, usize),
55    ) {
56        let (bw, bh) = blocks;
57        if base_q < 8 || bw * bh == 0 {
58            self.disable();
59            return;
60        }
61        let n = (bw * bh) as f64;
62        let mean_log = activity
63            .iter()
64            .map(|&a| ((a + 16) as f64).ln())
65            .sum::<f64>()
66            / n;
67        let base_step = qstep(base_q);
68        let mut q = [0i16; SEGMENTS];
69        let mut prev = i32::MIN;
70        for (s, scale) in SCALE.iter().enumerate() {
71            let qi = (q_index_for(base_step * scale) as i32)
72                .clamp(1, 255)
73                .max(prev);
74            prev = qi;
75            q[s] = (qi - base_q as i32) as i16;
76        }
77        let edges: [f64; SEGMENTS - 1] =
78            std::array::from_fn(|i| ((SCALE[i] * SCALE[i + 1]).sqrt().ln()) / STRENGTH);
79        let mut block_seg = std::mem::take(&mut self.block_seg);
80        block_seg.resize(bw * bh, 0);
81        for (seg, &a) in block_seg.iter_mut().zip(activity) {
82            let d = ((a + 16) as f64).ln() - mean_log;
83            *seg = edges.partition_point(|&e| e < d) as u8;
84        }
85        self.map.resize(mi_cols * mi_rows, 0);
86        let mut counts = [0u32; SEGMENTS];
87        for r in 0..mi_rows {
88            let by = (r / 2).min(bh - 1);
89            let row = &mut self.map[r * mi_cols..][..mi_cols];
90            for (c, m) in row.iter_mut().enumerate() {
91                let s = block_seg[by * bw + (c / 2).min(bw - 1)];
92                *m = s;
93                counts[s as usize] += 1;
94            }
95        }
96        self.block_seg = block_seg;
97        self.enabled = true;
98        self.update_map = true;
99        self.update_data = true;
100        self.abs_delta = false;
101        self.q = q;
102        self.lf = [0; SEGMENTS];
103        self.mi_cols = mi_cols;
104        self.mi_rows = mi_rows;
105        self.tree_probs = tree_probs(&counts);
106    }
107}
108
109pub fn tree_probs(c: &[u32; SEGMENTS]) -> [u8; 7] {
110    let p = |l: u32, r: u32| {
111        let t = l + r;
112        if t == 0 {
113            255
114        } else {
115            ((l as u64 * 256 + t as u64 / 2) / t as u64).clamp(1, 255) as u8
116        }
117    };
118    [
119        p(c[0] + c[1] + c[2] + c[3], c[4] + c[5] + c[6] + c[7]),
120        p(c[0] + c[1], c[2] + c[3]),
121        p(c[4] + c[5], c[6] + c[7]),
122        p(c[0], c[1]),
123        p(c[2], c[3]),
124        p(c[4], c[5]),
125        p(c[6], c[7]),
126    ]
127}
pito_codecs/vp9/enc/rc/aq.rs

pito_codecs/vp9/enc/rc/
aq.rs
