def test_rotate():
# fmt: off
expected = torch.Tensor([[[[3., 7., 11., 15.], [2., 6., 10., 14.],
[1., 5., 9., 13.], [0., 4., 8., 12.]]]])
# fmt: on
forward = F.Rotate().forward(INPUT, 90)
assert expected.allclose(forward)
assert INPUT.allclose(F.Rotate().backward(forward, 90))
class Augmentation(object):
transforms = {
"h_flip": F.HFlip(),
"v_flip": F.VFlip(),
"rotation": F.Rotate(),
"h_shift": F.HShift(),
"v_shift": F.VShift(),
# 'contrast': F.Contrast(),
"add": F.Add(),
"mul": F.Multiply(),
}
def __init__(self, **params):
super().__init__()
transforms = [Augmentation.transforms[k] for k in params.keys()]
transform_params = [params[k] for k in params.keys()]
# add identity parameters for all transforms and convert to list
transform_params = [
t.prepare(params)
for t, params in zip(transforms, transform_params)
]
# get all combinations of transforms params
transform_params = list(itertools.product(*transform_params))
self.forward_aug = [t.forward for t in transforms]
self.forward_params = transform_params
# reverse transforms
self.backward_aug = [t.backward for t in transforms[::-1]]
# reverse params
self.backward_params = [p[::-1] for p in transform_params]
self.n_transforms = len(transform_params)
print(f"Will merge {self.n_transforms} augmentations for each image.")
def forward(self, x):
self.bs = x.shape[0]
transformed_batches = []
for i, args in enumerate(self.forward_params):
batch = x
for f, arg in zip(self.forward_aug, args):
batch = f(batch, arg)
transformed_batches.append(batch)
# returns shape B*Aug x C x H x W
return torch.cat(transformed_batches, 0)
def backward(self, x):
# reshape to separate batches
x = x.reshape([-1, self.bs, *x.shape[1:]])
transformed_batches = []
for i, args in enumerate(self.backward_params):
batch = x[i]
for f, arg in zip(self.backward_aug, args):
batch = f(batch, arg)
transformed_batches.append(batch)
return torch.cat(transformed_batches, 0)