nihav-duck/src/codecs/vpenc/motion_est.rs

   1 use nihav_codec_support::codecs::{MV, ZERO_MV};
   2
   3 use std::str::FromStr;
   4
   5 #[derive(Debug,Clone,Copy,PartialEq,Default)]
   6 pub enum MVSearchMode {
   7     Full,
   8     SEA,
   9     Diamond,
  10     #[default]
  11     Hexagon,
  12     EPZS,
  13 }
  14
  15 pub struct ParseError{}
  16
  17 impl FromStr for MVSearchMode {
  18     type Err = ParseError;
  19
  20     fn from_str(s: &str) -> Result<Self, Self::Err> {
  21         match s {
  22             "full"  => Ok(MVSearchMode::Full),
  23             "sea"   => Ok(MVSearchMode::SEA),
  24             "dia"   => Ok(MVSearchMode::Diamond),
  25             "hex"   => Ok(MVSearchMode::Hexagon),
  26             "epzs"  => Ok(MVSearchMode::EPZS),
  27             _ => Err(ParseError{}),
  28         }
  29     }
  30 }
  31
  32 impl std::fmt::Display for MVSearchMode {
  33     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
  34         match *self {
  35             MVSearchMode::Full      => write!(f, "full"),
  36             MVSearchMode::SEA       => write!(f, "sea"),
  37             MVSearchMode::Diamond   => write!(f, "dia"),
  38             MVSearchMode::Hexagon   => write!(f, "hex"),
  39             MVSearchMode::EPZS      => write!(f, "epzs"),
  40         }
  41     }
  42 }
  43
  44 pub const DIA_PATTERN: [MV; 9] = [
  45     ZERO_MV,
  46     MV {x: -2, y:  0},
  47     MV {x: -1, y:  1},
  48     MV {x:  0, y:  2},
  49     MV {x:  1, y:  1},
  50     MV {x:  2, y:  0},
  51     MV {x:  1, y: -1},
  52     MV {x:  0, y: -2},
  53     MV {x: -1, y: -1}
  54 ];
  55
  56 pub const HEX_PATTERN: [MV; 7] = [
  57     ZERO_MV,
  58     MV {x: -2, y:  0},
  59     MV {x: -1, y:  2},
  60     MV {x:  1, y:  2},
  61     MV {x:  2, y:  0},
  62     MV {x:  1, y: -2},
  63     MV {x: -1, y: -2}
  64 ];
  65
  66 pub const REFINEMENT: [MV; 4] = [
  67     MV {x: -1, y:  0},
  68     MV {x:  0, y:  1},
  69     MV {x:  1, y:  0},
  70     MV {x:  0, y: -1}
  71 ];
  72
  73 #[macro_export]
  74 macro_rules! search_template {
  75     ($self: expr, $mv_est: expr, $cur_blk: expr, $mb_x: expr, $mb_y: expr, $sad_func: ident, $threshold: expr) => ({
  76         search_template!($self, $mv_est, $cur_blk, $mb_x, $mb_y, $sad_func, $threshold, ZERO_MV, MAX_DIST, true)
  77     });
  78     ($self: expr, $mv_est: expr, $cur_blk: expr, $mb_x: expr, $mb_y: expr, $sad_func: ident, $threshold: expr, $start_mv: expr, $best_dist: expr, $fullpel_stage: expr) => ({
  79         let mut best_dist = $best_dist;
  80         let mut best_mv = $start_mv;
  81
  82         let mut min_dist;
  83         let mut min_idx;
  84
  85         if $fullpel_stage {
  86             $self.reset();
  87             loop {
  88                 let mut cur_best_dist = best_dist;
  89                 for (dist, &point) in $self.dist.iter_mut().zip($self.point.iter()) {
  90                     if *dist == MAX_DIST {
  91                         *dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, point.from_pixels(), cur_best_dist);
  92                         cur_best_dist = cur_best_dist.min(*dist);
  93                         if *dist <= $threshold {
  94                             break;
  95                         }
  96                     }
  97                 }
  98                 min_dist = $self.dist[0];
  99                 min_idx = 0;
 100                 for (i, &dist) in $self.dist.iter().enumerate().skip(1) {
 101                     if dist < min_dist {
 102                         min_dist = dist;
 103                         min_idx = i;
 104                         if dist <= $threshold {
 105                             break;
 106                         }
 107                     }
 108                 }
 109                 if min_dist <= $threshold || min_idx == 0 || best_dist == min_dist || $self.point[min_idx].x.abs() >= $mv_est.mv_range || $self.point[min_idx].y.abs() >= $mv_est.mv_range {
 110                     break;
 111                 }
 112                 best_dist = min_dist;
 113                 $self.update($self.steps[min_idx]);
 114             }
 115             best_dist = min_dist;
 116             best_mv   = $self.point[min_idx];
 117             if best_dist <= $threshold {
 118                 return (best_mv.from_pixels(), best_dist);
 119             }
 120             for &step in REFINEMENT.iter() {
 121                 let mv = best_mv + step;
 122                 let dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, mv.from_pixels(), MAX_DIST);
 123                 if best_dist > dist {
 124                     best_dist = dist;
 125                     best_mv = mv;
 126                 }
 127             }
 128             best_mv = best_mv.from_pixels();
 129             if best_dist <= $threshold {
 130                 return (best_mv, best_dist);
 131             }
 132         }
 133
 134         // subpel refinement
 135         $self.set_new_point(best_mv, best_dist);
 136         loop {
 137             let mut cur_best_dist = best_dist;
 138             for (dist, &point) in $self.dist.iter_mut().zip($self.point.iter()) {
 139                 if *dist == MAX_DIST {
 140                     *dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, point, cur_best_dist);
 141                     cur_best_dist = cur_best_dist.min(*dist);
 142                     if *dist <= $threshold {
 143                         break;
 144                     }
 145                 }
 146             }
 147             min_dist = $self.dist[0];
 148             min_idx = 0;
 149             for (i, &dist) in $self.dist.iter().enumerate().skip(1) {
 150                 if dist < min_dist {
 151                     min_dist = dist;
 152                     min_idx = i;
 153                     if dist <= $threshold {
 154                         break;
 155                     }
 156                 }
 157             }
 158             if min_dist <= $threshold || min_idx == 0 || best_dist == min_dist || $self.point[min_idx].x.abs() >= $mv_est.mv_range * 8 || $self.point[min_idx].y.abs() >= $mv_est.mv_range * 8 {
 159                 break;
 160             }
 161             best_dist = min_dist;
 162             $self.update($self.steps[min_idx]);
 163         }
 164         best_dist = min_dist;
 165         best_mv   = $self.point[min_idx];
 166         if best_dist <= $threshold {
 167             return (best_mv, best_dist);
 168         }
 169         for &step in REFINEMENT.iter() {
 170             let mv = best_mv + step;
 171             let dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, mv, MAX_DIST);
 172             if best_dist > dist {
 173                 best_dist = dist;
 174                 best_mv = mv;
 175             }
 176         }
 177         (best_mv, best_dist)
 178     });
 179 }