nvdec: add frames_params support

1 files changed, 43 insertions, 31 deletions

diff --git a/libavcodec/nvdec.c b/libavcodec/nvdec.c
index db338accfa..e4babad43e 100644
--- a/libavcodec/nvdec.c
+++ b/libavcodec/nvdec.c
@@ -185,7 +185,7 @@ int ff_nvdec_decode_uninit(AVCodecContext *avctx)
     return 0;
 }
 
-int ff_nvdec_decode_init(AVCodecContext *avctx, unsigned int dpb_size)
+int ff_nvdec_decode_init(AVCodecContext *avctx)
 {
     NVDECContext *ctx = avctx->internal->hwaccel_priv_data;
 
@@ -214,37 +214,12 @@ int ff_nvdec_decode_init(AVCodecContext *avctx, unsigned int dpb_size)
         return AVERROR(ENOSYS);
     }
 
-    if (avctx->thread_type & FF_THREAD_FRAME)
-        dpb_size += avctx->thread_count;
-
     if (!avctx->hw_frames_ctx) {
-        AVHWFramesContext *frames_ctx;
-
-        if (!avctx->hw_device_ctx) {
-            av_log(avctx, AV_LOG_ERROR, "A hardware device or frames context "
-                   "is required for CUVID decoding.\n");
-            return AVERROR(EINVAL);
-        }
-
-        avctx->hw_frames_ctx = av_hwframe_ctx_alloc(avctx->hw_device_ctx);
-        if (!avctx->hw_frames_ctx)
-            return AVERROR(ENOMEM);
-        frames_ctx = (AVHWFramesContext*)avctx->hw_frames_ctx->data;
-
-        frames_ctx->format            = AV_PIX_FMT_CUDA;
-        frames_ctx->width             = avctx->coded_width;
-        frames_ctx->height            = avctx->coded_height;
-        frames_ctx->sw_format         = AV_PIX_FMT_NV12;
-        frames_ctx->sw_format         = sw_desc->comp[0].depth > 8 ?
-                                        AV_PIX_FMT_P010 : AV_PIX_FMT_NV12;
-        frames_ctx->initial_pool_size = dpb_size;
-
-        ret = av_hwframe_ctx_init(avctx->hw_frames_ctx);
-        if (ret < 0) {
-            av_log(avctx, AV_LOG_ERROR, "Error initializing internal frames context\n");
+        ret = ff_decode_get_hw_frames_ctx(avctx, AV_HWDEVICE_TYPE_CUDA);
+        if (ret < 0)
             return ret;
-        }
     }
+
     frames_ctx = (AVHWFramesContext*)avctx->hw_frames_ctx->data;
 
     params.ulWidth             = avctx->coded_width;
@@ -256,7 +231,7 @@ int ff_nvdec_decode_init(AVCodecContext *avctx, unsigned int dpb_size)
                                  cudaVideoSurfaceFormat_P016 : cudaVideoSurfaceFormat_NV12;
     params.CodecType           = cuvid_codec_type;
     params.ChromaFormat        = cuvid_chroma_format;
-    params.ulNumDecodeSurfaces = dpb_size;
+    params.ulNumDecodeSurfaces = frames_ctx->initial_pool_size;
     params.ulNumOutputSurfaces = 1;
 
     ret = nvdec_decoder_create(&ctx->decoder_ref, frames_ctx->device_ref, &params, avctx);
@@ -268,7 +243,7 @@ int ff_nvdec_decode_init(AVCodecContext *avctx, unsigned int dpb_size)
         ret = AVERROR(ENOMEM);
         goto fail;
     }
-    pool->dpb_size = dpb_size;
+    pool->dpb_size = frames_ctx->initial_pool_size;
 
     ctx->decoder_pool = av_buffer_pool_init2(sizeof(int), pool,
                                              nvdec_decoder_frame_alloc, av_free);
@@ -430,3 +405,40 @@ finish:
 
     return ret;
 }
+
+int ff_nvdec_frame_params(AVCodecContext *avctx,
+                          AVBufferRef *hw_frames_ctx,
+                          int dpb_size)
+{
+    AVHWFramesContext *frames_ctx = (AVHWFramesContext*)hw_frames_ctx->data;
+    const AVPixFmtDescriptor *sw_desc;
+    int cuvid_codec_type, cuvid_chroma_format;
+
+    sw_desc = av_pix_fmt_desc_get(avctx->sw_pix_fmt);
+    if (!sw_desc)
+        return AVERROR_BUG;
+
+    cuvid_codec_type = map_avcodec_id(avctx->codec_id);
+    if (cuvid_codec_type < 0) {
+        av_log(avctx, AV_LOG_ERROR, "Unsupported codec ID\n");
+        return AVERROR_BUG;
+    }
+
+    cuvid_chroma_format = map_chroma_format(avctx->sw_pix_fmt);
+    if (cuvid_chroma_format < 0) {
+        av_log(avctx, AV_LOG_VERBOSE, "Unsupported chroma format\n");
+        return AVERROR(EINVAL);
+    }
+
+    if (avctx->thread_type & FF_THREAD_FRAME)
+        dpb_size += avctx->thread_count;
+
+    frames_ctx->format            = AV_PIX_FMT_CUDA;
+    frames_ctx->width             = avctx->coded_width;
+    frames_ctx->height            = avctx->coded_height;
+    frames_ctx->sw_format         = sw_desc->comp[0].depth > 8 ?
+                                    AV_PIX_FMT_P010 : AV_PIX_FMT_NV12;
+    frames_ctx->initial_pool_size = dpb_size;
+
+    return 0;
+}