def _get_kappa_adv(adv_pc, ori_pc, ori_normal, k=2): b, _, n = adv_pc.size() # compute knn between advPC and oriPC to get normal n_p #intra_dis = ((adv_pc.unsqueeze(3) - ori_pc.unsqueeze(2))**2).sum(1) #intra_idx = torch.topk(intra_dis, 1, dim=2, largest=False, sorted=True)[1] #normal = torch.gather(ori_normal, 2, intra_idx.view(b,1,n).expand(b,3,n)) intra_KNN = knn_points(adv_pc.permute(0, 2, 1), ori_pc.permute(0, 2, 1), K=1) #[dists:[b,n,1], idx:[b,n,1]] normal = knn_gather(ori_normal.permute(0, 2, 1), intra_KNN.idx).permute( 0, 3, 1, 2).squeeze(3).contiguous() # [b, 3, n] # compute knn between advPC and itself to get \|q-p\|_2 #inter_dis = ((adv_pc.unsqueeze(3) - adv_pc.unsqueeze(2))**2).sum(1) #inter_idx = torch.topk(inter_dis, k+1, dim=2, largest=False, sorted=True)[1][:, :, 1:].contiguous() #nn_pts = torch.gather(adv_pc, 2, inter_idx.view(b,1,n*k).expand(b,3,n*k)).view(b,3,n,k) inter_KNN = knn_points(adv_pc.permute(0, 2, 1), adv_pc.permute(0, 2, 1), K=k + 1) #[dists:[b,n,k+1], idx:[b,n,k+1]] nn_pts = knn_gather(adv_pc.permute(0, 2, 1), inter_KNN.idx).permute( 0, 3, 1, 2)[:, :, :, 1:].contiguous() # [b, 3, n ,k] vectors = nn_pts - adv_pc.unsqueeze(3) vectors = _normalize(vectors) return torch.abs( (vectors * normal.unsqueeze(3)).sum(1)).mean(2), normal # [b, n], [b, 3, n]
def eval_batch(points_pred, points_gt): d_1, _, _ = knn_points(points_pred, points_gt) d_2, _, _ = knn_points(points_pred, points_gt) err_1, _ = d_1.squeeze(-1).max(dim=1) err_2, _ = d_2.squeeze(-1).max(dim=1) err = torch.cat((err_1.unsqueeze(-1), err_2.unsqueeze(-1)), dim=-1) e, _ = err.max(dim=1) return e
def chamfer_loss(adv_pc, ori_pc): # Chamfer distance (two sides) #intra_dis = ((adv_pc.unsqueeze(3) - ori_pc.unsqueeze(2))**2).sum(1) #dis_loss = intra_dis.min(2)[0].mean(1) + intra_dis.min(1)[0].mean(1) adv_KNN = knn_points(adv_pc.permute(0, 2, 1), ori_pc.permute(0, 2, 1), K=1) #[dists:[b,n,1], idx:[b,n,1]] ori_KNN = knn_points(ori_pc.permute(0, 2, 1), adv_pc.permute(0, 2, 1), K=1) #[dists:[b,n,1], idx:[b,n,1]] dis_loss = adv_KNN.dists.contiguous().squeeze(-1).mean( -1) + ori_KNN.dists.contiguous().squeeze(-1).mean(-1) #[b] return dis_loss
def build_edges(spatial, r_max, k_max, return_indices=False, target_spatial=None): if k_max > 200: if device == "cuda": res = faiss.StandardGpuResources() if target_spatial is None: D, I = faiss.knn_gpu(res, spatial, spatial, k_max) else: D, I = faiss.knn_gpu(res, spatial, target_spatial, k_max) elif device == "cpu": index = faiss.IndexFlatL2(spatial.shape[1]) index.add(spatial) if target_spatial is None: D, I = index.search(spatial, k_max) else: D, I = index.search(target_spatial, k_max) else: if target_spatial is None: knn_object = ops.knn_points(spatial.unsqueeze(0), spatial.unsqueeze(0), K=k_max, return_sorted=False) else: knn_object = ops.knn_points( spatial.unsqueeze(0), target_spatial.unsqueeze(0), K=k_max, return_sorted=False, ) I = knn_object.idx[0] D = knn_object.dists[0] # Overlay the "source" hit ID onto each neighbour ID (this is necessary as the FAISS algo does some shortcuts) ind = torch.Tensor.repeat(torch.arange(I.shape[0], device=device), (I.shape[1], 1), 1).T edge_list = torch.stack([ind[D <= r_max**2], I[D <= r_max**2]]) # Remove self-loops edge_list = edge_list[:, edge_list[0] != edge_list[1]] if return_indices: return edge_list, D, I, ind else: return edge_list
def hausdorff_loss(adv_pc, ori_pc): #dis = ((adv_pc.unsqueeze(3) - ori_pc.unsqueeze(2))**2).sum(1) #hd_loss = torch.max(torch.min(dis, dim=2)[0], dim=1)[0] adv_KNN = knn_points(adv_pc.permute(0, 2, 1), ori_pc.permute(0, 2, 1), K=1) #[dists:[b,n,1], idx:[b,n,1]] hd_loss = adv_KNN.dists.contiguous().squeeze(-1).max(-1)[0] #[b] return hd_loss
def forward(self, x, nsample, xyz): if nsample == 1: sampled_idx = None sampled_xyz = torch.mean(xyz, dim=-1, keepdim=True) _, sampled_idx, sampled_xyz = ops.knn_points(sampled_xyz.transpose( 1, 2), xyz.transpose(1, 2), return_nn=True) sampled_xyz = sampled_xyz.squeeze(2) sampled_idx = sampled_idx.squeeze(1) else: sampled_idx, sampled_xyz = furthest_point_sample(xyz, nsample, NCHW=True) sampled_x = gather_points(x, sampled_idx) for i, mlp in enumerate(self.mlps): if i == 0: y, idx = self.get_local_graph(sampled_x, x, k=self.k) sampled_x = sampled_x.unsqueeze(-1).expand(-1, -1, -1, self.k) y = torch.cat([nn.functional.relu_(mlp(y)), sampled_x], dim=1) elif i == (self.n - 1): y = torch.cat([mlp(y), y], dim=1) else: y = torch.cat([nn.functional.relu_(mlp(y)), y], dim=1) y, _ = torch.max(y, dim=-1) return y, sampled_xyz, sampled_idx
def get_local_graph(self, query, x, k, idx=None): """Construct edge feature [x, NN_i - x] for each point x :param x: (B, C, N) k: int idx: (B, N, k) :return edge features: (B, C, N, k) """ if idx is None: # BCN(K+1), BN(K+1) idx, knn_point = ops.knn_points(query.transpose(1, 2), x.transpose(1, 2), K=k + 1, return_nn=True) idx = idx[:, :, 1:] knn_point = knn_point.permute(0, 2, 3, 1) knn_point = knn_point[:, :, :, 1:] neighbor_center = torch.unsqueeze(query, dim=-1) neighbor_center = neighbor_center.expand_as(knn_point) edge_feature = torch.cat( [neighbor_center, knn_point - neighbor_center], dim=1) return edge_feature, idx
def pointUniformLaplacian(points, knn_idx=None, nn_size=3): """ Args: points: (B, N, 3) knn_idx: (B, N, K) Returns: laplacian: (B, N, 1) """ batch_size, num_points, _ = points.shape if knn_idx is None: # find neighborhood, (B,N,K,3), (B,N,K) _, knn_idx, group_points = ops.knn_points(points, points, K=nn_size+1, return_nn=True) knn_idx = knn_idx[:, :, 1:] group_points = group_points[:, :, 1:, :] else: points_expanded = points.unsqueeze(dim=1).expand( (-1, num_points, -1, -1)) # BxNxk -> BxNxNxC index_batch_expanded = knn_idx.unsqueeze(dim=-1).expand( (-1, -1, -1, points.size(-1))) # BxMxkxC group_points = torch.gather(points_expanded, 2, index_batch_expanded) lap = -torch.sum(group_points, dim=2)/knn_idx.shape[2] + points return lap, knn_idx
def _build_knn(self, point_clouds): """ search for KNN again set knn_tree and knn_mask attributes TODO(yifan): use a real Kd_tree library to be able to store the data tree and query at each forward pass? """ # Find local neighborhood to compute weights with torch.autograd.enable_grad(): points_padded = point_clouds.points_padded() lengths = point_clouds.num_points_per_cloud() knn_result = ops.knn_points(points_padded, points_padded, lengths, lengths, K=self.knn_k, return_nn=True) self.knn_mask = torch.full(knn_result.idx.shape, False, dtype=torch.bool, device=points_padded.device) # valid knn result for b in range(self.knn_mask.shape[0]): self.knn_mask[ b, :lengths[b], :min(self.knn_k, lengths[b].item())] = True assert (torch.all(knn_result.dists[b][~self.knn_mask[b]] == 0)) self.knn_tree = KNN(knn=knn_result.knn[:, :, 1:, :], dists=knn_result.dists[:, :, 1:], idx=knn_result.idx[:, :, 1:]) self.knn_mask = self.knn_mask[:, :, 1:] assert (self.knn_mask.shape == self.knn_tree.dists.shape)
def forward(self, points, knn_idx=None): batchSize, PN, _ = points.shape if knn_idx is None: distance2, knn_idx, knn_points = ops.knn_points(points, points, K=self.nn_size + 1, return_nn=True) knn_points = knn_points[:, :, 1:, :].contiguous().detach() knn_idx = knn_idx[:, :, 1:].contiguous() else: knn_points = torch.gather( points.unsqueeze(1).expand(-1, PN, -1, -1), 2, knn_idx.unsqueeze(-1).expand(-1, -1, -1, points.shape[-1])) knn_v = knn_points - points.unsqueeze(dim=2) distance2 = torch.sum(knn_v * knn_v, dim=-1) loss = 1 / torch.sqrt(distance2 + 1e-4) loss = torch.where(distance2 < self.radius2, loss, torch.zeros_like(loss)) if self.reduction == "mean": return loss.mean() elif self.reduction == "max": return torch.mean(torch.max(loss, dim=-1)[0]) elif self.reduction == "sum": return loss.mean(torch.sum(loss, dim=-1)) elif self.reduction == "none": return loss else: raise NotImplementedError return loss
def forward(self, points_ref, points): """ point1: (B,N,D) ref points (where connectivity is computed) point2: (B,N,D) pred points, uses connectivity of point1 """ # find neighborhood, (B,N,K,3), (B,N,K), (B,N,K) _, knn_idx, group_points_ref = ops.knn_points(points_ref, points_ref, K=self.nn_size + 1, return_nn=True) knn_idx = knn_idx[:, :, 1:] group_points_ref = group_points_ref[:, :, 1:, :] dist_ref = torch.norm(group_points_ref - points_ref.unsqueeze(2), dim=-1, p=2) group_points = torch.gather( points.unsqueeze(1).expand(-1, knn_idx.shape[1], -1, -1), 2, knn_idx.unsqueeze(-1).expand(-1, -1, -1, points.shape[-1])) dist = torch.norm(group_points - points.unsqueeze(2), dim=-1, p=2) stretch = torch.max(dist / (dist_ref + 1e-10) - 1, torch.zeros_like(dist)) if self.reduction == "mean": return torch.mean(stretch) elif self.reduction == "sum": return torch.mean(torch.sum(stretch, dim=-1)) elif self.reduction == "none": return stretch elif self.reduction == "max": return torch.mean(torch.max(stretch, dim=-1)[0]) else: raise NotImplementedError
def pseudo_chamfer_loss(adv_pc, ori_pc): # Chamfer pseudo distance (one side) #intra_dis = ((adv_pc.unsqueeze(3) - ori_pc.unsqueeze(2))**2).sum(1) #b*n*n #dis_loss = intra_dis.min(2)[0].mean(1) adv_KNN = knn_points(adv_pc.permute(0, 2, 1), ori_pc.permute(0, 2, 1), K=1) #[dists:[b,n,1], idx:[b,n,1]] dis_loss = adv_KNN.dists.contiguous().squeeze(-1).mean(-1) #[b] return dis_loss
def _compute_global_Vrk(self, pointclouds, refresh=True, **kwargs): """ determine variance scaler used in globally (see _compute_isotropic_Vrk) Args: pointclouds: pointclouds in object coorindates Returns: h_k: scaler S_k: local frame """ if not refresh and self._Vrk_h is not None: h_k = self._Vrk_h else: # compute average density with torch.autograd.enable_grad(): pts_world = pointclouds.points_padded() num_points_per_cloud = pointclouds.num_points_per_cloud() if self.frnn_radius <= 0: # use knn here # logger_py.info("vrk knn points") sq_dist, _, _ = ops3d.knn_points(pts_world, pts_world, num_points_per_cloud, num_points_per_cloud, K=7) else: sq_dist, _, _, _ = frnn.frnn_grid_points(pts_world, pts_world, num_points_per_cloud, num_points_per_cloud, K=7, r=self.frnn_radius) # logger_py.info("frnn and knn dist close: {}".format(torch.allclose(sq_dist, sq_dist2))) sq_dist = sq_dist[:, :, 1:] # knn search is unreliable, set sq_dist manually sq_dist[num_points_per_cloud < 7] = 1e-3 h_k = 0.5 * sq_dist.max(dim=-1, keepdim=True)[0] # prevent some outlier rendered be too large, or too small h_k = h_k.mean(dim=1, keepdim=True).clamp(5e-5, 1e-3) Vrk_h = gather_batch_to_packed(h_k, pointclouds.packed_to_cloud_idx()) # Sk, a transformation from 2D local surface frame to 3D world frame # Because isometry, two axis are equivalent, we can simply # find two 3d vectors perpendicular to the point normals # (totalP, 2, 3) with torch.autograd.enable_grad(): normals = pointclouds.normals_packed() u0 = F.normalize(torch.cross(normals, normals + torch.rand_like(normals)), dim=-1) u1 = F.normalize(torch.cross(normals, u0), dim=-1) Sk = torch.stack([u0, u1], dim=1) Vrk = Vrk_h.view(-1, 1, 1) * Sk.transpose(1, 2) @ Sk return Vrk, Sk
def output(): dists, idxs, nn = knn_points(points1, points2, lengths1, lengths2, K, version=-1, return_nn=False, return_sorted=True) torch.cuda.synchronize()
def uniform_loss(adv_pc, percentages=[0.004, 0.006, 0.008, 0.010, 0.012], radius=1.0, k=2): if adv_pc.size(1) == 3: adv_pc = adv_pc.permute(0, 2, 1).contiguous() b, n, _ = adv_pc.size() npoint = int(n * 0.05) for p in percentages: p = p * 4 nsample = int(n * p) r = math.sqrt(p * radius) disk_area = math.pi * (radius**2) * p / nsample expect_len = torch.sqrt(torch.Tensor([disk_area])).cuda() adv_pc_flipped = adv_pc.transpose(1, 2).contiguous() new_xyz = pointnet2_utils.gather_operation( adv_pc_flipped, pointnet2_utils.furthest_point_sample(adv_pc, npoint)).transpose( 1, 2).contiguous() # (batch_size, npoint, 3) idx = pointnet2_utils.ball_query( r, nsample, adv_pc, new_xyz) #(batch_size, npoint, nsample) grouped_pcd = pointnet2_utils.grouping_operation( adv_pc_flipped, idx).permute(0, 2, 3, 1).contiguous() # (batch_size, npoint, nsample, 3) grouped_pcd = torch.cat(torch.unbind(grouped_pcd, axis=1), axis=0) grouped_pcd = grouped_pcd.permute(0, 2, 1).contiguous() #dis = torch.sqrt(((grouped_pcd.unsqueeze(3) - grouped_pcd.unsqueeze(2))**2).sum(1)+1e-12) # (batch_size*npoint, nsample, nsample) #dists, _ = torch.topk(dis, k+1, dim=2, largest=False, sorted=True) # (batch_size*npoint, nsample, k+1) inter_KNN = knn_points(grouped_pcd.permute(0, 2, 1), grouped_pcd.permute(0, 2, 1), K=k + 1) #[dists:[b,n,k+1], idx:[b,n,k+1]] uniform_dis = inter_KNN.dists[:, :, 1:].contiguous() uniform_dis = torch.sqrt(torch.abs(uniform_dis) + 1e-12) uniform_dis = uniform_dis.mean(axis=[-1]) uniform_dis = (uniform_dis - expect_len)**2 / (expect_len + 1e-12) uniform_dis = torch.reshape(uniform_dis, [-1]) mean = uniform_dis.mean() mean = mean * math.pow(p * 100, 2) #nothing 4 try: loss = loss + mean except: loss = mean return loss / len(percentages)
def _get_kappa_ori(pc, normal, k=2): b, _, n = pc.size() #inter_dis = ((pc.unsqueeze(3) - pc.unsqueeze(2))**2).sum(1) #inter_idx = torch.topk(inter_dis, k+1, dim=2, largest=False, sorted=True)[1][:, :, 1:].contiguous() #nn_pts = torch.gather(pc, 2, inter_idx.view(b,1,n*k).expand(b,3,n*k)).view(b,3,n,k) inter_KNN = knn_points(pc.permute(0, 2, 1), pc.permute(0, 2, 1), K=k + 1) #[dists:[b,n,k+1], idx:[b,n,k+1]] nn_pts = knn_gather(pc.permute(0, 2, 1), inter_KNN.idx).permute( 0, 3, 1, 2)[:, :, :, 1:].contiguous() # [b, 3, n ,k] vectors = nn_pts - pc.unsqueeze(3) vectors = _normalize(vectors) return torch.abs((vectors * normal.unsqueeze(3)).sum(1)).mean(2) # [b, n]
def curvature_loss(adv_pc, ori_pc, adv_kappa, ori_kappa, k=2): b, _, n = adv_pc.size() # intra_dis = ((input_curr_iter.unsqueeze(3) - pc_ori.unsqueeze(2))**2).sum(1) # intra_idx = torch.topk(intra_dis, 1, dim=2, largest=False, sorted=True)[1] # knn_theta_normal = torch.gather(theta_normal, 1, intra_idx.view(b,n).expand(b,n)) # curv_loss = ((curv_loss - knn_theta_normal)**2).mean(-1) intra_KNN = knn_points(adv_pc.permute(0, 2, 1), ori_pc.permute(0, 2, 1), K=1) #[dists:[b,n,1], idx:[b,n,1]] onenn_ori_kappa = torch.gather( ori_kappa, 1, intra_KNN.idx.squeeze(-1)).contiguous() # [b, n] curv_loss = ((adv_kappa - onenn_ori_kappa)**2).mean(-1) return curv_loss
def kNN_smoothing_loss(adv_pc, k, threshold_coef=1.05): b, _, n = adv_pc.size() #dis = ((adv_pc.unsqueeze(3) - adv_pc.unsqueeze(2))**2).sum(1) #[b,n,n] #dis, idx = torch.topk(dis, k+1, dim=2, largest=False, sorted=True)#[b,n,k+1] inter_KNN = knn_points(adv_pc.permute(0, 2, 1), adv_pc.permute(0, 2, 1), K=k + 1) #[dists:[b,n,k+1], idx:[b,n,k+1]] knn_dis = inter_KNN.dists[:, :, 1:].contiguous().mean(-1) #[b,n] knn_dis_mean = knn_dis.mean(-1) #[b] knn_dis_std = knn_dis.std(-1) #[b] threshold = knn_dis_mean + threshold_coef * knn_dis_std #[b] condition = torch.gt(knn_dis, threshold.unsqueeze(1)).float() #[b,n] dis_mean = knn_dis * condition #[b,n] return dis_mean.mean(1) #[b]
def knn_group_withI(points1, points2, intensity2, k): ''' Input: points1: [B,3,N] points2: [B,3,N] intensity2: [B,1,N] ''' points1 = points1.permute(0,2,1).contiguous() points2 = points2.permute(0,2,1).contiguous() _, nn_idx, nn = knn_points(points1, points2, K=k, return_nn=True) points_resi = nn - points1.unsqueeze(2).repeat(1,1,k,1) # [B,M,k,3] grouped_dist = torch.norm(points_resi, dim=-1, keepdim=True) grouped_features = knn_gather(intensity2.permute(0,2,1), nn_idx) # [B,M,k,1] new_features = torch.cat([points_resi, grouped_dist], dim=-1) # [B,5,M,k], [B,3,M,k], [B,1,M,k] return new_features.permute(0,3,1,2).contiguous(), \ nn.permute(0,3,1,2).contiguous(), \ grouped_features.permute(0,3,1,2).contiguous()
def batch_normals(points, base=None, nn_size=20, NCHW=True, idx=None): """ compute normals vectors for batched points [B, C, M] If base is given, compute the normals of points using the neighborhood in base The direction of normal could flip. Args: points: (B,C,M) base: (B,C,N) idx (B,M,nn_size) Returns: normals: (B,C,M) """ if base is None: base = points if NCHW: points = points.transpose(2, 1).contiguous() base = base.transpose(2, 1).contiguous() assert(nn_size < base.shape[1]) batch_size, M, C = points.shape # B,M,k,C if idx is None: _, idx, grouped_points = ops.knn_points(points, base, K=nn_size, return_nn=True) else: grouped_points = torch.gather(base.unsqueeze(1).expand(-1,M,-1,-1), 2, idx.unsqueeze(-1).expand(-1,-1,-1,C)) group_center = torch.mean(grouped_points, dim=2, keepdim=True) points = grouped_points - group_center allpoints = points.view(-1, nn_size, C).contiguous() # MB,C,k U, S, V = batch_svd(allpoints) # V is MBxCxC, last_u MBxC normals = V[:, :, -1] normals = normals.view(batch_size, M, C) if NCHW: normals = normals.transpose(1, 2) return normals, idx
def forward(self, points_ref, points): """ point1: (B,N,D) ref points (where connectivity is computed) point2: (B,N,D) pred points, uses connectivity of point1 """ # find neighborhood, (B,N,K,3), (B,N,K) _, knn_idx, group_points = ops.knn_points(points_ref, points_ref, K=self.nn_size + 1, return_nn=True) knn_idx = knn_idx[:, :, 1:] group_points = group_points[:, :, 1:, :] dist_ref = torch.norm(group_points - points_ref.unsqueeze(2), dim=-1, p=2) # dist_ref = torch.sqrt(dist_ref) # B,N,K,D group_points = torch.gather( points.unsqueeze(1).expand(-1, knn_idx.shape[1], -1, -1), 2, knn_idx.unsqueeze(-1).expand(-1, -1, -1, points.shape[-1])) dist = torch.norm(group_points - points.unsqueeze(2), dim=-1, p=2) # print(group_points, group_points2) return self.metric(dist_ref, dist)
def estimate_normal_via_ori_normal(pc_adv, pc_ori, normal_ori, k): # pc_adv, pc_ori, normal_ori : [b,3,n] b, _, n = pc_adv.size() intra_KNN = knn_points(pc_adv.permute(0, 2, 1), pc_ori.permute(0, 2, 1), K=k) #[dists:[b,n,k], idx:[b,n,k]] inter_value = intra_KNN.dists[:, :, 0].contiguous() inter_idx = intra_KNN.idx.permute(0, 2, 1).contiguous() normal_pts = knn_gather(normal_ori.permute( 0, 2, 1), intra_KNN.idx).permute(0, 3, 1, 2).contiguous() # [b, 3, n ,k] normal_pts_avg = normal_pts.mean(dim=-1) normal_pts_avg = normal_pts_avg / (normal_pts_avg.norm(dim=1) + 1e-12) # If the points are not modified (distance = 0), use the normal directly from the original # one. Otherwise, use the mean of the normals of the k-nearest points. normal_ori_select = normal_pts[:, :, :, 0] condition = (inter_value < 1e-6).unsqueeze(1).expand_as(normal_ori_select) normals_estimated = torch.where(condition, normal_ori_select, normal_pts_avg) return normals_estimated
def knn_group(self, points1, points2, features2, k): ''' For each point in points1, query kNN points/features in points2/features2 Input: points1: [B,3,N] points2: [B,3,N] features2: [B,C,N] Output: new_features: [B,4,N] nn: [B,3,N] grouped_features: [B,C,N] ''' points1 = points1.permute(0,2,1).contiguous() points2 = points2.permute(0,2,1).contiguous() _, nn_idx, nn = knn_points(points1, points2, K=k, return_nn=True) points_resi = nn - points1.unsqueeze(2).repeat(1,1,k,1) grouped_dist = torch.norm(points_resi, dim=-1, keepdim=True) grouped_features = knn_gather(features2.permute(0,2,1), nn_idx) new_features = torch.cat([points_resi, grouped_dist], dim=-1) return new_features.permute(0,3,1,2).contiguous(),\ nn.permute(0,3,1,2).contiguous(),\ grouped_features.permute(0,3,1,2).contiguous()
def knn_pcl(pcl_1, pcl_2, n_neighbors=1, return_nn=False): dists, idxs, pcl_knn_to_1 = knn_points(pcl_1.points_padded(), pcl_2.points_padded(), pcl_1.num_points_per_cloud(), pcl_2.num_points_per_cloud(), K=n_neighbors, return_nn=return_nn) return dists, idxs, pcl_knn_to_1
def init_boundary_volume( batch_size: int, volume_size: Tuple[int, int, int], border_offset: int = 2, shape: str = "cube", volume_translation: torch.Tensor = ZERO_TRANSLATION, ): """ Generate a volume with sides colored with distinct colors. """ device = torch.device("cuda") # first center the volume for the purpose of generating the canonical shape volume_translation_tmp = (0.0, 0.0, 0.0) # set the voxel size to 1 / (volume_size-1) volume_voxel_size = 1 / (volume_size[0] - 1.0) # colors of the sides of the cube clr_sides = torch.tensor( [ [1.0, 1.0, 1.0], [1.0, 0.0, 0.0], [1.0, 0.0, 1.0], [1.0, 1.0, 0.0], [0.0, 1.0, 0.0], [0.0, 1.0, 1.0], ], dtype=torch.float32, device=device, ) # get the coord grid of the volume coord_grid = Volumes( densities=torch.zeros(1, 1, *volume_size, device=device), voxel_size=volume_voxel_size, volume_translation=volume_translation_tmp, ).get_coord_grid()[0] # extract the boundary points and their colors of the cube if shape == "cube": boundary_points, boundary_colors = [], [] for side, clr_side in enumerate(clr_sides): first = side % 2 dim = side // 2 slices = [slice(border_offset, -border_offset, 1)] * 3 slices[dim] = int(border_offset * (2 * first - 1)) slices.append(slice(0, 3, 1)) boundary_points_ = coord_grid[slices].reshape(-1, 3) boundary_points.append(boundary_points_) boundary_colors.append(clr_side[None].expand_as(boundary_points_)) # set the internal part of the volume to be completely opaque volume_densities = torch.zeros(*volume_size, device=device) volume_densities[[slice(border_offset, -border_offset, 1)] * 3] = 1.0 boundary_points, boundary_colors = [ torch.cat(p, dim=0) for p in [boundary_points, boundary_colors] ] # color the volume voxels with the nearest boundary points' color _, idx, _ = knn_points(coord_grid.view(1, -1, 3), boundary_points.view(1, -1, 3)) volume_colors = (boundary_colors[idx.view(-1)].view(*volume_size, 3).permute( 3, 0, 1, 2)) elif shape == "sphere": # set all voxels within a certain distance from the origin to be opaque volume_densities = (coord_grid.norm(dim=-1) <= 0.5 * volume_voxel_size * (volume_size[0] - border_offset)).float() # color each voxel with the standrd spherical color volume_colors = ( (torch.nn.functional.normalize(coord_grid, dim=-1) + 1.0) * 0.5).permute(3, 0, 1, 2) else: raise ValueError(shape) volume_voxel_size = torch.ones( (batch_size, 1), device=device) * volume_voxel_size volume_translation = volume_translation.expand(batch_size, 3) volumes = Volumes( densities=volume_densities[None, None].expand(batch_size, 1, *volume_size), features=volume_colors[None].expand(batch_size, 3, *volume_size), voxel_size=volume_voxel_size, volume_translation=volume_translation, ) return volumes, volume_voxel_size, volume_translation
def iterative_closest_point( X: Union[torch.Tensor, "Pointclouds"], Y: Union[torch.Tensor, "Pointclouds"], init_transform: Optional[SimilarityTransform] = None, max_iterations: int = 100, relative_rmse_thr: float = 1e-6, estimate_scale: bool = False, allow_reflection: bool = False, verbose: bool = False, ) -> ICPSolution: """ Executes the iterative closest point (ICP) algorithm [1, 2] in order to find a similarity transformation (rotation `R`, translation `T`, and optionally scale `s`) between two given differently-sized sets of `d`-dimensional points `X` and `Y`, such that: `s[i] X[i] R[i] + T[i] = Y[NN[i]]`, for all batch indices `i` in the least squares sense. Here, Y[NN[i]] stands for the indices of nearest neighbors from `Y` to each point in `X`. Note, however, that the solution is only a local optimum. Args: **X**: Batch of `d`-dimensional points of shape `(minibatch, num_points_X, d)` or a `Pointclouds` object. **Y**: Batch of `d`-dimensional points of shape `(minibatch, num_points_Y, d)` or a `Pointclouds` object. **init_transform**: A named-tuple `SimilarityTransform` of tensors `R`, `T, `s`, where `R` is a batch of orthonormal matrices of shape `(minibatch, d, d)`, `T` is a batch of translations of shape `(minibatch, d)` and `s` is a batch of scaling factors of shape `(minibatch,)`. **max_iterations**: The maximum number of ICP iterations. **relative_rmse_thr**: A threshold on the relative root mean squared error used to terminate the algorithm. **estimate_scale**: If `True`, also estimates a scaling component `s` of the transformation. Otherwise assumes the identity scale and returns a tensor of ones. **allow_reflection**: If `True`, allows the algorithm to return `R` which is orthonormal but has determinant==-1. **verbose**: If `True`, prints status messages during each ICP iteration. Returns: A named tuple `ICPSolution` with the following fields: **converged**: A boolean flag denoting whether the algorithm converged successfully (=`True`) or not (=`False`). **rmse**: Attained root mean squared error after termination of ICP. **Xt**: The point cloud `X` transformed with the final transformation (`R`, `T`, `s`). If `X` is a `Pointclouds` object, returns an instance of `Pointclouds`, otherwise returns `torch.Tensor`. **RTs**: A named tuple `SimilarityTransform` containing a batch of similarity transforms with fields: **R**: Batch of orthonormal matrices of shape `(minibatch, d, d)`. **T**: Batch of translations of shape `(minibatch, d)`. **s**: batch of scaling factors of shape `(minibatch, )`. **t_history**: A list of named tuples `SimilarityTransform` the transformation parameters after each ICP iteration. References: [1] Besl & McKay: A Method for Registration of 3-D Shapes. TPAMI, 1992. [2] https://en.wikipedia.org/wiki/Iterative_closest_point """ # make sure we convert input Pointclouds structures to # padded tensors of shape (N, P, 3) Xt, num_points_X = oputil.convert_pointclouds_to_tensor(X) Yt, num_points_Y = oputil.convert_pointclouds_to_tensor(Y) b, size_X, dim = Xt.shape if (Xt.shape[2] != Yt.shape[2]) or (Xt.shape[0] != Yt.shape[0]): raise ValueError( "Point sets X and Y have to have the same " + "number of batches and data dimensions." ) if ((num_points_Y < Yt.shape[1]).any() or (num_points_X < Xt.shape[1]).any()) and ( num_points_Y != num_points_X ).any(): # we have a heterogeneous input (e.g. because X/Y is # an instance of Pointclouds) mask_X = ( torch.arange(size_X, dtype=torch.int64, device=Xt.device)[None] < num_points_X[:, None] ).type_as(Xt) else: mask_X = Xt.new_ones(b, size_X) # clone the initial point cloud Xt_init = Xt.clone() if init_transform is not None: # parse the initial transform from the input and apply to Xt try: R, T, s = init_transform assert ( R.shape == torch.Size((b, dim, dim)) and T.shape == torch.Size((b, dim)) and s.shape == torch.Size((b,)) ) except Exception: raise ValueError( "The initial transformation init_transform has to be " "a named tuple SimilarityTransform with elements (R, T, s). " "R are dim x dim orthonormal matrices of shape " "(minibatch, dim, dim), T is a batch of dim-dimensional " "translations of shape (minibatch, dim) and s is a batch " "of scalars of shape (minibatch,)." ) # apply the init transform to the input point cloud Xt = _apply_similarity_transform(Xt, R, T, s) else: # initialize the transformation with identity R = oputil.eyes(dim, b, device=Xt.device, dtype=Xt.dtype) T = Xt.new_zeros((b, dim)) s = Xt.new_ones(b) prev_rmse = None rmse = None iteration = -1 converged = False # initialize the transformation history t_history = [] # the main loop over ICP iterations for iteration in range(max_iterations): Xt_nn_points = knn_points( Xt, Yt, lengths1=num_points_X, lengths2=num_points_Y, K=1, return_nn=True ).knn[:, :, 0, :] # get the alignment of the nearest neighbors from Yt with Xt_init R, T, s = corresponding_points_alignment( Xt_init, Xt_nn_points, weights=mask_X, estimate_scale=estimate_scale, allow_reflection=allow_reflection, ) # apply the estimated similarity transform to Xt_init Xt = _apply_similarity_transform(Xt_init, R, T, s) # add the current transformation to the history t_history.append(SimilarityTransform(R, T, s)) # compute the root mean squared error Xt_sq_diff = ((Xt - Xt_nn_points) ** 2).sum(2) rmse = oputil.wmean(Xt_sq_diff[:, :, None], mask_X).sqrt()[:, 0, 0] # compute the relative rmse if prev_rmse is None: relative_rmse = rmse.new_ones(b) else: relative_rmse = (prev_rmse - rmse) / prev_rmse if verbose: rmse_msg = ( f"ICP iteration {iteration}: mean/max rmse = " + f"{rmse.mean():1.2e}/{rmse.max():1.2e} " + f"; mean relative rmse = {relative_rmse.mean():1.2e}" ) print(rmse_msg) # check for convergence if (relative_rmse <= relative_rmse_thr).all(): converged = True break # update the previous rmse prev_rmse = rmse if verbose: if converged: print(f"ICP has converged in {iteration + 1} iterations.") else: print(f"ICP has not converged in {max_iterations} iterations.") if oputil.is_pointclouds(X): Xt = X.update_padded(Xt) # type: ignore return ICPSolution(converged, rmse, Xt, SimilarityTransform(R, T, s), t_history)
def mine_triplets(self, batch, true_edges, spatial, r_max, k_max): # -------- TRUTH torch_e = torch.sparse.FloatTensor( true_edges, torch.ones(true_edges.shape[1]).to(device), size=(len(spatial), len(spatial)), ) sparse_sum = torch.sparse.sum(torch_e, dim=0) num_true_torch = torch.zeros(len(spatial)).to(device).int() num_true_torch[sparse_sum.indices()] = sparse_sum.values().int() sorted_true_indices = torch.argsort(true_edges[0]) sorted_true_edges = true_edges[:, sorted_true_indices] # --------- HNM knn_object = ops.knn_points(spatial.unsqueeze(0), spatial.unsqueeze(0), K=k_max, return_sorted=False) I, D = knn_object.idx[0], knn_object.dists[0] # ---------- Shuffle shuffled_index = torch.randperm(I.shape[1]) shuffled_I = I[:, shuffled_index] shuffled_D = D[:, shuffled_index] ind = torch.Tensor.repeat( torch.arange(shuffled_I.shape[0], device=device), (shuffled_I.shape[1], 1), 1, ).T # ---------- Constraints shuffled_I[(shuffled_D > r_max) & (shuffled_D < (r_max / 2))] = -1 shuffled_I[batch.pid[ind] == batch.pid[shuffled_I]] = -1 shuffled_I[ind == shuffled_I] = -1 # ----------- Reshape with -1's squished_I = push_all_negs_back(shuffled_I.cpu().numpy()) squished_I = torch.from_numpy(squished_I).to(device) # ---------- Handle # pos > # neg pos_available = num_true_torch > 0 squished_I = torch.cat( [ squished_I, -1 * torch.ones( squished_I.shape[0], max(0, num_true_torch.max() - k_max), dtype=int, device=device, ), ], axis=-1, ) # ----------- Build Triplets selected_negatives = squished_I[pos_available][num_true_torch[ pos_available, None] > torch.arange(squished_I.shape[1], device=device)] triplets = torch.cat( [sorted_true_edges, selected_negatives.unsqueeze(0)], axis=0) triplets = triplets[:, triplets[2] != -1] return triplets
def _compute_isotropic_Vrk(self, pointclouds, refresh=True, **kwargs): """ determine the variance in the local surface frame h * Sk.T @ Sk, where Sk is 2x3 local surface coordinate to world coordinate. determine the h_k in V_k^r = h_k*Id using nearest neighbor heuristically h_k = mean(dist between points in a small neighbor) The larger h_k is, the larger the splat is NOTE: h_k in inverse to the definition in the paper, the larger h_k, the larger the splats Args: pointclouds: pointcloud in object coordinate Returns: h_k: [N,3,3] tensor for each point S_k: [N,2,3] local frame """ if not refresh and self._Vrk_h is not None and \ pointclouds.num_points_per_cloud().sum() == self._Vrk_h.shape[0]: pass else: with torch.autograd.enable_grad(): pts_world = pointclouds.points_padded() num_points_per_cloud = pointclouds.num_points_per_cloud() if self.frnn_radius <= 0: # logger_py.info("vrk knn points") sq_dist, _, _ = ops3d.knn_points(pts_world, pts_world, num_points_per_cloud, num_points_per_cloud, K=7) else: sq_dist, _, _, _ = frnn.frnn_grid_points(pts_world, pts_world, num_points_per_cloud, num_points_per_cloud, K=7, r=self.frnn_radius) sq_dist = sq_dist[:, :, 1:] # knn search is unreliable, set sq_dist manually sq_dist[num_points_per_cloud < 7] = 1e-3 # (totalP, knnK) sq_dist = ops3d.padded_to_packed( sq_dist, pointclouds.cloud_to_packed_first_idx(), num_points_per_cloud.sum().item()) # [totalP, ] h_k = 0.5 * sq_dist.max(dim=-1, keepdim=True)[0] # prevent some outlier rendered be too large, or too small self._Vrk_h = h_k.clamp(5e-5, 0.01) # Sk, a transformation from 2D local surface frame to 3D world frame # Because isometry, two axis are equivalent, we can simply # find two 3d vectors perpendicular to the point normals # (totalP, 2, 3) with torch.autograd.enable_grad(): normals = pointclouds.normals_packed() u0 = F.normalize(torch.cross(normals, normals + torch.rand_like(normals)), dim=-1) u1 = F.normalize(torch.cross(normals, u0), dim=-1) Sk = torch.stack([u0, u1], dim=1) Vrk = self._Vrk_h.view(-1, 1, 1) * Sk.transpose(1, 2) @ Sk return Vrk, Sk
def forward(self, body_model_output, camera, gt_joints, joints_conf, body_model_faces, joint_weights, use_vposer=False, pose_embedding=None, scan_tensor=None, visualize=False, scene_v=None, scene_vn=None, scene_f=None, ftov=None, **kwargs): batch_size = gt_joints.shape[0] forehead_vert_id = 336 # Patrick: replace head joint with a vertex on the head model_joints = body_model_output.joints model_joints[:, 0, :] = body_model_output.vertices[:, forehead_vert_id, :] projected_joints = camera(model_joints) # Calculate the weights for each joints weights = (joint_weights * joints_conf if self.use_joints_conf else joint_weights).unsqueeze(dim=-1) # Calculate the distance of the projected joints from # the ground truth 2D detections joint_err = gt_joints - projected_joints joint_err_sq = joint_err.pow(2) # joint_diff = self.robustifier(gt_joints - projected_joints) joint_loss = (torch.sum(weights**2 * joint_err_sq, dim=[1, 2]) * self.data_weight**2) # Calculate the loss from the Pose prior if use_vposer: pprior_loss = (pose_embedding.pow(2).sum() * self.body_pose_weight**2) else: pprior_loss = self.body_pose_prior( body_model_output.body_pose, body_model_output.betas) * self.body_pose_weight**2 shape_loss = self.shape_prior( body_model_output.betas) * self.shape_weight**2 # Patrick - weight and height loss tmp_betas = body_model_output.betas tmp_gender = self.gender_tensor # print('Gender tensor', tmp_gender) batch_weight_est, batch_height_est = self.betanet( tmp_gender, tmp_betas) batch_height_est = 100 * batch_height_est # print('Height {} est {}'.format(self.height, batch_height_est.detach().item())) # print('Weight {} est {}'.format(self.weight, batch_weight_est.detach().item())) # print('Cur gender flag', tmp_gender) global_vars.cur_weight = batch_weight_est.detach().cpu().numpy() global_vars.cur_height = batch_height_est.detach().cpu().numpy() d_weight = self.weight - batch_weight_est.squeeze() d_height = self.height - batch_height_est.squeeze() physical_loss = d_weight.pow(2) * self.weight_w + d_height.pow( 2) * self.height_w # Calculate the prior over the joint rotations. This a heuristic used # to prevent extreme rotation of the elbows and knees body_pose = body_model_output.full_pose[:, 3:66] # Patrick: turn off this loss angle_prior_loss = 0 * torch.sum( self.angle_prior(body_pose)) * self.bending_prior_weight**2 # Apply the prior on the pose space of the hand left_hand_prior_loss, right_hand_prior_loss = 0.0, 0.0 if self.use_hands and self.left_hand_prior is not None: left_hand_prior_loss = torch.sum( self.left_hand_prior( body_model_output.left_hand_pose)) * \ self.hand_prior_weight ** 2 if self.use_hands and self.right_hand_prior is not None: right_hand_prior_loss = torch.sum( self.right_hand_prior( body_model_output.right_hand_pose)) * \ self.hand_prior_weight ** 2 expression_loss = 0.0 jaw_prior_loss = 0.0 if self.use_face: expression_loss = torch.sum(self.expr_prior( body_model_output.expression)) * \ self.expr_prior_weight ** 2 if hasattr(self, 'jaw_prior'): jaw_prior_loss = torch.sum( self.jaw_prior( body_model_output.jaw_pose.mul(self.jaw_prior_weight))) pen_loss = 0.0 # Calculate the loss due to interpenetration if (self.interpenetration and self.coll_loss_weight.item() > 0): batch_size = projected_joints.shape[0] triangles = torch.index_select(body_model_output.vertices, 1, body_model_faces.view(-1)).view( batch_size, -1, 3, 3) with torch.no_grad(): collision_idxs = self.search_tree(triangles) # Remove unwanted collisions if self.tri_filtering_module is not None: collision_idxs = self.tri_filtering_module(collision_idxs) if collision_idxs.ge(0).sum().item() > 0: pen_loss = self.coll_loss_weight * self.pen_distance( triangles, collision_idxs) s2m_dist = 0.0 m2s_dist = 0.0 # calculate the scan2mesh and mesh2scan loss from the sparse point cloud if (self.s2m or self.m2s) and (self.s2m_weight > 0 or self.m2s_weight > 0) and scan_tensor is not None: # vertices_np = body_model_output.vertices.detach().cpu().numpy().squeeze() # body_faces_np = body_model_faces.detach().cpu().numpy().reshape(-1, 3) # m = Mesh(v=vertices_np, f=body_faces_np) # # (vis, n_dot) = visibility_compute(v=m.v, f=m.f, cams=np.array([[0.0, 0.0, 0.0]])) # vis = vis.squeeze() # TODO ignore body mask? # TODO ignore visibility map? Would help with points getting stuck inside body, only care about normal # Maybe need to do this in packed? in_mesh_verts = body_model_output.vertices in_mesh_faces = body_model_faces.unsqueeze(0).repeat( batch_size, 1, 1) body_mesh = pytorch3d.structures.Meshes(verts=in_mesh_verts, faces=in_mesh_faces) mesh_normals = body_mesh.verts_normals_padded() mesh_verts = body_mesh.verts_padded() camera_pos = torch.tensor([0.0, 0.0, 0.0], device=mesh_verts.device) vec_mesh_to_cam = camera_pos - mesh_verts towards_camera = torch.sum(mesh_normals * vec_mesh_to_cam, dim=2) > 0 num_mesh_verts_towards_camera = torch.sum(towards_camera, dim=1) mesh_verts_towards_camera = torch.zeros( [batch_size, num_mesh_verts_towards_camera.max(), 3], device=mesh_verts.device) for i in range( batch_size): # There's probably a cleaner way to do this mesh_verts_towards_camera[ i, 0:num_mesh_verts_towards_camera[i], :] = mesh_verts[ i, towards_camera[i, :], :] num_scan_points = scan_tensor.num_points_per_cloud( ) # Get number of points in mesh if self.s2m and self.s2m_weight > 0: # Note, returns squared distance s2m_dist = knn_points(scan_tensor.points_padded(), mesh_verts_towards_camera, lengths1=num_scan_points, lengths2=num_mesh_verts_towards_camera, K=1) # s2m_dist, _, _, _ = distChamfer(scan_tensor, body_model_output.vertices[:, np.where(vis > 0)[0], :]) s2m_dist = self.s2m_robustifier(s2m_dist.dists.sqrt()) s2m_dist = self.s2m_weight * s2m_dist.sum(dim=[1, 2]) if self.m2s and self.m2s_weight > 0: # _, m2s_dist, _, _ = distChamfer(scan_tensor, body_model_output.vertices[:, np.where(np.logical_and(vis > 0, self.body_mask))[0], :]) m2s_dist = knn_points(mesh_verts_towards_camera, scan_tensor.points_padded(), lengths2=num_scan_points, lengths1=num_mesh_verts_towards_camera, K=1) m2s_dist = self.m2s_robustifier(m2s_dist.dists.sqrt()) m2s_dist = self.m2s_weight * m2s_dist.sum(dim=[1, 2]) # Transform vertices to world coordinates if self.R is not None and self.t is not None: vertices = body_model_output.vertices.view(-1, 3) nv = vertices.shape[0] vertices = self.R.mm(vertices.t()).t() + self.t.repeat([nv, 1]) vertices = vertices.reshape(batch_size, -1, 3) # Compute scene penetration using signed distance field (SDF) # sdf_penetration_loss = 0.0 # if self.sdf_penetration and self.sdf_penetration_weight > 0: # grid_dim = self.sdf.shape[0] # sdf_ids = torch.round((vertices.squeeze() - self.grid_min) / self.voxel_size).to(dtype=torch.long) # Convert SMPL vertex to closest voxel ID # sdf_ids.clamp_(min=0, max=grid_dim-1) # Clamp to limits of grid # # norm_vertices = (vertices - self.grid_min) / (self.grid_max - self.grid_min) * 2 - 1 # Put SMPL verts into voxel space # body_sdf = F.grid_sample(self.sdf.view(1, 1, grid_dim, grid_dim, grid_dim), # norm_vertices[:, :, [2, 1, 0]].view(1, nv, 1, 1, 3), # padding_mode='border') # Calculate SDF for each SMPL vertex # sdf_normals = self.sdf_normals[sdf_ids[:,0], sdf_ids[:,1], sdf_ids[:,2]] # Find the SDF normal for each SMPL vertex # # if there are no penetrating vertices then set sdf_penetration_loss = 0 # if body_sdf.lt(0).sum().item() < 1: # sdf_penetration_loss = torch.tensor(0.0, dtype=joint_loss.dtype, device=joint_loss.device) # else: # if sdf_normals is None: # sdf_penetration_loss = self.sdf_penetration_weight * (body_sdf[body_sdf < 0].unsqueeze(dim=-1).abs()).pow(2).sum(dim=-1).sqrt().sum() # else: # sdf_penetration_loss = self.sdf_penetration_weight * (body_sdf[body_sdf < 0].unsqueeze(dim=-1).abs() * sdf_normals[body_sdf.view(-1) < 0, :]).pow(2).sum(dim=-1).sqrt().sum() sdf_penetration_loss = 0.0 if self.sdf_penetration and self.sdf_penetration_weight > 0: # Bed is at +z 2150, pointing down bed_height = 2.150 body_sdf = bed_height - vertices[:, :, 2] # Less than zero inside bed sdf_normals = torch.zeros_like(vertices) sdf_normals[:, :, 2] = -1 # if there are no penetrating vertices then set sdf_penetration_loss = 0 if body_sdf.lt(0).sum().item() < 1: sdf_penetration_loss = torch.tensor(0.0, dtype=joint_loss.dtype, device=joint_loss.device) else: sel_sdf = torch.max(body_sdf, torch.zeros_like(body_sdf)) sdf_dot = sel_sdf.unsqueeze(2) * sdf_normals sdf_penetration_loss = self.sdf_penetration_weight * sdf_dot.pow( 2).sum(dim=2).sqrt().sum(dim=1) # Compute the contact loss contact_loss = 0.0 if self.contact and self.contact_loss_weight > 0: # select contact vertices contact_body_vertices = vertices[:, self.contact_verts_ids, :] contact_dist, _, idx1, _ = distChamfer( contact_body_vertices.contiguous(), scene_v) body_triangles = torch.index_select(vertices, 1, body_model_faces).view( 1, -1, 3, 3) # Calculate the edges of the triangles # Size: BxFx3 edge0 = body_triangles[:, :, 1] - body_triangles[:, :, 0] edge1 = body_triangles[:, :, 2] - body_triangles[:, :, 0] # Compute the cross product of the edges to find the normal vector of # the triangle body_normals = torch.cross(edge0, edge1, dim=2) # Normalize the result to get a unit vector body_normals = body_normals / \ torch.norm(body_normals, 2, dim=2, keepdim=True) # compute the vertex normals body_v_normals = torch.mm(ftov, body_normals.squeeze()) body_v_normals = body_v_normals / \ torch.norm(body_v_normals, 2, dim=1, keepdim=True) # vertix normals of contact vertices contact_body_verts_normals = body_v_normals[ self.contact_verts_ids, :] # scene normals of the closest points on the scene surface to the contact vertices contact_scene_normals = scene_vn[:, idx1.squeeze().to(dtype=torch.long ), :].squeeze() # compute the angle between contact_verts normals and scene normals angles = torch.asin( torch.norm(torch.cross(contact_body_verts_normals, contact_scene_normals), 2, dim=1, keepdim=True)) * 180 / np.pi # consider only the vertices which their normals match valid_contact_mask = (angles.le(self.contact_angle) + angles.ge(180 - self.contact_angle)).ge(1) valid_contact_ids = valid_contact_mask.squeeze().nonzero().squeeze( ) contact_dist = self.contact_robustifier( contact_dist[:, valid_contact_ids].sqrt()) contact_loss = self.contact_loss_weight * contact_dist.mean() total_loss = (joint_loss + pprior_loss + shape_loss + angle_prior_loss + pen_loss + jaw_prior_loss + expression_loss + left_hand_prior_loss + right_hand_prior_loss + m2s_dist + s2m_dist + sdf_penetration_loss + contact_loss + physical_loss) loss_dict = { 'total': total_loss, 'joint': joint_loss, 's2m': s2m_dist, 'm2s': m2s_dist, 'pprior': pprior_loss, 'shape': shape_loss, 'angle_prior': angle_prior_loss, 'pen': pen_loss, 'sdf_penetration': sdf_penetration_loss, 'contact': contact_loss, 'physical': physical_loss } global_vars.cur_loss_dict = dict() for key, item in loss_dict.items(): if torch.is_tensor(item): global_vars.cur_loss_dict[key] = item.detach().cpu().numpy() else: global_vars.cur_loss_dict[key] = item if visualize: np.set_printoptions(precision=1) for key, value in global_vars.cur_loss_dict.items(): if isinstance(value, np.ndarray): if value.sum() == 0: continue else: if value == 0: continue print('{}:{}'.format(key, global_vars.cur_loss_dict[key]), end=' ') print('max_joint:{}'.format(global_vars.cur_max_joint)) return total_loss.sum()
def estimate_pointcloud_local_coord_frames( pointclouds: Union[torch.Tensor, Pointclouds], neighborhood_size: int = 50, disambiguate_directions: bool = True, return_knn_result: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor, Optional['KNN']]: """ Faster version of pytorch3d estimate_pointcloud_local_coord_frames Estimates the principal directions of curvature (which includes normals) of a batch of `pointclouds`. Returns: curvatures (N,P,3) ascending order local_frames (N,P,3,3) corresponding eigenvectors """ points_padded, num_points = convert_pointclouds_to_tensor(pointclouds) ba, N, dim = points_padded.shape if dim != 3: raise ValueError( "The pointclouds argument has to be of shape (minibatch, N, 3)") if (num_points <= neighborhood_size).any(): raise ValueError("The neighborhood_size argument has to be" + " >= size of each of the point clouds.") # undo global mean for stability # TODO: replace with tutil.wmean once landed pcl_mean = points_padded.sum(1) / num_points[:, None] points_centered = points_padded - pcl_mean[:, None, :] # get K nearest neighbor idx for each point in the point cloud knn_result = knn_points( points_padded, points_padded, lengths1=num_points, lengths2=num_points, K=neighborhood_size, return_nn=True, ) k_nearest_neighbors = knn_result.knn # obtain the mean of the neighborhood pt_mean = k_nearest_neighbors.mean(2, keepdim=True) # compute the diff of the neighborhood and the mean of the neighborhood # N,P,K,3 central_diff = k_nearest_neighbors - pt_mean per_pts_diff = central_diff.view(-1, neighborhood_size, 3) # S (NP,3) and local_coord_framds (NP,3,3) _, S, local_coord_frames = batch_svd(per_pts_diff) curvature = S * S / neighborhood_size local_coord_frames = local_coord_frames.view(ba, N, dim, dim) curvature = curvature.view(ba, N, dim) # flip to ascending order curvature = curvature.flip(-1) local_coord_frames = local_coord_frames.flip(-1) # disambiguate the directions of individual principal vectors if disambiguate_directions: # disambiguate normal n = _disambiguate_vector_directions(points_centered, k_nearest_neighbors, local_coord_frames[:, :, :, 0]) # disambiguate the main curvature z = _disambiguate_vector_directions(points_centered, k_nearest_neighbors, local_coord_frames[:, :, :, 2]) # the secondary curvature is just a cross between n and z y = torch.cross(n, z, dim=2) # cat to form the set of principal directions local_coord_frames = torch.stack((n, y, z), dim=3) if return_knn_result: return curvature, local_coord_frames, knn_result return curvature, local_coord_frames