def furthest_point_sample(xyz, npoint):
# type: (Any, torch.Tensor, int) -> torch.Tensor
r"""
Uses iterative furthest point sampling to select a set of npoint features that have the largest
minimum distance
Parameters
----------
xyz : torch.Tensor
(B, N, 3) tensor where N > npoint
npoint : int32
number of features in the sampled set
Returns
-------
torch.Tensor
(B, npoint) tensor containing the set
"""
if npoint > xyz.shape[1]:
raise ValueError(
"caanot sample %i points from an input set of %i points" %
(npoint, xyz.shape[1]))
if xyz.is_cuda:
return tpcuda.furthest_point_sampling(xyz, npoint)
else:
return tpcpu.fps(xyz, npoint, True)
def _multiscale_compute_fn(self,
batch,
collate_fn=None,
precompute_multi_scale=False,
num_scales=0,
sample_method='random'):
batch = collate_fn(batch)
if not precompute_multi_scale:
return batch
multiscale = []
pos = batch.pos # [B, N, 3]
for i in range(num_scales):
neighbor_idx = self._knn_search(pos, pos, self.kernel_size[i]) # [B, N, K]
sample_num = pos.shape[1] // self.ratio[i]
if sample_method.lower() == 'random':
choice = torch.randperm(pos.shape[1])[:sample_num]
sub_pos = pos[:, choice, :] # random sampled pos [B, S, 3]
sub_idx = neighbor_idx[:, choice, :] # the pool idx [B, S, K]
elif sample_method.lower() == 'fps':
choice = tpcuda.furthest_point_sampling(pos.cuda(), sample_num).to(torch.long).cpu()
sub_pos = pos.gather(dim=1, index=choice.unsqueeze(-1).repeat(1, 1, pos.shape[-1]))
sub_idx = neighbor_idx.gather(dim=1, index=choice.unsqueeze(-1).repeat(1, 1, neighbor_idx.shape[-1]))
else:
raise NotImplementedError('Only `random` or `fps` sampling method is implemented!')
up_idx = self._knn_search(sub_pos, pos, 1) # [B, N, 1]
multiscale.append(Data(pos=pos, neighbor_idx=neighbor_idx, sub_idx=sub_idx, up_idx=up_idx))
pos = sub_pos
return MultiScaleData(x=batch.x,
y=batch.y,
point_idx=batch.point_idx,
cloud_idx=batch.cloud_idx,
multiscale=multiscale)