shape_data/faust.py

import os.path as osp
import sys
import itertools
import math
import numpy as np
import torch
from torch.utils.data import Dataset
from pathlib import Path
import potpourri3d as pp3d 
import open3d as o3d 
from utils.geometry import get_operators, load_operators
from utils.surfaces import Surface
from utils.utils_func import may_create_folder
from utils.mesh import find_mesh_files

# def opt_rot_points(pts_1, pts_2, device="cuda:0"):
#     center_1 = pts_1.mean(dim=0)
#     pts_c1 = pts_1 - center_1
#     center_2 = pts_2.mean(dim=0)
#     pts_c2 = pts_2 - center_2 
#     to_sum = pts_c1[:, :, None] * pts_c2[:, None, :]
#     A = pts_c1.T @ pts_c2
#     #A = to_sum.sum(axis=0)
#     u, _, v = torch.linalg.svd(A)
#     a = torch.Tensor([[1, 0, 0], [0, 1, 0], [0, 0, torch.sign(torch.linalg.det(A))]]).float().to(device)
#     O = u @ a @ v
#     return O.T

def opt_rot_points(pts_1, pts_2):
    center_1 = pts_1.mean(axis=0)
    pts_c1 = pts_1 - center_1
    center_2 = pts_2.mean(axis=0)
    pts_c2 = pts_2 - center_2

    A = np.dot(pts_c1.T, pts_c2)
    u, _, v = np.linalg.svd(A)
    a = np.array([[1, 0, 0], [0, 1, 0], [0, 0, np.sign(np.linalg.det(A))]])
    O = u @ a @ v
    return O.T

def compute_vertex_normals(vertices, faces):
    mesh = o3d.geometry.TriangleMesh(o3d.utility.Vector3dVector(vertices), o3d.utility.Vector3iVector(faces))
    mesh.compute_vertex_normals()
    return np.asarray(mesh.vertex_normals, dtype=np.float32)


def numpy_to_open3d_mesh(V, F):
    # Create an empty TriangleMesh object
    mesh = o3d.geometry.TriangleMesh()
    # Set vertices
    mesh.vertices = o3d.utility.Vector3dVector(V)
    # Set triangles
    mesh.triangles = o3d.utility.Vector3iVector(F)
    return mesh

def open_mesh(path):
    """
    Tries to open a mesh.
    If it fails, try .ply, .obj, and .off alternatives.
    
    Parameters
    ----------
    path : str
        Path of the mesh
    Returns
    -------
    mesh or None
        Loaded mesh (V, F format) if successful, else None
    """
    p = Path(path)
    base, ext = p.with_suffix(""), p.suffix
    tried_exts = [ext, ".ply", ".obj", ".off"]
    for e in tried_exts:
        path = base.with_suffix(e)
        if Path.exists(path):
            try:
                temp = pp3d.read_mesh(str(path))
                return temp
            except Exception as err:
                print(f"Failed loading {path}: {err}")
    return None


KEYS = ['vertices', 'faces', 'frames', 'mass', 'L', 'evals', 'evecs', 'gradX', 'gradY', 'hks', 'wks', 'idx', 'name']


class ShapeDataset(Dataset):
    TRAIN_IDX = np.arange(0, 80)
    TEST_IDX = np.arange(80, 100)
    NAME = "FAUST"
    def __init__(self,
                 shape_dir,
                 cache_dir,
                 mode,
                 oriented=False,
                 rot_auto=False,
                 num_eigenbasis=256,
                 laplacian_type='mesh',
                 feature_type=None,
                 **kwargs):
        super().__init__()

        self.shape_dir = shape_dir
        self.cache_dir = cache_dir
        self.mode = mode
        self.oriented = oriented
        if self.oriented: 
            self.NAME = self.NAME + "_ori"
        self.num_eigenbasis = num_eigenbasis
        self.laplacian_type = laplacian_type
        self.feature_type = feature_type
        for k, w in kwargs.items():
            setattr(self, k, w)

        print(f'Loading {mode} data from {shape_dir}')
        self.shape_list = self._get_file_list()
        self._prepare()

        self.randg = np.random.RandomState(0)

    def _get_file_list(self):
        path_list = find_mesh_files(Path(self.shape_dir), alphanum_sort=True)
        file_list = [f.name for f in path_list]
        if self.mode.startswith('train'):
            assert self.TRAIN_IDX is not None
            shape_list = [file_list[idx] for idx in self.TRAIN_IDX]
        elif self.mode.startswith('test'):
            assert self.TEST_IDX is not None
            shape_list = [file_list[idx] for idx in self.TEST_IDX]
        else:
            raise RuntimeError(f'Mode {self.mode} is not supported.')
        return shape_list

    def _load_mesh(self, filepath, scale=True, return_vnormals=False):
        V, F = open_mesh(filepath)
        mesh = numpy_to_open3d_mesh(V, F)

        tmat = np.identity(4, dtype=np.float32)
        center = mesh.get_center()
        tmat[:3, 3] = -center
        if scale:
            smat = np.identity(4, dtype=np.float32)
            area = mesh.get_surface_area()
            smat[:3, :3] = np.identity(3, dtype=np.float32) / math.sqrt(area)
            tmat = smat @ tmat
        mesh.transform(tmat)

        vertices = np.asarray(mesh.vertices, dtype=np.float32)
        faces = np.asarray(mesh.triangles, dtype=np.int32)
        if return_vnormals:
            mesh.compute_vertex_normals()
            vnormals = np.asarray(mesh.vertex_normals, dtype=np.float32)
            return vertices, faces, vnormals
        else:
            return vertices, faces

    def _prepare(self):
        may_create_folder(self.cache_dir)
        for sid, sname in enumerate(self.shape_list):
            cache_prefix = osp.join(self.cache_dir, self.NAME, f'{sname[:-4]}_{self.laplacian_type}_{self.num_eigenbasis}k')
            cache_path = cache_prefix + '_0n.npz'
            if not Path(cache_path).is_file():
                vertices_np, faces_np, vertex_normals_np = self._load_mesh(osp.join(self.shape_dir, sname),
                                                                       scale=True,
                                                                       return_vnormals=True)

                if self.laplacian_type == 'mesh':
                    _ = get_operators(torch.from_numpy(vertices_np).float(), torch.from_numpy(faces_np).long(), self.num_eigenbasis, cache_path=cache_path)
                # elif self.laplacian_type == 'pcd':
                #     compute_operators(vertices_np, np.asarray([], dtype=np.int32), vertex_normals_np, self.num_eigenbasis,
                #                       cache_path)
                else:
                    raise RuntimeError(f'Basis type {self.laplacian_type} is not supported')

            # if self.aug_noise_type is not None and self.aug_noise_type != 'naive':
            #     max_magnitude, max_levels = self.aug_noise_args[:2]
            #     randg = np.random.RandomState(sid)
            #     for nlevel in range(1, max_levels + 1):
            #         cache_path = cache_prefix + f'_{nlevel}n.npz'
            #         if Path(cache_path).is_file():
            #             continue
            #         noise_mag = max_magnitude * nlevel / max_levels
            #         noise_mat = np.clip(noise_mag * randg.randn(vertices_np.shape[0], vertices_np.shape[1]), -noise_mag,
            #                             noise_mag)
            #         vertices_noise_np = vertices_np + noise_mat.astype(vertices_np.dtype)
            #         vertex_normals_noise_np = compute_vertex_normals(vertices_noise_np, faces_np)

            #         if self.laplacian_type == 'mesh':
            #             compute_operators(vertices_noise_np, faces_np, vertex_normals_noise_np, self.num_eigenbasis, cache_path)
            #         elif self.laplacian_type == 'pcd':
            #             compute_operators(vertices_noise_np, np.asarray([], dtype=np.int32), vertex_normals_noise_np,
            #                               self.num_eigenbasis, cache_path)
            #         else:
            #             raise RuntimeError(f'Basis type {self.laplacian_type} is not supported')

    def __getitem__(self, idx):
        sname = self.shape_list[idx]

        cache_prefix = osp.join(self.cache_dir, self.NAME, f'{sname[:-4]}_{self.laplacian_type}_{self.num_eigenbasis}k')
        cache_path = cache_prefix + '_0n.npz'

        assert Path(cache_path).is_file()

        sdict = load_operators(cache_path)
        sdict['idx'] = idx
        sdict['name'] = sname[:-4]

        if self.feature_type is not None:
            sdict['feats'] = np.concatenate([sdict[ft] for ft in self.feature_type.split('_')], axis=-1)
        vertices_np, _, _ = self._load_mesh(osp.join(self.shape_dir, sname), scale=True, return_vnormals=True)
        sdict['vertices'] = vertices_np
        sdict = self._centering(sdict)
        return sdict

    def __len__(self):
        return len(self.shape_list)

    def _centering(self, sdict):
        vertices, areas = sdict['vertices'], sdict["mass"]
        center = (vertices*areas[:, None]).sum()/areas.sum()
        sdict['vertices'] = vertices - center
        return sdict

    def _random_noise_naive(self, sdict, randg, args):
        vertices = sdict['vertices']
        dtype = vertices.dtype
        shape = vertices.shape
        std, clip = args

        noise = np.clip(std * randg.randn(*shape), -clip, clip)
        sdict['vertices'] = vertices + noise.astype(dtype)
        return sdict

    def _random_rotation(self, sdict, randg, axes, args):
        vertices = sdict['vertices']
        dtype = vertices.dtype

        max_x, max_y, max_z = args
        if 'x' in axes:
            anglex = randg.rand() * max_x * np.pi / 180.0
            cosx = np.cos(anglex)
            sinx = np.sin(anglex)
            Rx = np.asarray([[1, 0, 0], [0, cosx, -sinx], [0, sinx, cosx]], dtype=dtype)
        else:
            Rx = np.eye(3, dtype=dtype)

        if 'y' in axes:
            angley = randg.rand() * max_y * np.pi / 180.0
            cosy = np.cos(angley)
            siny = np.sin(angley)
            Ry = np.asarray([[cosy, 0, siny], [0, 1, 0], [-siny, 0, cosy]], dtype=dtype)
        else:
            Ry = np.eye(3, dtype=dtype)

        if 'z' in axes:
            anglez = randg.rand() * max_z * np.pi / 180.0
            cosz = np.cos(anglez)
            sinz = np.sin(anglez)
            Rz = np.asarray([[cosz, -sinz, 0], [sinz, cosz, 0], [0, 0, 1]], dtype=dtype)
        else:
            Rz = np.eye(3, dtype=dtype)

        Rxyz = randg.permutation(np.stack((Rx, Ry, Rz), axis=0))
        R = Rxyz[2] @ Rxyz[1] @ Rxyz[0]
        sdict['vertices'] = vertices @ R.T

        return sdict

    def _random_scaling(self, sdict, randg, args):
        scale_min, scale_max = args
        vertices = sdict['vertices']
        scale = scale_min + randg.rand(1, 3) * (scale_max - scale_min)
        sdict['vertices'] = vertices * scale
        return sdict

    def get_name_id_map(self):
        return {sname[:-4]: sid for sid, sname in enumerate(self.shape_list)}


class ShapePairDataset(Dataset):

    def __init__(self, corr_dir, mode, shape_data, rotate=False, **kwargs):
        super().__init__()
        self.corr_dir = corr_dir
        self.mode = mode
        self.shape_data = shape_data
        self.rotate = rotate
        if self.shape_data.oriented and self.rotate:
            self.rotate = False
        for k, w in kwargs.items():
            setattr(self, k, w)

        self._init()

        self.randg = np.random.RandomState(0)

    def _init(self):
        self.name_id_map = self.shape_data.get_name_id_map()
        self.pair_indices = list(itertools.combinations(range(len(self.shape_data)), 2))

    def __getitem__(self, idx):
        pidx = self.pair_indices[idx]
        sdict0 = self.shape_data[pidx[0]]
        sdict1 = self.shape_data[pidx[1]]
        return self._prepare_pair(sdict0, sdict1)

    def get_by_names(self, sname0, sname1):
        sdict0 = self.shape_data[self.name_id_map[sname0]]
        sdict1 = self.shape_data[self.name_id_map[sname1]]
        return self._prepare_pair(sdict0, sdict1)

    def _prepare_pair(self, sdict0, sdict1):
        corr_gt = self._load_corr_gt(sdict0, sdict1)
        # for fmap_size in self.fmap_sizes:
        #     fmap01_gt = pmap_to_fmap(sdict0['evecs'][:, :fmap_size], sdict1['evecs'][:, :fmap_size], corr_gt)
        #     pdict[f'fmap01_{fmap_size}_gt'] = fmap01_gt

        # for idx in range(2):
        #     indices_sel = farthest_point_sampling(pdict[f'vertices{idx}'], self.num_corrs, random_start=is_train)
        #     for k in ['vertices', 'evecs', 'feats']:
        #         kid = f'{k}{idx}'
        #         if kid in pdict:
        #             pdict[kid + '_sub'] = pdict[kid][indices_sel, :]
        #     if self.use_geodists:
        #         geodists = compute_geodesic_distance(pdict[f'vertices{idx}'], pdict[f'faces{idx}'], indices_sel)
        #         pdict[f'geodists{idx}_sub'] = geodists
        #     pdict[f'vindices{idx}_sub'] = indices_sel

        # fmap_size = self.fmap_sizes[-1]
        # corr_gt_sub = fmap_to_pmap(pdict['evecs0_sub'][:, :fmap_size], pdict['evecs1_sub'][:, :fmap_size],
        #                            pdict[f'fmap01_{fmap_size}_gt'])
        # pdict['corr_gt_sub'] = corr_gt_sub

        # if is_train:
        #     fmap_size = self.fmap_sizes[0]
        #     axis = self.randg.choice([0, 1]).item()
        #     max_bases = fmap_size // 2
        #     noise_ratio = 0.5
        #     if self.randg.rand() > 0.5:
        #         pdict[f'fmap01_{fmap_size}'] = self._random_scale(pdict[f'fmap01_{fmap_size}_gt'], self.randg, axis, max_bases)
        #     else:
        #         pdict[f'fmap01_{fmap_size}'] = self._random_noise(pdict[f'fmap01_{fmap_size}_gt'], self.randg, axis, max_bases,
        #                                                           noise_ratio)
        # else:
        #     if self.corr_loader is not None:
        #         corr_init = self.corr_loader.get_by_names(sdict0['name'], sdict1['name'])
        #         assert corr_init.ndim == 2 and len(corr_init) == len(sdict1['vertices'])
        #         fmap_size = self.fmap_sizes[0]
        #         fmap01_init = pmap_to_fmap(sdict0['evecs'][:, :fmap_size], sdict1['evecs'][:, :fmap_size], corr_init)
        #         pdict[f'fmap01_{fmap_size}'] = fmap01_init
        #         pdict['pmap10'] = corr_init[:, 0]

        vts_1, vts_2 = corr_gt[:, 0], corr_gt[:, 1]
        shape_dict, target_dict = sdict0, sdict1
        
        if self.rotate:
            pts_1, pts_2 = shape_dict['vertices'][vts_1], target_dict['vertices'][vts_2]
            rot = opt_rot_points(pts_1, pts_2).astype(np.float32)#, device="cuda")
            target_dict['vertices'] = target_dict['vertices'] @ rot
        target_surf = Surface(FV=[target_dict['faces'], target_dict['vertices']])
        target_normals = torch.from_numpy(target_surf.surfel/np.linalg.norm(target_surf.surfel, axis=-1, keepdims=True)).float().cuda()

        shape_surf = Surface(FV=[shape_dict['faces'], shape_dict['vertices']])
        map_info = (shape_dict['name'], vts_1, vts_2)
        return shape_dict, shape_surf, target_dict, target_surf, target_normals, map_info

    def _random_scale(self, fmap, randg, axis, max_bases):
        assert max_bases > 1
        assert axis in [0, 1]
        num_bases = randg.randint(1, max_bases)
        ids = randg.choice(fmap.shape[axis], num_bases, replace=False)
        fmap_out = np.copy(fmap)
        if axis == 0:
            fmap_out[ids, :] *= (randg.rand(num_bases, 1) * 2 - 1)
        else:
            fmap_out[:, ids] *= (randg.rand(1, num_bases) * 2 - 1)
        return fmap_out

    def _random_noise(self, fmap, randg, axis, max_bases, max_ratio):
        assert max_bases > 1
        assert axis in [0, 1]
        num_bases = randg.randint(1, max_bases)
        ids = randg.choice(fmap.shape[axis], num_bases, replace=False)
        fmap_out = np.copy(fmap)
        ratio = randg.rand() * max_ratio
        if axis == 0:
            maxvals = np.amax(np.abs(fmap_out[ids, :]), axis=1 - axis, keepdims=True)
            noise = ratio * maxvals * randg.randn(num_bases, fmap.shape[1 - axis])
            fmap_out[ids, :] += noise
        else:
            maxvals = np.amax(np.abs(fmap_out[:, ids]), axis=1 - axis, keepdims=True)
            noise = ratio * maxvals * randg.randn(fmap.shape[1 - axis], num_bases)
            fmap_out[:, ids] += noise
        return fmap_out

    def _load_corr_gt(self, sdict0, sdict1):
        corr0 = self._load_corr_file(sdict0['name'])
        corr1 = self._load_corr_file(sdict1['name'])
        corr_gt = np.stack((corr0, corr1), axis=1)
        return corr_gt

    def _load_corr_file(self, sname):
        corr_path = osp.join(self.corr_dir, f'{sname}.vts')
        corr = np.loadtxt(corr_path, dtype=np.int32)
        return corr - 1

    def __len__(self):
        return len(self.pair_indices)