diff --git a/test/test_transforms_tensor.py b/test/test_transforms_tensor.py index 2cd2f04831d..4e167d64b6a 100644 --- a/test/test_transforms_tensor.py +++ b/test/test_transforms_tensor.py @@ -18,7 +18,7 @@ _create_data_batch, _assert_equal_tensor_to_pil, _assert_approx_equal_tensor_to_pil, - cpu_and_gpu + cpu_and_gpu, ) from _assert_utils import assert_equal @@ -286,100 +286,6 @@ def test_resized_crop(self): with get_tmp_dir() as tmp_dir: s_transform.save(os.path.join(tmp_dir, "t_resized_crop.pt")) - def test_random_affine(self): - tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=self.device) - batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=self.device) - - def _test(**kwargs): - transform = T.RandomAffine(**kwargs) - s_transform = torch.jit.script(transform) - - _test_transform_vs_scripted(transform, s_transform, tensor) - _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors) - - return s_transform - - for interpolation in [NEAREST, BILINEAR]: - for shear in [15, 10.0, (5.0, 10.0), [-15, 15], [-10.0, 10.0, -11.0, 11.0]]: - _test(degrees=0.0, interpolation=interpolation, shear=shear) - - for scale in [(0.7, 1.2), [0.7, 1.2]]: - _test(degrees=0.0, interpolation=interpolation, scale=scale) - - for translate in [(0.1, 0.2), [0.2, 0.1]]: - _test(degrees=0.0, interpolation=interpolation, translate=translate) - - for degrees in [45, 35.0, (-45, 45), [-90.0, 90.0]]: - _test(degrees=degrees, interpolation=interpolation) - - for fill in [85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1, ], 1]: - _test(degrees=0.0, interpolation=interpolation, fill=fill) - - s_transform = _test(degrees=0.0) - with get_tmp_dir() as tmp_dir: - s_transform.save(os.path.join(tmp_dir, "t_random_affine.pt")) - - def test_random_rotate(self): - tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=self.device) - batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=self.device) - - for center in [(0, 0), [10, 10], None, (56, 44)]: - for expand in [True, False]: - for degrees in [45, 35.0, (-45, 45), [-90.0, 90.0]]: - for interpolation in [NEAREST, BILINEAR]: - for fill in [85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1, ], 1]: - transform = T.RandomRotation( - degrees=degrees, interpolation=interpolation, expand=expand, center=center, fill=fill - ) - s_transform = torch.jit.script(transform) - - _test_transform_vs_scripted(transform, s_transform, tensor) - _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors) - - with get_tmp_dir() as tmp_dir: - s_transform.save(os.path.join(tmp_dir, "t_random_rotate.pt")) - - def test_random_perspective(self): - tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=self.device) - batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=self.device) - - for distortion_scale in np.linspace(0.1, 1.0, num=20): - for interpolation in [NEAREST, BILINEAR]: - for fill in [85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1, ], 1]: - transform = T.RandomPerspective( - distortion_scale=distortion_scale, - interpolation=interpolation, - fill=fill - ) - s_transform = torch.jit.script(transform) - - _test_transform_vs_scripted(transform, s_transform, tensor) - _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors) - - with get_tmp_dir() as tmp_dir: - s_transform.save(os.path.join(tmp_dir, "t_perspective.pt")) - - def test_to_grayscale(self): - - meth_kwargs = {"num_output_channels": 1} - tol = 1.0 + 1e-10 - _test_class_op( - T.Grayscale, meth_kwargs=meth_kwargs, test_exact_match=False, device=self.device, - tol=tol, agg_method="max" - ) - - meth_kwargs = {"num_output_channels": 3} - _test_class_op( - T.Grayscale, meth_kwargs=meth_kwargs, test_exact_match=False, device=self.device, - tol=tol, agg_method="max" - ) - - meth_kwargs = {} - _test_class_op( - T.RandomGrayscale, meth_kwargs=meth_kwargs, test_exact_match=False, device=self.device, - tol=tol, agg_method="max" - ) - def test_normalize(self): fn = T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) tensor, _ = _create_data(26, 34, device=self.device) @@ -736,5 +642,125 @@ def setUp(self): self.device = "cuda" +def _test_random_affine_helper(device, **kwargs): + tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=device) + batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device) + transform = T.RandomAffine(**kwargs) + s_transform = torch.jit.script(transform) + + _test_transform_vs_scripted(transform, s_transform, tensor) + _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors) + + +@pytest.mark.parametrize('device', cpu_and_gpu()) +def test_random_affine(device): + transform = T.RandomAffine(degrees=45.0) + s_transform = torch.jit.script(transform) + with get_tmp_dir() as tmp_dir: + s_transform.save(os.path.join(tmp_dir, "t_random_affine.pt")) + + +@pytest.mark.parametrize('device', cpu_and_gpu()) +@pytest.mark.parametrize('interpolation', [NEAREST, BILINEAR]) +@pytest.mark.parametrize('shear', [15, 10.0, (5.0, 10.0), [-15, 15], [-10.0, 10.0, -11.0, 11.0]]) +def test_random_affine_shear(device, interpolation, shear): + _test_random_affine_helper(device, degrees=0.0, interpolation=interpolation, shear=shear) + + +@pytest.mark.parametrize('device', cpu_and_gpu()) +@pytest.mark.parametrize('interpolation', [NEAREST, BILINEAR]) +@pytest.mark.parametrize('scale', [(0.7, 1.2), [0.7, 1.2]]) +def test_random_affine_scale(device, interpolation, scale): + _test_random_affine_helper(device, degrees=0.0, interpolation=interpolation, scale=scale) + + +@pytest.mark.parametrize('device', cpu_and_gpu()) +@pytest.mark.parametrize('interpolation', [NEAREST, BILINEAR]) +@pytest.mark.parametrize('translate', [(0.1, 0.2), [0.2, 0.1]]) +def test_random_affine_translate(device, interpolation, translate): + _test_random_affine_helper(device, degrees=0.0, interpolation=interpolation, translate=translate) + + +@pytest.mark.parametrize('device', cpu_and_gpu()) +@pytest.mark.parametrize('interpolation', [NEAREST, BILINEAR]) +@pytest.mark.parametrize('degrees', [45, 35.0, (-45, 45), [-90.0, 90.0]]) +def test_random_affine_degrees(device, interpolation, degrees): + _test_random_affine_helper(device, degrees=degrees, interpolation=interpolation) + + +@pytest.mark.parametrize('device', cpu_and_gpu()) +@pytest.mark.parametrize('interpolation', [NEAREST, BILINEAR]) +@pytest.mark.parametrize('fill', [85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1, ], 1]) +def test_random_affine_fill(device, interpolation, fill): + _test_random_affine_helper(device, degrees=0.0, interpolation=interpolation, fill=fill) + + +@pytest.mark.parametrize('device', cpu_and_gpu()) +@pytest.mark.parametrize('center', [(0, 0), [10, 10], None, (56, 44)]) +@pytest.mark.parametrize('expand', [True, False]) +@pytest.mark.parametrize('degrees', [45, 35.0, (-45, 45), [-90.0, 90.0]]) +@pytest.mark.parametrize('interpolation', [NEAREST, BILINEAR]) +@pytest.mark.parametrize('fill', [85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1, ], 1]) +def test_random_rotate(device, center, expand, degrees, interpolation, fill): + tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=device) + batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device) + + transform = T.RandomRotation( + degrees=degrees, interpolation=interpolation, expand=expand, center=center, fill=fill + ) + s_transform = torch.jit.script(transform) + + _test_transform_vs_scripted(transform, s_transform, tensor) + _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors) + + +def test_random_rotate_save(): + transform = T.RandomRotation(degrees=45.0) + s_transform = torch.jit.script(transform) + with get_tmp_dir() as tmp_dir: + s_transform.save(os.path.join(tmp_dir, "t_random_rotate.pt")) + + +@pytest.mark.parametrize('device', cpu_and_gpu()) +@pytest.mark.parametrize('distortion_scale', np.linspace(0.1, 1.0, num=20)) +@pytest.mark.parametrize('interpolation', [NEAREST, BILINEAR]) +@pytest.mark.parametrize('fill', [85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1, ], 1]) +def test_random_perspective(device, distortion_scale, interpolation, fill): + tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=device) + batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device) + + transform = T.RandomPerspective( + distortion_scale=distortion_scale, + interpolation=interpolation, + fill=fill + ) + s_transform = torch.jit.script(transform) + + _test_transform_vs_scripted(transform, s_transform, tensor) + _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors) + + +def test_random_perspective_save(): + transform = T.RandomPerspective() + s_transform = torch.jit.script(transform) + with get_tmp_dir() as tmp_dir: + s_transform.save(os.path.join(tmp_dir, "t_perspective.pt")) + + +@pytest.mark.parametrize('device', cpu_and_gpu()) +@pytest.mark.parametrize('Klass, meth_kwargs', [ + (T.Grayscale, {"num_output_channels": 1}), + (T.Grayscale, {"num_output_channels": 3}), + (T.RandomGrayscale, {}) +]) +def test_to_grayscale(device, Klass, meth_kwargs): + + tol = 1.0 + 1e-10 + _test_class_op( + Klass, meth_kwargs=meth_kwargs, test_exact_match=False, device=device, + tol=tol, agg_method="max" + ) + + if __name__ == '__main__': unittest.main()