avcodec/nvenc: split interlaced output field slices

Signed-off-by: Timo Rothenpieler <timo@rothenpieler.org>

Author: BtbN

libavcodec/nvenc.c

@@ -21,6 +21,8 @@
 
 #include "config.h"
 
+#include "h2645_parse.h"
+#include "h264.h"
 #include "nvenc.h"
 
 #include "libavutil/hwcontext_cuda.h"
@@ -1073,6 +1075,15 @@ static av_cold int nvenc_setup_encoder(AVCodecContext *avctx)
     ctx->init_encode_params.frameRateNum = avctx->time_base.den;
     ctx->init_encode_params.frameRateDen = avctx->time_base.num * avctx->ticks_per_frame;
 
+    /* if (avctx->ticks_per_frame == 1) {
+       if(avctx->time_base.den < INT_MAX/2) {
+            avctx->time_base.den *= 2;
+        } else {
+            avctx->time_base.num /= 2;
+        }
+        avctx->ticks_per_frame = 2;
+    } */
+
     ctx->init_encode_params.enableEncodeAsync = 0;
     ctx->init_encode_params.enablePTD = 1;
 
@@ -1678,6 +1689,55 @@ static int nvenc_set_timestamp(AVCodecContext *avctx,
     return 0;
 }
 
+static int get_second_slice_offset(AVCodecContext *avctx, const uint8_t *buf, int length, int *offset)
+{
+    NvencContext *ctx = avctx->priv_data;
+
+    H2645Packet h2645_pkt = { 0 };
+    int idx = -1, count = 0, slice_count = 0;
+    int i, res;
+
+    // NVENC does not support HEVC interlaced encoding, so don't even bother.
+    if (avctx->codec->id != AV_CODEC_ID_H264 || ctx->encode_config.frameFieldMode != NV_ENC_PARAMS_FRAME_FIELD_MODE_FIELD) {
+        *offset = 0;
+        return 0;
+    }
+
+    res = ff_h2645_packet_split(&h2645_pkt, buf, length, avctx, 0, 0, avctx->codec->id, 1);
+    if (res < 0)
+        return res;
+
+    for (i = 0; i < h2645_pkt.nb_nals; i++)
+        if (h2645_pkt.nals[i].type == H264_NAL_SLICE || h2645_pkt.nals[i].type == H264_NAL_IDR_SLICE)
+            slice_count++;
+
+    if (slice_count < 2)
+        goto fail;
+
+    for (i = 0; i < h2645_pkt.nb_nals; i++) {
+        if (h2645_pkt.nals[i].type == H264_NAL_SLICE || h2645_pkt.nals[i].type == H264_NAL_IDR_SLICE) {
+            if (++count == slice_count / 2) {
+                idx = i;
+                break;
+            }
+        }
+    }
+
+    if (i >= h2645_pkt.nb_nals || idx < 0)
+        goto fail;
+
+    *offset = h2645_pkt.nals[idx + 1].raw_data - 4 - buf;
+
+    ff_h2645_packet_uninit(&h2645_pkt);
+
+    return 0;
+
+fail:
+    ff_h2645_packet_uninit(&h2645_pkt);
+    *offset = 0;
+    return 0;
+}
+
 static int process_output_surface(AVCodecContext *avctx, AVPacket *pkt, NvencSurface *tmpoutsurf)
 {
     NvencContext *ctx = avctx->priv_data;
@@ -1688,10 +1748,16 @@ static int process_output_surface(AVCodecContext *avctx, AVPacket *pkt, NvencSur
     uint32_t *slice_offsets = NULL;
     NV_ENC_LOCK_BITSTREAM lock_params = { 0 };
     NVENCSTATUS nv_status;
+    int main_size, second_slice_offset = -1;
     int res = 0;
 
     enum AVPictureType pict_type;
 
+    if (ctx->extra_packet) {
+        av_log(avctx, AV_LOG_ERROR, "extra_packet already allocated!\n");
+        return AVERROR_BUG;
+    }
+
     switch (avctx->codec->id) {
     case AV_CODEC_ID_H264:
         slice_mode_data = ctx->encode_config.encodeCodecConfig.h264Config.sliceModeData * 2;
@@ -1721,12 +1787,32 @@ static int process_output_surface(AVCodecContext *avctx, AVPacket *pkt, NvencSur
         goto error;
     }
 
-    if (res = ff_alloc_packet2(avctx, pkt, lock_params.bitstreamSizeInBytes, 0)) {
+    res = get_second_slice_offset(avctx, lock_params.bitstreamBufferPtr, lock_params.bitstreamSizeInBytes, &second_slice_offset);
+    if (res < 0)
+        goto error;
+
+    main_size = second_slice_offset > 0 ? second_slice_offset : lock_params.bitstreamSizeInBytes;
+
+    if (res = ff_alloc_packet2(avctx, pkt, main_size, 0)) {
         p_nvenc->nvEncUnlockBitstream(ctx->nvencoder, tmpoutsurf->output_surface);
         goto error;
     }
 
-    memcpy(pkt->data, lock_params.bitstreamBufferPtr, lock_params.bitstreamSizeInBytes);
+    memcpy(pkt->data, lock_params.bitstreamBufferPtr, main_size);
+
+    if (second_slice_offset > 0) {
+        ctx->extra_packet = av_packet_alloc();
+        if (!ctx->extra_packet) {
+            res = AVERROR(ENOMEM);
+            goto error;
+        }
+
+        res = av_new_packet(ctx->extra_packet, lock_params.bitstreamSizeInBytes - second_slice_offset);
+        if (res < 0)
+            goto error;
+
+        memcpy(ctx->extra_packet->data, (uint8_t*)lock_params.bitstreamBufferPtr + second_slice_offset, ctx->extra_packet->size);
+    }
 
     nv_status = p_nvenc->nvEncUnlockBitstream(ctx->nvencoder, tmpoutsurf->output_surface);
     if (nv_status != NV_ENC_SUCCESS)
@@ -1776,6 +1862,15 @@ FF_ENABLE_DEPRECATION_WARNINGS
     if (res < 0)
         goto error2;
 
+    if (ctx->extra_packet) {
+        //TODO: The timebase needs to be fixed(if ticks_per_frame == 1).
+        ctx->extra_packet->dts = AV_NOPTS_VALUE;
+        ctx->extra_packet->pts = AV_NOPTS_VALUE;
+
+        ff_side_data_set_encoder_stats(ctx->extra_packet,
+            (lock_params.frameAvgQP - 1) * FF_QP2LAMBDA, NULL, 0, pict_type);
+    }
+
     av_free(slice_offsets);
 
     return 0;
@@ -1933,7 +2028,11 @@ int ff_nvenc_receive_packet(AVCodecContext *avctx, AVPacket *pkt)
     if (!ctx->cu_context || !ctx->nvencoder)
         return AVERROR(EINVAL);
 
-    if (output_ready(avctx, ctx->encoder_flushing)) {
+    if (ctx->extra_packet) {
+        av_packet_move_ref(pkt, ctx->extra_packet);
+        av_packet_free(&ctx->extra_packet);
+        return 0;
+    } else if (output_ready(avctx, ctx->encoder_flushing)) {
         av_fifo_generic_read(ctx->output_surface_ready_queue, &tmp_out_surf, sizeof(tmp_out_surf), NULL);
 
         cu_res = dl_fn->cuda_dl->cuCtxPushCurrent(ctx->cu_context);

libavcodec/nvenc.h

@@ -118,6 +118,8 @@ typedef struct NvencContext
 
     int encoder_flushing;
 
+    AVPacket *extra_packet;
+
     struct {
         CUdeviceptr ptr;
         NV_ENC_REGISTERED_PTR regptr;