Пример #1
0
    def extract_mesh(self, occ_hat, z, c=None, stats_dict=dict()):
        ''' Extracts the mesh from the predicted occupancy grid.

        Args:
            occ_hat (tensor): value grid of occupancies
            z (tensor): latent code z
            c (tensor): latent conditioned code c
            stats_dict (dict): stats dictionary
        '''
        # Some short hands
        n_x, n_y, n_z = occ_hat.shape
        box_size = 1 + self.padding
        threshold = np.log(self.threshold) - np.log(1. - self.threshold)
        # Make sure that mesh is watertight
        t0 = time.time()
        occ_hat_padded = np.pad(occ_hat, 1, 'constant', constant_values=-1e6)
        vertices, triangles = libmcubes.marching_cubes(occ_hat_padded,
                                                       threshold)
        stats_dict['time (marching cubes)'] = time.time() - t0
        # Strange behaviour in libmcubes: vertices are shifted by 0.5
        vertices -= 0.5
        # Undo padding
        vertices -= 1
        # Normalize to bounding box
        vertices /= np.array([n_x - 1, n_y - 1, n_z - 1])
        vertices = box_size * (vertices - 0.5)

        # Estimate normals if needed
        if self.with_normals and not vertices.shape[0] == 0:
            t0 = time.time()
            normals = self.estimate_normals(vertices, z, c)
            stats_dict['time (normals)'] = time.time() - t0

        else:
            normals = None

        # Create mesh
        mesh = trimesh.Trimesh(vertices,
                               triangles,
                               vertex_normals=normals,
                               process=False)

        # Directly return if mesh is empty
        if vertices.shape[0] == 0:
            return mesh

        # TODO: normals are lost here
        if self.simplify_nfaces is not None:
            t0 = time.time()
            mesh = simplify_mesh(mesh, self.simplify_nfaces, 5.)
            stats_dict['time (simplify)'] = time.time() - t0

        # Refine mesh
        if self.refinement_step > 0:
            t0 = time.time()
            self.refine_mesh(mesh, occ_hat, z, c)
            stats_dict['time (refine)'] = time.time() - t0

        return mesh
Пример #2
0
    def extract_mesh(self,
                     occ_hat,
                     label_hat,
                     p_hat,
                     c,
                     stats_dict=dict(),
                     **kwargs):
        ''' Extracts the mesh from the predicted occupancy grid.

        Args:
            occ_hat (tensor): value grid of occupancies
            label_hat (tensor): value grid of predicted part labels
            p_hat (tensor): value grid of predicted locations in the A-pose
            c (tensor): latent conditioned code c
            stats_dict (dict): stats dictionary
        '''
        # Some short hands
        n_x, n_y, n_z = occ_hat.shape
        box_size = 1 + self.padding
        threshold = np.log(self.threshold) - np.log(1. - self.threshold)
        # Make sure that mesh is watertight
        t0 = time.time()
        occ_hat_padded = np.pad(occ_hat, 1, 'constant', constant_values=-1e6)
        vertices, triangles = libmcubes.marching_cubes(occ_hat_padded,
                                                       threshold)
        stats_dict['time (marching cubes)'] = time.time() - t0
        # Strange behaviour in libmcubes: vertices are shifted by 0.5
        vertices -= 0.5
        # Undo padding
        vertices -= 1
        # Construct part labels and A-pose vertices by sampling
        # occupancy grid and translation grid
        r_verts = np.round(vertices).astype('int32')
        labels = label_hat[r_verts[:, 0], r_verts[:, 1], r_verts[:, 2]]
        colors = self.colors[labels]
        # Normalize to bounding box
        vertices /= np.array([n_x - 1, n_y - 1, n_z - 1])
        if p_hat is not None:
            with torch.no_grad():
                v = torch.from_numpy(vertices).to(self.device).float()
                v = v * 2 - 1  # range: [-1, 1]
                v = v.unsqueeze(0).unsqueeze(1).unsqueeze(
                    1)  # 1 x 1 x 1 x n_pts x 3
                p_hat = torch.from_numpy(p_hat).to(self.device).float()
                p_hat = p_hat.permute(3, 2, 1,
                                      0).unsqueeze(0)  # 1 X C x D x H x W
                # v_rest is in [-1, 1]
                v_rest = torch.nn.functional.grid_sample(p_hat,
                                                         v,
                                                         align_corners=True)
                v_rest = v_rest.squeeze(0).squeeze(1).squeeze(1).transpose(
                    0, 1) / 1.5 * kwargs['scale']  # + kwargs['loc']

            vertices_rest = v_rest.detach().cpu().numpy()
        else:
            vertices_rest = None

        # vertices = box_size * (vertices - 0.5)
        vertices = 4 / 3 * kwargs['scale'].item() * (
            vertices - 0.5) + kwargs['loc'].cpu().numpy()

        # Estimate normals if needed
        if self.with_normals and not vertices.shape[0] == 0:
            t0 = time.time()
            normals = self.estimate_normals(vertices, c)
            stats_dict['time (normals)'] = time.time() - t0

        else:
            normals = None

        # Create mesh
        mesh = {}
        mesh['part_labels'] = labels
        mesh['posed'] = trimesh.Trimesh(vertices,
                                        triangles,
                                        vertex_normals=normals,
                                        vertex_colors=colors,
                                        process=False)
        if vertices_rest is not None:
            mesh['unposed'] = trimesh.Trimesh(vertices_rest,
                                              triangles,
                                              vertex_normals=normals,
                                              vertex_colors=colors,
                                              process=False)

        # Directly return if mesh is empty
        if vertices.shape[0] == 0:
            return mesh

        # TODO: normals are lost here
        if self.simplify_nfaces is not None:
            t0 = time.time()
            mesh = simplify_mesh(mesh, self.simplify_nfaces, 5.)
            stats_dict['time (simplify)'] = time.time() - t0

        # Refine mesh
        if self.refinement_step > 0:
            t0 = time.time()
            self.refine_mesh(mesh, occ_hat, c)
            stats_dict['time (refine)'] = time.time() - t0

        return mesh