[isp] Fix chroma coeff writing for ISP.

This commit is contained in:
siivonek 2022-08-31 14:54:30 +03:00 committed by Marko Viitanen
parent d8d206365c
commit 33cd44f11b
2 changed files with 36 additions and 22 deletions

View file

@ -527,7 +527,8 @@ static void encode_transform_unit(
bool only_chroma,
lcu_coeff_t* coeff,
enum uvg_tree_type tree_type,
bool last_split)
bool last_split,
const cu_loc_t *original_loc) // Original cu dimensions, before CU split
{
assert(depth >= 1 && depth <= MAX_PU_DEPTH);
@ -596,7 +597,8 @@ static void encode_transform_unit(
cbf_is_set(cur_pu->cbf, depth, COLOR_V);
if ((chroma_cbf_set || joint_chroma) && last_split) {
//Need to drop const to get lfnst constraints
encode_chroma_tu(state, cu_loc, depth, (cu_info_t*)cur_pu, &scan_idx, coeff, joint_chroma, tree_type);
// Use original dimensions instead of ISP split dimensions
encode_chroma_tu(state, original_loc, depth, (cu_info_t*)cur_pu, &scan_idx, coeff, joint_chroma, tree_type);
}
}
@ -620,7 +622,8 @@ static void encode_transform_coeff(
lcu_coeff_t* coeff,
enum uvg_tree_type tree_type,
bool last_split,
int *luma_cbf_ctx) // Always true except when writing sub partition coeffs (ISP)
int *luma_cbf_ctx, // Always true except when writing sub partition coeffs (ISP)
cu_loc_t *original_loc) // Original dimensions before ISP split
{
cabac_data_t * const cabac = &state->cabac;
const int x = cu_loc->x;
@ -707,7 +710,7 @@ static void encode_transform_coeff(
cu_loc_t loc;
uvg_cu_loc_ctor(&loc, (x + i * split_width), (y + j * split_height), width >> 1, height >> 1);
encode_transform_coeff(state, &loc, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type, true, luma_cbf_ctx);
encode_transform_coeff(state, &loc, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type, true, luma_cbf_ctx, &loc);
}
}
return;
@ -761,7 +764,7 @@ static void encode_transform_coeff(
CABAC_BIN(cabac, cur_pu->joint_cb_cr != 0, "tu_joint_cbcr_residual_flag");
}
encode_transform_unit(state, cu_loc, depth, only_chroma, coeff, tree_type, last_split);
encode_transform_unit(state, cu_loc, depth, only_chroma, coeff, tree_type, last_split, original_loc);
}
}
@ -1616,7 +1619,7 @@ void uvg_encode_coding_tree(
// Code (possible) coeffs to bitstream
if (cbf) {
int luma_cbf_ctx = 0;
encode_transform_coeff(state, &cu_loc, depth, 0, 0, 0, 0, coeff, tree_type, true, &luma_cbf_ctx);
encode_transform_coeff(state, &cu_loc, depth, 0, 0, 0, 0, coeff, tree_type, true, &luma_cbf_ctx, &cu_loc);
}
encode_mts_idx(state, cabac, cur_cu);
@ -1646,7 +1649,7 @@ void uvg_encode_coding_tree(
// Check if last split to write chroma
bool last_split = (i + 1) == split_limit;
encode_transform_coeff(state, &split_loc, depth, 0, 0, 0, 0, coeff, tree_type, last_split, &luma_cbf_ctx);
encode_transform_coeff(state, &split_loc, depth, 0, 0, 0, 0, coeff, tree_type, last_split, &luma_cbf_ctx, &cu_loc);
}
}
@ -1665,7 +1668,7 @@ void uvg_encode_coding_tree(
tmp->violates_lfnst_constrained_luma = false;
tmp->violates_lfnst_constrained_chroma = false;
tmp->lfnst_last_scan_pos = false;
encode_transform_coeff(state, &cu_loc, depth, 0, 0, 0, 1, coeff, tree_type, true, &luma_cbf_ctx);
encode_transform_coeff(state, &cu_loc, depth, 0, 0, 0, 1, coeff, tree_type, true, &luma_cbf_ctx, &cu_loc);
// Write LFNST only once for single tree structure
encode_lfnst_idx(state, cabac, tmp, x, y, depth, cu_width, cu_height, tree_type, COLOR_UV);
}

View file

@ -1148,20 +1148,7 @@ static double search_cu(
depth, &intra_search,
NULL,
lcu, tree_type,recon_luma,recon_chroma);
// Set isp split cbfs here
const int split_type = intra_search.pred_cu.intra.isp_mode;
const int split_num = split_type == ISP_MODE_NO_ISP ? 0 : uvg_get_isp_split_num(cu_width, cu_height, split_type);
for (int i = 0; i < split_num; ++i) {
cu_loc_t isp_loc;
uvg_get_isp_split_loc(&isp_loc, x, y, cu_width, cu_height, i, split_type);
//search_data->best_isp_cbfs |= cbf_is_set(cur_cu->cbf, depth, COLOR_Y) << (i++);
cu_info_t* split_cu = LCU_GET_CU_AT_PX(lcu, isp_loc.x % LCU_WIDTH, isp_loc.y % LCU_WIDTH);
bool cur_cbf = (intra_search.best_isp_cbfs >> i) & 1;
cbf_clear(&split_cu->cbf, depth, COLOR_Y);
if (cur_cbf) {
cbf_set(&split_cu->cbf, depth, COLOR_Y);
}
}
if(depth == 4 && x % 8 && y % 8 && tree_type != UVG_LUMA_T && state->encoder_control->chroma_format != UVG_CSP_400) {
intra_search.pred_cu.intra.mode_chroma = cur_cu->intra.mode_chroma;
@ -1173,6 +1160,30 @@ static double search_cu(
tree_type,false,true);
}
if (cur_cu->joint_cb_cr == 4) cur_cu->joint_cb_cr = 0;
// Set isp split cbfs here
const int split_type = intra_search.pred_cu.intra.isp_mode;
const int split_num = split_type == ISP_MODE_NO_ISP ? 0 : uvg_get_isp_split_num(cu_width, cu_height, split_type);
const int cbf_cb = cbf_is_set(cur_cu->cbf, depth, COLOR_U);
const int cbf_cr = cbf_is_set(cur_cu->cbf, depth, COLOR_V);
const int jccr = cur_cu->joint_cb_cr;
for (int i = 0; i < split_num; ++i) {
cu_loc_t isp_loc;
uvg_get_isp_split_loc(&isp_loc, x, y, cu_width, cu_height, i, split_type);
//search_data->best_isp_cbfs |= cbf_is_set(cur_cu->cbf, depth, COLOR_Y) << (i++);
cu_info_t* split_cu = LCU_GET_CU_AT_PX(lcu, isp_loc.x % LCU_WIDTH, isp_loc.y % LCU_WIDTH);
bool cur_cbf = (intra_search.best_isp_cbfs >> i) & 1;
cbf_clear(&split_cu->cbf, depth, COLOR_Y);
cbf_clear(&split_cu->cbf, depth, COLOR_U);
cbf_clear(&split_cu->cbf, depth, COLOR_V);
if (cur_cbf) {
cbf_set(&split_cu->cbf, depth, COLOR_Y);
}
if(cbf_cb) cbf_set(&split_cu->cbf, depth, COLOR_U);
if(cbf_cr) cbf_set(&split_cu->cbf, depth, COLOR_V);
split_cu->joint_cb_cr = jccr;
}
lcu_fill_cu_info(lcu, x_local, y_local, cu_width, cu_width, cur_cu);