me.rs 4.56 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11
// Copyright (c) 2017-2018, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use context::BlockOffset;
use context::BLOCK_TO_PLANE_SHIFT;
12
use context::MI_SIZE;
13 14 15 16
use partition::*;
use plane::*;
use FrameInvariants;
use FrameState;
17

18
#[inline(always)]
19 20 21 22
pub fn get_sad(
  plane_org: &mut PlaneSlice, plane_ref: &mut PlaneSlice, blk_h: usize,
  blk_w: usize
) -> u32 {
23 24 25 26 27 28
  let mut sum = 0 as u32;

  for _r in 0..blk_h {
    {
      let slice_org = plane_org.as_slice_w_width(blk_w);
      let slice_ref = plane_ref.as_slice_w_width(blk_w);
29 30 31 32 33
      sum += slice_org
        .iter()
        .zip(slice_ref)
        .map(|(&a, &b)| (a as i32 - b as i32).abs() as u32)
        .sum::<u32>();
34 35 36 37 38 39 40 41
    }
    plane_org.y += 1;
    plane_ref.y += 1;
  }

  sum
}

42 43 44 45
pub fn motion_estimation(
  fi: &FrameInvariants, fs: &mut FrameState, bsize: BlockSize,
  bo: &BlockOffset, ref_frame: usize
) -> MotionVector {
46 47
  match fi.rec_buffer.frames[fi.ref_frames[ref_frame - LAST_FRAME]] {
    Some(ref rec) => {
48 49 50 51
      let po = PlaneOffset {
        x: (bo.x as isize) << BLOCK_TO_PLANE_SHIFT,
        y: (bo.y as isize) << BLOCK_TO_PLANE_SHIFT
      };
52
      let range = 32 as isize;
53 54
      let blk_w = bsize.width();
      let blk_h = bsize.height();
55 56 57 58 59 60 61 62
      let border_w = 128 + blk_w as isize * 8;
      let border_h = 128 + blk_h as isize * 8;
      let cols = (rec.frame.planes[0].cfg.width + MI_SIZE - 1) / MI_SIZE;
      let rows = (rec.frame.planes[0].cfg.height + MI_SIZE - 1) / MI_SIZE;
      let x_min = -(bo.x as isize) * (8 * MI_SIZE) as isize - border_w;
      let x_max = (cols - bo.x - blk_w / MI_SIZE) as isize * (8 * MI_SIZE) as isize + border_w;
      let y_min = -(bo.y as isize) * (8 * MI_SIZE) as isize - border_h;
      let y_max = (rows - bo.y - blk_h / MI_SIZE) as isize * (8 * MI_SIZE) as isize + border_h;
Yushin Cho's avatar
Yushin Cho committed
63
      let x_lo = po.x + ((-range).max(x_min / 8));
64
      let x_hi = po.x + (range.min(x_max / 8));
Yushin Cho's avatar
Yushin Cho committed
65
      let y_lo = po.y + ((-range).max(y_min / 8));
66
      let y_hi = po.y + (range.min(y_max / 8));
67

68
      let mut lowest_sad = 128 * 128 * 4096 as u32;
69 70
      let mut best_mv = MotionVector { row: 0, col: 0 };

71 72
      for y in (y_lo..y_hi).step_by(2) {
        for x in (x_lo..x_hi).step_by(2) {
73
          let mut plane_org = fs.input.planes[0].slice(&po);
74
          let mut plane_ref = rec.frame.planes[0].slice(&PlaneOffset { x, y });
75

76
          let sad = get_sad(&mut plane_org, &mut plane_ref, blk_h, blk_w);
77 78 79

          if sad < lowest_sad {
            lowest_sad = sad;
80 81 82 83
            best_mv = MotionVector {
              row: 8 * (y as i16 - po.y as i16),
              col: 8 * (x as i16 - po.x as i16)
            }
84 85 86
          }
        }
      }
Frank Bossen's avatar
Frank Bossen committed
87 88

      let mode = PredictionMode::NEWMV;
fbossen's avatar
fbossen committed
89
      let mut tmp_plane = Plane::new(blk_w, blk_h, 0, 0, 0, 0);
Frank Bossen's avatar
Frank Bossen committed
90

91
      let mut steps = vec![8, 4, 2];
92 93 94 95 96
      if fi.allow_high_precision_mv {
        steps.push(1);
      }

      for step in steps {
Frank Bossen's avatar
Frank Bossen committed
97 98 99 100
        let center_mv_h = best_mv;
        for i in 0..3 {
          for j in 0..3 {
            // Skip the center point that was already tested
101 102 103
            if i == 1 && j == 1 {
              continue;
            }
Frank Bossen's avatar
Frank Bossen committed
104

105 106 107 108
            let cand_mv = MotionVector {
              row: center_mv_h.row + step * (i as i16 - 1),
              col: center_mv_h.col + step * (j as i16 - 1)
            };
Frank Bossen's avatar
Frank Bossen committed
109

110 111 112 113 114 115 116
            if (cand_mv.col as isize) < x_min || (cand_mv.col as isize) > x_max {
              continue;
            }
            if (cand_mv.row as isize) < y_min || (cand_mv.row as isize) > y_max {
              continue;
            }

Frank Bossen's avatar
Frank Bossen committed
117
            {
118 119
              let tmp_slice =
                &mut tmp_plane.mut_slice(&PlaneOffset { x: 0, y: 0 });
Frank Bossen's avatar
Frank Bossen committed
120

121 122 123
              mode.predict_inter(
                fi, 0, &po, tmp_slice, blk_w, blk_h, ref_frame, &cand_mv, 8,
              );
Frank Bossen's avatar
Frank Bossen committed
124 125 126
            }

            let mut plane_org = fs.input.planes[0].slice(&po);
127
            let mut plane_ref = tmp_plane.slice(&PlaneOffset { x: 0, y: 0 });
Frank Bossen's avatar
Frank Bossen committed
128

129
            let sad = get_sad(&mut plane_org, &mut plane_ref, blk_h, blk_w);
Frank Bossen's avatar
Frank Bossen committed
130 131 132 133 134 135 136 137 138

            if sad < lowest_sad {
              lowest_sad = sad;
              best_mv = cand_mv;
            }
          }
        }
      }

139
      best_mv
140
    }
141

142
    None => MotionVector { row: 0, col: 0 }
143 144
  }
}