[FFmpeg-devel] [PATCH 2/2] avcodec/nvenc: add master display and light level sei data for hdr10

[lance.lmwang at gmail.com]

From: Limin Wang <lance.lmwang at gmail.com>

---
 libavcodec/nvenc.c      | 73 +++++++++++++++++++++++++++++++++++++++++
 libavcodec/nvenc.h      |  1 +
 libavcodec/nvenc_hevc.c |  2 ++
 3 files changed, 76 insertions(+)

diff --git a/libavcodec/nvenc.c b/libavcodec/nvenc.c
index 77eb9918a4..55547d4246 100644
--- a/libavcodec/nvenc.c
+++ b/libavcodec/nvenc.c
@@ -22,6 +22,9 @@
 #include "config.h"
 
 #include "nvenc.h"
+#include "cbs_h265.h"
+#include "hevc_sei.h"
+#include "put_bits.h"
 
 #include "libavutil/hwcontext_cuda.h"
 #include "libavutil/hwcontext.h"
@@ -30,6 +33,7 @@
 #include "libavutil/avassert.h"
 #include "libavutil/mem.h"
 #include "libavutil/pixdesc.h"
+#include "libavutil/mastering_display_metadata.h"
 #include "internal.h"
 
 #define CHECK_CU(x) FF_CUDA_CHECK_DL(avctx, dl_fn->cuda_dl, x)
@@ -2110,6 +2114,75 @@ int ff_nvenc_send_frame(AVCodecContext *avctx, const AVFrame *frame)
             }
         }
 
+        if (ctx->hdr) {
+            AVFrameSideData *sd = av_frame_get_side_data(frame, AV_FRAME_DATA_MASTERING_DISPLAY_METADATA);
+
+            if (sd) {
+                AVMasteringDisplayMetadata *mdm = (AVMasteringDisplayMetadata *)sd->data;
+                // HEVC uses a g,b,r ordering, which we convert to a more natural r,g,b
+                const int mapping[3] = {2, 0, 1};
+                const int chroma_den = 50000;
+                const int luma_den = 10000;
+
+                if (mdm->has_primaries && mdm->has_luminance) {
+                    H265RawSEIMasteringDisplayColourVolume smd;
+
+                    for (i = 0; i < 3; i++) {
+                        const int j = mapping[i];
+                        smd.display_primaries_x[i] = chroma_den * av_q2d(mdm->display_primaries[j][0]);
+                        smd.display_primaries_y[i] = chroma_den * av_q2d( mdm->display_primaries[j][1]);
+                    }
+
+                    smd.white_point_x = chroma_den * av_q2d(mdm->white_point[0]);
+                    smd.white_point_y = chroma_den * av_q2d(mdm->white_point[1]);
+                    smd.max_display_mastering_luminance = luma_den * av_q2d(mdm->max_luminance);
+                    smd.min_display_mastering_luminance = luma_den * av_q2d(mdm->min_luminance);
+
+                    sei_data[sei_count].payloadSize = sizeof(H265RawSEIMasteringDisplayColourVolume);
+                    sei_data[sei_count].payloadType = HEVC_SEI_TYPE_MASTERING_DISPLAY_INFO;
+                    sei_data[sei_count].payload     = av_mallocz(sei_data[sei_count].payloadSize);
+                    if (sei_data[sei_count].payload) {
+                        PutBitContext pb;
+
+                        init_put_bits(&pb, sei_data[sei_count].payload, sei_data[sei_count].payloadSize);
+                        for (i = 0; i < 3; i++) {
+                            put_bits(&pb, 16, smd.display_primaries_x[i]);
+                            put_bits(&pb, 16, smd.display_primaries_y[i]);
+                        }
+                        put_bits(&pb, 16, smd.white_point_x);
+                        put_bits(&pb, 16, smd.white_point_y);
+                        put_bits(&pb, 32, smd.max_display_mastering_luminance);
+                        put_bits(&pb, 32, smd.min_display_mastering_luminance);
+                        flush_put_bits(&pb);
+
+                        sei_count ++;
+                    }
+                }
+            }
+
+            sd = av_frame_get_side_data(frame, AV_FRAME_DATA_CONTENT_LIGHT_LEVEL);
+            if (sd) {
+                AVContentLightMetadata *clm = (AVContentLightMetadata *)sd->data;
+                H265RawSEIContentLightLevelInfo clli;
+
+                clli.max_content_light_level     = FFMIN(clm->MaxCLL,  65535);
+                clli.max_pic_average_light_level = FFMIN(clm->MaxFALL, 65535);
+
+                sei_data[sei_count].payloadSize = sizeof(H265RawSEIContentLightLevelInfo);
+                sei_data[sei_count].payloadType = HEVC_SEI_TYPE_CONTENT_LIGHT_LEVEL_INFO;
+                sei_data[sei_count].payload     = av_mallocz(sei_data[sei_count].payloadSize);
+                if (sei_data[sei_count].payload) {
+                    PutBitContext pb;
+
+                    init_put_bits(&pb, sei_data[sei_count].payload, sei_data[sei_count].payloadSize);
+                    put_bits(&pb, 16, clli.max_content_light_level);
+                    put_bits(&pb, 16, clli.max_pic_average_light_level);
+                    flush_put_bits(&pb);
+
+                    sei_count ++;
+                }
+            }
+        }
         nvenc_codec_specific_pic_params(avctx, &pic_params, sei_data, sei_count);
     } else {
         pic_params.encodePicFlags = NV_ENC_PIC_FLAG_EOS;
diff --git a/libavcodec/nvenc.h b/libavcodec/nvenc.h
index ddd6168409..952e691147 100644
--- a/libavcodec/nvenc.h
+++ b/libavcodec/nvenc.h
@@ -192,6 +192,7 @@ typedef struct NvencContext
     int coder;
     int b_ref_mode;
     int a53_cc;
+    int hdr;
 } NvencContext;
 
 int ff_nvenc_encode_init(AVCodecContext *avctx);
diff --git a/libavcodec/nvenc_hevc.c b/libavcodec/nvenc_hevc.c
index d567d960ba..ea7580acb6 100644
--- a/libavcodec/nvenc_hevc.c
+++ b/libavcodec/nvenc_hevc.c
@@ -127,6 +127,8 @@ static const AVOption options[] = {
     { "each",         "",                                   0,                    AV_OPT_TYPE_CONST, { .i64 = 1 }, 0, 0,       VE, "b_ref_mode" },
     { "middle",       "",                                   0,                    AV_OPT_TYPE_CONST, { .i64 = 2 }, 0, 0,       VE, "b_ref_mode" },
 #endif
+    { "hdr","Include HDR metadata for mastering display colour volume and content light level information",
+                                                            OFFSET(hdr),AV_OPT_TYPE_INT,   { .i64 = 0 }, 0, 1, VE },
     { NULL }
 };
 
-- 
2.21.0