def test_small_faces_case(self): for device in [torch.device("cpu"), torch.device("cuda:0")]: mesh_vertices = torch.tensor( [ [-0.0021, -0.3769, 0.7146], [-0.0161, -0.3771, 0.7146], [-0.0021, -0.3771, 0.7147], ], dtype=torch.float32, device=device, ) mesh1_faces = torch.tensor([[0, 2, 1]], device=device) mesh2_faces = torch.tensor([[2, 0, 1]], device=device) pcd_points = torch.tensor([[-0.3623, -0.5340, 0.7727]], device=device) mesh1 = Meshes(verts=[mesh_vertices], faces=[mesh1_faces]) mesh2 = Meshes(verts=[mesh_vertices], faces=[mesh2_faces]) pcd = Pointclouds(points=[pcd_points]) loss1 = point_mesh_face_distance(mesh1, pcd) loss2 = point_mesh_face_distance(mesh2, pcd) self.assertClose(loss1, loss2)
def compute_loss(self, mesh, pcd=None): if pcd is None: pcd = self.pcd face_loss = pt3loss.point_mesh_face_distance(mesh, pcd) edge_loss = pt3loss.point_mesh_edge_distance(mesh, pcd) point_loss = pt3loss.chamfer_distance(mesh.verts_padded(), pcd)[0] length_loss = pt3loss.mesh_edge_loss(mesh) normal_loss = pt3loss.mesh_normal_consistency(mesh) mpcd = sample_points_from_meshes(mesh, 2 * pcd.points_padded()[0].shape[0]) sample_loss, _ = pt3loss.chamfer_distance(mpcd, pcd) losses = torch.tensor((face_loss, edge_loss, point_loss, length_loss, normal_loss, sample_loss), requires_grad=True).to(device='cuda') return losses
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 loss(): point_mesh_face_distance(meshes, pcls) torch.cuda.synchronize()
def test_point_mesh_face_distance(self): """ Test point_mesh_face_distance from pytorch3d.loss """ device = get_random_cuda_device() N, V, F, P = 4, 32, 16, 24 meshes, pcls = self.init_meshes_clouds(N, V, F, P, device=device) # clone and detach for another backward pass through the op verts_op = [verts.clone().detach() for verts in meshes.verts_list()] for i in range(N): verts_op[i].requires_grad = True faces_op = [faces.clone().detach() for faces in meshes.faces_list()] meshes_op = Meshes(verts=verts_op, faces=faces_op) points_op = [points.clone().detach() for points in pcls.points_list()] for i in range(N): points_op[i].requires_grad = True pcls_op = Pointclouds(points_op) # naive implementation loss_naive = torch.zeros(N, dtype=torch.float32, device=device) for i in range(N): points = pcls.points_list()[i] verts = meshes.verts_list()[i] faces = meshes.faces_list()[i] tris = verts[faces] num_p = points.shape[0] num_t = tris.shape[0] dists = torch.zeros((num_p, num_t), dtype=torch.float32, device=device) for p in range(num_p): for t in range(num_t): dist = self._point_to_tri_distance(points[p], tris[t]) dists[p, t] = dist min_dist_p, min_idx_p = dists.min(1) min_dist_t, min_idx_t = dists.min(0) loss_naive[i] = min_dist_p.mean() + min_dist_t.mean() loss_naive = loss_naive.mean() # Op loss_op = point_mesh_face_distance(meshes_op, pcls_op) # Compare forward pass self.assertClose(loss_op, loss_naive) # Compare backward pass rand_val = torch.rand(1).item() grad_dist = torch.tensor(rand_val, dtype=torch.float32, device=device) loss_naive.backward(grad_dist) loss_op.backward(grad_dist) # check verts grad for i in range(N): self.assertClose( meshes.verts_list()[i].grad, meshes_op.verts_list()[i].grad ) self.assertClose(pcls.points_list()[i].grad, pcls_op.points_list()[i].grad)