read jpeg linux

.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

.circleci/regenerate.py

@@ -21,10 +21,11 @@
 
 def workflows(prefix='', filter_branch=None, upload=False, indentation=6):
     w = []
-    for btype in ["wheel", "conda"]:
-        for os_type in ["linux", "macos"]:
-            for python_version in ["3.5", "3.6", "3.7", "3.8"]:
-                for cu_version in (["cpu", "cu92", "cu100", "cu101"] if os_type == "linux" else ["cpu"]):
+    for btype in ["conda"]:
+        for os_type in ["linux"]:# , "macos"]:
+            for python_version in ["3.5"]:# , "3.6", "3.7", "3.8"]:
+                # for cu_version in (["cpu", "cu92", "cu100", "cu101"] if os_type == "linux" else ["cpu"]):
+                for cu_version in ["cpu"]:
                     for unicode in ([False, True] if btype == "wheel" and python_version == "2.7" else [False]):
                         w += workflow_pair(
                             btype, os_type, python_version, cu_version,

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "third_party/libjpeg-turbo"]
+	path = third_party/libjpeg-turbo
+	url = https://github.com/libjpeg-turbo/libjpeg-turbo

.travis.yml

@@ -30,6 +30,7 @@ before_install:
   - bash miniconda.sh -b -p $HOME/miniconda
   - export PATH="$HOME/miniconda/bin:$PATH"
   - hash -r
+  - sudo apt-get install -y nasm
   - conda config --set always_yes yes --set changeps1 no
   # Useful for debugging any issues with conda
   - conda info -a
@@ -67,8 +68,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

@@ -11,8 +11,13 @@ if(WITH_CUDA)
 endif()
 
 find_package(Python3 COMPONENTS Development)
+
 find_package(Torch REQUIRED)
 
+if(Unix)
+  add_subdirectory("third_party/libjpeg-turbo")
+endif()
+
 file(GLOB HEADERS torchvision/csrc/*.h)
 file(GLOB OPERATOR_SOURCES torchvision/csrc/cpu/*.h torchvision/csrc/cpu/*.cpp torchvision/csrc/*.cpp)
 if(WITH_CUDA)

packaging/build_conda.sh

@@ -11,4 +11,5 @@ setup_conda_pytorch_constraint
 setup_conda_cudatoolkit_constraint
 setup_visual_studio_constraint
 setup_junit_results_folder
+conda install -c anaconda  -yq nasm
 conda build $CONDA_CHANNEL_FLAGS -c defaults -c conda-forge --no-anaconda-upload --python "$PYTHON_VERSION" packaging/torchvision

packaging/torchvision/meta.yaml

@@ -33,6 +33,12 @@ build:
     - NVCC_FLAGS
   features:
     {{ environ.get('CONDA_CPUONLY_FEATURE') }}
+  binary_relocation: False
+  rpaths:
+    - lib/
+  missing_dso_withelist:
+    - '**/libturbojpeg.so.0'
+
 
 test:
   imports:

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 = []
+
+    if sys.platform.startswith('linux'):
+        sources = sources + source_image_cpu
+        libraries.append('turbojpeg')
+        third_party_search_directories.append(os.path.join(cwd, "third_party/libjpeg-turbo"))
+        runtime_library_dirs = ['lib']
+
     extension = CppExtension
 
     compile_cpp_tests = os.getenv('WITH_CPP_MODELS_TEST', '0') == '1'
@@ -94,6 +108,7 @@ def get_extensions():
         models_dir = os.path.join(this_dir, 'torchvision', 'csrc', 'models')
         test_file = glob.glob(os.path.join(test_dir, '*.cpp'))
         source_models = glob.glob(os.path.join(models_dir, '*.cpp'))
+        extra_compile_args.setdefault('cxx', [])
 
         test_file = [os.path.join(test_dir, s) for s in test_file]
         source_models = [os.path.join(models_dir, s) for s in source_models]
@@ -142,9 +157,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 +215,42 @@ 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/libjpeg-turbo/")
+        throw_of_failure('cmake -DENABLE_STATIC=OFF -DREQUIRE_SIMD=TRUE  .')
+        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'):
+#if os.environ.get('BUILD_TYPE', "") == "conda":
+    data_files = [
+        ('lib/python3.5/site-packages/torchvision/lib', [
+            'third_party/libjpeg-turbo/libturbojpeg.so.0',
+            'third_party/libjpeg-turbo/libturbojpeg.so'])
+    ]
+    # else:
+        # data_files = [
+            # ('torchvision/lib', [
+                # 'third_party/libjpeg-turbo/libturbojpeg.so.0',
+                # 'third_party/libjpeg-turbo/libturbojpeg.so'])
+        # ]
+
 setup(
     # Metadata
     name=package_name,
@@ -218,7 +272,7 @@ 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,57 @@
+import os
+import unittest
+import sys
+
+import torch
+import torchvision
+if sys.platform.startswith('linux'):
+    from torchvision.io.image import read_jpeg, decode_jpeg
+    from PIL import Image
+import numpy as np
+
+
+IMAGE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets")
+
+
+def get_images(directory, img_ext):
+    assert os.path.isdir(directory)
+    for root, _, 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_jpeg(self):
+        for img_path in get_images(IMAGE_DIR, ".jpg"):
+            img_pil = torch.from_numpy(np.array(Image.open(img_path)))
+            img_ljpeg = read_jpeg(img_path)
+
+            err = torch.abs(img_ljpeg.flatten().float() - img_pil.flatten().float()).sum().float() / (img_ljpeg.shape[0] * img_ljpeg.shape[1] * img_ljpeg.shape[2] * 255)
+            self.assertLessEqual(err, 1e-2)
+
+    @unittest.skipUnless(sys.platform.startswith("linux"), "Support only available on linux for now.")
+    def test_decode_jpeg(self):
+        for img_path in get_images(IMAGE_DIR, ".jpg"):
+            img_pil = torch.from_numpy(np.array(Image.open(img_path)))
+            size = os.path.getsize(img_path)
+            img_ljpeg = decode_jpeg(torch.from_file(img_path, dtype=torch.uint8, size=size))
+
+            err = torch.abs(img_ljpeg.flatten().float() - img_pil.flatten().float()).sum().float() / (img_ljpeg.shape[0] * img_ljpeg.shape[1] * img_ljpeg.shape[2] * 255)
+
+            self.assertLessEqual(err, 1e-2)
+
+        with self.assertRaisesRegex(ValueError, "Expected a non empty 1-dimensional tensor."):
+            decode_jpeg(torch.empty((100, 1), dtype=torch.uint8))
+
+        with self.assertRaisesRegex(ValueError, "Expected a torch.uint8 tensor."):
+            decode_jpeg(torch.empty((100, ), dtype=torch.float16))
+
+        with self.assertRaisesRegex(ValueError, "Invalid jpeg input."):
+            decode_jpeg(torch.empty((100), dtype=torch.uint8))
+
+
+if __name__ == '__main__':
+    unittest.main()

torchvision/csrc/cpu/image/readjpeg_cpu.cpp

@@ -0,0 +1,37 @@
+#include "readjpeg_cpu.h"
+
+#include <setjmp.h>
+#include <string>
+#include <turbojpeg.h>
+
+torch::Tensor decodeJPEG(const torch::Tensor& data) {
+
+  tjhandle tjInstance = tjInitDecompress();
+  if (tjInstance == NULL) {
+    TORCH_CHECK(false, "libjpeg-turbo decompression initialization failed.");
+  }
+
+  auto datap = data.accessor<unsigned char, 1>().data();
+
+  int width, height;
+
+  if (tjDecompressHeader(tjInstance, datap, data.numel(), &width, &height) < 0) {
+    tjDestroy(tjInstance);
+    TORCH_CHECK(false, "Error while reading jpeg headers");
+  }
+  auto tensor =
+      torch::empty({int64_t(height), int64_t(width), int64_t(3)}, torch::kU8);
+
+  auto ptr = tensor.accessor<uint8_t, 3>().data();
+
+  int pixelFormat = TJPF_RGB;
+
+  auto ret = tjDecompress2(tjInstance, datap, data.numel(), ptr, width, 0, height,
+                      pixelFormat, NULL);
+  if(ret != 0){
+      tjDestroy(tjInstance);
+      TORCH_CHECK(false, "decompressing JPEG image");
+  }
+
+  return tensor;
+}

torchvision/csrc/cpu/image/readjpeg_cpu.h

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

torchvision/csrc/image.h

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

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_jpeg", &decodeJPEG)
+#endif
         .op("torchvision::_cuda_version", &_cuda_version);