]> git.nihav.org Git - nihav.git/blob - nihav-indeo/src/codecs/indeo3enc/mv.rs
d234e94f4aa41790c13dde2e0c45ddd38b95587d
[nihav.git] / nihav-indeo / src / codecs / indeo3enc / mv.rs
1 use super::{Indeo3Cell, Plane};
2
3 const MV_THRESHOLD: u16 = 64;
4 const FLAT_THRESHOLD: u16 = 8;
5
6 const DIA_LARGE: [(i8, i8); 5] = [(0, 0), (2, 0), (0, 2), (-2, 0), (0, -2)];
7 const DIA_SMALL: [(i8, i8); 5] = [(0, 0), (1, 0), (0, 1), (-1, 0), (0, -1)];
8 const SEARCH_RANGE: i8 = 16;
9
10 #[derive(Clone, Copy, PartialEq)]
11 pub struct MV {
12 pub x: i8,
13 pub y: i8
14 }
15
16 pub struct MotionEstimator {
17 pub mv_range: i8,
18 pub flat_thr: u16,
19 pub mv_thr: u16,
20 }
21
22 impl MotionEstimator {
23 pub fn new() -> Self {
24 Self {
25 mv_range: SEARCH_RANGE,
26 flat_thr: FLAT_THRESHOLD,
27 mv_thr: MV_THRESHOLD,
28 }
29 }
30 pub fn mv_search(&self, cur: &Plane, prev: &Plane, cell: Indeo3Cell) -> Option<(MV, bool)> {
31 let plane_w = prev.width as isize;
32 let plane_h = prev.height as isize;
33 let cell_w = cell.get_width() as isize;
34 let cell_h = cell.get_height() as isize;
35 let start_x = cell.get_x() as isize;
36 let start_y = cell.get_y() as isize;
37
38 let check_mv = |mv: MV| {
39 if mv.x.abs() < SEARCH_RANGE && mv.y.abs() < SEARCH_RANGE {
40 let new_x = start_x + isize::from(mv.x);
41 let new_y = start_y + isize::from(mv.y);
42 new_x >= 0 && new_x + cell_w <= plane_w && new_y >= 0 && new_y + cell_h <= plane_h
43 } else {
44 false
45 }
46 };
47
48 let area = (cell.get_width() * cell.get_height()) as u32;
49 let flat_thr = u32::from(self.flat_thr) * area;
50
51 let mut best_mv = MV{ x: 0, y: 0 };
52 let mut best_score = calc_mv_score(cur, prev, cell, best_mv);
53
54 if best_score < flat_thr {
55 return Some((best_mv, true));
56 }
57
58 let mut found_better = true;
59 while found_better {
60 found_better = false;
61 for step in DIA_LARGE.iter() {
62 let new_mv = MV{ x: best_mv.x + step.0, y: best_mv.y + step.1 };
63 if !check_mv(new_mv) {
64 continue;
65 }
66 let score = calc_mv_score(cur, prev, cell, new_mv);
67 if score < best_score {
68 best_mv = new_mv;
69 best_score = score;
70 found_better = true;
71 if best_score < flat_thr {
72 return Some((best_mv, true));
73 }
74 }
75 }
76 }
77 for step in DIA_SMALL.iter() {
78 let new_mv = MV{ x: best_mv.x + step.0, y: best_mv.y + step.1 };
79 if !check_mv(new_mv) {
80 continue;
81 }
82 let score = calc_mv_score(cur, prev, cell, new_mv);
83 if score < best_score {
84 best_mv = new_mv;
85 best_score = score;
86 if best_score < flat_thr {
87 break;
88 }
89 }
90 }
91 let score = (best_score / area) as u16;
92 if score < self.mv_thr {
93 Some((best_mv, false))
94 } else {
95 None
96 }
97 }
98 }
99
100 fn calc_cell_diff(src1: &[u8], src2: &[u8], stride: usize, width: usize, height: usize) -> u32 {
101 let mut score = 0;
102 for (line1, line2) in src1.chunks(stride).zip(src2.chunks(stride)).take(height) {
103 for (&a, &b) in line1.iter().zip(line2.iter()).take(width) {
104 let diff = if a >= b { u32::from(a - b) } else { u32::from(b - a) };
105 score += diff * diff;
106 }
107 }
108 score
109 }
110
111 fn calc_mv_score(cur: &Plane, prev: &Plane, cell: Indeo3Cell, mv: MV) -> u32 {
112 let xoff = (cell.get_x() as isize + isize::from(mv.x)) as usize;
113 let yoff = (cell.get_y() as isize + isize::from(mv.y)) as usize;
114
115 let cur_ptr = &cur.data[cell.get_x() + (cell.get_y() + 1) * cur.width..];
116 let ref_ptr = &prev.data[xoff + (yoff + 1) * prev.width..];
117
118 calc_cell_diff(cur_ptr, ref_ptr, cur.width, cell.get_width(), cell.get_height())
119 }
120
121 pub fn compact_mvs(mvs: &mut Vec<(MV, u16)>) {
122 mvs.sort_by(|a, b| b.1.cmp(&a.1));
123 mvs.truncate(256);
124 }
125
126 pub fn find_mv(mv: MV, mvs: &[(MV, u16)]) -> Option<u8> {
127 for (i, &(cand_mv, _)) in mvs.iter().enumerate() {
128 if cand_mv == mv {
129 return Some(i as u8);
130 }
131 }
132 None
133 }