Spaces:
Sleeping
Sleeping
| import torch | |
| from torch import nn | |
| from torch_cluster import knn_graph | |
| class DenseDilated(nn.Module): | |
| """ | |
| Find dilated neighbor from neighbor list | |
| edge_index: (2, batch_size, num_points, k) | |
| """ | |
| def __init__(self, k=9, dilation=1, stochastic=False, epsilon=0.0): | |
| super(DenseDilated, self).__init__() | |
| self.dilation = dilation | |
| self.stochastic = stochastic | |
| self.epsilon = epsilon | |
| self.k = k | |
| def forward(self, edge_index): | |
| if self.stochastic: | |
| if torch.rand(1) < self.epsilon and self.training: | |
| num = self.k * self.dilation | |
| randnum = torch.randperm(num)[:self.k] | |
| edge_index = edge_index[:, :, :, randnum] | |
| else: | |
| edge_index = edge_index[:, :, :, ::self.dilation] | |
| else: | |
| edge_index = edge_index[:, :, :, ::self.dilation] | |
| return edge_index | |
| def pairwise_distance(x): | |
| """ | |
| Compute pairwise distance of a point cloud. | |
| Args: | |
| x: tensor (batch_size, num_points, num_dims) | |
| Returns: | |
| pairwise distance: (batch_size, num_points, num_points) | |
| """ | |
| x_inner = -2*torch.matmul(x, x.transpose(2, 1)) | |
| x_square = torch.sum(torch.mul(x, x), dim=-1, keepdim=True) | |
| return x_square + x_inner + x_square.transpose(2, 1) | |
| def dense_knn_matrix(x, k=16): | |
| """Get KNN based on the pairwise distance. | |
| Args: | |
| x: (batch_size, num_dims, num_points, 1) | |
| k: int | |
| Returns: | |
| nearest neighbors: (batch_size, num_points ,k) (batch_size, num_points, k) | |
| """ | |
| with torch.no_grad(): | |
| x = x.transpose(2, 1).squeeze(-1) | |
| batch_size, n_points, n_dims = x.shape | |
| _, nn_idx = torch.topk(-pairwise_distance(x.detach()), k=k) | |
| center_idx = torch.arange(0, n_points, device=x.device).repeat(batch_size, k, 1).transpose(2, 1) | |
| return torch.stack((nn_idx, center_idx), dim=0) | |
| class DenseDilatedKnnGraph(nn.Module): | |
| """ | |
| Find the neighbors' indices based on dilated knn | |
| """ | |
| def __init__(self, k=9, dilation=1, stochastic=False, epsilon=0.0): | |
| super(DenseDilatedKnnGraph, self).__init__() | |
| self.dilation = dilation | |
| self.stochastic = stochastic | |
| self.epsilon = epsilon | |
| self.k = k | |
| self._dilated = DenseDilated(k, dilation, stochastic, epsilon) | |
| self.knn = dense_knn_matrix | |
| def forward(self, x): | |
| edge_index = self.knn(x, self.k * self.dilation) | |
| return self._dilated(edge_index) | |
| class DilatedKnnGraph(nn.Module): | |
| """ | |
| Find the neighbors' indices based on dilated knn | |
| """ | |
| def __init__(self, k=9, dilation=1, stochastic=False, epsilon=0.0): | |
| super(DilatedKnnGraph, self).__init__() | |
| self.dilation = dilation | |
| self.stochastic = stochastic | |
| self.epsilon = epsilon | |
| self.k = k | |
| self._dilated = DenseDilated(k, dilation, stochastic, epsilon) | |
| self.knn = knn_graph | |
| def forward(self, x): | |
| x = x.squeeze(-1) | |
| B, C, N = x.shape | |
| edge_index = [] | |
| for i in range(B): | |
| edgeindex = self.knn(x[i].contiguous().transpose(1, 0).contiguous(), self.k * self.dilation) | |
| edgeindex = edgeindex.view(2, N, self.k * self.dilation) | |
| edge_index.append(edgeindex) | |
| edge_index = torch.stack(edge_index, dim=1) | |
| return self._dilated(edge_index) | |