nvjpeg
copied
junjie.jiang
2 years ago
17 changed files with 755 additions and 0 deletions
@ -0,0 +1,80 @@ |
|||
cmake_minimum_required(VERSION 3.11 FATAL_ERROR) |
|||
|
|||
include(GNUInstallDirs) |
|||
|
|||
set(CMAKE_CXX_STANDARD 14) |
|||
set(CMAKE_CXX_STANDARD_REQUIRED ON) |
|||
set(CMAKE_CXX_EXTENSIONS OFF) |
|||
|
|||
SET(PROJECT_NAME pynvjpeg) |
|||
PROJECT(${PROJECT_NAME} VERSION 0.1 LANGUAGES CUDA CXX) |
|||
|
|||
if(NOT DEFINED CMAKE_CUDA_STANDARD) |
|||
set(CMAKE_CUDA_STANDARD 11) |
|||
set(CMAKE_CUDA_STANDARD_REQUIRED ON) |
|||
endif() |
|||
|
|||
include_directories(${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) |
|||
link_directories(${CMAKE_CUDA_INCLUDE_DIRS}) |
|||
|
|||
find_library(CUDART_LIBRARY cudart ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES}) |
|||
find_library(NVJPEG_LIBRARY nvjpeg ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES}) |
|||
|
|||
include_directories( |
|||
${PROJECT_SOURCE_DIR} |
|||
) |
|||
|
|||
include(FetchContent) |
|||
FetchContent_Declare( |
|||
pybind11_sources |
|||
GIT_REPOSITORY https://github.com/pybind/pybind11.git |
|||
GIT_TAG v2.9 |
|||
) |
|||
FetchContent_GetProperties(pybind11_sources) |
|||
if(NOT pybind11_sources_POPULATED) |
|||
FetchContent_Populate(pybind11_sources) |
|||
add_subdirectory( |
|||
${pybind11_sources_SOURCE_DIR} |
|||
${pybind11_sources_BINARY_DIR} |
|||
) |
|||
endif() |
|||
|
|||
if(CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT) |
|||
set(CMAKE_INSTALL_PREFIX ${CMAKE_BINARY_DIR} CACHE PATH "" FORCE) |
|||
endif(CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT) |
|||
|
|||
add_library(${PROJECT_NAME} |
|||
SHARED |
|||
nvjpeg_decoder.cpp |
|||
jpeg_image.cpp |
|||
cuda_util.cpp |
|||
python_nvjpeg_decoder.cpp |
|||
) |
|||
|
|||
target_link_libraries(${PROJECT_NAME} |
|||
PUBLIC |
|||
${CUDART_LIBRARY} |
|||
${NVJPEG_LIBRARY} |
|||
pybind11::module) |
|||
|
|||
set_target_properties(${PROJECT_NAME} |
|||
PROPERTIES |
|||
PREFIX "${PYTHON_MODULE_PREFIX}" |
|||
SUFFIX "${PYTHON_MODULE_EXTENSION}" |
|||
) |
|||
|
|||
add_executable(test_decode test.cpp) |
|||
target_link_libraries(test_decode ${PROJECT_NAME} pybind11::embed) |
|||
|
|||
include(CTest) |
|||
enable_testing() |
|||
|
|||
add_test( |
|||
NAME test_decode |
|||
COMMAND $<TARGET_FILE:test_decode> ${CMAKE_CURRENT_SOURCE_DIR}/cat.jpg |
|||
) |
|||
|
|||
add_test( |
|||
NAME python_test_short |
|||
COMMAND ${PYTHON_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test.py $<TARGET_FILE:${PROJECT_NAME}> ${CMAKE_CURRENT_SOURCE_DIR}/cat.jpg |
|||
) |
After Width: | Height: | Size: 3.0 KiB |
@ -0,0 +1,12 @@ |
|||
|
|||
#include "cuda_util.h" |
|||
|
|||
|
|||
namespace NVJpegDecoder { |
|||
|
|||
int DevMalloc(void **p, size_t s) { return (int)cudaMalloc(p, s); } |
|||
int DevFree(void *p) { return (int)cudaFree(p); } |
|||
int HostMalloc(void** p, size_t s, unsigned int f) { return (int)cudaHostAlloc(p, s, f); } |
|||
int HostFree(void* p) { return (int)cudaFreeHost(p); } |
|||
|
|||
} // namespace NVJpegDecoder
|
@ -0,0 +1,98 @@ |
|||
#pragma once |
|||
|
|||
#include <string> |
|||
#include <iostream> |
|||
#include <vector> |
|||
#include <cuda_runtime_api.h> |
|||
#include <nvjpeg.h> |
|||
|
|||
namespace NVJpegDecoder { |
|||
|
|||
|
|||
#define CHECK_CUDA(call) \ |
|||
do { \ |
|||
CudaStatus s(call); \ |
|||
if (!s.IsOk()) { \ |
|||
std::cout << "CUDA Runtime failure: '#" << s.Msg() << "' at " << __FILE__ << ":" << __LINE__ << std::endl; \ |
|||
return false; \ |
|||
} \ |
|||
} while (false) \ |
|||
|
|||
#define CHECK_NVJPEG(call) \ |
|||
do { \ |
|||
NvJpegStatus s(call); \ |
|||
if (!s.IsOk()) { \ |
|||
std::cout << "NVJPEG failure: '#" << s.Msg() << "' at " << __FILE__ << ":" << __LINE__ << std::endl; \ |
|||
return false; \ |
|||
} \ |
|||
} while (false) |
|||
|
|||
|
|||
int DevMalloc(void **p, size_t s); |
|||
int DevFree(void *p); |
|||
int HostMalloc(void** p, size_t s, unsigned int f); |
|||
int HostFree(void* p); |
|||
|
|||
|
|||
class CudaStatus { |
|||
public: |
|||
explicit CudaStatus(cudaError_t error) : mCode(error) {} |
|||
|
|||
CudaStatus(CudaStatus&) = default; |
|||
CudaStatus& operator=(CudaStatus&) = default; |
|||
|
|||
bool IsOk() { |
|||
return mCode == cudaSuccess; |
|||
} |
|||
|
|||
std::string Msg() { |
|||
return cudaGetErrorString(mCode); |
|||
} |
|||
|
|||
private: |
|||
cudaError_t mCode; |
|||
}; |
|||
|
|||
|
|||
class NvJpegStatus { |
|||
public: |
|||
explicit NvJpegStatus(nvjpegStatus_t error) : mCode(error) {} |
|||
|
|||
bool IsOk() { |
|||
return mCode == NVJPEG_STATUS_SUCCESS; |
|||
} |
|||
|
|||
inline nvjpegStatus_t Code() {return mCode;} |
|||
|
|||
std::string Msg() { |
|||
switch (mCode) { |
|||
case NVJPEG_STATUS_NOT_INITIALIZED: |
|||
return "NVJPEG_STATUS_NOT_INITIALIZED"; |
|||
case NVJPEG_STATUS_INVALID_PARAMETER: |
|||
return "NVJPEG_STATUS_INVALID_PARAMETER"; |
|||
case NVJPEG_STATUS_BAD_JPEG: |
|||
return "NVJPEG_STATUS_BAD_JPEG"; |
|||
case NVJPEG_STATUS_JPEG_NOT_SUPPORTED: |
|||
return "NVJPEG_STATUS_JPEG_NOT_SUPPORTED"; |
|||
case NVJPEG_STATUS_ALLOCATOR_FAILURE: |
|||
return "NVJPEG_STATUS_ALLOCATOR_FAILURE"; |
|||
case NVJPEG_STATUS_EXECUTION_FAILED: |
|||
return "NVJPEG_STATUS_EXECUTION_FAILED"; |
|||
case NVJPEG_STATUS_ARCH_MISMATCH: |
|||
return "NVJPEG_STATUS_ARCH_MISMATCH"; |
|||
case NVJPEG_STATUS_INTERNAL_ERROR: |
|||
return "NVJPEG_STATUS_INTERNAL_ERROR"; |
|||
case NVJPEG_STATUS_IMPLEMENTATION_NOT_SUPPORTED: |
|||
return "NVJPEG_STATUS_IMPLEMENTATION_NOT_SUPPORTED"; |
|||
case NVJPEG_STATUS_INCOMPLETE_BITSTREAM: |
|||
return "NVJPEG_STATUS_INCOMPLETE_BITSTREAM"; |
|||
default: |
|||
return "UNKNOWN NVJPEG ERROR"; |
|||
} |
|||
} |
|||
|
|||
private: |
|||
nvjpegStatus_t mCode; |
|||
}; |
|||
|
|||
} // namespace NVJpegDecoder
|
@ -0,0 +1,59 @@ |
|||
#include "jpeg_image.h" |
|||
#include "cuda_util.h" |
|||
|
|||
|
|||
namespace NVJpegDecoder { |
|||
|
|||
bool JpegImage::Init(int width, int height, int channels) { |
|||
mNvImage = std::make_unique<nvjpegImage_t>(); |
|||
unsigned char * pBuffer = nullptr; |
|||
CHECK_CUDA(cudaMalloc((void **)&pBuffer, height * width * channels)); |
|||
for(int i = 0; i < channels; i++) { |
|||
mNvImage->channel[i] = pBuffer + (height *width * i); |
|||
mNvImage->pitch[i] = (unsigned int)width; |
|||
} |
|||
|
|||
mNvImage->pitch[0] = (unsigned int)width * channels; |
|||
mWidth = width; |
|||
mHeight = height; |
|||
mChannels = channels; |
|||
return true; |
|||
} |
|||
|
|||
JpegImage::JpegImage(JpegImage&& rhs) { |
|||
mWidth = rhs.mWidth; |
|||
mHeight = rhs.mHeight; |
|||
mChannels = rhs.mChannels; |
|||
mSubsampling = rhs.mSubsampling; |
|||
mNvImage = std::move(rhs.mNvImage); |
|||
} |
|||
|
|||
JpegImage& JpegImage::operator=(JpegImage&& rhs) { |
|||
mWidth = rhs.mWidth; |
|||
mHeight = rhs.mHeight; |
|||
mChannels = rhs.mChannels; |
|||
mSubsampling = rhs.mSubsampling; |
|||
mNvImage = std::move(rhs.mNvImage); |
|||
return *this; |
|||
} |
|||
|
|||
unsigned char* JpegImage::Cpu() { |
|||
size_t size = mHeight * mWidth * mChannels; |
|||
auto buffer = std::make_unique<unsigned char[]>(size) ; |
|||
CudaStatus s(cudaMemcpy(buffer.get(), mNvImage->channel[0], size, cudaMemcpyDeviceToHost)); |
|||
|
|||
if (!s.IsOk()) { |
|||
std::cout << "Copy image from GPU to CPU failed: " << s.Msg() << std::endl; |
|||
return nullptr; |
|||
} |
|||
return buffer.release(); |
|||
} |
|||
|
|||
JpegImage::~JpegImage(){ |
|||
if (mNvImage != nullptr) { |
|||
cudaFree(mNvImage->channel[0]); |
|||
} |
|||
} |
|||
|
|||
|
|||
} // namespace NVJpegDecoder
|
@ -0,0 +1,39 @@ |
|||
#pragma once |
|||
#include <vector> |
|||
#include <memory> |
|||
#include <nvjpeg.h> |
|||
|
|||
namespace NVJpegDecoder { |
|||
|
|||
class JpegImage { |
|||
|
|||
public: |
|||
JpegImage() noexcept : mNvImage(nullptr) {} |
|||
virtual ~JpegImage(); |
|||
|
|||
JpegImage(const JpegImage&) = delete; |
|||
JpegImage& operator=(const JpegImage&) = delete; |
|||
|
|||
JpegImage(JpegImage&& rhs); |
|||
JpegImage& operator=(JpegImage&& rhs); |
|||
|
|||
bool Init(int width, int height, int channels); |
|||
|
|||
nvjpegImage_t* GetImagePoint() { |
|||
return mNvImage.get(); |
|||
} |
|||
|
|||
const std::vector<int64_t> Dims() { |
|||
return std::vector<int64_t>{mHeight, mWidth, mChannels}; |
|||
} |
|||
|
|||
unsigned char* Cpu(); |
|||
|
|||
private: |
|||
int mWidth = 0; |
|||
int mHeight = 0; |
|||
int mChannels = 0; |
|||
nvjpegChromaSubsampling_t mSubsampling; |
|||
std::unique_ptr<nvjpegImage_t> mNvImage; |
|||
}; |
|||
} // namespace NVJpegDecoder
|
@ -0,0 +1,23 @@ |
|||
#include <pybind11/pybind11.h> |
|||
#include <pybind11/embed.h> |
|||
#include <iostream> |
|||
|
|||
namespace py = pybind11; |
|||
|
|||
|
|||
void print(py::list l) { |
|||
for (auto item: l) { |
|||
std::cout << item.attr("__str__")().cast<std::string>() << std::endl; |
|||
} |
|||
} |
|||
|
|||
int main() { |
|||
py::scoped_interpreter guard{}; |
|||
py::print("Hello, World!"); |
|||
py::list data; |
|||
for (int i = 0; i < 10; i++) { |
|||
data.append(i); |
|||
} |
|||
print(data); |
|||
return 0; |
|||
} |
@ -0,0 +1,93 @@ |
|||
#include <string> |
|||
#include <iostream> |
|||
#include <fstream> |
|||
#include <cassert> |
|||
#include "cuda_util.h" |
|||
#include "nvjpeg_decoder.h" |
|||
|
|||
|
|||
namespace NVJpegDecoder { |
|||
|
|||
bool Decoder::BindDevice(int device_id) { |
|||
if (device_id == mDeviceId) { |
|||
return true; |
|||
} |
|||
|
|||
if (device_id < 0) { |
|||
std::cout<< "Device id must >= 0, the input is " |
|||
<< device_id << std::endl; |
|||
return false; |
|||
} |
|||
|
|||
Destroy(); |
|||
mDeviceId = device_id; |
|||
CHECK_CUDA(cudaSetDevice(mDeviceId)); |
|||
CHECK_NVJPEG(nvjpegCreateSimple(&mHandle)); |
|||
CHECK_NVJPEG(nvjpegJpegStateCreate(mHandle, &mState)); |
|||
CHECK_CUDA(cudaStreamCreateWithFlags(&mStream, cudaStreamNonBlocking)); |
|||
return true; |
|||
} |
|||
|
|||
bool Decoder::Read(const char* filename, JpegImage& image) { |
|||
std::ifstream input(filename); |
|||
if (!(input.is_open())) { |
|||
std::cout << "Open file " << filename << " failed" << std::endl; |
|||
return false; |
|||
} |
|||
std::string imagedata((std::istreambuf_iterator<char>(input)), std::istreambuf_iterator<char>()); |
|||
if (!Decode(imagedata, image)) { |
|||
return false; |
|||
} |
|||
return true; |
|||
} |
|||
|
|||
bool Decoder::Decode(std::string& imagedata , JpegImage& image) { |
|||
if (!PrepareJpegImage(imagedata, image)) { |
|||
return false; |
|||
} |
|||
CHECK_NVJPEG(nvjpegDecode( |
|||
mHandle, |
|||
mState, |
|||
(const unsigned char *)imagedata.data(), |
|||
imagedata.size(), |
|||
NVJPEG_OUTPUT_RGBI, |
|||
image.GetImagePoint(), |
|||
mStream)); |
|||
return true; |
|||
} |
|||
|
|||
bool Decoder::PrepareJpegImage(const std::string& image, JpegImage& output) { |
|||
int widths[NVJPEG_MAX_COMPONENT]; |
|||
int heights[NVJPEG_MAX_COMPONENT]; |
|||
int channels; |
|||
nvjpegChromaSubsampling_t subsampling; |
|||
|
|||
CHECK_NVJPEG(nvjpegGetImageInfo( |
|||
mHandle, (unsigned char *)image.data(), image.size(), |
|||
&channels, &subsampling, widths, heights)); |
|||
|
|||
if (NVJPEG_CSS_UNKNOWN == subsampling) { |
|||
std::cout << "Unknown chroma subsampling" << std::endl; |
|||
return false; |
|||
} |
|||
|
|||
if (!output.Init(widths[0], heights[0], channels)) { |
|||
return false; |
|||
} |
|||
return true; |
|||
} |
|||
|
|||
void Decoder::Destroy() { |
|||
if (mDeviceId >= 0) { |
|||
cudaStreamDestroy(mStream); |
|||
nvjpegJpegStateDestroy(mState); |
|||
nvjpegDestroy(mHandle); |
|||
mDeviceId = -1; |
|||
} |
|||
} |
|||
|
|||
Decoder::~Decoder() { |
|||
Destroy(); |
|||
} |
|||
|
|||
} // namespace NVJpegDecoder
|
@ -0,0 +1,39 @@ |
|||
#pragma once |
|||
|
|||
#include <vector> |
|||
#include <string> |
|||
#include <cuda_runtime_api.h> |
|||
#include <nvjpeg.h> |
|||
|
|||
#include "jpeg_image.h" |
|||
|
|||
namespace NVJpegDecoder { |
|||
|
|||
typedef std::vector<std::string> OriginJpegImages; |
|||
typedef std::vector<JpegImage> JpegImages; |
|||
|
|||
class Decoder { |
|||
public: |
|||
Decoder():mDeviceId(-1) {} |
|||
virtual ~Decoder(); |
|||
|
|||
Decoder(Decoder&) = delete; |
|||
Decoder& operator=(Decoder&) = delete; |
|||
|
|||
bool BindDevice(int device_id=0); |
|||
bool Read(const char* filename, JpegImage& image); |
|||
bool Decode(std::string& imagedata , JpegImage& image); |
|||
|
|||
private: |
|||
bool PrepareJpegImage(const std::string& image, JpegImage& output); |
|||
void Destroy(); |
|||
|
|||
private: |
|||
nvjpegHandle_t mHandle; |
|||
nvjpegJpegState_t mState; |
|||
|
|||
int mDeviceId; |
|||
cudaStream_t mStream; |
|||
}; |
|||
|
|||
} // namespace NVJpegDecoder
|
@ -0,0 +1,52 @@ |
|||
#include <string> |
|||
#include <iostream> |
|||
#include <pybind11/numpy.h> |
|||
#include <pybind11/stl_bind.h> |
|||
#include "python_nvjpeg_decoder.h" |
|||
|
|||
|
|||
namespace NVJpegDecoder { |
|||
|
|||
bool PythonDecoder::BindDevice(int device_id) { |
|||
py::gil_scoped_release release; |
|||
return mDecoder.BindDevice(device_id); |
|||
} |
|||
|
|||
py::object PythonDecoder::Read(std::string& filename) { |
|||
unsigned char* data = nullptr; |
|||
JpegImage image; |
|||
{ |
|||
py::gil_scoped_release release; |
|||
if (mDecoder.Read(filename.c_str(), image)) { |
|||
data = image.Cpu(); |
|||
} |
|||
} // gets gil
|
|||
|
|||
std::unique_ptr<unsigned char> ret(data); |
|||
if (ret != nullptr) { |
|||
return py::array(py::dtype(py::format_descriptor<uint8_t>::format()), image.Dims(), (void*)ret.get()); |
|||
} |
|||
return py::none(); |
|||
} |
|||
|
|||
py::object PythonDecoder::Decode(std::string& image_bytes) { |
|||
JpegImage image; |
|||
if (!mDecoder.Decode(image_bytes, image)) { |
|||
return py::none(); |
|||
} |
|||
std::unique_ptr<unsigned char> data(image.Cpu()); |
|||
if (nullptr == data) { |
|||
return py::none(); |
|||
} |
|||
return py::array(py::dtype(py::format_descriptor<uint8_t>::format()), image.Dims(), (void*)data.get()); |
|||
} |
|||
|
|||
PYBIND11_MODULE(pynvjpeg, m) { |
|||
py::class_<PythonDecoder, std::shared_ptr<PythonDecoder>>(m, "Decoder") |
|||
.def(py::init()) |
|||
.def("bind_device", &PythonDecoder::BindDevice) |
|||
.def("imread", &PythonDecoder::Read) |
|||
.def("imdecode", &PythonDecoder::Decode); |
|||
} |
|||
|
|||
} // namespace NVJpegDecoder
|
@ -0,0 +1,26 @@ |
|||
#pragma once |
|||
|
|||
#include <pybind11/pybind11.h> |
|||
#include <pybind11/stl.h> |
|||
#include "nvjpeg_decoder.h" |
|||
|
|||
namespace py = pybind11; |
|||
|
|||
namespace NVJpegDecoder { |
|||
|
|||
class PythonDecoder { |
|||
public: |
|||
PythonDecoder() = default; |
|||
|
|||
PythonDecoder(PythonDecoder&) = delete; |
|||
PythonDecoder& operator=(PythonDecoder&) = delete; |
|||
|
|||
bool BindDevice(int device_id=0); |
|||
py::object Read(std::string&); |
|||
py::object Decode(std::string&); |
|||
|
|||
private: |
|||
Decoder mDecoder; |
|||
}; |
|||
|
|||
} // namespace NVJpegDecoder
|
@ -0,0 +1,90 @@ |
|||
#include <string> |
|||
#include <sys/timeb.h> |
|||
#include <sstream> |
|||
#include <fstream> |
|||
#include <iostream> |
|||
#include "nvjpeg_decoder.h" |
|||
|
|||
|
|||
bool ReadImage(const char* filename, std::string& imagedata) { |
|||
std::ifstream input(filename); |
|||
if (!(input.is_open())) { |
|||
std::cout << "Open file " << filename << " failed" << std::endl; |
|||
return false; |
|||
} |
|||
|
|||
imagedata = std::string((std::istreambuf_iterator<char>(input)), std::istreambuf_iterator<char>()); |
|||
return true; |
|||
} |
|||
|
|||
int main(int argc, char *argv[]) { |
|||
const char* image_file = argv[1]; |
|||
NVJpegDecoder::Decoder decoder; |
|||
if (!decoder.BindDevice(0)) { |
|||
std::cout << "Init Failed" << std::endl; |
|||
return -1; |
|||
} |
|||
|
|||
// test deocde from bytes
|
|||
{ |
|||
std::string image_btyes; |
|||
if (!ReadImage(image_file, image_btyes)) { |
|||
return -1; |
|||
} |
|||
int count = 0; |
|||
timeb t1, t2; |
|||
ftime(&t1); |
|||
|
|||
while (count < 1000) { |
|||
NVJpegDecoder::JpegImage image; |
|||
if (!decoder.Decode(image_btyes, image)) { |
|||
return -1; |
|||
} |
|||
count++; |
|||
} |
|||
ftime(&t2); |
|||
std::cout << "--------- " << t2.time * 1000 + t2.millitm - t1.time * 1000 - t1.millitm << std::endl; |
|||
} |
|||
|
|||
// test decode from file
|
|||
{ |
|||
timeb t1, t2; |
|||
ftime(&t1); |
|||
{ |
|||
int count = 0; |
|||
while (count < 1000) { |
|||
NVJpegDecoder::JpegImage image; |
|||
if (!decoder.Read(image_file, image)) { |
|||
return -1; |
|||
} |
|||
unsigned char* d = image.Cpu(); |
|||
delete[] d; |
|||
count++; |
|||
} |
|||
} |
|||
ftime(&t2); |
|||
std::cout << "--------- " << t2.time * 1000 + t2.millitm - t1.time * 1000 - t1.millitm << std::endl; |
|||
} |
|||
|
|||
decoder.BindDevice(1); |
|||
{ |
|||
std::string image_btyes; |
|||
if (!ReadImage(image_file, image_btyes)) { |
|||
return -1; |
|||
} |
|||
int count = 0; |
|||
timeb t1, t2; |
|||
ftime(&t1); |
|||
|
|||
while (count < 1000) { |
|||
NVJpegDecoder::JpegImage image; |
|||
if (!decoder.Decode(image_btyes, image)) { |
|||
return -1; |
|||
} |
|||
count++; |
|||
} |
|||
ftime(&t2); |
|||
std::cout << "--------- " << t2.time * 1000 + t2.millitm - t1.time * 1000 - t1.millitm << std::endl; |
|||
} |
|||
return 0; |
|||
} |
@ -0,0 +1,26 @@ |
|||
import sys |
|||
import numpy as np |
|||
import cv2 |
|||
import time |
|||
sys.path.insert(0, '/'.join(sys.argv[1].split('/')[:-1])) |
|||
import pynvjpeg |
|||
|
|||
|
|||
if __name__ == '__main__': |
|||
decoder = pynvjpeg.Decoder() |
|||
assert(decoder.bind_device(0) is True) |
|||
filename = sys.argv[2] |
|||
image0 = decoder.imread(filename) |
|||
image0 = image0.astype(np.int32) |
|||
h, w, c = image0.shape |
|||
with open(filename, 'rb') as f: |
|||
data = f.read() |
|||
|
|||
|
|||
image1 = cv2.imread(filename) |
|||
image1 = cv2.cvtColor(image1, cv2.COLOR_BGR2RGB) |
|||
image1 = image1.astype(np.int32) |
|||
if np.sum((image1 - image0)) > h * w * c: |
|||
exit(-1) |
|||
exit(0) |
|||
|
@ -0,0 +1,22 @@ |
|||
#include <iostream> |
|||
#include <pybind11/embed.h> |
|||
#include "python_nvjpeg_decoder.h" |
|||
|
|||
|
|||
int main() { |
|||
py::scoped_interpreter guard{}; |
|||
NVJpegDecoder::PythonDecoder decoder; |
|||
if (!decoder.Init()) { |
|||
std::cout << "Init Failed" << std::endl; |
|||
return -1; |
|||
} |
|||
std::cout << 1 << std::endl;; |
|||
py::list data; |
|||
std::cout << 2 << std::endl;; |
|||
for (int i = 0; i < 10; i++) { |
|||
std::cout << i << ' '; |
|||
data.append(i); |
|||
} |
|||
decoder.BatchDecode(data); |
|||
return 0; |
|||
} |
@ -0,0 +1,96 @@ |
|||
# encoding=utf-8 |
|||
import os |
|||
import sys |
|||
import logging |
|||
import platform |
|||
|
|||
import cv2 |
|||
import requests |
|||
import numpy as np |
|||
from towhee.types import Image |
|||
from towhee.operator import PyOperator, SharedType |
|||
|
|||
logger = logging.getLogger() |
|||
|
|||
|
|||
def create_pynvjpeg(): |
|||
if platform.system().lower() != 'linux': |
|||
logger.warning('The GPU deocder only support linux') |
|||
return None |
|||
|
|||
py_version = platform.python_version() |
|||
version = ''.join(py_version.split('.')[0:2]) |
|||
cur_dir = os.path.dirname(os.path.abspath(__file__)) |
|||
lib_path = cur_dir + '/' + 'py' + version + 'nvjpeg' |
|||
sys.path.insert(0, lib_path) |
|||
try: |
|||
import pynvjpeg |
|||
decoder = pynvjpeg.Decoder() |
|||
if not decoder.bind_device(0): |
|||
logger.info('Init GPU deocder failed, use CPU decoder') |
|||
return None |
|||
else: |
|||
return decoder |
|||
except Exception as e: |
|||
logger.error('Create nvjpeg failed, use opencv decoder, errors: ', str(e)) |
|||
return None |
|||
|
|||
|
|||
class ImageDecoder(PyOperator): |
|||
def __init__(self): |
|||
self._nvjpeg = create_pynvjpeg() |
|||
|
|||
def to_device(self, device): |
|||
if self._nvjpeg is not None: |
|||
if not self._nvjpeg.bind_device(device): |
|||
self._nvjpeg = None |
|||
|
|||
def _download_image(self, image_url): |
|||
image_type = image_url.split('?')[0].split('.')[-1].lower() |
|||
is_jpeg = True if image_type == 'jpeg' or image_type == 'jpg' else False |
|||
r = requests.get(image_url, timeout=(20, 20)) |
|||
if r.status_code // 100 != 2: |
|||
logging.error('Download image from %s failed, error msg: %s, request code: %s ' % (image_url, |
|||
r.text, |
|||
r.status_code)) |
|||
return None |
|||
return is_jpeg, r.content |
|||
|
|||
def _read_image(self, image_path): |
|||
is_jpeg = image_path.lower().endswith('jpg') or image_path.lower().endswith('jpeg') |
|||
with open(image_path, 'rb') as f: |
|||
return is_jpeg, f.read() |
|||
|
|||
def load_image(self, image_path): |
|||
if image_path.startswith('http'): |
|||
return self._download_image(image_path) |
|||
else: |
|||
return self._read_image(image_path) |
|||
|
|||
def __call__(self, image_path: str): |
|||
is_jpeg, image_content = self.load_image(image_path) |
|||
if is_jpeg and self._nvjpeg is not None: |
|||
image = self._nvjpeg.imdecode(image_content) |
|||
else: |
|||
arr = np.asarray(bytearray(image_content), dtype=np.uint8) |
|||
image = cv2.imdecode(arr, -1) |
|||
if image is not None: |
|||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) |
|||
if image is None: |
|||
raise RuntimeError('Decode image %s failed' % image_path) |
|||
return Image(image, 'RGB') |
|||
|
|||
def input_schema(self): |
|||
return [(str, (1,))] |
|||
|
|||
def output_schema(self): |
|||
return [(Image, (-1, -1, 3))] |
|||
|
|||
@property |
|||
def shared_type(self): |
|||
return SharedType.NotShareable |
|||
|
|||
if __name__ == '__main__': |
|||
d = ImageDecoder() |
|||
print(d('/home/junjie.jiangjjj/images/1.png')) |
|||
print(d('/home/junjie.jiangjjj/images/towhee.jpeg')) |
Binary file not shown.
Loading…
Reference in new issue