def test_inverse_and_forward_return_transform(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.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"], ) out_inv = aug.inverse(inp, mask, bbox, keypoints) 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 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
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']).shape == mask.shape aug = K.AugmentationSequential( K.RandomAffine(360, p=1.0, return_transform=True)) assert aug.inverse(inp, data_keys=['input']).shape == inp.shape aug = K.AugmentationSequential( K.RandomAffine(360, p=1.0, return_transform=True)) assert aug.inverse(bbox, data_keys=['bbox']).shape == bbox.shape aug = K.AugmentationSequential( K.RandomAffine(360, p=1.0, return_transform=True)) assert aug.inverse(keypoints, data_keys=['keypoints']).shape == keypoints.shape aug = K.AugmentationSequential( K.RandomAffine(360, p=1.0, return_transform=True)) assert aug.inverse(mask, data_keys=['mask']).shape == mask.shape