| 302 | |
| 303 | |
| 304 | class PointNetFeaturePropagation(nn.Module): |
| 305 | def __init__(self, in_channel, mlp): |
| 306 | super(PointNetFeaturePropagation, self).__init__() |
| 307 | self.mlp_convs = nn.ModuleList() |
| 308 | self.mlp_bns = nn.ModuleList() |
| 309 | last_channel = in_channel |
| 310 | for out_channel in mlp: |
| 311 | self.mlp_convs.append(nn.Conv1d(last_channel, out_channel, 1)) |
| 312 | self.mlp_bns.append(nn.BatchNorm1d(out_channel)) |
| 313 | last_channel = out_channel |
| 314 | |
| 315 | def forward(self, xyz1, xyz2, points1, points2): |
| 316 | """ |
| 317 | Input: |
| 318 | xyz1: input points position data, [B, C, N] |
| 319 | xyz2: sampled input points position data, [B, C, S] |
| 320 | points1: input points data, [B, D, N] |
| 321 | points2: input points data, [B, D, S] |
| 322 | Return: |
| 323 | new_points: upsampled points data, [B, D', N] |
| 324 | """ |
| 325 | xyz1 = xyz1.permute(0, 2, 1) |
| 326 | xyz2 = xyz2.permute(0, 2, 1) |
| 327 | |
| 328 | points2 = points2.permute(0, 2, 1) |
| 329 | B, N, C = xyz1.shape |
| 330 | _, S, _ = xyz2.shape |
| 331 | |
| 332 | if S == 1: |
| 333 | interpolated_points = points2.repeat(1, N, 1) |
| 334 | else: |
| 335 | dists = square_distance(xyz1, xyz2) |
| 336 | dists, idx = dists.sort(dim=-1) |
| 337 | dists, idx = dists[:, :, :3], idx[:, :, :3] # [B, N, 3] |
| 338 | |
| 339 | dist_recip = 1.0 / (dists + 1e-8) |
| 340 | norm = torch.sum(dist_recip, dim=2, keepdim=True) |
| 341 | weight = dist_recip / norm |
| 342 | interpolated_points = torch.sum( |
| 343 | index_points(points2, idx) * weight.view(B, N, 3, 1), dim=2 |
| 344 | ) |
| 345 | |
| 346 | if points1 is not None: |
| 347 | points1 = points1.permute(0, 2, 1) |
| 348 | new_points = torch.cat([points1, interpolated_points], dim=-1) |
| 349 | else: |
| 350 | new_points = interpolated_points |
| 351 | |
| 352 | new_points = new_points.permute(0, 2, 1) |
| 353 | for i, conv in enumerate(self.mlp_convs): |
| 354 | bn = self.mlp_bns[i] |
| 355 | new_points = F.relu(bn(conv(new_points))) |
| 356 | return new_points |
nothing calls this directly
no outgoing calls
no test coverage detected