def _prepare_features(self, x, pos, new_pos, idx):
new_pos_trans = pos.transpose(1, 2).contiguous()
grouped_pos_absolute = tp.grouping_operation(
new_pos_trans, idx) # (B, 3, npoint, nsample)
centroids = new_pos.transpose(1, 2).unsqueeze(-1)
grouped_pos_normalized = grouped_pos_absolute - centroids
if x is not None:
grouped_features = tp.grouping_operation(x, idx)
if self.use_xyz:
new_features = torch.cat(
[
grouped_pos_absolute, grouped_pos_normalized,
grouped_features
],
dim=1) # (B, 3 + 3 + C, 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(
[grouped_pos_absolute, grouped_pos_normalized],
dim=1) # (B, 3 + 3 npoint, nsample)
return new_features, centroids
def _prepare_features(self, x, pos, new_pos, idx):
new_pos_trans = pos.transpose(1, 2).contiguous()
grouped_pos = tp.grouping_operation(new_pos_trans, idx) # (B, 3, npoint, nsample)
grouped_pos -= new_pos.transpose(1, 2).unsqueeze(-1)
if x is not None:
grouped_features = tp.grouping_operation(x, idx)
if self.use_xyz:
new_features = torch.cat([grouped_pos, grouped_features], dim=1) # (B, C + 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 = grouped_pos
return new_features
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
def test_simple(self):
features = torch.tensor([
[[0, 10, 0], [1, 11, 0], [2, 12, 0]],
[
[100, 110, 120], # x-coordinates
[101, 111, 121], # y-coordinates
[102, 112, 122], # z-coordinates
]
])
idx = torch.tensor([
[[1, 0], [0, 0]],
[[0, 1], [1, 2]]
])
expected = np.array([
[
[[10, 0], [0, 0]],
[[11, 1], [1, 1]],
[[12, 2], [2, 2]]
],
[ # 2nd batch
[ # x-coordinates
[100, 110], #x-coordinates of samples for point 0
[110, 120], #x-coordinates of samples for point 1
],
[[101, 111], [111, 121]], # y-coordinates
[[102, 112], [112, 122]], # z-coordinates
]
])
cpu_output = grouping_operation(features, idx).detach().cpu().numpy()
npt.assert_array_equal(expected, cpu_output)
if torch.cuda.is_available():
npt.assert_array_equal(
grouping_operation(
features.cuda(),
idx.cuda()
).detach().cpu().numpy(), expected)