lib/dataset/PIFuDataset.py

from lib.renderer.mesh import load_fit_body
from lib.dataset.hoppeMesh import HoppeMesh
from lib.dataset.body_model import TetraSMPLModel
from lib.common.render import Render
from lib.dataset.mesh_util import SMPLX, projection, cal_sdf_batch, get_visibility
from lib.pare.pare.utils.geometry import rotation_matrix_to_angle_axis
from termcolor import colored
import os.path as osp
import numpy as np
from PIL import Image
import random
import os
import trimesh
import torch
from kaolin.ops.mesh import check_sign
import torchvision.transforms as transforms
from huggingface_hub import hf_hub_download, cached_download


class PIFuDataset():
    def __init__(self, cfg, split='train', vis=False):

        self.split = split
        self.root = cfg.root
        self.bsize = cfg.batch_size
        self.overfit = cfg.overfit

        # for debug, only used in visualize_sampling3D
        self.vis = vis

        self.opt = cfg.dataset
        self.datasets = self.opt.types
        self.input_size = self.opt.input_size
        self.scales = self.opt.scales
        self.workers = cfg.num_threads
        self.prior_type = cfg.net.prior_type

        self.noise_type = self.opt.noise_type
        self.noise_scale = self.opt.noise_scale

        noise_joints = [4, 5, 7, 8, 13, 14, 16, 17, 18, 19, 20, 21]

        self.noise_smpl_idx = []
        self.noise_smplx_idx = []

        for idx in noise_joints:
            self.noise_smpl_idx.append(idx * 3)
            self.noise_smpl_idx.append(idx * 3 + 1)
            self.noise_smpl_idx.append(idx * 3 + 2)

            self.noise_smplx_idx.append((idx-1) * 3)
            self.noise_smplx_idx.append((idx-1) * 3 + 1)
            self.noise_smplx_idx.append((idx-1) * 3 + 2)

        self.use_sdf = cfg.sdf
        self.sdf_clip = cfg.sdf_clip

        # [(feat_name, channel_num),...]
        self.in_geo = [item[0] for item in cfg.net.in_geo]
        self.in_nml = [item[0] for item in cfg.net.in_nml]

        self.in_geo_dim = [item[1] for item in cfg.net.in_geo]
        self.in_nml_dim = [item[1] for item in cfg.net.in_nml]

        self.in_total = self.in_geo + self.in_nml
        self.in_total_dim = self.in_geo_dim + self.in_nml_dim

        if self.split == 'train':
            self.rotations = np.arange(
                0, 360, 360 / self.opt.rotation_num).astype(np.int32)
        else:
            self.rotations = range(0, 360, 120)

        self.datasets_dict = {}

        for dataset_id, dataset in enumerate(self.datasets):

            mesh_dir = None
            smplx_dir = None

            dataset_dir = osp.join(self.root, dataset)

            if dataset in ['thuman2']:
                mesh_dir = osp.join(dataset_dir, "scans")
                smplx_dir = osp.join(dataset_dir, "fits")
                smpl_dir = osp.join(dataset_dir, "smpl")

            self.datasets_dict[dataset] = {
                "subjects": np.loadtxt(osp.join(dataset_dir, "all.txt"), dtype=str),
                "smplx_dir": smplx_dir,
                "smpl_dir": smpl_dir,
                "mesh_dir": mesh_dir,
                "scale": self.scales[dataset_id]
            }

        self.subject_list = self.get_subject_list(split)
        self.smplx = SMPLX()

        # PIL to tensor
        self.image_to_tensor = transforms.Compose([
            transforms.Resize(self.input_size),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])

        # PIL to tensor
        self.mask_to_tensor = transforms.Compose([
            transforms.Resize(self.input_size),
            transforms.ToTensor(),
            transforms.Normalize((0.0, ), (1.0, ))
        ])

        self.device = torch.device(f"cuda:{cfg.gpus[0]}")
        self.render = Render(size=512, device=self.device)

    def render_normal(self, verts, faces):

        # render optimized mesh (normal, T_normal, image [-1,1])
        self.render.load_meshes(verts, faces)
        return self.render.get_rgb_image()

    def get_subject_list(self, split):

        subject_list = []

        for dataset in self.datasets:

            split_txt = osp.join(self.root, dataset, f'{split}.txt')

            if osp.exists(split_txt):
                print(f"load from {split_txt}")
                subject_list += np.loadtxt(split_txt, dtype=str).tolist()
            else:
                full_txt = osp.join(self.root, dataset, 'all.txt')
                print(f"split {full_txt} into train/val/test")

                full_lst = np.loadtxt(full_txt, dtype=str)
                full_lst = [dataset+"/"+item for item in full_lst]
                [train_lst, test_lst, val_lst] = np.split(
                    full_lst, [500, 500+5, ])

                np.savetxt(full_txt.replace(
                    "all", "train"), train_lst, fmt="%s")
                np.savetxt(full_txt.replace("all", "test"), test_lst, fmt="%s")
                np.savetxt(full_txt.replace("all", "val"), val_lst, fmt="%s")

                print(f"load from {split_txt}")
                subject_list += np.loadtxt(split_txt, dtype=str).tolist()

        if self.split != 'test':
            subject_list += subject_list[:self.bsize -
                                         len(subject_list) % self.bsize]
            print(colored(f"total: {len(subject_list)}", "yellow"))
            random.shuffle(subject_list)

        # subject_list = ["thuman2/0008"]
        return subject_list

    def __len__(self):
        return len(self.subject_list) * len(self.rotations)

    def __getitem__(self, index):

        # only pick the first data if overfitting
        if self.overfit:
            index = 0

        rid = index % len(self.rotations)
        mid = index // len(self.rotations)

        rotation = self.rotations[rid]
        subject = self.subject_list[mid].split("/")[1]
        dataset = self.subject_list[mid].split("/")[0]
        render_folder = "/".join([dataset +
                                 f"_{self.opt.rotation_num}views", subject])

        # setup paths
        data_dict = {
            'dataset': dataset,
            'subject': subject,
            'rotation': rotation,
            'scale': self.datasets_dict[dataset]["scale"],
            'mesh_path': osp.join(self.datasets_dict[dataset]["mesh_dir"], f"{subject}/{subject}.obj"),
            'smplx_path': osp.join(self.datasets_dict[dataset]["smplx_dir"], f"{subject}/smplx_param.pkl"),
            'smpl_path': osp.join(self.datasets_dict[dataset]["smpl_dir"], f"{subject}.pkl"),
            'calib_path': osp.join(self.root, render_folder, 'calib', f'{rotation:03d}.txt'),
            'vis_path': osp.join(self.root, render_folder, 'vis', f'{rotation:03d}.pt'),
            'image_path': osp.join(self.root, render_folder, 'render', f'{rotation:03d}.png')
        }

        # load training data
        data_dict.update(self.load_calib(data_dict))

        # image/normal/depth loader
        for name, channel in zip(self.in_total, self.in_total_dim):

            if f'{name}_path' not in data_dict.keys():
                data_dict.update({
                    f'{name}_path': osp.join(self.root, render_folder, name, f'{rotation:03d}.png')
                })

            # tensor update
            data_dict.update({
                name: self.imagepath2tensor(
                    data_dict[f'{name}_path'], channel, inv=False)
            })

        data_dict.update(self.load_mesh(data_dict))
        data_dict.update(self.get_sampling_geo(
            data_dict, is_valid=self.split == "val", is_sdf=self.use_sdf))
        data_dict.update(self.load_smpl(data_dict, self.vis))

        if self.prior_type == 'pamir':
            data_dict.update(self.load_smpl_voxel(data_dict))

        if (self.split != 'test') and (not self.vis):

            del data_dict['verts']
            del data_dict['faces']

        if not self.vis:
            del data_dict['mesh']

        path_keys = [
            key for key in data_dict.keys() if '_path' in key or '_dir' in key
        ]
        for key in path_keys:
            del data_dict[key]

        return data_dict

    def imagepath2tensor(self, path, channel=3, inv=False):

        rgba = Image.open(path).convert('RGBA')
        mask = rgba.split()[-1]
        image = rgba.convert('RGB')
        image = self.image_to_tensor(image)
        mask = self.mask_to_tensor(mask)
        image = (image * mask)[:channel]

        return (image * (0.5 - inv) * 2.0).float()

    def load_calib(self, data_dict):
        calib_data = np.loadtxt(data_dict['calib_path'], dtype=float)
        extrinsic = calib_data[:4, :4]
        intrinsic = calib_data[4:8, :4]
        calib_mat = np.matmul(intrinsic, extrinsic)
        calib_mat = torch.from_numpy(calib_mat).float()
        return {'calib': calib_mat}

    def load_mesh(self, data_dict):
        mesh_path = data_dict['mesh_path']
        scale = data_dict['scale']

        mesh_ori = trimesh.load(mesh_path,
                                skip_materials=True,
                                process=False,
                                maintain_order=True)
        verts = mesh_ori.vertices * scale
        faces = mesh_ori.faces

        vert_normals = np.array(mesh_ori.vertex_normals)
        face_normals = np.array(mesh_ori.face_normals)

        mesh = HoppeMesh(verts, faces, vert_normals, face_normals)

        return {
            'mesh': mesh,
            'verts': torch.as_tensor(mesh.verts).float(),
            'faces': torch.as_tensor(mesh.faces).long()
        }

    def add_noise(self,
                  beta_num,
                  smpl_pose,
                  smpl_betas,
                  noise_type,
                  noise_scale,
                  type,
                  hashcode):

        np.random.seed(hashcode)

        if type == 'smplx':
            noise_idx = self.noise_smplx_idx
        else:
            noise_idx = self.noise_smpl_idx

        if 'beta' in noise_type and noise_scale[noise_type.index("beta")] > 0.0:
            smpl_betas += (np.random.rand(beta_num) -
                           0.5) * 2.0 * noise_scale[noise_type.index("beta")]
            smpl_betas = smpl_betas.astype(np.float32)

        if 'pose' in noise_type and noise_scale[noise_type.index("pose")] > 0.0:
            smpl_pose[noise_idx] += (
                np.random.rand(len(noise_idx)) -
                0.5) * 2.0 * np.pi * noise_scale[noise_type.index("pose")]
            smpl_pose = smpl_pose.astype(np.float32)
        if type == 'smplx':
            return torch.as_tensor(smpl_pose[None, ...]), torch.as_tensor(smpl_betas[None, ...])
        else:
            return smpl_pose, smpl_betas

    def compute_smpl_verts(self, data_dict, noise_type=None, noise_scale=None):

        dataset = data_dict['dataset']
        smplx_dict = {}

        smplx_param = np.load(data_dict['smplx_path'], allow_pickle=True)
        smplx_pose = smplx_param["body_pose"]  # [1,63]
        smplx_betas = smplx_param["betas"]  # [1,10]
        smplx_pose, smplx_betas = self.add_noise(
            smplx_betas.shape[1],
            smplx_pose[0],
            smplx_betas[0],
            noise_type,
            noise_scale,
            type='smplx',
            hashcode=(hash(f"{data_dict['subject']}_{data_dict['rotation']}")) % (10**8))

        smplx_out, _ = load_fit_body(fitted_path=data_dict['smplx_path'],
                                     scale=self.datasets_dict[dataset]['scale'],
                                     smpl_type='smplx',
                                     smpl_gender='male',
                                     noise_dict=dict(betas=smplx_betas, body_pose=smplx_pose))

        smplx_dict.update({"type": "smplx",
                          "gender": 'male',
                           "body_pose": torch.as_tensor(smplx_pose),
                           "betas": torch.as_tensor(smplx_betas)})

        return smplx_out.vertices, smplx_dict

    def compute_voxel_verts(self,
                            data_dict,
                            noise_type=None,
                            noise_scale=None):

        smpl_param = np.load(data_dict['smpl_path'], allow_pickle=True)
        smplx_param = np.load(data_dict['smplx_path'], allow_pickle=True)

        smpl_pose = rotation_matrix_to_angle_axis(
            torch.as_tensor(smpl_param['full_pose'][0])).numpy()
        smpl_betas = smpl_param["betas"]

        smpl_path = cached_download(osp.join(self.smplx.model_dir, "smpl/SMPL_MALE.pkl"), use_auth_token=os.environ['ICON'])
        tetra_path = cached_download(osp.join(self.smplx.tedra_dir,
                              "tetra_male_adult_smpl.npz"), use_auth_token=os.environ['ICON'])

        smpl_model = TetraSMPLModel(smpl_path, tetra_path, 'adult')

        smpl_pose, smpl_betas = self.add_noise(
            smpl_model.beta_shape[0],
            smpl_pose.flatten(),
            smpl_betas[0],
            noise_type,
            noise_scale,
            type='smpl',
            hashcode=(hash(f"{data_dict['subject']}_{data_dict['rotation']}")) % (10**8))

        smpl_model.set_params(pose=smpl_pose.reshape(-1, 3),
                              beta=smpl_betas,
                              trans=smpl_param["transl"])
        
        verts = (np.concatenate([smpl_model.verts, smpl_model.verts_added],
                                axis=0) * smplx_param["scale"] + smplx_param["translation"]
                 ) * self.datasets_dict[data_dict['dataset']]['scale']
        faces = np.loadtxt(cached_download(osp.join(self.smplx.tedra_dir, "tetrahedrons_male_adult.txt"), use_auth_token=os.environ['ICON']),
                           dtype=np.int32) - 1

        pad_v_num = int(8000 - verts.shape[0])
        pad_f_num = int(25100 - faces.shape[0])

        verts = np.pad(verts, ((0, pad_v_num), (0, 0)),
                       mode='constant',
                       constant_values=0.0).astype(np.float32)
        faces = np.pad(faces, ((0, pad_f_num), (0, 0)),
                       mode='constant',
                       constant_values=0.0).astype(np.int32)
        

        return verts, faces, pad_v_num, pad_f_num

    def load_smpl(self, data_dict, vis=False):

        smplx_verts, smplx_dict = self.compute_smpl_verts(
            data_dict, self.noise_type,
            self.noise_scale)  # compute using smpl model

        smplx_verts = projection(smplx_verts, data_dict['calib']).float()
        smplx_faces = torch.as_tensor(self.smplx.faces).long()
        smplx_vis = torch.load(data_dict['vis_path']).float()
        smplx_cmap = torch.as_tensor(
            np.load(self.smplx.cmap_vert_path)).float()

        # get smpl_signs
        query_points = projection(data_dict['samples_geo'],
                                  data_dict['calib']).float()

        pts_signs = 2.0 * (check_sign(smplx_verts.unsqueeze(0),
                                      smplx_faces,
                                      query_points.unsqueeze(0)).float() - 0.5).squeeze(0)

        return_dict = {
            'smpl_verts': smplx_verts,
            'smpl_faces': smplx_faces,
            'smpl_vis': smplx_vis,
            'smpl_cmap': smplx_cmap,
            'pts_signs': pts_signs
        }
        if smplx_dict is not None:
            return_dict.update(smplx_dict)

        if vis:

            (xy, z) = torch.as_tensor(smplx_verts).to(
                self.device).split([2, 1], dim=1)
            smplx_vis = get_visibility(xy, z, torch.as_tensor(
                smplx_faces).to(self.device).long())

            T_normal_F, T_normal_B = self.render_normal(
                (smplx_verts*torch.tensor([1.0, -1.0, 1.0])).to(self.device),
                smplx_faces.to(self.device))

            return_dict.update({"T_normal_F": T_normal_F.squeeze(0),
                                "T_normal_B": T_normal_B.squeeze(0)})
            query_points = projection(data_dict['samples_geo'],
                                      data_dict['calib']).float()

            smplx_sdf, smplx_norm, smplx_cmap, smplx_vis = cal_sdf_batch(
                smplx_verts.unsqueeze(0).to(self.device),
                smplx_faces.unsqueeze(0).to(self.device),
                smplx_cmap.unsqueeze(0).to(self.device),
                smplx_vis.unsqueeze(0).to(self.device),
                query_points.unsqueeze(0).contiguous().to(self.device))

            return_dict.update({
                'smpl_feat':
                torch.cat(
                    (smplx_sdf[0].detach().cpu(),
                     smplx_cmap[0].detach().cpu(),
                     smplx_norm[0].detach().cpu(),
                     smplx_vis[0].detach().cpu()),
                    dim=1)
            })

        return return_dict

    def load_smpl_voxel(self, data_dict):

        smpl_verts, smpl_faces, pad_v_num, pad_f_num = self.compute_voxel_verts(
            data_dict, self.noise_type,
            self.noise_scale)  # compute using smpl model
        smpl_verts = projection(smpl_verts, data_dict['calib'])

        smpl_verts *= 0.5

        return {
            'voxel_verts': smpl_verts,
            'voxel_faces': smpl_faces,
            'pad_v_num': pad_v_num,
            'pad_f_num': pad_f_num
        }

    def get_sampling_geo(self, data_dict, is_valid=False, is_sdf=False):

        mesh = data_dict['mesh']
        calib = data_dict['calib']

        # Samples are around the true surface with an offset
        n_samples_surface = 4 * self.opt.num_sample_geo
        vert_ids = np.arange(mesh.verts.shape[0])
        thickness_sample_ratio = np.ones_like(vert_ids).astype(np.float32)

        thickness_sample_ratio /= thickness_sample_ratio.sum()

        samples_surface_ids = np.random.choice(vert_ids,
                                               n_samples_surface,
                                               replace=True,
                                               p=thickness_sample_ratio)

        samples_normal_ids = np.random.choice(vert_ids,
                                              self.opt.num_sample_geo // 2,
                                              replace=False,
                                              p=thickness_sample_ratio)

        surf_samples = mesh.verts[samples_normal_ids, :]
        surf_normals = mesh.vert_normals[samples_normal_ids, :]

        samples_surface = mesh.verts[samples_surface_ids, :]

        # Sampling offsets are random noise with constant scale (15cm - 20cm)
        offset = np.random.normal(scale=self.opt.sigma_geo,
                                  size=(n_samples_surface, 1))
        samples_surface += mesh.vert_normals[samples_surface_ids, :] * offset

        # Uniform samples in [-1, 1]
        calib_inv = np.linalg.inv(calib)
        n_samples_space = self.opt.num_sample_geo // 4
        samples_space_img = 2.0 * np.random.rand(n_samples_space, 3) - 1.0
        samples_space = projection(samples_space_img, calib_inv)

        # z-ray direction samples
        if self.opt.zray_type and not is_valid:
            n_samples_rayz = self.opt.ray_sample_num
            samples_surface_cube = projection(samples_surface, calib)
            samples_surface_cube_repeat = np.repeat(samples_surface_cube,
                                                    n_samples_rayz,
                                                    axis=0)

            thickness_repeat = np.repeat(0.5 *
                                         np.ones_like(samples_surface_ids),
                                         n_samples_rayz,
                                         axis=0)

            noise_repeat = np.random.normal(scale=0.40,
                                            size=(n_samples_surface *
                                                  n_samples_rayz, ))
            samples_surface_cube_repeat[:,
                                        -1] += thickness_repeat * noise_repeat
            samples_surface_rayz = projection(samples_surface_cube_repeat,
                                              calib_inv)

            samples = np.concatenate(
                [samples_surface, samples_space, samples_surface_rayz], 0)
        else:
            samples = np.concatenate([samples_surface, samples_space], 0)

        np.random.shuffle(samples)

        # labels: in->1.0; out->0.0.
        if is_sdf:
            sdfs = mesh.get_sdf(samples)
            inside_samples = samples[sdfs < 0]
            outside_samples = samples[sdfs >= 0]

            inside_sdfs = sdfs[sdfs < 0]
            outside_sdfs = sdfs[sdfs >= 0]
        else:
            inside = mesh.contains(samples)
            inside_samples = samples[inside >= 0.5]
            outside_samples = samples[inside < 0.5]

        nin = inside_samples.shape[0]

        if nin > self.opt.num_sample_geo // 2:
            inside_samples = inside_samples[:self.opt.num_sample_geo // 2]
            outside_samples = outside_samples[:self.opt.num_sample_geo // 2]
            if is_sdf:
                inside_sdfs = inside_sdfs[:self.opt.num_sample_geo // 2]
                outside_sdfs = outside_sdfs[:self.opt.num_sample_geo // 2]
        else:
            outside_samples = outside_samples[:(self.opt.num_sample_geo - nin)]
            if is_sdf:
                outside_sdfs = outside_sdfs[:(self.opt.num_sample_geo - nin)]

        if is_sdf:
            samples = np.concatenate(
                [inside_samples, outside_samples, surf_samples], 0)

            labels = np.concatenate([
                inside_sdfs, outside_sdfs, 0.0 * np.ones(surf_samples.shape[0])
            ])

            normals = np.zeros_like(samples)
            normals[-self.opt.num_sample_geo // 2:, :] = surf_normals

            # convert sdf from [-14, 130] to [0, 1]
            # outside: 0, inside: 1
            # Note: Marching cubes is defined on occupancy space (inside=1.0, outside=0.0)

            labels = -labels.clip(min=-self.sdf_clip, max=self.sdf_clip)
            labels += self.sdf_clip
            labels /= (self.sdf_clip * 2)

        else:
            samples = np.concatenate([inside_samples, outside_samples])
            labels = np.concatenate([
                np.ones(inside_samples.shape[0]),
                np.zeros(outside_samples.shape[0])
            ])

            normals = np.zeros_like(samples)

        samples = torch.from_numpy(samples).float()
        labels = torch.from_numpy(labels).float()
        normals = torch.from_numpy(normals).float()

        return {'samples_geo': samples, 'labels_geo': labels}