vaapi_h264: Add support for SEI messages

Send buffering_period and pic_timing messages when in modes
targetting bitrate.  Also adds NAL HRD parameters to VUI.

1 files changed, 206 insertions, 1 deletions

diff --git a/libavcodec/vaapi_encode_h264.c b/libavcodec/vaapi_encode_h264.c
index 139e6e942a..d991469b43 100644
--- a/libavcodec/vaapi_encode_h264.c
+++ b/libavcodec/vaapi_encode_h264.c
@@ -26,6 +26,7 @@
 
 #include "avcodec.h"
 #include "h264.h"
+#include "h264_sei.h"
 #include "internal.h"
 #include "vaapi_encode.h"
 #include "vaapi_encode_h26x.h"
@@ -84,6 +85,22 @@ typedef struct VAAPIEncodeH264MiscSequenceParams {
     char low_delay_hrd_flag;
     char pic_struct_present_flag;
     char bitstream_restriction_flag;
+
+    unsigned int cpb_cnt_minus1;
+    unsigned int bit_rate_scale;
+    unsigned int cpb_size_scale;
+    unsigned int bit_rate_value_minus1[32];
+    unsigned int cpb_size_value_minus1[32];
+    char cbr_flag[32];
+    unsigned int initial_cpb_removal_delay_length_minus1;
+    unsigned int cpb_removal_delay_length_minus1;
+    unsigned int dpb_output_delay_length_minus1;
+    unsigned int time_offset_length;
+
+    unsigned int initial_cpb_removal_delay;
+    unsigned int initial_cpb_removal_delay_offset;
+
+    unsigned int pic_struct;
 } VAAPIEncodeH264MiscSequenceParams;
 
 // This structure contains all possibly-useful per-slice syntax elements
@@ -126,7 +143,11 @@ typedef struct VAAPIEncodeH264Context {
     int next_frame_num;
     int64_t idr_pic_count;
 
+    int cpb_delay;
+    int dpb_delay;
+
     // Rate control configuration.
+    int send_timing_sei;
     struct {
         VAEncMiscParameterBuffer misc;
         VAEncMiscParameterRateControl rc;
@@ -183,6 +204,7 @@ static void vaapi_encode_h264_write_vui(PutBitContext *pbc,
     VAEncSequenceParameterBufferH264  *vseq = ctx->codec_sequence_params;
     VAAPIEncodeH264Context            *priv = ctx->priv_data;
     VAAPIEncodeH264MiscSequenceParams *mseq = &priv->misc_sequence_params;
+    int i;
 
     u(1, vvui_field(aspect_ratio_info_present_flag));
     if (vseq->vui_fields.bits.aspect_ratio_info_present_flag) {
@@ -224,7 +246,18 @@ static void vaapi_encode_h264_write_vui(PutBitContext *pbc,
 
     u(1, mseq_var(nal_hrd_parameters_present_flag));
     if (mseq->nal_hrd_parameters_present_flag) {
-        av_assert0(0 && "nal hrd parameters not supported");
+        ue(mseq_var(cpb_cnt_minus1));
+        u(4, mseq_var(bit_rate_scale));
+        u(4, mseq_var(cpb_size_scale));
+        for (i = 0; i <= mseq->cpb_cnt_minus1; i++) {
+            ue(mseq_var(bit_rate_value_minus1[i]));
+            ue(mseq_var(cpb_size_value_minus1[i]));
+            u(1, mseq_var(cbr_flag[i]));
+        }
+        u(5, mseq_var(initial_cpb_removal_delay_length_minus1));
+        u(5, mseq_var(cpb_removal_delay_length_minus1));
+        u(5, mseq_var(dpb_output_delay_length_minus1));
+        u(5, mseq_var(time_offset_length));
     }
     u(1, mseq_var(vcl_hrd_parameters_present_flag));
     if (mseq->vcl_hrd_parameters_present_flag) {
@@ -519,6 +552,110 @@ static void vaapi_encode_h264_write_slice_header2(PutBitContext *pbc,
     // No alignment - this need not be a byte boundary.
 }
 
+static void vaapi_encode_h264_write_buffering_period(PutBitContext *pbc,
+                                                     VAAPIEncodeContext *ctx,
+                                                     VAAPIEncodePicture *pic)
+{
+    VAAPIEncodeH264Context            *priv = ctx->priv_data;
+    VAAPIEncodeH264MiscSequenceParams *mseq = &priv->misc_sequence_params;
+    VAEncPictureParameterBufferH264   *vpic = pic->codec_picture_params;
+    int i;
+
+    ue(vpic_var(seq_parameter_set_id));
+
+    if (mseq->nal_hrd_parameters_present_flag) {
+        for (i = 0; i <= mseq->cpb_cnt_minus1; i++) {
+            u(mseq->initial_cpb_removal_delay_length_minus1 + 1,
+              mseq_var(initial_cpb_removal_delay));
+            u(mseq->initial_cpb_removal_delay_length_minus1 + 1,
+              mseq_var(initial_cpb_removal_delay_offset));
+        }
+    }
+    if (mseq->vcl_hrd_parameters_present_flag) {
+        av_assert0(0 && "vcl hrd parameters not supported");
+    }
+}
+
+static void vaapi_encode_h264_write_pic_timing(PutBitContext *pbc,
+                                               VAAPIEncodeContext *ctx,
+                                               VAAPIEncodePicture *pic)
+{
+    VAEncSequenceParameterBufferH264  *vseq = ctx->codec_sequence_params;
+    VAAPIEncodeH264Context            *priv = ctx->priv_data;
+    VAAPIEncodeH264MiscSequenceParams *mseq = &priv->misc_sequence_params;
+    int i, num_clock_ts;
+
+    if (mseq->nal_hrd_parameters_present_flag ||
+        mseq->vcl_hrd_parameters_present_flag) {
+        u(mseq->cpb_removal_delay_length_minus1 + 1,
+          2 * vseq->num_units_in_tick * priv->cpb_delay,
+          cpb_removal_delay);
+        u(mseq->dpb_output_delay_length_minus1 + 1,
+          2 * vseq->num_units_in_tick * priv->dpb_delay,
+          dpb_output_delay);
+    }
+    if (mseq->pic_struct_present_flag) {
+        u(4, mseq_var(pic_struct));
+        num_clock_ts = (mseq->pic_struct <= 2 ? 1 :
+                        mseq->pic_struct <= 4 ? 2 :
+                        mseq->pic_struct <= 8 ? 3 : 0);
+        for (i = 0; i < num_clock_ts; i++) {
+            u(1, 0, clock_timestamp_flag[i]);
+            // No full timestamp information.
+        }
+    }
+}
+
+static void vaapi_encode_h264_write_sei(PutBitContext *pbc,
+                                        VAAPIEncodeContext *ctx,
+                                        VAAPIEncodePicture *pic)
+{
+    VAAPIEncodeH264Context *priv = ctx->priv_data;
+    PutBitContext payload_bits;
+    char payload[256];
+    int payload_type, payload_size, i;
+    void (*write_payload)(PutBitContext *pbc,
+                          VAAPIEncodeContext *ctx,
+                          VAAPIEncodePicture *pic) = NULL;
+
+    vaapi_encode_h264_write_nal_header(pbc, NAL_SEI, 0);
+
+    for (payload_type = 0; payload_type < 64; payload_type++) {
+        switch (payload_type) {
+        case SEI_TYPE_BUFFERING_PERIOD:
+            if (!priv->send_timing_sei ||
+                pic->type != PICTURE_TYPE_IDR)
+                continue;
+            write_payload = &vaapi_encode_h264_write_buffering_period;
+            break;
+        case SEI_TYPE_PIC_TIMING:
+            if (!priv->send_timing_sei)
+                continue;
+            write_payload = &vaapi_encode_h264_write_pic_timing;
+            break;
+        default:
+            continue;
+        }
+
+        init_put_bits(&payload_bits, payload, sizeof(payload));
+        write_payload(&payload_bits, ctx, pic);
+        if (put_bits_count(&payload_bits) & 7) {
+            write_u(&payload_bits, 1, 1, bit_equal_to_one);
+            while (put_bits_count(&payload_bits) & 7)
+                write_u(&payload_bits, 1, 0, bit_equal_to_zero);
+        }
+        payload_size = put_bits_count(&payload_bits) / 8;
+        flush_put_bits(&payload_bits);
+
+        u(8, payload_type, last_payload_type_byte);
+        u(8, payload_size, last_payload_size_byte);
+        for (i = 0; i < payload_size; i++)
+            u(8, payload[i] & 0xff, sei_payload);
+    }
+
+    vaapi_encode_h264_write_trailing_rbsp(pbc);
+}
+
 static int vaapi_encode_h264_write_sequence_header(AVCodecContext *avctx,
                                                    char *data, size_t *data_len)
 {
@@ -580,6 +717,32 @@ static int vaapi_encode_h264_write_slice_header(AVCodecContext *avctx,
                                                         tmp, header_len);
 }
 
+static int vaapi_encode_h264_write_extra_header(AVCodecContext *avctx,
+                                                VAAPIEncodePicture *pic,
+                                                int index, int *type,
+                                                char *data, size_t *data_len)
+{
+    VAAPIEncodeContext *ctx = avctx->priv_data;
+    PutBitContext pbc;
+    char tmp[256];
+    size_t header_len;
+
+    if (index == 0 && ctx->va_rc_mode == VA_RC_CBR) {
+        *type = VAEncPackedHeaderH264_SEI;
+
+        init_put_bits(&pbc, tmp, sizeof(tmp));
+        vaapi_encode_h264_write_sei(&pbc, ctx, pic);
+        header_len = put_bits_count(&pbc);
+        flush_put_bits(&pbc);
+
+        return ff_vaapi_encode_h26x_nal_unit_to_byte_stream(data, data_len,
+                                                            tmp, header_len);
+
+    } else {
+        return AVERROR_EOF;
+    }
+}
+
 static int vaapi_encode_h264_init_sequence_params(AVCodecContext *avctx)
 {
     VAAPIEncodeContext                 *ctx = avctx->priv_data;
@@ -660,6 +823,43 @@ static int vaapi_encode_h264_init_sequence_params(AVCodecContext *avctx)
             mseq->fixed_frame_rate_flag = 0;
         }
 
+        if (ctx->va_rc_mode == VA_RC_CBR) {
+            priv->send_timing_sei = 1;
+            mseq->nal_hrd_parameters_present_flag = 1;
+
+            mseq->cpb_cnt_minus1 = 0;
+
+            // Try to scale these to a sensible range so that the
+            // golomb encode of the value is not overlong.
+            mseq->bit_rate_scale =
+                av_clip(av_log2(avctx->bit_rate) - 15, 0, 15);
+            mseq->bit_rate_value_minus1[0] =
+                (avctx->bit_rate >> mseq->bit_rate_scale) - 1;
+
+            mseq->cpb_size_scale =
+                av_clip(av_log2(priv->hrd_params.hrd.buffer_size) - 15, 0, 15);
+            mseq->cpb_size_value_minus1[0] =
+                (priv->hrd_params.hrd.buffer_size >> mseq->cpb_size_scale) - 1;
+
+            // CBR mode isn't actually available here, despite naming.
+            mseq->cbr_flag[0] = 0;
+
+            mseq->initial_cpb_removal_delay_length_minus1 = 23;
+            mseq->cpb_removal_delay_length_minus1         = 23;
+            mseq->dpb_output_delay_length_minus1          = 7;
+            mseq->time_offset_length = 0;
+
+            // This calculation can easily overflow 32 bits.
+            mseq->initial_cpb_removal_delay = 90000 *
+                (uint64_t)priv->hrd_params.hrd.initial_buffer_fullness /
+                priv->hrd_params.hrd.buffer_size;
+
+            mseq->initial_cpb_removal_delay_offset = 0;
+        } else {
+            priv->send_timing_sei = 0;
+            mseq->nal_hrd_parameters_present_flag = 0;
+        }
+
         vseq->intra_period     = ctx->p_per_i * (ctx->b_per_p + 1);
         vseq->intra_idr_period = vseq->intra_period;
         vseq->ip_period        = ctx->b_per_p + 1;
@@ -717,13 +917,16 @@ static int vaapi_encode_h264_init_picture_params(AVCodecContext *avctx,
         av_assert0(pic->display_order == pic->encode_order);
         vpic->frame_num = 0;
         priv->next_frame_num = 1;
+        priv->cpb_delay = 0;
     } else {
         vpic->frame_num = priv->next_frame_num;
         if (pic->type != PICTURE_TYPE_B) {
             // nal_ref_idc != 0
             ++priv->next_frame_num;
         }
+        ++priv->cpb_delay;
     }
+    priv->dpb_delay = pic->display_order - pic->encode_order + 1;
 
     vpic->frame_num = vpic->frame_num &
         ((1 << (4 + vseq->seq_fields.bits.log2_max_frame_num_minus4)) - 1);
@@ -1056,6 +1259,8 @@ static VAAPIEncodeType vaapi_encode_type_h264 = {
 
     .slice_header_type     = VAEncPackedHeaderH264_Slice,
     .write_slice_header    = &vaapi_encode_h264_write_slice_header,
+
+    .write_extra_header    = &vaapi_encode_h264_write_extra_header,
 };
 
 static av_cold int vaapi_encode_h264_init(AVCodecContext *avctx)