cbs_h265: read/write HEVC PREFIX SEI

Similar to H264, cbs_h265_{read, write}_nal_unit() can handle HEVC
prefix SEI NAL units. Currently mastering display colour volume SEI
message is added only, we may add more SEI message if needed later

Signed-off-by: Haihao Xiang <haihao.xiang@intel.com>

3 files changed, 239 insertions, 0 deletions

diff --git a/libavcodec/cbs_h2645.c b/libavcodec/cbs_h2645.c
index 1ce42c6..dd13d8a 100644
--- a/libavcodec/cbs_h2645.c
+++ b/libavcodec/cbs_h2645.c
@@ -29,6 +29,7 @@
 #include "h264_sei.h"
 #include "h2645_parse.h"
 #include "hevc.h"
+#include "hevc_sei.h"
 
 
 static int cbs_read_ue_golomb(CodedBitstreamContext *ctx, GetBitContext *gbc,
@@ -483,6 +484,26 @@ static void cbs_h265_free_slice(void *unit, uint8_t *content)
     av_freep(&content);
 }
 
+static void cbs_h265_free_sei_payload(H265RawSEIPayload *payload)
+{
+    switch (payload->payload_type) {
+    case HEVC_SEI_TYPE_MASTERING_DISPLAY_INFO:
+        break;
+    default:
+        av_buffer_unref(&payload->payload.other.data_ref);
+        break;
+    }
+}
+
+static void cbs_h265_free_sei(void *unit, uint8_t *content)
+{
+    H265RawSEI *sei = (H265RawSEI*)content;
+    int i;
+    for (i = 0; i < sei->payload_count; i++)
+        cbs_h265_free_sei_payload(&sei->payload[i]);
+    av_freep(&content);
+}
+
 static int cbs_h2645_fragment_add_nals(CodedBitstreamContext *ctx,
                                        CodedBitstreamFragment *frag,
                                        const H2645Packet *packet)
@@ -986,6 +1007,21 @@ static int cbs_h265_read_nal_unit(CodedBitstreamContext *ctx,
         }
         break;
 
+    case HEVC_NAL_SEI_PREFIX:
+        {
+            err = ff_cbs_alloc_unit_content(ctx, unit, sizeof(H265RawSEI),
+                                            &cbs_h265_free_sei);
+
+            if (err < 0)
+                return err;
+
+            err = cbs_h265_read_sei(ctx, &gbc, unit->content);
+
+            if (err < 0)
+                return err;
+        }
+        break;
+
     default:
         return AVERROR(ENOSYS);
     }
@@ -1226,6 +1262,15 @@ static int cbs_h265_write_nal_unit(CodedBitstreamContext *ctx,
         }
         break;
 
+    case HEVC_NAL_SEI_PREFIX:
+        {
+            err = cbs_h265_write_sei(ctx, pbc, unit->content);
+
+            if (err < 0)
+                return err;
+        }
+        break;
+
     default:
         av_log(ctx->log_ctx, AV_LOG_ERROR, "Write unimplemented for "
                "NAL unit type %"PRIu32".\n", unit->type);
diff --git a/libavcodec/cbs_h265.h b/libavcodec/cbs_h265.h
index 6d02979..983357e 100644
--- a/libavcodec/cbs_h265.h
+++ b/libavcodec/cbs_h265.h
@@ -25,6 +25,14 @@
 #include "cbs_h2645.h"
 #include "hevc.h"
 
+enum {
+    // This limit is arbitrary - it is sufficient for one message of each
+    // type plus some repeats, and will therefore easily cover all sane
+    // streams.  However, it is possible to make technically-valid streams
+    // for which it will fail (for example, by including a large number of
+    // user-data-unregistered messages).
+    H265_MAX_SEI_PAYLOADS = 64,
+};
 
 typedef struct H265RawNALUnitHeader {
     uint8_t forbidden_zero_bit;
@@ -516,6 +524,34 @@ typedef struct H265RawSlice {
     AVBufferRef *data_ref;
 } H265RawSlice;
 
+typedef struct H265RawSEIMasteringDisplayColourVolume {
+    uint16_t display_primaries_x[3];
+    uint16_t display_primaries_y[3];
+    uint16_t white_point_x;
+    uint16_t white_point_y;
+    uint32_t max_display_mastering_luminance;
+    uint32_t min_display_mastering_luminance;
+} H265RawSEIMasteringDisplayColourVolume;
+
+typedef struct H265RawSEIPayload {
+    uint32_t payload_type;
+    uint32_t payload_size;
+    union {
+        H265RawSEIMasteringDisplayColourVolume mastering_display;
+        struct {
+            uint8_t *data;
+            size_t data_length;
+            AVBufferRef *data_ref;
+        } other;
+    } payload;
+} H265RawSEIPayload;
+
+typedef struct H265RawSEI {
+    H265RawNALUnitHeader nal_unit_header;
+
+    H265RawSEIPayload payload[H265_MAX_SEI_PAYLOADS];
+    uint8_t payload_count;
+} H265RawSEI;
 
 typedef struct CodedBitstreamH265Context {
     // Reader/writer context in common with the H.264 implementation.
diff --git a/libavcodec/cbs_h265_syntax_template.c b/libavcodec/cbs_h265_syntax_template.c
index 9f13061..1c67271 100644
--- a/libavcodec/cbs_h265_syntax_template.c
+++ b/libavcodec/cbs_h265_syntax_template.c
@@ -1504,3 +1504,161 @@ static int FUNC(slice_segment_header)(CodedBitstreamContext *ctx, RWContext *rw,
 
     return 0;
 }
+
+static int FUNC(sei_mastering_display)(CodedBitstreamContext *ctx, RWContext *rw,
+                                       H265RawSEIMasteringDisplayColourVolume *current)
+{
+    int err, c;
+
+    for (c = 0; c < 3; c++) {
+        us(16, display_primaries_x[c], 0, 50000, 1, c);
+        us(16, display_primaries_y[c], 0, 50000, 1, c);
+    }
+
+    u(16, white_point_x, 0, 50000);
+    u(16, white_point_y, 0, 50000);
+
+    u(32, max_display_mastering_luminance,
+      1, MAX_UINT_BITS(32));
+    u(32, min_display_mastering_luminance,
+      0, current->max_display_mastering_luminance - 1);
+
+    return 0;
+}
+
+static int FUNC(sei_payload)(CodedBitstreamContext *ctx, RWContext *rw,
+                             H265RawSEIPayload *current)
+{
+    int err, i;
+    int start_position, end_position;
+
+#ifdef READ
+    start_position = get_bits_count(rw);
+#else
+    start_position = put_bits_count(rw);
+#endif
+
+    switch (current->payload_type) {
+    case HEVC_SEI_TYPE_MASTERING_DISPLAY_INFO:
+        CHECK(FUNC(sei_mastering_display)
+              (ctx, rw, &current->payload.mastering_display));
+
+        break;
+
+    default:
+        {
+#ifdef READ
+            current->payload.other.data_length = current->payload_size;
+#endif
+            allocate(current->payload.other.data, current->payload.other.data_length);
+
+            for (i = 0; i < current->payload_size; i++)
+                xu(8, payload_byte[i], current->payload.other.data[i], 0, 255,
+                   1, i);
+        }
+    }
+
+    if (byte_alignment(rw)) {
+        fixed(1, bit_equal_to_one, 1);
+        while (byte_alignment(rw))
+            fixed(1, bit_equal_to_zero, 0);
+    }
+
+#ifdef READ
+    end_position = get_bits_count(rw);
+    if (end_position < start_position + 8 * current->payload_size) {
+        av_log(ctx->log_ctx, AV_LOG_ERROR, "Incorrect SEI payload length: "
+               "header %"PRIu32" bits, actually %d bits.\n",
+               8 * current->payload_size,
+               end_position - start_position);
+        return AVERROR_INVALIDDATA;
+    }
+#else
+    end_position = put_bits_count(rw);
+    current->payload_size = (end_position - start_position) >> 3;
+#endif
+
+    return 0;
+}
+
+static int FUNC(sei)(CodedBitstreamContext *ctx, RWContext *rw,
+                     H265RawSEI *current)
+{
+    int err, k;
+
+    HEADER("Supplemental Enhancement Information");
+
+    CHECK(FUNC(nal_unit_header)(ctx, rw, &current->nal_unit_header,
+                                HEVC_NAL_SEI_PREFIX));
+
+#ifdef READ
+    for (k = 0; k < H265_MAX_SEI_PAYLOADS; k++) {
+        uint32_t payload_type = 0;
+        uint32_t payload_size = 0;
+        uint32_t tmp;
+
+        while (show_bits(rw, 8) == 0xff) {
+            fixed(8, ff_byte, 0xff);
+            payload_type += 255;
+        }
+        xu(8, last_payload_type_byte, tmp, 0, 254, 0);
+        payload_type += tmp;
+
+        while (show_bits(rw, 8) == 0xff) {
+            fixed(8, ff_byte, 0xff);
+            payload_size += 255;
+        }
+        xu(8, last_payload_size_byte, tmp, 0, 254, 0);
+        payload_size += tmp;
+
+        current->payload[k].payload_type = payload_type;
+        current->payload[k].payload_size = payload_size;
+
+        CHECK(FUNC(sei_payload)(ctx, rw, &current->payload[k]));
+
+        if (!cbs_h2645_read_more_rbsp_data(rw))
+            break;
+    }
+    if (k >= H265_MAX_SEI_PAYLOADS) {
+        av_log(ctx->log_ctx, AV_LOG_ERROR, "Too many payloads in "
+               "SEI message: found %d.\n", k);
+        return AVERROR_INVALIDDATA;
+    }
+    current->payload_count = k + 1;
+#else
+    for (k = 0; k < current->payload_count; k++) {
+        PutBitContext start_state;
+        uint32_t tmp;
+        int need_size, i;
+
+        // Somewhat clumsy: we write the payload twice when
+        // we don't know the size in advance.  This will mess
+        // with trace output, but is otherwise harmless.
+        start_state = *rw;
+        need_size = !current->payload[k].payload_size;
+        for (i = 0; i < 1 + need_size; i++) {
+            *rw = start_state;
+
+            tmp = current->payload[k].payload_type;
+            while (tmp >= 255) {
+                fixed(8, ff_byte, 0xff);
+                tmp -= 255;
+            }
+            xu(8, last_payload_type_byte, tmp, 0, 254, 0);
+
+            tmp = current->payload[k].payload_size;
+            while (tmp >= 255) {
+                fixed(8, ff_byte, 0xff);
+                tmp -= 255;
+            }
+            xu(8, last_payload_size_byte, tmp, 0, 254, 0);
+
+            CHECK(FUNC(sei_payload)(ctx, rw, &current->payload[k]));
+        }
+    }
+#endif
+
+    CHECK(FUNC(rbsp_trailing_bits)(ctx, rw));
+
+    return 0;
+}