Spaces:
Sleeping
Sleeping
File size: 3,393 Bytes
799e642 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
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)
|