多显卡设置

Hu Chunming
1 parent e41a52bb
Showing 9 changed files with 233 additions and 134 deletions
src/FFNvDecoder.cpp
src/FFNvDecoder.h
src/FFNvDecoderManager.cpp
src/FFNvDecoderManager.h
src/NV12ToRGB.cu
src/NvJpegEncoder.cpp
src/NvJpegEncoder.h
src/cuda_kernels.h
src/main.cpp
@@ -125,7 +125,6 @@ bool FFNvDecoder::init(const string&amp; path)
 		cout << "Failed to create specified HW device.";
 		return false;
 	}
-	avctx->hw_device_ctx = av_buffer_ref(hw_device_ctx);
 	// 打开解码器流
 	if (avcodec_open2(avctx, decoder, nullptr) < 0) {
@@ -136,6 +135,79 @@ bool FFNvDecoder::init(const string&amp; path)
 	return true;
 }
+bool FFNvDecoder::init2(FFDecConfig& cfg)
+{
+	m_cfg = cfg;
+
+	fstream infile(cfg.uri);
+	if (infile.is_open()){
+		m_bReal = false;
+		infile.close();
+	}else {
+		m_bReal = true;
+	}
+
+	
+
+	// 打开输入视频文件
+	AVDictionary *options = nullptr;
+	av_dict_set( &options, "bufsize", "655360", 0 );
+	av_dict_set( &options, "rtsp_transport", cfg.force_tcp ? "tcp" : "udp", 0 );
+	// av_dict_set( &options, "listen_timeout", "30", 0 ); // 单位为s
+	av_dict_set( &options, "stimeout", "3000000", 0 );
+
+	av_register_all();
+	avformat_network_init();
+	
+	fmt_ctx = avformat_alloc_context();
+	const char* input_file = cfg.uri.c_str();
+	if (avformat_open_input(&fmt_ctx, input_file, nullptr, &options) != 0) {
+		cout << "Cannot open input file" << input_file;
+		return false;
+	}
+
+	// 查找流信息
+	if (avformat_find_stream_info(fmt_ctx, nullptr) < 0) {
+		cout << "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) {
+		cout << "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;
+
+	avctx->opaque = this;
+	// 设置解码器管理器的像素格式回调函数
+	avctx->get_format = get_hw_format;
+
+	hw_pix_fmt = AV_PIX_FMT_CUDA;
+
+	// 打开解码器流
+	AVDictionary *op = nullptr;
+	av_dict_set( &op, "gpu", cfg.gpuid.c_str(), 0 );
+	if (avcodec_open2(avctx, vcodec, &op) < 0) {
+		cout << "Failed to open codec for stream" << stream_index;
+		return false;
+	}
+
+	return true;
+}
+
 void FFNvDecoder::start(){
 	m_bRunning = true;
@@ -210,7 +282,7 @@ void FFNvDecoder::decode_thread()
 			result = avcodec_send_packet(avctx, pkt);
 			if (result < 0)
 			{
-				cout << "Failed to send pkt" << result << endl;
+				cout << "Failed to send pkt:" << result << endl;
 				continue;
 			}
@@ -222,16 +294,6 @@ void FFNvDecoder::decode_thread()
 			}
 			mFrameQueue.addTail();
-
-			// if (gpuFrame->format == hw_pix_fmt) 
-			// {
-			// 	/* retrieve data from GPU to CPU */
-			// 	if (av_hwframe_transfer_data(pong_frame, gpuFrame, 0) < 0) 
-			// 	{
-			// 		cout << "Failed to transfer the data to system memory";
-			// 		return;
-			// 	}
-			// }
 		}
 		av_packet_unref(pkt);
 	}
@@ -29,11 +29,20 @@ using namespace std;
 **************************************************/
 typedef void(*POST_DECODE_CALLBACK)(const void * userPtr, AVFrame * gpuFrame);
+struct FFDecConfig
+{
+    string uri;
+    POST_DECODE_CALLBACK post_decoded_cbk;
+    string gpuid;
+    bool force_tcp{true};
+};
+
 class FFNvDecoder{
 public:
     FFNvDecoder();
     ~FFNvDecoder();
     bool init(const string& path);
+    bool init2(FFDecConfig& cfg);
     void close();
     void start();
     void pause();
@@ -55,6 +64,7 @@ private:
 public:
     POST_DECODE_CALLBACK post_decoded_cbk;
     const void * m_userPtr;
+    FFDecConfig m_cfg;
 private:
     AVStream* stream;
@@ -3,8 +3,8 @@
 using namespace std;
-FFNvDecoder* FFNvDecoderManager::createDecoder(const char * uri, POST_DECODE_CALLBACK post_decoded_cbk, string name){
-    int num = decoderMap.count(name);
+FFNvDecoder* FFNvDecoderManager::createDecoder(MgrDecConfig& config){
+    int num = decoderMap.count(config.name);
     if (num > 0)
     {
         cout << "已存在name所标记的解码器" << endl;
@@ -17,12 +17,12 @@ FFNvDecoder* FFNvDecoderManager::createDecoder(const char * uri, POST_DECODE_CAL
         return nullptr;
     }
-    bool bRet= dec->init(uri);
+    bool bRet= dec->init2(config.cfg);
     if (bRet)
     {
-        dec->setName(name) ;
-        dec->post_decoded_cbk = post_decoded_cbk;
-        decoderMap[name] = dec;
+        dec->setName(config.name) ;
+        dec->post_decoded_cbk = config.cfg.post_decoded_cbk;
+        decoderMap[config.name] = dec;
         return dec;
     }
@@ -5,6 +5,12 @@
 using namespace std;
+struct MgrDecConfig
+{
+    FFDecConfig cfg;
+    string name{""};
+};
+
 class FFNvDecoderManager{
 public:
     static FFNvDecoderManager* getInstance(){
@@ -20,7 +26,7 @@ public:
         closeAllDecoder();
     }
-    FFNvDecoder* createDecoder(const char * uri, POST_DECODE_CALLBACK post_decoded_cbk, string name = "");
+    FFNvDecoder* createDecoder(MgrDecConfig& config);
     void setUserPtr(string name, const void * userPtr);
     FFNvDecoder* getDecoderByName(const string name);
@@ -301,26 +301,6 @@ namespace cuda_common
 			fprintf(stderr, "cudaMemcpyToSymbol failed: %s\n", cudaGetErrorString(cudaStatus));
 		}
-		//showHueCSC();
-
-		return cudaStatus;
-	}
-
-	cudaError_t showHueCSC3()
-	{
- /*
-		cudaError_t cudaStatus;
-
-		float data[9];
-		cudaStatus = cudaMemcpyFromSymbol(data, constHueColorSpaceMat2, sizeof(data));
-		printf("hueCSC  ");
-		for (int i = 0; i < 9; i++)
-		{
-			printf("%f  ,", data[i]);
-		}
-		printf("\n");
-*/
-		cudaError_t cudaStatus;
 		return cudaStatus;
 	}
 #include "NvJpegEncoder.h"
-#include <nvjpeg.h>
-
 #include <fstream>
 #include <vector>
 #include <iostream>
+#define CHECK_NVJPEG(S) do {nvjpegStatus_t  status; \
+        status = S; \
+        if (status != NVJPEG_STATUS_SUCCESS ) std::cout << __LINE__ <<" CHECK_NVJPEG - status = " << status << std::endl; \
+        } while (false)
+
 int saveJpeg(const char * filepath, unsigned char* d_srcBGR, int width, int height)
 {
-	nvjpegHandle_t nvjpeg_handle;
-	nvjpegEncoderState_t encoder_state;
-	nvjpegEncoderParams_t encoder_params;
-
-	cudaEvent_t ev_start, ev_end;
-	cudaEventCreate(&ev_start);
-	cudaEventCreate(&ev_end);
-	
-
-	nvjpegImage_t input;
-	nvjpegInputFormat_t input_format = NVJPEG_INPUT_BGRI;
-	int image_width = width;
-	int image_height = height;
-
-	// int channel_size = image_width * image_height;
-	// for (int i = 0; i < 3; i++)
-	// {
-	// 	input.pitch[i] = image_width;
-	// 	(cudaMalloc((void**)&(input.channel[i]), channel_size));
-	// 	(cudaMemset(input.channel[i], 50 * 40 * i, channel_size));
-	// }
+    nvjpegHandle_t nvjpeg_handle;
+    nvjpegEncoderState_t encoder_state;
+    nvjpegEncoderParams_t encoder_params;
+
+    cudaEvent_t ev_start, ev_end;
+    cudaEventCreate(&ev_start);
+    cudaEventCreate(&ev_end);
+
+    nvjpegImage_t input;
+    nvjpegInputFormat_t input_format = NVJPEG_INPUT_BGRI;
+    int image_width = width;
+    int image_height = height;
+
+    // int channel_size = image_width * image_height;
+    // for (int i = 0; i < 3; i++)
+    // {
+    //     input.pitch[i] = image_width;
+    //     (cudaMalloc((void**)&(input.channel[i]), channel_size));
+    //     (cudaMemset(input.channel[i], 50 * 40 * i, channel_size));
+    // }
     input.channel[0] = d_srcBGR;
     input.pitch[0] = image_width * 3;
-	nvjpegBackend_t backend = NVJPEG_BACKEND_DEFAULT;
-
-	nvjpegCreate(backend, nullptr, &nvjpeg_handle);
-	
-	nvjpegEncoderParamsCreate(nvjpeg_handle, &encoder_params, NULL);
-	nvjpegEncoderStateCreate(nvjpeg_handle, &encoder_state, NULL);
-
-	// set params
-	nvjpegEncoderParamsSetEncoding(encoder_params, nvjpegJpegEncoding_t::NVJPEG_ENCODING_PROGRESSIVE_DCT_HUFFMAN, NULL);
-	nvjpegEncoderParamsSetOptimizedHuffman(encoder_params, 1, NULL);
-	nvjpegEncoderParamsSetQuality(encoder_params, 70, NULL);
-	nvjpegEncoderParamsSetSamplingFactors(encoder_params, nvjpegChromaSubsampling_t::NVJPEG_CSS_420, NULL);
-
-	cudaEventRecord(ev_start);
-	nvjpegEncodeImage(nvjpeg_handle, encoder_state, encoder_params, &input, input_format, image_width, image_height, NULL);
-	cudaEventRecord(ev_end);
-
-	std::vector<unsigned char> obuffer;
-	size_t length;
-	nvjpegEncodeRetrieveBitstream(
-		nvjpeg_handle,
-		encoder_state,
-		NULL,
-		&length,
-		NULL);
-
-	obuffer.resize(length);
-	nvjpegEncodeRetrieveBitstream(
-		nvjpeg_handle,
-		encoder_state,
-		obuffer.data(),
-		&length,
-		NULL);
-
-	std::ofstream outputFile(filepath, std::ios::out | std::ios::binary);
-	outputFile.write(reinterpret_cast<const char *>(obuffer.data()), static_cast<int>(length));
-
-	cudaEventSynchronize(ev_end);
-
-	float ms;
-	cudaEventElapsedTime(&ms, ev_start, ev_end);
-	std::cout << "time spend " << ms << " ms" << std::endl;
-	
-	return 0;
+    nvjpegBackend_t backend = NVJPEG_BACKEND_DEFAULT;
+
+    CHECK_NVJPEG(nvjpegCreate(backend, nullptr, &nvjpeg_handle));
+    
+    CHECK_NVJPEG(nvjpegEncoderParamsCreate(nvjpeg_handle, &encoder_params, NULL));
+    CHECK_NVJPEG(nvjpegEncoderStateCreate(nvjpeg_handle, &encoder_state, NULL));
+
+    // set params
+    CHECK_NVJPEG(nvjpegEncoderParamsSetEncoding(encoder_params, nvjpegJpegEncoding_t::NVJPEG_ENCODING_PROGRESSIVE_DCT_HUFFMAN, NULL));
+    CHECK_NVJPEG(nvjpegEncoderParamsSetOptimizedHuffman(encoder_params, 1, NULL));
+    CHECK_NVJPEG(nvjpegEncoderParamsSetQuality(encoder_params, 70, NULL));
+    CHECK_NVJPEG(nvjpegEncoderParamsSetSamplingFactors(encoder_params, nvjpegChromaSubsampling_t::NVJPEG_CSS_420, NULL));
+
+    cudaEventRecord(ev_start);
+    CHECK_NVJPEG(nvjpegEncodeImage(nvjpeg_handle, encoder_state, encoder_params, &input, input_format, image_width, image_height, NULL));
+    cudaEventRecord(ev_end);
+
+    std::vector<unsigned char> obuffer;
+    size_t length;
+    CHECK_NVJPEG(nvjpegEncodeRetrieveBitstream(
+        nvjpeg_handle,
+        encoder_state,
+        NULL,
+        &length,
+        NULL));
+
+    obuffer.resize(length);
+    CHECK_NVJPEG(nvjpegEncodeRetrieveBitstream(
+        nvjpeg_handle,
+        encoder_state,
+        obuffer.data(),
+        &length,
+        NULL));
+
+    cudaEventSynchronize(ev_end);
+
+    // 用完销毁，避免显存泄露
+    nvjpegEncoderParamsDestroy(encoder_params);
+    nvjpegEncoderStateDestroy(encoder_state);
+    nvjpegDestroy(nvjpeg_handle);
+
+    float ms;
+    cudaEventElapsedTime(&ms, ev_start, ev_end);
+    std::cout << "time spend " << ms << " ms" << std::endl;
+
+    std::ofstream outputFile(filepath, std::ios::out | std::ios::binary);
+    outputFile.write(reinterpret_cast<const char *>(obuffer.data()), static_cast<int>(length));
+    outputFile.close();
+    
+    return 0;
 }
 \ No newline at end of file
-
+#include <nvjpeg.h>
 int saveJpeg(const char * filepath, unsigned char* d_srcBGR, int width, int height);
 \ No newline at end of file
@@ -19,7 +19,6 @@ typedef enum
 namespace cuda_common
 {
 	cudaError_t setColorSpace2(e_ColorSpace CSC, float hue);
-	cudaError_t showHueCSC3();
 	cudaError_t NV12ToRGBnot(CUdeviceptr d_srcNV12, size_t nSourcePitch, unsigned char* d_dstRGB, int width, int height);
 	cudaError_t CUDAToBGR(CUdeviceptr dataY, CUdeviceptr dataUV, size_t pitchY, size_t pitchUV, unsigned char* d_dstRGB, int width, int height);
@@ -7,7 +7,6 @@
 unsigned char *pHwRgb = nullptr;
-
 void postDecoded(const void * userPtr, AVFrame * gpuFrame){
     FFNvDecoder* decoder = (FFNvDecoder*)userPtr;
     if (decoder!= nullptr)
@@ -19,42 +18,77 @@ void postDecoded(const void * userPtr, AVFrame * gpuFrame){
         cout << "decode successed ✿✿ヽ(°▽°)ノ✿ " << endl;
         if (gpuFrame->format == AV_PIX_FMT_CUDA)
-        {
-             cudaError_t cudaStatus;
-            if(pHwRgb == nullptr){
-                cudaStatus = cudaMalloc((void **)&pHwRgb, 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, gpuFrame->width, gpuFrame->height);
-            cudaDeviceSynchronize();
-            if (cudaStatus != cudaSuccess) {
-                cout << "CUDAToBGR failed !!!" << endl;
-                return;
-            }
-            saveJpeg("/home/cmhu/FFNvDecoder/a.jpg", pHwRgb, gpuFrame->width, gpuFrame->height);  // 验证 CUDAToRGB 
+        {   
+            cout << "gpuid = " << atoi(decoder->m_cfg.gpuid.c_str()) << endl;
+            // cudaSetDevice(atoi(decoder->m_cfg.gpuid.c_str()));
+            // cudaError_t cudaStatus;
+            // if(pHwRgb == nullptr){
+            //     cuda_common::setColorSpace2( ITU709, 0 );
+            //     cudaStatus = cudaMalloc((void **)&pHwRgb, 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, gpuFrame->width, gpuFrame->height);
+            // cudaDeviceSynchronize();
+            // if (cudaStatus != cudaSuccess) {
+            //     cout << "CUDAToBGR failed !!!" << endl;
+            //     return;
+            // }
+
+            // string path = "/home/cmhu/FFNvDecoder/" + decoder->m_cfg.gpuid + ".jpg";
+            // saveJpeg(path.c_str(), pHwRgb, gpuFrame->width, gpuFrame->height);  // 验证 CUDAToRGB 
         }
     }
 }
 int main(){
-    cuda_common::setColorSpace2( ITU709, 0 );
-
     FFNvDecoderManager* pDecManager = FFNvDecoderManager::getInstance();
-    // FFNvDecoder* decoder = pDecManager->createDecoder("rtmp://192.168.10.56:1935/objecteye/1",postDecoded,"dec1");
-    // if (!decoder)
-    // {
-    //     return 1;
-    // }
-    // pDecManager->setUserPtr("dec1", decoder);
-    // pDecManager->startDecodeByName("dec1");
-
-    FFNvDecoder* dec2 = pDecManager->createDecoder("/home/cmhu/data/duan1.avi",postDecoded,"dec2");
+    MgrDecConfig config;
+    config.name = "dec1";
+    config.cfg.uri = "rtmp://192.168.10.56:1935/objecteye/1";
+    config.cfg.post_decoded_cbk = postDecoded;
+    config.cfg.force_tcp = true;
+    config.cfg.gpuid = "1";
+    FFNvDecoder* decoder = pDecManager->createDecoder(config);
+    if (!decoder)
+    {
+        return 1;
+    }
+    pDecManager->setUserPtr(config.name, decoder);
+    pDecManager->startDecodeByName(config.name);
+
+    config.name = "dec2";
+    config.cfg.uri = "rtmp://192.168.10.56:1935/objecteye/1";
+    config.cfg.gpuid = "2";
+    FFNvDecoder* dec2 = pDecManager->createDecoder(config);
     if (!dec2)
     {
         return 1;
     }
-    pDecManager->setUserPtr("dec2", dec2);
-    pDecManager->startDecodeByName("dec2");
+    pDecManager->setUserPtr(config.name, dec2);
+    pDecManager->startDecodeByName(config.name);
+
+    
+    config.name = "dec0";
+    config.cfg.uri = "rtmp://192.168.10.56:1935/objecteye/1";
+    config.cfg.gpuid = "0";
+    FFNvDecoder* dec0 = pDecManager->createDecoder(config);
+    if (!dec0)
+    {
+        return 1;
+    }
+    pDecManager->setUserPtr(config.name, dec0);
+    pDecManager->startDecodeByName(config.name);
+
+    config.name = "dec01";
+    config.cfg.uri = "rtmp://192.168.10.56:1935/objecteye/1";
+    config.cfg.gpuid = "0";
+    FFNvDecoder* dec01 = pDecManager->createDecoder(config);
+    if (!dec01)
+    {
+        return 1;
+    }
+    pDecManager->setUserPtr(config.name, dec01);
+    pDecManager->startDecodeByName(config.name);
     // while (getchar() != 'q');