import pytest import torch from kornia.augmentation.random_generator import ( center_crop_generator, random_affine_generator, random_color_jitter_generator, random_crop_generator, random_crop_size_generator, random_cutmix_generator, random_mixup_generator, random_motion_blur_generator, random_perspective_generator, random_posterize_generator, random_prob_generator, random_rectangles_params_generator, random_rotation_generator, random_sharpness_generator, random_solarize_generator, ) from kornia.testing import assert_close from kornia.utils._compat import torch_version_geq class RandomGeneratorBaseTests: def test_valid_param_combinations(self, device, dtype): raise NotImplementedError def test_invalid_param_combinations(self, device, dtype): raise NotImplementedError def test_random_gen(self, device, dtype): raise NotImplementedError def test_same_on_batch(self, device, dtype): raise NotImplementedError class TestRandomProbGen(RandomGeneratorBaseTests): @pytest.mark.parametrize('p', [0.0, 0.5, 1.0]) @pytest.mark.parametrize('batch_size', [0, 1, 8]) @pytest.mark.parametrize('same_on_batch', [True, False]) def test_valid_param_combinations(self, p, batch_size, same_on_batch, device, dtype): random_prob_generator(batch_size=batch_size, p=p, same_on_batch=same_on_batch) @pytest.mark.parametrize( 'p', [ # Should be failed if p > 1. or p < 0. (-1.0), (2.0), ], ) def test_invalid_param_combinations(self, p, device, dtype): with pytest.raises(Exception): random_prob_generator(batch_size=8, p=p) @pytest.mark.parametrize( 'p,expected', [(0.0, [False] * 8), (0.5, [False, False, True, False, True, False, True, False]), (1.0, [True] * 8)], ) def test_random_gen(self, p, expected, device, dtype): torch.manual_seed(42) batch_size = 8 res = random_prob_generator(batch_size=batch_size, p=p) assert (res == torch.tensor(expected)).long().sum() == batch_size @pytest.mark.parametrize("seed,expected", [(42, [False] * 8), (0, [True] * 8)]) def test_same_on_batch(self, seed, expected, device, dtype): torch.manual_seed(seed) batch_size = 8 res = random_prob_generator(batch_size=batch_size, p=0.5, same_on_batch=True) assert (res == torch.tensor(expected)).long().sum() == batch_size class TestColorJitterGen(RandomGeneratorBaseTests): @pytest.mark.parametrize('brightness', [None, torch.tensor([0.8, 1.2])]) @pytest.mark.parametrize('contrast', [None, torch.tensor([0.8, 1.2])]) @pytest.mark.parametrize('saturation', [None, torch.tensor([0.8, 1.2])]) @pytest.mark.parametrize('hue', [None, torch.tensor([-0.1, 0.1])]) @pytest.mark.parametrize('batch_size', [0, 1, 8]) @pytest.mark.parametrize('same_on_batch', [True, False]) def test_valid_param_combinations( self, brightness, contrast, saturation, hue, batch_size, same_on_batch, device, dtype ): random_color_jitter_generator( batch_size, brightness.to(device=device, dtype=dtype) if brightness is not None else None, contrast.to(device=device, dtype=dtype) if contrast is not None else None, saturation.to(device=device, dtype=dtype) if saturation is not None else None, hue.to(device=device, dtype=dtype) if hue is not None else None, same_on_batch, ) @pytest.mark.parametrize( 'brightness,contrast,saturation,hue', [ # Should be failed if value out of bounds or tensor.shape != [1, 2] (torch.tensor([-1.0, 2.0]), None, None, None), (torch.tensor([0.0, 3.0]), None, None, None), (torch.tensor(0.0), None, None, None), (torch.tensor([0.0]), None, None, None), (torch.tensor([0.0, 1.0, 2.0]), None, None, None), (None, torch.tensor([-1.0, 2.0]), None, None), (None, torch.tensor(0.0), None, None), (None, torch.tensor([0.0]), None, None), (None, torch.tensor([0.0, 1.0, 2.0]), None, None), (None, None, torch.tensor([-1.0, 2.0]), None), (None, None, torch.tensor(0.0), None), (None, None, torch.tensor([0.0]), None), (None, None, torch.tensor([0.0, 1.0, 2.0]), None), (None, None, None, torch.tensor([-1.0, 0.0])), (None, None, None, torch.tensor([0, 1.0])), (None, None, None, torch.tensor(0.0)), (None, None, None, torch.tensor([0.0])), (None, None, None, torch.tensor([0.0, 1.0, 2.0])), ], ) def test_invalid_param_combinations(self, brightness, contrast, saturation, hue, device, dtype): with pytest.raises(Exception): random_color_jitter_generator( 8, brightness.to(device=device, dtype=dtype) if brightness is not None else None, contrast.to(device=device, dtype=dtype) if contrast is not None else None, saturation.to(device=device, dtype=dtype) if saturation is not None else None, hue.to(device=device, dtype=dtype) if hue is not None else None, ) def test_random_gen(self, device, dtype): # TODO(jian): crashes with pytorch 1.10, cuda and fp64 if torch_version_geq(1, 10) and "cuda" in str(device): pytest.skip("AssertionError: Tensor-likes are not close!") torch.manual_seed(42) batch_size = 8 jitter_params = random_color_jitter_generator( batch_size, brightness=torch.tensor([0.8, 1.2], device=device, dtype=dtype), contrast=torch.tensor([0.7, 1.3], device=device, dtype=dtype), saturation=torch.tensor([0.6, 1.4], device=device, dtype=dtype), hue=torch.tensor([-0.1, 0.1], device=device, dtype=dtype), ) expected_jitter_params = { 'brightness_factor': torch.tensor( [1.1529, 1.1660, 0.9531, 1.1837, 0.9562, 1.0404, 0.9026, 1.1175], device=device, dtype=dtype ), 'contrast_factor': torch.tensor( [1.2645, 0.7799, 1.2608, 1.0561, 1.2216, 1.0406, 1.1447, 0.9576], device=device, dtype=dtype ), 'hue_factor': torch.tensor( [0.0771, 0.0148, -0.0467, 0.0255, -0.0461, -0.0117, -0.0406, 0.0663], device=device, dtype=dtype ), 'saturation_factor': torch.tensor( [0.6843, 0.8156, 0.8871, 0.7595, 1.0378, 0.6049, 1.3612, 0.6602], device=device, dtype=dtype ), 'order': torch.tensor([3, 2, 0, 1], device=device, dtype=dtype), } assert set(list(jitter_params.keys())) == { 'brightness_factor', 'contrast_factor', 'hue_factor', 'saturation_factor', 'order', }, "Redundant keys found apart from \ 'brightness_factor', 'contrast_factor', 'hue_factor', 'saturation_factor', 'order'" assert_close( jitter_params['brightness_factor'], expected_jitter_params['brightness_factor'], rtol=1e-4, atol=1e-4 ) assert_close(jitter_params['contrast_factor'], expected_jitter_params['contrast_factor'], rtol=1e-4, atol=1e-4) assert_close(jitter_params['hue_factor'], expected_jitter_params['hue_factor'], rtol=1e-4, atol=1e-4) assert_close( jitter_params['saturation_factor'], expected_jitter_params['saturation_factor'], rtol=1e-4, atol=1e-4 ) assert_close(jitter_params['order'].to(dtype), expected_jitter_params['order'], rtol=1e-4, atol=1e-4) def test_same_on_batch(self, device, dtype): torch.manual_seed(42) batch_size = 8 jitter_params = random_color_jitter_generator( batch_size, brightness=torch.tensor([0.8, 1.2], device=device, dtype=dtype), contrast=torch.tensor([0.7, 1.3], device=device, dtype=dtype), saturation=torch.tensor([0.6, 1.4], device=device, dtype=dtype), hue=torch.tensor([-0.1, 0.1], device=device, dtype=dtype), same_on_batch=True, ) expected_res = { 'brightness_factor': torch.tensor([1.1529] * batch_size, device=device, dtype=dtype), 'contrast_factor': torch.tensor([1.2490] * batch_size, device=device, dtype=dtype), 'hue_factor': torch.tensor([-0.0234] * batch_size, device=device, dtype=dtype), 'saturation_factor': torch.tensor([1.3674] * batch_size, device=device, dtype=dtype), 'order': torch.tensor([2, 3, 0, 1], device=device, dtype=dtype), } assert_close(jitter_params['brightness_factor'], expected_res['brightness_factor'], rtol=1e-4, atol=1e-4) assert_close(jitter_params['contrast_factor'], expected_res['contrast_factor'], rtol=1e-4, atol=1e-4) assert_close(jitter_params['hue_factor'], expected_res['hue_factor'], rtol=1e-4, atol=1e-4) assert_close(jitter_params['saturation_factor'], expected_res['saturation_factor'], rtol=1e-4, atol=1e-4) assert_close(jitter_params['order'].to(dtype), expected_res['order'], rtol=1e-4, atol=1e-4) class TestRandomPerspectiveGen(RandomGeneratorBaseTests): @pytest.mark.parametrize('height,width', [(200, 200)]) @pytest.mark.parametrize('distortion_scale', [torch.tensor(0.0), torch.tensor(0.5), torch.tensor(1.0)]) @pytest.mark.parametrize('batch_size', [0, 1, 8]) @pytest.mark.parametrize('same_on_batch', [True, False]) def test_valid_param_combinations(self, height, width, distortion_scale, batch_size, same_on_batch, device, dtype): random_perspective_generator( batch_size=8, height=height, width=width, distortion_scale=distortion_scale.to(device=device, dtype=dtype), same_on_batch=same_on_batch, ) @pytest.mark.parametrize( 'height,width,distortion_scale', [ # Should be failed if distortion_scale > 1. or distortion_scale < 0. (-100, 100, torch.tensor(0.5)), (100, -100, torch.tensor(0.5)), (100, 100, torch.tensor(-0.5)), (100, 100, torch.tensor(1.5)), (100, 100, torch.tensor([0.0, 0.5])), ], ) def test_invalid_param_combinations(self, height, width, distortion_scale, device, dtype): with pytest.raises(Exception): random_perspective_generator( batch_size=8, height=height, width=width, distortion_scale=distortion_scale.to(device=device, dtype=dtype), ) def test_random_gen(self, device, dtype): torch.manual_seed(42) batch_size = 2 res = random_perspective_generator(batch_size, 200, 200, torch.tensor(0.5, device=device, dtype=dtype)) expected = dict( start_points=torch.tensor( [ [[0.0, 0.0], [199.0, 0.0], [199.0, 199.0], [0.0, 199.0]], [[0.0, 0.0], [199.0, 0.0], [199.0, 199.0], [0.0, 199.0]], ], device=device, dtype=dtype, ), end_points=torch.tensor( [ [[44.1135, 45.7502], [179.8568, 47.9653], [179.4776, 168.9552], [12.8286, 159.3179]], [[47.0386, 6.6593], [152.2701, 29.6790], [155.5298, 170.6142], [37.0547, 177.5298]], ], device=device, dtype=dtype, ), ) assert res.keys() == expected.keys() assert_close(res['start_points'], expected['start_points']) assert_close(res['end_points'], expected['end_points']) def test_same_on_batch(self, device, dtype): torch.manual_seed(42) batch_size = 2 res = random_perspective_generator( batch_size, 200, 200, torch.tensor(0.5, device=device, dtype=dtype), same_on_batch=True ) expected = dict( start_points=torch.tensor( [[[0.0, 0.0], [199.0, 0.0], [199.0, 199.0], [0.0, 199.0]]], device=device, dtype=dtype ).repeat(2, 1, 1), end_points=torch.tensor( [[[44.1135, 45.7502], [179.8568, 47.9653], [179.4776, 168.9552], [12.8286, 159.3179]]], device=device, dtype=dtype, ).repeat(2, 1, 1), ) assert res.keys() == expected.keys() assert_close(res['start_points'], expected['start_points']) assert_close(res['end_points'], expected['end_points']) class TestRandomAffineGen(RandomGeneratorBaseTests): @pytest.mark.parametrize('batch_size', [0, 1, 4]) @pytest.mark.parametrize('height', [200]) @pytest.mark.parametrize('width', [300]) @pytest.mark.parametrize('degrees', [torch.tensor([0, 30])]) @pytest.mark.parametrize('translate', [None, torch.tensor([0.1, 0.1])]) @pytest.mark.parametrize('scale', [None, torch.tensor([0.7, 1.2])]) @pytest.mark.parametrize('shear', [None, torch.tensor([[0, 20], [0, 20]])]) @pytest.mark.parametrize('same_on_batch', [True, False]) def test_valid_param_combinations( self, batch_size, height, width, degrees, translate, scale, shear, same_on_batch, device, dtype ): random_affine_generator( batch_size=batch_size, height=height, width=width, degrees=degrees.to(device=device, dtype=dtype), translate=translate.to(device=device, dtype=dtype) if translate is not None else None, scale=scale.to(device=device, dtype=dtype) if scale is not None else None, shear=shear.to(device=device, dtype=dtype) if shear is not None else None, same_on_batch=same_on_batch, ) @pytest.mark.parametrize( 'height,width,degrees,translate,scale,shear', [ (-100, 100, torch.tensor([10, 20]), None, None, None), (100, -100, torch.tensor([10, 20]), None, None, None), (100, 100, 0.5, None, None, None), (100, 100, torch.tensor([10, 20, 30]), None, None, None), (100, 100, torch.tensor([10, 20]), torch.tensor([0.1]), None, None), (10, 10, torch.tensor([1, 2]), torch.tensor([0.1, 0.2, 0.3]), None, None), (100, 100, torch.tensor([10, 20]), None, torch.tensor([1]), None), (100, 100, torch.tensor([10, 20]), None, torch.tensor([1, 2, 3]), None), (100, 100, torch.tensor([10, 20]), None, None, torch.tensor([1])), (100, 100, torch.tensor([10, 20]), None, None, torch.tensor([1, 2])), (10, 10, torch.tensor([1, 2]), None, None, torch.tensor([1, 2, 3])), (10, 10, torch.tensor([1, 2]), None, None, torch.tensor([1, 2, 3, 4])), (10, 10, torch.tensor([1, 2]), None, None, torch.tensor([1, 2, 3, 4, 5])), ], ) def test_invalid_param_combinations(self, height, width, degrees, translate, scale, shear, device, dtype): with pytest.raises(Exception): random_affine_generator( batch_size=8, height=height, width=width, degrees=degrees.to(device=device, dtype=dtype), translate=translate.to(device=device, dtype=dtype) if translate is not None else None, scale=scale.to(device=device, dtype=dtype) if scale is not None else None, shear=shear.to(device=device, dtype=dtype) if shear is not None else None, ) def test_random_gen(self, device, dtype): torch.manual_seed(42) degrees = torch.tensor([10, 20], device=device, dtype=dtype) translate = torch.tensor([0.1, 0.1], device=device, dtype=dtype) scale = torch.tensor([0.7, 1.2], device=device, dtype=dtype) shear = torch.tensor([[10, 20], [10, 20]], device=device, dtype=dtype) res = random_affine_generator( batch_size=2, height=200, width=200, degrees=degrees, translate=translate, scale=scale, shear=shear, same_on_batch=False, ) expected = dict( translations=torch.tensor([[-4.3821, -9.7371], [4.0358, 11.7457]], device=device, dtype=dtype), center=torch.tensor([[99.5000, 99.5000], [99.5000, 99.5000]], device=device, dtype=dtype), scale=torch.tensor([[0.8914, 0.8914], [1.1797, 1.1797]], device=device, dtype=dtype), angle=torch.tensor([18.8227, 19.1500], device=device, dtype=dtype), sx=torch.tensor([19.4077, 11.3319], device=device, dtype=dtype), sy=torch.tensor([19.3460, 15.9358], device=device, dtype=dtype), ) assert res.keys() == expected.keys() assert_close(res['translations'], expected['translations'], rtol=1e-4, atol=1e-4) assert_close(res['center'], expected['center'], rtol=1e-4, atol=1e-4) assert_close(res['scale'], expected['scale'], rtol=1e-4, atol=1e-4) assert_close(res['angle'], expected['angle'], rtol=1e-4, atol=1e-4) assert_close(res['sx'], expected['sx'], rtol=1e-4, atol=1e-4) assert_close(res['sy'], expected['sy'], rtol=1e-4, atol=1e-4) def test_same_on_batch(self, device, dtype): torch.manual_seed(42) degrees = torch.tensor([10, 20], device=device, dtype=dtype) translate = torch.tensor([0.1, 0.1], device=device, dtype=dtype) scale = torch.tensor([0.7, 1.2], device=device, dtype=dtype) shear = torch.tensor([[10, 20], [10, 20]], device=device, dtype=dtype) res = random_affine_generator( batch_size=2, height=200, width=200, degrees=degrees, translate=translate, scale=scale, shear=shear, same_on_batch=True, ) expected = dict( translations=torch.tensor([[-4.6854, 18.3722], [-4.6854, 18.3722]], device=device, dtype=dtype), center=torch.tensor([[99.5000, 99.5000], [99.5000, 99.5000]], device=device, dtype=dtype), scale=torch.tensor([[1.1575, 1.1575], [1.1575, 1.1575]], device=device, dtype=dtype), angle=torch.tensor([18.8227, 18.8227], device=device, dtype=dtype), sx=torch.tensor([13.9045, 13.9045], device=device, dtype=dtype), sy=torch.tensor([16.0090, 16.0090], device=device, dtype=dtype), ) assert res.keys() == expected.keys() assert_close(res['translations'], expected['translations'], rtol=1e-4, atol=1e-4) assert_close(res['center'], expected['center'], rtol=1e-4, atol=1e-4) assert_close(res['scale'], expected['scale'], rtol=1e-4, atol=1e-4) assert_close(res['angle'], expected['angle'], rtol=1e-4, atol=1e-4) assert_close(res['sx'], expected['sx'], rtol=1e-4, atol=1e-4) assert_close(res['sy'], expected['sy'], rtol=1e-4, atol=1e-4) class TestRandomRotationGen(RandomGeneratorBaseTests): @pytest.mark.parametrize('batch_size', [0, 1, 8]) @pytest.mark.parametrize('degrees', [torch.tensor([0, 30])]) @pytest.mark.parametrize('same_on_batch', [True, False]) def test_valid_param_combinations(self, batch_size, degrees, same_on_batch, device, dtype): random_rotation_generator( batch_size=batch_size, degrees=degrees.to(device=device, dtype=dtype), same_on_batch=same_on_batch ) @pytest.mark.parametrize('degrees', [(torch.tensor(10)), (torch.tensor([10])), (torch.tensor([10, 20, 30]))]) def test_invalid_param_combinations(self, degrees, device, dtype): batch_size = 8 with pytest.raises(Exception): random_rotation_generator(batch_size=batch_size, degrees=degrees.to(device=device, dtype=dtype)) def test_random_gen(self, device, dtype): torch.manual_seed(42) degrees = torch.tensor([10, 20]) res = random_rotation_generator( batch_size=2, degrees=degrees.to(device=device, dtype=dtype), same_on_batch=False ) expected = dict(degrees=torch.tensor([18.8227, 19.1500], device=device, dtype=dtype)) assert res.keys() == expected.keys() assert_close(res['degrees'], expected['degrees']) def test_same_on_batch(self, device, dtype): torch.manual_seed(42) degrees = torch.tensor([10, 20]) res = random_rotation_generator( batch_size=2, degrees=degrees.to(device=device, dtype=dtype), same_on_batch=True ) expected = dict(degrees=torch.tensor([18.8227, 18.8227], device=device, dtype=dtype)) assert res.keys() == expected.keys() assert_close(res['degrees'], expected['degrees']) class TestRandomCropGen(RandomGeneratorBaseTests): @pytest.mark.parametrize('batch_size', [0, 2]) @pytest.mark.parametrize('input_size', [(200, 200)]) @pytest.mark.parametrize('size', [(100, 100), torch.tensor([50, 50])]) @pytest.mark.parametrize('resize_to', [None, (100, 100)]) @pytest.mark.parametrize('same_on_batch', [True, False]) def test_valid_param_combinations(self, batch_size, input_size, size, resize_to, same_on_batch, device, dtype): if isinstance(size, torch.Tensor): size = size.repeat(batch_size, 1).to(device=device, dtype=dtype) random_crop_generator( batch_size=batch_size, input_size=input_size, size=size, resize_to=resize_to, same_on_batch=same_on_batch, device=device, dtype=dtype, ) @pytest.mark.parametrize( 'input_size,size,resize_to', [((-300, 300), (200, 200), (100, 100)), ((200, 200), torch.tensor([50, 50]), (100, 100))], ) def test_invalid_param_combinations(self, input_size, size, resize_to, device, dtype): batch_size = 2 with pytest.raises(Exception): random_crop_generator( batch_size=batch_size, input_size=input_size, size=size.to(device=device, dtype=dtype) if isinstance(size, torch.Tensor) else size, resize_to=resize_to, ) def test_random_gen(self, device, dtype): torch.manual_seed(42) res = random_crop_generator( batch_size=2, input_size=(100, 100), size=torch.tensor([[50, 60], [70, 80]], device=device, dtype=dtype), resize_to=(200, 200), ) expected = dict( src=torch.tensor( [[[36, 19], [95, 19], [95, 68], [36, 68]], [[19, 29], [98, 29], [98, 98], [19, 98]]], device=device, dtype=dtype, ), dst=torch.tensor( [[[0, 0], [199, 0], [199, 199], [0, 199]], [[0, 0], [199, 0], [199, 199], [0, 199]]], device=device, dtype=dtype, ), input_size=torch.tensor([[100, 100], [100, 100]], device=device, dtype=torch.long), ) assert res.keys() == expected.keys() assert_close(res['src'], expected['src']) assert_close(res['dst'], expected['dst']) def test_same_on_batch(self, device, dtype): torch.manual_seed(42) res = random_crop_generator( batch_size=2, input_size=(100, 100), size=torch.tensor([[50, 60], [70, 80]], device=device, dtype=dtype), resize_to=(200, 200), same_on_batch=True, ) expected = dict( src=torch.tensor( [[[36, 46], [95, 46], [95, 95], [36, 95]], [[36, 46], [115, 46], [115, 115], [36, 115]]], device=device, dtype=dtype, ), dst=torch.tensor( [[[0, 0], [199, 0], [199, 199], [0, 199]], [[0, 0], [199, 0], [199, 199], [0, 199]]], device=device, dtype=dtype, ), input_size=torch.tensor([[100, 100], [100, 100]], device=device, dtype=torch.long), ) assert res.keys() == expected.keys() assert_close(res['src'], expected['src']) assert_close(res['dst'], expected['dst']) class TestRandomCropSizeGen(RandomGeneratorBaseTests): @pytest.mark.parametrize('batch_size', [0, 1, 8]) @pytest.mark.parametrize('size', [(200, 200)]) @pytest.mark.parametrize('scale', [torch.tensor([0.7, 1.3])]) @pytest.mark.parametrize('ratio', [torch.tensor([0.9, 1.1])]) @pytest.mark.parametrize('same_on_batch', [True, False]) def test_valid_param_combinations(self, batch_size, size, scale, ratio, same_on_batch, device, dtype): random_crop_size_generator( batch_size=batch_size, size=size, scale=scale.to(device=device, dtype=dtype), ratio=ratio.to(device=device, dtype=dtype), same_on_batch=same_on_batch, ) @pytest.mark.parametrize( 'size,scale,ratio', [ ((100), torch.tensor([0.7, 1.3]), torch.tensor([0.9, 1.1])), ((100, 100, 100), torch.tensor([0.7, 1.3]), torch.tensor([0.9, 1.1])), ((100, 100), torch.tensor([0.7]), torch.tensor([0.9, 1.1])), ((100, 100), torch.tensor([0.7, 1.3, 1.5]), torch.tensor([0.9, 1.1])), ((100, 100), torch.tensor([0.7, 1.3]), torch.tensor([0.9])), ((100, 100), torch.tensor([0.7, 1.3]), torch.tensor([0.9, 1.1, 1.3])), ], ) def test_invalid_param_combinations(self, size, scale, ratio, device, dtype): batch_size = 2 with pytest.raises(Exception): random_crop_size_generator( batch_size=batch_size, size=size, scale=scale.to(device=device, dtype=dtype), ratio=ratio.to(device=device, dtype=dtype), same_on_batch=same_on_batch, # noqa: F821 raises NameError ) def test_random_gen(self, device, dtype): torch.manual_seed(42) res = random_crop_size_generator( batch_size=8, size=(100, 100), scale=torch.tensor([0.7, 1.3], device=device, dtype=dtype), ratio=torch.tensor([0.9, 1.1], device=device, dtype=dtype), same_on_batch=False, ) expected = dict( size=torch.tensor( [[99, 94], [91, 95], [90, 96], [87, 86], [94, 98], [87, 81], [85, 93], [83, 90]], device=device, dtype=dtype, ) ) assert res.keys() == expected.keys() assert_close(res['size'], expected['size']) res = random_crop_size_generator( batch_size=100, size=(100, 100), scale=torch.tensor([0.999, 1.0], device=device, dtype=dtype), ratio=torch.tensor([1.0, 1.0], device=device, dtype=dtype), same_on_batch=False, ) expected = dict(size=torch.tensor([[100, 100]], device=device, dtype=dtype).repeat(100, 1)) assert res.keys() == expected.keys() assert_close(res['size'], expected['size']) def test_same_on_batch(self, device, dtype): torch.manual_seed(42) res = random_crop_size_generator( batch_size=8, size=(100, 100), scale=torch.tensor([0.7, 1.3], device=device, dtype=dtype), ratio=torch.tensor([0.9, 1.1], device=device, dtype=dtype), same_on_batch=True, ) expected = dict( size=torch.tensor( [[99, 95], [99, 95], [99, 95], [99, 95], [99, 95], [99, 95], [99, 95], [99, 95]], device=device, dtype=dtype, ) ) assert res.keys() == expected.keys() assert_close(res['size'], expected['size']) class TestRandomRectangleGen(RandomGeneratorBaseTests): @pytest.mark.parametrize('batch_size', [0, 1, 8]) @pytest.mark.parametrize('height', [200]) @pytest.mark.parametrize('width', [300]) @pytest.mark.parametrize('scale', [torch.tensor([0.7, 1.1])]) @pytest.mark.parametrize('ratio', [torch.tensor([0.7, 1.1])]) @pytest.mark.parametrize('value', [0]) @pytest.mark.parametrize('same_on_batch', [True, False]) def test_valid_param_combinations( self, batch_size, height, width, scale, ratio, value, same_on_batch, device, dtype ): random_rectangles_params_generator( batch_size=batch_size, height=height, width=width, scale=scale.to(device=device, dtype=dtype), ratio=ratio.to(device=device, dtype=dtype), value=value, same_on_batch=same_on_batch, ) @pytest.mark.parametrize( 'height,width,scale,ratio,value', [ (-100, 100, torch.tensor([0.7, 1.3]), torch.tensor([0.7, 1.3]), 0), (100, -100, torch.tensor([0.7, 1.3]), torch.tensor([0.7, 1.3]), 0), (100, -100, torch.tensor([0.7]), torch.tensor([0.7, 1.3]), 0), (100, 100, torch.tensor([0.7, 1.3, 1.5]), torch.tensor([0.7, 1.3]), 0), (100, 100, torch.tensor([0.7, 1.3]), torch.tensor([0.7]), 0), (100, 100, torch.tensor([0.7, 1.3]), torch.tensor([0.7, 1.3, 1.5]), 0), (100, 100, torch.tensor([0.7, 1.3]), torch.tensor([0.7, 1.3]), -1), (100, 100, torch.tensor([0.7, 1.3]), torch.tensor([0.7, 1.3]), 2), (100, 100, torch.tensor([0.5, 0.7]), torch.tensor([0.7, 0.9]), torch.tensor(0.5)), ], ) def test_invalid_param_combinations(self, height, width, scale, ratio, value, device, dtype): batch_size = 8 with pytest.raises(Exception): random_rectangles_params_generator( batch_size=batch_size, height=height, width=width, scale=scale.to(device=device, dtype=dtype), ratio=ratio.to(device=device, dtype=dtype), value=value, same_on_batch=same_on_batch, # noqa: F821 raises NameError ) def test_random_gen(self, device, dtype): torch.manual_seed(42) width, height = 100, 150 scale = torch.tensor([0.7, 1.3], device=device, dtype=dtype) ratio = torch.tensor([0.7, 1.3], device=device, dtype=dtype) value = 0.5 res = random_rectangles_params_generator( batch_size=2, height=height, width=width, scale=scale, ratio=ratio, value=value, same_on_batch=False ) expected = dict( widths=torch.tensor([100, 100], device=device, dtype=dtype), heights=torch.tensor([0, 0], device=device, dtype=dtype), xs=torch.tensor([0, 0], device=device, dtype=dtype), ys=torch.tensor([6, 8], device=device, dtype=dtype), values=torch.tensor([0.5000, 0.5000], device=device, dtype=dtype), ) assert res.keys() == expected.keys() assert_close(res['widths'], expected['widths']) assert_close(res['widths'], expected['widths']) assert_close(res['xs'], expected['xs']) assert_close(res['ys'], expected['ys']) assert_close(res['values'], expected['values']) def test_same_on_batch(self, device, dtype): torch.manual_seed(42) width, height = 100, 150 scale = torch.tensor([0.7, 1.3], device=device, dtype=dtype) ratio = torch.tensor([0.7, 1.3], device=device, dtype=dtype) value = 0.5 res = random_rectangles_params_generator( batch_size=2, height=height, width=width, scale=scale, ratio=ratio, value=value, same_on_batch=True ) expected = dict( widths=torch.tensor([100, 100], device=device, dtype=dtype), heights=torch.tensor([0, 0], device=device, dtype=dtype), xs=torch.tensor([0, 0], device=device, dtype=dtype), ys=torch.tensor([10, 10], device=device, dtype=dtype), values=torch.tensor([0.5000, 0.5000], device=device, dtype=dtype), ) assert res.keys() == expected.keys() assert_close(res['widths'], expected['widths']) assert_close(res['widths'], expected['widths']) assert_close(res['xs'], expected['xs']) assert_close(res['ys'], expected['ys']) assert_close(res['values'], expected['values']) class TestCenterCropGen(RandomGeneratorBaseTests): @pytest.mark.parametrize('batch_size', [0, 2]) @pytest.mark.parametrize('height', [200]) @pytest.mark.parametrize('width', [200]) @pytest.mark.parametrize('size', [(100, 100)]) def test_valid_param_combinations(self, batch_size, height, width, size, device, dtype): center_crop_generator(batch_size=batch_size, height=height, width=width, size=size) @pytest.mark.parametrize( 'height,width,size', [ (200, -200, (100, 100)), (-200, 200, (100, 100)), (100, 100, (120, 120)), (150, 100, (120, 120)), (100, 150, (120, 120)), ], ) def test_invalid_param_combinations(self, height, width, size, device, dtype): batch_size = 2 with pytest.raises(Exception): center_crop_generator(batch_size=batch_size, height=height, width=width, size=size) def test_random_gen(self, device, dtype): torch.manual_seed(42) res = center_crop_generator(batch_size=2, height=200, width=200, size=(120, 150)) expected = dict( src=torch.tensor( [[[25, 40], [174, 40], [174, 159], [25, 159]], [[25, 40], [174, 40], [174, 159], [25, 159]]], device=device, dtype=torch.long, ), dst=torch.tensor( [[[0, 0], [149, 0], [149, 119], [0, 119]], [[0, 0], [149, 0], [149, 119], [0, 119]]], device=device, dtype=torch.long, ), input_size=torch.tensor([[200, 200], [200, 200]], device=device, dtype=torch.long), ) assert res.keys() == expected.keys() assert_close(res['src'].to(device=device), expected['src']) assert_close(res['dst'].to(device=device), expected['dst']) def test_same_on_batch(self, device, dtype): pass class TestRandomMotionBlur(RandomGeneratorBaseTests): @pytest.mark.parametrize('batch_size', [0, 1, 8]) @pytest.mark.parametrize('kernel_size', [3, (3, 5)]) @pytest.mark.parametrize('angle', [torch.tensor([10, 30])]) @pytest.mark.parametrize('direction', [torch.tensor([-1, -1]), torch.tensor([1, 1])]) @pytest.mark.parametrize('same_on_batch', [True, False]) def test_valid_param_combinations(self, batch_size, kernel_size, angle, direction, same_on_batch, device, dtype): random_motion_blur_generator( batch_size=batch_size, kernel_size=kernel_size, angle=angle.to(device=device, dtype=dtype), direction=direction.to(device=device, dtype=dtype), same_on_batch=same_on_batch, ) @pytest.mark.parametrize( 'kernel_size,angle,direction', [ (4, torch.tensor([30, 100]), torch.tensor([-1, 1])), (1, torch.tensor([30, 100]), torch.tensor([-1, 1])), ((1, 2, 3), torch.tensor([30, 100]), torch.tensor([-1, 1])), (3, torch.tensor([30, 100]), torch.tensor([-2, 1])), (3, torch.tensor([30, 100]), torch.tensor([-1, 2])), ], ) def test_invalid_param_combinations(self, kernel_size, angle, direction, device, dtype): with pytest.raises(Exception): random_motion_blur_generator( batch_size=8, kernel_size=kernel_size, angle=angle.to(device=device, dtype=dtype), direction=direction.to(device=device, dtype=dtype), ) def test_random_gen(self, device, dtype): torch.manual_seed(42) angle = torch.tensor([30, 90]) direction = torch.tensor([-1, 1]) res = random_motion_blur_generator( batch_size=2, kernel_size=3, angle=angle.to(device=device, dtype=dtype), direction=direction.to(device=device, dtype=dtype), same_on_batch=False, ) expected = dict( ksize_factor=torch.tensor([3, 3], device=device, dtype=torch.int32), angle_factor=torch.tensor([82.9362, 84.9002], device=device, dtype=dtype), direction_factor=torch.tensor([-0.2343, 0.9186], device=device, dtype=dtype), ) assert res.keys() == expected.keys() assert_close(res['ksize_factor'], expected['ksize_factor'], rtol=1e-4, atol=1e-4) assert_close(res['angle_factor'], expected['angle_factor'], rtol=1e-4, atol=1e-4) assert_close(res['direction_factor'], expected['direction_factor'], rtol=1e-4, atol=1e-4) def test_same_on_batch(self, device, dtype): torch.manual_seed(42) angle = torch.tensor([30, 90]) direction = torch.tensor([-1, 1]) res = random_motion_blur_generator( batch_size=2, kernel_size=3, angle=angle.to(device=device, dtype=dtype), direction=direction.to(device=device, dtype=dtype), same_on_batch=True, ) expected = dict( ksize_factor=torch.tensor([3, 3], device=device, dtype=torch.int32), angle_factor=torch.tensor([82.9362, 82.9362], device=device, dtype=dtype), direction_factor=torch.tensor([0.8300, 0.8300], device=device, dtype=dtype), ) assert res.keys() == expected.keys() assert_close(res['ksize_factor'], expected['ksize_factor'], rtol=1e-4, atol=1e-4) assert_close(res['angle_factor'], expected['angle_factor'], rtol=1e-4, atol=1e-4) assert_close(res['direction_factor'], expected['direction_factor'], rtol=1e-4, atol=1e-4) class TestRandomSolarizeGen(RandomGeneratorBaseTests): @pytest.mark.parametrize('batch_size', [0, 1, 8]) @pytest.mark.parametrize('thresholds', [torch.tensor([0, 1]), torch.tensor([0.4, 0.6])]) @pytest.mark.parametrize('additions', [torch.tensor([-0.5, 0.5])]) @pytest.mark.parametrize('same_on_batch', [True, False]) def test_valid_param_combinations(self, batch_size, thresholds, additions, same_on_batch, device, dtype): random_solarize_generator( batch_size=batch_size, thresholds=thresholds.to(device=device, dtype=dtype), additions=additions.to(device=device, dtype=dtype), same_on_batch=same_on_batch, ) @pytest.mark.parametrize( 'thresholds,additions', [ (torch.tensor([0, 2]), torch.tensor([-0.5, 0.5])), (torch.tensor([-1, 1]), torch.tensor([-0.5, 0.5])), ([0, 1], torch.tensor([-0.5, 0.5])), (torch.tensor([0, 1]), torch.tensor([-0.5, 1])), (torch.tensor([0, 1]), torch.tensor([-1, 0.5])), (torch.tensor([0, 1]), [-0.5, 0.5]), ], ) def test_invalid_param_combinations(self, thresholds, additions, device, dtype): with pytest.raises(Exception): random_solarize_generator( batch_size=batch_size, # noqa: F821 raises NameError thresholds=thresholds.to(device=device, dtype=dtype), additions=additions.to(device=device, dtype=dtype), ) def test_random_gen(self, device, dtype): torch.manual_seed(42) batch_size = 8 res = random_solarize_generator( batch_size=batch_size, thresholds=torch.tensor([0, 1], device=device, dtype=dtype), additions=torch.tensor([-0.5, 0.5], device=device, dtype=dtype), same_on_batch=False, ) expected = dict( thresholds_factor=torch.tensor( [0.8823, 0.9150, 0.3829, 0.9593, 0.3904, 0.6009, 0.2566, 0.7936], device=device, dtype=dtype ), additions_factor=torch.tensor( [0.4408, -0.3668, 0.4346, 0.0936, 0.3694, 0.0677, 0.2411, -0.0706], device=device, dtype=dtype ), ) assert res.keys() == expected.keys() assert_close(res['thresholds_factor'], expected['thresholds_factor'], rtol=1e-4, atol=1e-4) assert_close(res['additions_factor'], expected['additions_factor'], rtol=1e-4, atol=1e-4) def test_same_on_batch(self, device, dtype): torch.manual_seed(42) batch_size = 8 res = random_solarize_generator( batch_size=batch_size, thresholds=torch.tensor([0, 1], device=device, dtype=dtype), additions=torch.tensor([-0.5, 0.5], device=device, dtype=dtype), same_on_batch=True, ) expected = dict( thresholds_factor=torch.tensor( [0.8823, 0.8823, 0.8823, 0.8823, 0.8823, 0.8823, 0.8823, 0.8823], device=device, dtype=dtype ), additions_factor=torch.tensor( [0.4150, 0.4150, 0.4150, 0.4150, 0.4150, 0.4150, 0.4150, 0.4150], device=device, dtype=dtype ), ) assert res.keys() == expected.keys() assert_close(res['thresholds_factor'], expected['thresholds_factor'], rtol=1e-4, atol=1e-4) assert_close(res['additions_factor'], expected['additions_factor'], rtol=1e-4, atol=1e-4) class TestRandomPosterizeGen(RandomGeneratorBaseTests): @pytest.mark.parametrize('batch_size', [0, 1, 8]) @pytest.mark.parametrize('bits', [torch.tensor([0, 8])]) @pytest.mark.parametrize('same_on_batch', [True, False]) def test_valid_param_combinations(self, batch_size, bits, same_on_batch, device, dtype): random_posterize_generator( batch_size=batch_size, bits=bits.to(device=device, dtype=dtype), same_on_batch=same_on_batch ) @pytest.mark.parametrize('bits', [(torch.tensor([-1, 1])), (torch.tensor([0, 9])), (torch.tensor([3])), ([0, 8])]) def test_invalid_param_combinations(self, bits, device, dtype): with pytest.raises(Exception): random_posterize_generator( batch_size=batch_size, # noqa: F821 raises NameError bits=bits.to(device=device, dtype=dtype), same_on_batch=same_on_batch, # noqa: F821 raises NameError ) def test_random_gen(self, device, dtype): torch.manual_seed(9) batch_size = 8 res = random_posterize_generator( batch_size=batch_size, bits=torch.tensor([0, 8], device=device, dtype=dtype), same_on_batch=False ) expected = dict(bits_factor=torch.tensor([5, 2, 3, 6, 7, 7, 2, 7], device=device, dtype=torch.int32)) assert res.keys() == expected.keys() assert_close(res['bits_factor'], expected['bits_factor'], rtol=1e-4, atol=1e-4) def test_same_on_batch(self, device, dtype): torch.manual_seed(9) batch_size = 8 res = random_posterize_generator( batch_size=batch_size, bits=torch.tensor([0, 8], device=device, dtype=dtype), same_on_batch=True ) expected = dict(bits_factor=torch.tensor([5, 5, 5, 5, 5, 5, 5, 5], device=device, dtype=torch.int32)) assert res.keys() == expected.keys() assert_close(res['bits_factor'], expected['bits_factor'], rtol=1e-4, atol=1e-4) class TestRandomSharpnessGen(RandomGeneratorBaseTests): @pytest.mark.parametrize('batch_size', [0, 1, 8]) @pytest.mark.parametrize('sharpness', [torch.tensor([0.0, 1.0])]) @pytest.mark.parametrize('same_on_batch', [True, False]) def test_valid_param_combinations(self, batch_size, sharpness, same_on_batch, device, dtype): random_sharpness_generator( batch_size=batch_size, sharpness=sharpness.to(device=device, dtype=dtype), same_on_batch=same_on_batch ) @pytest.mark.parametrize('sharpness', [(torch.tensor([-1, 5])), (torch.tensor([3])), ([0, 1.0])]) def test_invalid_param_combinations(self, sharpness, device, dtype): with pytest.raises(Exception): random_sharpness_generator( batch_size=batch_size, # noqa: F821 raises NameError sharpness=sharpness.to(device=device, dtype=dtype), same_on_batch=same_on_batch, # noqa: F821 raises NameError ) def test_random_gen(self, device, dtype): torch.manual_seed(42) batch_size = 8 res = random_sharpness_generator( batch_size=batch_size, sharpness=torch.tensor([0.0, 1.0], device=device, dtype=dtype), same_on_batch=False ) expected = dict( sharpness_factor=torch.tensor( [0.8823, 0.9150, 0.3829, 0.9593, 0.3904, 0.6009, 0.2566, 0.7936], device=device, dtype=dtype ) ) assert res.keys() == expected.keys() assert_close(res['sharpness_factor'], expected['sharpness_factor'], rtol=1e-4, atol=1e-4) def test_same_on_batch(self, device, dtype): torch.manual_seed(42) batch_size = 8 res = random_sharpness_generator( batch_size=batch_size, sharpness=torch.tensor([0.0, 1.0], device=device, dtype=dtype), same_on_batch=True ) expected = dict( sharpness_factor=torch.tensor( [0.8823, 0.8823, 0.8823, 0.8823, 0.8823, 0.8823, 0.8823, 0.8823], device=device, dtype=dtype ) ) assert res.keys() == expected.keys() assert_close(res['sharpness_factor'], expected['sharpness_factor'], rtol=1e-4, atol=1e-4) class TestRandomMixUpGen(RandomGeneratorBaseTests): @pytest.mark.parametrize('batch_size', [0, 1, 8]) @pytest.mark.parametrize('p', [0.0, 0.5, 1.0]) @pytest.mark.parametrize('lambda_val', [None, torch.tensor([0.0, 1.0])]) @pytest.mark.parametrize('same_on_batch', [True, False]) def test_valid_param_combinations(self, batch_size, p, lambda_val, same_on_batch, device, dtype): random_mixup_generator( batch_size=batch_size, p=p, lambda_val=lambda_val.to(device=device, dtype=dtype) if isinstance(lambda_val, (torch.Tensor)) else lambda_val, same_on_batch=same_on_batch, ) @pytest.mark.parametrize( 'lambda_val', [(torch.tensor([-1, 1])), (torch.tensor([0, 2])), (torch.tensor([0, 0.5, 1])), ([0.0, 1.0])] ) def test_invalid_param_combinations(self, lambda_val, device, dtype): with pytest.raises(Exception): random_mixup_generator(batch_size=8, lambda_val=lambda_val.to(device=device, dtype=dtype)) def test_random_gen(self, device, dtype): torch.manual_seed(42) batch_size = 8 res = random_mixup_generator( batch_size=batch_size, p=0.5, lambda_val=torch.tensor([0.0, 1.0], device=device, dtype=dtype), same_on_batch=False, ) expected = dict( mixup_pairs=torch.tensor([6, 1, 0, 7, 2, 5, 3, 4], device=device, dtype=torch.long), mixup_lambdas=torch.tensor( [0.0000, 0.0000, 0.5739, 0.0000, 0.6274, 0.0000, 0.4414, 0.0000], device=device, dtype=dtype ), ) assert res.keys() == expected.keys() assert_close(res['mixup_pairs'], expected['mixup_pairs'], rtol=1e-4, atol=1e-4) assert_close(res['mixup_lambdas'], expected['mixup_lambdas'], rtol=1e-4, atol=1e-4) def test_same_on_batch(self, device, dtype): torch.manual_seed(9) batch_size = 8 res = random_mixup_generator( batch_size=batch_size, p=0.9999, lambda_val=torch.tensor([0.0, 1.0], device=device, dtype=dtype), same_on_batch=True, ) expected = dict( mixup_pairs=torch.tensor([4, 6, 7, 5, 0, 1, 3, 2], device=device, dtype=torch.long), mixup_lambdas=torch.tensor( [0.3804, 0.3804, 0.3804, 0.3804, 0.3804, 0.3804, 0.3804, 0.3804], device=device, dtype=dtype ), ) assert res.keys() == expected.keys() assert_close(res['mixup_pairs'], expected['mixup_pairs'], rtol=1e-4, atol=1e-4) assert_close(res['mixup_lambdas'], expected['mixup_lambdas'], rtol=1e-4, atol=1e-4) class TestRandomCutMixGen(RandomGeneratorBaseTests): @pytest.mark.parametrize('batch_size', [0, 1, 8]) @pytest.mark.parametrize('p', [0, 0.5, 1.0]) @pytest.mark.parametrize('width,height', [(200, 200)]) @pytest.mark.parametrize('num_mix', [1, 3]) @pytest.mark.parametrize('beta', [None, torch.tensor(1e-15), torch.tensor(1.0)]) @pytest.mark.parametrize('cut_size', [None, torch.tensor([0.0, 1.0]), torch.tensor([0.3, 0.6])]) @pytest.mark.parametrize('same_on_batch', [True, False]) def test_valid_param_combinations( self, batch_size, p, width, height, num_mix, beta, cut_size, same_on_batch, device, dtype ): random_cutmix_generator( batch_size=batch_size, p=p, width=width, height=height, num_mix=num_mix, beta=beta.to(device=device, dtype=dtype) if isinstance(beta, (torch.Tensor)) else beta, cut_size=cut_size.to(device=device, dtype=dtype) if isinstance(cut_size, (torch.Tensor)) else cut_size, same_on_batch=same_on_batch, ) @pytest.mark.parametrize( 'width,height,num_mix,beta,cut_size', [ (200, -200, 1, None, None), (-200, 200, 1, None, None), (200, 200, 0, None, None), (200, 200, 1.5, None, None), (200, 200, 1, torch.tensor([0.0, 1.0]), None), (200, 200, 1, None, torch.tensor([-1.0, 1.0])), (200, 200, 1, None, torch.tensor([0.0, 2.0])), ], ) @pytest.mark.parametrize('same_on_batch', [True, False]) def test_invalid_param_combinations(self, width, height, num_mix, beta, cut_size, same_on_batch, device, dtype): with pytest.raises(Exception): random_cutmix_generator( batch_size=8, p=0.5, width=width, height=height, num_mix=num_mix, beta=beta.to(device=device, dtype=dtype) if isinstance(beta, (torch.Tensor)) else beta, cut_size=beta.to(device=device, dtype=dtype) if isinstance(cut_size, (torch.Tensor)) else cut_size, same_on_batch=same_on_batch, ) def test_random_gen(self, device, dtype): torch.manual_seed(42) batch_size = 2 res = random_cutmix_generator( batch_size=batch_size, width=200, height=200, p=0.5, num_mix=1, beta=torch.tensor(1.0, device=device, dtype=dtype), cut_size=torch.tensor([0.0, 1.0], device=device, dtype=dtype), same_on_batch=False, ) expected = dict( mix_pairs=torch.tensor([[0, 1]], device=device, dtype=torch.long), crop_src=torch.tensor( [[[[71, 108], [70, 108], [70, 107], [71, 107]], [[39, 1], [38, 1], [38, 0], [39, 0]]]], device=device, dtype=dtype, ), ) assert res.keys() == expected.keys() assert_close(res['mix_pairs'], expected['mix_pairs'], rtol=1e-4, atol=1e-4) assert_close(res['crop_src'], expected['crop_src'], rtol=1e-4, atol=1e-4) def test_same_on_batch(self, device, dtype): torch.manual_seed(42) batch_size = 2 res = random_cutmix_generator( batch_size=batch_size, width=200, height=200, p=0.5, num_mix=1, beta=torch.tensor(1.0, device=device, dtype=dtype), cut_size=torch.tensor([0.0, 1.0], device=device, dtype=dtype), same_on_batch=True, ) expected = dict( mix_pairs=torch.tensor([[1, 0]], device=device, dtype=torch.long), crop_src=torch.tensor( [[[[114, 53], [113, 53], [113, 52], [114, 52]], [[114, 53], [113, 53], [113, 52], [114, 52]]]], device=device, dtype=dtype, ), ) assert res.keys() == expected.keys() assert_close(res['mix_pairs'], expected['mix_pairs'], rtol=1e-4, atol=1e-4) assert_close(res['crop_src'], expected['crop_src'], rtol=1e-4, atol=1e-4)