def three_nn(xyz1, xyz2):
'''
:param xyz1: shape=(B, N1, 3)
:param xyz2: shape=(B, N2, 3)
:return: dists: shape=(B, N1, 3), inds: shape=(B, N1, 3)
'''
dists = get_dists(xyz1, xyz2)
dists, inds = torch.sort(dists, dim=-1)
dists, inds = dists[:, :, :3], inds[:, :, :3]
return dists, inds
def ball_query(xyz, new_xyz, radius, K):
device = xyz.device
B, N, C = xyz.shape
M = new_xyz.shape[1]
grouped_inds = torch.arange(0, N, dtype=torch.long).to(device).view(
1, 1, N).repeat(B, M, 1)
dists = get_dists(new_xyz, xyz)
grouped_inds[dists > radius] = N
grouped_inds = torch.sort(grouped_inds, dim=-1)[0][:, :, :K]
grouped_min_inds = grouped_inds[:, :, 0:1].repeat(1, 1, K)
grouped_inds[grouped_inds == N] = grouped_min_inds[grouped_inds == N]
return grouped_inds
def fps(xyz, M):
device = xyz.device
B, N, C = xyz.shape
centroids = torch.zeros(size=(B, M), dtype=torch.long).to(device)
dists = torch.ones(B, N).to(device) * 1e5
inds = torch.randint(0, N, size=(B, ), dtype=torch.long).to(device)
batchlists = torch.arange(0, B, dtype=torch.long).to(device)
for i in range(M):
centroids[:, i] = inds
cur_point = xyz[batchlists, inds, :] # (B, 3)
cur_dist = torch.squeeze(get_dists(torch.unsqueeze(cur_point, 1), xyz),
dim=1)
dists[cur_dist < dists] = cur_dist[cur_dist < dists]
inds = torch.max(dists, dim=1)[1]
return centroids
def fps(xyz, M):
'''
Sample M points from points according to farthest point sampling (FPS) algorithm.
:param xyz: shape=(B, N, 3)
:return: inds: shape=(B, M)
'''
device = xyz.device
B, N, C = xyz.shape
centroids = torch.zeros(size=(B, M), dtype=torch.long).to(device)
dists = torch.ones(B, N).to(device) * 1e5
inds = torch.randint(0, N, size=(B, ), dtype=torch.long).to(device)
batchlists = torch.arange(0, B, dtype=torch.long).to(device)
for i in range(M):
centroids[:, i] = inds
cur_point = xyz[batchlists, inds, :] # (B, 3)
cur_dist = torch.squeeze(get_dists(torch.unsqueeze(cur_point, 1), xyz))
dists[cur_dist < dists] = cur_dist[cur_dist < dists]
inds = torch.max(dists, dim=1)[1]
return centroids
def ball_query(xyz, new_xyz, radius, K):
'''
:param xyz: shape=(B, N, 3)
:param new_xyz: shape=(B, M, 3)
:param radius: int
:param K: int, an upper limit samples
:return: shape=(B, M, K)
'''
device = xyz.device
B, N, C = xyz.shape
M = new_xyz.shape[1]
grouped_inds = torch.arange(0, N, dtype=torch.long).to(device).view(1, 1, N).repeat(B, M, 1)
dists = get_dists(new_xyz, xyz)
grouped_inds[dists > radius] = N
grouped_inds = torch.sort(grouped_inds, dim=-1)[0][:, :, :K]
grouped_min_inds = grouped_inds[:, :, 0:1].repeat(1, 1, K)
grouped_inds[grouped_inds == N] = grouped_min_inds[grouped_inds == N]
return grouped_inds
def three_nn(xyz1, xyz2):
dists = get_dists(xyz1, xyz2)
dists, inds = torch.sort(dists, dim=-1)
dists, inds = dists[:, :, :3], inds[:, :, :3]
return dists, inds