Adds readpng features

Author: r-zenine

.circleci/config.yml

@@ -128,7 +128,7 @@ jobs:
               ca-certificates \
               curl \
               gnupg-agent \
-              software-properties-common
+              software-properties-common 
 
           curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo apt-key add -
 
@@ -219,7 +219,8 @@ jobs:
             curl -o conda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh
             sh conda.sh -b
             source $HOME/miniconda3/bin/activate
-            packaging/build_wheel.sh
+            conda install -yq libpng
+            packaging/build_wheel_macos.sh
       - store_artifacts:
           path: dist
       - persist_to_workspace:
@@ -239,6 +240,7 @@ jobs:
             sh conda.sh -b
             source $HOME/miniconda3/bin/activate
             conda install -yq conda-build
+            conda install -yq libpng
             packaging/build_conda.sh
       - store_artifacts:
           path: /Users/distiller/miniconda3/conda-bld/osx-64

.circleci/config.yml.in

@@ -128,7 +128,7 @@ jobs:
               ca-certificates \
               curl \
               gnupg-agent \
-              software-properties-common
+              software-properties-common 
 
           curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo apt-key add -
 
@@ -219,7 +219,8 @@ jobs:
             curl -o conda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh
             sh conda.sh -b
             source $HOME/miniconda3/bin/activate
-            packaging/build_wheel.sh
+            conda install -yq libpng
+            packaging/build_wheel_macos.sh
       - store_artifacts:
           path: dist
       - persist_to_workspace:
@@ -239,6 +240,7 @@ jobs:
             sh conda.sh -b
             source $HOME/miniconda3/bin/activate
             conda install -yq conda-build
+            conda install -yq libpng
             packaging/build_conda.sh
       - store_artifacts:
           path: /Users/distiller/miniconda3/conda-bld/osx-64

.travis.yml

@@ -34,6 +34,7 @@ matrix:
 
 before_install:
   - sudo apt-get update
+  - sudo apt-get install -y libpng16-16 libpng16-dev
   - wget https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh;
   - bash miniconda.sh -b -p $HOME/miniconda
   - export PATH="$HOME/miniconda/bin:$PATH"

CMakeLists.txt

@@ -11,8 +11,8 @@ 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)
 
-file(GLOB IMAGE_HEADERS torchvision/csrc/image/*.h)
-file(GLOB IMAGE_SOURCES torchvision/csrc/image/*.h torchvision/csrc/image/*.cpp)
+file(GLOB IMAGE_HEADERS torchvision/csrc/image.h)
+file(GLOB IMAGE_SOURCES torchvision/csrc/cpu/image/*.h torchvision/csrc/cpu/image/*.cpp)
 
 add_library(${PROJECT_NAME} SHARED ${MODELS_SOURCES} ${IMAGE_SOURCES})
 target_link_libraries(${PROJECT_NAME} PUBLIC "${PNG_LIBRARY}")

packaging/build_conda.sh

@@ -10,4 +10,5 @@ export SOURCE_ROOT_DIR="$PWD"
 setup_conda_pytorch_constraint
 setup_conda_cudatoolkit_constraint
 setup_visual_studio_constraint
+conda install -yq libpng
 conda build $CONDA_CHANNEL_FLAGS -c defaults -c conda-forge --no-anaconda-upload --python "$PYTHON_VERSION" packaging/torchvision

packaging/build_wheel.sh

@@ -7,6 +7,7 @@ script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 export BUILD_TYPE=wheel
 setup_env 0.5.0
 setup_wheel_python
+yum install -y libpng libpng-devel
 pip_install numpy pyyaml future ninja
 # TODO remove after https://github.com/pytorch/pytorch/pull/27282 gets merged
 pip_install six

packaging/build_wheel_macos.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+set -ex
+
+script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
+. "$script_dir/pkg_helpers.bash"
+
+export BUILD_TYPE=wheel
+setup_env 0.5.0
+setup_wheel_python
+pip_install numpy pyyaml future ninja
+# TODO remove after https://github.com/pytorch/pytorch/pull/27282 gets merged
+pip_install six
+setup_pip_pytorch_version
+python setup.py clean
+IS_WHEEL=1 python setup.py bdist_wheel

packaging/torchvision/meta.yaml

@@ -12,6 +12,7 @@ requirements:
   host:
     - python
     - setuptools
+    - libpng >=1.6.37
     {{ environ.get('CONDA_PYTORCH_BUILD_CONSTRAINT') }}
     {{ environ.get('CONDA_CUDATOOLKIT_CONSTRAINT') }}
     {{ environ.get('CONDA_CPUONLY_FEATURE') }}
@@ -20,6 +21,7 @@ requirements:
     - python
     - pillow >=4.1.1
     - numpy >=1.11
+    - libpng >=1.6.37
     - six
     {{ environ.get('CONDA_PYTORCH_CONSTRAINT') }}
     {{ environ.get('CONDA_CUDATOOLKIT_CONSTRAINT') }}

packaging/wheel/linux_manywheel.sh

@@ -32,7 +32,7 @@ rm -rf vision
 git clone https://github.com/pytorch/vision
 
 cd /tmp/vision
-
+yum install -y libpng libpng-devel
 for PYDIR in "${python_installations[@]}"; do
     export PATH=$PYDIR/bin:$OLD_PATH
     pip install --upgrade pip

setup.py

@@ -83,9 +83,10 @@ def get_extensions():
 
     main_file = glob.glob(os.path.join(extensions_dir, '*.cpp'))
     source_cpu = glob.glob(os.path.join(extensions_dir, 'cpu', '*.cpp'))
+    source_image_cpu = glob.glob(os.path.join(extensions_dir, 'cpu', 'image', '*.cpp'))
     source_cuda = glob.glob(os.path.join(extensions_dir, 'cuda', '*.cu'))
 
-    sources = main_file + source_cpu
+    sources = main_file + source_cpu + source_image_cpu
     extension = CppExtension
 
     compile_cpp_tests = os.getenv('WITH_CPP_MODELS_TEST', '0') == '1'
@@ -142,6 +143,7 @@ def get_extensions():
         extension(
             'torchvision._C',
             sources,
+            libraries=['png'],
             include_dirs=include_dirs,
             define_macros=define_macros,
             extra_compile_args=extra_compile_args,

test/test_image.py

@@ -0,0 +1,37 @@
+import os
+import unittest
+
+import torch
+from PIL import Image
+from torchvision.io.image import read_png, decode_png
+import numpy as np
+
+IMAGE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets", "fakedata", "imagefolder")
+
+
+def get_png_images(directory):
+    assert os.path.isdir(directory)
+    for root, dir, files in os.walk(directory):
+        for fl in files:
+            _, ext = os.path.splitext(fl)
+            if ext == ".png":
+                yield os.path.join(root, fl)
+
+
+class ImageTester(unittest.TestCase):
+    def test_read_png(self):
+        for img_path in get_png_images(IMAGE_DIR):
+            img_pil = torch.from_numpy(np.array(Image.open(img_path)))
+            img_lpng = read_png(img_path)
+            self.assertTrue(torch.all(img_lpng == img_pil))
+
+    def test_decode_png(self):
+        for img_path in get_png_images(IMAGE_DIR):
+            img_pil = torch.from_numpy(np.array(Image.open(img_path)))
+            size = os.path.getsize(img_path)
+            img_lpng = decode_png(torch.from_file(img_path, dtype=torch.uint8, size=size))
+            self.assertTrue(torch.all(img_lpng == img_pil))
+
+
+if __name__ == '__main__':
+    unittest.main()

torchvision/csrc/image/readpng.cpp renamed to torchvision/csrc/cpu/image/readpng_cpu.cpp

@@ -1,48 +1,43 @@
-#include "readpng.h"
-
-#include <sxx/sxx.h>
+#include "readpng_cpu.h"
 
 #include <png.h>
+#include <setjmp.h>
+#include <string>
 
-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;
-      };
-
+torch::Tensor decodePNG(const torch::Tensor& data) {
   auto png_ptr =
       png_create_read_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr);
-  if (!png_ptr)
-    return torch::tensor({0});
-
+  TORCH_CHECK(png_ptr, "libpng read structure allocation failed!")
   auto info_ptr = png_create_info_struct(png_ptr);
   if (!info_ptr) {
     png_destroy_read_struct(&png_ptr, nullptr, nullptr);
-    return torch::tensor({0});
+    // Seems redundant with the if statement. done here to avoid leaking memory.
+    TORCH_CHECK(info_ptr, "libpng info structure allocation failed!")
   }
 
+  auto datap = data.accessor<unsigned char, 1>().data();
+
   if (setjmp(png_jmpbuf(png_ptr)) != 0) {
     png_destroy_read_struct(&png_ptr, &info_ptr, nullptr);
-    return torch::tensor({0});
+    auto errptr = static_cast<png_const_charp>(png_get_error_ptr(png_ptr));
+    TORCH_CHECK(false, "Internal error.");
   }
+  auto is_png = !png_sig_cmp(datap, 0, 8);
+  TORCH_CHECK(is_png, "Content is not png!")
 
-  png_set_read_fn(png_ptr, &reader, read_callback);
+  struct Reader {
+    png_const_bytep ptr;
+  } reader;
+  reader.ptr = png_const_bytep(datap) + 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;
+      };
   png_set_sig_bytes(png_ptr, 8);
+  png_set_read_fn(png_ptr, &reader, read_callback);
   png_read_info(png_ptr, info_ptr);
 
   png_uint_32 width, height;
@@ -58,26 +53,24 @@ torch::Tensor readpng(const void* data) {
       nullptr,
       nullptr);
 
-  if (retval != 1 || color_type != PNG_COLOR_TYPE_RGB) {
+  if (retval != 1) {
+    png_destroy_read_struct(&png_ptr, &info_ptr, nullptr);
+    TORCH_CHECK(retval == 1, "Could read image metadata from content.")
+  }
+  if (color_type != PNG_COLOR_TYPE_RGB) {
     png_destroy_read_struct(&png_ptr, &info_ptr, nullptr);
-    return torch::tensor({0});
+    TORCH_CHECK(
+        color_type == PNG_COLOR_TYPE_RGB, "Non RGB images are not supported.")
   }
 
   auto tensor =
       torch::empty({int64_t(height), int64_t(width), int64_t(3)}, torch::kU8);
-  auto ptr = tensor.data<uint8_t>();
+  auto ptr = tensor.accessor<uint8_t, 3>().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

torchvision/csrc/cpu/image/readpng_cpu.h

@@ -0,0 +1,6 @@
+#pragma once
+
+#include <torch/torch.h>
+#include <string>
+
+torch::Tensor decodePNG(const torch::Tensor& data);

torchvision/csrc/image.h

@@ -0,0 +1,3 @@
+#pragma once
+
+#include "cpu/image/readpng_cpu.h"

torchvision/csrc/image/image.h

@@ -11,6 +11,7 @@
 #include "ROIAlign.h"
 #include "ROIPool.h"
 #include "empty_tensor_op.h"
+#include "image.h"
 #include "nms.h"
 
 // If we are in a Windows environment, we need to define
@@ -49,4 +50,5 @@ static auto registry =
         .op("torchvision::ps_roi_align", &ps_roi_align)
         .op("torchvision::ps_roi_pool", &ps_roi_pool)
         .op("torchvision::deform_conv2d", &deform_conv2d)
+        .op("torchvision::decode_png", &decodePNG)
         .op("torchvision::_cuda_version", &_cuda_version);

torchvision/csrc/image/readpng.h

@@ -0,0 +1,47 @@
+import torch
+from torch import nn, Tensor
+import os
+
+
+def decode_png(input):
+    # type: (Tensor) -> Tensor
+    """
+    Decodes a PNG image into a 3 dimensional RGB Tensor.
+    The values of the output tensor are uint8 between 0 and 255.
+
+    Arguments:
+        input (Tensor[1]): a one dimensional int8 tensor containing
+    the raw bytes of the PNG image.
+
+    Returns:
+        output (Tensor[image_width, image_height, 3])
+    """
+    if not isinstance(input, torch.Tensor) or len(input) == 0:
+        raise ValueError("Expected a non empty 1-dimensional tensor.")
+
+    if not input.dtype == torch.uint8:
+        raise ValueError("Expected a torch.uint8 tensor.")
+    output = torch.ops.torchvision.decode_png(input)
+    return output
+
+
+def read_png(path):
+    # type: (str) -> Tensor
+    """
+    Reads a PNG image into a 3 dimensional RGB Tensor.
+    The values of the output tensor are uint8 between 0 and 255.
+
+    Arguments:
+        path (str): path of the PNG image.
+
+    Returns:
+        output (Tensor[image_width, image_height, 3])
+    """
+    if not os.path.isfile(path):
+        raise ValueError("Excepted a valid file path.")
+
+    size = os.path.getsize(path)
+    if size == 0:
+        raise ValueError("Excepted a non empty file.")
+    data = torch.from_file(path, dtype=torch.uint8, size=size)
+    return decode_png(data)