def test_save_load_with_normals(self): points = torch.tensor([[0, 0, 0], [0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=torch.float32) normals = torch.tensor([[0, 1, 0], [1, 0, 0], [1, 4, 1], [1, 0, 0]], dtype=torch.float32) features = torch.rand_like(points) for do_features, do_normals in itertools.product([True, False], [True, False]): cloud = Pointclouds( points=[points], features=[features] if do_features else None, normals=[normals] if do_normals else None, ) device = torch.device("cuda:0") io = IO() with NamedTemporaryFile(mode="w", suffix=".ply") as f: io.save_pointcloud(cloud.cuda(), f.name) f.flush() cloud2 = io.load_pointcloud(f.name, device=device) self.assertEqual(cloud2.device, device) cloud2 = cloud2.cpu() self.assertClose(cloud2.points_padded(), cloud.points_padded()) if do_normals: self.assertClose(cloud2.normals_padded(), cloud.normals_padded()) else: self.assertIsNone(cloud.normals_padded()) self.assertIsNone(cloud2.normals_padded()) if do_features: self.assertClose(cloud2.features_packed(), features) else: self.assertIsNone(cloud2.features_packed())
def eval_one_dir(exp_dir, n_pts=50000): """ Function for one directory """ device = torch.device('cuda:0') cfg = config.load_config(os.path.join(exp_dir, 'config.yaml')) dataset = config.create_dataset(cfg.data, mode='val') meshes_gt = dataset.get_meshes().to(device) val_gt_pts_file = os.path.join(cfg.data.data_dir, 'val%d.ply' % n_pts) if os.path.isfile(val_gt_pts_file): points, normals = np.split(read_ply(val_gt_pts_file), 2, axis=1) pcl_gt = Pointclouds( torch.from_numpy(points[None, ...]).float(), torch.from_numpy(normals[None, ...]).float()).to(device) else: pcl_gt = dataset.get_pointclouds(n_pts).to(device) trimesh.Trimesh(pcl_gt.points_packed().cpu().numpy(), vertex_normals=pcl_gt.normals_packed().cpu().numpy(), process=False).export(val_gt_pts_file, vertex_normal=True) # load vis directories vis_dir = os.path.join(exp_dir, 'vis') vis_files = sorted(get_filenames(vis_dir, '_mesh.ply')) iters = [int(os.path.basename(v).split('_')[0]) for v in vis_files] best_dict = defaultdict(lambda: float('inf')) vis_eval_csv = os.path.join(vis_dir, "evaluation_n%d.csv" % n_pts) if not os.path.isfile(vis_eval_csv): with open(os.path.join(vis_dir, "evaluation_n%d.csv" % n_pts), "w") as f: fieldnames = ['mtime', 'it', 'chamfer_p', 'chamfer_n', 'pf_dist'] writer = csv.DictWriter(f, fieldnames=fieldnames, restval="-", extrasaction="ignore") writer.writeheader() mtime0 = None for it, vis_file in zip(iters, vis_files): eval_dict = OrderedDict() mtime = os.path.getmtime(vis_file) if mtime0 is None: mtime0 = mtime eval_dict['it'] = it eval_dict['mtime'] = mtime - mtime0 val_pts_file = os.path.join( vis_dir, os.path.basename(vis_file).replace('_mesh', '_val%d' % n_pts)) if os.path.isfile(val_pts_file): points, normals = np.split(read_ply(val_pts_file), 2, axis=1) points = torch.from_numpy(points).float().to( device=device).view(1, -1, 3) normals = torch.from_numpy(normals).float().to( device=device).view(1, -1, 3) else: mesh = trimesh.load(vis_file, process=False) # points, normals = pcu.sample_mesh_poisson_disk( # mesh.vertices, mesh.faces, # mesh.vertex_normals.ravel().reshape(-1, 3), n_pts, use_geodesic_distance=True) # p_idx = np.random.permutation(points.shape[0])[:n_pts] # points = points[p_idx, ...] # normals = normals[p_idx, ...] # points = torch.from_numpy(points).float().to( # device=device).view(1, -1, 3) # normals = torch.from_numpy(normals).float().to( # device=device).view(1, -1, 3) meshes = Meshes( torch.from_numpy(mesh.vertices[None, ...]).float(), torch.from_numpy(mesh.faces[None, ...]).float()).to(device) points, normals = sample_points_from_meshes( meshes, n_pts, return_normals=True) trimesh.Trimesh(points.cpu().numpy()[0], vertex_normals=normals.cpu().numpy()[0], process=False).export(val_pts_file, vertex_normal=True) pcl = Pointclouds(points, normals) chamfer_p, chamfer_n = chamfer_distance( points, pcl_gt.points_padded(), x_normals=normals, y_normals=pcl_gt.normals_padded(), ) eval_dict['chamfer_p'] = chamfer_p.item() eval_dict['chamfer_n'] = chamfer_n.item() pf_dist = point_mesh_face_distance(meshes_gt, pcl) eval_dict['pf_dist'] = pf_dist.item() writer.writerow(eval_dict) for k, v in eval_dict.items(): if v < best_dict[k]: best_dict[k] = v print('best {} so far ({}): {:.4g}'.format( k, vis_file, v)) # generation dictories gen_dir = os.path.join(exp_dir, 'generation') if not os.path.isdir(gen_dir): return final_file = os.path.join(gen_dir, 'mesh.ply') val_pts_file = final_file[:-4] + '_val%d' % n_pts + '.ply' if not os.path.isfile(final_file): return gen_file_csv = os.path.join(gen_dir, "evaluation_n%d.csv" % n_pts) if not os.path.isfile(gen_file_csv): with open(os.path.join(gen_dir, "evaluation_n%d.csv" % n_pts), "w") as f: fieldnames = ['chamfer_p', 'chamfer_n', 'pf_dist'] writer = csv.DictWriter(f, fieldnames=fieldnames, restval="-", extrasaction="ignore") writer.writeheader() eval_dict = OrderedDict() mesh = trimesh.load(final_file) # points, normals = pcu.sample_mesh_poisson_disk( # mesh.vertices, mesh.faces, # mesh.vertex_normals.ravel().reshape(-1, 3), n_pts, use_geodesic_distance=True) # p_idx = np.random.permutation(points.shape[0])[:n_pts] # points = points[p_idx, ...] # normals = normals[p_idx, ...] # points = torch.from_numpy(points).float().to( # device=device).view(1, -1, 3) # normals = torch.from_numpy(normals).float().to( # device=device).view(1, -1, 3) meshes = Meshes( torch.from_numpy(mesh.vertices[None, ...]).float(), torch.from_numpy(mesh.faces[None, ...]).float()).to(device) points, normals = sample_points_from_meshes(meshes, n_pts, return_normals=True) trimesh.Trimesh(points.cpu().numpy()[0], vertex_normals=normals.cpu().numpy()[0], process=False).export(val_pts_file, vertex_normal=True) pcl = Pointclouds(points, normals) chamfer_p, chamfer_n = chamfer_distance( points, pcl_gt.points_padded(), x_normals=normals, y_normals=pcl_gt.normals_padded(), ) eval_dict['chamfer_p'] = chamfer_p.item() eval_dict['chamfer_n'] = chamfer_n.item() pf_dist = point_mesh_face_distance(meshes_gt, pcl) eval_dict['pf_dist'] = pf_dist.item() writer.writerow(eval_dict) for k, v in eval_dict.items(): if v < best_dict[k]: best_dict[k] = v print('best {} so far ({}): {:.4g}'.format( k, final_file, v))
def forward( self, camera: CamerasBase, image_rgb: torch.Tensor, depth_map: torch.Tensor, fg_probability: torch.Tensor, frame_type: List[str], **kwargs, ) -> Dict[str, Any]: # TODO: return a namedtuple or dataclass """ Given a set of input source cameras images and depth maps, unprojects all RGBD maps to a colored point cloud and renders into the target views. Args: camera: A batch of `N` PyTorch3D cameras. image_rgb: A batch of `N` images of shape `(N, 3, H, W)`. depth_map: A batch of `N` depth maps of shape `(N, 1, H, W)`. fg_probability: A batch of `N` foreground probability maps of shape `(N, 1, H, W)`. frame_type: A list of `N` strings containing frame type indicators which specify target and source views. Returns: preds: A dict with the following fields: nvs_prediction: The rendered colors, depth and mask of the target views. point_cloud: The point cloud of the scene. It's renders are stored in `nvs_prediction`. """ is_known = is_known_frame(frame_type) is_known_idx = torch.where(is_known)[0] mask_fg = (fg_probability > 0.5).type_as(image_rgb) point_cloud = get_rgbd_point_cloud( camera[is_known_idx], image_rgb[is_known_idx], depth_map[is_known_idx], mask_fg[is_known_idx], ) pcl_size = int(point_cloud.num_points_per_cloud()) if (self.max_points > 0) and (pcl_size > self.max_points): prm = torch.randperm(pcl_size)[:self.max_points] point_cloud = Pointclouds( point_cloud.points_padded()[:, prm, :], # pyre-fixme[16]: Optional type has no attribute `__getitem__`. features=point_cloud.features_padded()[:, prm, :], ) is_target_idx = torch.where(~is_known)[0] depth_render, image_render, mask_render = [], [], [] # render into target frames in a for loop to save memory for tgt_idx in is_target_idx: _image_render, _mask_render, _depth_render = render_point_cloud_pytorch3d( camera[int(tgt_idx)], point_cloud, render_size=(self.image_size, self.image_size), point_radius=1e-2, topk=10, bg_color=self.bg_color, ) _image_render = _image_render.clamp(0.0, 1.0) # the mask is the set of pixels with opacity bigger than eps _mask_render = (_mask_render > 1e-4).float() depth_render.append(_depth_render) image_render.append(_image_render) mask_render.append(_mask_render) nvs_prediction = NewViewSynthesisPrediction( **{ k: torch.cat(v, dim=0) for k, v in zip( ["depth_render", "image_render", "mask_render"], [depth_render, image_render, mask_render], ) }) preds = { "nvs_prediction": nvs_prediction, "point_cloud": point_cloud, } return preds
def test_ndc_grid_sample_rendering(self): """ Use PyTorch3D point renderer to render a colored point cloud, then sample the image at the locations of the point projections with `ndc_grid_sample`. Finally, assert that the sampled colors are equal to the original point cloud colors. Note that, in order to ensure correctness, we use a nearest-neighbor assignment point renderer (i.e. no soft splatting). """ # generate a bunch of 3D points on a regular grid lying in the z-plane n_grid_pts = 10 grid_scale = 0.9 z_plane = 2.0 image_size = [128, 128] point_radius = 0.015 n_pts = n_grid_pts * n_grid_pts pts = torch.stack( meshgrid_ij([torch.linspace(-grid_scale, grid_scale, n_grid_pts)] * 2, ), dim=-1, ) pts = torch.cat([pts, z_plane * torch.ones_like(pts[..., :1])], dim=-1) pts = pts.reshape(1, n_pts, 3) # color the points randomly pts_colors = torch.rand(1, n_pts, 3) # make trivial rendering cameras cameras = PerspectiveCameras( R=eyes(dim=3, N=1), device=pts.device, T=torch.zeros(1, 3, dtype=torch.float32, device=pts.device), ) # render the point cloud pcl = Pointclouds(points=pts, features=pts_colors) renderer = NearestNeighborPointsRenderer( rasterizer=PointsRasterizer( cameras=cameras, raster_settings=PointsRasterizationSettings( image_size=image_size, radius=point_radius, points_per_pixel=1, ), ), compositor=AlphaCompositor(), ) im_render = renderer(pcl) # sample the render at projected pts pts_proj = cameras.transform_points(pcl.points_padded())[..., :2] pts_colors_sampled = ndc_grid_sample( im_render, pts_proj, mode="nearest", align_corners=False, ).permute(0, 2, 1) # assert that the samples are the same as original points self.assertClose(pts_colors, pts_colors_sampled, atol=1e-4)
def cleanup_eval_depth( point_cloud: Pointclouds, camera: CamerasBase, depth: torch.Tensor, mask: torch.Tensor, sigma: float = 0.01, image=None, ): ba, _, H, W = depth.shape pcl = point_cloud.points_padded() n_pts = point_cloud.num_points_per_cloud() pcl_mask = ( torch.arange(pcl.shape[1], dtype=torch.int64, device=pcl.device)[None] < n_pts[:, None] ).type_as(pcl) pcl_proj = camera.transform_points(pcl, eps=1e-2)[..., :-1] pcl_depth = camera.get_world_to_view_transform().transform_points(pcl)[..., -1] depth_and_idx = torch.cat( ( depth, torch.arange(H * W).view(1, 1, H, W).expand(ba, 1, H, W).type_as(depth), ), dim=1, ) depth_and_idx_sampled = Fu.grid_sample( depth_and_idx, -pcl_proj[:, None], mode="nearest" )[:, :, 0].view(ba, 2, -1) depth_sampled, idx_sampled = depth_and_idx_sampled.split([1, 1], dim=1) df = (depth_sampled[:, 0] - pcl_depth).abs() # the threshold is a sigma-multiple of the standard deviation of the depth mu = wmean(depth.view(ba, -1, 1), mask.view(ba, -1)).view(ba, 1) std = ( wmean((depth.view(ba, -1) - mu).view(ba, -1, 1) ** 2, mask.view(ba, -1)) .clamp(1e-4) .sqrt() .view(ba, -1) ) good_df_thr = std * sigma good_depth = (df <= good_df_thr).float() * pcl_mask perc_kept = good_depth.sum(dim=1) / pcl_mask.sum(dim=1).clamp(1) # print(f'Kept {100.0 * perc_kept.mean():1.3f} % points') good_depth_raster = torch.zeros_like(depth).view(ba, -1) # pyre-ignore[16]: scatter_add_ good_depth_raster.scatter_add_(1, torch.round(idx_sampled[:, 0]).long(), good_depth) good_depth_mask = (good_depth_raster.view(ba, 1, H, W) > 0).float() # if float(torch.rand(1)) > 0.95: # depth_ok = depth * good_depth_mask # # visualize # visdom_env = 'depth_cleanup_dbg' # from visdom import Visdom # # from tools.vis_utils import make_depth_image # from pytorch3d.vis.plotly_vis import plot_scene # viz = Visdom() # show_pcls = { # 'pointclouds': point_cloud, # } # for d, nm in zip( # (depth, depth_ok), # ('pointclouds_unproj', 'pointclouds_unproj_ok'), # ): # pointclouds_unproj = get_rgbd_point_cloud( # camera, image, d, # ) # if int(pointclouds_unproj.num_points_per_cloud()) > 0: # show_pcls[nm] = pointclouds_unproj # scene_dict = {'1': { # **show_pcls, # 'cameras': camera, # }} # scene = plot_scene( # scene_dict, # pointcloud_max_points=5000, # pointcloud_marker_size=1.5, # camera_scale=1.0, # ) # viz.plotlyplot(scene, env=visdom_env, win='scene') # # depth_image_ok = make_depth_image(depths_ok, masks) # # viz.images(depth_image_ok, env=visdom_env, win='depth_ok') # # depth_image = make_depth_image(depths, masks) # # viz.images(depth_image, env=visdom_env, win='depth') # # # viz.images(rgb_rendered, env=visdom_env, win='images_render') # # viz.images(images, env=visdom_env, win='images') # import pdb; pdb.set_trace() return good_depth_mask