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 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 initialize_mu(self, K, features): X = TensorList() for TV in features: Xi = np.random.randn(TV[0].shape[0], K).astype(np.float32) Xi = torch.from_numpy(Xi).to(TV[0].device) Xi = Xi / torch.norm(Xi, dim=0, keepdim=True) X.append(Xi.permute(1, 0)) return X
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 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 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 initialize(self, K, features): init_method = self.params.get("init_method", "dirichlet") if init_method == "uniform": init_distr_list = [] for i in range(len(features)): init_distr_list.append( torch.ones(K, self.params.num_feature_clusters).float() / self.params.num_feature_clusters) elif init_method == "dirichlet": init_distr_list = TensorList() relative_distr = features.sum(dim=1).sum_list() relative_distr = relative_distr / relative_distr.sum() for r in relative_distr: dir = torch.distributions.dirichlet.Dirichlet(r) init_distr_list.append(dir.sample((K, ))) else: init_distr_list = [] for i in range(len(features)): init_distr_list.append( torch.ones(K, self.params.num_feature_clusters).float() / self.params.num_feature_clusters) self.distr = TensorListList(init_distr_list, repeat=self.repeat_list)
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): 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 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