diff --git a/src/encode_coding_tree.c b/src/encode_coding_tree.c index fd5b6ce6..9b0bf884 100644 --- a/src/encode_coding_tree.c +++ b/src/encode_coding_tree.c @@ -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); } diff --git a/src/search.c b/src/search.c index 66eac299..5dad1b3b 100644 --- a/src/search.c +++ b/src/search.c @@ -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);