Add support for png decoding on linux

.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,10 +230,12 @@ 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
             source $HOME/miniconda3/bin/activate
+            conda install -yq cmake
             packaging/build_wheel.sh
       - store_artifacts:
           path: dist
@@ -238,11 +251,12 @@ 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
             source $HOME/miniconda3/bin/activate
-            conda install -yq conda-build
+            conda install -yq conda-build cmake
             packaging/build_conda.sh
       - store_artifacts:
           path: /Users/distiller/miniconda3/conda-bld/osx-64

.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,10 +230,12 @@ 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
             source $HOME/miniconda3/bin/activate
+            conda install -yq cmake
             packaging/build_wheel.sh
       - store_artifacts:
           path: dist
@@ -238,11 +251,12 @@ 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
             source $HOME/miniconda3/bin/activate
-            conda install -yq conda-build
+            conda install -yq conda-build cmake
             packaging/build_conda.sh
       - store_artifacts:
           path: /Users/distiller/miniconda3/conda-bld/osx-64

.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

.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

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

packaging/pkg_helpers.bash

@@ -153,6 +153,7 @@ setup_wheel_python() {
     conda env remove -n "env$PYTHON_VERSION" || true
     conda create -yn "env$PYTHON_VERSION" python="$PYTHON_VERSION"
     conda activate "env$PYTHON_VERSION"
+    export PATH="$PATH:/Users/distiller/miniconda3/bin"
   else
     case "$PYTHON_VERSION" in
       2.7)

packaging/torchvision/meta.yaml

@@ -7,7 +7,8 @@ source:
 
 requirements:
   build:
-    - {{ compiler('c') }} # [win]
+    - {{ compiler('c') }}
+    - {{ compiler('cxx') }}
 
   host:
     - python

setup.py

@@ -8,6 +8,7 @@
 import subprocess
 import distutils.command.clean
 import distutils.spawn
+import multiprocessing
 import glob
 import shutil
 
@@ -83,9 +84,22 @@ 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 = []
+
+    runtime_library_dirs = None
+    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"))
+        runtime_library_dirs = ['lib']
+
     extension = CppExtension
 
     compile_cpp_tests = os.getenv('WITH_CPP_MODELS_TEST', '0') == '1'
@@ -142,9 +156,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=runtime_library_dirs
         )
     ]
     if compile_cpp_tests:
@@ -197,6 +214,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")
+        if sys.platform.startswith("linux"):
+            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 +272,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
 )

test/test_image.py

@@ -0,0 +1,50 @@
+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.assertTrue(torch.equal(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.assertTrue(torch.equal(img_lpng, img_pil))
+
+        with self.assertRaisesRegex(ValueError, "Expected a non empty 1-dimensional tensor."):
+            decode_png(torch.empty((100, 1), dtype=torch.uint8))
+
+        with self.assertRaisesRegex(ValueError, "Expected a torch.uint8 tensor."):
+            decode_png(torch.empty((100, ), dtype=torch.float16))
+
+        with self.assertRaisesRegex(ValueError, "Invalid png input."):
+            decode_png(torch.empty((100), dtype=torch.uint8))
+
+
+if __name__ == '__main__':
+    unittest.main()