| 19 | return idx |
| 20 | |
| 21 | def get_graph_feature(x, k=20, idx=None, x_coord=None): |
| 22 | batch_size = x.size(0) |
| 23 | num_points = x.size(3) |
| 24 | x = x.view(batch_size, -1, num_points) |
| 25 | if idx is None: |
| 26 | if x_coord is None: # dynamic knn graph |
| 27 | idx = knn(x, k=k) |
| 28 | else: # fixed knn graph with input point coordinates |
| 29 | idx = knn(x_coord, k=k) |
| 30 | device = torch.device('cuda') |
| 31 | |
| 32 | idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points |
| 33 | |
| 34 | idx = idx + idx_base |
| 35 | |
| 36 | idx = idx.view(-1) |
| 37 | |
| 38 | _, num_dims, _ = x.size() |
| 39 | num_dims = num_dims // 3 |
| 40 | |
| 41 | x = x.transpose(2, 1).contiguous() |
| 42 | feature = x.view(batch_size * num_points, -1)[idx, :] |
| 43 | feature = feature.view(batch_size, num_points, k, num_dims, 3) |
| 44 | x = x.view(batch_size, num_points, 1, num_dims, 3).repeat(1, 1, k, 1, 1) |
| 45 | |
| 46 | feature = torch.cat((feature - x, x), dim=3).permute(0, 3, 4, 1, 2).contiguous() |
| 47 | |
| 48 | return feature |
| 49 | |
| 50 | class VN_DGCNN(nn.Module): |
| 51 | |