From 1a02e726f50ae2ce007854ee473a9f69617e2947 Mon Sep 17 00:00:00 2001
From: Jonas Karlman <jonas@kwiboo.se>
Date: Sat, 22 Jun 2019 16:08:25 +0000
Subject: [PATCH] rkmppdec: set mastering display and content light side data

---
 libavcodec/rkmppdec.c | 62 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 62 insertions(+)

diff --git a/libavcodec/rkmppdec.c b/libavcodec/rkmppdec.c
index 48016bc681..74938443df 100644
--- a/libavcodec/rkmppdec.c
+++ b/libavcodec/rkmppdec.c
@@ -35,6 +35,7 @@
 #include "libavutil/frame.h"
 #include "libavutil/hwcontext.h"
 #include "libavutil/hwcontext_drm.h"
+#include "libavutil/mastering_display_metadata.h"
 #include "libavutil/imgutils.h"
 #include "libavutil/log.h"
 
@@ -319,6 +320,59 @@ static void rkmpp_release_frame(void *opaque, uint8_t *data)
     av_free(desc);
 }
 
+static void rkmpp_hevc_mastering_display(AVFrame *frame, MppFrameMasteringDisplayMetadata mastering_display)
+{
+    // HEVC uses a g,b,r ordering, which we convert to a more natural r,g,b
+    AVMasteringDisplayMetadata *metadata;
+    const int mapping[3] = {2, 0, 1};
+    const int chroma_den = 50000;
+    const int luma_den = 10000;
+    int i;
+
+    AVFrameSideData *sd = av_frame_get_side_data(frame, AV_FRAME_DATA_MASTERING_DISPLAY_METADATA);
+    if (sd)
+        metadata = (AVMasteringDisplayMetadata *)sd->data;
+    else
+        metadata = av_mastering_display_metadata_create_side_data(frame);
+    if (!metadata)
+        return AVERROR(ENOMEM);
+
+    for (i = 0; i < 3; i++) {
+        const int j = mapping[i];
+        metadata->display_primaries[i][0].num = mastering_display.display_primaries[j][0];
+        metadata->display_primaries[i][0].den = chroma_den;
+        metadata->display_primaries[i][1].num = mastering_display.display_primaries[j][1];
+        metadata->display_primaries[i][1].den = chroma_den;
+    }
+    metadata->white_point[0].num = mastering_display.white_point[0];
+    metadata->white_point[0].den = chroma_den;
+    metadata->white_point[1].num = mastering_display.white_point[1];
+    metadata->white_point[1].den = chroma_den;
+
+    metadata->max_luminance.num = mastering_display.max_luminance;
+    metadata->max_luminance.den = luma_den;
+    metadata->min_luminance.num = mastering_display.min_luminance;
+    metadata->min_luminance.den = luma_den;
+    metadata->has_luminance = 1;
+    metadata->has_primaries = 1;
+}
+
+static void rkmpp_hevc_content_light(AVFrame *frame, MppFrameContentLightMetadata content_light)
+{
+    AVContentLightMetadata *metadata;
+
+    AVFrameSideData *sd = av_frame_get_side_data(frame, AV_FRAME_DATA_CONTENT_LIGHT_LEVEL);
+    if (sd)
+        metadata = (AVContentLightMetadata *)sd->data;
+    else
+        metadata = av_content_light_metadata_create_side_data(frame);
+    if (!metadata)
+        return AVERROR(ENOMEM);
+
+    metadata->MaxCLL  = content_light.MaxCLL;
+    metadata->MaxFALL = content_light.MaxFALL;
+}
+
 static int rkmpp_retrieve_frame(AVCodecContext *avctx, AVFrame *frame)
 {
     RKMPPDecodeContext *rk_context = avctx->priv_data;
@@ -352,6 +406,11 @@ static int rkmpp_retrieve_frame(AVCodecContext *avctx, AVFrame *frame)
             avctx->width = mpp_frame_get_width(mppframe);
             avctx->height = mpp_frame_get_height(mppframe);
 
+            avctx->color_range = mpp_frame_get_color_range(mppframe);
+            avctx->color_primaries = mpp_frame_get_color_primaries(mppframe);
+            avctx->color_trc = mpp_frame_get_color_trc(mppframe);
+            avctx->colorspace = mpp_frame_get_colorspace(mppframe);
+
             decoder->mpi->control(decoder->ctx, MPP_DEC_SET_INFO_CHANGE_READY, NULL);
 
             av_buffer_unref(&decoder->frames_ref);
@@ -410,6 +469,9 @@ static int rkmpp_retrieve_frame(AVCodecContext *avctx, AVFrame *frame)
             MppFrameRational dar = mpp_frame_get_sar(mppframe);
             frame->sample_aspect_ratio = av_div_q((AVRational) { dar.num, dar.den },
                                                   (AVRational) { frame->width, frame->height });
+        } else if (avctx->codec_id == AV_CODEC_ID_HEVC && frame->color_trc > AVCOL_TRC_UNSPECIFIED) {
+            rkmpp_hevc_mastering_display(frame, mpp_frame_get_mastering_display(mppframe));
+            rkmpp_hevc_content_light(frame, mpp_frame_get_content_light(mppframe));
         }
 
         mode = mpp_frame_get_mode(mppframe);
