Spaces:
Build error
Build error
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| from typing import Optional, Tuple | |
| import torch | |
| from torch.autograd import Function | |
| from ..utils import ext_loader | |
| ext_module = ext_loader.load_ext( | |
| '_ext', ['ball_query_forward', 'stack_ball_query_forward']) | |
| class BallQuery(Function): | |
| """Find nearby points in spherical space.""" | |
| def forward( | |
| ctx, | |
| min_radius: float, | |
| max_radius: float, | |
| sample_num: int, | |
| xyz: torch.Tensor, | |
| center_xyz: torch.Tensor, | |
| xyz_batch_cnt: Optional[torch.Tensor] = None, | |
| center_xyz_batch_cnt: Optional[torch.Tensor] = None | |
| ) -> torch.Tensor: | |
| """ | |
| Args: | |
| min_radius (float): minimum radius of the balls. | |
| max_radius (float): maximum radius of the balls. | |
| sample_num (int): maximum number of features in the balls. | |
| xyz (torch.Tensor): (B, N, 3) xyz coordinates of the features, | |
| or staked input (N1 + N2 ..., 3). | |
| center_xyz (torch.Tensor): (B, npoint, 3) centers of the ball | |
| query, or staked input (M1 + M2 ..., 3). | |
| xyz_batch_cnt: (batch_size): Stacked input xyz coordinates nums in | |
| each batch, just like (N1, N2, ...). Defaults to None. | |
| New in version 1.7.0. | |
| center_xyz_batch_cnt: (batch_size): Stacked centers coordinates | |
| nums in each batch, just line (M1, M2, ...). Defaults to None. | |
| New in version 1.7.0. | |
| Returns: | |
| torch.Tensor: (B, npoint, nsample) tensor with the indices of the | |
| features that form the query balls. | |
| """ | |
| assert center_xyz.is_contiguous() | |
| assert xyz.is_contiguous() | |
| assert min_radius < max_radius | |
| if xyz_batch_cnt is not None and center_xyz_batch_cnt is not None: | |
| assert xyz_batch_cnt.dtype == torch.int | |
| assert center_xyz_batch_cnt.dtype == torch.int | |
| idx = center_xyz.new_zeros((center_xyz.shape[0], sample_num), | |
| dtype=torch.int32) | |
| ext_module.stack_ball_query_forward( | |
| center_xyz, | |
| center_xyz_batch_cnt, | |
| xyz, | |
| xyz_batch_cnt, | |
| idx, | |
| max_radius=max_radius, | |
| nsample=sample_num, | |
| ) | |
| else: | |
| B, N, _ = xyz.size() | |
| npoint = center_xyz.size(1) | |
| idx = xyz.new_zeros(B, npoint, sample_num, dtype=torch.int32) | |
| ext_module.ball_query_forward( | |
| center_xyz, | |
| xyz, | |
| idx, | |
| b=B, | |
| n=N, | |
| m=npoint, | |
| min_radius=min_radius, | |
| max_radius=max_radius, | |
| nsample=sample_num) | |
| if torch.__version__ != 'parrots': | |
| ctx.mark_non_differentiable(idx) | |
| return idx | |
| def backward(ctx, a=None) -> Tuple[None, None, None, None]: | |
| return None, None, None, None | |
| ball_query = BallQuery.apply | |