diff --git a/.circleci/config.yml b/.circleci/config.yml index 364cdd5f0b9..9668b59915f 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -21,6 +21,9 @@ commands: description: "checkout merge branch" steps: - checkout + - run: + name: initialize submodules + command: git submodule update --init --recursive # - run: # name: Checkout merge branch # command: | @@ -83,7 +86,9 @@ jobs: resource_class: 2xlarge+ steps: - checkout_merge - - run: packaging/build_wheel.sh + - run: + command: packaging/build_wheel.sh + no_output_timeout: 30m - store_artifacts: path: dist - persist_to_workspace: @@ -98,7 +103,10 @@ jobs: resource_class: 2xlarge+ steps: - checkout_merge - - run: packaging/build_conda.sh + - run: + command: packaging/build_conda.sh + no_output_timeout: 30m + - store_artifacts: path: /opt/conda/conda-bld/linux-64 - persist_to_workspace: @@ -167,6 +175,7 @@ jobs: - run: name: Build and run tests + no_output_timeout: 30m command: | set -e @@ -191,6 +200,7 @@ jobs: conda activate base conda install -yq conda-build "conda-package-handling!=1.5.0" bash packaging/build_conda.sh + no_output_timeout: 30m shell: powershell.exe - store_test_results: path: build_results/ @@ -207,6 +217,7 @@ jobs: conda activate base conda install -yq conda-build "conda-package-handling!=1.5.0" bash packaging/build_conda.sh + no_output_timeout: 30m shell: powershell.exe binary_macos_wheel: @@ -219,6 +230,7 @@ jobs: # Cannot easily deduplicate this as source'ing activate # will set environment variables which we need to propagate # to build_wheel.sh + no_output_timeout: 30m command: | curl -o conda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh sh conda.sh -b @@ -238,6 +250,7 @@ jobs: steps: - checkout_merge - run: + no_output_timeout: 30m command: | curl -o conda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh sh conda.sh -b diff --git a/.circleci/config.yml.in b/.circleci/config.yml.in index de5ec111684..6517daaca37 100644 --- a/.circleci/config.yml.in +++ b/.circleci/config.yml.in @@ -21,6 +21,9 @@ commands: description: "checkout merge branch" steps: - checkout + - run: + name: initialize submodules + command: git submodule update --init --recursive # - run: # name: Checkout merge branch # command: | @@ -83,7 +86,9 @@ jobs: resource_class: 2xlarge+ steps: - checkout_merge - - run: packaging/build_wheel.sh + - run: + command: packaging/build_wheel.sh + no_output_timeout: 30m - store_artifacts: path: dist - persist_to_workspace: @@ -98,7 +103,10 @@ jobs: resource_class: 2xlarge+ steps: - checkout_merge - - run: packaging/build_conda.sh + - run: + command: packaging/build_conda.sh + no_output_timeout: 30m + - store_artifacts: path: /opt/conda/conda-bld/linux-64 - persist_to_workspace: @@ -167,6 +175,7 @@ jobs: - run: name: Build and run tests + no_output_timeout: 30m command: | set -e @@ -191,6 +200,7 @@ jobs: conda activate base conda install -yq conda-build "conda-package-handling!=1.5.0" bash packaging/build_conda.sh + no_output_timeout: 30m shell: powershell.exe - store_test_results: path: build_results/ @@ -207,6 +217,7 @@ jobs: conda activate base conda install -yq conda-build "conda-package-handling!=1.5.0" bash packaging/build_conda.sh + no_output_timeout: 30m shell: powershell.exe binary_macos_wheel: @@ -219,6 +230,7 @@ jobs: # Cannot easily deduplicate this as source'ing activate # will set environment variables which we need to propagate # to build_wheel.sh + no_output_timeout: 30m command: | curl -o conda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh sh conda.sh -b @@ -238,6 +250,7 @@ jobs: steps: - checkout_merge - run: + no_output_timeout: 30m command: | curl -o conda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh sh conda.sh -b diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 00000000000..55b372801aa --- /dev/null +++ b/.gitmodules @@ -0,0 +1,6 @@ +[submodule "third_party/libpng"] + path = third_party/libpng + url = https://github.com/glennrp/libpng +[submodule "third_party/zlib"] + path = third_party/zlib + url = https://github.com/madler/zlib diff --git a/.travis.yml b/.travis.yml index 1b6ecb7a65b..51b518833ea 100644 --- a/.travis.yml +++ b/.travis.yml @@ -67,8 +67,8 @@ install: cd - script: - - pytest --cov-config .coveragerc --cov torchvision --cov $TV_INSTALL_PATH -k 'not TestVideoReader and not TestVideoTransforms' test - - pytest test/test_hub.py + - travis_wait 60 pytest --cov-config .coveragerc --cov torchvision --cov $TV_INSTALL_PATH -k 'not TestVideoReader and not TestVideoTransforms' test + - travis_wait 60 pytest test/test_hub.py after_success: # Necessary to run coverage combine to rewrite paths from diff --git a/CMakeLists.txt b/CMakeLists.txt index fcedf41e5ab..fd4baaa96e1 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -10,6 +10,10 @@ if(WITH_CUDA) add_definitions(-D__CUDA_NO_HALF_OPERATORS__) endif() +if(Unix) + add_subdirectory("third_party/libpng") +endif() + find_package(Python3 COMPONENTS Development) find_package(Torch REQUIRED) @@ -21,8 +25,17 @@ endif() 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} ${OPERATOR_SOURCES}) -target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES} Python3::Python) +file(GLOB IMAGE_HEADERS torchvision/csrc/image.h) +file(GLOB IMAGE_SOURCES torchvision/csrc/cpu/image/*.h torchvision/csrc/cpu/image/*.cpp) + +if(Unix) + add_library(${PROJECT_NAME} SHARED ${MODELS_SOURCES} ${OPERATOR_SOURCES} {IMAGE_SOURCES}) + target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES} Python3::Python "${PNG_LIBRARIES}") +else() + add_library(${PROJECT_NAME} SHARED ${MODELS_SOURCES} ${OPERATOR_SOURCES}) + target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES} Python3::Python) +endif() + set_target_properties(${PROJECT_NAME} PROPERTIES EXPORT_NAME TorchVision) target_include_directories(${PROJECT_NAME} INTERFACE diff --git a/setup.py b/setup.py index 71d420573ed..285e9e952cc 100644 --- a/setup.py +++ b/setup.py @@ -8,6 +8,7 @@ import subprocess import distutils.command.clean import distutils.spawn +import multiprocessing import glob import shutil @@ -83,9 +84,20 @@ 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 + + libraries = [] + extra_compile_args = {} + third_party_search_directories = [] + + if sys.platform.startswith('linux'): + sources = sources + source_image_cpu + libraries.append('png') + third_party_search_directories.append(os.path.join(cwd, "third_party/libpng")) + extension = CppExtension compile_cpp_tests = os.getenv('WITH_CPP_MODELS_TEST', '0') == '1' @@ -142,9 +154,12 @@ def get_extensions(): extension( 'torchvision._C', sources, - include_dirs=include_dirs, + libraries=libraries, + library_dirs=third_party_search_directories, + include_dirs=include_dirs + third_party_search_directories, define_macros=define_macros, extra_compile_args=extra_compile_args, + runtime_library_dirs=['lib'] ) ] if compile_cpp_tests: @@ -197,6 +212,43 @@ def run(self): distutils.command.clean.clean.run(self) +def throw_of_failure(command): + ret = os.system(command) + if ret != 0: + raise Exception("{} failed".format(command)) + + +def build_deps(): + this_dir = os.path.dirname(os.path.abspath(__file__)) + if sys.platform.startswith('linux'): + cpu_count = multiprocessing.cpu_count() + os.chdir("third_party/zlib/") + throw_of_failure('cmake .') + throw_of_failure("cmake --build . -- -j {}".format(cpu_count)) + os.chdir(this_dir) + + zlib_path = os.path.join(this_dir, "third_party/zlib") + libpng_cmake_options = "-DPNG_BUILD_ZLIB=ON -DPNG_STATIC=OFF -DZLIB_INCLUDE_DIR:PATH={zlib_path} -DZLIB_LIBRARY:FILEPATH={zlib_path}/libz.so".format(zlib_path=zlib_path) + os.chdir("third_party/libpng/") + os.system('cmake {} .'.format(libpng_cmake_options)) + throw_of_failure("cmake --build . -- -j {}".format(cpu_count)) + os.chdir(this_dir) + + +def build_ext_with_dependencies(self): + build_deps() + return BuildExtension.with_options(no_python_abi_suffix=True)(self) + + +data_files = [] +if sys.platform.startswith('linux'): + data_files = [ + ('torchvision/lib', [ + 'third_party/zlib/libz.so', + 'third_party/libpng/libpng.so']) + ] + + setup( # Metadata name=package_name, @@ -218,7 +270,8 @@ def run(self): }, ext_modules=get_extensions(), cmdclass={ - 'build_ext': BuildExtension.with_options(no_python_abi_suffix=True), + 'build_ext': build_ext_with_dependencies, 'clean': clean, - } + }, + data_files=data_files ) diff --git a/test/test_image.py b/test/test_image.py new file mode 100644 index 00000000000..2dee7912aa7 --- /dev/null +++ b/test/test_image.py @@ -0,0 +1,44 @@ +import os +import unittest +import sys + +import torch +from PIL import Image +if sys.platform.startswith('linux'): + 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_images(directory, img_ext): + assert os.path.isdir(directory) + for root, dir, files in os.walk(directory): + for fl in files: + _, ext = os.path.splitext(fl) + if ext == img_ext: + yield os.path.join(root, fl) + + +class ImageTester(unittest.TestCase): + @unittest.skipUnless(sys.platform.startswith("linux"), "Support only available on linux for now.") + def test_read_png(self): + for img_path in get_images(IMAGE_DIR, "png"): + img_pil = torch.from_numpy(np.array(Image.open(img_path))) + img_lpng = read_png(img_path) + self.assertEqual(img_lpng, img_pil) + + @unittest.skipUnless(sys.platform.startswith("linux"), "Support only available on linux for now.") + def test_decode_png(self): + for img_path in get_images(IMAGE_DIR, "png"): + 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.assertEqual(img_lpng, img_pil) + + self.assertEqual(decode_png(torch.empty()), torch.empty()) + self.assertEqual(decode_png(torch.randint(3, 5, (300,))), torch.empty()) + + +if __name__ == '__main__': + unittest.main() diff --git a/third_party/libpng b/third_party/libpng new file mode 160000 index 00000000000..a40189cf881 --- /dev/null +++ b/third_party/libpng @@ -0,0 +1 @@ +Subproject commit a40189cf881e9f0db80511c382292a5604c3c3d1 diff --git a/third_party/zlib b/third_party/zlib new file mode 160000 index 00000000000..cacf7f1d4e3 --- /dev/null +++ b/third_party/zlib @@ -0,0 +1 @@ +Subproject commit cacf7f1d4e3d44d871b605da3b647f07d718623f diff --git a/torchvision/csrc/cpu/image/readpng_cpu.cpp b/torchvision/csrc/cpu/image/readpng_cpu.cpp new file mode 100644 index 00000000000..c6581f168b1 --- /dev/null +++ b/torchvision/csrc/cpu/image/readpng_cpu.cpp @@ -0,0 +1,75 @@ +#include "readpng_cpu.h" + +#include +#include +#include + +torch::Tensor decodePNG(const torch::Tensor& data) { + auto png_ptr = + png_create_read_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr); + 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); + // 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); + TORCH_CHECK(false, "Internal error."); + } + auto is_png = !png_sig_cmp(datap, 0, 8); + TORCH_CHECK(is_png, "Content is not png!") + + 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; + 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) { + 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); + 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.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; +} 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/vision.cpp b/torchvision/csrc/vision.cpp index 8d8699ecc26..baa439e0d9a 100644 --- a/torchvision/csrc/vision.cpp +++ b/torchvision/csrc/vision.cpp @@ -11,6 +11,9 @@ #include "ROIAlign.h" #include "ROIPool.h" #include "empty_tensor_op.h" +#ifdef __linux__ +#include "image.h" +#endif #include "nms.h" // If we are in a Windows environment, we need to define @@ -49,4 +52,7 @@ 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) +#ifdef __linux__ + .op("torchvision::decode_png", &decodePNG) +#endif .op("torchvision::_cuda_version", &_cuda_version); diff --git a/torchvision/io/image.py b/torchvision/io/image.py new file mode 100644 index 00000000000..bbf7470b097 --- /dev/null +++ b/torchvision/io/image.py @@ -0,0 +1,48 @@ +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("Expected a valid file path.") + + size = os.path.getsize(path) + if size == 0: + raise ValueError("Expected a non empty file.") + data = torch.from_file(path, dtype=torch.uint8, size=size) + return decode_png(data) +