Add support for png decoding on linux

- Pipelines are timing out, increasing to 30 minutes. - Vendors zlib and add's libpng as a dependency. - Ship libpng.so and zlib.so with package on linux
pytorch · Feb 17, 2020 · 7e94649 · 7e94649
1 parent 6c2cda6
commit 7e94649
Show file tree

Hide file tree

Showing 14 changed files with 293 additions and 11 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
diff --git a/third_party/zlib b/third_party/zlib