me.rs 3.75 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12
// 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 std::cmp;
use FrameInvariants;
use FrameState;
Frank Bossen's avatar
Frank Bossen committed
13
use partition::*;
14
use context::BlockOffset;
Frank Bossen's avatar
Frank Bossen committed
15
use plane::*;
16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31
use context::BLOCK_TO_PLANE_SHIFT;

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) => {
      let po = PlaneOffset { x: bo.x << BLOCK_TO_PLANE_SHIFT, y: bo.y << BLOCK_TO_PLANE_SHIFT };
      let range = 16 as usize;
      let blk_w = bsize.width();
      let blk_h = bsize.height();
      let x_lo = cmp::max(0, po.x as isize - range as isize) as usize;
      let x_hi = cmp::min(fs.input.planes[0].cfg.width - blk_w, po.x + range);
      let y_lo = cmp::max(0, po.y as isize - range as isize) as usize;
      let y_hi = cmp::min(fs.input.planes[0].cfg.height - blk_h, po.y + range);

32
      let mut lowest_sad = 128*128*4096 as u32;
33 34
      let mut best_mv = MotionVector { row: 0, col: 0 };

35 36
      for y in (y_lo..y_hi).step_by(1) {
        for x in (x_lo..x_hi).step_by(1) {
37 38 39
          let mut sad = 0 as u32;
          let mut plane_org = fs.input.planes[0].slice(&po);
          let mut plane_ref = rec.planes[0].slice(&PlaneOffset { x: x, y: y });
40

41 42 43 44 45
          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);
              sad += slice_org.iter().zip(slice_ref).map(|(&a, &b)| (a as i32 - b as i32).abs() as u32).sum::<u32>();
46
            }
47 48
            plane_org.y += 1;
            plane_ref.y += 1;
49 50 51 52 53 54 55 56 57
          }

          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
58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99

      let mode = PredictionMode::NEWMV;
      let mut tmp_plane = Plane::new(blk_w, blk_h, 0, 0);

      for step in [4, 2].iter() {
        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 sad = 0 as u32;
            let mut plane_org = fs.input.planes[0].slice(&po);
            let mut plane_ref = tmp_plane.slice(&PlaneOffset { x:0, y:0 });

            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);
                sad += 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;
            }

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

100
      best_mv
101
    },
102 103 104 105

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