nvjpeg
              
                 
                
            
          copied
				 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
	
	