def _infer_flow(self, flow_f_1, flow_f_2):
# Normalize the features
if self.normalize_feature:
flow_f_1 = ME.SparseTensor(
flow_f_1.F / torch.norm(flow_f_1.F, p=2, dim=1, keepdim=True),
coordinate_map_key=flow_f_1.coordinate_map_key,
coordinate_manager=flow_f_1.coordinate_manager)
flow_f_2 = ME.SparseTensor(
flow_f_2.F / torch.norm(flow_f_2.F, p=2, dim=1, keepdim=True),
coordinate_map_key=flow_f_2.coordinate_map_key,
coordinate_manager=flow_f_2.coordinate_manager)
# Extract the coarse flow based on the feature correspondences
coarse_flow = []
# Iterate over the examples in the batch
for b_idx in range(len(flow_f_1.decomposed_coordinates)):
feat_s = flow_f_1.F[flow_f_1.C[:, 0] == b_idx]
feat_t = flow_f_2.F[flow_f_2.C[:, 0] == b_idx]
coor_s = flow_f_1.C[flow_f_1.C[:, 0] == b_idx, 1:].to(
self.device) * self.voxel_size
coor_t = flow_f_2.C[flow_f_2.C[:, 0] == b_idx, 1:].to(
self.device) * self.voxel_size
# Squared l2 distance between points points of both point clouds
coor_s, coor_t = coor_s.unsqueeze(0), coor_t.unsqueeze(0)
feat_s, feat_t = feat_s.unsqueeze(0), feat_t.unsqueeze(0)
# Force transport to be zero for points further than 10 m apart
support = (pairwise_distance(coor_s, coor_t, normalized=False) < 10
**2).float()
# Transport cost matrix
C = pairwise_distance(feat_s, feat_t)
K = torch.exp(
-C / (torch.exp(self.epsilon) + self.tau_offset)) * support
row_sum = K.sum(-1, keepdim=True)
# Estimate flow
corr_flow = (K @ coor_t) / (row_sum + 1e-8) - coor_s
coarse_flow.append(corr_flow.squeeze(0))
coarse_flow = torch.cat(coarse_flow, dim=0)
st_cf = ME.SparseTensor(features=coarse_flow,
coordinate_manager=flow_f_1.coordinate_manager,
coordinate_map_key=flow_f_1.coordinate_map_key)
self.inferred_values['coarse_flow'] = st_cf.F
# Refine the flow with the second network
refined_flow = self.flow_refiner(st_cf)
self.inferred_values['refined_flow'] = refined_flow.F
def forward(self, x_f, y_f, y_c):
""" Computes the correspondences in the feature space based on the selected parameters.
Args:
x_f (torch.tensor): infered features of points x [b,n,c]
y_f (torch.tensor): infered features of points y [b,m,c]
y_c (torch.tensor): coordinates of point y [b,m,3]
Returns:
x_corr (torch.tensor): coordinates of the feature based correspondences of points x [b,n,3]
"""
dist = pairwise_distance(x_f, y_f).detach()
if self.corr_type == 'soft':
y_soft = torch.softmax(-dist / (self.get_temp()), dim=2)
if self.st:
# Straight through.
index = y_soft.max(dim=2, keepdim=True)[1]
y_hard = torch.zeros_like(y_soft).scatter_(dim=2,
index=index,
value=1.0)
ret = y_hard - y_soft.detach() + y_soft
else:
ret = y_soft
elif self.corr_type == 'soft_gumbel':
if self.st:
# Straight through.
ret = F.gumbel_softmax(-dist, tau=self.get_temp(), hard=True)
else:
ret = F.gumbel_softmax(-dist, tau=self.get_temp(), hard=False)
else:
index = dist.min(dim=2, keepdim=True)[1]
ret = torch.zeros_like(dist).scatter_(dim=2,
index=index,
value=1.0)
# Compute corresponding coordinates
x_corr = torch.matmul(ret, y_c)
return x_corr
def _infer_ego_motion(self, flow_f_1, flow_f_2, sem_label_s, sem_label_t):
ego_motion_R = []
ego_motion_t = []
ego_motion_perm = []
run_b_len_s = 0
run_b_len_t = 0
# Normalize the features
if self.normalize_feature:
flow_f_1 = ME.SparseTensor(
flow_f_1.F / torch.norm(flow_f_1.F, p=2, dim=1, keepdim=True),
coordinate_map_key=flow_f_1.coordinate_map_key,
coordinate_manager=flow_f_1.coordinate_manager)
flow_f_2 = ME.SparseTensor(
flow_f_2.F / torch.norm(flow_f_2.F, p=2, dim=1, keepdim=True),
coordinate_map_key=flow_f_2.coordinate_map_key,
coordinate_manager=flow_f_2.coordinate_manager)
for b_idx in range(len(flow_f_1.decomposed_coordinates)):
feat_s = flow_f_1.F[flow_f_1.C[:, 0] == b_idx]
feat_t = flow_f_2.F[flow_f_2.C[:, 0] == b_idx]
coor_s = flow_f_1.C[flow_f_1.C[:, 0] == b_idx, 1:].to(
self.device) * self.voxel_size
coor_t = flow_f_2.C[flow_f_2.C[:, 0] == b_idx, 1:].to(
self.device) * self.voxel_size
# Get the number of points in the current b_idx
b_len_s = feat_s.shape[0]
b_len_t = feat_t.shape[0]
# Extract the semantic labels for the current b_idx (0 are the background points)
mask_s = (sem_label_s[run_b_len_s:(run_b_len_s + b_len_s)] == 0)
mask_t = (sem_label_t[run_b_len_t:(run_b_len_t + b_len_t)] == 0)
# Update the running number of points
run_b_len_s += b_len_s
run_b_len_t += b_len_t
# Squared l2 distance between points points of both point clouds
coor_s, coor_t = coor_s[mask_s, :].unsqueeze(0), coor_t[
mask_t, :].unsqueeze(0)
feat_s, feat_t = feat_s[mask_s, :].unsqueeze(0), feat_t[
mask_t, :].unsqueeze(0)
# Sample the points randomly (to keep the computation memory tracktable)
idx_ego_s = torch.randperm(coor_s.shape[1])[:self.ego_n_points]
idx_ego_t = torch.randperm(coor_t.shape[1])[:self.ego_n_points]
coor_s_ego = coor_s[:, idx_ego_s, :]
coor_t_ego = coor_t[:, idx_ego_t, :]
feat_s_ego = feat_s[:, idx_ego_s, :]
feat_t_ego = feat_t[:, idx_ego_t, :]
# Force transport to be zero for points further than 10 m apart
support_ego = (pairwise_distance(
coor_s_ego, coor_t_ego, normalized=False) < 10**2).float()
# Cost matrix in the feature space
feat_dist = pairwise_distance(feat_s_ego, feat_t_ego)
R_est, t_est, perm_matrix = self.ego_motion_decoder(
feat_dist, support_ego, coor_s_ego, coor_t_ego)
ego_motion_R.append(R_est)
ego_motion_t.append(t_est)
ego_motion_perm.append(perm_matrix)
# Save ego motion results
self.inferred_values['R_est'] = torch.cat(ego_motion_R, dim=0)
self.inferred_values['t_est'] = torch.cat(ego_motion_t, dim=0)
self.inferred_values['permutation'] = ego_motion_perm