def test_input_unchanged(self):
x = torch.tensor(
[
[1, 0, 0],
[0, 1, 0],
[0, 0, 1],
]
).float()
x_copy = x.clone()
center_crop(x, 0.5, 0.5, 0.5)
assert torch.allclose(x, x_copy)
def test_optional_sizes(self):
x = torch.tensor(
[
[1, 0, 0],
[0, 1, 0],
[0, 0, 1],
]
).float()
mask = center_crop(x, 0.5)
assert not mask[0]
assert mask.sum() == 2
def test_masked_range(self):
x = torch.tensor(
[
[1, 0, 0],
[0, 2, 0],
[0, 0, 3],
]
).float()
mask = center_crop(x, 2.0, 1.0, 6.1)
assert mask[0]
assert not mask[1]
assert mask[2]
def test_apply_mask_to_coords(self, batched, dtype):
x = torch.tensor(
[
[1, 0, 0],
[0, 1, 0],
[0, 0, 1],
]
)
if batched:
x.unsqueeze_(0).repeat(2, 1, 1)
x = x.long() if dtype == "long" else x.float()
mask = center_crop(x, 100, 100, 0.5)
cropped = x[mask]
assert torch.allclose(x[..., :2, :], cropped)
def test_crop_big_window_points_unchanged(self, batched, dtype):
x = torch.tensor(
[
[1, 0, 0],
[0, 1, 0],
[0, 0, 1],
]
)
if batched:
x.unsqueeze_(0)
x = x.long() if dtype == "long" else x.float()
mask = center_crop(x, 100, 100, 100)
assert mask.ndim == x.ndim - 1
assert mask.all()
assert not batched or mask.shape[0] == 1
def projection_mapping(
coords: Tensor,
resolution: float = 1.0,
image_size: Optional[Tuple[int, int]] = None,
) -> Tuple[Tensor, Tensor, Tuple[int, int]]:
r"""Performs a parallel projection of Cartesian coordinates at a given resolution, computing a mapping of each
3D point to a 2D pixel. Returns a mapping of points within the image area to height and witdth pixel coordinates,
along with a mask of points that lie within the projected window and the height and width of the 2D projection.
Args:
coords (Tensor):
Cartesian coordinates comprising a point cloud, in XYZ order.
resolution (float):
The resolution at which to construct the 2D projection
image_size (optional, 2-tuple of ints):
By default, the size of the resultant projection is determined based on ``resolution`` and the
range of x-y coordinates in ``coords``. An optional image size can be provided to override
this behavior.
Returns:
Tuple of the form ``mapping, mask (height, width)``
Shapes
* ``coords`` - :math:`(N, 3)` or :math:`(1, N, 3)`
* Mapping - :math:`(N, 2)` where the final dimension gives the height and width coordinate
* Mask - :math:`(N)`
"""
# apply batch dim if not present
coords = coords.view(coords.shape[-2], coords.shape[-1])
mask = torch.empty(1, device=coords.device).type_as(coords)
x, y, z = coords[..., 0], coords[..., 1], coords[..., 2]
# find xy mins/maxes
min_x = coords[..., 0].min()
max_x = coords[..., 0].max()
min_y = coords[..., 1].min()
max_y = coords[..., 1].max()
mins = torch.stack([min_x, min_y], dim=0)
maxes = torch.stack([max_x, max_y], dim=0)
if image_size is not None:
height, width = image_size
# crop point cloud based on resolution and image size
crop_height = float(height)
crop_width = float(width)
mask = center_crop(coords, crop_height, crop_width)
assert mask.any()
coords = coords[mask]
else:
# calculate size one from min/max
height = int(max_y.sub(min_y).floor_divide_(resolution).long().item())
width = int(max_x.sub(min_x).floor_divide_(resolution).long().item())
mask = torch.tensor([True],
device=coords.device).expand(coords.shape[0])
assert height > 0
assert width > 0
x, y, z = coords[..., 0], coords[..., 1], coords[..., 2]
# recompute mins/maxes after cropping
mins = torch.min(coords[..., :2], dim=0).values
maxes = torch.max(coords[..., :2], dim=0).values
# map each point to a height/width in the 2d grid
x, y, z = coords[..., 0], coords[..., 1], coords[..., 2]
mapping = coords[..., :2].roll(1, dims=-1).sub(mins.roll(
1, dims=-1)).floor_divide(resolution)
mapping[..., 0].clamp_max_(height - 1)
mapping[..., 1].clamp_max_(width - 1)
return mapping, mask, (height, width)
def func():
center_crop(coords, 500, 500)