summaryrefslogtreecommitdiffstats
path: root/nihav-indeo/src/codecs/indeo3enc/mv.rs
blob: d234e94f4aa41790c13dde2e0c45ddd38b95587d (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
use super::{Indeo3Cell, Plane};

const MV_THRESHOLD: u16 = 64;
const FLAT_THRESHOLD: u16 = 8;

const DIA_LARGE: [(i8, i8); 5] = [(0, 0), (2, 0), (0, 2), (-2, 0), (0, -2)];
const DIA_SMALL: [(i8, i8); 5] = [(0, 0), (1, 0), (0, 1), (-1, 0), (0, -1)];
const SEARCH_RANGE: i8 = 16;

#[derive(Clone, Copy, PartialEq)]
pub struct MV {
    pub x: i8,
    pub y: i8
}

pub struct MotionEstimator {
    pub mv_range:   i8,
    pub flat_thr:   u16,
    pub mv_thr:     u16,
}

impl MotionEstimator {
    pub fn new() -> Self {
        Self {
            mv_range:       SEARCH_RANGE,
            flat_thr:       FLAT_THRESHOLD,
            mv_thr:         MV_THRESHOLD,
        }
    }
    pub fn mv_search(&self, cur: &Plane, prev: &Plane, cell: Indeo3Cell) -> Option<(MV, bool)> {
        let plane_w = prev.width  as isize;
        let plane_h = prev.height as isize;
        let cell_w = cell.get_width()  as isize;
        let cell_h = cell.get_height() as isize;
        let start_x = cell.get_x() as isize;
        let start_y = cell.get_y() as isize;

        let check_mv = |mv: MV| {
                if mv.x.abs() < SEARCH_RANGE && mv.y.abs() < SEARCH_RANGE {
                    let new_x = start_x + isize::from(mv.x);
                    let new_y = start_y + isize::from(mv.y);
                    new_x >= 0 && new_x + cell_w <= plane_w && new_y >= 0 && new_y + cell_h <= plane_h
                } else {
                    false
                }
            };

        let area = (cell.get_width() * cell.get_height()) as u32;
        let flat_thr = u32::from(self.flat_thr) * area;

        let mut best_mv = MV{ x: 0, y: 0 };
        let mut best_score = calc_mv_score(cur, prev, cell, best_mv);

        if best_score < flat_thr {
            return Some((best_mv, true));
        }

        let mut found_better = true;
        while found_better {
            found_better = false;
            for step in DIA_LARGE.iter() {
                let new_mv = MV{ x: best_mv.x + step.0, y: best_mv.y + step.1 };
                if !check_mv(new_mv) {
                    continue;
                }
                let score = calc_mv_score(cur, prev, cell, new_mv);
                if score < best_score {
                    best_mv = new_mv;
                    best_score = score;
                    found_better = true;
                    if best_score < flat_thr {
                        return Some((best_mv, true));
                    }
                }
            }
        }
        for step in DIA_SMALL.iter() {
            let new_mv = MV{ x: best_mv.x + step.0, y: best_mv.y + step.1 };
            if !check_mv(new_mv) {
                continue;
            }
            let score = calc_mv_score(cur, prev, cell, new_mv);
            if score < best_score {
                best_mv = new_mv;
                best_score = score;
                if best_score < flat_thr {
                    break;
                }
            }
        }
        let score = (best_score / area) as u16;
        if score < self.mv_thr {
            Some((best_mv, false))
        } else {
            None
        }
    }
}

fn calc_cell_diff(src1: &[u8], src2: &[u8], stride: usize, width: usize, height: usize) -> u32 {
    let mut score = 0;
    for (line1, line2) in src1.chunks(stride).zip(src2.chunks(stride)).take(height) {
        for (&a, &b) in line1.iter().zip(line2.iter()).take(width) {
            let diff = if a >= b { u32::from(a - b) } else { u32::from(b - a) };
            score += diff * diff;
        }
    }
    score
}

fn calc_mv_score(cur: &Plane, prev: &Plane, cell: Indeo3Cell, mv: MV) -> u32 {
    let xoff = (cell.get_x() as isize + isize::from(mv.x)) as usize;
    let yoff = (cell.get_y() as isize + isize::from(mv.y)) as usize;

    let cur_ptr = &cur.data[cell.get_x() + (cell.get_y() + 1) * cur.width..];
    let ref_ptr = &prev.data[xoff + (yoff + 1) * prev.width..];

    calc_cell_diff(cur_ptr, ref_ptr, cur.width, cell.get_width(), cell.get_height())
}

pub fn compact_mvs(mvs: &mut Vec<(MV, u16)>) {
    mvs.sort_by(|a, b| b.1.cmp(&a.1));
    mvs.truncate(256);
}

pub fn find_mv(mv: MV, mvs: &[(MV, u16)]) -> Option<u8> {
    for (i, &(cand_mv, _)) in mvs.iter().enumerate() {
        if cand_mv == mv {
            return Some(i as u8);
        }
    }
    None
}