nvjpeg
copied
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
Readme
Files and versions
52 lines
1.4 KiB
52 lines
1.4 KiB
#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
|