]>
Commit | Line | Data |
---|---|---|
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 | } |