NvJpegEncoder.cpp
2.89 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
#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)
{
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, 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;
}