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, 499], [660, 499], [660, 259], [355, 259]]], device=device, dtype=dtype) expected_bbox_horizontal_flip = torch.tensor( [[[664, 10], [359, 10], [359, 250], [664, 250]]], device=device, dtype=dtype) aug_ver = K.AugmentationSequential(K.RandomVerticalFlip(p=1.0), data_keys=["input", "bbox"], return_transform=False, same_on_batch=False) aug_hor = K.AugmentationSequential(K.RandomHorizontalFlip(p=1.0), data_keys=["input", "bbox"], return_transform=False, 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_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), 1000, 500)[:, None].float() aug = K.AugmentationSequential( K.RandomAffine(360, p=1.0, return_transform=False), data_keys=['input', 'mask', 'bbox', 'keypoints']) reproducibility_test((inp, mask, bbox, keypoints), aug) aug = K.AugmentationSequential( K.RandomAffine(360, p=1.0, return_transform=True)) assert aug(inp, data_keys=['input'])[0].shape == inp.shape aug = K.AugmentationSequential( K.RandomAffine(360, p=1.0, return_transform=False)) 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_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_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.ColorJitter(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0, return_transform=True)), K.ColorJitter(0.1, 0.1, 0.1, 0.1, p=1.0, return_transform=True), K.RandomAffine(360, p=1.0, return_transform=True), data_keys=["input", "mask", "bbox", "keypoints"], random_apply=random_apply, ) out = aug(inp, mask, bbox, keypoints) assert out[0][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_video(self, device, dtype): input = torch.randn(2, 3, 5, 6, device=device, dtype=dtype)[None] bbox = torch.tensor([[ [1., 1.], [2., 1.], [2., 2.], [1., 2.], ]], device=device, dtype=dtype).expand(2, -1, -1)[None] points = torch.tensor([[[1., 1.]]], device=device, dtype=dtype).expand(2, -1, -1)[None] aug_list = K.AugmentationSequential( K.VideoSequential( kornia.augmentation.ColorJitter(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_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.ColorJitter(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_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_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_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_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_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_mixup(self, inp, return_transform, random_apply, same_on_batch, device, dtype): inp = torch.as_tensor(inp, device=device, dtype=dtype) aug = K.AugmentationSequential( K.ColorJitter(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, return_transform=return_transform, same_on_batch=same_on_batch, ) out = aug(inp) if aug.return_label: out, label = out if return_transform and isinstance(out, (tuple, list)): out = out[0] assert out.shape[-3:] == inp.shape[-3:] reproducibility_test(inp, aug)
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)
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)