Added gray image support to adjust_saturation function (#4480)

* update channels parameter to every calling to check_functional_vs_PIL_vs_scripted

* update adjust_saturation

* update docstrings for functional transformations

* parametrize channels

* update docstring of ColorJitter class

* move channels to class's parameter

* remove testing channels for geometric transforms

* revert redundant changes

* revert redundant changes

* update grayscale test cases for randaugment, autoaugment, trivialaugment

* update docstrings of randaugment, autoaugment, trivialaugment

* update docstring of ColorJitter

* fix adjust_hue's docstring

* change test equal tolerance

* refactor grayscale tests

* make get_grayscale_test_image private

Author: VinhLoiIT

test/common_utils.py

@@ -148,7 +148,7 @@ def _create_data_batch(height=3, width=3, channels=3, num_samples=4, device="cpu
     return batch_tensor
 
 
-assert_equal = functools.partial(torch.testing.assert_close, rtol=0, atol=0)
+assert_equal = functools.partial(torch.testing.assert_close, rtol=0, atol=1e-6)
 
 
 def get_list_of_videos(tmpdir, num_videos=5, sizes=None, fps=None):

test/test_functional_tensor.py

@@ -681,57 +681,65 @@ def check_functional_vs_PIL_vs_scripted(
 @pytest.mark.parametrize('device', cpu_and_gpu())
 @pytest.mark.parametrize('dtype', (None, torch.float32, torch.float64))
 @pytest.mark.parametrize('config', [{"brightness_factor": f} for f in (0.1, 0.5, 1.0, 1.34, 2.5)])
-def test_adjust_brightness(device, dtype, config):
+@pytest.mark.parametrize('channels', [1, 3])
+def test_adjust_brightness(device, dtype, config, channels):
     check_functional_vs_PIL_vs_scripted(
         F.adjust_brightness,
         F_pil.adjust_brightness,
         F_t.adjust_brightness,
         config,
         device,
         dtype,
+        channels,
     )
 
 
 @pytest.mark.parametrize('device', cpu_and_gpu())
 @pytest.mark.parametrize('dtype', (None, torch.float32, torch.float64))
-def test_invert(device, dtype):
+@pytest.mark.parametrize('channels', [1, 3])
+def test_invert(device, dtype, channels):
     check_functional_vs_PIL_vs_scripted(
         F.invert,
         F_pil.invert,
         F_t.invert,
         {},
         device,
         dtype,
+        channels,
         tol=1.0,
         agg_method="max"
     )
 
 
 @pytest.mark.parametrize('device', cpu_and_gpu())
 @pytest.mark.parametrize('config', [{"bits": bits} for bits in range(0, 8)])
-def test_posterize(device, config):
+@pytest.mark.parametrize('channels', [1, 3])
+def test_posterize(device, config, channels):
     check_functional_vs_PIL_vs_scripted(
         F.posterize,
         F_pil.posterize,
         F_t.posterize,
         config,
         device,
         dtype=None,
+        channels=channels,
         tol=1.0,
         agg_method="max",
     )
 
 
 @pytest.mark.parametrize('device', cpu_and_gpu())
 @pytest.mark.parametrize('config', [{"threshold": threshold} for threshold in [0, 64, 128, 192, 255]])
-def test_solarize1(device, config):
+@pytest.mark.parametrize('channels', [1, 3])
+def test_solarize1(device, config, channels):
     check_functional_vs_PIL_vs_scripted(
         F.solarize,
         F_pil.solarize,
         F_t.solarize,
         config,
         device,
         dtype=None,
+        channels=channels,
         tol=1.0,
         agg_method="max",
     )
@@ -740,14 +748,16 @@ def test_solarize1(device, config):
 @pytest.mark.parametrize('device', cpu_and_gpu())
 @pytest.mark.parametrize('dtype', (torch.float32, torch.float64))
 @pytest.mark.parametrize('config', [{"threshold": threshold} for threshold in [0.0, 0.25, 0.5, 0.75, 1.0]])
-def test_solarize2(device, dtype, config):
+@pytest.mark.parametrize('channels', [1, 3])
+def test_solarize2(device, dtype, config, channels):
     check_functional_vs_PIL_vs_scripted(
         F.solarize,
         lambda img, threshold: F_pil.solarize(img, 255 * threshold),
         F_t.solarize,
         config,
         device,
         dtype,
+        channels,
         tol=1.0,
         agg_method="max",
     )
@@ -756,34 +766,39 @@ def test_solarize2(device, dtype, config):
 @pytest.mark.parametrize('device', cpu_and_gpu())
 @pytest.mark.parametrize('dtype', (None, torch.float32, torch.float64))
 @pytest.mark.parametrize('config', [{"sharpness_factor": f} for f in [0.2, 0.5, 1.0, 1.5, 2.0]])
-def test_adjust_sharpness(device, dtype, config):
+@pytest.mark.parametrize('channels', [1, 3])
+def test_adjust_sharpness(device, dtype, config, channels):
     check_functional_vs_PIL_vs_scripted(
         F.adjust_sharpness,
         F_pil.adjust_sharpness,
         F_t.adjust_sharpness,
         config,
         device,
         dtype,
+        channels,
     )
 
 
 @pytest.mark.parametrize('device', cpu_and_gpu())
 @pytest.mark.parametrize('dtype', (None, torch.float32, torch.float64))
-def test_autocontrast(device, dtype):
+@pytest.mark.parametrize('channels', [1, 3])
+def test_autocontrast(device, dtype, channels):
     check_functional_vs_PIL_vs_scripted(
         F.autocontrast,
         F_pil.autocontrast,
         F_t.autocontrast,
         {},
         device,
         dtype,
+        channels,
         tol=1.0,
         agg_method="max"
     )
 
 
 @pytest.mark.parametrize('device', cpu_and_gpu())
-def test_equalize(device):
+@pytest.mark.parametrize('channels', [1, 3])
+def test_equalize(device, channels):
     torch.use_deterministic_algorithms(False)
     check_functional_vs_PIL_vs_scripted(
         F.equalize,
@@ -792,6 +807,7 @@ def test_equalize(device):
         {},
         device,
         dtype=None,
+        channels=channels,
         tol=1.0,
         agg_method="max",
     )
@@ -809,35 +825,39 @@ def test_adjust_contrast(device, dtype, config, channels):
         config,
         device,
         dtype,
-        channels=channels
+        channels
     )
 
 
 @pytest.mark.parametrize('device', cpu_and_gpu())
 @pytest.mark.parametrize('dtype', (None, torch.float32, torch.float64))
 @pytest.mark.parametrize('config', [{"saturation_factor": f} for f in [0.5, 0.75, 1.0, 1.5, 2.0]])
-def test_adjust_saturation(device, dtype, config):
+@pytest.mark.parametrize('channels', [1, 3])
+def test_adjust_saturation(device, dtype, config, channels):
     check_functional_vs_PIL_vs_scripted(
         F.adjust_saturation,
         F_pil.adjust_saturation,
         F_t.adjust_saturation,
         config,
         device,
-        dtype
+        dtype,
+        channels
     )
 
 
 @pytest.mark.parametrize('device', cpu_and_gpu())
 @pytest.mark.parametrize('dtype', (None, torch.float32, torch.float64))
 @pytest.mark.parametrize('config', [{"hue_factor": f} for f in [-0.45, -0.25, 0.0, 0.25, 0.45]])
-def test_adjust_hue(device, dtype, config):
+@pytest.mark.parametrize('channels', [1, 3])
+def test_adjust_hue(device, dtype, config, channels):
     check_functional_vs_PIL_vs_scripted(
         F.adjust_hue,
         F_pil.adjust_hue,
         F_t.adjust_hue,
         config,
         device,
         dtype,
+        channels,
         tol=16.1,
         agg_method="max"
     )
@@ -846,14 +866,16 @@ def test_adjust_hue(device, dtype, config):
 @pytest.mark.parametrize('device', cpu_and_gpu())
 @pytest.mark.parametrize('dtype', (None, torch.float32, torch.float64))
 @pytest.mark.parametrize('config', [{"gamma": g1, "gain": g2} for g1, g2 in zip([0.8, 1.0, 1.2], [0.7, 1.0, 1.3])])
-def test_adjust_gamma(device, dtype, config):
+@pytest.mark.parametrize('channels', [1, 3])
+def test_adjust_gamma(device, dtype, config, channels):
     check_functional_vs_PIL_vs_scripted(
         F.adjust_gamma,
         F_pil.adjust_gamma,
         F_t.adjust_gamma,
         config,
         device,
         dtype,
+        channels,
     )
 
 

test/test_transforms.py

@@ -26,6 +26,12 @@
     os.path.dirname(os.path.abspath(__file__)), 'assets', 'encode_jpeg', 'grace_hopper_517x606.jpg')
 
 
+def _get_grayscale_test_image(img, fill=None):
+    img = img.convert('L')
+    fill = (fill[0], ) if isinstance(fill, tuple) else fill
+    return img, fill
+
+
 class TestConvertImageDtype:
     @pytest.mark.parametrize('input_dtype, output_dtype', cycle_over(float_dtypes()))
     def test_float_to_float(self, input_dtype, output_dtype):
@@ -1482,9 +1488,12 @@ def test_five_crop(single_dim):
 
 @pytest.mark.parametrize('policy', transforms.AutoAugmentPolicy)
 @pytest.mark.parametrize('fill', [None, 85, (128, 128, 128)])
-def test_autoaugment(policy, fill):
+@pytest.mark.parametrize('grayscale', [True, False])
+def test_autoaugment(policy, fill, grayscale):
     random.seed(42)
     img = Image.open(GRACE_HOPPER)
+    if grayscale:
+        img, fill = _get_grayscale_test_image(img, fill)
     transform = transforms.AutoAugment(policy=policy, fill=fill)
     for _ in range(100):
         img = transform(img)
@@ -1494,9 +1503,12 @@ def test_autoaugment(policy, fill):
 @pytest.mark.parametrize('num_ops', [1, 2, 3])
 @pytest.mark.parametrize('magnitude', [7, 9, 11])
 @pytest.mark.parametrize('fill', [None, 85, (128, 128, 128)])
-def test_randaugment(num_ops, magnitude, fill):
+@pytest.mark.parametrize('grayscale', [True, False])
+def test_randaugment(num_ops, magnitude, fill, grayscale):
     random.seed(42)
     img = Image.open(GRACE_HOPPER)
+    if grayscale:
+        img, fill = _get_grayscale_test_image(img, fill)
     transform = transforms.RandAugment(num_ops=num_ops, magnitude=magnitude, fill=fill)
     for _ in range(100):
         img = transform(img)
@@ -1505,9 +1517,12 @@ def test_randaugment(num_ops, magnitude, fill):
 
 @pytest.mark.parametrize('fill', [None, 85, (128, 128, 128)])
 @pytest.mark.parametrize('num_magnitude_bins', [10, 13, 30])
-def test_trivialaugmentwide(fill, num_magnitude_bins):
+@pytest.mark.parametrize('grayscale', [True, False])
+def test_trivialaugmentwide(fill, num_magnitude_bins, grayscale):
     random.seed(42)
     img = Image.open(GRACE_HOPPER)
+    if grayscale:
+        img, fill = _get_grayscale_test_image(img, fill)
     transform = transforms.TrivialAugmentWide(fill=fill, num_magnitude_bins=num_magnitude_bins)
     for _ in range(100):
         img = transform(img)