av_scenechange/data/
motion.rs

1use std::{
2    marker::PhantomData,
3    ops,
4    ops::{Index, IndexMut},
5    slice,
6    sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard},
7};
8
9use arrayvec::ArrayVec;
10use v_frame::{frame::Frame, pixel::Pixel, plane::Plane};
11
12const MV_IN_USE_BITS: usize = 14;
13pub const MV_UPP: i32 = 1 << MV_IN_USE_BITS;
14pub const MV_LOW: i32 = -(1 << MV_IN_USE_BITS);
15
16#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
17pub struct MotionVector {
18    pub row: i16,
19    pub col: i16,
20}
21
22impl MotionVector {
23    pub const fn quantize_to_fullpel(self) -> Self {
24        Self {
25            row: (self.row / 8) * 8,
26            col: (self.col / 8) * 8,
27        }
28    }
29}
30
31impl ops::Mul<i16> for MotionVector {
32    type Output = MotionVector;
33
34    fn mul(self, rhs: i16) -> MotionVector {
35        MotionVector {
36            row: self.row * rhs,
37            col: self.col * rhs,
38        }
39    }
40}
41
42impl ops::Mul<u16> for MotionVector {
43    type Output = MotionVector;
44
45    fn mul(self, rhs: u16) -> MotionVector {
46        MotionVector {
47            row: self.row * rhs as i16,
48            col: self.col * rhs as i16,
49        }
50    }
51}
52
53impl ops::Shr<u8> for MotionVector {
54    type Output = MotionVector;
55
56    fn shr(self, rhs: u8) -> MotionVector {
57        MotionVector {
58            row: self.row >> rhs,
59            col: self.col >> rhs,
60        }
61    }
62}
63
64impl ops::Shl<u8> for MotionVector {
65    type Output = MotionVector;
66
67    fn shl(self, rhs: u8) -> MotionVector {
68        MotionVector {
69            row: self.row << rhs,
70            col: self.col << rhs,
71        }
72    }
73}
74
75impl ops::Add<MotionVector> for MotionVector {
76    type Output = MotionVector;
77
78    fn add(self, rhs: MotionVector) -> MotionVector {
79        MotionVector {
80            row: self.row + rhs.row,
81            col: self.col + rhs.col,
82        }
83    }
84}
85
86#[derive(Debug, Copy, Clone, Default)]
87pub struct MEStats {
88    pub mv: MotionVector,
89    /// SAD value on the scale of a 128x128 block
90    pub normalized_sad: u32,
91}
92
93#[derive(Debug, Clone)]
94pub struct FrameMEStats {
95    stats: Box<[MEStats]>,
96    pub cols: usize,
97    pub rows: usize,
98}
99
100pub const REF_FRAMES_LOG2: usize = 3;
101pub const REF_FRAMES: usize = 1 << REF_FRAMES_LOG2;
102pub type RefMEStats = Arc<RwLock<[FrameMEStats; REF_FRAMES]>>;
103pub type ReadGuardMEStats<'a> = RwLockReadGuard<'a, [FrameMEStats; REF_FRAMES]>;
104pub type WriteGuardMEStats<'a> = RwLockWriteGuard<'a, [FrameMEStats; REF_FRAMES]>;
105
106impl FrameMEStats {
107    pub fn new(cols: usize, rows: usize) -> Self {
108        Self {
109            // dynamic allocation: once per frame
110            stats: vec![MEStats::default(); cols * rows].into_boxed_slice(),
111            cols,
112            rows,
113        }
114    }
115
116    pub fn new_arc_array(cols: usize, rows: usize) -> RefMEStats {
117        Arc::new(RwLock::new([
118            FrameMEStats::new(cols, rows),
119            FrameMEStats::new(cols, rows),
120            FrameMEStats::new(cols, rows),
121            FrameMEStats::new(cols, rows),
122            FrameMEStats::new(cols, rows),
123            FrameMEStats::new(cols, rows),
124            FrameMEStats::new(cols, rows),
125            FrameMEStats::new(cols, rows),
126        ]))
127    }
128}
129
130impl Index<usize> for FrameMEStats {
131    type Output = [MEStats];
132
133    fn index(&self, index: usize) -> &Self::Output {
134        &self.stats[index * self.cols..(index + 1) * self.cols]
135    }
136}
137
138impl IndexMut<usize> for FrameMEStats {
139    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
140        &mut self.stats[index * self.cols..(index + 1) * self.cols]
141    }
142}
143
144/// Tiled view of `FrameMEStats`
145#[derive(Debug)]
146pub struct TileMEStats<'a> {
147    data: *const MEStats,
148    // expressed in mi blocks
149    // private to guarantee borrowing rules
150    x: usize,
151    y: usize,
152    cols: usize,
153    rows: usize,
154    /// number of cols in the underlying `FrameMEStats`
155    stride: usize,
156    phantom: PhantomData<&'a MotionVector>,
157}
158
159/// Mutable tiled view of `FrameMEStats`
160#[derive(Debug)]
161pub struct TileMEStatsMut<'a> {
162    data: *mut MEStats,
163    // expressed in mi blocks
164    // private to guarantee borrowing rules
165    x: usize,
166    y: usize,
167    cols: usize,
168    rows: usize,
169    /// number of cols in the underlying `FrameMEStats`
170    stride: usize,
171    phantom: PhantomData<&'a mut MotionVector>,
172}
173
174// common impl for TileMotionVectors and TileMotionVectorsMut
175macro_rules! tile_me_stats_common {
176  // $name: TileMEStats or TileMEStatsMut
177  // $opt_mut: nothing or mut
178  ($name:ident $(,$opt_mut:tt)?) => {
179    impl<'a> $name<'a> {
180
181      /// # Panics
182      ///
183      /// - If the requested dimensions are larger than the frame MV size
184      #[allow(dead_code)]
185      pub fn new(
186        frame_mvs: &'a $($opt_mut)? FrameMEStats,
187        x: usize,
188        y: usize,
189        cols: usize,
190        rows: usize,
191      ) -> Self {
192        assert!(x + cols <= frame_mvs.cols);
193        assert!(y + rows <= frame_mvs.rows);
194        Self {
195          data: & $($opt_mut)? frame_mvs[y][x],
196          x,
197          y,
198          cols,
199          rows,
200          stride: frame_mvs.cols,
201          phantom: PhantomData,
202        }
203      }
204
205      #[allow(dead_code)]
206      pub const fn x(&self) -> usize {
207        self.x
208      }
209
210      #[allow(dead_code)]
211      pub const fn y(&self) -> usize {
212        self.y
213      }
214
215      #[allow(dead_code)]
216      pub const fn cols(&self) -> usize {
217        self.cols
218      }
219
220      #[allow(dead_code)]
221      pub const fn rows(&self) -> usize {
222        self.rows
223      }
224    }
225
226    unsafe impl Send for $name<'_> {}
227    unsafe impl Sync for $name<'_> {}
228
229    impl Index<usize> for $name<'_> {
230      type Output = [MEStats];
231
232      fn index(&self, index: usize) -> &Self::Output {
233        assert!(index < self.rows);
234        // SAFETY: The above assert ensures we do not access OOB data.
235        unsafe {
236          let ptr = self.data.add(index * self.stride);
237          slice::from_raw_parts(ptr, self.cols)
238        }
239      }
240    }
241  }
242}
243
244tile_me_stats_common!(TileMEStats);
245tile_me_stats_common!(TileMEStatsMut, mut);
246
247impl TileMEStatsMut<'_> {
248    pub const fn as_const(&self) -> TileMEStats<'_> {
249        TileMEStats {
250            data: self.data,
251            x: self.x,
252            y: self.y,
253            cols: self.cols,
254            rows: self.rows,
255            stride: self.stride,
256            phantom: PhantomData,
257        }
258    }
259}
260
261impl IndexMut<usize> for TileMEStatsMut<'_> {
262    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
263        assert!(index < self.rows);
264        // SAFETY: The above assert ensures we do not access OOB data.
265        unsafe {
266            let ptr = self.data.add(index * self.stride);
267            slice::from_raw_parts_mut(ptr, self.cols)
268        }
269    }
270}
271
272#[derive(Debug, Copy, Clone, Eq, PartialEq)]
273#[allow(clippy::upper_case_acronyms)]
274pub enum MVSamplingMode {
275    INIT,
276    CORNER { right: bool, bottom: bool },
277}
278
279#[derive(Debug, Clone)]
280pub struct ReferenceFrame<T: Pixel> {
281    pub frame: Arc<Frame<T>>,
282    pub input_hres: Arc<Plane<T>>,
283    pub input_qres: Arc<Plane<T>>,
284    pub frame_me_stats: RefMEStats,
285}
286
287#[derive(Debug, Clone, Default)]
288pub struct ReferenceFramesSet<T: Pixel> {
289    pub frames: [Option<Arc<ReferenceFrame<T>>>; REF_FRAMES],
290}
291
292impl<T: Pixel> ReferenceFramesSet<T> {
293    pub fn new() -> Self {
294        Self {
295            frames: Default::default(),
296            // deblock: Default::default()
297        }
298    }
299}
300
301pub struct MotionEstimationSubsets {
302    pub min_sad: u32,
303    pub median: Option<MotionVector>,
304    pub subset_b: ArrayVec<MotionVector, 5>,
305    pub subset_c: ArrayVec<MotionVector, 5>,
306}
307
308impl MotionEstimationSubsets {
309    pub fn all_mvs(&self) -> ArrayVec<MotionVector, 11> {
310        let mut all = ArrayVec::new();
311        if let Some(median) = self.median {
312            all.push(median);
313        }
314
315        all.extend(self.subset_b.iter().copied());
316        all.extend(self.subset_c.iter().copied());
317
318        all
319    }
320}
av_scenechange/data/motion.rs

av_scenechange/data/
motion.rs
