Hu Chunming / FFNvDecoder

#include "FFNvDecoder.h"

#include <chrono>
#include <thread>
#include <fstream>

#include "FFCuContextManager.h"

using namespace std;

// 参考博客： https://blog.csdn.net/qq_40116098/article/details/120704340

static AVPixelFormat get_hw_format(AVCodecContext *avctx, const AVPixelFormat *pix_fmts)
{
	FFNvDecoder* _this = (FFNvDecoder*)avctx->opaque;

	const AVPixelFormat *p;

	for (p = pix_fmts; *p != -1; p++) {
		if (*p == _this->getHwPixFmt())
			return *p;
	}

	av_log(nullptr, AV_LOG_ERROR, "Failed to get HW surface format. \n");
	return AV_PIX_FMT_NONE;
}

FFNvDecoder::FFNvDecoder()
{
	// 初始化解码对象
	fmt_ctx = nullptr;
	avctx = nullptr;
	m_bRunning = false;

	stream = nullptr;
    stream_index = -1;
    hw_pix_fmt = AV_PIX_FMT_NONE;
    name = "";

	m_bPause = false;
	m_bReal = true;

	m_decode_thread = 0;
	m_post_decode_thread = 0;

	m_bFinished = false;
	m_dec_keyframe = false;
}

FFNvDecoder::~FFNvDecoder()
{
	m_dec_keyframe = false;
}

bool FFNvDecoder::init(FFDecConfig& cfg)
{
	m_cfg = cfg;

	fstream infile(cfg.uri);
	if (infile.is_open()){
		m_bReal = false;
		infile.close();
	}else {
		m_bReal = true;
	}

	return init(cfg.uri.c_str(), cfg.gpuid.c_str(),cfg.force_tcp);
}

bool FFNvDecoder::init(const char* uri, const char* gpuid, bool force_tcp)
{
	// av_log_set_level(AV_LOG_DEBUG);

	avformat_network_init();

	// 打开输入视频文件
	AVDictionary *options = nullptr;
	av_dict_set( &options, "bufsize", "655360", 0 );
	av_dict_set( &options, "rtsp_transport", force_tcp ? "tcp" : "udp", 0 );
	// av_dict_set( &options, "listen_timeout", "30", 0 ); // 单位为s
	av_dict_set( &options, "stimeout", "30000000", 0 ); // 单位为 百万分之一秒
	
	fmt_ctx = avformat_alloc_context();
	const char* input_file = uri;
	if (avformat_open_input(&fmt_ctx, input_file, nullptr, &options) != 0) {
		av_log(nullptr, AV_LOG_ERROR, "Cannot open input file: %s \n", input_file);
		return false;
	}

	// 查找流信息
	if (avformat_find_stream_info(fmt_ctx, nullptr) < 0) {
		av_log(nullptr, AV_LOG_ERROR, "Cannot find input stream information ! \n");
		return false;
	}

	// 查找视频流信息
	AVCodec *decoder = nullptr;
	stream_index = av_find_best_stream(fmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, &decoder, 0);
	if (stream_index < 0) {
		av_log(nullptr, AV_LOG_ERROR, "Cannot find a video stream in the input file ! \n");
		return false;
	}

	string cuvid_dec_name = string(decoder->name) + "_cuvid";
	AVCodec *vcodec = avcodec_find_decoder_by_name(cuvid_dec_name.c_str());
	if (!(avctx = avcodec_alloc_context3(vcodec)))
		return (bool)AVERROR(ENOMEM);

	// 得到视频流对象
	stream = fmt_ctx->streams[stream_index];
	if (avcodec_parameters_to_context(avctx, stream->codecpar) < 0)
		return false;

	avctx->opaque = this;
	// 设置解码器管理器的像素格式回调函数
	avctx->get_format = get_hw_format;

	hw_pix_fmt = AV_PIX_FMT_CUDA;

	FFCuContextManager* pCtxMgr = FFCuContextManager::getInstance();

	AVBufferRef *hw_device_ctx = pCtxMgr->getCuCtx(gpuid);
	if(nullptr == hw_device_ctx){
		av_log(nullptr, AV_LOG_ERROR, "create CUDA context failed ! \n");
		return false;
	}
	avctx->hw_device_ctx = av_buffer_ref(hw_device_ctx);
	if (nullptr == avctx->hw_device_ctx)
	{
		return false;
	}

	// 打开解码器流
	AVDictionary *op = nullptr;
	av_dict_set( &op, "gpu", gpuid, 0 );
	// av_dict_set( &op, "surfaces", "10", 0 );
	if (avcodec_open2(avctx, vcodec, &op) < 0) {
		av_log(nullptr, AV_LOG_ERROR, "Failed to open codec for stream ! \n");
		return false;
	}

	return true;
}

bool FFNvDecoder::isSurport(FFDecConfig& cfg)
{
	bool bRet = init(cfg);
    decode_finished();
    return bRet;
}

void FFNvDecoder::start(){

	m_bRunning = true;

	pthread_create(&m_decode_thread,0,
        [](void* arg)
        {
            FFNvDecoder* a=(FFNvDecoder*)arg;
            a->decode_thread();
            return (void*)0;
        }
    ,this);
}

static long long get_cur_time(){
    // 获取操作系统当前时间点（精确到微秒）
    chrono::time_point<chrono::system_clock, chrono::milliseconds> tpMicro
        = chrono::time_point_cast<chrono::milliseconds>(chrono::system_clock::now());
    // (微秒精度的)时间点 => (微秒精度的)时间戳
    time_t currentTime = tpMicro.time_since_epoch().count();

    return (long long )currentTime;
}

void FFNvDecoder::decode_thread()
{
	AVPacket* pkt ;
	pkt = av_packet_alloc();
	av_init_packet( pkt );

	pthread_create(&m_post_decode_thread,0,
        [](void* arg)
        {
            FFNvDecoder* a=(FFNvDecoder*)arg;
            a->post_decode_thread();
            return (void*)0;
        }
    ,this);

	// long start_time = get_cur_time();

	while (m_bRunning)
	{
		if (!m_bReal)
		{
			if (m_bPause)
			{
				std::this_thread::sleep_for(std::chrono::milliseconds(3));
				continue;
			}
		}

		AVFrame * gpuFrame = mFrameQueue.getTail();
		if (gpuFrame == nullptr)
		{
			std::this_thread::sleep_for(std::chrono::milliseconds(1));
			continue;
		}
		
		int result = av_read_frame(fmt_ctx, pkt);
		if (result == AVERROR_EOF || result < 0)
		{
			av_log(nullptr, AV_LOG_ERROR, "Failed to read frame! \n");
			break;
		}

		if (m_dec_keyframe && !(pkt->flags & AV_PKT_FLAG_KEY)) {
			av_packet_unref(pkt);
			continue;
		}

		if (m_bReal)
		{
			if (m_bPause)
			{
				av_packet_unref(pkt);
				std::this_thread::sleep_for(std::chrono::milliseconds(3));
				continue;
			}
		}

		if (stream_index == pkt->stream_index){
			result = avcodec_send_packet(avctx, pkt);
			if (result < 0){
				av_log(nullptr, AV_LOG_ERROR, "Failed to send pkt: %d \n",result);
				continue;
			}

			result = avcodec_receive_frame(avctx, gpuFrame);
			if ((result == AVERROR(EAGAIN) || result == AVERROR_EOF) || result < 0){
				av_log(nullptr, AV_LOG_ERROR, "Failed to receive frame: %d \n",result);
				continue;
			}

			mFrameQueue.addTail();
		}
		av_packet_unref(pkt);
	}

	// 队列中没有数据了再结束
	while (mFrameQueue.length() > 0){
		if(!m_bRunning){
			break;
		}
		std::this_thread::sleep_for(std::chrono::milliseconds(10));
	}

	m_bRunning = false;

	// long end_time = get_cur_time();

	// cout << "解码用时：" << end_time - start_time << endl;

	if (m_post_decode_thread != 0)
	{
		pthread_join(m_post_decode_thread,0);
	}

	decode_finished_cbk(m_userPtr);

	decode_finished();

	av_log(nullptr, AV_LOG_INFO, "decode thread exited. \n");
}

void FFNvDecoder::decode_finished()
{
	if (avctx)
	{
		avcodec_free_context(&avctx);
	}
	
	if (fmt_ctx)
	{
		avformat_close_input(&fmt_ctx);
	}

	m_bFinished = true;
	m_dec_keyframe = false;
}

void FFNvDecoder::post_decode_thread()
{
	 while (m_bRunning || mFrameQueue.length() > 0)
	 {
		AVFrame * gpuFrame = mFrameQueue.getHead();
		if (gpuFrame == nullptr)
		{
			std::this_thread::sleep_for(std::chrono::milliseconds(3));
			continue;
		}

		post_decoded_cbk(m_userPtr, gpuFrame);

		mFrameQueue.addHead();
	 }

	 av_log(nullptr, AV_LOG_INFO, "post decode thread exited. \n");
}

void FFNvDecoder::close()
{
	m_bRunning=false;
	if(m_decode_thread != 0){
		pthread_join(m_decode_thread,0);
	}
	m_dec_keyframe = false;
}

void FFNvDecoder::setName(string nm){
	name = nm;
}

string FFNvDecoder::getName(){
	return name;
}

AVPixelFormat FFNvDecoder::getHwPixFmt()
{
	return hw_pix_fmt;
}

bool FFNvDecoder::isRunning()
{
	return m_bRunning;
}

bool FFNvDecoder::isFinished()
{
	return m_bFinished;
}

bool FFNvDecoder::isPausing(){
	return m_bPause;
}

bool FFNvDecoder::getResolution( int &width, int &height )
{
	if (stream != nullptr && stream->codecpar != nullptr)
	{
		width = stream->codecpar->width;
		height = stream->codecpar->height;
		return true;
	}
	
	return false;
}

void FFNvDecoder::pause()
{
	m_bPause = true;
}

void FFNvDecoder::resume()
{
	m_bPause = false;
}

void FFNvDecoder::setDecKeyframe(bool bKeyframe)
{
	m_dec_keyframe = bKeyframe;
}

int FFNvDecoder::getCachedQueueLength(){
	return mFrameQueue.length();
}

FFImgInfo* FFNvDecoder::snapshot(const string& uri){
 
    AVFormatContext* ifmt_ctx = nullptr;
	AVCodecContext* codec_ctx = nullptr;
	AVCodec* codec = nullptr;
	AVPacket* pkt = nullptr;
	AVFrame *frame = nullptr;
	AVFrame *pFrameRGB = nullptr;	
	int video_index = -1;
	AVStream* st = nullptr;
	SwsContext *img_convert_ctx = nullptr;
	uint8_t *buffer = nullptr;
    int numBytes = 0;

	FFImgInfo* imgInfo = nullptr;


	avformat_network_init();

	// 打开输入视频文件
	AVDictionary *options = nullptr;
	av_dict_set( &options, "bufsize", "655360", 0 );
	av_dict_set( &options, "rtsp_transport", "tcp", 0 );
	// av_dict_set( &options, "listen_timeout", "30", 0 ); // 单位为s
	av_dict_set( &options, "stimeout", "30000000", 0 ); // 单位为 百万分之一秒
	
	///打开输入的流
	ifmt_ctx = avformat_alloc_context();
	int ret = avformat_open_input(&ifmt_ctx, uri.c_str(), nullptr, &options);
	if (ret != 0){
		av_log(nullptr, AV_LOG_ERROR, "Couldn't open input stream ! \n");
		goto end_flag ;
	}
 
	//查找流信息
	if (avformat_find_stream_info(ifmt_ctx, nullptr) < 0){
		av_log(nullptr, AV_LOG_ERROR, "Couldn't find stream information ! \n");
		goto end_flag ;
	}
 
	//找到视频流索引
    video_index = av_find_best_stream(ifmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, nullptr, 0);
 
    st = ifmt_ctx->streams[video_index];
    
    //找到解码器
    codec = avcodec_find_decoder(st->codecpar->codec_id);
    if (!codec){
		av_log(nullptr, AV_LOG_ERROR, "Codec not found ! \n");
        goto end_flag ;
    }
 
    //申请AVCodecContext
    codec_ctx = avcodec_alloc_context3(codec);
    if (!codec_ctx){
        goto end_flag ;
    }
 
	avcodec_parameters_to_context(codec_ctx, ifmt_ctx->streams[video_index]->codecpar);
 
    //打开解码器
    if ((ret = avcodec_open2(codec_ctx, codec, nullptr) < 0)){
        goto end_flag ;
    }
	
    // 计算解码后原始数据所需缓冲区大小，并分配内存空间 Determine required buffer size and allocate buffer
    numBytes = av_image_get_buffer_size(AV_PIX_FMT_RGB24, codec_ctx->width, codec_ctx->height, 1);
    buffer = (uint8_t *)av_malloc(numBytes * sizeof(uint8_t));
    
	pFrameRGB = av_frame_alloc();
    av_image_fill_arrays(pFrameRGB->data, pFrameRGB->linesize, buffer, AV_PIX_FMT_BGR24, codec_ctx->width, codec_ctx->height, 1);

	img_convert_ctx = sws_getContext(codec_ctx->width, codec_ctx->height,codec_ctx->pix_fmt, codec_ctx->width, codec_ctx->height, AV_PIX_FMT_BGR24, 
		   SWS_BICUBIC, nullptr, nullptr, nullptr);
 
	pkt = av_packet_alloc();
	frame = av_frame_alloc();
	while (av_read_frame(ifmt_ctx, pkt) >= 0){
		if (pkt->stream_index == video_index){
			int ret = avcodec_send_packet(codec_ctx, pkt);
			if (ret >= 0){
				ret = avcodec_receive_frame(codec_ctx, frame);
				if ((ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) || ret < 0){
					av_log(nullptr, AV_LOG_ERROR, "Failed to receive frame: %d \n", ret);
					av_packet_unref(pkt);
					continue;
				}

				sws_scale(img_convert_ctx, (const unsigned char* const*)frame->data, frame->linesize, 0, codec_ctx->height, pFrameRGB->data, pFrameRGB->linesize);

				imgInfo = new FFImgInfo();
				imgInfo->pData = buffer;
				imgInfo->height = codec_ctx->height;
				imgInfo->width = codec_ctx->width;
				break;
			}
		}

		av_packet_unref(pkt);
	}

end_flag:
	if (codec_ctx != nullptr){
		avcodec_close(codec_ctx);
		avcodec_free_context(&codec_ctx);
	}
	
	if (ifmt_ctx != nullptr){
		avformat_close_input(&ifmt_ctx);
	}
	
	if (frame != nullptr){
		av_frame_free(&frame);
	}

	if (pFrameRGB != nullptr){
		av_frame_free(&pFrameRGB);
	}

	if (pkt != nullptr){
		av_packet_free(&pkt);
	}

	return imgInfo;
}

void FFNvDecoder::releaseFFImgInfo(FFImgInfo* info){
	if(nullptr != info){
		if(info->pData != nullptr){
			av_free(info->pData);
			info->pData = nullptr;
		}
		delete info;
		info = nullptr;
	}
}