def forward(self, xyz1, xyz2, points1, points2):
"""
Args:
xyz1: Tensor, (B, 3, N)
xyz2: Tensor, (B, 3, M)
points1: Tensor, (B, in_channel, N)
points2: Tensor, (B, in_channel, M)
Returns:
new_points: Tensor, (B, mlp[-1], N)
"""
dist, idx = three_nn(
xyz1.permute(0, 2, 1).contiguous(),
xyz2.permute(0, 2, 1).contiguous())
dist = torch.clamp_min(dist, 1e-10) # (B, N, 3)
recip_dist = 1.0 / dist
norm = torch.sum(recip_dist, 2, keepdim=True).repeat((1, 1, 3))
weight = recip_dist / norm
interpolated_points = three_interpolate(points2, idx,
weight) # B, in_channel, N
if self.use_points1:
new_points = torch.cat([interpolated_points, points1], 1)
else:
new_points = interpolated_points
new_points = self.mlp_conv(new_points)
return new_points
def forward(self, unknown, known, unknow_feats, known_feats):
# type: (PointnetFPModule, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor) -> torch.Tensor
r"""
Parameters
----------
unknown : torch.Tensor
(B, n, 3) tensor of the xyz positions of the unknown features
known : torch.Tensor
(B, m, 3) tensor of the xyz positions of the known features
unknow_feats : torch.Tensor
(B, C1, n) tensor of the features to be propigated to
known_feats : torch.Tensor
(B, C2, m) tensor of features to be propigated
Returns
-------
new_features : torch.Tensor
(B, mlp[-1], n) tensor of the features of the unknown features
"""
# known_feats are features at positions specified at known
# unknow_feats are features at positions specified at unknown
# unknow_feats features at the positons unknown obtained in the encoder
# we first interpolate known_feats to interpolated_feats at positions unknown
# then concat interpolated_feats with unknow_feats to get new_features
# then new_features is put through a pointnet
if known is not None:
dist, idx = pointnet2_utils.three_nn(unknown, known)
dist_recip = 1.0 / (dist + 1e-8)
norm = torch.sum(dist_recip, dim=2, keepdim=True)
weight = dist_recip / norm
interpolated_feats = pointnet2_utils.three_interpolate(
known_feats, idx, weight)
else:
interpolated_feats = known_feats.expand(*(known_feats.size()[0:2] +
[unknown.size(1)]))
if unknow_feats is not None:
new_features = torch.cat([interpolated_feats, unknow_feats],
dim=1) # (B, C2 + C1, n)
else:
new_features = interpolated_feats
new_features = new_features.unsqueeze(-1)
new_features = self.mlp(new_features)
return new_features.squeeze(-1)
def forward(self, unknown, known, unknow_feats, known_feats):
# type: (PointnetFPModule, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor) -> torch.Tensor
r"""
Parameters
----------
unknown : torch.Tensor
(B, n, 3) tensor of the xyz positions of the unknown features
known : torch.Tensor
(B, m, 3) tensor of the xyz positions of the known features
unknow_feats : torch.Tensor
(B, C1, n) tensor of the features to be propigated to
known_feats : torch.Tensor
(B, C2, m) tensor of features to be propigated
Returns
-------
new_features : torch.Tensor
(B, mlp[-1], n) tensor of the features of the unknown features
"""
if known is not None:
dist, idx = pointnet2_utils.three_nn(unknown, known)
dist_recip = 1.0 / (dist + 1e-8)
norm = torch.sum(dist_recip, dim=2, keepdim=True)
weight = dist_recip / norm
interpolated_feats = pointnet2_utils.three_interpolate(
known_feats, idx, weight
)
else:
interpolated_feats = known_feats.expand(
(known_feats.size(0), known_feats.size(1), unknown.size(1))
)
if unknow_feats is not None:
new_features = torch.cat(
[interpolated_feats, unknow_feats], dim=1
) # (B, C2 + C1, n)
else:
new_features = interpolated_feats
new_features = new_features.unsqueeze(-1)
new_features = self.mlp(new_features)
return new_features.squeeze(-1)