FFNvDecoder/main.cpp

#include "FFNvDecoderManager.h"
#include <iostream>
#include <thread>
#include <chrono>
#include "check_tool.h"
#include "cuda_kernels.h"
#include "NvJpegEncoder.h"
#include "opencv2\opencv.hpp"
using namespace std;
using namespace cv;
unsigned char *pHwRgb[2] = {nullptr, nullptr};
int sum1 = 0;
int sum2 = 0;
mutex m_mutex;
void saveFrame(AVFrame * gpuFrame, string file_name) {
	std::lock_guard<std::mutex> l(m_mutex);


	unsigned char *pHwData = nullptr;
	cudaError_t cudaStatus = cudaMalloc((void **)&pHwData, 3 * gpuFrame->width * gpuFrame->height * sizeof(unsigned char));

	cuda_common::setColorSpace(ITU709, 0);
	cudaStatus = cuda_common::CUDAToBGR((CUdeviceptr)gpuFrame->data[0], (CUdeviceptr)gpuFrame->data[1], gpuFrame->linesize[0], gpuFrame->linesize[1], pHwData, gpuFrame->width, gpuFrame->height);
	cudaDeviceSynchronize();
	if (cudaStatus != cudaSuccess) {
		cout << "CUDAToBGR failed !!!" << endl;
		return;
	}

	string path = file_name + ".jpg";
	saveJpeg(path.c_str(), pHwData, gpuFrame->width, gpuFrame->height, nullptr);  // 验证 CUDAToRGB 

	cudaDeviceSynchronize();
	cudaFree(pHwData);
	pHwData = nullptr;
}
mutex m_mutex_show;
void showFrame(AVFrame * gpuFrame) {
	std::lock_guard<std::mutex> l(m_mutex_show);
	unsigned char *pHwData = nullptr;
	cudaError_t cudaStatus = cudaMalloc((void **)&pHwData, 3 * gpuFrame->width * gpuFrame->height * sizeof(unsigned char));

	cuda_common::setColorSpace(ITU709, 0);
	cudaStatus = cuda_common::CUDAToBGR((CUdeviceptr)gpuFrame->data[0], (CUdeviceptr)gpuFrame->data[1], gpuFrame->linesize[0], gpuFrame->linesize[1], pHwData, gpuFrame->width, gpuFrame->height);
	cudaDeviceSynchronize();
	if (cudaStatus != cudaSuccess) {
		cout << "CUDAToBGR failed !!!" << endl;
		return;
	}


	unsigned char * pHwRgb = pHwData;
	int channel = 3;
	int width = gpuFrame->width;
	int height = gpuFrame->height;

	if (pHwRgb != nullptr && channel > 0 && width > 0 && height > 0) {
		int nSize = channel * height * width;
		unsigned char* cpu_data = new unsigned char[nSize];

		cudaMemcpy(cpu_data, pHwRgb, nSize * sizeof(unsigned char), cudaMemcpyDeviceToHost);
		cudaDeviceSynchronize();

		cv::Mat img_(height, width, CV_8UC3, cpu_data);
		bool bWrite = cv::imwrite("dec0.jpg", img_);

		imshow("show", img_);
		waitKey(0);

		delete[] cpu_data;
		cpu_data = nullptr;

	}

	cudaFree(pHwData);
	pHwData = nullptr;
}
/**
 * 注意： gpuFrame 在解码器设置的显卡上，后续操作要十分注意这一点，尤其是多线程情况
 * */
void postDecoded(const void * userPtr, AVFrame * gpuFrame){
    FFNvDecoder* decoder = (FFNvDecoder*)userPtr;
    if (decoder!= nullptr)
    {
        // cout << "decode name: " << decoder->getName() << endl;
        // const char* gpu_pixfmt = av_get_pix_fmt_name((AVPixelFormat)gpuFrame->format);
        // cout << "pixfmt: " << gpu_pixfmt << endl;
        // cout << "keyframe: " << gpuFrame->key_frame << " width: " << gpuFrame->width << " height: "<< gpuFrame->height << endl;
        // cout << "decode successed ✿✿ヽ(°▽°)ノ✿ " << endl;
		if (gpuFrame->format == AV_PIX_FMT_CUDA)
		{
			// cout << "gpuid = " << atoi(decoder->m_cfg.gpuid.c_str()) << endl;
			cudaSetDevice(atoi(decoder->m_cfg.gpuid.c_str()));

			saveFrame(gpuFrame, decoder->getName());
			showFrame(gpuFrame);
		}
    }
}
long long start_time = 0;
long long end_time = 0;
bool count_flag = false;
int count_num = 0;
int count_std = 100;
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();
}
static int suming = 0;
unsigned char *pHwData = nullptr;
void postDecoded0(const void * userPtr, AVFrame * gpuFrame){
    // std::this_thread::sleep_for(std::chrono::milliseconds(30000));
    FFNvDecoder* decoder = (FFNvDecoder*)userPtr;
    if (decoder!= nullptr)
    {
        // cout << "decode name: " << decoder->getName() << endl;
        if (decoder->getName() == "dec")
        {
            if (! count_flag)
            {
                count_flag = true;
                count_num = 0;
                end_time = start_time = get_cur_time();
            }
            count_num++;
            suming ++ ;
            if (count_num >= count_std)
            {
                // end_time = get_cur_time();
                // long time_using = end_time - start_time;
                // double time_per_frame = double(time_using)/count_std ;
                // cout << count_std << "帧用时:" << time_using << "ms 每帧用时：" << time_per_frame << "ms" << endl;
                cout << "keyframe: " << gpuFrame->key_frame << " width: " << gpuFrame->width << " height: "<< gpuFrame->height << endl;
                count_flag = false;
            }
            cout << "帧数：" << suming << endl;
        }
    }
}
void decode_finished_cbk(const void* userPtr){
    cout << "decode_finish timestamp: " << get_cur_time() << endl;
}
// string test_uri = "rtmp://192.168.10.56:1935/objecteye/1";
// string test_uri = "/home/cmhu/data/output_800x480.mp4";
// string test_uri = "/home/cmhu/data/output_1920x1080.mp4";
// string test_uri = "rtsp://176.10.0.2:8554/stream";
// string test_uri = "/mnt/f/fiss/test_data/h265.mp4";
string test_uri = "rtsp://176.10.0.4:8554/stream";
void createDecode(int index){
    FFNvDecoderManager* pDecManager = FFNvDecoderManager::getInstance();
    MgrDecConfig config;
    config.name = "dec" + to_string(index);
    config.cfg.uri = test_uri;
    config.cfg.post_decoded_cbk = postDecoded;
    config.cfg.decode_finished_cbk = decode_finished_cbk;
    config.cfg.force_tcp = true;
    if (index % 2 == 0)
    {
        config.cfg.gpuid = "0";
    }
    else
    {
        config.cfg.gpuid = "0";
    }
    
    FFNvDecoder* decoder = pDecManager->createDecoder(config);
    if (!decoder)
    {
        return ;
    }
    pDecManager->setUserPtr(config.name, decoder);
    pDecManager->startDecodeByName(config.name);
}
void logFF(void *, int level, const char *fmt, va_list ap)
{
    vfprintf(stdout, fmt, ap);
}
int main(int argc, char* argv[]) {
    printf("start \n");
    if (argc != 3) {
        fprintf(stderr, "./xxx uri gpu_id\n");
        return -1;
    }
    char* uri = argv[1];
    char* gpuid = argv[2];
    cout << av_version_info() << endl;
    //evalQuality(uri, gpuid);
    // av_log_set_callback(&logFF);
    FFNvDecoderManager* pDecManager = FFNvDecoderManager::getInstance();
    // int count = 99;
    // for (size_t i = 0; i < count ; i++)
    // {
    //     createDecode(i);
    // }
    MgrDecConfig config;
    config.name = "dec";
    config.cfg.uri = uri;
    config.cfg.post_decoded_cbk = postDecoded;
    config.cfg.decode_finished_cbk = decode_finished_cbk;
    config.cfg.force_tcp = true;
    config.cfg.gpuid = "0";
    FFNvDecoder* dec2 = pDecManager->createDecoder(config);
    if (!dec2)
    {
        return 1;
    }
    pDecManager->setUserPtr(config.name, dec2);
    // pDecManager->setDecKeyframe(config.name, true);
    pDecManager->startDecodeByName(config.name);
    int w,h;
    pDecManager->getResolution(config.name, w,h);
    printf( "%s : %dx%d\n", config.name.c_str() , w,h );
    thread* m_thread = new thread([](void* arg)
        {
            while (true)
            {
                std::this_thread::sleep_for(std::chrono::milliseconds(5000));
                FFNvDecoderManager* pDecManager = FFNvDecoderManager::getInstance();
                int count = pDecManager->count();
                cout << "当前运行路数： " << pDecManager->count() << endl;
            }
            return (void*)0;
        }
    , nullptr);
    while (getchar() != 'q');
    cout << "总共帧数：" << sum << endl;
    pDecManager->closeAllDecoder();
}