]> git.nihav.org Git - nihav.git/blob - nihav-indeo/src/codecs/indeo3enc/tree.rs
indeo3enc: rework the logic with motion cells
[nihav.git] / nihav-indeo / src / codecs / indeo3enc / tree.rs
1 use super::Indeo3Writer;
2 use super::mv::*;
3 use super::cell::{CellEncoder, MAX_CELL_SIZE};
4 use std::ops::DerefMut;
5
6 pub enum Indeo3PrimaryTree {
7 VSplit(Box<Indeo3PrimaryTree>, Box<Indeo3PrimaryTree>),
8 HSplit(Box<Indeo3PrimaryTree>, Box<Indeo3PrimaryTree>),
9 RelFill(MV, Box<Indeo3SecondaryTree>),
10 AbsFill(Box<Indeo3SecondaryTree>),
11 }
12
13 impl Indeo3PrimaryTree {
14 pub fn print(&self) {
15 println!("Plane tree:");
16 self.print1(1);
17 }
18 fn print1(&self, depth: u8) {
19 for _ in 0..depth {
20 print!(" ");
21 }
22 match self {
23 Indeo3PrimaryTree::VSplit(t1, t2) => {
24 println!("vertical split");
25 t1.print1(depth + 1);
26 t2.print1(depth + 1);
27 },
28 Indeo3PrimaryTree::HSplit(t1, t2) => {
29 println!("horizontal split");
30 t1.print1(depth + 1);
31 t2.print1(depth + 1);
32 },
33 Indeo3PrimaryTree::RelFill(mv, sec) => {
34 println!("relative fill {},{}", mv.x, mv.y);
35 sec.print1(depth + 1);
36 },
37 Indeo3PrimaryTree::AbsFill(sec) => {
38 println!("absolute fill");
39 sec.print1(depth + 1);
40 }
41 }
42 }
43 }
44
45 pub enum Indeo3SecondaryTree {
46 VSplit(Box<Indeo3SecondaryTree>, Box<Indeo3SecondaryTree>),
47 HSplit(Box<Indeo3SecondaryTree>, Box<Indeo3SecondaryTree>),
48 VQData(u8),
49 VQNull(u8),
50 }
51
52 impl Indeo3SecondaryTree {
53 fn print1(&self, depth: u8) {
54 for _ in 0..depth {
55 print!(" ");
56 }
57 match self {
58 Indeo3SecondaryTree::VSplit(t1, t2) => {
59 println!("vertical split");
60 t1.print1(depth + 1);
61 t2.print1(depth + 1);
62 },
63 Indeo3SecondaryTree::HSplit(t1, t2) => {
64 println!("horizontal split");
65 t1.print1(depth + 1);
66 t2.print1(depth + 1);
67 },
68 Indeo3SecondaryTree::VQData(mode) => {
69 println!("VQ data ({})", mode);
70 },
71 Indeo3SecondaryTree::VQNull(mode) => {
72 println!("VQ Null ({})", mode);
73 }
74 }
75 }
76 }
77
78 const THRESHOLD: u32 = 64;
79
80 #[derive(Clone, Copy)]
81 pub struct Indeo3Cell {
82 x: u8,
83 y: u8,
84 w: u8,
85 h: u8,
86 intra: bool,
87 }
88
89 impl Indeo3Cell {
90 pub fn new(width: usize, height: usize, intra: bool) -> Self {
91 Self {
92 x: 0,
93 y: 0,
94 w: (width / 4) as u8,
95 h: (height / 4) as u8,
96 intra,
97 }
98 }
99
100 pub fn get_x(&self) -> usize { usize::from(self.x) * 4 }
101 pub fn get_y(&self) -> usize { usize::from(self.y) * 4 }
102 pub fn get_width(&self) -> usize { usize::from(self.w) * 4 }
103 pub fn get_height(&self) -> usize { usize::from(self.h) * 4 }
104 pub fn is_intra(&self) -> bool { self.intra }
105
106 fn split_h(&self) -> (Self, Self) {
107 let h1 = if self.h > 2 { ((self.h + 2) >> 2) << 1 } else { 1 };
108 let h2 = self.h - h1;
109 let mut cell1 = *self;
110 cell1.h = h1;
111 let mut cell2 = *self;
112 cell2.y += h1;
113 cell2.h = h2;
114 (cell1, cell2)
115 }
116 fn split_v(&self, stripw: u8) -> (Self, Self) {
117 let w1 = if self.w > stripw {
118 if self.w > stripw * 2 { stripw * 2 } else { stripw }
119 } else {
120 if self.w > 2 { ((self.w + 2) >> 2) << 1 } else { 1 }
121 };
122 let w2 = self.w - w1;
123 let mut cell1 = *self;
124 cell1.w = w1;
125 let mut cell2 = *self;
126 cell2.x += w1;
127 cell2.w = w2;
128 (cell1, cell2)
129 }
130 }
131
132 impl std::fmt::Display for Indeo3Cell {
133 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
134 write!(f, "[{}x{} @ {},{}{}]", self.get_width(), self.get_height(), self.get_x(), self.get_y(), if self.intra { " intra" } else { "" })
135 }
136 }
137
138 #[derive(Default)]
139 pub struct Plane {
140 pub data: Vec<u8>,
141 pub width: usize,
142 pub height: usize,
143 pub stripw: u8,
144 pub mvs: Vec<(MV, u16)>,
145 }
146
147 fn ssd(a: u8, b: u8) -> u32 {
148 let diff = i32::from(a) - i32::from(b);
149 (diff * diff) as u32
150 }
151
152 impl Plane {
153 pub fn alloc(&mut self, width: usize, height: usize, stripw: u8) {
154 self.data.resize(width * height, 0);
155 self.width = width;
156 self.height = height;
157 self.stripw = stripw;
158 if self.mvs.capacity() < 256 {
159 self.mvs.reserve(256);
160 }
161 }
162 pub fn fill(&mut self, src: &[u8], stride: usize) {
163 for (dline, sline) in self.data.chunks_mut(self.width).zip(src.chunks(stride)) {
164 for (dst, &src) in dline.iter_mut().zip(sline.iter()) {
165 *dst = src >> 1;
166 }
167 }
168 }
169 pub fn clear_mvs(&mut self){
170 self.mvs.clear();
171 }
172 pub fn checksum(&self) -> u16 {
173 let xors = self.data.chunks(2).fold([0u8; 2], |acc, pair| [acc[0] ^ pair[0], acc[1] ^ pair[1]]);
174 u16::from(xors[0]) | (u16::from(xors[1]) * 256)
175 }
176 pub fn find_cells(&mut self, is_intra: bool, pplane: &Plane, mv_est: &MotionEstimator) -> Box<Indeo3PrimaryTree> {
177 let cell = Indeo3Cell::new(self.width, self.height, is_intra);
178 self.split_pri(cell, pplane, mv_est)
179 }
180 fn split_pri(&mut self, cell: Indeo3Cell, pplane: &Plane, mv_est: &MotionEstimator) -> Box<Indeo3PrimaryTree> {
181 let width = cell.get_width();
182 let height = cell.get_height();
183 if width * height > MAX_CELL_SIZE {
184 let (hsplit, vsplit) = if width != height {
185 (width > (self.stripw as usize) * 4 || width > height, height > width)
186 } else {
187 let (hdiff, vdiff) = self.calculate_diffs(cell);
188 (vdiff > THRESHOLD && vdiff > hdiff,
189 hdiff > THRESHOLD && hdiff > vdiff)
190 };
191 match (hsplit, vsplit) {
192 (true, _) => {
193 let (cell1, cell2) = cell.split_v(self.stripw);
194 let tree1 = self.split_pri(cell1, pplane, mv_est);
195 let tree2 = self.split_pri(cell2, pplane, mv_est);
196 Box::new(Indeo3PrimaryTree::VSplit(tree1, tree2))
197 },
198 (_, true) => {
199 let (cell1, cell2) = cell.split_h();
200 let tree1 = self.split_pri(cell1, pplane, mv_est);
201 let tree2 = self.split_pri(cell2, pplane, mv_est);
202 Box::new(Indeo3PrimaryTree::HSplit(tree1, tree2))
203 },
204 (false, false) => {
205 let sec = self.split_sec(cell);
206 Box::new(Indeo3PrimaryTree::AbsFill(sec))
207 },
208 }
209 } else {
210 if !cell.intra {
211 if let Some((mv, flat)) = mv_est.mv_search(self, pplane, cell) {
212 return self.add_mv_tree(mv, flat, cell);
213 }
214
215 // try splitting once to see if it improves situation
216 if width >= 16 && height >= 16 {
217 let vsplit = width > height;
218 let (cell1, cell2) = if vsplit {
219 cell.split_v(self.stripw)
220 } else {
221 cell.split_h()
222 };
223 let search1 = mv_est.mv_search(self, pplane, cell1);
224 let search2 = mv_est.mv_search(self, pplane, cell2);
225 if search1.is_some() || search2.is_some() {
226 let tree1 = if let Some((mv, flat)) = search1 {
227 self.add_mv_tree(mv, flat, cell1)
228 } else {
229 let sec = self.split_sec(cell1);
230 Box::new(Indeo3PrimaryTree::AbsFill(sec))
231 };
232 let tree2 = if let Some((mv, flat)) = search2 {
233 self.add_mv_tree(mv, flat, cell2)
234 } else {
235 let sec = self.split_sec(cell2);
236 Box::new(Indeo3PrimaryTree::AbsFill(sec))
237 };
238 return if vsplit {
239 Box::new(Indeo3PrimaryTree::VSplit(tree1, tree2))
240 } else {
241 Box::new(Indeo3PrimaryTree::HSplit(tree1, tree2))
242 }
243 }
244 }
245 }
246 let sec = self.split_sec(cell);
247 Box::new(Indeo3PrimaryTree::AbsFill(sec))
248 }
249 }
250 fn add_mv_tree(&mut self, mv: MV, flat: bool, cell: Indeo3Cell) -> Box<Indeo3PrimaryTree> {
251 let sec = if flat {
252 Box::new(Indeo3SecondaryTree::VQNull(0))
253 } else {
254 Box::new(Indeo3SecondaryTree::VQData(0))
255 };
256
257 let mut found = false;
258 for (ref cmv, ref mut count) in self.mvs.iter_mut() {
259 if cmv == &mv {
260 *count += u16::from(cell.w) * u16::from(cell.h);
261 found = true;
262 break;
263 }
264 }
265 if !found {
266 self.mvs.push((mv, 1));
267 }
268
269 Box::new(Indeo3PrimaryTree::RelFill(mv, sec))
270 }
271 fn split_sec(&mut self, cell: Indeo3Cell) -> Box<Indeo3SecondaryTree> {
272 let (hdiff, vdiff) = self.calculate_diffs(cell);
273 if hdiff == 0 && vdiff == 0 {
274 if !cell.intra {
275 return Box::new(Indeo3SecondaryTree::VQNull(0));
276 } else {
277 return Box::new(Indeo3SecondaryTree::VQData(0));
278 }
279 }
280 if cell.get_width() > 16 && cell.get_height() > 16 {
281 let hsplit = vdiff > THRESHOLD && vdiff > hdiff * 2;
282 let vsplit = hdiff > THRESHOLD && hdiff > vdiff * 2;
283 match (vsplit, hsplit) {
284 (true, _) => {
285 let (cell1, cell2) = cell.split_v(self.stripw);
286 let tree1 = self.split_sec(cell1);
287 let tree2 = self.split_sec(cell2);
288 Box::new(Indeo3SecondaryTree::VSplit(tree1, tree2))
289 },
290 (_, true) => {
291 let (cell1, cell2) = cell.split_h();
292 let tree1 = self.split_sec(cell1);
293 let tree2 = self.split_sec(cell2);
294 Box::new(Indeo3SecondaryTree::HSplit(tree1, tree2))
295 },
296 _ => {
297 Box::new(Indeo3SecondaryTree::VQData(0))
298 },
299 }
300 } else {
301 let is_w8 = (cell.get_width() & 7) == 0;
302 let is_h8 = (cell.get_height() & 7) == 0;
303 let mode = match (hdiff > THRESHOLD, vdiff > THRESHOLD) {
304 (false, false) if is_w8 && is_h8 => 10,
305 (_, true) if is_h8 => 3,
306 _ => 0,
307 };
308 Box::new(Indeo3SecondaryTree::VQData(mode))
309 }
310 }
311 fn calculate_diffs(&self, cell: Indeo3Cell) -> (u32, u32) {
312 let offset = cell.get_x() + cell.get_y() * self.width;
313 let mut w = cell.get_width();
314 if cell.get_x() + w == self.width { w -= 1; }
315 let mut h = cell.get_height();
316 if cell.get_y() + h == self.height { h -= 1; }
317
318 let mut vdiff = 0;
319 let mut hdiff = 0;
320 let src0 = &self.data[offset..];
321 let src1 = &self.data[offset + self.width..];
322 for (line0, line1) in src0.chunks(self.width).zip(src1.chunks(self.width)).take(h) {
323 for ((&cur, &right), &bottom) in line0.iter().zip(line0[1..].iter()).zip(line1.iter()).take(w) {
324 hdiff += ssd(cur, right);
325 vdiff += ssd(cur, bottom);
326 }
327 }
328 let area = (w * h) as u32;
329 (hdiff * 16 / area, vdiff * 16 / area)
330 }
331 pub fn prune_extra_mvs(&mut self, tree: &mut Box<Indeo3PrimaryTree>) {
332 let cell = Indeo3Cell::new(self.width, self.height, true);
333 self.prune_pri(cell, tree)
334 }
335 fn prune_pri(&mut self, cell: Indeo3Cell, tree: &mut Box<Indeo3PrimaryTree>) {
336 match tree.deref_mut() {
337 Indeo3PrimaryTree::HSplit(ref mut tree1, ref mut tree2) => {
338 let (cell1, cell2) = cell.split_h();
339 self.prune_pri(cell1, tree1);
340 self.prune_pri(cell2, tree2);
341 },
342 Indeo3PrimaryTree::VSplit(ref mut tree1, ref mut tree2) => {
343 let (cell1, cell2) = cell.split_v(self.stripw);
344 self.prune_pri(cell1, tree1);
345 self.prune_pri(cell2, tree2);
346 },
347 Indeo3PrimaryTree::AbsFill(_) => {},
348 Indeo3PrimaryTree::RelFill(ref mv, ref _sec) => {
349 if find_mv(*mv, &self.mvs).is_none() {
350 let sec = self.split_sec(cell);
351 *tree = Box::new(Indeo3PrimaryTree::AbsFill(sec));
352 }
353 },
354 }
355 }
356 pub fn encode_tree(&mut self, iw: &mut Indeo3Writer, tree: &Indeo3PrimaryTree, cenc: &mut CellEncoder, is_intra: bool, refp: &Plane) {
357 let cell = Indeo3Cell::new(self.width, self.height, is_intra);
358 self.encode_pri(iw, cell, tree, cenc, refp);
359 }
360 fn encode_pri(&mut self, iw: &mut Indeo3Writer, mut cell: Indeo3Cell, tree: &Indeo3PrimaryTree, cenc: &mut CellEncoder, refp: &Plane) {
361 match tree {
362 Indeo3PrimaryTree::HSplit(t1, t2) => {
363 iw.put_2bits(0);
364 let (cell1, cell2) = cell.split_h();
365 self.encode_pri(iw, cell1, t1, cenc, refp);
366 self.encode_pri(iw, cell2, t2, cenc, refp);
367 },
368 Indeo3PrimaryTree::VSplit(t1, t2) => {
369 iw.put_2bits(1);
370 let (cell1, cell2) = cell.split_v(self.stripw);
371 self.encode_pri(iw, cell1, t1, cenc, refp);
372 self.encode_pri(iw, cell2, t2, cenc, refp);
373 },
374 Indeo3PrimaryTree::AbsFill(sec) => {
375 iw.put_2bits(2);
376 cell.intra = true;
377 self.encode_sec(iw, cell, sec, cenc);
378 }
379 Indeo3PrimaryTree::RelFill(mv, sec) => {
380 if let Some(mv_idx) = find_mv(*mv, &self.mvs) {
381 iw.put_2bits(3);
382 iw.put_byte(mv_idx);
383 cell.intra = false;
384 let real_mv = self.mvs[usize::from(mv_idx)].0;
385 self.encode_sec_inter(iw, cell, sec, cenc, real_mv, refp);
386 } else {
387 iw.put_2bits(2);
388 cell.intra = true;
389 self.encode_sec(iw, cell, sec, cenc);
390 }
391 },
392 }
393 }
394 fn encode_sec(&mut self, iw: &mut Indeo3Writer, cell: Indeo3Cell, tree: &Indeo3SecondaryTree, cenc: &mut CellEncoder) {
395 match tree {
396 Indeo3SecondaryTree::HSplit(t1, t2) => {
397 iw.put_2bits(0);
398 let (cell1, cell2) = cell.split_h();
399 self.encode_sec(iw, cell1, t1, cenc);
400 self.encode_sec(iw, cell2, t2, cenc);
401 },
402 Indeo3SecondaryTree::VSplit(t1, t2) => {
403 iw.put_2bits(1);
404 let (cell1, cell2) = cell.split_v(self.stripw);
405 self.encode_sec(iw, cell1, t1, cenc);
406 self.encode_sec(iw, cell2, t2, cenc);
407 },
408 Indeo3SecondaryTree::VQNull(mode) => {
409 iw.put_2bits(2);
410 iw.put_2bits(*mode);
411 },
412 Indeo3SecondaryTree::VQData(mode) => {
413 iw.put_2bits(3);
414 self.encode_cell_data_intra(iw, cell, cenc, *mode);
415 },
416 }
417 }
418 fn encode_sec_inter(&mut self, iw: &mut Indeo3Writer, cell: Indeo3Cell, tree: &Indeo3SecondaryTree, cenc: &mut CellEncoder, mv: MV, refp: &Plane) {
419 match tree {
420 Indeo3SecondaryTree::HSplit(_t1, _t2) => {
421 unimplemented!();
422 },
423 Indeo3SecondaryTree::VSplit(_t1, _t2) => {
424 unimplemented!();
425 },
426 Indeo3SecondaryTree::VQNull(mode) => {
427 iw.put_2bits(2);
428 iw.put_2bits(*mode);
429 cenc.read_mv_buffer(refp, cell, mv);
430 cenc.null_mv();
431 cenc.put_buffer(self);
432 },
433 Indeo3SecondaryTree::VQData(_mode) => {
434 iw.put_2bits(3);
435 self.encode_cell_data_inter(iw, cell, cenc, mv, refp);
436 },
437 }
438 }
439 fn encode_cell_data_intra(&mut self, iw: &mut Indeo3Writer, cell: Indeo3Cell, cenc: &mut CellEncoder, mode: u8) {
440 cenc.read_buffer(self, cell);
441 cenc.gen_diffs_intra();
442 cenc.compress_intra(mode);
443 cenc.put_buffer(self);
444 for &b in cenc.out[..cenc.osize].iter() {
445 iw.put_byte(b);
446 }
447 }
448 fn encode_cell_data_inter(&mut self, iw: &mut Indeo3Writer, cell: Indeo3Cell, cenc: &mut CellEncoder, mv: MV, refp: &Plane) {
449 cenc.read_buffer(self, cell);
450 cenc.read_mv_buffer(refp, cell, mv);
451 cenc.gen_diffs_inter();
452 cenc.compress_inter();
453 cenc.put_buffer(self);
454 for &b in cenc.out[..cenc.osize].iter() {
455 iw.put_byte(b);
456 }
457 }
458 }