def test_forward(self, shape, padding, patchwise_apply, same_on_batch, keepdim, random_apply, device, dtype): torch.manual_seed(11) try: # skip wrong param settings. seq = K.PatchSequential( K.color.RgbToBgr(), K.ColorJiggle(0.1, 0.1, 0.1, 0.1), K.ImageSequential( K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=0.5), K.RandomPerspective(0.2, p=0.5), K.RandomSolarize(0.1, 0.1, p=0.5), ), K.RandomMixUp(p=1.0), grid_size=(2, 2), padding=padding, patchwise_apply=patchwise_apply, same_on_batch=same_on_batch, keepdim=keepdim, random_apply=random_apply, ) # TODO: improve me and remove the exception. except Exception: return input = torch.randn(*shape, device=device, dtype=dtype) out = seq(input) if seq.return_label: out, _ = out assert out.shape[-3:] == input.shape[-3:] reproducibility_test(input, seq)
def test_inverse_and_forward_return_transform(self, random_apply, device, dtype): inp = torch.randn(1, 3, 1000, 500, device=device, dtype=dtype) bbox = torch.tensor([[[355, 10], [660, 10], [660, 250], [355, 250]]], device=device, dtype=dtype) bbox_2 = [ # torch.tensor([[[355, 10], [660, 10], [660, 250], [355, 250]]], device=device, dtype=dtype), torch.tensor( [[[355, 10], [660, 10], [660, 250], [355, 250]], [[355, 10], [660, 10], [660, 250], [355, 250]]], device=device, dtype=dtype, ) ] bbox_wh = torch.tensor([[[30, 40, 100, 100]]], device=device, dtype=dtype) bbox_wh_2 = [ # torch.tensor([[30, 40, 100, 100]], device=device, dtype=dtype), torch.tensor([[30, 40, 100, 100], [30, 40, 100, 100]], device=device, dtype=dtype) ] keypoints = torch.tensor([[[465, 115], [545, 116]]], device=device, dtype=dtype) mask = bbox_to_mask( torch.tensor([[[155, 0], [900, 0], [900, 400], [155, 400]]], device=device, dtype=dtype), 1000, 500)[:, None].float() aug = K.AugmentationSequential( K.ImageSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0)), K.AugmentationSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0)), K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0), data_keys=[ "input", "mask", "bbox", "keypoints", "bbox", "BBOX_XYWH", "BBOX_XYWH" ], random_apply=random_apply, ) with pytest.raises( Exception): # No parameters available for inversing. aug.inverse(inp, mask, bbox, keypoints, bbox_2, bbox_wh, bbox_wh_2) out = aug(inp, mask, bbox, keypoints, bbox_2, bbox_wh, bbox_wh_2) assert out[0].shape == inp.shape assert out[1].shape == mask.shape assert out[2].shape == bbox.shape assert out[3].shape == keypoints.shape reproducibility_test( (inp, mask, bbox, keypoints, bbox_2, bbox_wh, bbox_wh_2), aug)
def test_individual_forward_and_inverse(self, device, dtype): inp = torch.randn(1, 3, 1000, 500, device=device, dtype=dtype) bbox = torch.tensor([[[355, 10], [660, 10], [660, 250], [355, 250]]], device=device, dtype=dtype) keypoints = torch.tensor([[[465, 115], [545, 116]]], device=device, dtype=dtype) mask = bbox_to_mask( torch.tensor([[[155, 0], [900, 0], [900, 400], [155, 400]]], device=device, dtype=dtype), 500, 1000)[:, None].float() crop_size = (200, 200) aug = K.AugmentationSequential( K.ImageSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0)), K.AugmentationSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0)), K.RandomAffine(360, p=1.0), K.RandomCrop(crop_size, padding=1, cropping_mode='resample', fill=0), data_keys=['input', 'mask', 'bbox', 'keypoints'], ) reproducibility_test((inp, mask, bbox, keypoints), aug) out = aug(inp, mask, bbox, keypoints) assert out[0].shape == (*inp.shape[:2], *crop_size) assert out[1].shape == (*mask.shape[:2], *crop_size) assert out[2].shape == bbox.shape assert out[3].shape == keypoints.shape out_inv = aug.inverse(*out) assert out_inv[0].shape == inp.shape assert out_inv[1].shape == mask.shape assert out_inv[2].shape == bbox.shape assert out_inv[3].shape == keypoints.shape aug = K.AugmentationSequential(K.RandomAffine(360, p=1.0)) assert aug(inp, data_keys=['input']).shape == inp.shape aug = K.AugmentationSequential(K.RandomAffine(360, p=1.0)) assert aug(inp, data_keys=['input']).shape == inp.shape assert aug(mask, data_keys=['mask'], params=aug._params).shape == mask.shape assert aug.inverse(inp, data_keys=['input']).shape == inp.shape assert aug.inverse(bbox, data_keys=['bbox']).shape == bbox.shape assert aug.inverse(keypoints, data_keys=['keypoints']).shape == keypoints.shape assert aug.inverse(mask, data_keys=['mask']).shape == mask.shape
def test_intensity_only(self): seq = K.PatchSequential( K.ImageSequential( K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=0.5), K.RandomPerspective(0.2, p=0.5), K.RandomSolarize(0.1, 0.1, p=0.5), ), K.ColorJiggle(0.1, 0.1, 0.1, 0.1), K.ImageSequential( K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=0.5), K.RandomPerspective(0.2, p=0.5), K.RandomSolarize(0.1, 0.1, p=0.5), ), K.ColorJiggle(0.1, 0.1, 0.1, 0.1), grid_size=(2, 2), ) assert not seq.is_intensity_only() seq = K.PatchSequential( K.ImageSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=0.5)), K.ColorJiggle(0.1, 0.1, 0.1, 0.1), K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=0.5), K.ColorJiggle(0.1, 0.1, 0.1, 0.1), grid_size=(2, 2), ) assert seq.is_intensity_only()
def test_exception(self, shape, data_format, device, dtype): aug_list = K.VideoSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1), data_format=data_format, same_on_frame=True) with pytest.raises(AssertionError): img = torch.randn(*shape, device=device, dtype=dtype) aug_list(img)
def test_jit(self, device, dtype): B, C, D, H, W = 2, 3, 5, 4, 4 img = torch.ones(B, C, D, H, W, device=device, dtype=dtype) op = K.VideoSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1), same_on_frame=True) op_jit = torch.jit.script(op) assert_close(op(img), op_jit(img))
def test_exception(self, error_param): with pytest.raises(Exception): # AssertError and NotImplementedError K.PatchSequential( K.ImageSequential( K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=0.5), K.RandomPerspective(0.2, p=0.5), K.RandomSolarize(0.1, 0.1, p=0.5), ), K.ColorJiggle(0.1, 0.1, 0.1, 0.1), K.ImageSequential( K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=0.5), K.RandomPerspective(0.2, p=0.5), K.RandomSolarize(0.1, 0.1, p=0.5), ), K.ColorJiggle(0.1, 0.1, 0.1, 0.1), **error_param, )
def test_jit(self, device, dtype): B, C, H, W = 2, 3, 4, 4 img = torch.ones(B, C, H, W, device=device, dtype=dtype) op = K.AugmentationSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0), same_on_batch=True) op_jit = torch.jit.script(op) assert_close(op(img), op_jit(img))
def test_mixup(self, inp, random_apply, same_on_batch, device, dtype): inp = torch.as_tensor(inp, device=device, dtype=dtype) aug = K.AugmentationSequential( K.ImageSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0)), K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0), K.RandomMixUp(p=1.0), data_keys=["input"], random_apply=random_apply, same_on_batch=same_on_batch, ) out = aug(inp) if aug.return_label: out, _ = out assert out.shape[-3:] == inp.shape[-3:] reproducibility_test(inp, aug)
def test_forward(self, random_apply, device, dtype): inp = torch.randn(1, 3, 30, 30, device=device, dtype=dtype) aug = K.ImageSequential( K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), kornia.filters.MedianBlur((3, 3)), K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.ImageSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0)), K.ImageSequential(K.RandomAffine(360, p=1.0)), K.RandomAffine(360, p=1.0), K.RandomMixUp(p=1.0), random_apply=random_apply, ) out = aug(inp) if aug.return_label: out, _ = out assert out.shape == inp.shape aug.inverse(inp) reproducibility_test(inp, aug)
def test_forward_and_inverse_return_transform(self, random_apply, device, dtype): inp = torch.randn(1, 3, 1000, 500, device=device, dtype=dtype) bbox = torch.tensor([[[355, 10], [660, 10], [660, 250], [355, 250]]], device=device, dtype=dtype) keypoints = torch.tensor([[[465, 115], [545, 116]]], device=device, dtype=dtype) mask = bbox_to_mask( torch.tensor([[[155, 0], [900, 0], [900, 400], [155, 400]]], device=device, dtype=dtype), 1000, 500)[:, None].float() aug = K.AugmentationSequential( K.ImageSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0)), K.AugmentationSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0)), K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0), data_keys=["input", "mask", "bbox", "keypoints"], random_apply=random_apply, ) out = aug(inp, mask, bbox, keypoints) assert out[0].shape == inp.shape assert out[1].shape == mask.shape assert out[2].shape == bbox.shape assert out[3].shape == keypoints.shape reproducibility_test((inp, mask, bbox, keypoints), aug) # TODO(jian): we sometimes throw the following error # AttributeError: 'tuple' object has no attribute 'shape' out_inv = aug.inverse(*out) assert out_inv[0].shape == inp.shape assert out_inv[1].shape == mask.shape assert out_inv[2].shape == bbox.shape assert out_inv[3].shape == keypoints.shape
def test_construction(self, same_on_batch, keepdim, random_apply): aug = K.ImageSequential( K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0), K.RandomMixUp(p=1.0), same_on_batch=same_on_batch, keepdim=keepdim, random_apply=random_apply, ) c = 0 for a in aug.get_forward_sequence(): if isinstance(a, (MixAugmentationBase, )): c += 1 assert c < 2 aug.same_on_batch = True aug.keepdim = True for m in aug.children(): assert m.same_on_batch is True, m.same_on_batch assert m.keepdim is True, m.keepdim
class TestVideoSequential: @pytest.mark.parametrize('shape', [(3, 4), (2, 3, 4), (2, 3, 5, 6), (2, 3, 4, 5, 6, 7)]) @pytest.mark.parametrize('data_format', ["BCTHW", "BTCHW"]) def test_exception(self, shape, data_format, device, dtype): aug_list = K.VideoSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1), data_format=data_format, same_on_frame=True) with pytest.raises(AssertionError): img = torch.randn(*shape, device=device, dtype=dtype) aug_list(img) @pytest.mark.parametrize( 'augmentation', [ K.RandomAffine(360, p=1.0), K.CenterCrop((3, 3), p=1.0), K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomCrop((5, 5), p=1.0), K.RandomErasing(p=1.0), K.RandomGrayscale(p=1.0), K.RandomHorizontalFlip(p=1.0), K.RandomVerticalFlip(p=1.0), K.RandomPerspective(p=1.0), K.RandomResizedCrop((5, 5), p=1.0), K.RandomRotation(360.0, p=1.0), K.RandomSolarize(p=1.0), K.RandomPosterize(p=1.0), K.RandomSharpness(p=1.0), K.RandomEqualize(p=1.0), K.RandomMotionBlur(3, 35.0, 0.5, p=1.0), K.Normalize(torch.tensor([0.5, 0.5, 0.5]), torch.tensor([0.5, 0.5, 0.5]), p=1.0), K.Denormalize(torch.tensor([0.5, 0.5, 0.5]), torch.tensor([0.5, 0.5, 0.5]), p=1.0), ], ) @pytest.mark.parametrize('data_format', ["BCTHW", "BTCHW"]) def test_augmentation(self, augmentation, data_format, device, dtype): input = torch.randint(255, (1, 3, 3, 5, 6), device=device, dtype=dtype).repeat(2, 1, 1, 1, 1) / 255.0 torch.manual_seed(21) aug_list = K.VideoSequential(augmentation, data_format=data_format, same_on_frame=True) reproducibility_test(input, aug_list) @pytest.mark.parametrize( 'augmentations', [ [ K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0) ], [ K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0) ], [K.RandomAffine(360, p=1.0), kornia.color.BgrToRgb()], [ K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=0.0), K.RandomAffine(360, p=0.0) ], [K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=0.0)], [K.RandomAffine(360, p=0.0)], [ K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0), K.RandomMixUp(p=1.0) ], ], ) @pytest.mark.parametrize('data_format', ["BCTHW", "BTCHW"]) @pytest.mark.parametrize('random_apply', [1, (1, 1), (1, ), 10, True, False]) def test_same_on_frame(self, augmentations, data_format, random_apply, device, dtype): aug_list = K.VideoSequential(*augmentations, data_format=data_format, same_on_frame=True, random_apply=random_apply) if data_format == 'BCTHW': input = torch.randn(2, 3, 1, 5, 6, device=device, dtype=dtype).repeat(1, 1, 4, 1, 1) output = aug_list(input) if aug_list.return_label: output, _ = output assert (output[:, :, 0] == output[:, :, 1]).all() assert (output[:, :, 1] == output[:, :, 2]).all() assert (output[:, :, 2] == output[:, :, 3]).all() if data_format == 'BTCHW': input = torch.randn(2, 1, 3, 5, 6, device=device, dtype=dtype).repeat(1, 4, 1, 1, 1) output = aug_list(input) if aug_list.return_label: output, _ = output assert (output[:, 0] == output[:, 1]).all() assert (output[:, 1] == output[:, 2]).all() assert (output[:, 2] == output[:, 3]).all() reproducibility_test(input, aug_list) @pytest.mark.parametrize( 'augmentations', [ [K.RandomAffine(360, p=1.0)], [K.RandomCrop((2, 2), padding=2)], [K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0)], [ K.RandomAffine(360, p=0.0), K.ImageSequential(K.RandomAffine(360, p=0.0)) ], ], ) @pytest.mark.parametrize('data_format', ["BCTHW", "BTCHW"]) def test_against_sequential(self, augmentations, data_format, device, dtype): aug_list_1 = K.VideoSequential(*augmentations, data_format=data_format, same_on_frame=False) aug_list_2 = torch.nn.Sequential(*augmentations) if data_format == 'BCTHW': input = torch.randn(2, 3, 1, 5, 6, device=device, dtype=dtype).repeat(1, 1, 4, 1, 1) if data_format == 'BTCHW': input = torch.randn(2, 1, 3, 5, 6, device=device, dtype=dtype).repeat(1, 4, 1, 1, 1) torch.manual_seed(0) output_1 = aug_list_1(input) torch.manual_seed(0) if data_format == 'BCTHW': input = input.transpose(1, 2) output_2 = aug_list_2(input.reshape(-1, 3, 5, 6)) if any(isinstance(a, K.RandomCrop) for a in augmentations): output_2 = output_2.view(2, 4, 3, 2, 2) else: output_2 = output_2.view(2, 4, 3, 5, 6) if data_format == 'BCTHW': output_2 = output_2.transpose(1, 2) assert (output_1 == output_2).all(), dict(aug_list_1._params) @pytest.mark.jit @pytest.mark.skip(reason="turn off due to Union Type") def test_jit(self, device, dtype): B, C, D, H, W = 2, 3, 5, 4, 4 img = torch.ones(B, C, D, H, W, device=device, dtype=dtype) op = K.VideoSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1), same_on_frame=True) op_jit = torch.jit.script(op) assert_close(op(img), op_jit(img))
class TestAugmentationSequential: @pytest.mark.parametrize('data_keys', [ "input", ["mask", "input"], ["input", "bbox_yxyx"], [0, 10], [BorderType.REFLECT] ]) @pytest.mark.parametrize("augmentation_list", [K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0)]) def test_exception(self, augmentation_list, data_keys, device, dtype): with pytest.raises(Exception): # AssertError and NotImplementedError K.AugmentationSequential(augmentation_list, data_keys=data_keys) @pytest.mark.parametrize('same_on_batch', [True, False]) @pytest.mark.parametrize('random_apply', [1, (2, 2), (1, 2), (2, ), 10, True, False]) @pytest.mark.parametrize( 'inp', [torch.randn(1, 3, 1000, 500), torch.randn(3, 1000, 500)]) def test_mixup(self, inp, random_apply, same_on_batch, device, dtype): inp = torch.as_tensor(inp, device=device, dtype=dtype) aug = K.AugmentationSequential( K.ImageSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0)), K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0), K.RandomMixUp(p=1.0), data_keys=["input"], random_apply=random_apply, same_on_batch=same_on_batch, ) out = aug(inp) if aug.return_label: out, _ = out assert out.shape[-3:] == inp.shape[-3:] reproducibility_test(inp, aug) def test_video(self, device, dtype): input = torch.randn(2, 3, 5, 6, device=device, dtype=dtype)[None] bbox = torch.tensor([[[1.0, 1.0], [2.0, 1.0], [2.0, 2.0], [1.0, 2.0]]], device=device, dtype=dtype).expand(2, -1, -1)[None] points = torch.tensor([[[1.0, 1.0]]], device=device, dtype=dtype).expand(2, -1, -1)[None] aug_list = K.AugmentationSequential( K.VideoSequential( kornia.augmentation.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), kornia.augmentation.RandomAffine(360, p=1.0)), data_keys=["input", "mask", "bbox", "keypoints"], ) out = aug_list(input, input, bbox, points) assert out[0].shape == input.shape assert out[1].shape == input.shape assert out[2].shape == bbox.shape assert out[3].shape == points.shape out_inv = aug_list.inverse(*out) assert out_inv[0].shape == input.shape assert out_inv[1].shape == input.shape assert out_inv[2].shape == bbox.shape assert out_inv[3].shape == points.shape def test_3d_augmentations(self, device, dtype): input = torch.randn(2, 2, 3, 5, 6, device=device, dtype=dtype) aug_list = K.AugmentationSequential( K.RandomAffine3D(360., p=1.), K.RandomHorizontalFlip3D(p=1.), data_keys=["input"], ) out = aug_list(input) assert out.shape == input.shape def test_random_flips(self, device, dtype): inp = torch.randn(1, 3, 510, 1020, device=device, dtype=dtype) bbox = torch.tensor([[[355, 10], [660, 10], [660, 250], [355, 250]]], device=device, dtype=dtype) expected_bbox_vertical_flip = torch.tensor( [[[355, 259], [660, 259], [660, 499], [355, 499]]], device=device, dtype=dtype) expected_bbox_horizontal_flip = torch.tensor( [[[359, 10], [664, 10], [664, 250], [359, 250]]], device=device, dtype=dtype) aug_ver = K.AugmentationSequential(K.RandomVerticalFlip(p=1.0), data_keys=["input", "bbox"], same_on_batch=False) aug_hor = K.AugmentationSequential(K.RandomHorizontalFlip(p=1.0), data_keys=["input", "bbox"], same_on_batch=False) out_ver = aug_ver(inp.clone(), bbox.clone()) out_hor = aug_hor(inp.clone(), bbox.clone()) assert_close(out_ver[1], expected_bbox_vertical_flip) assert_close(out_hor[1], expected_bbox_horizontal_flip) def test_random_crops_and_flips(self, device, dtype): width, height = 100, 100 crop_width, crop_height = 3, 3 input = torch.randn(3, 3, width, height, device=device, dtype=dtype) bbox = torch.tensor([[[1.0, 1.0, 2.0, 2.0], [0.0, 0.0, 1.0, 2.0], [0.0, 0.0, 2.0, 1.0]]], device=device, dtype=dtype).expand(3, -1, -1) aug = K.AugmentationSequential( K.RandomCrop((crop_width, crop_height), padding=1, cropping_mode='resample', fill=0), K.RandomHorizontalFlip(p=1.0), data_keys=["input", "bbox_xyxy"], ) reproducibility_test((input, bbox), aug) _params = aug.forward_parameters(input.shape) # specifying the crop locations allows us to compute by hand the expected outputs crop_locations = torch.tensor( [[1.0, 2.0], [1.0, 1.0], [2.0, 0.0]], device=_params[0].data['src'].device, dtype=_params[0].data['src'].dtype, ) crops = crop_locations.expand(4, -1, -1).permute(1, 0, 2).clone() crops[:, 1:3, 0] += crop_width - 1 crops[:, 2:4, 1] += crop_height - 1 _params[0].data['src'] = crops # expected output bboxes after crop for specified crop locations and crop size (3,3) expected_out_bbox = torch.tensor( [ [[1.0, 0.0, 2.0, 1.0], [0.0, -1.0, 1.0, 1.0], [0.0, -1.0, 2.0, 0.0]], [[1.0, 1.0, 2.0, 2.0], [0.0, 0.0, 1.0, 2.0], [0.0, 0.0, 2.0, 1.0]], [[0.0, 2.0, 1.0, 3.0], [-1.0, 1.0, 0.0, 3.0], [-1.0, 1.0, 1.0, 2.0]], ], device=device, dtype=dtype, ) # horizontally flip boxes based on crop width xmins = expected_out_bbox[..., 0].clone() xmaxs = expected_out_bbox[..., 2].clone() expected_out_bbox[..., 0] = crop_width - xmaxs expected_out_bbox[..., 2] = crop_width - xmins out = aug(input, bbox, params=_params) assert out[1].shape == bbox.shape assert_close(out[1], expected_out_bbox, atol=1e-4, rtol=1e-4) out_inv = aug.inverse(*out) assert out_inv[1].shape == bbox.shape assert_close(out_inv[1], bbox, atol=1e-4, rtol=1e-4) def test_random_erasing(self, device, dtype): fill_value = 0.5 input = torch.randn(3, 3, 100, 100, device=device, dtype=dtype) aug = K.AugmentationSequential( K.RandomErasing(p=1., value=fill_value), data_keys=["input", "mask"], ) reproducibility_test((input, input), aug) out = aug(input, input) assert torch.all(out[1][out[0] == fill_value] == 0.) def test_random_crops(self, device, dtype): torch.manual_seed(233) input = torch.randn(3, 3, 3, 3, device=device, dtype=dtype) bbox = torch.tensor([[[1.0, 1.0, 2.0, 2.0], [0.0, 0.0, 1.0, 2.0], [0.0, 0.0, 2.0, 1.0]]], device=device, dtype=dtype).expand(3, -1, -1) points = torch.tensor([[[0.0, 0.0], [1.0, 1.0]]], device=device, dtype=dtype).expand(3, -1, -1) aug = K.AugmentationSequential( K.RandomCrop((3, 3), padding=1, cropping_mode='resample', fill=0), K.RandomAffine((360., 360.), p=1.), data_keys=["input", "mask", "bbox_xyxy", "keypoints"], ) reproducibility_test((input, input, bbox, points), aug) _params = aug.forward_parameters(input.shape) # specifying the crops allows us to compute by hand the expected outputs _params[0].data['src'] = torch.tensor( [ [[1.0, 2.0], [3.0, 2.0], [3.0, 4.0], [1.0, 4.0]], [[1.0, 1.0], [3.0, 1.0], [3.0, 3.0], [1.0, 3.0]], [[2.0, 0.0], [4.0, 0.0], [4.0, 2.0], [2.0, 2.0]], ], device=_params[0].data['src'].device, dtype=_params[0].data['src'].dtype, ) expected_out_bbox = torch.tensor( [ [[1.0, 0.0, 2.0, 1.0], [0.0, -1.0, 1.0, 1.0], [0.0, -1.0, 2.0, 0.0]], [[1.0, 1.0, 2.0, 2.0], [0.0, 0.0, 1.0, 2.0], [0.0, 0.0, 2.0, 1.0]], [[0.0, 2.0, 1.0, 3.0], [-1.0, 1.0, 0.0, 3.0], [-1.0, 1.0, 1.0, 2.0]], ], device=device, dtype=dtype, ) expected_out_points = torch.tensor( [[[0.0, -1.0], [1.0, 0.0]], [[0.0, 0.0], [1.0, 1.0]], [[-1.0, 1.0], [0.0, 2.0]]], device=device, dtype=dtype) out = aug(input, input, bbox, points, params=_params) assert out[0].shape == (3, 3, 3, 3) assert_close(out[0], out[1], atol=1e-4, rtol=1e-4) assert out[2].shape == bbox.shape assert_close(out[2], expected_out_bbox, atol=1e-4, rtol=1e-4) assert out[3].shape == points.shape assert_close(out[3], expected_out_points, atol=1e-4, rtol=1e-4) out_inv = aug.inverse(*out) assert out_inv[0].shape == input.shape assert_close(out_inv[0], out_inv[1], atol=1e-4, rtol=1e-4) assert out_inv[2].shape == bbox.shape assert_close(out_inv[2], bbox, atol=1e-4, rtol=1e-4) assert out_inv[3].shape == points.shape assert_close(out_inv[3], points, atol=1e-4, rtol=1e-4) @pytest.mark.parametrize('random_apply', [1, (2, 2), (1, 2), (2, ), 10, True, False]) def test_forward_and_inverse(self, random_apply, device, dtype): inp = torch.randn(1, 3, 1000, 500, device=device, dtype=dtype) bbox = torch.tensor([[[355, 10], [660, 10], [660, 250], [355, 250]]], device=device, dtype=dtype) keypoints = torch.tensor([[[465, 115], [545, 116]]], device=device, dtype=dtype) mask = bbox_to_mask( torch.tensor([[[155, 0], [900, 0], [900, 400], [155, 400]]], device=device, dtype=dtype), 1000, 500)[:, None].float() aug = K.AugmentationSequential( K.ImageSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0)), K.AugmentationSequential( K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0), K.RandomAffine(360, p=1.0), data_keys=["input", "mask", "bbox", "keypoints"]), K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0), data_keys=["input", "mask", "bbox", "keypoints"], random_apply=random_apply, ) out = aug(inp, mask, bbox, keypoints) assert out[0].shape == inp.shape assert out[1].shape == mask.shape assert out[2].shape == bbox.shape assert out[3].shape == keypoints.shape reproducibility_test((inp, mask, bbox, keypoints), aug) out_inv = aug.inverse(*out) assert out_inv[0].shape == inp.shape assert out_inv[1].shape == mask.shape assert out_inv[2].shape == bbox.shape assert out_inv[3].shape == keypoints.shape def test_individual_forward_and_inverse(self, device, dtype): inp = torch.randn(1, 3, 1000, 500, device=device, dtype=dtype) bbox = torch.tensor([[[355, 10], [660, 10], [660, 250], [355, 250]]], device=device, dtype=dtype) keypoints = torch.tensor([[[465, 115], [545, 116]]], device=device, dtype=dtype) mask = bbox_to_mask( torch.tensor([[[155, 0], [900, 0], [900, 400], [155, 400]]], device=device, dtype=dtype), 500, 1000)[:, None].float() crop_size = (200, 200) aug = K.AugmentationSequential( K.ImageSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0)), K.AugmentationSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0)), K.RandomAffine(360, p=1.0), K.RandomCrop(crop_size, padding=1, cropping_mode='resample', fill=0), data_keys=['input', 'mask', 'bbox', 'keypoints'], ) reproducibility_test((inp, mask, bbox, keypoints), aug) out = aug(inp, mask, bbox, keypoints) assert out[0].shape == (*inp.shape[:2], *crop_size) assert out[1].shape == (*mask.shape[:2], *crop_size) assert out[2].shape == bbox.shape assert out[3].shape == keypoints.shape out_inv = aug.inverse(*out) assert out_inv[0].shape == inp.shape assert out_inv[1].shape == mask.shape assert out_inv[2].shape == bbox.shape assert out_inv[3].shape == keypoints.shape aug = K.AugmentationSequential(K.RandomAffine(360, p=1.0)) assert aug(inp, data_keys=['input']).shape == inp.shape aug = K.AugmentationSequential(K.RandomAffine(360, p=1.0)) assert aug(inp, data_keys=['input']).shape == inp.shape assert aug(mask, data_keys=['mask'], params=aug._params).shape == mask.shape assert aug.inverse(inp, data_keys=['input']).shape == inp.shape assert aug.inverse(bbox, data_keys=['bbox']).shape == bbox.shape assert aug.inverse(keypoints, data_keys=['keypoints']).shape == keypoints.shape assert aug.inverse(mask, data_keys=['mask']).shape == mask.shape @pytest.mark.parametrize('random_apply', [2, (1, 1), (2, ), 10, True, False]) def test_forward_and_inverse_return_transform(self, random_apply, device, dtype): inp = torch.randn(1, 3, 1000, 500, device=device, dtype=dtype) bbox = torch.tensor([[[355, 10], [660, 10], [660, 250], [355, 250]]], device=device, dtype=dtype) keypoints = torch.tensor([[[465, 115], [545, 116]]], device=device, dtype=dtype) mask = bbox_to_mask( torch.tensor([[[155, 0], [900, 0], [900, 400], [155, 400]]], device=device, dtype=dtype), 1000, 500)[:, None].float() aug = K.AugmentationSequential( K.ImageSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0)), K.AugmentationSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0)), K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0), data_keys=["input", "mask", "bbox", "keypoints"], random_apply=random_apply, ) out = aug(inp, mask, bbox, keypoints) assert out[0].shape == inp.shape assert out[1].shape == mask.shape assert out[2].shape == bbox.shape assert out[3].shape == keypoints.shape reproducibility_test((inp, mask, bbox, keypoints), aug) # TODO(jian): we sometimes throw the following error # AttributeError: 'tuple' object has no attribute 'shape' out_inv = aug.inverse(*out) assert out_inv[0].shape == inp.shape assert out_inv[1].shape == mask.shape assert out_inv[2].shape == bbox.shape assert out_inv[3].shape == keypoints.shape @pytest.mark.parametrize('random_apply', [1, (2, 2), (1, 2), (2, ), 10, True, False]) def test_inverse_and_forward_return_transform(self, random_apply, device, dtype): inp = torch.randn(1, 3, 1000, 500, device=device, dtype=dtype) bbox = torch.tensor([[[355, 10], [660, 10], [660, 250], [355, 250]]], device=device, dtype=dtype) bbox_2 = [ # torch.tensor([[[355, 10], [660, 10], [660, 250], [355, 250]]], device=device, dtype=dtype), torch.tensor( [[[355, 10], [660, 10], [660, 250], [355, 250]], [[355, 10], [660, 10], [660, 250], [355, 250]]], device=device, dtype=dtype, ) ] bbox_wh = torch.tensor([[[30, 40, 100, 100]]], device=device, dtype=dtype) bbox_wh_2 = [ # torch.tensor([[30, 40, 100, 100]], device=device, dtype=dtype), torch.tensor([[30, 40, 100, 100], [30, 40, 100, 100]], device=device, dtype=dtype) ] keypoints = torch.tensor([[[465, 115], [545, 116]]], device=device, dtype=dtype) mask = bbox_to_mask( torch.tensor([[[155, 0], [900, 0], [900, 400], [155, 400]]], device=device, dtype=dtype), 1000, 500)[:, None].float() aug = K.AugmentationSequential( K.ImageSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0)), K.AugmentationSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0)), K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0), data_keys=[ "input", "mask", "bbox", "keypoints", "bbox", "BBOX_XYWH", "BBOX_XYWH" ], random_apply=random_apply, ) with pytest.raises( Exception): # No parameters available for inversing. aug.inverse(inp, mask, bbox, keypoints, bbox_2, bbox_wh, bbox_wh_2) out = aug(inp, mask, bbox, keypoints, bbox_2, bbox_wh, bbox_wh_2) assert out[0].shape == inp.shape assert out[1].shape == mask.shape assert out[2].shape == bbox.shape assert out[3].shape == keypoints.shape reproducibility_test( (inp, mask, bbox, keypoints, bbox_2, bbox_wh, bbox_wh_2), aug) @pytest.mark.jit @pytest.mark.skip(reason="turn off due to Union Type") def test_jit(self, device, dtype): B, C, H, W = 2, 3, 4, 4 img = torch.ones(B, C, H, W, device=device, dtype=dtype) op = K.AugmentationSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0), same_on_batch=True) op_jit = torch.jit.script(op) assert_close(op(img), op_jit(img))
def test_exception(self, random_apply_weights, device, dtype): inp = torch.randn(1, 3, 30, 30, device=device, dtype=dtype) with pytest.raises(Exception): # AssertError and NotImplementedError K.ImageSequential( K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), random_apply_weights=random_apply_weights).inverse(inp)