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