优化demo

Hu Chunming
1 parent f40cc409
Showing 4 changed files with 105 additions and 45 deletions
.vscode/settings.json
src/NvJpegEncoder.cpp
src/NvJpegEncoder.h
src/main.cpp
@@ -5,6 +5,10 @@
         "array": "cpp",
         "string": "cpp",
         "string_view": "cpp",
-        "iosfwd": "cpp"
+        "iosfwd": "cpp",
+        "typeinfo": "cpp",
+        "iostream": "cpp",
+        "thread": "cpp",
+        "nvdec.h": "c"
     }
 }
 \ No newline at end of file
@@ -11,7 +11,7 @@
         } while (false)
-int saveJpeg(const char * filepath, unsigned char* d_srcBGR, int width, int height)
+int saveJpeg(const char * filepath, unsigned char* d_srcBGR, int width, int height, cudaStream_t stream)
 {
     nvjpegHandle_t nvjpeg_handle;
     nvjpegEncoderState_t encoder_state;
@@ -41,17 +41,17 @@ int saveJpeg(const char * filepath, unsigned char* d_srcBGR, int width, int heig
     CHECK_NVJPEG(nvjpegCreate(backend, nullptr, &nvjpeg_handle));
-    CHECK_NVJPEG(nvjpegEncoderParamsCreate(nvjpeg_handle, &encoder_params, NULL));
-    CHECK_NVJPEG(nvjpegEncoderStateCreate(nvjpeg_handle, &encoder_state, NULL));
+    CHECK_NVJPEG(nvjpegEncoderParamsCreate(nvjpeg_handle, &encoder_params, stream));
+    CHECK_NVJPEG(nvjpegEncoderStateCreate(nvjpeg_handle, &encoder_state, stream));
     // 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));
+    CHECK_NVJPEG(nvjpegEncoderParamsSetEncoding(encoder_params, nvjpegJpegEncoding_t::NVJPEG_ENCODING_PROGRESSIVE_DCT_HUFFMAN, stream));
+    CHECK_NVJPEG(nvjpegEncoderParamsSetOptimizedHuffman(encoder_params, 1, stream));
+    CHECK_NVJPEG(nvjpegEncoderParamsSetQuality(encoder_params, 70, stream));
+    CHECK_NVJPEG(nvjpegEncoderParamsSetSamplingFactors(encoder_params, nvjpegChromaSubsampling_t::NVJPEG_CSS_420, stream));
     cudaEventRecord(ev_start);
-    CHECK_NVJPEG(nvjpegEncodeImage(nvjpeg_handle, encoder_state, encoder_params, &input, input_format, image_width, image_height, NULL));
+    CHECK_NVJPEG(nvjpegEncodeImage(nvjpeg_handle, encoder_state, encoder_params, &input, input_format, image_width, image_height, stream));
     cudaEventRecord(ev_end);
     std::vector<unsigned char> obuffer;
@@ -61,7 +61,7 @@ int saveJpeg(const char * filepath, unsigned char* d_srcBGR, int width, int heig
         encoder_state,
         NULL,
         &length,
-        NULL));
+        stream));
     obuffer.resize(length);
     CHECK_NVJPEG(nvjpegEncodeRetrieveBitstream(
@@ -69,7 +69,7 @@ int saveJpeg(const char * filepath, unsigned char* d_srcBGR, int width, int heig
         encoder_state,
         obuffer.data(),
         &length,
-        NULL));
+        stream));
     cudaEventSynchronize(ev_end);
@@ -80,7 +80,7 @@ int saveJpeg(const char * filepath, unsigned char* d_srcBGR, int width, int heig
     float ms;
     cudaEventElapsedTime(&ms, ev_start, ev_end);
-    std::cout << "time spend " << ms << " ms" << std::endl;
+    // 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));
 #include <nvjpeg.h>
-int saveJpeg(const char * filepath, unsigned char* d_srcBGR, int width, int height);
 \ No newline at end of file
+int saveJpeg(const char * filepath, unsigned char* d_srcBGR, int width, int height, cudaStream_t stream);
 \ No newline at end of file
@@ -10,7 +10,12 @@
 #include <chrono>
-unsigned char *pHwRgb = nullptr;
+unsigned char *pHwRgb[2] = {nullptr, nullptr};
+
+int sum1 = 0;
+int sum2 = 0;
+
+cudaStream_t stream[2];
 /**
  * 注意： gpuFrame 在解码器设置的显卡上，后续操作要十分注意这一点，尤其是多线程情况
@@ -21,35 +26,63 @@ void postDecoded(const void * userPtr, AVFrame * gpuFrame){
     {
         // cout << "decode name: " << decoder->getName() << endl;
-        if (decoder->getName() == "dec1")
-        {
-            /* code */
-        }
-        
-        // 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()));
-            // 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 
-        }
+            // 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::setColorSpace2( 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 = "/home/cmhu/data/" + decoder->getName() + "/" + to_string(sum) + ".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::setColorSpace2( 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 = "/home/cmhu/data/" + decoder->getName() + "/" + to_string(sum) + ".jpg";
+                    saveJpeg(path.c_str(), pHwRgb[1], gpuFrame->width, gpuFrame->height, stream[1]);  // 验证 CUDAToRGB 
+                }
+            }
     }
 }
@@ -72,6 +105,8 @@ long long get_cur_time(){
 }
 int sum = 0;
+unsigned char *pHwData = nullptr;
+
 void postDecoded0(const void * userPtr, AVFrame * gpuFrame){
     FFNvDecoder* decoder = (FFNvDecoder*)userPtr;
     if (decoder!= nullptr)
@@ -98,12 +133,33 @@ void postDecoded0(const void * userPtr, AVFrame * gpuFrame){
                 count_flag = false;
             }
+
+            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::setColorSpace2( 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 = "/home/cmhu/data/test/" + to_string(sum) + ".jpg";
+                saveJpeg(path.c_str(), pHwData, gpuFrame->width, gpuFrame->height, nullptr);  // 验证 CUDAToRGB 
+            }
         }
     }
 }
 // string test_uri = "rtmp://192.168.10.56:1935/objecteye/1";
-string test_uri = "/home/cmhu/data/test.mp4";
+// string test_uri = "/home/cmhu/data/output_800x480.mp4";
+string test_uri = "/home/cmhu/data/output_1920x1080.mp4";
 void createDecode(int index){
     FFNvDecoderManager* pDecManager = FFNvDecoderManager::getInstance();
@@ -170,7 +226,7 @@ int main(){
     FFNvDecoderManager* pDecManager = FFNvDecoderManager::getInstance();
-    int count = 105;
+    int count = 99;
     for (size_t i = 0; i < count ; i++)
     {
         createDecode(i);