Use libjpeg-turbo

.github/scripts/setup-env.sh

@@ -40,9 +40,10 @@ conda create \
   --quiet --yes \
   python="${PYTHON_VERSION}" pip \
   ninja cmake \
-  libpng jpeg \
+  libpng \
   'ffmpeg<4.3'
 conda activate ci
+conda install libjpeg-turbo -c pytorch -y
 pip install --progress-bar=off --upgrade setuptools
 
 # See https://github.com/pytorch/vision/issues/6790

.github/scripts/unittest.sh

@@ -11,4 +11,5 @@ echo '::group::Install testing utilities'
 pip install --progress-bar=off pytest pytest-mock pytest-cov
 echo '::endgroup::'
 
+python test/smoke_test.py
 pytest --junit-xml="${RUNNER_TEST_RESULTS_DIR}/test-results.xml" -v --durations=25

packaging/pre_build_script.sh

@@ -13,8 +13,8 @@ fi
 
 if [[ "$(uname)" == Darwin || "$OSTYPE" == "msys" ]]; then
   # Install libpng from Anaconda (defaults)
-  conda install libpng "jpeg<=9b" -yq
-  conda install -yq ffmpeg=4.2 -c pytorch
+  conda install libpng -yq
+  conda install -yq ffmpeg=4.2 libjpeg-turbo -c pytorch
 
   # Copy binaries to be included in the wheel distribution
   if [[ "$OSTYPE" == "msys" ]]; then

packaging/torchvision/meta.yaml

@@ -10,7 +10,7 @@ requirements:
   build:
     - {{ compiler('c') }} # [win]
     - libpng
-    - jpeg
+    - libjpeg-turbo
     # NOTE: The only ffmpeg version that we build is actually 4.2
     - ffmpeg >=4.2  # [not win]
 
@@ -28,7 +28,7 @@ requirements:
     - requests
     - libpng
     - ffmpeg >=4.2  # [not win]
-    - jpeg
+    - libjpeg-turbo
     - pillow >=5.3.0, !=8.3.*
     - pytorch-mutex 1.0 {{ build_variant }}  # [not osx ]
     {{ environ.get('CONDA_PYTORCH_CONSTRAINT') }}
@@ -62,7 +62,7 @@ test:
   requires:
     - pytest
     - scipy
-    - jpeg
+    - libjpeg-turbo
     - ca-certificates
 
 

test/smoke_test.py

@@ -78,6 +78,8 @@ def smoke_test_torchvision_resnet50_classify(device: str = "cpu") -> None:
 def main() -> None:
     print(f"torchvision: {torchvision.__version__}")
     print(f"torch.cuda.is_available: {torch.cuda.is_available()}")
+    print(f"{torch.ops.image._jpeg_version() = }")
+    assert torch.ops.image._is_compiled_against_turbo()
     smoke_test_torchvision()
     smoke_test_torchvision_read_decode()
     smoke_test_torchvision_resnet50_classify()

torchvision/csrc/io/image/cpu/decode_jpeg.cpp

@@ -152,8 +152,23 @@ torch::Tensor decode_jpeg(const torch::Tensor& data, ImageReadMode mode) {
   jpeg_destroy_decompress(&cinfo);
   return tensor.permute({2, 0, 1});
 }
+#endif // #if !JPEG_FOUND
 
+int64_t _jpeg_version() {
+#ifdef JPEG_FOUND
+  return JPEG_LIB_VERSION;
+#else
+  return -1;
+#endif
+}
+
+bool _is_compiled_against_turbo() {
+#ifdef LIBJPEG_TURBO_VERSION
+  return true;
+#else
+  return false;
 #endif
+}
 
 } // namespace image
 } // namespace vision

torchvision/csrc/io/image/cpu/decode_jpeg.h

@@ -10,5 +10,8 @@ C10_EXPORT torch::Tensor decode_jpeg(
     const torch::Tensor& data,
     ImageReadMode mode = IMAGE_READ_MODE_UNCHANGED);
 
+C10_EXPORT int64_t _jpeg_version();
+C10_EXPORT bool _is_compiled_against_turbo();
+
 } // namespace image
 } // namespace vision

torchvision/csrc/io/image/image.cpp

@@ -19,15 +19,18 @@ PyMODINIT_FUNC PyInit_image(void) {
 namespace vision {
 namespace image {
 
-static auto registry = torch::RegisterOperators()
-                           .op("image::decode_png", &decode_png)
-                           .op("image::encode_png", &encode_png)
-                           .op("image::decode_jpeg", &decode_jpeg)
-                           .op("image::encode_jpeg", &encode_jpeg)
-                           .op("image::read_file", &read_file)
-                           .op("image::write_file", &write_file)
-                           .op("image::decode_image", &decode_image)
-                           .op("image::decode_jpeg_cuda", &decode_jpeg_cuda);
+static auto registry =
+    torch::RegisterOperators()
+        .op("image::decode_png", &decode_png)
+        .op("image::encode_png", &encode_png)
+        .op("image::decode_jpeg", &decode_jpeg)
+        .op("image::encode_jpeg", &encode_jpeg)
+        .op("image::read_file", &read_file)
+        .op("image::write_file", &write_file)
+        .op("image::decode_image", &decode_image)
+        .op("image::decode_jpeg_cuda", &decode_jpeg_cuda)
+        .op("image::_jpeg_version", &_jpeg_version)
+        .op("image::_is_compiled_against_turbo", &_is_compiled_against_turbo);
 
 } // namespace image
 } // namespace vision