Refactor test_transforms.py into classes (#4030)

test/test_transforms.py

@@ -195,8 +195,10 @@ class TestAccImage:
         torch.testing.assert_close(output, expected_output)
 
 
-@pytest.mark.parametrize('channels', [1, 3, 4])
-def test_to_tensor(channels):
+class TestToTensor:
+
+    @pytest.mark.parametrize('channels', [1, 3, 4])
+    def test_to_tensor(self, channels):
         height, width = 4, 4
         trans = transforms.ToTensor()
 
@@ -221,8 +223,7 @@ def test_to_tensor(channels):
         output = trans(img)
         torch.testing.assert_close(input_data, output, check_dtype=False, check_stride=False)
 
-
-def test_to_tensor_errors():
+    def test_to_tensor_errors(self):
         height, width = 4, 4
         trans = transforms.ToTensor()
 
@@ -235,9 +236,8 @@ def test_to_tensor_errors():
         with pytest.raises(ValueError):
             trans(np.random.rand(1, 1, height, width))
 
-
-@pytest.mark.parametrize('dtype', [torch.float16, torch.float, torch.double])
-def test_to_tensor_with_other_default_dtypes(dtype):
+    @pytest.mark.parametrize('dtype', [torch.float16, torch.float, torch.double])
+    def test_to_tensor_with_other_default_dtypes(self, dtype):
         current_def_dtype = torch.get_default_dtype()
 
         t = transforms.ToTensor()
@@ -250,9 +250,8 @@ def test_to_tensor_with_other_default_dtypes(dtype):
 
         torch.set_default_dtype(current_def_dtype)
 
-
-@pytest.mark.parametrize('channels', [1, 3, 4])
-def test_pil_to_tensor(channels):
+    @pytest.mark.parametrize('channels', [1, 3, 4])
+    def test_pil_to_tensor(self, channels):
         height, width = 4, 4
         trans = transforms.PILToTensor()
 
@@ -279,8 +278,7 @@ def test_pil_to_tensor(channels):
         output = trans(img).view(torch.uint8).bool().to(torch.uint8)
         torch.testing.assert_close(input_data, output, check_stride=False)
 
-
-def test_pil_to_tensor_errors():
+    def test_pil_to_tensor_errors(self):
         height, width = 4, 4
         trans = transforms.PILToTensor()
 
@@ -537,7 +535,9 @@ def test_randomness(fn, trans, config, p):
     assert p_value > 0.0001
 
 
-def _get_1_channel_tensor_various_types():
+class TestToPil:
+
+    def _get_1_channel_tensor_various_types():
         img_data_float = torch.Tensor(1, 4, 4).uniform_()
         expected_output = img_data_float.mul(255).int().float().div(255).numpy()
         yield img_data_float, expected_output, 'L'
@@ -554,10 +554,26 @@ def _get_1_channel_tensor_various_types():
         expected_output = img_data_int.numpy()
         yield img_data_int, expected_output, 'I'
 
+    def _get_2d_tensor_various_types():
+        img_data_float = torch.Tensor(4, 4).uniform_()
+        expected_output = img_data_float.mul(255).int().float().div(255).numpy()
+        yield img_data_float, expected_output, 'L'
+
+        img_data_byte = torch.ByteTensor(4, 4).random_(0, 255)
+        expected_output = img_data_byte.float().div(255.0).numpy()
+        yield img_data_byte, expected_output, 'L'
+
+        img_data_short = torch.ShortTensor(4, 4).random_()
+        expected_output = img_data_short.numpy()
+        yield img_data_short, expected_output, 'I;16'
+
+        img_data_int = torch.IntTensor(4, 4).random_()
+        expected_output = img_data_int.numpy()
+        yield img_data_int, expected_output, 'I'
 
-@pytest.mark.parametrize('with_mode', [False, True])
-@pytest.mark.parametrize('img_data, expected_output, expected_mode', _get_1_channel_tensor_various_types())
-def test_1_channel_tensor_to_pil_image(with_mode, img_data, expected_output, expected_mode):
+    @pytest.mark.parametrize('with_mode', [False, True])
+    @pytest.mark.parametrize('img_data, expected_output, expected_mode', _get_1_channel_tensor_various_types())
+    def test_1_channel_tensor_to_pil_image(self, with_mode, img_data, expected_output, expected_mode):
         transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage()
         to_tensor = transforms.ToTensor()
 
@@ -565,8 +581,7 @@ def test_1_channel_tensor_to_pil_image(with_mode, img_data, expected_output, exp
         assert img.mode == expected_mode
         torch.testing.assert_close(expected_output, to_tensor(img).numpy(), check_stride=False)
 
-
-def test_1_channel_float_tensor_to_pil_image():
+    def test_1_channel_float_tensor_to_pil_image(self):
         img_data = torch.Tensor(1, 4, 4).uniform_()
         # 'F' mode for torch.FloatTensor
         img_F_mode = transforms.ToPILImage(mode='F')(img_data)
@@ -575,15 +590,14 @@ def test_1_channel_float_tensor_to_pil_image():
             np.array(Image.fromarray(img_data.squeeze(0).numpy(), mode='F')), np.array(img_F_mode)
         )
 
-
-@pytest.mark.parametrize('with_mode', [False, True])
-@pytest.mark.parametrize('img_data, expected_mode', [
+    @pytest.mark.parametrize('with_mode', [False, True])
+    @pytest.mark.parametrize('img_data, expected_mode', [
         (torch.Tensor(4, 4, 1).uniform_().numpy(), 'F'),
         (torch.ByteTensor(4, 4, 1).random_(0, 255).numpy(), 'L'),
         (torch.ShortTensor(4, 4, 1).random_().numpy(), 'I;16'),
         (torch.IntTensor(4, 4, 1).random_().numpy(), 'I'),
-])
-def test_1_channel_ndarray_to_pil_image(with_mode, img_data, expected_mode):
+    ])
+    def test_1_channel_ndarray_to_pil_image(self, with_mode, img_data, expected_mode):
         transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage()
         img = transform(img_data)
         assert img.mode == expected_mode
@@ -591,9 +605,8 @@ def test_1_channel_ndarray_to_pil_image(with_mode, img_data, expected_mode):
         # and otherwise assert_close wouldn't be able to construct a tensor from the uint16 array
         torch.testing.assert_close(img_data[:, :, 0], np.asarray(img).astype(img_data.dtype))
 
-
-@pytest.mark.parametrize('expected_mode', [None, 'LA'])
-def test_2_channel_ndarray_to_pil_image(expected_mode):
+    @pytest.mark.parametrize('expected_mode', [None, 'LA'])
+    def test_2_channel_ndarray_to_pil_image(self, expected_mode):
         img_data = torch.ByteTensor(4, 4, 2).random_(0, 255).numpy()
 
         if expected_mode is None:
@@ -606,8 +619,7 @@ def test_2_channel_ndarray_to_pil_image(expected_mode):
         for i in range(2):
             torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i]), check_stride=False)
 
-
-def test_2_channel_ndarray_to_pil_image_error():
+    def test_2_channel_ndarray_to_pil_image_error(self):
         img_data = torch.ByteTensor(4, 4, 2).random_(0, 255).numpy()
         transforms.ToPILImage().__repr__()
 
@@ -619,9 +631,8 @@ def test_2_channel_ndarray_to_pil_image_error():
         with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"):
             transforms.ToPILImage(mode='RGB')(img_data)
 
-
-@pytest.mark.parametrize('expected_mode', [None, 'LA'])
-def test_2_channel_tensor_to_pil_image(expected_mode):
+    @pytest.mark.parametrize('expected_mode', [None, 'LA'])
+    def test_2_channel_tensor_to_pil_image(self, expected_mode):
         img_data = torch.Tensor(2, 4, 4).uniform_()
         expected_output = img_data.mul(255).int().float().div(255)
         if expected_mode is None:
@@ -635,8 +646,7 @@ def test_2_channel_tensor_to_pil_image(expected_mode):
         for i in range(2):
             torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy())
 
-
-def test_2_channel_tensor_to_pil_image_error():
+    def test_2_channel_tensor_to_pil_image_error(self):
         img_data = torch.Tensor(2, 4, 4).uniform_()
 
         # should raise if we try a mode for 4 or 1 or 3 channel images
@@ -647,28 +657,9 @@ def test_2_channel_tensor_to_pil_image_error():
         with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"):
             transforms.ToPILImage(mode='RGB')(img_data)
 
-
-def _get_2d_tensor_various_types():
-    img_data_float = torch.Tensor(4, 4).uniform_()
-    expected_output = img_data_float.mul(255).int().float().div(255).numpy()
-    yield img_data_float, expected_output, 'L'
-
-    img_data_byte = torch.ByteTensor(4, 4).random_(0, 255)
-    expected_output = img_data_byte.float().div(255.0).numpy()
-    yield img_data_byte, expected_output, 'L'
-
-    img_data_short = torch.ShortTensor(4, 4).random_()
-    expected_output = img_data_short.numpy()
-    yield img_data_short, expected_output, 'I;16'
-
-    img_data_int = torch.IntTensor(4, 4).random_()
-    expected_output = img_data_int.numpy()
-    yield img_data_int, expected_output, 'I'
-
-
-@pytest.mark.parametrize('with_mode', [False, True])
-@pytest.mark.parametrize('img_data, expected_output, expected_mode', _get_2d_tensor_various_types())
-def test_2d_tensor_to_pil_image(with_mode, img_data, expected_output, expected_mode):
+    @pytest.mark.parametrize('with_mode', [False, True])
+    @pytest.mark.parametrize('img_data, expected_output, expected_mode', _get_2d_tensor_various_types())
+    def test_2d_tensor_to_pil_image(self, with_mode, img_data, expected_output, expected_mode):
         transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage()
         to_tensor = transforms.ToTensor()
 
@@ -676,23 +667,21 @@ def test_2d_tensor_to_pil_image(with_mode, img_data, expected_output, expected_m
         assert img.mode == expected_mode
         torch.testing.assert_close(expected_output, to_tensor(img).numpy()[0])
 
-
-@pytest.mark.parametrize('with_mode', [False, True])
-@pytest.mark.parametrize('img_data, expected_mode', [
+    @pytest.mark.parametrize('with_mode', [False, True])
+    @pytest.mark.parametrize('img_data, expected_mode', [
         (torch.Tensor(4, 4).uniform_().numpy(), 'F'),
         (torch.ByteTensor(4, 4).random_(0, 255).numpy(), 'L'),
         (torch.ShortTensor(4, 4).random_().numpy(), 'I;16'),
         (torch.IntTensor(4, 4).random_().numpy(), 'I'),
-])
-def test_2d_ndarray_to_pil_image(with_mode, img_data, expected_mode):
+    ])
+    def test_2d_ndarray_to_pil_image(self, with_mode, img_data, expected_mode):
         transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage()
         img = transform(img_data)
         assert img.mode == expected_mode
         np.testing.assert_allclose(img_data, img)
 
-
-@pytest.mark.parametrize('expected_mode', [None, 'RGB', 'HSV', 'YCbCr'])
-def test_3_channel_tensor_to_pil_image(expected_mode):
+    @pytest.mark.parametrize('expected_mode', [None, 'RGB', 'HSV', 'YCbCr'])
+    def test_3_channel_tensor_to_pil_image(self, expected_mode):
         img_data = torch.Tensor(3, 4, 4).uniform_()
         expected_output = img_data.mul(255).int().float().div(255)
 
@@ -706,8 +695,7 @@ def test_3_channel_tensor_to_pil_image(expected_mode):
         for i in range(3):
             torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy())
 
-
-def test_3_channel_tensor_to_pil_image_error():
+    def test_3_channel_tensor_to_pil_image_error(self):
         img_data = torch.Tensor(3, 4, 4).uniform_()
         error_message_3d = r"Only modes \['RGB', 'YCbCr', 'HSV'\] are supported for 3D inputs"
         # should raise if we try a mode for 4 or 1 or 2 channel images
@@ -721,9 +709,8 @@ def test_3_channel_tensor_to_pil_image_error():
         with pytest.raises(ValueError, match=r'pic should be 2/3 dimensional. Got \d+ dimensions.'):
             transforms.ToPILImage()(torch.Tensor(1, 3, 4, 4).uniform_())
 
-
-@pytest.mark.parametrize('expected_mode', [None, 'RGB', 'HSV', 'YCbCr'])
-def test_3_channel_ndarray_to_pil_image(expected_mode):
+    @pytest.mark.parametrize('expected_mode', [None, 'RGB', 'HSV', 'YCbCr'])
+    def test_3_channel_ndarray_to_pil_image(self, expected_mode):
         img_data = torch.ByteTensor(4, 4, 3).random_(0, 255).numpy()
 
         if expected_mode is None:
@@ -736,8 +723,7 @@ def test_3_channel_ndarray_to_pil_image(expected_mode):
         for i in range(3):
             torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i]), check_stride=False)
 
-
-def test_3_channel_ndarray_to_pil_image_error():
+    def test_3_channel_ndarray_to_pil_image_error(self):
         img_data = torch.ByteTensor(4, 4, 3).random_(0, 255).numpy()
 
         # Checking if ToPILImage can be printed as string
@@ -752,9 +738,8 @@ def test_3_channel_ndarray_to_pil_image_error():
         with pytest.raises(ValueError, match=error_message_3d):
             transforms.ToPILImage(mode='LA')(img_data)
 
-
-@pytest.mark.parametrize('expected_mode', [None, 'RGBA', 'CMYK', 'RGBX'])
-def test_4_channel_tensor_to_pil_image(expected_mode):
+    @pytest.mark.parametrize('expected_mode', [None, 'RGBA', 'CMYK', 'RGBX'])
+    def test_4_channel_tensor_to_pil_image(self, expected_mode):
         img_data = torch.Tensor(4, 4, 4).uniform_()
         expected_output = img_data.mul(255).int().float().div(255)
 
@@ -769,8 +754,7 @@ def test_4_channel_tensor_to_pil_image(expected_mode):
         for i in range(4):
             torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy())
 
-
-def test_4_channel_tensor_to_pil_image_error():
+    def test_4_channel_tensor_to_pil_image_error(self):
         img_data = torch.Tensor(4, 4, 4).uniform_()
 
         error_message_4d = r"Only modes \['RGBA', 'CMYK', 'RGBX'\] are supported for 4D inputs"
@@ -782,9 +766,8 @@ def test_4_channel_tensor_to_pil_image_error():
         with pytest.raises(ValueError, match=error_message_4d):
             transforms.ToPILImage(mode='LA')(img_data)
 
-
-@pytest.mark.parametrize('expected_mode', [None, 'RGBA', 'CMYK', 'RGBX'])
-def test_4_channel_ndarray_to_pil_image(expected_mode):
+    @pytest.mark.parametrize('expected_mode', [None, 'RGBA', 'CMYK', 'RGBX'])
+    def test_4_channel_ndarray_to_pil_image(self, expected_mode):
         img_data = torch.ByteTensor(4, 4, 4).random_(0, 255).numpy()
 
         if expected_mode is None:
@@ -797,8 +780,7 @@ def test_4_channel_ndarray_to_pil_image(expected_mode):
         for i in range(4):
             torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i]), check_stride=False)
 
-
-def test_4_channel_ndarray_to_pil_image_error():
+    def test_4_channel_ndarray_to_pil_image_error(self):
         img_data = torch.ByteTensor(4, 4, 4).random_(0, 255).numpy()
 
         error_message_4d = r"Only modes \['RGBA', 'CMYK', 'RGBX'\] are supported for 4D inputs"
@@ -810,8 +792,7 @@ def test_4_channel_ndarray_to_pil_image_error():
         with pytest.raises(ValueError, match=error_message_4d):
             transforms.ToPILImage(mode='LA')(img_data)
 
-
-def test_ndarray_bad_types_to_pil_image():
+    def test_ndarray_bad_types_to_pil_image(self):
         trans = transforms.ToPILImage()
         reg_msg = r'Input type \w+ is not supported'
         with pytest.raises(TypeError, match=reg_msg):
@@ -828,8 +809,7 @@ def test_ndarray_bad_types_to_pil_image():
         with pytest.raises(ValueError, match=r'pic should not have > 4 channels. Got \d+ channels.'):
             transforms.ToPILImage()(np.ones([4, 4, 6]))
 
-
-def test_tensor_bad_types_to_pil_image():
+    def test_tensor_bad_types_to_pil_image(self):
         with pytest.raises(ValueError, match=r'pic should be 2/3 dimensional. Got \d+ dimensions.'):
             transforms.ToPILImage()(torch.ones(1, 3, 4, 4))
         with pytest.raises(ValueError, match=r'pic should not have > 4 channels. Got \d+ channels.'):