2 files changed, 49 insertions, 5 deletions

diff --git a/libavcodec/nvenc.c b/libavcodec/nvenc.c
index e23ec8868e..2dfea9431e 100644
--- a/libavcodec/nvenc.c
+++ b/libavcodec/nvenc.c
@@ -663,7 +663,6 @@ static int nvenc_setup_encoder(AVCodecContext *avctx)
 
     if (avctx->gop_size > 0) {
         if (avctx->max_b_frames > 0) {
-            ctx->last_dts = -2;
             /* 0 is intra-only,
              * 1 is I/P only,
              * 2 is one B Frame,
@@ -681,6 +680,9 @@ static int nvenc_setup_encoder(AVCodecContext *avctx)
     if (ctx->config.frameIntervalP > 1)
         avctx->max_b_frames = ctx->config.frameIntervalP - 1;
 
+    ctx->initial_pts[0] = AV_NOPTS_VALUE;
+    ctx->initial_pts[1] = AV_NOPTS_VALUE;
+
     nvenc_setup_rate_control(avctx);
 
     if (avctx->flags & AV_CODEC_FLAG_INTERLACED_DCT) {
@@ -1049,13 +1051,35 @@ static inline int nvenc_dequeue_surface(AVFifoBuffer *f,
     return av_fifo_generic_read(f, surf, sizeof(*surf), NULL);
 }
 
-static int nvenc_set_timestamp(NVENCContext *ctx,
+static int nvenc_set_timestamp(AVCodecContext *avctx,
                                NV_ENC_LOCK_BITSTREAM *params,
                                AVPacket *pkt)
 {
+    NVENCContext *ctx = avctx->priv_data;
+
     pkt->pts      = params->outputTimeStamp;
     pkt->duration = params->outputDuration;
 
+    /* generate the first dts by linearly extrapolating the
+     * first two pts values to the past */
+    if (avctx->max_b_frames > 0 && !ctx->first_packet_output &&
+        ctx->initial_pts[1] != AV_NOPTS_VALUE) {
+        int64_t ts0 = ctx->initial_pts[0], ts1 = ctx->initial_pts[1];
+        int64_t delta;
+
+        if ((ts0 < 0 && ts1 > INT64_MAX + ts0) ||
+            (ts0 > 0 && ts1 < INT64_MIN + ts0))
+            return AVERROR(ERANGE);
+        delta = ts1 - ts0;
+
+        if ((delta < 0 && ts0 > INT64_MAX + delta) ||
+            (delta > 0 && ts0 < INT64_MIN + delta))
+            return AVERROR(ERANGE);
+        pkt->dts = ts0 - delta;
+
+        ctx->first_packet_output = 1;
+        return 0;
+    }
     return nvenc_dequeue_timestamp(ctx->timestamps, &pkt->dts);
 }
 
@@ -1090,7 +1114,7 @@ static int nvenc_get_frame(AVCodecContext *avctx, AVPacket *pkt)
 
     out->busy = out->in->locked = 0;
 
-    ret = nvenc_set_timestamp(ctx, &params, pkt);
+    ret = nvenc_set_timestamp(avctx, &params, pkt);
     if (ret < 0)
         return ret;
 
@@ -1119,6 +1143,18 @@ FF_ENABLE_DEPRECATION_WARNINGS
     return 0;
 }
 
+static int output_ready(AVCodecContext *avctx, int flush)
+{
+    NVENCContext *ctx = avctx->priv_data;
+
+    /* when B-frames are enabled, we wait for two initial timestamps to
+     * calculate the first dts */
+    if (!flush && avctx->max_b_frames > 0 &&
+        (ctx->initial_pts[0] == AV_NOPTS_VALUE || ctx->initial_pts[1] == AV_NOPTS_VALUE))
+        return 0;
+    return av_fifo_size(ctx->ready) > 0;
+}
+
 int ff_nvenc_encode_frame(AVCodecContext *avctx, AVPacket *pkt,
                           const AVFrame *frame, int *got_packet)
 {
@@ -1162,6 +1198,11 @@ int ff_nvenc_encode_frame(AVCodecContext *avctx, AVPacket *pkt,
         ret = nvenc_enqueue_timestamp(ctx->timestamps, frame->pts);
         if (ret < 0)
             return ret;
+
+        if (ctx->initial_pts[0] == AV_NOPTS_VALUE)
+            ctx->initial_pts[0] = frame->pts;
+        else if (ctx->initial_pts[1] == AV_NOPTS_VALUE)
+            ctx->initial_pts[1] = frame->pts;
     } else {
         params.encodePicFlags = NV_ENC_PIC_FLAG_EOS;
     }
@@ -1185,7 +1226,7 @@ int ff_nvenc_encode_frame(AVCodecContext *avctx, AVPacket *pkt,
         }
     }
 
-    if (av_fifo_size(ctx->ready) > 0) {
+    if (output_ready(avctx, !frame)) {
         ret = nvenc_get_frame(avctx, pkt);
         if (ret < 0)
             return ret;
diff --git a/libavcodec/nvenc.h b/libavcodec/nvenc.h
index 8819b3ca32..32d3345414 100644
--- a/libavcodec/nvenc.h
+++ b/libavcodec/nvenc.h
@@ -112,7 +112,10 @@ typedef struct NVENCContext {
     AVFifoBuffer *timestamps;
     AVFifoBuffer *pending, *ready;
 
-    int64_t last_dts;
+    /* timestamps of the first two frames, for computing the first dts
+     * when b-frames are present */
+    int64_t initial_pts[2];
+    int first_packet_output;
 
     void *nvenc_ctx;