Skip to content

Instantly share code, notes, and snippets.

@asigatchov

asigatchov/ffmpeg_nvmpi_with_YUVJ420P.patch

Created March 2, 2021 13:39
Show Gist options
  • Save asigatchov/a91f4fecb8b0fb41c4ee3939c53c07fa to your computer and use it in GitHub Desktop.
Save asigatchov/a91f4fecb8b0fb41c4ee3939c53c07fa to your computer and use it in GitHub Desktop.
Download ZIP
Raw
ffmpeg_nvmpi_with_YUVJ420P.patch
diff --git a/configure b/configure
index 6a7a85c..5810ab5 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) && (avctx->pix_fmt != AV_PIX_FMT_YUVJ420P)){
+ av_log(avctx, AV_LOG_ERROR, "Invalid Pix_FMT for NVMPI: Only YUV420P and YUVJ420P are 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..0e1390e
--- /dev/null
+++ b/libavcodec/nvmpi_enc.c
@@ -0,0 +1,272 @@
+#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;
+}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 nvmpi_encode_frame(AVCodecContext *avctx, AVPacket *pkt,const AVFrame *frame, int *got_packet){
+
+ nvmpiEncodeContext * nvmpi_context = avctx->priv_data;
+ nvFrame _nvframe={0};
+ nvPacket packet={0};
+ int res;
+
+ 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;
+ }
+
+
+ if(nvmpi_encoder_get_packet(nvmpi_context->ctx,&packet)<0)
+ return 0;
+
+
+ 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;
+
+
+ *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, \
+ .encode2 = 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);
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment