Spaces:
Running
on
A10G
Running
on
A10G
File size: 2,550 Bytes
251e479 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import torch
from torch.autograd import Function
from ..utils import ext_loader
ext_module = ext_loader.load_ext('_ext', [
'furthest_point_sampling_forward',
'furthest_point_sampling_with_dist_forward'
])
class FurthestPointSampling(Function):
"""Uses iterative furthest point sampling to select a set of features whose
corresponding points have the furthest distance."""
@staticmethod
def forward(ctx, points_xyz: torch.Tensor,
num_points: int) -> torch.Tensor:
"""
Args:
points_xyz (Tensor): (B, N, 3) where N > num_points.
num_points (int): Number of points in the sampled set.
Returns:
Tensor: (B, num_points) indices of the sampled points.
"""
assert points_xyz.is_contiguous()
B, N = points_xyz.size()[:2]
output = torch.cuda.IntTensor(B, num_points)
temp = torch.cuda.FloatTensor(B, N).fill_(1e10)
ext_module.furthest_point_sampling_forward(
points_xyz,
temp,
output,
b=B,
n=N,
m=num_points,
)
if torch.__version__ != 'parrots':
ctx.mark_non_differentiable(output)
return output
@staticmethod
def backward(xyz, a=None):
return None, None
class FurthestPointSamplingWithDist(Function):
"""Uses iterative furthest point sampling to select a set of features whose
corresponding points have the furthest distance."""
@staticmethod
def forward(ctx, points_dist: torch.Tensor,
num_points: int) -> torch.Tensor:
"""
Args:
points_dist (Tensor): (B, N, N) Distance between each point pair.
num_points (int): Number of points in the sampled set.
Returns:
Tensor: (B, num_points) indices of the sampled points.
"""
assert points_dist.is_contiguous()
B, N, _ = points_dist.size()
output = points_dist.new_zeros([B, num_points], dtype=torch.int32)
temp = points_dist.new_zeros([B, N]).fill_(1e10)
ext_module.furthest_point_sampling_with_dist_forward(
points_dist, temp, output, b=B, n=N, m=num_points)
if torch.__version__ != 'parrots':
ctx.mark_non_differentiable(output)
return output
@staticmethod
def backward(xyz, a=None):
return None, None
furthest_point_sample = FurthestPointSampling.apply
furthest_point_sample_with_dist = FurthestPointSamplingWithDist.apply
|