Commit 33c298d3 authored by Frank Bossen's avatar Frank Bossen Committed by Thomas Daede

Add rate constraint to ME

parent 50967bdf
......@@ -14,6 +14,7 @@ use partition::*;
use plane::*;
use FrameInvariants;
use FrameState;
use rdo::get_lambda_sqrt;
#[cfg(all(target_arch = "x86_64", not(windows), feature = "nasm"))]
mod nasm {
......@@ -147,7 +148,7 @@ fn get_mv_range(fi: &FrameInvariants, bo: &BlockOffset, blk_w: usize, blk_h: usi
pub fn motion_estimation(
fi: &FrameInvariants, fs: &FrameState, bsize: BlockSize, bo: &BlockOffset,
ref_frame: usize, pmv: MotionVector, bit_depth: usize
ref_frame: usize, cmv: MotionVector, bit_depth: usize, pmv: &[MotionVector; 2]
) -> MotionVector {
match fi.rec_buffer.frames[fi.ref_frames[ref_frame - LAST_FRAME] as usize] {
Some(ref rec) => {
......@@ -159,14 +160,17 @@ pub fn motion_estimation(
let blk_w = bsize.width();
let blk_h = bsize.height();
let (mvx_min, mvx_max, mvy_min, mvy_max) = get_mv_range(fi, bo, blk_w, blk_h);
let x_lo = po.x + ((-range + (pmv.col / 8) as isize).max(mvx_min / 8));
let x_hi = po.x + ((range + (pmv.col / 8) as isize).min(mvx_max / 8));
let y_lo = po.y + ((-range + (pmv.row / 8) as isize).max(mvy_min / 8));
let y_hi = po.y + ((range + (pmv.row / 8) as isize).min(mvy_max / 8));
let x_lo = po.x + ((-range + (cmv.col / 8) as isize).max(mvx_min / 8));
let x_hi = po.x + ((range + (cmv.col / 8) as isize).min(mvx_max / 8));
let y_lo = po.y + ((-range + (cmv.row / 8) as isize).max(mvy_min / 8));
let y_hi = po.y + ((range + (cmv.row / 8) as isize).min(mvy_max / 8));
let mut lowest_sad = 128 * 128 * 4096 as u32;
let mut lowest_cost = std::u32::MAX;
let mut best_mv = MotionVector { row: 0, col: 0 };
// 0.5 is a fudge factor
let lambda = (get_lambda_sqrt(fi, bit_depth) * 256.0 * 0.5) as u32;
full_search(
x_lo,
x_hi,
......@@ -177,10 +181,13 @@ pub fn motion_estimation(
&fs.input.planes[0],
&rec.frame.planes[0],
&mut best_mv,
&mut lowest_sad,
&mut lowest_cost,
&po,
2,
bit_depth
bit_depth,
lambda,
pmv,
fi.allow_high_precision_mv
);
let mode = PredictionMode::NEWMV;
......@@ -234,8 +241,13 @@ pub fn motion_estimation(
let sad = get_sad(&plane_org, &plane_ref, blk_h, blk_w, bit_depth);
if sad < lowest_sad {
lowest_sad = sad;
let rate1 = get_mv_rate(cand_mv, pmv[0], fi.allow_high_precision_mv);
let rate2 = get_mv_rate(cand_mv, pmv[1], fi.allow_high_precision_mv);
let rate = rate1.min(rate2 + 1);
let cost = 256 * sad + rate * lambda;
if cost < lowest_cost {
lowest_cost = cost;
best_mv = cand_mv;
}
}
......@@ -252,7 +264,8 @@ pub fn motion_estimation(
fn full_search(
x_lo: isize, x_hi: isize, y_lo: isize, y_hi: isize, blk_h: usize,
blk_w: usize, p_org: &Plane, p_ref: &Plane, best_mv: &mut MotionVector,
lowest_sad: &mut u32, po: &PlaneOffset, step: usize, bit_depth: usize
lowest_cost: &mut u32, po: &PlaneOffset, step: usize, bit_depth: usize,
lambda: u32, pmv: &[MotionVector; 2], allow_high_precision_mv: bool
) {
for y in (y_lo..y_hi).step_by(step) {
for x in (x_lo..x_hi).step_by(step) {
......@@ -261,12 +274,19 @@ fn full_search(
let sad = get_sad(&plane_org, &plane_ref, blk_h, blk_w, bit_depth);
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)
}
let mv = MotionVector {
row: 8 * (y as i16 - po.y as i16),
col: 8 * (x as i16 - po.x as i16)
};
let rate1 = get_mv_rate(mv, pmv[0], allow_high_precision_mv);
let rate2 = get_mv_rate(mv, pmv[1], allow_high_precision_mv);
let rate = rate1.min(rate2 + 1);
let cost = 256 * sad + rate * lambda;
if cost < *lowest_cost {
*lowest_cost = cost;
*best_mv = mv;
}
}
}
......@@ -280,6 +300,19 @@ fn adjust_bo(bo: &BlockOffset, fi: &FrameInvariants, blk_w: usize, blk_h: usize)
}
}
fn get_mv_rate(a: MotionVector, b: MotionVector, allow_high_precision_mv: bool) -> u32 {
fn diff_to_rate(diff: i16, allow_high_precision_mv: bool) -> u32 {
let d = if allow_high_precision_mv { diff } else { diff >> 1 };
if d == 0 {
0
} else {
2 * (16 - d.abs().leading_zeros())
}
}
diff_to_rate(a.row - b.row, allow_high_precision_mv) + diff_to_rate(a.col - b.col, allow_high_precision_mv)
}
pub fn estimate_motion_ss4(
fi: &FrameInvariants, fs: &FrameState, bsize: BlockSize, ref_idx: usize,
bo: &BlockOffset, bit_depth: usize
......@@ -299,9 +332,12 @@ pub fn estimate_motion_ss4(
let y_lo = po.y + (((-range).max(mvy_min / 8)) >> 2);
let y_hi = po.y + (((range).min(mvy_max / 8)) >> 2);
let mut lowest_sad = ((blk_w >> 2) * (blk_h >> 2) * 4096) as u32;
let mut lowest_cost = std::u32::MAX;
let mut best_mv = MotionVector { row: 0, col: 0 };
// Divide by 16 to account for subsampling, 0.125 is a fudge factor
let lambda = (get_lambda_sqrt(fi, bit_depth) * 256.0 / 16.0 * 0.125) as u32;
full_search(
x_lo,
x_hi,
......@@ -312,10 +348,13 @@ pub fn estimate_motion_ss4(
&fs.input_qres,
&rec.input_qres,
&mut best_mv,
&mut lowest_sad,
&mut lowest_cost,
&po,
1,
bit_depth
bit_depth,
lambda,
&[MotionVector { row: 0, col: 0 }; 2],
fi.allow_high_precision_mv
);
Some(MotionVector { row: best_mv.row * 4, col: best_mv.col * 4 })
......@@ -339,9 +378,12 @@ pub fn estimate_motion_ss2(
let range = 16;
let (mvx_min, mvx_max, mvy_min, mvy_max) = get_mv_range(fi, &bo_adj, blk_w, blk_h);
let mut lowest_sad = ((blk_w >> 1) * (blk_h >> 1) * 4096) as u32;
let mut lowest_cost = std::u32::MAX;
let mut best_mv = MotionVector { row: 0, col: 0 };
// Divide by 4 to account for subsampling, 0.125 is a fudge factor
let lambda = (get_lambda_sqrt(fi, bit_depth) * 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)) >> 1);
......@@ -359,10 +401,13 @@ pub fn estimate_motion_ss2(
&fs.input_hres,
&rec.input_hres,
&mut best_mv,
&mut lowest_sad,
&mut lowest_cost,
&po,
1,
bit_depth
bit_depth,
lambda,
&[MotionVector { row: 0, col: 0 }; 2],
fi.allow_high_precision_mv
);
}
}
......
......@@ -150,6 +150,17 @@ pub fn get_lambda(fi: &FrameInvariants, bit_depth: usize) -> f64 {
q0 * q0 * std::f64::consts::LN_2 / 6.0
}
pub fn get_lambda_sqrt(fi: &FrameInvariants, bit_depth: usize) -> f64 {
let q = dc_q(fi.base_q_idx, bit_depth) as f64;
// Convert q into Q0 precision, given that libaom quantizers are Q3
let q0 = q / 8.0_f64;
// Lambda formula from doc/theoretical_results.lyx in the daala repo
// Use Q0 quantizer since lambda will be applied to Q0 pixel domain
q0 * (std::f64::consts::LN_2 / 6.0).sqrt()
}
// Compute the rate-distortion cost for an encode
fn compute_rd_cost(
fi: &FrameInvariants, fs: &FrameState, w_y: usize, h_y: usize,
......@@ -385,11 +396,6 @@ pub fn rdo_mode_decision(
ref_frames_set.push([i, NONE_FRAME]);
let slot_idx = fi.ref_frames[i - LAST_FRAME];
ref_slot_set.push(slot_idx);
let pmv = pmvs[slot_idx as usize].unwrap();
mvs_from_me.push([
motion_estimation(fi, fs, bsize, bo, i, pmv, seq.bit_depth),
MotionVector { row: 0, col: 0 }
]);
}
}
assert!(ref_frames_set.len() != 0);
......@@ -404,6 +410,15 @@ pub fn rdo_mode_decision(
mode_contexts.push(cw.find_mvrefs(bo, ref_frames, &mut mv_stack, bsize, false, fi, false));
if fi.frame_type == FrameType::INTER {
let mut pmv = [MotionVector{ row: 0, col: 0 }; 2];
if mv_stack.len() > 0 { pmv[0] = mv_stack[0].this_mv; }
if mv_stack.len() > 1 { pmv[1] = mv_stack[1].this_mv; }
let cmv = pmvs[ref_slot_set[i] as usize].unwrap();
mvs_from_me.push([
motion_estimation(fi, fs, bsize, bo, ref_frames[0], cmv, seq.bit_depth, &pmv),
MotionVector { row: 0, col: 0 }
]);
for &x in RAV1E_INTER_MODES_MINIMAL {
mode_set.push((x, i));
}
......
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