def test_simple(self):
a = torch.tensor([[[0, 0, 0], [1, 0, 0], [2, 0,
0]]]).to(torch.float).cuda()
b = torch.tensor([[[0, 0, 0]]]).to(torch.float).cuda()
npt.assert_array_equal(
ball_query(1, 2, a, b).detach().cpu().numpy(), np.array([[[0,
0]]]))
def _variance_estimator_dense(r, pos, f):
nei_idx = tp.ball_query(r, _MAX_NEIGHBOURS, pos, pos,
sort=True)[0].reshape(pos.shape[0],
-1).long() # [B,N * nei]
f_neighboors = f.gather(1, nei_idx).reshape(f.shape[0], f.shape[1],
-1) # [B,N , nei]
gradient = (f.unsqueeze(-1).repeat(1, 1, f_neighboors.shape[-1]) -
f_neighboors)**2 # [B,N,nei]
return gradient.sum(-1)
def find_neighbours(self, x, y, scale_idx=0):
if scale_idx >= self.num_scales:
raise ValueError("Scale %i is out of bounds %i" %
(scale_idx, self.num_scales))
num_neighbours = self._max_num_neighbors[scale_idx]
neighbours = tp.ball_query(self._radius[scale_idx], num_neighbours, x,
y)[0]
if DEBUGGING_VARS["FIND_NEIGHBOUR_DIST"]:
for i in range(neighbours.shape[0]):
start = neighbours[i, :, 0]
valid_neighbours = (neighbours[i, :, 1:] != start.view(
(-1, 1)).repeat(1, num_neighbours - 1)).sum(1) + 1
self._dist_meters[scale_idx].add_valid_neighbours(
valid_neighbours)
return neighbours
def find_neighbours(self, x, y, batch_x=None, batch_y=None):
if self._conv_type == ConvolutionFormat.MESSAGE_PASSING.value:
return radius(x,
y,
self._radius,
batch_x,
batch_y,
max_num_neighbors=self._max_num_neighbors)
elif self._conv_type == ConvolutionFormat.DENSE.value or ConvolutionFormat.PARTIAL_DENSE.value:
return tp.ball_query(self._radius,
self._max_num_neighbors,
x,
y,
mode=self._conv_type,
batch_x=batch_x,
batch_y=batch_y)
else:
raise NotImplementedError
def find_neighbours(self, x, y, scale_idx=0):
if scale_idx >= self.num_scales:
raise ValueError("Scale %i is out of bounds %i" %
(scale_idx, self.num_scales))
num_neighbours = self._max_num_neighbors[scale_idx]
neighbours = tp.ball_query(self._radius[scale_idx], num_neighbours, x,
y)
# mean_count = (
# (neighbours[0, :, :] != neighbours[0, :, 0].view((-1, 1)).repeat(1, num_neighbours)).sum(1).float().mean()
# )
# for i in range(1, neighbours.shape[0]):
# start = neighbours[i, :, 0]
# valid_neighbours = (neighbours[i, :, :] != start.view((-1, 1)).repeat(1, num_neighbours)).sum(1)
# mean_count += valid_neighbours.float().mean()
# mean_count = mean_count / neighbours.shape[0]
# print(
# "Radius: %f, Num_neighbours %i, actual, %f" % (self._radius[scale_idx], num_neighbours, mean_count.item())
# )
return neighbours
def forward(
self,
xyz: torch.Tensor,
new_xyz: torch.Tensor,
features: torch.Tensor = None,
fps_idx: torch.IntTensor = None
) -> Tuple[torch.Tensor]:
r"""
Parameters
----------
xyz : torch.Tensor
xyz coordinates of the features (B, N, 3)
new_xyz : torch.Tensor
centriods (B, npoint, 3)
features : torch.Tensor
Descriptors of the features (B, C, N)
Returns
-------
new_features : torch.Tensor
(B, 3 + C, npoint, nsample) tensor
"""
idx = tp.ball_query(self.radius, self.nsample, xyz, new_xyz, mode='dense')
xyz_trans = xyz.transpose(1, 2).contiguous()
grouped_xyz = tp.grouping_operation(
xyz_trans, idx
) # (B, 3, npoint, nsample)
raw_grouped_xyz = grouped_xyz
grouped_xyz -= new_xyz.transpose(1, 2).unsqueeze(-1)
if features is not None:
grouped_features = tp.grouping_operation(features, idx)
if self.use_xyz:
new_features = torch.cat([raw_grouped_xyz, grouped_xyz, grouped_features],
dim=1) # (B, C + 3 + 3, npoint, nsample)
else:
new_features = grouped_features
else:
assert self.use_xyz, "Cannot have not features and not use xyz as a feature!"
new_features = torch.cat([raw_grouped_xyz, grouped_xyz], dim = 1)
return new_features
