Skip to main content

av_scenechange/data/
prediction.rs

1use v_frame::{
2    pixel::Pixel,
3    plane::{Plane, PlaneConfig, PlaneOffset, PlaneSlice},
4};
5
6use crate::{
7    cpu::CpuFeatureLevel,
8    data::{
9        frame::{FrameInvariants, RefType},
10        mc::put_8tap,
11        motion::MotionVector,
12        plane::PlaneRegionMut,
13        tile::TileRect,
14    },
15};
16
17// There are more modes than in the spec because every allowed
18// drl index for NEAR modes is considered its own mode.
19#[allow(non_camel_case_types)]
20#[allow(clippy::upper_case_acronyms)]
21#[allow(dead_code)]
22#[derive(Copy, Clone, Debug, PartialEq, PartialOrd, Eq, Ord, Default)]
23pub enum PredictionMode {
24    #[default]
25    DC_PRED, // Average of above and left pixels
26    V_PRED,      // Vertical
27    H_PRED,      // Horizontal
28    D45_PRED,    // Directional 45  degree
29    D135_PRED,   // Directional 135 degree
30    D113_PRED,   // Directional 113 degree
31    D157_PRED,   // Directional 157 degree
32    D203_PRED,   // Directional 203 degree
33    D67_PRED,    // Directional 67  degree
34    SMOOTH_PRED, // Combination of horizontal and vertical interpolation
35    SMOOTH_V_PRED,
36    SMOOTH_H_PRED,
37    PAETH_PRED,
38    UV_CFL_PRED,
39    NEARESTMV,
40    NEAR0MV,
41    NEAR1MV,
42    NEAR2MV,
43    GLOBALMV,
44    NEWMV,
45    // Compound ref compound modes
46    NEAREST_NEARESTMV,
47    NEAR_NEAR0MV,
48    NEAR_NEAR1MV,
49    NEAR_NEAR2MV,
50    NEAREST_NEWMV,
51    NEW_NEARESTMV,
52    NEAR_NEW0MV,
53    NEAR_NEW1MV,
54    NEAR_NEW2MV,
55    NEW_NEAR0MV,
56    NEW_NEAR1MV,
57    NEW_NEAR2MV,
58    GLOBAL_GLOBALMV,
59    NEW_NEWMV,
60}
61
62impl PredictionMode {
63    pub fn is_intra(self) -> bool {
64        self < PredictionMode::NEARESTMV
65    }
66
67    /// Inter prediction with a single reference (i.e. not compound mode)
68    ///
69    /// # Panics
70    ///
71    /// - If called on an intra `PredictionMode`
72    #[allow(clippy::too_many_arguments)]
73    pub fn predict_inter_single<T: Pixel>(
74        self,
75        fi: &FrameInvariants<T>,
76        tile_rect: TileRect,
77        p: usize,
78        po: PlaneOffset,
79        dst: &mut PlaneRegionMut<'_, T>,
80        width: usize,
81        height: usize,
82        ref_frame: RefType,
83        mv: MotionVector,
84        bit_depth: usize,
85        cpu_feature_level: CpuFeatureLevel,
86    ) {
87        assert!(!self.is_intra());
88        let frame_po = tile_rect.to_frame_plane_offset(po);
89
90        if let Some(ref rec) = fi.rec_buffer.frames[fi.ref_frames[ref_frame.to_index()] as usize] {
91            let (row_frac, col_frac, src) =
92                PredictionMode::get_mv_params(&rec.frame.planes[p], frame_po, mv);
93            put_8tap(
94                dst,
95                src,
96                width,
97                height,
98                col_frac,
99                row_frac,
100                bit_depth,
101                cpu_feature_level,
102            );
103        }
104    }
105
106    // Used by inter prediction to extract the fractional component of a mv and
107    // obtain the correct PlaneSlice to operate on.
108    fn get_mv_params<T: Pixel>(
109        rec_plane: &Plane<T>,
110        po: PlaneOffset,
111        mv: MotionVector,
112    ) -> (i32, i32, PlaneSlice<T>) {
113        let &PlaneConfig { xdec, ydec, .. } = &rec_plane.cfg;
114        let row_offset = mv.row as i32 >> (3 + ydec);
115        let col_offset = mv.col as i32 >> (3 + xdec);
116        let row_frac = ((mv.row as i32) << (1 - ydec)) & 0xf;
117        let col_frac = ((mv.col as i32) << (1 - xdec)) & 0xf;
118        let qo = PlaneOffset {
119            x: po.x + col_offset as isize - 3,
120            y: po.y + row_offset as isize - 3,
121        };
122        (
123            row_frac,
124            col_frac,
125            rec_plane.slice(qo).clamp().subslice(3, 3),
126        )
127    }
128}
129
130#[derive(Copy, Clone, Debug)]
131#[allow(clippy::upper_case_acronyms)]
132pub enum PredictionVariant {
133    NONE,
134    LEFT,
135    TOP,
136    BOTH,
137}
138
139impl PredictionVariant {
140    pub const fn new(x: usize, y: usize) -> Self {
141        match (x, y) {
142            (0, 0) => PredictionVariant::NONE,
143            (_, 0) => PredictionVariant::LEFT,
144            (0, _) => PredictionVariant::TOP,
145            _ => PredictionVariant::BOTH,
146        }
147    }
148}