def extract_label_points(pcd_file, label_file):
points = np.asarray(open3d.io.read_point_cloud(pcd_file).points)
colors = np.zeros(points.shape, dtype=int)
points_label = np.zeros(points.shape[0], dtype=np.uint8)
xyz = torch.from_numpy(np.expand_dims(points,
axis=0)).type(torch.FloatTensor)
xyz = xyz.to("cuda", non_blocking=True)
labels = load_label_file(label_file)
for label in labels:
if (int(label[0]) == 1): # vehicle
center = np.array(
[float(label[2]),
float(label[3]),
float(label[7]) * 0.5]) #float(label[4])])
size = np.array(
[float(label[5]),
float(label[6]),
float(label[7])])
heading = float(label[8])
center_shift = np.array([
size[0] * np.cos(heading) * 0.5,
size[0] * np.sin(heading) * 0.5, 0
])
proposal_centers = np.asarray(
[[center + center_shift, center - center_shift]])
idx = pointnet_utils.ball_query(
size[2] * 0.5, args.nsample, xyz,
torch.from_numpy(proposal_centers).type(torch.Tensor).to(
"cuda", non_blocking=True))
idx = np.squeeze(idx.cpu().numpy())
idx1 = np.unique(idx[0])
idx2 = np.unique(idx[1])
if idx1.size > idx2.size:
# colors[idx1] = [255, 0, 0]
points_label[idx1] = 1
else:
# colors[idx2] = [255, 0, 0]
points_label[idx2] = 1
return points, colors, points_label
def ball_query(p: torch.Tensor, q: torch.Tensor, r: float, k: int) -> torch.Tensor:
"""Point cloud ball (sphere) query.
Args:
p: Reference point cloud of shape [batch_size, dim, num_point].
q: Query Point cloud of shape [batch_size, dim, num_query].
r (float): Ball radius.
k (int): Maximum group size.
Returns:
Indices tensor of shape [batch_size, num_query, k].
"""
p_t = p.transpose(1, 2).contiguous()
q_t = q.transpose(1, 2).contiguous()
idx = _PU.ball_query(r, k, p_t, q_t)
return idx
def forward(self, pos1, pos2, feature1, feature2):
"""
Feature propagation from xyz2 (less points) to xyz1 (more points)
Inputs:
xyz1: (batch_size, 3, npoint1)
xyz2: (batch_size, 3, npoint2)
feat1: (batch_size, channel1, npoint1) features for xyz1 points (earlier layers, more points)
feat2: (batch_size, channel1, npoint2) features for xyz2 points
Output:
feat1_new: (batch_size, npoint2, mlp[-1] or mlp2[-1] or channel1+3)
TODO: Add support for skip links. Study how delta(XYZ) plays a role in feature updating.
"""
pos1_t = pos1.permute(0, 2, 1).contiguous()
pos2_t = pos2.permute(0, 2, 1).contiguous()
B, C, N = pos1.shape
if self.knn:
_, idx = pointutils.knn(self.nsample, pos1_t, pos2_t) # [B, N1, S]
else:
idx = pointutils.ball_query(self.radius, self.nsample, pos2_t,
pos1_t)
pos2_grouped = pointutils.grouping_operation(pos2, idx)
pos_diff = pos2_grouped - pos1.view(B, -1, N, 1) # [B, 3, N1, S]
feat2_grouped = pointutils.grouping_operation(feature2, idx)
feat_new = torch.cat([feat2_grouped, pos_diff],
dim=1) # [B, C1+3, N1, S]
for conv in self.mlp1_convs:
feat_new = conv(feat_new)
# max pooling
feat_new = feat_new.max(-1)[0] # [B, mlp1[-1], N1]
# concatenate feature in early layer
if feature1 is not None:
feat_new = torch.cat([feat_new, feature1],
dim=1) # [B, mlp1[-1]+feat1_channel, N1]
for conv in self.mlp2_convs:
feat_new = conv(feat_new)
return feat_new
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, pos1, pos2, feature1, feature2):
"""
Input:
xyz1: (batch_size, 3, npoint)
xyz2: (batch_size, 3, npoint)
feat1: (batch_size, channel, npoint)
feat2: (batch_size, channel, npoint)
Output:
xyz1: (batch_size, 3, npoint)
feat1_new: (batch_size, mlp[-1], npoint)
"""
pos1_t = pos1.permute(0, 2, 1).contiguous()
pos2_t = pos2.permute(0, 2, 1).contiguous()
B, N, C = pos1_t.shape
if self.knn:
_, idx = pointutils.knn(self.nsample, pos1_t, pos2_t) # [B, N, S]
else:
idx = pointutils.ball_query(self.radius, self.nsample, pos2_t,
pos1_t)
pos2_grouped = pointutils.grouping_operation(pos2, idx) # [B, 3, N, S]
pos_diff = pos2_grouped - pos1.view(B, -1, N, 1) # [B, 3, N, S]
feat2_grouped = pointutils.grouping_operation(feature2,
idx) # [B, C, N, S]
if self.corr_func == 'concat':
feat_diff = torch.cat([
feat2_grouped,
feature1.view(B, -1, N, 1).repeat(1, 1, 1, self.nsample)
],
dim=1) # [B, 2*C, N, S]
feat1_new = torch.cat([pos_diff, feat_diff], dim=1) # [B, 2*C+3, N, S]
for i, conv in enumerate(self.mlp_convs):
bn = self.mlp_bns[i]
feat1_new = F.relu(bn(conv(feat1_new)))
feat1_new = torch.max(feat1_new, -1)[0] # [B, mlp[-1], npoint]
return pos1, feat1_new