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)