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