From f6e5417d3c997688e1c9841ed007b5a4db1b6754 Mon Sep 17 00:00:00 2001 From: Shahriar Date: Sat, 10 Aug 2019 15:37:19 +0430 Subject: [PATCH 1/2] Added PNG reading functions --- CMakeLists.txt | 10 +++- torchvision/csrc/image/image.h | 6 +++ torchvision/csrc/image/readpng.cpp | 83 ++++++++++++++++++++++++++++++ torchvision/csrc/image/readpng.h | 19 +++++++ 4 files changed, 117 insertions(+), 1 deletion(-) create mode 100644 torchvision/csrc/image/image.h create mode 100644 torchvision/csrc/image/readpng.cpp create mode 100644 torchvision/csrc/image/readpng.h diff --git a/CMakeLists.txt b/CMakeLists.txt index df77482c870..3e4b065e83e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -2,15 +2,23 @@ cmake_minimum_required(VERSION 2.8) project(torchvision) set(CMAKE_CXX_STANDARD 11) +find_package(PNG REQUIRED) find_package(Torch REQUIRED) +include_directories(${PNG_INCLUDE_DIR}) file(GLOB HEADERS torchvision/csrc/vision.h) + file(GLOB MODELS_HEADERS torchvision/csrc/models/*.h) file(GLOB MODELS_SOURCES torchvision/csrc/models/*.h torchvision/csrc/models/*.cpp) -add_library (${PROJECT_NAME} SHARED ${MODELS_SOURCES}) +file(GLOB IMAGE_HEADERS torchvision/csrc/image/*.h) +file(GLOB IMAGE_SOURCES torchvision/csrc/image/*.h torchvision/csrc/image/*.cpp) + +add_library(${PROJECT_NAME} SHARED ${MODELS_SOURCES} ${IMAGE_SOURCES}) +target_link_libraries(${PROJECT_NAME} PUBLIC "${PNG_LIBRARY}") target_link_libraries(${PROJECT_NAME} PUBLIC "${TORCH_LIBRARIES}") install(TARGETS ${PROJECT_NAME} DESTINATION ${CMAKE_INSTALL_PREFIX}/lib) install(FILES ${HEADERS} DESTINATION ${CMAKE_INSTALL_PREFIX}/include/${PROJECT_NAME}) install(FILES ${MODELS_HEADERS} DESTINATION ${CMAKE_INSTALL_PREFIX}/include/${PROJECT_NAME}/models) +install(FILES ${IMAGE_HEADERS} DESTINATION ${CMAKE_INSTALL_PREFIX}/include/${PROJECT_NAME}/image) diff --git a/torchvision/csrc/image/image.h b/torchvision/csrc/image/image.h new file mode 100644 index 00000000000..f3b6973ecf7 --- /dev/null +++ b/torchvision/csrc/image/image.h @@ -0,0 +1,6 @@ +#ifndef IMAGE_H +#define IMAGE_H + +#include "readpng.h" + +#endif // IMAGE_H diff --git a/torchvision/csrc/image/readpng.cpp b/torchvision/csrc/image/readpng.cpp new file mode 100644 index 00000000000..d589d60912e --- /dev/null +++ b/torchvision/csrc/image/readpng.cpp @@ -0,0 +1,83 @@ +#include "readpng.h" + +#include + +#include + +namespace torch { +namespace vision { +namespace image { +namespace impl { +bool is_png(const void* data) { + return png_sig_cmp(png_const_bytep(data), 0, 8) == 0; +} + +torch::Tensor readpng(const void* data) { + struct Reader { + png_const_bytep ptr; + } reader; + + reader.ptr = png_const_bytep(data) + 8; + auto read_callback = + [](png_structp png_ptr, png_bytep output, png_size_t bytes) { + auto reader = static_cast(png_get_io_ptr(png_ptr)); + std::copy(reader->ptr, reader->ptr + bytes, output); + reader->ptr += bytes; + }; + + auto png_ptr = + png_create_read_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr); + if (!png_ptr) + return torch::tensor({0}); + + auto info_ptr = png_create_info_struct(png_ptr); + if (!info_ptr) { + png_destroy_read_struct(&png_ptr, nullptr, nullptr); + return torch::tensor({0}); + } + + if (setjmp(png_jmpbuf(png_ptr)) != 0) { + png_destroy_read_struct(&png_ptr, &info_ptr, nullptr); + return torch::tensor({0}); + } + + png_set_read_fn(png_ptr, &reader, read_callback); + png_set_sig_bytes(png_ptr, 8); + png_read_info(png_ptr, info_ptr); + + png_uint_32 width, height; + int bit_depth, color_type; + auto retval = png_get_IHDR( + png_ptr, + info_ptr, + &width, + &height, + &bit_depth, + &color_type, + nullptr, + nullptr, + nullptr); + + if (retval != 1 || color_type != PNG_COLOR_TYPE_RGB) { + png_destroy_read_struct(&png_ptr, &info_ptr, nullptr); + return torch::tensor({0}); + } + + auto tensor = + torch::empty({int64_t(height), int64_t(width), int64_t(3)}, torch::kU8); + auto ptr = tensor.data(); + auto bytes = png_get_rowbytes(png_ptr, info_ptr); + + for (decltype(height) i = 0; i < height; ++i) { + png_read_row(png_ptr, ptr, nullptr); + ptr += bytes; + } + + png_destroy_read_struct(&png_ptr, &info_ptr, nullptr); + return tensor; +} + +} // namespace impl +} // namespace image +} // namespace vision +} // namespace torch diff --git a/torchvision/csrc/image/readpng.h b/torchvision/csrc/image/readpng.h new file mode 100644 index 00000000000..135f07e3f47 --- /dev/null +++ b/torchvision/csrc/image/readpng.h @@ -0,0 +1,19 @@ +#ifndef READPNG_H +#define READPNG_H + +#include + +namespace torch { +namespace vision { +namespace image { +namespace impl { + +bool is_png(const void* data); +torch::Tensor readpng(const void* data); + +} // namespace impl +} // namespace image +} // namespace vision +} // namespace torch + +#endif // READPNG_H From 73c09aa118b9e20db8f7c24f2fddfb78c5ac8a31 Mon Sep 17 00:00:00 2001 From: Shahriar Date: Sat, 10 Aug 2019 15:48:28 +0430 Subject: [PATCH 2/2] Fixed some stuff --- torchvision/csrc/image/readpng.cpp | 10 +++++----- torchvision/csrc/image/readpng.h | 2 +- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/torchvision/csrc/image/readpng.cpp b/torchvision/csrc/image/readpng.cpp index d589d60912e..b92d05ab998 100644 --- a/torchvision/csrc/image/readpng.cpp +++ b/torchvision/csrc/image/readpng.cpp @@ -12,7 +12,7 @@ bool is_png(const void* data) { return png_sig_cmp(png_const_bytep(data), 0, 8) == 0; } -torch::Tensor readpng(const void* data) { +torch::Tensor read_png(const void* data) { struct Reader { png_const_bytep ptr; } reader; @@ -28,17 +28,17 @@ torch::Tensor readpng(const void* data) { auto png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr); if (!png_ptr) - return torch::tensor({0}); + return torch::empty({}); auto info_ptr = png_create_info_struct(png_ptr); if (!info_ptr) { png_destroy_read_struct(&png_ptr, nullptr, nullptr); - return torch::tensor({0}); + return torch::empty({}); } if (setjmp(png_jmpbuf(png_ptr)) != 0) { png_destroy_read_struct(&png_ptr, &info_ptr, nullptr); - return torch::tensor({0}); + return torch::empty({}); } png_set_read_fn(png_ptr, &reader, read_callback); @@ -60,7 +60,7 @@ torch::Tensor readpng(const void* data) { if (retval != 1 || color_type != PNG_COLOR_TYPE_RGB) { png_destroy_read_struct(&png_ptr, &info_ptr, nullptr); - return torch::tensor({0}); + return torch::empty({}); } auto tensor = diff --git a/torchvision/csrc/image/readpng.h b/torchvision/csrc/image/readpng.h index 135f07e3f47..62ef5afbb43 100644 --- a/torchvision/csrc/image/readpng.h +++ b/torchvision/csrc/image/readpng.h @@ -9,7 +9,7 @@ namespace image { namespace impl { bool is_png(const void* data); -torch::Tensor readpng(const void* data); +torch::Tensor read_png(const void* data); } // namespace impl } // namespace image