nvmpi_dec.cpp

#include "nvmpi.h"
#include "NvVideoDecoder.h"

#include "nvbufsurface.h"
#include "nvbufsurftransform.h"
#include "NvBufSurface.h"
#include <vector>
#include <iostream>
#include <thread>
#include <unistd.h>
#include <queue>
#include <mutex>
#include <condition_variable>

#define CHUNK_SIZE 4000000
#define MAX_BUFFERS 32

#define TEST_ERROR(condition, message, errorCode)    \
	if (condition)                               \
{                                                    \
	std::cout<< message;			     \
}

using namespace std;

struct nvmpictx
{
	NvVideoDecoder *dec{nullptr};
	bool eos{false};
	bool got_res_event{false};
	int index{0};
	unsigned int coded_width{0};
	unsigned int coded_height{0};
	int dst_dma_fd{0};
	int numberCaptureBuffers{0};
	int dmaBufferFileDescriptor[MAX_BUFFERS];
	nvPixFormat out_pixfmt;
	unsigned int decoder_pixfmt{0};
	std::thread * dec_capture_loop{nullptr};
	std::mutex* mutex{nullptr};
	std::condition_variable* has_frame_cv{nullptr};
	std::queue<int> * frame_pools{nullptr};
	unsigned char * bufptr_0[MAX_BUFFERS];
	unsigned char * bufptr_1[MAX_BUFFERS];
	unsigned char * bufptr_2[MAX_BUFFERS];
	unsigned int frame_size[NVBUF_MAX_PLANES];
	unsigned int frame_linesize[NVBUF_MAX_PLANES];
	unsigned long long timestamp[MAX_BUFFERS];
};

void respondToResolutionEvent(v4l2_format &format, v4l2_crop &crop,nvmpictx* ctx){
	int32_t minimumDecoderCaptureBuffers;
	int ret=0;
        NvBufSurf::NvCommonAllocateParams params = {0};
	NvBufSurf::NvCommonAllocateParams capParams = {0};

	ret = ctx->dec->capture_plane.getFormat(format);
	TEST_ERROR(ret < 0, "Error: Could not get format from decoder capture plane", ret);

	ret = ctx->dec->capture_plane.getCrop(crop);
	TEST_ERROR(ret < 0, "Error: Could not get crop from decoder capture plane", ret);

	ctx->coded_width=crop.c.width;
	ctx->coded_height=crop.c.height;

	if(ctx->dst_dma_fd != -1)
	{
        	ret = NvBufSurf::NvDestroy(ctx->dst_dma_fd);
        	ctx->dst_dma_fd = -1;
	}

	params.memType = NVBUF_MEM_SURFACE_ARRAY;
	params.width = crop.c.width;
	params.height = crop.c.height;
	params.layout = NVBUF_LAYOUT_PITCH;
	params.colorFormat = ctx->out_pixfmt==NV_PIX_NV12?NVBUF_COLOR_FORMAT_NV12:NVBUF_COLOR_FORMAT_YUV420;
	params.memtag = NvBufSurfaceTag_VIDEO_DEC;

	ctx->dec->capture_plane.deinitPlane();

	for (int index = 0; index < ctx->numberCaptureBuffers; index++)
	{
		if (ctx->dmaBufferFileDescriptor[index] != 0)
		{
			ret = NvBufSurf::NvDestroy(ctx->dmaBufferFileDescriptor[index]);
			TEST_ERROR(ret < 0, "Failed to Destroy NvBuffer", ret);
		}

	}

	ret=ctx->dec->setCapturePlaneFormat(format.fmt.pix_mp.pixelformat,format.fmt.pix_mp.width,format.fmt.pix_mp.height);
	TEST_ERROR(ret < 0, "Error in setting decoder capture plane format", ret);

	ret=ctx->dec->setMaxPerfMode(1);
	TEST_ERROR(ret < 0, "Error in setting decoder maximum performance mode", ret);

	ctx->dec->getMinimumCapturePlaneBuffers(minimumDecoderCaptureBuffers);
	TEST_ERROR(ret < 0, "Error while getting value of minimum capture plane buffers",ret);

	ctx->numberCaptureBuffers = minimumDecoderCaptureBuffers + 5;

	switch (format.fmt.pix_mp.colorspace)
	{
		case V4L2_COLORSPACE_SMPTE170M:
			if (format.fmt.pix_mp.quantization == V4L2_QUANTIZATION_DEFAULT)
			{
				// "Decoder colorspace ITU-R BT.601 with standard range luma (16-235)"
				capParams.colorFormat = NVBUF_COLOR_FORMAT_NV12;
			}
			else
			{
				//"Decoder colorspace ITU-R BT.601 with extended range luma (0-255)";
				capParams.colorFormat = NVBUF_COLOR_FORMAT_NV12_ER;
			}
			break;
		case V4L2_COLORSPACE_REC709:
			if (format.fmt.pix_mp.quantization == V4L2_QUANTIZATION_DEFAULT)
			{
				//"Decoder colorspace ITU-R BT.709 with standard range luma (16-235)";
				capParams.colorFormat = NVBUF_COLOR_FORMAT_NV12_709;
			}
			else
			{
				//"Decoder colorspace ITU-R BT.709 with extended range luma (0-255)";
				capParams.colorFormat = NVBUF_COLOR_FORMAT_NV12_709_ER;
			}
			break;
		case V4L2_COLORSPACE_BT2020:
			{
				//"Decoder colorspace ITU-R BT.2020";
				capParams.colorFormat = NVBUF_COLOR_FORMAT_NV12_2020;
			}
			break;
		default:
			if (format.fmt.pix_mp.quantization == V4L2_QUANTIZATION_DEFAULT)
			{
				//"Decoder colorspace ITU-R BT.601 with standard range luma (16-235)";
				capParams.colorFormat = NVBUF_COLOR_FORMAT_NV12;
			}
			else
			{
				//"Decoder colorspace ITU-R BT.601 with extended range luma (0-255)";
				capParams.colorFormat = NVBUF_COLOR_FORMAT_NV12_ER;
			}
			break;
	}




        ret = NvBufSurf::NvAllocate(&params, 1, &ctx->dst_dma_fd);
        TEST_ERROR(ret == -1, "create dst_dmabuf failed", error);

        capParams.width = crop.c.width;
        capParams.height = crop.c.height;
        capParams.layout = NVBUF_LAYOUT_BLOCK_LINEAR;;
        capParams.memType = NVBUF_MEM_SURFACE_ARRAY;
        capParams.memtag = NvBufSurfaceTag_VIDEO_DEC;
        ret = NvBufSurf::NvAllocate(&capParams, ctx->numberCaptureBuffers, ctx->dmaBufferFileDescriptor);
	TEST_ERROR(ret < 0, "Failed to create buffers", ret);

	ctx->dec->capture_plane.reqbufs(V4L2_MEMORY_DMABUF, ctx->numberCaptureBuffers);
	TEST_ERROR(ret < 0, "Error in decoder capture plane streamon", ret);

	ctx->dec->capture_plane.setStreamStatus(true);
	TEST_ERROR(ret < 0, "Error in decoder capture plane streamon", ret);


	for (uint32_t i = 0; i < ctx->dec->capture_plane.getNumBuffers(); i++)
	{
		struct v4l2_buffer v4l2_buf;
		struct v4l2_plane planes[MAX_PLANES];

		memset(&v4l2_buf, 0, sizeof(v4l2_buf));
		memset(planes, 0, sizeof(planes));

		v4l2_buf.index = i;
		v4l2_buf.m.planes = planes;
		v4l2_buf.type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
		v4l2_buf.memory = V4L2_MEMORY_DMABUF;
		v4l2_buf.m.planes[0].m.fd = ctx->dmaBufferFileDescriptor[i];

		ret = ctx->dec->capture_plane.qBuffer(v4l2_buf, NULL);
		TEST_ERROR(ret < 0, "Error Qing buffer at output plane", ret);
	}

	ctx->got_res_event = true;
}

void dec_capture_loop_fcn(void *arg){
	nvmpictx* ctx=(nvmpictx*)arg;

	struct v4l2_format v4l2Format;
	struct v4l2_crop v4l2Crop;
	struct v4l2_event v4l2Event;
	int ret,buf_index=0;

	while (!(ctx->dec->isInError()||ctx->eos)){
		NvBuffer *dec_buffer;

		ret = ctx->dec->dqEvent(v4l2Event, ctx->got_res_event ? 0 : 500);
		if (ret == 0)
		{
			switch (v4l2Event.type)
			{
				case V4L2_EVENT_RESOLUTION_CHANGE:
					respondToResolutionEvent(v4l2Format, v4l2Crop,ctx);
					continue;
			}
		}

		if (!ctx->got_res_event) {
			continue;
		}

		while(!ctx->eos){
			struct v4l2_buffer v4l2_buf;
			struct v4l2_plane planes[MAX_PLANES];
			v4l2_buf.m.planes = planes;

			if (ctx->dec->capture_plane.dqBuffer(v4l2_buf, &dec_buffer, NULL, 0)){
				if (errno == EAGAIN)
				{
					usleep(1000);
				}
				else
				{

					ERROR_MSG("Error while calling dequeue at capture plane");
					ctx->eos=true;
				}
				break;

			}

			dec_buffer->planes[0].fd = ctx->dmaBufferFileDescriptor[v4l2_buf.index];

			NvBufSurf::NvCommonTransformParams transform_params;

                        transform_params.src_top = 0;
			transform_params.src_left = 0;
			transform_params.src_width = ctx->coded_width;
			transform_params.src_height = ctx->coded_height;

			transform_params.dst_top = 0;
			transform_params.dst_left = 0;
			transform_params.dst_width = ctx->coded_width;
			transform_params.dst_height = ctx->coded_height;


			transform_params.flag = NVBUFSURF_TRANSFORM_FILTER;
			transform_params.flip = NvBufSurfTransform_None;
			transform_params.filter = NvBufSurfTransformInter_Algo3;

			ctx->mutex->lock();

			if(!ctx->eos){
                                ret = NvBufSurf::NvTransform(&transform_params, dec_buffer->planes[0].fd, ctx->dst_dma_fd);
				TEST_ERROR(ret==-1, "Transform failed",ret);

                                NvBufSurface* buf_surf = 0;
                                NvBufSurfaceFromFd(ctx->dst_dma_fd, (void**)(&buf_surf));
				NvBufSurfaceParams surParm = buf_surf->surfaceList[0];
				NvBufSurfacePlaneParams parm = surParm.planeParams;

                                // cout <<  "batchSize:" << buf_surf->batchSize << endl;
                                // cout <<  "memType:" << buf_surf->memType << endl;
                                // cout <<  "width:" << surParm.width << endl;
                                // cout <<  "height:" << surParm.height << endl;
                                // cout <<  "dataSize:" << surParm.dataSize << endl;
                                // cout <<  "num_p:" << parm.num_planes << endl;
                                // cout <<  "psize0:" << parm.psize[0] << endl;
                                // cout <<  "psize1:" << parm.psize[1] << endl;
                                // cout <<  "psize2:" << parm.psize[2] << endl;
                                // cout <<  "bytesPerPix:" << parm.bytesPerPix << endl;


				if(!ctx->frame_size[0]){
					for(int index=0;index<MAX_BUFFERS;index++){
                                          ctx->bufptr_0[index]=new unsigned char[parm.psize[0]];//Y
                                          ctx->bufptr_1[index]=new unsigned char[parm.psize[1]];//UV or UU
                                          ctx->bufptr_2[index]=new unsigned char[parm.psize[2]];//VV
					}
				}


				ctx->frame_linesize[0]=parm.width[0];
				ctx->frame_size[0]=parm.psize[0];

				ctx->frame_linesize[1]=parm.width[1];
				ctx->frame_size[1]=parm.psize[1];
				ctx->frame_linesize[2]=parm.width[2];
				ctx->frame_size[2]=parm.psize[2];

                                ret=NvBufSurface2Raw(buf_surf, 0, 0, parm.width[0], parm.height[0],ctx->bufptr_0[buf_index]);
                                TEST_ERROR(ret < 0, "Copy plane0 to mem error", ret);
                                ret=NvBufSurface2Raw(buf_surf, 0, 1, parm.width[1], parm.height[1],ctx->bufptr_1[buf_index]);
                                TEST_ERROR(ret < 0, "Copy plane1 to mem error", ret);
				if(ctx->out_pixfmt==NV_PIX_YUV420)
                                        ret=NvBufSurface2Raw(buf_surf, 0, 2, parm.width[2], parm.height[2],ctx->bufptr_2[buf_index]);
                                        TEST_ERROR(ret < 0, "Copy plane2 to mem error", ret);

				ctx->frame_pools->push(buf_index);
				ctx->timestamp[buf_index]= (v4l2_buf.timestamp.tv_usec % 1000000) + (v4l2_buf.timestamp.tv_sec * 1000000UL);

				buf_index=(buf_index+1)%MAX_BUFFERS;

			}

			ctx->mutex->unlock();

			if (ctx->eos) {
				break;
			}

			ctx->has_frame_cv->notify_one();

			v4l2_buf.m.planes[0].m.fd = ctx->dmaBufferFileDescriptor[v4l2_buf.index];
			if (ctx->dec->capture_plane.qBuffer(v4l2_buf, NULL) < 0){
				ERROR_MSG("Error while queueing buffer at decoder capture plane");
			}
		}
	}

	// Wake all waiting threads at EOS or decoder error
	ctx->has_frame_cv->notify_all();
}

nvmpictx* nvmpi_create_decoder(nvCodingType codingType,nvPixFormat pixFormat){
	int ret;
	log_level = LOG_LEVEL_INFO;

	nvmpictx* ctx=new nvmpictx;

	ctx->dec = NvVideoDecoder::createVideoDecoder("dec0");
	TEST_ERROR(!ctx->dec, "Could not create decoder",ret);

	ret=ctx->dec->subscribeEvent(V4L2_EVENT_RESOLUTION_CHANGE, 0, 0);
	TEST_ERROR(ret < 0, "Could not subscribe to V4L2_EVENT_RESOLUTION_CHANGE", ret);

	switch(codingType){
		case NV_VIDEO_CodingH264:
			ctx->decoder_pixfmt=V4L2_PIX_FMT_H264;
			break;
		case NV_VIDEO_CodingHEVC:
			ctx->decoder_pixfmt=V4L2_PIX_FMT_H265;
			break;
		case NV_VIDEO_CodingMPEG4:
			ctx->decoder_pixfmt=V4L2_PIX_FMT_MPEG4;
			break;
		case NV_VIDEO_CodingMPEG2:
			ctx->decoder_pixfmt=V4L2_PIX_FMT_MPEG2;
			break;
		case NV_VIDEO_CodingVP8:
			ctx->decoder_pixfmt=V4L2_PIX_FMT_VP8;
			break;
		case NV_VIDEO_CodingVP9:
			ctx->decoder_pixfmt=V4L2_PIX_FMT_VP9;
			break;
		default:
			ctx->decoder_pixfmt=V4L2_PIX_FMT_H264;
			break;
	}

	ret=ctx->dec->setOutputPlaneFormat(ctx->decoder_pixfmt, CHUNK_SIZE);
	TEST_ERROR(ret < 0, "Could not set output plane format", ret);

	ret=ctx->dec->setMaxPerfMode(1);
	TEST_ERROR(ret < 0, "Error in setting decoder maximum performance mode", ret);

	//ctx->nalu_parse_buffer = new char[CHUNK_SIZE];
	ret = ctx->dec->setFrameInputMode(0);
	TEST_ERROR(ret < 0, "Error in decoder setFrameInputMode for NALU", ret);

	ret = ctx->dec->output_plane.setupPlane(V4L2_MEMORY_USERPTR, 10, false, true);
	TEST_ERROR(ret < 0, "Error while setting up output plane", ret);

	ctx->dec->output_plane.setStreamStatus(true);
	TEST_ERROR(ret < 0, "Error in output plane stream on", ret);

	ctx->out_pixfmt=pixFormat;
	ctx->dst_dma_fd=-1;
	ctx->eos=false;
	ctx->got_res_event=false;
	ctx->index=0;
	ctx->frame_size[0]=0;
	ctx->frame_pools=new std::queue<int>;
	ctx->mutex = new std::mutex();
	ctx->has_frame_cv = new std::condition_variable();
	for(int index=0;index<MAX_BUFFERS;index++)
		ctx->dmaBufferFileDescriptor[index]=0;
	for(int index=0;index<MAX_BUFFERS;index++){
		ctx->bufptr_0[index] = nullptr;
		ctx->bufptr_1[index] = nullptr;
		ctx->bufptr_2[index] = nullptr;
	}
	ctx->numberCaptureBuffers=0;
	ctx->dec_capture_loop=new thread(dec_capture_loop_fcn,ctx);

	return ctx;
}




int nvmpi_decoder_put_packet(nvmpictx* ctx,nvPacket* packet){
	int ret;

	struct v4l2_buffer v4l2_buf;
	struct v4l2_plane planes[MAX_PLANES];
	NvBuffer *nvBuffer;

	memset(&v4l2_buf, 0, sizeof(v4l2_buf));
	memset(planes, 0, sizeof(planes));

	v4l2_buf.m.planes = planes;

	if (ctx->index < (int)ctx->dec->output_plane.getNumBuffers()) {
		nvBuffer = ctx->dec->output_plane.getNthBuffer(ctx->index);
	} else {
		ret = ctx->dec->output_plane.dqBuffer(v4l2_buf, &nvBuffer, NULL, -1);
		if (ret < 0) {
			cout << "Error DQing buffer at output plane" << std::endl;
			return false;
		}
	}

	memcpy(nvBuffer->planes[0].data,packet->payload,packet->payload_size);
	nvBuffer->planes[0].bytesused=packet->payload_size;



	if (ctx->index < ctx->dec->output_plane.getNumBuffers())
	{
		v4l2_buf.index = ctx->index ;
		v4l2_buf.m.planes = planes;
	}

	v4l2_buf.m.planes[0].bytesused = nvBuffer->planes[0].bytesused;

	v4l2_buf.flags |= V4L2_BUF_FLAG_TIMESTAMP_COPY;
	v4l2_buf.timestamp.tv_sec = packet->pts / 1000000;
	v4l2_buf.timestamp.tv_usec = packet->pts % 1000000;


	ret = ctx->dec->output_plane.qBuffer(v4l2_buf, NULL);
	if (ret < 0) {
		std::cout << "Error Qing buffer at output plane" << std::endl;
		return false;
	}

	if (ctx->index < ctx->dec->output_plane.getNumBuffers())
		ctx->index++;

	if (v4l2_buf.m.planes[0].bytesused == 0) {
		ctx->eos=true;
		std::cout << "Input file read complete" << std::endl;
	}


	return 0;

}


int nvmpi_decoder_get_frame(nvmpictx* ctx,nvFrame* frame,bool wait){
	int ret,picture_index;
	std::unique_lock<std::mutex> lock(*ctx->mutex);

	if (wait) {
		while (ctx->frame_pools->empty() && !ctx->eos && !ctx->dec->isInError()) {
			ctx->has_frame_cv->wait(lock);
		}
	}

	if (ctx->frame_pools->empty()) {
		return -1;
	}

	picture_index=ctx->frame_pools->front();
	ctx->frame_pools->pop();
	frame->width=ctx->coded_width;
	frame->height=ctx->coded_height;

	frame->linesize[0]=ctx->frame_linesize[0];
	frame->linesize[1]=ctx->frame_linesize[1];
	frame->linesize[2]=ctx->frame_linesize[2];

	frame->payload[0]=ctx->bufptr_0[picture_index];
	frame->payload[1]=ctx->bufptr_1[picture_index];
	frame->payload[2]=ctx->bufptr_2[picture_index];

	frame->payload_size[0]=ctx->frame_size[0];
	frame->payload_size[1]=ctx->frame_size[1];
	frame->payload_size[2]=ctx->frame_size[2];
	frame->timestamp=ctx->timestamp[picture_index];

	return 0;

}

int nvmpi_decoder_close(nvmpictx* ctx){
	ctx->mutex->lock();
	ctx->eos=true;
	ctx->mutex->unlock();

	ctx->dec->capture_plane.setStreamStatus(false);

	if (ctx->dec_capture_loop) {
		ctx->dec_capture_loop->join();
		delete ctx->dec_capture_loop;
		ctx->dec_capture_loop = nullptr;
	}

	if(ctx->dst_dma_fd != -1)
	{
		NvBufSurf::NvDestroy(ctx->dst_dma_fd);
		ctx->dst_dma_fd = -1;
	}

	for (int index = 0; index < ctx->numberCaptureBuffers; index++)
	{
		if (ctx->dmaBufferFileDescriptor[index] != 0)
		{
			int ret = NvBufSurf::NvDestroy(ctx->dmaBufferFileDescriptor[index]);
			TEST_ERROR(ret < 0, "Failed to Destroy NvBuffer", ret);
		}

	}

	delete ctx->dec; ctx->dec = nullptr;

	for(int index=0;index<MAX_BUFFERS;index++){
		delete[] ctx->bufptr_0[index];
		delete[] ctx->bufptr_1[index];
		delete[] ctx->bufptr_2[index];
	}

	delete ctx->mutex; ctx->mutex = nullptr;
	delete ctx->has_frame_cv; ctx->has_frame_cv = nullptr;
	delete ctx->frame_pools; ctx->frame_pools = nullptr;

	delete ctx; ctx = nullptr;

	return 0;
}