pito_codecs/vp9/enc/pred/
mvcost.rs

1use super::Mv;
2
3pub use crate::vp9::tables::{MV_CLASS_TREE, MV_FR_TREE, MV_JOINT_TREE};
4
5const TABLE: i32 = 1024;
6
7pub use super::super::bool::cost as bit_cost;
8
9pub fn tree_cost(tree: &[i8], probs: &[u8], value: u8) -> u32 {
10    fn walk(tree: &[i8], probs: &[u8], node: usize, value: u8, acc: u32) -> Option<u32> {
11        for b in 0..2usize {
12            let t = tree[node + b];
13            let c = acc + bit_cost(b == 1, probs[node >> 1]);
14            if t <= 0 {
15                if (-t) as u8 == value {
16                    return Some(c);
17                }
18            } else if let Some(r) = walk(tree, probs, t as usize, value, c) {
19                return Some(r);
20            }
21        }
22        None
23    }
24    walk(tree, probs, 0, value, 0).unwrap_or(0)
25}
26
27#[derive(Clone, Copy, Debug)]
28pub struct MvProbs {
29    pub joints: [u8; 3],
30    pub sign: [u8; 2],
31    pub classes: [[u8; 10]; 2],
32    pub class0: [u8; 2],
33    pub bits: [[u8; 10]; 2],
34    pub class0_fr: [[[u8; 3]; 2]; 2],
35    pub fr: [[u8; 3]; 2],
36    pub class0_hp: [u8; 2],
37    pub hp: [u8; 2],
38}
39
40impl Default for MvProbs {
41    fn default() -> Self {
42        MvProbs {
43            joints: [32, 64, 96],
44            sign: [128, 128],
45            classes: [
46                [224, 144, 192, 168, 192, 176, 192, 198, 198, 245],
47                [216, 128, 176, 160, 176, 176, 192, 198, 198, 208],
48            ],
49            class0: [216, 208],
50            bits: [
51                [136, 140, 148, 160, 176, 192, 224, 234, 234, 240],
52                [136, 140, 148, 160, 176, 192, 224, 234, 234, 240],
53            ],
54            class0_fr: [
55                [[128, 128, 64], [96, 112, 64]],
56                [[128, 128, 64], [96, 112, 64]],
57            ],
58            fr: [[64, 96, 64], [64, 96, 64]],
59            class0_hp: [160, 160],
60            hp: [128, 128],
61        }
62    }
63}
64
65pub struct MvCost {
66    joint: [u32; 4],
67    comp: [[Vec<u32>; 2]; 2],
68    probs: MvProbs,
69}
70
71fn mv_class(z: u32) -> u32 {
72    if z < 16 {
73        0
74    } else {
75        (31 - (z >> 3).leading_zeros()).min(10)
76    }
77}
78
79fn component_cost(p: &MvProbs, comp: usize, v: i32, usehp: bool) -> u32 {
80    if v == 0 {
81        return 0;
82    }
83    let mut c = bit_cost(v < 0, p.sign[comp]);
84    let z = v.unsigned_abs() - 1;
85    let class = mv_class(z);
86    c += tree_cost(&MV_CLASS_TREE, &p.classes[comp], class as u8);
87    if class == 0 {
88        let d = (z >> 3) as usize;
89        c += bit_cost(d == 1, p.class0[comp]);
90        c += tree_cost(&MV_FR_TREE, &p.class0_fr[comp][d], ((z >> 1) & 3) as u8);
91        if usehp {
92            c += bit_cost(z & 1 == 1, p.class0_hp[comp]);
93        }
94    } else {
95        let offset = z - (2 << (class + 2));
96        let d = offset >> 3;
97        for i in 0..class as usize {
98            c += bit_cost((d >> i) & 1 == 1, p.bits[comp][i]);
99        }
100        c += tree_cost(&MV_FR_TREE, &p.fr[comp], ((offset >> 1) & 3) as u8);
101        if usehp {
102            c += bit_cost(offset & 1 == 1, p.hp[comp]);
103        }
104    }
105    c
106}
107
108impl MvCost {
109    pub fn new(probs: MvProbs) -> Self {
110        let joint = std::array::from_fn(|j| tree_cost(&MV_JOINT_TREE, &probs.joints, j as u8));
111        let comp = std::array::from_fn(|comp| {
112            std::array::from_fn(|hp| {
113                (-TABLE..=TABLE)
114                    .map(|v| component_cost(&probs, comp, v, hp == 1))
115                    .collect()
116            })
117        });
118        MvCost { joint, comp, probs }
119    }
120
121    #[inline(always)]
122    fn comp(&self, comp: usize, v: i32, usehp: bool) -> u32 {
123        if v.abs() <= TABLE {
124            self.comp[comp][usehp as usize][(v + TABLE) as usize]
125        } else {
126            component_cost(&self.probs, comp, v, usehp)
127        }
128    }
129
130    pub fn rate(&self, mv: Mv, reference: Mv, usehp: bool) -> u32 {
131        let (dr, dc) = (
132            mv.row as i32 - reference.row as i32,
133            mv.col as i32 - reference.col as i32,
134        );
135        let j = ((dr != 0) as usize) << 1 | (dc != 0) as usize;
136        self.joint[j] + self.comp(0, dr, usehp) + self.comp(1, dc, usehp)
137    }
138}
pito_codecs/vp9/enc/pred/mvcost.rs

pito_codecs/vp9/enc/pred/
mvcost.rs
