Hu Chunming / FFNvDecoder

  #include "NvJpegEncoder.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, cudaStream_t stream)

  {

      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;
  
      CHECK_NVJPEG(nvjpegCreate(backend, nullptr, &nvjpeg_handle));
      

      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, 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, stream));

      cudaEventRecord(ev_end);
  
      std::vector<unsigned char> obuffer;
      size_t length;
      CHECK_NVJPEG(nvjpegEncodeRetrieveBitstream(
          nvjpeg_handle,
          encoder_state,
          NULL,
          &length,

          stream));

  
      obuffer.resize(length);
      CHECK_NVJPEG(nvjpegEncodeRetrieveBitstream(
          nvjpeg_handle,
          encoder_state,
          obuffer.data(),
          &length,

          stream));

  
      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;

  }