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) -> Self {
  91         Self {
  92             x:      0,
  93             y:      0,
  94             w:      (width / 4) as u8,
  95             h:      (height / 4) as u8,
  96             intra:  false,
  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     #[allow(clippy::collapsible_else_if)]
 117     fn split_v(&self, stripw: u8) -> (Self, Self) {
 118         let w1 = if self.w > stripw {
 119                 if self.w > stripw * 2 { stripw * 2 } else { stripw }
 120             } else {
 121                 if self.w > 2 { ((self.w + 2) >> 2) << 1 } else { 1 }
 122             };
 123         let w2 = self.w - w1;
 124         let mut cell1 = *self;
 125         cell1.w  = w1;
 126         let mut cell2 = *self;
 127         cell2.x += w1;
 128         cell2.w  = w2;
 129         (cell1, cell2)
 130     }
 131 }
 132
 133 impl std::fmt::Display for Indeo3Cell {
 134     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
 135         write!(f, "[{}x{} @ {},{}{}]", self.get_width(), self.get_height(), self.get_x(), self.get_y(), if self.intra { " intra" } else { "" })
 136     }
 137 }
 138
 139 #[derive(Default)]
 140 pub struct Plane {
 141     pub data:   Vec<u8>,
 142     pub width:  usize,
 143     pub height: usize,
 144     pub stripw: u8,
 145     pub mvs:    Vec<(MV, u16)>,
 146     pub zero:   bool,
 147 }
 148
 149 fn ssd(a: u8, b: u8) -> u32 {
 150     let diff = i32::from(a) - i32::from(b);
 151     (diff * diff) as u32
 152 }
 153
 154 impl Plane {
 155     pub fn alloc(&mut self, width: usize, height: usize, stripw: u8) {
 156         self.data.resize(width * height, 0);
 157         self.width  = width;
 158         self.height = height;
 159         self.stripw = stripw;
 160         if self.mvs.capacity() < 256 {
 161             self.mvs.reserve(256);
 162         }
 163     }
 164     pub fn fill(&mut self, src: &[u8], stride: usize) {
 165         for (dline, sline) in self.data.chunks_mut(self.width).zip(src.chunks(stride)) {
 166             for (dst, &src) in dline.iter_mut().zip(sline.iter()) {
 167                 *dst = src >> 1;
 168             }
 169         }
 170     }
 171     pub fn clear_mvs(&mut self){
 172         self.mvs.clear();
 173     }
 174     pub fn checksum(&self) -> u16 {
 175         let xors = self.data.chunks(2).fold([0u8; 2], |acc, pair| [acc[0] ^ pair[0], acc[1] ^ pair[1]]);
 176         u16::from(xors[0]) | (u16::from(xors[1]) * 256)
 177     }
 178     pub fn find_cells(&mut self, is_intra: bool, pplane: &Plane, mv_est: &MotionEstimator) -> Box<Indeo3PrimaryTree> {
 179         let cell = Indeo3Cell::new(self.width, self.height);
 180         self.split_pri(cell, pplane, mv_est, is_intra)
 181     }
 182     fn split_pri(&mut self, mut cell: Indeo3Cell, pplane: &Plane, mv_est: &MotionEstimator, is_intra: bool) -> Box<Indeo3PrimaryTree> {
 183         let width  = cell.get_width();
 184         let height = cell.get_height();
 185         if width * height > MAX_CELL_SIZE {
 186             let (hsplit, vsplit) = if width != height {
 187                     (width > (self.stripw as usize) * 4 || width > height, height > width)
 188                 } else {
 189                     let (hdiff, vdiff) = self.calculate_diffs(cell);
 190                     (vdiff > THRESHOLD && vdiff > hdiff,
 191                      hdiff > THRESHOLD && hdiff > vdiff)
 192                 };
 193             match (hsplit, vsplit) {
 194                 (true, _) => {
 195                     let (cell1, cell2) = cell.split_v(self.stripw);
 196                     let tree1 = self.split_pri(cell1, pplane, mv_est, is_intra);
 197                     let tree2 = self.split_pri(cell2, pplane, mv_est, is_intra);
 198                     Box::new(Indeo3PrimaryTree::VSplit(tree1, tree2))
 199                 },
 200                 (_, true) => {
 201                     let (cell1, cell2) = cell.split_h();
 202                     let tree1 = self.split_pri(cell1, pplane, mv_est, is_intra);
 203                     let tree2 = self.split_pri(cell2, pplane, mv_est, is_intra);
 204                     Box::new(Indeo3PrimaryTree::HSplit(tree1, tree2))
 205                 },
 206                 (false, false) => {
 207                     let sec = self.split_sec(cell);
 208                     Box::new(Indeo3PrimaryTree::AbsFill(sec))
 209                 },
 210             }
 211         } else {
 212             if !is_intra {
 213                 if let Some((mv, flat)) = mv_est.mv_search(self, pplane, cell) {
 214                     return self.add_mv_tree(mv, flat, cell);
 215                 }
 216
 217                 // try splitting once to see if it improves situation
 218                 if width >= 16 && height >= 16 {
 219                     let vsplit = width > height;
 220                     let (mut cell1, mut cell2) = if vsplit {
 221                             cell.split_v(self.stripw)
 222                         } else {
 223                             cell.split_h()
 224                         };
 225                     let search1 = mv_est.mv_search(self, pplane, cell1);
 226                     let search2 = mv_est.mv_search(self, pplane, cell2);
 227                     if search1.is_some() || search2.is_some() {
 228                         let tree1 = if let Some((mv, flat)) = search1 {
 229                                 self.add_mv_tree(mv, flat, cell1)
 230                             } else {
 231                                 cell1.intra = true;
 232                                 let sec = self.split_sec(cell1);
 233                                 Box::new(Indeo3PrimaryTree::AbsFill(sec))
 234                             };
 235                         let tree2 = if let Some((mv, flat)) = search2 {
 236                                 self.add_mv_tree(mv, flat, cell2)
 237                             } else {
 238                                 cell2.intra = true;
 239                                 let sec = self.split_sec(cell2);
 240                                 Box::new(Indeo3PrimaryTree::AbsFill(sec))
 241                             };
 242                         return if vsplit {
 243                                 Box::new(Indeo3PrimaryTree::VSplit(tree1, tree2))
 244                             } else {
 245                                 Box::new(Indeo3PrimaryTree::HSplit(tree1, tree2))
 246                             }
 247                     }
 248                 }
 249             }
 250             cell.intra = true;
 251             let sec = self.split_sec(cell);
 252             Box::new(Indeo3PrimaryTree::AbsFill(sec))
 253         }
 254     }
 255     fn add_mv_tree(&mut self, mv: MV, flat: bool, cell: Indeo3Cell) -> Box<Indeo3PrimaryTree> {
 256         let sec = if flat {
 257                 Box::new(Indeo3SecondaryTree::VQNull(0))
 258             } else {
 259                 Box::new(Indeo3SecondaryTree::VQData(0))
 260             };
 261
 262         let mut found = false;
 263         for (ref cmv, ref mut count) in self.mvs.iter_mut() {
 264             if cmv == &mv {
 265                 *count += u16::from(cell.w) * u16::from(cell.h);
 266                 found = true;
 267                 break;
 268             }
 269         }
 270         if !found {
 271             self.mvs.push((mv, 1));
 272         }
 273
 274         Box::new(Indeo3PrimaryTree::RelFill(mv, sec))
 275     }
 276     fn split_sec(&mut self, cell: Indeo3Cell) -> Box<Indeo3SecondaryTree> {
 277         let (hdiff, vdiff) = self.calculate_diffs(cell);
 278         if hdiff == 0 && vdiff == 0 {
 279             if !cell.intra {
 280                 return Box::new(Indeo3SecondaryTree::VQNull(0));
 281             } else {
 282                 return Box::new(Indeo3SecondaryTree::VQData(0));
 283             }
 284         }
 285         if cell.get_width() > 16 && cell.get_height() > 16 {
 286             let hsplit = vdiff > THRESHOLD && vdiff > hdiff * 2;
 287             let vsplit = hdiff > THRESHOLD && hdiff > vdiff * 2;
 288             match (vsplit, hsplit) {
 289                 (true, _) => {
 290                     let (cell1, cell2) = cell.split_v(self.stripw);
 291                     let tree1 = self.split_sec(cell1);
 292                     let tree2 = self.split_sec(cell2);
 293                     Box::new(Indeo3SecondaryTree::VSplit(tree1, tree2))
 294                 },
 295                 (_, true) => {
 296                     let (cell1, cell2) = cell.split_h();
 297                     let tree1 = self.split_sec(cell1);
 298                     let tree2 = self.split_sec(cell2);
 299                     Box::new(Indeo3SecondaryTree::HSplit(tree1, tree2))
 300                 },
 301                 _ => {
 302                     Box::new(Indeo3SecondaryTree::VQData(0))
 303                 },
 304             }
 305         } else {
 306             let is_w8 = (cell.get_width()  & 7) == 0;
 307             let is_h8 = (cell.get_height() & 7) == 0;
 308             let mode = match (hdiff > THRESHOLD, vdiff > THRESHOLD) {
 309                     (false, false) if is_w8 && is_h8 => 10,
 310                     (_, true)      if is_h8          => 3,
 311                     _                                => 0,
 312                 };
 313             Box::new(Indeo3SecondaryTree::VQData(mode))
 314         }
 315     }
 316     fn calculate_diffs(&self, cell: Indeo3Cell) -> (u32, u32) {
 317         let offset = cell.get_x() + cell.get_y() * self.width;
 318         let mut w = cell.get_width();
 319         if cell.get_x() + w == self.width { w -= 1; }
 320         let mut h = cell.get_height();
 321         if cell.get_y() + h == self.height { h -= 1; }
 322
 323         let mut vdiff = 0;
 324         let mut hdiff = 0;
 325         let src0 = &self.data[offset..];
 326         let src1 = &self.data[offset + self.width..];
 327         for (line0, line1) in src0.chunks(self.width).zip(src1.chunks(self.width)).take(h) {
 328             for ((&cur, &right), &bottom) in line0.iter().zip(line0[1..].iter()).zip(line1.iter()).take(w) {
 329                 hdiff += ssd(cur, right);
 330                 vdiff += ssd(cur, bottom);
 331             }
 332         }
 333         let area = (w * h) as u32;
 334         (hdiff * 16 / area, vdiff * 16 / area)
 335     }
 336     pub fn prune_extra_mvs(&mut self, tree: &mut Box<Indeo3PrimaryTree>) {
 337         let cell = Indeo3Cell::new(self.width, self.height);
 338         self.prune_pri(cell, tree)
 339     }
 340     fn prune_pri(&mut self, cell: Indeo3Cell, tree: &mut Box<Indeo3PrimaryTree>) {
 341         match tree.deref_mut() {
 342             Indeo3PrimaryTree::HSplit(ref mut tree1, ref mut tree2) => {
 343                 let (cell1, cell2) = cell.split_h();
 344                 self.prune_pri(cell1, tree1);
 345                 self.prune_pri(cell2, tree2);
 346             },
 347             Indeo3PrimaryTree::VSplit(ref mut tree1, ref mut tree2) => {
 348                 let (cell1, cell2) = cell.split_v(self.stripw);
 349                 self.prune_pri(cell1, tree1);
 350                 self.prune_pri(cell2, tree2);
 351             },
 352             Indeo3PrimaryTree::AbsFill(_) => {},
 353             Indeo3PrimaryTree::RelFill(ref mv, ref _sec) => {
 354                 if find_mv(*mv, &self.mvs).is_none() {
 355                     let sec = self.split_sec(cell);
 356                     *tree = Box::new(Indeo3PrimaryTree::AbsFill(sec));
 357                 }
 358             },
 359         }
 360     }
 361     pub fn encode_tree(&mut self, iw: &mut Indeo3Writer, tree: &Indeo3PrimaryTree, cenc: &mut CellEncoder, refp: &Plane) {
 362         let cell = Indeo3Cell::new(self.width, self.height);
 363         self.zero = true;
 364         self.encode_pri(iw, cell, tree, cenc, refp);
 365     }
 366     fn encode_pri(&mut self, iw: &mut Indeo3Writer, mut cell: Indeo3Cell, tree: &Indeo3PrimaryTree, cenc: &mut CellEncoder, refp: &Plane) {
 367         match tree {
 368             Indeo3PrimaryTree::HSplit(t1, t2) => {
 369                 iw.put_2bits(0);
 370                 let (cell1, cell2) = cell.split_h();
 371                 self.encode_pri(iw, cell1, t1, cenc, refp);
 372                 self.encode_pri(iw, cell2, t2, cenc, refp);
 373             },
 374             Indeo3PrimaryTree::VSplit(t1, t2) => {
 375                 iw.put_2bits(1);
 376                 let (cell1, cell2) = cell.split_v(self.stripw);
 377                 self.encode_pri(iw, cell1, t1, cenc, refp);
 378                 self.encode_pri(iw, cell2, t2, cenc, refp);
 379             },
 380             Indeo3PrimaryTree::AbsFill(sec) => {
 381                 iw.put_2bits(2);
 382                 cell.intra = true;
 383                 self.zero = false;
 384                 self.encode_sec(iw, cell, sec, cenc);
 385             }
 386             Indeo3PrimaryTree::RelFill(mv, sec) => {
 387                 if let Some(mv_idx) = find_mv(*mv, &self.mvs) {
 388                     iw.put_2bits(3);
 389                     iw.put_byte(mv_idx);
 390                     cell.intra = false;
 391                     let real_mv = self.mvs[usize::from(mv_idx)].0;
 392                     self.encode_sec_inter(iw, cell, sec, cenc, real_mv, refp);
 393                 } else {
 394                     iw.put_2bits(2);
 395                     cell.intra = true;
 396                     self.zero = false;
 397                     self.encode_sec(iw, cell, sec, cenc);
 398                 }
 399             },
 400         }
 401     }
 402     fn encode_sec(&mut self, iw: &mut Indeo3Writer, cell: Indeo3Cell, tree: &Indeo3SecondaryTree, cenc: &mut CellEncoder) {
 403         match tree {
 404             Indeo3SecondaryTree::HSplit(t1, t2) => {
 405                 iw.put_2bits(0);
 406                 let (cell1, cell2) = cell.split_h();
 407                 self.encode_sec(iw, cell1, t1, cenc);
 408                 self.encode_sec(iw, cell2, t2, cenc);
 409             },
 410             Indeo3SecondaryTree::VSplit(t1, t2) => {
 411                 iw.put_2bits(1);
 412                 let (cell1, cell2) = cell.split_v(self.stripw);
 413                 self.encode_sec(iw, cell1, t1, cenc);
 414                 self.encode_sec(iw, cell2, t2, cenc);
 415             },
 416             Indeo3SecondaryTree::VQNull(mode) => {
 417                 iw.put_2bits(2);
 418                 iw.put_2bits(*mode);
 419             },
 420             Indeo3SecondaryTree::VQData(mode) => {
 421                 iw.put_2bits(3);
 422                 self.zero = false;
 423                 self.encode_cell_data_intra(iw, cell, cenc, *mode);
 424             },
 425         }
 426     }
 427     fn encode_sec_inter(&mut self, iw: &mut Indeo3Writer, cell: Indeo3Cell, tree: &Indeo3SecondaryTree, cenc: &mut CellEncoder, mv: MV, refp: &Plane) {
 428         match tree {
 429             Indeo3SecondaryTree::HSplit(_t1, _t2) => {
 430                 unimplemented!();
 431             },
 432             Indeo3SecondaryTree::VSplit(_t1, _t2) => {
 433                 unimplemented!();
 434             },
 435             Indeo3SecondaryTree::VQNull(mode) => {
 436                 iw.put_2bits(2);
 437                 iw.put_2bits(*mode);
 438                 cenc.read_mv_buffer(refp, cell, mv);
 439                 if mv.x != 0 || mv.y != 0 {
 440                     self.zero = false;
 441                 }
 442                 cenc.null_mv();
 443                 cenc.put_buffer(self);
 444             },
 445             Indeo3SecondaryTree::VQData(_mode) => {
 446                 iw.put_2bits(3);
 447                 self.zero = false;
 448                 self.encode_cell_data_inter(iw, cell, cenc, mv, refp);
 449             },
 450         }
 451     }
 452     fn encode_cell_data_intra(&mut self, iw: &mut Indeo3Writer, cell: Indeo3Cell, cenc: &mut CellEncoder, mode: u8) {
 453         cenc.read_buffer(self, cell);
 454         cenc.gen_diffs_intra();
 455         cenc.compress_intra(mode);
 456         cenc.put_buffer(self);
 457         for &b in cenc.out[..cenc.osize].iter() {
 458             iw.put_byte(b);
 459         }
 460     }
 461     fn encode_cell_data_inter(&mut self, iw: &mut Indeo3Writer, cell: Indeo3Cell, cenc: &mut CellEncoder, mv: MV, refp: &Plane) {
 462         cenc.read_buffer(self, cell);
 463         cenc.read_mv_buffer(refp, cell, mv);
 464         cenc.gen_diffs_inter();
 465         cenc.compress_inter();
 466         cenc.put_buffer(self);
 467         for &b in cenc.out[..cenc.osize].iter() {
 468             iw.put_byte(b);
 469         }
 470     }
 471 }