0001-POC-for-new-FFmpeg-nvmpi-API-implementation.patch

From 437f8c556b55a9e8722d0100f2aab41ce7bfea10 Mon Sep 17 00:00:00 2001
From: Bartosz Meglicki <bmegli@op.pl>
Date: Mon, 20 Mar 2023 17:29:46 +0100
Subject: [PATCH] POC for new FFmpeg nvmpi API implementation

- implement send_frame, receive_packet FFmpeg API
- this should allow low latency streaming
  - previous encode2 API may return 1 packet per frame
  - new API allows "draining" pending packets while streaming

EOS:
- current implementation doesn't make sure that all packets are drained from hardware
- so you might miss some last few packets

Related to:
- https://github.com/jocover/jetson-ffmpeg/issues/131
---
 configure              |  13 ++
 libavcodec/Makefile    |   8 ++
 libavcodec/allcodecs.c |   8 ++
 libavcodec/nvmpi_dec.c | 171 +++++++++++++++++++++++
 libavcodec/nvmpi_enc.c | 309 +++++++++++++++++++++++++++++++++++++++++
 5 files changed, 509 insertions(+)
 create mode 100644 libavcodec/nvmpi_dec.c
 create mode 100644 libavcodec/nvmpi_enc.c

diff --git a/configure b/configure
index 927857f..70433f4 100755
--- a/configure
+++ b/configure
@@ -340,6 +340,7 @@ External library support:
   --disable-vaapi          disable Video Acceleration API (mainly Unix/Intel) code [autodetect]
   --disable-vdpau          disable Nvidia Video Decode and Presentation API for Unix code [autodetect]
   --disable-videotoolbox   disable VideoToolbox code [autodetect]
+  --enable-nvmpi           enable nvmpi code
 
 Toolchain options:
   --arch=ARCH              select architecture [$arch]
@@ -1851,6 +1852,7 @@ HWACCEL_LIBRARY_LIST="
     mmal
     omx
     opencl
+    nvmpi
 "
 
 DOCUMENT_LIST="
@@ -3014,11 +3016,14 @@ h264_mediacodec_decoder_deps="mediacodec"
 h264_mediacodec_decoder_select="h264_mp4toannexb_bsf h264_parser"
 h264_mmal_decoder_deps="mmal"
 h264_nvenc_encoder_deps="nvenc"
+h264_nvmpi_encoder_deps="nvmpi"
 h264_omx_encoder_deps="omx"
 h264_qsv_decoder_select="h264_mp4toannexb_bsf h264_parser qsvdec"
 h264_qsv_encoder_select="qsvenc"
 h264_rkmpp_decoder_deps="rkmpp"
 h264_rkmpp_decoder_select="h264_mp4toannexb_bsf"
+h264_nvmpi_decoder_deps="nvmpi"
+h264_nvmpi_decoder_select="h264_mp4toannexb_bsf"
 h264_vaapi_encoder_select="cbs_h264 vaapi_encode"
 h264_v4l2m2m_decoder_deps="v4l2_m2m h264_v4l2_m2m"
 h264_v4l2m2m_decoder_select="h264_mp4toannexb_bsf"
@@ -3029,10 +3034,13 @@ hevc_cuvid_decoder_select="hevc_mp4toannexb_bsf"
 hevc_mediacodec_decoder_deps="mediacodec"
 hevc_mediacodec_decoder_select="hevc_mp4toannexb_bsf hevc_parser"
 hevc_nvenc_encoder_deps="nvenc"
+hevc_nvmpi_encoder_deps="nvmpi"
 hevc_qsv_decoder_select="hevc_mp4toannexb_bsf hevc_parser qsvdec"
 hevc_qsv_encoder_select="hevcparse qsvenc"
 hevc_rkmpp_decoder_deps="rkmpp"
 hevc_rkmpp_decoder_select="hevc_mp4toannexb_bsf"
+hevc_nvmpi_decoder_deps="nvmpi"
+hevc_nvmpi_decoder_select="hevc_mp4toannexb_bsf"
 hevc_vaapi_encoder_deps="VAEncPictureParameterBufferHEVC"
 hevc_vaapi_encoder_select="cbs_h265 vaapi_encode"
 hevc_v4l2m2m_decoder_deps="v4l2_m2m hevc_v4l2_m2m"
@@ -3047,6 +3055,7 @@ mpeg1_cuvid_decoder_deps="cuvid"
 mpeg1_v4l2m2m_decoder_deps="v4l2_m2m mpeg1_v4l2_m2m"
 mpeg2_crystalhd_decoder_select="crystalhd"
 mpeg2_cuvid_decoder_deps="cuvid"
+mpeg2_nvmpi_decoder_deps="nvmpi"
 mpeg2_mmal_decoder_deps="mmal"
 mpeg2_mediacodec_decoder_deps="mediacodec"
 mpeg2_qsv_decoder_select="qsvdec mpegvideo_parser"
@@ -3055,6 +3064,7 @@ mpeg2_vaapi_encoder_select="cbs_mpeg2 vaapi_encode"
 mpeg2_v4l2m2m_decoder_deps="v4l2_m2m mpeg2_v4l2_m2m"
 mpeg4_crystalhd_decoder_select="crystalhd"
 mpeg4_cuvid_decoder_deps="cuvid"
+mpeg4_nvmpi_decoder_deps="nvmpi"
 mpeg4_mediacodec_decoder_deps="mediacodec"
 mpeg4_mmal_decoder_deps="mmal"
 mpeg4_omx_encoder_deps="omx"
@@ -3069,6 +3079,7 @@ vc1_mmal_decoder_deps="mmal"
 vc1_qsv_decoder_select="qsvdec vc1_parser"
 vc1_v4l2m2m_decoder_deps="v4l2_m2m vc1_v4l2_m2m"
 vp8_cuvid_decoder_deps="cuvid"
+vp8_nvmpi_decoder_deps="nvmpi"
 vp8_mediacodec_decoder_deps="mediacodec"
 vp8_qsv_decoder_select="qsvdec vp8_parser"
 vp8_rkmpp_decoder_deps="rkmpp"
@@ -3077,6 +3088,7 @@ vp8_vaapi_encoder_select="vaapi_encode"
 vp8_v4l2m2m_decoder_deps="v4l2_m2m vp8_v4l2_m2m"
 vp8_v4l2m2m_encoder_deps="v4l2_m2m vp8_v4l2_m2m"
 vp9_cuvid_decoder_deps="cuvid"
+vp9_nvmpi_decoder_deps="nvmpi"
 vp9_mediacodec_decoder_deps="mediacodec"
 vp9_rkmpp_decoder_deps="rkmpp"
 vp9_vaapi_encoder_deps="VAEncPictureParameterBufferVP9"
@@ -6366,6 +6378,7 @@ enabled rkmpp             && { require_pkg_config rkmpp rockchip_mpp  rockchip/r
                                  die "ERROR: rkmpp requires --enable-libdrm"; }
                              }
 enabled vapoursynth       && require_pkg_config vapoursynth "vapoursynth-script >= 42" VSScript.h vsscript_init
+enabled nvmpi		  && require_pkg_config nvmpi nvmpi nvmpi.h nvmpi_create_decoder
 
 
 if enabled gcrypt; then
diff --git a/libavcodec/Makefile b/libavcodec/Makefile
index 3cd73fb..c3ed5cc 100644
--- a/libavcodec/Makefile
+++ b/libavcodec/Makefile
@@ -354,6 +354,8 @@ OBJS-$(CONFIG_H264_MMAL_DECODER)       += mmaldec.o
 OBJS-$(CONFIG_H264_NVENC_ENCODER)      += nvenc_h264.o
 OBJS-$(CONFIG_NVENC_ENCODER)           += nvenc_h264.o
 OBJS-$(CONFIG_NVENC_H264_ENCODER)      += nvenc_h264.o
+OBJS-$(CONFIG_H264_NVMPI_DECODER)      += nvmpi_dec.o
+OBJS-$(CONFIG_H264_NVMPI_ENCODER)      += nvmpi_enc.o
 OBJS-$(CONFIG_H264_OMX_ENCODER)        += omx.o
 OBJS-$(CONFIG_H264_QSV_DECODER)        += qsvdec_h2645.o
 OBJS-$(CONFIG_H264_QSV_ENCODER)        += qsvenc_h264.o
@@ -379,6 +381,8 @@ OBJS-$(CONFIG_HEVC_QSV_ENCODER)        += qsvenc_hevc.o hevc_ps_enc.o       \
 OBJS-$(CONFIG_HEVC_RKMPP_DECODER)      += rkmppdec.o
 OBJS-$(CONFIG_HEVC_VAAPI_ENCODER)      += vaapi_encode_h265.o h265_profile_level.o
 OBJS-$(CONFIG_HEVC_V4L2M2M_DECODER)    += v4l2_m2m_dec.o
+OBJS-$(CONFIG_HEVC_NVMPI_DECODER)      += nvmpi_dec.o
+OBJS-$(CONFIG_HEVC_NVMPI_ENCODER)      += nvmpi_enc.o
 OBJS-$(CONFIG_HEVC_V4L2M2M_ENCODER)    += v4l2_m2m_enc.o
 OBJS-$(CONFIG_HNM4_VIDEO_DECODER)      += hnm4video.o
 OBJS-$(CONFIG_HQ_HQA_DECODER)          += hq_hqa.o hq_hqadata.o hq_hqadsp.o \
@@ -464,11 +468,13 @@ OBJS-$(CONFIG_MPEG2_QSV_ENCODER)       += qsvenc_mpeg2.o
 OBJS-$(CONFIG_MPEG2VIDEO_DECODER)      += mpeg12dec.o mpeg12.o mpeg12data.o
 OBJS-$(CONFIG_MPEG2VIDEO_ENCODER)      += mpeg12enc.o mpeg12.o
 OBJS-$(CONFIG_MPEG2_CUVID_DECODER)     += cuviddec.o
+OBJS-$(CONFIG_MPEG2_NVMPI_DECODER)     += nvmpi_dec.o
 OBJS-$(CONFIG_MPEG2_MEDIACODEC_DECODER) += mediacodecdec.o
 OBJS-$(CONFIG_MPEG2_VAAPI_ENCODER)     += vaapi_encode_mpeg2.o
 OBJS-$(CONFIG_MPEG2_V4L2M2M_DECODER)   += v4l2_m2m_dec.o
 OBJS-$(CONFIG_MPEG4_DECODER)           += xvididct.o
 OBJS-$(CONFIG_MPEG4_CUVID_DECODER)     += cuviddec.o
+OBJS-$(CONFIG_MPEG4_NVMPI_DECODER)     += nvmpi_dec.o
 OBJS-$(CONFIG_MPEG4_MEDIACODEC_DECODER) += mediacodecdec.o
 OBJS-$(CONFIG_MPEG4_OMX_ENCODER)       += omx.o
 OBJS-$(CONFIG_MPEG4_V4L2M2M_DECODER)   += v4l2_m2m_dec.o
@@ -669,6 +675,7 @@ OBJS-$(CONFIG_VP6_DECODER)             += vp6.o vp56.o vp56data.o \
 OBJS-$(CONFIG_VP7_DECODER)             += vp8.o vp56rac.o
 OBJS-$(CONFIG_VP8_DECODER)             += vp8.o vp56rac.o
 OBJS-$(CONFIG_VP8_CUVID_DECODER)       += cuviddec.o
+OBJS-$(CONFIG_VP8_NVMPI_DECODER)       += nvmpi_dec.o
 OBJS-$(CONFIG_VP8_MEDIACODEC_DECODER)  += mediacodecdec.o
 OBJS-$(CONFIG_VP8_QSV_DECODER)         += qsvdec_other.o
 OBJS-$(CONFIG_VP8_RKMPP_DECODER)       += rkmppdec.o
@@ -679,6 +686,7 @@ OBJS-$(CONFIG_VP9_DECODER)             += vp9.o vp9data.o vp9dsp.o vp9lpf.o vp9r
                                           vp9block.o vp9prob.o vp9mvs.o vp56rac.o \
                                           vp9dsp_8bpp.o vp9dsp_10bpp.o vp9dsp_12bpp.o
 OBJS-$(CONFIG_VP9_CUVID_DECODER)       += cuviddec.o
+OBJS-$(CONFIG_VP9_NVMPI_DECODER)       += nvmpi_dec.o
 OBJS-$(CONFIG_VP9_MEDIACODEC_DECODER)  += mediacodecdec.o
 OBJS-$(CONFIG_VP9_RKMPP_DECODER)       += rkmppdec.o
 OBJS-$(CONFIG_VP9_VAAPI_ENCODER)       += vaapi_encode_vp9.o
diff --git a/libavcodec/allcodecs.c b/libavcodec/allcodecs.c
index d2f9a39..04dc62b 100644
--- a/libavcodec/allcodecs.c
+++ b/libavcodec/allcodecs.c
@@ -143,11 +143,15 @@ extern AVCodec ff_h264_mediacodec_decoder;
 extern AVCodec ff_h264_mmal_decoder;
 extern AVCodec ff_h264_qsv_decoder;
 extern AVCodec ff_h264_rkmpp_decoder;
+extern AVCodec ff_h264_nvmpi_decoder;
+extern AVCodec ff_h264_nvmpi_encoder;
 extern AVCodec ff_hap_encoder;
 extern AVCodec ff_hap_decoder;
 extern AVCodec ff_hevc_decoder;
 extern AVCodec ff_hevc_qsv_decoder;
 extern AVCodec ff_hevc_rkmpp_decoder;
+extern AVCodec ff_hevc_nvmpi_decoder;
+extern AVCodec ff_hevc_nvmpi_encoder;
 extern AVCodec ff_hevc_v4l2m2m_decoder;
 extern AVCodec ff_hnm4_video_decoder;
 extern AVCodec ff_hq_hqa_decoder;
@@ -766,18 +770,22 @@ extern AVCodec ff_mjpeg_qsv_encoder;
 extern AVCodec ff_mjpeg_vaapi_encoder;
 extern AVCodec ff_mpeg1_cuvid_decoder;
 extern AVCodec ff_mpeg2_cuvid_decoder;
+extern AVCodec ff_mpeg2_nvmpi_decoder;
 extern AVCodec ff_mpeg2_qsv_encoder;
 extern AVCodec ff_mpeg2_vaapi_encoder;
 extern AVCodec ff_mpeg4_cuvid_decoder;
+extern AVCodec ff_mpeg4_nvmpi_decoder;
 extern AVCodec ff_mpeg4_mediacodec_decoder;
 extern AVCodec ff_mpeg4_v4l2m2m_encoder;
 extern AVCodec ff_vc1_cuvid_decoder;
 extern AVCodec ff_vp8_cuvid_decoder;
+extern AVCodec ff_vp8_nvmpi_decoder;
 extern AVCodec ff_vp8_mediacodec_decoder;
 extern AVCodec ff_vp8_qsv_decoder;
 extern AVCodec ff_vp8_v4l2m2m_encoder;
 extern AVCodec ff_vp8_vaapi_encoder;
 extern AVCodec ff_vp9_cuvid_decoder;
+extern AVCodec ff_vp9_nvmpi_decoder;
 extern AVCodec ff_vp9_mediacodec_decoder;
 extern AVCodec ff_vp9_vaapi_encoder;
 
diff --git a/libavcodec/nvmpi_dec.c b/libavcodec/nvmpi_dec.c
new file mode 100644
index 0000000..f82aa61
--- /dev/null
+++ b/libavcodec/nvmpi_dec.c
@@ -0,0 +1,171 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <sys/time.h>
+
+#include <nvmpi.h>
+#include "avcodec.h"
+#include "decode.h"
+#include "internal.h"
+#include "libavutil/buffer.h"
+#include "libavutil/common.h"
+#include "libavutil/frame.h"
+#include "libavutil/hwcontext.h"
+#include "libavutil/hwcontext_drm.h"
+#include "libavutil/imgutils.h"
+#include "libavutil/log.h"
+
+
+
+
+typedef struct {
+	char eos_reached;
+	nvmpictx* ctx;
+	AVClass *av_class;
+} nvmpiDecodeContext;
+
+static nvCodingType nvmpi_get_codingtype(AVCodecContext *avctx)
+{
+	switch (avctx->codec_id) {
+		case AV_CODEC_ID_H264:          return NV_VIDEO_CodingH264;
+		case AV_CODEC_ID_HEVC:          return NV_VIDEO_CodingHEVC;
+		case AV_CODEC_ID_VP8:           return NV_VIDEO_CodingVP8;
+		case AV_CODEC_ID_VP9:           return NV_VIDEO_CodingVP9;
+		case AV_CODEC_ID_MPEG4:		return NV_VIDEO_CodingMPEG4;
+		case AV_CODEC_ID_MPEG2VIDEO:    return NV_VIDEO_CodingMPEG2;
+		default:                        return NV_VIDEO_CodingUnused;
+	}
+};
+
+
+static int nvmpi_init_decoder(AVCodecContext *avctx){
+
+	nvmpiDecodeContext *nvmpi_context = avctx->priv_data;
+	nvCodingType codectype=NV_VIDEO_CodingUnused;
+
+	codectype =nvmpi_get_codingtype(avctx);
+	if (codectype == NV_VIDEO_CodingUnused) {
+		av_log(avctx, AV_LOG_ERROR, "Unknown codec type (%d).\n", avctx->codec_id);
+		return AVERROR_UNKNOWN;
+	}
+
+	//Workaround for default pix_fmt not being set, so check if it isnt set and set it,
+   //or if it is set, but isnt set to something we can work with.
+
+	if(avctx->pix_fmt ==AV_PIX_FMT_NONE){
+		 avctx->pix_fmt=AV_PIX_FMT_YUV420P;
+	}else if(avctx-> pix_fmt != AV_PIX_FMT_YUV420P){
+		av_log(avctx, AV_LOG_ERROR, "Invalid Pix_FMT for NVMPI Only yuv420p is supported\n");
+		return AVERROR_INVALIDDATA;
+	}
+
+	nvmpi_context->ctx=nvmpi_create_decoder(codectype,NV_PIX_YUV420);
+
+	if(!nvmpi_context->ctx){
+		av_log(avctx, AV_LOG_ERROR, "Failed to nvmpi_create_decoder (code = %d).\n", AVERROR_EXTERNAL);
+		return AVERROR_EXTERNAL;
+	}
+   return 0;
+
+}
+
+
+
+static int nvmpi_close(AVCodecContext *avctx){
+
+	nvmpiDecodeContext *nvmpi_context = avctx->priv_data;
+	return nvmpi_decoder_close(nvmpi_context->ctx);
+
+}
+
+
+
+static int nvmpi_decode(AVCodecContext *avctx,void *data,int *got_frame, AVPacket *avpkt){
+
+	nvmpiDecodeContext *nvmpi_context = avctx->priv_data;
+	AVFrame *frame = data;
+	nvFrame _nvframe={0};
+	nvPacket packet;
+	uint8_t* ptrs[3];
+	int res,linesize[3];
+
+	if(avpkt->size){
+		packet.payload_size=avpkt->size;
+		packet.payload=avpkt->data;
+		packet.pts=avpkt->pts;
+
+		res=nvmpi_decoder_put_packet(nvmpi_context->ctx,&packet);
+	}
+
+	res=nvmpi_decoder_get_frame(nvmpi_context->ctx,&_nvframe,avctx->flags & AV_CODEC_FLAG_LOW_DELAY);
+
+	if(res<0)
+		return avpkt->size;
+
+	if (ff_get_buffer(avctx, frame, 0) < 0) {
+		return AVERROR(ENOMEM);
+
+	}
+
+	linesize[0]=_nvframe.linesize[0];
+	linesize[1]=_nvframe.linesize[1];
+	linesize[2]=_nvframe.linesize[2];
+
+	ptrs[0]=_nvframe.payload[0];
+	ptrs[1]=_nvframe.payload[1];
+	ptrs[2]=_nvframe.payload[2];
+
+	av_image_copy(frame->data, frame->linesize, (const uint8_t **) ptrs, linesize, avctx->pix_fmt, _nvframe.width,_nvframe.height);
+
+	frame->width=_nvframe.width;
+	frame->height=_nvframe.height;
+
+	frame->format=AV_PIX_FMT_YUV420P;
+	frame->pts=_nvframe.timestamp;
+	frame->pkt_dts = AV_NOPTS_VALUE;
+
+	avctx->coded_width=_nvframe.width;
+	avctx->coded_height=_nvframe.height;
+	avctx->width=_nvframe.width;
+	avctx->height=_nvframe.height;
+
+	*got_frame = 1;
+
+	return avpkt->size;
+}
+
+
+
+
+#define NVMPI_DEC_CLASS(NAME) \
+	static const AVClass nvmpi_##NAME##_dec_class = { \
+		.class_name = "nvmpi_" #NAME "_dec", \
+		.version    = LIBAVUTIL_VERSION_INT, \
+	};
+
+#define NVMPI_DEC(NAME, ID, BSFS) \
+	NVMPI_DEC_CLASS(NAME) \
+	AVCodec ff_##NAME##_nvmpi_decoder = { \
+		.name           = #NAME "_nvmpi", \
+		.long_name      = NULL_IF_CONFIG_SMALL(#NAME " (nvmpi)"), \
+		.type           = AVMEDIA_TYPE_VIDEO, \
+		.id             = ID, \
+		.priv_data_size = sizeof(nvmpiDecodeContext), \
+		.init           = nvmpi_init_decoder, \
+		.close          = nvmpi_close, \
+		.decode         = nvmpi_decode, \
+		.priv_class     = &nvmpi_##NAME##_dec_class, \
+		.capabilities   = AV_CODEC_CAP_DELAY | AV_CODEC_CAP_AVOID_PROBING | AV_CODEC_CAP_HARDWARE, \
+		.pix_fmts	=(const enum AVPixelFormat[]){AV_PIX_FMT_YUV420P,AV_PIX_FMT_NV12,AV_PIX_FMT_NONE},\
+		.bsfs           = BSFS, \
+		.wrapper_name   = "nvmpi", \
+	};
+
+
+
+NVMPI_DEC(h264,  AV_CODEC_ID_H264,"h264_mp4toannexb");
+NVMPI_DEC(hevc,  AV_CODEC_ID_HEVC,"hevc_mp4toannexb");
+NVMPI_DEC(mpeg2, AV_CODEC_ID_MPEG2VIDEO,NULL);
+NVMPI_DEC(mpeg4, AV_CODEC_ID_MPEG4,NULL);
+NVMPI_DEC(vp9,  AV_CODEC_ID_VP9,NULL);
+NVMPI_DEC(vp8, AV_CODEC_ID_VP8,NULL);
+
diff --git a/libavcodec/nvmpi_enc.c b/libavcodec/nvmpi_enc.c
new file mode 100644
index 0000000..53fca62
--- /dev/null
+++ b/libavcodec/nvmpi_enc.c
@@ -0,0 +1,309 @@
+#include <nvmpi.h>
+#include "avcodec.h"
+#include "internal.h"
+#include <stdio.h>
+#include "libavutil/avstring.h"
+#include "libavutil/avutil.h"
+#include "libavutil/common.h"
+#include "libavutil/imgutils.h"
+#include "libavutil/log.h"
+#include "libavutil/opt.h"
+
+
+typedef struct {
+	const AVClass *class;
+	nvmpictx* ctx;
+	int num_capture_buffers;
+	int profile;
+	int level;
+	int rc;
+	int preset;
+	int encoder_flushing;
+}nvmpiEncodeContext;
+
+static av_cold int nvmpi_encode_init(AVCodecContext *avctx){
+
+	nvmpiEncodeContext * nvmpi_context = avctx->priv_data;
+
+	nvEncParam param={0};
+
+	param.width=avctx->width;
+	param.height=avctx->height;
+	param.bitrate=avctx->bit_rate;
+	param.mode_vbr=0;
+	param.idr_interval=60;
+	param.iframe_interval=30;
+	param.peak_bitrate=0;
+	param.fps_n=avctx->framerate.num;
+	param.fps_d=avctx->framerate.den;
+	param.profile=nvmpi_context->profile& ~FF_PROFILE_H264_INTRA;
+	param.level=nvmpi_context->level;
+	param.capture_num=nvmpi_context->num_capture_buffers;
+	param.hw_preset_type=nvmpi_context->preset;
+	param.insert_spspps_idr=(avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER)?0:1;
+
+	if(nvmpi_context->rc==1){
+		param.mode_vbr=1;
+	}
+
+	if(avctx->qmin >= 0 && avctx->qmax >= 0){
+		param.qmin=avctx->qmin;
+		param.qmax=avctx->qmax;
+	}
+
+	if (avctx->refs >= 0){
+		param.refs=avctx->refs;
+
+	}
+
+	if(avctx->max_b_frames > 0 && avctx->max_b_frames < 3){
+		param.max_b_frames=avctx->max_b_frames;
+	}
+
+	if(avctx->gop_size>0){
+		param.idr_interval=param.iframe_interval=avctx->gop_size;
+
+	}
+
+
+	if ((avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) && (avctx->codec->id == AV_CODEC_ID_H264)){
+
+		uint8_t *dst[4];
+		int linesize[4];
+		nvFrame _nvframe={0};
+		nvPacket packet={0};
+		int i;
+		int ret;
+		nvmpictx* _ctx;
+		av_image_alloc(dst, linesize,avctx->width,avctx->height,avctx->pix_fmt,1);
+
+		_ctx=nvmpi_create_encoder(NV_VIDEO_CodingH264,&param);
+		i=0;
+
+		while(1){
+
+			_nvframe.payload[0]=dst[0];
+			_nvframe.payload[1]=dst[1];
+			_nvframe.payload[2]=dst[2];
+			_nvframe.payload_size[0]=linesize[0]*avctx->height;
+			_nvframe.payload_size[1]=linesize[1]*avctx->height/2;
+			_nvframe.payload_size[2]=linesize[2]*avctx->height/2;
+
+			nvmpi_encoder_put_frame(_ctx,&_nvframe);
+
+			ret=nvmpi_encoder_get_packet(_ctx,&packet);
+
+			if(ret<0)
+				continue;
+
+			//find idr index 0x0000000165
+			while((packet.payload[i]!=0||packet.payload[i+1]!=0||packet.payload[i+2]!=0||packet.payload[i+3]!=0x01||packet.payload[i+4]!=0x65)){
+				i++;
+
+			}
+
+			avctx->extradata_size=i;
+			avctx->extradata	= av_mallocz( avctx->extradata_size + AV_INPUT_BUFFER_PADDING_SIZE );
+			memcpy( avctx->extradata, packet.payload,avctx->extradata_size);
+			memset( avctx->extradata + avctx->extradata_size, 0, AV_INPUT_BUFFER_PADDING_SIZE );
+
+			break;
+
+		}
+
+		nvmpi_encoder_close(_ctx);
+
+
+	}
+
+	if(avctx->codec->id == AV_CODEC_ID_H264)
+		nvmpi_context->ctx=nvmpi_create_encoder(NV_VIDEO_CodingH264,&param);
+	else if(avctx->codec->id == AV_CODEC_ID_HEVC){
+		nvmpi_context->ctx=nvmpi_create_encoder(NV_VIDEO_CodingHEVC,&param);
+	}
+
+
+	return 0;
+}
+
+static int ff_nvmpi_send_frame(AVCodecContext *avctx,const AVFrame *frame)
+{
+	nvmpiEncodeContext * nvmpi_context = avctx->priv_data;
+	nvFrame _nvframe={0};
+	int res;
+
+	if (nvmpi_context->encoder_flushing)
+		return AVERROR_EOF;
+
+	if(frame){
+
+		_nvframe.payload[0]=frame->data[0];
+		_nvframe.payload[1]=frame->data[1];
+		_nvframe.payload[2]=frame->data[2];
+
+		_nvframe.payload_size[0]=frame->linesize[0]*frame->height;
+		_nvframe.payload_size[1]=frame->linesize[1]*frame->height/2;
+		_nvframe.payload_size[2]=frame->linesize[2]*frame->height/2;
+
+		_nvframe.linesize[0]=frame->linesize[0];
+		_nvframe.linesize[1]=frame->linesize[1];
+		_nvframe.linesize[2]=frame->linesize[2];
+
+		_nvframe.timestamp=frame->pts;
+
+		res=nvmpi_encoder_put_frame(nvmpi_context->ctx,&_nvframe);
+
+		if(res<0)
+			return res;
+	}
+	else
+		nvmpi_context->encoder_flushing = 1;
+
+	return 0;
+}
+
+static int ff_nvmpi_receive_packet(AVCodecContext *avctx, AVPacket *pkt)
+{
+	nvmpiEncodeContext * nvmpi_context = avctx->priv_data;
+	nvPacket packet={0};
+
+	if(nvmpi_encoder_get_packet(nvmpi_context->ctx,&packet)<0)
+	{
+		if (nvmpi_context->encoder_flushing)
+			return AVERROR_EOF; //we don't really wait so we may miss some last packets
+
+		return AVERROR(EAGAIN); //nvmpi get_packet returns -1 if no packets are pending
+	}
+
+	ff_alloc_packet2(avctx,pkt,packet.payload_size,packet.payload_size);
+
+	memcpy(pkt->data,packet.payload,packet.payload_size);
+	pkt->dts=pkt->pts=packet.pts;
+
+	if(packet.flags& AV_PKT_FLAG_KEY)
+		pkt->flags = AV_PKT_FLAG_KEY;
+
+	return 0;
+}
+
+static int ff_nvmpi_encode_frame(AVCodecContext *avctx, AVPacket *pkt,const AVFrame *frame, int *got_packet)
+{
+	int res;
+	*got_packet = 0;
+
+	res = ff_nvmpi_send_frame(avctx, frame);
+
+	if(res < 0)
+		return res;
+
+	res = ff_nvmpi_receive_packet(avctx, pkt);
+
+	if (res == AVERROR(EAGAIN) || res == AVERROR_EOF)
+		return 0;
+
+	if (res < 0)
+		return res;
+
+	*got_packet = 1;
+
+	return 0;
+}
+
+static av_cold int nvmpi_encode_close(AVCodecContext *avctx){
+
+	nvmpiEncodeContext *nvmpi_context = avctx->priv_data;
+	nvmpi_encoder_close(nvmpi_context->ctx);
+
+	return 0;
+}
+
+static const AVCodecDefault defaults[] = {
+	{ "b", "2M" },
+	{ "qmin", "-1" },
+	{ "qmax", "-1" },
+	{ "qdiff", "-1" },
+	{ "qblur", "-1" },
+	{ "qcomp", "-1" },
+	{ "g", "50" },
+	{ "bf", "0" },
+	{ "refs", "0" },
+	{ NULL },
+};
+
+
+#define OFFSET(x) offsetof(nvmpiEncodeContext, x)
+#define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
+
+static const AVOption options[] = {
+	{ "num_capture_buffers", "Number of buffers in the capture context", OFFSET(num_capture_buffers), AV_OPT_TYPE_INT, {.i64 = 10 }, 1, 32, AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM },
+	/// Profile,
+
+	{ "profile",  "Set the encoding profile", OFFSET(profile), AV_OPT_TYPE_INT,   { .i64 = FF_PROFILE_UNKNOWN },       FF_PROFILE_UNKNOWN, FF_PROFILE_H264_HIGH, VE, "profile" },
+	{ "baseline", "",                         0,               AV_OPT_TYPE_CONST, { .i64 = FF_PROFILE_H264_BASELINE }, 0, 0, VE, "profile" },
+	{ "main",     "",                         0,               AV_OPT_TYPE_CONST, { .i64 = FF_PROFILE_H264_MAIN },     0, 0, VE, "profile" },
+	{ "high",     "",                         0,               AV_OPT_TYPE_CONST, { .i64 = FF_PROFILE_H264_HIGH },     0, 0, VE, "profile" },
+
+	/// Profile Level
+	{ "level",          "Profile Level",        OFFSET(level),  AV_OPT_TYPE_INT,   { .i64 = 0  }, 0, 62, VE, "level" },
+	{ "auto",           "",                     0,              AV_OPT_TYPE_CONST, { .i64 = 0  }, 0, 0,  VE, "level" },
+	{ "1.0",            "",                     0,              AV_OPT_TYPE_CONST, { .i64 = 10 }, 0, 0,  VE, "level" },
+	{ "1.1",            "",                     0,              AV_OPT_TYPE_CONST, { .i64 = 11 }, 0, 0,  VE, "level" },
+	{ "1.2",            "",                     0,              AV_OPT_TYPE_CONST, { .i64 = 12 }, 0, 0,  VE, "level" },
+	{ "1.3",            "",                     0,              AV_OPT_TYPE_CONST, { .i64 = 13 }, 0, 0,  VE, "level" },
+	{ "2.0",            "",                     0,              AV_OPT_TYPE_CONST, { .i64 = 20 }, 0, 0,  VE, "level" },
+	{ "2.1",            "",                     0,              AV_OPT_TYPE_CONST, { .i64 = 21 }, 0, 0,  VE, "level" },
+	{ "2.2",            "",                     0,              AV_OPT_TYPE_CONST, { .i64 = 22 }, 0, 0,  VE, "level" },
+	{ "3.0",            "",                     0,              AV_OPT_TYPE_CONST, { .i64 = 30 }, 0, 0,  VE, "level" },
+	{ "3.1",            "",                     0,              AV_OPT_TYPE_CONST, { .i64 = 31 }, 0, 0,  VE, "level" },
+	{ "3.2",            "",                     0,              AV_OPT_TYPE_CONST, { .i64 = 32 }, 0, 0,  VE, "level" },
+	{ "4.0",            "",                     0,              AV_OPT_TYPE_CONST, { .i64 = 40 }, 0, 0,  VE, "level" },
+	{ "4.1",            "",                     0,              AV_OPT_TYPE_CONST, { .i64 = 41 }, 0, 0,  VE, "level" },
+	{ "4.2",            "",                     0,              AV_OPT_TYPE_CONST, { .i64 = 42 }, 0, 0,  VE, "level" },
+	{ "5.0",            "",                     0,              AV_OPT_TYPE_CONST, { .i64 = 50 }, 0, 0,  VE, "level" },
+	{ "5.1",            "",                     0,              AV_OPT_TYPE_CONST, { .i64 = 51 }, 0, 0,  VE, "level" },
+
+	{ "rc",           "Override the preset rate-control",   OFFSET(rc),           AV_OPT_TYPE_INT,   { .i64 = -1 },                                  -1, INT_MAX, VE, "rc" },
+	{ "cbr",          "Constant bitrate mode",              0,                    AV_OPT_TYPE_CONST, { .i64 = 0 },                       0, 0, VE, "rc" },
+	{ "vbr",          "Variable bitrate mode",              0,                    AV_OPT_TYPE_CONST, { .i64 = 1 },                       0, 0, VE, "rc" },
+
+	{ "preset",          "Set the encoding preset",            OFFSET(preset),       AV_OPT_TYPE_INT,   { .i64 = 3 }, 1, 4, VE, "preset" },
+	{ "default",         "",                                   0,                    AV_OPT_TYPE_CONST, { .i64 = 3 }, 0, 0, VE, "preset" },
+	{ "slow",            "",                        0,                    AV_OPT_TYPE_CONST, { .i64 = 4 },            0, 0, VE, "preset" },
+	{ "medium",          "",                        0,                    AV_OPT_TYPE_CONST, { .i64 = 3 },            0, 0, VE, "preset" },
+	{ "fast",            "",                        0,                    AV_OPT_TYPE_CONST, { .i64 = 2 },            0, 0, VE, "preset" },
+	{ "ultrafast",       "",                        0,                    AV_OPT_TYPE_CONST, { .i64 = 1 },            0, 0, VE, "preset" },
+	{ NULL }
+};
+
+
+#define NVMPI_ENC_CLASS(NAME) \
+	static const AVClass nvmpi_ ## NAME ## _enc_class = { \
+		.class_name = #NAME "_nvmpi_encoder", \
+		.item_name  = av_default_item_name, \
+		.option     = options, \
+		.version    = LIBAVUTIL_VERSION_INT, \
+	};
+
+
+#define NVMPI_ENC(NAME, LONGNAME, CODEC) \
+	NVMPI_ENC_CLASS(NAME) \
+	AVCodec ff_ ## NAME ## _nvmpi_encoder = { \
+		.name           = #NAME "_nvmpi" , \
+		.long_name      = NULL_IF_CONFIG_SMALL("nvmpi " LONGNAME " encoder wrapper"), \
+		.type           = AVMEDIA_TYPE_VIDEO, \
+		.id             = CODEC , \
+		.priv_data_size = sizeof(nvmpiEncodeContext), \
+		.priv_class     = &nvmpi_ ## NAME ##_enc_class, \
+		.init           = nvmpi_encode_init, \
+		.send_frame     = ff_nvmpi_send_frame, \
+		.receive_packet = ff_nvmpi_receive_packet, \
+		.encode2        = ff_nvmpi_encode_frame, \
+		.close          = nvmpi_encode_close, \
+		.pix_fmts       = (const enum AVPixelFormat[]) { AV_PIX_FMT_YUV420P, AV_PIX_FMT_NONE },\
+		.capabilities   = AV_CODEC_CAP_HARDWARE | AV_CODEC_CAP_DELAY, \
+		.defaults       = defaults,\
+		.wrapper_name   = "nvmpi", \
+	};
+
+NVMPI_ENC(h264, "H.264", AV_CODEC_ID_H264);
+NVMPI_ENC(hevc, "HEVC", AV_CODEC_ID_HEVC);
-- 
2.25.1