Commit 5ecf002b authored by Yushin Cho's avatar Yushin Cho Committed by Thomas Daede

Enable partition : part 1, split down to 8x8 blocks

This is part 1 of ongoing work of enabling partition encoding.
Every SuperBlock is partitioned down to 8x8 partitions,
where pre-existing RDO mode decision decides intra modes
(among three, DC_PRED, HORZ_PRED. VERT_PRED) for the 8x8 partition.

If input image size is not multiple of 8 pixels,
the coded image is padded to be multiple of 8 pixels in both width and height.

TEST: subset1 images.

cargo build --bin rav1e --release
rm test.ivf test_rec.y4m test_dec.y4m
./target/release/rav1e /home/yushin/sequences/subset1-y4m/125_-_Québec_-_Pont_de_Québec_de_nuit_-_Septembre_2009.y4m -o test.ivf  -r test_rec.y4m --quantizer 50
./aom_test/aomdec test.ivf -o test_dec.y4m -v
mpv --keep-open test_dec.y4m &
mpv --keep-open test_rec.y4m &

TODO:
1. Fix partition down to 4x4 blocks, which will include fixing
   has_chroma() function
2. RDO-based block size decision
parent 5ce8ef5c
......@@ -135,7 +135,6 @@ fn write_b_bench(b: &mut Bencher) {
bc: bc,
};
let mode = PredictionMode::DC_PRED;
let tx_type = TxType::DCT_DCT;
let sbx = 0;
......@@ -148,7 +147,9 @@ fn write_b_bench(b: &mut Bencher) {
for by in 0..8 {
for bx in 0..8 {
let bo = sbo.block_offset(bx, by);
write_b(&mut cw, &mut fi, &mut fs, p, &bo, mode, tx_type);
let tx_bo = BlockOffset{x: bo.x + bx, y: bo.y + by};
let po = tx_bo.plane_offset(&fs.input.planes[p].cfg);
encode_tx_block(&mut fi, &mut fs, &mut cw, p, &bo, mode, tx_type, &po);
}
}
}
......
......@@ -13,9 +13,10 @@ fn main() {
Some(rec_file) => Some(y4m::encode(width, height, framerate).write_header(rec_file).unwrap()),
None => None
};
let mut fi = FrameInvariants::new(width, height, files.quantizer);
let sequence = Sequence::new();
write_ivf_header(&mut files.output_file, fi.sb_width*64, fi.sb_height*64, framerate.num, framerate.den);
write_ivf_header(&mut files.output_file, fi.padded_w, fi.padded_h, framerate.num, framerate.den);
let mut last_rec: Option<Frame> = None;
loop {
......
......@@ -19,7 +19,7 @@ fn main() {
};
let mut fi = FrameInvariants::new(width, height, files.quantizer);
let sequence = Sequence::new();
write_ivf_header(&mut files.output_file, fi.sb_width*64, fi.sb_height*64, framerate.num, framerate.den);
write_ivf_header(&mut files.output_file, fi.padded_w, fi.padded_h, framerate.num, framerate.den);
let mut rl = Editor::<()>::new();
let _ = rl.load_history(".rav1e-history");
......
......@@ -12,13 +12,13 @@ use plane::*;
const PLANES: usize = 3;
const PARTITION_PLOFFSET: usize = 4;
const PARTITION_CONTEXTS: usize = 16;
const PARTITION_TYPES: usize = 4;
const PARTITION_CONTEXTS: usize = 20;
pub const PARTITION_TYPES: usize = 4;
const MI_SIZE_LOG2: usize = 2;
pub const MI_SIZE_LOG2: usize = 2;
const MI_SIZE: usize = (1 << MI_SIZE_LOG2);
const MAX_MIB_SIZE_LOG2: usize = (MAX_SB_SIZE_LOG2 - MI_SIZE_LOG2);
const MAX_MIB_SIZE: usize = (1 << MAX_MIB_SIZE_LOG2);
pub const MAX_MIB_SIZE: usize = (1 << MAX_MIB_SIZE_LOG2);
const MAX_MIB_MASK: usize = (MAX_MIB_SIZE - 1);
const MAX_SB_SIZE_LOG2: usize = 6;
......@@ -28,14 +28,20 @@ const MAX_SB_SQUARE: usize = (MAX_SB_SIZE * MAX_SB_SIZE);
const INTRA_MODES: usize = 13;
const UV_INTRA_MODES: usize = 13;
static mi_size_wide: [u8; BLOCK_SIZES_ALL] = [
pub static mi_size_wide: [u8; BLOCK_SIZES_ALL] = [
1, 1, 2, 2, 2, 4, 4, 4, 8, 8, 8, 16, 16, 1, 4, 2, 8, 4, 16];
static mi_size_high: [u8; BLOCK_SIZES_ALL] = [
pub static mi_size_high: [u8; BLOCK_SIZES_ALL] = [
1, 2, 1, 2, 4, 2, 4, 8, 4, 8, 16, 8, 16, 4, 1, 8, 2, 16, 4];
static b_width_log2_lookup: [u8; BLOCK_SIZES_ALL] = [0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 0, 2, 1, 3, 2, 4];
static b_height_log2_lookup: [u8; BLOCK_SIZES_ALL] = [0, 1, 0, 1, 2, 1, 2, 3, 2, 3, 4, 3, 4, 2, 0, 3, 1, 4, 2];
static tx_size_wide_log2: [usize; TX_SIZES_ALL] = [2, 3, 4, 5, 2, 3, 3, 4, 4, 5, 2, 4, 3, 5];
static tx_size_high_log2: [usize; TX_SIZES_ALL] = [2, 3, 4, 5, 3, 2, 4, 3, 5, 4, 4, 2, 5, 3];
pub static b_width_log2_lookup: [u8; BLOCK_SIZES_ALL] = [0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 0, 2, 1, 3, 2, 4];
pub static b_height_log2_lookup: [u8; BLOCK_SIZES_ALL] = [0, 1, 0, 1, 2, 1, 2, 3, 2, 3, 4, 3, 4, 2, 0, 3, 1, 4, 2];
pub static tx_size_wide_log2: [usize; TX_SIZES_ALL] = [2, 3, 4, 5, 2, 3, 3, 4, 4, 5, 2, 4, 3, 5];
pub static tx_size_high_log2: [usize; TX_SIZES_ALL] = [2, 3, 4, 5, 3, 2, 4, 3, 5, 4, 4, 2, 5, 3];
// Width/height lookup tables in units of various block sizes
pub static block_size_wide: [u8; BLOCK_SIZES_ALL] = [
4, 4, 8, 8, 8, 16, 16, 16, 32, 32, 32, 64, 64, 4, 16, 8, 32, 16, 64 ];
pub static block_size_high: [u8; BLOCK_SIZES_ALL] = [
4, 8, 4, 8, 16, 8, 16, 32, 16, 32, 64, 32, 64, 16,4, 32, 8, 64, 16 ];
const EXT_TX_SIZES: usize = 4;
const EXT_TX_SET_TYPES: usize = 6;
......@@ -324,7 +330,7 @@ impl CDFContext {
const SUPERBLOCK_TO_PLANE_SHIFT: usize = MAX_SB_SIZE_LOG2;
const SUPERBLOCK_TO_BLOCK_SHIFT: usize = MAX_MIB_SIZE_LOG2;
const BLOCK_TO_PLANE_SHIFT: usize = MI_SIZE_LOG2;
const LOCAL_BLOCK_MASK: usize = (1 << SUPERBLOCK_TO_BLOCK_SHIFT) - 1;
pub const LOCAL_BLOCK_MASK: usize = (1 << SUPERBLOCK_TO_BLOCK_SHIFT) - 1;
/// Absolute offset in superblocks inside a plane, where a superblock is defined
/// to be an N*N square where N = (1 << SUPERBLOCK_TO_PLANE_SHIFT).
......@@ -386,6 +392,8 @@ impl BlockOffset {
#[derive(Copy,Clone)]
pub struct Block {
pub mode: PredictionMode,
pub bsize: BlockSize,
pub partition: PartitionType,
pub skip: bool,
}
......@@ -393,18 +401,20 @@ impl Block {
pub fn default() -> Block {
Block {
mode: PredictionMode::DC_PRED,
bsize: BlockSize::BLOCK_64X64,
partition: PartitionType::PARTITION_NONE,
skip: false,
}
}
pub fn is_inter(&self) -> bool {
false
self.mode >= PredictionMode::NEARESTMV
}
}
#[derive(Clone, Default)]
pub struct BlockContext {
cols: usize,
rows: usize,
pub cols: usize,
pub rows: usize,
above_partition_context: Vec<u8>,
left_partition_context: [u8; MAX_MIB_SIZE],
above_coeff_context: [Vec<u8>; PLANES],
......@@ -483,6 +493,14 @@ impl BlockContext {
//TODO(anyone): Call reset_left_tx_context() here.
}
pub fn set_mode(&mut self, bo: &BlockOffset, mode: PredictionMode) {
self.blocks[bo.y][bo.x].mode = mode;
}
pub fn get_mode(&mut self, bo: &BlockOffset) -> PredictionMode {
self.blocks[bo.y][bo.x].mode
}
fn partition_plane_context(&self, bo: &BlockOffset,
bsize: BlockSize) -> usize {
// TODO: this should be way simpler without sub8x8
......@@ -539,10 +557,49 @@ pub struct ContextWriter {
}
impl ContextWriter {
pub fn write_partition(&mut self, p: PartitionType, bsize: BlockSize) {
let bo = BlockOffset { x: 0, y: 0 };
fn cdf_element_prob(cdf: &[u16], element: usize) -> u16 {
return if element > 0 { cdf[element - 1] } else { 32768 } - cdf[element];
}
fn partition_gather_horz_alike(out: &mut [u16; 2], cdf_in: &[u16], _bsize: BlockSize) {
out[0] = 32768;
out[0] -= ContextWriter::cdf_element_prob(cdf_in, PartitionType::PARTITION_HORZ as usize);
out[0] -= ContextWriter::cdf_element_prob(cdf_in, PartitionType::PARTITION_SPLIT as usize);
out[0] = 32768 - out[0];
out[1] = 0;
}
fn partition_gather_vert_alike(out: &mut [u16; 2], cdf_in: &[u16], _bsize: BlockSize) {
out[0] = 32768;
out[0] -= ContextWriter::cdf_element_prob(cdf_in, PartitionType::PARTITION_VERT as usize);
out[0] -= ContextWriter::cdf_element_prob(cdf_in, PartitionType::PARTITION_SPLIT as usize);
out[0] = 32768 - out[0];
out[1] = 0;
}
pub fn write_partition(&mut self, bo: &BlockOffset, p: PartitionType, bsize: BlockSize) {
let hbs = (mi_size_wide[bsize as usize] / 2) as usize;
let has_cols = (bo.x + hbs) < self.bc.cols;
let has_rows = (bo.y + hbs) < self.bc.rows;
let ctx = self.bc.partition_plane_context(&bo, bsize);
self.w.symbol(p as u32, &mut self.fc.partition_cdf[ctx], PARTITION_TYPES);
assert!(ctx < PARTITION_CONTEXTS);
let partition_cdf = &mut self.fc.partition_cdf[ctx];
if !has_rows && !has_cols {
return;
}
if has_rows && has_cols {
self.w.symbol(p as u32, partition_cdf, PARTITION_TYPES);
} else if !has_rows && has_cols {
let mut cdf = [0u16; 2];
ContextWriter::partition_gather_vert_alike(&mut cdf, partition_cdf, bsize);
self.w.cdf((p == PartitionType::PARTITION_SPLIT) as u32, &cdf);
} else {
let mut cdf = [0u16; 2];
ContextWriter::partition_gather_horz_alike(&mut cdf, partition_cdf, bsize);
self.w.cdf((p == PartitionType::PARTITION_SPLIT) as u32, &cdf);
}
}
pub fn write_intra_mode_kf(&mut self, bo: &BlockOffset, mode: PredictionMode) {
let above_mode = self.bc.above_of(bo).mode as usize;
......
......@@ -71,8 +71,8 @@ pub struct FrameState {
impl FrameState {
pub fn new(fi: &FrameInvariants) -> FrameState {
FrameState {
input: Frame::new(fi.sb_width*64, fi.sb_height*64),
rec: Frame::new(fi.sb_width*64, fi.sb_height*64),
input: Frame::new(fi.padded_w, fi.padded_h),
rec: Frame::new(fi.padded_w, fi.padded_h),
}
}
}
......@@ -84,8 +84,12 @@ pub struct FrameInvariants {
pub qindex: usize,
pub width: usize,
pub height: usize,
pub padded_w: usize,
pub padded_h: usize,
pub sb_width: usize,
pub sb_height: usize,
pub w_in_b: usize,
pub h_in_b: usize,
pub number: u64,
pub ftype: FrameType,
pub show_existing_frame: bool,
......@@ -97,8 +101,12 @@ impl FrameInvariants {
qindex: qindex,
width: width,
height: height,
padded_w: ((width+7)/8)*8,
padded_h: ((height+7)/8)*8,
sb_width: (width+63)/64,
sb_height: (height+63)/64,
w_in_b: ((width+7)/8)*8 >> MI_SIZE_LOG2,
h_in_b: ((height+7)/8)*8 >> MI_SIZE_LOG2,
number: 0,
ftype: FrameType::KEY,
show_existing_frame: false,
......@@ -228,8 +236,8 @@ fn write_uncompressed_header(packet: &mut Write, sequence: &Sequence, fi: &Frame
//uch.write(8+7,0)?; // frame id
uch.write(3,0)?; // colorspace
uch.write(1,0)?; // color range
uch.write(16,(fi.sb_width*64-1) as u16)?; // width
uch.write(16,(fi.sb_height*64-1) as u16)?; // height
uch.write(16,(fi.padded_w-1) as u16)?; // width
uch.write(16,(fi.padded_h-1) as u16)?; // height
uch.write_bit(false)?; // scaling active
uch.write_bit(false)?; // screen content tools
uch.write(3,0x0)?; // frame context
......@@ -255,7 +263,7 @@ fn write_uncompressed_header(packet: &mut Write, sequence: &Sequence, fi: &Frame
//uch.write_bit(false)?; // use hybrid pred
//uch.write_bit(false)?; // use compound pred
uch.write_bit(true)?; // reduced tx
if fi.sb_width*64-1 > 256 {
if fi.padded_w > 256 {
uch.write(1,0)?; // tile cols
}
uch.write(1,0)?; // tile rows
......@@ -275,14 +283,29 @@ fn diff_4x4(dst: &mut [i16; 16], src1: &PlaneSlice, src2: &PlaneSlice) {
}
}
pub fn write_b(cw: &mut ContextWriter, fi: &FrameInvariants, fs: &mut FrameState, p: usize, bo: &BlockOffset, mode: PredictionMode, tx_type: TxType) {
fn has_chroma(bo: &BlockOffset, bsize: BlockSize,
subsampling_x: usize, subsampling_y: usize) -> bool {
let bw = mi_size_wide[bsize as usize] as u8;
let bh = mi_size_high[bsize as usize] as u8;
let ref_pos = ((bo.x & 0x01) == 1 || (bw & 0x01) == 0 || subsampling_x == 0) &&
((bo.y & 0x01) == 1 || (bh & 0x01) == 0 || subsampling_y == 0);
ref_pos
}
// For a transform block,
// predict, transform, quantize, write coefficients to a bitstream,
// dequantize, inverse-transform.
pub fn encode_tx_block(fi: &FrameInvariants, fs: &mut FrameState, cw: &mut ContextWriter,
p: usize, bo: &BlockOffset, mode: PredictionMode, tx_type: TxType,
po: &PlaneOffset) {
let stride = fs.input.planes[p].cfg.stride;
let rec = &mut fs.rec.planes[p];
let po = bo.plane_offset(&fs.input.planes[p].cfg);
mode.predict_4x4(&mut rec.mut_slice(&po));
let mut residual = [0 as i16; 16];
diff_4x4(&mut residual,
&fs.input.planes[p].slice(&po),
&rec.slice(&po));
......@@ -299,92 +322,181 @@ pub fn write_b(cw: &mut ContextWriter, fi: &FrameInvariants, fs: &mut FrameState
iht4x4_add(&mut rcoeffs, &mut rec.mut_slice(&po).as_mut_slice(), stride, tx_type);
}
fn write_sb(cw: &mut ContextWriter, fi: &FrameInvariants, fs: &mut FrameState,
sbo: &SuperBlockOffset, mode: PredictionMode, bsize: BlockSize) {
cw.write_partition(PartitionType::PARTITION_NONE, bsize);
fn encode_block(fi: &FrameInvariants, fs: &mut FrameState, cw: &mut ContextWriter,
mode: PredictionMode, bsize: BlockSize, bo: &BlockOffset) {
cw.write_skip(bo, false);
cw.write_intra_mode_kf(bo, mode);
// TODO(yushin): If partition type is PARTITION_SPLIT, recursively call write_sb here.
// Otherwise, call write_b for each parition
let xdec = fs.input.planes[1].cfg.xdec;
let ydec = fs.input.planes[1].cfg.ydec;
// TODO(yushin): Factor out new function which handles four different types of a partition.
// The partition offset is represented using a BlockOffset
let po = sbo.block_offset(0, 0);
cw.write_skip(&po, false);
cw.write_intra_mode_kf(&po, mode);
let uv_mode = mode;
cw.write_intra_uv_mode(uv_mode, mode);
if has_chroma(bo, bsize, xdec, ydec) {
cw.write_intra_uv_mode(uv_mode, mode);
}
let tx_type = TxType::DCT_DCT;
cw.write_tx_type(tx_type, mode);
let bw = mi_size_wide[bsize as usize] as usize;
let bh = mi_size_high[bsize as usize] as usize;
// FIXME(you): Loop for TX blocks. For now, fixed as a 4x4 TX only,
// but consider factor out as write_tx_blocks()
for p in 0..1 {
for by in 0..16 {
for bx in 0..16 {
let bo = sbo.block_offset(bx, by);
cw.bc.at(&bo).mode = mode;
write_b(cw, fi, fs, p, &bo, mode, tx_type);
for by in 0..bh {
for bx in 0..bw {
let tx_bo = BlockOffset{x: bo.x + bx, y: bo.y + by};
let po = tx_bo.plane_offset(&fs.input.planes[p].cfg);
// FIXME(anyone): Another way to do this is
// to set mode for each MI block after mode decision is done.
// It works now because we set mode for each tx position, where tx_size is 4x4.
cw.bc.set_mode(&tx_bo, mode);
encode_tx_block(fi, fs, cw, p, &tx_bo, mode, tx_type, &po);
}
}
}
let uv_tx_type = exported_intra_mode_to_tx_type_context[uv_mode as usize];
for p in 1..3 {
for by in 0..8 {
for bx in 0..8 {
let bo = sbo.block_offset(bx, by);
write_b(cw, fi, fs, p, &bo, uv_mode, uv_tx_type);
if has_chroma(bo, bsize, xdec, ydec) {
let uv_tx_type = exported_intra_mode_to_tx_type_context[uv_mode as usize];
let partition_x = (bo.x & LOCAL_BLOCK_MASK) >> xdec << MI_SIZE_LOG2;
let partition_y = (bo.y & LOCAL_BLOCK_MASK) >> ydec << MI_SIZE_LOG2;
for p in 1..3 {
let sb_offset = bo.sb_offset().plane_offset(&fs.input.planes[p].cfg);
for by in 0..bh >> ydec {
for bx in 0..bw >> xdec {
let tx_bo = BlockOffset{x: bo.x + (bx << xdec), y: bo.y + (by << ydec)};
let po = PlaneOffset {
x: sb_offset.x + partition_x + (bx << MI_SIZE_LOG2),
y: sb_offset.y + partition_y + (by << MI_SIZE_LOG2) };
encode_tx_block(fi, fs, cw, p, &tx_bo, uv_mode, uv_tx_type, &po);
}
}
}
}
}
// RDO-based mode decision
fn rdo_mode_decision(fi: &FrameInvariants, fs: &mut FrameState,
cw: &mut ContextWriter,
bsize: BlockSize, bo: &BlockOffset) -> RDOOutput {
let q = dc_q(fi.qindex) as f64;
let q0 = q / 8.0_f64; // Convert q into Q0 precision, given thatn libaom quantizers are Q3.
// Lambda formula from doc/theoretical_results.lyx in the daala repo
let lambda = q0*q0*2.0_f64.log2()/6.0; // Use Q0 quantizer since lambda will be applied to Q0 pixel domain
let mut best_mode = PredictionMode::DC_PRED;
let mut best_rd = std::f64::MAX;
let tell = cw.w.tell_frac();
let w = block_size_wide[bsize as usize];
let h = block_size_high[bsize as usize];
for &mode in RAV1E_INTRA_MODES {
let checkpoint = cw.checkpoint();
encode_block(fi, fs, cw, mode, bsize, bo);
let po = bo.plane_offset(&fs.input.planes[0].cfg);
let d = sse_wxh(&fs.input.planes[0].slice(&po), &fs.rec.planes[0].slice(&po),
w as usize, h as usize);
let r = ((cw.w.tell_frac() - tell) as f64)/8.0;
let rd = (d as f64) + lambda*r;
if rd < best_rd {
best_rd = rd;
best_mode = mode;
}
cw.rollback(checkpoint.clone());
}
// Update partition context
let subsize = get_subsize(bsize, PartitionType::PARTITION_NONE);
assert!(best_rd as i64 >= 0);
cw.bc.update_partition_context(&po, subsize, bsize);
let rdo_output = RDOOutput { rd_cost: best_rd as u64,
pred_mode: best_mode};
rdo_output
}
fn encode_partition(fi: &FrameInvariants, fs: &mut FrameState, cw: &mut ContextWriter,
bsize: BlockSize, bo: &BlockOffset) {
if bo.x >= cw.bc.cols || bo.y >= cw.bc.rows {
return;
}
// TODO(anyone): Until we have RDO-based block size decision,
// split all the way down to 8x8 blocks, then do rdo_mode_decision() for each 8x8 block.
// FIXIME(anyone): Enable partition into 4x4 blocks.
let partition = if bsize > BlockSize::BLOCK_8X8 {
PartitionType::PARTITION_SPLIT }
else { PartitionType::PARTITION_NONE };
assert!(mi_size_wide[bsize as usize] == mi_size_high[bsize as usize]);
assert!(PartitionType::PARTITION_NONE <= partition &&
partition < PartitionType::PARTITION_INVALID);
let bs = mi_size_wide[bsize as usize];
let hbs = bs >> 1; // Half the block size in blocks
let subsize = get_subsize(bsize, partition);
if bsize >= BlockSize::BLOCK_8X8 {
cw.write_partition(bo, partition, bsize);
}
match partition {
PartitionType::PARTITION_NONE => {
// TODO(anyone): Until we have RDO-based block size decision,
// call rdo_mode_decision() for a partition.
let rdo_none = rdo_mode_decision(fi, fs, cw, bsize, bo);
// FIXME(anyone): Instead of calling set_mode() in encode_block() for each 4x4tx position,
// it would be better to call set_mode() for each MI block position here.
cw.bc.set_mode(bo, rdo_none.pred_mode);
encode_block(fi, fs, cw, rdo_none.pred_mode, bsize, bo);
},
PartitionType::PARTITION_SPLIT => {
assert!(subsize != BlockSize::BLOCK_INVALID);
encode_partition(fi, fs, cw, subsize, bo);
encode_partition(fi, fs, cw, subsize, &BlockOffset{x: bo.x + hbs as usize, y: bo.y});
encode_partition(fi, fs, cw, subsize, &BlockOffset{x: bo.x, y: bo.y + hbs as usize});
encode_partition(fi, fs, cw, subsize, &BlockOffset{x: bo.x + hbs as usize, y: bo.y + hbs as usize});
},
_ => { assert!(false); },
}
if bsize >= BlockSize::BLOCK_8X8 &&
(bsize == BlockSize::BLOCK_8X8 || partition != PartitionType::PARTITION_SPLIT) {
cw.bc.update_partition_context(bo, subsize, bsize);
}
}
fn encode_tile(fi: &FrameInvariants, fs: &mut FrameState) -> Vec<u8> {
let w = ec::Writer::new();
let fc = CDFContext::new();
let bc = BlockContext::new(fi.sb_width*16, fi.sb_height*16);
let bc = BlockContext::new(fi.w_in_b, fi.h_in_b);
let mut cw = ContextWriter {
w: w,
fc: fc,
bc: bc,
};
let q = dc_q(fi.qindex) as f64;
let q0 = q / 8.0_f64; // Convert q into Q0 precision, given thatn libaom quantizers are Q3.
// Lambda formula from doc/theoretical_results.lyx in the daala repo
let lambda = q0*q0*2.0_f64.log2()/6.0; // Use Q0 quantizer since lambda will be applied to Q0 pixel domain
for sby in 0..fi.sb_height {
cw.bc.reset_left_contexts();
for sbx in 0..fi.sb_width {
let sbo = SuperBlockOffset { x: sbx, y: sby };
let po = sbo.plane_offset(&fs.input.planes[0].cfg);
let tell = cw.w.tell_frac();
let bo = sbo.block_offset(0, 0);
let mut best_mode = PredictionMode::DC_PRED;
let mut best_rd = std::f64::MAX;
for &mode in RAV1E_INTRA_MODES {
let checkpoint = cw.checkpoint();
write_sb(&mut cw, fi, fs, &sbo, mode, BlockSize::BLOCK_64X64);
let d = sse_64x64(&fs.input.planes[0].slice(&po), &fs.rec.planes[0].slice(&po));
let r = ((cw.w.tell_frac() - tell) as f64)/8.0;
let rd = (d as f64) + lambda*r;
if rd < best_rd {
best_rd = rd;
best_mode = mode;
}
cw.rollback(checkpoint.clone());
}
// TODO(anyone):
// RDO-based block size decision for each SuperBlock can be done here.
write_sb(&mut cw, fi, fs, &sbo, best_mode, BlockSize::BLOCK_64X64);
// Encode SuperBlock
encode_partition(fi, fs, &mut cw, BlockSize::BLOCK_64X64, &bo);
}
}
let mut h = cw.w.done();
......
#![allow(non_camel_case_types)]
#![allow(dead_code)]
#[derive(Copy,Clone,PartialEq)]
#[derive(Copy,Clone,PartialEq,PartialOrd)]
pub enum PartitionType {
PARTITION_NONE,
PARTITION_HORZ,
PARTITION_VERT,
PARTITION_SPLIT,
PARTITION_INVALID = 255
PARTITION_INVALID
}
pub const BLOCK_SIZES_ALL: usize = 19;
#[derive(Copy,Clone)]
#[derive(Copy,Clone,PartialEq,PartialOrd)]
pub enum BlockSize {
BLOCK_4X4,
BLOCK_4X8,
......@@ -33,7 +33,7 @@ pub enum BlockSize {
BLOCK_32X8,
BLOCK_16X64,
BLOCK_64X16,
BLOCK_INVALID = 255
BLOCK_INVALID
}
pub const TX_SIZES: usize = 4;
......@@ -80,7 +80,7 @@ pub enum TxType {
H_FLIPADST = 15,
}
#[derive(Copy, Clone, Debug, PartialEq)]
#[derive(Copy, Clone, Debug, PartialEq, PartialOrd)]
pub enum PredictionMode {
DC_PRED, // Average of above and left pixels
V_PRED, // Vertical
......
......@@ -11,12 +11,20 @@
#![allow(non_camel_case_types)]
use plane::*;
use partition::PredictionMode;
//use std::io::prelude::*;
// Sum of Squared Error for a 64x64 block
pub fn sse_64x64(src1: &PlaneSlice, src2: &PlaneSlice) -> u64 {
#[derive(Copy,Clone)]
pub struct RDOOutput {
pub rd_cost: u64,
pub pred_mode: PredictionMode,
}
// Sum of Squared Error for a wxh block
pub fn sse_wxh(src1: &PlaneSlice, src2: &PlaneSlice, w: usize, h: usize) -> u64 {
let mut sse: u64 = 0;
for j in 0..64 {
for i in 0..64 {
for j in 0..h {
for i in 0..w {
let dist = (src1.p(i, j) as i16 - src2.p(i, j) as i16) as i64;
sse += (dist * dist) as u64;
}
......
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