diff --git a/src/intra.c b/src/intra.c index 01e3eccd..df9662b5 100644 --- a/src/intra.c +++ b/src/intra.c @@ -1417,6 +1417,30 @@ void uvg_intra_predict( } } +// This function works on luma coordinates +const cu_info_t* kvz_get_co_located_luma_cu( + int x, + int y, + int width, + int height, + const lcu_t* const lcu, + const cu_array_t* const cu_array, + enum kvz_tree_type tree_type) +{ + assert((cu_array || lcu) && !(cu_array && lcu)); + assert(tree_type != KVZ_LUMA_T && "Luma only CU shouldn't need colocated luma CU"); + if(tree_type == KVZ_CHROMA_T) { + x += width >> 1; + y += height >> 1; + } + if(cu_array) { + return uvg_cu_array_at_const(cu_array, x, y); + } + else { + return LCU_GET_CU_AT_PX(lcu, SUB_SCU(x), SUB_SCU(y)); + } +} + static void intra_recon_tb_leaf( encoder_state_t* const state, diff --git a/src/intra.h b/src/intra.h index 6357bc29..3d48e16c 100644 --- a/src/intra.h +++ b/src/intra.h @@ -146,4 +146,13 @@ void uvg_intra_recon_cu( lcu_t *lcu, enum kvz_tree_type tree_type); +const cu_info_t* kvz_get_co_located_luma_cu( + int x, + int y, + int width, + int height, + const lcu_t* const lcu, + const cu_array_t* const cu_array, + enum kvz_tree_type tree_type); + int uvg_get_mip_flag_context(int x, int y, int width, int height, const lcu_t* lcu, cu_array_t* const cu_a); diff --git a/src/search.c b/src/search.c index cec50a97..f94ca0fd 100644 --- a/src/search.c +++ b/src/search.c @@ -944,6 +944,9 @@ static double search_cu( // rd2. Possibly because the luma mode search already takes chroma // into account, so there is less of a chanse of luma mode being // really bad for chroma. + if(tree_type == KVZ_CHROMA_T) { + intra_search.pred_cu.intra = kvz_get_co_located_luma_cu(x, y, luma_width, luma_width, lcu, NULL, KVZ_CHROMA_T)->intra; + } intra_search.pred_cu.intra.mode_chroma = intra_search.pred_cu.intra.mode; if (ctrl->cfg.rdo >= 3 || ctrl->cfg.jccr || ctrl->cfg.lfnst) { uvg_search_cu_intra_chroma(state, x, y, depth, lcu, &intra_search, tree_type);