| import torch |
| from torch.autograd import Function |
|
|
| from pointops._C import farthest_point_sampling_cuda |
|
|
|
|
| class FarthestPointSampling(Function): |
| @staticmethod |
| def forward(ctx, xyz, offset, new_offset): |
| """ |
| input: coords: (n, 3), offset: (b), new_offset: (b) |
| output: idx: (m) |
| """ |
| assert xyz.is_contiguous() |
| n, b, n_max = xyz.shape[0], offset.shape[0], offset[0] |
| for i in range(1, b): |
| n_max = max(offset[i] - offset[i - 1], n_max) |
| idx = torch.zeros(new_offset[b - 1].item(), dtype=torch.int, device=xyz.device) |
| tmp = torch.ones(n, dtype=torch.float, device=xyz.device) * 1e10 |
| farthest_point_sampling_cuda( |
| b, n_max, xyz, offset.int(), new_offset.int(), tmp, idx |
| ) |
| del tmp |
| return idx |
|
|
|
|
| farthest_point_sampling = FarthestPointSampling.apply |
|
|
