[FFmpeg-devel] [PATCH] avcodec/libaomenc: fix encoding of sRGB streams

[James Almer]

Signed-off-by: James Almer <jamrial at gmail.com>
---
 libavcodec/libaomenc.c | 25 ++++++++++++++++++++++---
 1 file changed, 22 insertions(+), 3 deletions(-)

diff --git a/libavcodec/libaomenc.c b/libavcodec/libaomenc.c
index 8ebdcc20e3..fb9d60527b 100644
--- a/libavcodec/libaomenc.c
+++ b/libavcodec/libaomenc.c
@@ -67,6 +67,7 @@ typedef struct AOMEncoderContext {
     int static_thresh;
     int drop_threshold;
     int noise_sensitivity;
+    int is_rgb;
 } AOMContext;
 
 static const char *const ctlidstr[] = {
@@ -221,6 +222,7 @@ static int set_pix_fmt(AVCodecContext *avctx, aom_codec_caps_t codec_caps,
         *img_fmt = AOM_IMG_FMT_I422;
         return 0;
     case AV_PIX_FMT_GBRP:
+        ctx->is_rgb = 1;
     case AV_PIX_FMT_YUV444P:
         enccfg->g_profile = FF_PROFILE_AV1_HIGH;
         *img_fmt = AOM_IMG_FMT_I444;
@@ -250,6 +252,7 @@ static int set_pix_fmt(AVCodecContext *avctx, aom_codec_caps_t codec_caps,
         break;
     case AV_PIX_FMT_GBRP10:
     case AV_PIX_FMT_GBRP12:
+        ctx->is_rgb = 1;
     case AV_PIX_FMT_YUV444P10:
     case AV_PIX_FMT_YUV444P12:
         if (codec_caps & AOM_CODEC_CAP_HIGHBITDEPTH) {
@@ -270,6 +273,24 @@ static int set_pix_fmt(AVCodecContext *avctx, aom_codec_caps_t codec_caps,
     return AVERROR_INVALIDDATA;
 }
 
+static void set_colorspace(AVCodecContext *avctx)
+{
+    AOMContext *ctx = avctx->priv_data;
+    enum aom_color_primaries aom_cp = avctx->color_primaries;
+    enum aom_transfer_characteristics aom_tc = avctx->color_trc;
+    enum aom_matrix_coefficients aom_mc = avctx->colorspace;
+
+    if (ctx->is_rgb) {
+        // Forces the proper sRGB codepath in libaom.
+        aom_cp = AOM_CICP_CP_BT_709;
+        aom_tc = AOM_CICP_TC_SRGB;
+        aom_mc = AOM_CICP_MC_IDENTITY;
+    }
+    codecctl_int(avctx, AV1E_SET_COLOR_PRIMARIES, aom_cp);
+    codecctl_int(avctx, AV1E_SET_TRANSFER_CHARACTERISTICS, aom_tc);
+    codecctl_int(avctx, AV1E_SET_MATRIX_COEFFICIENTS, aom_mc);
+}
+
 static void set_color_range(AVCodecContext *avctx)
 {
     enum aom_color_range aom_cr;
@@ -451,9 +472,7 @@ static av_cold int aom_init(AVCodecContext *avctx,
     if (ctx->crf >= 0)
         codecctl_int(avctx, AOME_SET_CQ_LEVEL,          ctx->crf);
 
-    codecctl_int(avctx, AV1E_SET_COLOR_PRIMARIES, avctx->color_primaries);
-    codecctl_int(avctx, AV1E_SET_MATRIX_COEFFICIENTS, avctx->colorspace);
-    codecctl_int(avctx, AV1E_SET_TRANSFER_CHARACTERISTICS, avctx->color_trc);
+    set_colorspace(avctx);
     set_color_range(avctx);
 
     // provide dummy value to initialize wrapper, values will be updated each _encode()
-- 
2.16.2