Upload 8 files

lib/dataset/Evaluator.py

@@ -0,0 +1,264 @@
+
+# -*- 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: ps-license@tuebingen.mpg.de
+
+
+from lib.renderer.gl.normal_render import NormalRender
+from lib.dataset.mesh_util import projection
+from lib.common.render import Render
+from PIL import Image
+import numpy as np
+import torch
+from torch import nn
+import trimesh
+import os.path as osp
+from PIL import Image
+
+
+class Evaluator:
+
+    _normal_render = None
+
+    @staticmethod
+    def init_gl():
+        Evaluator._normal_render = NormalRender(width=512, height=512)
+
+    def __init__(self, device):
+        self.device = device
+        self.render = Render(size=512, device=self.device)
+        self.error_term = nn.MSELoss()
+
+        self.offset = 0.0
+        self.scale_factor = None
+
+    def set_mesh(self, result_dict, scale_factor=1.0, offset=0.0):
+
+        for key in result_dict.keys():
+            if torch.is_tensor(result_dict[key]):
+                result_dict[key] = result_dict[key].detach().cpu().numpy()
+
+        for k, v in result_dict.items():
+            setattr(self, k, v)
+
+        self.scale_factor = scale_factor
+        self.offset = offset
+
+    def _render_normal(self, mesh, deg, norms=None):
+        view_mat = np.identity(4)
+        rz = deg / 180.0 * np.pi
+        model_mat = np.identity(4)
+        model_mat[:3, :3] = self._normal_render.euler_to_rot_mat(0, rz, 0)
+        model_mat[1, 3] = self.offset
+        view_mat[2, 2] *= -1
+
+        self._normal_render.set_matrices(view_mat, model_mat)
+        if norms is None:
+            norms = mesh.vertex_normals
+        self._normal_render.set_normal_mesh(self.scale_factor * mesh.vertices,
+                                            mesh.faces, norms, mesh.faces)
+        self._normal_render.draw()
+        normal_img = self._normal_render.get_color()
+        return normal_img
+
+    def render_mesh_list(self, mesh_lst):
+
+        self.offset = 0.0
+        self.scale_factor = 1.0
+
+        full_list = []
+        for mesh in mesh_lst:
+            row_lst = []
+            for deg in np.arange(0, 360, 90):
+                normal = self._render_normal(mesh, deg)
+                row_lst.append(normal)
+            full_list.append(np.concatenate(row_lst, axis=1))
+
+        res_array = np.concatenate(full_list, axis=0)
+
+        return res_array
+
+    def _get_reproj_normal_error(self, deg):
+
+        tgt_normal = self._render_normal(self.tgt_mesh, deg)
+        src_normal = self._render_normal(self.src_mesh, deg)
+        error = (((src_normal[:, :, :3] -
+                   tgt_normal[:, :, :3])**2).sum(axis=2).mean(axis=(0, 1)))
+
+        return error, [src_normal, tgt_normal]
+
+    def render_normal(self, verts, faces):
+
+        verts = verts[0].detach().cpu().numpy()
+        faces = faces[0].detach().cpu().numpy()
+
+        mesh_F = trimesh.Trimesh(verts * np.array([1.0, -1.0, 1.0]), faces)
+        mesh_B = trimesh.Trimesh(verts * np.array([1.0, -1.0, -1.0]), faces)
+
+        self.scale_factor = 1.0
+
+        normal_F = self._render_normal(mesh_F, 0)
+        normal_B = self._render_normal(mesh_B,
+                                       0,
+                                       norms=mesh_B.vertex_normals *
+                                       np.array([-1.0, -1.0, 1.0]))
+
+        mask = normal_F[:, :, 3:4]
+        normal_F = (torch.as_tensor(2.0 * (normal_F - 0.5) * mask).permute(
+            2, 0, 1)[:3, :, :].float().unsqueeze(0).to(self.device))
+        normal_B = (torch.as_tensor(2.0 * (normal_B - 0.5) * mask).permute(
+            2, 0, 1)[:3, :, :].float().unsqueeze(0).to(self.device))
+
+        return {"T_normal_F": normal_F, "T_normal_B": normal_B}
+
+    def calculate_normal_consist(
+        self,
+        frontal=True,
+        back=True,
+        left=True,
+        right=True,
+        save_demo_img=None,
+        return_demo=False,
+    ):
+
+        # reproj error
+        # if save_demo_img is not None, save a visualization at the given path (etc, "./test.png")
+        if self._normal_render is None:
+            print(
+                "In order to use normal render, "
+                "you have to call init_gl() before initialing any evaluator objects."
+            )
+            return -1
+
+        side_cnt = 0
+        total_error = 0
+        demo_list = []
+
+        if frontal:
+            side_cnt += 1
+            error, normal_lst = self._get_reproj_normal_error(0)
+            total_error += error
+            demo_list.append(np.concatenate(normal_lst, axis=0))
+        if back:
+            side_cnt += 1
+            error, normal_lst = self._get_reproj_normal_error(180)
+            total_error += error
+            demo_list.append(np.concatenate(normal_lst, axis=0))
+        if left:
+            side_cnt += 1
+            error, normal_lst = self._get_reproj_normal_error(90)
+            total_error += error
+            demo_list.append(np.concatenate(normal_lst, axis=0))
+        if right:
+            side_cnt += 1
+            error, normal_lst = self._get_reproj_normal_error(270)
+            total_error += error
+            demo_list.append(np.concatenate(normal_lst, axis=0))
+        if save_demo_img is not None:
+            res_array = np.concatenate(demo_list, axis=1)
+            res_img = Image.fromarray((res_array * 255).astype(np.uint8))
+            res_img.save(save_demo_img)
+
+        if return_demo:
+            res_array = np.concatenate(demo_list, axis=1)
+            return res_array
+        else:
+            return total_error
+
+    def space_transfer(self):
+
+        # convert from GT to SDF
+        self.verts_pr -= self.recon_size / 2.0
+        self.verts_pr /= self.recon_size / 2.0
+
+        self.verts_gt = projection(self.verts_gt, self.calib)
+        self.verts_gt[:, 1] *= -1
+
+        self.tgt_mesh = trimesh.Trimesh(self.verts_gt, self.faces_gt)
+        self.src_mesh = trimesh.Trimesh(self.verts_pr, self.faces_pr)
+
+        # (self.tgt_mesh+self.src_mesh).show()
+
+    def export_mesh(self, dir, name):
+        self.tgt_mesh.visual.vertex_colors = np.array([255, 0, 0])
+        self.src_mesh.visual.vertex_colors = np.array([0, 255, 0])
+
+        (self.tgt_mesh + self.src_mesh).export(
+            osp.join(dir, f"{name}_gt_pr.obj"))
+
+    def calculate_chamfer_p2s(self, sampled_points=1000):
+        """calculate the geometry metrics [chamfer, p2s, chamfer_H, p2s_H]
+
+        Args:
+            verts_gt (torch.cuda.tensor): [N, 3]
+            faces_gt (torch.cuda.tensor): [M, 3]
+            verts_pr (torch.cuda.tensor): [N', 3]
+            faces_pr (torch.cuda.tensor): [M', 3]
+            sampled_points (int, optional): use smaller number for faster testing. Defaults to 1000.
+
+        Returns:
+            tuple: chamfer, p2s, chamfer_H, p2s_H
+        """
+
+        gt_surface_pts, _ = trimesh.sample.sample_surface_even(
+            self.tgt_mesh, sampled_points)
+        pred_surface_pts, _ = trimesh.sample.sample_surface_even(
+            self.src_mesh, sampled_points)
+
+        _, dist_pred_gt, _ = trimesh.proximity.closest_point(
+            self.src_mesh, gt_surface_pts)
+        _, dist_gt_pred, _ = trimesh.proximity.closest_point(
+            self.tgt_mesh, pred_surface_pts)
+
+        dist_pred_gt[np.isnan(dist_pred_gt)] = 0
+        dist_gt_pred[np.isnan(dist_gt_pred)] = 0
+        chamfer_dist = 0.5 * (dist_pred_gt.mean() +
+                              dist_gt_pred.mean()).item() * 100
+        p2s_dist = dist_pred_gt.mean().item() * 100
+
+        return chamfer_dist, p2s_dist
+
+    def calc_acc(self, output, target, thres=0.5, use_sdf=False):
+
+        # # remove the surface points with thres
+        # non_surf_ids = (target != thres)
+        # output = output[non_surf_ids]
+        # target = target[non_surf_ids]
+
+        with torch.no_grad():
+            output = output.masked_fill(output < thres, 0.0)
+            output = output.masked_fill(output > thres, 1.0)
+
+            if use_sdf:
+                target = target.masked_fill(target < thres, 0.0)
+                target = target.masked_fill(target > thres, 1.0)
+
+            acc = output.eq(target).float().mean()
+
+            # iou, precison, recall
+            output = output > thres
+            target = target > thres
+
+            union = output | target
+            inter = output & target
+
+            _max = torch.tensor(1.0).to(output.device)
+
+            union = max(union.sum().float(), _max)
+            true_pos = max(inter.sum().float(), _max)
+            vol_pred = max(output.sum().float(), _max)
+            vol_gt = max(target.sum().float(), _max)
+
+            return acc, true_pos / union, true_pos / vol_pred, true_pos / vol_gt

lib/dataset/NormalDataset.py

@@ -0,0 +1,212 @@
+
+# -*- 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: ps-license@tuebingen.mpg.de
+
+import os.path as osp
+import numpy as np
+from PIL import Image
+import torchvision.transforms as transforms
+
+
+class NormalDataset():
+    def __init__(self, cfg, split='train'):
+
+        self.split = split
+        self.root = cfg.root
+        self.overfit = cfg.overfit
+
+        self.opt = cfg.dataset
+        self.datasets = self.opt.types
+        self.input_size = self.opt.input_size
+        self.set_splits = self.opt.set_splits
+        self.scales = self.opt.scales
+        self.pifu = self.opt.pifu
+
+        # input data types and dimensions
+        self.in_nml = [item[0] for item in cfg.net.in_nml]
+        self.in_nml_dim = [item[1] for item in cfg.net.in_nml]
+        self.in_total = self.in_nml + ['normal_F', 'normal_B']
+        self.in_total_dim = self.in_nml_dim + [3, 3]
+
+        if self.split != 'train':
+            self.rotations = range(0, 360, 120)
+        else:
+            self.rotations = np.arange(0, 360, 360 /
+                                       self.opt.rotation_num).astype(np.int)
+
+        self.datasets_dict = {}
+        for dataset_id, dataset in enumerate(self.datasets):
+            dataset_dir = osp.join(self.root, dataset, "smplx")
+            self.datasets_dict[dataset] = {
+                "subjects":
+                np.loadtxt(osp.join(self.root, dataset, "all.txt"), dtype=str),
+                "path":
+                dataset_dir,
+                "scale":
+                self.scales[dataset_id]
+            }
+
+        self.subject_list = self.get_subject_list(split)
+
+        # 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, ))
+        ])
+
+    def get_subject_list(self, split):
+
+        subject_list = []
+
+        for dataset in self.datasets:
+
+            if self.pifu:
+                txt = osp.join(self.root, dataset, f'{split}_pifu.txt')
+            else:
+                txt = osp.join(self.root, dataset, f'{split}.txt')
+
+            if osp.exists(txt):
+                print(f"load from {txt}")
+                subject_list += sorted(np.loadtxt(txt, dtype=str).tolist())
+
+                if self.pifu:
+                    miss_pifu = sorted(
+                        np.loadtxt(osp.join(self.root, dataset,
+                                            "miss_pifu.txt"),
+                                   dtype=str).tolist())
+                    subject_list = [
+                        subject for subject in subject_list
+                        if subject not in miss_pifu
+                    ]
+                    subject_list = [
+                        "renderpeople/" + subject for subject in subject_list
+                    ]
+
+            else:
+                train_txt = osp.join(self.root, dataset, 'train.txt')
+                val_txt = osp.join(self.root, dataset, 'val.txt')
+                test_txt = osp.join(self.root, dataset, 'test.txt')
+
+                print(
+                    f"generate lists of [train, val, test] \n {train_txt} \n {val_txt} \n {test_txt} \n"
+                )
+
+                split_txt = osp.join(self.root, dataset, f'{split}.txt')
+
+                subjects = self.datasets_dict[dataset]['subjects']
+                train_split = int(len(subjects) * self.set_splits[0])
+                val_split = int(
+                    len(subjects) * self.set_splits[1]) + train_split
+
+                with open(train_txt, "w") as f:
+                    f.write("\n".join(dataset + "/" + item
+                                      for item in subjects[:train_split]))
+                with open(val_txt, "w") as f:
+                    f.write("\n".join(
+                        dataset + "/" + item
+                        for item in subjects[train_split:val_split]))
+                with open(test_txt, "w") as f:
+                    f.write("\n".join(dataset + "/" + item
+                                      for item in subjects[val_split:]))
+
+                subject_list += sorted(
+                    np.loadtxt(split_txt, dtype=str).tolist())
+
+        bug_list = sorted(
+            np.loadtxt(osp.join(self.root, 'bug.txt'), dtype=str).tolist())
+
+        subject_list = [
+            subject for subject in subject_list if (subject not in bug_list)
+        ]
+
+        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]
+
+        # choose specific test sets
+        subject = self.subject_list[mid]
+
+        subject_render = "/".join(
+            [subject.split("/")[0] + "_12views",
+             subject.split("/")[1]])
+
+        # setup paths
+        data_dict = {
+            'dataset':
+            subject.split("/")[0],
+            'subject':
+            subject,
+            'rotation':
+            rotation,
+            'image_path':
+            osp.join(self.root, subject_render, 'render',
+                     f'{rotation:03d}.png')
+        }
+
+        # image/normal/depth loader
+        for name, channel in zip(self.in_total, self.in_total_dim):
+
+            if name != 'image':
+                data_dict.update({
+                    f'{name}_path':
+                    osp.join(self.root, subject_render, name,
+                             f'{rotation:03d}.png')
+                })
+            data_dict.update({
+                name:
+                self.imagepath2tensor(data_dict[f'{name}_path'],
+                                      channel,
+                                      inv='depth_B' in name)
+            })
+
+        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()

lib/dataset/NormalModule.py

@@ -0,0 +1,94 @@
+
+# -*- 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: ps-license@tuebingen.mpg.de
+
+import numpy as np
+from torch.utils.data import DataLoader
+from .NormalDataset import NormalDataset
+
+# pytorch lightning related libs
+import pytorch_lightning as pl
+
+
+class NormalModule(pl.LightningDataModule):
+    def __init__(self, cfg):
+        super(NormalModule, self).__init__()
+        self.cfg = cfg
+        self.overfit = self.cfg.overfit
+
+        if self.overfit:
+            self.batch_size = 1
+        else:
+            self.batch_size = self.cfg.batch_size
+
+        self.data_size = {}
+
+    def prepare_data(self):
+
+        pass
+
+    @staticmethod
+    def worker_init_fn(worker_id):
+        np.random.seed(np.random.get_state()[1][0] + worker_id)
+
+    def setup(self, stage):
+
+        if stage == 'fit' or stage is None:
+            self.train_dataset = NormalDataset(cfg=self.cfg, split="train")
+            self.val_dataset = NormalDataset(cfg=self.cfg, split="val")
+            self.data_size = {
+                'train': len(self.train_dataset),
+                'val': len(self.val_dataset)
+            }
+
+        if stage == 'test' or stage is None:
+            self.test_dataset = NormalDataset(cfg=self.cfg, split="test")
+
+    def train_dataloader(self):
+
+        train_data_loader = DataLoader(self.train_dataset,
+                                       batch_size=self.batch_size,
+                                       shuffle=not self.overfit,
+                                       num_workers=self.cfg.num_threads,
+                                       pin_memory=True,
+                                       worker_init_fn=self.worker_init_fn)
+
+        return train_data_loader
+
+    def val_dataloader(self):
+
+        if self.overfit:
+            current_dataset = self.train_dataset
+        else:
+            current_dataset = self.val_dataset
+
+        val_data_loader = DataLoader(current_dataset,
+                                     batch_size=self.batch_size,
+                                     shuffle=False,
+                                     num_workers=self.cfg.num_threads,
+                                     pin_memory=True)
+
+        return val_data_loader
+
+    def test_dataloader(self):
+
+        test_data_loader = DataLoader(self.test_dataset,
+                                      batch_size=1,
+                                      shuffle=False,
+                                      num_workers=self.cfg.num_threads,
+                                      pin_memory=True)
+
+        return test_data_loader

lib/dataset/PIFuDataModule.py

@@ -0,0 +1,71 @@
+import numpy as np
+from torch.utils.data import DataLoader
+from .PIFuDataset import PIFuDataset
+import pytorch_lightning as pl
+
+
+class PIFuDataModule(pl.LightningDataModule):
+    def __init__(self, cfg):
+        super(PIFuDataModule, self).__init__()
+        self.cfg = cfg
+        self.overfit = self.cfg.overfit
+
+        if self.overfit:
+            self.batch_size = 1
+        else:
+            self.batch_size = self.cfg.batch_size
+
+        self.data_size = {}
+
+    def prepare_data(self):
+
+        pass
+
+    @staticmethod
+    def worker_init_fn(worker_id):
+        np.random.seed(np.random.get_state()[1][0] + worker_id)
+
+    def setup(self, stage):
+
+        if stage == 'fit':
+            self.train_dataset = PIFuDataset(cfg=self.cfg, split="train")
+            self.val_dataset = PIFuDataset(cfg=self.cfg, split="val")
+            self.data_size = {'train': len(self.train_dataset),
+                              'val': len(self.val_dataset)}
+
+        if stage == 'test':
+            self.test_dataset = PIFuDataset(cfg=self.cfg, split="test")
+
+    def train_dataloader(self):
+
+        train_data_loader = DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size, shuffle=True,
+            num_workers=self.cfg.num_threads, pin_memory=True,
+            worker_init_fn=self.worker_init_fn)
+
+        return train_data_loader
+
+    def val_dataloader(self):
+
+        if self.overfit:
+            current_dataset = self.train_dataset
+        else:
+            current_dataset = self.val_dataset
+
+        val_data_loader = DataLoader(
+            current_dataset,
+            batch_size=1, shuffle=False,
+            num_workers=self.cfg.num_threads, pin_memory=True,
+            worker_init_fn=self.worker_init_fn)
+
+        return val_data_loader
+
+    def test_dataloader(self):
+
+        test_data_loader = DataLoader(
+            self.test_dataset,
+            batch_size=1, shuffle=False,
+            num_workers=self.cfg.num_threads, pin_memory=True)
+
+        return test_data_loader

lib/dataset/PIFuDataset.py

@@ -0,0 +1,589 @@
+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}

lib/dataset/TestDataset.py

@@ -0,0 +1,254 @@
+
+# -*- 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: ps-license@tuebingen.mpg.de
+
+import os
+
+import lib.smplx as smplx
+from lib.pymaf.utils.geometry import rotation_matrix_to_angle_axis, batch_rodrigues
+from lib.pymaf.utils.imutils import process_image
+from lib.pymaf.core import path_config
+from lib.pymaf.models import pymaf_net
+from lib.common.config import cfg
+from lib.common.render import Render
+from lib.dataset.body_model import TetraSMPLModel
+from lib.dataset.mesh_util import get_visibility, SMPLX
+import os.path as osp
+import torch
+import numpy as np
+import random
+from termcolor import colored
+from PIL import ImageFile
+from huggingface_hub import cached_download
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+class TestDataset():
+    def __init__(self, cfg, device):
+
+        random.seed(1993)
+
+        self.image_path = cfg['image_path']
+        self.seg_dir = cfg['seg_dir']
+        self.has_det = cfg['has_det']
+        self.hps_type = cfg['hps_type']
+        self.smpl_type = 'smpl' if cfg['hps_type'] != 'pixie' else 'smplx'
+        self.smpl_gender = 'neutral'
+
+        self.device = device
+
+        self.subject_list = [self.image_path]
+
+        # smpl related
+        self.smpl_data = SMPLX()
+
+        self.get_smpl_model = lambda smpl_type, smpl_gender: smplx.create(
+            model_path=self.smpl_data.model_dir,
+            gender=smpl_gender,
+            model_type=smpl_type,
+            ext='npz')
+
+        # Load SMPL model
+        self.smpl_model = self.get_smpl_model(
+            self.smpl_type, self.smpl_gender).to(self.device)
+        self.faces = self.smpl_model.faces
+
+        self.hps = pymaf_net(path_config.SMPL_MEAN_PARAMS,
+                                pretrained=True).to(self.device)
+        self.hps.load_state_dict(torch.load(
+            path_config.CHECKPOINT_FILE)['model'],
+            strict=True)
+        self.hps.eval()
+
+        print(colored(f"Using {self.hps_type} as HPS Estimator\n", "green"))
+
+        self.render = Render(size=512, device=device)
+
+    def __len__(self):
+        return len(self.subject_list)
+
+    def compute_vis_cmap(self, smpl_verts, smpl_faces):
+
+        (xy, z) = torch.as_tensor(smpl_verts).split([2, 1], dim=1)
+        smpl_vis = get_visibility(xy, -z, torch.as_tensor(smpl_faces).long())
+        if self.smpl_type == 'smpl':
+            smplx_ind = self.smpl_data.smpl2smplx(np.arange(smpl_vis.shape[0]))
+        else:
+            smplx_ind = np.arange(smpl_vis.shape[0])
+        smpl_cmap = self.smpl_data.get_smpl_mat(smplx_ind)
+
+        return {
+            'smpl_vis': smpl_vis.unsqueeze(0).to(self.device),
+            'smpl_cmap': smpl_cmap.unsqueeze(0).to(self.device),
+            'smpl_verts': smpl_verts.unsqueeze(0)
+        }
+
+    def compute_voxel_verts(self, body_pose, global_orient, betas, trans,
+                            scale):
+
+        smpl_path = cached_download(osp.join(self.smpl_data.model_dir, "smpl/SMPL_NEUTRAL.pkl"), use_auth_token=os.environ['ICON'])
+        tetra_path = cached_download(osp.join(self.smpl_data.tedra_dir,
+                              'tetra_neutral_adult_smpl.npz'), use_auth_token=os.environ['ICON'])
+        smpl_model = TetraSMPLModel(smpl_path, tetra_path, 'adult')
+
+        pose = torch.cat([global_orient[0], body_pose[0]], dim=0)
+        smpl_model.set_params(rotation_matrix_to_angle_axis(pose),
+                              beta=betas[0])
+
+        verts = np.concatenate(
+            [smpl_model.verts, smpl_model.verts_added],
+            axis=0) * scale.item() + trans.detach().cpu().numpy()
+        faces = np.loadtxt(cached_download(osp.join(self.smpl_data.tedra_dir,
+                                    'tetrahedrons_neutral_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) * 0.5
+        faces = np.pad(faces, ((0, pad_f_num), (0, 0)),
+                       mode='constant',
+                       constant_values=0.0).astype(np.int32)
+
+        verts[:, 2] *= -1.0
+
+        voxel_dict = {
+            'voxel_verts':
+            torch.from_numpy(verts).to(self.device).unsqueeze(0).float(),
+            'voxel_faces':
+            torch.from_numpy(faces).to(self.device).unsqueeze(0).long(),
+            'pad_v_num':
+            torch.tensor(pad_v_num).to(self.device).unsqueeze(0).long(),
+            'pad_f_num':
+            torch.tensor(pad_f_num).to(self.device).unsqueeze(0).long()
+        }
+
+        return voxel_dict
+
+    def __getitem__(self, index):
+
+        img_path = self.subject_list[index]
+        img_name = img_path.split("/")[-1].rsplit(".", 1)[0]
+
+        if self.seg_dir is None:
+            img_icon, img_hps, img_ori, img_mask, uncrop_param = process_image(
+                img_path, self.hps_type, 512, self.device)
+
+            data_dict = {
+                'name': img_name,
+                'image': img_icon.to(self.device).unsqueeze(0),
+                'ori_image': img_ori,
+                'mask': img_mask,
+                'uncrop_param': uncrop_param
+            }
+
+        else:
+            img_icon, img_hps, img_ori, img_mask, uncrop_param, segmentations = process_image(
+                img_path, self.hps_type, 512, self.device,
+                seg_path=os.path.join(self.seg_dir, f'{img_name}.json'))
+            data_dict = {
+                'name': img_name,
+                'image': img_icon.to(self.device).unsqueeze(0),
+                'ori_image': img_ori,
+                'mask': img_mask,
+                'uncrop_param': uncrop_param,
+                'segmentations': segmentations
+            }
+
+        with torch.no_grad():
+            # import ipdb; ipdb.set_trace()
+            preds_dict = self.hps.forward(img_hps)
+
+        data_dict['smpl_faces'] = torch.Tensor(
+            self.faces.astype(np.int16)).long().unsqueeze(0).to(
+                self.device)
+
+        if self.hps_type == 'pymaf':
+            output = preds_dict['smpl_out'][-1]
+            scale, tranX, tranY = output['theta'][0, :3]
+            data_dict['betas'] = output['pred_shape']
+            data_dict['body_pose'] = output['rotmat'][:, 1:]
+            data_dict['global_orient'] = output['rotmat'][:, 0:1]
+            data_dict['smpl_verts'] = output['verts']
+
+        elif self.hps_type == 'pare':
+            data_dict['body_pose'] = preds_dict['pred_pose'][:, 1:]
+            data_dict['global_orient'] = preds_dict['pred_pose'][:, 0:1]
+            data_dict['betas'] = preds_dict['pred_shape']
+            data_dict['smpl_verts'] = preds_dict['smpl_vertices']
+            scale, tranX, tranY = preds_dict['pred_cam'][0, :3]
+
+        elif self.hps_type == 'pixie':
+            data_dict.update(preds_dict)
+            data_dict['body_pose'] = preds_dict['body_pose']
+            data_dict['global_orient'] = preds_dict['global_pose']
+            data_dict['betas'] = preds_dict['shape']
+            data_dict['smpl_verts'] = preds_dict['vertices']
+            scale, tranX, tranY = preds_dict['cam'][0, :3]
+
+        elif self.hps_type == 'hybrik':
+            data_dict['body_pose'] = preds_dict['pred_theta_mats'][:, 1:]
+            data_dict['global_orient'] = preds_dict['pred_theta_mats'][:, [0]]
+            data_dict['betas'] = preds_dict['pred_shape']
+            data_dict['smpl_verts'] = preds_dict['pred_vertices']
+            scale, tranX, tranY = preds_dict['pred_camera'][0, :3]
+            scale = scale * 2
+
+        elif self.hps_type == 'bev':
+            data_dict['betas'] = torch.from_numpy(preds_dict['smpl_betas'])[
+                [0], :10].to(self.device).float()
+            pred_thetas = batch_rodrigues(torch.from_numpy(
+                preds_dict['smpl_thetas'][0]).reshape(-1, 3)).float()
+            data_dict['body_pose'] = pred_thetas[1:][None].to(self.device)
+            data_dict['global_orient'] = pred_thetas[[0]][None].to(self.device)
+            data_dict['smpl_verts'] = torch.from_numpy(
+                preds_dict['verts'][[0]]).to(self.device).float()
+            tranX = preds_dict['cam_trans'][0, 0]
+            tranY = preds_dict['cam'][0, 1] + 0.28
+            scale = preds_dict['cam'][0, 0] * 1.1
+
+        data_dict['scale'] = scale
+        data_dict['trans'] = torch.tensor(
+            [tranX, tranY, 0.0]).to(self.device).float()
+
+        # data_dict info (key-shape):
+        # scale, tranX, tranY - tensor.float
+        # betas - [1,10] / [1, 200]
+        # body_pose - [1, 23, 3, 3] / [1, 21, 3, 3]
+        # global_orient - [1, 1, 3, 3]
+        # smpl_verts - [1, 6890, 3] / [1, 10475, 3]
+        
+        # from rot_mat to rot_6d for better optimization
+        N_body = data_dict["body_pose"].shape[1]
+        data_dict["body_pose"] =  data_dict["body_pose"][:, :, :, :2].reshape(1, N_body,-1)
+        data_dict["global_orient"] =  data_dict["global_orient"][:, :, :, :2].reshape(1, 1,-1)
+
+        return data_dict
+
+    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 render_depth(self, verts, faces):
+    
+        # render optimized mesh (normal, T_normal, image [-1,1])
+        self.render.load_meshes(verts, faces)
+        return self.render.get_depth_map(cam_ids=[0, 2])

lib/dataset/body_model.py

@@ -0,0 +1,494 @@
+
+# -*- 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: ps-license@tuebingen.mpg.de
+
+import numpy as np
+import pickle
+import torch
+import os
+
+
+class SMPLModel():
+    def __init__(self, model_path, age):
+        """
+        SMPL model.
+
+        Parameter:
+        ---------
+        model_path: Path to the SMPL model parameters, pre-processed by
+        `preprocess.py`.
+
+        """
+        with open(model_path, 'rb') as f:
+            params = pickle.load(f, encoding='latin1')
+
+            self.J_regressor = params['J_regressor']
+            self.weights = np.asarray(params['weights'])
+            self.posedirs = np.asarray(params['posedirs'])
+            self.v_template = np.asarray(params['v_template'])
+            self.shapedirs = np.asarray(params['shapedirs'])
+            self.faces = np.asarray(params['f'])
+            self.kintree_table = np.asarray(params['kintree_table'])
+
+        self.pose_shape = [24, 3]
+        self.beta_shape = [10]
+        self.trans_shape = [3]
+
+        if age == 'kid':
+            v_template_smil = np.load(
+                os.path.join(os.path.dirname(model_path),
+                             "smpl/smpl_kid_template.npy"))
+            v_template_smil -= np.mean(v_template_smil, axis=0)
+            v_template_diff = np.expand_dims(v_template_smil - self.v_template,
+                                             axis=2)
+            self.shapedirs = np.concatenate(
+                (self.shapedirs[:, :, :self.beta_shape[0]], v_template_diff),
+                axis=2)
+            self.beta_shape[0] += 1
+
+        id_to_col = {
+            self.kintree_table[1, i]: i
+            for i in range(self.kintree_table.shape[1])
+        }
+        self.parent = {
+            i: id_to_col[self.kintree_table[0, i]]
+            for i in range(1, self.kintree_table.shape[1])
+        }
+
+        self.pose = np.zeros(self.pose_shape)
+        self.beta = np.zeros(self.beta_shape)
+        self.trans = np.zeros(self.trans_shape)
+
+        self.verts = None
+        self.J = None
+        self.R = None
+        self.G = None
+
+        self.update()
+
+    def set_params(self, pose=None, beta=None, trans=None):
+        """
+        Set pose, shape, and/or translation parameters of SMPL model. Verices of the
+        model will be updated and returned.
+
+        Prameters:
+        ---------
+        pose: Also known as 'theta', a [24,3] matrix indicating child joint rotation
+        relative to parent joint. For root joint it's global orientation.
+        Represented in a axis-angle format.
+
+        beta: Parameter for model shape. A vector of shape [10]. Coefficients for
+        PCA component. Only 10 components were released by MPI.
+
+        trans: Global translation of shape [3].
+
+        Return:
+        ------
+        Updated vertices.
+
+        """
+        if pose is not None:
+            self.pose = pose
+        if beta is not None:
+            self.beta = beta
+        if trans is not None:
+            self.trans = trans
+        self.update()
+        return self.verts
+
+    def update(self):
+        """
+        Called automatically when parameters are updated.
+
+        """
+        # how beta affect body shape
+        v_shaped = self.shapedirs.dot(self.beta) + self.v_template
+        # joints location
+        self.J = self.J_regressor.dot(v_shaped)
+        pose_cube = self.pose.reshape((-1, 1, 3))
+        # rotation matrix for each joint
+        self.R = self.rodrigues(pose_cube)
+        I_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0),
+                                 (self.R.shape[0] - 1, 3, 3))
+        lrotmin = (self.R[1:] - I_cube).ravel()
+        # how pose affect body shape in zero pose
+        v_posed = v_shaped + self.posedirs.dot(lrotmin)
+        # world transformation of each joint
+        G = np.empty((self.kintree_table.shape[1], 4, 4))
+        G[0] = self.with_zeros(
+            np.hstack((self.R[0], self.J[0, :].reshape([3, 1]))))
+        for i in range(1, self.kintree_table.shape[1]):
+            G[i] = G[self.parent[i]].dot(
+                self.with_zeros(
+                    np.hstack([
+                        self.R[i],
+                        ((self.J[i, :] - self.J[self.parent[i], :]).reshape(
+                            [3, 1]))
+                    ])))
+        # remove the transformation due to the rest pose
+        G = G - self.pack(
+            np.matmul(
+                G,
+                np.hstack([self.J, np.zeros([24, 1])]).reshape([24, 4, 1])))
+        # transformation of each vertex
+        T = np.tensordot(self.weights, G, axes=[[1], [0]])
+        rest_shape_h = np.hstack((v_posed, np.ones([v_posed.shape[0], 1])))
+        v = np.matmul(T, rest_shape_h.reshape([-1, 4, 1])).reshape([-1,
+                                                                    4])[:, :3]
+        self.verts = v + self.trans.reshape([1, 3])
+        self.G = G
+
+    def rodrigues(self, r):
+        """
+        Rodrigues' rotation formula that turns axis-angle vector into rotation
+        matrix in a batch-ed manner.
+
+        Parameter:
+        ----------
+        r: Axis-angle rotation vector of shape [batch_size, 1, 3].
+
+        Return:
+        -------
+        Rotation matrix of shape [batch_size, 3, 3].
+
+        """
+        theta = np.linalg.norm(r, axis=(1, 2), keepdims=True)
+        # avoid zero divide
+        theta = np.maximum(theta, np.finfo(np.float64).tiny)
+        r_hat = r / theta
+        cos = np.cos(theta)
+        z_stick = np.zeros(theta.shape[0])
+        m = np.dstack([
+            z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
+            -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick
+        ]).reshape([-1, 3, 3])
+        i_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0),
+                                 [theta.shape[0], 3, 3])
+        A = np.transpose(r_hat, axes=[0, 2, 1])
+        B = r_hat
+        dot = np.matmul(A, B)
+        R = cos * i_cube + (1 - cos) * dot + np.sin(theta) * m
+        return R
+
+    def with_zeros(self, x):
+        """
+        Append a [0, 0, 0, 1] vector to a [3, 4] matrix.
+
+        Parameter:
+        ---------
+        x: Matrix to be appended.
+
+        Return:
+        ------
+        Matrix after appending of shape [4,4]
+
+        """
+        return np.vstack((x, np.array([[0.0, 0.0, 0.0, 1.0]])))
+
+    def pack(self, x):
+        """
+        Append zero matrices of shape [4, 3] to vectors of [4, 1] shape in a batched
+        manner.
+
+        Parameter:
+        ----------
+        x: Matrices to be appended of shape [batch_size, 4, 1]
+
+        Return:
+        ------
+        Matrix of shape [batch_size, 4, 4] after appending.
+
+        """
+        return np.dstack((np.zeros((x.shape[0], 4, 3)), x))
+
+    def save_to_obj(self, path):
+        """
+        Save the SMPL model into .obj file.
+
+        Parameter:
+        ---------
+        path: Path to save.
+
+        """
+        with open(path, 'w') as fp:
+            for v in self.verts:
+                fp.write('v %f %f %f\n' % (v[0], v[1], v[2]))
+            for f in self.faces + 1:
+                fp.write('f %d %d %d\n' % (f[0], f[1], f[2]))
+
+
+class TetraSMPLModel():
+    def __init__(self,
+                 model_path,
+                 model_addition_path,
+                 age='adult',
+                 v_template=None):
+        """
+        SMPL model.
+
+        Parameter:
+        ---------
+        model_path: Path to the SMPL model parameters, pre-processed by
+        `preprocess.py`.
+
+        """
+        with open(model_path, 'rb') as f:
+            params = pickle.load(f, encoding='latin1')
+
+            self.J_regressor = params['J_regressor']
+            self.weights = np.asarray(params['weights'])
+            self.posedirs = np.asarray(params['posedirs'])
+
+            if v_template is not None:
+                self.v_template = v_template
+            else:
+                self.v_template = np.asarray(params['v_template'])
+
+            self.shapedirs = np.asarray(params['shapedirs'])
+            self.faces = np.asarray(params['f'])
+            self.kintree_table = np.asarray(params['kintree_table'])
+
+        params_added = np.load(model_addition_path)
+        self.v_template_added = params_added['v_template_added']
+        self.weights_added = params_added['weights_added']
+        self.shapedirs_added = params_added['shapedirs_added']
+        self.posedirs_added = params_added['posedirs_added']
+        self.tetrahedrons = params_added['tetrahedrons']
+
+        id_to_col = {
+            self.kintree_table[1, i]: i
+            for i in range(self.kintree_table.shape[1])
+        }
+        self.parent = {
+            i: id_to_col[self.kintree_table[0, i]]
+            for i in range(1, self.kintree_table.shape[1])
+        }
+
+        self.pose_shape = [24, 3]
+        self.beta_shape = [10]
+        self.trans_shape = [3]
+
+        if age == 'kid':
+            v_template_smil = np.load(
+                os.path.join(os.path.dirname(model_path),
+                             "smpl/smpl_kid_template.npy"))
+            v_template_smil -= np.mean(v_template_smil, axis=0)
+            v_template_diff = np.expand_dims(v_template_smil - self.v_template,
+                                             axis=2)
+            self.shapedirs = np.concatenate(
+                (self.shapedirs[:, :, :self.beta_shape[0]], v_template_diff),
+                axis=2)
+            self.beta_shape[0] += 1
+
+        self.pose = np.zeros(self.pose_shape)
+        self.beta = np.zeros(self.beta_shape)
+        self.trans = np.zeros(self.trans_shape)
+
+        self.verts = None
+        self.verts_added = None
+        self.J = None
+        self.R = None
+        self.G = None
+
+        self.update()
+
+    def set_params(self, pose=None, beta=None, trans=None):
+        """
+        Set pose, shape, and/or translation parameters of SMPL model. Verices of the
+        model will be updated and returned.
+
+        Prameters:
+        ---------
+        pose: Also known as 'theta', a [24,3] matrix indicating child joint rotation
+        relative to parent joint. For root joint it's global orientation.
+        Represented in a axis-angle format.
+
+        beta: Parameter for model shape. A vector of shape [10]. Coefficients for
+        PCA component. Only 10 components were released by MPI.
+
+        trans: Global translation of shape [3].
+
+        Return:
+        ------
+        Updated vertices.
+
+        """
+
+        if torch.is_tensor(pose):
+            pose = pose.detach().cpu().numpy()
+        if torch.is_tensor(beta):
+            beta = beta.detach().cpu().numpy()
+
+        if pose is not None:
+            self.pose = pose
+        if beta is not None:
+            self.beta = beta
+        if trans is not None:
+            self.trans = trans
+        self.update()
+        return self.verts
+
+    def update(self):
+        """
+        Called automatically when parameters are updated.
+
+        """
+        # how beta affect body shape
+        v_shaped = self.shapedirs.dot(self.beta) + self.v_template
+        v_shaped_added = self.shapedirs_added.dot(
+            self.beta) + self.v_template_added
+        # joints location
+        self.J = self.J_regressor.dot(v_shaped)
+        pose_cube = self.pose.reshape((-1, 1, 3))
+        # rotation matrix for each joint
+        self.R = self.rodrigues(pose_cube)
+        I_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0),
+                                 (self.R.shape[0] - 1, 3, 3))
+        lrotmin = (self.R[1:] - I_cube).ravel()
+        # how pose affect body shape in zero pose
+        v_posed = v_shaped + self.posedirs.dot(lrotmin)
+        v_posed_added = v_shaped_added + self.posedirs_added.dot(lrotmin)
+        # world transformation of each joint
+        G = np.empty((self.kintree_table.shape[1], 4, 4))
+        G[0] = self.with_zeros(
+            np.hstack((self.R[0], self.J[0, :].reshape([3, 1]))))
+        for i in range(1, self.kintree_table.shape[1]):
+            G[i] = G[self.parent[i]].dot(
+                self.with_zeros(
+                    np.hstack([
+                        self.R[i],
+                        ((self.J[i, :] - self.J[self.parent[i], :]).reshape(
+                            [3, 1]))
+                    ])))
+        # remove the transformation due to the rest pose
+        G = G - self.pack(
+            np.matmul(
+                G,
+                np.hstack([self.J, np.zeros([24, 1])]).reshape([24, 4, 1])))
+        self.G = G
+        # transformation of each vertex
+        T = np.tensordot(self.weights, G, axes=[[1], [0]])
+        rest_shape_h = np.hstack((v_posed, np.ones([v_posed.shape[0], 1])))
+        v = np.matmul(T, rest_shape_h.reshape([-1, 4, 1])).reshape([-1,
+                                                                    4])[:, :3]
+        self.verts = v + self.trans.reshape([1, 3])
+        T_added = np.tensordot(self.weights_added, G, axes=[[1], [0]])
+        rest_shape_added_h = np.hstack(
+            (v_posed_added, np.ones([v_posed_added.shape[0], 1])))
+        v_added = np.matmul(T_added,
+                            rest_shape_added_h.reshape([-1, 4,
+                                                        1])).reshape([-1, 4
+                                                                      ])[:, :3]
+        self.verts_added = v_added + self.trans.reshape([1, 3])
+
+    def rodrigues(self, r):
+        """
+        Rodrigues' rotation formula that turns axis-angle vector into rotation
+        matrix in a batch-ed manner.
+
+        Parameter:
+        ----------
+        r: Axis-angle rotation vector of shape [batch_size, 1, 3].
+
+        Return:
+        -------
+        Rotation matrix of shape [batch_size, 3, 3].
+
+        """
+        theta = np.linalg.norm(r, axis=(1, 2), keepdims=True)
+        # avoid zero divide
+        theta = np.maximum(theta, np.finfo(np.float64).tiny)
+        r_hat = r / theta
+        cos = np.cos(theta)
+        z_stick = np.zeros(theta.shape[0])
+        m = np.dstack([
+            z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
+            -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick
+        ]).reshape([-1, 3, 3])
+        i_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0),
+                                 [theta.shape[0], 3, 3])
+        A = np.transpose(r_hat, axes=[0, 2, 1])
+        B = r_hat
+        dot = np.matmul(A, B)
+        R = cos * i_cube + (1 - cos) * dot + np.sin(theta) * m
+        return R
+
+    def with_zeros(self, x):
+        """
+        Append a [0, 0, 0, 1] vector to a [3, 4] matrix.
+
+        Parameter:
+        ---------
+        x: Matrix to be appended.
+
+        Return:
+        ------
+        Matrix after appending of shape [4,4]
+
+        """
+        return np.vstack((x, np.array([[0.0, 0.0, 0.0, 1.0]])))
+
+    def pack(self, x):
+        """
+        Append zero matrices of shape [4, 3] to vectors of [4, 1] shape in a batched
+        manner.
+
+        Parameter:
+        ----------
+        x: Matrices to be appended of shape [batch_size, 4, 1]
+
+        Return:
+        ------
+        Matrix of shape [batch_size, 4, 4] after appending.
+
+        """
+        return np.dstack((np.zeros((x.shape[0], 4, 3)), x))
+
+    def save_mesh_to_obj(self, path):
+        """
+        Save the SMPL model into .obj file.
+
+        Parameter:
+        ---------
+        path: Path to save.
+
+        """
+        with open(path, 'w') as fp:
+            for v in self.verts:
+                fp.write('v %f %f %f\n' % (v[0], v[1], v[2]))
+            for f in self.faces + 1:
+                fp.write('f %d %d %d\n' % (f[0], f[1], f[2]))
+
+    def save_tetrahedron_to_obj(self, path):
+        """
+        Save the tetrahedron SMPL model into .obj file.
+
+        Parameter:
+        ---------
+        path: Path to save.
+
+        """
+
+        with open(path, 'w') as fp:
+            for v in self.verts:
+                fp.write('v %f %f %f 1 0 0\n' % (v[0], v[1], v[2]))
+            for va in self.verts_added:
+                fp.write('v %f %f %f 0 0 1\n' % (va[0], va[1], va[2]))
+            for t in self.tetrahedrons + 1:
+                fp.write('f %d %d %d\n' % (t[0], t[2], t[1]))
+                fp.write('f %d %d %d\n' % (t[0], t[3], t[2]))
+                fp.write('f %d %d %d\n' % (t[0], t[1], t[3]))
+                fp.write('f %d %d %d\n' % (t[1], t[2], t[3]))

lib/dataset/hoppeMesh.py

@@ -0,0 +1,116 @@
+
+# -*- 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: ps-license@tuebingen.mpg.de
+
+import numpy as np
+from scipy.spatial import cKDTree
+import trimesh
+
+import logging
+
+logging.getLogger("trimesh").setLevel(logging.ERROR)
+
+
+def save_obj_mesh(mesh_path, verts, faces):
+    file = open(mesh_path, 'w')
+    for v in verts:
+        file.write('v %.4f %.4f %.4f\n' % (v[0], v[1], v[2]))
+    for f in faces:
+        f_plus = f + 1
+        file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2]))
+    file.close()
+
+
+def save_obj_mesh_with_color(mesh_path, verts, faces, colors):
+    file = open(mesh_path, 'w')
+
+    for idx, v in enumerate(verts):
+        c = colors[idx]
+        file.write('v %.4f %.4f %.4f %.4f %.4f %.4f\n' %
+                   (v[0], v[1], v[2], c[0], c[1], c[2]))
+    for f in faces:
+        f_plus = f + 1
+        file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2]))
+    file.close()
+
+
+def save_ply(mesh_path, points, rgb):
+    '''
+    Save the visualization of sampling to a ply file.
+    Red points represent positive predictions.
+    Green points represent negative predictions.
+    :param mesh_path: File name to save
+    :param points: [N, 3] array of points
+    :param rgb: [N, 3] array of rgb values in the range [0~1]
+    :return:
+    '''
+    to_save = np.concatenate([points, rgb * 255], axis=-1)
+    return np.savetxt(
+        mesh_path,
+        to_save,
+        fmt='%.6f %.6f %.6f %d %d %d',
+        comments='',
+        header=(
+            'ply\nformat ascii 1.0\nelement vertex {:d}\n' +
+            'property float x\nproperty float y\nproperty float z\n' +
+            'property uchar red\nproperty uchar green\nproperty uchar blue\n' +
+            'end_header').format(points.shape[0]))
+
+
+class HoppeMesh:
+    def __init__(self, verts, faces, vert_normals, face_normals):
+        '''
+        The HoppeSDF calculates signed distance towards a predefined oriented point cloud
+        http://hhoppe.com/recon.pdf
+        For clean and high-resolution pcl data, this is the fastest and accurate approximation of sdf
+        :param points: pts
+        :param normals: normals
+        '''
+        self.verts = verts  # [n, 3]
+        self.faces = faces  # [m, 3]
+        self.vert_normals = vert_normals  # [n, 3]
+        self.face_normals = face_normals  # [m, 3]
+
+        self.kd_tree = cKDTree(self.verts)
+        self.len = len(self.verts)
+
+    def query(self, points):
+        dists, idx = self.kd_tree.query(points, n_jobs=1)
+        # FIXME: because the eyebows are removed, cKDTree around eyebows
+        # are not accurate. Cause a few false-inside labels here.
+        dirs = points - self.verts[idx]
+        signs = (dirs * self.vert_normals[idx]).sum(axis=1)
+        signs = (signs > 0) * 2 - 1
+        return signs * dists
+
+    def contains(self, points):
+
+        labels = trimesh.Trimesh(vertices=self.verts,
+                                 faces=self.faces).contains(points)
+        return labels
+
+    def export(self, path):
+        if self.colors is not None:
+            save_obj_mesh_with_color(path, self.verts, self.faces,
+                                     self.colors[:, 0:3] / 255.0)
+        else:
+            save_obj_mesh(path, self.verts, self.faces)
+
+    def export_ply(self, path):
+        save_ply(path, self.verts, self.colors[:, 0:3] / 255.0)
+
+    def triangles(self):
+        return self.verts[self.faces]  # [n, 3, 3]