diff --git a/src/cu.h b/src/cu.h index b61ed2f2..762467b1 100644 --- a/src/cu.h +++ b/src/cu.h @@ -168,6 +168,7 @@ typedef struct uint8_t violates_lfnst_constrained_chroma : 1; // Two types, luma and chroma. Luma index is 0. uint8_t lfnst_last_scan_pos : 1; uint8_t lfnst_idx : 2; + uint8_t cr_lfnst_idx : 2; union { struct { diff --git a/src/encode_coding_tree.c b/src/encode_coding_tree.c index 00ca218f..0bb6efeb 100644 --- a/src/encode_coding_tree.c +++ b/src/encode_coding_tree.c @@ -243,7 +243,7 @@ static bool encode_lfnst_idx( // VTM seems to force explicit dual tree structure for small 4x4 blocks bool is_separate_tree = depth == 4; // TODO: if/when separate/dual tree structure is implemented, get proper value for this - const int lfnst_index = pred_cu->lfnst_idx; + const int lfnst_index = color == COLOR_Y ? pred_cu->lfnst_idx : pred_cu->cr_lfnst_idx; assert((lfnst_index >= 0 && lfnst_index < 3) && "Invalid LFNST index."); uint16_t ctx_idx = 0; diff --git a/src/search_intra.c b/src/search_intra.c index ae4f4c7e..56065ab5 100644 --- a/src/search_intra.c +++ b/src/search_intra.c @@ -1434,6 +1434,11 @@ int8_t uvg_search_intra_chroma_rdo( const int offset = ((lcu_px.x & ~7) >> 1) + ((lcu_px.y & ~7) >> 1)* LCU_WIDTH_C; int lfnst_modes_to_check[3]; + if(depth == 4) { + for (int i = 0; i < 3; ++i) { + lfnst_modes_to_check[i] = i; + } + } if(chroma_data->pred_cu.lfnst_idx) { lfnst_modes_to_check[0] = chroma_data->pred_cu.lfnst_idx; lfnst_modes_to_check[1] = -1; @@ -1461,6 +1466,7 @@ int8_t uvg_search_intra_chroma_rdo( if (lfnst == -1) { continue; } + pred_cu->cr_lfnst_idx = lfnst; chroma_data[mode_i].lfnst_costs[lfnst] += mode_bits * state->lambda; if (pred_cu->tr_depth == pred_cu->depth) { uvg_intra_predict( @@ -1555,7 +1561,7 @@ int8_t uvg_search_intra_chroma_rdo( memcpy(&state->search_cabac, &temp_cabac, sizeof(cabac_data_t)); } } - pred_cu->lfnst_idx = best_lfnst_index; + pred_cu->cr_lfnst_idx = best_lfnst_index; } sort_modes(chroma_data, num_modes); @@ -1603,7 +1609,9 @@ int8_t uvg_search_cu_intra_chroma(encoder_state_t * const state, chroma_data[i].pred_cu.intra.mode_chroma = num_modes == 1 ? intra_mode : modes[i]; chroma_data[i].pred_cu.intra.mode = -1; chroma_data[i].cost = 0; - memcpy(chroma_data[i].lfnst_costs, search_data->lfnst_costs, sizeof(double) * 3); + if(depth != 4) { + memcpy(chroma_data[i].lfnst_costs, search_data->lfnst_costs, sizeof(double) * 3); + } } // Don't do rough mode search if all modes are selected. // FIXME: It might make more sense to only disable rough search if diff --git a/src/strategies/generic/quant-generic.c b/src/strategies/generic/quant-generic.c index 1cbc5d4e..f103acc9 100644 --- a/src/strategies/generic/quant-generic.c +++ b/src/strategies/generic/quant-generic.c @@ -476,7 +476,7 @@ int uvg_quantize_residual_generic(encoder_state_t *const state, uvg_transform2d(state->encoder_control, residual, coeff, width, color, cur_cu); } - const uint16_t lfnst_index = cur_cu->lfnst_idx; + const uint16_t lfnst_index = color == COLOR_Y ? cur_cu->lfnst_idx : cur_cu->cr_lfnst_idx; if (state->encoder_control->cfg.lfnst && cur_cu->type == CU_INTRA) { // Forward low frequency non-separable transform diff --git a/src/transform.c b/src/transform.c index 4b069d4f..2a58e6ad 100644 --- a/src/transform.c +++ b/src/transform.c @@ -501,7 +501,7 @@ void uvg_chroma_transform_search( const int can_use_tr_skip = state->encoder_control->cfg.trskip_enable && (1 << state->encoder_control->cfg.trskip_max_size) >= width && state->encoder_control->cfg.chroma_trskip_enable && - pred_cu->lfnst_idx == 0; + pred_cu->cr_lfnst_idx == 0 ; if (can_use_tr_skip) { uvg_transformskip(state->encoder_control, u_resi, u_coeff + num_transforms * trans_offset, width); @@ -552,16 +552,16 @@ void uvg_chroma_transform_search( scan_order, &u_has_coeffs, &v_has_coeffs, - pred_cu->lfnst_idx); - /* - if(pred_cu->type == CU_INTRA && transforms[i] != CHROMA_TS) { + pred_cu->cr_lfnst_idx); + + if(pred_cu->type == CU_INTRA && transforms[i] != CHROMA_TS && depth == 4) { bool constraints[2] = { false, false }; uvg_derive_lfnst_constraints(pred_cu, depth, constraints, &u_coeff[i * trans_offset], width, height); if(!IS_JCCR_MODE(transforms[i])) { uvg_derive_lfnst_constraints(pred_cu, depth, constraints, &v_coeff[i * trans_offset], width, height); } - if ((constraints[0] || !constraints[1]) && pred_cu->lfnst_idx != 0) continue; - }*/ + if ((constraints[0] || !constraints[1]) && pred_cu->cr_lfnst_idx != 0) continue; + } if (IS_JCCR_MODE(transforms[i]) && !u_has_coeffs) continue;