Repository: pytorch/vision

Issue:
BUG: ColorJitter in torchvision.transforms
## 🐛 Bug

ColorJitter is supposed to be called on an image of type PIL or Tensor, but can only be called on images of type PIL. 

## To Reproduce

Steps to reproduce the behavior:

1. Load the data by specifically composing a `ToTensor()` transformation followed by a `ColorJitter()` one.
2. Create a DataLoader using that dataset
3. Try to loop through the loader

Code example that reproduces this bug:

```python
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision.datasets import CIFAR10
root = 'path/to/cifar/data'

color_jitter = transforms.ColorJitter(0.8, 0.8, 0.8, 0.2)
ds_train = CIFAR10(root, download=True, train=True, transform=transforms.Compose([transforms.ToTensor(), color_jitter]))

train_loader = DataLoader(ds_train, batch_size=128, num_workers=4, drop_last=True, shuffle=False)

for (images, label) in train_loader:
    print("There is no Bug!")
```

Error message:

```
Traceback (most recent call last):
  File "train.py", line 11, in <module>
    for (images, label) in train_loader:
  File "/path/to/home/dir/.conda/envs/contrastive/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 400, in __next__
    data = self._next_data()
  File "/path/to/home/dir/.conda/envs/contrastive/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 1032, in _next_data
    return self._process_data(data)
  File "/path/to/home/dir/.conda/envs/contrastive/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 1058, in _process_data
    data.reraise()
  File "/path/to/home/dir/.conda/envs/contrastive/lib/python3.7/site-packages/torch/_utils.py", line 420, in reraise
    raise self.exc_type(msg)
TypeError: Caught TypeError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/path/to/home/dir/.conda/envs/contrastive/lib/python3.7/site-packages/torch/utils/data/_utils/worker.py", line 185, in _worker_loop
    data = fetcher.fetch(index)
  File "/path/to/home/dir/.conda/envs/contrastive/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py", line 44, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/path/to/home/dir/.conda/envs/contrastive/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py", line 44, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/path/to/home/dir/.conda/envs/contrastive/lib/python3.7/site-packages/torchvision/datasets/cifar.py", line 120, in __getitem__
    img = self.transform(img)
  File "/path/to/home/dir/.conda/envs/contrastive/lib/python3.7/site-packages/torchvision/transforms/transforms.py", line 54, in __call__
    img = t(img)
  File "/path/to/home/dir/.conda/envs/contrastive/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/path/to/home/dir/.conda/envs/contrastive/lib/python3.7/site-packages/torchvision/transforms/transforms.py", line 1092, in forward
    img = F.adjust_hue(img, hue_factor)
  File "/path/to/home/dir/.conda/envs/contrastive/lib/python3.7/site-packages/torchvision/transforms/functional.py", line 728, in adjust_hue
    raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
TypeError: img should be PIL Image. Got <class 'torch.Tensor'>
```

## Expected behavior

Should be able to run without errors.

## Environment

```
PyTorch version: 1.7.0.dev20200807
Is debug build: No
CUDA used to build PyTorch: 10.2

OS: Ubuntu 18.04.3 LTS
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
CMake version: version 3.10.2

Python version: 3.7
Is CUDA available: Yes
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: GeForce RTX 2080 Ti
GPU 1: GeForce RTX 2080 Ti
GPU 2: GeForce RTX 2080 Ti
GPU 3: GeForce RTX 2080 Ti

Nvidia driver version: 440.33.01
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5

Versions of relevant libraries:
[pip3] numpy==1.19.1
[pip3] torch==1.7.0.dev20200807
[pip3] torchfile==0.1.0
[pip3] torchnet==0.0.4
[pip3] torchvision==0.8.0.dev20200807
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               10.2.89              hfd86e86_1  
[conda] mkl                       2020.1                      217  
[conda] mkl-service               2.3.0            py37he904b0f_0  
[conda] mkl_fft                   1.1.0            py37h23d657b_0  
[conda] mkl_random                1.1.1            py37h0573a6f_0  
[conda] numpy                     1.19.1           py37hbc911f0_0  
[conda] numpy-base                1.19.1           py37hfa32c7d_0  
[conda] pytorch                   1.7.0.dev20200807 py3.7_cuda10.2.89_cudnn7.6.5_0    pytorch-nightly
[conda] torchfile                 0.1.0                    pypi_0    pypi
[conda] torchnet                  0.0.4                    pypi_0    pypi
[conda] torchvision               0.8.0.dev20200807      py37_cu102    pytorch-nightly
```
## Additional context

By looking into the code of `transforms/transforms.py` and `transforms/functional.py` of the master branch of this repo, the description of `forward` at line 1064 (in `class ColorJitter`) of `transforms.py` says:

```python
    def forward(self, img):
        """
        Args:
            img (PIL Image or Tensor): Input image.
        Returns:
            PIL Image or Tensor: Color jittered image.
        """
```
which contradicts the fact that `ColorJitter` was placed in the "Transforms on PIL Image" category in the official documentation at https://pytorch.org/docs/stable/torchvision/transforms.html. Then `forward` proceeds to call, on lines 1077, 1082, 1087 and 1092, the following functions in this order:
`F.adjust_brightness(img, brightness_factor)`, `F.adjust_contrast(img, contrast_factor)`, `F.adjust_saturation(img, saturation_factor)` and `F.adjust_hue(img, hue_factor)`.

Now, if we look into `transforms/functional.py`, the first three functions handle both Tensor and PIL type cases. For example, at line 675, we have (in the case of `adjust_brightness`):

```python
    if not isinstance(img, torch.Tensor):
        return F_pil.adjust_brightness(img, brightness_factor)

    return F_t.adjust_brightness(img, brightness_factor)
```

However, it seems to have been forgotten to update the last one as well. Indeed, at line 742, in `adjust_hue`, we have:

```python
    if not isinstance(img, torch.Tensor):
        return F_pil.adjust_hue(img, hue_factor)

    raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
```

It would also be nice if other transformations, like RandomGrayscale, could support types of both Tensor and PIL as inputs (with the types of their outputs being consistent).

## Pull Request

I opened a pull request (#2566) to attempt to address the reported bug.



Discussion:

Resolving patch:
diff --git a/torchvision/transforms/functional.py b/torchvision/transforms/functional.py
--- a/torchvision/transforms/functional.py
+++ b/torchvision/transforms/functional.py
@@ -736,7 +736,7 @@ def adjust_hue(img: Tensor, hue_factor: float) -> Tensor:
     .. _Hue: https://en.wikipedia.org/wiki/Hue
 
     Args:
-        img (PIL Image): PIL Image to be adjusted.
+        img (PIL Image or Tensor): Image to be adjusted.
         hue_factor (float):  How much to shift the hue channel. Should be in
             [-0.5, 0.5]. 0.5 and -0.5 give complete reversal of hue channel in
             HSV space in positive and negative direction respectively.
@@ -744,12 +744,12 @@ def adjust_hue(img: Tensor, hue_factor: float) -> Tensor:
             with complementary colors while 0 gives the original image.
 
     Returns:
-        PIL Image: Hue adjusted image.
+        PIL Image or Tensor: Hue adjusted image.
     """
     if not isinstance(img, torch.Tensor):
         return F_pil.adjust_hue(img, hue_factor)
 
-    raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
+    return F_t.adjust_hue(img, hue_factor)
 
 
 def adjust_gamma(img: Tensor, gamma: float, gain: float = 1) -> Tensor:
diff --git a/torchvision/transforms/functional_tensor.py b/torchvision/transforms/functional_tensor.py
--- a/torchvision/transforms/functional_tensor.py
+++ b/torchvision/transforms/functional_tensor.py
@@ -157,7 +157,7 @@ def adjust_contrast(img: Tensor, contrast_factor: float) -> Tensor:
     return _blend(img, mean, contrast_factor)
 
 
-def adjust_hue(img, hue_factor):
+def adjust_hue(img: Tensor, hue_factor: float) -> Tensor:
     """Adjust hue of an image.
 
     The image hue is adjusted by converting the image to HSV and
@@ -185,8 +185,8 @@ def adjust_hue(img, hue_factor):
     if not (-0.5 <= hue_factor <= 0.5):
         raise ValueError('hue_factor ({}) is not in [-0.5, 0.5].'.format(hue_factor))
 
-    if not _is_tensor_a_torch_image(img):
-        raise TypeError('tensor is not a torch image.')
+    if not (isinstance(img, torch.Tensor) and _is_tensor_a_torch_image(img)):
+        raise TypeError('img should be Tensor image. Got {}'.format(type(img)))
 
     orig_dtype = img.dtype
     if img.dtype == torch.uint8:
@@ -194,8 +194,7 @@ def adjust_hue(img, hue_factor):
 
     img = _rgb2hsv(img)
     h, s, v = img.unbind(0)
-    h += hue_factor
-    h = h % 1.0
+    h = (h + hue_factor) % 1.0
     img = torch.stack((h, s, v))
     img_hue_adj = _hsv2rgb(img)
 
@@ -408,6 +407,8 @@ def _blend(img1: Tensor, img2: Tensor, ratio: float) -> Tensor:
 def _rgb2hsv(img):
     r, g, b = img.unbind(0)
 
+    # Implementation is based on https://github.com/python-pillow/Pillow/blob/4174d4267616897df3746d315d5a2d0f82c656ee/
+    # src/libImaging/Convert.c#L330
     maxc = torch.max(img, dim=0).values
     minc = torch.min(img, dim=0).values
 


Test patch:
diff --git a/test/test_functional_tensor.py b/test/test_functional_tensor.py
--- a/test/test_functional_tensor.py
+++ b/test/test_functional_tensor.py
@@ -217,11 +217,9 @@ def test_pad(self):
         with self.assertRaises(ValueError, msg="Padding can not be negative for symmetric padding_mode"):
             F_t.pad(tensor, (-2, -3), padding_mode="symmetric")
 
-    def _test_adjust_fn(self, fn, fn_pil, fn_t, configs):
+    def _test_adjust_fn(self, fn, fn_pil, fn_t, configs, tol=2.0 + 1e-10, agg_method="max"):
         script_fn = torch.jit.script(fn)
-
         torch.manual_seed(15)
-
         tensor, pil_img = self._create_data(26, 34, device=self.device)
 
         for dt in [None, torch.float32, torch.float64]:
@@ -230,7 +228,6 @@ def _test_adjust_fn(self, fn, fn_pil, fn_t, configs):
                 tensor = F.convert_image_dtype(tensor, dt)
 
             for config in configs:
-
                 adjusted_tensor = fn_t(tensor, **config)
                 adjusted_pil = fn_pil(pil_img, **config)
                 scripted_result = script_fn(tensor, **config)
@@ -245,9 +242,12 @@ def _test_adjust_fn(self, fn, fn_pil, fn_t, configs):
 
                 # Check that max difference does not exceed 2 in [0, 255] range
                 # Exact matching is not possible due to incompatibility convert_image_dtype and PIL results
-                tol = 2.0 + 1e-10
-                self.approxEqualTensorToPIL(rbg_tensor.float(), adjusted_pil, tol, msg=msg, agg_method="max")
-                self.assertTrue(adjusted_tensor.allclose(scripted_result), msg=msg)
+                self.approxEqualTensorToPIL(rbg_tensor.float(), adjusted_pil, tol=tol, msg=msg, agg_method=agg_method)
+
+                atol = 1e-6
+                if adjusted_tensor.dtype == torch.uint8 and "cuda" in torch.device(self.device).type:
+                    atol = 1.0
+                self.assertTrue(adjusted_tensor.allclose(scripted_result, atol=atol), msg=msg)
 
     def test_adjust_brightness(self):
         self._test_adjust_fn(
@@ -273,6 +273,16 @@ def test_adjust_saturation(self):
             [{"saturation_factor": f} for f in [0.5, 0.75, 1.0, 1.5, 2.0]]
         )
 
+    def test_adjust_hue(self):
+        self._test_adjust_fn(
+            F.adjust_hue,
+            F_pil.adjust_hue,
+            F_t.adjust_hue,
+            [{"hue_factor": f} for f in [-0.45, -0.25, 0.0, 0.25, 0.45]],
+            tol=0.1,
+            agg_method="mean"
+        )
+
     def test_adjust_gamma(self):
         self._test_adjust_fn(
             F.adjust_gamma,
diff --git a/test/test_transforms_tensor.py b/test/test_transforms_tensor.py
--- a/test/test_transforms_tensor.py
+++ b/test/test_transforms_tensor.py
@@ -60,24 +60,36 @@ def test_random_vertical_flip(self):
     def test_color_jitter(self):
 
         tol = 1.0 + 1e-10
-        for f in [0.1, 0.5, 1.0, 1.34]:
+        for f in [0.1, 0.5, 1.0, 1.34, (0.3, 0.7), [0.4, 0.5]]:
             meth_kwargs = {"brightness": f}
             self._test_class_op(
                 "ColorJitter", meth_kwargs=meth_kwargs, test_exact_match=False, tol=tol, agg_method="max"
             )
 
-        for f in [0.2, 0.5, 1.0, 1.5]:
+        for f in [0.2, 0.5, 1.0, 1.5, (0.3, 0.7), [0.4, 0.5]]:
             meth_kwargs = {"contrast": f}
             self._test_class_op(
                 "ColorJitter", meth_kwargs=meth_kwargs, test_exact_match=False, tol=tol, agg_method="max"
             )
 
-        for f in [0.5, 0.75, 1.0, 1.25]:
+        for f in [0.5, 0.75, 1.0, 1.25, (0.3, 0.7), [0.3, 0.4]]:
             meth_kwargs = {"saturation": f}
             self._test_class_op(
                 "ColorJitter", meth_kwargs=meth_kwargs, test_exact_match=False, tol=tol, agg_method="max"
             )
 
+        for f in [0.2, 0.5, (-0.2, 0.3), [-0.4, 0.5]]:
+            meth_kwargs = {"hue": f}
+            self._test_class_op(
+                "ColorJitter", meth_kwargs=meth_kwargs, test_exact_match=False, tol=0.1, agg_method="mean"
+            )
+
+        # All 4 parameters together
+        meth_kwargs = {"brightness": 0.2, "contrast": 0.2, "saturation": 0.2, "hue": 0.2}
+        self._test_class_op(
+            "ColorJitter", meth_kwargs=meth_kwargs, test_exact_match=False, tol=0.1, agg_method="mean"
+        )
+
     def test_pad(self):
 
         # Test functional.pad (PIL and Tensor) with padding as single int