Skip to content

Commit

Permalink
New types for distortion
Browse files Browse the repository at this point in the history
  • Loading branch information
barrbrain authored and tdaede committed Oct 11, 2019
1 parent 26f2656 commit 4a0a0b1
Show file tree
Hide file tree
Showing 2 changed files with 105 additions and 77 deletions.
55 changes: 16 additions & 39 deletions src/encoder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1086,7 +1086,7 @@ pub fn encode_tx_block<T: Pixel>(
alpha: i16,
rdo_type: RDOType,
need_recon_pixel: bool,
) -> (bool, i64) {
) -> (bool, ScaledDistortion) {
let qidx = get_qidx(fi, ts, cw, tile_bo);
assert_ne!(qidx, 0); // lossless is not yet supported
let PlaneConfig { xdec, ydec, .. } = ts.input.planes[p].cfg;
Expand Down Expand Up @@ -1129,7 +1129,7 @@ pub fn encode_tx_block<T: Pixel>(
}

if skip {
return (false, -1);
return (false, ScaledDistortion::zero());
}

let mut residual_storage: AlignedArray<[i16; 64 * 64]> =
Expand Down Expand Up @@ -1195,7 +1195,7 @@ pub fn encode_tx_block<T: Pixel>(
fi.ac_delta_q[p],
);

let mut tx_dist: i64 = -1;
let mut tx_dist: u64 = 0;

if !fi.use_tx_domain_distortion || need_recon_pixel {
inverse_transform_add(
Expand All @@ -1215,29 +1215,24 @@ pub fn encode_tx_block<T: Pixel>(
let c = *a as i32 - *b as i32;
(c * c) as u64
})
.sum::<u64>() as i64;
.sum::<u64>();

let tx_dist_scale_bits = 2 * (3 - get_log_tx_scale(tx_size));
let tx_dist_scale_rounding_offset = 1 << (tx_dist_scale_bits - 1);
tx_dist = (tx_dist + tx_dist_scale_rounding_offset) >> tx_dist_scale_bits;
}
if fi.config.train_rdo {
ts.rdo.add_rate(
fi.base_q_idx,
tx_size,
tx_dist as u64,
cost_coeffs as u64,
);
ts.rdo.add_rate(fi.base_q_idx, tx_size, tx_dist, cost_coeffs as u64);
}

if rdo_type == RDOType::TxDistEstRate {
// look up rate and distortion in table
let estimated_rate = estimate_rate(fi.base_q_idx, tx_size, tx_dist as u64);
let estimated_rate = estimate_rate(fi.base_q_idx, tx_size, tx_dist);
w.add_bits_frac(estimated_rate as u32);
}
let bias =
compute_distortion_bias(fi, ts.to_frame_block_offset(tile_bo), bsize);
(has_coeff, (tx_dist as f64 * bias * fi.dist_scale[p]) as i64)
(has_coeff, RawDistortion::new(tx_dist) * bias * fi.dist_scale[p])
}

pub fn motion_compensate<T: Pixel>(
Expand Down Expand Up @@ -1491,7 +1486,7 @@ pub fn encode_block_post_cdef<T: Pixel>(
skip: bool, cfl: CFLParams, tx_size: TxSize, tx_type: TxType,
mode_context: usize, mv_stack: &[CandidateMV], rdo_type: RDOType,
need_recon_pixel: bool, record_stats: bool,
) -> (bool, i64) {
) -> (bool, ScaledDistortion) {
let is_inter = !luma_mode.is_intra();
if is_inter {
assert!(luma_mode == chroma_mode);
Expand Down Expand Up @@ -1764,15 +1759,15 @@ pub fn write_tx_blocks<T: Pixel>(
chroma_mode: PredictionMode, tile_bo: TileBlockOffset, bsize: BlockSize,
tx_size: TxSize, tx_type: TxType, skip: bool, cfl: CFLParams,
luma_only: bool, rdo_type: RDOType, need_recon_pixel: bool,
) -> (bool, i64) {
) -> (bool, ScaledDistortion) {
let bw = bsize.width_mi() / tx_size.width_mi();
let bh = bsize.height_mi() / tx_size.height_mi();
let qidx = get_qidx(fi, ts, cw, tile_bo);

let PlaneConfig { xdec, ydec, .. } = ts.input.planes[1].cfg;
let mut ac: AlignedArray<[i16; 32 * 32]> = AlignedArray::uninitialized();
let mut partition_has_coeff: bool = false;
let mut tx_dist: i64 = 0;
let mut tx_dist = ScaledDistortion::zero();
let do_chroma = has_chroma(tile_bo, bsize, xdec, ydec);

ts.qc.update(
Expand Down Expand Up @@ -1814,9 +1809,6 @@ pub fn write_tx_blocks<T: Pixel>(
need_recon_pixel,
);
partition_has_coeff |= has_coeff;
assert!(
!fi.use_tx_domain_distortion || need_recon_pixel || skip || dist >= 0
);
tx_dist += dist;
}
}
Expand Down Expand Up @@ -1900,12 +1892,6 @@ pub fn write_tx_blocks<T: Pixel>(
need_recon_pixel,
);
partition_has_coeff |= has_coeff;
assert!(
!fi.use_tx_domain_distortion
|| need_recon_pixel
|| skip
|| dist >= 0
);
tx_dist += dist;
}
}
Expand All @@ -1923,17 +1909,16 @@ pub fn write_tx_tree<T: Pixel>(
tile_bo: TileBlockOffset, bsize: BlockSize, tx_size: TxSize,
tx_type: TxType, skip: bool, luma_only: bool, rdo_type: RDOType,
need_recon_pixel: bool,
) -> (bool, i64) {
) -> (bool, ScaledDistortion) {
if skip {
return (false, -1);
return (false, ScaledDistortion::zero());
}
let bw = bsize.width_mi() / tx_size.width_mi();
let bh = bsize.height_mi() / tx_size.height_mi();
let qidx = get_qidx(fi, ts, cw, tile_bo);

let PlaneConfig { xdec, ydec, .. } = ts.input.planes[1].cfg;
let ac = &[0i16; 0];
let mut tx_dist: i64 = 0;
let mut partition_has_coeff: bool = false;

ts.qc.update(
Expand Down Expand Up @@ -1968,10 +1953,7 @@ pub fn write_tx_tree<T: Pixel>(
need_recon_pixel,
);
partition_has_coeff |= has_coeff;
assert!(
!fi.use_tx_domain_distortion || need_recon_pixel || skip || dist >= 0
);
tx_dist += dist;
let mut tx_dist = dist;

if luma_only {
return (partition_has_coeff, tx_dist);
Expand Down Expand Up @@ -2044,12 +2026,6 @@ pub fn write_tx_tree<T: Pixel>(
need_recon_pixel,
);
partition_has_coeff |= has_coeff;
assert!(
!fi.use_tx_domain_distortion
|| need_recon_pixel
|| skip
|| dist >= 0
);
tx_dist += dist;
}
}
Expand Down Expand Up @@ -2176,7 +2152,7 @@ fn encode_partition_bottomup<T: Pixel, W: Writer>(
let w: &mut W = if cw.bc.cdef_coded { w_post_cdef } else { w_pre_cdef };
let tell = w.tell_frac();
cw.write_partition(w, tile_bo, PartitionType::PARTITION_NONE, bsize);
compute_rd_cost(fi, w.tell_frac() - tell, 0)
compute_rd_cost(fi, w.tell_frac() - tell, ScaledDistortion::zero())
} else {
0.0
};
Expand Down Expand Up @@ -2276,7 +2252,8 @@ fn encode_partition_bottomup<T: Pixel, W: Writer>(
if cw.bc.cdef_coded { w_post_cdef } else { w_pre_cdef };
let tell = w.tell_frac();
cw.write_partition(w, tile_bo, partition, bsize);
rd_cost = compute_rd_cost(fi, w.tell_frac() - tell, 0);
rd_cost =
compute_rd_cost(fi, w.tell_frac() - tell, ScaledDistortion::zero());
}

let four_partitions = [
Expand Down
Loading

0 comments on commit 4a0a0b1

Please sign in to comment.