def assemble_matrix(self, **data) -> torch.Tensor:
"""
Assembles the matrix (and takes care of batching and having it on the
right device and in the correct dtype and dimensionality).
Args:
**data: the data to be transformed. Will be used to determine
batchsize, dimensionality, dtype and device
Returns:
torch.Tensor: the (batched) transformation matrix
"""
if self.matrix is None:
raise ValueError("Matrix needs to be initialized or overwritten.")
if not torch.is_tensor(self.matrix):
self.matrix = torch.tensor(self.matrix)
self.matrix = self.matrix.to(data[self.keys[0]])
batchsize = data[self.keys[0]].shape[0]
ndim = len(data[self.keys[0]].shape) - 2 # channel and batch dim
# batch dimension missing -> Replicate for each sample in batch
if len(self.matrix.shape) == 2:
self.matrix = self.matrix[None].expand(batchsize, -1, -1).clone()
if self.matrix.shape == (batchsize, ndim, ndim + 1):
return self.matrix
elif self.matrix.shape == (batchsize, ndim, ndim):
return matrix_to_homogeneous(self.matrix)[:, :-1]
elif self.matrix.shape == (batchsize, ndim + 1, ndim + 1):
return matrix_to_cartesian(self.matrix)
raise ValueError(
"Invalid Shape for affine transformation matrix. "
"Got %s but expected %s" % (str(tuple(self.matrix.shape)), str((batchsize, ndim, ndim + 1)))
)
def test_matrix_to_cartesian(self):
expectations = [
torch.tensor([[[1, 2], [3, 4]]]), # single sample, no trans, 2d
torch.tensor([[[1, 2, 5], [3, 4, 6]]]), # single sample, trans, 2d
torch.tensor([[[1, 2, 3], [4, 5, 6], [7, 8, 9]]]), # single sample, no trans, 3d
torch.tensor([[[1, 2, 3, 10], [4, 5, 6, 11], [7, 8, 9, 12]]]), # single sample, trans, 3d
torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]), # multiple samples, no trans, 2d
torch.tensor([[[1, 2, 10], [3, 4, 11]], [[5, 6, 12], [7, 8, 13]]]), # multiple samples, trans, 2d
# multiple samples, trans, 3d
torch.tensor([[[1, 2, 3], [4, 5, 6], [7, 8, 9]], [[10, 11, 12], [13, 14, 15], [16, 17, 18]]]),
# multiple samples, trans, 3d
torch.tensor([[[1, 2, 3, 21], [4, 5, 6, 22], [7, 8, 9, 23]],
[[10, 11, 12, 24], [13, 14, 15, 25], [16, 17, 18, 26]]])
]
inputs = [
torch.tensor([[[1, 2, 0], [3, 4, 0], [0, 0, 1]]]),
torch.tensor([[[1, 2, 5], [3, 4, 6], [0, 0, 1]]]),
torch.tensor([[[1, 2, 3, 0], [4, 5, 6, 0], [7, 8, 9, 0], [0, 0, 0, 1]]]),
torch.tensor([[[1, 2, 3, 10], [4, 5, 6, 11], [7, 8, 9, 12], [0, 0, 0, 1]]]),
torch.tensor([[[1, 2, 0], [3, 4, 0], [0, 0, 1]], [[5, 6, 0], [7, 8, 0], [0, 0, 1]]]),
torch.tensor([[[1, 2, 10], [3, 4, 11], [0, 0, 1]], [[5, 6, 12], [7, 8, 13], [0, 0, 1]]]),
torch.tensor([[[1, 2, 3, 0], [4, 5, 6, 0], [7, 8, 9, 0], [0, 0, 0, 1]],
[[10, 11, 12, 0], [13, 14, 15, 0], [16, 17, 18, 0], [0, 0, 0, 1]]]),
torch.tensor([[[1, 2, 3, 21], [4, 5, 6, 22], [7, 8, 9, 23], [0, 0, 0, 1]],
[[10, 11, 12, 24], [13, 14, 15, 25], [16, 17, 18, 26], [0, 0, 0, 1]]])
]
keep_square = True
for inp, exp in zip(inputs, expectations):
with self.subTest(input=inp, expected=exp):
self.assertTrue(torch.allclose(matrix_to_cartesian(inp, keep_square=keep_square), exp))
keep_square = not keep_square
def test_stacked_transformation_assembly(self):
first_matrix = torch.tensor([[[2.0, 0.0, 1.0], [0.0, 3.0, 2.0]]])
second_matrix = torch.tensor([[[4.0, 0.0, 3.0], [0.0, 5.0, 4.0]]])
trafo = StackedAffine([first_matrix, second_matrix])
sample = {"data": torch.rand(1, 3, 25, 25)}
matrix = trafo.assemble_matrix(**sample)
target_matrix = matrix_to_cartesian(
torch.bmm(matrix_to_homogeneous(first_matrix),
matrix_to_homogeneous(second_matrix)))
self.assertTrue(torch.allclose(matrix, target_matrix))
def assemble_matrix(self, **data) -> torch.Tensor:
"""
Handles the matrix assembly and stacking
Args:
**data: the data to be transformed.
Will be used to determine batchsize,
dimensionality, dtype and device
Returns:
torch.Tensor: the (batched) transformation matrix
"""
whole_trafo = None
for trafo in self.transforms:
matrix = matrix_to_homogeneous(trafo.assemble_matrix(**data))
if whole_trafo is None:
whole_trafo = matrix
else:
whole_trafo = torch.bmm(whole_trafo, matrix)
return matrix_to_cartesian(whole_trafo)
def test_matrix_parametrization(self):
inputs = [
{
"scale": None,
"translation": None,
"rotation": None,
"batchsize": 2,
"ndim": 2,
"dtype": torch.float,
"image_transform": False,
},
{
"scale": [4, 5],
"translation": [2, 3],
"rotation": 90,
"batchsize": 1,
"ndim": 2,
"dtype": torch.float,
"degree": True,
"image_transform": False,
},
{
"scale": [4, 5],
"translation": [2, 3],
"rotation": 90,
"batchsize": 1,
"ndim": 2,
"dtype": torch.float,
"degree": True,
"image_transform": True,
},
{
"scale": [2, 5],
"translation": [9, 18, 27],
"rotation": [180, 0, 180],
"degree": True,
"batchsize": 3,
"ndim": 2,
"dtype": torch.float,
"image_transform": False,
},
]
expectations = [
torch.tensor([
[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
]),
torch.tensor([[[0.0, -4.0, -12.0], [5.0, 0.0, 10.0],
[0.0, 0.0, 1.0]]]),
torch.tensor([[[0.0, 1 / 5.0, -2.0], [-1 / 4.0, 0.0, -3.0],
[0.0, 0.0, 1.0]]]),
torch.bmm(
torch.bmm(
torch.tensor([
[[2.0, 0.0, 0], [0.0, 5.0, 0.0], [0.0, 0.0, 1.0]],
[[2.0, 0.0, 0.0], [0.0, 5.0, 0.0], [0.0, 0.0, 1.0]],
[[2.0, 0.0, 0.0], [0.0, 5.0, 0.0], [0.0, 0.0, 1.0]],
]),
torch.tensor([
[[-1.0, 0.0, 0.0], [0.0, -1.0, 0.0], [0.0, 0.0, 1.0]],
[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
[[-1.0, 0.0, 0.0], [0.0, -1.0, 0.0], [0.0, 0.0, -1.0]],
]),
),
torch.tensor([
[[1.0, 0.0, 9.0], [0.0, 1.0, 9.0], [0.0, 0.0, 1.0]],
[[1.0, 0.0, 18.0], [0.0, 1.0, 18.0], [0.0, 0.0, 1.0]],
[[1.0, 0.0, 27.0], [0.0, 1.0, 27.0], [0.0, 0.0, 1.0]],
]),
),
]
for inp, exp in zip(inputs, expectations):
with self.subTest(input=inp, expected=exp):
res = parametrize_matrix(**inp).to(exp.dtype)
self.assertTrue(
torch.allclose(res, matrix_to_cartesian(exp), atol=1e-6))
def test_check_image_size(self):
images = [
torch.rand(11, 2, 3, 4, 5),
torch.rand(11, 2, 3, 4),
torch.rand(11, 2, 3, 3),
]
img_sizes = [[3, 4, 5], [3, 4], 3]
scales = [
torch.tensor([2.0, 3.0, 4.0]),
torch.tensor([2.0, 3.0]),
torch.tensor([2.0, 3.0])
]
rots = [[45.0, 90.0, 135.0], [45.0], [45.0]]
trans = [[0.0, 10.0, 20.0], [10.0, 20.0], [10.0, 20.0]]
edges = [
[
[0.0, 0.0, 0.0, 1.0],
[0.0, 0.0, 5.0, 1.0],
[0.0, 4.0, 0.0, 1.0],
[0.0, 4.0, 5.0, 1.0],
[3.0, 0.0, 0.0, 1.0],
[3.0, 0.0, 5.0, 1.0],
[3.0, 4.0, 0.0, 1.0],
[3.0, 4.0, 5.0, 1.0],
],
[[0.0, 0.0, 1.0], [0.0, 4.0, 1.0], [3.0, 0.0, 1.0],
[3.0, 4.0, 1.0]],
[[0.0, 0.0, 1.0], [0.0, 3.0, 1.0], [3.0, 0.0, 1.0],
[3.0, 3.0, 1.0]],
]
for img, size, scale, rot, tran, edge_pts in zip(
images, img_sizes, scales, rots, trans, edges):
ndim = scale.size(-1)
with self.subTest(ndim=ndim):
affine = matrix_to_homogeneous(
parametrize_matrix(
scale=scale,
rotation=rot,
translation=tran,
degree=True,
batchsize=1,
ndim=ndim,
dtype=torch.float,
image_transform=False,
))
edge_pts = torch.tensor(edge_pts, dtype=torch.float)
img = img.to(torch.float)
new_edges = torch.bmm(edge_pts.unsqueeze(0),
affine.clone().permute(0, 2, 1))
img_size_zero_border = new_edges.max(dim=1)[0][0]
img_size_non_zero_border = (new_edges.max(dim=1)[0] -
new_edges.min(dim=1)[0])[0]
fn_result_zero_border = _check_new_img_size(
size,
matrix_to_cartesian(
affine.expand(img.size(0), -1, -1).clone()),
zero_border=True,
)
fn_result_non_zero_border = _check_new_img_size(
size,
matrix_to_cartesian(
affine.expand(img.size(0), -1, -1).clone()),
zero_border=False,
)
self.assertTrue(
torch.allclose(img_size_zero_border[:-1],
fn_result_zero_border))
self.assertTrue(
torch.allclose(img_size_non_zero_border[:-1],
fn_result_non_zero_border))
def test_matrix_parametrization(self):
inputs = [{
'scale': None,
'translation': None,
'rotation': None,
'batchsize': 2,
'ndim': 2,
'dtype': torch.float,
"image_transform": False
}, {
'scale': [4, 5],
'translation': [2, 3],
'rotation': 90,
'batchsize': 1,
'ndim': 2,
'dtype': torch.float,
"degree": True,
"image_transform": False
}, {
'scale': [4, 5],
'translation': [2, 3],
'rotation': 90,
'batchsize': 1,
'ndim': 2,
'dtype': torch.float,
"degree": True,
"image_transform": True
}, {
'scale': [2, 5],
'translation': [9, 18, 27],
'rotation': [180, 0, 180],
'degree': True,
'batchsize': 3,
'ndim': 2,
'dtype': torch.float,
"image_transform": False
}]
expectations = [
torch.tensor([[[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]],
[[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]]),
torch.tensor([[[0., -4., -12.], [5., 0., 10.], [0., 0., 1.]]]),
torch.tensor([[[0., 1 / 5., -2.], [-1 / 4., 0., -3.], [0., 0.,
1.]]]),
torch.bmm(
torch.bmm(
torch.tensor([[[2., 0., 0], [0., 5., 0.], [0., 0., 1.]],
[[2., 0., 0.], [0., 5., 0.], [0., 0., 1.]],
[[2., 0., 0.], [0., 5., 0.], [0., 0., 1.]]]),
torch.tensor([[[-1., 0., 0.], [0., -1., 0.], [0., 0., 1.]],
[[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]],
[[-1., 0., 0.], [0., -1., 0.], [0., 0.,
-1.]]])),
torch.tensor([[[1., 0., 9.], [0., 1., 9.], [0., 0., 1.]],
[[1., 0., 18.], [0., 1., 18.], [0., 0., 1.]],
[[1., 0., 27.], [0., 1., 27.], [0., 0., 1.]]]))
]
for inp, exp in zip(inputs, expectations):
with self.subTest(input=inp, expected=exp):
res = parametrize_matrix(**inp).to(exp.dtype)
self.assertTrue(
torch.allclose(res, matrix_to_cartesian(exp), atol=1e-6))
