def test_invalid_inputs_shapes(self, device="cuda:0"): with self.assertRaisesRegex( ValueError, "input can only be 2-dimensional." ): values = torch.rand((100, 50, 2), device=device) first_idxs = torch.tensor([0, 80], dtype=torch.int64, device=device) packed_to_padded(values, first_idxs, 100) with self.assertRaisesRegex( ValueError, "input can only be 3-dimensional." ): values = torch.rand((100,), device=device) first_idxs = torch.tensor([0, 80], dtype=torch.int64, device=device) padded_to_packed(values, first_idxs, 20) with self.assertRaisesRegex( ValueError, "input can only be 3-dimensional." ): values = torch.rand((100, 50, 2, 2), device=device) first_idxs = torch.tensor([0, 80], dtype=torch.int64, device=device) padded_to_packed(values, first_idxs, 20)
def _test_padded_to_packed_helper(self, D, device): """ Check the results from packed_to_padded and PyTorch implementations are the same. """ meshes = self.init_meshes(16, 100, 300, device=device) mesh_to_faces_packed_first_idx = meshes.mesh_to_faces_packed_first_idx() num_faces_per_mesh = meshes.num_faces_per_mesh() max_faces = num_faces_per_mesh.max().item() if D == 0: values = torch.rand((len(meshes), max_faces), device=device) else: values = torch.rand((len(meshes), max_faces, D), device=device) for i, num in enumerate(num_faces_per_mesh): values[i, num:] = 0 values.requires_grad = True values_torch = values.detach().clone() values_torch.requires_grad = True values_packed = padded_to_packed( values, mesh_to_faces_packed_first_idx, num_faces_per_mesh.sum().item(), ) values_packed_torch = TestPackedToPadded.padded_to_packed_python( values_torch, mesh_to_faces_packed_first_idx, num_faces_per_mesh.sum().item(), device, ) # check forward self.assertClose(values_packed, values_packed_torch) # check backward if D == 0: grad_inputs = torch.rand( (num_faces_per_mesh.sum().item()), device=device ) else: grad_inputs = torch.rand( (num_faces_per_mesh.sum().item(), D), device=device ) values_packed.backward(grad_inputs) grad_outputs = values.grad values_packed_torch.backward(grad_inputs) grad_outputs_torch1 = values_torch.grad grad_outputs_torch2 = TestPackedToPadded.packed_to_padded_python( grad_inputs, mesh_to_faces_packed_first_idx, values.size(1), device=device, ) self.assertClose(grad_outputs, grad_outputs_torch1) self.assertClose(grad_outputs, grad_outputs_torch2)
def compute(self, pointclouds: PointClouds3D, neighborhood_size=None): if neighborhood_size is None: neighborhood_size = self.neighborhood_size num_points = pointclouds.num_points_per_cloud() normals_packed = pointclouds.normals_packed() assert (normals_packed is not None) normals_ref = estimate_pointcloud_normals( pointclouds, neighborhood_size=neighborhood_size) normals_ref = padded_to_packed(normals_ref, pointclouds.cloud_to_packed_first_idx(), num_points.sum().item()) return 1 - F.cosine_similarity(normals_packed, normals_ref).abs()
def _compute_anisotropic_Vrk(self, pointclouds, **kwargs): """ determine the variance in the local surface frame h * Sk.T @ Sk based on curvature. Args: points_packed (Pointclouds3D): point clouds in object coordinates Returns: Vr (N, 3, 3): V_k^r matrix packed """ with torch.autograd.enable_grad(): points_padded = pointclouds.points_padded() num_points = pointclouds.num_points_per_cloud() # eigenvectors (=principal directions) in an ascending order of their # corresponding eigenvalues, while the smallest eigenvalue's eigenvector # corresponds to the normal direction curvatures, local_frame = estimate_pointcloud_local_coord_frames( pointclouds, neighborhood_size=8, disambiguate_directions=False) local_frame = ops3d.padded_to_packed( local_frame.reshape(local_frame.shape[:2] + (-1, )), pointclouds.cloud_to_packed_first_idx(), num_points.sum().item()) curvatures = ops3d.padded_to_packed( curvatures, pointclouds.cloud_to_packed_first_idx(), num_points.sum().item()) local_frame = local_frame.view(-1, 3, 3)[:, :, 1:] # curvature only determines ratio of the two principle axis # the actual size is based on a global max_size curvatures = curvatures.view(-1, 3)[:, 1:] curvature_ratios = curvatures / curvatures[:, -1:] # TODO: compute density curvatures = curvatures Vr = local_frame @ torch.diag_embed( curvatures) @ local_frame.transpose(1, 2) return Vr, local_frame.transpose(1, 2)
def _denoise_normals(self, point_clouds, weights, point_clouds_filter=None): """ robust normal mollification (Sec 4.4), i.e. replace normals with a weighted average from neighboring normals do this only for invisible points (?) Args: weights (tensors): (N,max_P,K) """ lengths = point_clouds.num_points_per_cloud() P_total = lengths.sum().item() normals = point_clouds.normals_padded() batch_size, max_P, _ = normals.shape knn_normals = ops.knn_gather(normals, self.knn_tree.idx, lengths) normals_denoised = torch.sum(knn_normals * weights[:, :, :, None], dim=-2) / \ eps_denom(torch.sum(weights, dim=-1, keepdim=True)) # get point visibility so that we update only the non-visible or out-of-mask normals if point_clouds_filter is not None: try: reliable_normals_mask = point_clouds_filter.visibility & point_clouds_filter.inmask if len(point_clouds) != reliable_normals_mask.shape[0]: if len(point_clouds ) == 1 and reliable_normals_mask.shape[0] > 1: reliable_normals_mask = reliable_normals_mask.any( dim=0, keepdim=True) else: ValueError( "Incompatible point clouds {} and mask {}".format( len(point_clouds), reliable_normals_mask.shape)) # found visibility 0/1 as the last dimension of the features # reset visible points normals to its original ones normals_denoised[pts_reliable_normals_mask == 1] = normals[ reliable_normals_mask == 1] except KeyError as e: pass normals_packed = point_clouds.normals_packed() normals_denoised_packed = ops.padded_to_packed( normals_denoised, point_clouds.cloud_to_packed_first_idx(), P_total) point_clouds.update_normals_(normals_denoised_packed) return point_clouds
def _filter_points_with_invalid_depth(self, point_clouds, **kwargs): self.cameras = kwargs.get('cameras', self.cameras) lengths = point_clouds.num_points_per_cloud() points_padded = point_clouds.points_padded() with torch.autograd.no_grad(): to_view = self.cameras.get_world_to_view_transform() points = to_view.transform_points(points_padded) znear = getattr(self.cameras, 'znear', kwargs.get('znear', 1.0)) zfar = getattr(self.cameras, 'zfar', kwargs.get('zfar', 100.0)) mask = (points[..., 2] >= znear) & (points[..., 2] <= zfar) mask_packed = ops3d.padded_to_packed( mask.float(), point_clouds.cloud_to_packed_first_idx(), lengths.sum().item()).bool() if torch.all(mask_packed): return point_clouds, mask_packed # create point clouds again from packed points_padded = point_clouds.points_padded() normals_padded = point_clouds.normals_padded() features_padded = point_clouds.features_padded() # tuple to list, since pointclouds class doesn't accept tuples points_list = [ points_padded[b][mask[b]] for b in range(points_padded.shape[0]) ] normals_list = features_list = None if normals_padded is not None: normals_list = [ normals_padded[b][mask[b]] for b in range(normals_padded.shape[0]) ] if features_padded is not None: features_list = [ features_padded[b][mask[b]] for b in range(features_padded.shape[0]) ] new_point_clouds = point_clouds.__class__(points=points_list, normals=normals_list, features=features_list) # for k in per_point_info: # per_point_info[k] = per_point_info[k][mask_packed] return new_point_clouds, mask_packed
def _filter_backface_points(self, point_clouds, **kwargs): with torch.autograd.no_grad(): normals_view = self.transform_normals(point_clouds, **kwargs) # mask = (normals_view[:, :, 2] < 1e-3) # error buffer mask = normals_view[:, :, 2] < 0 lengths = point_clouds.num_points_per_cloud() mask_packed = ops3d.padded_to_packed( mask.float(), point_clouds.cloud_to_packed_first_idx(), lengths.sum().item()).bool() if torch.all(mask_packed): return point_clouds, mask_packed # create point clouds again from packed points_padded = point_clouds.points_padded() normals_padded = point_clouds.normals_padded() features_padded = point_clouds.features_padded() # tuple to list, since pointclouds class doesn't accept tuples points_list = [ points_padded[b][mask[b]] for b in range(points_padded.shape[0]) ] normals_list = [ normals_padded[b][mask[b]] for b in range(normals_padded.shape[0]) ] if features_padded is not None: features_list = [ features_padded[b][mask[b]] for b in range(features_padded.shape[0]) ] new_point_clouds = point_clouds.__class__(points=points_list, normals=normals_list, features=features_list) else: new_point_clouds = point_clouds.__class__(points=points_list, normals=normals_list) # for k in per_point_info: # per_point_info[k] = per_point_info[k][mask_packed] return new_point_clouds, mask_packed
def backward(ctx, idx_grad, zbuf_grad, qvalue_grad, occ_grad): # idx_grad and zbuf_grad are None (unless maybe we make weights depend on z? i.e. volumetric splatting) grad_radii = None grad_cloud_to_packed_first_idx = None grad_num_points_per_cloud = None grad_cutoff_thres = None grad_depth_merging_thres = None grad_image_size = None grad_points_per_pixel = None grad_bin_size = None grad_max_points_per_bin = None grad_radii_s = None grad_backward_rbf = None grads = (grad_cutoff_thres, grad_radii, grad_cloud_to_packed_first_idx, grad_num_points_per_cloud, grad_depth_merging_thres, grad_image_size, grad_points_per_pixel, grad_bin_size, grad_max_points_per_bin, grad_radii_s, grad_backward_rbf) radii_s = ctx.radii_backward_scaler # either use OccRBFBackward or use OccBackward pts_screen, ellipse_param, cutoff_threshold, radii, idx, zbuf0, \ cloud_to_packed_first_idx, num_points_per_cloud, \ = ctx.saved_tensors depth_merging_threshold = ctx.depth_merging_threshold backward_occ_fast = True if not backward_occ_fast: device = pts_screen.device grads_input_xy = pts_screen.new_zeros((pts_screen.shape[0], 2)) grads_input_z = pts_screen.new_zeros((pts_screen.shape[0], 1)) mask = (idx[..., 0] >= 0).bool() # float pts_visibility = torch.full((pts_screen.shape[0], ), False, dtype=torch.bool, device=pts_screen.device) # all rendered points (indices in packed points) visible_idx = idx[mask].unique().long().view(-1) visible_idx = visible_idx[visible_idx >= 0] pts_visibility[visible_idx] = True num_points_per_cloud = torch.stack([ x.sum() for x in torch.split( pts_visibility, num_points_per_cloud.tolist(), dim=0) ]) cloud_to_packed_first_idx = num_points_2_cloud_to_packed_first_idx( num_points_per_cloud) pts_screen = pts_screen[pts_visibility] radii = radii[pts_visibility] grad_visible = _C._splat_points_occ_backward( pts_screen, radii, occ_grad, cloud_to_packed_first_idx, num_points_per_cloud, radii_s, depth_merging_threshold) if torch.isnan(grad_visible).any( ) or not torch.isfinite(grad_visible).all(): print('invalid grad_visible') assert (pts_visibility.sum() == grad_visible.shape[0]) grads_input_xy[pts_visibility] = grad_visible _C._backward_zbuf(idx, zbuf_grad, grads_input_z) # TODO necessary to concatenate grads_input = torch.cat([grads_input_xy, grads_input_z], dim=-1) else: """ We only care about rasterized points (visible points) 1. Filter [P,*] data to [P_visible,*] data 2. Fast backward cuda 2a. call FRNN insertion 2b. count_sort """ device = pts_screen.device mask = (idx[..., 0] >= 0).bool() # float pts_visibility = torch.full((pts_screen.shape[0], ), False, dtype=torch.bool, device=pts_screen.device) # all rendered points (indices in packed points) visible_idx = idx[mask].unique().long().view(-1) visible_idx = visible_idx[visible_idx >= 0] pts_visibility[visible_idx] = True num_points_per_cloud = torch.stack([ x.sum() for x in torch.split( pts_visibility, num_points_per_cloud.tolist(), dim=0) ]) cloud_to_packed_first_idx = num_points_2_cloud_to_packed_first_idx( num_points_per_cloud) pts_screen_visible = pts_screen[pts_visibility] radii_visible = radii[pts_visibility] ##################################### # 2a. call FRNN insertion ##################################### N = num_points_per_cloud.shape[0] P = pts_screen_visible.shape[0] assert (num_points_per_cloud.sum().item() == P) # from frnn.frnn import GRID_PARAMS_SIZE, MAX_RES, prefix_sum_cuda # imported from from prefix_sum import prefix_sum_cuda GRID_2D_PARAMS_SIZE = 6 GRID_2D_MAX_RES = 1024 GRID_2D_DELTA = 2 GRID_2D_TOTAL = 5 RADIUS_CELL_RATIO = 2 # first convert to padded max_P = num_points_per_cloud.max().item() pts_padded = ops3d.packed_to_padded(pts_screen_visible, cloud_to_packed_first_idx, max_P) radii_padded = ops3d.packed_to_padded(radii_visible, cloud_to_packed_first_idx, max_P) # determine search radius as max(radii)*radii_s search_radius = torch.tensor([ radii_padded[i, :num_points_per_cloud[i]].median() * radii_s for i in range(N) ], dtype=torch.float, device=device) # create grid from scratch # setup grid params grid_params_cuda = torch.zeros((N, GRID_2D_PARAMS_SIZE), dtype=torch.float, device=pts_padded.device) G = -1 pts_padded_2D = pts_padded[:, :, :2].clone().contiguous() for i in range(N): # 0-2 grid_min; 3 grid_delta; 4-6 grid_res; 7 grid_total grid_min = pts_padded_2D[i, :num_points_per_cloud[i]].min( dim=0)[0] grid_max = pts_padded_2D[i, :num_points_per_cloud[i]].max( dim=0)[0] grid_params_cuda[i, :GRID_2D_DELTA] = grid_min grid_size = grid_max - grid_min cell_size = search_radius[i].item() / RADIUS_CELL_RATIO if cell_size < grid_size.min() / GRID_2D_MAX_RES: cell_size = grid_size.min() / GRID_2D_MAX_RES grid_params_cuda[i, GRID_2D_DELTA] = 1 / cell_size grid_params_cuda[i, GRID_2D_DELTA + 1:GRID_2D_TOTAL] = torch.floor( grid_size / cell_size) + 1 grid_params_cuda[i, GRID_2D_TOTAL] = torch.prod( grid_params_cuda[i, GRID_2D_DELTA + 1:GRID_2D_TOTAL]) if G < grid_params_cuda[i, GRID_2D_TOTAL]: G = int(grid_params_cuda[i, GRID_2D_TOTAL].item()) # insert points into the grid pc_grid_cnt = torch.zeros((N, G), dtype=torch.int, device=device) pc_grid_cell = torch.full((N, max_P), -1, dtype=torch.int, device=device) pc_grid_idx = torch.full((N, max_P), -1, dtype=torch.int, device=device) frnn._C.insert_points_cuda(pts_padded_2D, num_points_per_cloud, grid_params_cuda, pc_grid_cnt, pc_grid_cell, pc_grid_idx, G) # use prefix_sum from Matt Dean grid_params = grid_params_cuda.cpu() pc_grid_off = torch.full((N, G), 0, dtype=torch.int, device=device) for i in range(N): prefix_sum_cuda(pc_grid_cnt[i], grid_params[i, GRID_2D_TOTAL], pc_grid_off[i]) # sort points according to their grid positions and insertion orders # sort based on x, y first. Then we will use points_sorted_idxs to recover the points_sorted with Z points_sorted = torch.zeros((N, max_P, 2), dtype=torch.float, device=device) points_sorted_idxs = torch.full((N, max_P), -1, dtype=torch.int, device=device) frnn._C.counting_sort_cuda( pts_padded_2D, num_points_per_cloud, pc_grid_cell, pc_grid_idx, pc_grid_off, points_sorted, # (N,P,2) points_sorted_idxs # (N,P) ) new_points_sorted = torch.zeros_like(pts_padded) for i in range(N): points_sorted_idxs_i = points_sorted_idxs[ i, :num_points_per_cloud[i]].long().unsqueeze(1).expand( -1, 3) new_points_sorted[i, :num_points_per_cloud[i]] = torch.gather( pts_padded[i], 0, points_sorted_idxs_i) # print(points_sorted[i, :10]) # print(new_points_sorted[i, :10]) # new_points_sorted = torch.gather(pts_padded, 1, points_sorted_idxs.long().unsqueeze(2).expand(-1, -1, 3)) assert (new_points_sorted is not None and pc_grid_off is not None and points_sorted_idxs is not None and grid_params_cuda is not None) # convert sorted_points and sorted_points_idxs to packed (P, ) points_sorted = ops3d.padded_to_packed(new_points_sorted, cloud_to_packed_first_idx, P) # padded_to_packed only supports torch.float32... shifted_points_sorted_idxs = points_sorted_idxs + cloud_to_packed_first_idx.float( ).unsqueeze(1) points_sorted_idxs = ops3d.padded_to_packed( shifted_points_sorted_idxs, cloud_to_packed_first_idx, P) points_sorted_idxs_2D = points_sorted_idxs.long().unsqueeze( 1).expand(-1, 2) radii_sorted = torch.gather(radii_visible, 0, points_sorted_idxs_2D) pc_grid_off += cloud_to_packed_first_idx.unsqueeze(1) grad_sorted = _C._splat_points_occ_fast_cuda_backward( points_sorted, radii_sorted, search_radius, occ_grad, num_points_per_cloud, cloud_to_packed_first_idx, pc_grid_off, grid_params_cuda) # grad_sorted_slow = _C._splat_points_occ_backward(points_sorted, radii_sorted, # occ_grad, cloud_to_packed_first_idx, num_points_per_cloud, # radii_s, depth_merging_threshold) # breakpoint() # points_sorted_idxs_3D = points_sorted_idxs.long().unsqueeze(1).expand(-1, 3) # print(points_sorted_idxs_3D.max(), grad_sorted.shape[0]) grad_visible = torch.zeros_like(grad_sorted).scatter_( 0, points_sorted_idxs_2D, grad_sorted) # grad_visible_slow = _C._splat_points_occ_backward(pts_screen[pts_visibility], radii[pts_visibility], # occ_grad, cloud_to_packed_first_idx, num_points_per_cloud, # radii_s, depth_merging_threshold) # breakpoint() if torch.isnan(grad_visible).any( ) or not torch.isfinite(grad_visible).all(): print('invalid grad_visible') assert (pts_visibility.sum() == grad_visible.shape[0]) grads_input_xy = pts_screen.new_zeros(pts_screen.shape[0], 2) grads_input_z = pts_screen.new_zeros(pts_screen.shape[0], 1) # print("1") grads_input_xy[pts_visibility] = grad_visible _C._backward_zbuf(idx, zbuf_grad, grads_input_z) grads_input = torch.cat([grads_input_xy, grads_input_z], dim=-1) # print("2") pts_grad = grads_input return (pts_grad, None) + grads
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 compute(self, point_clouds: PointClouds3D, points_filters=None, rebuild_knn=True, **kwargs): self.knn_tree = kwargs.get('knn_tree', self.knn_tree) self.knn_mask = kwargs.get('knn_mask', self.knn_mask) lengths = point_clouds.num_points_per_cloud() P_total = lengths.sum().item() points_padded = point_clouds.points_padded() # Compute necessary weights to project points to local plane # TODO(yifan): This part is same as ProjectionLoss # how can we at best save repetitive computation with torch.autograd.no_grad(): if rebuild_knn or self.knn_tree is None or points_padded.shape[: 2] != self.knn_tree.shape[: 2]: self._build_knn(point_clouds) phi = self.get_phi(point_clouds, **kwargs) self._denoise_normals(point_clouds, phi, points_filters) # compute wn and wr # TODO(yifan): visibility weight? normal_w = self.get_normal_w(point_clouds, **kwargs) # update normals for a second iteration (?) Eq.(10) point_clouds = self._denoise_normals(point_clouds, phi * normal_w, points_filters) # compose weights weights = phi * normal_w weights[~self.knn_mask] = 0 # outside filter_scale*local_point_spacing weights mask_ball_query = self.knn_tree.dists > ( self.filter_scale * self.knn_tree.dists[:, :, :1] * 2.0) weights[mask_ball_query] = 0.0 # project the point to a local surface knn_normals = ops.knn_gather(point_clouds.normals_padded(), self.knn_tree.idx, lengths) dist_to_surface = torch.sum( (self.knn_tree.knn.detach() - points_padded.unsqueeze(-2)) * knn_normals, dim=-1) deltap = torch.sum( dist_to_surface[..., None] * weights[..., None] * knn_normals, dim=-2) / eps_denom(torch.sum(weights, dim=-1, keepdim=True)) points_projected = points_padded + deltap if get_debugging_mode(): # points_padded.requires_grad_(True) def save_grad(): lengths = point_clouds.num_points_per_cloud() def _save_grad(grad): dbg_tensor = get_debugging_tensor() if dbg_tensor is None: logger_py.error("dbg_tensor is None") if grad is None: logger_py.error('grad is None') # a dict of list of tensors dbg_tensor.pts_world_grad['repel'] = [ grad[b, :lengths[b]].detach().cpu() for b in range(grad.shape[0]) ] return _save_grad dbg_tensor = get_debugging_tensor() dbg_tensor.pts_world['repel'] = [ points_padded[b, :lengths[b]].detach().cpu() for b in range(points_padded.shape[0]) ] handle = points_padded.register_hook(save_grad()) self.hooks.append(handle) with torch.autograd.no_grad(): spatial_w = self.get_spatial_w(point_clouds, points_projected) # density_w = self.get_density_w(point_clouds) # density weight is actually spatial_w + 1 density_w = torch.sum(spatial_w, dim=-1, keepdim=True) + 1.0 weights = normal_w * spatial_w * density_w weights[~self.knn_mask] = 0 weights[mask_ball_query] = 0 deltap = points_projected[:, :, None, :] - self.knn_tree.knn.detach() point_to_point_dist = torch.sum(deltap * deltap, dim=-1) # convert everything to packed weights = ops.padded_to_packed( weights, point_clouds.cloud_to_packed_first_idx(), P_total) point_to_point_dist = ops.padded_to_packed( point_to_point_dist, point_clouds.cloud_to_packed_first_idx(), P_total) # we want to maximize this, so negative sign point_to_point_dist = -torch.sum(point_to_point_dist * weights, dim=1) / eps_denom( torch.sum(weights, dim=1)) return point_to_point_dist
def compute(self, point_clouds: PointClouds3D, points_filters=None, rebuild_knn=False, **kwargs): """ Args: point_clouds (optional) knn_tree: output from ops.knn_points excluding the query point itself (optional) knn_mask: mask valid knn results Returns: (P, N) """ self.sharpness_sigma = kwargs.get('sharpness_sigma', self.sharpness_sigma) self.filter_scale = kwargs.get('filter_scale', self.filter_scale) self.knn_tree = kwargs.get('knn_tree', self.knn_tree) self.knn_mask = kwargs.get('knn_mask', self.knn_mask) lengths = point_clouds.num_points_per_cloud() P_total = lengths.sum().item() points = point_clouds.points_padded() # - determine phi spatial with using local point spacing (i.e. 2*dist_to_nn) # - denoise normals # - determine w_normal # - mask out values outside ballneighbor i.e. d > filterSpatialScale * localPointSpacing # - projected distance dot(ni, x-xi) # - multiply and normalize the weights with torch.autograd.no_grad(): if rebuild_knn or self.knn_tree is None or self.knn_tree.idx.shape[: 2] != points.shape[: 2]: self._build_knn(point_clouds) phi = self.get_phi(point_clouds, **kwargs) # robust normal mollification (Sec 4.4), i.e. replace normals with a weighted average # from neighboring normals Eq.(11) point_clouds = self._denoise_normals(point_clouds, phi, points_filters) # compute wn and wr # TODO(yifan): visibility weight? normal_w = self.get_normal_w(point_clouds, **kwargs) spatial_w = self.get_spatial_w(point_clouds, **kwargs) # update normals for a second iteration (?) Eq.(10) point_clouds = self._denoise_normals(point_clouds, phi * normal_w, points_filters) # compose weights weights = phi * spatial_w * normal_w weights[~self.knn_mask] = 0 # outside filter_scale*local_point_spacing weights mask_ball_query = self.knn_tree.dists > ( self.filter_scale * self.knn_tree.dists[:, :, :1] * 2.0) weights[mask_ball_query] = 0.0 # (B, P, k), dot product distance to surface # (we need to gather again because the normals have been changed in the denoising step) knn_normals = ops.knn_gather(point_clouds.normals_padded(), self.knn_tree.idx, lengths) # if points.requires_grad: # from DSS.core.rasterizer import _dbg_tensor # def save_grad(name): # def _save_grad(grad): # _dbg_tensor[name] = grad.detach().cpu() # return _save_grad # points.register_hook(save_grad('proj_grad')) dist_to_surface = torch.sum( (self.knn_tree.knn.detach() - points.unsqueeze(-2)) * knn_normals, dim=-1) if get_debugging_mode(): # points.requires_grad_(True) def save_grad(): lengths = point_clouds.num_points_per_cloud() def _save_grad(grad): dbg_tensor = get_debugging_tensor() if dbg_tensor is None: logger_py.error("dbg_tensor is None") if grad is None: logger_py.error('grad is None') # a dict of list of tensors dbg_tensor.pts_world_grad['proj'] = [ grad[b, :lengths[b]].detach().cpu() for b in range(grad.shape[0]) ] return _save_grad dbg_tensor = get_debugging_tensor() dbg_tensor.pts_world['proj'] = [ points[b, :lengths[b]].detach().cpu() for b in range(points.shape[0]) ] handle = points.register_hook(save_grad()) self.hooks.append(handle) # convert everything to packed weights = ops.padded_to_packed( weights, point_clouds.cloud_to_packed_first_idx(), P_total) dist_to_surface = ops.padded_to_packed( dist_to_surface, point_clouds.cloud_to_packed_first_idx(), P_total) # compute weighted signed distance to surface dist_to_surface = torch.sum(weights * dist_to_surface, dim=-1) / eps_denom( torch.sum(weights, dim=-1)) loss = dist_to_surface * dist_to_surface return loss
def compute(self, point_clouds: PointClouds3D, points_filter=None, rebuild_knn=True, **kwargs): self.knn_tree = kwargs.get('knn_tree', self.knn_tree) self.knn_mask = kwargs.get('knn_mask', self.knn_mask) lengths = point_clouds.num_points_per_cloud() P_total = lengths.sum().item() points_padded = point_clouds.points_padded() if not points_padded.requires_grad: logger_py.warn( 'Computing repulsion loss, but points_padded is not differentiable.' ) # Compute necessary weights to project points to local plane # TODO(yifan): This part is same as ProjectionLoss # how can we at best save repetitive computation with torch.autograd.no_grad(): if rebuild_knn or self.knn_tree is None or points_padded.shape[: 2] != self.knn_tree.shape[: 2]: self._build_knn(point_clouds) phi = self.get_phi(point_clouds, **kwargs) point_clouds = self._denoise_normals(point_clouds, phi, points_filter, inplace=False) # project the point to a local surface knn_diff = points_padded.unsqueeze(-2) - self.knn_tree.knn.detach() knn_normals = ops.knn_gather(point_clouds.normals_padded(), self.knn_tree.idx, lengths) pts_diff_proj = knn_diff - \ (knn_diff * knn_normals).sum(dim=-1, keepdim=True) * knn_normals if get_debugging_mode(): # points_padded.requires_grad_(True) def save_grad(): lengths = point_clouds.num_points_per_cloud() def _save_grad(grad): dbg_tensor = get_debugging_tensor() if dbg_tensor is None: logger_py.error("dbg_tensor is None") if grad is None: logger_py.error('grad is None') # a dict of list of tensors dbg_tensor.pts_world_grad['repel'] = [ grad[b, :lengths[b]].detach().cpu() for b in range(grad.shape[0]) ] return _save_grad if points_padded.requires_grad: dbg_tensor = get_debugging_tensor() dbg_tensor.pts_world['repel'] = [ points_padded[b, :lengths[b]].detach().cpu() for b in range(points_padded.shape[0]) ] handle = points_padded.register_hook(save_grad()) self.hooks.append(handle) with torch.autograd.no_grad(): spatial_w = self.get_spatial_w(point_clouds, **kwargs) # set far neighbors' spatial_w to 0 normal_w = self.get_normal_w(point_clouds, **kwargs) density_w = torch.sum(spatial_w, dim=-1, keepdim=True) + 1.0 weights = spatial_w * normal_w # convert everything to packed weights = ops.padded_to_packed( weights, point_clouds.cloud_to_packed_first_idx(), P_total) pts_diff_proj = ops.padded_to_packed( pts_diff_proj.contiguous().view(pts_diff_proj.shape[0], pts_diff_proj.shape[1], -1), point_clouds.cloud_to_packed_first_idx(), P_total).view(P_total, -1, 3) density_w = ops.padded_to_packed( density_w, point_clouds.cloud_to_packed_first_idx(), P_total) # we want to maximize this, so negative sign repel_vec = torch.sum(pts_diff_proj * weights.unsqueeze(-1), dim=1) / eps_denom( torch.sum(weights, dim=1).unsqueeze(-1)) repel_vec = repel_vec * density_w loss = torch.exp(-repel_vec.abs()) # if get_debugging_mode(): # # save to dbg folder as normal # from ..utils.io import save_ply # save_ply('./dbg_repel_diff.ply', point_clouds.points_packed().cpu().detach(), normals=repel_vec.cpu().detach()) return loss
def compute(self, point_clouds: PointClouds3D, points_filter=None, rebuild_knn=False, **kwargs): """ Args: point_clouds (optional) knn_tree: output from ops.knn_points excluding the query point itself (optional) knn_mask: mask valid knn results Returns: (P, N) """ self.sharpness_sigma = kwargs.get('sharpness_sigma', self.sharpness_sigma) self.filter_scale = kwargs.get('filter_scale', self.filter_scale) self.knn_tree = kwargs.get('knn_tree', self.knn_tree) self.knn_mask = kwargs.get('knn_mask', self.knn_mask) lengths = point_clouds.num_points_per_cloud() P_total = lengths.sum().item() points = point_clouds.points_padded() # - determine phi spatial with using local point spacing (i.e. 2*dist_to_nn) # - denoise normals # - determine w_normal # - mask out values outside ballneighbor i.e. d > filterSpatialScale * localPointSpacing # - projected distance dot(ni, x-xi) # - multiply and normalize the weights with torch.autograd.no_grad(): if rebuild_knn or self.knn_tree is None or self.knn_tree.idx.shape[: 2] != points.shape[: 2]: self._build_knn(point_clouds) phi = self.get_phi(point_clouds, **kwargs) # robust normal mollification (Sec 4.4), i.e. replace normals with a weighted average # from neighboring normals Eq.(11) point_clouds = self._denoise_normals(point_clouds, phi, points_filter, inplace=False) # compute wn and wr normal_w = self.get_normal_w(point_clouds, **kwargs) # visibility weight visibility_nb = ops.knn_gather( points_filter.visibility.unsqueeze(-1), self.knn_tree.idx, lengths) visibility_w = visibility_nb.float() visibility_w[~visibility_nb] = 0.1 # compose weights weights = phi * normal_w * visibility_w.squeeze(-1) # (B, P, k), dot product distance to surface knn_normals = ops.knn_gather(point_clouds.normals_padded(), self.knn_tree.idx, lengths) if get_debugging_mode(): # points.requires_grad_(True) def save_grad(): lengths = point_clouds.num_points_per_cloud() def _save_grad(grad): dbg_tensor = get_debugging_tensor() if dbg_tensor is None: logger_py.error("dbg_tensor is None") if grad is None: logger_py.error('grad is None') # a dict of list of tensors dbg_tensor.pts_world_grad['proj'] = [ grad[b, :lengths[b]].detach().cpu() for b in range(grad.shape[0]) ] return _save_grad if points.requires_grad: dbg_tensor = get_debugging_tensor() dbg_tensor.pts_world['proj'] = [ points[b, :lengths[b]].detach().cpu() for b in range(points.shape[0]) ] handle = points.register_hook(save_grad()) self.hooks.append(handle) sdf = torch.sum( (self.knn_tree.knn.detach() - points.unsqueeze(-2)) * knn_normals, dim=-1) # convert everything to packed weights = ops.padded_to_packed( weights, point_clouds.cloud_to_packed_first_idx(), P_total) sdf = ops.padded_to_packed(sdf, point_clouds.cloud_to_packed_first_idx(), P_total) # if get_debugging_mode(): # # save to dbg folder as normal # from ..utils.io import save_ply # save_ply('./dbg_repel_diff.ply', point_clouds.points_packed().cpu().detach(), normals=repel_vec.cpu().detach()) distance_to_face = sdf * sdf # compute weighted signed distance to surface loss = torch.sum(weights * distance_to_face, dim=-1) / eps_denom( torch.sum(weights, dim=-1)) return loss
def test_dataset(self): # 1. rerender input point clouds / meshes using the saved camera_mat # compare mask image with saved mask image # 2. backproject masked points to space with dense depth map, # fuse all views and save batch_size = 1 device = torch.device('cuda:0') data_dir = 'data/synthetic/cube_mesh' output_dir = os.path.join('tests', 'outputs', 'test_data') if not os.path.isdir(output_dir): os.makedirs(output_dir) # dataset dataset = MVRDataset(data_dir=data_dir, load_dense_depth=True, mode="train") data_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, num_workers=0, shuffle=False) meshes = load_objs_as_meshes([os.path.join(data_dir, 'mesh.obj')]).to(device) cams = dataset.get_cameras().to(device) image_size = imageio.imread(dataset.image_files[0]).shape[0] # initialize rasterizer, we check mask pngs only, so no need to create lights and shaders etc raster_settings = RasterizationSettings( image_size=image_size, blur_radius=0.0, faces_per_pixel=5, bin_size= None, # this setting controls whether naive or coarse-to-fine rasterization is used max_faces_per_bin=None # this setting is for coarse rasterization ) rasterizer = MeshRasterizer(cameras=None, raster_settings=raster_settings) # render with loaded cameras positions and training tranformation functions pixel_world_all = [] for idx, data in enumerate(data_loader): # get datas img = data.get('img.rgb').to(device) assert (img.min() >= 0 and img.max() <= 1 ), "Image must be a floating number between 0 and 1." mask_gt = data.get('img.mask').to(device).permute(0, 2, 3, 1) camera_mat = data['camera_mat'].to(device) cams.R, cams.T = decompose_to_R_and_t(camera_mat) cams._N = cams.R.shape[0] cams.to(device) self.assertTrue( torch.equal(cams.get_world_to_view_transform().get_matrix(), camera_mat)) # transform to view and rerender with non-rotated camera verts_padded = transform_to_camera_space(meshes.verts_padded(), cams) meshes_in_view = meshes.offset_verts( -meshes.verts_packed() + padded_to_packed( verts_padded, meshes.mesh_to_verts_packed_first_idx(), meshes.verts_packed().shape[0])) fragments = rasterizer(meshes_in_view, cameras=dataset.get_cameras().to(device)) # compare mask mask = fragments.pix_to_face[..., :1] >= 0 imageio.imwrite(os.path.join(output_dir, "mask_%06d.png" % idx), mask[0, ...].cpu().to(dtype=torch.uint8) * 255) # allow 5 pixels difference self.assertTrue(torch.sum(mask_gt != mask) < 5) # check dense maps # backproject points to the world pixel range (-1, 1) pixels = arange_pixels((image_size, image_size), batch_size)[1].to(device) depth_img = data.get('img.depth').to(device) # get the depth and mask at the sampled pixel position depth_gt = get_tensor_values(depth_img, pixels, squeeze_channel_dim=True) mask_gt = get_tensor_values(mask.permute(0, 3, 1, 2).float(), pixels, squeeze_channel_dim=True).bool() # get pixels and depth inside the masked area pixels_packed = pixels[mask_gt] depth_gt_packed = depth_gt[mask_gt] first_idx = torch.zeros((pixels.shape[0], ), device=device, dtype=torch.long) num_pts_in_mask = mask_gt.sum(dim=1) first_idx[1:] = num_pts_in_mask.cumsum(dim=0)[:-1] pixels_padded = packed_to_padded(pixels_packed, first_idx, num_pts_in_mask.max().item()) depth_gt_padded = packed_to_padded(depth_gt_packed, first_idx, num_pts_in_mask.max().item()) # backproject to world coordinates # contains nan and infinite values due to depth_gt_padded containing 0.0 pixel_world_padded = transform_to_world(pixels_padded, depth_gt_padded[..., None], cams) # transform back to list, containing no padded values split_size = num_pts_in_mask[..., None].repeat(1, 2) split_size[:, 1] = 3 pixel_world_list = padded_to_list(pixel_world_padded, split_size) pixel_world_all.extend(pixel_world_list) idx += 1 if idx >= 10: break pixel_world_all = torch.cat(pixel_world_all, dim=0) mesh = trimesh.Trimesh(vertices=pixel_world_all.cpu(), faces=None, process=False) mesh.export(os.path.join(output_dir, 'pixel_to_world.ply'))