def _check_new_img_size(curr_img_size,
matrix: torch.Tensor,
zero_border: bool = False) -> torch.Tensor:
"""
Calculates the image size so that the whole image content fits the image.
The resulting size will be the maximum size of the batch, so that the
images can remain batched.
Args:
curr_img_size: the size of the current image.
If int, it will be used as size for all image dimensions
matrix: a batch of affine matrices with shape [N, NDIM, NDIM+1]
zero_border: whether or not to have a fixed image border at zero
Returns:
torch.Tensor: the new image size
"""
n_dim = matrix.size(-1) - 1
if check_scalar(curr_img_size):
curr_img_size = [curr_img_size] * n_dim
possible_points = unit_box(n_dim, torch.tensor(curr_img_size)).to(matrix)
transformed_edges = affine_point_transform(
possible_points[None].expand(matrix.size(0),
*[-1 for _ in possible_points.shape
]).clone(), matrix)
if zero_border:
substr = 0
else:
substr = transformed_edges.min(1)[0]
return (transformed_edges.max(1)[0] - substr).max(0)[0]
def test_unit_box_3d(self):
curr_img_size = torch.tensor([2, 3, 4])
box = torch.tensor([
[0.0, 0.0, 0.0],
[0.0, 0.0, curr_img_size[2]],
[0.0, curr_img_size[1], 0],
[0.0, curr_img_size[1], curr_img_size[2]],
[curr_img_size[0], 0.0, 0.0],
[curr_img_size[0], 0.0, curr_img_size[2]],
[curr_img_size[0], curr_img_size[1], 0.0],
curr_img_size,
])
created_box = unit_box(3, curr_img_size).to(box)
self.compare_points_unordered(box, created_box)
def test_unit_box_2d(self):
curr_img_size = torch.tensor([2, 3])
box = torch.tensor([[0., 0.], [0., curr_img_size[1]],
[curr_img_size[0], 0], curr_img_size])
created_box = unit_box(2, curr_img_size).to(box)
self.compare_points_unordered(box, created_box)