me.rs 3.47 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 FrameInvariants;
use FrameState;
Frank Bossen's avatar
Frank Bossen committed
12
use partition::*;
13
use context::BlockOffset;
Frank Bossen's avatar
Frank Bossen committed
14
use plane::*;
15 16
use context::BLOCK_TO_PLANE_SHIFT;

17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33
#[inline(always)]
pub fn get_sad(plane_org: &mut PlaneSlice, plane_ref: &mut PlaneSlice, blk_h: usize, blk_w: usize) -> u32 {
  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);
      sum += slice_org.iter().zip(slice_ref).map(|(&a, &b)| (a as i32 - b as i32).abs() as u32).sum::<u32>();
    }
    plane_org.y += 1;
    plane_ref.y += 1;
  }

  sum
}

34 35 36 37 38
pub fn motion_estimation(fi: &FrameInvariants, fs: &mut FrameState, bsize: BlockSize,
                         bo: &BlockOffset, ref_frame: usize) -> MotionVector {

  match fi.rec_buffer.frames[fi.ref_frames[ref_frame - LAST_FRAME]] {
    Some(ref rec) => {
fbossen's avatar
fbossen committed
39 40
      let po = PlaneOffset { x: (bo.x as isize) << BLOCK_TO_PLANE_SHIFT, y: (bo.y as isize) << BLOCK_TO_PLANE_SHIFT };
      let range = 16 as isize;
41 42
      let blk_w = bsize.width();
      let blk_h = bsize.height();
fbossen's avatar
fbossen committed
43 44 45 46
      let x_lo = po.x - range;
      let x_hi = po.x + range;
      let y_lo = po.y - range;
      let y_hi = po.y + range;
47

48
      let mut lowest_sad = 128*128*4096 as u32;
49 50
      let mut best_mv = MotionVector { row: 0, col: 0 };

51 52
      for y in (y_lo..y_hi).step_by(1) {
        for x in (x_lo..x_hi).step_by(1) {
53
          let mut plane_org = fs.input.planes[0].slice(&po);
54
          let mut plane_ref = rec.frame.planes[0].slice(&PlaneOffset { x: x, y: y });
55

56
          let sad = get_sad(&mut plane_org, &mut plane_ref, blk_h, blk_w);
57 58 59 60 61 62 63 64

          if sad < lowest_sad {
            lowest_sad = sad;
            best_mv = MotionVector { row: 8*(y as i16 - po.y as i16), col: 8*(x as i16 - po.x as i16) }
          }

        }
      }
Frank Bossen's avatar
Frank Bossen committed
65 66

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

69 70 71 72 73 74
      let mut steps = vec![4, 2];
      if fi.allow_high_precision_mv {
        steps.push(1);
      }

      for step in steps {
Frank Bossen's avatar
Frank Bossen committed
75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92
        let center_mv_h = best_mv;
        for i in 0..3 {
          for j in 0..3 {
            // Skip the center point that was already tested
            if i == 1 && j == 1 { continue; }

            let cand_mv = MotionVector { row: center_mv_h.row + step*(i as i16 - 1),
            col: center_mv_h.col + step*(j as i16 - 1) };

            {
              let tmp_slice = &mut tmp_plane.mut_slice(&PlaneOffset { x:0, y:0 });

              mode.predict_inter(fi, 0, &po, tmp_slice, blk_w, blk_h, ref_frame, &cand_mv, 8);
            }

            let mut plane_org = fs.input.planes[0].slice(&po);
            let mut plane_ref = tmp_plane.slice(&PlaneOffset { x:0, y:0 });

93
            let sad = get_sad(&mut plane_org, &mut plane_ref, blk_h, blk_w);
Frank Bossen's avatar
Frank Bossen committed
94 95 96 97 98 99 100 101 102

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

103
      best_mv
104
    },
105 106 107 108

    None => MotionVector { row: 0, col : 0 }
  }
}