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 test_simple(self):
features = torch.tensor([
[[0, 10, 0], [1, 11, 0], [2, 12, 0]],
[
[100, 110, 120],
[101, 111, 121],
[102, 112, 122],
], # x-coordinates # y-coordinates # z-coordinates
]).type(torch.float)
idx = torch.tensor([[[1, 0], [0, 0]], [[0, 1], [1,
2]]]).type(torch.long)
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,
)
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 _prepare_features(self,
xyz: torch.Tensor,
new_xyz: torch.Tensor,
features: torch.Tensor = None,
idx: torch.Tensor = None) -> Tuple[torch.Tensor]:
"""
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
"""
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, 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([raw_grouped_xyz, grouped_xyz], dim=1)
return new_features