Commit 6503e40c authored by Raphael Zumer's avatar Raphael Zumer Committed by Raphaël Zumer

Map U/V transform size in partition.rs

Support all chroma sampling formats.
parent 678afd62
......@@ -139,7 +139,7 @@ fn cfl_rdo_bench(b: &mut Bencher, bsize: BlockSize) {
let fi = FrameInvariants::new(1024, 1024, config);
let mut fs = FrameState::new(&fi);
let offset = BlockOffset { x: 1, y: 1 };
b.iter(|| rdo_cfl_alpha(&mut fs, &offset, bsize, 8))
b.iter(|| rdo_cfl_alpha(&mut fs, &offset, bsize, 8, Default::default()))
}
criterion_group!(intra_prediction, predict::pred_bench,);
......
......@@ -1940,9 +1940,9 @@ pub fn encode_block_b(seq: &Sequence, fi: &FrameInvariants, fs: &mut FrameState,
motion_compensate(fi, fs, cw, luma_mode, ref_frames, mvs, bsize, bo, bit_depth, false);
if is_inter {
write_tx_tree(fi, fs, cw, w, luma_mode, bo, bsize, tx_size, tx_type, skip, bit_depth, false, for_rdo_use)
write_tx_tree(fi, fs, cw, w, luma_mode, bo, bsize, tx_size, tx_type, skip, bit_depth, seq.chroma_sampling, false, for_rdo_use)
} else {
write_tx_blocks(fi, fs, cw, w, luma_mode, chroma_mode, bo, bsize, tx_size, tx_type, skip, bit_depth, cfl, false, for_rdo_use)
write_tx_blocks(fi, fs, cw, w, luma_mode, chroma_mode, bo, bsize, tx_size, tx_type, skip, bit_depth, seq.chroma_sampling, cfl, false, for_rdo_use)
}
}
......@@ -1986,7 +1986,8 @@ pub fn write_tx_blocks(fi: &FrameInvariants, fs: &mut FrameState,
cw: &mut ContextWriter, w: &mut dyn Writer,
luma_mode: PredictionMode, chroma_mode: PredictionMode, bo: &BlockOffset,
bsize: BlockSize, tx_size: TxSize, tx_type: TxType, skip: bool, bit_depth: usize,
cfl: CFLParams, luma_only: bool, for_rdo_use: bool) -> i64 {
chroma_sampling: ChromaSampling, cfl: CFLParams, luma_only: bool,
for_rdo_use: bool) -> i64 {
let bw = bsize.width_mi() / tx_size.width_mi();
let bh = bsize.height_mi() / tx_size.height_mi();
let qidx = get_qidx(fi, fs, cw, bo);
......@@ -2018,13 +2019,7 @@ pub fn write_tx_blocks(fi: &FrameInvariants, fs: &mut FrameState,
if luma_only { return tx_dist };
// TODO: these are only valid for 4:2:0
let uv_tx_size = match bsize {
BlockSize::BLOCK_4X4 | BlockSize::BLOCK_8X8 => TxSize::TX_4X4,
BlockSize::BLOCK_16X16 => TxSize::TX_8X8,
BlockSize::BLOCK_32X32 => TxSize::TX_16X16,
_ => TxSize::TX_32X32
};
let uv_tx_size = bsize.largest_uv_tx_size(chroma_sampling);
let mut bw_uv = (bw * tx_size.width_mi()) >> xdec;
let mut bh_uv = (bh * tx_size.height_mi()) >> ydec;
......@@ -2084,7 +2079,7 @@ pub fn write_tx_blocks(fi: &FrameInvariants, fs: &mut FrameState,
pub fn write_tx_tree(fi: &FrameInvariants, fs: &mut FrameState, cw: &mut ContextWriter, w: &mut dyn Writer,
luma_mode: PredictionMode, bo: &BlockOffset,
bsize: BlockSize, tx_size: TxSize, tx_type: TxType, skip: bool, bit_depth: usize,
luma_only: bool, for_rdo_use: bool) -> i64 {
chroma_sampling: ChromaSampling, luma_only: bool, for_rdo_use: bool) -> i64 {
let bw = bsize.width_mi() / tx_size.width_mi();
let bh = bsize.height_mi() / tx_size.height_mi();
let qidx = get_qidx(fi, fs, cw, bo);
......@@ -2105,13 +2100,7 @@ pub fn write_tx_tree(fi: &FrameInvariants, fs: &mut FrameState, cw: &mut Context
if luma_only { return tx_dist };
// TODO: these are only valid for 4:2:0
let uv_tx_size = match bsize {
BlockSize::BLOCK_4X4 | BlockSize::BLOCK_8X8 => TxSize::TX_4X4,
BlockSize::BLOCK_16X16 => TxSize::TX_8X8,
BlockSize::BLOCK_32X32 => TxSize::TX_16X16,
_ => TxSize::TX_32X32
};
let uv_tx_size = bsize.largest_uv_tx_size(chroma_sampling);
let mut bw_uv = (bw * tx_size.width_mi()) >> xdec;
let mut bh_uv = (bh * tx_size.height_mi()) >> ydec;
......
......@@ -12,7 +12,7 @@
use self::BlockSize::*;
use self::TxSize::*;
use encoder::FrameInvariants;
use encoder::{ChromaSampling, FrameInvariants};
pub const NONE_FRAME: usize = 8;
pub const INTRA_FRAME: usize = 0;
......@@ -133,6 +133,73 @@ impl BlockSize {
self.height() >> MI_SIZE_LOG2
}
pub fn tx_size(self) -> TxSize {
match self {
BLOCK_4X4 => TX_4X4,
BLOCK_4X8 => TX_4X8,
BLOCK_8X4 => TX_8X4,
BLOCK_8X8 => TX_8X8,
BLOCK_8X16 => TX_8X16,
BLOCK_16X8 => TX_16X8,
BLOCK_16X16 => TX_16X16,
BLOCK_16X32 => TX_16X32,
BLOCK_32X16 => TX_32X16,
BLOCK_32X32 => TX_32X32,
BLOCK_32X64 => TX_32X64,
BLOCK_64X32 => TX_64X32,
BLOCK_4X16 => TX_4X16,
BLOCK_16X4 => TX_16X4,
BLOCK_8X32 => TX_8X32,
BLOCK_32X8 => TX_32X8,
BLOCK_16X64 => TX_16X64,
BLOCK_64X16 => TX_64X16,
BLOCK_INVALID => unreachable!(),
_ => TX_64X64
}
}
pub fn largest_uv_tx_size(self, chroma_sampling: ChromaSampling) -> TxSize {
match chroma_sampling {
ChromaSampling::Cs444 => match self {
BLOCK_64X64 | BLOCK_64X32 | BLOCK_32X64 |
BLOCK_128X128 | BLOCK_128X64 | BLOCK_64X128 => TX_32X32,
BLOCK_64X16 => TX_32X16,
BLOCK_16X64 => TX_16X32,
_ => self.tx_size()
},
ChromaSampling::Cs422 => match self {
BLOCK_4X4 | BLOCK_8X4 => TX_4X4,
BLOCK_8X8 => TX_4X8,
BLOCK_16X8 => TX_8X8,
BLOCK_16X16 => TX_8X16,
BLOCK_32X16 => TX_16X16,
BLOCK_32X32 => TX_16X32,
BLOCK_64X32 | BLOCK_64X64 |
BLOCK_128X64 | BLOCK_128X128 => TX_32X32,
BLOCK_16X4 => TX_8X4,
BLOCK_32X8 => TX_16X8,
BLOCK_64X16 => TX_32X16,
_ => unreachable!() // vertical splits are illegal in 4:2:2
},
ChromaSampling::Cs420 => match self {
BLOCK_4X4 | BLOCK_8X4 | BLOCK_4X8 | BLOCK_8X8 => TX_4X4,
BLOCK_8X16 | BLOCK_4X16 => TX_4X8,
BLOCK_16X8 | BLOCK_16X4 => TX_8X4,
BLOCK_16X16 => TX_8X8,
BLOCK_16X32 => TX_8X16,
BLOCK_32X16 => TX_16X8,
BLOCK_32X32 => TX_16X16,
BLOCK_32X64 => TX_16X32,
BLOCK_64X32 => TX_32X16,
BLOCK_8X32 => TX_4X16,
BLOCK_32X8 => TX_16X4,
BLOCK_16X64 => TX_8X32,
BLOCK_64X16 => TX_32X8,
_ => TX_32X32
}
}
}
pub fn is_sqr(self) -> bool {
self.width_log2() == self.height_log2()
}
......
......@@ -37,7 +37,7 @@ use FrameState;
use FrameType;
use Tune;
use Sequence;
use encoder::ReferenceMode;
use encoder::{ChromaSampling, ReferenceMode};
use api::PredictionModesSetting;
#[derive(Clone)]
......@@ -313,7 +313,8 @@ pub fn rdo_tx_size_type(
bo,
tx_size,
tx_set,
seq.bit_depth
seq.bit_depth,
seq.chroma_sampling
)
} else {
TxType::DCT_DCT
......@@ -617,12 +618,13 @@ pub fn rdo_mode_decision(
best.tx_type,
false,
seq.bit_depth,
seq.chroma_sampling,
CFLParams::new(),
true,
false
);
cw.rollback(&cw_checkpoint);
if let Some(cfl) = rdo_cfl_alpha(fs, bo, bsize, seq.bit_depth) {
if let Some(cfl) = rdo_cfl_alpha(fs, bo, bsize, seq.bit_depth, seq.chroma_sampling) {
let mut wr: &mut dyn Writer = &mut WriterCounter::new();
let tell = wr.tell_frac();
......@@ -700,15 +702,9 @@ pub fn rdo_mode_decision(
}
pub fn rdo_cfl_alpha(
fs: &mut FrameState, bo: &BlockOffset, bsize: BlockSize, bit_depth: usize
) -> Option<CFLParams> {
// TODO: these are only valid for 4:2:0
let uv_tx_size = match bsize {
BlockSize::BLOCK_4X4 | BlockSize::BLOCK_8X8 => TxSize::TX_4X4,
BlockSize::BLOCK_16X16 => TxSize::TX_8X8,
BlockSize::BLOCK_32X32 => TxSize::TX_16X16,
_ => TxSize::TX_32X32
};
fs: &mut FrameState, bo: &BlockOffset, bsize: BlockSize, bit_depth: usize,
chroma_sampling: ChromaSampling) -> Option<CFLParams> {
let uv_tx_size = bsize.largest_uv_tx_size(chroma_sampling);
let mut ac = [0i16; 32 * 32];
luma_ac(&mut ac, fs, bo, bsize);
......@@ -746,7 +742,7 @@ pub fn rdo_cfl_alpha(
pub fn rdo_tx_type_decision(
fi: &FrameInvariants, fs: &mut FrameState, cw: &mut ContextWriter,
mode: PredictionMode, ref_frames: [usize; 2], mvs: [MotionVector; 2], bsize: BlockSize, bo: &BlockOffset, tx_size: TxSize,
tx_set: TxSet, bit_depth: usize
tx_set: TxSet, bit_depth: usize, chroma_sampling: ChromaSampling
) -> TxType {
let mut best_type = TxType::DCT_DCT;
let mut best_rd = std::f64::MAX;
......@@ -774,12 +770,12 @@ pub fn rdo_tx_type_decision(
let tell = wr.tell_frac();
let tx_dist = if is_inter {
write_tx_tree(
fi, fs, cw, wr, mode, bo, bsize, tx_size, tx_type, false, bit_depth, true, true
fi, fs, cw, wr, mode, bo, bsize, tx_size, tx_type, false, bit_depth, chroma_sampling, true, true
)
} else {
let cfl = CFLParams::new(); // Unused
write_tx_blocks(
fi, fs, cw, wr, mode, mode, bo, bsize, tx_size, tx_type, false, bit_depth, cfl, true, true
fi, fs, cw, wr, mode, mode, bo, bsize, tx_size, tx_type, false, bit_depth, chroma_sampling, cfl, true, true
)
};
......@@ -1076,4 +1072,3 @@ pub fn rdo_cdef_decision(sbo: &SuperBlockOffset, fi: &FrameInvariants,
}
best_index
}
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