| 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 | fn get_mv_diff(mv1: MV, mv2: MV) -> i8 { |
| 122 | (mv1.x - mv2.x).abs() + (mv1.y - mv2.y).abs() |
| 123 | } |
| 124 | |
| 125 | pub fn compact_mvs(mvs: &mut Vec<(MV, u16)>) { |
| 126 | mvs.sort_by(|a, b| a.1.cmp(&b.1)); |
| 127 | while mvs.len() > 256 { |
| 128 | let (mv, _) = mvs.pop().unwrap(); |
| 129 | if let Some(idx) = find_mv(mv, mvs) { |
| 130 | mvs[usize::from(idx)].1 += 1; |
| 131 | } |
| 132 | } |
| 133 | } |
| 134 | |
| 135 | pub fn find_mv(mv: MV, mvs: &[(MV, u16)]) -> Option<u8> { |
| 136 | let mut best_idx = None; |
| 137 | let mut best_cand_diff = 42; |
| 138 | let mut best_cand_count = 0; |
| 139 | for (i, &(cand_mv, count)) in mvs.iter().enumerate() { |
| 140 | if cand_mv == mv { |
| 141 | return Some(i as u8); |
| 142 | } |
| 143 | let diff = get_mv_diff(mv, cand_mv); |
| 144 | if diff <= 2 { |
| 145 | if diff < best_cand_diff || (diff == best_cand_diff && count > best_cand_count) { |
| 146 | best_idx = Some(i as u8); |
| 147 | best_cand_diff = diff; |
| 148 | best_cand_count = count; |
| 149 | } |
| 150 | } |
| 151 | } |
| 152 | best_idx |
| 153 | } |