init

.gitignore

@@ -0,0 +1,17 @@
+data/*/*
+data/thuman*
+!data/tbfo.ttf
+__pycache__
+debug/
+log/
+results/*
+.vscode
+!.gitignore
+force_push.sh
+.idea
+smplx/
+human_det/
+kaolin/
+neural_voxelization_layer/
+pytorch3d/
+force_push.sh

LICENSE

@@ -0,0 +1,53 @@
+License
+
+Software Copyright License for non-commercial scientific research purposes
+Please read carefully the following terms and conditions and any accompanying documentation before you download and/or use the ICON model, data and software, (the "Data & Software"), including 3D meshes, images, videos, textures, software, scripts, and animations. By downloading and/or using the Data & Software (including downloading, cloning, installing, and any other use of the corresponding github repository), you acknowledge that you have read these terms and conditions, understand them, and agree to be bound by them. If you do not agree with these terms and conditions, you must not download and/or use the Data & Software. Any infringement of the terms of this agreement will automatically terminate your rights under this License
+
+Ownership / Licensees
+The Software and the associated materials has been developed at the Max Planck Institute for Intelligent Systems (hereinafter "MPI"). Any copyright or patent right is owned by and proprietary material of the Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (hereinafter “MPG”; MPI and MPG hereinafter collectively “Max-Planck”) hereinafter the “Licensor”.
+
+License Grant
+Licensor grants you (Licensee) personally a single-user, non-exclusive, non-transferable, free of charge right:
+
+	• To install the Model & Software on computers owned, leased or otherwise controlled by you and/or your organization;
+	• To use the Model & Software for the sole purpose of performing peaceful non-commercial scientific research, non-commercial education, or non-commercial artistic projects;
+	• To modify, adapt, translate or create derivative works based upon the Model & Software.
+
+Any other use, in particular any use for commercial, pornographic, military, or surveillance, purposes is prohibited. This includes, without limitation, incorporation in a commercial product, use in a commercial service, or production of other artifacts for commercial purposes. The Data & Software may not be used to create fake, libelous, misleading, or defamatory content of any kind excluding analyses in peer-reviewed scientific research. The Data & Software may not be reproduced, modified and/or made available in any form to any third party without Max-Planck’s prior written permission.
+
+The Data & Software may not be used for pornographic purposes or to generate pornographic material whether commercial or not. This license also prohibits the use of the Software to train methods/algorithms/neural networks/etc. for commercial, pornographic, military, surveillance, or defamatory use of any kind. By downloading the Data & Software, you agree not to reverse engineer it.
+
+No Distribution
+The Data & Software and the license herein granted shall not be copied, shared, distributed, re-sold, offered for re-sale, transferred or sub-licensed in whole or in part except that you may make one copy for archive purposes only.
+
+Disclaimer of Representations and Warranties
+You expressly acknowledge and agree that the Data & Software results from basic research, is provided “AS IS”, may contain errors, and that any use of the Data & Software is at your sole risk. LICENSOR MAKES NO REPRESENTATIONS OR WARRANTIES OF ANY KIND CONCERNING THE DATA & SOFTWARE, NEITHER EXPRESS NOR IMPLIED, AND THE ABSENCE OF ANY LEGAL OR ACTUAL DEFECTS, WHETHER DISCOVERABLE OR NOT. Specifically, and not to limit the foregoing, licensor makes no representations or warranties (i) regarding the merchantability or fitness for a particular purpose of the Data & Software, (ii) that the use of the Data & Software will not infringe any patents, copyrights or other intellectual property rights of a third party, and (iii) that the use of the Data & Software will not cause any damage of any kind to you or a third party.
+
+Limitation of Liability
+Because this Data & Software License Agreement qualifies as a donation, according to Section 521 of the German Civil Code (Bürgerliches Gesetzbuch – BGB) Licensor as a donor is liable for intent and gross negligence only. If the Licensor fraudulently conceals a legal or material defect, they are obliged to compensate the Licensee for the resulting damage.
+Licensor shall be liable for loss of data only up to the amount of typical recovery costs which would have arisen had proper and regular data backup measures been taken. For the avoidance of doubt Licensor shall be liable in accordance with the German Product Liability Act in the event of product liability. The foregoing applies also to Licensor’s legal representatives or assistants in performance. Any further liability shall be excluded.
+Patent claims generated through the usage of the Data & Software cannot be directed towards the copyright holders.
+The Data & Software is provided in the state of development the licensor defines. If modified or extended by Licensee, the Licensor makes no claims about the fitness of the Data & Software and is not responsible for any problems such modifications cause.
+
+No Maintenance Services
+You understand and agree that Licensor is under no obligation to provide either maintenance services, update services, notices of latent defects, or corrections of defects with regard to the Data & Software. Licensor nevertheless reserves the right to update, modify, or discontinue the Data & Software at any time.
+
+Defects of the Data & Software must be notified in writing to the Licensor with a comprehensible description of the error symptoms. The notification of the defect should enable the reproduction of the error. The Licensee is encouraged to communicate any use, results, modification or publication.
+
+Publications using the Model & Software
+You acknowledge that the Data & Software is a valuable scientific resource and agree to appropriately reference the following paper in any publication making use of the Data & Software.
+
+Citation:
+
+@inproceedings{xiu2022icon,
+  title={{ICON}: {I}mplicit {C}lothed humans {O}btained from {N}ormals},
+  author={Xiu, Yuliang and Yang, Jinlong and Tzionas, Dimitrios and Black, Michael J.},
+  booktitle={IEEE/CVF Conf.~on Computer Vision and Pattern Recognition (CVPR)},
+  month = jun,
+  year={2022}
+}
+
+Commercial licensing opportunities
+For commercial uses of the Model & Software, please send email to ps-license@tue.mpg.de
+
+This Agreement shall be governed by the laws of the Federal Republic of Germany except for the UN Sales Convention.

README.md

@@ -1,13 +1,12 @@
 ---
 title: ICON
-emoji: 🌖
+metaTitle: "Image2Human by Yuliang Xiu"
+emoji: 🤼
 colorFrom: indigo
-colorTo: indigo
+colorTo: yellow
 sdk: gradio
 sdk_version: 3.1.1
-app_file: app.py
-pinned: false
-license: other
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+app_file: ./apps/app.py
+pinned: true
+python_version: 3.8
+---

apps/ICON.py

@@ -0,0 +1,762 @@
+# -*- 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.common.seg3d_lossless import Seg3dLossless
+from lib.dataset.Evaluator import Evaluator
+from lib.net import HGPIFuNet
+from lib.common.train_util import *
+from lib.renderer.gl.init_gl import initialize_GL_context
+from lib.common.render import Render
+from lib.dataset.mesh_util import SMPLX, update_mesh_shape_prior_losses, get_visibility
+import warnings
+import logging
+import torch
+import smplx
+import numpy as np
+from torch import nn
+from skimage.transform import resize
+import pytorch_lightning as pl
+
+torch.backends.cudnn.benchmark = True
+
+logging.getLogger("lightning").setLevel(logging.ERROR)
+
+warnings.filterwarnings("ignore")
+
+
+class ICON(pl.LightningModule):
+    def __init__(self, cfg):
+        super(ICON, self).__init__()
+
+        self.cfg = cfg
+        self.batch_size = self.cfg.batch_size
+        self.lr_G = self.cfg.lr_G
+
+        self.use_sdf = cfg.sdf
+        self.prior_type = cfg.net.prior_type
+        self.mcube_res = cfg.mcube_res
+        self.clean_mesh_flag = cfg.clean_mesh
+
+        self.netG = HGPIFuNet(
+            self.cfg,
+            self.cfg.projection_mode,
+            error_term=nn.SmoothL1Loss() if self.use_sdf else nn.MSELoss(),
+        )
+
+        # TODO: replace the renderer from opengl to pytorch3d
+        self.evaluator = Evaluator(
+            device=torch.device(f"cuda:{self.cfg.gpus[0]}"))
+
+        self.resolutions = (
+            np.logspace(
+                start=5,
+                stop=np.log2(self.mcube_res),
+                base=2,
+                num=int(np.log2(self.mcube_res) - 4),
+                endpoint=True,
+            )
+            + 1.0
+        )
+        self.resolutions = self.resolutions.astype(np.int16).tolist()
+
+        self.icon_keys = ["smpl_verts", "smpl_faces", "smpl_vis", "smpl_cmap"]
+        self.pamir_keys = ["voxel_verts",
+                           "voxel_faces", "pad_v_num", "pad_f_num"]
+
+        self.reconEngine = Seg3dLossless(
+            query_func=query_func,
+            b_min=[[-1.0, 1.0, -1.0]],
+            b_max=[[1.0, -1.0, 1.0]],
+            resolutions=self.resolutions,
+            align_corners=True,
+            balance_value=0.50,
+            device=torch.device(f"cuda:{self.cfg.test_gpus[0]}"),
+            visualize=False,
+            debug=False,
+            use_cuda_impl=False,
+            faster=True,
+        )
+
+        self.render = Render(
+            size=512, device=torch.device(f"cuda:{self.cfg.test_gpus[0]}")
+        )
+        self.smpl_data = SMPLX()
+
+        self.get_smpl_model = lambda smpl_type, gender, age, v_template: smplx.create(
+            self.smpl_data.model_dir,
+            kid_template_path=osp.join(
+                osp.realpath(self.smpl_data.model_dir),
+                f"{smpl_type}/{smpl_type}_kid_template.npy",
+            ),
+            model_type=smpl_type,
+            gender=gender,
+            age=age,
+            v_template=v_template,
+            use_face_contour=False,
+            ext="pkl",
+        )
+
+        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_total = self.in_geo + self.in_nml
+        self.smpl_dim = cfg.net.smpl_dim
+
+        self.export_dir = None
+        self.result_eval = {}
+
+    def get_progress_bar_dict(self):
+        tqdm_dict = super().get_progress_bar_dict()
+        if "v_num" in tqdm_dict:
+            del tqdm_dict["v_num"]
+        return tqdm_dict
+
+    # Training related
+    def configure_optimizers(self):
+
+        # set optimizer
+        weight_decay = self.cfg.weight_decay
+        momentum = self.cfg.momentum
+
+        optim_params_G = [
+            {"params": self.netG.if_regressor.parameters(), "lr": self.lr_G}
+        ]
+
+        if self.cfg.net.use_filter:
+            optim_params_G.append(
+                {"params": self.netG.F_filter.parameters(), "lr": self.lr_G}
+            )
+
+        if self.cfg.net.prior_type == "pamir":
+            optim_params_G.append(
+                {"params": self.netG.ve.parameters(), "lr": self.lr_G}
+            )
+
+        if self.cfg.optim == "Adadelta":
+
+            optimizer_G = torch.optim.Adadelta(
+                optim_params_G, lr=self.lr_G, weight_decay=weight_decay
+            )
+
+        elif self.cfg.optim == "Adam":
+
+            optimizer_G = torch.optim.Adam(
+                optim_params_G, lr=self.lr_G, weight_decay=weight_decay
+            )
+
+        elif self.cfg.optim == "RMSprop":
+
+            optimizer_G = torch.optim.RMSprop(
+                optim_params_G,
+                lr=self.lr_G,
+                weight_decay=weight_decay,
+                momentum=momentum,
+            )
+
+        else:
+            raise NotImplementedError
+
+        # set scheduler
+        scheduler_G = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer_G, milestones=self.cfg.schedule, gamma=self.cfg.gamma
+        )
+
+        return [optimizer_G], [scheduler_G]
+
+    def training_step(self, batch, batch_idx):
+
+        if not self.cfg.fast_dev:
+            export_cfg(self.logger, self.cfg)
+
+        self.netG.train()
+
+        in_tensor_dict = {
+            "sample": batch["samples_geo"].permute(0, 2, 1),
+            "calib": batch["calib"],
+            "label": batch["labels_geo"].unsqueeze(1),
+        }
+
+        for name in self.in_total:
+            in_tensor_dict.update({name: batch[name]})
+
+        if self.prior_type == "icon":
+            for key in self.icon_keys:
+                in_tensor_dict.update({key: batch[key]})
+        elif self.prior_type == "pamir":
+            for key in self.pamir_keys:
+                in_tensor_dict.update({key: batch[key]})
+        else:
+            pass
+
+        preds_G, error_G = self.netG(in_tensor_dict)
+
+        acc, iou, prec, recall = self.evaluator.calc_acc(
+            preds_G.flatten(),
+            in_tensor_dict["label"].flatten(),
+            0.5,
+            use_sdf=self.cfg.sdf,
+        )
+
+        # metrics processing
+        metrics_log = {
+            "train_loss": error_G.item(),
+            "train_acc": acc.item(),
+            "train_iou": iou.item(),
+            "train_prec": prec.item(),
+            "train_recall": recall.item(),
+        }
+
+        tf_log = tf_log_convert(metrics_log)
+        bar_log = bar_log_convert(metrics_log)
+
+        if batch_idx % int(self.cfg.freq_show_train) == 0:
+
+            with torch.no_grad():
+                self.render_func(in_tensor_dict, dataset="train")
+
+        metrics_return = {
+            k.replace("train_", ""): torch.tensor(v) for k, v in metrics_log.items()
+        }
+
+        metrics_return.update(
+            {"loss": error_G, "log": tf_log, "progress_bar": bar_log})
+
+        return metrics_return
+
+    def training_epoch_end(self, outputs):
+
+        if [] in outputs:
+            outputs = outputs[0]
+
+        # metrics processing
+        metrics_log = {
+            "train_avgloss": batch_mean(outputs, "loss"),
+            "train_avgiou": batch_mean(outputs, "iou"),
+            "train_avgprec": batch_mean(outputs, "prec"),
+            "train_avgrecall": batch_mean(outputs, "recall"),
+            "train_avgacc": batch_mean(outputs, "acc"),
+        }
+
+        tf_log = tf_log_convert(metrics_log)
+
+        return {"log": tf_log}
+
+    def validation_step(self, batch, batch_idx):
+
+        self.netG.eval()
+        self.netG.training = False
+
+        in_tensor_dict = {
+            "sample": batch["samples_geo"].permute(0, 2, 1),
+            "calib": batch["calib"],
+            "label": batch["labels_geo"].unsqueeze(1),
+        }
+
+        for name in self.in_total:
+            in_tensor_dict.update({name: batch[name]})
+
+        if self.prior_type == "icon":
+            for key in self.icon_keys:
+                in_tensor_dict.update({key: batch[key]})
+        elif self.prior_type == "pamir":
+            for key in self.pamir_keys:
+                in_tensor_dict.update({key: batch[key]})
+        else:
+            pass
+
+        preds_G, error_G = self.netG(in_tensor_dict)
+
+        acc, iou, prec, recall = self.evaluator.calc_acc(
+            preds_G.flatten(),
+            in_tensor_dict["label"].flatten(),
+            0.5,
+            use_sdf=self.cfg.sdf,
+        )
+
+        if batch_idx % int(self.cfg.freq_show_val) == 0:
+            with torch.no_grad():
+                self.render_func(in_tensor_dict, dataset="val", idx=batch_idx)
+
+        metrics_return = {
+            "val_loss": error_G,
+            "val_acc": acc,
+            "val_iou": iou,
+            "val_prec": prec,
+            "val_recall": recall,
+        }
+
+        return metrics_return
+
+    def validation_epoch_end(self, outputs):
+
+        # metrics processing
+        metrics_log = {
+            "val_avgloss": batch_mean(outputs, "val_loss"),
+            "val_avgacc": batch_mean(outputs, "val_acc"),
+            "val_avgiou": batch_mean(outputs, "val_iou"),
+            "val_avgprec": batch_mean(outputs, "val_prec"),
+            "val_avgrecall": batch_mean(outputs, "val_recall"),
+        }
+
+        tf_log = tf_log_convert(metrics_log)
+
+        return {"log": tf_log}
+
+    def compute_vis_cmap(self, smpl_type, 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 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),
+        }
+
+    @torch.enable_grad()
+    def optim_body(self, in_tensor_dict, batch):
+
+        smpl_model = self.get_smpl_model(
+            batch["type"][0], batch["gender"][0], batch["age"][0], None
+        ).to(self.device)
+        in_tensor_dict["smpl_faces"] = (
+            torch.tensor(smpl_model.faces.astype(np.int))
+            .long()
+            .unsqueeze(0)
+            .to(self.device)
+        )
+
+        # The optimizer and variables
+        optimed_pose = torch.tensor(
+            batch["body_pose"][0], device=self.device, requires_grad=True
+        )  # [1,23,3,3]
+        optimed_trans = torch.tensor(
+            batch["transl"][0], device=self.device, requires_grad=True
+        )  # [3]
+        optimed_betas = torch.tensor(
+            batch["betas"][0], device=self.device, requires_grad=True
+        )  # [1,10]
+        optimed_orient = torch.tensor(
+            batch["global_orient"][0], device=self.device, requires_grad=True
+        )  # [1,1,3,3]
+
+        optimizer_smpl = torch.optim.SGD(
+            [optimed_pose, optimed_trans, optimed_betas, optimed_orient],
+            lr=1e-3,
+            momentum=0.9,
+        )
+        scheduler_smpl = torch.optim.lr_scheduler.ReduceLROnPlateau(
+            optimizer_smpl, mode="min", factor=0.5, verbose=0, min_lr=1e-5, patience=5
+        )
+        loop_smpl = range(50)
+        for i in loop_smpl:
+
+            optimizer_smpl.zero_grad()
+
+            # prior_loss, optimed_pose = dataset.vposer_prior(optimed_pose)
+            smpl_out = smpl_model(
+                betas=optimed_betas,
+                body_pose=optimed_pose,
+                global_orient=optimed_orient,
+                transl=optimed_trans,
+                return_verts=True,
+            )
+
+            smpl_verts = smpl_out.vertices[0] * 100.0
+            smpl_verts = projection(
+                smpl_verts, batch["calib"][0], format="tensor")
+            smpl_verts[:, 1] *= -1
+            # render optimized mesh (normal, T_normal, image [-1,1])
+            self.render.load_meshes(
+                smpl_verts, in_tensor_dict["smpl_faces"])
+            (
+                in_tensor_dict["T_normal_F"],
+                in_tensor_dict["T_normal_B"],
+            ) = self.render.get_rgb_image()
+
+            T_mask_F, T_mask_B = self.render.get_silhouette_image()
+
+            with torch.no_grad():
+                (
+                    in_tensor_dict["normal_F"],
+                    in_tensor_dict["normal_B"],
+                ) = self.netG.normal_filter(in_tensor_dict)
+
+            # mask = torch.abs(in_tensor['T_normal_F']).sum(dim=0, keepdims=True) > 0.0
+            diff_F_smpl = torch.abs(
+                in_tensor_dict["T_normal_F"] - in_tensor_dict["normal_F"]
+            )
+            diff_B_smpl = torch.abs(
+                in_tensor_dict["T_normal_B"] - in_tensor_dict["normal_B"]
+            )
+            loss = (diff_F_smpl + diff_B_smpl).mean()
+
+            # silhouette loss
+            smpl_arr = torch.cat([T_mask_F, T_mask_B], dim=-1)[0]
+            gt_arr = torch.cat(
+                [in_tensor_dict["normal_F"][0], in_tensor_dict["normal_B"][0]], dim=2
+            ).permute(1, 2, 0)
+            gt_arr = ((gt_arr + 1.0) * 0.5).to(self.device)
+            bg_color = (
+                torch.Tensor([0.5, 0.5, 0.5]).unsqueeze(
+                    0).unsqueeze(0).to(self.device)
+            )
+            gt_arr = ((gt_arr - bg_color).sum(dim=-1) != 0.0).float()
+            loss += torch.abs(smpl_arr - gt_arr).mean()
+
+            # Image.fromarray(((in_tensor_dict['T_normal_F'][0].permute(1,2,0)+1.0)*0.5*255.0).detach().cpu().numpy().astype(np.uint8)).show()
+
+            # loop_smpl.set_description(f"smpl = {loss:.3f}")
+
+            loss.backward(retain_graph=True)
+            optimizer_smpl.step()
+            scheduler_smpl.step(loss)
+            in_tensor_dict["smpl_verts"] = smpl_verts.unsqueeze(0)
+
+        in_tensor_dict.update(
+            self.compute_vis_cmap(
+                batch["type"][0],
+                in_tensor_dict["smpl_verts"][0],
+                in_tensor_dict["smpl_faces"][0],
+            )
+        )
+
+        features, inter = self.netG.filter(in_tensor_dict, return_inter=True)
+
+        return features, inter, in_tensor_dict
+
+    @torch.enable_grad()
+    def optim_cloth(self, verts_pr, faces_pr, inter):
+
+        # convert from GT to SDF
+        verts_pr -= (self.resolutions[-1] - 1) / 2.0
+        verts_pr /= (self.resolutions[-1] - 1) / 2.0
+
+        losses = {
+            "cloth": {"weight": 5.0, "value": 0.0},
+            "edge": {"weight": 100.0, "value": 0.0},
+            "normal": {"weight": 0.2, "value": 0.0},
+            "laplacian": {"weight": 100.0, "value": 0.0},
+            "smpl": {"weight": 1.0, "value": 0.0},
+            "deform": {"weight": 20.0, "value": 0.0},
+        }
+
+        deform_verts = torch.full(
+            verts_pr.shape, 0.0, device=self.device, requires_grad=True
+        )
+        optimizer_cloth = torch.optim.SGD(
+            [deform_verts], lr=1e-1, momentum=0.9)
+        scheduler_cloth = torch.optim.lr_scheduler.ReduceLROnPlateau(
+            optimizer_cloth, mode="min", factor=0.1, verbose=0, min_lr=1e-3, patience=5
+        )
+        # cloth optimization
+        loop_cloth = range(100)
+
+        for i in loop_cloth:
+
+            optimizer_cloth.zero_grad()
+
+            self.render.load_meshes(
+                verts_pr.unsqueeze(0).to(self.device),
+                faces_pr.unsqueeze(0).to(self.device).long(),
+                deform_verts,
+            )
+            P_normal_F, P_normal_B = self.render.get_rgb_image()
+
+            update_mesh_shape_prior_losses(self.render.mesh, losses)
+            diff_F_cloth = torch.abs(P_normal_F[0] - inter[:3])
+            diff_B_cloth = torch.abs(P_normal_B[0] - inter[3:])
+            losses["cloth"]["value"] = (diff_F_cloth + diff_B_cloth).mean()
+            losses["deform"]["value"] = torch.topk(
+                torch.abs(deform_verts.flatten()), 30
+            )[0].mean()
+
+            # Weighted sum of the losses
+            cloth_loss = torch.tensor(0.0, device=self.device)
+            pbar_desc = ""
+
+            for k in losses.keys():
+                if k != "smpl":
+                    cloth_loss_per_cls = losses[k]["value"] * \
+                        losses[k]["weight"]
+                    pbar_desc += f"{k}: {cloth_loss_per_cls:.3f} | "
+                    cloth_loss += cloth_loss_per_cls
+
+            # loop_cloth.set_description(pbar_desc)
+            cloth_loss.backward(retain_graph=True)
+            optimizer_cloth.step()
+            scheduler_cloth.step(cloth_loss)
+
+        # convert from GT to SDF
+        deform_verts = deform_verts.flatten().detach()
+        deform_verts[torch.topk(torch.abs(deform_verts), 30)[
+            1]] = deform_verts.mean()
+        deform_verts = deform_verts.view(-1, 3).cpu()
+
+        verts_pr += deform_verts
+        verts_pr *= (self.resolutions[-1] - 1) / 2.0
+        verts_pr += (self.resolutions[-1] - 1) / 2.0
+
+        return verts_pr
+
+    def test_step(self, batch, batch_idx):
+
+        # dict_keys(['dataset', 'subject', 'rotation', 'scale', 'calib',
+        #            'normal_F', 'normal_B', 'image', 'T_normal_F', 'T_normal_B',
+        #            'z-trans', 'verts', 'faces', 'samples_geo', 'labels_geo',
+        #            'smpl_verts', 'smpl_faces', 'smpl_vis', 'smpl_cmap', 'pts_signs',
+        #            'type', 'gender', 'age', 'body_pose', 'global_orient', 'betas', 'transl'])
+
+        if self.evaluator._normal_render is None:
+            self.evaluator.init_gl()
+
+        self.netG.eval()
+        self.netG.training = False
+        in_tensor_dict = {}
+
+        # export paths
+        mesh_name = batch["subject"][0]
+        mesh_rot = batch["rotation"][0].item()
+        ckpt_dir = self.cfg.name
+
+        for kid, key in enumerate(self.cfg.dataset.noise_type):
+            ckpt_dir += f"_{key}_{self.cfg.dataset.noise_scale[kid]}"
+
+        if self.cfg.optim_cloth:
+            ckpt_dir += "_optim_cloth"
+        if self.cfg.optim_body:
+            ckpt_dir += "_optim_body"
+
+        self.export_dir = osp.join(self.cfg.results_path, ckpt_dir, mesh_name)
+        os.makedirs(self.export_dir, exist_ok=True)
+
+        for name in self.in_total:
+            if name in batch.keys():
+                in_tensor_dict.update({name: batch[name]})
+
+        # update the new T_normal_F/B
+        in_tensor_dict.update(
+            self.evaluator.render_normal(
+                batch["smpl_verts"], batch["smpl_faces"])
+        )
+
+        # update the new smpl_vis
+        (xy, z) = batch["smpl_verts"][0].split([2, 1], dim=1)
+        smpl_vis = get_visibility(
+            xy,
+            z,
+            torch.as_tensor(self.smpl_data.faces).type_as(
+                batch["smpl_verts"]).long(),
+        )
+        in_tensor_dict.update({"smpl_vis": smpl_vis.unsqueeze(0)})
+
+        if self.prior_type == "icon":
+            for key in self.icon_keys:
+                in_tensor_dict.update({key: batch[key]})
+        elif self.prior_type == "pamir":
+            for key in self.pamir_keys:
+                in_tensor_dict.update({key: batch[key]})
+        else:
+            pass
+
+        with torch.no_grad():
+            if self.cfg.optim_body:
+                features, inter, in_tensor_dict = self.optim_body(
+                    in_tensor_dict, batch)
+            else:
+                features, inter = self.netG.filter(
+                    in_tensor_dict, return_inter=True)
+            sdf = self.reconEngine(
+                opt=self.cfg, netG=self.netG, features=features, proj_matrix=None
+            )
+
+        # save inter results
+        image = (
+            in_tensor_dict["image"][0].permute(
+                1, 2, 0).detach().cpu().numpy() + 1.0
+        ) * 0.5
+        smpl_F = (
+            in_tensor_dict["T_normal_F"][0].permute(
+                1, 2, 0).detach().cpu().numpy()
+            + 1.0
+        ) * 0.5
+        smpl_B = (
+            in_tensor_dict["T_normal_B"][0].permute(
+                1, 2, 0).detach().cpu().numpy()
+            + 1.0
+        ) * 0.5
+        image_inter = np.concatenate(
+            self.tensor2image(512, inter[0]) + [smpl_F, smpl_B, image], axis=1
+        )
+        Image.fromarray((image_inter * 255.0).astype(np.uint8)).save(
+            osp.join(self.export_dir, f"{mesh_rot}_inter.png")
+        )
+
+        verts_pr, faces_pr = self.reconEngine.export_mesh(sdf)
+
+        if self.clean_mesh_flag:
+            verts_pr, faces_pr = clean_mesh(verts_pr, faces_pr)
+
+        if self.cfg.optim_cloth:
+            verts_pr = self.optim_cloth(verts_pr, faces_pr, inter[0].detach())
+
+        verts_gt = batch["verts"][0]
+        faces_gt = batch["faces"][0]
+
+        self.result_eval.update(
+            {
+                "verts_gt": verts_gt,
+                "faces_gt": faces_gt,
+                "verts_pr": verts_pr,
+                "faces_pr": faces_pr,
+                "recon_size": (self.resolutions[-1] - 1.0),
+                "calib": batch["calib"][0],
+            }
+        )
+
+        self.evaluator.set_mesh(self.result_eval, scale_factor=1.0)
+        self.evaluator.space_transfer()
+
+        chamfer, p2s = self.evaluator.calculate_chamfer_p2s(
+            sampled_points=1000)
+        normal_consist = self.evaluator.calculate_normal_consist(
+            save_demo_img=osp.join(self.export_dir, f"{mesh_rot}_nc.png")
+        )
+
+        test_log = {"chamfer": chamfer, "p2s": p2s, "NC": normal_consist}
+
+        return test_log
+
+    def test_epoch_end(self, outputs):
+
+        # make_test_gif("/".join(self.export_dir.split("/")[:-2]))
+
+        accu_outputs = accumulate(
+            outputs,
+            rot_num=3,
+            split={
+                "thuman2": (0, 5),
+            },
+        )
+
+        print(colored(self.cfg.name, "green"))
+        print(colored(self.cfg.dataset.noise_scale, "green"))
+
+        self.logger.experiment.add_hparams(
+            hparam_dict={"lr_G": self.lr_G, "bsize": self.batch_size},
+            metric_dict=accu_outputs,
+        )
+
+        np.save(
+            osp.join(self.export_dir, "../test_results.npy"),
+            accu_outputs,
+            allow_pickle=True,
+        )
+
+        return accu_outputs
+
+    def tensor2image(self, height, inter):
+
+        all = []
+        for dim in self.in_geo_dim:
+            img = resize(
+                np.tile(
+                    ((inter[:dim].cpu().numpy() + 1.0) /
+                     2.0).transpose(1, 2, 0),
+                    (1, 1, int(3 / dim)),
+                ),
+                (height, height),
+                anti_aliasing=True,
+            )
+
+            all.append(img)
+            inter = inter[dim:]
+
+        return all
+
+    def render_func(self, in_tensor_dict, dataset="title", idx=0):
+
+        for name in in_tensor_dict.keys():
+            in_tensor_dict[name] = in_tensor_dict[name][0:1]
+
+        self.netG.eval()
+        features, inter = self.netG.filter(in_tensor_dict, return_inter=True)
+        sdf = self.reconEngine(
+            opt=self.cfg, netG=self.netG, features=features, proj_matrix=None
+        )
+
+        if sdf is not None:
+            render = self.reconEngine.display(sdf)
+
+            image_pred = np.flip(render[:, :, ::-1], axis=0)
+            height = image_pred.shape[0]
+
+            image_gt = resize(
+                ((in_tensor_dict["image"].cpu().numpy()[0] + 1.0) / 2.0).transpose(
+                    1, 2, 0
+                ),
+                (height, height),
+                anti_aliasing=True,
+            )
+            image_inter = self.tensor2image(height, inter[0])
+            image = np.concatenate(
+                [image_pred, image_gt] + image_inter, axis=1)
+
+            step_id = self.global_step if dataset == "train" else self.global_step + idx
+            self.logger.experiment.add_image(
+                tag=f"Occupancy-{dataset}/{step_id}",
+                img_tensor=image.transpose(2, 0, 1),
+                global_step=step_id,
+            )
+
+    def test_single(self, batch):
+
+        self.netG.eval()
+        self.netG.training = False
+        in_tensor_dict = {}
+
+        for name in self.in_total:
+            if name in batch.keys():
+                in_tensor_dict.update({name: batch[name]})
+
+        if self.prior_type == "icon":
+            for key in self.icon_keys:
+                in_tensor_dict.update({key: batch[key]})
+        elif self.prior_type == "pamir":
+            for key in self.pamir_keys:
+                in_tensor_dict.update({key: batch[key]})
+        else:
+            pass
+
+        features, inter = self.netG.filter(in_tensor_dict, return_inter=True)
+        sdf = self.reconEngine(
+            opt=self.cfg, netG=self.netG, features=features, proj_matrix=None
+        )
+
+        verts_pr, faces_pr = self.reconEngine.export_mesh(sdf)
+
+        if self.clean_mesh_flag:
+            verts_pr, faces_pr = clean_mesh(verts_pr, faces_pr)
+
+        verts_pr -= (self.resolutions[-1] - 1) / 2.0
+        verts_pr /= (self.resolutions[-1] - 1) / 2.0
+
+        return verts_pr, faces_pr, inter

apps/Normal.py

@@ -0,0 +1,213 @@
+from lib.net import NormalNet
+from lib.common.train_util import *
+import logging
+import torch
+import numpy as np
+from torch import nn
+from skimage.transform import resize
+import pytorch_lightning as pl
+
+torch.backends.cudnn.benchmark = True
+
+logging.getLogger("lightning").setLevel(logging.ERROR)
+
+
+class Normal(pl.LightningModule):
+    def __init__(self, cfg):
+        super(Normal, self).__init__()
+        self.cfg = cfg
+        self.batch_size = self.cfg.batch_size
+        self.lr_N = self.cfg.lr_N
+
+        self.schedulers = []
+
+        self.netG = NormalNet(self.cfg, error_term=nn.SmoothL1Loss())
+
+        self.in_nml = [item[0] for item in cfg.net.in_nml]
+
+    def get_progress_bar_dict(self):
+        tqdm_dict = super().get_progress_bar_dict()
+        if "v_num" in tqdm_dict:
+            del tqdm_dict["v_num"]
+        return tqdm_dict
+
+    # Training related
+    def configure_optimizers(self):
+
+        # set optimizer
+        weight_decay = self.cfg.weight_decay
+        momentum = self.cfg.momentum
+
+        optim_params_N_F = [
+            {"params": self.netG.netF.parameters(), "lr": self.lr_N}]
+        optim_params_N_B = [
+            {"params": self.netG.netB.parameters(), "lr": self.lr_N}]
+
+        optimizer_N_F = torch.optim.Adam(
+            optim_params_N_F, lr=self.lr_N, weight_decay=weight_decay
+        )
+
+        optimizer_N_B = torch.optim.Adam(
+            optim_params_N_B, lr=self.lr_N, weight_decay=weight_decay
+        )
+
+        scheduler_N_F = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer_N_F, milestones=self.cfg.schedule, gamma=self.cfg.gamma
+        )
+
+        scheduler_N_B = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer_N_B, milestones=self.cfg.schedule, gamma=self.cfg.gamma
+        )
+
+        self.schedulers = [scheduler_N_F, scheduler_N_B]
+        optims = [optimizer_N_F, optimizer_N_B]
+
+        return optims, self.schedulers
+
+    def render_func(self, render_tensor):
+
+        height = render_tensor["image"].shape[2]
+        result_list = []
+
+        for name in render_tensor.keys():
+            result_list.append(
+                resize(
+                    ((render_tensor[name].cpu().numpy()[0] + 1.0) / 2.0).transpose(
+                        1, 2, 0
+                    ),
+                    (height, height),
+                    anti_aliasing=True,
+                )
+            )
+        result_array = np.concatenate(result_list, axis=1)
+
+        return result_array
+
+    def training_step(self, batch, batch_idx, optimizer_idx):
+
+        export_cfg(self.logger, self.cfg)
+
+        # retrieve the data
+        in_tensor = {}
+        for name in self.in_nml:
+            in_tensor[name] = batch[name]
+
+        FB_tensor = {"normal_F": batch["normal_F"],
+                     "normal_B": batch["normal_B"]}
+
+        self.netG.train()
+
+        preds_F, preds_B = self.netG(in_tensor)
+        error_NF, error_NB = self.netG.get_norm_error(
+            preds_F, preds_B, FB_tensor)
+
+        (opt_nf, opt_nb) = self.optimizers()
+
+        opt_nf.zero_grad()
+        opt_nb.zero_grad()
+
+        self.manual_backward(error_NF, opt_nf)
+        self.manual_backward(error_NB, opt_nb)
+
+        opt_nf.step()
+        opt_nb.step()
+
+        if batch_idx > 0 and batch_idx % int(self.cfg.freq_show_train) == 0:
+
+            self.netG.eval()
+            with torch.no_grad():
+                nmlF, nmlB = self.netG(in_tensor)
+                in_tensor.update({"nmlF": nmlF, "nmlB": nmlB})
+                result_array = self.render_func(in_tensor)
+
+                self.logger.experiment.add_image(
+                    tag=f"Normal-train/{self.global_step}",
+                    img_tensor=result_array.transpose(2, 0, 1),
+                    global_step=self.global_step,
+                )
+
+        # metrics processing
+        metrics_log = {
+            "train_loss-NF": error_NF.item(),
+            "train_loss-NB": error_NB.item(),
+        }
+
+        tf_log = tf_log_convert(metrics_log)
+        bar_log = bar_log_convert(metrics_log)
+
+        return {
+            "loss": error_NF + error_NB,
+            "loss-NF": error_NF,
+            "loss-NB": error_NB,
+            "log": tf_log,
+            "progress_bar": bar_log,
+        }
+
+    def training_epoch_end(self, outputs):
+
+        if [] in outputs:
+            outputs = outputs[0]
+
+        # metrics processing
+        metrics_log = {
+            "train_avgloss": batch_mean(outputs, "loss"),
+            "train_avgloss-NF": batch_mean(outputs, "loss-NF"),
+            "train_avgloss-NB": batch_mean(outputs, "loss-NB"),
+        }
+
+        tf_log = tf_log_convert(metrics_log)
+
+        tf_log["lr-NF"] = self.schedulers[0].get_last_lr()[0]
+        tf_log["lr-NB"] = self.schedulers[1].get_last_lr()[0]
+
+        return {"log": tf_log}
+
+    def validation_step(self, batch, batch_idx):
+
+        # retrieve the data
+        in_tensor = {}
+        for name in self.in_nml:
+            in_tensor[name] = batch[name]
+
+        FB_tensor = {"normal_F": batch["normal_F"],
+                     "normal_B": batch["normal_B"]}
+
+        self.netG.train()
+
+        preds_F, preds_B = self.netG(in_tensor)
+        error_NF, error_NB = self.netG.get_norm_error(
+            preds_F, preds_B, FB_tensor)
+
+        if (batch_idx > 0 and batch_idx % int(self.cfg.freq_show_train) == 0) or (
+            batch_idx == 0
+        ):
+
+            with torch.no_grad():
+                nmlF, nmlB = self.netG(in_tensor)
+                in_tensor.update({"nmlF": nmlF, "nmlB": nmlB})
+                result_array = self.render_func(in_tensor)
+
+                self.logger.experiment.add_image(
+                    tag=f"Normal-val/{self.global_step}",
+                    img_tensor=result_array.transpose(2, 0, 1),
+                    global_step=self.global_step,
+                )
+
+        return {
+            "val_loss": error_NF + error_NB,
+            "val_loss-NF": error_NF,
+            "val_loss-NB": error_NB,
+        }
+
+    def validation_epoch_end(self, outputs):
+
+        # metrics processing
+        metrics_log = {
+            "val_avgloss": batch_mean(outputs, "val_loss"),
+            "val_avgloss-NF": batch_mean(outputs, "val_loss-NF"),
+            "val_avgloss-NB": batch_mean(outputs, "val_loss-NB"),
+        }
+
+        tf_log = tf_log_convert(metrics_log)
+
+        return {"log": tf_log}

apps/app.py

@@ -0,0 +1,21 @@
+# install
+
+import os
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"  
+os.environ["CUDA_VISIBLE_DEVICES"]="0"
+try:
+    os.system("bash install.sh")
+except Exception as e:
+    print(e)
+    
+
+# running
+ 
+import gradio as gr
+
+def image_classifier(inp):
+    return {'cat': 0.3, 'dog': 0.7}
+
+demo = gr.Interface(fn=image_classifier, inputs="image", outputs="label")
+demo.launch(auth=("icon@tue.mpg.de", "icon_2022"),
+            auth_message="Register at icon.is.tue.mpg.de/download to get the username and password.")

apps/infer.py

@@ -0,0 +1,616 @@
+# -*- 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 logging
+from lib.common.render import query_color, image2vid
+from lib.common.config import cfg
+from lib.common.cloth_extraction import extract_cloth
+from lib.dataset.mesh_util import (
+    load_checkpoint,
+    update_mesh_shape_prior_losses,
+    get_optim_grid_image,
+    blend_rgb_norm,
+    unwrap,
+    remesh,
+    tensor2variable,
+    normal_loss
+)
+
+from lib.dataset.TestDataset import TestDataset
+from lib.net.local_affine import LocalAffine
+from pytorch3d.structures import Meshes
+from apps.ICON import ICON
+
+import os
+from termcolor import colored
+import argparse
+import numpy as np
+from PIL import Image
+import trimesh
+import pickle
+import numpy as np
+
+import torch
+torch.backends.cudnn.benchmark = True
+
+logging.getLogger("trimesh").setLevel(logging.ERROR)
+
+
+if __name__ == "__main__":
+
+    # loading cfg file
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("-gpu", "--gpu_device", type=int, default=0)
+    parser.add_argument("-colab", action="store_true")
+    parser.add_argument("-loop_smpl", "--loop_smpl", type=int, default=100)
+    parser.add_argument("-patience", "--patience", type=int, default=5)
+    parser.add_argument("-vis_freq", "--vis_freq", type=int, default=10)
+    parser.add_argument("-loop_cloth", "--loop_cloth", type=int, default=200)
+    parser.add_argument("-hps_type", "--hps_type", type=str, default="pymaf")
+    parser.add_argument("-export_video", action="store_true")
+    parser.add_argument("-in_dir", "--in_dir", type=str, default="./examples")
+    parser.add_argument("-out_dir", "--out_dir",
+                        type=str, default="./results")
+    parser.add_argument('-seg_dir', '--seg_dir', type=str, default=None)
+    parser.add_argument(
+        "-cfg", "--config", type=str, default="./configs/icon-filter.yaml"
+    )
+
+    args = parser.parse_args()
+
+    # cfg read and merge
+    cfg.merge_from_file(args.config)
+    cfg.merge_from_file("./lib/pymaf/configs/pymaf_config.yaml")
+
+    cfg_show_list = [
+        "test_gpus",
+        [args.gpu_device],
+        "mcube_res",
+        256,
+        "clean_mesh",
+        True,
+    ]
+
+    cfg.merge_from_list(cfg_show_list)
+    cfg.freeze()
+
+    os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
+    device = torch.device(f"cuda:{args.gpu_device}")
+
+    if args.colab:
+        print(colored("colab environment...", "red"))
+        from tqdm.notebook import tqdm
+    else:
+        print(colored("normal environment...", "red"))
+        from tqdm import tqdm
+
+    # load model and dataloader
+    model = ICON(cfg)
+    model = load_checkpoint(model, cfg)
+
+    dataset_param = {
+        'image_dir': args.in_dir,
+        'seg_dir': args.seg_dir,
+        'has_det': True,            # w/ or w/o detection
+        'hps_type': args.hps_type   # pymaf/pare/pixie
+    }
+
+    if args.hps_type == "pixie" and "pamir" in args.config:
+        print(colored("PIXIE isn't compatible with PaMIR, thus switch to PyMAF", "red"))
+        dataset_param["hps_type"] = "pymaf"
+
+    dataset = TestDataset(dataset_param, device)
+
+    print(colored(f"Dataset Size: {len(dataset)}", "green"))
+
+    pbar = tqdm(dataset)
+
+    for data in pbar:
+
+        pbar.set_description(f"{data['name']}")
+
+        in_tensor = {"smpl_faces": data["smpl_faces"], "image": data["image"]}
+
+        # The optimizer and variables
+        optimed_pose = torch.tensor(
+            data["body_pose"], device=device, requires_grad=True
+        )  # [1,23,3,3]
+        optimed_trans = torch.tensor(
+            data["trans"], device=device, requires_grad=True
+        )  # [3]
+        optimed_betas = torch.tensor(
+            data["betas"], device=device, requires_grad=True
+        )  # [1,10]
+        optimed_orient = torch.tensor(
+            data["global_orient"], device=device, requires_grad=True
+        )  # [1,1,3,3]
+
+        optimizer_smpl = torch.optim.SGD(
+            [optimed_pose, optimed_trans, optimed_betas, optimed_orient],
+            lr=1e-3,
+            momentum=0.9,
+        )
+        scheduler_smpl = torch.optim.lr_scheduler.ReduceLROnPlateau(
+            optimizer_smpl,
+            mode="min",
+            factor=0.5,
+            verbose=0,
+            min_lr=1e-5,
+            patience=args.patience,
+        )
+
+        losses = {
+            "cloth": {"weight": 1e1, "value": 0.0},             # Cloth: Normal_recon - Normal_pred
+            "stiffness": {"weight": 1e5, "value": 0.0},         # Cloth: [RT]_v1 - [RT]_v2 (v1-edge-v2)
+            "rigid": {"weight": 1e5, "value": 0.0},             # Cloth: det(R) = 1
+            "edge": {"weight": 0, "value": 0.0},                # Cloth: edge length
+            "nc": {"weight": 0, "value": 0.0},                  # Cloth: normal consistency
+            "laplacian": {"weight": 1e2, "value": 0.0},         # Cloth: laplacian smoonth
+            "normal": {"weight": 1e0, "value": 0.0},            # Body: Normal_pred - Normal_smpl
+            "silhouette": {"weight": 1e1, "value": 0.0},        # Body: Silhouette_pred - Silhouette_smpl
+        }
+
+        # smpl optimization
+
+        loop_smpl = tqdm(
+            range(args.loop_smpl if cfg.net.prior_type != "pifu" else 1))
+
+        per_data_lst = []
+
+        for i in loop_smpl:
+
+            per_loop_lst = []
+
+            optimizer_smpl.zero_grad()
+
+            if dataset_param["hps_type"] != "pixie":
+                smpl_out = dataset.smpl_model(
+                    betas=optimed_betas,
+                    body_pose=optimed_pose,
+                    global_orient=optimed_orient,
+                    pose2rot=False,
+                )
+
+                smpl_verts = ((smpl_out.vertices) +
+                              optimed_trans) * data["scale"]
+            else:
+                smpl_verts, _, _ = dataset.smpl_model(
+                    shape_params=optimed_betas,
+                    expression_params=tensor2variable(data["exp"], device),
+                    body_pose=optimed_pose,
+                    global_pose=optimed_orient,
+                    jaw_pose=tensor2variable(data["jaw_pose"], device),
+                    left_hand_pose=tensor2variable(
+                        data["left_hand_pose"], device),
+                    right_hand_pose=tensor2variable(
+                        data["right_hand_pose"], device),
+                )
+
+                smpl_verts = (smpl_verts + optimed_trans) * data["scale"]
+
+            # render optimized mesh (normal, T_normal, image [-1,1])
+            in_tensor["T_normal_F"], in_tensor["T_normal_B"] = dataset.render_normal(
+                smpl_verts *
+                torch.tensor([1.0, -1.0, -1.0]
+                             ).to(device), in_tensor["smpl_faces"]
+            )
+            T_mask_F, T_mask_B = dataset.render.get_silhouette_image()
+
+            with torch.no_grad():
+                in_tensor["normal_F"], in_tensor["normal_B"] = model.netG.normal_filter(
+                    in_tensor
+                )
+
+            diff_F_smpl = torch.abs(
+                in_tensor["T_normal_F"] - in_tensor["normal_F"])
+            diff_B_smpl = torch.abs(
+                in_tensor["T_normal_B"] - in_tensor["normal_B"])
+
+            loss_F_smpl = normal_loss(
+                in_tensor["T_normal_F"], in_tensor["normal_F"])
+            loss_B_smpl = normal_loss(
+                in_tensor["T_normal_B"], in_tensor["normal_B"])
+
+            losses["normal"]["value"] = (loss_F_smpl + loss_B_smpl).mean()
+
+            # silhouette loss
+            smpl_arr = torch.cat([T_mask_F, T_mask_B], dim=-1)[0]
+            gt_arr = torch.cat(
+                [in_tensor["normal_F"][0], in_tensor["normal_B"][0]], dim=2
+            ).permute(1, 2, 0)
+            gt_arr = ((gt_arr + 1.0) * 0.5).to(device)
+            bg_color = (
+                torch.Tensor([0.5, 0.5, 0.5]).unsqueeze(
+                    0).unsqueeze(0).to(device)
+            )
+            gt_arr = ((gt_arr - bg_color).sum(dim=-1) != 0.0).float()
+            diff_S = torch.abs(smpl_arr - gt_arr)
+            losses["silhouette"]["value"] = diff_S.mean()
+
+            # Weighted sum of the losses
+            smpl_loss = 0.0
+            pbar_desc = "Body Fitting --- "
+            for k in ["normal", "silhouette"]:
+                pbar_desc += f"{k}: {losses[k]['value'] * losses[k]['weight']:.3f} | "
+                smpl_loss += losses[k]["value"] * losses[k]["weight"]
+            pbar_desc += f"Total: {smpl_loss:.3f}"
+            loop_smpl.set_description(pbar_desc)
+
+            if i % args.vis_freq == 0:
+
+                per_loop_lst.extend(
+                    [
+                        in_tensor["image"],
+                        in_tensor["T_normal_F"],
+                        in_tensor["normal_F"],
+                        diff_F_smpl / 2.0,
+                        diff_S[:, :512].unsqueeze(
+                            0).unsqueeze(0).repeat(1, 3, 1, 1),
+                    ]
+                )
+                per_loop_lst.extend(
+                    [
+                        in_tensor["image"],
+                        in_tensor["T_normal_B"],
+                        in_tensor["normal_B"],
+                        diff_B_smpl / 2.0,
+                        diff_S[:, 512:].unsqueeze(
+                            0).unsqueeze(0).repeat(1, 3, 1, 1),
+                    ]
+                )
+                per_data_lst.append(
+                    get_optim_grid_image(
+                        per_loop_lst, None, nrow=5, type="smpl")
+                )
+
+            smpl_loss.backward()
+            optimizer_smpl.step()
+            scheduler_smpl.step(smpl_loss)
+            in_tensor["smpl_verts"] = smpl_verts * \
+                torch.tensor([1.0, 1.0, -1.0]).to(device)
+
+        # visualize the optimization process
+        # 1. SMPL Fitting
+        # 2. Clothes Refinement
+
+        os.makedirs(os.path.join(args.out_dir, cfg.name,
+                    "refinement"), exist_ok=True)
+
+        # visualize the final results in self-rotation mode
+        os.makedirs(os.path.join(args.out_dir, cfg.name, "vid"), exist_ok=True)
+
+        # final results rendered as image
+        # 1. Render the final fitted SMPL (xxx_smpl.png)
+        # 2. Render the final reconstructed clothed human (xxx_cloth.png)
+        # 3. Blend the original image with predicted cloth normal (xxx_overlap.png)
+
+        os.makedirs(os.path.join(args.out_dir, cfg.name, "png"), exist_ok=True)
+
+        # final reconstruction meshes
+        # 1. SMPL mesh (xxx_smpl.obj)
+        # 2. SMPL params (xxx_smpl.npy)
+        # 3. clohted mesh (xxx_recon.obj)
+        # 4. remeshed clothed mesh (xxx_remesh.obj)
+        # 5. refined clothed mesh (xxx_refine.obj)
+
+        os.makedirs(os.path.join(args.out_dir, cfg.name, "obj"), exist_ok=True)
+
+        if cfg.net.prior_type != "pifu":
+
+            per_data_lst[0].save(
+                os.path.join(
+                    args.out_dir, cfg.name, f"refinement/{data['name']}_smpl.gif"
+                ),
+                save_all=True,
+                append_images=per_data_lst[1:],
+                duration=500,
+                loop=0,
+            )
+
+            if args.vis_freq == 1:
+                image2vid(
+                    per_data_lst,
+                    os.path.join(
+                        args.out_dir, cfg.name, f"refinement/{data['name']}_smpl.avi"
+                    ),
+                )
+
+            per_data_lst[-1].save(
+                os.path.join(args.out_dir, cfg.name,
+                             f"png/{data['name']}_smpl.png")
+            )
+
+        norm_pred = (
+            ((in_tensor["normal_F"][0].permute(1, 2, 0) + 1.0) * 255.0 / 2.0)
+            .detach()
+            .cpu()
+            .numpy()
+            .astype(np.uint8)
+        )
+
+        norm_orig = unwrap(norm_pred, data)
+        mask_orig = unwrap(
+            np.repeat(
+                data["mask"].permute(1, 2, 0).detach().cpu().numpy(), 3, axis=2
+            ).astype(np.uint8),
+            data,
+        )
+        rgb_norm = blend_rgb_norm(data["ori_image"], norm_orig, mask_orig)
+
+        Image.fromarray(
+            np.concatenate(
+                [data["ori_image"].astype(np.uint8), rgb_norm], axis=1)
+        ).save(os.path.join(args.out_dir, cfg.name, f"png/{data['name']}_overlap.png"))
+
+        smpl_obj = trimesh.Trimesh(
+            in_tensor["smpl_verts"].detach().cpu()[0] *
+            torch.tensor([1.0, -1.0, 1.0]),
+            in_tensor['smpl_faces'].detach().cpu()[0],
+            process=False,
+            maintains_order=True
+        )
+        smpl_obj.export(
+            f"{args.out_dir}/{cfg.name}/obj/{data['name']}_smpl.obj")
+
+        smpl_info = {'betas': optimed_betas,
+                     'pose': optimed_pose,
+                     'orient': optimed_orient,
+                     'trans': optimed_trans}
+
+        np.save(
+            f"{args.out_dir}/{cfg.name}/obj/{data['name']}_smpl.npy", smpl_info, allow_pickle=True)
+
+        # ------------------------------------------------------------------------------------------------------------------
+
+        # cloth optimization
+
+        per_data_lst = []
+
+        # cloth recon
+        in_tensor.update(
+            dataset.compute_vis_cmap(
+                in_tensor["smpl_verts"][0], in_tensor["smpl_faces"][0]
+            )
+        )
+
+        if cfg.net.prior_type == "pamir":
+            in_tensor.update(
+                dataset.compute_voxel_verts(
+                    optimed_pose,
+                    optimed_orient,
+                    optimed_betas,
+                    optimed_trans,
+                    data["scale"],
+                )
+            )
+
+        with torch.no_grad():
+            verts_pr, faces_pr, _ = model.test_single(in_tensor)
+
+        recon_obj = trimesh.Trimesh(
+            verts_pr, faces_pr, process=False, maintains_order=True
+        )
+        recon_obj.export(
+            os.path.join(args.out_dir, cfg.name,
+                         f"obj/{data['name']}_recon.obj")
+        )
+
+        # Isotropic Explicit Remeshing for better geometry topology
+        verts_refine, faces_refine = remesh(os.path.join(args.out_dir, cfg.name,
+                                                         f"obj/{data['name']}_recon.obj"), 0.5, device)
+
+        # define local_affine deform verts
+        mesh_pr = Meshes(verts_refine, faces_refine).to(device)
+        local_affine_model = LocalAffine(
+            mesh_pr.verts_padded().shape[1], mesh_pr.verts_padded().shape[0], mesh_pr.edges_packed()).to(device)
+        optimizer_cloth = torch.optim.Adam(
+            [{'params': local_affine_model.parameters()}], lr=1e-4, amsgrad=True)
+
+        scheduler_cloth = torch.optim.lr_scheduler.ReduceLROnPlateau(
+            optimizer_cloth,
+            mode="min",
+            factor=0.1,
+            verbose=0,
+            min_lr=1e-5,
+            patience=args.patience,
+        )
+
+        with torch.no_grad():
+            per_loop_lst = []
+            rotate_recon_lst = dataset.render.get_rgb_image(cam_ids=[
+                0, 1, 2, 3])
+            per_loop_lst.extend(rotate_recon_lst)
+            per_data_lst.append(get_optim_grid_image(
+                per_loop_lst, None, type="cloth"))
+
+        final = None
+
+        if args.loop_cloth > 0:
+
+            loop_cloth = tqdm(range(args.loop_cloth))
+
+            for i in loop_cloth:
+
+                per_loop_lst = []
+
+                optimizer_cloth.zero_grad()
+
+                deformed_verts, stiffness, rigid = local_affine_model(
+                    verts_refine.to(device), return_stiff=True)
+                mesh_pr = mesh_pr.update_padded(deformed_verts)
+
+                # losses for laplacian, edge, normal consistency
+                update_mesh_shape_prior_losses(mesh_pr, losses)
+
+                in_tensor["P_normal_F"], in_tensor["P_normal_B"] = dataset.render_normal(
+                    mesh_pr.verts_padded(), mesh_pr.faces_padded())
+
+                diff_F_cloth = torch.abs(
+                    in_tensor["P_normal_F"] - in_tensor["normal_F"])
+                diff_B_cloth = torch.abs(
+                    in_tensor["P_normal_B"] - in_tensor["normal_B"])
+
+                losses["cloth"]["value"] = (diff_F_cloth + diff_B_cloth).mean()
+                losses["stiffness"]["value"] = torch.mean(stiffness)
+                losses["rigid"]["value"] = torch.mean(rigid)
+
+                # Weighted sum of the losses
+                cloth_loss = torch.tensor(0.0, requires_grad=True).to(device)
+                pbar_desc = "Cloth Refinement --- "
+
+                for k in losses.keys():
+                    if k not in ["normal", "silhouette"] and losses[k]["weight"] > 0.0:
+                        cloth_loss = cloth_loss + \
+                            losses[k]["value"] * losses[k]["weight"]
+                        pbar_desc += f"{k}:{losses[k]['value']* losses[k]['weight']:.5f} | "
+
+                pbar_desc += f"Total: {cloth_loss:.5f}"
+                loop_cloth.set_description(pbar_desc)
+
+                # update params
+                cloth_loss.backward(retain_graph=True)
+                optimizer_cloth.step()
+                scheduler_cloth.step(cloth_loss)
+
+                # for vis
+                with torch.no_grad():
+                    if i % args.vis_freq == 0:
+
+                        rotate_recon_lst = dataset.render.get_rgb_image(cam_ids=[
+                            0, 1, 2, 3])
+
+                        per_loop_lst.extend(
+                            [
+                                in_tensor["image"],
+                                in_tensor["P_normal_F"],
+                                in_tensor["normal_F"],
+                                diff_F_cloth / 2.0,
+                            ]
+                        )
+                        per_loop_lst.extend(
+                            [
+                                in_tensor["image"],
+                                in_tensor["P_normal_B"],
+                                in_tensor["normal_B"],
+                                diff_B_cloth / 2.0,
+                            ]
+                        )
+                        per_loop_lst.extend(rotate_recon_lst)
+                        per_data_lst.append(
+                            get_optim_grid_image(
+                                per_loop_lst, None, type="cloth")
+                        )
+
+            # gif for optimization
+            per_data_lst[1].save(
+                os.path.join(
+                    args.out_dir, cfg.name, f"refinement/{data['name']}_cloth.gif"
+                ),
+                save_all=True,
+                append_images=per_data_lst[2:],
+                duration=500,
+                loop=0,
+            )
+
+            if args.vis_freq == 1:
+                image2vid(
+                    per_data_lst,
+                    os.path.join(
+                        args.out_dir, cfg.name, f"refinement/{data['name']}_cloth.avi"
+                    ),
+                )
+
+            final = trimesh.Trimesh(
+                mesh_pr.verts_packed().detach().squeeze(0).cpu(),
+                mesh_pr.faces_packed().detach().squeeze(0).cpu(),
+                process=False, maintains_order=True
+            )
+            final_colors = query_color(
+                mesh_pr.verts_packed().detach().squeeze(0).cpu(),
+                mesh_pr.faces_packed().detach().squeeze(0).cpu(),
+                in_tensor["image"],
+                device=device,
+            )
+            final.visual.vertex_colors = final_colors
+            final.export(
+                f"{args.out_dir}/{cfg.name}/obj/{data['name']}_refine.obj")
+
+        # always export visualized png regardless of the cloth refinment
+        per_data_lst[-1].save(
+            os.path.join(args.out_dir, cfg.name,
+                         f"png/{data['name']}_cloth.png")
+        )
+
+        # always export visualized video regardless of the cloth refinment
+        if args.export_video:
+            if final is not None:
+                verts_lst = [verts_pr, final.vertices]
+                faces_lst = [faces_pr, final.faces]
+            else:
+                verts_lst = [verts_pr]
+                faces_lst = [faces_pr]
+
+            # self-rotated video
+            dataset.render.load_meshes(
+                verts_lst, faces_lst)
+            dataset.render.get_rendered_video(
+                [data["ori_image"], rgb_norm],
+                os.path.join(args.out_dir, cfg.name,
+                             f"vid/{data['name']}_cloth.mp4"),
+            )
+
+        # garment extraction from deepfashion images
+        if not (args.seg_dir is None):
+            if final is not None:
+                recon_obj = final.copy()
+
+            os.makedirs(os.path.join(
+                args.out_dir, cfg.name, "clothes"), exist_ok=True)
+            os.makedirs(os.path.join(args.out_dir, cfg.name,
+                        "clothes", "info"), exist_ok=True)
+            for seg in data['segmentations']:
+                # These matrices work for PyMaf, not sure about the other hps type
+                K = np.array([[1.0000,  0.0000,  0.0000,  0.0000],
+                              [0.0000,  1.0000,  0.0000,  0.0000],
+                              [0.0000,  0.0000, -0.5000,  0.0000],
+                              [-0.0000, -0.0000,  0.5000,  1.0000]]).T
+
+                R = np.array([[-1.,  0.,  0.],
+                              [0.,  1.,  0.],
+                              [0.,  0., -1.]])
+
+                t = np.array([[-0.,  -0., 100.]])
+                clothing_obj = extract_cloth(recon_obj, seg, K, R, t, smpl_obj)
+                if clothing_obj is not None:
+                    cloth_type = seg['type'].replace(' ', '_')
+                    cloth_info = {
+                        'betas': optimed_betas,
+                        'body_pose': optimed_pose,
+                        'global_orient': optimed_orient,
+                        'pose2rot': False,
+                        'clothing_type': cloth_type,
+                    }
+
+                    file_id = f"{data['name']}_{cloth_type}"
+                    with open(os.path.join(args.out_dir, cfg.name, "clothes", "info", f"{file_id}_info.pkl"), 'wb') as fp:
+                        pickle.dump(cloth_info, fp)
+
+                    clothing_obj.export(os.path.join(
+                        args.out_dir, cfg.name, "clothes", f"{file_id}.obj"))
+                else:
+                    print(
+                        f"Unable to extract clothing of type {seg['type']} from image {data['name']}")

configs/icon-filter.yaml

@@ -0,0 +1,25 @@
+name: icon-filter
+ckpt_dir: "./data/ckpt/"
+resume_path: "./data/ckpt/icon-filter.ckpt"
+normal_path: "./data/ckpt/normal.ckpt"
+
+test_mode: True
+batch_size: 1
+
+net:
+  mlp_dim: [256, 512, 256, 128, 1]
+  res_layers: [2,3,4]
+  num_stack: 2
+  prior_type: "icon" # icon/pamir/icon
+  use_filter: True
+  in_geo: (('normal_F',3), ('normal_B',3))
+  in_nml: (('image',3), ('T_normal_F',3), ('T_normal_B',3))
+  smpl_feats: ['sdf', 'norm', 'vis', 'cmap']
+  gtype: 'HGPIFuNet'
+  norm_mlp: 'batch'
+  hourglass_dim: 6
+  smpl_dim: 7
+
+# user defined
+mcube_res: 512    # occupancy field resolution, higher --> more details
+clean_mesh: False # if True, will remove floating pieces

configs/icon-nofilter.yaml

@@ -0,0 +1,25 @@
+name: icon-nofilter
+ckpt_dir: "./data/ckpt/"
+resume_path: "./data/ckpt/icon-nofilter.ckpt"
+normal_path: "./data/ckpt/normal.ckpt"
+
+test_mode: True
+batch_size: 1
+
+net:
+  mlp_dim: [256, 512, 256, 128, 1]
+  res_layers: [2,3,4]
+  num_stack: 2
+  prior_type: "icon" # icon/pamir/icon
+  use_filter: False
+  in_geo: (('normal_F',3), ('normal_B',3))
+  in_nml: (('image',3), ('T_normal_F',3), ('T_normal_B',3))
+  smpl_feats: ['sdf', 'norm', 'vis', 'cmap']
+  gtype: 'HGPIFuNet'
+  norm_mlp: 'batch'
+  hourglass_dim: 6
+  smpl_dim: 7
+
+# user defined
+mcube_res: 512    # occupancy field resolution, higher --> more details
+clean_mesh: False # if True, will remove floating pieces

configs/pamir.yaml

@@ -0,0 +1,24 @@
+name: pamir
+ckpt_dir: "./data/ckpt/"
+resume_path: "./data/ckpt/pamir.ckpt"
+normal_path: "./data/ckpt/normal.ckpt"
+
+test_mode: True
+batch_size: 1
+
+net:
+  mlp_dim: [256, 512, 256, 128, 1]
+  res_layers: [2,3,4]
+  num_stack: 2
+  prior_type: "pamir" # icon/pamir/icon
+  use_filter: True
+  in_geo: (('image',3), ('normal_F',3), ('normal_B',3))
+  in_nml: (('image',3), ('T_normal_F',3), ('T_normal_B',3))
+  gtype: 'HGPIFuNet'
+  norm_mlp: 'batch'
+  hourglass_dim: 6
+  voxel_dim: 7
+
+# user defined
+mcube_res: 512    # occupancy field resolution, higher --> more details
+clean_mesh: False # if True, will remove floating pieces

configs/pifu.yaml

@@ -0,0 +1,24 @@
+name: pifu
+ckpt_dir: "./data/ckpt/"
+resume_path: "./data/ckpt/pifu.ckpt"
+normal_path: "./data/ckpt/normal.ckpt"
+
+test_mode: True
+batch_size: 1
+
+net:
+  mlp_dim: [256, 512, 256, 128, 1]
+  res_layers: [2,3,4]
+  num_stack: 2
+  prior_type: "pifu" # icon/pamir/icon
+  use_filter: True
+  in_geo: (('image',3), ('normal_F',3), ('normal_B',3))
+  in_nml: (('image',3), ('T_normal_F',3), ('T_normal_B',3))
+  gtype: 'HGPIFuNet'
+  norm_mlp: 'batch'
+  hourglass_dim: 12
+
+
+# user defined
+mcube_res: 512    # occupancy field resolution, higher --> more details
+clean_mesh: False # if True, will remove floating pieces

environment.yaml

@@ -0,0 +1,16 @@
+name: icon
+channels:
+  - pytorch-lts
+  - nvidia
+  - conda-forge
+  - fvcore
+  - iopath
+  - bottler
+  - defaults
+dependencies:
+  - pytorch
+  - torchvision
+  - fvcore
+  - iopath
+  - nvidiacub
+  - pyembree

examples/segmentation/003883.json

@@ -0,0 +1,136 @@
+{
+  "item2": {
+    "segmentation": [
+      [
+        232.29572649572654, 34.447388414055126, 237.0364672364673,
+        40.57084520417861, 244.9377018043686, 47.089363722697165,
+        252.04881291547974, 49.65726495726508, 262.5179487179489,
+        51.43504273504287, 269.233998100665, 50.447388414055204,
+        277.5446343779678, 49.12725546058881, 285.64339981006657,
+        46.16429249762584, 294.9273504273506, 41.22602089268754,
+        299.9377967711301, 36.514245014245084, 304.67853751187084,
+        30.588319088319132, 306.0612535612536, 25.65004748338083,
+        307.64150047483383, 23.477207977207982, 311.19705603038943,
+        24.859924026590704, 317.12298195631536, 28.020417853751216,
+        323.04890788224134, 29.008072174738874, 331.34520417853764,
+        30.193257359924065, 339.4439696106365, 34.7364672364673,
+        346.75261158594515, 39.279677113010536, 350.11063627730323,
+        44.61301044634389, 355.00541310541314, 61.422317188983875,
+        358.9560303893638, 77.6198480531815, 362.1165242165243,
+        90.26182336182353, 364.88195631528976, 103.29886039886063,
+        367.6473884140552, 118.11367521367552, 369.42516619183294,
+        129.37293447293484, 369.2324786324788, 132.60550807217476,
+        365.6769230769232, 134.77834757834762, 359.15840455840464,
+        138.3339031339032, 353.43000949667623, 140.70427350427357,
+        351.4547008547009, 141.4943969610637, 351.25716999050337,
+        138.5314339981007, 351.05963912630585, 136.75365622032294,
+        345.7263057929725, 137.34624881291552, 337.8250712250712,
+        139.51908831908838, 331.5040835707502, 141.09933523266864,
+        324.7880341880341, 143.66723646723653, 322.2201329534662,
+        146.43266856600198, 322.2201329534662, 151.5684710351378,
+        323.0102564102563, 160.6548907882243, 324.95185185185176,
+        173.44615384615395, 325.34691358024685, 190.23627730294416,
+        325.93950617283946, 205.64368471035164, 325.93950617283946,
+        215.71775878442577, 325.93950617283946, 220.06343779677147,
+        322.7790123456789, 223.22393162393197, 315.0753086419752,
+        228.55726495726532, 309.34691358024673, 230.53257359924066,
+        290.1866096866098, 230.87929724596398, 263.91500474833805,
+        229.6941120607788, 236.45821462488112, 229.29905033238373,
+        218.48290598290572, 226.73114909781583, 202.65650522317188,
+        224.82811016144353, 197.71823361823357, 221.07502374169044,
+        195.15033238366567, 214.55650522317188, 195.74292497625825,
+        200.53181386514711, 197.125641025641, 180.5811965811964,
+        197.33285849952523, 164.68736942070285, 198.51804368471042,
+        154.21823361823365, 198.51804368471042, 138.61329534662863,
+        193.5797720797721, 136.4404558404558, 185.08594491927823,
+        133.08243114909774, 177.77730294396957, 128.73675213675205,
+        174.41927825261152, 128.53922127255453, 173.82668566001894,
+        133.2799620132953, 174.02421652421646, 136.24292497625825,
+        172.83903133903127, 137.03304843304838, 167.11063627730283,
+        134.86020892687554, 159.9995251661917, 130.51452991452985,
+        159.01187084520404, 129.1318138651471, 159.60446343779662,
+        123.60094966761622, 162.6012345679013, 111.57578347578357,
+        165.95925925925934, 98.53874643874646, 170.30493827160504,
+        82.7362773029439, 173.92307692307693, 70.05584045584048,
+        177.08357075023744, 54.84596391263053, 180.58129154795822,
+        41.73190883190885, 183.14919278252614, 34.423266856600165,
+        188.51623931623936, 30.279962013295354, 195.6273504273505,
+        25.539221272554588, 201.75080721747398, 22.971320037986676,
+        211.23228869895553, 22.37872744539408, 221.10883190883212,
+        20.996011396011355, 224.8619183285852, 20.996011396011355,
+        226.04710351377042, 23.56391263057927, 229.01006647673339,
+        30.279962013295354
+      ]
+    ],
+    "category_id": 1,
+    "category_name": "short sleeve top"
+  },
+  "item1": {
+    "segmentation": [
+      [
+        201.51804815682925, 224.7401022799914, 218.41555508203712,
+        227.23317707223518, 236.42109524824218, 228.89522693373104,
+        256.91971020669104, 229.44924355422967, 280.188408267633,
+        230.2802684849776, 296.53189857234224, 230.2802684849776,
+        313.7064138077994, 229.72625186447897, 315.32667803111013,
+        236.8076070743661, 317.8197528233539, 240.96273172810572,
+        318.65077775410185, 246.2258896228426, 321.4208608565949,
+        253.15109737907534, 322.8059024078415, 265.0624547197956,
+        324.74496057958663, 273.6497123375242, 325.9612827615598,
+        284.4076070743661, 325.40726614106114, 299.9200724483274,
+        324.29923290006394, 316.8175793735353, 322.0831664180694,
+        325.9588536117625, 320.16803750266354, 336.5366716386107,
+        316.0129128489239, 344.01589601534204, 315.18188791817596,
+        357.86631152780745, 312.4118048156829, 368.1156190070319,
+        308.5336884721926, 378.64193479650567, 306.31762199019806,
+        385.29013424248905, 305.76360536969946, 398.3095248242066,
+        305.48659705945016, 409.6668655444283, 304.94393777967184,
+        419.3418708715109, 302.7278712976774, 427.0981035584915,
+        301.3428297464308, 433.74630300447495, 301.3428297464308,
+        445.3806520349459, 300.5118048156829, 461.72414233965515,
+        299.89735776688684, 467.352311953974, 297.9582995951417,
+        477.60161943319844, 295.1882164926486, 491.7290432559132,
+        293.52616663115276, 497.2692094608994, 291.8641167696569,
+        503.36339228638417, 291.3101001491583, 510.8426166631155,
+        289.37104197741314, 513.8897080758579, 287.4433411463882,
+        519.2043682079693, 283.0112081823993, 519.7583848284679,
+        275.5319838056679, 519.4813765182186, 270.26882591093107,
+        518.096334966972, 265.8366929469421, 513.6642020029831,
+        263.62062646494763, 509.78608565949276, 264.7286597059449,
+        498.9827615597697, 265.2826763264435, 478.76115491157015,
+        266.1137012571914, 467.1268058810992, 266.1137012571914,
+        454.6614319198803, 264.17464308544623, 441.64204133816276,
+        263.06660984444903, 424.19051779245626, 261.5834221180482,
+        407.2581504368212, 259.92137225655233, 396.45482633709815,
+        257.1512891540592, 380.1113360323889, 257.42829746430857,
+        359.05870445344146, 256.8742808438099, 338.56008949499255,
+        256.8742808438099, 321.3855742595354, 254.10419774131685,
+        320.5545493287875, 251.05710632857443, 326.6487321542723,
+        249.39505646707858, 339.1141061154912, 249.11804815682927,
+        356.28862135094835, 248.28702322608135, 372.3551033454083,
+        245.23993181333896, 387.59056040912026, 243.5766673769444,
+        409.1404219049649, 241.91461751544855, 424.92989558917554,
+        240.52957596420202, 440.4423609631369, 238.86752610270617,
+        455.40080971659955, 238.86752610270617, 470.91327509056083,
+        238.31350948220754, 486.42574046452216, 238.81966759002768,
+        501.19639889196685, 239.6506925207756, 511.168698060942,
+        236.0495844875346, 515.6008310249309, 229.40138504155118,
+        519.4789473684212, 221.6451523545705, 520.3099722991692,
+        216.65900277008296, 517.2628808864267, 213.33490304709125,
+        509.50664819944615, 208.3487534626037, 491.50110803324105,
+        205.8556786703599, 475.1576177285318, 203.63961218836545,
+        460.75318559556774, 203.63961218836545, 443.3016620498613,
+        203.63961218836545, 421.9720221606645, 200.59252077562303,
+        415.60083102493036, 197.5052844662264, 406.9847858512679,
+        195.28921798423193, 392.0263370978052, 193.35015981248677,
+        370.97370551885774, 190.857085020243, 343.82689111442545,
+        187.8099936075006, 322.77425953547794, 187.0028979330919,
+        309.89237161730256, 186.17187300234397, 291.33281483059886,
+        188.11093117408916, 266.67907521841033, 191.15802258683155,
+        250.3355849137011, 196.69818879181773, 234.82311953973982
+      ]
+    ],
+    "category_id": 8,
+    "category_name": "trousers"
+  }
+}

examples/segmentation/028009.json

@@ -0,0 +1,191 @@
+{
+  "item2": {
+    "segmentation": [
+      [
+        314.7474747474744, 204.84848484848482, 328.9696969696967,
+        209.7373737373737, 342.74747474747454, 211.95959595959593,
+        360.0808080808079, 211.07070707070704, 375.19191919191906,
+        210.18181818181816, 384.5252525252524, 207.07070707070704,
+        390.30303030303025, 204.84848484848482, 396.080808080808,
+        201.29292929292924, 402.3030303030303, 204.40404040404036,
+        412.969696969697, 203.9595959595959, 425.8585858585859,
+        206.18181818181813, 434.3030303030304, 211.95959595959593,
+        439.63636363636374, 223.0707070707071, 444.0808080808082,
+        234.18181818181824, 448.52525252525265, 250.62626262626276,
+        449.41414141414157, 260.848484848485, 452.08080808080825,
+        279.0707070707073, 456.08080808080825, 300.84848484848516,
+        457.858585858586, 308.40404040404076, 460.5252525252526,
+        315.7575757575756, 460.96969696969705, 329.97979797979787,
+        460.5252525252526, 345.9797979797979, 456.969696969697,
+        363.75757575757575, 453.41414141414145, 373.5353535353536,
+        450.3030303030303, 385.97979797979804, 447.1919191919192,
+        393.9797979797981, 443.6363636363636, 401.9797979797981,
+        438.3030303030303, 403.7575757575759, 433.85858585858585,
+        401.09090909090924, 430.7474747474747, 393.0909090909092,
+        426.7474747474747, 383.3131313131314, 424.9696969696969,
+        374.8686868686869, 424.9696969696969, 369.0909090909091,
+        423.63636363636357, 363.3131313131313, 423.63636363636357,
+        359.3131313131313, 423.63636363636357, 352.6464646464646,
+        420.9696969696969, 350.86868686868684, 422.74747474747466,
+        345.53535353535347, 422.74747474747466, 340.64646464646455,
+        422.74747474747466, 332.2020202020201, 421.8585858585858,
+        321.53535353535335, 418.74747474747466, 313.0909090909089,
+        416.5252525252524, 306.4242424242422, 412.9696969696969,
+        314.8686868686867, 410.3030303030302, 320.20202020202004,
+        411.6363636363635, 327.3131313131312, 414.74747474747466,
+        336.2020202020201, 418.74747474747466, 351.7575757575757,
+        420.9696969696969, 365.0909090909091, 423.1919191919191,
+        377.0909090909091, 423.1919191919191, 385.0909090909092,
+        424.5252525252525, 398.42424242424255, 396.0808080808079,
+        398.42424242424255, 374.7474747474745, 400.6464646464648,
+        354.7474747474744, 400.6464646464648, 331.6363636363632,
+        400.6464646464648, 313.41414141414094, 400.6464646464648,
+        305.4141414141409, 399.3131313131314, 297.4141414141409,
+        396.6464646464648, 284.525252525252, 396.2020202020203,
+        282.8686868686866, 391.59595959595964, 282.42424242424215,
+        373.81818181818176, 282.42424242424215, 358.26262626262616,
+        281.09090909090884, 334.70707070707056, 281.5353535353533,
+        313.37373737373713, 283.31313131313107, 297.3737373737371,
+        282.8686868686866, 283.1515151515148, 280.6464646464644,
+        266.7070707070703, 271.313131313131, 253.3737373737369,
+        264.6464646464643, 246.70707070707022, 257.5353535353532,
+        239.59595959595907, 249.9797979797976, 228.9292929292924,
+        242.42424242424204, 220.92929292929236, 233.17171717171723,
+        209.01010101010093, 225.1717171717172, 194.78787878787867,
+        222.06060606060606, 185.4545454545453, 224.2828282828283,
+        179.6767676767675, 230.0606060606061, 171.67676767676747,
+        232.72727272727278, 169.89898989898967, 243.83838383838392,
+        167.67676767676744, 256.2828282828284, 165.4545454545452,
+        274.06060606060623, 165.4545454545452, 291.8383838383841,
+        167.67676767676744, 302.5050505050508, 168.1212121212119,
+        310.94949494949526, 177.0101010101008, 314.0606060606064,
+        181.45454545454527, 314.94949494949526, 187.2323232323231,
+        312.7272727272731, 193.01010101010087, 307.8383838383842,
+        191.2323232323231, 302.94949494949526, 193.45454545454533,
+        292.727272727273, 193.45454545454533, 290.50505050505075,
+        195.67676767676755, 287.39393939393966, 197.45454545454533,
+        285.61616161616183, 197.45454545454533, 283.3939393939396,
+        193.89898989898978, 278.94949494949515, 197.45454545454533,
+        274.94949494949515, 199.67676767676755, 279.83838383838406,
+        201.45454545454535, 286.50505050505075, 201.45454545454535,
+        291.8383838383841, 201.8989898989898, 296.2828282828286,
+        202.7878787878787, 303.3939393939397, 202.34343434343424
+      ]
+    ],
+    "category_id": 2,
+    "category_name": "long sleeve top"
+  },
+  "item1": {
+    "segmentation": [
+      [
+        346.9494949494949, 660.6868686868687, 397.6161616161618,
+        661.5757575757576, 398.06060606060623, 674.0202020202021,
+        398.94949494949515, 691.3535353535356, 397.6161616161618,
+        710.0202020202022, 395.838383838384, 726.0202020202023,
+        393.1717171717173, 742.0202020202023, 346.9494949494949,
+        738.9090909090912, 346.50505050505046, 724.2424242424245,
+        347.3939393939394, 713.5757575757578, 348.72727272727275,
+        706.0202020202022, 349.17171717171715, 686.0202020202022,
+        348.72727272727275, 675.7979797979799, 347.3939393939394,
+        667.7979797979799
+      ],
+      [
+        283.71717171717165, 396.68686868686876, 289.9393939393939,
+        396.68686868686876, 303.27272727272725, 397.1313131313132,
+        312.16161616161617, 399.7979797979799, 334.3838383838385,
+        400.68686868686876, 351.7171717171719, 400.68686868686876,
+        361.93939393939417, 401.5757575757577, 376.60606060606085,
+        401.5757575757577, 390.82828282828314, 398.46464646464653,
+        410.3838383838388, 397.5757575757577, 425.0505050505055,
+        394.46464646464653, 431.71717171717216, 422.9090909090911,
+        434.38383838383885, 447.79797979798, 430.38383838383885,
+        478.0202020202024, 423.2727272727277, 507.79797979798025,
+        418.3838383838388, 530.0202020202025, 411.8787878787878,
+        557.3333333333333, 403.43434343434336, 590.6666666666666,
+        400.7676767676767, 611.5555555555557, 399.8787878787878,
+        619.1111111111112, 399.8787878787878, 630.6666666666669,
+        398.10101010101, 635.1111111111113, 399.43434343434336,
+        641.7777777777779, 399.43434343434336, 656.4444444444447,
+        398.10101010101, 662.666666666667, 347.4343434343432, 662.666666666667,
+        346.1010101010098, 637.7777777777779, 347.4343434343432,
+        610.6666666666667, 349.21212121212096, 576.4444444444445,
+        350.98989898989873, 556.4444444444443, 349.6565656565654,
+        541.3333333333331, 348.32323232323205, 535.9999999999998,
+        348.32323232323205, 523.5555555555553, 349.21212121212096,
+        505.33333333333303, 342.5454545454543, 511.5555555555553,
+        338.9898989898987, 516.8888888888887, 334.5454545454542,
+        523.5555555555553, 325.6565656565653, 543.111111111111,
+        319.87878787878753, 556.4444444444443, 314.1010101010097,
+        568.4444444444443, 307.8787878787875, 583.1111111111111,
+        300.3232323232319, 608.0000000000001, 298.10101010100965,
+        617.7777777777778, 298.5454545454541, 624.0000000000001,
+        295.43434343434296, 628.0000000000001, 293.2121212121208,
+        628.0000000000001, 293.6565656565652, 632.4444444444446,
+        291.43434343434296, 638.6666666666669, 290.54545454545405,
+        644.4444444444447, 292.3232323232319, 648.8888888888891,
+        303.8787878787875, 667.1111111111114, 313.65656565656525,
+        684.0000000000003, 319.87878787878753, 700.8888888888893,
+        322.54545454545416, 712.8888888888894, 324.323232323232,
+        720.0000000000005, 327.87878787878753, 731.5555555555561,
+        330.9898989898987, 738.6666666666672, 331.87878787878753,
+        743.1111111111117, 334.5454545454542, 745.7777777777783,
+        336.3232323232325, 749.1313131313133, 338.54545454545473,
+        754.0202020202022, 338.54545454545473, 757.5757575757577,
+        341.6565656565658, 760.6868686868688, 344.76767676767696,
+        767.3535353535356, 345.2121212121214, 770.9090909090911,
+        346.9898989898992, 754.0202020202022, 347.43434343434365,
+        738.909090909091, 393.2121212121216, 740.6868686868687,
+        389.65656565656604, 764.6868686868688, 386.5454545454549,
+        784.2424242424245, 384.3232323232327, 806.9090909090912,
+        382.54545454545485, 812.686868686869, 381.13131313131316,
+        818.7070707070708, 378.020202020202, 828.4848484848485,
+        375.35353535353534, 839.5959595959597, 374.9090909090909,
+        854.2626262626264, 373.1313131313131, 856.9292929292931,
+        376.24242424242425, 864.9292929292931, 372.24242424242425,
+        874.2626262626264, 366.4646464646464, 880.9292929292932,
+        357.13131313131305, 872.9292929292932, 345.13131313131305,
+        868.0404040404043, 337.131313131313, 867.1515151515154,
+        337.131313131313, 856.0404040404042, 338.4646464646463,
+        850.7070707070709, 336.2424242424241, 846.2626262626264,
+        335.3535353535352, 841.3737373737375, 338.4646464646463,
+        827.5959595959597, 342.0202020202019, 815.5959595959596,
+        344.6868686868686, 809.3737373737374, 344.6868686868686,
+        796.4848484848484, 344.6868686868686, 786.7070707070707,
+        346.0202020202019, 779.151515151515, 344.24242424242414,
+        776.0404040404039, 343.3535353535352, 786.2626262626262,
+        342.0202020202019, 796.0404040404039, 338.90909090909076,
+        801.8181818181818, 333.57575757575745, 809.3737373737374,
+        326.02020202020185, 813.8181818181819, 320.242424242424,
+        812.4848484848485, 318.02020202020185, 810.7070707070707,
+        317.13131313131294, 807.1515151515151, 315.79797979797956,
+        803.5959595959596, 313.57575757575734, 799.5959595959596,
+        311.3535353535351, 793.8181818181818, 306.90909090909065,
+        791.1515151515151, 305.57575757575734, 787.5959595959595,
+        304.242424242424, 782.7070707070706, 302.02020202020174,
+        776.4848484848484, 298.90909090909065, 773.8181818181816,
+        294.90909090909065, 771.151515151515, 290.34343434343435,
+        758.909090909091, 284.5656565656566, 742.020202020202,
+        278.78787878787875, 729.5757575757575, 270.3434343434343,
+        713.131313131313, 257.8989898989898, 689.1313131313129,
+        247.2323232323231, 669.1313131313128, 239.23232323232307,
+        657.5757575757573, 233.89898989898973, 642.9090909090905,
+        233.0101010101008, 634.0202020202016, 233.45454545454527,
+        630.0202020202016, 235.23232323232304, 611.7979797979793,
+        241.93939393939402, 583.0707070707073, 245.93939393939405,
+        567.5151515151516, 251.2727272727274, 540.4040404040404,
+        256.1616161616163, 518.6262626262626, 260.60606060606074,
+        501.2929292929292, 263.7171717171719, 493.7373737373736,
+        268.16161616161634, 481.73737373737356, 270.38383838383857,
+        469.73737373737356, 272.6060606060608, 462.18181818181796,
+        276.1616161616164, 457.7373737373735, 276.1616161616164,
+        454.1818181818179, 277.05050505050525, 450.1818181818179,
+        278.828282828283, 433.292929292929, 278.3838383838386,
+        419.0707070707067, 278.828282828283, 417.29292929292893,
+        281.0505050505053, 414.1818181818178, 281.93939393939417,
+        404.8484848484844, 283.71717171717194, 401.2929292929289
+      ]
+    ],
+    "category_id": 8,
+    "category_name": "trousers"
+  }
+}

fetch_data.sh

@@ -0,0 +1,60 @@
+#!/bin/bash
+urle () { [[ "${1}" ]] || return 1; local LANG=C i x; for (( i = 0; i < ${#1}; i++ )); do x="${1:i:1}"; [[ "${x}" == [a-zA-Z0-9.~-] ]] && echo -n "${x}" || printf '%%%02X' "'${x}"; done; echo; }
+
+mkdir -p data/smpl_related/models
+
+# username and password input
+echo -e "\nYou need to register at https://icon.is.tue.mpg.de/, according to Installation Instruction."
+read -p "Username (ICON):" username
+read -p "Password (ICON):" password
+username=$(urle $username)
+password=$(urle $password)
+
+# SMPL (Male, Female)
+echo -e "\nDownloading SMPL..."
+wget --post-data "username=$username&password=$password" 'https://download.is.tue.mpg.de/download.php?domain=smpl&sfile=SMPL_python_v.1.0.0.zip&resume=1' -O './data/smpl_related/models/SMPL_python_v.1.0.0.zip' --no-check-certificate --continue
+unzip data/smpl_related/models/SMPL_python_v.1.0.0.zip -d data/smpl_related/models
+mv data/smpl_related/models/smpl/models/basicModel_f_lbs_10_207_0_v1.0.0.pkl data/smpl_related/models/smpl/SMPL_FEMALE.pkl
+mv data/smpl_related/models/smpl/models/basicmodel_m_lbs_10_207_0_v1.0.0.pkl data/smpl_related/models/smpl/SMPL_MALE.pkl
+cd data/smpl_related/models
+rm -rf *.zip __MACOSX smpl/models smpl/smpl_webuser
+cd ../../..
+
+# SMPL (Neutral, from SMPLIFY)
+echo -e "\nDownloading SMPLify..."
+wget --post-data "username=$username&password=$password" 'https://download.is.tue.mpg.de/download.php?domain=smplify&sfile=mpips_smplify_public_v2.zip&resume=1' -O './data/smpl_related/models/mpips_smplify_public_v2.zip' --no-check-certificate --continue
+unzip data/smpl_related/models/mpips_smplify_public_v2.zip -d data/smpl_related/models
+mv data/smpl_related/models/smplify_public/code/models/basicModel_neutral_lbs_10_207_0_v1.0.0.pkl data/smpl_related/models/smpl/SMPL_NEUTRAL.pkl
+cd data/smpl_related/models
+rm -rf *.zip smplify_public 
+cd ../../..
+
+# ICON
+echo -e "\nDownloading ICON..."
+wget --post-data "username=$username&password=$password" 'https://download.is.tue.mpg.de/download.php?domain=icon&sfile=icon_data.zip&resume=1' -O './data/icon_data.zip' --no-check-certificate --continue
+cd data && unzip icon_data.zip
+mv smpl_data smpl_related/
+rm -f icon_data.zip
+cd ..
+
+function download_for_training () {
+    
+    # SMPL-X (optional)
+    echo -e "\nDownloading SMPL-X..."
+    wget --post-data "username=$1&password=$2" 'https://download.is.tue.mpg.de/download.php?domain=smplx&sfile=models_smplx_v1_1.zip&resume=1' -O './data/smpl_related/models/models_smplx_v1_1.zip' --no-check-certificate --continue
+    unzip data/smpl_related/models/models_smplx_v1_1.zip -d data/smpl_related
+    rm -f data/smpl_related/models/models_smplx_v1_1.zip
+
+    # SMIL (optional)
+    echo -e "\nDownloading SMIL..."
+    wget --post-data "username=$1&password=$2" 'https://download.is.tue.mpg.de/download.php?domain=agora&sfile=smpl_kid_template.npy&resume=1' -O './data/smpl_related/models/smpl/smpl_kid_template.npy' --no-check-certificate --continue
+    wget --post-data "username=$1&password=$2" 'https://download.is.tue.mpg.de/download.php?domain=agora&sfile=smplx_kid_template.npy&resume=1' -O './data/smpl_related/models/smplx/smplx_kid_template.npy' --no-check-certificate --continue
+}
+
+
+read -p "(optional) Download models used for training (y/n)?" choice
+case "$choice" in 
+  y|Y ) download_for_training $username $password;;
+  n|N ) echo "Great job! Try the demo for now!";;
+  * ) echo "Invalid input! Please use y|Y or n|N";;
+esac

install.sh

@@ -0,0 +1,16 @@
+# # conda installation
+# wget https://repo.anaconda.com/miniconda/Miniconda3-py38_4.10.3-Linux-x86_64.sh
+# chmod +x Miniconda3-py38_4.10.3-Linux-x86_64.sh
+# bash Miniconda3-py38_4.10.3-Linux-x86_64.sh -b -f -p /home/user/.local
+# rm Miniconda3-py38_4.10.3-Linux-x86_64.sh
+# conda config --env --set always_yes true
+# conda update -n base -c defaults conda -y
+
+# # conda environment setup
+# conda env create -f environment.yaml
+# conda init bash
+# source /home/user/.bashrc
+# source activate icon
+nvidia-smi
+pip install torch==1.8.2 torchvision==0.9.2 --extra-index-url https://download.pytorch.org/whl/lts/1.8/cu111
+pip install -r requirement.txt

lib/common/cloth_extraction.py

@@ -0,0 +1,170 @@
+import numpy as np
+import json
+import os
+import itertools
+import trimesh
+from matplotlib.path import Path
+from collections import Counter
+from sklearn.neighbors import KNeighborsClassifier
+
+
+def load_segmentation(path, shape):
+    """
+    Get a segmentation mask for a given image
+    Arguments:
+        path: path to the segmentation json file
+        shape: shape of the output mask
+    Returns:
+        Returns a segmentation mask
+    """
+    with open(path) as json_file:
+        dict = json.load(json_file)
+        segmentations = []
+        for key, val in dict.items():
+            if not key.startswith('item'):
+                continue
+
+            # Each item can have multiple polygons. Combine them to one
+            # segmentation_coord = list(itertools.chain.from_iterable(val['segmentation']))
+            # segmentation_coord = np.round(np.array(segmentation_coord)).astype(int)
+
+            coordinates = []
+            for segmentation_coord in val['segmentation']:
+                # The format before is [x1,y1, x2, y2, ....]
+                x = segmentation_coord[::2]
+                y = segmentation_coord[1::2]
+                xy = np.vstack((x, y)).T
+                coordinates.append(xy)
+
+            segmentations.append(
+                {'type': val['category_name'], 'type_id': val['category_id'], 'coordinates': coordinates})
+
+        return segmentations
+
+
+def smpl_to_recon_labels(recon, smpl, k=1):
+    """
+    Get the bodypart labels for the recon object by using the labels from the corresponding smpl object
+    Arguments:
+        recon: trimesh object (fully clothed model)
+        shape: trimesh object (smpl model)
+        k: number of nearest neighbours to use
+    Returns:
+        Returns a dictionary containing the bodypart and the corresponding indices
+    """
+    smpl_vert_segmentation = json.load(
+        open(os.path.join(os.path.dirname(__file__), 'smpl_vert_segmentation.json')))
+    n = smpl.vertices.shape[0]
+    y = np.array([None] * n)
+    for key, val in smpl_vert_segmentation.items():
+        y[val] = key
+
+    classifier = KNeighborsClassifier(n_neighbors=1)
+    classifier.fit(smpl.vertices, y)
+
+    y_pred = classifier.predict(recon.vertices)
+
+    recon_labels = {}
+    for key in smpl_vert_segmentation.keys():
+        recon_labels[key] = list(np.argwhere(
+            y_pred == key).flatten().astype(int))
+
+    return recon_labels
+
+
+def extract_cloth(recon, segmentation, K, R, t, smpl=None):
+    """
+    Extract a portion of a mesh using 2d segmentation coordinates
+    Arguments:
+        recon: fully clothed mesh
+        seg_coord: segmentation coordinates in 2D (NDC)
+        K: intrinsic matrix of the projection
+        R: rotation matrix of the projection
+        t: translation vector of the projection
+    Returns:
+        Returns a submesh using the segmentation coordinates
+    """
+    seg_coord = segmentation['coord_normalized']
+    mesh = trimesh.Trimesh(recon.vertices, recon.faces)
+    extrinsic = np.zeros((3, 4))
+    extrinsic[:3, :3] = R
+    extrinsic[:, 3] = t
+    P = K[:3, :3] @ extrinsic
+
+    P_inv = np.linalg.pinv(P)
+
+    # Each segmentation can contain multiple polygons
+    # We need to check them separately
+    points_so_far = []
+    faces = recon.faces
+    for polygon in seg_coord:
+        n = len(polygon)
+        coords_h = np.hstack((polygon, np.ones((n, 1))))
+        # Apply the inverse projection on homogeneus 2D coordinates to get the corresponding 3d Coordinates
+        XYZ = P_inv @ coords_h[:, :, None]
+        XYZ = XYZ.reshape((XYZ.shape[0], XYZ.shape[1]))
+        XYZ = XYZ[:, :3] / XYZ[:, 3, None]
+
+        p = Path(XYZ[:, :2])
+
+        grid = p.contains_points(recon.vertices[:, :2])
+        indeces = np.argwhere(grid == True)
+        points_so_far += list(indeces.flatten())
+
+    if smpl is not None:
+        num_verts = recon.vertices.shape[0]
+        recon_labels = smpl_to_recon_labels(recon, smpl)
+        body_parts_to_remove = ['rightHand', 'leftToeBase', 'leftFoot', 'rightFoot', 'head',
+                                'leftHandIndex1', 'rightHandIndex1', 'rightToeBase', 'leftHand', 'rightHand']
+        type = segmentation['type_id']
+
+        # Remove additional bodyparts that are most likely not part of the segmentation but might intersect (e.g. hand in front of torso)
+        # https://github.com/switchablenorms/DeepFashion2
+        # Short sleeve clothes
+        if type == 1 or type == 3 or type == 10:
+            body_parts_to_remove += ['leftForeArm', 'rightForeArm']
+        # No sleeves at all or lower body clothes
+        elif type == 5 or type == 6 or type == 12 or type == 13 or type == 8 or type == 9:
+            body_parts_to_remove += ['leftForeArm',
+                                     'rightForeArm', 'leftArm', 'rightArm']
+        # Shorts
+        elif type == 7:
+            body_parts_to_remove += ['leftLeg', 'rightLeg',
+                                     'leftForeArm', 'rightForeArm', 'leftArm', 'rightArm']
+
+        verts_to_remove = list(itertools.chain.from_iterable(
+            [recon_labels[part] for part in body_parts_to_remove]))
+
+        label_mask = np.zeros(num_verts, dtype=bool)
+        label_mask[verts_to_remove] = True
+
+        seg_mask = np.zeros(num_verts, dtype=bool)
+        seg_mask[points_so_far] = True
+
+        # Remove points that belong to other bodyparts
+        # If a vertice in pointsSoFar is included in the bodyparts to remove, then these points should be removed
+        extra_verts_to_remove = np.array(list(seg_mask) and list(label_mask))
+
+        combine_mask = np.zeros(num_verts, dtype=bool)
+        combine_mask[points_so_far] = True
+        combine_mask[extra_verts_to_remove] = False
+
+        all_indices = np.argwhere(combine_mask == True).flatten()
+
+    i_x = np.where(np.in1d(faces[:, 0], all_indices))[0]
+    i_y = np.where(np.in1d(faces[:, 1], all_indices))[0]
+    i_z = np.where(np.in1d(faces[:, 2], all_indices))[0]
+
+    faces_to_keep = np.array(list(set(i_x).union(i_y).union(i_z)))
+    mask = np.zeros(len(recon.faces), dtype=bool)
+    if len(faces_to_keep) > 0:
+        mask[faces_to_keep] = True
+
+        mesh.update_faces(mask)
+        mesh.remove_unreferenced_vertices()
+
+        # mesh.rezero()
+
+        return mesh
+
+    return None

lib/common/config.py

@@ -0,0 +1,218 @@
+
+# -*- 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 yacs.config import CfgNode as CN
+import os
+
+_C = CN(new_allowed=True)
+
+# needed by trainer
+_C.name = 'default'
+_C.gpus = [0]
+_C.test_gpus = [1]
+_C.root = "./data/"
+_C.ckpt_dir = './data/ckpt/'
+_C.resume_path = ''
+_C.normal_path = ''
+_C.corr_path = ''
+_C.results_path = './data/results/'
+_C.projection_mode = 'orthogonal'
+_C.num_views = 1
+_C.sdf = False
+_C.sdf_clip = 5.0
+
+_C.lr_G = 1e-3
+_C.lr_C = 1e-3
+_C.lr_N = 2e-4
+_C.weight_decay = 0.0
+_C.momentum = 0.0
+_C.optim = 'RMSprop'
+_C.schedule = [5, 10, 15]
+_C.gamma = 0.1
+
+_C.overfit = False
+_C.resume = False
+_C.test_mode = False
+_C.test_uv = False
+_C.draw_geo_thres = 0.60
+_C.num_sanity_val_steps = 2
+_C.fast_dev = 0
+_C.get_fit = False
+_C.agora = False
+_C.optim_cloth = False
+_C.optim_body = False
+_C.mcube_res = 256
+_C.clean_mesh = True
+_C.remesh = False
+
+_C.batch_size = 4
+_C.num_threads = 8
+
+_C.num_epoch = 10
+_C.freq_plot = 0.01
+_C.freq_show_train = 0.1
+_C.freq_show_val = 0.2
+_C.freq_eval = 0.5
+_C.accu_grad_batch = 4
+
+_C.test_items = ['sv', 'mv', 'mv-fusion', 'hybrid', 'dc-pred', 'gt']
+
+_C.net = CN()
+_C.net.gtype = 'HGPIFuNet'
+_C.net.ctype = 'resnet18'
+_C.net.classifierIMF = 'MultiSegClassifier'
+_C.net.netIMF = 'resnet18'
+_C.net.norm = 'group'
+_C.net.norm_mlp = 'group'
+_C.net.norm_color = 'group'
+_C.net.hg_down = 'ave_pool'
+_C.net.num_views = 1
+
+# kernel_size, stride, dilation, padding
+
+_C.net.conv1 = [7, 2, 1, 3]
+_C.net.conv3x3 = [3, 1, 1, 1]
+
+_C.net.num_stack = 4
+_C.net.num_hourglass = 2
+_C.net.hourglass_dim = 256
+_C.net.voxel_dim = 32
+_C.net.resnet_dim = 120
+_C.net.mlp_dim = [320, 1024, 512, 256, 128, 1]
+_C.net.mlp_dim_knn = [320, 1024, 512, 256, 128, 3]
+_C.net.mlp_dim_color = [513, 1024, 512, 256, 128, 3]
+_C.net.mlp_dim_multiseg = [1088, 2048, 1024, 500]
+_C.net.res_layers = [2, 3, 4]
+_C.net.filter_dim = 256
+_C.net.smpl_dim = 3
+
+_C.net.cly_dim = 3
+_C.net.soft_dim = 64
+_C.net.z_size = 200.0
+_C.net.N_freqs = 10
+_C.net.geo_w = 0.1
+_C.net.norm_w = 0.1
+_C.net.dc_w = 0.1
+_C.net.C_cat_to_G = False
+
+_C.net.skip_hourglass = True
+_C.net.use_tanh = True
+_C.net.soft_onehot = True
+_C.net.no_residual = True
+_C.net.use_attention = False
+
+_C.net.prior_type = "sdf"
+_C.net.smpl_feats = ['sdf', 'cmap', 'norm', 'vis']
+_C.net.use_filter = True
+_C.net.use_cc = False
+_C.net.use_PE = False
+_C.net.use_IGR = False
+_C.net.in_geo = ()
+_C.net.in_nml = ()
+
+_C.dataset = CN()
+_C.dataset.root = ''
+_C.dataset.set_splits = [0.95, 0.04]
+_C.dataset.types = [
+    "3dpeople", "axyz", "renderpeople", "renderpeople_p27", "humanalloy"
+]
+_C.dataset.scales = [1.0, 100.0, 1.0, 1.0, 100.0 / 39.37]
+_C.dataset.rp_type = "pifu900"
+_C.dataset.th_type = 'train'
+_C.dataset.input_size = 512
+_C.dataset.rotation_num = 3
+_C.dataset.num_precomp = 10  # Number of segmentation classifiers
+_C.dataset.num_multiseg = 500  # Number of categories per classifier
+_C.dataset.num_knn = 10  # for loss/error
+_C.dataset.num_knn_dis = 20  # for accuracy
+_C.dataset.num_verts_max = 20000
+_C.dataset.zray_type = False
+_C.dataset.online_smpl = False
+_C.dataset.noise_type = ['z-trans', 'pose', 'beta']
+_C.dataset.noise_scale = [0.0, 0.0, 0.0]
+_C.dataset.num_sample_geo = 10000
+_C.dataset.num_sample_color = 0
+_C.dataset.num_sample_seg = 0
+_C.dataset.num_sample_knn = 10000
+
+_C.dataset.sigma_geo = 5.0
+_C.dataset.sigma_color = 0.10
+_C.dataset.sigma_seg = 0.10
+_C.dataset.thickness_threshold = 20.0
+_C.dataset.ray_sample_num = 2
+_C.dataset.semantic_p = False
+_C.dataset.remove_outlier = False
+
+_C.dataset.train_bsize = 1.0
+_C.dataset.val_bsize = 1.0
+_C.dataset.test_bsize = 1.0
+
+
+def get_cfg_defaults():
+    """Get a yacs CfgNode object with default values for my_project."""
+    # Return a clone so that the defaults will not be altered
+    # This is for the "local variable" use pattern
+    return _C.clone()
+
+
+# Alternatively, provide a way to import the defaults as
+# a global singleton:
+cfg = _C  # users can `from config import cfg`
+
+# cfg = get_cfg_defaults()
+# cfg.merge_from_file('./configs/example.yaml')
+
+# # Now override from a list (opts could come from the command line)
+# opts = ['dataset.root', './data/XXXX', 'learning_rate', '1e-2']
+# cfg.merge_from_list(opts)
+
+
+def update_cfg(cfg_file):
+    # cfg = get_cfg_defaults()
+    _C.merge_from_file(cfg_file)
+    # return cfg.clone()
+    return _C
+
+
+def parse_args(args):
+    cfg_file = args.cfg_file
+    if args.cfg_file is not None:
+        cfg = update_cfg(args.cfg_file)
+    else:
+        cfg = get_cfg_defaults()
+
+    # if args.misc is not None:
+    #     cfg.merge_from_list(args.misc)
+
+    return cfg
+
+
+def parse_args_extend(args):
+    if args.resume:
+        if not os.path.exists(args.log_dir):
+            raise ValueError(
+                'Experiment are set to resume mode, but log directory does not exist.'
+            )
+
+        # load log's cfg
+        cfg_file = os.path.join(args.log_dir, 'cfg.yaml')
+        cfg = update_cfg(cfg_file)
+
+        if args.misc is not None:
+            cfg.merge_from_list(args.misc)
+    else:
+        parse_args(args)

lib/common/render.py

@@ -0,0 +1,387 @@
+# -*- 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 pytorch3d.renderer import (
+    BlendParams,
+    blending,
+    look_at_view_transform,
+    FoVOrthographicCameras,
+    PointLights,
+    RasterizationSettings,
+    PointsRasterizationSettings,
+    PointsRenderer,
+    AlphaCompositor,
+    PointsRasterizer,
+    MeshRenderer,
+    MeshRasterizer,
+    SoftPhongShader,
+    SoftSilhouetteShader,
+    TexturesVertex,
+)
+from pytorch3d.renderer.mesh import TexturesVertex
+from pytorch3d.structures import Meshes
+from lib.dataset.mesh_util import SMPLX, get_visibility
+
+import lib.common.render_utils as util
+import torch
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+import os
+import cv2
+import math
+from termcolor import colored
+
+
+def image2vid(images, vid_path):
+
+    w, h = images[0].size
+    videodims = (w, h)
+    fourcc = cv2.VideoWriter_fourcc(*'XVID')
+    video = cv2.VideoWriter(vid_path, fourcc, 30, videodims)
+    for image in images:
+        video.write(cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR))
+    video.release()
+
+
+def query_color(verts, faces, image, device):
+    """query colors from points and image
+
+    Args:
+        verts ([B, 3]): [query verts]
+        faces ([M, 3]): [query faces]
+        image ([B, 3, H, W]): [full image]
+
+    Returns:
+        [np.float]: [return colors]
+    """
+
+    verts = verts.float().to(device)
+    faces = faces.long().to(device)
+
+    (xy, z) = verts.split([2, 1], dim=1)
+    visibility = get_visibility(xy, z, faces[:, [0, 2, 1]]).flatten()
+    uv = xy.unsqueeze(0).unsqueeze(2)  # [B, N, 2]
+    uv = uv * torch.tensor([1.0, -1.0]).type_as(uv)
+    colors = (torch.nn.functional.grid_sample(image, uv, align_corners=True)[
+              0, :, :, 0].permute(1, 0) + 1.0) * 0.5 * 255.0
+    colors[visibility == 0.0] = ((Meshes(verts.unsqueeze(0), faces.unsqueeze(
+        0)).verts_normals_padded().squeeze(0) + 1.0) * 0.5 * 255.0)[visibility == 0.0]
+
+    return colors.detach().cpu()
+
+
+class cleanShader(torch.nn.Module):
+    def __init__(self, device="cpu", cameras=None, blend_params=None):
+        super().__init__()
+        self.cameras = cameras
+        self.blend_params = blend_params if blend_params is not None else BlendParams()
+
+    def forward(self, fragments, meshes, **kwargs):
+        cameras = kwargs.get("cameras", self.cameras)
+        if cameras is None:
+            msg = "Cameras must be specified either at initialization \
+                or in the forward pass of TexturedSoftPhongShader"
+
+            raise ValueError(msg)
+
+        # get renderer output
+        blend_params = kwargs.get("blend_params", self.blend_params)
+        texels = meshes.sample_textures(fragments)
+        images = blending.softmax_rgb_blend(
+            texels, fragments, blend_params, znear=-256, zfar=256
+        )
+
+        return images
+
+
+class Render:
+    def __init__(self, size=512, device=torch.device("cuda:0")):
+        self.device = device
+        self.mesh_y_center = 100.0
+        self.dis = 100.0
+        self.scale = 1.0
+        self.size = size
+        self.cam_pos = [(0, 100, 100)]
+
+        self.mesh = None
+        self.deform_mesh = None
+        self.pcd = None
+        self.renderer = None
+        self.meshRas = None
+        self.type = None
+        self.knn = None
+        self.knn_inverse = None
+
+        self.smpl_seg = None
+        self.smpl_cmap = None
+
+        self.smplx = SMPLX()
+
+        self.uv_rasterizer = util.Pytorch3dRasterizer(self.size)
+
+    def get_camera(self, cam_id):
+
+        R, T = look_at_view_transform(
+            eye=[self.cam_pos[cam_id]],
+            at=((0, self.mesh_y_center, 0),),
+            up=((0, 1, 0),),
+        )
+
+        camera = FoVOrthographicCameras(
+            device=self.device,
+            R=R,
+            T=T,
+            znear=100.0,
+            zfar=-100.0,
+            max_y=100.0,
+            min_y=-100.0,
+            max_x=100.0,
+            min_x=-100.0,
+            scale_xyz=(self.scale * np.ones(3),),
+        )
+
+        return camera
+
+    def init_renderer(self, camera, type="clean_mesh", bg="gray"):
+
+        if "mesh" in type:
+
+            # rasterizer
+            self.raster_settings_mesh = RasterizationSettings(
+                image_size=self.size,
+                blur_radius=np.log(1.0 / 1e-4) * 1e-7,
+                faces_per_pixel=30,
+            )
+            self.meshRas = MeshRasterizer(
+                cameras=camera, raster_settings=self.raster_settings_mesh
+            )
+
+        if bg == "black":
+            blendparam = BlendParams(1e-4, 1e-4, (0.0, 0.0, 0.0))
+        elif bg == "white":
+            blendparam = BlendParams(1e-4, 1e-8, (1.0, 1.0, 1.0))
+        elif bg == "gray":
+            blendparam = BlendParams(1e-4, 1e-8, (0.5, 0.5, 0.5))
+
+        if type == "ori_mesh":
+
+            lights = PointLights(
+                device=self.device,
+                ambient_color=((0.8, 0.8, 0.8),),
+                diffuse_color=((0.2, 0.2, 0.2),),
+                specular_color=((0.0, 0.0, 0.0),),
+                location=[[0.0, 200.0, 0.0]],
+            )
+
+            self.renderer = MeshRenderer(
+                rasterizer=self.meshRas,
+                shader=SoftPhongShader(
+                    device=self.device,
+                    cameras=camera,
+                    lights=lights,
+                    blend_params=blendparam,
+                ),
+            )
+
+        if type == "silhouette":
+            self.raster_settings_silhouette = RasterizationSettings(
+                image_size=self.size,
+                blur_radius=np.log(1.0 / 1e-4 - 1.0) * 5e-5,
+                faces_per_pixel=50,
+                cull_backfaces=True,
+            )
+
+            self.silhouetteRas = MeshRasterizer(
+                cameras=camera, raster_settings=self.raster_settings_silhouette
+            )
+            self.renderer = MeshRenderer(
+                rasterizer=self.silhouetteRas, shader=SoftSilhouetteShader()
+            )
+
+        if type == "pointcloud":
+            self.raster_settings_pcd = PointsRasterizationSettings(
+                image_size=self.size, radius=0.006, points_per_pixel=10
+            )
+
+            self.pcdRas = PointsRasterizer(
+                cameras=camera, raster_settings=self.raster_settings_pcd
+            )
+            self.renderer = PointsRenderer(
+                rasterizer=self.pcdRas,
+                compositor=AlphaCompositor(background_color=(0, 0, 0)),
+            )
+
+        if type == "clean_mesh":
+
+            self.renderer = MeshRenderer(
+                rasterizer=self.meshRas,
+                shader=cleanShader(
+                    device=self.device, cameras=camera, blend_params=blendparam
+                ),
+            )
+
+    def VF2Mesh(self, verts, faces):
+
+        if not torch.is_tensor(verts):
+            verts = torch.tensor(verts)
+        if not torch.is_tensor(faces):
+            faces = torch.tensor(faces)
+
+        if verts.ndimension() == 2:
+            verts = verts.unsqueeze(0).float()
+        if faces.ndimension() == 2:
+            faces = faces.unsqueeze(0).long()
+
+        verts = verts.to(self.device)
+        faces = faces.to(self.device)
+
+        mesh = Meshes(verts, faces).to(self.device)
+
+        mesh.textures = TexturesVertex(
+            verts_features=(mesh.verts_normals_padded() + 1.0) * 0.5
+        )
+
+        return mesh
+
+    def load_meshes(self, verts, faces):
+        """load mesh into the pytorch3d renderer
+
+        Args:
+            verts ([N,3]): verts
+            faces ([N,3]): faces
+            offset ([N,3]): offset
+        """
+
+        # camera setting
+        self.scale = 100.0
+        self.mesh_y_center = 0.0
+
+        self.cam_pos = [
+            (0, self.mesh_y_center, 100.0),
+            (100.0, self.mesh_y_center, 0),
+            (0, self.mesh_y_center, -100.0),
+            (-100.0, self.mesh_y_center, 0),
+        ]
+
+        self.type = "color"
+
+        if isinstance(verts, list):
+            self.meshes = []
+            for V, F in zip(verts, faces):
+                self.meshes.append(self.VF2Mesh(V, F))
+        else:
+            self.meshes = [self.VF2Mesh(verts, faces)]
+
+    def get_depth_map(self, cam_ids=[0, 2]):
+
+        depth_maps = []
+        for cam_id in cam_ids:
+            self.init_renderer(self.get_camera(cam_id), "clean_mesh", "gray")
+            fragments = self.meshRas(self.meshes[0])
+            depth_map = fragments.zbuf[..., 0].squeeze(0)
+            if cam_id == 2:
+                depth_map = torch.fliplr(depth_map)
+            depth_maps.append(depth_map)
+
+        return depth_maps
+
+    def get_rgb_image(self, cam_ids=[0, 2]):
+
+        images = []
+        for cam_id in range(len(self.cam_pos)):
+            if cam_id in cam_ids:
+                self.init_renderer(self.get_camera(
+                    cam_id), "clean_mesh", "gray")
+                if len(cam_ids) == 4:
+                    rendered_img = (
+                        self.renderer(self.meshes[0])[
+                            0:1, :, :, :3].permute(0, 3, 1, 2)
+                        - 0.5
+                    ) * 2.0
+                else:
+                    rendered_img = (
+                        self.renderer(self.meshes[0])[
+                            0:1, :, :, :3].permute(0, 3, 1, 2)
+                        - 0.5
+                    ) * 2.0
+                if cam_id == 2 and len(cam_ids) == 2:
+                    rendered_img = torch.flip(rendered_img, dims=[3])
+                images.append(rendered_img)
+
+        return images
+
+    def get_rendered_video(self, images, save_path):
+        
+        self.cam_pos = []
+        for angle in range(360):
+            self.cam_pos.append(
+                (
+                    100.0 * math.cos(np.pi / 180 * angle),
+                    self.mesh_y_center,
+                    100.0 * math.sin(np.pi / 180 * angle),
+                )
+            )
+
+        old_shape = np.array(images[0].shape[:2])
+        new_shape = np.around(
+            (self.size / old_shape[0]) * old_shape).astype(np.int)
+
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+        video = cv2.VideoWriter(
+            save_path, fourcc, 30, (self.size * len(self.meshes) +
+                                    new_shape[1] * len(images), self.size)
+        )
+
+        pbar = tqdm(range(len(self.cam_pos)))
+        pbar.set_description(colored(f"exporting video {os.path.basename(save_path)}...", "blue"))
+        for cam_id in pbar:
+            self.init_renderer(self.get_camera(cam_id), "clean_mesh", "gray")
+
+            img_lst = [
+                np.array(Image.fromarray(img).resize(new_shape[::-1])).astype(np.uint8)[
+                    :, :, [2, 1, 0]
+                ]
+                for img in images
+            ]
+
+            for mesh in self.meshes:
+                rendered_img = (
+                    (self.renderer(mesh)[0, :, :, :3] * 255.0)
+                    .detach()
+                    .cpu()
+                    .numpy()
+                    .astype(np.uint8)
+                )
+
+                img_lst.append(rendered_img)
+            final_img = np.concatenate(img_lst, axis=1)
+            video.write(final_img)
+
+        video.release()
+
+    def get_silhouette_image(self, cam_ids=[0, 2]):
+
+        images = []
+        for cam_id in range(len(self.cam_pos)):
+            if cam_id in cam_ids:
+                self.init_renderer(self.get_camera(cam_id), "silhouette")
+                rendered_img = self.renderer(self.meshes[0])[0:1, :, :, 3]
+                if cam_id == 2 and len(cam_ids) == 2:
+                    rendered_img = torch.flip(rendered_img, dims=[2])
+                images.append(rendered_img)
+
+        return images

lib/common/render_utils.py

@@ -0,0 +1,221 @@
+
+# -*- 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 torch
+from torch import nn
+import trimesh
+import math
+from typing import NewType
+from pytorch3d.structures import Meshes
+from pytorch3d.renderer.mesh import rasterize_meshes
+
+Tensor = NewType('Tensor', torch.Tensor)
+
+
+def solid_angles(points: Tensor,
+                 triangles: Tensor,
+                 thresh: float = 1e-8) -> Tensor:
+    ''' Compute solid angle between the input points and triangles
+        Follows the method described in:
+        The Solid Angle of a Plane Triangle
+        A. VAN OOSTEROM AND J. STRACKEE
+        IEEE TRANSACTIONS ON BIOMEDICAL ENGINEERING,
+        VOL. BME-30, NO. 2, FEBRUARY 1983
+        Parameters
+        -----------
+            points: BxQx3
+                Tensor of input query points
+            triangles: BxFx3x3
+                Target triangles
+            thresh: float
+                float threshold
+        Returns
+        -------
+            solid_angles: BxQxF
+                A tensor containing the solid angle between all query points
+                and input triangles
+    '''
+    # Center the triangles on the query points. Size should be BxQxFx3x3
+    centered_tris = triangles[:, None] - points[:, :, None, None]
+
+    # BxQxFx3
+    norms = torch.norm(centered_tris, dim=-1)
+
+    # Should be BxQxFx3
+    cross_prod = torch.cross(centered_tris[:, :, :, 1],
+                             centered_tris[:, :, :, 2],
+                             dim=-1)
+    # Should be BxQxF
+    numerator = (centered_tris[:, :, :, 0] * cross_prod).sum(dim=-1)
+    del cross_prod
+
+    dot01 = (centered_tris[:, :, :, 0] * centered_tris[:, :, :, 1]).sum(dim=-1)
+    dot12 = (centered_tris[:, :, :, 1] * centered_tris[:, :, :, 2]).sum(dim=-1)
+    dot02 = (centered_tris[:, :, :, 0] * centered_tris[:, :, :, 2]).sum(dim=-1)
+    del centered_tris
+
+    denominator = (norms.prod(dim=-1) + dot01 * norms[:, :, :, 2] +
+                   dot02 * norms[:, :, :, 1] + dot12 * norms[:, :, :, 0])
+    del dot01, dot12, dot02, norms
+
+    # Should be BxQ
+    solid_angle = torch.atan2(numerator, denominator)
+    del numerator, denominator
+
+    torch.cuda.empty_cache()
+
+    return 2 * solid_angle
+
+
+def winding_numbers(points: Tensor,
+                    triangles: Tensor,
+                    thresh: float = 1e-8) -> Tensor:
+    ''' Uses winding_numbers to compute inside/outside
+        Robust inside-outside segmentation using generalized winding numbers
+        Alec Jacobson,
+        Ladislav Kavan,
+        Olga Sorkine-Hornung
+        Fast Winding Numbers for Soups and Clouds SIGGRAPH 2018
+        Gavin Barill
+        NEIL G. Dickson
+        Ryan Schmidt
+        David I.W. Levin
+        and Alec Jacobson
+        Parameters
+        -----------
+            points: BxQx3
+                Tensor of input query points
+            triangles: BxFx3x3
+                Target triangles
+            thresh: float
+                float threshold
+        Returns
+        -------
+            winding_numbers: BxQ
+                A tensor containing the Generalized winding numbers
+    '''
+    # The generalized winding number is the sum of solid angles of the point
+    # with respect to all triangles.
+    return 1 / (4 * math.pi) * solid_angles(points, triangles,
+                                            thresh=thresh).sum(dim=-1)
+
+
+def batch_contains(verts, faces, points):
+
+    B = verts.shape[0]
+    N = points.shape[1]
+
+    verts = verts.detach().cpu()
+    faces = faces.detach().cpu()
+    points = points.detach().cpu()
+    contains = torch.zeros(B, N)
+
+    for i in range(B):
+        contains[i] = torch.as_tensor(
+            trimesh.Trimesh(verts[i], faces[i]).contains(points[i]))
+
+    return 2.0 * (contains - 0.5)
+
+
+def dict2obj(d):
+    # if isinstance(d, list):
+    #     d = [dict2obj(x) for x in d]
+    if not isinstance(d, dict):
+        return d
+
+    class C(object):
+        pass
+
+    o = C()
+    for k in d:
+        o.__dict__[k] = dict2obj(d[k])
+    return o
+
+
+def face_vertices(vertices, faces):
+    """ 
+    :param vertices: [batch size, number of vertices, 3]
+    :param faces: [batch size, number of faces, 3]
+    :return: [batch size, number of faces, 3, 3]
+    """
+
+    bs, nv = vertices.shape[:2]
+    bs, nf = faces.shape[:2]
+    device = vertices.device
+    faces = faces + (torch.arange(bs, dtype=torch.int32).to(device) *
+                     nv)[:, None, None]
+    vertices = vertices.reshape((bs * nv, vertices.shape[-1]))
+
+    return vertices[faces.long()]
+
+
+class Pytorch3dRasterizer(nn.Module):
+    """  Borrowed from https://github.com/facebookresearch/pytorch3d
+    Notice:
+        x,y,z are in image space, normalized
+        can only render squared image now
+    """
+
+    def __init__(self, image_size=224):
+        """
+        use fixed raster_settings for rendering faces
+        """
+        super().__init__()
+        raster_settings = {
+            'image_size': image_size,
+            'blur_radius': 0.0,
+            'faces_per_pixel': 1,
+            'bin_size': None,
+            'max_faces_per_bin': None,
+            'perspective_correct': True,
+            'cull_backfaces': True,
+        }
+        raster_settings = dict2obj(raster_settings)
+        self.raster_settings = raster_settings
+
+    def forward(self, vertices, faces, attributes=None):
+        fixed_vertices = vertices.clone()
+        fixed_vertices[..., :2] = -fixed_vertices[..., :2]
+        meshes_screen = Meshes(verts=fixed_vertices.float(),
+                               faces=faces.long())
+        raster_settings = self.raster_settings
+        pix_to_face, zbuf, bary_coords, dists = rasterize_meshes(
+            meshes_screen,
+            image_size=raster_settings.image_size,
+            blur_radius=raster_settings.blur_radius,
+            faces_per_pixel=raster_settings.faces_per_pixel,
+            bin_size=raster_settings.bin_size,
+            max_faces_per_bin=raster_settings.max_faces_per_bin,
+            perspective_correct=raster_settings.perspective_correct,
+        )
+        vismask = (pix_to_face > -1).float()
+        D = attributes.shape[-1]
+        attributes = attributes.clone()
+        attributes = attributes.view(attributes.shape[0] * attributes.shape[1],
+                                     3, attributes.shape[-1])
+        N, H, W, K, _ = bary_coords.shape
+        mask = pix_to_face == -1
+        pix_to_face = pix_to_face.clone()
+        pix_to_face[mask] = 0
+        idx = pix_to_face.view(N * H * W * K, 1, 1).expand(N * H * W * K, 3, D)
+        pixel_face_vals = attributes.gather(0, idx).view(N, H, W, K, 3, D)
+        pixel_vals = (bary_coords[..., None] * pixel_face_vals).sum(dim=-2)
+        pixel_vals[mask] = 0  # Replace masked values in output.
+        pixel_vals = pixel_vals[:, :, :, 0].permute(0, 3, 1, 2)
+        pixel_vals = torch.cat(
+            [pixel_vals, vismask[:, :, :, 0][:, None, :, :]], dim=1)
+        return pixel_vals

lib/common/seg3d_lossless.py

@@ -0,0 +1,604 @@
+
+# -*- 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 .seg3d_utils import (
+    create_grid3D,
+    plot_mask3D,
+    SmoothConv3D,
+)
+
+import torch
+import torch.nn as nn
+import numpy as np
+import torch.nn.functional as F
+import mcubes
+from kaolin.ops.conversions import voxelgrids_to_trianglemeshes
+import logging
+
+logging.getLogger("lightning").setLevel(logging.ERROR)
+
+
+class Seg3dLossless(nn.Module):
+    def __init__(self,
+                 query_func,
+                 b_min,
+                 b_max,
+                 resolutions,
+                 channels=1,
+                 balance_value=0.5,
+                 align_corners=False,
+                 visualize=False,
+                 debug=False,
+                 use_cuda_impl=False,
+                 faster=False,
+                 use_shadow=False,
+                 **kwargs):
+        """
+        align_corners: same with how you process gt. (grid_sample / interpolate) 
+        """
+        super().__init__()
+        self.query_func = query_func
+        self.register_buffer(
+            'b_min',
+            torch.tensor(b_min).float().unsqueeze(1))  # [bz, 1, 3]
+        self.register_buffer(
+            'b_max',
+            torch.tensor(b_max).float().unsqueeze(1))  # [bz, 1, 3]
+
+        # ti.init(arch=ti.cuda)
+        # self.mciso_taichi = MCISO(dim=3, N=resolutions[-1]-1)
+
+        if type(resolutions[0]) is int:
+            resolutions = torch.tensor([(res, res, res)
+                                        for res in resolutions])
+        else:
+            resolutions = torch.tensor(resolutions)
+        self.register_buffer('resolutions', resolutions)
+        self.batchsize = self.b_min.size(0)
+        assert self.batchsize == 1
+        self.balance_value = balance_value
+        self.channels = channels
+        assert self.channels == 1
+        self.align_corners = align_corners
+        self.visualize = visualize
+        self.debug = debug
+        self.use_cuda_impl = use_cuda_impl
+        self.faster = faster
+        self.use_shadow = use_shadow
+
+        for resolution in resolutions:
+            assert resolution[0] % 2 == 1 and resolution[1] % 2 == 1, \
+                f"resolution {resolution} need to be odd becuase of align_corner."
+
+        # init first resolution
+        init_coords = create_grid3D(0,
+                                    resolutions[-1] - 1,
+                                    steps=resolutions[0])  # [N, 3]
+        init_coords = init_coords.unsqueeze(0).repeat(self.batchsize, 1,
+                                                      1)  # [bz, N, 3]
+        self.register_buffer('init_coords', init_coords)
+
+        # some useful tensors
+        calculated = torch.zeros(
+            (self.resolutions[-1][2], self.resolutions[-1][1],
+             self.resolutions[-1][0]),
+            dtype=torch.bool)
+        self.register_buffer('calculated', calculated)
+
+        gird8_offsets = torch.stack(
+            torch.meshgrid([
+                torch.tensor([-1, 0, 1]),
+                torch.tensor([-1, 0, 1]),
+                torch.tensor([-1, 0, 1])
+            ])).int().view(3, -1).t()  # [27, 3]
+        self.register_buffer('gird8_offsets', gird8_offsets)
+
+        # smooth convs
+        self.smooth_conv3x3 = SmoothConv3D(in_channels=1,
+                                           out_channels=1,
+                                           kernel_size=3)
+        self.smooth_conv5x5 = SmoothConv3D(in_channels=1,
+                                           out_channels=1,
+                                           kernel_size=5)
+        self.smooth_conv7x7 = SmoothConv3D(in_channels=1,
+                                           out_channels=1,
+                                           kernel_size=7)
+        self.smooth_conv9x9 = SmoothConv3D(in_channels=1,
+                                           out_channels=1,
+                                           kernel_size=9)
+
+    def batch_eval(self, coords, **kwargs):
+        """
+        coords: in the coordinates of last resolution
+        **kwargs: for query_func
+        """
+        coords = coords.detach()
+        # normalize coords to fit in [b_min, b_max]
+        if self.align_corners:
+            coords2D = coords.float() / (self.resolutions[-1] - 1)
+        else:
+            step = 1.0 / self.resolutions[-1].float()
+            coords2D = coords.float() / self.resolutions[-1] + step / 2
+        coords2D = coords2D * (self.b_max - self.b_min) + self.b_min
+        # query function
+        occupancys = self.query_func(**kwargs, points=coords2D)
+        if type(occupancys) is list:
+            occupancys = torch.stack(occupancys)  # [bz, C, N]
+        assert len(occupancys.size()) == 3, \
+            "query_func should return a occupancy with shape of [bz, C, N]"
+        return occupancys
+
+    def forward(self, **kwargs):
+        if self.faster:
+            return self._forward_faster(**kwargs)
+        else:
+            return self._forward(**kwargs)
+
+    def _forward_faster(self, **kwargs):
+        """
+        In faster mode, we make following changes to exchange accuracy for speed:
+        1. no conflict checking: 4.88 fps -> 6.56 fps
+        2. smooth_conv9x9 ~ smooth_conv3x3 for different resolution
+        3. last step no examine
+        """
+        final_W = self.resolutions[-1][0]
+        final_H = self.resolutions[-1][1]
+        final_D = self.resolutions[-1][2]
+
+        for resolution in self.resolutions:
+            W, H, D = resolution
+            stride = (self.resolutions[-1] - 1) / (resolution - 1)
+
+            # first step
+            if torch.equal(resolution, self.resolutions[0]):
+                coords = self.init_coords.clone()  # torch.long
+                occupancys = self.batch_eval(coords, **kwargs)
+                occupancys = occupancys.view(self.batchsize, self.channels, D,
+                                             H, W)
+                if (occupancys > 0.5).sum() == 0:
+                    # return F.interpolate(
+                    #     occupancys, size=(final_D, final_H, final_W),
+                    #     mode="linear", align_corners=True)
+                    return None
+
+                if self.visualize:
+                    self.plot(occupancys, coords, final_D, final_H, final_W)
+
+                with torch.no_grad():
+                    coords_accum = coords / stride
+
+            # last step
+            elif torch.equal(resolution, self.resolutions[-1]):
+
+                with torch.no_grad():
+                    # here true is correct!
+                    valid = F.interpolate(
+                        (occupancys > self.balance_value).float(),
+                        size=(D, H, W),
+                        mode="trilinear",
+                        align_corners=True)
+
+                # here true is correct!
+                occupancys = F.interpolate(occupancys.float(),
+                                           size=(D, H, W),
+                                           mode="trilinear",
+                                           align_corners=True)
+
+                # is_boundary = (valid > 0.0) & (valid < 1.0)
+                is_boundary = valid == 0.5
+
+            # next steps
+            else:
+                coords_accum *= 2
+
+                with torch.no_grad():
+                    # here true is correct!
+                    valid = F.interpolate(
+                        (occupancys > self.balance_value).float(),
+                        size=(D, H, W),
+                        mode="trilinear",
+                        align_corners=True)
+
+                # here true is correct!
+                occupancys = F.interpolate(occupancys.float(),
+                                           size=(D, H, W),
+                                           mode="trilinear",
+                                           align_corners=True)
+
+                is_boundary = (valid > 0.0) & (valid < 1.0)
+
+                with torch.no_grad():
+                    if torch.equal(resolution, self.resolutions[1]):
+                        is_boundary = (self.smooth_conv9x9(is_boundary.float())
+                                       > 0)[0, 0]
+                    elif torch.equal(resolution, self.resolutions[2]):
+                        is_boundary = (self.smooth_conv7x7(is_boundary.float())
+                                       > 0)[0, 0]
+                    else:
+                        is_boundary = (self.smooth_conv3x3(is_boundary.float())
+                                       > 0)[0, 0]
+
+                    coords_accum = coords_accum.long()
+                    is_boundary[coords_accum[0, :, 2], coords_accum[0, :, 1],
+                                coords_accum[0, :, 0]] = False
+                    point_coords = is_boundary.permute(
+                        2, 1, 0).nonzero(as_tuple=False).unsqueeze(0)
+                    point_indices = (point_coords[:, :, 2] * H * W +
+                                     point_coords[:, :, 1] * W +
+                                     point_coords[:, :, 0])
+
+                    R, C, D, H, W = occupancys.shape
+
+                    # inferred value
+                    coords = point_coords * stride
+
+                if coords.size(1) == 0:
+                    continue
+                occupancys_topk = self.batch_eval(coords, **kwargs)
+
+                # put mask point predictions to the right places on the upsampled grid.
+                R, C, D, H, W = occupancys.shape
+                point_indices = point_indices.unsqueeze(1).expand(-1, C, -1)
+                occupancys = (occupancys.reshape(R, C, D * H * W).scatter_(
+                    2, point_indices, occupancys_topk).view(R, C, D, H, W))
+
+                with torch.no_grad():
+                    voxels = coords / stride
+                    coords_accum = torch.cat([voxels, coords_accum],
+                                             dim=1).unique(dim=1)
+
+        return occupancys[0, 0]
+
+    def _forward(self, **kwargs):
+        """
+        output occupancy field would be:
+        (bz, C, res, res)
+        """
+        final_W = self.resolutions[-1][0]
+        final_H = self.resolutions[-1][1]
+        final_D = self.resolutions[-1][2]
+
+        calculated = self.calculated.clone()
+
+        for resolution in self.resolutions:
+            W, H, D = resolution
+            stride = (self.resolutions[-1] - 1) / (resolution - 1)
+
+            if self.visualize:
+                this_stage_coords = []
+
+            # first step
+            if torch.equal(resolution, self.resolutions[0]):
+                coords = self.init_coords.clone()  # torch.long
+                occupancys = self.batch_eval(coords, **kwargs)
+                occupancys = occupancys.view(self.batchsize, self.channels, D,
+                                             H, W)
+
+                if self.visualize:
+                    self.plot(occupancys, coords, final_D, final_H, final_W)
+
+                with torch.no_grad():
+                    coords_accum = coords / stride
+                    calculated[coords[0, :, 2], coords[0, :, 1],
+                               coords[0, :, 0]] = True
+
+            # next steps
+            else:
+                coords_accum *= 2
+
+                with torch.no_grad():
+                    # here true is correct!
+                    valid = F.interpolate(
+                        (occupancys > self.balance_value).float(),
+                        size=(D, H, W),
+                        mode="trilinear",
+                        align_corners=True)
+
+                # here true is correct!
+                occupancys = F.interpolate(occupancys.float(),
+                                           size=(D, H, W),
+                                           mode="trilinear",
+                                           align_corners=True)
+
+                is_boundary = (valid > 0.0) & (valid < 1.0)
+
+                with torch.no_grad():
+                    # TODO
+                    if self.use_shadow and torch.equal(resolution,
+                                                       self.resolutions[-1]):
+                        # larger z means smaller depth here
+                        depth_res = resolution[2].item()
+                        depth_index = torch.linspace(0,
+                                                     depth_res - 1,
+                                                     steps=depth_res).type_as(
+                                                         occupancys.device)
+                        depth_index_max = torch.max(
+                            (occupancys > self.balance_value) *
+                            (depth_index + 1),
+                            dim=-1,
+                            keepdim=True)[0] - 1
+                        shadow = depth_index < depth_index_max
+                        is_boundary[shadow] = False
+                        is_boundary = is_boundary[0, 0]
+                    else:
+                        is_boundary = (self.smooth_conv3x3(is_boundary.float())
+                                       > 0)[0, 0]
+                        # is_boundary = is_boundary[0, 0]
+
+                    is_boundary[coords_accum[0, :, 2], coords_accum[0, :, 1],
+                                coords_accum[0, :, 0]] = False
+                    point_coords = is_boundary.permute(
+                        2, 1, 0).nonzero(as_tuple=False).unsqueeze(0)
+                    point_indices = (point_coords[:, :, 2] * H * W +
+                                     point_coords[:, :, 1] * W +
+                                     point_coords[:, :, 0])
+
+                    R, C, D, H, W = occupancys.shape
+                    # interpolated value
+                    occupancys_interp = torch.gather(
+                        occupancys.reshape(R, C, D * H * W), 2,
+                        point_indices.unsqueeze(1))
+
+                    # inferred value
+                    coords = point_coords * stride
+
+                if coords.size(1) == 0:
+                    continue
+                occupancys_topk = self.batch_eval(coords, **kwargs)
+                if self.visualize:
+                    this_stage_coords.append(coords)
+
+                # put mask point predictions to the right places on the upsampled grid.
+                R, C, D, H, W = occupancys.shape
+                point_indices = point_indices.unsqueeze(1).expand(-1, C, -1)
+                occupancys = (occupancys.reshape(R, C, D * H * W).scatter_(
+                    2, point_indices, occupancys_topk).view(R, C, D, H, W))
+
+                with torch.no_grad():
+                    # conflicts
+                    conflicts = ((occupancys_interp - self.balance_value) *
+                                 (occupancys_topk - self.balance_value) < 0)[0,
+                                                                             0]
+
+                    if self.visualize:
+                        self.plot(occupancys, coords, final_D, final_H,
+                                  final_W)
+
+                    voxels = coords / stride
+                    coords_accum = torch.cat([voxels, coords_accum],
+                                             dim=1).unique(dim=1)
+                    calculated[coords[0, :, 2], coords[0, :, 1],
+                               coords[0, :, 0]] = True
+
+                while conflicts.sum() > 0:
+                    if self.use_shadow and torch.equal(resolution,
+                                                       self.resolutions[-1]):
+                        break
+
+                    with torch.no_grad():
+                        conflicts_coords = coords[0, conflicts, :]
+
+                        if self.debug:
+                            self.plot(occupancys,
+                                      conflicts_coords.unsqueeze(0),
+                                      final_D,
+                                      final_H,
+                                      final_W,
+                                      title='conflicts')
+
+                        conflicts_boundary = (conflicts_coords.int() +
+                                              self.gird8_offsets.unsqueeze(1) *
+                                              stride.int()).reshape(
+                                                  -1, 3).long().unique(dim=0)
+                        conflicts_boundary[:, 0] = (
+                            conflicts_boundary[:, 0].clamp(
+                                0,
+                                calculated.size(2) - 1))
+                        conflicts_boundary[:, 1] = (
+                            conflicts_boundary[:, 1].clamp(
+                                0,
+                                calculated.size(1) - 1))
+                        conflicts_boundary[:, 2] = (
+                            conflicts_boundary[:, 2].clamp(
+                                0,
+                                calculated.size(0) - 1))
+
+                        coords = conflicts_boundary[calculated[
+                            conflicts_boundary[:, 2], conflicts_boundary[:, 1],
+                            conflicts_boundary[:, 0]] == False]
+
+                        if self.debug:
+                            self.plot(occupancys,
+                                      coords.unsqueeze(0),
+                                      final_D,
+                                      final_H,
+                                      final_W,
+                                      title='coords')
+
+                        coords = coords.unsqueeze(0)
+                        point_coords = coords / stride
+                        point_indices = (point_coords[:, :, 2] * H * W +
+                                         point_coords[:, :, 1] * W +
+                                         point_coords[:, :, 0])
+
+                        R, C, D, H, W = occupancys.shape
+                        # interpolated value
+                        occupancys_interp = torch.gather(
+                            occupancys.reshape(R, C, D * H * W), 2,
+                            point_indices.unsqueeze(1))
+
+                        # inferred value
+                        coords = point_coords * stride
+
+                    if coords.size(1) == 0:
+                        break
+                    occupancys_topk = self.batch_eval(coords, **kwargs)
+                    if self.visualize:
+                        this_stage_coords.append(coords)
+
+                    with torch.no_grad():
+                        # conflicts
+                        conflicts = ((occupancys_interp - self.balance_value) *
+                                     (occupancys_topk - self.balance_value) <
+                                     0)[0, 0]
+
+                    # put mask point predictions to the right places on the upsampled grid.
+                    point_indices = point_indices.unsqueeze(1).expand(
+                        -1, C, -1)
+                    occupancys = (occupancys.reshape(R, C, D * H * W).scatter_(
+                        2, point_indices, occupancys_topk).view(R, C, D, H, W))
+
+                    with torch.no_grad():
+                        voxels = coords / stride
+                        coords_accum = torch.cat([voxels, coords_accum],
+                                                 dim=1).unique(dim=1)
+                        calculated[coords[0, :, 2], coords[0, :, 1],
+                                   coords[0, :, 0]] = True
+
+                if self.visualize:
+                    this_stage_coords = torch.cat(this_stage_coords, dim=1)
+                    self.plot(occupancys, this_stage_coords, final_D, final_H,
+                              final_W)
+
+        return occupancys[0, 0]
+
+    def plot(self,
+             occupancys,
+             coords,
+             final_D,
+             final_H,
+             final_W,
+             title='',
+             **kwargs):
+        final = F.interpolate(occupancys.float(),
+                              size=(final_D, final_H, final_W),
+                              mode="trilinear",
+                              align_corners=True)  # here true is correct!
+        x = coords[0, :, 0].to("cpu")
+        y = coords[0, :, 1].to("cpu")
+        z = coords[0, :, 2].to("cpu")
+
+        plot_mask3D(final[0, 0].to("cpu"), title, (x, y, z), **kwargs)
+
+    def find_vertices(self, sdf, direction="front"):
+        '''
+            - direction: "front" | "back" | "left" | "right"
+        '''
+        resolution = sdf.size(2)
+        if direction == "front":
+            pass
+        elif direction == "left":
+            sdf = sdf.permute(2, 1, 0)
+        elif direction == "back":
+            inv_idx = torch.arange(sdf.size(2) - 1, -1, -1).long()
+            sdf = sdf[inv_idx, :, :]
+        elif direction == "right":
+            inv_idx = torch.arange(sdf.size(2) - 1, -1, -1).long()
+            sdf = sdf[:, :, inv_idx]
+            sdf = sdf.permute(2, 1, 0)
+
+        inv_idx = torch.arange(sdf.size(2) - 1, -1, -1).long()
+        sdf = sdf[inv_idx, :, :]
+        sdf_all = sdf.permute(2, 1, 0)
+
+        # shadow
+        grad_v = (sdf_all > 0.5) * torch.linspace(
+            resolution, 1, steps=resolution).to(sdf.device)
+        grad_c = torch.ones_like(sdf_all) * torch.linspace(
+            0, resolution - 1, steps=resolution).to(sdf.device)
+        max_v, max_c = grad_v.max(dim=2)
+        shadow = grad_c > max_c.view(resolution, resolution, 1)
+        keep = (sdf_all > 0.5) & (~shadow)
+
+        p1 = keep.nonzero(as_tuple=False).t()  # [3, N]
+        p2 = p1.clone()  # z
+        p2[2, :] = (p2[2, :] - 2).clamp(0, resolution)
+        p3 = p1.clone()  # y
+        p3[1, :] = (p3[1, :] - 2).clamp(0, resolution)
+        p4 = p1.clone()  # x
+        p4[0, :] = (p4[0, :] - 2).clamp(0, resolution)
+
+        v1 = sdf_all[p1[0, :], p1[1, :], p1[2, :]]
+        v2 = sdf_all[p2[0, :], p2[1, :], p2[2, :]]
+        v3 = sdf_all[p3[0, :], p3[1, :], p3[2, :]]
+        v4 = sdf_all[p4[0, :], p4[1, :], p4[2, :]]
+
+        X = p1[0, :].long()  # [N,]
+        Y = p1[1, :].long()  # [N,]
+        Z = p2[2, :].float() * (0.5 - v1) / (v2 - v1) + \
+            p1[2, :].float() * (v2 - 0.5) / (v2 - v1)  # [N,]
+        Z = Z.clamp(0, resolution)
+
+        # normal
+        norm_z = v2 - v1
+        norm_y = v3 - v1
+        norm_x = v4 - v1
+        # print (v2.min(dim=0)[0], v2.max(dim=0)[0], v3.min(dim=0)[0], v3.max(dim=0)[0])
+
+        norm = torch.stack([norm_x, norm_y, norm_z], dim=1)
+        norm = norm / torch.norm(norm, p=2, dim=1, keepdim=True)
+
+        return X, Y, Z, norm
+
+    def render_normal(self, resolution, X, Y, Z, norm):
+        image = torch.ones((1, 3, resolution, resolution),
+                           dtype=torch.float32).to(norm.device)
+        color = (norm + 1) / 2.0
+        color = color.clamp(0, 1)
+        image[0, :, Y, X] = color.t()
+        return image
+
+    def display(self, sdf):
+
+        # render
+        X, Y, Z, norm = self.find_vertices(sdf, direction="front")
+        image1 = self.render_normal(self.resolutions[-1, -1], X, Y, Z, norm)
+        X, Y, Z, norm = self.find_vertices(sdf, direction="left")
+        image2 = self.render_normal(self.resolutions[-1, -1], X, Y, Z, norm)
+        X, Y, Z, norm = self.find_vertices(sdf, direction="right")
+        image3 = self.render_normal(self.resolutions[-1, -1], X, Y, Z, norm)
+        X, Y, Z, norm = self.find_vertices(sdf, direction="back")
+        image4 = self.render_normal(self.resolutions[-1, -1], X, Y, Z, norm)
+
+        image = torch.cat([image1, image2, image3, image4], axis=3)
+        image = image.detach().cpu().numpy()[0].transpose(1, 2, 0) * 255.0
+
+        return np.uint8(image)
+
+    def export_mesh(self, occupancys):
+
+        final = occupancys[1:, 1:, 1:].contiguous()
+
+        if final.shape[0] > 256:
+            # for voxelgrid larger than 256^3, the required GPU memory will be > 9GB
+            # thus we use CPU marching_cube to avoid "CUDA out of memory"
+            occu_arr = final.detach().cpu().numpy()                 # non-smooth surface
+            # occu_arr = mcubes.smooth(final.detach().cpu().numpy())  # smooth surface
+            vertices, triangles = mcubes.marching_cubes(
+                occu_arr, self.balance_value)
+            verts = torch.as_tensor(vertices[:, [2, 1, 0]])
+            faces = torch.as_tensor(triangles.astype(
+                np.long), dtype=torch.long)[:, [0, 2, 1]]
+        else:
+            torch.cuda.empty_cache()
+            vertices, triangles = voxelgrids_to_trianglemeshes(
+                final.unsqueeze(0))
+            verts = vertices[0][:, [2, 1, 0]].cpu()
+            faces = triangles[0][:, [0, 2, 1]].cpu()
+
+        return verts, faces

lib/common/seg3d_utils.py

@@ -0,0 +1,392 @@
+
+# -*- 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 torch
+import torch.nn as nn
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
+
+
+def plot_mask2D(mask,
+                title="",
+                point_coords=None,
+                figsize=10,
+                point_marker_size=5):
+    '''
+    Simple plotting tool to show intermediate mask predictions and points 
+    where PointRend is applied.
+
+    Args:
+    mask (Tensor): mask prediction of shape HxW
+    title (str): title for the plot
+    point_coords ((Tensor, Tensor)): x and y point coordinates
+    figsize (int): size of the figure to plot
+    point_marker_size (int): marker size for points
+    '''
+
+    H, W = mask.shape
+    plt.figure(figsize=(figsize, figsize))
+    if title:
+        title += ", "
+    plt.title("{}resolution {}x{}".format(title, H, W), fontsize=30)
+    plt.ylabel(H, fontsize=30)
+    plt.xlabel(W, fontsize=30)
+    plt.xticks([], [])
+    plt.yticks([], [])
+    plt.imshow(mask.detach(),
+               interpolation="nearest",
+               cmap=plt.get_cmap('gray'))
+    if point_coords is not None:
+        plt.scatter(x=point_coords[0],
+                    y=point_coords[1],
+                    color="red",
+                    s=point_marker_size,
+                    clip_on=True)
+    plt.xlim(-0.5, W - 0.5)
+    plt.ylim(H - 0.5, -0.5)
+    plt.show()
+
+
+def plot_mask3D(mask=None,
+                title="",
+                point_coords=None,
+                figsize=1500,
+                point_marker_size=8,
+                interactive=True):
+    '''
+    Simple plotting tool to show intermediate mask predictions and points 
+    where PointRend is applied.
+
+    Args:
+    mask (Tensor): mask prediction of shape DxHxW
+    title (str): title for the plot
+    point_coords ((Tensor, Tensor, Tensor)): x and y and z point coordinates
+    figsize (int): size of the figure to plot
+    point_marker_size (int): marker size for points
+    '''
+    import trimesh
+    import vtkplotter
+    from skimage import measure
+
+    vp = vtkplotter.Plotter(title=title, size=(figsize, figsize))
+    vis_list = []
+
+    if mask is not None:
+        mask = mask.detach().to("cpu").numpy()
+        mask = mask.transpose(2, 1, 0)
+
+        # marching cube to find surface
+        verts, faces, normals, values = measure.marching_cubes_lewiner(
+            mask, 0.5, gradient_direction='ascent')
+
+        # create a mesh
+        mesh = trimesh.Trimesh(verts, faces)
+        mesh.visual.face_colors = [200, 200, 250, 100]
+        vis_list.append(mesh)
+
+    if point_coords is not None:
+        point_coords = torch.stack(point_coords, 1).to("cpu").numpy()
+
+        # import numpy as np
+        # select_x = np.logical_and(point_coords[:, 0] >= 16, point_coords[:, 0] <= 112)
+        # select_y = np.logical_and(point_coords[:, 1] >= 48, point_coords[:, 1] <= 272)
+        # select_z = np.logical_and(point_coords[:, 2] >= 16, point_coords[:, 2] <= 112)
+        # select = np.logical_and(np.logical_and(select_x, select_y), select_z)
+        # point_coords = point_coords[select, :]
+
+        pc = vtkplotter.Points(point_coords, r=point_marker_size, c='red')
+        vis_list.append(pc)
+
+    vp.show(*vis_list,
+            bg="white",
+            axes=1,
+            interactive=interactive,
+            azimuth=30,
+            elevation=30)
+
+
+def create_grid3D(min, max, steps):
+    if type(min) is int:
+        min = (min, min, min)  # (x, y, z)
+    if type(max) is int:
+        max = (max, max, max)  # (x, y)
+    if type(steps) is int:
+        steps = (steps, steps, steps)  # (x, y, z)
+    arrangeX = torch.linspace(min[0], max[0], steps[0]).long()
+    arrangeY = torch.linspace(min[1], max[1], steps[1]).long()
+    arrangeZ = torch.linspace(min[2], max[2], steps[2]).long()
+    gridD, girdH, gridW = torch.meshgrid([arrangeZ, arrangeY, arrangeX])
+    coords = torch.stack([gridW, girdH,
+                          gridD])  # [2, steps[0], steps[1], steps[2]]
+    coords = coords.view(3, -1).t()  # [N, 3]
+    return coords
+
+
+def create_grid2D(min, max, steps):
+    if type(min) is int:
+        min = (min, min)  # (x, y)
+    if type(max) is int:
+        max = (max, max)  # (x, y)
+    if type(steps) is int:
+        steps = (steps, steps)  # (x, y)
+    arrangeX = torch.linspace(min[0], max[0], steps[0]).long()
+    arrangeY = torch.linspace(min[1], max[1], steps[1]).long()
+    girdH, gridW = torch.meshgrid([arrangeY, arrangeX])
+    coords = torch.stack([gridW, girdH])  # [2, steps[0], steps[1]]
+    coords = coords.view(2, -1).t()  # [N, 2]
+    return coords
+
+
+class SmoothConv2D(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=3):
+        super().__init__()
+        assert kernel_size % 2 == 1, "kernel_size for smooth_conv must be odd: {3, 5, ...}"
+        self.padding = (kernel_size - 1) // 2
+
+        weight = torch.ones(
+            (in_channels, out_channels, kernel_size, kernel_size),
+            dtype=torch.float32) / (kernel_size**2)
+        self.register_buffer('weight', weight)
+
+    def forward(self, input):
+        return F.conv2d(input, self.weight, padding=self.padding)
+
+
+class SmoothConv3D(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=3):
+        super().__init__()
+        assert kernel_size % 2 == 1, "kernel_size for smooth_conv must be odd: {3, 5, ...}"
+        self.padding = (kernel_size - 1) // 2
+
+        weight = torch.ones(
+            (in_channels, out_channels, kernel_size, kernel_size, kernel_size),
+            dtype=torch.float32) / (kernel_size**3)
+        self.register_buffer('weight', weight)
+
+    def forward(self, input):
+        return F.conv3d(input, self.weight, padding=self.padding)
+
+
+def build_smooth_conv3D(in_channels=1,
+                        out_channels=1,
+                        kernel_size=3,
+                        padding=1):
+    smooth_conv = torch.nn.Conv3d(in_channels=in_channels,
+                                  out_channels=out_channels,
+                                  kernel_size=kernel_size,
+                                  padding=padding)
+    smooth_conv.weight.data = torch.ones(
+        (in_channels, out_channels, kernel_size, kernel_size, kernel_size),
+        dtype=torch.float32) / (kernel_size**3)
+    smooth_conv.bias.data = torch.zeros(out_channels)
+    return smooth_conv
+
+
+def build_smooth_conv2D(in_channels=1,
+                        out_channels=1,
+                        kernel_size=3,
+                        padding=1):
+    smooth_conv = torch.nn.Conv2d(in_channels=in_channels,
+                                  out_channels=out_channels,
+                                  kernel_size=kernel_size,
+                                  padding=padding)
+    smooth_conv.weight.data = torch.ones(
+        (in_channels, out_channels, kernel_size, kernel_size),
+        dtype=torch.float32) / (kernel_size**2)
+    smooth_conv.bias.data = torch.zeros(out_channels)
+    return smooth_conv
+
+
+def get_uncertain_point_coords_on_grid3D(uncertainty_map, num_points,
+                                         **kwargs):
+    """
+    Find `num_points` most uncertain points from `uncertainty_map` grid.
+    Args:
+        uncertainty_map (Tensor): A tensor of shape (N, 1, H, W, D) that contains uncertainty
+            values for a set of points on a regular H x W x D grid.
+        num_points (int): The number of points P to select.
+    Returns:
+        point_indices (Tensor): A tensor of shape (N, P) that contains indices from
+            [0, H x W x D) of the most uncertain points.
+        point_coords (Tensor): A tensor of shape (N, P, 3) that contains [0, 1] x [0, 1] normalized
+            coordinates of the most uncertain points from the H x W x D grid.
+    """
+    R, _, D, H, W = uncertainty_map.shape
+    # h_step = 1.0 / float(H)
+    # w_step = 1.0 / float(W)
+    # d_step = 1.0 / float(D)
+
+    num_points = min(D * H * W, num_points)
+    point_scores, point_indices = torch.topk(uncertainty_map.view(
+        R, D * H * W),
+        k=num_points,
+        dim=1)
+    point_coords = torch.zeros(R,
+                               num_points,
+                               3,
+                               dtype=torch.float,
+                               device=uncertainty_map.device)
+    # point_coords[:, :, 0] = h_step / 2.0 + (point_indices // (W * D)).to(torch.float) * h_step
+    # point_coords[:, :, 1] = w_step / 2.0 + (point_indices % (W * D) // D).to(torch.float) * w_step
+    # point_coords[:, :, 2] = d_step / 2.0 + (point_indices % D).to(torch.float) * d_step
+    point_coords[:, :, 0] = (point_indices % W).to(torch.float)  # x
+    point_coords[:, :, 1] = (point_indices % (H * W) // W).to(torch.float)  # y
+    point_coords[:, :, 2] = (point_indices // (H * W)).to(torch.float)  # z
+    print(f"resolution {D} x {H} x {W}", point_scores.min(),
+          point_scores.max())
+    return point_indices, point_coords
+
+
+def get_uncertain_point_coords_on_grid3D_faster(uncertainty_map, num_points,
+                                                clip_min):
+    """
+    Find `num_points` most uncertain points from `uncertainty_map` grid.
+    Args:
+        uncertainty_map (Tensor): A tensor of shape (N, 1, H, W, D) that contains uncertainty
+            values for a set of points on a regular H x W x D grid.
+        num_points (int): The number of points P to select.
+    Returns:
+        point_indices (Tensor): A tensor of shape (N, P) that contains indices from
+            [0, H x W x D) of the most uncertain points.
+        point_coords (Tensor): A tensor of shape (N, P, 3) that contains [0, 1] x [0, 1] normalized
+            coordinates of the most uncertain points from the H x W x D grid.
+    """
+    R, _, D, H, W = uncertainty_map.shape
+    # h_step = 1.0 / float(H)
+    # w_step = 1.0 / float(W)
+    # d_step = 1.0 / float(D)
+
+    assert R == 1, "batchsize > 1 is not implemented!"
+    uncertainty_map = uncertainty_map.view(D * H * W)
+    indices = (uncertainty_map >= clip_min).nonzero().squeeze(1)
+    num_points = min(num_points, indices.size(0))
+    point_scores, point_indices = torch.topk(uncertainty_map[indices],
+                                             k=num_points,
+                                             dim=0)
+    point_indices = indices[point_indices].unsqueeze(0)
+
+    point_coords = torch.zeros(R,
+                               num_points,
+                               3,
+                               dtype=torch.float,
+                               device=uncertainty_map.device)
+    # point_coords[:, :, 0] = h_step / 2.0 + (point_indices // (W * D)).to(torch.float) * h_step
+    # point_coords[:, :, 1] = w_step / 2.0 + (point_indices % (W * D) // D).to(torch.float) * w_step
+    # point_coords[:, :, 2] = d_step / 2.0 + (point_indices % D).to(torch.float) * d_step
+    point_coords[:, :, 0] = (point_indices % W).to(torch.float)  # x
+    point_coords[:, :, 1] = (point_indices % (H * W) // W).to(torch.float)  # y
+    point_coords[:, :, 2] = (point_indices // (H * W)).to(torch.float)  # z
+    # print (f"resolution {D} x {H} x {W}", point_scores.min(), point_scores.max())
+    return point_indices, point_coords
+
+
+def get_uncertain_point_coords_on_grid2D(uncertainty_map, num_points,
+                                         **kwargs):
+    """
+    Find `num_points` most uncertain points from `uncertainty_map` grid.
+    Args:
+        uncertainty_map (Tensor): A tensor of shape (N, 1, H, W) that contains uncertainty
+            values for a set of points on a regular H x W grid.
+        num_points (int): The number of points P to select.
+    Returns:
+        point_indices (Tensor): A tensor of shape (N, P) that contains indices from
+            [0, H x W) of the most uncertain points.
+        point_coords (Tensor): A tensor of shape (N, P, 2) that contains [0, 1] x [0, 1] normalized
+            coordinates of the most uncertain points from the H x W grid.
+    """
+    R, _, H, W = uncertainty_map.shape
+    # h_step = 1.0 / float(H)
+    # w_step = 1.0 / float(W)
+
+    num_points = min(H * W, num_points)
+    point_scores, point_indices = torch.topk(uncertainty_map.view(R, H * W),
+                                             k=num_points,
+                                             dim=1)
+    point_coords = torch.zeros(R,
+                               num_points,
+                               2,
+                               dtype=torch.long,
+                               device=uncertainty_map.device)
+    # point_coords[:, :, 0] = w_step / 2.0 + (point_indices % W).to(torch.float) * w_step
+    # point_coords[:, :, 1] = h_step / 2.0 + (point_indices // W).to(torch.float) * h_step
+    point_coords[:, :, 0] = (point_indices % W).to(torch.long)
+    point_coords[:, :, 1] = (point_indices // W).to(torch.long)
+    # print (point_scores.min(), point_scores.max())
+    return point_indices, point_coords
+
+
+def get_uncertain_point_coords_on_grid2D_faster(uncertainty_map, num_points,
+                                                clip_min):
+    """
+    Find `num_points` most uncertain points from `uncertainty_map` grid.
+    Args:
+        uncertainty_map (Tensor): A tensor of shape (N, 1, H, W) that contains uncertainty
+            values for a set of points on a regular H x W grid.
+        num_points (int): The number of points P to select.
+    Returns:
+        point_indices (Tensor): A tensor of shape (N, P) that contains indices from
+            [0, H x W) of the most uncertain points.
+        point_coords (Tensor): A tensor of shape (N, P, 2) that contains [0, 1] x [0, 1] normalized
+            coordinates of the most uncertain points from the H x W grid.
+    """
+    R, _, H, W = uncertainty_map.shape
+    # h_step = 1.0 / float(H)
+    # w_step = 1.0 / float(W)
+
+    assert R == 1, "batchsize > 1 is not implemented!"
+    uncertainty_map = uncertainty_map.view(H * W)
+    indices = (uncertainty_map >= clip_min).nonzero().squeeze(1)
+    num_points = min(num_points, indices.size(0))
+    point_scores, point_indices = torch.topk(uncertainty_map[indices],
+                                             k=num_points,
+                                             dim=0)
+    point_indices = indices[point_indices].unsqueeze(0)
+
+    point_coords = torch.zeros(R,
+                               num_points,
+                               2,
+                               dtype=torch.long,
+                               device=uncertainty_map.device)
+    # point_coords[:, :, 0] = w_step / 2.0 + (point_indices % W).to(torch.float) * w_step
+    # point_coords[:, :, 1] = h_step / 2.0 + (point_indices // W).to(torch.float) * h_step
+    point_coords[:, :, 0] = (point_indices % W).to(torch.long)
+    point_coords[:, :, 1] = (point_indices // W).to(torch.long)
+    # print (point_scores.min(), point_scores.max())
+    return point_indices, point_coords
+
+
+def calculate_uncertainty(logits, classes=None, balance_value=0.5):
+    """
+    We estimate uncerainty as L1 distance between 0.0 and the logit prediction in 'logits' for the
+        foreground class in `classes`.
+    Args:
+        logits (Tensor): A tensor of shape (R, C, ...) or (R, 1, ...) for class-specific or
+            class-agnostic, where R is the total number of predicted masks in all images and C is
+            the number of foreground classes. The values are logits.
+        classes (list): A list of length R that contains either predicted of ground truth class
+            for eash predicted mask.
+    Returns:
+        scores (Tensor): A tensor of shape (R, 1, ...) that contains uncertainty scores with
+            the most uncertain locations having the highest uncertainty score.
+    """
+    if logits.shape[1] == 1:
+        gt_class_logits = logits
+    else:
+        gt_class_logits = logits[
+            torch.arange(logits.shape[0], device=logits.device),
+            classes].unsqueeze(1)
+    return -torch.abs(gt_class_logits - balance_value)

lib/common/train_util.py

@@ -0,0 +1,597 @@
+
+# -*- 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 yaml
+import os.path as osp
+import torch
+import numpy as np
+import torch.nn.functional as F
+from ..dataset.mesh_util import *
+from ..net.geometry import orthogonal
+from pytorch3d.renderer.mesh import rasterize_meshes
+from .render_utils import Pytorch3dRasterizer
+from pytorch3d.structures import Meshes
+import cv2
+from PIL import Image
+from tqdm import tqdm
+import os
+from termcolor import colored
+
+
+def reshape_sample_tensor(sample_tensor, num_views):
+    if num_views == 1:
+        return sample_tensor
+    # Need to repeat sample_tensor along the batch dim num_views times
+    sample_tensor = sample_tensor.unsqueeze(dim=1)
+    sample_tensor = sample_tensor.repeat(1, num_views, 1, 1)
+    sample_tensor = sample_tensor.view(
+        sample_tensor.shape[0] * sample_tensor.shape[1],
+        sample_tensor.shape[2], sample_tensor.shape[3])
+    return sample_tensor
+
+
+def gen_mesh_eval(opt, net, cuda, data, resolution=None):
+    resolution = opt.resolution if resolution is None else resolution
+    image_tensor = data['img'].to(device=cuda)
+    calib_tensor = data['calib'].to(device=cuda)
+
+    net.filter(image_tensor)
+
+    b_min = data['b_min']
+    b_max = data['b_max']
+    try:
+        verts, faces, _, _ = reconstruction_faster(net,
+                                                   cuda,
+                                                   calib_tensor,
+                                                   resolution,
+                                                   b_min,
+                                                   b_max,
+                                                   use_octree=False)
+
+    except Exception as e:
+        print(e)
+        print('Can not create marching cubes at this time.')
+        verts, faces = None, None
+    return verts, faces
+
+
+def gen_mesh(opt, net, cuda, data, save_path, resolution=None):
+    resolution = opt.resolution if resolution is None else resolution
+    image_tensor = data['img'].to(device=cuda)
+    calib_tensor = data['calib'].to(device=cuda)
+
+    net.filter(image_tensor)
+
+    b_min = data['b_min']
+    b_max = data['b_max']
+    try:
+        save_img_path = save_path[:-4] + '.png'
+        save_img_list = []
+        for v in range(image_tensor.shape[0]):
+            save_img = (np.transpose(image_tensor[v].detach().cpu().numpy(),
+                                     (1, 2, 0)) * 0.5 +
+                        0.5)[:, :, ::-1] * 255.0
+            save_img_list.append(save_img)
+        save_img = np.concatenate(save_img_list, axis=1)
+        Image.fromarray(np.uint8(save_img[:, :, ::-1])).save(save_img_path)
+
+        verts, faces, _, _ = reconstruction_faster(net, cuda, calib_tensor,
+                                                   resolution, b_min, b_max)
+        verts_tensor = torch.from_numpy(
+            verts.T).unsqueeze(0).to(device=cuda).float()
+        xyz_tensor = net.projection(verts_tensor, calib_tensor[:1])
+        uv = xyz_tensor[:, :2, :]
+        color = netG.index(image_tensor[:1], uv).detach().cpu().numpy()[0].T
+        color = color * 0.5 + 0.5
+        save_obj_mesh_with_color(save_path, verts, faces, color)
+    except Exception as e:
+        print(e)
+        print('Can not create marching cubes at this time.')
+        verts, faces, color = None, None, None
+    return verts, faces, color
+
+
+def gen_mesh_color(opt, netG, netC, cuda, data, save_path, use_octree=True):
+    image_tensor = data['img'].to(device=cuda)
+    calib_tensor = data['calib'].to(device=cuda)
+
+    netG.filter(image_tensor)
+    netC.filter(image_tensor)
+    netC.attach(netG.get_im_feat())
+
+    b_min = data['b_min']
+    b_max = data['b_max']
+    try:
+        save_img_path = save_path[:-4] + '.png'
+        save_img_list = []
+        for v in range(image_tensor.shape[0]):
+            save_img = (np.transpose(image_tensor[v].detach().cpu().numpy(),
+                                     (1, 2, 0)) * 0.5 +
+                        0.5)[:, :, ::-1] * 255.0
+            save_img_list.append(save_img)
+        save_img = np.concatenate(save_img_list, axis=1)
+        Image.fromarray(np.uint8(save_img[:, :, ::-1])).save(save_img_path)
+
+        verts, faces, _, _ = reconstruction_faster(netG,
+                                                   cuda,
+                                                   calib_tensor,
+                                                   opt.resolution,
+                                                   b_min,
+                                                   b_max,
+                                                   use_octree=use_octree)
+
+        # Now Getting colors
+        verts_tensor = torch.from_numpy(
+            verts.T).unsqueeze(0).to(device=cuda).float()
+        verts_tensor = reshape_sample_tensor(verts_tensor, opt.num_views)
+        color = np.zeros(verts.shape)
+        interval = 10000
+        for i in range(len(color) // interval):
+            left = i * interval
+            right = i * interval + interval
+            if i == len(color) // interval - 1:
+                right = -1
+            netC.query(verts_tensor[:, :, left:right], calib_tensor)
+            rgb = netC.get_preds()[0].detach().cpu().numpy() * 0.5 + 0.5
+            color[left:right] = rgb.T
+
+        save_obj_mesh_with_color(save_path, verts, faces, color)
+    except Exception as e:
+        print(e)
+        print('Can not create marching cubes at this time.')
+        verts, faces, color = None, None, None
+    return verts, faces, color
+
+
+def adjust_learning_rate(optimizer, epoch, lr, schedule, gamma):
+    """Sets the learning rate to the initial LR decayed by schedule"""
+    if epoch in schedule:
+        lr *= gamma
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = lr
+    return lr
+
+
+def compute_acc(pred, gt, thresh=0.5):
+    '''
+    return:
+        IOU, precision, and recall
+    '''
+    with torch.no_grad():
+        vol_pred = pred > thresh
+        vol_gt = gt > thresh
+
+        union = vol_pred | vol_gt
+        inter = vol_pred & vol_gt
+
+        true_pos = inter.sum().float()
+
+        union = union.sum().float()
+        if union == 0:
+            union = 1
+        vol_pred = vol_pred.sum().float()
+        if vol_pred == 0:
+            vol_pred = 1
+        vol_gt = vol_gt.sum().float()
+        if vol_gt == 0:
+            vol_gt = 1
+        return true_pos / union, true_pos / vol_pred, true_pos / vol_gt
+
+
+# def calc_metrics(opt, net, cuda, dataset, num_tests,
+#                  resolution=128, sampled_points=1000, use_kaolin=True):
+#     if num_tests > len(dataset):
+#         num_tests = len(dataset)
+#     with torch.no_grad():
+#         chamfer_arr, p2s_arr = [], []
+#         for idx in tqdm(range(num_tests)):
+#             data = dataset[idx * len(dataset) // num_tests]
+
+#             verts, faces = gen_mesh_eval(opt, net, cuda, data, resolution)
+#             if verts is None:
+#                 continue
+
+#             mesh_gt = trimesh.load(data['mesh_path'])
+#             mesh_gt = mesh_gt.split(only_watertight=False)
+#             comp_num = [mesh.vertices.shape[0] for mesh in mesh_gt]
+#             mesh_gt = mesh_gt[comp_num.index(max(comp_num))]
+
+#             mesh_pred = trimesh.Trimesh(verts, faces)
+
+#             gt_surface_pts, _ = trimesh.sample.sample_surface_even(
+#                 mesh_gt, sampled_points)
+#             pred_surface_pts, _ = trimesh.sample.sample_surface_even(
+#                 mesh_pred, sampled_points)
+
+#             if use_kaolin and has_kaolin:
+#                 kal_mesh_gt = kal.rep.TriangleMesh.from_tensors(
+#                         torch.tensor(mesh_gt.vertices).float().to(device=cuda),
+#                         torch.tensor(mesh_gt.faces).long().to(device=cuda))
+#                 kal_mesh_pred = kal.rep.TriangleMesh.from_tensors(
+#                     torch.tensor(mesh_pred.vertices).float().to(device=cuda),
+#                     torch.tensor(mesh_pred.faces).long().to(device=cuda))
+
+#                 kal_distance_0 = kal.metrics.mesh.point_to_surface(
+#                     torch.tensor(pred_surface_pts).float().to(device=cuda), kal_mesh_gt)
+#                 kal_distance_1 = kal.metrics.mesh.point_to_surface(
+#                     torch.tensor(gt_surface_pts).float().to(device=cuda), kal_mesh_pred)
+
+#                 dist_gt_pred = torch.sqrt(kal_distance_0).cpu().numpy()
+#                 dist_pred_gt = torch.sqrt(kal_distance_1).cpu().numpy()
+#             else:
+#                 try:
+#                     _, dist_pred_gt, _ = trimesh.proximity.closest_point(mesh_pred, gt_surface_pts)
+#                     _, dist_gt_pred, _ = trimesh.proximity.closest_point(mesh_gt, pred_surface_pts)
+#                 except Exception as e:
+#                     print (e)
+#                     continue
+
+#             chamfer_dist = 0.5 * (dist_pred_gt.mean() + dist_gt_pred.mean())
+#             p2s_dist = dist_pred_gt.mean()
+
+#             chamfer_arr.append(chamfer_dist)
+#             p2s_arr.append(p2s_dist)
+
+#     return np.average(chamfer_arr), np.average(p2s_arr)
+
+
+def calc_error(opt, net, cuda, dataset, num_tests):
+    if num_tests > len(dataset):
+        num_tests = len(dataset)
+    with torch.no_grad():
+        erorr_arr, IOU_arr, prec_arr, recall_arr = [], [], [], []
+        for idx in tqdm(range(num_tests)):
+            data = dataset[idx * len(dataset) // num_tests]
+            # retrieve the data
+            image_tensor = data['img'].to(device=cuda)
+            calib_tensor = data['calib'].to(device=cuda)
+            sample_tensor = data['samples'].to(device=cuda).unsqueeze(0)
+            if opt.num_views > 1:
+                sample_tensor = reshape_sample_tensor(sample_tensor,
+                                                      opt.num_views)
+            label_tensor = data['labels'].to(device=cuda).unsqueeze(0)
+
+            res, error = net.forward(image_tensor,
+                                     sample_tensor,
+                                     calib_tensor,
+                                     labels=label_tensor)
+
+            IOU, prec, recall = compute_acc(res, label_tensor)
+
+            # print(
+            #     '{0}/{1} | Error: {2:06f} IOU: {3:06f} prec: {4:06f} recall: {5:06f}'
+            #         .format(idx, num_tests, error.item(), IOU.item(), prec.item(), recall.item()))
+            erorr_arr.append(error.item())
+            IOU_arr.append(IOU.item())
+            prec_arr.append(prec.item())
+            recall_arr.append(recall.item())
+
+    return np.average(erorr_arr), np.average(IOU_arr), np.average(
+        prec_arr), np.average(recall_arr)
+
+
+def calc_error_color(opt, netG, netC, cuda, dataset, num_tests):
+    if num_tests > len(dataset):
+        num_tests = len(dataset)
+    with torch.no_grad():
+        error_color_arr = []
+
+        for idx in tqdm(range(num_tests)):
+            data = dataset[idx * len(dataset) // num_tests]
+            # retrieve the data
+            image_tensor = data['img'].to(device=cuda)
+            calib_tensor = data['calib'].to(device=cuda)
+            color_sample_tensor = data['color_samples'].to(
+                device=cuda).unsqueeze(0)
+
+            if opt.num_views > 1:
+                color_sample_tensor = reshape_sample_tensor(
+                    color_sample_tensor, opt.num_views)
+
+            rgb_tensor = data['rgbs'].to(device=cuda).unsqueeze(0)
+
+            netG.filter(image_tensor)
+            _, errorC = netC.forward(image_tensor,
+                                     netG.get_im_feat(),
+                                     color_sample_tensor,
+                                     calib_tensor,
+                                     labels=rgb_tensor)
+
+            # print('{0}/{1} | Error inout: {2:06f} | Error color: {3:06f}'
+            #       .format(idx, num_tests, errorG.item(), errorC.item()))
+            error_color_arr.append(errorC.item())
+
+    return np.average(error_color_arr)
+
+
+# pytorch lightning training related fucntions
+
+
+def query_func(opt, netG, features, points, proj_matrix=None):
+    '''
+        - points: size of (bz, N, 3)
+        - proj_matrix: size of (bz, 4, 4)
+    return: size of (bz, 1, N)
+    '''
+    assert len(points) == 1
+    samples = points.repeat(opt.num_views, 1, 1)
+    samples = samples.permute(0, 2, 1)  # [bz, 3, N]
+
+    # view specific query
+    if proj_matrix is not None:
+        samples = orthogonal(samples, proj_matrix)
+
+    calib_tensor = torch.stack([torch.eye(4).float()], dim=0).type_as(samples)
+
+    preds = netG.query(features=features,
+                       points=samples,
+                       calibs=calib_tensor,
+                       regressor=netG.if_regressor)
+
+    if type(preds) is list:
+        preds = preds[0]
+
+    return preds
+
+
+def isin(ar1, ar2):
+    return (ar1[..., None] == ar2).any(-1)
+
+
+def in1d(ar1, ar2):
+    mask = ar2.new_zeros((max(ar1.max(), ar2.max()) + 1, ), dtype=torch.bool)
+    mask[ar2.unique()] = True
+    return mask[ar1]
+
+
+def get_visibility(xy, z, faces):
+    """get the visibility of vertices
+
+    Args:
+        xy (torch.tensor): [N,2]
+        z (torch.tensor): [N,1]
+        faces (torch.tensor): [N,3]
+        size (int): resolution of rendered image
+    """
+
+    xyz = torch.cat((xy, -z), dim=1)
+    xyz = (xyz + 1.0) / 2.0
+    faces = faces.long()
+
+    rasterizer = Pytorch3dRasterizer(image_size=2**12)
+    meshes_screen = Meshes(verts=xyz[None, ...], faces=faces[None, ...])
+    raster_settings = rasterizer.raster_settings
+
+    pix_to_face, zbuf, bary_coords, dists = rasterize_meshes(
+        meshes_screen,
+        image_size=raster_settings.image_size,
+        blur_radius=raster_settings.blur_radius,
+        faces_per_pixel=raster_settings.faces_per_pixel,
+        bin_size=raster_settings.bin_size,
+        max_faces_per_bin=raster_settings.max_faces_per_bin,
+        perspective_correct=raster_settings.perspective_correct,
+        cull_backfaces=raster_settings.cull_backfaces,
+    )
+
+    vis_vertices_id = torch.unique(faces[torch.unique(pix_to_face), :])
+    vis_mask = torch.zeros(size=(z.shape[0], 1))
+    vis_mask[vis_vertices_id] = 1.0
+
+    # print("------------------------\n")
+    # print(f"keep points : {vis_mask.sum()/len(vis_mask)}")
+
+    return vis_mask
+
+
+def batch_mean(res, key):
+    # recursive mean for multilevel dicts
+    return torch.stack([
+        x[key] if isinstance(x, dict) else batch_mean(x, key) for x in res
+    ]).mean()
+
+
+def tf_log_convert(log_dict):
+    new_log_dict = log_dict.copy()
+    for k, v in log_dict.items():
+        new_log_dict[k.replace("_", "/")] = v
+        del new_log_dict[k]
+
+    return new_log_dict
+
+
+def bar_log_convert(log_dict, name=None, rot=None):
+    from decimal import Decimal
+
+    new_log_dict = {}
+
+    if name is not None:
+        new_log_dict['name'] = name[0]
+    if rot is not None:
+        new_log_dict['rot'] = rot[0]
+
+    for k, v in log_dict.items():
+        color = "yellow"
+        if 'loss' in k:
+            color = "red"
+            k = k.replace("loss", "L")
+        elif 'acc' in k:
+            color = "green"
+            k = k.replace("acc", "A")
+        elif 'iou' in k:
+            color = "green"
+            k = k.replace("iou", "I")
+        elif 'prec' in k:
+            color = "green"
+            k = k.replace("prec", "P")
+        elif 'recall' in k:
+            color = "green"
+            k = k.replace("recall", "R")
+
+        if 'lr' not in k:
+            new_log_dict[colored(k.split("_")[1],
+                                 color)] = colored(f"{v:.3f}", color)
+        else:
+            new_log_dict[colored(k.split("_")[1],
+                                 color)] = colored(f"{Decimal(str(v)):.1E}",
+                                                   color)
+
+    if 'loss' in new_log_dict.keys():
+        del new_log_dict['loss']
+
+    return new_log_dict
+
+
+def accumulate(outputs, rot_num, split):
+
+    hparam_log_dict = {}
+
+    metrics = outputs[0].keys()
+    datasets = split.keys()
+
+    for dataset in datasets:
+        for metric in metrics:
+            keyword = f"hparam/{dataset}-{metric}"
+            if keyword not in hparam_log_dict.keys():
+                hparam_log_dict[keyword] = 0
+            for idx in range(split[dataset][0] * rot_num,
+                             split[dataset][1] * rot_num):
+                hparam_log_dict[keyword] += outputs[idx][metric]
+            hparam_log_dict[keyword] /= (split[dataset][1] -
+                                         split[dataset][0]) * rot_num
+
+    print(colored(hparam_log_dict, "green"))
+
+    return hparam_log_dict
+
+
+def calc_error_N(outputs, targets):
+    """calculate the error of normal (IGR)
+
+    Args:
+        outputs (torch.tensor): [B, 3, N]
+        target (torch.tensor): [B, N, 3]
+
+    # manifold loss and grad_loss in IGR paper
+    grad_loss = ((nonmnfld_grad.norm(2, dim=-1) - 1) ** 2).mean()
+    normals_loss = ((mnfld_grad - normals).abs()).norm(2, dim=1).mean()
+
+    Returns:
+        torch.tensor: error of valid normals on the surface
+    """
+    # outputs = torch.tanh(-outputs.permute(0,2,1).reshape(-1,3))
+    outputs = -outputs.permute(0, 2, 1).reshape(-1, 1)
+    targets = targets.reshape(-1, 3)[:, 2:3]
+    with_normals = targets.sum(dim=1).abs() > 0.0
+
+    # eikonal loss
+    grad_loss = ((outputs[with_normals].norm(2, dim=-1) - 1)**2).mean()
+    # normals loss
+    normal_loss = (outputs - targets)[with_normals].abs().norm(2, dim=1).mean()
+
+    return grad_loss * 0.0 + normal_loss
+
+
+def calc_knn_acc(preds, carn_verts, labels, pick_num):
+    """calculate knn accuracy
+
+    Args:
+        preds (torch.tensor): [B, 3, N]
+        carn_verts (torch.tensor): [SMPLX_V_num, 3]
+        labels (torch.tensor): [B, N_knn, N]
+    """
+    N_knn_full = labels.shape[1]
+    preds = preds.permute(0, 2, 1).reshape(-1, 3)
+    labels = labels.permute(0, 2, 1).reshape(-1, N_knn_full)  # [BxN, num_knn]
+    labels = labels[:, :pick_num]
+
+    dist = torch.cdist(preds, carn_verts, p=2)  # [BxN, SMPL_V_num]
+    knn = dist.topk(k=pick_num, dim=1, largest=False)[1]  # [BxN, num_knn]
+    cat_mat = torch.sort(torch.cat((knn, labels), dim=1))[0]
+    bool_col = torch.zeros_like(cat_mat)[:, 0]
+    for i in range(pick_num * 2 - 1):
+        bool_col += cat_mat[:, i] == cat_mat[:, i + 1]
+    acc = (bool_col > 0).sum() / len(bool_col)
+
+    return acc
+
+
+def calc_acc_seg(output, target, num_multiseg):
+    from pytorch_lightning.metrics import Accuracy
+    return Accuracy()(output.reshape(-1, num_multiseg).cpu(),
+                      target.flatten().cpu())
+
+
+def add_watermark(imgs, titles):
+
+    # Write some Text
+
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    bottomLeftCornerOfText = (350, 50)
+    bottomRightCornerOfText = (800, 50)
+    fontScale = 1
+    fontColor = (1.0, 1.0, 1.0)
+    lineType = 2
+
+    for i in range(len(imgs)):
+
+        title = titles[i + 1]
+        cv2.putText(imgs[i], title, bottomLeftCornerOfText, font, fontScale,
+                    fontColor, lineType)
+
+        if i == 0:
+            cv2.putText(imgs[i], str(titles[i][0]), bottomRightCornerOfText,
+                        font, fontScale, fontColor, lineType)
+
+    result = np.concatenate(imgs, axis=0).transpose(2, 0, 1)
+
+    return result
+
+
+def make_test_gif(img_dir):
+
+    if img_dir is not None and len(os.listdir(img_dir)) > 0:
+        for dataset in os.listdir(img_dir):
+            for subject in sorted(os.listdir(osp.join(img_dir, dataset))):
+                img_lst = []
+                im1 = None
+                for file in sorted(
+                        os.listdir(osp.join(img_dir, dataset, subject))):
+                    if file[-3:] not in ['obj', 'gif']:
+                        img_path = os.path.join(img_dir, dataset, subject,
+                                                file)
+                        if im1 == None:
+                            im1 = Image.open(img_path)
+                        else:
+                            img_lst.append(Image.open(img_path))
+
+                print(os.path.join(img_dir, dataset, subject, "out.gif"))
+                im1.save(os.path.join(img_dir, dataset, subject, "out.gif"),
+                         save_all=True,
+                         append_images=img_lst,
+                         duration=500,
+                         loop=0)
+
+
+def export_cfg(logger, cfg):
+
+    cfg_export_file = osp.join(logger.save_dir, logger.name,
+                               f"version_{logger.version}", "cfg.yaml")
+
+    if not osp.exists(cfg_export_file):
+        os.makedirs(osp.dirname(cfg_export_file), exist_ok=True)
+        with open(cfg_export_file, "w+") as file:
+            _ = yaml.dump(cfg, file)

lib/dataloader_demo.py

@@ -0,0 +1,58 @@
+import argparse
+from lib.common.config import get_cfg_defaults
+from lib.dataset.PIFuDataset import PIFuDataset
+
+if __name__ == '__main__':
+        
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-v',
+                        '--show',
+                        action='store_true',
+                        help='vis sampler 3D')
+    parser.add_argument('-s',
+                        '--speed',
+                        action='store_true',
+                        help='vis sampler 3D')
+    parser.add_argument('-l',
+                        '--list',
+                        action='store_true',
+                        help='vis sampler 3D')
+    parser.add_argument('-c',
+                        '--config',
+                        default='./configs/train/icon-filter.yaml',
+                        help='vis sampler 3D')
+    args_c = parser.parse_args()
+
+    args = get_cfg_defaults()
+    args.merge_from_file(args_c.config)
+
+    dataset = PIFuDataset(args, split='train', vis=args_c.show)
+    print(f"Number of subjects :{len(dataset.subject_list)}")
+    data_dict = dataset[0]
+
+    if args_c.list:
+        for k in data_dict.keys():
+            if not hasattr(data_dict[k], "shape"):
+                print(f"{k}: {data_dict[k]}")
+            else:
+                print(f"{k}: {data_dict[k].shape}")
+
+    if args_c.show:
+        # for item in dataset:
+        item = dataset[0]
+        dataset.visualize_sampling3D(item, mode='occ')
+
+    if args_c.speed:
+        # original: 2 it/s
+        # smpl online compute: 2 it/s
+        # normal online compute: 1.5 it/s
+        from tqdm import tqdm
+        for item in tqdm(dataset):
+            # pass
+            for k in item.keys():
+                if 'voxel' in k:
+                    if not hasattr(item[k], "shape"):
+                        print(f"{k}: {item[k]}")
+                    else:
+                        print(f"{k}: {item[k].shape}")
+            print("--------------------")

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 os
+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,662 @@
+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 trimesh
+import torch
+import vedo
+from kaolin.ops.mesh import check_sign
+import torchvision.transforms as transforms
+from ipdb import set_trace
+
+
+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 = osp.join(self.smplx.model_dir, "smpl/SMPL_MALE.pkl")
+        tetra_path = osp.join(self.smplx.tedra_dir,
+                              "tetra_male_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"])
+        
+        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(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 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}
+
+    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]
+        for img_id, img_key in enumerate(['normal_F', 'image', 'T_normal_B']):
+            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]
+            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=15, c=np.float32(colors))
+        vis_list.append(pc)
+
+        vp.show(*vis_list, bg="white", axes=1.0, interactive=True)

lib/dataset/TestDataset.py

@@ -0,0 +1,342 @@
+
+# -*- 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 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 os
+import torch
+import glob
+import numpy as np
+import random
+import human_det
+from termcolor import colored
+from PIL import ImageFile
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+class TestDataset():
+    def __init__(self, cfg, device):
+
+        random.seed(1993)
+
+        self.image_dir = cfg['image_dir']
+        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
+
+        if self.has_det:
+            self.det = human_det.Detection()
+        else:
+            self.det = None
+
+        keep_lst = sorted(glob.glob(f"{self.image_dir}/*"))
+        img_fmts = ['jpg', 'png', 'jpeg', "JPG", 'bmp']
+        keep_lst = [
+            item for item in keep_lst if item.split(".")[-1] in img_fmts
+        ]
+
+        self.subject_list = sorted(
+            [item for item in keep_lst if item.split(".")[-1] in img_fmts])
+
+        # 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 = osp.join(self.smpl_data.model_dir, "smpl/SMPL_NEUTRAL.pkl")
+        tetra_path = osp.join(self.smpl_data.tedra_dir,
+                              'tetra_neutral_adult_smpl.npz')
+        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(osp.join(self.smpl_data.tedra_dir,
+                                    'tetrahedrons_neutral_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) * 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.det, 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.det, 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]
+
+        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])
+
+    def visualize_alignment(self, data):
+
+        import vedo
+        import trimesh
+
+        if self.hps_type != 'pixie':
+            smpl_out = self.smpl_model(betas=data['betas'],
+                                       body_pose=data['body_pose'],
+                                       global_orient=data['global_orient'],
+                                       pose2rot=False)
+            smpl_verts = (
+                (smpl_out.vertices + data['trans']) * data['scale']).detach().cpu().numpy()[0]
+        else:
+            smpl_verts, _, _ = self.smpl_model(shape_params=data['betas'],
+                                               expression_params=data['exp'],
+                                               body_pose=data['body_pose'],
+                                               global_pose=data['global_orient'],
+                                               jaw_pose=data['jaw_pose'],
+                                               left_hand_pose=data['left_hand_pose'],
+                                               right_hand_pose=data['right_hand_pose'])
+
+            smpl_verts = (
+                (smpl_verts + data['trans']) * data['scale']).detach().cpu().numpy()[0]
+
+        smpl_verts *= np.array([1.0, -1.0, -1.0])
+        faces = data['smpl_faces'][0].detach().cpu().numpy()
+
+        image_P = data['image']
+        image_F, image_B = self.render_normal(smpl_verts, faces)
+
+        # create plot
+        vp = vedo.Plotter(title="", size=(1500, 1500))
+        vis_list = []
+
+        image_F = (
+            0.5 * (1.0 + image_F[0].permute(1, 2, 0).detach().cpu().numpy()) * 255.0)
+        image_B = (
+            0.5 * (1.0 + image_B[0].permute(1, 2, 0).detach().cpu().numpy()) * 255.0)
+        image_P = (
+            0.5 * (1.0 + image_P[0].permute(1, 2, 0).detach().cpu().numpy()) * 255.0)
+
+        vis_list.append(vedo.Picture(image_P*0.5+image_F *
+                        0.5).scale(2.0/image_P.shape[0]).pos(-1.0, -1.0, 1.0))
+        vis_list.append(vedo.Picture(image_F).scale(
+            2.0/image_F.shape[0]).pos(-1.0, -1.0, -0.5))
+        vis_list.append(vedo.Picture(image_B).scale(
+            2.0/image_B.shape[0]).pos(-1.0, -1.0, -1.0))
+
+        # create a mesh
+        mesh = trimesh.Trimesh(smpl_verts, faces, process=False)
+        mesh.visual.vertex_colors = [200, 200, 0]
+        vis_list.append(mesh)
+
+        vp.show(*vis_list, bg="white", axes=1, interactive=True)
+
+
+if __name__ == '__main__':
+
+    cfg.merge_from_file("./configs/icon-filter.yaml")
+    cfg.merge_from_file('./lib/pymaf/configs/pymaf_config.yaml')
+
+    cfg_show_list = [
+        'test_gpus', ['0'], 'mcube_res', 512, 'clean_mesh', False
+    ]
+
+    cfg.merge_from_list(cfg_show_list)
+    cfg.freeze()
+
+    os.environ['CUDA_VISIBLE_DEVICES'] = "0"
+    device = torch.device('cuda:0')
+
+    dataset = TestDataset(
+        {
+            'image_dir': "./examples",
+            'has_det': True,    # w/ or w/o detection
+            'hps_type': 'bev'  # pymaf/pare/pixie/hybrik/bev
+        }, device)
+
+    for i in range(len(dataset)):
+        dataset.visualize_alignment(dataset[i])