Hu Chunming / FFNvDecoder

#include "FFNvDecoderManager.h"
#include <iostream>

#include "cuda_kernels.h"

#include "NvJpegEncoder.h"

#include <pthread.h>
#include <thread>

#include <chrono>



#include <unistd.h>


#ifdef _WIN32
#include "Winsock2.h"
#pragma comment(lib, "ws2_32.lib")
#endif

#ifdef __linux__
#include "arpa/inet.h"
#endif


#define MIN_RTP_PORT		10000
#define MAX_RTP_PORT		60000

// ȡ MIN_RTP_PORT(10000)~MAX_RTP_PORT(60000)֮�������˿�(ż���������������˿ڿ���)
int allocRtpPort() {

	static int s_rtpPort = MIN_RTP_PORT;
	if (MIN_RTP_PORT == s_rtpPort)
	{
		srand((unsigned int)time(NULL));
		s_rtpPort = MIN_RTP_PORT + (rand() % MIN_RTP_PORT);
	}

	if (s_rtpPort % 2)
		++s_rtpPort;

	while (true)
	{
		s_rtpPort += 2;
		s_rtpPort = s_rtpPort >= MAX_RTP_PORT ? MIN_RTP_PORT : s_rtpPort;

		int i = 0;
		for (; i < 2; i++)
		{
			sockaddr_in sRecvAddr;
			int s = socket(AF_INET, SOCK_DGRAM, 0);

			sRecvAddr.sin_family = AF_INET;        
			sRecvAddr.sin_addr.s_addr = htonl(INADDR_ANY);    
			sRecvAddr.sin_port = htons(s_rtpPort + i); 

			int nResult = bind(s, (sockaddr *)&sRecvAddr, sizeof(sRecvAddr));
			if (nResult != 0)
			{
				break;
			}

			nResult = close(s);
			if (nResult != 0)
			{
				printf("closesocket failed:%d\n", nResult);
				break;
			}
		}

		if (i == 2)
			break;
	}

	return s_rtpPort;
}





unsigned char *pHwRgb[2] = {nullptr, nullptr};

int sum1 = 0;
int sum2 = 0;

cudaStream_t stream[2];

/**
 * 注意： gpuFrame 在解码器设置的显卡上，后续操作要十分注意这一点，尤其是多线程情况
 * */
void postDecoded(const void * userPtr, AVFrame * gpuFrame){
    AbstractDecoder* decoder = (AbstractDecoder*)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;

            int sum = sum1;
            if (decoder->getName() == "dec1")
            {
                sum1 ++ ;
                sum = sum1;

                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()));
                    cudaError_t cudaStatus;
                    if(pHwRgb[0] == nullptr){
                        // cudaStreamCreate(&stream[0]);
                        cuda_common::setColorSpace( ITU709, 0 );
                        cudaStatus = cudaMalloc((void **)&pHwRgb[0], 3 * gpuFrame->width * gpuFrame->height * sizeof(unsigned char));
                    }
                    cudaStatus = cuda_common::CUDAToBGR((CUdeviceptr)gpuFrame->data[0],(CUdeviceptr)gpuFrame->data[1], gpuFrame->linesize[0], gpuFrame->linesize[1], pHwRgb[0], gpuFrame->width, gpuFrame->height);
                    cudaDeviceSynchronize();
                    if (cudaStatus != cudaSuccess) {
                        cout << "CUDAToBGR failed !!!" << endl;
                        return;
                    }

                    string path = "/mnt/f/fiss/data/" + decoder->getName() + ".jpg";
                    saveJpeg(path.c_str(), pHwRgb[0], gpuFrame->width, gpuFrame->height, stream[0]);  // 验证 CUDAToRGB 
                }
            } else if (decoder->getName() == "dec2") 
            {
                sum2 ++ ;
                sum = sum2;

                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()));
                    cudaError_t cudaStatus;
                    if(pHwRgb[1] == nullptr){
                        // cudaStreamCreate(&stream[1]);
                        cuda_common::setColorSpace( ITU709, 0 );
                        cudaStatus = cudaMalloc((void **)&pHwRgb[1], 3 * gpuFrame->width * gpuFrame->height * sizeof(unsigned char));
                    }
                    cudaStatus = cuda_common::CUDAToBGR((CUdeviceptr)gpuFrame->data[0],(CUdeviceptr)gpuFrame->data[1], gpuFrame->linesize[0], gpuFrame->linesize[1], pHwRgb[1], gpuFrame->width, gpuFrame->height);
                    cudaDeviceSynchronize();
                    if (cudaStatus != cudaSuccess) {
                        cout << "CUDAToBGR failed !!!" << endl;
                        return;
                    }

                    string path = "/mnt/f/fiss/data/" + decoder->getName() + ".jpg";
                    saveJpeg(path.c_str(), pHwRgb[1], gpuFrame->width, gpuFrame->height, stream[1]);  // 验证 CUDAToRGB 
                }
            }
    }
}

long start_time = 0;
long end_time = 0;
bool count_flag = false;
int count = 0;
int count_std = 100;

static long long get_cur_time(){
    // 获取操作系统当前时间点（精确到ms）
    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 sum = 0;
unsigned char *pHwData = nullptr;

void postDecoded0(const void * userPtr, AVFrame * gpuFrame){
    // std::this_thread::sleep_for(std::chrono::milliseconds(30000));

    AbstractDecoder* decoder = (AbstractDecoder*)userPtr;
    if (decoder!= nullptr)
    {
        // cout << "decode name: " << decoder->getName() << endl;
        if (decoder->getName() == "dec")
        {
            if (! count_flag)
            {
                count_flag = true;
                count = 0;
                end_time = start_time = get_cur_time();
            }
            count++;
            sum ++ ;
            if (count >= 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;
                // cout << gpuFrame->pts << endl;

                count_flag = false;
            }
            // cout << "帧数：" << sum << endl;

            if (gpuFrame->format == AV_PIX_FMT_CUDA)
            {   
                cudaSetDevice(atoi(decoder->m_cfg.gpuid.c_str()));
                // cout << "gpu id : " << decoder->m_cfg.gpuid.c_str() << endl;
                cudaError_t cudaStatus;
                if(pHwData == nullptr){
                    cuda_common::setColorSpace( ITU709, 0 );
                    cudaStatus = cudaMalloc((void **)&pHwData, 3 * gpuFrame->width * gpuFrame->height * sizeof(unsigned char));
                }
                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 = "/mnt/f/fiss/data/" + decoder->getName() + ".jpg";
                saveJpeg(path.c_str(), pHwData, gpuFrame->width, gpuFrame->height, nullptr);  // 验证 CUDAToRGB 
            }
        }
    }
}

void decode_finished_cbk(const void* userPtr){
    cout << "当前时间戳: " << 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 = "ws://127.0.0.1:10000/sms/34020000002020000001/flv/hls/34020000001110005555_34020000001310005554.flv";
// string test_uri = "rtsp://176.10.0.4:8554/stream";

char* gpu_id = "0";

void createDecode(int index, char* gpuid){
    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;
    config.dec_type = DECODER_TYPE_FFMPEG;

    config.cfg.gpuid = gpuid;
    // if (index % 2 == 0)
    // {
    //     config.cfg.gpuid = "0";
    // }
    // else
    // {
    //     config.cfg.gpuid = "0";
    // }
    
    AbstractDecoder* decoder = pDecManager->createDecoder(config);
    if (!decoder)
    {
        return ;
    }
    pDecManager->setPostDecArg(config.name, decoder);
    pDecManager->setFinishedDecArg(config.name, decoder);
    pDecManager->startDecodeByName(config.name);
}

void createGB28181Decode(int index, char* gpuid){
    FFNvDecoderManager* pDecManager = FFNvDecoderManager::getInstance();
    MgrDecConfig config;
    config.name = "dec" + to_string(index);
    config.cfg.uri = config.name;
    config.cfg.post_decoded_cbk = postDecoded;
    config.cfg.decode_finished_cbk = decode_finished_cbk;
    config.cfg.force_tcp = true;

    config.dec_type = DECODER_TYPE_GB28181;
    config.cfg.port = 30012;//allocRtpPort();

    config.cfg.gpuid = gpuid;
    
    AbstractDecoder* decoder = pDecManager->createDecoder(config);
    if (!decoder)
    {
        return ;
    }
    pDecManager->setPostDecArg(config.name, decoder);
    pDecManager->setFinishedDecArg(config.name, decoder);
    pDecManager->startDecodeByName(config.name);
}

#define checkCudaErrors(S) do {CUresult  status; \
        status = S; \
        if (status != CUDA_SUCCESS ) std::cout << __LINE__ <<" checkCudaErrors - status = " << status << std::endl; \
        } while (false)

int CheckCUDAProperty( int devId )
{
    cuInit(0);

	CUdevice dev = devId;
	size_t memSize = 0;
	char devName[256] = {0};
	int major = 0, minor = 0;
	CUresult rlt = CUDA_SUCCESS;

    rlt = cuDeviceGetAttribute(&major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, dev);
    checkCudaErrors( rlt );

    rlt = cuDeviceGetAttribute(&minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, dev);
	checkCudaErrors( rlt );

	rlt = cuDeviceGetName( devName, sizeof( devName ), dev );
	checkCudaErrors( rlt );

	printf( "Using GPU Device %d: %s has SM %d.%d compute capability\n",
		    dev, devName, major, minor );

	rlt = cuDeviceTotalMem( &memSize, dev );
	checkCudaErrors( rlt );

	printf( "Total amount of global memory:   %4.4f MB\n",
		   (float)memSize / ( 1024 * 1024 ) );

	return 0;
}

void logFF(void *, int level, const char *fmt, va_list ap)
{
    vfprintf(stdout, fmt, ap);
}


int main(int argc, char* argv[]){

    // av_log_set_callback(&logFF);

    CheckCUDAProperty(0);

    test_uri = argv[1];
    gpu_id = argv[2];
    cout << test_uri << "   gpu_id:" << gpu_id << endl;



    pthread_t m_decode_thread;
    pthread_create(&m_decode_thread,0,
        [](void* arg)
        {
            while (true)
            {
                std::this_thread::sleep_for(std::chrono::milliseconds(5000));
                FFNvDecoderManager* pDecManager = FFNvDecoderManager::getInstance();
                int count = pDecManager->count();
                cout << "当前时间：" << get_cur_time() << "  当前运行路数： " << pDecManager->count() << endl;
            }  

            return (void*)0;
        }
    ,nullptr);

    
    

    FFNvDecoderManager* pDecManager = FFNvDecoderManager::getInstance();
    int i = 0;

    while (true)
    {
        int ch = getchar();
        if (ch == 'q')
        {
            break;
        }
        

        switch (ch)
        {
        case 'f':
        case 'F':
            createDecode(i, gpu_id);
            i++;
            break;
        case 'g':
        case 'G':
            createGB28181Decode(i, gpu_id);
            i++;
            break;
        case 'r':
        case 'R':
            pDecManager->resumeDecoder("dec0");
            break;
        case 'p':
        case 'P':
            pDecManager->pauseDecoder("dec0");
            break;

        case 'c':
        case 'C':
            i--;
            pDecManager->closeDecoderByName("dec" + to_string(i));
            break;

        case 'i':
        case 'I':
        {
            int w,h;
            pDecManager->getResolution("dec0", w,h);
            printf( "%s : %dx%d\n", "dec0" , w,h );
        }
            break;

        default:
            break;
        }

        /* code */
    }

    cout << "总共帧数：" << sum << endl;
    pDecManager->closeAllDecoder();
}