def test_create_control_grid(self):
with self.assertRaisesRegex(TypeError, ''):
create_control_grid(None, None)
with self.assertRaisesRegex(TypeError, ''):
create_control_grid((1, 1), 2.)
g = create_control_grid((1., 1.), (1., 1.))
expected = np.array([
[[-1., -1., -1.], [0., 0., 0.], [1., 1., 1.]],
[[-1., 0., 1.], [-1., 0., 1.], [-1., 0., 1.]],
[[1., 1., 1.], [1., 1., 1.], [1., 1., 1.]],
])
np.testing.assert_allclose(g, expected)
g = create_control_grid((1., 1.), (2., 2.))
expected = np.array([
[[-2., -2., -2.], [0., 0., 0.], [2., 2., 2.]],
[[-2., 0., 2.], [-2., 0., 2.], [-2., 0., 2.]],
[[1., 1., 1.], [1., 1., 1.], [1., 1., 1.]],
])
np.testing.assert_allclose(g, expected)
g = create_control_grid((2., 2.), (1., 1.))
expected = np.array([
[[-1.5, -1.5, -1.5, -1.5], [-0.5, -0.5, -0.5, -0.5],
[0.5, 0.5, 0.5, 0.5], [1.5, 1.5, 1.5, 1.5]],
[[-1.5, -0.5, 0.5, 1.5], [-1.5, -0.5, 0.5, 1.5],
[-1.5, -0.5, 0.5, 1.5], [-1.5, -0.5, 0.5, 1.5]],
[[1., 1., 1., 1.], [1., 1., 1., 1.], [1., 1., 1., 1.],
[1., 1., 1., 1.]],
])
np.testing.assert_allclose(g, expected)
g = create_control_grid((2., 2.), (2., 2.))
expected = np.array([
[[-3., -3., -3., -3.], [-1., -1., -1., -1.], [1., 1., 1., 1.],
[3., 3., 3., 3.]],
[[-3., -1., 1., 3.], [-3., -1., 1., 3.], [-3., -1., 1., 3.],
[-3., -1., 1., 3.]],
[[1., 1., 1., 1.], [1., 1., 1., 1.], [1., 1., 1., 1.],
[1., 1., 1., 1.]],
])
np.testing.assert_allclose(g, expected)
g = create_control_grid((1., 1., 1.), (2., 2., 2.), homogeneous=False)
expected = np.array([[[[-2., -2., -2.], [-2., -2., -2.],
[-2., -2., -2.]],
[[0., 0., 0.], [0., 0., 0.], [0., 0., 0.]],
[[2., 2., 2.], [2., 2., 2.], [2., 2., 2.]]],
[[[-2., -2., -2.], [0., 0., 0.], [2., 2., 2.]],
[[-2., -2., -2.], [0., 0., 0.], [2., 2., 2.]],
[[-2., -2., -2.], [0., 0., 0.], [2., 2., 2.]]],
[[[-2., 0., 2.], [-2., 0., 2.], [-2., 0., 2.]],
[[-2., 0., 2.], [-2., 0., 2.], [-2., 0., 2.]],
[[-2., 0., 2.], [-2., 0., 2.], [-2., 0., 2.]]]])
np.testing.assert_allclose(g, expected)
def __call__(self, spatial_size):
control_grid = create_control_grid(spatial_size, self.spacing)
self.randomize(control_grid.shape[1:])
control_grid[: len(spatial_size)] += self.rand_mag * self.random_offset
if self.as_tensor_output:
control_grid = torch.as_tensor(np.ascontiguousarray(control_grid), device=self.device)
return control_grid
def __call__(
self,
spatial_size: Sequence[int]) -> Union[np.ndarray, torch.Tensor]:
"""
Args:
spatial_size: spatial size of the grid.
"""
self.spacing = fall_back_tuple(self.spacing,
(1.0, ) * len(spatial_size))
control_grid = create_control_grid(spatial_size, self.spacing)
self.randomize()
dist1 = np.sqrt(
np.sum(np.square(control_grid - self.center1)[:2], axis=0) + 1e-6)
dist2 = np.sqrt(
np.sum(np.square(control_grid - self.center2)[:2], axis=0) + 1e-6)
deform_mag = self.rand_mag / (dist1 + dist2)
center = (self.center1 + self.center2) / 2.0
deform = (control_grid - center)[:2] * deform_mag
control_grid[:len(spatial_size)] -= deform
if self.as_tensor_output:
control_grid = torch.as_tensor(np.ascontiguousarray(control_grid),
device=self.device)
return control_grid