Commit 2c83dfd6 authored by Adrien Maglo's avatar Adrien Maglo Committed by Luca Barbato

Use diamond search for the half resolution motion estimation

parent fef34fda
......@@ -1112,14 +1112,16 @@ pub fn motion_compensate<T: Pixel>(
pub fn save_block_motion<T: Pixel>(
fs: &mut FrameState<T>,
w_in_b: usize, _h_in_b: usize,
w_in_b: usize, h_in_b: usize,
bsize: BlockSize, bo: &BlockOffset,
ref_frame: usize, mv: MotionVector,
) {
let frame_mvs = &mut fs.frame_mvs;
for mi_y in (bo.y)..(bo.y + bsize.height_mi()) {
for mi_x in (bo.x)..(bo.x + bsize.width_mi()) {
frame_mvs[ref_frame][mi_y][mi_x] = mv;
let frame_mvs = &mut fs.frame_mvs[ref_frame];
let bo_x_end = (bo.x + bsize.width_mi()).min(w_in_b);
let bo_y_end = (bo.y + bsize.height_mi()).min(h_in_b);
for mi_y in bo.y..bo_y_end {
for mi_x in bo.x..bo_x_end {
frame_mvs[mi_y][mi_x] = mv;
}
}
}
......@@ -2044,7 +2046,8 @@ fn encode_tile<T: Pixel>(fi: &FrameInvariants<T>, fs: &mut FrameState<T>) -> Vec
BlockSize::BLOCK_32X32,
r,
&sbo.block_offset(0, 0),
&[Some(pmv), pmv_w, pmv_n]
&[Some(pmv), pmv_w, pmv_n],
i
)
});
s.spawn(|_| {
......@@ -2054,7 +2057,8 @@ fn encode_tile<T: Pixel>(fi: &FrameInvariants<T>, fs: &mut FrameState<T>) -> Vec
BlockSize::BLOCK_32X32,
r,
&sbo.block_offset(8, 0),
&[Some(pmv), pmv_e, pmv_n]
&[Some(pmv), pmv_e, pmv_n],
i
)
});
s.spawn(|_| {
......@@ -2064,7 +2068,8 @@ fn encode_tile<T: Pixel>(fi: &FrameInvariants<T>, fs: &mut FrameState<T>) -> Vec
BlockSize::BLOCK_32X32,
r,
&sbo.block_offset(0, 8),
&[Some(pmv), pmv_w, pmv_s]
&[Some(pmv), pmv_w, pmv_s],
i
)
});
s.spawn(|_| {
......@@ -2074,7 +2079,8 @@ fn encode_tile<T: Pixel>(fi: &FrameInvariants<T>, fs: &mut FrameState<T>) -> Vec
BlockSize::BLOCK_32X32,
r,
&sbo.block_offset(8, 8),
&[Some(pmv), pmv_e, pmv_s]
&[Some(pmv), pmv_e, pmv_s],
i
)
});
});
......@@ -2083,6 +2089,19 @@ fn encode_tile<T: Pixel>(fi: &FrameInvariants<T>, fs: &mut FrameState<T>) -> Vec
pmvs[2][r] = pmvs2;
pmvs[3][r] = pmvs3;
pmvs[4][r] = pmvs4;
if let Some(mv1) = pmvs1 {
save_block_motion(fs, fi.w_in_b, fi.h_in_b, BlockSize::BLOCK_32X32, &sbo.block_offset(0, 0), i, mv1);
}
if let Some(mv2) = pmvs2 {
save_block_motion(fs, fi.w_in_b, fi.h_in_b, BlockSize::BLOCK_32X32, &sbo.block_offset(8, 0), i, mv2);
}
if let Some(mv3) = pmvs3 {
save_block_motion(fs, fi.w_in_b, fi.h_in_b, BlockSize::BLOCK_32X32, &sbo.block_offset(0, 8), i, mv3);
}
if let Some(mv4) = pmvs4 {
save_block_motion(fs, fi.w_in_b, fi.h_in_b, BlockSize::BLOCK_32X32, &sbo.block_offset(8, 8), i, mv4);
}
}
}
}
......
......@@ -263,7 +263,8 @@ fn get_mv_range(
}
pub fn get_subset_predictors<T: Pixel>(
fi: &FrameInvariants<T>, bo: &BlockOffset, cmv: MotionVector,
bo: &BlockOffset, cmv: MotionVector,
w_in_b: usize, h_in_b: usize,
frame_mvs: &FrameMotionVectors, frame_ref_opt: &Option<Arc<ReferenceFrame<T>>>,
ref_frame_id: usize
) -> (Vec<MotionVector>) {
......@@ -288,7 +289,7 @@ pub fn get_subset_predictors<T: Pixel>(
median_preds.push(top);
if !top.is_zero() { predictors.push(top); }
if bo.x < fi.w_in_b - 1 {
if bo.x < w_in_b - 1 {
let top_right = frame_mvs[bo.y - 1][bo.x + 1];
median_preds.push(top_right);
if !top_right.is_zero() { predictors.push(top_right); }
......@@ -318,11 +319,11 @@ pub fn get_subset_predictors<T: Pixel>(
let top = prev_frame_mvs[bo.y - 1][bo.x];
if !top.is_zero() { predictors.push(top); }
}
if bo.x < fi.w_in_b - 1 {
if bo.x < w_in_b - 1 {
let right = prev_frame_mvs[bo.y][bo.x + 1];
if !right.is_zero() { predictors.push(right); }
}
if bo.y < fi.h_in_b - 1 {
if bo.y < h_in_b - 1 {
let bottom = prev_frame_mvs[bo.y + 1][bo.x];
if !bottom.is_zero() { predictors.push(bottom); }
}
......@@ -869,12 +870,13 @@ pub fn estimate_motion_ss4<T: Pixel>(
pub fn estimate_motion_ss2<T: Pixel>(
fi: &FrameInvariants<T>, fs: &FrameState<T>, bsize: BlockSize, ref_idx: usize,
bo: &BlockOffset, pmvs: &[Option<MotionVector>; 3]
bo: &BlockOffset, pmvs: &[Option<MotionVector>; 3], ref_frame: usize
) -> Option<MotionVector> {
if let Some(ref rec) = fi.rec_buffer.frames[ref_idx] {
let blk_w = bsize.width();
let blk_h = bsize.height();
let bo_adj = adjust_bo(bo, fi, blk_w, blk_h);
let bo_adj_h = BlockOffset{x: bo_adj.x >> 1, y: bo_adj.y >> 1};
let po = PlaneOffset {
x: (bo_adj.x as isize) << BLOCK_TO_PLANE_SHIFT >> 1,
y: (bo_adj.y as isize) << BLOCK_TO_PLANE_SHIFT >> 1
......@@ -882,37 +884,69 @@ pub fn estimate_motion_ss2<T: Pixel>(
let range = 16;
let (mvx_min, mvx_max, mvy_min, mvy_max) = get_mv_range(fi.w_in_b, fi.h_in_b, &bo_adj, blk_w, blk_h);
let global_mv = [MotionVector{row: 0, col: 0}; 2];
let frame_mvs = &fs.frame_mvs[ref_frame];
let frame_ref = &fi.rec_buffer.frames[fi.ref_frames[0] as usize];
let mut lowest_cost = std::u64::MAX;
let mut best_mv = MotionVector::default();
// Divide by 4 to account for subsampling, 0.125 is a fudge factor
let lambda = (fi.me_lambda * 256.0 / 4.0 * 0.125) as u32;
for omv in pmvs.iter() {
if let Some(pmv) = omv {
let x_lo = po.x + (((pmv.col as isize / 8 - range).max(mvx_min / 8).min(mvx_max / 8)) >> 1);
let x_hi = po.x + (((pmv.col as isize / 8 + range).max(mvx_min / 8).min(mvx_max / 8)) >> 1);
let y_lo = po.y + (((pmv.row as isize / 8 - range).max(mvy_min / 8).min(mvy_max / 8)) >> 1);
let y_hi = po.y + (((pmv.row as isize / 8 + range).max(mvy_min / 8).min(mvy_max / 8)) >> 1);
full_search(
x_lo,
x_hi,
y_lo,
y_hi,
blk_h >> 1,
blk_w >> 1,
&fs.input_hres,
&rec.input_hres,
&mut best_mv,
&mut lowest_cost,
&po,
1,
fi.sequence.bit_depth,
lambda,
[MotionVector::default(); 2],
fi.allow_high_precision_mv
);
if fi.config.speed_settings.diamond_me {
for omv in pmvs.iter() {
if let Some(pmv) = omv {
let mut predictors = get_subset_predictors::<T>(
&bo_adj_h,
MotionVector{row: pmv.row, col: pmv.col},
fi.w_in_b, fi.h_in_b,
&frame_mvs, frame_ref, 0
);
for predictor in &mut predictors {
predictor.row >>= 1;
predictor.col >>= 1;
}
diamond_me_search(
fi, &po,
&fs.input_hres, &rec.input_hres,
&predictors, fi.sequence.bit_depth,
global_mv, lambda,
mvx_min >> 1, mvx_max >> 1, mvy_min >> 1, mvy_max >> 1,
blk_w >> 1, blk_h >> 1,
&mut best_mv, &mut lowest_cost,
&mut None, 0
);
}
}
} else {
for omv in pmvs.iter() {
if let Some(pmv) = omv {
let x_lo = po.x + (((pmv.col as isize / 8 - range).max(mvx_min / 8).min(mvx_max / 8)) >> 1);
let x_hi = po.x + (((pmv.col as isize / 8 + range).max(mvx_min / 8).min(mvx_max / 8)) >> 1);
let y_lo = po.y + (((pmv.row as isize / 8 - range).max(mvy_min / 8).min(mvy_max / 8)) >> 1);
let y_hi = po.y + (((pmv.row as isize / 8 + range).max(mvy_min / 8).min(mvy_max / 8)) >> 1);
full_search(
x_lo,
x_hi,
y_lo,
y_hi,
blk_h >> 1,
blk_w >> 1,
&fs.input_hres,
&rec.input_hres,
&mut best_mv,
&mut lowest_cost,
&po,
1,
fi.sequence.bit_depth,
lambda,
[MotionVector::default(); 2],
fi.allow_high_precision_mv
);
}
}
}
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment