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;
......
This diff is collapsed.
#![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