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