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) { ...@@ -135,7 +135,6 @@ fn write_b_bench(b: &mut Bencher) {
bc: bc, bc: bc,
}; };
let mode = PredictionMode::DC_PRED;
let tx_type = TxType::DCT_DCT; let tx_type = TxType::DCT_DCT;
let sbx = 0; let sbx = 0;
...@@ -148,7 +147,9 @@ fn write_b_bench(b: &mut Bencher) { ...@@ -148,7 +147,9 @@ fn write_b_bench(b: &mut Bencher) {
for by in 0..8 { for by in 0..8 {
for bx in 0..8 { for bx in 0..8 {
let bo = sbo.block_offset(bx, by); 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() { ...@@ -13,9 +13,10 @@ fn main() {
Some(rec_file) => Some(y4m::encode(width, height, framerate).write_header(rec_file).unwrap()), Some(rec_file) => Some(y4m::encode(width, height, framerate).write_header(rec_file).unwrap()),
None => None None => None
}; };
let mut fi = FrameInvariants::new(width, height, files.quantizer); let mut fi = FrameInvariants::new(width, height, files.quantizer);
let sequence = Sequence::new(); 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; let mut last_rec: Option<Frame> = None;
loop { loop {
......
...@@ -19,7 +19,7 @@ fn main() { ...@@ -19,7 +19,7 @@ fn main() {
}; };
let mut fi = FrameInvariants::new(width, height, files.quantizer); let mut fi = FrameInvariants::new(width, height, files.quantizer);
let sequence = Sequence::new(); 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 mut rl = Editor::<()>::new();
let _ = rl.load_history(".rav1e-history"); let _ = rl.load_history(".rav1e-history");
......
...@@ -12,13 +12,13 @@ use plane::*; ...@@ -12,13 +12,13 @@ use plane::*;
const PLANES: usize = 3; const PLANES: usize = 3;
const PARTITION_PLOFFSET: usize = 4; const PARTITION_PLOFFSET: usize = 4;
const PARTITION_CONTEXTS: usize = 16; const PARTITION_CONTEXTS: usize = 20;
const PARTITION_TYPES: usize = 4; 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 MI_SIZE: usize = (1 << MI_SIZE_LOG2);
const MAX_MIB_SIZE_LOG2: usize = (MAX_SB_SIZE_LOG2 - 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_MIB_MASK: usize = (MAX_MIB_SIZE - 1);
const MAX_SB_SIZE_LOG2: usize = 6; const MAX_SB_SIZE_LOG2: usize = 6;
...@@ -28,14 +28,20 @@ const MAX_SB_SQUARE: usize = (MAX_SB_SIZE * MAX_SB_SIZE); ...@@ -28,14 +28,20 @@ const MAX_SB_SQUARE: usize = (MAX_SB_SIZE * MAX_SB_SIZE);
const INTRA_MODES: usize = 13; const INTRA_MODES: usize = 13;
const UV_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]; 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]; 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]; 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];
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 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]; pub 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 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_SIZES: usize = 4;
const EXT_TX_SET_TYPES: usize = 6; const EXT_TX_SET_TYPES: usize = 6;
...@@ -324,7 +330,7 @@ impl CDFContext { ...@@ -324,7 +330,7 @@ impl CDFContext {
const SUPERBLOCK_TO_PLANE_SHIFT: usize = MAX_SB_SIZE_LOG2; const SUPERBLOCK_TO_PLANE_SHIFT: usize = MAX_SB_SIZE_LOG2;
const SUPERBLOCK_TO_BLOCK_SHIFT: usize = MAX_MIB_SIZE_LOG2; const SUPERBLOCK_TO_BLOCK_SHIFT: usize = MAX_MIB_SIZE_LOG2;
const BLOCK_TO_PLANE_SHIFT: usize = MI_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 /// 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). /// to be an N*N square where N = (1 << SUPERBLOCK_TO_PLANE_SHIFT).
...@@ -386,6 +392,8 @@ impl BlockOffset { ...@@ -386,6 +392,8 @@ impl BlockOffset {
#[derive(Copy,Clone)] #[derive(Copy,Clone)]
pub struct Block { pub struct Block {
pub mode: PredictionMode, pub mode: PredictionMode,
pub bsize: BlockSize,
pub partition: PartitionType,
pub skip: bool, pub skip: bool,
} }
...@@ -393,18 +401,20 @@ impl Block { ...@@ -393,18 +401,20 @@ impl Block {
pub fn default() -> Block { pub fn default() -> Block {
Block { Block {
mode: PredictionMode::DC_PRED, mode: PredictionMode::DC_PRED,
bsize: BlockSize::BLOCK_64X64,
partition: PartitionType::PARTITION_NONE,
skip: false, skip: false,
} }
} }
pub fn is_inter(&self) -> bool { pub fn is_inter(&self) -> bool {
false self.mode >= PredictionMode::NEARESTMV
} }
} }
#[derive(Clone, Default)] #[derive(Clone, Default)]
pub struct BlockContext { pub struct BlockContext {
cols: usize, pub cols: usize,
rows: usize, pub rows: usize,
above_partition_context: Vec<u8>, above_partition_context: Vec<u8>,
left_partition_context: [u8; MAX_MIB_SIZE], left_partition_context: [u8; MAX_MIB_SIZE],
above_coeff_context: [Vec<u8>; PLANES], above_coeff_context: [Vec<u8>; PLANES],
...@@ -483,6 +493,14 @@ impl BlockContext { ...@@ -483,6 +493,14 @@ impl BlockContext {
//TODO(anyone): Call reset_left_tx_context() here. //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, fn partition_plane_context(&self, bo: &BlockOffset,
bsize: BlockSize) -> usize { bsize: BlockSize) -> usize {
// TODO: this should be way simpler without sub8x8 // TODO: this should be way simpler without sub8x8
...@@ -539,10 +557,49 @@ pub struct ContextWriter { ...@@ -539,10 +557,49 @@ pub struct ContextWriter {
} }
impl ContextWriter { impl ContextWriter {
pub fn write_partition(&mut self, p: PartitionType, bsize: BlockSize) { fn cdf_element_prob(cdf: &[u16], element: usize) -> u16 {
let bo = BlockOffset { x: 0, y: 0 }; 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); 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) { pub fn write_intra_mode_kf(&mut self, bo: &BlockOffset, mode: PredictionMode) {
let above_mode = self.bc.above_of(bo).mode as usize; let above_mode = self.bc.above_of(bo).mode as usize;
......
This diff is collapsed.
#![allow(non_camel_case_types)] #![allow(non_camel_case_types)]
#![allow(dead_code)] #![allow(dead_code)]
#[derive(Copy,Clone,PartialEq)] #[derive(Copy,Clone,PartialEq,PartialOrd)]
pub enum PartitionType { pub enum PartitionType {
PARTITION_NONE, PARTITION_NONE,
PARTITION_HORZ, PARTITION_HORZ,
PARTITION_VERT, PARTITION_VERT,
PARTITION_SPLIT, PARTITION_SPLIT,
PARTITION_INVALID = 255 PARTITION_INVALID
} }
pub const BLOCK_SIZES_ALL: usize = 19; pub const BLOCK_SIZES_ALL: usize = 19;
#[derive(Copy,Clone)] #[derive(Copy,Clone,PartialEq,PartialOrd)]
pub enum BlockSize { pub enum BlockSize {
BLOCK_4X4, BLOCK_4X4,
BLOCK_4X8, BLOCK_4X8,
...@@ -33,7 +33,7 @@ pub enum BlockSize { ...@@ -33,7 +33,7 @@ pub enum BlockSize {
BLOCK_32X8, BLOCK_32X8,
BLOCK_16X64, BLOCK_16X64,
BLOCK_64X16, BLOCK_64X16,
BLOCK_INVALID = 255 BLOCK_INVALID
} }
pub const TX_SIZES: usize = 4; pub const TX_SIZES: usize = 4;
...@@ -80,7 +80,7 @@ pub enum TxType { ...@@ -80,7 +80,7 @@ pub enum TxType {
H_FLIPADST = 15, H_FLIPADST = 15,
} }
#[derive(Copy, Clone, Debug, PartialEq)] #[derive(Copy, Clone, Debug, PartialEq, PartialOrd)]
pub enum PredictionMode { pub enum PredictionMode {
DC_PRED, // Average of above and left pixels DC_PRED, // Average of above and left pixels
V_PRED, // Vertical V_PRED, // Vertical
......
...@@ -11,12 +11,20 @@ ...@@ -11,12 +11,20 @@
#![allow(non_camel_case_types)] #![allow(non_camel_case_types)]
use plane::*; use plane::*;
use partition::PredictionMode;
//use std::io::prelude::*;
// Sum of Squared Error for a 64x64 block #[derive(Copy,Clone)]
pub fn sse_64x64(src1: &PlaneSlice, src2: &PlaneSlice) -> u64 { 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; let mut sse: u64 = 0;
for j in 0..64 { for j in 0..h {
for i in 0..64 { for i in 0..w {
let dist = (src1.p(i, j) as i16 - src2.p(i, j) as i16) as i64; let dist = (src1.p(i, j) as i16 - src2.p(i, j) as i16) as i64;
sse += (dist * dist) as u64; 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