Wang's blog

使用nvJPEG加速解码JPG图像

Published on

背景

在读取JPG图像时，需要先对图像进行解码。一般的图像库（如opencv）采用CPU进行解码，如果对速度有很高要求则需要另外实现。

nvJPEG是英伟达提供的可以使用GPU加速的JPG解码库，它包含在较高版本的CUDA（>10.0）中，如果使用低版本CUDA，则需要另外安装。

步骤

1. 假设待解码的本地文件为test.jpg，首先将其二进制数据读入内存

    // 以二进制方式打开文件
    std::ifstream in_fs("test.jpg", std::ifstream::binary);

    // 获取文件大小
    in_fs.seekg(0, std::ios::end);
    auto in_size = in_fs.tellg();

    // 读入全部数据（使用vector避免手动申请/释放空间）
    std::vector<uchar> in_buf(in_size);
    in_fs.seekg(0);
    in_fs.read((char*)in_buf.data(), in_size);

2. 初始化nvJPEG，建立handle与state

    nvjpegHandle_t handle;
    nvjpegJpegState_t state;
    nvjpegCreateSimple(&handle);
    nvjpegJpegStateCreate(handle, &state);

3. 获取图像长宽、通道数、subsampling等信息

    int widths[NVJPEG_MAX_COMPONENT];
    int heights[NVJPEG_MAX_COMPONENT];
    int channels;
    nvjpegChromaSubsampling_t subsampling;
    nvjpegGetImageInfo(handle, in_buf.data(), in_size, &channels, &subsampling, widths, heights);

4. 设置输出参数并解码

解码时，除了需要传入输入数据及大小外，还需要指定输出格式（nvjpegOutputFormat_t）及输出图像信息（nvjpegImage_t，包含每个通道的行宽度pitch，以及每个通道的输出缓存地址channel）

此处指定输出格式为NVJPEG_OUTPUT_BGRI，该格式将所有输出写入channel[0]中，格式为"BGRBGRBGR…"，因此picth[0]为图像宽度×3，channel[0]缓存大小至少应为图像长度×图像宽度×3。如果需要使用其它格式请参考官方文档

    // 计算输出行宽度与总大小
    int mul = 3;
    int out_step = widths[0] * mul;
    int out_size = out_step * heights[0];

    // 设置输出信息，并在显存上申请缓存
    nvjpegImage_t out_buf;
    out_buf.pitch[0] = out_step;
    cudaMalloc((void**)&out_buf.channel[0], out_size);

    // 解码
    auto ret = nvjpegDecode(handle, state, in_buf.data(), in_size, NVJPEG_OUTPUT_BGRI, &out_buf, nullptr);

5. 输出格式NVJPEG_OUTPUT_BGRI与opencv的默认格式完全相同，因此可以直接拷贝到cv::Mat中使用。另外也可以直接使用缓存建立cv::cuda::GpuMat，无需拷贝即可使用

    // 根据长宽建立opencv图像
    cv::Mat image(heights[0], widths[0], CV_8UC3);
    // 将数据拷贝后即可直接使用image
    cudaMemcpy(image.data, out_buf.channel[0], out_size, cudaMemcpyDeviceToHost);

6. 释放缓存与nvJPEG的handle与state

    cudaFree(out_buf.channel[0]);
    nvjpegJpegStateDestroy(state);
    nvjpegDestroy(handle);

参考