PR: Read JPEG images directly

.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 .

.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 .

.circleci/unittest/linux/scripts/environment.yml

@@ -7,6 +7,7 @@ dependencies:
   - codecov
   - pip
   - libpng
+  - jpeg
   - ca-certificates
   - pip:
     - future

.circleci/unittest/windows/scripts/environment.yml

@@ -7,6 +7,7 @@ dependencies:
   - codecov
   - pip
   - libpng
+  - jpeg
   - ca-certificates
   - pip:
     - future

.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"

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)

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
 =======

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

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

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') }}

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,

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"):

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);

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"

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

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/io/__init__.py

@@ -30,5 +30,5 @@
     "_read_video_clip_from_memory",
     "_read_video_meta_data",
     "VideoMetaData",
-    "Timebase",
+    "Timebase"
 ]