def GARPointList(t1_list, t2_list, R1_list, R2_list, V, c, alpha): err = TensorListList() for t1, t2, R1, R2, Vb in zip(t1_list, t2_list, R1_list, R2_list, V): M = len(t1) err1 = TensorList() for ind1 in range(M - 1): for ind2 in range(ind1 + 1, M): Rtest = R1[ind1].permute(1, 0) @ R1[ind2] # estimated relative translation from ind2 to ind1 test = R1[ind1].permute(1, 0) @ t1[ind2] - R1[ind1].permute( 1, 0) @ t1[ind1] Rgt = R2[ind1].permute(1, 0) @ R2[ind2] # ground truth translation from ind2 to ind1 tgt = R2[ind1].permute(1, 0) @ t2[ind2] - R2[ind1].permute( 1, 0) @ t2[ind1] diff = Rtest @ Vb[ind2] + test - (Rgt @ Vb[ind2] + tgt) err_i = (diff * diff).sum(dim=-2).sqrt() err1.append(lossfun(err_i, alpha=alpha, scale=c).mean()) err.append(err1) return err
def empirical_estimate(points, num_neighbors): ps, batch = points.permute(1, 0).cat_tensors() N = ps.shape[0] val = knn(ps.contiguous(), ps.contiguous(), batch_x=batch, batch_y=batch, k=num_neighbors) A = ps[val[1, :]].reshape(N, num_neighbors, 3) A = A - A.mean(dim=1, keepdim=True) Asqr = A.permute(0, 2, 1).bmm(A) sigma, _ = Asqr.cpu().symeig() w = (sigma[:, 2] * sigma[:, 1]).sqrt() val = val[1, :].reshape(N, num_neighbors) w, _ = torch.median(w[val].to(ps.device), dim=1, keepdim=True) weights = TensorListList() bi = 0 for point_list in points: ww = TensorList() for p in point_list: ww.append(w[batch == bi]) bi = bi + 1 weights.append(ww) return weights
def global_to_relative(R, t, ref_index): R_rel = TensorListList() t_rel = TensorListList() for Ri, ti in zip(R, t): R_ref = Ri[ref_index] t_ref = ti[ref_index] R_reli = R_ref.permute(1, 0) @ Ri t_reli = R_ref.permute(1, 0) @ ti - R_ref.permute(1, 0) @ t_ref R_rel.append(R_reli) t_rel.append(t_reli) return R_rel, t_rel
def L2sqrList(R1_list, R2_list): err = TensorListList() for R1, R2 in zip(R1_list, R2_list): M = len(R1) err1 = TensorList() for ind1 in range(M - 1): for ind2 in range(ind1 + 1, M): Rdiff = R1[ind1].permute(1, 0) @ R1[ind2] - R2[ind1].permute( 1, 0) @ R2[ind2] err1.append( (Rdiff * Rdiff).sum(dim=-1).sum(dim=-1).unsqueeze(0)) err.append(err1) return err
def register_point_sets(self, x): point_clouds = x["coordinates"].clone() if isinstance(point_clouds, TensorListList): Rs = TensorListList() ts = TensorListList() else: # B, P, N, M = point_clouds.shape # assert P == 2 Rs = [] ts = [] target_R = torch.eye(3,3) target_t = torch.zeros(3,1) for pcds in point_clouds: source = o3d.geometry.PointCloud() target = o3d.geometry.PointCloud() source.points = o3d.utility.Vector3dVector(pcds[0].cpu().numpy().T) target.points = o3d.utility.Vector3dVector(pcds[1].cpu().numpy().T) source_down, target_down, source_fpfh, target_fpfh = \ self.prepare_dataset(source, target, self.params.voxel_size) result_fast = self.execute_global_registration(source_down, target_down, source_fpfh, target_fpfh, self.params.voxel_size) T = result_fast.transformation R = T[0:3, 0:3] t = T[0:3, 3:] if self.params.get("refine", True): radius_normal = self.params.voxel_size * 2 if self.params.metric == "p2pl": source.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=radius_normal, max_nn=30)) target.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=radius_normal, max_nn=30)) #o3d.geometry.estimate_normals(source) #o3d.geometry.estimate_normals(target) distance_threshold = self.params.voxel_size * 0.5 reg_p2p = o3d.registration.registration_icp( source, target, max_correspondence_distance=distance_threshold, estimation_method=o3d.registration.TransformationEstimationPointToPlane(), init=result_fast.transformation) T = reg_p2p.transformation R = T[0:3,0:3] t = T[0:3,3:] if isinstance(point_clouds, TensorListList): Rs.append(TensorList([torch.from_numpy(R).float(), target_R])) ts.append(TensorList([torch.from_numpy(t).float(), target_t])) else: Rs.append(torch.stack([torch.from_numpy(R).float(), target_R])) ts.append(torch.stack([torch.from_numpy(t).float(), target_t])) if isinstance(point_clouds, TensorListList): return dict(Rs=Rs, ts=ts, Vs=point_clouds) else: return dict(Rs=torch.stack(Rs), ts=torch.stack(ts), Vs=point_clouds)
def L2sqrTransList(t1_list, t2_list, R1_list, R2_list): err = TensorListList() for t1, t2, R1, R2 in zip(t1_list, t2_list, R1_list, R2_list): M = len(t1) err1 = TensorList() for ind1 in range(M - 1): for ind2 in range(ind1 + 1, M): # estimated relative translation from ind2 to ind1 test = R1[ind1].permute(1, 0) @ t1[ind2] - R1[ind1].permute( 1, 0) @ t1[ind1] # ground truth translation from ind2 to ind1 tgt = R2[ind1].permute(1, 0) @ t2[ind2] - R2[ind1].permute( 1, 0) @ t2[ind1] diff = test - tgt err1.append((diff * diff).sum(dim=-2)) err.append(err1) return err
def cluster_features(self, features, num_clusters): feature_labels_LL = TensorListList() for f in features: feature_labels_L = TensorList() fcat = torch.cat([fi for fi in f]) fcat = fcat.to("cpu").numpy() kmeans = KMeans(n_clusters=num_clusters, random_state=0).fit(fcat) labels = torch.from_numpy(kmeans.labels_) onehot = torch.nn.functional.one_hot( labels.long(), num_clusters).to(self.params.device) cnt = 0 for fi in f: feature_labels_L.append(onehot[cnt:cnt + fi.shape[0]]) cnt += fi.shape[0] feature_labels_LL.append(feature_labels_L.permute(1, 0).float()) return feature_labels_LL
def extract_features(self, x): coords = x['coordinates'] sinput, inds_list = self.create_sparse_tensors(coords) features_dict = self.feature_extractor(sinput) features = features_dict["features"] if "attention" in features_dict.keys(): att = features_dict["attention"] batch_indices = list(features.coords_man.get_batch_indices()) features_LL = TensorListList() att_LL = TensorListList() coords_ds_LL = TensorListList() inds_LL = [] ind_cnt = 0 for coords_b in coords: features_L = TensorList() att_L = TensorList() coords_ds_L = TensorList() inds_L = [] for coords_s in coords_b: mask = features.C[:, 0] == batch_indices[ind_cnt] # hacky way of finding batch channel dimension if mask.int().sum() != inds_list[ind_cnt].shape[0]: mask = features.C[:, -1] == batch_indices[ind_cnt] f = features.F[mask] assert f.shape[0] == inds_list[ind_cnt].shape[0] if "attention" in features_dict.keys(): a = att.F[mask] assert a.shape[0] == inds_list[ind_cnt].shape[0] att_L.append(a) features_L.append(f.permute(1, 0)) coords_ds_L.append(coords_s[:, inds_list[ind_cnt]]) inds_L.append(inds_list[ind_cnt]) ind_cnt = ind_cnt + 1 features_LL.append(features_L) att_LL.append(att_L) coords_ds_LL.append(coords_ds_L) inds_LL.append(inds_L) out = dict() out['features'] = features_LL out['att'] = att_LL out['coordinates_ds'] = coords_ds_LL.to(self.params.device) out['indices'] = inds_LL out['coordinates'] = x['coordinates'] return out
def register_point_sets(self, x): point_clouds = x["coordinates"].clone() t_tot = time.time() if isinstance(point_clouds, TensorListList): Rs = TensorListList() ts = TensorListList() else: # B, P, N, M = point_clouds.shape # assert P == 2 Rs = [] ts = [] target_R = torch.eye(3, 3) target_t = torch.zeros(3, 1) for pcds in point_clouds: source = o3d.geometry.PointCloud() target = o3d.geometry.PointCloud() source.points = o3d.utility.Vector3dVector(pcds[0].cpu().numpy().T) target.points = o3d.utility.Vector3dVector(pcds[1].cpu().numpy().T) T = self.dgr.register(source, target) R = T[0:3, 0:3] t = T[0:3, 3:] if isinstance(point_clouds, TensorListList): Rs.append(TensorList([torch.from_numpy(R).float(), target_R])) ts.append(TensorList([torch.from_numpy(t).float(), target_t])) else: Rs.append(torch.stack([torch.from_numpy(R).float(), target_R])) ts.append(torch.stack([torch.from_numpy(t).float(), target_t])) tot_time = time.time() - t_tot if isinstance(point_clouds, TensorListList): return dict(Rs=Rs, ts=ts, Vs=Rs @ point_clouds + ts, time=tot_time) else: return dict(Rs=torch.stack(Rs), ts=torch.stack(ts), Vs=point_clouds, time=tot_time)
def get_init_transformation_list_dgr(pcds, dgr_init_model): if isinstance(pcds, TensorListList): Rs = TensorListList() ts = TensorListList() else: # B, P, N, M = point_clouds.shape # assert P == 2 Rs = [] ts = [] target_R = torch.eye(3, 3).to(pcds[0][0].device) target_t = torch.zeros(3, 1).to(pcds[0][0].device) for pc in pcds: assert len(pc) == 2 R,t=dgr_init_model.register_point_sets(pc[0].permute(1,0), pc[1].permute(1,0)) if isinstance(pcds, TensorListList): Rs.append(TensorList([R, target_R])) ts.append(TensorList([t.unsqueeze(dim=1), target_t])) else: Rs.append(torch.stack([R, target_R])) ts.append(torch.stack([t.unsqueeze(dim=1), target_t])) return Rs, ts
def forward(self, x): coords = x['coordinates'].clone() t_tot = time.time() features_LL = TensorListList() att_LL = TensorListList() coords_ds_LL = TensorListList() inds_LL = [] ind_cnt = 0 for coords_b in coords: features_L = TensorList() att_L = TensorList() coords_ds_L = TensorList() inds_L = [] for coords_s in coords_b: pcd_down, f = self.preprocess_point_cloud( coords_s, self.params.voxel_size) features_L.append(f) coords_ds_L.append(pcd_down) ind_cnt = ind_cnt + 1 features_LL.append(features_L) att_LL.append(att_L) coords_ds_LL.append(coords_ds_L) inds_LL.append(inds_L) x = dict() x['features'] = self.cluster_features( features_LL, self.params.feature_distr_parameters.num_feature_clusters) x['att'] = att_LL x['coordinates_ds'] = coords_ds_LL.to(self.params.device) x['indices'] = inds_LL out = self.registration(x) tot_time = time.time() - t_tot print("tot time: %.1f ms" % (tot_time * 1000)) out["time"] = tot_time out["features"] = features_LL out["indices"] = inds_LL out["coordinates_ds"] = coords_ds_LL return out
def register_point_sets(self, x): Vs = x["coordinates_ds"] features = x["features"] features_w = x["att"] repeat_list = [len(V) for V in Vs] init_R, init_t = get_init_transformation_list( Vs, self.params.get("mean_init", True)) TVs = init_R @ Vs + init_t X = TensorList() Q = TensorList() mu = TensorList() for TV, Fs in zip(TVs, features): if self.params.cluster_init == "box": Xi = get_randn_box_cluster_means_list(TV, self.params.K) else: Xi = get_randn_sphere_cluster_means_list( TV, self.params.K, self.params.get("cluster_mean_scale", 1.0)) Q.append( get_scaled_cluster_precisions_list( TV, Xi, self.params.cluster_precision_scale)) X.append(Xi.T) feature_distr = feature_models.MultinomialModel( self.params.feature_distr_parameters, self.params.K, features, repeat_list=repeat_list) feature_distr.to(self.params.device) X = TensorListList(X, repeat=repeat_list) self.betas = get_default_beta(Q, self.params.gamma) Vs = Vs TVs = TVs # Compute the observation weights if self.params.use_dare_weighting: observation_weights = empirical_estimate(Vs, self.params.ow_args) ow_reg_factor = 8.0 ow_mean = observation_weights.mean(dim=0, keepdim=True) for idx in range(len(observation_weights)): for idxx in range(len(observation_weights[idx])): observation_weights[idx][idxx][observation_weights[idx][idxx] > ow_reg_factor * ow_mean[idx][idxx]] \ = ow_reg_factor * ow_mean[idx][idxx] else: observation_weights = 1.0 ds = TVs.permute(1, 0).sqe(X).permute(1, 0) if self.params.debug: self.visdom.register( dict(pcds=Vs[0].cpu(), X=X[0][0].cpu(), c=None), 'point_clouds', 2, 'init') time.sleep(1) Rs = init_R.to(self.params.device) ts = init_t.to(self.params.device) self.betas = TensorListList(self.betas, repeat=repeat_list) QL = TensorListList(Q, repeat=repeat_list) Riter = TensorListList() titer = TensorListList() TVs_iter = TensorListList() for i in range(self.params.num_iters): Qt = QL.permute(1, 0) ap = (-0.5 * ds * QL).exp() * QL.pow(1.5) if i < 1000: pyz_feature = feature_distr.posteriors(features) else: pyz_feature = 1.0 a = ap * pyz_feature ac_den = a.sum(dim=-1, keepdim=True) + self.betas a = a / ac_den # normalize row-wise a = a * observation_weights L = a.sum(dim=-2, keepdim=True).permute(1, 0) W = (Vs @ a) * QL b = L * Qt # weights, b mW = W.sum(dim=-1, keepdim=True) mX = (b.permute(1, 0) @ X).permute(1, 0) z = L.permute(1, 0) @ Qt P = (W @ X).permute(1, 0) - mX @ mW.permute(1, 0) / z # Compute R and t svd_list_list = P.cpu().svd() Rs = TensorListList() for svd_list in svd_list_list: Rs_list = TensorList() for svd in svd_list: uu, vv = svd.U, svd.V vvt = vv.permute(1, 0) detuvt = uu @ vvt detuvt = detuvt.det() S = torch.ones(1, 3) S[:, -1] = detuvt Rs_list.append((uu * S) @ vvt) Rs.append(Rs_list) Rs = Rs.to(self.params.device) Riter.append(Rs) ts = (mX - Rs @ mW) / z titer.append(ts) TVs = Rs @ Vs + ts TVs_iter.append(TVs.clone()) if self.params.debug: self.visdom.register( dict(pcds=TVs[0].cpu(), X=X[0][0].cpu(), c=None), 'point_clouds', 2, 'registration-iter') time.sleep(0.2) # Update X den = L.sum_list() if self.params.fix_cluster_pos_iter < i: X = (TVs @ a).permute(1, 0) X = TensorListList(X.sum_list() / den, repeat_list) # Update Q ds = TVs.permute(1, 0).sqe(X).permute(1, 0) wn = (a * ds).sum(dim=-2, keepdim=True).sum_list() Q = (3 * den / (wn.permute(1, 0) + 3 * den * self.params.epsilon)).permute( 1, 0) QL = TensorListList(Q, repeat=repeat_list) feature_distr.maximize(ap=ap, ow=observation_weights, y=features, den=ac_den) out = dict(Rs=Rs, ts=ts, X=X, Riter=Riter[:-1], titer=titer[:-1], Vs=TVs, Vs_iter=TVs_iter[:-1], ow=observation_weights) return out
def forward(self, x): coords = x['coordinates'].clone() t_tot = time.time() sinput, inds_list = self.create_sparse_tensors(coords) features_dict = self.feature_extractor(sinput) features = features_dict["features"] if "attention" in features_dict.keys(): att = features_dict["attention"] if torch.isnan(features.feats).any(): print("nans in features!") batch_indices = list(features.coords_man.get_batch_indices()) features_LL = TensorListList() att_LL = TensorListList() coords_ds_LL = TensorListList() inds_LL = [] ind_cnt = 0 for coords_b in coords: features_L = TensorList() att_L = TensorList() coords_ds_L = TensorList() inds_L = [] for coords_s in coords_b: mask = features.C[:, 0] == batch_indices[ind_cnt] if mask.int().sum() != inds_list[ind_cnt].shape[0]: mask = features.C[:, -1] == batch_indices[ind_cnt] f = features.F[mask] assert f.shape[0] == inds_list[ind_cnt].shape[0] if "attention" in features_dict.keys(): a = att.F[mask] assert a.shape[0] == inds_list[ind_cnt].shape[0] att_L.append(a) features_L.append(f) coords_ds_L.append(coords_s[:, inds_list[ind_cnt]]) inds_L.append(inds_list[ind_cnt]) ind_cnt = ind_cnt + 1 features_LL.append(features_L) att_LL.append(att_L) coords_ds_LL.append(coords_ds_L) inds_LL.append(inds_L) x = dict() x['features'] = self.cluster_features( features_LL, self.params.feature_distr_parameters.num_feature_clusters) x['att'] = att_LL x['coordinates_ds'] = coords_ds_LL.to(self.params.device) x['indices'] = inds_LL out = self.registration(x) tot_time = time.time() - t_tot print("tot time: %.1f ms" % (tot_time * 1000)) out["time"] = tot_time out["features"] = features_LL out["indices"] = inds_LL out["coordinates_ds"] = coords_ds_LL return out
def register_point_sets(self, x): point_clouds = x["coordinates"].clone() t_tot = time.time() if isinstance(point_clouds, TensorListList): Rs = TensorListList() ts = TensorListList() else: #B, P, N, M = point_clouds.shape #assert P == 2 Rs = [] ts = [] target_R = torch.eye(3, 3) target_t = torch.zeros(3, 1) for pcds in point_clouds: if self.params.get("mean_init", True): m_target = pcds[1].mean(dim=1) m_source = pcds[0].mean(dim=1) init_t = -(m_target - m_source).cpu().numpy() else: init_t = [0, 0, 0] trans_init = np.asarray([[1., 0., 0., init_t[0]], [0., 1., 0., init_t[1]], [0., 0., 1., init_t[2]], [0.0, 0.0, 0.0, 1.0]]) source = o3d.geometry.PointCloud() target = o3d.geometry.PointCloud() source.points = o3d.utility.Vector3dVector(pcds[0].cpu().numpy().T) target.points = o3d.utility.Vector3dVector(pcds[1].cpu().numpy().T) voxel_size = self.params.voxel_size max_correspondence_distance = self.params.threshold radius_normal = float(self.params.radius_normal) if self.params.metric == "p2pl": source.estimate_normals( search_param=o3d.geometry.KDTreeSearchParamHybrid( radius=radius_normal, max_nn=30)) target.estimate_normals( search_param=o3d.geometry.KDTreeSearchParamHybrid( radius=radius_normal, max_nn=30)) # o3d.geometry.estimate_normals(source) # o3d.geometry.estimate_normals(target) if self.params.get("downsample", True): source = source.voxel_down_sample(voxel_size=voxel_size) target = target.voxel_down_sample(voxel_size=voxel_size) reg_p2p = o3d.registration.registration_icp( source, target, max_correspondence_distance=max_correspondence_distance, estimation_method=self.metric, init=trans_init) T = reg_p2p.transformation R = T[0:3, 0:3] t = T[0:3, 3:] if isinstance(point_clouds, TensorListList): Rs.append(TensorList([torch.from_numpy(R).float(), target_R])) ts.append(TensorList([torch.from_numpy(t).float(), target_t])) else: Rs.append(torch.stack([torch.from_numpy(R).float(), target_R])) ts.append(torch.stack([torch.from_numpy(t).float(), target_t])) tot_time = time.time() - t_tot print("ICP tot time: %.1f ms" % (tot_time * 1000)) if isinstance(point_clouds, TensorListList): return dict(Rs=Rs, ts=ts, Vs=Rs @ point_clouds + ts, time=tot_time) else: return dict(Rs=torch.stack(Rs), ts=torch.stack(ts), Vs=point_clouds, time=tot_time)
def forward(self, x_in): coords = x_in['coordinates'].clone() t_tot = time.time() sinput, inds_list = self.create_sparse_tensors(coords) resample_time = time.time() - t_tot print("resample time: %.1f ms" % ((resample_time) * 1000)) if not self.params.feature_distr_parameters.model=='none': features_dict = self.feature_extractor(sinput) features = features_dict["features"] if "attention" in features_dict.keys(): att = features_dict["attention"] extract_time=time.time()-t_tot print("extract time: %.1f ms" % ((extract_time) * 1000)) batch_indices = list(features.coords_man.get_batch_indices()) else: features_dict=None extract_time=0 time_preprocess = time.time() features_LL = TensorListList() att_LL = TensorListList() coords_ds_LL = TensorListList() inds_LL = [] ind_cnt = 0 for coords_b in coords: features_L = TensorList() att_L = TensorList() coords_ds_L = TensorList() inds_L = [] for coords_s in coords_b: if not features_dict is None: mask = features.C[:, 0] == batch_indices[ind_cnt] # hacky way of finding batch channel dimension if mask.int().sum() != inds_list[ind_cnt].shape[0]: mask = features.C[:, -1] == batch_indices[ind_cnt] f = features.F[mask] assert f.shape[0] == inds_list[ind_cnt].shape[0] if "attention" in features_dict.keys(): a = att.F[mask] assert a.shape[0] == inds_list[ind_cnt].shape[0] att_L.append(a) features_L.append(f.permute(1, 0)) coords_ds_L.append(coords_s[:, inds_list[ind_cnt]]) inds_L.append(inds_list[ind_cnt]) ind_cnt = ind_cnt + 1 features_LL.append(features_L) att_LL.append(att_L) coords_ds_LL.append(coords_ds_L) inds_LL.append(inds_L) x = dict() x['features'] = features_LL x['att'] = att_LL x['coordinates_ds'] = coords_ds_LL.to(self.params.device) x['indices'] = inds_LL x['coordinates'] = x_in['coordinates'] print("preprocess time: %.1f ms" % ((time.time()-time_preprocess) * 1000)) reg_time=time.time() out = self.registration(x) reg_time2 = time.time() - reg_time print("reg time: %.1f ms" % ((time.time() - reg_time) * 1000)) tot_time = time.time() - t_tot print("tot time: %.1f ms" % (tot_time * 1000)) out["time"] = tot_time out["reg_time"] = reg_time2 out["extract_time"] =extract_time out["resample_time"] = resample_time out["coordinates_ds"] = coords_ds_LL return out