|
|
|
import torch |
|
from torch.autograd import Function |
|
|
|
from ..utils import ext_loader |
|
|
|
ext_module = ext_loader.load_ext('_ext', ['ball_query_forward']) |
|
|
|
|
|
class BallQuery(Function): |
|
"""Find nearby points in spherical space.""" |
|
|
|
@staticmethod |
|
def forward(ctx, min_radius: float, max_radius: float, sample_num: int, |
|
xyz: torch.Tensor, center_xyz: torch.Tensor) -> 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 (Tensor): (B, N, 3) xyz coordinates of the features. |
|
center_xyz (Tensor): (B, npoint, 3) centers of the ball query. |
|
|
|
Returns: |
|
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 |
|
|
|
B, N, _ = xyz.size() |
|
npoint = center_xyz.size(1) |
|
idx = xyz.new_zeros(B, npoint, sample_num, dtype=torch.int) |
|
|
|
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 |
|
|
|
@staticmethod |
|
def backward(ctx, a=None): |
|
return None, None, None, None |
|
|
|
|
|
ball_query = BallQuery.apply |
|
|
