def forward(ctx, radius: float, nsample: int, xyz: torch.Tensor, new_xyz: torch.Tensor) -> torch.Tensor: """ :param ctx: :param radius: float, radius of the balls :param nsample: int, maximum number of features in the balls :param xyz: (B, N, 3) xyz coordinates of the features :param new_xyz: (B, npoint, 3) centers of the ball query :return: idx: (B, npoint, nsample) tensor with the indicies of the features that form the query balls """ assert new_xyz.is_contiguous() assert xyz.is_contiguous() B, N, _ = xyz.size() npoint = new_xyz.size(1) idx = torch.cuda.IntTensor(B, npoint, nsample).zero_() pointnet2.ball_query_wrapper(B, N, npoint, radius, nsample, new_xyz, xyz, idx) return idx
def forward(ctx, radius: float, nsample: int, xyz: torch.Tensor, new_xyz: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """ :param ctx: :param radius: float, radius of the balls :param nsample: int, maximum number of features in the balls :param xyz: (B, N, 3) xyz coordinates of the features :param new_xyz: (B, npoint, 3) centers of the ball query :return: idx: (B, npoint, nsample) tensor with the indicies of the features that form the query balls """ # assert new_xyz.is_contiguous() # assert xyz.is_contiguous() B, N, _ = xyz.size() npoint = new_xyz.size(1) idx = torch.cuda.IntTensor(B, npoint, nsample).zero_() his = torch.cuda.FloatTensor(B, npoint, int(nsample/3)).zero_() # idx = torch.zeros(B, npoint, nsample).type(torch.cuda.IntTensor) # his = torch.zeros(B, npoint, nsample).type(torch.cuda.FloatTensor) # print("1",idx) pointnet2.ball_query_wrapper(B, N, npoint, radius, nsample, new_xyz, xyz, idx, his) # print("2",idx) return idx, his