Unverified Commit bc5a464b authored by fbossen's avatar fbossen Committed by GitHub

Hierarchical motion estimation (#673)

* Generate downsampled versions of input frames and store them with reference frames

* Add hierarchical ME, where initial ME is done using 4x subsampled 64x64 blocks. MVs are refined using 2x subsampled 32x32 blocks. In the refinement step, search is done around MV found for colocated 64x64 block, and also around MVs found for neighboring 64x64 blocks
parent 27b67a5b
......@@ -7,7 +7,6 @@
// 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::CDFContext;
use encoder::*;
use partition::*;
......@@ -291,13 +290,7 @@ impl Context {
if self.fi.show_existing_frame {
self.idx = self.idx + 1;
let mut fs = FrameState {
input: Arc::new(Frame::new(self.fi.padded_w, self.fi.padded_h)), // dummy
rec: Frame::new(self.fi.padded_w, self.fi.padded_h),
qc: Default::default(),
cdfs: CDFContext::new(0),
deblock: Default::default(),
};
let mut fs = FrameState::new(&self.fi);
let data = encode_frame(&mut self.seq, &mut self.fi, &mut fs);
......@@ -310,13 +303,7 @@ impl Context {
self.idx = self.idx + 1;
if let Some(frame) = f {
let mut fs = FrameState {
input: frame,
rec: Frame::new(self.fi.padded_w, self.fi.padded_h),
qc: Default::default(),
cdfs: CDFContext::new(0),
deblock: Default::default(),
};
let mut fs = FrameState::new_with_frame(&self.fi, frame);
let data = encode_frame(&mut self.seq, &mut self.fi, &mut fs);
self.packet_data.extend(data);
......
......@@ -52,7 +52,7 @@ pub const MAX_MIB_SIZE: usize = (1 << MAX_MIB_SIZE_LOG2);
pub const MAX_MIB_MASK: usize = (MAX_MIB_SIZE - 1);
const MAX_SB_SIZE_LOG2: usize = 6;
const MAX_SB_SIZE: usize = (1 << MAX_SB_SIZE_LOG2);
pub const MAX_SB_SIZE: usize = (1 << MAX_SB_SIZE_LOG2);
const MAX_SB_SQUARE: usize = (MAX_SB_SIZE * MAX_SB_SIZE);
pub const MAX_TX_SIZE: usize = 32;
......
......@@ -20,6 +20,7 @@ use rdo::*;
use std::fmt;
use transform::*;
use util::*;
use me::*;
use bitstream_io::{BitWriter, BigEndian, LittleEndian};
use std;
......@@ -37,13 +38,27 @@ pub struct Frame {
pub planes: [Plane; 3]
}
const FRAME_MARGIN: usize = 16 + SUBPEL_FILTER_SIZE;
impl Frame {
pub fn new(width: usize, height:usize) -> Frame {
Frame {
planes: [
Plane::new(width, height, 0, 0, 128+8, 128+8),
Plane::new(width/2, height/2, 1, 1, 64+8, 64+8),
Plane::new(width/2, height/2, 1, 1, 64+8, 64+8)
Plane::new(
width, height,
0, 0,
MAX_SB_SIZE + FRAME_MARGIN, MAX_SB_SIZE + FRAME_MARGIN
),
Plane::new(
width/2, height/2,
1, 1,
MAX_SB_SIZE/2 + FRAME_MARGIN, MAX_SB_SIZE/2 + FRAME_MARGIN
),
Plane::new(
width/2, height/2,
1, 1,
MAX_SB_SIZE/2 + FRAME_MARGIN, MAX_SB_SIZE/2 + FRAME_MARGIN
)
]
}
}
......@@ -69,6 +84,8 @@ impl Frame {
pub struct ReferenceFrame {
pub order_hint: u32,
pub frame: Frame,
pub input_hres: Plane,
pub input_qres: Plane,
pub cdfs: CDFContext
}
......@@ -313,6 +330,8 @@ use std::sync::Arc;
#[derive(Debug)]
pub struct FrameState {
pub input: Arc<Frame>,
pub input_hres: Plane, // half-resolution version of input luma
pub input_qres: Plane, // quarter-resolution version of input luma
pub rec: Frame,
pub qc: QuantizationContext,
pub cdfs: CDFContext,
......@@ -321,8 +340,22 @@ pub struct FrameState {
impl FrameState {
pub fn new(fi: &FrameInvariants) -> FrameState {
FrameState::new_with_frame(fi, Arc::new(Frame::new(fi.padded_w, fi.padded_h)))
}
pub fn new_with_frame(fi: &FrameInvariants, frame: Arc<Frame>) -> FrameState {
FrameState {
input: Arc::new(Frame::new(fi.padded_w, fi.padded_h)),
input: frame,
input_hres: Plane::new(
fi.padded_w/2, fi.padded_h/2,
1, 1,
(MAX_SB_SIZE + FRAME_MARGIN) / 2, (MAX_SB_SIZE + FRAME_MARGIN) / 2
),
input_qres: Plane::new(
fi.padded_w/4, fi.padded_h/4,
2, 2,
(MAX_SB_SIZE + FRAME_MARGIN) / 4, (MAX_SB_SIZE + FRAME_MARGIN) / 4
),
rec: Frame::new(fi.padded_w, fi.padded_h),
qc: Default::default(),
cdfs: CDFContext::new(0),
......@@ -333,6 +366,8 @@ impl FrameState {
pub fn window(&self, sbo: &SuperBlockOffset) -> FrameState {
FrameState {
input: Arc::new(self.input.window(sbo)),
input_hres: self.input_hres.window(&sbo.plane_offset(&self.input_hres.cfg)),
input_qres: self.input_qres.window(&sbo.plane_offset(&self.input_qres.cfg)),
rec: self.rec.window(sbo),
qc: self.qc.clone(),
cdfs: self.cdfs.clone(),
......@@ -499,17 +534,6 @@ impl FrameInvariants {
use_tx_domain_distortion: use_tx_domain_distortion,
}
}
pub fn new_frame_state(&self) -> FrameState {
FrameState {
input: Arc::new(Frame::new(self.padded_w, self.padded_h)),
rec: Frame::new(self.padded_w, self.padded_h),
qc: Default::default(),
cdfs: CDFContext::new(0),
deblock: Default::default(),
}
}
}
impl fmt::Display for FrameInvariants{
......@@ -1795,7 +1819,8 @@ pub fn write_tx_tree(fi: &FrameInvariants, fs: &mut FrameState, cw: &mut Context
fn encode_partition_bottomup(seq: &Sequence, fi: &FrameInvariants, fs: &mut FrameState,
cw: &mut ContextWriter, w_pre_cdef: &mut dyn Writer, w_post_cdef: &mut dyn Writer,
bsize: BlockSize, bo: &BlockOffset, pmv: &MotionVector) -> f64 {
bsize: BlockSize, bo: &BlockOffset, pmvs: &[[Option<MotionVector>; REF_FRAMES]; 5]
) -> f64 {
let mut rd_cost = std::f64::MAX;
if bo.x >= cw.bc.cols || bo.y >= cw.bc.rows {
......@@ -1845,7 +1870,15 @@ fn encode_partition_bottomup(seq: &Sequence, fi: &FrameInvariants, fs: &mut Fram
cw.write_partition(w, bo, partition, bsize);
cost = (w.tell_frac() - tell) as f64 * get_lambda(fi, seq.bit_depth)/ ((1 << OD_BITRES) as f64);
}
let mode_decision = rdo_mode_decision(seq, fi, fs, cw, bsize, bo, pmv).part_modes[0].clone();
let pmv_idx = if bsize > BlockSize::BLOCK_32X32 {
0
} else {
((bo.x & 32) >> 5) + ((bo.y & 32) >> 4) + 1
};
let spmvs = &pmvs[pmv_idx];
let mode_decision = rdo_mode_decision(seq, fi, fs, cw, bsize, bo, spmvs).part_modes[0].clone();
let (mode_luma, mode_chroma) = (mode_decision.pred_mode_luma, mode_decision.pred_mode_chroma);
let cfl = mode_decision.pred_cfl_params;
{
......@@ -1910,7 +1943,7 @@ fn encode_partition_bottomup(seq: &Sequence, fi: &FrameInvariants, fs: &mut Fram
w_post_cdef,
subsize,
offset,
&best_decision.mvs[0]
pmvs//&best_decision.mvs[0]
)
}).sum::<f64>();
......@@ -1966,7 +1999,9 @@ fn encode_partition_bottomup(seq: &Sequence, fi: &FrameInvariants, fs: &mut Fram
fn encode_partition_topdown(seq: &Sequence, fi: &FrameInvariants, fs: &mut FrameState,
cw: &mut ContextWriter, w_pre_cdef: &mut dyn Writer, w_post_cdef: &mut dyn Writer,
bsize: BlockSize, bo: &BlockOffset, block_output: &Option<RDOOutput>) {
bsize: BlockSize, bo: &BlockOffset, block_output: &Option<RDOOutput>,
pmvs: &[[Option<MotionVector>; REF_FRAMES]; 5]
) {
if bo.x >= cw.bc.cols || bo.y >= cw.bc.rows {
return;
......@@ -1991,7 +2026,7 @@ fn encode_partition_topdown(seq: &Sequence, fi: &FrameInvariants, fs: &mut Frame
partition = PartitionType::PARTITION_SPLIT;
} else if bsize > fi.min_partition_size {
// Blocks of sizes within the supported range are subjected to a partitioning decision
rdo_output = rdo_partition_decision(seq, fi, fs, cw, bsize, bo, &rdo_output);
rdo_output = rdo_partition_decision(seq, fi, fs, cw, bsize, bo, &rdo_output, pmvs);
partition = rdo_output.part_type;
} else {
// Blocks of sizes below the supported range are encoded directly
......@@ -2004,7 +2039,6 @@ fn encode_partition_topdown(seq: &Sequence, fi: &FrameInvariants, fs: &mut Frame
let hbs = bs >> 1; // Half the block size in blocks
let subsize = bsize.subsize(partition);
let pmv = MotionVector { row: 0, col: 0 };
if bsize >= BlockSize::BLOCK_8X8 {
let w: &mut dyn Writer = if cw.bc.cdef_coded {w_post_cdef} else {w_pre_cdef};
......@@ -2017,8 +2051,15 @@ fn encode_partition_topdown(seq: &Sequence, fi: &FrameInvariants, fs: &mut Frame
// The optimal prediction mode is known from a previous iteration
rdo_output.part_modes[0].clone()
} else {
let pmv_idx = if bsize > BlockSize::BLOCK_32X32 {
0
} else {
((bo.x & 32) >> 5) + ((bo.y & 32) >> 4) + 1
};
let spmvs = &pmvs[pmv_idx];
// Make a prediction mode decision for blocks encoded with no rdo_partition_decision call (e.g. edges)
rdo_mode_decision(seq, fi, fs, cw, bsize, bo, &pmv).part_modes[0].clone()
rdo_mode_decision(seq, fi, fs, cw, bsize, bo, spmvs).part_modes[0].clone()
};
let mut mode_luma = part_decision.pred_mode_luma;
......@@ -2099,7 +2140,7 @@ fn encode_partition_topdown(seq: &Sequence, fi: &FrameInvariants, fs: &mut Frame
&Some(RDOOutput {
rd_cost: mode.rd_cost,
part_type: PartitionType::PARTITION_NONE,
part_modes: vec![mode] }));
part_modes: vec![mode] }), pmvs);
}
}
else {
......@@ -2119,7 +2160,8 @@ fn encode_partition_topdown(seq: &Sequence, fi: &FrameInvariants, fs: &mut Frame
w_post_cdef,
subsize,
offset,
&None
&None,
pmvs
);
});
}
......@@ -2150,6 +2192,26 @@ fn encode_tile(sequence: &mut Sequence, fi: &FrameInvariants, fs: &mut FrameStat
let rc = RestorationContext::new(fi.sb_width, fi.sb_height);
let mut cw = ContextWriter::new(fc, bc, rc);
// initial coarse ME loop
let mut frame_pmvs = Vec::new();
for sby in 0..fi.sb_height {
for sbx in 0..fi.sb_width {
let sbo = SuperBlockOffset { x: sbx, y: sby };
let bo = sbo.block_offset(0, 0);
let mut pmvs: [Option<MotionVector>; REF_FRAMES] = [None; REF_FRAMES];
for i in 0..INTER_REFS_PER_FRAME {
let r = fi.ref_frames[i] as usize;
if pmvs[r].is_none() {
assert!(!sequence.use_128x128_superblock);
pmvs[r] = estimate_motion_ss4(fi, fs, BlockSize::BLOCK_64X64, r, &bo);
}
}
frame_pmvs.push(pmvs);
}
}
// main loop
for sby in 0..fi.sb_height {
cw.bc.reset_left_contexts();
......@@ -2162,17 +2224,61 @@ fn encode_tile(sequence: &mut Sequence, fi: &FrameInvariants, fs: &mut FrameStat
cw.bc.cdef_coded = false;
cw.bc.code_deltas = fi.delta_q_present;
// Do subsampled ME
let mut pmvs: [[Option<MotionVector>; REF_FRAMES]; 5] = [[None; REF_FRAMES]; 5];
for i in 0..INTER_REFS_PER_FRAME {
let r = fi.ref_frames[i] as usize;
if pmvs[0][r].is_none() {
pmvs[0][r] = frame_pmvs[sby * fi.sb_width + sbx][r];
if let Some(pmv) = pmvs[0][r] {
let pmv_w = if sbx > 0 {
frame_pmvs[sby * fi.sb_width + sbx - 1][r]
} else {
None
};
let pmv_e = if sbx < fi.sb_width - 1 {
frame_pmvs[sby * fi.sb_width + sbx + 1][r]
} else {
None
};
let pmv_n = if sby > 0 {
frame_pmvs[sby * fi.sb_width + sbx - fi.sb_width][r]
} else {
None
};
let pmv_s = if sby < fi.sb_height - 1 {
frame_pmvs[sby * fi.sb_width + sbx + fi.sb_width][r]
} else {
None
};
assert!(!sequence.use_128x128_superblock);
pmvs[1][r] = estimate_motion_ss2(
fi, fs, BlockSize::BLOCK_32X32, r, &sbo.block_offset(0, 0), &[Some(pmv), pmv_w, pmv_n]
);
pmvs[2][r] = estimate_motion_ss2(
fi, fs, BlockSize::BLOCK_32X32, r, &sbo.block_offset(8, 0), &[Some(pmv), pmv_e, pmv_n]
);
pmvs[3][r] = estimate_motion_ss2(
fi, fs, BlockSize::BLOCK_32X32, r, &sbo.block_offset(0, 8), &[Some(pmv), pmv_w, pmv_s]
);
pmvs[4][r] = estimate_motion_ss2(
fi, fs, BlockSize::BLOCK_32X32, r, &sbo.block_offset(8, 8), &[Some(pmv), pmv_e, pmv_s]
);
}
}
}
// Encode SuperBlock
if fi.config.speed == 0 {
let pmv = MotionVector { row: 0, col: 0 };
encode_partition_bottomup(sequence, fi, fs, &mut cw,
&mut w_pre_cdef, &mut w_post_cdef,
BlockSize::BLOCK_64X64, &bo, &pmv);
BlockSize::BLOCK_64X64, &bo, &pmvs);
}
else {
encode_partition_topdown(sequence, fi, fs, &mut cw,
&mut w_pre_cdef, &mut w_post_cdef,
BlockSize::BLOCK_64X64, &bo, &None);
BlockSize::BLOCK_64X64, &bo, &None, &pmvs);
}
// CDEF has to be decisded before loop restoration, but coded after
......@@ -2252,6 +2358,11 @@ pub fn encode_frame(sequence: &mut Sequence, fi: &mut FrameInvariants, fs: &mut
}
}
fs.input_hres.downsample_from(&fs.input.planes[0]);
fs.input_hres.pad();
fs.input_qres.downsample_from(&fs.input_hres);
fs.input_qres.pad();
let bit_depth = sequence.bit_depth;
let tile = encode_tile(sequence, fi, fs, bit_depth); // actually tile group
......@@ -2284,7 +2395,15 @@ pub fn encode_frame(sequence: &mut Sequence, fi: &mut FrameInvariants, fs: &mut
}
pub fn update_rec_buffer(fi: &mut FrameInvariants, fs: FrameState) {
let rfs = Rc::new(ReferenceFrame { order_hint: fi.order_hint, frame: fs.rec, cdfs: fs.cdfs } );
let rfs = Rc::new(
ReferenceFrame {
order_hint: fi.order_hint,
frame: fs.rec,
input_hres: fs.input_hres,
input_qres: fs.input_qres,
cdfs: fs.cdfs
}
);
for i in 0..(REF_FRAMES as usize) {
if (fi.refresh_frame_flags & (1 << i)) != 0 {
fi.rec_buffer.frames[i] = Some(Rc::clone(&rfs));
......
......@@ -35,6 +35,18 @@ pub fn get_sad(
sum
}
fn get_mv_range(fi: &FrameInvariants, bo: &BlockOffset, blk_w: usize, blk_h: usize) -> (isize, isize, isize, isize) {
let border_w = 128 + blk_w as isize * 8;
let border_h = 128 + blk_h as isize * 8;
let mvx_min = -(bo.x as isize) * (8 * MI_SIZE) as isize - border_w;
let mvx_max = (fi.w_in_b - bo.x - blk_w / MI_SIZE) as isize * (8 * MI_SIZE) as isize + border_w;
let mvy_min = -(bo.y as isize) * (8 * MI_SIZE) as isize - border_h;
let mvy_max = (fi.h_in_b - bo.y - blk_h / MI_SIZE) as isize * (8 * MI_SIZE) as isize + border_h;
(mvx_min, mvx_max, mvy_min, mvy_max)
}
pub fn motion_estimation(
fi: &FrameInvariants, fs: &FrameState, bsize: BlockSize,
bo: &BlockOffset, ref_frame: usize, pmv: &MotionVector
......@@ -45,15 +57,10 @@ pub fn motion_estimation(
x: (bo.x as isize) << BLOCK_TO_PLANE_SHIFT,
y: (bo.y as isize) << BLOCK_TO_PLANE_SHIFT
};
let range = 32 * fi.me_range_scale as isize;
let range = 16;
let blk_w = bsize.width();
let blk_h = bsize.height();
let border_w = 128 + blk_w as isize * 8;
let border_h = 128 + blk_h as isize * 8;
let mvx_min = -(bo.x as isize) * (8 * MI_SIZE) as isize - border_w;
let mvx_max = (fi.w_in_b - bo.x - blk_w / MI_SIZE) as isize * (8 * MI_SIZE) as isize + border_w;
let mvy_min = -(bo.y as isize) * (8 * MI_SIZE) as isize - border_h;
let mvy_max = (fi.h_in_b - bo.y - blk_h / MI_SIZE) as isize * (8 * MI_SIZE) as isize + border_h;
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));
......@@ -62,22 +69,10 @@ pub fn motion_estimation(
let mut lowest_sad = 128 * 128 * 4096 as u32;
let mut best_mv = MotionVector { row: 0, col: 0 };
for y in (y_lo..y_hi).step_by(2) {
for x in (x_lo..x_hi).step_by(2) {
let plane_org = fs.input.planes[0].slice(&po);
let plane_ref = rec.frame.planes[0].slice(&PlaneOffset { x, y });
let sad = get_sad(&plane_org, &plane_ref, blk_h, blk_w);
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)
}
}
}
}
full_search(
x_lo, x_hi, y_lo, y_hi, blk_h, blk_w,
&fs.input.planes[0], &rec.frame.planes[0], &mut best_mv, &mut lowest_sad, &po, 2
);
let mode = PredictionMode::NEWMV;
let mut tmp_plane = Plane::new(blk_w, blk_h, 0, 0, 0, 0);
......@@ -138,6 +133,104 @@ 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) {
for y in (y_lo..y_hi).step_by(step) {
for x in (x_lo..x_hi).step_by(step) {
let plane_org = p_org.slice(po);
let plane_ref = p_ref.slice(&PlaneOffset { x, y });
let sad = get_sad(&plane_org, &plane_ref, blk_h, blk_w);
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)
}
}
}
}
}
// Adjust block offset such that entire block lies within frame boundaries
fn adjust_bo(bo: &BlockOffset, fi: &FrameInvariants, blk_w: usize, blk_h: usize) -> BlockOffset {
BlockOffset {
x: (bo.x as isize).min(fi.w_in_b as isize - blk_w as isize / 4).max(0) as usize,
y: (bo.y as isize).min(fi.h_in_b as isize - blk_h as isize / 4).max(0) as usize
}
}
pub fn estimate_motion_ss4(
fi: &FrameInvariants, fs: &FrameState, bsize: BlockSize, ref_idx: usize, bo: &BlockOffset
) -> 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 po = PlaneOffset {
x: (bo_adj.x as isize) << BLOCK_TO_PLANE_SHIFT >> 2,
y: (bo_adj.y as isize) << BLOCK_TO_PLANE_SHIFT >> 2
};
let range = 64 * fi.me_range_scale as isize;
let (mvx_min, mvx_max, mvy_min, mvy_max) = get_mv_range(fi, &bo_adj, blk_w, blk_h);
let x_lo = po.x + (((-range).max(mvx_min / 8)) >> 2);
let x_hi = po.x + (((range).min(mvx_max / 8)) >> 2);
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 best_mv = MotionVector { row: 0, col: 0 };
full_search(
x_lo, x_hi, y_lo, y_hi, blk_h >> 2, blk_w >> 2,
&fs.input_qres, &rec.input_qres, &mut best_mv, &mut lowest_sad, &po, 1
);
Some(MotionVector { row: best_mv.row * 4, col: best_mv.col * 4 })
} else {
None
}
}
pub fn estimate_motion_ss2(
fi: &FrameInvariants, fs: &FrameState, bsize: BlockSize, ref_idx: usize, bo: &BlockOffset, pmvs: &[Option<MotionVector>; 3]
) -> 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 po = PlaneOffset {
x: (bo_adj.x as isize) << BLOCK_TO_PLANE_SHIFT >> 1,
y: (bo_adj.y as isize) << BLOCK_TO_PLANE_SHIFT >> 1
};
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 best_mv = MotionVector { row: 0, col: 0 };
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);
let x_hi = po.x + (((pmv.col as isize / 8 + range).min(mvx_max / 8)) >> 1);
let y_lo = po.y + (((pmv.row as isize / 8 - range).max(mvy_min / 8)) >> 1);
let y_hi = po.y + (((pmv.row as isize / 8 + range).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_sad, &po, 1
);
}
}
Some(MotionVector { row: best_mv.row * 2, col: best_mv.col * 2 })
} else {
None
}
}
#[cfg(test)]
pub mod test {
use super::*;
......
......@@ -689,7 +689,9 @@ pub enum MvSubpelPrecision {
MV_SUBPEL_HIGH_PRECISION
}
const SUBPEL_FILTERS: [[[i32; 8]; 16]; 6] = [
pub const SUBPEL_FILTER_SIZE: usize = 8;
const SUBPEL_FILTERS: [[[i32; SUBPEL_FILTER_SIZE]; 16]; 6] = [
[
[0, 0, 0, 128, 0, 0, 0, 0],
[0, 2, -6, 126, 8, -2, 0, 0],
......
......@@ -218,6 +218,29 @@ impl Plane {
}
}
}
pub fn downsample_from(&mut self, src: &Plane) {
let width = self.cfg.width;
let height = self.cfg.height;
assert!(width * 2 == src.cfg.width);
assert!(height * 2 == src.cfg.height);
for row in 0..height {
let mut dst_slice = self.mut_slice(&PlaneOffset{ x: 0, y: row as isize });
let mut dst = dst_slice.as_mut_slice();
for col in 0..width {
let mut sum = 0;
sum = sum + src.p(2*col, 2*row);
sum = sum + src.p(2*col+1, 2*row);
sum = sum + src.p(2*col, 2*row+1);
sum = sum + src.p(2*col+1, 2*row+1);
let avg = (sum + 2) >> 2;
dst[col] = avg;
}
}
}
}
#[derive(Clone, Copy)]
......
......@@ -340,7 +340,7 @@ impl Default for EncodingSettings {
pub fn rdo_mode_decision(
seq: &Sequence, fi: &FrameInvariants, fs: &mut FrameState,
cw: &mut ContextWriter, bsize: BlockSize, bo: &BlockOffset,
pmv: &MotionVector
pmvs: &[Option<MotionVector>]
) -> RDOOutput {
let mut best = EncodingSettings::default();
......@@ -381,8 +381,10 @@ pub fn rdo_mode_decision(
bwdref = Some(ref_frames_set.len());
}
ref_frames_set.push([i, NONE_FRAME]);
ref_slot_set.push(fi.ref_frames[i - LAST_FRAME]);
mvs_from_me.push([motion_estimation(fi, fs, bsize, bo, i, pmv), MotionVector { row: 0, col: 0 }]);
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), MotionVector { row: 0, col: 0 }]);
}
}
assert!(ref_frames_set.len() != 0);
......@@ -794,7 +796,7 @@ pub fn rdo_tx_type_decision(
pub fn rdo_partition_decision(
seq: &Sequence, fi: &FrameInvariants, fs: &mut FrameState,
cw: &mut ContextWriter, bsize: BlockSize, bo: &BlockOffset,
cached_block: &RDOOutput
cached_block: &RDOOutput, pmvs: &[[Option<MotionVector>; REF_FRAMES]; 5]
) -> RDOOutput {
let max_rd = std::f64::MAX;
......@@ -812,7 +814,6 @@ pub fn rdo_partition_decision(
let mut rd: f64;
let mut child_modes = std::vec::Vec::new();
let mut pmv = MotionVector { row: 0, col: 0 };
match partition {
PartitionType::PARTITION_NONE => {
......@@ -820,11 +821,14 @@ pub fn rdo_partition_decision(
continue;
}
let pmv_idx = ((bo.x & 32) >> 5) + ((bo.y & 32) >> 4) + 1;
let spmvs = &pmvs[pmv_idx];
let mode_decision = cached_block
.part_modes
.get(0)
.unwrap_or(
&rdo_mode_decision(seq, fi, fs, cw, bsize, bo, &pmv).part_modes[0]
&rdo_mode_decision(seq, fi, fs, cw, bsize, bo, spmvs).part_modes[0]
).clone();
child_modes.push(mode_decision);
}
......@@ -834,7 +838,7 @@ pub fn rdo_partition_decision(
if subsize == BlockSize::BLOCK_INVALID {
continue;
}
pmv = best_pred_modes[0].mvs[0];
//pmv = best_pred_modes[0].mvs[0];
assert!(best_pred_modes.len() <= 4);
let bs = bsize.width_mi();
......@@ -845,11 +849,21 @@ pub fn rdo_partition_decision(
&BlockOffset{ x: bo.x, y: bo.y + hbs as usize },
&BlockOffset{ x: bo.x + hbs as usize, y: bo.y + hbs as usize }
];
let pmv_idxs = partitions.iter().map(|&offset| {
if subsize > BlockSize::BLOCK_32X32 {
0
} else {
((offset.x & 32) >> 5) + ((offset.y & 32) >> 4) + 1
}
}).collect::<Vec<_>>();
child_modes.extend(
partitions
.iter()
.map(|&offset| {
rdo_mode_decision(seq, fi, fs, cw, subsize, &offset, &pmv)
.iter().zip(pmv_idxs)
.map(|(&offset, pmv_idx)| {
rdo_mode_decision(seq, fi, fs, cw, subsize, &offset,
&pmvs[pmv_idx])
.part_modes[0]
.clone()
}).collect::<Vec<_>>()
...