def test_prob_k_spatial_axes(self):
rotate = Rotate90(k=2, spatial_axes=(0, 1))
rotated = rotate(self.imt[0])
expected = list()
for channel in self.imt[0]:
expected.append(np.rot90(channel, 2, (0, 1)))
expected = np.stack(expected)
self.assertTrue(np.allclose(rotated, expected))
def test_rotate90_default(self):
rotate = Rotate90()
rotated = rotate(self.imt[0])
expected = list()
for channel in self.imt[0]:
expected.append(np.rot90(channel, 1, (0, 1)))
expected = np.stack(expected)
self.assertTrue(np.allclose(rotated, expected))
def __call__(self, data):
self.randomize()
if not self._do_transform:
return data
rotator = Rotate90(self._rand_k, self.spatial_axes)
d = dict(data)
for key in self.keys:
d[key] = rotator(d[key])
return d
def __init__(self, keys, k=1, spatial_axes=(0, 1)):
"""
Args:
k (int): number of times to rotate by 90 degrees.
spatial_axes (2 ints): defines the plane to rotate with 2 spatial axes.
Default: (0, 1), this is the first two axis in spatial dimensions.
"""
MapTransform.__init__(self, keys)
self.k = k
self.spatial_axes = spatial_axes
self.rotator = Rotate90(self.k, self.spatial_axes)