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