diff --git a/.circleci/config.yml b/.circleci/config.yml index 5bdb0e69e72..1ff6aeaac78 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -128,7 +128,9 @@ jobs: ca-certificates \ curl \ gnupg-agent \ - software-properties-common + software-properties-common \ + libpng16-16 \ + libpng-dev curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo apt-key add - @@ -219,7 +221,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 +242,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 diff --git a/.circleci/config.yml.in b/.circleci/config.yml.in index e3747134c6f..7f72519d7da 100644 --- a/.circleci/config.yml.in +++ b/.circleci/config.yml.in @@ -128,7 +128,9 @@ jobs: ca-certificates \ curl \ gnupg-agent \ - software-properties-common + software-properties-common \ + libpng16-16 \ + libpng-dev curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo apt-key add - @@ -219,7 +221,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 +242,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 diff --git a/.travis.yml b/.travis.yml index 3e667b6fc1e..d9ddbf6c2ee 100644 --- a/.travis.yml +++ b/.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" diff --git a/CMakeLists.txt b/CMakeLists.txt index 3e4b065e83e..16727531ef8 100644 --- a/CMakeLists.txt +++ b/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}") diff --git a/packaging/build_wheel.sh b/packaging/build_wheel.sh index 7d37239563d..4f986534243 100755 --- a/packaging/build_wheel.sh +++ b/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 diff --git a/packaging/build_wheel_macos.sh b/packaging/build_wheel_macos.sh new file mode 100755 index 00000000000..7d37239563d --- /dev/null +++ b/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 diff --git a/packaging/torchvision/meta.yaml b/packaging/torchvision/meta.yaml index da075ff03cb..c7c3fd3d4ed 100644 --- a/packaging/torchvision/meta.yaml +++ b/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') }} diff --git a/packaging/wheel/linux_manywheel.sh b/packaging/wheel/linux_manywheel.sh index d04e334d237..a3bb22d7f88 100644 --- a/packaging/wheel/linux_manywheel.sh +++ b/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 diff --git a/setup.py b/setup.py index 8ece63ce739..a01bd120383 100644 --- a/setup.py +++ b/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, diff --git a/test/test_image.py b/test/test_image.py new file mode 100644 index 00000000000..35dfb1884ba --- /dev/null +++ b/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() diff --git a/torchvision/csrc/image/readpng.cpp b/torchvision/csrc/cpu/image/readpng_cpu.cpp similarity index 59% rename from torchvision/csrc/image/readpng.cpp rename to torchvision/csrc/cpu/image/readpng_cpu.cpp index d589d60912e..a395e5ed663 100644 --- a/torchvision/csrc/image/readpng.cpp +++ b/torchvision/csrc/cpu/image/readpng_cpu.cpp @@ -1,48 +1,43 @@ -#include "readpng.h" - -#include +#include "readpng_cpu.h" #include +#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; - }; - +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().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_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(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(); + auto ptr = tensor.accessor().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/cpu/image/readpng_cpu.h b/torchvision/csrc/cpu/image/readpng_cpu.h new file mode 100644 index 00000000000..d2151a43aa9 --- /dev/null +++ b/torchvision/csrc/cpu/image/readpng_cpu.h @@ -0,0 +1,6 @@ +#pragma once + +#include +#include + +torch::Tensor decodePNG(const torch::Tensor& data); diff --git a/torchvision/csrc/image.h b/torchvision/csrc/image.h new file mode 100644 index 00000000000..01d2063564d --- /dev/null +++ b/torchvision/csrc/image.h @@ -0,0 +1,3 @@ +#pragma once + +#include "cpu/image/readpng_cpu.h" diff --git a/torchvision/csrc/image/image.h b/torchvision/csrc/image/image.h deleted file mode 100644 index f3b6973ecf7..00000000000 --- a/torchvision/csrc/image/image.h +++ /dev/null @@ -1,6 +0,0 @@ -#ifndef IMAGE_H -#define IMAGE_H - -#include "readpng.h" - -#endif // IMAGE_H diff --git a/torchvision/csrc/image/readpng.h b/torchvision/csrc/image/readpng.h deleted file mode 100644 index 135f07e3f47..00000000000 --- a/torchvision/csrc/image/readpng.h +++ /dev/null @@ -1,19 +0,0 @@ -#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 diff --git a/torchvision/csrc/vision.cpp b/torchvision/csrc/vision.cpp index b761dc88710..d785d4ec5c6 100644 --- a/torchvision/csrc/vision.cpp +++ b/torchvision/csrc/vision.cpp @@ -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); diff --git a/torchvision/io/image.py b/torchvision/io/image.py new file mode 100644 index 00000000000..ded604fdd60 --- /dev/null +++ b/torchvision/io/image.py @@ -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)