[FFmpeg-devel,v1] avcodec: add hdr sei support for hevc_nvenc

diff --git a/libavcodec/cbs.c b/libavcodec/cbs.c
index f98531e131..b3db4749bf 100644
--- a/libavcodec/cbs.c
+++ b/libavcodec/cbs.c
@@ -26,6 +26,7 @@ 
 
 #include "cbs.h"
 #include "cbs_internal.h"
+#include "cbs_h265.h"
 
 
 static const CodedBitstreamType *cbs_type_table[] = {
diff --git a/libavcodec/cbs_h2645.c b/libavcodec/cbs_h2645.c
index 434322492c..ca55c4f298 100644
--- a/libavcodec/cbs_h2645.c
+++ b/libavcodec/cbs_h2645.c
@@ -1538,6 +1538,66 @@  const CodedBitstreamType ff_cbs_type_h265 = {
     .close             = &cbs_h265_close,
 };
 
+int ff_cbs_h265_add_sei_message(CodedBitstreamFragment *au,
+  H265RawSEIPayload *payload)
+{
+  H265RawSEI *sei = NULL;
+  int err, i, k;
+
+  // Use single SEI NAL unit to store HDR info.
+  for (i = 0; i < au->nb_units; i++) {
+    if (au->units[i].type == HEVC_NAL_SEI_PREFIX || au->units[i].type == HEVC_NAL_SEI_SUFFIX) {
+      H265RawSEI * raw_sei = (H265RawSEI *)au->units[i].content;
+      if (raw_sei) {
+        for (k = 0; k < raw_sei->payload_count; k++) {
+          if (raw_sei->payload[k].payload_type == payload->payload_type) {
+            sei = au->units[i].content;
+            if (sei->payload_count < H265_MAX_SEI_PAYLOADS)
+              break;
+            sei = NULL;
+          }
+        }
+      }
+    }
+  }
+
+  if (!sei) {
+    AVBufferRef *sei_ref;
+
+    sei = av_mallocz(sizeof(*sei));
+    if (!sei) {
+      err = AVERROR(ENOMEM);
+      goto fail;
+    }
+
+    sei->nal_unit_header.nal_unit_type = HEVC_NAL_SEI_PREFIX;
+    sei->nal_unit_header.nuh_layer_id = 0;
+    sei->nal_unit_header.nuh_temporal_id_plus1 = 1;
+    sei_ref = av_buffer_create((uint8_t*)sei, sizeof(*sei),
+      &cbs_h265_free_sei, NULL, 0);
+    if (!sei_ref) {
+      av_freep(&sei);
+      err = AVERROR(ENOMEM);
+      goto fail;
+    }
+
+    err = ff_cbs_insert_unit_content(au, i, HEVC_NAL_SEI_PREFIX,
+      sei, sei_ref);
+    av_buffer_unref(&sei_ref);
+    if (err < 0)
+      goto fail;
+  }
+
+  memcpy(&sei->payload[sei->payload_count], payload, sizeof(*payload));
+  ++sei->payload_count;
+
+  return 0;
+fail:
+  cbs_h265_free_sei_payload(payload);
+  return err;
+}
+
+
 int ff_cbs_h264_add_sei_message(CodedBitstreamFragment *au,
                                 H264RawSEIPayload *payload)
 {
diff --git a/libavcodec/cbs_h265.h b/libavcodec/cbs_h265.h
index 15b22bbfd4..cd3824c7fa 100644
--- a/libavcodec/cbs_h265.h
+++ b/libavcodec/cbs_h265.h
@@ -24,6 +24,7 @@ 
 
 #include "cbs_h2645.h"
 #include "hevc.h"
+#include "cbs.h"
 
 enum {
     // This limit is arbitrary - it is sufficient for one message of each
@@ -746,5 +747,6 @@  typedef struct CodedBitstreamH265Context {
     const H265RawPPS *active_pps;
 } CodedBitstreamH265Context;
 
+int ff_cbs_h265_add_sei_message(CodedBitstreamFragment *au, H265RawSEIPayload *payload);
 
 #endif /* AVCODEC_CBS_H265_H */
diff --git a/libavcodec/h265_metadata_bsf.c b/libavcodec/h265_metadata_bsf.c
index 504a75dac2..91e5df05d2 100644
--- a/libavcodec/h265_metadata_bsf.c
+++ b/libavcodec/h265_metadata_bsf.c
@@ -25,6 +25,7 @@ 
 #include "cbs_h265.h"
 #include "hevc.h"
 #include "h265_profile_level.h"
+#include "hevc_sei.h"
 
 enum {
     PASS,
@@ -70,6 +71,11 @@  typedef struct H265MetadataContext {
     int level;
     int level_guess;
     int level_warned;
+
+    char * max_cll;
+    char * mastering_display_color_volume;
+    int done_first_au;
+
 } H265MetadataContext;
 
 
@@ -393,7 +399,7 @@  static int h265_metadata_filter(AVBSFContext *bsf, AVPacket *pkt)
 {
     H265MetadataContext *ctx = bsf->priv_data;
     CodedBitstreamFragment *au = &ctx->access_unit;
-    int err, i;
+    int err, i, has_sps;
 
     err = ff_bsf_get_packet_ref(bsf, pkt);
     if (err < 0)
@@ -471,15 +477,64 @@  static int h265_metadata_filter(AVBSFContext *bsf, AVPacket *pkt)
             err = h265_metadata_update_sps(bsf, au->units[i].content);
             if (err < 0)
                 goto fail;
+            has_sps = 1;
         }
     }
 
+     if (has_sps || !ctx->done_first_au) {
+      if (ctx->max_cll) {
+        H265RawSEIPayload payload = {
+             .payload_type = HEVC_SEI_TYPE_CONTENT_LIGHT_LEVEL_INFO,
+        };
+        payload.payload_size = 4;
+
+        if (sscanf(ctx->max_cll, "%hu|%hu",
+          &payload.payload.content_light_level.max_content_light_level,
+          &payload.payload.content_light_level.max_pic_average_light_level) == 2) {
+          err = ff_cbs_h265_add_sei_message(au, &payload);
+          if (err < 0) {
+            av_log(bsf, AV_LOG_ERROR, "Failed to add HEVC_SEI_TYPE_CONTENT_LIGHT_LEVEL_INFO SEI "
+              "message to access unit.\n");
+            goto fail;
+          }
+        }
+      }
+
+      if (ctx->mastering_display_color_volume) {
+        H265RawSEIPayload payload = {
+            .payload_type = HEVC_SEI_TYPE_MASTERING_DISPLAY_INFO,
+        };
+        payload.payload_size = 24;
+
+        if (sscanf(ctx->mastering_display_color_volume,
+          "G(%hu|%hu)B(%hu|%hu)R(%hu|%hu)WP(%hu|%hu)L(%u|%u)",
+          &payload.payload.mastering_display.display_primaries_x[0],
+          &payload.payload.mastering_display.display_primaries_y[0],
+          &payload.payload.mastering_display.display_primaries_x[1],
+          &payload.payload.mastering_display.display_primaries_y[1],
+          &payload.payload.mastering_display.display_primaries_x[2],
+          &payload.payload.mastering_display.display_primaries_y[2],
+          &payload.payload.mastering_display.white_point_x,
+          &payload.payload.mastering_display.white_point_y,
+          &payload.payload.mastering_display.max_display_mastering_luminance,
+          &payload.payload.mastering_display.min_display_mastering_luminance) == 10) {
+          err = ff_cbs_h265_add_sei_message(au, &payload);
+          if (err < 0) {
+            av_log(bsf, AV_LOG_ERROR, "Failed to add HEVC_SEI_TYPE_MASTERING_DISPLAY_INFO SEI "
+              "message to access unit.\n");
+            goto fail;
+          }
+        }
+      }
+    }
+
     err = ff_cbs_write_packet(ctx->output, pkt, au);
     if (err < 0) {
         av_log(bsf, AV_LOG_ERROR, "Failed to write packet.\n");
         goto fail;
     }
 
+    ctx->done_first_au = 1;
     err = 0;
 fail:
     ff_cbs_fragment_reset(au);
@@ -607,6 +662,13 @@  static const AVOption h265_metadata_options[] = {
         OFFSET(crop_bottom), AV_OPT_TYPE_INT,
         { .i64 = -1 }, -1, HEVC_MAX_HEIGHT, FLAGS },
 
+    { "max_cll", "Content light level for hdr. example 1000|800",
+        OFFSET(max_cll), AV_OPT_TYPE_STRING,
+        { .str = NULL }, .flags = FLAGS },
+    { "master_display", "Mastering display color volume for hdr.",
+        OFFSET(mastering_display_color_volume), AV_OPT_TYPE_STRING,
+        { .str = NULL }, .flags = FLAGS },
+
     { "level", "Set level (tables A.6 and A.7)",
         OFFSET(level), AV_OPT_TYPE_INT,
         { .i64 = LEVEL_UNSET }, LEVEL_UNSET, 0xff, FLAGS, "level" },