[FFmpeg-devel] [PATCH v3 3/3] libavcodec: write out user data unregistered SEI for nvenc

[Brad Hards]

Signed-off-by: Brad Hards <bradh at frogmouth.net>
---
 libavcodec/nvenc.c | 64 ++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 53 insertions(+), 11 deletions(-)

diff --git a/libavcodec/nvenc.c b/libavcodec/nvenc.c
index 0dcd93a99c..9bf2728cba 100644
--- a/libavcodec/nvenc.c
+++ b/libavcodec/nvenc.c
@@ -2170,9 +2170,10 @@ static int nvenc_send_frame(AVCodecContext *avctx, const AVFrame *frame)
     NVENCSTATUS nv_status;
     NvencSurface *tmp_out_surf, *in_surf;
     int res, res2;
-    NV_ENC_SEI_PAYLOAD sei_data[8];
+    NV_ENC_SEI_PAYLOAD *sei_data = 0;
     int sei_count = 0;
     int i;
+    int total_unregistered_sei = 0;
 
     NvencContext *ctx = avctx->priv_data;
     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
@@ -2185,6 +2186,8 @@ static int nvenc_send_frame(AVCodecContext *avctx, const AVFrame *frame)
         return AVERROR(EINVAL);
 
     if (frame && frame->buf[0]) {
+        void *a53_data = NULL;
+        void *tc_data = NULL;
         in_surf = get_free_frame(ctx);
         if (!in_surf)
             return AVERROR(EAGAIN);
@@ -2230,7 +2233,6 @@ static int nvenc_send_frame(AVCodecContext *avctx, const AVFrame *frame)
         pic_params.inputTimeStamp = frame->pts;
 
         if (ctx->a53_cc && av_frame_get_side_data(frame, AV_FRAME_DATA_A53_CC)) {
-            void *a53_data = NULL;
             size_t a53_size = 0;
 
             if (ff_alloc_a53_sei(frame, 0, (void**)&a53_data, &a53_size) < 0) {
@@ -2238,15 +2240,21 @@ static int nvenc_send_frame(AVCodecContext *avctx, const AVFrame *frame)
             }
 
             if (a53_data) {
-                sei_data[sei_count].payloadSize = (uint32_t)a53_size;
-                sei_data[sei_count].payloadType = 4;
-                sei_data[sei_count].payload = (uint8_t*)a53_data;
-                sei_count ++;
+                sei_data = av_realloc_array(sei_data, sei_count + 1, sizeof(NV_ENC_SEI_PAYLOAD));
+                if (!sei_data) {
+                    av_log(ctx, AV_LOG_ERROR, "Not enough memory for closed captions SEI, skipping\n");
+                    av_free(a53_data);
+                    sei_count = 0;
+                } else {
+                    sei_data[sei_count].payloadSize = (uint32_t)a53_size;
+                    sei_data[sei_count].payloadType = 4;
+                    sei_data[sei_count].payload = (uint8_t*)a53_data;
+                    sei_count ++;
+                }
             }
         }
 
         if (ctx->s12m_tc && av_frame_get_side_data(frame, AV_FRAME_DATA_S12M_TIMECODE)) {
-            void *tc_data = NULL;
             size_t tc_size = 0;
 
             if (ff_alloc_timecode_sei(frame, avctx->framerate, 0, (void**)&tc_data, &tc_size) < 0) {
@@ -2254,13 +2262,46 @@ static int nvenc_send_frame(AVCodecContext *avctx, const AVFrame *frame)
             }
 
             if (tc_data) {
-                sei_data[sei_count].payloadSize = (uint32_t)tc_size;
-                sei_data[sei_count].payloadType = SEI_TYPE_TIME_CODE;
-                sei_data[sei_count].payload = (uint8_t*)tc_data;
-                sei_count ++;
+                sei_data = av_realloc_array(sei_data, sei_count + 1, sizeof(NV_ENC_SEI_PAYLOAD));
+                if (!sei_data) {
+                    av_log(ctx, AV_LOG_ERROR, "Not enough memory for timecode SEI, skipping\n");
+                    av_free(a53_data);
+                    av_free(tc_data);
+                    sei_count = 0;
+                } else {
+                    sei_data[sei_count].payloadSize = (uint32_t)tc_size;
+                    sei_data[sei_count].payloadType = SEI_TYPE_TIME_CODE;
+                    sei_data[sei_count].payload = (uint8_t*)tc_data;
+                    sei_count ++;
+                }
             }
         }
 
+        for (int j = 0; j < frame->nb_side_data; j++)
+            if (frame->side_data[j]->type == AV_FRAME_DATA_SEI_UNREGISTERED)
+                total_unregistered_sei++;
+        if (total_unregistered_sei > 0) {
+            sei_data = av_realloc_array(sei_data,
+                                        sei_count + total_unregistered_sei,
+                                        sizeof(NV_ENC_SEI_PAYLOAD));
+            if (!sei_data) {
+                av_log(ctx, AV_LOG_ERROR, "Not enough memory for unregistered SEI, skipping\n");
+                av_free(a53_data);
+                av_free(tc_data);
+                sei_count = 0;
+            } else    
+                for (int j = 0; j < frame->nb_side_data; j++) {
+                    AVFrameSideData *side_data = frame->side_data[j];
+                    if (side_data->type == AV_FRAME_DATA_SEI_UNREGISTERED) {
+                        sei_data[sei_count].payloadSize = side_data->size;
+                        sei_data[sei_count].payloadType = SEI_TYPE_USER_DATA_UNREGISTERED;
+                        sei_data[sei_count].payload = av_memdup(side_data->data, side_data->size);
+                        if (sei_data[sei_count].payload)
+                            sei_count ++;
+                    }
+                }
+        }
+
         nvenc_codec_specific_pic_params(avctx, &pic_params, sei_data, sei_count);
     } else {
         pic_params.encodePicFlags = NV_ENC_PIC_FLAG_EOS;
@@ -2274,6 +2315,7 @@ static int nvenc_send_frame(AVCodecContext *avctx, const AVFrame *frame)
 
     for ( i = 0; i < sei_count; i++)
         av_freep(&sei_data[i].payload);
+    av_free(sei_data);
 
     res = nvenc_pop_context(avctx);
     if (res < 0)
-- 
2.27.0