PR: Read JPEG images directly (#2388)

andfoy · r-zenine · web-flow · commit 4246abc897aa · 2020-07-08T19:58:56.000+02:00
* Add libpng requirement into conda recipe

* Try to install libjpeg-turbo

* Add PNG reading capabilities

* Remove newline

* Add image extension to compilation instructions

* Include png functions as part of the main library

* Update CMakeLists

* Detect if building on conda-build

* Debug

* More debug messages

* Print globbed libreries

* Print globbed libreries

* Point to correct PNG path

* Remove libJPEG preventively

* Debug extension loading

* Link libpng explicitly

* Link with PNG

* Add PNG reading capabilities

* Add libpng requirement into conda recipe

* Try to install libjpeg-turbo

* Remove newline

* Add image extension to compilation instructions

* Include png functions as part of the main library

* Update CMakeLists

* Detect if building on conda-build

* Debug

* More debug messages

* Print globbed libreries

* Print globbed libreries

* Point to correct PNG path

* Remove libJPEG preventively

* Debug extension loading

* Link libpng explicitly

* Link with PNG

* Install libpng on conda-based wheel distributions

* Add -y flag

* Add -y flag to yum

* Locate LibPNG on windows conda

* Remove empty else

* Copy libpng16.so

* Copy dylib on Mac

* Improve check on Windows

* Try to install ninja using conda on windows

* Use libpng on Windows

* Package lib on windows wheel

* Point library to the correct place

* Include binaries as part of wheel

* Copy libpng.so on linux

* Look for png.h on Windows when using conda-build

* Do not skip png tests on Mac/Win

* Restore libjpeg-turbo

* Install jpeg-turbo on wheel distributions

* Install libjpeg-turbo from conda-forge on wheel distributions

* Do not pull av on conda-build

* Add pillow disclaimer

* Vendors libjpeg-turbo 2.0.4

* Merge JPEG work

* Remove submodules

* Regenerate circle config

* Fix style issues

* Fix C++ style issues

* More style corrections

* Add JPEG-turbo to linking libraries

* More style corrections

* More style corrections

* More style corrections

* Install libjpeg-turbo-devel

* Install libturbo-jpeg on typing pipeline

* Update Circle template

* Windows and Unix turbojpeg have the same linking name

* Install turbojpeg-devel instead of libjpeg-turbo

* Copy TurboJPEG binaries to wheel

* Move test image

* Move back test image

* Update JPEG test path

* Remove dot from extension

* Move image functions to extension

* Use stdout arg in subprocess

* Disable image extension if libpng or turbojpeg are not found

* Append libpng stdout

* Prevent list appending on lists

* Minor path correction

* Minor error correction

* Add linking flags

* Style issues correction

* Address minor review corrections

* Refactor library search

* Restore access index

* Fix JPEG tests

* Update libpng version in Travis

* Add -y flag

* Remove dot

* Update libpng using apt

* Check libpng version

* Change libturbojpeg binary

* Update import

* Change call

* Restore av in conda recipe

* Minor error correction

* Remove unused comment in travis.yml

* Update README

* Fix missing links

* Remove fixes for 16.04

* Enable JPEG support using libjpeg directly

* Install libjpeg-turbo8 on Travis

* Fix styling issues

* Do not append to paths if library found in standard library locations

* Add macro flag

* Use custom error handler

* Use 3?

* Use short-lived buffer

* Return TRUE instead of true

* Assert RuntimeError

* Use .jpg extension

* Remove conda-forge

* Use data_ptr instead of accessor

* Use assertTrue for jpeg verification

* Remove unnecessary memcpy

* Debug test on Windows

* Remove PIL from jpeg verification

Co-authored-by: Ryad ZENINE &lt;r.zenine@gmail.com&gt;
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -107,6 +107,8 @@ jobs:
       - checkout
       - run:
           command: |
+            sudo apt-get update -y
+            sudo apt install -y libturbojpeg-dev
             pip install --user --progress-bar off numpy mypy
             pip install --user --progress-bar off --pre torch -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
             pip install --user --progress-bar off --editable .
diff --git a/.circleci/config.yml.in b/.circleci/config.yml.in
@@ -107,6 +107,8 @@ jobs:
       - checkout
       - run:
           command: |
+            sudo apt-get update -y
+            sudo apt install -y libturbojpeg-dev
             pip install --user --progress-bar off numpy mypy
             pip install --user --progress-bar off --pre torch -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
             pip install --user --progress-bar off --editable .
diff --git a/.circleci/unittest/linux/scripts/environment.yml b/.circleci/unittest/linux/scripts/environment.yml
@@ -7,6 +7,7 @@ dependencies:
   - codecov
   - pip
   - libpng
+  - jpeg
   - ca-certificates
   - pip:
     - future
diff --git a/.circleci/unittest/windows/scripts/environment.yml b/.circleci/unittest/windows/scripts/environment.yml
@@ -7,6 +7,7 @@ dependencies:
   - codecov
   - pip
   - libpng
+  - jpeg
   - ca-certificates
   - pip:
     - future
diff --git a/.travis.yml b/.travis.yml
@@ -13,7 +13,7 @@ jobs:
 
 before_install:
   - sudo apt-get update
-  - sudo apt-get install -y libpng-dev
+  - sudo apt-get install -y libpng-dev libjpeg-turbo8-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
@@ -14,6 +14,7 @@ find_package(Python3 COMPONENTS Development)
 
 find_package(Torch REQUIRED)
 find_package(PNG REQUIRED)
+find_package(JPEG REQUIRED)
 
 
 file(GLOB HEADERS torchvision/csrc/*.h)
@@ -28,12 +29,12 @@ 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} ${IMAGE_SOURCES})
-target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES} ${PNG_LIBRARY} Python3::Python)
+target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES} ${PNG_LIBRARY} ${JPEG_LIBRARIES} Python3::Python)
 # target_link_libraries(${PROJECT_NAME} PRIVATE ${PNG_LIBRARY} Python3::Python)
 set_target_properties(${PROJECT_NAME} PROPERTIES EXPORT_NAME TorchVision)
 
 target_include_directories(${PROJECT_NAME} INTERFACE
-  $<BUILD_INTERFACE:${HEADERS}:${PNG_INCLUDE_DIR}>
+  $<BUILD_INTERFACE:${HEADERS}:${PNG_INCLUDE_DIR}:${JPEG_INCLUDE_DIRS}>
   $<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>)
 
 include(GNUInstallDirs)
diff --git a/README.rst b/README.rst
@@ -80,13 +80,17 @@ Torchvision currently supports the following image backends:
 
 * `libpng`_ - can be installed via conda :code:`conda install libpng` or any of the package managers for debian-based and RHEL-based Linux distributions.
 
-**Notes:** ``libpng`` must be available at compilation time in order to be available. Make sure that it is available on the standard library locations,
+* `libjpeg`_ - can be installed via conda :code:`conda install jpeg` or any of the package managers for debian-based and RHEL-based Linux distributions. `libjpeg-turbo`_ can be used as well.
+
+**Notes:** ``libpng`` and ``libjpeg`` must be available at compilation time in order to be available. Make sure that it is available on the standard library locations,
 otherwise, add the include and library paths in the environment variables ``TORCHVISION_INCLUDE`` and ``TORCHVISION_LIBRARY``, respectively.
 
 .. _libpng : http://www.libpng.org/pub/png/libpng.html
 .. _Pillow : https://python-pillow.org/
 .. _Pillow-SIMD : https://github.com/uploadcare/pillow-simd
 .. _accimage: https://github.com/pytorch/accimage
+.. _libjpeg: http://ijg.org/
+.. _libjpeg-turbo: https://libjpeg-turbo.org/
 
 C++ API
 =======
diff --git a/packaging/build_wheel.sh b/packaging/build_wheel.sh
@@ -19,12 +19,17 @@ if [[ "$(uname)" == Darwin || "$OSTYPE" == "msys" ]]; then
     if [[ "$(uname)" == Darwin ]]; then
         # Include LibPNG
         cp "$env_path/lib/libpng16.dylib" torchvision
+        # Include LibJPEG
+        cp "$env_path/lib/libjpeg.dylib" torchvision
     else
         cp "$bin_path/Library/bin/libpng16.dll" torchvision
+        cp "$bin_path/Library/bin/libjpeg.dll" torchvision
     fi
 else
     # Include LibPNG
     cp "/usr/lib64/libpng.so" torchvision
+    # Include LibJPEG
+    cp "/usr/lib64/libjpeg.so" torchvision
 fi
 
 if [[ "$OSTYPE" == "msys" ]]; then
diff --git a/packaging/pkg_helpers.bash b/packaging/pkg_helpers.bash
@@ -171,10 +171,10 @@ setup_wheel_python() {
     conda create -yn "env$PYTHON_VERSION" python="$PYTHON_VERSION"
     conda activate "env$PYTHON_VERSION"
     # Install libpng from Anaconda (defaults)
-    conda install libpng -y
+    conda install libpng jpeg -y
   else
     # Install native CentOS libPNG
-    yum install -y libpng-devel
+    yum install -y libpng-devel libjpeg-turbo-devel
     case "$PYTHON_VERSION" in
       2.7)
         if [[ -n "$UNICODE_ABI" ]]; then
diff --git a/packaging/torchvision/meta.yaml b/packaging/torchvision/meta.yaml
@@ -9,6 +9,7 @@ requirements:
   build:
     - {{ compiler('c') }} # [win]
     - libpng
+    - jpeg
 
   host:
     - python
@@ -20,6 +21,7 @@ requirements:
   run:
     - python
     - libpng
+    - jpeg
     - pillow >=4.1.1
     - numpy >=1.11
     {{ environ.get('CONDA_PYTORCH_CONSTRAINT') }}
diff --git a/setup.py b/setup.py
@@ -327,10 +327,23 @@ def get_extensions():
             image_include += [png_include]
             image_link_flags.append('libpng')
 
+    # Locating libjpeg
+    (jpeg_found, jpeg_conda,
+     jpeg_include, jpeg_lib) = find_library('jpeglib', vision_include)
+
+    print('JPEG found: {0}'.format(jpeg_found))
+    image_macros += [('JPEG_FOUND', str(int(jpeg_found)))]
+    if jpeg_found:
+        print('Building torchvision with JPEG image support')
+        image_link_flags.append('jpeg')
+        if jpeg_conda:
+            image_library += [jpeg_lib]
+            image_include += [jpeg_include]
+
     image_path = os.path.join(extensions_dir, 'cpu', 'image')
     image_src = glob.glob(os.path.join(image_path, '*.cpp'))
 
-    if png_found:
+    if png_found or jpeg_found:
         ext_modules.append(extension(
             'torchvision.image',
             image_src,
diff --git a/test/assets/grace_hopper_517x606.pth b/test/assets/grace_hopper_517x606.pth
diff --git a/test/test_image.py b/test/test_image.py
@@ -5,7 +5,7 @@
 import torch
 import torchvision
 from PIL import Image
-from torchvision.io.image import read_png, decode_png
+from torchvision.io.image import read_png, decode_png, read_jpeg, decode_jpeg
 import numpy as np
 
 IMAGE_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets")
@@ -22,6 +22,28 @@ def get_images(directory, img_ext):
 
 
 class ImageTester(unittest.TestCase):
+    def test_read_jpeg(self):
+        for img_path in get_images(IMAGE_ROOT, ".jpg"):
+            img_pil = torch.load(img_path.replace('jpg', 'pth'))
+            img_ljpeg = read_jpeg(img_path)
+            self.assertTrue(img_ljpeg.equal(img_pil))
+
+    def test_decode_jpeg(self):
+        for img_path in get_images(IMAGE_ROOT, ".jpg"):
+            img_pil = torch.load(img_path.replace('jpg', 'pth'))
+            size = os.path.getsize(img_path)
+            img_ljpeg = decode_jpeg(torch.from_file(img_path, dtype=torch.uint8, size=size))
+            self.assertTrue(img_ljpeg.equal(img_pil))
+
+        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.assertRaises(RuntimeError):
+            decode_jpeg(torch.empty((100), dtype=torch.uint8))
+
     def test_read_png(self):
         # Check across .png
         for img_path in get_images(IMAGE_DIR, ".png"):
diff --git a/torchvision/csrc/cpu/image/image.cpp b/torchvision/csrc/cpu/image/image.cpp
@@ -12,5 +12,6 @@ PyMODINIT_FUNC PyInit_image(void) {
 }
 #endif
 
-static auto registry =
-    torch::RegisterOperators().op("image::decode_png", &decodePNG);
+static auto registry = torch::RegisterOperators()
+                           .op("image::decode_png", &decodePNG)
+                           .op("image::decode_jpeg", &decodeJPEG);
diff --git a/torchvision/csrc/cpu/image/image.h b/torchvision/csrc/cpu/image/image.h
@@ -4,4 +4,5 @@
 // Comment
 #include <torch/script.h>
 #include <torch/torch.h>
+#include "readjpeg_cpu.h"
 #include "readpng_cpu.h"
diff --git a/torchvision/csrc/cpu/image/readjpeg_cpu.cpp b/torchvision/csrc/cpu/image/readjpeg_cpu.cpp
@@ -0,0 +1,140 @@
+#include "readjpeg_cpu.h"
+
+#include <ATen/ATen.h>
+#include <setjmp.h>
+#include <string>
+
+#if !JPEG_FOUND
+
+torch::Tensor decodeJPEG(const torch::Tensor& data) {
+  AT_ERROR("decodeJPEG: torchvision not compiled with libjpeg support");
+}
+
+#else
+#include <jpeglib.h>
+
+const static JOCTET EOI_BUFFER[1] = {JPEG_EOI};
+char jpegLastErrorMsg[JMSG_LENGTH_MAX];
+
+struct torch_jpeg_error_mgr {
+  struct jpeg_error_mgr pub; /* "public" fields */
+  jmp_buf setjmp_buffer; /* for return to caller */
+};
+
+typedef struct torch_jpeg_error_mgr* torch_jpeg_error_ptr;
+
+void torch_jpeg_error_exit(j_common_ptr cinfo) {
+  /* cinfo->err really points to a torch_jpeg_error_mgr struct, so coerce
+   * pointer */
+  torch_jpeg_error_ptr myerr = (torch_jpeg_error_ptr)cinfo->err;
+
+  /* Always display the message. */
+  /* We could postpone this until after returning, if we chose. */
+  // (*cinfo->err->output_message)(cinfo);
+  /* Create the message */
+  (*(cinfo->err->format_message))(cinfo, jpegLastErrorMsg);
+
+  /* Return control to the setjmp point */
+  longjmp(myerr->setjmp_buffer, 1);
+}
+
+struct torch_jpeg_mgr {
+  struct jpeg_source_mgr pub;
+  const JOCTET* data;
+  size_t len;
+};
+
+static void torch_jpeg_init_source(j_decompress_ptr cinfo) {}
+
+static boolean torch_jpeg_fill_input_buffer(j_decompress_ptr cinfo) {
+  torch_jpeg_mgr* src = (torch_jpeg_mgr*)cinfo->src;
+  // No more data.  Probably an incomplete image;  just output EOI.
+  src->pub.next_input_byte = EOI_BUFFER;
+  src->pub.bytes_in_buffer = 1;
+  return TRUE;
+}
+
+static void torch_jpeg_skip_input_data(j_decompress_ptr cinfo, long num_bytes) {
+  torch_jpeg_mgr* src = (torch_jpeg_mgr*)cinfo->src;
+  if (src->pub.bytes_in_buffer < num_bytes) {
+    // Skipping over all of remaining data;  output EOI.
+    src->pub.next_input_byte = EOI_BUFFER;
+    src->pub.bytes_in_buffer = 1;
+  } else {
+    // Skipping over only some of the remaining data.
+    src->pub.next_input_byte += num_bytes;
+    src->pub.bytes_in_buffer -= num_bytes;
+  }
+}
+
+static void torch_jpeg_term_source(j_decompress_ptr cinfo) {}
+
+static void torch_jpeg_set_source_mgr(
+    j_decompress_ptr cinfo,
+    const unsigned char* data,
+    size_t len) {
+  torch_jpeg_mgr* src;
+  if (cinfo->src == 0) { // if this is first time;  allocate memory
+    cinfo->src = (struct jpeg_source_mgr*)(*cinfo->mem->alloc_small)(
+        (j_common_ptr)cinfo, JPOOL_PERMANENT, sizeof(torch_jpeg_mgr));
+  }
+  src = (torch_jpeg_mgr*)cinfo->src;
+  src->pub.init_source = torch_jpeg_init_source;
+  src->pub.fill_input_buffer = torch_jpeg_fill_input_buffer;
+  src->pub.skip_input_data = torch_jpeg_skip_input_data;
+  src->pub.resync_to_restart = jpeg_resync_to_restart; // default
+  src->pub.term_source = torch_jpeg_term_source;
+  // fill the buffers
+  src->data = (const JOCTET*)data;
+  src->len = len;
+  src->pub.bytes_in_buffer = len;
+  src->pub.next_input_byte = src->data;
+}
+
+torch::Tensor decodeJPEG(const torch::Tensor& data) {
+  struct jpeg_decompress_struct cinfo;
+  struct torch_jpeg_error_mgr jerr;
+
+  auto datap = data.data_ptr<uint8_t>();
+  // Setup decompression structure
+  cinfo.err = jpeg_std_error(&jerr.pub);
+  jerr.pub.error_exit = torch_jpeg_error_exit;
+  /* Establish the setjmp return context for my_error_exit to use. */
+  if (setjmp(jerr.setjmp_buffer)) {
+    /* If we get here, the JPEG code has signaled an error.
+     * We need to clean up the JPEG object.
+     */
+    jpeg_destroy_decompress(&cinfo);
+    AT_ERROR(jpegLastErrorMsg);
+  }
+
+  jpeg_create_decompress(&cinfo);
+  torch_jpeg_set_source_mgr(&cinfo, datap, data.numel());
+
+  // read info from header.
+  jpeg_read_header(&cinfo, TRUE);
+  jpeg_start_decompress(&cinfo);
+
+  int height = cinfo.output_height;
+  int width = cinfo.output_width;
+  int components = cinfo.output_components;
+
+  auto stride = width * components;
+  auto tensor = torch::empty(
+      {int64_t(height), int64_t(width), int64_t(components)}, torch::kU8);
+  auto ptr = tensor.data_ptr<uint8_t>();
+  while (cinfo.output_scanline < cinfo.output_height) {
+    /* jpeg_read_scanlines expects an array of pointers to scanlines.
+     * Here the array is only one element long, but you could ask for
+     * more than one scanline at a time if that's more convenient.
+     */
+    jpeg_read_scanlines(&cinfo, &ptr, 1);
+    ptr += stride;
+  }
+
+  jpeg_finish_decompress(&cinfo);
+  jpeg_destroy_decompress(&cinfo);
+  return tensor;
+}
+
+#endif // JPEG_FOUND
diff --git a/torchvision/csrc/cpu/image/readjpeg_cpu.h b/torchvision/csrc/cpu/image/readjpeg_cpu.h
@@ -0,0 +1,5 @@
+#pragma once
+
+#include <torch/torch.h>
+
+torch::Tensor decodeJPEG(const torch::Tensor& data);
diff --git a/torchvision/io/__init__.py b/torchvision/io/__init__.py
@@ -30,5 +30,5 @@
     "_read_video_clip_from_memory",
     "_read_video_meta_data",
     "VideoMetaData",
-    "Timebase",
+    "Timebase"
 ]
diff --git a/torchvision/io/image.py b/torchvision/io/image.py

-Original file line number
+Diff line change
   - codecov
   - pip
   - libpng
 +  - jpeg
   - ca-certificates
   - pip:
     - future