lib/dataset/PIFuDataset.py

# -*- coding: utf-8 -*-

# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: [email protected]

from lib.renderer.mesh import load_fit_body, compute_normal_batch
from lib.dataset.body_model import TetraSMPLModel
from lib.common.render import Render
from lib.dataset.mesh_util import *
from lib.pare.pare.utils.geometry import rotation_matrix_to_angle_axis
from lib.net.nerf_util import sample_ray_h36m, get_wsampling_points
from termcolor import colored
import os.path as osp
import numpy as np
from PIL import Image
import random
import os, cv2
import trimesh
import torch
import vedo
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import trimesh
os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1"

cape_gender = {
    "male": [
        '00032', '00096', '00122', '00127', '00145', '00215', '02474', '03284',
        '03375', '03394'
    ],
    "female": ['00134', '00159', '03223', '03331', '03383']
}



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

        self.base_keys = ["smpl_verts", "smpl_faces"]
        self.feat_names = cfg.net.smpl_feats

        self.feat_keys = self.base_keys + [f"smpl_{feat_name}" for feat_name in self.feat_names]

        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)

            mesh_dir = osp.join(dataset_dir, "scans")
            smplx_dir = osp.join(dataset_dir, "smplx")
            smpl_dir = osp.join(dataset_dir, "smpl")

            self.datasets_dict[dataset] = {
                "smplx_dir": smplx_dir,
                "smpl_dir": smpl_dir,
                "mesh_dir": mesh_dir,
                "scale": self.scales[dataset_id]
            }

            if split == 'train':
                self.datasets_dict[dataset].update(
                    {"subjects": np.loadtxt(osp.join(dataset_dir, "all.txt"), dtype=str)})
            else:
                self.datasets_dict[dataset].update(
                    {"subjects": np.loadtxt(osp.join(dataset_dir, "test.txt"), dtype=str)})

        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)

        self.UV_RENDER='/sdb/zzc/zzc/paper_models/PIFu-master/PIFu-master/training_data/UV_RENDER'
        self.UV_MASK='/sdb/zzc/zzc/paper_models/PIFu-master/PIFu-master/training_data/UV_MASK'
        self.UV_POS='/sdb/zzc/zzc/paper_models/PIFu-master/PIFu-master/training_data/UV_POS'
        self.UV_NORMAL='/sdb/zzc/zzc/paper_models/PIFu-master/PIFu-master/training_data/UV_NORMAL'
        self.IMAGE_MASK='/sdb/zzc/zzc/paper_models/PIFu-master/PIFu-master/training_data/MASK'
        self.PARAM='/sdb/zzc/zzc/paper_models/PIFu-master/PIFu-master/training_data/PARAM'
        self.depth='./data/thuman2_36views'

    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])
        if dataset=='thuman2':
            old_folder="/".join([dataset + f"_{self.opt.rotation_num}views_nosideview", subject])
        else:
            old_folder=render_folder
        # add uv map path
        # use pifu dataset
        
        uv_render_path=os.path.join(self.UV_RENDER,subject,'%d_%d_%02d.jpg'%(rotation,0,0))
        uv_mask_path = os.path.join(self.UV_MASK, subject, '%02d.png' % (0))
        uv_pos_path = os.path.join(self.UV_POS, subject, '%02d.exr' % (0))
        uv_normal_path = os.path.join(self.UV_NORMAL, subject, '%02d.png' % (0))
        image_mask_path=os.path.join(self.IMAGE_MASK,subject,'%d_%d_%02d.png'%(rotation,0,0))
        param_path=os.path.join(self.PARAM, subject,'%d_%d_%02d.npy'%(rotation,0,0))
        depth_path=os.path.join(self.depth,subject,"depth_F",'%03d.png'%(rotation))



        # setup paths
        data_dict = {
            'dataset': dataset,
            'subject': subject,
            'rotation': rotation,
            'scale': self.datasets_dict[dataset]["scale"],
            'calib_path': osp.join(self.root, render_folder, 'calib', f'{rotation:03d}.txt'),
            'image_path': osp.join(self.root, render_folder, 'render', f'{rotation:03d}.png'),
            'smpl_path': osp.join(self.datasets_dict[dataset]["smpl_dir"], f"{subject}.obj"),
            'vis_path': osp.join(self.root, old_folder, 'vis', f'{rotation:03d}.pt'),
            'uv_render_path': osp.join(self.root, render_folder, 'uv_color', f'{rotation:03d}.png'),
            'uv_mask_path': uv_mask_path,
            'uv_pos_path': uv_pos_path,
            'uv_normal_path': osp.join(self.root, render_folder, 'uv_normal', '%02d.png' % (0)),
            'image_mask_path':image_mask_path,
            'param_path':param_path,
            'depth_path':depth_path,
        }

        if dataset == 'thuman2':
            data_dict.update({
                '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}.obj"),
                    osp.join("./data/thuman2/smplx/", f"{subject}.obj"),
                'smpl_param':
                    osp.join(self.datasets_dict[dataset]["smpl_dir"], f"{subject}.pkl"),
                'smplx_param':
                    osp.join(self.datasets_dict[dataset]["smplx_dir"], f"{subject}.pkl"),
            })
        elif dataset == 'cape':
            data_dict.update({
                'mesh_path': osp.join(self.datasets_dict[dataset]["mesh_dir"], f"{subject}.obj"),
                'smpl_param': osp.join(self.datasets_dict[dataset]["smpl_dir"], f"{subject}.npz"),
            })

        # 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
            if os.path.exists(data_dict[f'{name}_path']):
                data_dict.update(
                    {name: self.imagepath2tensor(data_dict[f'{name}_path'], channel, inv=False)})
                # if name=='normal_F' and dataset == 'thuman2':
                #     right_angle=(rotation+270)%360
                #     left_angle=(rotation+90)%360
                #     normal_right_path=osp.join(self.root, render_folder, name, f'{right_angle:03d}.png')
                #     normal_left_path=osp.join(self.root, render_folder, name, f'{left_angle:03d}.png')
                #     data_dict.update(
                #         {'normal_R': self.imagepath2tensor(normal_right_path, channel, inv=False)})
                #     data_dict.update(
                #         {'normal_L': self.imagepath2tensor(normal_left_path, channel, inv=False)})
                    
                if name=='T_normal_F' and dataset == 'thuman2':
                    
                    normal_right_path=osp.join(self.root, render_folder, "T_normal_R", f'{rotation:03d}.png')
                    normal_left_path=osp.join(self.root, render_folder, "T_normal_L", f'{rotation:03d}.png')
                    data_dict.update(
                        {'T_normal_R': self.imagepath2tensor(normal_right_path, channel, inv=False)})
                    data_dict.update(
                        {'T_normal_L': self.imagepath2tensor(normal_left_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')

        # remove CAPE's noisy outliers using OpenCV's inpainting
        if "cape" in path and 'T_' not in path:
            mask = (cv2.imread(path.replace(path.split("/")[-2], "mask"), 0) > 1)
            img = np.asarray(rgba)[:, :, :3]
            fill_mask = ((mask & (img.sum(axis=2) == 0))).astype(np.uint8)
            image = Image.fromarray(
                cv2.inpaint(img * mask[..., None], fill_mask, 3, cv2.INPAINT_TELEA))
            mask = Image.fromarray(mask)
        else:
            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']

        verts, faces = obj_loader(mesh_path)

        mesh = HoppeMesh(verts * scale, faces)

        return {
            'mesh': mesh,
            'verts': torch.as_tensor(verts * scale).float(),
            'faces': torch.as_tensor(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_param'], 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_param'],
                                     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_param"], allow_pickle=True)

        if data_dict['dataset'] == 'cape':
            pid = data_dict['subject'].split("-")[0]
            gender = "male" if pid in cape_gender["male"] else "female"
            smpl_pose = smpl_param['pose'].flatten()
            smpl_betas = np.zeros((1, 10))
        else:
            gender = 'male'
            smpl_pose = rotation_matrix_to_angle_axis(torch.as_tensor(
                smpl_param["full_pose"][0])).numpy()
            smpl_betas = smpl_param["betas"]

        smpl_path = osp.join(self.smplx.model_dir, f"smpl/SMPL_{gender.upper()}.pkl")
        tetra_path = osp.join(self.smplx.tedra_dir, f"tetra_{gender}_adult_smpl.npz")

        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"])
        if data_dict['dataset'] == 'cape':
            verts = np.concatenate([smpl_model.verts, smpl_model.verts_added], axis=0) * 100.0
        else:
            verts = (np.concatenate([smpl_model.verts, smpl_model.verts_added], axis=0) *
                     smpl_param["scale"] +
                     smpl_param["translation"]) * self.datasets_dict[data_dict["dataset"]]["scale"]

        faces = (np.loadtxt(
            osp.join(self.smplx.tedra_dir, "tetrahedrons_male_adult.txt"),
            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 densely_sample(self, verts, faces):
        # TODO: subdivided the triangular mesh
        new_vertices,new_faces,index=trimesh.remesh.subdivide(verts,faces)
        
        return new_vertices, torch.LongTensor(new_faces)
        ...

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

        smpl_type = "smplx" if ('smplx_path' in data_dict.keys() and
                                os.path.exists(data_dict['smplx_path'])) else "smpl"
        
        return_dict = {}

        if 'smplx_param' in data_dict.keys() and \
            os.path.exists(data_dict['smplx_param']) and \
                sum(self.noise_scale) > 0.0:
            smplx_verts, smplx_dict = self.compute_smpl_verts(data_dict, self.noise_type,
                                                              self.noise_scale)
            smplx_faces = torch.as_tensor(self.smplx.smplx_faces).long()
            smplx_cmap = torch.as_tensor(np.load(self.smplx.cmap_vert_path)).float()

        else:
            smplx_vis = torch.load(data_dict['vis_path']).float()
            return_dict.update({'smpl_vis': smplx_vis})

            # noise_factor = 20
            # noise = np.random.normal(loc=0, scale=noise_factor)
            smplx_verts = rescale_smpl(data_dict[f"{smpl_type}_path"], scale=100.0)
            smplx_faces = torch.as_tensor(getattr(self.smplx, f"{smpl_type}_faces")).long()
            smplx_cmap = self.smplx.cmap_smpl_vids(smpl_type)

        if densely_sample:
            smplx_verts,smplx_faces=self.densely_sample(smplx_verts,smplx_faces)

        smplx_verts = projection(smplx_verts, data_dict['calib']).float()

        # get smpl_vis
        if "smpl_vis" not in return_dict.keys() and "smpl_vis" in self.feat_keys:
            (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())
            return_dict['smpl_vis'] = smplx_vis

        if "smpl_norm" not in return_dict.keys() and "smpl_norm" in self.feat_keys:
            # get smpl_norms
            smplx_norms = compute_normal_batch(smplx_verts.unsqueeze(0),
                                               smplx_faces.unsqueeze(0))[0]
            return_dict["smpl_norm"] = smplx_norms

        if "smpl_cmap" not in return_dict.keys() and "smpl_cmap" in self.feat_keys:
            return_dict["smpl_cmap"] = smplx_cmap

        sample_num=smplx_verts.shape[0]
        verts_ids=np.arange(smplx_verts.shape[0])
       
        sample_ids=torch.LongTensor(verts_ids)
        return_dict.update({
            'smpl_verts': smplx_verts,
            'smpl_faces': smplx_faces,
            'smpl_cmap': smplx_cmap,
            'smpl_sample_id':sample_ids,
        })

        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(np.array([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):

        #assert 0
        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])

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

        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

        # samples=np.concatenate([samples_surface], 0)
        # np.random.shuffle(samples)
        # 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)

        samples = np.concatenate([samples_surface, samples_space], 0)
        np.random.shuffle(samples)

        # labels: in->1.0; out->0.0.
        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]
        else:
            outside_samples = outside_samples[:(self.opt.num_sample_geo - nin)]

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

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

        # sample color
        if not self.datasets[0]=='cape':
            # get color
            uv_render_path = data_dict['uv_render_path']
            uv_mask_path = data_dict['uv_mask_path']
            uv_pos_path = data_dict['uv_pos_path']
            uv_normal_path = data_dict['uv_normal_path']

            # Segmentation mask for the uv render.
            # [H, W] bool
            uv_mask = cv2.imread(uv_mask_path)
            uv_mask = uv_mask[:, :, 0] != 0
            # UV render. each pixel is the color of the point.
            # [H, W, 3] 0 ~ 1 float
            uv_render = cv2.imread(uv_render_path)
            uv_render = cv2.cvtColor(uv_render, cv2.COLOR_BGR2RGB) / 255.0

            # Normal render. each pixel is the surface normal of the point.
            # [H, W, 3] -1 ~ 1 float
            uv_normal = cv2.imread(uv_normal_path)
            uv_normal = cv2.cvtColor(uv_normal, cv2.COLOR_BGR2RGB) / 255.0
            uv_normal = 2.0 * uv_normal - 1.0

            # Position render. each pixel is the xyz coordinates of the point
            uv_pos = cv2.imread(uv_pos_path, 2 | 4)[:, :, ::-1]
            
            ### In these few lines we flattern the masks, positions, and normals
            uv_mask = uv_mask.reshape((-1))          # 512*512
            uv_pos = uv_pos.reshape((-1, 3))
            uv_render = uv_render.reshape((-1, 3))   # 512*512,3
            uv_normal = uv_normal.reshape((-1, 3))

            surface_points = uv_pos[uv_mask]
            surface_colors = uv_render[uv_mask]
            surface_normal = uv_normal[uv_mask]

            # Samples are around the true surface with an offset
            n_samples_surface = self.opt.num_sample_color

            if n_samples_space>surface_points.shape[0]:
                print(surface_points.shape[0])
                print( uv_pos_path)
                assert 0
            sample_list = random.sample(range(0, surface_points.shape[0] - 1), n_samples_surface)
            surface_points=surface_points[sample_list].T
            surface_colors=surface_colors[sample_list].T
            surface_normal=surface_normal[sample_list].T

            # Samples are around the true surface with an offset
            normal = torch.Tensor(surface_normal).float()
            samples_surface = torch.Tensor(surface_points).float() \
                    + torch.normal(mean=torch.zeros((1, normal.size(1))), std=self.opt.sigma_color).expand_as(normal) * normal

            sample_color=samples_surface.T
            rgbs_color=(surface_colors-0.5)*2   # range -1 - 1
            rgbs_color=rgbs_color.T
            colors = torch.from_numpy(rgbs_color).float()

            

          
        # center.unsqueeze(0).float()

            return {'samples_geo': samples, 'labels_geo': labels,"samples_color":sample_color,"color_labels":colors}
        else:
            return {'samples_geo': samples, 'labels_geo': labels}
    def get_param(self,data_dict):
        W=512
        H=512
        
        param_path = data_dict['param_path']
        # loading calibration data
        param = np.load(param_path, allow_pickle=True)
        # pixel unit / world unit
        ortho_ratio = param.item().get('ortho_ratio')
        # world unit / model unit
        scale = param.item().get('scale')
        # camera center world coordinate
        center = param.item().get('center')
        # model rotation
        R = param.item().get('R')
        translate = -np.matmul(R, center).reshape(3, 1) 
        extrinsic = np.concatenate([R, translate], axis=1)
        extrinsic = np.concatenate([extrinsic, np.array([0, 0, 0, 1]).reshape(1, 4)], 0)

       
        # Match camera space to image pixel space
        scale_intrinsic = np.identity(4)
        scale_intrinsic[0, 0] = scale / ortho_ratio
        scale_intrinsic[1, 1] = -scale / ortho_ratio
        scale_intrinsic[2, 2] = scale / ortho_ratio
        # Match image pixel space to image uv space
        uv_intrinsic = np.identity(4)
        uv_intrinsic[0, 0] = 1.0 / float(W // 2)
        uv_intrinsic[1, 1] = 1.0 / float(W // 2)
        uv_intrinsic[2, 2] = 1.0 / float(W // 2)
       

        return scale_intrinsic[:3,:3],R,translate.numpy(),center

    def get_extrinsics(self,data_dict):
        calib_data = np.loadtxt(data_dict['calib_path'], dtype=float)
        extrinsic = calib_data[:4, :4]
        intrinsic = calib_data[4:8, :4]
        return intrinsic.astype(np.float32)

    def visualize_sampling3D(self, data_dict, mode='vis'):

        # create plot
        vp = vedo.Plotter(title="", size=(1500, 1500), axes=0, bg='white')
        vis_list = []

        assert mode in ['vis', 'sdf', 'normal', 'cmap', 'occ']

        # sdf-1 cmap-3 norm-3 vis-1
        if mode == 'vis':
            labels = data_dict[f'smpl_feat'][:, [-1]]  # visibility
            colors = np.concatenate([labels, labels, labels], axis=1)
        elif mode == 'occ':
            labels = data_dict[f'labels_geo'][..., None]  # occupancy
            colors = np.concatenate([labels, labels, labels], axis=1)
        elif mode == 'sdf':
            labels = data_dict[f'smpl_feat'][:, [0]]  # sdf
            labels -= labels.min()
            labels /= labels.max()
            colors = np.concatenate([labels, labels, labels], axis=1)
        elif mode == 'normal':
            labels = data_dict[f'smpl_feat'][:, -4:-1]  # normal
            colors = (labels + 1.0) * 0.5
        elif mode == 'cmap':
            labels = data_dict[f'smpl_feat'][:, -7:-4]  # colormap
            colors = np.array(labels)

        points = projection(data_dict['samples_geo'], data_dict['calib'])
        verts = projection(data_dict['verts'], data_dict['calib'])
        points[:, 1] *= -1
        
        verts[:, 1] *= -1
        
        # create a mesh
        mesh = trimesh.Trimesh(verts, data_dict['faces'], process=True)
        mesh.visual.vertex_colors = [128.0, 128.0, 128.0, 255.0]
        vis_list.append(mesh)

        if 'voxel_verts' in data_dict.keys():
            print(colored("voxel verts", "green"))
            voxel_verts = data_dict['voxel_verts'] * 2.0
            voxel_faces = data_dict['voxel_faces']
            voxel_verts[:, 1] *= -1
            voxel = trimesh.Trimesh(voxel_verts,
                                    voxel_faces[:, [0, 2, 1]],
                                    process=False,
                                    maintain_order=True)
            voxel.visual.vertex_colors = [0.0, 128.0, 0.0, 255.0]
            vis_list.append(voxel)

        if 'smpl_verts' in data_dict.keys():
            print(colored("smpl verts", "green"))
            smplx_verts = data_dict['smpl_verts']
            smplx_faces = data_dict['smpl_faces']
            smplx_verts[:, 1] *= -1
            smplx = trimesh.Trimesh(smplx_verts,
                                    smplx_faces[:, [0, 2, 1]],
                                    process=False,
                                    maintain_order=True)
            smplx.visual.vertex_colors = [128.0, 128.0, 0.0, 255.0]
            vis_list.append(smplx)

        # create a picure
        img_pos = [1.0, 0.0, -1.0,-1.0,1.0]
        for img_id, img_key in enumerate(['normal_F', 'image', 'T_normal_B','T_normal_L','T_normal_R']):
            image_arr = (data_dict[img_key].detach().cpu().permute(1, 2, 0).numpy() +
                         1.0) * 0.5 * 255.0
            image_dim = image_arr.shape[0]

            if img_id==3:
                image=vedo.Picture(image_arr).scale(2.0 / image_dim).pos(-1.0, -1.0, -1.0).rotateY(90)
            elif img_id==4:
                image=vedo.Picture(image_arr).scale(2.0 / image_dim).pos(-1.0, -1.0, 1.0).rotateY(90)
            else:
                image = vedo.Picture(image_arr).scale(2.0 / image_dim).pos(-1.0, -1.0, img_pos[img_id])
            vis_list.append(image)

        # create a pointcloud
        pc = vedo.Points(points, r=1)
        vis_list.append(pc)

        vp.show(*vis_list, bg="white", axes=1.0, interactive=True)