From b1ef906fb7892c93a0a099dfade3e8695ae07f17 Mon Sep 17 00:00:00 2001
From: Shahriar <shahriar25.ss@gmail.com>
Date: Sat, 10 Aug 2019 15:37:19 +0430
Subject: [PATCH] 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 <sxx/sxx.h>
+
+#include <png.h>
+
+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<Reader*>(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<uint8_t>();
+  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 <torch/torch.h>
+
+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