indeo3enc: rework cell intra flag logic

[nihav.git] / nihav-indeo / src / codecs / indeo3enc / tree.rs

diff --git a/nihav-indeo/src/codecs/indeo3enc/tree.rs b/nihav-indeo/src/codecs/indeo3enc/tree.rs
index 8c51ad74bd2ad0cc703287841d89332a6684253a..921d2d9a4467d1ed7cb095e56e9aaa613c269634 100644 (file)
--- a/nihav-indeo/src/codecs/indeo3enc/tree.rs
+++ b/nihav-indeo/src/codecs/indeo3enc/tree.rs
@@ -1,6 +1,7 @@
 use super::Indeo3Writer;
 use super::mv::*;
 use super::cell::{CellEncoder, MAX_CELL_SIZE};
+use std::ops::DerefMut;
 
 pub enum Indeo3PrimaryTree {
     VSplit(Box<Indeo3PrimaryTree>, Box<Indeo3PrimaryTree>),
@@ -86,13 +87,13 @@ pub struct Indeo3Cell {
 }
 
 impl Indeo3Cell {
-    pub fn new(width: usize, height: usize, intra: bool) -> Self {
+    pub fn new(width: usize, height: usize) -> Self {
         Self {
             x:      0,
             y:      0,
             w:      (width / 4) as u8,
             h:      (height / 4) as u8,
-            intra,
+            intra:  false,
         }
     }
 
@@ -169,14 +170,14 @@ impl Plane {
         self.mvs.clear();
     }
     pub fn checksum(&self) -> u16 {
-        let xors = self.data[self.width..].chunks(2).fold([0u8; 2], |acc, pair| [acc[0] ^ pair[0], acc[1] ^ pair[1]]);
+        let xors = self.data.chunks(2).fold([0u8; 2], |acc, pair| [acc[0] ^ pair[0], acc[1] ^ pair[1]]);
         u16::from(xors[0]) | (u16::from(xors[1]) * 256)
     }
     pub fn find_cells(&mut self, is_intra: bool, pplane: &Plane, mv_est: &MotionEstimator) -> Box<Indeo3PrimaryTree> {
-        let cell = Indeo3Cell::new(self.width, self.height, is_intra);
-        self.split_pri(cell, pplane, mv_est)
+        let cell = Indeo3Cell::new(self.width, self.height);
+        self.split_pri(cell, pplane, mv_est, is_intra)
     }
-    fn split_pri(&mut self, cell: Indeo3Cell, pplane: &Plane, mv_est: &MotionEstimator) -> Box<Indeo3PrimaryTree> {
+    fn split_pri(&mut self, mut cell: Indeo3Cell, pplane: &Plane, mv_est: &MotionEstimator, is_intra: bool) -> Box<Indeo3PrimaryTree> {
         let width  = cell.get_width();
         let height = cell.get_height();
         if width * height > MAX_CELL_SIZE {
@@ -190,14 +191,14 @@ impl Plane {
             match (hsplit, vsplit) {
                 (true, _) => {
                     let (cell1, cell2) = cell.split_v(self.stripw);
-                    let tree1 = self.split_pri(cell1, pplane, mv_est);
-                    let tree2 = self.split_pri(cell2, pplane, mv_est);
+                    let tree1 = self.split_pri(cell1, pplane, mv_est, is_intra);
+                    let tree2 = self.split_pri(cell2, pplane, mv_est, is_intra);
                     Box::new(Indeo3PrimaryTree::VSplit(tree1, tree2))
                 },
                 (_, true) => {
                     let (cell1, cell2) = cell.split_h();
-                    let tree1 = self.split_pri(cell1, pplane, mv_est);
-                    let tree2 = self.split_pri(cell2, pplane, mv_est);
+                    let tree1 = self.split_pri(cell1, pplane, mv_est, is_intra);
+                    let tree2 = self.split_pri(cell2, pplane, mv_est, is_intra);
                     Box::new(Indeo3PrimaryTree::HSplit(tree1, tree2))
                 },
                 (false, false) => {
@@ -206,15 +207,15 @@ impl Plane {
                 },
             }
         } else {
-            if !cell.intra {
+            if !is_intra {
                 if let Some((mv, flat)) = mv_est.mv_search(self, pplane, cell) {
-                    return self.add_mv_tree(mv, flat);
+                    return self.add_mv_tree(mv, flat, cell);
                 }
 
                 // try splitting once to see if it improves situation
                 if width >= 16 && height >= 16 {
                     let vsplit = width > height;
-                    let (cell1, cell2) = if vsplit {
+                    let (mut cell1, mut cell2) = if vsplit {
                             cell.split_v(self.stripw)
                         } else {
                             cell.split_h()
@@ -223,14 +224,16 @@ impl Plane {
                     let search2 = mv_est.mv_search(self, pplane, cell2);
                     if search1.is_some() || search2.is_some() {
                         let tree1 = if let Some((mv, flat)) = search1 {
-                                self.add_mv_tree(mv, flat)
+                                self.add_mv_tree(mv, flat, cell1)
                             } else {
+                                cell1.intra = true;
                                 let sec = self.split_sec(cell1);
                                 Box::new(Indeo3PrimaryTree::AbsFill(sec))
                             };
                         let tree2 = if let Some((mv, flat)) = search2 {
-                                self.add_mv_tree(mv, flat)
+                                self.add_mv_tree(mv, flat, cell2)
                             } else {
+                                cell2.intra = true;
                                 let sec = self.split_sec(cell2);
                                 Box::new(Indeo3PrimaryTree::AbsFill(sec))
                             };
@@ -242,11 +245,12 @@ impl Plane {
                     }
                 }
             }
+            cell.intra = true;
             let sec = self.split_sec(cell);
             Box::new(Indeo3PrimaryTree::AbsFill(sec))
         }
     }
-    fn add_mv_tree(&mut self, mv: MV, flat: bool) -> Box<Indeo3PrimaryTree> {
+    fn add_mv_tree(&mut self, mv: MV, flat: bool, cell: Indeo3Cell) -> Box<Indeo3PrimaryTree> {
         let sec = if flat {
                 Box::new(Indeo3SecondaryTree::VQNull(0))
             } else {
@@ -256,7 +260,7 @@ impl Plane {
         let mut found = false;
         for (ref cmv, ref mut count) in self.mvs.iter_mut() {
             if cmv == &mv {
-                *count += 1;
+                *count += u16::from(cell.w) * u16::from(cell.h);
                 found = true;
                 break;
             }
@@ -327,8 +331,33 @@ impl Plane {
         let area = (w * h) as u32;
         (hdiff * 16 / area, vdiff * 16 / area)
     }
-    pub fn encode_tree(&mut self, iw: &mut Indeo3Writer, tree: &Indeo3PrimaryTree, cenc: &mut CellEncoder, is_intra: bool, refp: &Plane) {
-        let cell = Indeo3Cell::new(self.width, self.height, is_intra);
+    pub fn prune_extra_mvs(&mut self, tree: &mut Box<Indeo3PrimaryTree>) {
+        let cell = Indeo3Cell::new(self.width, self.height);
+        self.prune_pri(cell, tree)
+    }
+    fn prune_pri(&mut self, cell: Indeo3Cell, tree: &mut Box<Indeo3PrimaryTree>) {
+        match tree.deref_mut() {
+            Indeo3PrimaryTree::HSplit(ref mut tree1, ref mut tree2) => {
+                let (cell1, cell2) = cell.split_h();
+                self.prune_pri(cell1, tree1);
+                self.prune_pri(cell2, tree2);
+            },
+            Indeo3PrimaryTree::VSplit(ref mut tree1, ref mut tree2) => {
+                let (cell1, cell2) = cell.split_v(self.stripw);
+                self.prune_pri(cell1, tree1);
+                self.prune_pri(cell2, tree2);
+            },
+            Indeo3PrimaryTree::AbsFill(_) => {},
+            Indeo3PrimaryTree::RelFill(ref mv, ref _sec) => {
+                if find_mv(*mv, &self.mvs).is_none() {
+                    let sec = self.split_sec(cell);
+                    *tree = Box::new(Indeo3PrimaryTree::AbsFill(sec));
+                }
+            },
+        }
+    }
+    pub fn encode_tree(&mut self, iw: &mut Indeo3Writer, tree: &Indeo3PrimaryTree, cenc: &mut CellEncoder, refp: &Plane) {
+        let cell = Indeo3Cell::new(self.width, self.height);
         self.encode_pri(iw, cell, tree, cenc, refp);
     }
     fn encode_pri(&mut self, iw: &mut Indeo3Writer, mut cell: Indeo3Cell, tree: &Indeo3PrimaryTree, cenc: &mut CellEncoder, refp: &Plane) {