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)
Пример #2
0
    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
Пример #3
0
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))
Пример #4
0
 def loss():
     point_mesh_face_distance(meshes, pcls)
     torch.cuda.synchronize()
Пример #5
0
    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)