def forward(
self,
points_xyz: torch.Tensor,
features: torch.Tensor = None,
indices: torch.Tensor = None,
target_xyz: torch.Tensor = None,
) -> (torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor):
"""forward.
Args:
points_xyz (Tensor): (B, N, 3) xyz coordinates of the features.
features (Tensor): (B, C, N) features of each point.
Default: None.
indices (Tensor): (B, num_point) Index of the features.
Default: None.
target_xyz (Tensor): (B, M, 3) new_xyz coordinates of the outputs.
Returns:
Tensor: (B, M, 3) where M is the number of points.
New features xyz.
Tensor: (B, M, sum_k(mlps[k][-1])) where M is the number
of points. New feature descriptors.
Tensor: (B, M) where M is the number of points.
Index of the features.
"""
new_features_list = []
xyz_flipped = points_xyz.transpose(1, 2).contiguous()
if indices is not None:
assert (indices.shape[1] == self.num_point[0])
new_xyz = gather_points(xyz_flipped, indices).transpose(
1, 2).contiguous() if self.num_point is not None else None
elif target_xyz is not None:
new_xyz = target_xyz.contiguous()
else:
indices = self.points_sampler(points_xyz, features)
new_xyz = gather_points(xyz_flipped, indices).transpose(
1, 2).contiguous() if self.num_point is not None else None
for i in range(len(self.groupers)):
# (B, C, num_point, nsample)
new_features = self.groupers[i](points_xyz, new_xyz, features)
# (B, mlp[-1], num_point, nsample)
new_features = self.mlps[i](new_features)
if self.pool_mod == 'max':
# (B, mlp[-1], num_point, 1)
new_features = F.max_pool2d(
new_features, kernel_size=[1, new_features.size(3)])
elif self.pool_mod == 'avg':
# (B, mlp[-1], num_point, 1)
new_features = F.avg_pool2d(
new_features, kernel_size=[1, new_features.size(3)])
else:
raise NotImplementedError
new_features = new_features.squeeze(-1) # (B, mlp[-1], num_point)
new_features_list.append(new_features)
return new_xyz, torch.cat(new_features_list, dim=1), indices
def _sample_points(self, points_xyz, features, indices, target_xyz):
"""Perform point sampling based on inputs.
If `indices` is specified, directly sample corresponding points.
Else if `target_xyz` is specified, use is as sampled points.
Otherwise sample points using `self.points_sampler`.
Args:
points_xyz (Tensor): (B, N, 3) xyz coordinates of the features.
features (Tensor): (B, C, N) features of each point.
Default: None.
indices (Tensor): (B, num_point) Index of the features.
Default: None.
target_xyz (Tensor): (B, M, 3) new_xyz coordinates of the outputs.
Returns:
Tensor: (B, num_point, 3) sampled xyz coordinates of points.
Tensor: (B, num_point) sampled points' index.
"""
xyz_flipped = points_xyz.transpose(1, 2).contiguous()
if indices is not None:
assert (indices.shape[1] == self.num_point[0])
new_xyz = gather_points(xyz_flipped, indices).transpose(
1, 2).contiguous() if self.num_point is not None else None
elif target_xyz is not None:
new_xyz = target_xyz.contiguous()
else:
indices = self.points_sampler(points_xyz, features)
new_xyz = gather_points(xyz_flipped, indices).transpose(
1, 2).contiguous() if self.num_point is not None else None
return new_xyz, indices
def test_gather_points():
if not torch.cuda.is_available():
pytest.skip()
features = torch.tensor([[[
-1.6095, -0.1029, -0.8876, -1.2447, -2.4031, 0.3708, -1.1586, -1.4967,
-0.4800, 0.2252
],
[
1.9138, 3.4979, 1.6854, 1.5631, 3.6776,
3.1154, 2.1705, 2.5221, 2.0411, 3.1446
],
[
-1.4173, 0.3073, -1.4339, -1.4340, -1.2770,
-0.2867, -1.4162, -1.4044, -1.4245, -1.4074
]],
[[
0.2160, 0.0842, 0.3661, -0.2749, -0.4909,
-0.6066, -0.8773, -0.0745, -0.9496, 0.1434
],
[
1.3644, 1.8087, 1.6855, 1.9563, 1.2746,
1.9662, 0.9566, 1.8778, 1.1437, 1.3639
],
[
-0.7172, 0.1692, 0.2241, 0.0721, -0.7540,
0.0462, -0.6227, 0.3223, -0.6944, -0.5294
]]]).cuda()
idx = torch.tensor([[0, 1, 4, 0, 0, 0], [0, 5, 6, 0, 0, 0]]).int().cuda()
output = gather_points(features, idx)
expected_output = torch.tensor(
[[[-1.6095, -0.1029, -2.4031, -1.6095, -1.6095, -1.6095],
[1.9138, 3.4979, 3.6776, 1.9138, 1.9138, 1.9138],
[-1.4173, 0.3073, -1.2770, -1.4173, -1.4173, -1.4173]],
[[0.2160, -0.6066, -0.8773, 0.2160, 0.2160, 0.2160],
[1.3644, 1.9662, 0.9566, 1.3644, 1.3644, 1.3644],
[-0.7172, 0.0462, -0.6227, -0.7172, -0.7172, -0.7172]]]).cuda()
assert torch.allclose(output, expected_output)
output_half = gather_points(features.half(), idx)
assert torch.allclose(output_half, expected_output.half())
def forward(self, xyz, features, sample_inds):
"""Forward pass.
Args:
xyz： (B, N, 3) the coordinates of the features.
features (Tensor): (B, C, N) features to sample.
sample_inds (Tensor): (B, M) the given index,
where M is the number of points.
Returns:
Tensor: (B, M, 3) coordinates of sampled features
Tensor: (B, C, M) the sampled features.
Tensor: (B, M) the given index.
"""
xyz_t = xyz.transpose(1, 2).contiguous()
new_xyz = gather_points(xyz_t, sample_inds).transpose(1,
2).contiguous()
new_features = gather_points(features, sample_inds).contiguous()
return new_xyz, new_features, sample_inds
