[WIP] Package torchvision with libjpeg-turbo instead of libjpeg

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

@@ -1,13 +1,14 @@
 channels:
   - pytorch
   - defaults
+  - conda-forge
 dependencies:
   - pytest
   - pytest-cov
   - pytest-mock
   - pip
   - libpng
-  - jpeg
+  - libjpeg-turbo
   - ca-certificates
   - h5py
   - pip:

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

@@ -1,13 +1,14 @@
 channels:
   - pytorch
   - defaults
+  - conda-forge
 dependencies:
   - pytest
   - pytest-cov
   - pytest-mock
   - pip
   - libpng
-  - jpeg
+  - libjpeg-turbo
   - ca-certificates
   - hdf5
   - setuptools

.github/workflows/build-m1-binaries.yml

@@ -52,7 +52,7 @@ jobs:
             setup_build_version
           fi
 
-          conda create -yp ${ENV_NAME} python=${PY_VERS} numpy libpng jpeg wheel pkg-config
+          conda create -yp ${ENV_NAME} python=${PY_VERS} numpy libpng wheel pkg-config libjpeg-turbo -c conda-forge
           conda run -p ${ENV_NAME} python3 -mpip install torch --pre --extra-index-url=https://download.pytorch.org/whl/${CHANNEL}
           conda run -p ${ENV_NAME} python3 -mpip install delocate
           conda run -p ${ENV_NAME} python3 setup.py bdist_wheel

.github/workflows/test-m1.yml

@@ -34,7 +34,7 @@ jobs:
           # Needed for JPEG library detection as setup.py detects conda presence by running `shutil.which('conda')`
           export PATH=~/miniconda3/bin:$PATH
           set -ex
-          conda create -yp ${ENV_NAME} python=${PY_VERS} numpy libpng jpeg scipy
+          conda create -yp ${ENV_NAME} python=${PY_VERS} numpy libpng scipy libjpeg-turbo -c conda-forge
           conda run -p ${ENV_NAME} python3 -mpip install --pre torch --extra-index-url=https://download.pytorch.org/whl/${CHANNEL}
           conda run -p ${ENV_NAME} python3 setup.py develop
           conda run -p ${ENV_NAME} python3 -mpip install pytest pytest-mock av

packaging/build_cmake.sh

@@ -46,7 +46,8 @@ conda install -yq \pytorch=$PYTORCH_VERSION $CONDA_CUDATOOLKIT_CONSTRAINT $PYTOR
 TORCH_PATH=$(dirname $(python -c "import torch; print(torch.__file__)"))
 
 if [[ "$(uname)" == Darwin || "$OSTYPE" == "msys" ]]; then
-    conda install -yq libpng jpeg
+    conda install -yq libpng
+    conda install -yq libjpeg-turbo -c conda-forge
 else
     yum install -y libpng-devel libjpeg-turbo-devel
 fi

packaging/build_conda.sh

@@ -13,4 +13,7 @@ setup_visual_studio_constraint
 setup_junit_results_folder
 export CUDATOOLKIT_CHANNEL="nvidia"
 
-conda build -c $CUDATOOLKIT_CHANNEL $CONDA_CHANNEL_FLAGS --no-anaconda-upload --python "$PYTHON_VERSION" packaging/torchvision
+conda install conda-libmamba-solver
+export CONDA_EXPERIMENTAL_SOLVER="libmamba"
+
+conda build -c $CUDATOOLKIT_CHANNEL $CONDA_CHANNEL_FLAGS -c conda-forge --no-anaconda-upload --python "$PYTHON_VERSION" packaging/torchvision

packaging/build_wheel.sh

@@ -22,7 +22,7 @@ if [[ "$(uname)" == Darwin || "$OSTYPE" == "msys" ]]; then
         pip_install "delocate>=0.9"
     else
         cp "$bin_path/Library/bin/libpng16.dll" torchvision
-        cp "$bin_path/Library/bin/libjpeg.dll" torchvision
+        # cp "$bin_path/Library/bin/libjpeg.dll" torchvision
     fi
 else
     # Install auditwheel to get some inspection utilities

packaging/pkg_helpers.bash

@@ -172,7 +172,9 @@ setup_wheel_python() {
     conda create ${CONDA_CHANNEL_FLAGS} -yn "env$PYTHON_VERSION" python="$PYTHON_VERSION"
     conda activate "env$PYTHON_VERSION"
     # Install libpng from Anaconda (defaults)
-    conda install ${CONDA_CHANNEL_FLAGS} libpng "jpeg<=9b" -y
+    conda install ${CONDA_CHANNEL_FLAGS} libpng -y
+    # Install libjpeg-turbo from conda-forge
+    conda install -c conda-forge libjpeg-turbo -y
   else
     # Install native CentOS libJPEG, freetype and GnuTLS
     yum install -y libjpeg-turbo-devel freetype gnutls

packaging/torchvision/meta.yaml

@@ -1,4 +1,6 @@
 {% set build_variant = environ.get('CONDA_BUILD_VARIANT', 'cpu') %}
+{% set cuda_version = environ.get('CU_VERSION', 'cpu').replace('cpu', 'cu-1')[2:] | int %}
+
 package:
   name: torchvision
   version: "{{ environ.get('BUILD_VERSION') }}"
@@ -10,7 +12,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]
 
@@ -27,7 +29,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') }}
@@ -38,7 +40,10 @@ requirements:
     - cpuonly
   {% elif not osx %}
   run_constrained:
-     - cpuonly <0
+    - cpuonly <0
+    {% if cuda_version < 116 %}
+    - cudatoolkit <0
+    {% endif %}
   {% endif %}
 
 build:
@@ -61,7 +66,7 @@ test:
   requires:
     - pytest
     - scipy
-    - jpeg
+    - libjpeg-turbo
     - ca-certificates
 
 

test/test_image.py

@@ -95,10 +95,10 @@ def test_decode_jpeg(img_path, pil_mode, mode):
     data = read_file(img_path)
     img_ljpeg = decode_image(data, mode=mode)
 
-    # Permit a small variation on pixel values to account for implementation
-    # differences between Pillow and LibJPEG.
+    # Permit a small variation on pixel values to account for differences
+    # between Pillow's libjpeg and ours.
     abs_mean_diff = (img_ljpeg.type(torch.float32) - img_pil).abs().mean().item()
-    assert abs_mean_diff < 2
+    assert abs_mean_diff < 1
 
 
 def test_decode_jpeg_errors():
@@ -412,77 +412,6 @@ def test_encode_jpeg_errors():
         encode_jpeg(torch.empty((100, 100), dtype=torch.uint8))
 
 
-def _collect_if(cond):
-    # TODO: remove this once test_encode_jpeg_reference and test_write_jpeg_reference
-    # are removed
-    def _inner(test_func):
-        if cond:
-            return test_func
-        else:
-            return pytest.mark.dont_collect(test_func)
-
-    return _inner
-
-
-@_collect_if(cond=False)
-@pytest.mark.parametrize(
-    "img_path",
-    [pytest.param(jpeg_path, id=_get_safe_image_name(jpeg_path)) for jpeg_path in get_images(ENCODE_JPEG, ".jpg")],
-)
-def test_encode_jpeg_reference(img_path):
-    # This test is *wrong*.
-    # It compares a torchvision-encoded jpeg with a PIL-encoded jpeg (the reference), but it
-    # starts encoding the torchvision version from an image that comes from
-    # decode_jpeg, which can yield different results from pil.decode (see
-    # test_decode... which uses a high tolerance).
-    # Instead, we should start encoding from the exact same decoded image, for a
-    # valid comparison. This is done in test_encode_jpeg, but unfortunately
-    # these more correct tests fail on windows (probably because of a difference
-    # in libjpeg) between torchvision and PIL.
-    # FIXME: make the correct tests pass on windows and remove this.
-    dirname = os.path.dirname(img_path)
-    filename, _ = os.path.splitext(os.path.basename(img_path))
-    write_folder = os.path.join(dirname, "jpeg_write")
-    expected_file = os.path.join(write_folder, f"{filename}_pil.jpg")
-    img = decode_jpeg(read_file(img_path))
-
-    with open(expected_file, "rb") as f:
-        pil_bytes = f.read()
-        pil_bytes = torch.as_tensor(list(pil_bytes), dtype=torch.uint8)
-    for src_img in [img, img.contiguous()]:
-        # PIL sets jpeg quality to 75 by default
-        jpeg_bytes = encode_jpeg(src_img, quality=75)
-        assert_equal(jpeg_bytes, pil_bytes)
-
-
-@_collect_if(cond=False)
-@pytest.mark.parametrize(
-    "img_path",
-    [pytest.param(jpeg_path, id=_get_safe_image_name(jpeg_path)) for jpeg_path in get_images(ENCODE_JPEG, ".jpg")],
-)
-def test_write_jpeg_reference(img_path, tmpdir):
-    # FIXME: Remove this eventually, see test_encode_jpeg_reference
-    data = read_file(img_path)
-    img = decode_jpeg(data)
-
-    basedir = os.path.dirname(img_path)
-    filename, _ = os.path.splitext(os.path.basename(img_path))
-    torch_jpeg = os.path.join(tmpdir, f"{filename}_torch.jpg")
-    pil_jpeg = os.path.join(basedir, "jpeg_write", f"{filename}_pil.jpg")
-
-    write_jpeg(img, torch_jpeg, quality=75)
-
-    with open(torch_jpeg, "rb") as f:
-        torch_bytes = f.read()
-
-    with open(pil_jpeg, "rb") as f:
-        pil_bytes = f.read()
-
-    assert_equal(torch_bytes, pil_bytes)
-
-
-# TODO: Remove the skip. See https://github.com/pytorch/vision/issues/5162.
-@pytest.mark.skip("this test fails because PIL uses libjpeg-turbo")
 @pytest.mark.parametrize(
     "img_path",
     [pytest.param(jpeg_path, id=_get_safe_image_name(jpeg_path)) for jpeg_path in get_images(ENCODE_JPEG, ".jpg")],
@@ -501,8 +430,6 @@ def test_encode_jpeg(img_path):
         assert_equal(encoded_jpeg_torch, encoded_jpeg_pil)
 
 
-# TODO: Remove the skip. See https://github.com/pytorch/vision/issues/5162.
-@pytest.mark.skip("this test fails because PIL uses libjpeg-turbo")
 @pytest.mark.parametrize(
     "img_path",
     [pytest.param(jpeg_path, id=_get_safe_image_name(jpeg_path)) for jpeg_path in get_images(ENCODE_JPEG, ".jpg")],