-
Notifications
You must be signed in to change notification settings - Fork 5
/
Copy pathvn_dgcnn_util.py
121 lines (87 loc) · 3.78 KB
/
vn_dgcnn_util.py
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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
import torch
def knn(x, k):
inner = -2 * torch.matmul(x.transpose(2, 1), x)
xx = torch.sum(x ** 2, dim=1, keepdim=True)
pairwise_distance = -xx - inner - xx.transpose(2, 1)
idx = pairwise_distance.topk(k=k, dim=-1)[1] # (batch_size, num_points, k)
return idx
def get_graph_feature(x, k=20, idx=None, x_coord=None):
device = x.device
batch_size = x.size(0)
num_points = x.size(3)
x = x.view(batch_size, -1, num_points)
if idx is None:
if x_coord is None: # dynamic knn graph
idx = knn(x, k=k)
else: # fixed knn graph with input point coordinates
idx = knn(x_coord, k=k)
idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points
idx = idx + idx_base
idx = idx.view(-1)
_, num_dims, _ = x.size()
num_dims = num_dims // 3
x = x.transpose(2, 1).contiguous()
feature = x.view(batch_size * num_points, -1)[idx, :]
feature = feature.view(batch_size, num_points, k, num_dims, 3)
x = x.view(batch_size, num_points, 1, num_dims, 3).repeat(1, 1, k, 1, 1)
feature = torch.cat((feature - x, x), dim=3).permute(0, 3, 4, 1, 2).contiguous()
return feature
def get_graph_offset_old(x, k=20, idx=None, x_coord=None):
device = x.device
batch_size = x.size(0)
num_points = x.size(2)
x = x.view(batch_size, -1, num_points)
if idx is None:
if x_coord is None: # dynamic knn graph
idx = knn(x, k=k)
else: # fixed knn graph with input point coordinates
idx = knn(x_coord, k=k)
idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points
idx = idx + idx_base
idx = idx.view(-1)
_, num_dims, _ = x.size()
x = x.transpose(2, 1).contiguous() # (batch_size, num_points, num_dims) -> (batch_size*num_points, num_dims) # batch_size * num_points * k + range(0, batch_size*num_points)
feature = x.view(batch_size * num_points, -1)[idx, :]
feature = feature.view(batch_size, num_points, k, num_dims)
x = x.view(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1)
feature = torch.cat((feature - x, x), dim=3).permute(0, 3, 1, 2).contiguous()
return feature
def get_graph_offset(x, k=20, idx=None, x_coord=None):
device = x.device
batch_size = x.size(0)
num_points = x.size(2)
x = x.view(batch_size, -1, num_points)
if idx is None:
if x_coord is None: # dynamic knn graph
idx = knn(x, k=k)
else: # fixed knn graph with input point coordinates
idx = knn(x_coord, k=k)
idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points
idx = idx + idx_base
idx = idx.view(-1)
_, num_dims, _ = x.size()
num_dims = num_dims // 3
x = x.transpose(2, 1).contiguous()
feature = x.view(batch_size * num_points, -1)[idx, :]
feature = feature.view(batch_size, num_points, k, num_dims, 3)
x = x.view(batch_size, num_points, 1, num_dims, 3).repeat(1, 1, k, 1, 1)
return (feature - x)
def get_graph_feature_cross(x, k=20, idx=None):
device = x.device
batch_size = x.size(0)
num_points = x.size(3)
x = x.view(batch_size, -1, num_points)
if idx is None:
idx = knn(x, k=k)
idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points
idx = idx + idx_base
idx = idx.view(-1)
_, num_dims, _ = x.size()
num_dims = num_dims // 3
x = x.transpose(2, 1).contiguous()
feature = x.view(batch_size * num_points, -1)[idx, :]
feature = feature.view(batch_size, num_points, k, num_dims, 3)
x = x.view(batch_size, num_points, 1, num_dims, 3).repeat(1, 1, k, 1, 1)
cross = torch.cross(feature, x, dim=-1)
feature = torch.cat((feature - x, x, cross), dim=3).permute(0, 3, 4, 1, 2).contiguous()
return feature