done

.gitattributes

@@ -0,0 +1,37 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.obj filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text
+*.glb filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,15 @@
+data/*/*
+data/thuman*
+!data/tbfo.ttf
+__pycache__
+debug/
+log/
+.vscode
+!.gitignore
+force_push.sh
+.idea
+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 [email protected]
+
+This Agreement shall be governed by the laws of the Federal Republic of Germany except for the UN Sales Convention.

README.md

@@ -0,0 +1,12 @@
+---
+title: ICON - Clothed Human Digitization
+metaTitle: "Making yourself an ICON, by Yuliang Xiu"
+emoji: 🤼
+colorFrom: indigo
+colorTo: yellow
+sdk: gradio
+sdk_version: 3.1.1
+app_file: app.py
+pinned: true
+python_version: 3.8.13
+---

app.py

@@ -0,0 +1,129 @@
+# install
+
+
+import glob
+import gradio as gr
+import os, random
+
+import subprocess
+
+if os.getenv('SYSTEM') == 'spaces':
+    subprocess.run('pip install pyembree'.split())
+    subprocess.run('pip install rembg'.split())
+    subprocess.run(
+        'pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html'.split())
+    subprocess.run(
+        'pip install git+https://github.com/YuliangXiu/kaolin.git'.split())
+    # subprocess.run('pip install https://download.is.tue.mpg.de/icon/kaolin-0.11.0-cp38-cp38-linux_x86_64.whl'.split())
+    subprocess.run('pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py38_cu113_pyt1110/download.html'.split())
+    subprocess.run(
+        'pip install git+https://github.com/Project-Splinter/human_det.git'.split())
+    subprocess.run(
+        'pip install git+https://github.com/YuliangXiu/neural_voxelization_layer.git'.split())
+
+from apps.infer import generate_model
+
+# running
+
+description = '''
+# ICON Clothed Human Digitization
+### ICON: Implicit Clothed humans Obtained from Normals (CVPR 2022)
+
+<table style="width:26%; padding:0; margin:0;">
+<tr>
+<th><iframe src="https://ghbtns.com/github-btn.html?user=yuliangxiu&repo=ICON&type=star&count=true&v=2&size=small" frameborder="0" scrolling="0" width="100" height="20"></iframe></th>
+<th><img alt="Twitter Follow" src="https://img.shields.io/twitter/follow/yuliangxiu?style=social"></th>
+<th><img alt="YouTube Video Views" src="https://img.shields.io/youtube/views/hZd6AYin2DE?style=social"></th>
+</tr>
+</table>
+
+#### Acknowledgments:
+
+- [StyleGAN-Human, ECCV 2022](https://stylegan-human.github.io/)
+- [nagolinc/styleGanHuman_and_PIFu](https://huggingface.co/spaces/nagolinc/styleGanHuman_and_PIFu)
+- [radames/PIFu-Clothed-Human-Digitization](https://huggingface.co/spaces/radames/PIFu-Clothed-Human-Digitization)
+
+#### The reconstruction + refinement + video take about 80 seconds for single image.
+
+<details>
+
+<summary>More</summary>
+
+#### Image Credits
+
+* [Pinterest](https://www.pinterest.com/search/pins/?q=parkour&rs=sitelinks_searchbox)
+* [Qianli Ma](https://qianlim.github.io/)
+
+#### Related works
+
+* [ICON @ MPI](https://icon.is.tue.mpg.de/)
+* [MonoPort @ USC](https://xiuyuliang.cn/monoport)
+* [Phorhum @ Google](https://phorhum.github.io/)
+* [PIFuHD @ Meta](https://shunsukesaito.github.io/PIFuHD/)
+* [PaMIR @ Tsinghua](http://www.liuyebin.com/pamir/pamir.html)
+
+</details>
+'''
+
+
+def generate_image(seed, psi):
+    iface = gr.Interface.load("spaces/hysts/StyleGAN-Human")
+    img = iface(seed, psi)
+    return img
+
+random.seed(1993)
+model_types = ['icon-filter', 'pifu', 'pamir']
+examples = [[item, random.choice(model_types)] for item in sorted(glob.glob('examples/*.png'))]
+
+with gr.Blocks() as demo:
+    gr.Markdown(description)
+    
+    out_lst = []
+    with gr.Row():
+        with gr.Column():
+            with gr.Row():
+                with gr.Column():
+                    seed = gr.inputs.Slider(
+                        0, 100, step=1, default=0, label='Seed (For Image Generation)')
+                    psi = gr.inputs.Slider(
+                        0, 2, step=0.05, default=0.7, label='Truncation psi (For Image Generation)')
+                    radio_choice = gr.Radio(model_types, label='Method (For Reconstruction)', value='icon-filter')
+                inp = gr.Image(type="filepath", label="Input Image")
+            with gr.Row():
+                btn_sample = gr.Button("Sample Image")
+                btn_submit = gr.Button("Submit Image")
+
+            gr.Examples(examples=examples,
+                        inputs=[inp, radio_choice],
+                        cache_examples=True,
+                        fn=generate_model,
+                        outputs=out_lst)
+
+            out_vid_download = gr.File(label="Download Video, welcome share on Twitter with #ICON")
+
+        with gr.Column():
+            overlap_inp = gr.Image(type="filepath", label="Image Normal Overlap")
+            out_smpl = gr.Model3D(
+                clear_color=[0.0, 0.0, 0.0, 0.0],  label="SMPL")
+            out_smpl_download = gr.File(label="Download SMPL mesh")
+            out_smpl_npy_download = gr.File(label="Download SMPL params")
+            out_recon = gr.Model3D(
+                clear_color=[0.0, 0.0, 0.0, 0.0],  label="ICON")
+            out_recon_download = gr.File(label="Download clothed human mesh")
+            out_final = gr.Model3D(
+                clear_color=[0.0, 0.0, 0.0, 0.0],  label="ICON++")
+            out_final_download = gr.File(label="Download refined clothed human mesh")
+
+    out_lst = [out_smpl, out_smpl_download, out_smpl_npy_download, out_recon, out_recon_download,
+               out_final, out_final_download, out_vid_download, overlap_inp]
+
+    btn_submit.click(fn=generate_model, inputs=[inp, radio_choice], outputs=out_lst)
+    btn_sample.click(fn=generate_image, inputs=[seed, psi], outputs=inp)
+
+if __name__ == "__main__":
+
+    # demo.launch(debug=False, enable_queue=False,
+    #             auth=("[email protected]", "icon_2022"),
+    #             auth_message="Register at icon.is.tue.mpg.de to get HuggingFace username and password.")
+
+    demo.launch(debug=True, enable_queue=True)

apps/ICON.py

@@ -0,0 +1,765 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+
+import os
+
+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.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 lib.smplx as smplx
+import numpy as np
+from torch import nn
+import os.path as osp
+
+from skimage.transform import resize
+import pytorch_lightning as pl
+from huggingface_hub import cached_download
+
+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=cached_download(osp.join(self.smpl_data.model_dir,
+                f"{smpl_type}/{smpl_type}_kid_template.npy"), use_auth_token=os.environ['ICON']),
+            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/infer.py

@@ -0,0 +1,467 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+import os
+
+import logging
+from lib.common.render import query_color
+from lib.common.config import cfg
+from lib.dataset.mesh_util import (
+    load_checkpoint,
+    update_mesh_shape_prior_losses,
+    blend_rgb_norm,
+    unwrap,
+    remesh,
+    tensor2variable,
+)
+
+from lib.dataset.TestDataset import TestDataset
+from lib.net.local_affine import LocalAffine
+from pytorch3d.structures import Meshes
+from apps.ICON import ICON
+
+from termcolor import colored
+import numpy as np
+from PIL import Image
+import trimesh
+import numpy as np
+from tqdm import tqdm
+
+import torch
+torch.backends.cudnn.benchmark = True
+
+logging.getLogger("trimesh").setLevel(logging.ERROR)
+
+
+def generate_model(in_path, model_type):
+    
+    torch.cuda.empty_cache()
+
+    config_dict = {'loop_smpl': 50,
+                   'loop_cloth': 100,
+                   'patience': 5,
+                   'vis_freq': 10,
+                   'out_dir': './results',
+                   'hps_type': 'pymaf',
+                   'config': f"./configs/{model_type}.yaml"}
+
+    # cfg read and merge
+    cfg.merge_from_file(config_dict['config'])
+    cfg.merge_from_file("./lib/pymaf/configs/pymaf_config.yaml")
+    
+    os.makedirs(config_dict['out_dir'], exist_ok=True)
+
+    cfg_show_list = [
+        "test_gpus",
+        [0],
+        "mcube_res",
+        256,
+        "clean_mesh",
+        True,
+    ]
+
+    cfg.merge_from_list(cfg_show_list)
+    cfg.freeze()
+
+    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+    device = torch.device(f"cuda:0")
+
+    # load model and dataloader
+    model = ICON(cfg)
+    model = load_checkpoint(model, cfg)
+
+    dataset_param = {
+        'image_path': in_path,
+        'seg_dir': None,
+        'has_det': True,            # w/ or w/o detection
+        'hps_type': 'pymaf'   # pymaf/pare/pixie
+    }
+
+    if config_dict['hps_type'] == "pixie" and "pamir" in config_dict['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=config_dict['patience'],
+        )
+
+        losses = {
+            # Cloth: Normal_recon - Normal_pred
+            "cloth": {"weight": 1e1, "value": 0.0},
+            # Cloth: [RT]_v1 - [RT]_v2 (v1-edge-v2)
+            "stiffness": {"weight": 1e5, "value": 0.0},
+            # Cloth: det(R) = 1
+            "rigid": {"weight": 1e5, "value": 0.0},
+            # Cloth: edge length
+            "edge": {"weight": 0, "value": 0.0},
+            # Cloth: normal consistency
+            "nc": {"weight": 0, "value": 0.0},
+            # Cloth: laplacian smoonth
+            "laplacian": {"weight": 1e2, "value": 0.0},
+            # Body: Normal_pred - Normal_smpl
+            "normal": {"weight": 1e0, "value": 0.0},
+            # Body: Silhouette_pred - Silhouette_smpl
+            "silhouette": {"weight": 1e0, "value": 0.0},
+        }
+
+        # smpl optimization
+
+        loop_smpl = tqdm(
+            range(config_dict['loop_smpl'] if cfg.net.prior_type != "pifu" else 1))
+
+        for i in loop_smpl:
+
+            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"])
+
+            losses["normal"]["value"] = (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["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)
+
+            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(config_dict['out_dir'], cfg.name,
+                    "refinement"), exist_ok=True)
+
+        # visualize the final results in self-rotation mode
+        os.makedirs(os.path.join(config_dict['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(config_dict['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(config_dict['out_dir'],
+                    cfg.name, "obj"), exist_ok=True)
+
+
+        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(config_dict['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"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_smpl.glb")
+
+        smpl_info = {'betas': optimed_betas,
+                     'pose': optimed_pose,
+                     'orient': optimed_orient,
+                     'trans': optimed_trans}
+
+        np.save(
+            f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_smpl.npy", smpl_info, allow_pickle=True)
+
+        # ------------------------------------------------------------------------------------------------------------------
+
+        # cloth optimization
+
+        # 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(config_dict['out_dir'], cfg.name,
+                         f"obj/{data['name']}_recon.obj")
+        )
+        
+        recon_obj.export(
+            os.path.join(config_dict['out_dir'], cfg.name,
+                         f"obj/{data['name']}_recon.glb")
+        )
+
+        # Isotropic Explicit Remeshing for better geometry topology
+        verts_refine, faces_refine = remesh(os.path.join(config_dict['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=config_dict['patience'],
+        )
+
+        final = None
+
+        if config_dict['loop_cloth'] > 0:
+
+            loop_cloth = tqdm(range(config_dict['loop_cloth']))
+
+            for i in loop_cloth:
+
+                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)
+
+
+            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"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_refine.glb")
+
+
+        # always export visualized video regardless of the cloth refinment
+        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(config_dict['out_dir'], cfg.name,
+                         f"vid/{data['name']}_cloth.mp4"),
+        )
+    
+    smpl_path = f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_smpl.glb"
+    smpl_npy_path = f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_smpl.npy"
+    recon_path = f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_recon.glb"
+    refine_path = f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_refine.glb"
+    
+    video_path = os.path.join(config_dict['out_dir'], cfg.name, f"vid/{data['name']}_cloth.mp4")
+    overlap_path = os.path.join(config_dict['out_dir'], cfg.name, f"png/{data['name']}_overlap.png")
+    
+    torch.cuda.empty_cache()
+    del model
+    del dataset
+    del local_affine_model
+    del optimizer_smpl
+    del optimizer_cloth
+    del scheduler_smpl
+    del scheduler_cloth
+    del losses
+    del in_tensor
+    
+    return [smpl_path, smpl_path, smpl_npy_path, recon_path, recon_path, refine_path, refine_path, video_path, overlap_path]

configs/icon-filter.yaml

@@ -0,0 +1,25 @@
+name: icon-filter
+ckpt_dir: "./data/ckpt/"
+resume_path: "https://huggingface.co/Yuliang/ICON/resolve/main/icon-filter.ckpt"
+normal_path: "https://huggingface.co/Yuliang/ICON/resolve/main/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: "https://huggingface.co/Yuliang/ICON/resolve/main/icon-nofilter.ckpt"
+normal_path: "https://huggingface.co/Yuliang/ICON/resolve/main/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: "https://huggingface.co/Yuliang/ICON/resolve/main/pamir.ckpt"
+normal_path: "https://huggingface.co/Yuliang/ICON/resolve/main/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: "https://huggingface.co/Yuliang/ICON/resolve/main/pifu.ckpt"
+normal_path: "https://huggingface.co/Yuliang/ICON/resolve/main/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

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

lib/dataset/NormalDataset.py

@@ -0,0 +1,212 @@
+
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+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: [email protected]
+
+import numpy as np
+from torch.utils.data import DataLoader
+from .NormalDataset import NormalDataset
+
+# pytorch lightning related libs
+import pytorch_lightning as pl
+
+
+class NormalModule(pl.LightningDataModule):
+    def __init__(self, cfg):
+        super(NormalModule, self).__init__()
+        self.cfg = cfg
+        self.overfit = self.cfg.overfit
+
+        if self.overfit:
+            self.batch_size = 1
+        else:
+            self.batch_size = self.cfg.batch_size
+
+        self.data_size = {}
+
+    def prepare_data(self):
+
+        pass
+
+    @staticmethod
+    def worker_init_fn(worker_id):
+        np.random.seed(np.random.get_state()[1][0] + worker_id)
+
+    def setup(self, stage):
+
+        if stage == 'fit' or stage is None:
+            self.train_dataset = NormalDataset(cfg=self.cfg, split="train")
+            self.val_dataset = NormalDataset(cfg=self.cfg, split="val")
+            self.data_size = {
+                'train': len(self.train_dataset),
+                'val': len(self.val_dataset)
+            }
+
+        if stage == 'test' or stage is None:
+            self.test_dataset = NormalDataset(cfg=self.cfg, split="test")
+
+    def train_dataloader(self):
+
+        train_data_loader = DataLoader(self.train_dataset,
+                                       batch_size=self.batch_size,
+                                       shuffle=not self.overfit,
+                                       num_workers=self.cfg.num_threads,
+                                       pin_memory=True,
+                                       worker_init_fn=self.worker_init_fn)
+
+        return train_data_loader
+
+    def val_dataloader(self):
+
+        if self.overfit:
+            current_dataset = self.train_dataset
+        else:
+            current_dataset = self.val_dataset
+
+        val_data_loader = DataLoader(current_dataset,
+                                     batch_size=self.batch_size,
+                                     shuffle=False,
+                                     num_workers=self.cfg.num_threads,
+                                     pin_memory=True)
+
+        return val_data_loader
+
+    def test_dataloader(self):
+
+        test_data_loader = DataLoader(self.test_dataset,
+                                      batch_size=1,
+                                      shuffle=False,
+                                      num_workers=self.cfg.num_threads,
+                                      pin_memory=True)
+
+        return test_data_loader

lib/dataset/PIFuDataModule.py

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

lib/dataset/PIFuDataset.py

@@ -0,0 +1,589 @@
+from lib.renderer.mesh import load_fit_body
+from lib.dataset.hoppeMesh import HoppeMesh
+from lib.dataset.body_model import TetraSMPLModel
+from lib.common.render import Render
+from lib.dataset.mesh_util import SMPLX, projection, cal_sdf_batch, get_visibility
+from lib.pare.pare.utils.geometry import rotation_matrix_to_angle_axis
+from termcolor import colored
+import os.path as osp
+import numpy as np
+from PIL import Image
+import random
+import os
+import trimesh
+import torch
+from kaolin.ops.mesh import check_sign
+import torchvision.transforms as transforms
+from huggingface_hub import hf_hub_download, cached_download
+
+
+class PIFuDataset():
+    def __init__(self, cfg, split='train', vis=False):
+
+        self.split = split
+        self.root = cfg.root
+        self.bsize = cfg.batch_size
+        self.overfit = cfg.overfit
+
+        # for debug, only used in visualize_sampling3D
+        self.vis = vis
+
+        self.opt = cfg.dataset
+        self.datasets = self.opt.types
+        self.input_size = self.opt.input_size
+        self.scales = self.opt.scales
+        self.workers = cfg.num_threads
+        self.prior_type = cfg.net.prior_type
+
+        self.noise_type = self.opt.noise_type
+        self.noise_scale = self.opt.noise_scale
+
+        noise_joints = [4, 5, 7, 8, 13, 14, 16, 17, 18, 19, 20, 21]
+
+        self.noise_smpl_idx = []
+        self.noise_smplx_idx = []
+
+        for idx in noise_joints:
+            self.noise_smpl_idx.append(idx * 3)
+            self.noise_smpl_idx.append(idx * 3 + 1)
+            self.noise_smpl_idx.append(idx * 3 + 2)
+
+            self.noise_smplx_idx.append((idx-1) * 3)
+            self.noise_smplx_idx.append((idx-1) * 3 + 1)
+            self.noise_smplx_idx.append((idx-1) * 3 + 2)
+
+        self.use_sdf = cfg.sdf
+        self.sdf_clip = cfg.sdf_clip
+
+        # [(feat_name, channel_num),...]
+        self.in_geo = [item[0] for item in cfg.net.in_geo]
+        self.in_nml = [item[0] for item in cfg.net.in_nml]
+
+        self.in_geo_dim = [item[1] for item in cfg.net.in_geo]
+        self.in_nml_dim = [item[1] for item in cfg.net.in_nml]
+
+        self.in_total = self.in_geo + self.in_nml
+        self.in_total_dim = self.in_geo_dim + self.in_nml_dim
+
+        if self.split == 'train':
+            self.rotations = np.arange(
+                0, 360, 360 / self.opt.rotation_num).astype(np.int32)
+        else:
+            self.rotations = range(0, 360, 120)
+
+        self.datasets_dict = {}
+
+        for dataset_id, dataset in enumerate(self.datasets):
+
+            mesh_dir = None
+            smplx_dir = None
+
+            dataset_dir = osp.join(self.root, dataset)
+
+            if dataset in ['thuman2']:
+                mesh_dir = osp.join(dataset_dir, "scans")
+                smplx_dir = osp.join(dataset_dir, "fits")
+                smpl_dir = osp.join(dataset_dir, "smpl")
+
+            self.datasets_dict[dataset] = {
+                "subjects": np.loadtxt(osp.join(dataset_dir, "all.txt"), dtype=str),
+                "smplx_dir": smplx_dir,
+                "smpl_dir": smpl_dir,
+                "mesh_dir": mesh_dir,
+                "scale": self.scales[dataset_id]
+            }
+
+        self.subject_list = self.get_subject_list(split)
+        self.smplx = SMPLX()
+
+        # PIL to tensor
+        self.image_to_tensor = transforms.Compose([
+            transforms.Resize(self.input_size),
+            transforms.ToTensor(),
+            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+        ])
+
+        # PIL to tensor
+        self.mask_to_tensor = transforms.Compose([
+            transforms.Resize(self.input_size),
+            transforms.ToTensor(),
+            transforms.Normalize((0.0, ), (1.0, ))
+        ])
+
+        self.device = torch.device(f"cuda:{cfg.gpus[0]}")
+        self.render = Render(size=512, device=self.device)
+
+    def render_normal(self, verts, faces):
+
+        # render optimized mesh (normal, T_normal, image [-1,1])
+        self.render.load_meshes(verts, faces)
+        return self.render.get_rgb_image()
+
+    def get_subject_list(self, split):
+
+        subject_list = []
+
+        for dataset in self.datasets:
+
+            split_txt = osp.join(self.root, dataset, f'{split}.txt')
+
+            if osp.exists(split_txt):
+                print(f"load from {split_txt}")
+                subject_list += np.loadtxt(split_txt, dtype=str).tolist()
+            else:
+                full_txt = osp.join(self.root, dataset, 'all.txt')
+                print(f"split {full_txt} into train/val/test")
+
+                full_lst = np.loadtxt(full_txt, dtype=str)
+                full_lst = [dataset+"/"+item for item in full_lst]
+                [train_lst, test_lst, val_lst] = np.split(
+                    full_lst, [500, 500+5, ])
+
+                np.savetxt(full_txt.replace(
+                    "all", "train"), train_lst, fmt="%s")
+                np.savetxt(full_txt.replace("all", "test"), test_lst, fmt="%s")
+                np.savetxt(full_txt.replace("all", "val"), val_lst, fmt="%s")
+
+                print(f"load from {split_txt}")
+                subject_list += np.loadtxt(split_txt, dtype=str).tolist()
+
+        if self.split != 'test':
+            subject_list += subject_list[:self.bsize -
+                                         len(subject_list) % self.bsize]
+            print(colored(f"total: {len(subject_list)}", "yellow"))
+            random.shuffle(subject_list)
+
+        # subject_list = ["thuman2/0008"]
+        return subject_list
+
+    def __len__(self):
+        return len(self.subject_list) * len(self.rotations)
+
+    def __getitem__(self, index):
+
+        # only pick the first data if overfitting
+        if self.overfit:
+            index = 0
+
+        rid = index % len(self.rotations)
+        mid = index // len(self.rotations)
+
+        rotation = self.rotations[rid]
+        subject = self.subject_list[mid].split("/")[1]
+        dataset = self.subject_list[mid].split("/")[0]
+        render_folder = "/".join([dataset +
+                                 f"_{self.opt.rotation_num}views", subject])
+
+        # setup paths
+        data_dict = {
+            'dataset': dataset,
+            'subject': subject,
+            'rotation': rotation,
+            'scale': self.datasets_dict[dataset]["scale"],
+            'mesh_path': osp.join(self.datasets_dict[dataset]["mesh_dir"], f"{subject}/{subject}.obj"),
+            'smplx_path': osp.join(self.datasets_dict[dataset]["smplx_dir"], f"{subject}/smplx_param.pkl"),
+            'smpl_path': osp.join(self.datasets_dict[dataset]["smpl_dir"], f"{subject}.pkl"),
+            'calib_path': osp.join(self.root, render_folder, 'calib', f'{rotation:03d}.txt'),
+            'vis_path': osp.join(self.root, render_folder, 'vis', f'{rotation:03d}.pt'),
+            'image_path': osp.join(self.root, render_folder, 'render', f'{rotation:03d}.png')
+        }
+
+        # load training data
+        data_dict.update(self.load_calib(data_dict))
+
+        # image/normal/depth loader
+        for name, channel in zip(self.in_total, self.in_total_dim):
+
+            if f'{name}_path' not in data_dict.keys():
+                data_dict.update({
+                    f'{name}_path': osp.join(self.root, render_folder, name, f'{rotation:03d}.png')
+                })
+
+            # tensor update
+            data_dict.update({
+                name: self.imagepath2tensor(
+                    data_dict[f'{name}_path'], channel, inv=False)
+            })
+
+        data_dict.update(self.load_mesh(data_dict))
+        data_dict.update(self.get_sampling_geo(
+            data_dict, is_valid=self.split == "val", is_sdf=self.use_sdf))
+        data_dict.update(self.load_smpl(data_dict, self.vis))
+
+        if self.prior_type == 'pamir':
+            data_dict.update(self.load_smpl_voxel(data_dict))
+
+        if (self.split != 'test') and (not self.vis):
+
+            del data_dict['verts']
+            del data_dict['faces']
+
+        if not self.vis:
+            del data_dict['mesh']
+
+        path_keys = [
+            key for key in data_dict.keys() if '_path' in key or '_dir' in key
+        ]
+        for key in path_keys:
+            del data_dict[key]
+
+        return data_dict
+
+    def imagepath2tensor(self, path, channel=3, inv=False):
+
+        rgba = Image.open(path).convert('RGBA')
+        mask = rgba.split()[-1]
+        image = rgba.convert('RGB')
+        image = self.image_to_tensor(image)
+        mask = self.mask_to_tensor(mask)
+        image = (image * mask)[:channel]
+
+        return (image * (0.5 - inv) * 2.0).float()
+
+    def load_calib(self, data_dict):
+        calib_data = np.loadtxt(data_dict['calib_path'], dtype=float)
+        extrinsic = calib_data[:4, :4]
+        intrinsic = calib_data[4:8, :4]
+        calib_mat = np.matmul(intrinsic, extrinsic)
+        calib_mat = torch.from_numpy(calib_mat).float()
+        return {'calib': calib_mat}
+
+    def load_mesh(self, data_dict):
+        mesh_path = data_dict['mesh_path']
+        scale = data_dict['scale']
+
+        mesh_ori = trimesh.load(mesh_path,
+                                skip_materials=True,
+                                process=False,
+                                maintain_order=True)
+        verts = mesh_ori.vertices * scale
+        faces = mesh_ori.faces
+
+        vert_normals = np.array(mesh_ori.vertex_normals)
+        face_normals = np.array(mesh_ori.face_normals)
+
+        mesh = HoppeMesh(verts, faces, vert_normals, face_normals)
+
+        return {
+            'mesh': mesh,
+            'verts': torch.as_tensor(mesh.verts).float(),
+            'faces': torch.as_tensor(mesh.faces).long()
+        }
+
+    def add_noise(self,
+                  beta_num,
+                  smpl_pose,
+                  smpl_betas,
+                  noise_type,
+                  noise_scale,
+                  type,
+                  hashcode):
+
+        np.random.seed(hashcode)
+
+        if type == 'smplx':
+            noise_idx = self.noise_smplx_idx
+        else:
+            noise_idx = self.noise_smpl_idx
+
+        if 'beta' in noise_type and noise_scale[noise_type.index("beta")] > 0.0:
+            smpl_betas += (np.random.rand(beta_num) -
+                           0.5) * 2.0 * noise_scale[noise_type.index("beta")]
+            smpl_betas = smpl_betas.astype(np.float32)
+
+        if 'pose' in noise_type and noise_scale[noise_type.index("pose")] > 0.0:
+            smpl_pose[noise_idx] += (
+                np.random.rand(len(noise_idx)) -
+                0.5) * 2.0 * np.pi * noise_scale[noise_type.index("pose")]
+            smpl_pose = smpl_pose.astype(np.float32)
+        if type == 'smplx':
+            return torch.as_tensor(smpl_pose[None, ...]), torch.as_tensor(smpl_betas[None, ...])
+        else:
+            return smpl_pose, smpl_betas
+
+    def compute_smpl_verts(self, data_dict, noise_type=None, noise_scale=None):
+
+        dataset = data_dict['dataset']
+        smplx_dict = {}
+
+        smplx_param = np.load(data_dict['smplx_path'], allow_pickle=True)
+        smplx_pose = smplx_param["body_pose"]  # [1,63]
+        smplx_betas = smplx_param["betas"]  # [1,10]
+        smplx_pose, smplx_betas = self.add_noise(
+            smplx_betas.shape[1],
+            smplx_pose[0],
+            smplx_betas[0],
+            noise_type,
+            noise_scale,
+            type='smplx',
+            hashcode=(hash(f"{data_dict['subject']}_{data_dict['rotation']}")) % (10**8))
+
+        smplx_out, _ = load_fit_body(fitted_path=data_dict['smplx_path'],
+                                     scale=self.datasets_dict[dataset]['scale'],
+                                     smpl_type='smplx',
+                                     smpl_gender='male',
+                                     noise_dict=dict(betas=smplx_betas, body_pose=smplx_pose))
+
+        smplx_dict.update({"type": "smplx",
+                          "gender": 'male',
+                           "body_pose": torch.as_tensor(smplx_pose),
+                           "betas": torch.as_tensor(smplx_betas)})
+
+        return smplx_out.vertices, smplx_dict
+
+    def compute_voxel_verts(self,
+                            data_dict,
+                            noise_type=None,
+                            noise_scale=None):
+
+        smpl_param = np.load(data_dict['smpl_path'], allow_pickle=True)
+        smplx_param = np.load(data_dict['smplx_path'], allow_pickle=True)
+
+        smpl_pose = rotation_matrix_to_angle_axis(
+            torch.as_tensor(smpl_param['full_pose'][0])).numpy()
+        smpl_betas = smpl_param["betas"]
+
+        smpl_path = cached_download(osp.join(self.smplx.model_dir, "smpl/SMPL_MALE.pkl"), use_auth_token=os.environ['ICON'])
+        tetra_path = cached_download(osp.join(self.smplx.tedra_dir,
+                              "tetra_male_adult_smpl.npz"), use_auth_token=os.environ['ICON'])
+
+        smpl_model = TetraSMPLModel(smpl_path, tetra_path, 'adult')
+
+        smpl_pose, smpl_betas = self.add_noise(
+            smpl_model.beta_shape[0],
+            smpl_pose.flatten(),
+            smpl_betas[0],
+            noise_type,
+            noise_scale,
+            type='smpl',
+            hashcode=(hash(f"{data_dict['subject']}_{data_dict['rotation']}")) % (10**8))
+
+        smpl_model.set_params(pose=smpl_pose.reshape(-1, 3),
+                              beta=smpl_betas,
+                              trans=smpl_param["transl"])
+        
+        verts = (np.concatenate([smpl_model.verts, smpl_model.verts_added],
+                                axis=0) * smplx_param["scale"] + smplx_param["translation"]
+                 ) * self.datasets_dict[data_dict['dataset']]['scale']
+        faces = np.loadtxt(cached_download(osp.join(self.smplx.tedra_dir, "tetrahedrons_male_adult.txt"), use_auth_token=os.environ['ICON']),
+                           dtype=np.int32) - 1
+
+        pad_v_num = int(8000 - verts.shape[0])
+        pad_f_num = int(25100 - faces.shape[0])
+
+        verts = np.pad(verts, ((0, pad_v_num), (0, 0)),
+                       mode='constant',
+                       constant_values=0.0).astype(np.float32)
+        faces = np.pad(faces, ((0, pad_f_num), (0, 0)),
+                       mode='constant',
+                       constant_values=0.0).astype(np.int32)
+        
+
+        return verts, faces, pad_v_num, pad_f_num
+
+    def load_smpl(self, data_dict, vis=False):
+
+        smplx_verts, smplx_dict = self.compute_smpl_verts(
+            data_dict, self.noise_type,
+            self.noise_scale)  # compute using smpl model
+
+        smplx_verts = projection(smplx_verts, data_dict['calib']).float()
+        smplx_faces = torch.as_tensor(self.smplx.faces).long()
+        smplx_vis = torch.load(data_dict['vis_path']).float()
+        smplx_cmap = torch.as_tensor(
+            np.load(self.smplx.cmap_vert_path)).float()
+
+        # get smpl_signs
+        query_points = projection(data_dict['samples_geo'],
+                                  data_dict['calib']).float()
+
+        pts_signs = 2.0 * (check_sign(smplx_verts.unsqueeze(0),
+                                      smplx_faces,
+                                      query_points.unsqueeze(0)).float() - 0.5).squeeze(0)
+
+        return_dict = {
+            'smpl_verts': smplx_verts,
+            'smpl_faces': smplx_faces,
+            'smpl_vis': smplx_vis,
+            'smpl_cmap': smplx_cmap,
+            'pts_signs': pts_signs
+        }
+        if smplx_dict is not None:
+            return_dict.update(smplx_dict)
+
+        if vis:
+
+            (xy, z) = torch.as_tensor(smplx_verts).to(
+                self.device).split([2, 1], dim=1)
+            smplx_vis = get_visibility(xy, z, torch.as_tensor(
+                smplx_faces).to(self.device).long())
+
+            T_normal_F, T_normal_B = self.render_normal(
+                (smplx_verts*torch.tensor([1.0, -1.0, 1.0])).to(self.device),
+                smplx_faces.to(self.device))
+
+            return_dict.update({"T_normal_F": T_normal_F.squeeze(0),
+                                "T_normal_B": T_normal_B.squeeze(0)})
+            query_points = projection(data_dict['samples_geo'],
+                                      data_dict['calib']).float()
+
+            smplx_sdf, smplx_norm, smplx_cmap, smplx_vis = cal_sdf_batch(
+                smplx_verts.unsqueeze(0).to(self.device),
+                smplx_faces.unsqueeze(0).to(self.device),
+                smplx_cmap.unsqueeze(0).to(self.device),
+                smplx_vis.unsqueeze(0).to(self.device),
+                query_points.unsqueeze(0).contiguous().to(self.device))
+
+            return_dict.update({
+                'smpl_feat':
+                torch.cat(
+                    (smplx_sdf[0].detach().cpu(),
+                     smplx_cmap[0].detach().cpu(),
+                     smplx_norm[0].detach().cpu(),
+                     smplx_vis[0].detach().cpu()),
+                    dim=1)
+            })
+
+        return return_dict
+
+    def load_smpl_voxel(self, data_dict):
+
+        smpl_verts, smpl_faces, pad_v_num, pad_f_num = self.compute_voxel_verts(
+            data_dict, self.noise_type,
+            self.noise_scale)  # compute using smpl model
+        smpl_verts = projection(smpl_verts, data_dict['calib'])
+
+        smpl_verts *= 0.5
+
+        return {
+            'voxel_verts': smpl_verts,
+            'voxel_faces': smpl_faces,
+            'pad_v_num': pad_v_num,
+            'pad_f_num': pad_f_num
+        }
+
+    def get_sampling_geo(self, data_dict, is_valid=False, is_sdf=False):
+
+        mesh = data_dict['mesh']
+        calib = data_dict['calib']
+
+        # Samples are around the true surface with an offset
+        n_samples_surface = 4 * self.opt.num_sample_geo
+        vert_ids = np.arange(mesh.verts.shape[0])
+        thickness_sample_ratio = np.ones_like(vert_ids).astype(np.float32)
+
+        thickness_sample_ratio /= thickness_sample_ratio.sum()
+
+        samples_surface_ids = np.random.choice(vert_ids,
+                                               n_samples_surface,
+                                               replace=True,
+                                               p=thickness_sample_ratio)
+
+        samples_normal_ids = np.random.choice(vert_ids,
+                                              self.opt.num_sample_geo // 2,
+                                              replace=False,
+                                              p=thickness_sample_ratio)
+
+        surf_samples = mesh.verts[samples_normal_ids, :]
+        surf_normals = mesh.vert_normals[samples_normal_ids, :]
+
+        samples_surface = mesh.verts[samples_surface_ids, :]
+
+        # Sampling offsets are random noise with constant scale (15cm - 20cm)
+        offset = np.random.normal(scale=self.opt.sigma_geo,
+                                  size=(n_samples_surface, 1))
+        samples_surface += mesh.vert_normals[samples_surface_ids, :] * offset
+
+        # Uniform samples in [-1, 1]
+        calib_inv = np.linalg.inv(calib)
+        n_samples_space = self.opt.num_sample_geo // 4
+        samples_space_img = 2.0 * np.random.rand(n_samples_space, 3) - 1.0
+        samples_space = projection(samples_space_img, calib_inv)
+
+        # z-ray direction samples
+        if self.opt.zray_type and not is_valid:
+            n_samples_rayz = self.opt.ray_sample_num
+            samples_surface_cube = projection(samples_surface, calib)
+            samples_surface_cube_repeat = np.repeat(samples_surface_cube,
+                                                    n_samples_rayz,
+                                                    axis=0)
+
+            thickness_repeat = np.repeat(0.5 *
+                                         np.ones_like(samples_surface_ids),
+                                         n_samples_rayz,
+                                         axis=0)
+
+            noise_repeat = np.random.normal(scale=0.40,
+                                            size=(n_samples_surface *
+                                                  n_samples_rayz, ))
+            samples_surface_cube_repeat[:,
+                                        -1] += thickness_repeat * noise_repeat
+            samples_surface_rayz = projection(samples_surface_cube_repeat,
+                                              calib_inv)
+
+            samples = np.concatenate(
+                [samples_surface, samples_space, samples_surface_rayz], 0)
+        else:
+            samples = np.concatenate([samples_surface, samples_space], 0)
+
+        np.random.shuffle(samples)
+
+        # labels: in->1.0; out->0.0.
+        if is_sdf:
+            sdfs = mesh.get_sdf(samples)
+            inside_samples = samples[sdfs < 0]
+            outside_samples = samples[sdfs >= 0]
+
+            inside_sdfs = sdfs[sdfs < 0]
+            outside_sdfs = sdfs[sdfs >= 0]
+        else:
+            inside = mesh.contains(samples)
+            inside_samples = samples[inside >= 0.5]
+            outside_samples = samples[inside < 0.5]
+
+        nin = inside_samples.shape[0]
+
+        if nin > self.opt.num_sample_geo // 2:
+            inside_samples = inside_samples[:self.opt.num_sample_geo // 2]
+            outside_samples = outside_samples[:self.opt.num_sample_geo // 2]
+            if is_sdf:
+                inside_sdfs = inside_sdfs[:self.opt.num_sample_geo // 2]
+                outside_sdfs = outside_sdfs[:self.opt.num_sample_geo // 2]
+        else:
+            outside_samples = outside_samples[:(self.opt.num_sample_geo - nin)]
+            if is_sdf:
+                outside_sdfs = outside_sdfs[:(self.opt.num_sample_geo - nin)]
+
+        if is_sdf:
+            samples = np.concatenate(
+                [inside_samples, outside_samples, surf_samples], 0)
+
+            labels = np.concatenate([
+                inside_sdfs, outside_sdfs, 0.0 * np.ones(surf_samples.shape[0])
+            ])
+
+            normals = np.zeros_like(samples)
+            normals[-self.opt.num_sample_geo // 2:, :] = surf_normals
+
+            # convert sdf from [-14, 130] to [0, 1]
+            # outside: 0, inside: 1
+            # Note: Marching cubes is defined on occupancy space (inside=1.0, outside=0.0)
+
+            labels = -labels.clip(min=-self.sdf_clip, max=self.sdf_clip)
+            labels += self.sdf_clip
+            labels /= (self.sdf_clip * 2)
+
+        else:
+            samples = np.concatenate([inside_samples, outside_samples])
+            labels = np.concatenate([
+                np.ones(inside_samples.shape[0]),
+                np.zeros(outside_samples.shape[0])
+            ])
+
+            normals = np.zeros_like(samples)
+
+        samples = torch.from_numpy(samples).float()
+        labels = torch.from_numpy(labels).float()
+        normals = torch.from_numpy(normals).float()
+
+        return {'samples_geo': samples, 'labels_geo': labels}

lib/dataset/TestDataset.py

@@ -0,0 +1,256 @@
+
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+import os
+
+import lib.smplx as smplx
+from lib.pymaf.utils.geometry import rotation_matrix_to_angle_axis, batch_rodrigues
+from lib.pymaf.utils.imutils import process_image
+from lib.pymaf.core import path_config
+from lib.pymaf.models import pymaf_net
+from lib.common.config import cfg
+from lib.common.render import Render
+from lib.dataset.body_model import TetraSMPLModel
+from lib.dataset.mesh_util import get_visibility, SMPLX
+import os.path as osp
+import torch
+import numpy as np
+import random
+import human_det
+from termcolor import colored
+from PIL import ImageFile
+from huggingface_hub import cached_download
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+class TestDataset():
+    def __init__(self, cfg, device):
+
+        random.seed(1993)
+
+        self.image_path = cfg['image_path']
+        self.seg_dir = cfg['seg_dir']
+        self.has_det = cfg['has_det']
+        self.hps_type = cfg['hps_type']
+        self.smpl_type = 'smpl' if cfg['hps_type'] != 'pixie' else 'smplx'
+        self.smpl_gender = 'neutral'
+
+        self.device = device
+
+        if self.has_det:
+            self.det = human_det.Detection()
+        else:
+            self.det = None
+
+
+        self.subject_list = [self.image_path]
+
+        # smpl related
+        self.smpl_data = SMPLX()
+
+        self.get_smpl_model = lambda smpl_type, smpl_gender: smplx.create(
+            model_path=self.smpl_data.model_dir,
+            gender=smpl_gender,
+            model_type=smpl_type,
+            ext='npz')
+
+        # Load SMPL model
+        self.smpl_model = self.get_smpl_model(
+            self.smpl_type, self.smpl_gender).to(self.device)
+        self.faces = self.smpl_model.faces
+
+        self.hps = pymaf_net(path_config.SMPL_MEAN_PARAMS,
+                                pretrained=True).to(self.device)
+        self.hps.load_state_dict(torch.load(
+            path_config.CHECKPOINT_FILE)['model'],
+            strict=True)
+        self.hps.eval()
+
+        print(colored(f"Using {self.hps_type} as HPS Estimator\n", "green"))
+
+        self.render = Render(size=512, device=device)
+
+    def __len__(self):
+        return len(self.subject_list)
+
+    def compute_vis_cmap(self, smpl_verts, smpl_faces):
+
+        (xy, z) = torch.as_tensor(smpl_verts).split([2, 1], dim=1)
+        smpl_vis = get_visibility(xy, -z, torch.as_tensor(smpl_faces).long())
+        if self.smpl_type == 'smpl':
+            smplx_ind = self.smpl_data.smpl2smplx(np.arange(smpl_vis.shape[0]))
+        else:
+            smplx_ind = np.arange(smpl_vis.shape[0])
+        smpl_cmap = self.smpl_data.get_smpl_mat(smplx_ind)
+
+        return {
+            'smpl_vis': smpl_vis.unsqueeze(0).to(self.device),
+            'smpl_cmap': smpl_cmap.unsqueeze(0).to(self.device),
+            'smpl_verts': smpl_verts.unsqueeze(0)
+        }
+
+    def compute_voxel_verts(self, body_pose, global_orient, betas, trans,
+                            scale):
+
+        smpl_path = cached_download(osp.join(self.smpl_data.model_dir, "smpl/SMPL_NEUTRAL.pkl"), use_auth_token=os.environ['ICON'])
+        tetra_path = cached_download(osp.join(self.smpl_data.tedra_dir,
+                              'tetra_neutral_adult_smpl.npz'), use_auth_token=os.environ['ICON'])
+        smpl_model = TetraSMPLModel(smpl_path, tetra_path, 'adult')
+
+        pose = torch.cat([global_orient[0], body_pose[0]], dim=0)
+        smpl_model.set_params(rotation_matrix_to_angle_axis(pose),
+                              beta=betas[0])
+
+        verts = np.concatenate(
+            [smpl_model.verts, smpl_model.verts_added],
+            axis=0) * scale.item() + trans.detach().cpu().numpy()
+        faces = np.loadtxt(cached_download(osp.join(self.smpl_data.tedra_dir,
+                                    'tetrahedrons_neutral_adult.txt'), use_auth_token=os.environ['ICON']),
+                           dtype=np.int32) - 1
+
+        pad_v_num = int(8000 - verts.shape[0])
+        pad_f_num = int(25100 - faces.shape[0])
+
+        verts = np.pad(verts, ((0, pad_v_num), (0, 0)),
+                       mode='constant',
+                       constant_values=0.0).astype(np.float32) * 0.5
+        faces = np.pad(faces, ((0, pad_f_num), (0, 0)),
+                       mode='constant',
+                       constant_values=0.0).astype(np.int32)
+
+        verts[:, 2] *= -1.0
+
+        voxel_dict = {
+            'voxel_verts':
+            torch.from_numpy(verts).to(self.device).unsqueeze(0).float(),
+            'voxel_faces':
+            torch.from_numpy(faces).to(self.device).unsqueeze(0).long(),
+            'pad_v_num':
+            torch.tensor(pad_v_num).to(self.device).unsqueeze(0).long(),
+            'pad_f_num':
+            torch.tensor(pad_f_num).to(self.device).unsqueeze(0).long()
+        }
+
+        return voxel_dict
+
+    def __getitem__(self, index):
+
+        img_path = self.subject_list[index]
+        img_name = img_path.split("/")[-1].rsplit(".", 1)[0]
+
+        if self.seg_dir is None:
+            img_icon, img_hps, img_ori, img_mask, uncrop_param = process_image(
+                img_path, self.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])

lib/dataset/body_model.py

@@ -0,0 +1,494 @@
+
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+import numpy as np
+import pickle
+import torch
+import os
+
+
+class SMPLModel():
+    def __init__(self, model_path, age):
+        """
+        SMPL model.
+
+        Parameter:
+        ---------
+        model_path: Path to the SMPL model parameters, pre-processed by
+        `preprocess.py`.
+
+        """
+        with open(model_path, 'rb') as f:
+            params = pickle.load(f, encoding='latin1')
+
+            self.J_regressor = params['J_regressor']
+            self.weights = np.asarray(params['weights'])
+            self.posedirs = np.asarray(params['posedirs'])
+            self.v_template = np.asarray(params['v_template'])
+            self.shapedirs = np.asarray(params['shapedirs'])
+            self.faces = np.asarray(params['f'])
+            self.kintree_table = np.asarray(params['kintree_table'])
+
+        self.pose_shape = [24, 3]
+        self.beta_shape = [10]
+        self.trans_shape = [3]
+
+        if age == 'kid':
+            v_template_smil = np.load(
+                os.path.join(os.path.dirname(model_path),
+                             "smpl/smpl_kid_template.npy"))
+            v_template_smil -= np.mean(v_template_smil, axis=0)
+            v_template_diff = np.expand_dims(v_template_smil - self.v_template,
+                                             axis=2)
+            self.shapedirs = np.concatenate(
+                (self.shapedirs[:, :, :self.beta_shape[0]], v_template_diff),
+                axis=2)
+            self.beta_shape[0] += 1
+
+        id_to_col = {
+            self.kintree_table[1, i]: i
+            for i in range(self.kintree_table.shape[1])
+        }
+        self.parent = {
+            i: id_to_col[self.kintree_table[0, i]]
+            for i in range(1, self.kintree_table.shape[1])
+        }
+
+        self.pose = np.zeros(self.pose_shape)
+        self.beta = np.zeros(self.beta_shape)
+        self.trans = np.zeros(self.trans_shape)
+
+        self.verts = None
+        self.J = None
+        self.R = None
+        self.G = None
+
+        self.update()
+
+    def set_params(self, pose=None, beta=None, trans=None):
+        """
+        Set pose, shape, and/or translation parameters of SMPL model. Verices of the
+        model will be updated and returned.
+
+        Prameters:
+        ---------
+        pose: Also known as 'theta', a [24,3] matrix indicating child joint rotation
+        relative to parent joint. For root joint it's global orientation.
+        Represented in a axis-angle format.
+
+        beta: Parameter for model shape. A vector of shape [10]. Coefficients for
+        PCA component. Only 10 components were released by MPI.
+
+        trans: Global translation of shape [3].
+
+        Return:
+        ------
+        Updated vertices.
+
+        """
+        if pose is not None:
+            self.pose = pose
+        if beta is not None:
+            self.beta = beta
+        if trans is not None:
+            self.trans = trans
+        self.update()
+        return self.verts
+
+    def update(self):
+        """
+        Called automatically when parameters are updated.
+
+        """
+        # how beta affect body shape
+        v_shaped = self.shapedirs.dot(self.beta) + self.v_template
+        # joints location
+        self.J = self.J_regressor.dot(v_shaped)
+        pose_cube = self.pose.reshape((-1, 1, 3))
+        # rotation matrix for each joint
+        self.R = self.rodrigues(pose_cube)
+        I_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0),
+                                 (self.R.shape[0] - 1, 3, 3))
+        lrotmin = (self.R[1:] - I_cube).ravel()
+        # how pose affect body shape in zero pose
+        v_posed = v_shaped + self.posedirs.dot(lrotmin)
+        # world transformation of each joint
+        G = np.empty((self.kintree_table.shape[1], 4, 4))
+        G[0] = self.with_zeros(
+            np.hstack((self.R[0], self.J[0, :].reshape([3, 1]))))
+        for i in range(1, self.kintree_table.shape[1]):
+            G[i] = G[self.parent[i]].dot(
+                self.with_zeros(
+                    np.hstack([
+                        self.R[i],
+                        ((self.J[i, :] - self.J[self.parent[i], :]).reshape(
+                            [3, 1]))
+                    ])))
+        # remove the transformation due to the rest pose
+        G = G - self.pack(
+            np.matmul(
+                G,
+                np.hstack([self.J, np.zeros([24, 1])]).reshape([24, 4, 1])))
+        # transformation of each vertex
+        T = np.tensordot(self.weights, G, axes=[[1], [0]])
+        rest_shape_h = np.hstack((v_posed, np.ones([v_posed.shape[0], 1])))
+        v = np.matmul(T, rest_shape_h.reshape([-1, 4, 1])).reshape([-1,
+                                                                    4])[:, :3]
+        self.verts = v + self.trans.reshape([1, 3])
+        self.G = G
+
+    def rodrigues(self, r):
+        """
+        Rodrigues' rotation formula that turns axis-angle vector into rotation
+        matrix in a batch-ed manner.
+
+        Parameter:
+        ----------
+        r: Axis-angle rotation vector of shape [batch_size, 1, 3].
+
+        Return:
+        -------
+        Rotation matrix of shape [batch_size, 3, 3].
+
+        """
+        theta = np.linalg.norm(r, axis=(1, 2), keepdims=True)
+        # avoid zero divide
+        theta = np.maximum(theta, np.finfo(np.float64).tiny)
+        r_hat = r / theta
+        cos = np.cos(theta)
+        z_stick = np.zeros(theta.shape[0])
+        m = np.dstack([
+            z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
+            -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick
+        ]).reshape([-1, 3, 3])
+        i_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0),
+                                 [theta.shape[0], 3, 3])
+        A = np.transpose(r_hat, axes=[0, 2, 1])
+        B = r_hat
+        dot = np.matmul(A, B)
+        R = cos * i_cube + (1 - cos) * dot + np.sin(theta) * m
+        return R
+
+    def with_zeros(self, x):
+        """
+        Append a [0, 0, 0, 1] vector to a [3, 4] matrix.
+
+        Parameter:
+        ---------
+        x: Matrix to be appended.
+
+        Return:
+        ------
+        Matrix after appending of shape [4,4]
+
+        """
+        return np.vstack((x, np.array([[0.0, 0.0, 0.0, 1.0]])))
+
+    def pack(self, x):
+        """
+        Append zero matrices of shape [4, 3] to vectors of [4, 1] shape in a batched
+        manner.
+
+        Parameter:
+        ----------
+        x: Matrices to be appended of shape [batch_size, 4, 1]
+
+        Return:
+        ------
+        Matrix of shape [batch_size, 4, 4] after appending.
+
+        """
+        return np.dstack((np.zeros((x.shape[0], 4, 3)), x))
+
+    def save_to_obj(self, path):
+        """
+        Save the SMPL model into .obj file.
+
+        Parameter:
+        ---------
+        path: Path to save.
+
+        """
+        with open(path, 'w') as fp:
+            for v in self.verts:
+                fp.write('v %f %f %f\n' % (v[0], v[1], v[2]))
+            for f in self.faces + 1:
+                fp.write('f %d %d %d\n' % (f[0], f[1], f[2]))
+
+
+class TetraSMPLModel():
+    def __init__(self,
+                 model_path,
+                 model_addition_path,
+                 age='adult',
+                 v_template=None):
+        """
+        SMPL model.
+
+        Parameter:
+        ---------
+        model_path: Path to the SMPL model parameters, pre-processed by
+        `preprocess.py`.
+
+        """
+        with open(model_path, 'rb') as f:
+            params = pickle.load(f, encoding='latin1')
+
+            self.J_regressor = params['J_regressor']
+            self.weights = np.asarray(params['weights'])
+            self.posedirs = np.asarray(params['posedirs'])
+
+            if v_template is not None:
+                self.v_template = v_template
+            else:
+                self.v_template = np.asarray(params['v_template'])
+
+            self.shapedirs = np.asarray(params['shapedirs'])
+            self.faces = np.asarray(params['f'])
+            self.kintree_table = np.asarray(params['kintree_table'])
+
+        params_added = np.load(model_addition_path)
+        self.v_template_added = params_added['v_template_added']
+        self.weights_added = params_added['weights_added']
+        self.shapedirs_added = params_added['shapedirs_added']
+        self.posedirs_added = params_added['posedirs_added']
+        self.tetrahedrons = params_added['tetrahedrons']
+
+        id_to_col = {
+            self.kintree_table[1, i]: i
+            for i in range(self.kintree_table.shape[1])
+        }
+        self.parent = {
+            i: id_to_col[self.kintree_table[0, i]]
+            for i in range(1, self.kintree_table.shape[1])
+        }
+
+        self.pose_shape = [24, 3]
+        self.beta_shape = [10]
+        self.trans_shape = [3]
+
+        if age == 'kid':
+            v_template_smil = np.load(
+                os.path.join(os.path.dirname(model_path),
+                             "smpl/smpl_kid_template.npy"))
+            v_template_smil -= np.mean(v_template_smil, axis=0)
+            v_template_diff = np.expand_dims(v_template_smil - self.v_template,
+                                             axis=2)
+            self.shapedirs = np.concatenate(
+                (self.shapedirs[:, :, :self.beta_shape[0]], v_template_diff),
+                axis=2)
+            self.beta_shape[0] += 1
+
+        self.pose = np.zeros(self.pose_shape)
+        self.beta = np.zeros(self.beta_shape)
+        self.trans = np.zeros(self.trans_shape)
+
+        self.verts = None
+        self.verts_added = None
+        self.J = None
+        self.R = None
+        self.G = None
+
+        self.update()
+
+    def set_params(self, pose=None, beta=None, trans=None):
+        """
+        Set pose, shape, and/or translation parameters of SMPL model. Verices of the
+        model will be updated and returned.
+
+        Prameters:
+        ---------
+        pose: Also known as 'theta', a [24,3] matrix indicating child joint rotation
+        relative to parent joint. For root joint it's global orientation.
+        Represented in a axis-angle format.
+
+        beta: Parameter for model shape. A vector of shape [10]. Coefficients for
+        PCA component. Only 10 components were released by MPI.
+
+        trans: Global translation of shape [3].
+
+        Return:
+        ------
+        Updated vertices.
+
+        """
+
+        if torch.is_tensor(pose):
+            pose = pose.detach().cpu().numpy()
+        if torch.is_tensor(beta):
+            beta = beta.detach().cpu().numpy()
+
+        if pose is not None:
+            self.pose = pose
+        if beta is not None:
+            self.beta = beta
+        if trans is not None:
+            self.trans = trans
+        self.update()
+        return self.verts
+
+    def update(self):
+        """
+        Called automatically when parameters are updated.
+
+        """
+        # how beta affect body shape
+        v_shaped = self.shapedirs.dot(self.beta) + self.v_template
+        v_shaped_added = self.shapedirs_added.dot(
+            self.beta) + self.v_template_added
+        # joints location
+        self.J = self.J_regressor.dot(v_shaped)
+        pose_cube = self.pose.reshape((-1, 1, 3))
+        # rotation matrix for each joint
+        self.R = self.rodrigues(pose_cube)
+        I_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0),
+                                 (self.R.shape[0] - 1, 3, 3))
+        lrotmin = (self.R[1:] - I_cube).ravel()
+        # how pose affect body shape in zero pose
+        v_posed = v_shaped + self.posedirs.dot(lrotmin)
+        v_posed_added = v_shaped_added + self.posedirs_added.dot(lrotmin)
+        # world transformation of each joint
+        G = np.empty((self.kintree_table.shape[1], 4, 4))
+        G[0] = self.with_zeros(
+            np.hstack((self.R[0], self.J[0, :].reshape([3, 1]))))
+        for i in range(1, self.kintree_table.shape[1]):
+            G[i] = G[self.parent[i]].dot(
+                self.with_zeros(
+                    np.hstack([
+                        self.R[i],
+                        ((self.J[i, :] - self.J[self.parent[i], :]).reshape(
+                            [3, 1]))
+                    ])))
+        # remove the transformation due to the rest pose
+        G = G - self.pack(
+            np.matmul(
+                G,
+                np.hstack([self.J, np.zeros([24, 1])]).reshape([24, 4, 1])))
+        self.G = G
+        # transformation of each vertex
+        T = np.tensordot(self.weights, G, axes=[[1], [0]])
+        rest_shape_h = np.hstack((v_posed, np.ones([v_posed.shape[0], 1])))
+        v = np.matmul(T, rest_shape_h.reshape([-1, 4, 1])).reshape([-1,
+                                                                    4])[:, :3]
+        self.verts = v + self.trans.reshape([1, 3])
+        T_added = np.tensordot(self.weights_added, G, axes=[[1], [0]])
+        rest_shape_added_h = np.hstack(
+            (v_posed_added, np.ones([v_posed_added.shape[0], 1])))
+        v_added = np.matmul(T_added,
+                            rest_shape_added_h.reshape([-1, 4,
+                                                        1])).reshape([-1, 4
+                                                                      ])[:, :3]
+        self.verts_added = v_added + self.trans.reshape([1, 3])
+
+    def rodrigues(self, r):
+        """
+        Rodrigues' rotation formula that turns axis-angle vector into rotation
+        matrix in a batch-ed manner.
+
+        Parameter:
+        ----------
+        r: Axis-angle rotation vector of shape [batch_size, 1, 3].
+
+        Return:
+        -------
+        Rotation matrix of shape [batch_size, 3, 3].
+
+        """
+        theta = np.linalg.norm(r, axis=(1, 2), keepdims=True)
+        # avoid zero divide
+        theta = np.maximum(theta, np.finfo(np.float64).tiny)
+        r_hat = r / theta
+        cos = np.cos(theta)
+        z_stick = np.zeros(theta.shape[0])
+        m = np.dstack([
+            z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
+            -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick
+        ]).reshape([-1, 3, 3])
+        i_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0),
+                                 [theta.shape[0], 3, 3])
+        A = np.transpose(r_hat, axes=[0, 2, 1])
+        B = r_hat
+        dot = np.matmul(A, B)
+        R = cos * i_cube + (1 - cos) * dot + np.sin(theta) * m
+        return R
+
+    def with_zeros(self, x):
+        """
+        Append a [0, 0, 0, 1] vector to a [3, 4] matrix.
+
+        Parameter:
+        ---------
+        x: Matrix to be appended.
+
+        Return:
+        ------
+        Matrix after appending of shape [4,4]
+
+        """
+        return np.vstack((x, np.array([[0.0, 0.0, 0.0, 1.0]])))
+
+    def pack(self, x):
+        """
+        Append zero matrices of shape [4, 3] to vectors of [4, 1] shape in a batched
+        manner.
+
+        Parameter:
+        ----------
+        x: Matrices to be appended of shape [batch_size, 4, 1]
+
+        Return:
+        ------
+        Matrix of shape [batch_size, 4, 4] after appending.
+
+        """
+        return np.dstack((np.zeros((x.shape[0], 4, 3)), x))
+
+    def save_mesh_to_obj(self, path):
+        """
+        Save the SMPL model into .obj file.
+
+        Parameter:
+        ---------
+        path: Path to save.
+
+        """
+        with open(path, 'w') as fp:
+            for v in self.verts:
+                fp.write('v %f %f %f\n' % (v[0], v[1], v[2]))
+            for f in self.faces + 1:
+                fp.write('f %d %d %d\n' % (f[0], f[1], f[2]))
+
+    def save_tetrahedron_to_obj(self, path):
+        """
+        Save the tetrahedron SMPL model into .obj file.
+
+        Parameter:
+        ---------
+        path: Path to save.
+
+        """
+
+        with open(path, 'w') as fp:
+            for v in self.verts:
+                fp.write('v %f %f %f 1 0 0\n' % (v[0], v[1], v[2]))
+            for va in self.verts_added:
+                fp.write('v %f %f %f 0 0 1\n' % (va[0], va[1], va[2]))
+            for t in self.tetrahedrons + 1:
+                fp.write('f %d %d %d\n' % (t[0], t[2], t[1]))
+                fp.write('f %d %d %d\n' % (t[0], t[3], t[2]))
+                fp.write('f %d %d %d\n' % (t[0], t[1], t[3]))
+                fp.write('f %d %d %d\n' % (t[1], t[2], t[3]))

lib/dataset/hoppeMesh.py

@@ -0,0 +1,116 @@
+
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+import numpy as np
+from scipy.spatial import cKDTree
+import trimesh
+
+import logging
+
+logging.getLogger("trimesh").setLevel(logging.ERROR)
+
+
+def save_obj_mesh(mesh_path, verts, faces):
+    file = open(mesh_path, 'w')
+    for v in verts:
+        file.write('v %.4f %.4f %.4f\n' % (v[0], v[1], v[2]))
+    for f in faces:
+        f_plus = f + 1
+        file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2]))
+    file.close()
+
+
+def save_obj_mesh_with_color(mesh_path, verts, faces, colors):
+    file = open(mesh_path, 'w')
+
+    for idx, v in enumerate(verts):
+        c = colors[idx]
+        file.write('v %.4f %.4f %.4f %.4f %.4f %.4f\n' %
+                   (v[0], v[1], v[2], c[0], c[1], c[2]))
+    for f in faces:
+        f_plus = f + 1
+        file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2]))
+    file.close()
+
+
+def save_ply(mesh_path, points, rgb):
+    '''
+    Save the visualization of sampling to a ply file.
+    Red points represent positive predictions.
+    Green points represent negative predictions.
+    :param mesh_path: File name to save
+    :param points: [N, 3] array of points
+    :param rgb: [N, 3] array of rgb values in the range [0~1]
+    :return:
+    '''
+    to_save = np.concatenate([points, rgb * 255], axis=-1)
+    return np.savetxt(
+        mesh_path,
+        to_save,
+        fmt='%.6f %.6f %.6f %d %d %d',
+        comments='',
+        header=(
+            'ply\nformat ascii 1.0\nelement vertex {:d}\n' +
+            'property float x\nproperty float y\nproperty float z\n' +
+            'property uchar red\nproperty uchar green\nproperty uchar blue\n' +
+            'end_header').format(points.shape[0]))
+
+
+class HoppeMesh:
+    def __init__(self, verts, faces, vert_normals, face_normals):
+        '''
+        The HoppeSDF calculates signed distance towards a predefined oriented point cloud
+        http://hhoppe.com/recon.pdf
+        For clean and high-resolution pcl data, this is the fastest and accurate approximation of sdf
+        :param points: pts
+        :param normals: normals
+        '''
+        self.verts = verts  # [n, 3]
+        self.faces = faces  # [m, 3]
+        self.vert_normals = vert_normals  # [n, 3]
+        self.face_normals = face_normals  # [m, 3]
+
+        self.kd_tree = cKDTree(self.verts)
+        self.len = len(self.verts)
+
+    def query(self, points):
+        dists, idx = self.kd_tree.query(points, n_jobs=1)
+        # FIXME: because the eyebows are removed, cKDTree around eyebows
+        # are not accurate. Cause a few false-inside labels here.
+        dirs = points - self.verts[idx]
+        signs = (dirs * self.vert_normals[idx]).sum(axis=1)
+        signs = (signs > 0) * 2 - 1
+        return signs * dists
+
+    def contains(self, points):
+
+        labels = trimesh.Trimesh(vertices=self.verts,
+                                 faces=self.faces).contains(points)
+        return labels
+
+    def export(self, path):
+        if self.colors is not None:
+            save_obj_mesh_with_color(path, self.verts, self.faces,
+                                     self.colors[:, 0:3] / 255.0)
+        else:
+            save_obj_mesh(path, self.verts, self.faces)
+
+    def export_ply(self, path):
+        save_ply(path, self.verts, self.colors[:, 0:3] / 255.0)
+
+    def triangles(self):
+        return self.verts[self.faces]  # [n, 3, 3]

lib/dataset/mesh_util.py

@@ -0,0 +1,894 @@
+
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+import numpy as np
+import cv2
+import pymeshlab
+import torch
+import torchvision
+import trimesh
+from pytorch3d.io import load_obj
+from termcolor import colored
+from scipy.spatial import cKDTree
+
+from pytorch3d.structures import Meshes
+import torch.nn.functional as F
+
+import os
+from lib.pymaf.utils.imutils import uncrop
+from lib.common.render_utils import Pytorch3dRasterizer, face_vertices
+
+from pytorch3d.renderer.mesh import rasterize_meshes
+from PIL import Image, ImageFont, ImageDraw
+from kaolin.ops.mesh import check_sign
+from kaolin.metrics.trianglemesh import point_to_mesh_distance
+
+from pytorch3d.loss import (
+    mesh_laplacian_smoothing,
+    mesh_normal_consistency
+)
+
+from huggingface_hub import hf_hub_download, hf_hub_url, cached_download
+
+
+def tensor2variable(tensor, device):
+    # [1,23,3,3]
+    return torch.tensor(tensor, device=device, requires_grad=True)
+
+
+def normal_loss(vec1, vec2):
+
+    # vec1_mask = vec1.sum(dim=1) != 0.0
+    # vec2_mask = vec2.sum(dim=1) != 0.0
+    # union_mask = vec1_mask * vec2_mask
+    vec_sim = torch.nn.CosineSimilarity(dim=1, eps=1e-6)(vec1, vec2)
+    # vec_diff = ((vec_sim-1.0)**2)[union_mask].mean()
+    vec_diff = ((vec_sim-1.0)**2).mean()
+
+    return vec_diff
+
+
+class GMoF(torch.nn.Module):
+    def __init__(self, rho=1):
+        super(GMoF, self).__init__()
+        self.rho = rho
+
+    def extra_repr(self):
+        return 'rho = {}'.format(self.rho)
+
+    def forward(self, residual):
+        dist = torch.div(residual, residual + self.rho ** 2)
+        return self.rho ** 2 * dist
+
+
+def mesh_edge_loss(meshes, target_length: float = 0.0):
+    """
+    Computes mesh edge length regularization loss averaged across all meshes
+    in a batch. Each mesh contributes equally to the final loss, regardless of
+    the number of edges per mesh in the batch by weighting each mesh with the
+    inverse number of edges. For example, if mesh 3 (out of N) has only E=4
+    edges, then the loss for each edge in mesh 3 should be multiplied by 1/E to
+    contribute to the final loss.
+
+    Args:
+        meshes: Meshes object with a batch of meshes.
+        target_length: Resting value for the edge length.
+
+    Returns:
+        loss: Average loss across the batch. Returns 0 if meshes contains
+        no meshes or all empty meshes.
+    """
+    if meshes.isempty():
+        return torch.tensor(
+            [0.0], dtype=torch.float32, device=meshes.device, requires_grad=True
+        )
+
+    N = len(meshes)
+    edges_packed = meshes.edges_packed()  # (sum(E_n), 3)
+    verts_packed = meshes.verts_packed()  # (sum(V_n), 3)
+    edge_to_mesh_idx = meshes.edges_packed_to_mesh_idx()  # (sum(E_n), )
+    num_edges_per_mesh = meshes.num_edges_per_mesh()  # N
+
+    # Determine the weight for each edge based on the number of edges in the
+    # mesh it corresponds to.
+    # TODO (nikhilar) Find a faster way of computing the weights for each edge
+    # as this is currently a bottleneck for meshes with a large number of faces.
+    weights = num_edges_per_mesh.gather(0, edge_to_mesh_idx)
+    weights = 1.0 / weights.float()
+
+    verts_edges = verts_packed[edges_packed]
+    v0, v1 = verts_edges.unbind(1)
+    loss = ((v0 - v1).norm(dim=1, p=2) - target_length) ** 2.0
+    loss_vertex = loss * weights
+    # loss_outlier = torch.topk(loss, 100)[0].mean()
+    # loss_all = (loss_vertex.sum() + loss_outlier.mean()) / N
+    loss_all = loss_vertex.sum() / N
+
+    return loss_all
+
+
+def remesh(obj_path, perc, device):
+
+    ms = pymeshlab.MeshSet()
+    ms.load_new_mesh(obj_path)
+    ms.laplacian_smooth()
+    ms.remeshing_isotropic_explicit_remeshing(
+        targetlen=pymeshlab.Percentage(perc), adaptive=True)
+    ms.save_current_mesh(obj_path.replace("recon", "remesh"))
+    polished_mesh = trimesh.load_mesh(obj_path.replace("recon", "remesh"))
+    verts_pr = torch.tensor(polished_mesh.vertices).float().unsqueeze(0).to(device)
+    faces_pr = torch.tensor(polished_mesh.faces).long().unsqueeze(0).to(device)
+
+    return verts_pr, faces_pr
+
+
+def possion(mesh, obj_path):
+
+    mesh.export(obj_path)
+    ms = pymeshlab.MeshSet()
+    ms.load_new_mesh(obj_path)
+    ms.surface_reconstruction_screened_poisson(depth=10)
+    ms.set_current_mesh(1)
+    ms.save_current_mesh(obj_path)
+
+    return trimesh.load(obj_path)
+
+
+def get_mask(tensor, dim):
+
+    mask = torch.abs(tensor).sum(dim=dim, keepdims=True) > 0.0
+    mask = mask.type_as(tensor)
+
+    return mask
+
+
+def blend_rgb_norm(rgb, norm, mask):
+
+    # [0,0,0] or [127,127,127] should be marked as mask
+    final = rgb * (1-mask) + norm * (mask)
+
+    return final.astype(np.uint8)
+
+
+def unwrap(image, data):
+
+    img_uncrop = uncrop(np.array(Image.fromarray(image).resize(data['uncrop_param']['box_shape'][:2])),
+                        data['uncrop_param']['center'],
+                        data['uncrop_param']['scale'],
+                        data['uncrop_param']['crop_shape'])
+
+    img_orig = cv2.warpAffine(img_uncrop,
+                              np.linalg.inv(data['uncrop_param']['M'])[:2, :],
+                              data['uncrop_param']['ori_shape'][::-1][1:],
+                              flags=cv2.INTER_CUBIC)
+
+    return img_orig
+
+
+# Losses to smooth / regularize the mesh shape
+def update_mesh_shape_prior_losses(mesh, losses):
+
+    # and (b) the edge length of the predicted mesh
+    losses["edge"]['value'] = mesh_edge_loss(mesh)
+    # mesh normal consistency
+    losses["nc"]['value'] = mesh_normal_consistency(mesh)
+    # mesh laplacian smoothing
+    losses["laplacian"]['value'] = mesh_laplacian_smoothing(
+        mesh, method="uniform")
+
+
+def rename(old_dict, old_name, new_name):
+    new_dict = {}
+    for key, value in zip(old_dict.keys(), old_dict.values()):
+        new_key = key if key != old_name else new_name
+        new_dict[new_key] = old_dict[key]
+    return new_dict
+
+
+def load_checkpoint(model, cfg):
+
+    model_dict = model.state_dict()
+    main_dict = {}
+    normal_dict = {}
+
+    device = torch.device(f"cuda:{cfg['test_gpus'][0]}")
+
+    main_dict = torch.load(cached_download(cfg.resume_path, use_auth_token=os.environ['ICON']),
+                            map_location=device)['state_dict']
+
+    main_dict = {
+        k: v
+        for k, v in main_dict.items()
+        if k in model_dict and v.shape == model_dict[k].shape and (
+            'reconEngine' not in k) and ("normal_filter" not in k) and (
+                'voxelization' not in k)
+    }
+    print(colored(f"Resume MLP weights from {cfg.resume_path}", 'green'))
+
+    normal_dict = torch.load(cached_download(cfg.normal_path, use_auth_token=os.environ['ICON']),
+                                map_location=device)['state_dict']
+
+    for key in normal_dict.keys():
+        normal_dict = rename(normal_dict, key,
+                                key.replace("netG", "netG.normal_filter"))
+
+    normal_dict = {
+        k: v
+        for k, v in normal_dict.items()
+        if k in model_dict and v.shape == model_dict[k].shape
+    }
+    print(colored(f"Resume normal model from {cfg.normal_path}", 'green'))
+
+    model_dict.update(main_dict)
+    model_dict.update(normal_dict)
+    model.load_state_dict(model_dict)
+
+    model.netG = model.netG.to(device)
+    model.reconEngine = model.reconEngine.to(device)
+
+    model.netG.training = False
+    model.netG.eval()
+
+    del main_dict
+    del normal_dict
+    del model_dict
+
+    return model
+
+
+def read_smpl_constants(folder):
+    """Load smpl vertex code"""
+    smpl_vtx_std = np.loadtxt(cached_download(os.path.join(folder, 'vertices.txt'), use_auth_token=os.environ['ICON']))
+    min_x = np.min(smpl_vtx_std[:, 0])
+    max_x = np.max(smpl_vtx_std[:, 0])
+    min_y = np.min(smpl_vtx_std[:, 1])
+    max_y = np.max(smpl_vtx_std[:, 1])
+    min_z = np.min(smpl_vtx_std[:, 2])
+    max_z = np.max(smpl_vtx_std[:, 2])
+
+    smpl_vtx_std[:, 0] = (smpl_vtx_std[:, 0] - min_x) / (max_x - min_x)
+    smpl_vtx_std[:, 1] = (smpl_vtx_std[:, 1] - min_y) / (max_y - min_y)
+    smpl_vtx_std[:, 2] = (smpl_vtx_std[:, 2] - min_z) / (max_z - min_z)
+    smpl_vertex_code = np.float32(np.copy(smpl_vtx_std))
+    """Load smpl faces & tetrahedrons"""
+    smpl_faces = np.loadtxt(cached_download(os.path.join(folder, 'faces.txt'), use_auth_token=os.environ['ICON']),
+                            dtype=np.int32) - 1
+    smpl_face_code = (smpl_vertex_code[smpl_faces[:, 0]] +
+                      smpl_vertex_code[smpl_faces[:, 1]] +
+                      smpl_vertex_code[smpl_faces[:, 2]]) / 3.0
+    smpl_tetras = np.loadtxt(cached_download(os.path.join(folder, 'tetrahedrons.txt'), use_auth_token=os.environ['ICON']),
+                             dtype=np.int32) - 1
+
+    return smpl_vertex_code, smpl_face_code, smpl_faces, smpl_tetras
+
+
+def feat_select(feat, select):
+
+    # feat [B, featx2, N]
+    # select [B, 1, N]
+    # return [B, feat, N]
+
+    dim = feat.shape[1] // 2
+    idx = torch.tile((1-select), (1, dim, 1))*dim + \
+        torch.arange(0, dim).unsqueeze(0).unsqueeze(2).type_as(select)
+    feat_select = torch.gather(feat, 1, idx.long())
+
+    return feat_select
+
+
+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 barycentric_coordinates_of_projection(points, vertices):
+    ''' https://github.com/MPI-IS/mesh/blob/master/mesh/geometry/barycentric_coordinates_of_projection.py
+    '''
+    """Given a point, gives projected coords of that point to a triangle
+    in barycentric coordinates.
+    See
+        **Heidrich**, Computing the Barycentric Coordinates of a Projected Point, JGT 05
+        at http://www.cs.ubc.ca/~heidrich/Papers/JGT.05.pdf
+    
+    :param p: point to project. [B, 3]
+    :param v0: first vertex of triangles. [B, 3]
+    :returns: barycentric coordinates of ``p``'s projection in triangle defined by ``q``, ``u``, ``v``
+            vectorized so ``p``, ``q``, ``u``, ``v`` can all be ``3xN``
+    """
+    #(p, q, u, v)
+    v0, v1, v2 = vertices[:, 0], vertices[:, 1], vertices[:, 2]
+    p = points
+
+    q = v0
+    u = v1 - v0
+    v = v2 - v0
+    n = torch.cross(u, v)
+    s = torch.sum(n * n, dim=1)
+    # If the triangle edges are collinear, cross-product is zero,
+    # which makes "s" 0, which gives us divide by zero. So we
+    # make the arbitrary choice to set s to epsv (=numpy.spacing(1)),
+    # the closest thing to zero
+    s[s == 0] = 1e-6
+    oneOver4ASquared = 1.0 / s
+    w = p - q
+    b2 = torch.sum(torch.cross(u, w) * n, dim=1) * oneOver4ASquared
+    b1 = torch.sum(torch.cross(w, v) * n, dim=1) * oneOver4ASquared
+    weights = torch.stack((1 - b1 - b2, b1, b2), dim=-1)
+    # check barycenric weights
+    # p_n = v0*weights[:,0:1] + v1*weights[:,1:2] + v2*weights[:,2:3]
+    return weights
+
+
+def cal_sdf_batch(verts, faces, cmaps, vis, points):
+
+    # verts [B, N_vert, 3]
+    # faces [B, N_face, 3]
+    # triangles [B, N_face, 3, 3]
+    # points [B, N_point, 3]
+    # cmaps [B, N_vert, 3]
+
+    Bsize = points.shape[0]
+
+    normals = Meshes(verts, faces).verts_normals_padded()
+
+    triangles = face_vertices(verts, faces)
+    normals = face_vertices(normals, faces)
+    cmaps = face_vertices(cmaps, faces)
+    vis = face_vertices(vis, faces)
+
+    residues, pts_ind, _ = point_to_mesh_distance(points, triangles)
+    closest_triangles = torch.gather(
+        triangles, 1, pts_ind[:, :, None, None].expand(-1, -1, 3, 3)).view(-1, 3, 3)
+    closest_normals = torch.gather(
+        normals, 1, pts_ind[:, :, None, None].expand(-1, -1, 3, 3)).view(-1, 3, 3)
+    closest_cmaps = torch.gather(
+        cmaps, 1, pts_ind[:, :, None, None].expand(-1, -1, 3, 3)).view(-1, 3, 3)
+    closest_vis = torch.gather(
+        vis, 1, pts_ind[:, :, None, None].expand(-1, -1, 3, 1)).view(-1, 3, 1)
+    bary_weights = barycentric_coordinates_of_projection(
+        points.view(-1, 3), closest_triangles)
+
+    pts_cmap = (closest_cmaps*bary_weights[:, :, None]).sum(1).unsqueeze(0)
+    pts_vis = (closest_vis*bary_weights[:,
+               :, None]).sum(1).unsqueeze(0).ge(1e-1)
+    pts_norm = (closest_normals*bary_weights[:, :, None]).sum(
+        1).unsqueeze(0) * torch.tensor([-1.0, 1.0, -1.0]).type_as(normals)
+    pts_dist = torch.sqrt(residues) / torch.sqrt(torch.tensor(3))
+
+    pts_signs = 2.0 * (check_sign(verts, faces[0], points).float() - 0.5)
+    pts_sdf = (pts_dist * pts_signs).unsqueeze(-1)
+
+    return pts_sdf.view(Bsize, -1, 1), pts_norm.view(Bsize, -1, 3), pts_cmap.view(Bsize, -1, 3), pts_vis.view(Bsize, -1, 1)
+
+
+def orthogonal(points, calibrations, transforms=None):
+    '''
+    Compute the orthogonal projections of 3D points into the image plane by given projection matrix
+    :param points: [B, 3, N] Tensor of 3D points
+    :param calibrations: [B, 3, 4] Tensor of projection matrix
+    :param transforms: [B, 2, 3] Tensor of image transform matrix
+    :return: xyz: [B, 3, N] Tensor of xyz coordinates in the image plane
+    '''
+    rot = calibrations[:, :3, :3]
+    trans = calibrations[:, :3, 3:4]
+    pts = torch.baddbmm(trans, rot, points)  # [B, 3, N]
+    if transforms is not None:
+        scale = transforms[:2, :2]
+        shift = transforms[:2, 2:3]
+        pts[:, :2, :] = torch.baddbmm(shift, scale, pts[:, :2, :])
+    return pts
+
+
+def projection(points, calib, format='numpy'):
+    if format == 'tensor':
+        return torch.mm(calib[:3, :3], points.T).T + calib[:3, 3]
+    else:
+        return np.matmul(calib[:3, :3], points.T).T + calib[:3, 3]
+
+
+def load_calib(calib_path):
+    calib_data = np.loadtxt(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_mat
+
+
+def load_obj_mesh_for_Hoppe(mesh_file):
+    vertex_data = []
+    face_data = []
+
+    if isinstance(mesh_file, str):
+        f = open(mesh_file, "r")
+    else:
+        f = mesh_file
+    for line in f:
+        if isinstance(line, bytes):
+            line = line.decode("utf-8")
+        if line.startswith('#'):
+            continue
+        values = line.split()
+        if not values:
+            continue
+
+        if values[0] == 'v':
+            v = list(map(float, values[1:4]))
+            vertex_data.append(v)
+
+        elif values[0] == 'f':
+            # quad mesh
+            if len(values) > 4:
+                f = list(map(lambda x: int(x.split('/')[0]), values[1:4]))
+                face_data.append(f)
+                f = list(
+                    map(lambda x: int(x.split('/')[0]),
+                        [values[3], values[4], values[1]]))
+                face_data.append(f)
+            # tri mesh
+            else:
+                f = list(map(lambda x: int(x.split('/')[0]), values[1:4]))
+                face_data.append(f)
+
+    vertices = np.array(vertex_data)
+    faces = np.array(face_data)
+    faces[faces > 0] -= 1
+
+    normals, _ = compute_normal(vertices, faces)
+
+    return vertices, normals, faces
+
+
+def load_obj_mesh_with_color(mesh_file):
+    vertex_data = []
+    color_data = []
+    face_data = []
+
+    if isinstance(mesh_file, str):
+        f = open(mesh_file, "r")
+    else:
+        f = mesh_file
+    for line in f:
+        if isinstance(line, bytes):
+            line = line.decode("utf-8")
+        if line.startswith('#'):
+            continue
+        values = line.split()
+        if not values:
+            continue
+
+        if values[0] == 'v':
+            v = list(map(float, values[1:4]))
+            vertex_data.append(v)
+            c = list(map(float, values[4:7]))
+            color_data.append(c)
+
+        elif values[0] == 'f':
+            # quad mesh
+            if len(values) > 4:
+                f = list(map(lambda x: int(x.split('/')[0]), values[1:4]))
+                face_data.append(f)
+                f = list(
+                    map(lambda x: int(x.split('/')[0]),
+                        [values[3], values[4], values[1]]))
+                face_data.append(f)
+            # tri mesh
+            else:
+                f = list(map(lambda x: int(x.split('/')[0]), values[1:4]))
+                face_data.append(f)
+
+    vertices = np.array(vertex_data)
+    colors = np.array(color_data)
+    faces = np.array(face_data)
+    faces[faces > 0] -= 1
+
+    return vertices, colors, faces
+
+
+def load_obj_mesh(mesh_file, with_normal=False, with_texture=False):
+    vertex_data = []
+    norm_data = []
+    uv_data = []
+
+    face_data = []
+    face_norm_data = []
+    face_uv_data = []
+
+    if isinstance(mesh_file, str):
+        f = open(mesh_file, "r")
+    else:
+        f = mesh_file
+    for line in f:
+        if isinstance(line, bytes):
+            line = line.decode("utf-8")
+        if line.startswith('#'):
+            continue
+        values = line.split()
+        if not values:
+            continue
+
+        if values[0] == 'v':
+            v = list(map(float, values[1:4]))
+            vertex_data.append(v)
+        elif values[0] == 'vn':
+            vn = list(map(float, values[1:4]))
+            norm_data.append(vn)
+        elif values[0] == 'vt':
+            vt = list(map(float, values[1:3]))
+            uv_data.append(vt)
+
+        elif values[0] == 'f':
+            # quad mesh
+            if len(values) > 4:
+                f = list(map(lambda x: int(x.split('/')[0]), values[1:4]))
+                face_data.append(f)
+                f = list(
+                    map(lambda x: int(x.split('/')[0]),
+                        [values[3], values[4], values[1]]))
+                face_data.append(f)
+            # tri mesh
+            else:
+                f = list(map(lambda x: int(x.split('/')[0]), values[1:4]))
+                face_data.append(f)
+
+            # deal with texture
+            if len(values[1].split('/')) >= 2:
+                # quad mesh
+                if len(values) > 4:
+                    f = list(map(lambda x: int(x.split('/')[1]), values[1:4]))
+                    face_uv_data.append(f)
+                    f = list(
+                        map(lambda x: int(x.split('/')[1]),
+                            [values[3], values[4], values[1]]))
+                    face_uv_data.append(f)
+                # tri mesh
+                elif len(values[1].split('/')[1]) != 0:
+                    f = list(map(lambda x: int(x.split('/')[1]), values[1:4]))
+                    face_uv_data.append(f)
+            # deal with normal
+            if len(values[1].split('/')) == 3:
+                # quad mesh
+                if len(values) > 4:
+                    f = list(map(lambda x: int(x.split('/')[2]), values[1:4]))
+                    face_norm_data.append(f)
+                    f = list(
+                        map(lambda x: int(x.split('/')[2]),
+                            [values[3], values[4], values[1]]))
+                    face_norm_data.append(f)
+                # tri mesh
+                elif len(values[1].split('/')[2]) != 0:
+                    f = list(map(lambda x: int(x.split('/')[2]), values[1:4]))
+                    face_norm_data.append(f)
+
+    vertices = np.array(vertex_data)
+    faces = np.array(face_data)
+    faces[faces > 0] -= 1
+
+    if with_texture and with_normal:
+        uvs = np.array(uv_data)
+        face_uvs = np.array(face_uv_data)
+        face_uvs[face_uvs > 0] -= 1
+        norms = np.array(norm_data)
+        if norms.shape[0] == 0:
+            norms, _ = compute_normal(vertices, faces)
+            face_normals = faces
+        else:
+            norms = normalize_v3(norms)
+            face_normals = np.array(face_norm_data)
+            face_normals[face_normals > 0] -= 1
+        return vertices, faces, norms, face_normals, uvs, face_uvs
+
+    if with_texture:
+        uvs = np.array(uv_data)
+        face_uvs = np.array(face_uv_data) - 1
+        return vertices, faces, uvs, face_uvs
+
+    if with_normal:
+        norms = np.array(norm_data)
+        norms = normalize_v3(norms)
+        face_normals = np.array(face_norm_data) - 1
+        return vertices, faces, norms, face_normals
+
+    return vertices, faces
+
+
+def normalize_v3(arr):
+    ''' Normalize a numpy array of 3 component vectors shape=(n,3) '''
+    lens = np.sqrt(arr[:, 0]**2 + arr[:, 1]**2 + arr[:, 2]**2)
+    eps = 0.00000001
+    lens[lens < eps] = eps
+    arr[:, 0] /= lens
+    arr[:, 1] /= lens
+    arr[:, 2] /= lens
+    return arr
+
+
+def compute_normal(vertices, faces):
+    # Create a zeroed array with the same type and shape as our vertices i.e., per vertex normal
+    vert_norms = np.zeros(vertices.shape, dtype=vertices.dtype)
+    # Create an indexed view into the vertex array using the array of three indices for triangles
+    tris = vertices[faces]
+    # Calculate the normal for all the triangles, by taking the cross product of the vectors v1-v0, and v2-v0 in each triangle
+    face_norms = np.cross(tris[::, 1] - tris[::, 0], tris[::, 2] - tris[::, 0])
+    # n is now an array of normals per triangle. The length of each normal is dependent the vertices,
+    # we need to normalize these, so that our next step weights each normal equally.
+    normalize_v3(face_norms)
+    # now we have a normalized array of normals, one per triangle, i.e., per triangle normals.
+    # But instead of one per triangle (i.e., flat shading), we add to each vertex in that triangle,
+    # the triangles' normal. Multiple triangles would then contribute to every vertex, so we need to normalize again afterwards.
+    # The cool part, we can actually add the normals through an indexed view of our (zeroed) per vertex normal array
+    vert_norms[faces[:, 0]] += face_norms
+    vert_norms[faces[:, 1]] += face_norms
+    vert_norms[faces[:, 2]] += face_norms
+    normalize_v3(vert_norms)
+
+    return vert_norms, face_norms
+
+
+def save_obj_mesh(mesh_path, verts, faces):
+    file = open(mesh_path, 'w')
+    for v in verts:
+        file.write('v %.4f %.4f %.4f\n' % (v[0], v[1], v[2]))
+    for f in faces:
+        f_plus = f + 1
+        file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2]))
+    file.close()
+
+
+def save_obj_mesh_with_color(mesh_path, verts, faces, colors):
+    file = open(mesh_path, 'w')
+
+    for idx, v in enumerate(verts):
+        c = colors[idx]
+        file.write('v %.4f %.4f %.4f %.4f %.4f %.4f\n' %
+                   (v[0], v[1], v[2], c[0], c[1], c[2]))
+    for f in faces:
+        f_plus = f + 1
+        file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2]))
+    file.close()
+
+
+def calculate_mIoU(outputs, labels):
+
+    SMOOTH = 1e-6
+
+    outputs = outputs.int()
+    labels = labels.int()
+
+    intersection = (
+        outputs
+        & labels).float().sum()  # Will be zero if Truth=0 or Prediction=0
+    union = (outputs | labels).float().sum()  # Will be zzero if both are 0
+
+    iou = (intersection + SMOOTH) / (union + SMOOTH
+                                     )  # We smooth our devision to avoid 0/0
+
+    thresholded = torch.clamp(
+        20 * (iou - 0.5), 0,
+        10).ceil() / 10  # This is equal to comparing with thresolds
+
+    return thresholded.mean().detach().cpu().numpy(
+    )  # Or thresholded.mean() if you are interested in average across the batch
+
+
+def mask_filter(mask, number=1000):
+    """only keep {number} True items within a mask
+
+    Args:
+        mask (bool array): [N, ]
+        number (int, optional): total True item. Defaults to 1000.
+    """
+    true_ids = np.where(mask)[0]
+    keep_ids = np.random.choice(true_ids, size=number)
+    filter_mask = np.isin(np.arange(len(mask)), keep_ids)
+
+    return filter_mask
+
+
+def query_mesh(path):
+
+    verts, faces_idx, _ = load_obj(path)
+
+    return verts, faces_idx.verts_idx
+
+
+def add_alpha(colors, alpha=0.7):
+
+    colors_pad = np.pad(colors, ((0, 0), (0, 1)),
+                        mode='constant',
+                        constant_values=alpha)
+
+    return colors_pad
+
+
+def get_optim_grid_image(per_loop_lst, loss=None, nrow=4, type='smpl'):
+
+    font_path = os.path.join(os.path.dirname(__file__), "tbfo.ttf")
+    font = ImageFont.truetype(font_path, 30)
+    grid_img = torchvision.utils.make_grid(torch.cat(per_loop_lst, dim=0),
+                                           nrow=nrow)
+    grid_img = Image.fromarray(
+        ((grid_img.permute(1, 2, 0).detach().cpu().numpy() + 1.0) * 0.5 *
+         255.0).astype(np.uint8))
+
+    # add text
+    draw = ImageDraw.Draw(grid_img)
+    grid_size = 512
+    if loss is not None:
+        draw.text((10, 5), f"error: {loss:.3f}", (255, 0, 0), font=font)
+
+    if type == 'smpl':
+        for col_id, col_txt in enumerate(
+                ['image', 'smpl-norm(render)', 'cloth-norm(pred)', 'diff-norm', 'diff-mask']):
+            draw.text((10+(col_id*grid_size), 5),
+                      col_txt, (255, 0, 0), font=font)
+    elif type == 'cloth':
+        for col_id, col_txt in enumerate(
+                ['image', 'cloth-norm(recon)', 'cloth-norm(pred)', 'diff-norm']):
+            draw.text((10+(col_id*grid_size), 5),
+                      col_txt, (255, 0, 0), font=font)
+        for col_id, col_txt in enumerate(
+                ['0', '90', '180', '270']):
+            draw.text((10+(col_id*grid_size), grid_size*2+5),
+                      col_txt, (255, 0, 0), font=font)
+    else:
+        print(f"{type} should be 'smpl' or 'cloth'")
+
+    grid_img = grid_img.resize((grid_img.size[0], grid_img.size[1]),
+                               Image.ANTIALIAS)
+
+    return grid_img
+
+
+def clean_mesh(verts, faces):
+
+    device = verts.device
+
+    mesh_lst = trimesh.Trimesh(verts.detach().cpu().numpy(),
+                               faces.detach().cpu().numpy())
+    mesh_lst = mesh_lst.split(only_watertight=False)
+    comp_num = [mesh.vertices.shape[0] for mesh in mesh_lst]
+    mesh_clean = mesh_lst[comp_num.index(max(comp_num))]
+
+    final_verts = torch.as_tensor(mesh_clean.vertices).float().to(device)
+    final_faces = torch.as_tensor(mesh_clean.faces).int().to(device)
+
+    return final_verts, final_faces
+
+
+def merge_mesh(verts_A, faces_A, verts_B, faces_B, color=False):
+
+    sep_mesh = trimesh.Trimesh(np.concatenate([verts_A, verts_B], axis=0),
+                               np.concatenate(
+                                   [faces_A, faces_B + faces_A.max() + 1],
+                                   axis=0),
+                               maintain_order=True,
+                               process=False)
+    if color:
+        colors = np.ones_like(sep_mesh.vertices)
+        colors[:verts_A.shape[0]] *= np.array([255.0, 0.0, 0.0])
+        colors[verts_A.shape[0]:] *= np.array([0.0, 255.0, 0.0])
+        sep_mesh.visual.vertex_colors = colors
+
+    # union_mesh = trimesh.boolean.union([trimesh.Trimesh(verts_A, faces_A),
+    #                                     trimesh.Trimesh(verts_B, faces_B)], engine='blender')
+
+    return sep_mesh
+
+
+def mesh_move(mesh_lst, step, scale=1.0):
+
+    trans = np.array([1.0, 0.0, 0.0]) * step
+
+    resize_matrix = trimesh.transformations.scale_and_translate(
+        scale=(scale), translate=trans)
+
+    results = []
+
+    for mesh in mesh_lst:
+        mesh.apply_transform(resize_matrix)
+        results.append(mesh)
+
+    return results
+
+
+class SMPLX():
+    def __init__(self):
+        
+        REPO_ID = "Yuliang/SMPL"
+
+        self.smpl_verts_path = hf_hub_download(REPO_ID, filename='smpl_data/smpl_verts.npy', use_auth_token=os.environ['ICON'])
+        self.smplx_verts_path = hf_hub_download(REPO_ID, filename='smpl_data/smplx_verts.npy', use_auth_token=os.environ['ICON'])
+        self.faces_path = hf_hub_download(REPO_ID, filename='smpl_data/smplx_faces.npy', use_auth_token=os.environ['ICON'])
+        self.cmap_vert_path = hf_hub_download(REPO_ID, filename='smpl_data/smplx_cmap.npy', use_auth_token=os.environ['ICON'])
+
+        self.faces = np.load(self.faces_path)
+        self.verts = np.load(self.smplx_verts_path)
+        self.smpl_verts = np.load(self.smpl_verts_path)
+
+        self.model_dir = hf_hub_url(REPO_ID, filename='models')
+        self.tedra_dir = hf_hub_url(REPO_ID, filename='tedra_data')
+        
+    def get_smpl_mat(self, vert_ids):
+
+        mat = torch.as_tensor(np.load(self.cmap_vert_path)).float()
+        return mat[vert_ids, :]
+
+    def smpl2smplx(self, vert_ids=None):
+        """convert vert_ids in smpl to vert_ids in smplx
+
+        Args:
+            vert_ids ([int.array]): [n, knn_num]
+        """
+        smplx_tree = cKDTree(self.verts, leafsize=1)
+        _, ind = smplx_tree.query(self.smpl_verts, k=1)  # ind: [smpl_num, 1]
+
+        if vert_ids is not None:
+            smplx_vert_ids = ind[vert_ids]
+        else:
+            smplx_vert_ids = ind
+
+        return smplx_vert_ids
+
+    def smplx2smpl(self, vert_ids=None):
+        """convert vert_ids in smplx to vert_ids in smpl
+
+        Args:
+            vert_ids ([int.array]): [n, knn_num]
+        """
+        smpl_tree = cKDTree(self.smpl_verts, leafsize=1)
+        _, ind = smpl_tree.query(self.verts, k=1)  # ind: [smplx_num, 1]
+        if vert_ids is not None:
+            smpl_vert_ids = ind[vert_ids]
+        else:
+            smpl_vert_ids = ind
+
+        return smpl_vert_ids

lib/net/BasePIFuNet.py

@@ -0,0 +1,84 @@
+
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+import torch.nn as nn
+import pytorch_lightning as pl
+
+from .geometry import index, orthogonal, perspective
+
+
+class BasePIFuNet(pl.LightningModule):
+    def __init__(
+            self,
+            projection_mode='orthogonal',
+            error_term=nn.MSELoss(),
+    ):
+        """
+        :param projection_mode:
+        Either orthogonal or perspective.
+        It will call the corresponding function for projection.
+        :param error_term:
+        nn Loss between the predicted [B, Res, N] and the label [B, Res, N]
+        """
+        super(BasePIFuNet, self).__init__()
+        self.name = 'base'
+
+        self.error_term = error_term
+
+        self.index = index
+        self.projection = orthogonal if projection_mode == 'orthogonal' else perspective
+
+    def forward(self, points, images, calibs, transforms=None):
+        '''
+        :param points: [B, 3, N] world space coordinates of points
+        :param images: [B, C, H, W] input images
+        :param calibs: [B, 3, 4] calibration matrices for each image
+        :param transforms: Optional [B, 2, 3] image space coordinate transforms
+        :return: [B, Res, N] predictions for each point
+        '''
+        features = self.filter(images)
+        preds = self.query(features, points, calibs, transforms)
+        return preds
+
+    def filter(self, images):
+        '''
+        Filter the input images
+        store all intermediate features.
+        :param images: [B, C, H, W] input images
+        '''
+        return None
+
+    def query(self, features, points, calibs, transforms=None):
+        '''
+        Given 3D points, query the network predictions for each point.
+        Image features should be pre-computed before this call.
+        store all intermediate features.
+        query() function may behave differently during training/testing.
+        :param points: [B, 3, N] world space coordinates of points
+        :param calibs: [B, 3, 4] calibration matrices for each image
+        :param transforms: Optional [B, 2, 3] image space coordinate transforms
+        :param labels: Optional [B, Res, N] gt labeling
+        :return: [B, Res, N] predictions for each point
+        '''
+        return None
+
+    def get_error(self, preds, labels):
+        '''
+        Get the network loss from the last query
+        :return: loss term
+        '''
+        return self.error_term(preds, labels)

lib/net/FBNet.py

@@ -0,0 +1,387 @@
+'''
+Copyright (C) 2019 NVIDIA Corporation. Ting-Chun Wang, Ming-Yu Liu, Jun-Yan Zhu.
+BSD License. All rights reserved. 
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE, INCLUDING ALL 
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR ANY PARTICULAR PURPOSE. 
+IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL 
+DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, 
+WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING 
+OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+'''
+import torch
+import torch.nn as nn
+import functools
+import numpy as np
+import pytorch_lightning as pl
+
+
+###############################################################################
+# Functions
+###############################################################################
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        m.weight.data.normal_(0.0, 0.02)
+    elif classname.find('BatchNorm2d') != -1:
+        m.weight.data.normal_(1.0, 0.02)
+        m.bias.data.fill_(0)
+
+
+def get_norm_layer(norm_type='instance'):
+    if norm_type == 'batch':
+        norm_layer = functools.partial(nn.BatchNorm2d, affine=True)
+    elif norm_type == 'instance':
+        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False)
+    else:
+        raise NotImplementedError('normalization layer [%s] is not found' %
+                                  norm_type)
+    return norm_layer
+
+
+def define_G(input_nc,
+             output_nc,
+             ngf,
+             netG,
+             n_downsample_global=3,
+             n_blocks_global=9,
+             n_local_enhancers=1,
+             n_blocks_local=3,
+             norm='instance',
+             gpu_ids=[],
+             last_op=nn.Tanh()):
+    norm_layer = get_norm_layer(norm_type=norm)
+    if netG == 'global':
+        netG = GlobalGenerator(input_nc,
+                               output_nc,
+                               ngf,
+                               n_downsample_global,
+                               n_blocks_global,
+                               norm_layer,
+                               last_op=last_op)
+    elif netG == 'local':
+        netG = LocalEnhancer(input_nc, output_nc, ngf, n_downsample_global,
+                             n_blocks_global, n_local_enhancers,
+                             n_blocks_local, norm_layer)
+    elif netG == 'encoder':
+        netG = Encoder(input_nc, output_nc, ngf, n_downsample_global,
+                       norm_layer)
+    else:
+        raise ('generator not implemented!')
+    # print(netG)
+    if len(gpu_ids) > 0:
+        assert (torch.cuda.is_available())
+        netG.cuda(gpu_ids[0])
+    netG.apply(weights_init)
+    return netG
+
+
+def print_network(net):
+    if isinstance(net, list):
+        net = net[0]
+    num_params = 0
+    for param in net.parameters():
+        num_params += param.numel()
+    print(net)
+    print('Total number of parameters: %d' % num_params)
+
+
+##############################################################################
+# Generator
+##############################################################################
+class LocalEnhancer(pl.LightningModule):
+    def __init__(self,
+                 input_nc,
+                 output_nc,
+                 ngf=32,
+                 n_downsample_global=3,
+                 n_blocks_global=9,
+                 n_local_enhancers=1,
+                 n_blocks_local=3,
+                 norm_layer=nn.BatchNorm2d,
+                 padding_type='reflect'):
+        super(LocalEnhancer, self).__init__()
+        self.n_local_enhancers = n_local_enhancers
+
+        ###### global generator model #####
+        ngf_global = ngf * (2**n_local_enhancers)
+        model_global = GlobalGenerator(input_nc, output_nc, ngf_global,
+                                       n_downsample_global, n_blocks_global,
+                                       norm_layer).model
+        model_global = [model_global[i] for i in range(len(model_global) - 3)
+                        ]  # get rid of final convolution layers
+        self.model = nn.Sequential(*model_global)
+
+        ###### local enhancer layers #####
+        for n in range(1, n_local_enhancers + 1):
+            # downsample
+            ngf_global = ngf * (2**(n_local_enhancers - n))
+            model_downsample = [
+                nn.ReflectionPad2d(3),
+                nn.Conv2d(input_nc, ngf_global, kernel_size=7, padding=0),
+                norm_layer(ngf_global),
+                nn.ReLU(True),
+                nn.Conv2d(ngf_global,
+                          ngf_global * 2,
+                          kernel_size=3,
+                          stride=2,
+                          padding=1),
+                norm_layer(ngf_global * 2),
+                nn.ReLU(True)
+            ]
+            # residual blocks
+            model_upsample = []
+            for i in range(n_blocks_local):
+                model_upsample += [
+                    ResnetBlock(ngf_global * 2,
+                                padding_type=padding_type,
+                                norm_layer=norm_layer)
+                ]
+
+            # upsample
+            model_upsample += [
+                nn.ConvTranspose2d(ngf_global * 2,
+                                   ngf_global,
+                                   kernel_size=3,
+                                   stride=2,
+                                   padding=1,
+                                   output_padding=1),
+                norm_layer(ngf_global),
+                nn.ReLU(True)
+            ]
+
+            # final convolution
+            if n == n_local_enhancers:
+                model_upsample += [
+                    nn.ReflectionPad2d(3),
+                    nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0),
+                    nn.Tanh()
+                ]
+
+            setattr(self, 'model' + str(n) + '_1',
+                    nn.Sequential(*model_downsample))
+            setattr(self, 'model' + str(n) + '_2',
+                    nn.Sequential(*model_upsample))
+
+        self.downsample = nn.AvgPool2d(3,
+                                       stride=2,
+                                       padding=[1, 1],
+                                       count_include_pad=False)
+
+    def forward(self, input):
+        # create input pyramid
+        input_downsampled = [input]
+        for i in range(self.n_local_enhancers):
+            input_downsampled.append(self.downsample(input_downsampled[-1]))
+
+        # output at coarest level
+        output_prev = self.model(input_downsampled[-1])
+        # build up one layer at a time
+        for n_local_enhancers in range(1, self.n_local_enhancers + 1):
+            model_downsample = getattr(self,
+                                       'model' + str(n_local_enhancers) + '_1')
+            model_upsample = getattr(self,
+                                     'model' + str(n_local_enhancers) + '_2')
+            input_i = input_downsampled[self.n_local_enhancers -
+                                        n_local_enhancers]
+            output_prev = model_upsample(
+                model_downsample(input_i) + output_prev)
+        return output_prev
+
+
+class GlobalGenerator(pl.LightningModule):
+    def __init__(self,
+                 input_nc,
+                 output_nc,
+                 ngf=64,
+                 n_downsampling=3,
+                 n_blocks=9,
+                 norm_layer=nn.BatchNorm2d,
+                 padding_type='reflect',
+                 last_op=nn.Tanh()):
+        assert (n_blocks >= 0)
+        super(GlobalGenerator, self).__init__()
+        activation = nn.ReLU(True)
+
+        model = [
+            nn.ReflectionPad2d(3),
+            nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0),
+            norm_layer(ngf), activation
+        ]
+        # downsample
+        for i in range(n_downsampling):
+            mult = 2**i
+            model += [
+                nn.Conv2d(ngf * mult,
+                          ngf * mult * 2,
+                          kernel_size=3,
+                          stride=2,
+                          padding=1),
+                norm_layer(ngf * mult * 2), activation
+            ]
+
+        # resnet blocks
+        mult = 2**n_downsampling
+        for i in range(n_blocks):
+            model += [
+                ResnetBlock(ngf * mult,
+                            padding_type=padding_type,
+                            activation=activation,
+                            norm_layer=norm_layer)
+            ]
+
+        # upsample
+        for i in range(n_downsampling):
+            mult = 2**(n_downsampling - i)
+            model += [
+                nn.ConvTranspose2d(ngf * mult,
+                                   int(ngf * mult / 2),
+                                   kernel_size=3,
+                                   stride=2,
+                                   padding=1,
+                                   output_padding=1),
+                norm_layer(int(ngf * mult / 2)), activation
+            ]
+        model += [
+            nn.ReflectionPad2d(3),
+            nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)
+        ]
+        if last_op is not None:
+            model += [last_op]
+        self.model = nn.Sequential(*model)
+
+    def forward(self, input):
+        return self.model(input)
+
+
+# Define a resnet block
+class ResnetBlock(pl.LightningModule):
+    def __init__(self,
+                 dim,
+                 padding_type,
+                 norm_layer,
+                 activation=nn.ReLU(True),
+                 use_dropout=False):
+        super(ResnetBlock, self).__init__()
+        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer,
+                                                activation, use_dropout)
+
+    def build_conv_block(self, dim, padding_type, norm_layer, activation,
+                         use_dropout):
+        conv_block = []
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' %
+                                      padding_type)
+
+        conv_block += [
+            nn.Conv2d(dim, dim, kernel_size=3, padding=p),
+            norm_layer(dim), activation
+        ]
+        if use_dropout:
+            conv_block += [nn.Dropout(0.5)]
+
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' %
+                                      padding_type)
+        conv_block += [
+            nn.Conv2d(dim, dim, kernel_size=3, padding=p),
+            norm_layer(dim)
+        ]
+
+        return nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        out = x + self.conv_block(x)
+        return out
+
+
+class Encoder(pl.LightningModule):
+    def __init__(self,
+                 input_nc,
+                 output_nc,
+                 ngf=32,
+                 n_downsampling=4,
+                 norm_layer=nn.BatchNorm2d):
+        super(Encoder, self).__init__()
+        self.output_nc = output_nc
+
+        model = [
+            nn.ReflectionPad2d(3),
+            nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0),
+            norm_layer(ngf),
+            nn.ReLU(True)
+        ]
+        # downsample
+        for i in range(n_downsampling):
+            mult = 2**i
+            model += [
+                nn.Conv2d(ngf * mult,
+                          ngf * mult * 2,
+                          kernel_size=3,
+                          stride=2,
+                          padding=1),
+                norm_layer(ngf * mult * 2),
+                nn.ReLU(True)
+            ]
+
+        # upsample
+        for i in range(n_downsampling):
+            mult = 2**(n_downsampling - i)
+            model += [
+                nn.ConvTranspose2d(ngf * mult,
+                                   int(ngf * mult / 2),
+                                   kernel_size=3,
+                                   stride=2,
+                                   padding=1,
+                                   output_padding=1),
+                norm_layer(int(ngf * mult / 2)),
+                nn.ReLU(True)
+            ]
+
+        model += [
+            nn.ReflectionPad2d(3),
+            nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0),
+            nn.Tanh()
+        ]
+        self.model = nn.Sequential(*model)
+
+    def forward(self, input, inst):
+        outputs = self.model(input)
+
+        # instance-wise average pooling
+        outputs_mean = outputs.clone()
+        inst_list = np.unique(inst.cpu().numpy().astype(int))
+        for i in inst_list:
+            for b in range(input.size()[0]):
+                indices = (inst[b:b + 1] == int(i)).nonzero()  # n x 4
+                for j in range(self.output_nc):
+                    output_ins = outputs[indices[:, 0] + b, indices[:, 1] + j,
+                                         indices[:, 2], indices[:, 3]]
+                    mean_feat = torch.mean(output_ins).expand_as(output_ins)
+                    outputs_mean[indices[:, 0] + b, indices[:, 1] + j,
+                                 indices[:, 2], indices[:, 3]] = mean_feat
+        return outputs_mean

lib/net/HGFilters.py

@@ -0,0 +1,197 @@
+
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+from lib.net.net_util import *
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class HourGlass(nn.Module):
+    def __init__(self, num_modules, depth, num_features, opt):
+        super(HourGlass, self).__init__()
+        self.num_modules = num_modules
+        self.depth = depth
+        self.features = num_features
+        self.opt = opt
+
+        self._generate_network(self.depth)
+
+    def _generate_network(self, level):
+        self.add_module('b1_' + str(level),
+                        ConvBlock(self.features, self.features, self.opt))
+
+        self.add_module('b2_' + str(level),
+                        ConvBlock(self.features, self.features, self.opt))
+
+        if level > 1:
+            self._generate_network(level - 1)
+        else:
+            self.add_module('b2_plus_' + str(level),
+                            ConvBlock(self.features, self.features, self.opt))
+
+        self.add_module('b3_' + str(level),
+                        ConvBlock(self.features, self.features, self.opt))
+
+    def _forward(self, level, inp):
+        # Upper branch
+        up1 = inp
+        up1 = self._modules['b1_' + str(level)](up1)
+
+        # Lower branch
+        low1 = F.avg_pool2d(inp, 2, stride=2)
+        low1 = self._modules['b2_' + str(level)](low1)
+
+        if level > 1:
+            low2 = self._forward(level - 1, low1)
+        else:
+            low2 = low1
+            low2 = self._modules['b2_plus_' + str(level)](low2)
+
+        low3 = low2
+        low3 = self._modules['b3_' + str(level)](low3)
+
+        # NOTE: for newer PyTorch (1.3~), it seems that training results are degraded due to implementation diff in F.grid_sample
+        # if the pretrained model behaves weirdly, switch with the commented line.
+        # NOTE: I also found that "bicubic" works better.
+        up2 = F.interpolate(low3,
+                            scale_factor=2,
+                            mode='bicubic',
+                            align_corners=True)
+        # up2 = F.interpolate(low3, scale_factor=2, mode='nearest)
+
+        return up1 + up2
+
+    def forward(self, x):
+        return self._forward(self.depth, x)
+
+
+class HGFilter(nn.Module):
+    def __init__(self, opt, num_modules, in_dim):
+        super(HGFilter, self).__init__()
+        self.num_modules = num_modules
+
+        self.opt = opt
+        [k, s, d, p] = self.opt.conv1
+
+        # self.conv1 = nn.Conv2d(in_dim, 64, kernel_size=7, stride=2, padding=3)
+        self.conv1 = nn.Conv2d(in_dim,
+                               64,
+                               kernel_size=k,
+                               stride=s,
+                               dilation=d,
+                               padding=p)
+
+        if self.opt.norm == 'batch':
+            self.bn1 = nn.BatchNorm2d(64)
+        elif self.opt.norm == 'group':
+            self.bn1 = nn.GroupNorm(32, 64)
+
+        if self.opt.hg_down == 'conv64':
+            self.conv2 = ConvBlock(64, 64, self.opt)
+            self.down_conv2 = nn.Conv2d(64,
+                                        128,
+                                        kernel_size=3,
+                                        stride=2,
+                                        padding=1)
+        elif self.opt.hg_down == 'conv128':
+            self.conv2 = ConvBlock(64, 128, self.opt)
+            self.down_conv2 = nn.Conv2d(128,
+                                        128,
+                                        kernel_size=3,
+                                        stride=2,
+                                        padding=1)
+        elif self.opt.hg_down == 'ave_pool':
+            self.conv2 = ConvBlock(64, 128, self.opt)
+        else:
+            raise NameError('Unknown Fan Filter setting!')
+
+        self.conv3 = ConvBlock(128, 128, self.opt)
+        self.conv4 = ConvBlock(128, 256, self.opt)
+
+        # Stacking part
+        for hg_module in range(self.num_modules):
+            self.add_module('m' + str(hg_module),
+                            HourGlass(1, opt.num_hourglass, 256, self.opt))
+
+            self.add_module('top_m_' + str(hg_module),
+                            ConvBlock(256, 256, self.opt))
+            self.add_module(
+                'conv_last' + str(hg_module),
+                nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0))
+            if self.opt.norm == 'batch':
+                self.add_module('bn_end' + str(hg_module), nn.BatchNorm2d(256))
+            elif self.opt.norm == 'group':
+                self.add_module('bn_end' + str(hg_module),
+                                nn.GroupNorm(32, 256))
+
+            self.add_module(
+                'l' + str(hg_module),
+                nn.Conv2d(256,
+                          opt.hourglass_dim,
+                          kernel_size=1,
+                          stride=1,
+                          padding=0))
+
+            if hg_module < self.num_modules - 1:
+                self.add_module(
+                    'bl' + str(hg_module),
+                    nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0))
+                self.add_module(
+                    'al' + str(hg_module),
+                    nn.Conv2d(opt.hourglass_dim,
+                              256,
+                              kernel_size=1,
+                              stride=1,
+                              padding=0))
+
+    def forward(self, x):
+        x = F.relu(self.bn1(self.conv1(x)), True)
+        tmpx = x
+        if self.opt.hg_down == 'ave_pool':
+            x = F.avg_pool2d(self.conv2(x), 2, stride=2)
+        elif self.opt.hg_down in ['conv64', 'conv128']:
+            x = self.conv2(x)
+            x = self.down_conv2(x)
+        else:
+            raise NameError('Unknown Fan Filter setting!')
+
+        x = self.conv3(x)
+        x = self.conv4(x)
+
+        previous = x
+
+        outputs = []
+        for i in range(self.num_modules):
+            hg = self._modules['m' + str(i)](previous)
+
+            ll = hg
+            ll = self._modules['top_m_' + str(i)](ll)
+
+            ll = F.relu(
+                self._modules['bn_end' + str(i)](
+                    self._modules['conv_last' + str(i)](ll)), True)
+
+            # Predict heatmaps
+            tmp_out = self._modules['l' + str(i)](ll)
+            outputs.append(tmp_out)
+
+            if i < self.num_modules - 1:
+                ll = self._modules['bl' + str(i)](ll)
+                tmp_out_ = self._modules['al' + str(i)](tmp_out)
+                previous = previous + ll + tmp_out_
+
+        return outputs