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