Files
FastDeploy/fastdeploy/vision/common/image_decoder/image_decoder.cc
Wang Xinyu efa46563f3 [nvJPEG] Integrate nvJPEG decoder (#1288)
* nvjpeg cmake

* add common decoder, nvjpeg decoder and add image name predict api

* ppclas support nvjpeg decoder

* remove useless comments

* image decoder support opencv

* nvjpeg decode fallback to opencv

* fdtensor add nbytes_allocated

* single image decode api

* fix bug

* add pybind

* ignore nvjpeg on jetson

* fix cmake in

* predict on fdmat

* remove image names predict api, add image decoder tutorial

* Update __init__.py

* fix pybind
2023-02-17 10:27:05 +08:00

113 lines
3.5 KiB
C++

// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "fastdeploy/vision/common/image_decoder/image_decoder.h"
#include "opencv2/imgcodecs.hpp"
namespace fastdeploy {
namespace vision {
ImageDecoder::ImageDecoder(ImageDecoderLib lib) {
if (lib == ImageDecoderLib::NVJPEG) {
#ifdef ENABLE_NVJPEG
nvjpeg::init_decoder(nvjpeg_params_);
#endif
}
lib_ = lib;
}
ImageDecoder::~ImageDecoder() {
if (lib_ == ImageDecoderLib::NVJPEG) {
#ifdef ENABLE_NVJPEG
nvjpeg::destroy_decoder(nvjpeg_params_);
#endif
}
}
bool ImageDecoder::Decode(const std::string& img_name, FDMat* mat) {
std::vector<FDMat> mats(1);
mats[0] = std::move(*mat);
if (!BatchDecode({img_name}, &mats)) {
return false;
}
*mat = std::move(mats[0]);
return true;
}
bool ImageDecoder::BatchDecode(const std::vector<std::string>& img_names,
std::vector<FDMat>* mats) {
if (lib_ == ImageDecoderLib::OPENCV) {
return ImplByOpenCV(img_names, mats);
} else if (lib_ == ImageDecoderLib::NVJPEG) {
return ImplByNvJpeg(img_names, mats);
}
return true;
}
bool ImageDecoder::ImplByOpenCV(const std::vector<std::string>& img_names,
std::vector<FDMat>* mats) {
for (size_t i = 0; i < img_names.size(); ++i) {
cv::Mat im = cv::imread(img_names[i]);
(*mats)[i].SetMat(im);
(*mats)[i].layout = Layout::HWC;
(*mats)[i].SetWidth(im.cols);
(*mats)[i].SetHeight(im.rows);
(*mats)[i].SetChannels(im.channels());
}
return true;
}
bool ImageDecoder::ImplByNvJpeg(const std::vector<std::string>& img_names,
std::vector<FDMat>* mats) {
#ifdef ENABLE_NVJPEG
nvjpeg_params_.batch_size = img_names.size();
std::vector<nvjpegImage_t> output_imgs(nvjpeg_params_.batch_size);
std::vector<int> widths(nvjpeg_params_.batch_size);
std::vector<int> heights(nvjpeg_params_.batch_size);
// TODO(wangxinyu): support other output format
nvjpeg_params_.fmt = NVJPEG_OUTPUT_BGRI;
double total;
nvjpeg_params_.stream = (*mats)[0].Stream();
std::vector<FDTensor*> output_buffers;
for (size_t i = 0; i < mats->size(); ++i) {
FDASSERT((*mats)[i].output_cache != nullptr,
"The output_cache of FDMat was not set.");
output_buffers.push_back((*mats)[i].output_cache);
}
if (nvjpeg::process_images(img_names, nvjpeg_params_, total, output_imgs,
output_buffers, widths, heights)) {
// If nvJPEG decode failed, will fallback to OpenCV,
// e.g. png format is not supported by nvJPEG
FDWARNING << "nvJPEG decode failed, falling back to OpenCV for this batch"
<< std::endl;
return ImplByOpenCV(img_names, mats);
}
for (size_t i = 0; i < mats->size(); ++i) {
(*mats)[i].mat_type = ProcLib::CUDA;
(*mats)[i].layout = Layout::HWC;
(*mats)[i].SetTensor(output_buffers[i]);
}
#else
FDASSERT(false, "FastDeploy didn't compile with NVJPEG.");
#endif
return true;
}
} // namespace vision
} // namespace fastdeploy