[FFmpeg-devel] [PATCH v2 3/3] avcodec/nvenc: add more sei support for SDR HLG

[lance.lmwang at gmail.com]

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

The testing command for the atc HLG output with nvenc:

./ffmpeg_g -y -i 4K.mp4 -c:v hevc_nvenc -g 7 -color_primaries bt2020 -colorspace bt2020_ncl -color_trc arib-std-b67 -sei hlg test.ts

Signed-off-by: Limin Wang <lance.lmwang at gmail.com>
---
 libavcodec/nvenc.c      | 22 ++++++++++++++++++++++
 libavcodec/nvenc.h      |  1 +
 libavcodec/nvenc_hevc.c |  3 +++
 3 files changed, 26 insertions(+)

diff --git a/libavcodec/nvenc.c b/libavcodec/nvenc.c
index 30cf4df928..72a7b7f217 100644
--- a/libavcodec/nvenc.c
+++ b/libavcodec/nvenc.c
@@ -1018,6 +1018,10 @@ static av_cold int nvenc_setup_hevc_config(AVCodecContext *avctx)
     vui->colourMatrix = avctx->colorspace;
     vui->colourPrimaries = avctx->color_primaries;
     vui->transferCharacteristics = avctx->color_trc;
+    if (ctx->sei & SEI_ALTER_TRANSFER_CHAR) {
+        if (IS_10BIT(ctx->data_pix_fmt))
+            vui->transferCharacteristics = AVCOL_TRC_BT2020_10;
+    }
     vui->videoFullRangeFlag = (avctx->color_range == AVCOL_RANGE_JPEG
         || ctx->data_pix_fmt == AV_PIX_FMT_YUVJ420P || ctx->data_pix_fmt == AV_PIX_FMT_YUVJ422P || ctx->data_pix_fmt == AV_PIX_FMT_YUVJ444P);
 
@@ -2224,6 +2228,24 @@ int ff_nvenc_send_frame(AVCodecContext *avctx, const AVFrame *frame)
             }
         }
 
+        if (ctx->sei & SEI_ALTER_TRANSFER_CHAR) {
+            if (avctx->codec->id == AV_CODEC_ID_HEVC) {
+                size_t transfer_sei_size = 1;
+
+                sei_data[sei_count].payloadSize = transfer_sei_size;
+                sei_data[sei_count].payloadType = HEVC_SEI_TYPE_ALTERNATIVE_TRANSFER_CHARACTERISTICS;
+                sei_data[sei_count].payload = av_malloc(transfer_sei_size);
+                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, 8, AVCOL_TRC_ARIB_STD_B67);
+                    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 583c48d090..8be0d341ea 100644
--- a/libavcodec/nvenc.h
+++ b/libavcodec/nvenc.h
@@ -55,6 +55,7 @@ typedef void ID3D11Device;
 #endif
 
 enum {
+    SEI_ALTER_TRANSFER_CHAR     = 0x04,
     SEI_MASTERING_DISPLAY       = 0x08,
     SEI_CONTENT_LIGHT_LEVEL     = 0x10,
 };
diff --git a/libavcodec/nvenc_hevc.c b/libavcodec/nvenc_hevc.c
index 74ebd03d8e..23ba593439 100644
--- a/libavcodec/nvenc_hevc.c
+++ b/libavcodec/nvenc_hevc.c
@@ -134,6 +134,9 @@ static const AVOption options[] = {
     { "hdr10","Include HDR metadata for mastering display colour volume and content light level information",
                                                             0, AV_OPT_TYPE_CONST,   { .i64 = SEI_MASTERING_DISPLAY | SEI_CONTENT_LIGHT_LEVEL },
                                                             0, 0, VE, "sei" },
+    { "hlg","Include HLG metadata for alternative transfer characteristics",
+                                                            0, AV_OPT_TYPE_CONST,   { .i64 = SEI_ALTER_TRANSFER_CHAR },
+                                                            0, 0, VE, "sei" },
     { "master_display",     "SMPTE ST 2086 master display color volume info SEI (HDR), the string format is: \"G(x,y)B(x,y)R(x,y)WP(x,y)L(max,min)\"",
                                                             OFFSET(master_display), AV_OPT_TYPE_STRING,   { .str = NULL }, 0, 0, VE },
     { "max_cll",             "content light level info, the string format is: \"cll, fall\"",
-- 
2.21.0