avcodec/prores_ks : use official quant_matrix (for proxy and xq codec luma and chroma quant matrix is not the same)

disable the use of the official luma xq matrix for now (output
appear to be desaturate)

1 files changed, 83 insertions, 16 deletions

diff --git a/libavcodec/proresenc_kostya.c b/libavcodec/proresenc_kostya.c
index 06d5aa31e8..81f3865ea6 100644
--- a/libavcodec/proresenc_kostya.c
+++ b/libavcodec/proresenc_kostya.c
@@ -51,9 +51,11 @@ enum {
 
 enum {
     QUANT_MAT_PROXY = 0,
+    QUANT_MAT_PROXY_CHROMA,
     QUANT_MAT_LT,
     QUANT_MAT_STANDARD,
     QUANT_MAT_HQ,
+    QUANT_MAT_XQ_LUMA,
     QUANT_MAT_DEFAULT,
 };
 
@@ -68,6 +70,16 @@ static const uint8_t prores_quant_matrices[][64] = {
         13, 63, 63, 63, 63, 63, 63, 63,
         63, 63, 63, 63, 63, 63, 63, 63,
     },
+    { // proxy chromas
+        4,  7,  9, 11, 13, 14, 63, 63,
+        7,  7, 11, 12, 14, 63, 63, 63,
+        9, 11, 13, 14, 63, 63, 63, 63,
+        11, 11, 13, 14, 63, 63, 63, 63,
+        11, 13, 14, 63, 63, 63, 63, 63,
+        13, 14, 63, 63, 63, 63, 63, 63,
+        13, 63, 63, 63, 63, 63, 63, 63,
+        63, 63, 63, 63, 63, 63, 63, 63
+    },
     { // LT
          4,  5,  6,  7,  9, 11, 13, 15,
          5,  5,  7,  8, 11, 13, 15, 17,
@@ -98,6 +110,16 @@ static const uint8_t prores_quant_matrices[][64] = {
          4,  4,  4,  4,  5,  5,  6,  7,
          4,  4,  4,  4,  5,  6,  7,  7,
     },
+    { // XQ luma
+        2,  2,  2,  2,  2,  2,  2,  2,
+        2,  2,  2,  2,  2,  2,  2,  2,
+        2,  2,  2,  2,  2,  2,  2,  2,
+        2,  2,  2,  2,  2,  2,  2,  3,
+        2,  2,  2,  2,  2,  2,  3,  3,
+        2,  2,  2,  2,  2,  3,  3,  3,
+        2,  2,  2,  2,  3,  3,  3,  4,
+        2,  2,  2,  2,  3,  3,  4,  4,
+    },
     { // codec default
          4,  4,  4,  4,  4,  4,  4,  4,
          4,  4,  4,  4,  4,  4,  4,  4,
@@ -125,6 +147,7 @@ static const struct prores_profile {
     int         max_quant;
     int         br_tab[NUM_MB_LIMITS];
     int         quant;
+    int         quant_chroma;
 } prores_profile_info[6] = {
     {
         .full_name = "proxy",
@@ -133,6 +156,7 @@ static const struct prores_profile {
         .max_quant = 8,
         .br_tab    = { 300, 242, 220, 194 },
         .quant     = QUANT_MAT_PROXY,
+        .quant_chroma = QUANT_MAT_PROXY_CHROMA,
     },
     {
         .full_name = "LT",
@@ -141,6 +165,7 @@ static const struct prores_profile {
         .max_quant = 9,
         .br_tab    = { 720, 560, 490, 440 },
         .quant     = QUANT_MAT_LT,
+        .quant_chroma = QUANT_MAT_LT,
     },
     {
         .full_name = "standard",
@@ -149,6 +174,7 @@ static const struct prores_profile {
         .max_quant = 6,
         .br_tab    = { 1050, 808, 710, 632 },
         .quant     = QUANT_MAT_STANDARD,
+        .quant_chroma = QUANT_MAT_STANDARD,
     },
     {
         .full_name = "high quality",
@@ -157,6 +183,7 @@ static const struct prores_profile {
         .max_quant = 6,
         .br_tab    = { 1566, 1216, 1070, 950 },
         .quant     = QUANT_MAT_HQ,
+        .quant_chroma = QUANT_MAT_HQ,
     },
     {
         .full_name = "4444",
@@ -165,6 +192,7 @@ static const struct prores_profile {
         .max_quant = 6,
         .br_tab    = { 2350, 1828, 1600, 1425 },
         .quant     = QUANT_MAT_HQ,
+        .quant_chroma = QUANT_MAT_HQ,
     },
     {
         .full_name = "4444XQ",
@@ -172,7 +200,8 @@ static const struct prores_profile {
         .min_quant = 1,
         .max_quant = 6,
         .br_tab    = { 3525, 2742, 2400, 2137 },
-        .quant     = QUANT_MAT_HQ,
+        .quant     = QUANT_MAT_HQ, /* Fix me : use QUANT_MAT_XQ_LUMA */
+        .quant_chroma = QUANT_MAT_HQ,
     }
 };
 
@@ -200,8 +229,10 @@ typedef struct ProresContext {
     DECLARE_ALIGNED(16, int16_t, blocks)[MAX_PLANES][64 * 4 * MAX_MBS_PER_SLICE];
     DECLARE_ALIGNED(16, uint16_t, emu_buf)[16*16];
     int16_t quants[MAX_STORED_Q][64];
+    int16_t quants_chroma[MAX_STORED_Q][64];
     int16_t custom_q[64];
     const uint8_t *quant_mat;
+    const uint8_t *quant_chroma_mat;
     const uint8_t *scantable;
 
     void (*fdct)(FDCTDSPContext *fdsp, const uint16_t *src,
@@ -527,6 +558,7 @@ static int encode_slice(AVCodecContext *avctx, const AVFrame *pic,
     ptrdiff_t linesize;
     int plane_factor, is_chroma;
     uint16_t *qmat;
+    uint16_t *qmat_chroma;
 
     if (ctx->pictures_per_frame == 1)
         line_add = 0;
@@ -535,12 +567,17 @@ static int encode_slice(AVCodecContext *avctx, const AVFrame *pic,
 
     if (ctx->force_quant) {
         qmat = ctx->quants[0];
+        qmat_chroma = ctx->quants_chroma[0];
     } else if (quant < MAX_STORED_Q) {
         qmat = ctx->quants[quant];
+        qmat_chroma = ctx->quants_chroma[quant];
     } else {
         qmat = ctx->custom_q;
-        for (i = 0; i < 64; i++)
+        qmat_chroma = ctx->custom_q;
+        for (i = 0; i < 64; i++) {
             qmat[i] = ctx->quant_mat[i] * quant;
+            qmat_chroma[i] = ctx->quant_chroma_mat[i] * quant;
+        }
     }
 
     for (i = 0; i < ctx->num_planes; i++) {
@@ -569,10 +606,17 @@ static int encode_slice(AVCodecContext *avctx, const AVFrame *pic,
                            pwidth, avctx->height / ctx->pictures_per_frame,
                            ctx->blocks[0], ctx->emu_buf,
                            mbs_per_slice, num_cblocks, is_chroma);
-            sizes[i] = encode_slice_plane(ctx, pb, src, linesize,
-                                          mbs_per_slice, ctx->blocks[0],
-                                          num_cblocks, plane_factor,
-                                          qmat);
+            if (!is_chroma) {/* luma quant */
+                sizes[i] = encode_slice_plane(ctx, pb, src, linesize,
+                                              mbs_per_slice, ctx->blocks[0],
+                                              num_cblocks, plane_factor,
+                                              qmat);
+            } else { /* chroma plane */
+                sizes[i] = encode_slice_plane(ctx, pb, src, linesize,
+                                              mbs_per_slice, ctx->blocks[0],
+                                              num_cblocks, plane_factor,
+                                              qmat_chroma);
+            }
         } else {
             get_alpha_data(ctx, src, linesize, xp, yp,
                            pwidth, avctx->height / ctx->pictures_per_frame,
@@ -771,6 +815,7 @@ static int find_slice_quant(AVCodecContext *avctx,
     int slice_bits[TRELLIS_WIDTH], slice_score[TRELLIS_WIDTH];
     int overquant;
     uint16_t *qmat;
+    uint16_t *qmat_chroma;
     int linesize[4], line_add;
     int alpha_bits = 0;
 
@@ -825,12 +870,17 @@ static int find_slice_quant(AVCodecContext *avctx,
     for (q = min_quant; q <= max_quant; q++) {
         bits  = alpha_bits;
         error = 0;
-        for (i = 0; i < ctx->num_planes - !!ctx->alpha_bits; i++) {
+        bits += estimate_slice_plane(ctx, &error, 0,
+                                     src, linesize[0],
+                                     mbs_per_slice,
+                                     num_cblocks[0], plane_factor[0],
+                                     ctx->quants[q], td); /* estimate luma plane */
+        for (i = 1; i < ctx->num_planes - !!ctx->alpha_bits; i++) { /* estimate chroma plane */
             bits += estimate_slice_plane(ctx, &error, i,
                                          src, linesize[i],
                                          mbs_per_slice,
                                          num_cblocks[i], plane_factor[i],
-                                         ctx->quants[q], td);
+                                         ctx->quants_chroma[q], td);
         }
         if (bits > 65000 * 8)
             error = SCORE_LIMIT;
@@ -848,17 +898,26 @@ static int find_slice_quant(AVCodecContext *avctx,
             error = 0;
             if (q < MAX_STORED_Q) {
                 qmat = ctx->quants[q];
+                qmat_chroma = ctx->quants_chroma[q];
             } else {
                 qmat = td->custom_q;
-                for (i = 0; i < 64; i++)
+                qmat_chroma = td->custom_q;
+                for (i = 0; i < 64; i++) {
                     qmat[i] = ctx->quant_mat[i] * q;
+                    qmat_chroma[i] = ctx->quant_chroma_mat[i] * q;
+                }
             }
-            for (i = 0; i < ctx->num_planes - !!ctx->alpha_bits; i++) {
+            bits += estimate_slice_plane(ctx, &error, 0,
+                                         src, linesize[0],
+                                         mbs_per_slice,
+                                         num_cblocks[0], plane_factor[0],
+                                         qmat, td);/* estimate luma plane */
+            for (i = 1; i < ctx->num_planes - !!ctx->alpha_bits; i++) { /* estimate chroma plane */
                 bits += estimate_slice_plane(ctx, &error, i,
                                              src, linesize[i],
                                              mbs_per_slice,
                                              num_cblocks[i], plane_factor[i],
-                                             qmat, td);
+                                             qmat_chroma, td);
             }
             if (bits <= ctx->bits_per_mb * mbs_per_slice)
                 break;
@@ -1194,10 +1253,13 @@ FF_ENABLE_DEPRECATION_WARNINGS
     ctx->slices_per_picture = ctx->mb_height * ctx->slices_width;
     ctx->pictures_per_frame = 1 + interlaced;
 
-    if (ctx->quant_sel == -1)
+    if (ctx->quant_sel == -1) {
         ctx->quant_mat = prores_quant_matrices[ctx->profile_info->quant];
-    else
+        ctx->quant_chroma_mat = prores_quant_matrices[ctx->profile_info->quant_chroma];
+    } else {
         ctx->quant_mat = prores_quant_matrices[ctx->quant_sel];
+        ctx->quant_chroma_mat = prores_quant_matrices[ctx->quant_sel];
+    }
 
     if (strlen(ctx->vendor) != 4) {
         av_log(avctx, AV_LOG_ERROR, "vendor ID should be 4 bytes\n");
@@ -1222,8 +1284,10 @@ FF_ENABLE_DEPRECATION_WARNINGS
         min_quant = ctx->profile_info->min_quant;
         max_quant = ctx->profile_info->max_quant;
         for (i = min_quant; i < MAX_STORED_Q; i++) {
-            for (j = 0; j < 64; j++)
+            for (j = 0; j < 64; j++) {
                 ctx->quants[i][j] = ctx->quant_mat[j] * i;
+                ctx->quants_chroma[i][j] = ctx->quant_chroma_mat[j] * i;
+            }
         }
 
         ctx->slice_q = av_malloc(ctx->slices_per_picture * sizeof(*ctx->slice_q));
@@ -1254,6 +1318,7 @@ FF_ENABLE_DEPRECATION_WARNINGS
         }
     } else {
         int ls = 0;
+        int ls_chroma = 0;
 
         if (ctx->force_quant > 64) {
             av_log(avctx, AV_LOG_ERROR, "too large quantiser, maximum is 64\n");
@@ -1262,12 +1327,14 @@ FF_ENABLE_DEPRECATION_WARNINGS
 
         for (j = 0; j < 64; j++) {
             ctx->quants[0][j] = ctx->quant_mat[j] * ctx->force_quant;
+            ctx->quants_chroma[0][j] = ctx->quant_chroma_mat[j] * ctx->force_quant;
             ls += av_log2((1 << 11)  / ctx->quants[0][j]) * 2 + 1;
+            ls_chroma += av_log2((1 << 11)  / ctx->quants_chroma[0][j]) * 2 + 1;
         }
 
-        ctx->bits_per_mb = ls * 8;
+        ctx->bits_per_mb = ls * 4 + ls_chroma * 4;
         if (ctx->chroma_factor == CFACTOR_Y444)
-            ctx->bits_per_mb += ls * 4;
+            ctx->bits_per_mb += ls_chroma * 4;
     }
 
     ctx->frame_size_upper_bound = (ctx->pictures_per_frame *