Hu Chunming / FFNvDecoder

  #include "FFNvDecoder.h"

  
  #include <chrono>
  #include <thread>

  #include <fstream>

  #include <chrono>
  

  #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(NULL, 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(NULL, AV_LOG_ERROR, "Cannot open input file: %s \n", input_file);

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

  		av_log(NULL, 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(NULL, 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();
  	avctx->hw_device_ctx = av_buffer_ref(pCtxMgr->getCuCtx(gpuid));
  	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(NULL, 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::microseconds> tpMicro
          = chrono::time_point_cast<chrono::microseconds>(chrono::system_clock::now());
      // (微秒精度的)时间点 => (微秒精度的)时间戳
      time_t totalMicroSeconds = tpMicro.time_since_epoch().count();
  
  	long long currentTime = ((long long)totalMicroSeconds)/1000;
  
      return 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(NULL, 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(NULL, 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(NULL, 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(NULL, 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(NULL, 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 (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(){
  	return mFrameQueue.length();
  }
  
  FFImgInfo* FFNvDecoder::snapshot(const string& uri){
   
      AVFormatContext* ifmt_ctx = NULL;
  	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 = NULL;
      int numBytes = 0;
  
  	FFImgInfo* imgInfo = nullptr;
   
   	//av_register_all();
  	avformat_network_init();
  
  	///打开输入的流
  	int ret = avformat_open_input(&ifmt_ctx, uri.c_str(), NULL, NULL);
  	if (ret != 0){
  		av_log(NULL, AV_LOG_ERROR, "Couldn't open input stream ! \n");
  		goto end_flag ;
  	}
   
  	//查找流信息
  	if (avformat_find_stream_info(ifmt_ctx, NULL) < 0){
  		av_log(NULL, 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, NULL, 0);
   
      st = ifmt_ctx->streams[video_index];
      
      //找到解码器
      codec = avcodec_find_decoder(st->codecpar->codec_id);
      if (!codec){
  		av_log(NULL, 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, NULL) < 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, NULL, NULL, NULL);
   
  	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(NULL, AV_LOG_ERROR, "Failed to receive frame: %d \n", ret);
  					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;
  			}
  		}
  	}
  
  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;
  	}

  }