def forward(self, xyz, feats, new_xyz=None):
'''
:param pcd: B, C_in, N
:return:
new_pcd: B, C_out, np
'''
if new_xyz is None:
assert self.npoint is not None
xyz_flipped = xyz.transpose(1, 2).contiguous() # B,3,npoint
idx = pointnet2_utils.furthest_point_sample(
xyz, self.npoint) # B,npoint
new_xyz_flipped = pointnet2_utils.gather_operation(
xyz_flipped, idx) # B,3,npoint
new_xyz = new_xyz_flipped.transpose(1,
2).contiguous() # B,npoint,3
idx = pointnet2_utils.ball_query(self.radius, self.nsample, xyz,
new_xyz)
gped_feats = pointnet2_utils.grouping_operation(feats,
idx) # B,C,np,ns
gped_feats = F.max_pool2d(gped_feats,
kernel_size=[1, self.nsample]) # B,C,np,1
gped_feats = gped_feats.squeeze(-1) # B,C,np
return self.conv(gped_feats)
def forward(self, xyz, points, global_features, select_idx, use_sample):
"""
Input:
xyz: input points position data, [B, N, C]
points: input points data, [B, D, N]
Return:
new_xyz: sampled points position data, [B, C, S]
new_points_concat: sample points feature data, [B, D', S]
"""
B, N, C = xyz.shape
if use_sample == 1:
num_points = select_idx.shape[1]
new_xyz = torch.from_numpy(select_idx).transpose(1,
2).float().cuda()
else:
new_xyz = pointnet2_utils.gather_operation(
xyz.permute(0, 2, 1).contiguous(), select_idx.int())
num_points = select_idx.shape[-1]
new_radius_list = []
for i, radius in enumerate(self.radius_list):
K = self.nsample_list[i]
new_radius = torch.zeros(
[select_idx.shape[0], select_idx.shape[1]])
query_and_group = pointnet2_utils.QueryAndGroup(radius,
K,
use_xyz=True)
grouped_points = query_and_group(
xyz.contiguous(),
new_xyz.permute(0, 2, 1).contiguous(),
points.permute(0, 2, 1).contiguous())
for j in range(len(self.conv_blocks[i])):
conv = self.conv_blocks[i][j]
bn = self.bn_blocks[i][j]
grouped_points = F.leaky_relu(bn(conv(grouped_points)),
negative_slope=0.2)
new_radius = torch.max(grouped_points, 3)[0] # [B, D', S]
new_radius_list.append(new_radius)
new_xyz = new_xyz.permute(0, 2, 1)
new_radius_concat = torch.cat(new_radius_list, dim=1)
new_radius_concat = F.leaky_relu(self.conv1(new_radius_concat),
negative_slope=0.2)
num_point = new_radius_concat.shape[2]
new_radius_concat = F.leaky_relu(self.conv2(new_radius_concat),
negative_slope=0.2)
new_radius_concat = self.conv3(new_radius_concat).transpose(1, 2).view(
-1, num_point, 6)
return new_radius_concat
def get_emd_loss(self, pred, gt, pcd_radius):
idx, _ = auction_match(pred, gt)
matched_out = pn2_utils.gather_operation(
gt.transpose(1, 2).contiguous(), idx)
matched_out = matched_out.transpose(1, 2).contiguous()
dist2 = (pred - matched_out)**2
dist2 = dist2.view(dist2.shape[0], -1) # <-- ???
dist2 = torch.mean(dist2, dim=1, keepdims=True) # B,
dist2 /= pcd_radius
return torch.mean(dist2)
def sample_keypoints(self, points):
"""
fps expects points shape (B, N, 3)
fps returns indices shape (B, K)
gather expects features shape (B, C, N)
"""
points = points[..., :3].contiguous()
indices = furthest_point_sample(points, self.cfg.NUM_KEYPOINTS)
keypoints = gather_operation(points.transpose(1, 2).contiguous(), indices)
keypoints = keypoints.transpose(1, 2).contiguous()
return keypoints
def index_points_gather(points, fps_idx):
"""
Input:
points: input points data, [B, N, C]
idx: sample index data, [B, S]
Return:
new_points:, indexed points data, [B, S, C]
"""
points_flipped = points.permute(0, 2, 1).contiguous()
new_points = pointnet2_utils.gather_operation(points_flipped, fps_idx)
return new_points.permute(0, 2, 1).contiguous()
def get_emd_loss(self, pred, gt, radius=1.0):
'''
pred and gt is B N 3
'''
idx, _ = auction_match(pred.contiguous(), gt.contiguous())
#gather operation has to be B 3 N
#print(gt.transpose(1,2).shape)
matched_out = pn2_utils.gather_operation(
gt.transpose(1, 2).contiguous(), idx)
matched_out = matched_out.transpose(1, 2).contiguous()
dist2 = (pred - matched_out)**2
dist2 = dist2.view(dist2.shape[0], -1) # <-- ???
dist2 = torch.mean(dist2, dim=1, keepdims=True) # B,
dist2 /= radius
return torch.mean(dist2)
def forward(self,
xyz: torch.Tensor,
features: torch.Tensor = None,
inds: torch.Tensor = None) -> (torch.Tensor, torch.Tensor):
r"""
Parameters
----------
xyz : torch.Tensor
(B, N, 3) tensor of the xyz coordinates of the features
features : torch.Tensor
(B, C, C) tensor of the descriptors of the the features
inds : torch.Tensor
(B, npoint) tensor that stores index to the xyz points (values in 0-N-1)
Returns
-------
new_xyz : torch.Tensor
(B, npoint, 3) tensor of the new features' xyz
new_features : torch.Tensor
(B, \sum_k(mlps[k][-1]), npoint) tensor of the new_features descriptors
inds: torch.Tensor
(B, npoint) tensor of the inds
"""
new_features_list = []
xyz_flipped = xyz.transpose(1, 2).contiguous()
if inds is None:
inds = pointnet2_utils.furthest_point_sample(xyz, self.npoint)
new_xyz = pointnet2_utils.gather_operation(
xyz_flipped, inds).transpose(
1, 2).contiguous() if self.npoint is not None else None
for i in range(len(self.groupers)):
new_features = self.groupers[i](
xyz, new_xyz, features) # (B, C, npoint, nsample)
new_features = self.mlps[i](
new_features) # (B, mlp[-1], npoint, nsample)
new_features = F.max_pool2d(new_features,
kernel_size=[
1, new_features.size(3)
]) # (B, mlp[-1], npoint, 1)
new_features = new_features.squeeze(-1) # (B, mlp[-1], npoint)
new_features_list.append(new_features)
return new_xyz, torch.cat(new_features_list, dim=1), inds
def forward(self, xyz, features=None):
# type: (_PointnetSAModuleBase, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]
r"""
Parameters
----------
xyz : torch.Tensor
(B, N, 3) tensor of the xyz coordinates of the features
features : torch.Tensor
(B, N, C) tensor of the descriptors of the the features
Returns
-------
new_xyz : torch.Tensor
(B, npoint, 3) tensor of the new features' xyz
new_features : torch.Tensor
(B, \sum_k(mlps[k][-1]), npoint) tensor of the new_features descriptors
"""
new_features_list = []
xyz_flipped = xyz.transpose(1, 2).contiguous()
new_xyz = (pointnet2_utils.gather_operation(
xyz_flipped, pointnet2_utils.furthest_point_sample(
xyz, self.npoint)).transpose(1, 2).contiguous()
if self.npoint is not None else None)
for i in range(len(self.groupers)):
new_features = self.groupers[i](
xyz, new_xyz, features) # (B, C, npoint, nsample)
new_features = self.mlps[i](
new_features) # (B, mlp[-1], npoint, nsample)
new_features = F.max_pool2d(new_features,
kernel_size=[
1, new_features.size(3)
]) # (B, mlp[-1], npoint, 1)
new_features = new_features.squeeze(-1) # (B, mlp[-1], npoint)
new_features_list.append(new_features)
return new_xyz, torch.cat(new_features_list, dim=1)
def get_uniform_loss(self,
pcd,
percentage=[0.004, 0.006, 0.008, 0.010, 0.012],
radius=1.0):
B, N, C = pcd.shape[0], pcd.shape[1], pcd.shape[2]
npoint = int(N * 0.05)
loss = 0
further_point_idx = pn2_utils.furthest_point_sample(
pcd.permute(0, 2, 1).contiguous(), npoint)
new_xyz = pn2_utils.gather_operation(
pcd.permute(0, 2, 1).contiguous(), further_point_idx) # B,C,N
for p in percentage:
nsample = int(N * p)
r = math.sqrt(p * radius)
disk_area = math.pi * (radius**2) / N
idx = pn2_utils.ball_query(r, nsample, pcd.contiguous(),
new_xyz.permute(
0, 2, 1).contiguous()) #b N nsample
expect_len = math.sqrt(disk_area)
grouped_pcd = pn2_utils.grouping_operation(
pcd.permute(0, 2, 1).contiguous(), idx) #B C N nsample
grouped_pcd = grouped_pcd.permute(0, 2, 3, 1) #B N nsample C
grouped_pcd = torch.cat(torch.unbind(grouped_pcd, dim=1),
dim=0) #B*N nsample C
dist, _ = self.knn_uniform(grouped_pcd, grouped_pcd)
#print(dist.shape)
uniform_dist = dist[:, :, 1:] #B*N nsample 1
uniform_dist = torch.abs(uniform_dist + 1e-8)
uniform_dist = torch.mean(uniform_dist, dim=1)
uniform_dist = (uniform_dist - expect_len)**2 / (expect_len + 1e-8)
mean_loss = torch.mean(uniform_dist)
mean_loss = mean_loss * math.pow(p * 100, 2)
loss += mean_loss
return loss / len(percentage)
def forward(self, xyz: torch.Tensor, features: torch.Tensor = None, new_xyz=None) -> (torch.Tensor, torch.Tensor):
"""
:param xyz: (B, N, 3) tensor of the xyz coordinates of the features
:param features: (B, N, C) tensor of the descriptors of the the features
:param new_xyz:
:return:
new_xyz: (B, npoint, 3) tensor of the new features' xyz
new_features: (B, npoint, \sum_k(mlps[k][-1])) tensor of the new_features descriptors
"""
new_features_list = []
xyz_flipped = xyz.transpose(1, 2).contiguous()
if new_xyz is None:
new_xyz = pointnet2_utils.gather_operation(
xyz_flipped,
pointnet2_utils.furthest_point_sample(xyz, self.npoint)
).transpose(1, 2).contiguous() if self.npoint is not None else None
for i in range(len(self.groupers)):
new_features = self.groupers[i](xyz, new_xyz, features) # (B, C, npoint, nsample)
new_features = self.mlps[i](new_features) # (B, mlp[-1], npoint, nsample)
if self.pool_method == 'max_pool':
new_features = F.max_pool2d(
new_features, kernel_size=[1, new_features.size(3)]
) # (B, mlp[-1], npoint, 1)
elif self.pool_method == 'avg_pool':
new_features = F.avg_pool2d(
new_features, kernel_size=[1, new_features.size(3)]
) # (B, mlp[-1], npoint, 1)
else:
raise NotImplementedError
new_features = new_features.squeeze(-1) # (B, mlp[-1], npoint)
new_features_list.append(new_features)
return new_xyz, torch.cat(new_features_list, dim=1)
def forward(self,
xyz: torch.Tensor,
features: torch.Tensor = None,
inds: torch.Tensor = None) -> (torch.Tensor, torch.Tensor):
r"""
Parameters
----------
xyz : torch.Tensor
(B, N, 3) tensor of the xyz coordinates of the features
features : torch.Tensor
(B, C, N) tensor of the descriptors of the the features
inds : torch.Tensor
(B, npoint) tensor that stores index to the xyz points (values in 0-N-1)
Returns
-------
new_xyz : torch.Tensor
(B, npoint, 3) tensor of the new features' xyz
new_features : torch.Tensor
(B, \sum_k(mlps[k][-1]), npoint) tensor of the new_features descriptors
inds: torch.Tensor
(B, npoint) tensor of the inds
"""
xyz_flipped = xyz.transpose(1, 2).contiguous()
if inds is None:
inds = pointnet2_utils.furthest_point_sample(xyz, self.npoint)
else:
assert (inds.shape[1] == self.npoint)
new_xyz = pointnet2_utils.gather_operation(
xyz_flipped, inds).transpose(
1, 2).contiguous() if self.npoint is not None else None
if not self.ret_unique_cnt:
grouped_features, grouped_xyz = self.grouper(
xyz, new_xyz, features) # (B, C, npoint, nsample)
else:
grouped_features, grouped_xyz, unique_cnt = self.grouper(
xyz, new_xyz, features
) # (B, C, npoint, nsample), (B,3,npoint,nsample), (B,npoint)
new_features = self.mlp_module(
grouped_features) # (B, mlp[-1], npoint, nsample)
if self.pooling == 'max':
new_features = F.max_pool2d(new_features,
kernel_size=[
1, new_features.size(3)
]) # (B, mlp[-1], npoint, 1)
elif self.pooling == 'avg':
new_features = F.avg_pool2d(new_features,
kernel_size=[
1, new_features.size(3)
]) # (B, mlp[-1], npoint, 1)
elif self.pooling == 'rbf':
# Use radial basis function kernel for weighted sum of features (normalized by nsample and sigma)
# Ref: https://en.wikipedia.org/wiki/Radial_basis_function_kernel
rbf = torch.exp(-1 * grouped_xyz.pow(2).sum(1, keepdim=False) /
(self.sigma**2) / 2) # (B, npoint, nsample)
new_features = torch.sum(
new_features * rbf.unsqueeze(1), -1, keepdim=True) / float(
self.nsample) # (B, mlp[-1], npoint, 1)
new_features = new_features.squeeze(-1) # (B, mlp[-1], npoint)
if not self.ret_unique_cnt:
return new_xyz, new_features, inds
else:
return new_xyz, new_features, inds, unique_cnt
