diff --git a/src/encode_coding_tree.c b/src/encode_coding_tree.c index ba74215f..63cf4d66 100644 --- a/src/encode_coding_tree.c +++ b/src/encode_coding_tree.c @@ -525,7 +525,8 @@ static void encode_transform_unit( int depth, bool only_chroma, lcu_coeff_t* coeff, - enum uvg_tree_type tree_type) + enum uvg_tree_type tree_type, + bool last_split) { assert(depth >= 1 && depth <= MAX_PU_DEPTH); @@ -586,7 +587,7 @@ static void encode_transform_unit( bool chroma_cbf_set = cbf_is_set(cur_pu->cbf, depth, COLOR_U) || cbf_is_set(cur_pu->cbf, depth, COLOR_V); - if (chroma_cbf_set || joint_chroma) { + if ((chroma_cbf_set || joint_chroma) && last_split) { //Need to drop const to get lfnst constraints encode_chroma_tu(state, x, y, depth, width_c, (cu_info_t*)cur_pu, &scan_idx, coeff, joint_chroma, tree_type); } @@ -611,7 +612,8 @@ static void encode_transform_coeff( uint8_t parent_coeff_v, bool only_chroma, lcu_coeff_t* coeff, - enum uvg_tree_type tree_type) + enum uvg_tree_type tree_type, + bool last_split) // Always true except when writing sub partition coeffs (ISP) { cabac_data_t * const cabac = &state->cabac; //const encoder_control_t *const ctrl = state->encoder_control; @@ -642,8 +644,6 @@ static void encode_transform_coeff( int8_t split = (LCU_WIDTH >> depth > TR_MAX_WIDTH); - - const int cb_flag_y = tree_type != UVG_CHROMA_T ? cbf_is_set(cur_pu->cbf, depth, COLOR_Y) : 0; const int cb_flag_u = tree_type != UVG_LUMA_T ?( cur_pu->joint_cb_cr ? (cur_pu->joint_cb_cr >> 1) & 1 : cbf_is_set(cur_cu->cbf, depth, COLOR_U)) : 0; const int cb_flag_v = tree_type != UVG_LUMA_T ? (cur_pu->joint_cb_cr ? cur_pu->joint_cb_cr & 1 : cbf_is_set(cur_cu->cbf, depth, COLOR_V)) : 0; @@ -671,7 +671,7 @@ static void encode_transform_coeff( // - they have already been signaled to 0 previously // When they are not present they are inferred to be 0, except for size 4 // when the flags from previous level are used. - if (state->encoder_control->chroma_format != UVG_CSP_400 && (depth != 4 || only_chroma) && tree_type != UVG_LUMA_T) { + if (state->encoder_control->chroma_format != UVG_CSP_400 && (depth != 4 || only_chroma) && tree_type != UVG_LUMA_T && last_split) { if (!split) { if (true) { @@ -690,10 +690,10 @@ static void encode_transform_coeff( uint8_t offset = LCU_WIDTH >> (depth + 1); int x2 = x + offset; int y2 = y + offset; - encode_transform_coeff(state, x, y, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type); - encode_transform_coeff(state, x2, y, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type); - encode_transform_coeff(state, x, y2, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type); - encode_transform_coeff(state, x2, y2, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type); + encode_transform_coeff(state, x, y, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type, true); + encode_transform_coeff(state, x2, y, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type, true); + encode_transform_coeff(state, x, y2, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type, true); + encode_transform_coeff(state, x2, y2, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type, true); return; } @@ -737,12 +737,13 @@ static void encode_transform_coeff( || (cb_flag_u && cb_flag_v)) && (depth != 4 || only_chroma || tree_type == UVG_CHROMA_T) && state->encoder_control->cfg.jccr + && last_split ) { assert(cur_pu->joint_cb_cr < 4 && "JointCbCr is in search state."); cabac->cur_ctx = &cabac->ctx.joint_cb_cr[cb_flag_u * 2 + cb_flag_v - 1]; CABAC_BIN(cabac, cur_pu->joint_cb_cr != 0, "tu_joint_cbcr_residual_flag"); } - encode_transform_unit(state, x, y, depth, only_chroma, coeff, tree_type); + encode_transform_unit(state, x, y, depth, only_chroma, coeff, tree_type, last_split); } } @@ -1671,7 +1672,7 @@ void uvg_encode_coding_tree( // Code (possible) coeffs to bitstream if (cbf) { - encode_transform_coeff(state, x, y, depth, 0, 0, 0, 0, coeff, tree_type); + encode_transform_coeff(state, x, y, depth, 0, 0, 0, 0, coeff, tree_type, true); } encode_mts_idx(state, cabac, cur_cu); @@ -1688,7 +1689,25 @@ void uvg_encode_coding_tree( } if (tree_type != UVG_CHROMA_T) { - encode_transform_coeff(state, x, y, depth, 0, 0, 0, 0, coeff, tree_type); + // Cycle through sub partitions if ISP enabled. + // ISP split is done horizontally or vertically depending on ISP mode, 2 or 4 times depending on block dimensions. + // Small blocks are split only twice. + int split_type = cur_cu->intra.isp_mode; + + int part_dim = cu_width; + if (split_type != ISP_MODE_NO_ISP) { + part_dim = uvg_get_isp_split_dim(cu_width, cu_height, split_type); + } + int limit = split_type == ISP_MODE_HOR ? cu_height : cu_width; + + for (int part = 0; part < limit; part += part_dim) { + const int part_x = split_type == ISP_MODE_HOR ? x : x + part; + const int part_y = split_type == ISP_MODE_HOR ? y + part : y; + + // Check if last split to write chroma + bool last_split = (part + part_dim) == limit; + encode_transform_coeff(state, part_x, part_y, depth, 0, 0, 0, 0, coeff, tree_type, last_split); + } } if (tree_type != UVG_CHROMA_T) { @@ -1706,7 +1725,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, x, y, depth, 0, 0, 0, 1, coeff, tree_type); + encode_transform_coeff(state, x, y, depth, 0, 0, 0, 1, coeff, tree_type, true); // 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/intra.c b/src/intra.c index 1a1e3817..e65acfea 100644 --- a/src/intra.c +++ b/src/intra.c @@ -1481,8 +1481,10 @@ const cu_info_t* uvg_get_co_located_luma_cu( * \param height Block height. * \param split_type Horizontal or vertical split. */ -static int get_isp_split_dim(const int width, const int height, const int split_type) +int uvg_get_isp_split_dim(const int width, const int height, const int split_type) { + assert(split_type != ISP_MODE_NO_ISP && "Cannot calculate split dimension if no split type is set. Make sure this function is not called in this case."); + bool divide_in_rows = split_type == SPLIT_TYPE_HOR; int split_dim_size, non_split_dim_size, partition_size, div_shift = 2; @@ -1666,9 +1668,9 @@ void uvg_intra_recon_cu( // ISP split is done horizontally or vertically depending on ISP mode, 2 or 4 times depending on block dimensions. // Small blocks are split only twice. int split_type = search_data->pred_cu.intra.isp_mode; - int part_dim = get_isp_split_dim(width, height, split_type); + int part_dim = uvg_get_isp_split_dim(width, height, split_type); int limit = split_type == ISP_MODE_HOR ? height : width; - for (int part = 0; part < limit; part + part_dim) { + for (int part = 0; part < limit; part += part_dim) { const int part_x = split_type == ISP_MODE_HOR ? x : x + part; const int part_y = split_type == ISP_MODE_HOR ? y + part: y; const int part_w = split_type == ISP_MODE_HOR ? part_dim : width; @@ -1749,8 +1751,8 @@ bool uvg_can_use_isp_with_lfnst(const int width, const int height, const int isp return false; } - const int tu_width = (isp_split_type == ISP_MODE_HOR) ? width : get_isp_split_dim(width, height, SPLIT_TYPE_VER); - const int tu_height = (isp_split_type == ISP_MODE_HOR) ? get_isp_split_dim(width, height, SPLIT_TYPE_HOR) : height; + const int tu_width = (isp_split_type == ISP_MODE_HOR) ? width : uvg_get_isp_split_dim(width, height, SPLIT_TYPE_VER); + const int tu_height = (isp_split_type == ISP_MODE_HOR) ? uvg_get_isp_split_dim(width, height, SPLIT_TYPE_HOR) : height; if (!(tu_width >= TR_MIN_WIDTH && tu_height >= TR_MIN_WIDTH)) {