def forward(self, center_coords, coords, features):
knn_indexes, _ = k_nearest_neighbors(center_coords, coords,
self.k * self.dilation)
knn_indexes = knn_indexes[:, :, ::self.dilation]
knn_coords = index2points(coords, knn_indexes)
knn_local_coords = localize(center_coords, knn_coords) # (B,C,N,k)
feature_d = self.mlp_d(knn_local_coords)
if self.point_feature_size == 0:
feature_a = feature_d
else:
knn_features = index2points(features, knn_indexes)
feature_a = torch.cat([feature_d, knn_features],
dim=1) # [B, C+add_C, N, k]
if self.use_x_transformation:
trans = self.x_trans(feature_a)
fx = self.transform(feature_a, trans)
else:
fx = feature_a
fx = self.conv1(fx)
fx = torch.squeeze(fx, dim=3).contiguous()
if self.with_global:
fts_global = self.linear1(center_coords)
res = torch.can([fts_global, fx], dim=1)
else:
res = fx
return res
def forward(self, coords, knn_indices):
"""
Parameters
----------
coords: torch.tensor (B, C, N)
"""
# Get spaces between points.
knn_coords = index2points(coords, knn_indices)
coords_space = localize(coords, knn_coords) * -1
# Use theta.
outputs = self.mlp_theta(coords_space)
return outputs
def forward(self, x, coords):
# Get p and x of fps indices
# fps_indices = furthest_point_sampling(coords, self.num_samples)
# fps_coords = index2points(coords, fps_indices) # p
# Get knn indices
knn_indices, _ = k_nearest_neighbors(coords, coords, self.k)
# Get local coords
knn_coords = index2points(coords, knn_indices)
local_coods = localize(coords, knn_coords) * -1
# get knn_features and use MLP
knn_x = index2points(x, knn_indices)
knn_mlp_x = self.mlp(knn_x)
# Use local max pooling. y = pooled features
y, _ = torch.max(knn_mlp_x, dim=-1)
return y, knn_mlp_x, local_coods
def forward(self, pooled_features, coords, knn_features, local_coords):
"""
Parameters
----------
pooled_features: torch.tensor (B, C, N)
coords: torch.tensor (B, 3, N)
knn_features: torch.tensor (B, C, N, k)
local_coords: torch.tensor (B, 3, N, k)
"""
# Get knn indexes.
# knn_indices, _ = k_nearest_neighbors(coords, coords, self.k) # get (B, N, k)
# Get delta.
outputs_delta = self.pe_delta(coords, local_coords) # get (B, C, N, k)
# Get pointwise feature.
# outputs_phi, outputs_psi, outputs_alpha = torch.chunk(
# self.input_linear(features), chunks=3, dim=1) # to (B, C, N) x 3
# Get weights.
outputs_phi = self.phi(pooled_features)
outputs_psi = self.psi(knn_features)
# outputs_psi = index2points(outputs_psi, knn_indices) # to (B, C, N, k)
inputs_gamma = localize(
outputs_phi, outputs_psi) * -1 + outputs_delta # get (B, C, N, k)
outputs_gamma = self.mlp_gamma(inputs_gamma)
outputs_rho = self.normalization_rho(outputs_gamma)
# \alpha(x_j) + \delta
outputs_alpha = self.alpha(knn_features)
# outputs_alpha = index2points(outputs_alpha, knn_indices) # to (B, C, N, k)
outputs_alpha_delta = outputs_alpha + outputs_delta
# outputs_alpha_delta = outputs_alpha
# compute value with hadamard product and aggregation
outputs_hp = outputs_rho * outputs_alpha_delta
outputs_aggregation = torch.sum(outputs_hp, dim=-1) # get (B, C, N)
return outputs_aggregation