[FFmpeg-devel] [PATCH 2/2] vaapi_encode_h265: Insert mastering display colour colume if needed

[Haihao Xiang]

'-sei xxx' is added to control SEI insertion, so far only mastering
display colour colume is available for testing. Another patch is
required to change mastering display settings later.

Signed-off-by: Haihao Xiang <haihao.xiang at intel.com>
---
 libavcodec/vaapi_encode_h265.c | 94 +++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 93 insertions(+), 1 deletion(-)

diff --git a/libavcodec/vaapi_encode_h265.c b/libavcodec/vaapi_encode_h265.c
index 5203c6871d..a8cb6a6d8b 100644
--- a/libavcodec/vaapi_encode_h265.c
+++ b/libavcodec/vaapi_encode_h265.c
@@ -29,10 +29,14 @@
 #include "cbs.h"
 #include "cbs_h265.h"
 #include "hevc.h"
+#include "hevc_sei.h"
 #include "internal.h"
 #include "put_bits.h"
 #include "vaapi_encode.h"
 
+enum {
+    SEI_MASTERING_DISPLAY       = 0x01,
+};
 
 typedef struct VAAPIEncodeH265Context {
     unsigned int ctu_width;
@@ -47,6 +51,9 @@ typedef struct VAAPIEncodeH265Context {
     H265RawSPS sps;
     H265RawPPS pps;
     H265RawSlice slice;
+    H265RawSEI sei;
+
+    H265RawSEIMasteringDiplayColorVolume mastering_display;
 
     int64_t last_idr_frame;
     int pic_order_cnt;
@@ -58,6 +65,7 @@ typedef struct VAAPIEncodeH265Context {
     CodedBitstreamContext *cbc;
     CodedBitstreamFragment current_access_unit;
     int aud_needed;
+    int sei_needed;
 } VAAPIEncodeH265Context;
 
 typedef struct VAAPIEncodeH265Options {
@@ -65,6 +73,7 @@ typedef struct VAAPIEncodeH265Options {
     int aud;
     int profile;
     int level;
+    int sei;
 } VAAPIEncodeH265Options;
 
 
@@ -175,6 +184,61 @@ fail:
     return err;
 }
 
+static int vaapi_encode_h265_write_extra_header(AVCodecContext *avctx,
+                                                VAAPIEncodePicture *pic,
+                                                int index, int *type,
+                                                char *data, size_t *data_len)
+{
+    VAAPIEncodeContext      *ctx = avctx->priv_data;
+    VAAPIEncodeH265Context *priv = ctx->priv_data;
+    VAAPIEncodeH265Options  *opt = ctx->codec_options;
+    CodedBitstreamFragment   *au = &priv->current_access_unit;
+    int err, i;
+
+    if (priv->sei_needed) {
+        if (priv->aud_needed) {
+            err = vaapi_encode_h265_add_nal(avctx, au, &priv->aud);
+            if (err < 0)
+                goto fail;
+            priv->aud_needed = 0;
+        }
+
+        memset(&priv->sei, 0, sizeof(priv->sei));
+        priv->sei.nal_unit_header.nal_unit_type = HEVC_NAL_SEI_PREFIX;
+        priv->sei.nal_unit_header.nuh_temporal_id_plus1 = 1;
+        i = 0;
+
+        if (opt->sei & SEI_MASTERING_DISPLAY && pic->type == PICTURE_TYPE_IDR) {
+            priv->sei.payload[i].payload_type = HEVC_SEI_TYPE_MASTERING_DISPLAY_INFO;
+            priv->sei.payload[i].payload.mastering_display = priv->mastering_display;
+            ++i;
+        }
+
+        priv->sei.payload_count = i;
+        av_assert0(priv->sei.payload_count > 0);
+
+        err = vaapi_encode_h265_add_nal(avctx, au, &priv->sei);
+        if (err < 0)
+            goto fail;
+        priv->sei_needed = 0;
+
+        err = vaapi_encode_h265_write_access_unit(avctx, data, data_len, au);
+        if (err < 0)
+            goto fail;
+
+        ff_cbs_fragment_uninit(priv->cbc, au);
+
+        *type = VAEncPackedHeaderRawData;
+        return 0;
+    } else {
+        return AVERROR_EOF;
+    }
+
+fail:
+    ff_cbs_fragment_uninit(priv->cbc, au);
+    return err;
+}
+
 static int vaapi_encode_h265_init_sequence_params(AVCodecContext *avctx)
 {
     VAAPIEncodeContext                *ctx = avctx->priv_data;
@@ -587,6 +651,23 @@ static int vaapi_encode_h265_init_picture_params(AVCodecContext *avctx,
         priv->aud_needed = 0;
     }
 
+    if ((opt->sei & SEI_MASTERING_DISPLAY) &&
+        (pic->type == PICTURE_TYPE_I || pic->type == PICTURE_TYPE_IDR)) {
+        // hard-coded value for testing, change it later
+        for (i = 0; i < 3; i++) {
+            priv->mastering_display.display_primaries[i].x = 50000;
+            priv->mastering_display.display_primaries[i].y = 50000;
+        }
+
+        priv->mastering_display.white_point_x = 50000;
+        priv->mastering_display.white_point_y = 50000;
+
+        priv->mastering_display.max_display_mastering_luminance = 5000;
+        priv->mastering_display.min_display_mastering_luminance = 1000;
+
+        priv->sei_needed = 1;
+    }
+
     vpic->decoded_curr_pic = (VAPictureHEVC) {
         .picture_id    = pic->recon_surface,
         .pic_order_cnt = priv->pic_order_cnt,
@@ -895,6 +976,8 @@ static const VAAPIEncodeType vaapi_encode_type_h265 = {
 
     .slice_header_type     = VAEncPackedHeaderHEVC_Slice,
     .write_slice_header    = &vaapi_encode_h265_write_slice_header,
+
+    .write_extra_header    = &vaapi_encode_h265_write_extra_header,
 };
 
 static av_cold int vaapi_encode_h265_init(AVCodecContext *avctx)
@@ -943,7 +1026,8 @@ static av_cold int vaapi_encode_h265_init(AVCodecContext *avctx)
 
     ctx->va_packed_headers =
         VA_ENC_PACKED_HEADER_SEQUENCE | // VPS, SPS and PPS.
-        VA_ENC_PACKED_HEADER_SLICE;     // Slice headers.
+        VA_ENC_PACKED_HEADER_SLICE |    // Slice headers.
+        VA_ENC_PACKED_HEADER_MISC;      // SEI
 
     ctx->surface_width  = FFALIGN(avctx->width,  16);
     ctx->surface_height = FFALIGN(avctx->height, 16);
@@ -1003,6 +1087,14 @@ static const AVOption vaapi_encode_h265_options[] = {
     { LEVEL("6.2", 186) },
 #undef LEVEL
 
+    { "sei", "Set SEI to include",
+      OFFSET(sei), AV_OPT_TYPE_FLAGS,
+      { .i64 = 0 },
+      0, INT_MAX, FLAGS, "sei" },
+    { "mastering_display", "Include mastering display colour volumne",
+      0, AV_OPT_TYPE_CONST, { .i64 = SEI_MASTERING_DISPLAY },
+      INT_MIN, INT_MAX, FLAGS, "sei" },
+
     { NULL },
 };
 
-- 
2.14.1