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)