model/pvcnn/modules/functional/ball_query.py

from torch.autograd import Function

from .backend import _backend

__all__ = ['ball_query']


def ball_query(centers_coords, points_coords, radius, num_neighbors):
        """
        :param centers_coords: coordinates of centers, FloatTensor[B, 3, M]
        :param points_coords: coordinates of points, FloatTensor[B, 3, N]
        :param radius: float, radius of ball query
        :param num_neighbors: int, maximum number of neighbors
        :return:
            neighbor_indices: indices of neighbors, IntTensor[B, M, U]
        """
        centers_coords = centers_coords.contiguous()
        points_coords = points_coords.contiguous()
        return _backend.ball_query(centers_coords, points_coords, radius, num_neighbors)