[FFmpeg-devel] [PATCH] lavc/vaapi_encode_h264: add option to insert AUD nal to AVC stream

[Jun Zhao]

From e39078e9fe02c8d77c5e28436aef4d80a2d7b3a0 Mon Sep 17 00:00:00 2001
From: Jun Zhao <jun.zhao at intel.com>
Date: Tue, 2 May 2017 10:36:55 +0800
Subject: [PATCH] lavc/vaapi_encode_h264: add option to insert AUD nal to AVC
 stream.

Add aud option to support insert AUD nal in AVC stream. Default is
disable.

Signed-off-by: Jun Zhao <jun.zhao at intel.com>
Signed-off-by: Yi A Wang <yi.a.wang at intel.com>
---
 libavcodec/vaapi_encode.c      | 15 ++++++++++++++
 libavcodec/vaapi_encode.h      |  4 ++++
 libavcodec/vaapi_encode_h264.c | 45 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 64 insertions(+)

diff --git a/libavcodec/vaapi_encode.c b/libavcodec/vaapi_encode.c
index 7e9c00f51d..77a10f98a7 100644
--- a/libavcodec/vaapi_encode.c
+++ b/libavcodec/vaapi_encode.c
@@ -236,6 +236,21 @@ static int vaapi_encode_issue(AVCodecContext *avctx,
             goto fail;
     }
 
+    if (ctx->va_packed_headers & VA_ENC_PACKED_HEADER_RAW_DATA &&
+        ctx->codec->write_aud_header) {
+        bit_len = 8 * sizeof(data);
+        err = ctx->codec->write_aud_header(avctx, pic, data, &bit_len);
+        if (err < 0) {
+            av_log(avctx, AV_LOG_ERROR, "Failed to write aud "
+                   "header %d: %d.\n", err);
+            goto fail;
+        }
+        err = vaapi_encode_make_packed_header(avctx, pic, VAEncPackedHeaderRawData,
+                                              data, bit_len);
+        if (err < 0)
+            goto fail;
+    }
+
     if (pic->type == PICTURE_TYPE_IDR) {
         if (ctx->va_packed_headers & VA_ENC_PACKED_HEADER_SEQUENCE &&
             ctx->codec->write_sequence_header) {
diff --git a/libavcodec/vaapi_encode.h b/libavcodec/vaapi_encode.h
index 0edf27e4cb..09a5d87f7d 100644
--- a/libavcodec/vaapi_encode.h
+++ b/libavcodec/vaapi_encode.h
@@ -267,6 +267,10 @@ typedef struct VAAPIEncodeType {
                                  VAAPIEncodePicture *pic,
                                  int index, int *type,
                                  char *data, size_t *data_len);
+     // Write an AU packed header, called by AVC encoder to insert AUD
+     int    (*write_aud_header)(AVCodecContext *avctx,
+                                VAAPIEncodePicture *pic,
+                                char *data, size_t *data_len);
 } VAAPIEncodeType;
 
 
diff --git a/libavcodec/vaapi_encode_h264.c b/libavcodec/vaapi_encode_h264.c
index 47d0c9496a..7fa9ca70e0 100644
--- a/libavcodec/vaapi_encode_h264.c
+++ b/libavcodec/vaapi_encode_h264.c
@@ -168,6 +168,7 @@ typedef struct VAAPIEncodeH264Options {
     int qp;
     int quality;
     int low_power;
+    int aud;
 } VAAPIEncodeH264Options;
 
 
@@ -750,6 +751,41 @@ static int vaapi_encode_h264_write_slice_header(AVCodecContext *avctx,
                                                         tmp, header_len);
 }
 
+static int vaapi_encode_h264_write_aud_header(AVCodecContext *avctx,
+                                              VAAPIEncodePicture *pic,
+                                              char *data, size_t *data_len)
+{
+    VAAPIEncodeContext *ctx = avctx->priv_data;
+    PutBitContext pbc;
+    char tmp[256];
+    size_t header_len;
+    int primary_pic_type;
+
+    init_put_bits(&pbc, tmp, sizeof(tmp));
+    vaapi_encode_h264_write_nal_header(&pbc, H264_NAL_AUD, 0);
+    switch (pic->type) {
+        case PICTURE_TYPE_IDR:
+        case PICTURE_TYPE_I:
+            primary_pic_type = 0;
+            break;
+        case PICTURE_TYPE_P:
+            primary_pic_type = 1;
+            break;
+        case PICTURE_TYPE_B:
+            primary_pic_type = 2;
+            break;
+        default:
+            av_assert0(0 && "unknown pic type");
+            break;
+    }
+    write_u(&pbc, 3, primary_pic_type, primary_pic_type);
+    vaapi_encode_h264_write_trailing_rbsp(&pbc);
+    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);
+}
+
 static int vaapi_encode_h264_write_extra_header(AVCodecContext *avctx,
                                                 VAAPIEncodePicture *pic,
                                                 int index, int *type,
@@ -1180,6 +1216,8 @@ static const VAAPIEncodeType vaapi_encode_type_h264 = {
     .write_slice_header    = &vaapi_encode_h264_write_slice_header,
 
     .write_extra_header    = &vaapi_encode_h264_write_extra_header,
+
+    .write_aud_header      = &vaapi_encode_h264_write_aud_header,
 };
 
 static av_cold int vaapi_encode_h264_init(AVCodecContext *avctx)
@@ -1265,6 +1303,11 @@ static av_cold int vaapi_encode_h264_init(AVCodecContext *avctx)
         VA_ENC_PACKED_HEADER_SLICE    | // Slice headers.
         VA_ENC_PACKED_HEADER_MISC;      // SEI.
 
+    if (opt->aud == 1) {
+        ctx->va_packed_headers |=
+            VA_ENC_PACKED_HEADER_RAW_DATA;
+    }
+
     ctx->surface_width  = FFALIGN(avctx->width,  16);
     ctx->surface_height = FFALIGN(avctx->height, 16);
 
@@ -1282,6 +1325,8 @@ static const AVOption vaapi_encode_h264_options[] = {
     { "low_power", "Use low-power encoding mode (experimental: only supported "
       "on some platforms, does not support all features)",
       OFFSET(low_power), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, 1, FLAGS },
+    { "aud", "Use access unit delimiters",
+      OFFSET(aud), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, 1, FLAGS },
     { NULL },
 };
 
-- 
2.11.0