Hu Chunming / FFNvDecoder

  #include "FFNvDecoder.h"

  
  #include <chrono>
  #include <thread>

  #include <fstream>

  #include <chrono>
  

  #include "FFCuContextManager.h"
  

  #include "common_header.h"

  #include "GpuRgbMemory.hpp"
  #include "cuda_kernels.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;
  	}
  

  	LOG_ERROR("Failed to get HW surface format");

  	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;

      m_dec_name = "";

  
  	m_bPause = false;

  	m_bReal = true;

  
  	m_decode_thread = 0;
  	m_post_decode_thread = 0;

  
  	m_bFinished = false;

  	m_dec_keyframe = false;

  	m_fps = 0.0;

  }
  
  FFNvDecoder::~FFNvDecoder()
  {

  	m_dec_keyframe = false;

  }
  

  bool FFNvDecoder::init(FFDecConfig& cfg)

  {
  	m_cfg = cfg;

  	m_dec_name = cfg.dec_name;

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

  	post_decoded_cbk = cfg.post_decoded_cbk;
      decode_finished_cbk = cfg.decode_finished_cbk;
  

  	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) {

  		LOG_ERROR("Cannot open input file:{}",input_file);

  		return false;
  	}
  
  	// 查找流信息
  	if (avformat_find_stream_info(fmt_ctx, nullptr) < 0) {

  		LOG_ERROR("Cannot find input stream information");

  		return false;
  	}
  
  	// 查找视频流信息
  	AVCodec *decoder = nullptr;
  	stream_index = av_find_best_stream(fmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, &decoder, 0);
  	if (stream_index < 0) {

  		LOG_ERROR("Cannot find a video stream in the input file");

  		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;
  

  	m_fps = av_q2d(stream ->avg_frame_rate);
  

  	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", "5", 0 );

  	if (avcodec_open2(avctx, vcodec, &op) < 0) {

  		LOG_ERROR("Failed to open codec for stream");

  		return false;
  	}

  	return true;
  }
  

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

  bool FFNvDecoder::start(){

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

  
  	return true;

  }
  

  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 = UtilTools::get_cur_time_ms();

  	while (m_bRunning)
  	{

  		if (!m_bReal)

  		{

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

  		}

  		
  		int result = av_read_frame(fmt_ctx, pkt);

  		if (result == AVERROR_EOF || result < 0)

  		{

  			LOG_ERROR("Failed to read frame!");

  			break;
  		}
  

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

  		if (stream_index == pkt->stream_index){

  			result = avcodec_send_packet(avctx, pkt);

  			if (result < 0){

  				av_packet_unref(pkt);

  				LOG_ERROR("{} - Failed to send pkt: {}", m_dec_name, result);

  				continue;
  			}
  

  			AVFrame* gpuFrame = av_frame_alloc();

  			result = avcodec_receive_frame(avctx, gpuFrame);

  			if ((result == AVERROR(EAGAIN) || result == AVERROR_EOF) || result < 0){

  				LOG_ERROR("{} - Failed to receive frame: {}", m_dec_name, result);

  				av_frame_free(&gpuFrame); 
  				av_packet_unref(pkt);

  				continue;
  			}

  			av_packet_unref(pkt);

  			if (m_bReal){
  				if (m_bPause){
  					av_frame_free(&gpuFrame); 
  					std::this_thread::sleep_for(std::chrono::milliseconds(3));
  					continue;
  				}
  			}
  

  			if(gpuFrame != nullptr){
  				m_queue_mutex.lock();
  				if(mFrameQueue.size() <= 10){
  					mFrameQueue.push(gpuFrame);
  				}else{
  					av_frame_free(&gpuFrame); 
  				}
  				m_queue_mutex.unlock();
  			}

  		}
  		av_packet_unref(pkt);
  	}
  

  	m_bRunning = false;

  	av_packet_free(&pkt);

  	// long end_time = UtilTools::get_cur_time_ms();

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

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

  	decode_finished_cbk(m_finishedDecArg);

  	decode_finished();
  

  	// 清空队列
  	while(mFrameQueue.size() > 0){
  		AVFrame * gpuFrame = mFrameQueue.front();
  		av_frame_free(&gpuFrame); 
  		mFrameQueue.pop();
  	}
  

  	LOG_INFO("{} - decode thread exited.", m_dec_name);

  }
  

  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(){
  	int skip_frame = m_cfg.skip_frame;
  	if (skip_frame <= 0){
  		skip_frame = 1;
  	}
  	
  	int index = 0;

  	while (m_bRunning)

  	{

  		if(mFrameQueue.size() > 0){
  			std::lock_guard<std::mutex> l(m_snapshot_mutex);
  			// 取队头数据
  			m_queue_mutex.lock();
  			AVFrame * gpuFrame = mFrameQueue.front();
  			mFrameQueue.pop();
  			m_queue_mutex.unlock();
  			// 跳帧
  			if (skip_frame == 1 || index % skip_frame == 0){

  				post_decoded_cbk(m_postDecArg, convert2bgr(gpuFrame));

  				index = 0;
  			}

  			av_frame_free(&gpuFrame); 

  			index++;
  		}

  	}
  
  	LOG_INFO("post decode thread exited.");

  }
  

  void FFNvDecoder::close(){

  	m_bRunning=false;

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

  	m_dec_keyframe = false;

  }
  

  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 (avctx != nullptr)

  	{

  		width = avctx->width;
  		height = avctx->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(){

  	m_queue_mutex.lock();
  	int queue_size = mFrameQueue.size(); 
  	m_queue_mutex.lock();
  	return queue_size;

  }
  

  float FFNvDecoder::fps(){
  	return m_fps;

  }

  void FFNvDecoder::setPostDecArg(const void* postDecArg){
  	m_postDecArg = postDecArg;
  }
  
  void FFNvDecoder::setFinishedDecArg(const void* finishedDecArg){
  	m_finishedDecArg = finishedDecArg;
  }
  
  DeviceRgbMemory* FFNvDecoder::convert2bgr(AVFrame * gpuFrame){
  	if (gpuFrame != nullptr && gpuFrame->format == AV_PIX_FMT_CUDA ){
  		LOG_DEBUG("decode task: gpuid: {}  width: {} height: {}", m_cfg.gpuid, gpuFrame->width, gpuFrame->height);
  		GpuRgbMemory* gpuMem = new GpuRgbMemory(3, gpuFrame->width, gpuFrame->height, getName(), m_cfg.gpuid, false, true);
  
  		do{
  			if (gpuMem->getMem() == nullptr){
  				LOG_ERROR("new GpuRgbMemory failed !!!");
  				break;
  			}
  			
  			cudaSetDevice(atoi(m_cfg.gpuid.c_str()));
  			cuda_common::setColorSpace( ITU_709, 0 );
  			cudaError_t cudaStatus = cuda_common::CUDAToBGR((CUdeviceptr)gpuFrame->data[0],(CUdeviceptr)gpuFrame->data[1], gpuFrame->linesize[0], gpuFrame->linesize[1], gpuMem->getMem(), gpuFrame->width, gpuFrame->height);
  			cudaDeviceSynchronize();
  			if (cudaStatus != cudaSuccess) {
  				LOG_ERROR("CUDAToBGR failed failed !!!");
  				break;
  			}
  
  			return gpuMem;
  		}while(0);
  
  		delete gpuMem;
  		gpuMem = nullptr;
  	}
  
  	return nullptr;
  }
  

  FFImgInfo* FFNvDecoder::snapshot(){
  
  	// 锁住停止队列消耗
  	std::lock_guard<std::mutex> l(m_snapshot_mutex);
  
  	AVFrame * gpuFrame = nullptr;
  
  	bool bFirst = true;
  	while(true){
  		m_queue_mutex.lock();
  		if(mFrameQueue.size() <= 0){
  			m_queue_mutex.unlock();
  			if(bFirst){
  				std::this_thread::sleep_for(std::chrono::milliseconds(100));
  				bFirst = false;
  				continue;
  			}else{
  				// 再进来说明前面已经等了 100 ms
  				// 100 ms都没有等到解码数据，则退出
  				return nullptr;
  			}
  		}
  
  		// 队列中数据大于1 
  		gpuFrame = mFrameQueue.front();
  		m_queue_mutex.unlock();
  		break;
  	}
  
  	if (gpuFrame != nullptr && gpuFrame->format == AV_PIX_FMT_CUDA ){
  		LOG_DEBUG("decode task: gpuid: {}  width: {} height: {}", m_cfg.gpuid, gpuFrame->width, gpuFrame->height);

  		GpuRgbMemory* gpuMem = new GpuRgbMemory(3, gpuFrame->width, gpuFrame->height, getName(), m_cfg.gpuid, false, true);

  
  		if (gpuMem->getMem() == nullptr){
  			LOG_ERROR("new GpuRgbMemory failed !!!");
  			return nullptr;
  		}
  		
  		cudaSetDevice(atoi(m_cfg.gpuid.c_str()));
  		cuda_common::setColorSpace( ITU_709, 0 );
  		cudaError_t cudaStatus = cuda_common::CUDAToBGR((CUdeviceptr)gpuFrame->data[0],(CUdeviceptr)gpuFrame->data[1], gpuFrame->linesize[0], gpuFrame->linesize[1], gpuMem->getMem(), gpuFrame->width, gpuFrame->height);
  		cudaDeviceSynchronize();
  		if (cudaStatus != cudaSuccess) {
  			LOG_ERROR("CUDAToBGR failed failed !!!");
  			return nullptr;
  		}
  
  		unsigned char * pHwRgb = gpuMem->getMem();
  		int channel = gpuMem->getChannel();
  		int width = gpuMem->getWidth();
  		int height = gpuMem->getHeight();
  
  		if (pHwRgb != nullptr && channel > 0 && width > 0 && height > 0){
  			int nSize = channel * height * width;
  
  			LOG_INFO("channel:{} height:{} width:{}", channel, height, width);
  			// unsigned char* cpu_data = new unsigned char[nSize];
  
              unsigned char* cpu_data = (unsigned char *)av_malloc(nSize * sizeof(unsigned char));
  
  			cudaMemcpy(cpu_data, pHwRgb, nSize * sizeof(unsigned char), cudaMemcpyDeviceToHost);
  			cudaDeviceSynchronize();
  
  			delete gpuMem;
  			gpuMem = nullptr;
  
  			FFImgInfo* imgInfo = new FFImgInfo();
  			imgInfo->dec_name = m_dec_name;
  			imgInfo->pData = cpu_data;
  			imgInfo->height = height;
  			imgInfo->width = width;
  			imgInfo->timestamp = UtilTools::get_cur_time_ms();
  			imgInfo->index = m_index;
  
  			m_index++;
  
  			return imgInfo;
  		}
  
  		delete gpuMem;
  		gpuMem = nullptr;
  	}
  
  	return nullptr;
  }