diff --git a/.circleci/config.yml b/.circleci/config.yml index 29b5fc77aab..f3fc23b7c92 100644 --- 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 index d9bd257eae6..f63c3f408ba 100644 --- 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 index 2b3604ee1c8..96b66319ed6 100644 --- 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 index ddbf7445a92..49795f73bc3 100644 --- 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 index d8c45c2defe..f5656f926f1 100644 --- 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 index 2d28bc8a4c5..5d2e86291f6 100644 --- 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 - $ + $ $) include(GNUInstallDirs) diff --git a/README.rst b/README.rst index 0de404c74b3..1d55e9387c3 100644 --- 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 index 1a6e1b1761a..043d2ed7ea9 100755 --- 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 index 128d0b51913..ff7abdab7fb 100644 --- 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 index b3fc2a2e9df..1b61464f01e 100644 --- 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 index 198c14bfbf1..e1587996362 100644 --- 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 new file mode 100644 index 00000000000..54b39dc0cd7 Binary files /dev/null and b/test/assets/grace_hopper_517x606.pth differ diff --git a/test/test_image.py b/test/test_image.py index a7f660127b8..09a591d19ef 100644 --- 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 index 0dc82a69827..5efe53b02b7 100644 --- 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 index f5b86cf683b..324ecea8a28 100644 --- a/torchvision/csrc/cpu/image/image.h +++ b/torchvision/csrc/cpu/image/image.h @@ -4,4 +4,5 @@ // Comment #include #include +#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 new file mode 100644 index 00000000000..4954b2c2474 --- /dev/null +++ b/torchvision/csrc/cpu/image/readjpeg_cpu.cpp @@ -0,0 +1,140 @@ +#include "readjpeg_cpu.h" + +#include +#include +#include + +#if !JPEG_FOUND + +torch::Tensor decodeJPEG(const torch::Tensor& data) { + AT_ERROR("decodeJPEG: torchvision not compiled with libjpeg support"); +} + +#else +#include + +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(); + // 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(); + 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 new file mode 100644 index 00000000000..40404df29b5 --- /dev/null +++ b/torchvision/csrc/cpu/image/readjpeg_cpu.h @@ -0,0 +1,5 @@ +#pragma once + +#include + +torch::Tensor decodeJPEG(const torch::Tensor& data); diff --git a/torchvision/io/__init__.py b/torchvision/io/__init__.py index cbbf560412e..4c47d8a51d5 100644 --- 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 index 1ad13ed27ad..8d5da4899ca 100644 --- a/torchvision/io/image.py +++ b/torchvision/io/image.py @@ -66,3 +66,44 @@ def read_png(path): raise ValueError("Expected a non empty file.") data = torch.from_file(path, dtype=torch.uint8, size=size) return decode_png(data) + + +def decode_jpeg(input): + # type: (Tensor) -> Tensor + """ + Decodes a JPEG image into a 3 dimensional RGB Tensor. + The values of the output tensor are uint8 between 0 and 255. + Arguments: + input (Tensor[1]): a one dimensional int8 tensor containing + the raw bytes of the JPEG image. + Returns: + output (Tensor[image_width, image_height, 3]) + """ + if not isinstance(input, torch.Tensor) or len(input) == 0 or input.ndim != 1: + raise ValueError("Expected a non empty 1-dimensional tensor.") + + if not input.dtype == torch.uint8: + raise ValueError("Expected a torch.uint8 tensor.") + + output = torch.ops.image.decode_jpeg(input) + return output + + +def read_jpeg(path): + # type: (str) -> Tensor + """ + Reads a JPEG image into a 3 dimensional RGB Tensor. + The values of the output tensor are uint8 between 0 and 255. + Arguments: + path (str): path of the JPEG image. + Returns: + output (Tensor[image_width, image_height, 3]) + """ + if not os.path.isfile(path): + raise ValueError("Expected a valid file path.") + + size = os.path.getsize(path) + if size == 0: + raise ValueError("Expected a non empty file.") + data = torch.from_file(path, dtype=torch.uint8, size=size) + return decode_jpeg(data)