check_tool.cpp 5.15 KB
#include "check_tool.h"

#include "FFCuContextManager.h"

extern "C"
{
#include <libavcodec/avcodec.h> 
#include <libavdevice/avdevice.h> 
#include <libavformat/avformat.h> 
#include <libavfilter/avfilter.h> 
#include <libavutil/avutil.h> 
#include <libavutil/pixdesc.h> 
#include <libswscale/swscale.h>
#include <libavutil/imgutils.h>
}

#include <chrono>

static AVPixelFormat get_hw_format(AVCodecContext* avctx, const AVPixelFormat* pix_fmts) {

	const AVPixelFormat* p;

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

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

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());

	return tpMicro.time_since_epoch().count();
}

void evalQuality(const string& uri, const char* gpuid) {
	AVFormatContext* ifmt_ctx = nullptr;
	AVCodecContext* codec_ctx = nullptr;
	const AVCodec* decoder = nullptr;
	const AVCodec* codec = nullptr;
	AVPacket* pkt = nullptr;
	int video_index = -1;
	AVStream* st = nullptr;
	SwsContext* img_convert_ctx = nullptr;
	uint8_t* buffer = nullptr;
	int numBytes = 0;

	long long start_time = 0;
	long long end_time = 0;
	long long s_time = 0;
	long long e_time = 0;
	long long sum = 0;
	double avg_time = 0.0;

	int sep = 0;

	string cuvid_dec_name = "";
	FFCuContextManager* pCtxMgr = FFCuContextManager::getInstance();
	AVDictionary* op = nullptr;
	AVBufferRef* hw_device_ctx = 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, &decoder, 0);
	if (video_index < 0) {
		av_log(nullptr, AV_LOG_ERROR, "Cannot find a video stream in the input file ! \n");
		goto end_flag;
	}

	cuvid_dec_name = string(decoder->name) + "_cuvid";
	codec = avcodec_find_decoder_by_name(cuvid_dec_name.c_str());
	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;
	}

	st = ifmt_ctx->streams[video_index];
	if (avcodec_parameters_to_context(codec_ctx, st->codecpar) < 0) {
		goto end_flag;
	}

	{
		av_log(nullptr, AV_LOG_INFO, "path: %s \n", ifmt_ctx->url);
		av_log(nullptr, AV_LOG_INFO, "format: %s \n", ifmt_ctx->iformat->name);
		const char* profile = avcodec_profile_name(codec_ctx->codec_id, codec_ctx->profile);
		av_log(nullptr, AV_LOG_INFO, "codec: %s(%s) \n", decoder->name, profile);
		const char* pix_fmt_name = av_get_pix_fmt_name(codec_ctx->pix_fmt);
		av_log(nullptr, AV_LOG_INFO, "pix_fmt_name: %s \n", pix_fmt_name);
		av_log(nullptr, AV_LOG_INFO, "width x height: %dX%d \n", st->codecpar->width, st->codecpar->height);

		av_log(NULL, AV_LOG_INFO, "frame_rate: %1.0f ", av_q2d(st->r_frame_rate));
		av_log(NULL, AV_LOG_INFO, "  -->  reference PPS: %f ms\n", 1 / av_q2d(st->r_frame_rate) * 1000);

	}

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

	hw_device_ctx = pCtxMgr->getCuCtx(gpuid);
	if (nullptr == hw_device_ctx) {
		av_log(nullptr, AV_LOG_ERROR, "create CUDA context failed ! \n");
		goto end_flag;
	}
	codec_ctx->hw_device_ctx = av_buffer_ref(hw_device_ctx);
	if (nullptr == codec_ctx->hw_device_ctx) {
		goto end_flag;
	}

	// 打开解码器流
	av_dict_set(&op, "gpu", gpuid, 0);
	if (avcodec_open2(codec_ctx, codec, &op) < 0) {
		av_log(nullptr, AV_LOG_ERROR, "Failed to open codec for stream ! \n");
		goto end_flag;
	}

	s_time = get_cur_time();
	start_time = get_cur_time();
	pkt = av_packet_alloc();
	while (av_read_frame(ifmt_ctx, pkt) >= 0) {
		end_time = get_cur_time();
		sum++;
		sep++;
		if (end_time - start_time > 1000) {
			avg_time = double(end_time - start_time) / sep;
			start_time = get_cur_time();
			sep = 0;
			av_log(nullptr, AV_LOG_INFO, "PPS: %f ms\n", avg_time);
		}
		av_packet_unref(pkt);
	}

	e_time = get_cur_time();
	avg_time = double(e_time - s_time) / sum;
	av_log(nullptr, AV_LOG_INFO, "TOOTAL PPS: %f ms\n", avg_time);

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 (pkt != nullptr) {
		av_packet_free(&pkt);
	}
}