init

.gitignore

@@ -0,0 +1,18 @@
+__pycache__
+debug/
+log/
+results/*
+results
+.vscode
+!.gitignore
+.idea
+cluster/
+cluster
+*.zip
+data/
+data
+wandb
+build
+dist
+*egg-info
+*.so

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 ECON 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{xiu2022econ,
+  title={{ECON}: {E}xplicit {C}lothed humans {O}btained from {N}ormals},
+  author={Xiu, Yuliang and Yang, Jinlong and Cao, Xu and Tzionas, Dimitrios and Black, Michael J.},
+  booktitle=arXiv,
+  month = Dec,
+  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,165 @@
+<!-- PROJECT LOGO -->
+
+<p align="center">
+
+  <h1 align="center">ECON: Explicit Clothed humans Obtained from Normals</h1>
+  <p align="center">
+    <a href="https://ps.is.tuebingen.mpg.de/person/yxiu"><strong>Yuliang Xiu</strong></a>
+    ·
+    <a href="https://ps.is.tuebingen.mpg.de/person/jyang"><strong>Jinlong Yang</strong></a>
+    ·
+    <a href="https://hoshino042.github.io/homepage/"><strong>Xu Cao</strong></a>
+    ·
+    <a href="https://ps.is.mpg.de/~dtzionas"><strong>Dimitrios Tzionas</strong></a>
+    ·
+    <a href="https://ps.is.tuebingen.mpg.de/person/black"><strong>Michael J. Black</strong></a>
+  </p>
+  <h2 align="center">arXiv 2022</h2>
+  <div align="center">
+    <img src="./assets/teaser.gif" alt="Logo" width="100%">
+  </div>
+
+  <p align="center">
+  <br>
+    <a href="https://pytorch.org/get-started/locally/"><img alt="PyTorch" src="https://img.shields.io/badge/PyTorch-ee4c2c?logo=pytorch&logoColor=white"></a>
+    <a href="https://pytorchlightning.ai/"><img alt="Lightning" src="https://img.shields.io/badge/-Lightning-792ee5?logo=pytorchlightning&logoColor=white"></a>
+    <br></br>
+    <a href=''>
+      <img src='https://img.shields.io/badge/Paper-PDF-green?style=for-the-badge&logo=arXiv&logoColor=green' alt='Paper PDF'>
+    </a>
+    <a href="https://discord.gg/Vqa7KBGRyk"><img src="https://img.shields.io/discord/940240966844035082?color=7289DA&labelColor=4a64bd&logo=discord&logoColor=white&style=for-the-badge"></a>
+    <a href="https://youtu.be/j5hw4tsWpoY"><img alt="youtube views" title="Subscribe to my YouTube channel" src="https://img.shields.io/youtube/views/j5hw4tsWpoY?logo=youtube&labelColor=ce4630&style=for-the-badge"/></a>
+  </p>
+</p>
+
+<br/>
+
+ECON is designed for **"Human digitization from a color image"**, which combines the best properties of implicit and explicit representations, to infer high-fidelity 3D clothed humans from in-the-wild images, even with **loose clothing** or in **challenging poses**. ECON also supports batch reconstruction from **multi-person** photos.
+<br/>
+<br/>
+
+## News :triangular_flag_on_post:
+- [2022/03/05] <a href="">arXiv</a> and <a href="#demo">demo</a> are available.
+
+## TODO
+- [ ] Blender add-on for FBX export
+- [ ] Full RGB texture generation
+
+<br>
+
+<!-- TABLE OF CONTENTS -->
+<details open="open" style='padding: 10px; border-radius:5px 30px 30px 5px; border-style: solid; border-width: 1px;'>
+  <summary>Table of Contents</summary>
+  <ol>
+    <li>
+      <a href="#instructions">Instructions</a>
+    </li>
+    <li>
+      <a href="#demo">Demo</a>
+    </li>
+    <li>
+      <a href="#tricks">Tricks</a>
+    </li>
+    <li>
+      <a href="#citation">Citation</a>
+    </li>
+  </ol>
+</details>
+
+<br/>
+
+## Instructions
+
+- See [docs/installation.md](docs/installation.md) to install all the required packages and setup the models
+
+
+## Demo
+
+```bash
+# For image-based reconstruction
+python -m apps.infer -cfg ./configs/econ.yaml -in_dir ./examples -out_dir ./results
+
+# For video rendering
+python -m apps.multi_render -n {filename}
+```
+
+## Tricks
+### Some adjustable parameters in *config/econ.yaml*
+- `use_ifnet`
+  - True: use IF-Nets+ for mesh completion ( $\text{ECON}_\text{IF}$ - Better quality)
+  - False: use SMPL-X for mesh completion ( $\text{ECON}_\text{EX}$ - Faster speed)
+- `use_smpl`
+  - [ ]: don't use either hands or face parts from SMPL-X
+  - ["hand"]: only use the **visible** hands from SMPL-X
+  - ["hand", "face"]: use both **visible** hands and face from SMPL-X
+- `thickness` (default 2cm)
+  - could be increased accordingly in case **xx_full.obj** looks flat
+- `hps_type`
+  - "pixie": more accurate for face and hands
+  - "pymafx": more robust for challenging poses
+
+<br/>
+
+## More Qualitative Results
+
+|![OOD Poses](assets/OOD-poses.jpg)|
+| :----------------------: |
+|_Challenging Poses_|
+|![OOD Clothes](assets/OOD-outfits.jpg)|
+|_Loose Clothes_|
+|![SHHQ](assets/SHHQ.gif)|
+|_ECON Results on [SHHQ Dataset](https://github.com/stylegan-human/StyleGAN-Human)_|
+|![crowd](assets/crowd.gif)|
+|_ECON Results on Multi-Person Image_|
+
+
+<br/>
+<br/>
+
+## Citation
+
+```bibtex
+@inproceedings{xiu2022econ,
+  title     = {{ECON}: {E}xplicit {C}lothed humans {O}btained from {N}ormals},
+  author    = {Xiu, Yuliang and Yang, Jinlong and Cao, Xu and Tzionas, Dimitrios and Black, Michael J.},
+  booktitle = arXiv,
+  month     = {Dec},
+  year      = {2022},
+}
+```
+<br/>
+
+## Acknowledgments
+
+We thank [Lea Hering](https://is.mpg.de/person/lhering) and [Radek Daněček](https://is.mpg.de/person/rdanecek) for proof reading, [Yao Feng](https://ps.is.mpg.de/person/yfeng), [Haven Feng](https://is.mpg.de/person/hfeng), and [Weiyang Liu](https://wyliu.com/) for their feedback and discussions, [Tsvetelina Alexiadis](https://ps.is.mpg.de/person/talexiadis) for her help with the AMT perceptual study.
+
+Here are some great resources we benefit from:
+
+- [ICON](https://github.com/YuliangXiu/ICON) for Body Fitting
+- [MonoPortDataset](https://github.com/Project-Splinter/MonoPortDataset) for Data Processing
+- [rembg](https://github.com/danielgatis/rembg) for Human Segmentation
+- [smplx](https://github.com/vchoutas/smplx), [PyMAF-X](https://www.liuyebin.com/pymaf-x/), [PIXIE](https://github.com/YadiraF/PIXIE) for Human Pose & Shape Estimation
+- [CAPE](https://github.com/qianlim/CAPE) and [THuman](https://github.com/ZhengZerong/DeepHuman/tree/master/THUmanDataset) for Dataset
+- [PyTorch3D](https://github.com/facebookresearch/pytorch3d) for Differential Rendering
+
+Some images used in the qualitative examples come from [pinterest.com](https://www.pinterest.com/).
+
+This project has received funding from the European Union’s Horizon 2020 research and innovation programme under the Marie Skłodowska-Curie grant agreement No.860768 ([CLIPE Project](https://www.clipe-itn.eu)).
+
+--------------
+
+<br>
+
+## License
+
+This code and model are available for non-commercial scientific research purposes as defined in the [LICENSE](LICENSE) file. By downloading and using the code and model you agree to the terms in the [LICENSE](LICENSE).
+
+## Disclosure
+
+MJB has received research gift funds from Adobe, Intel, Nvidia, Meta/Facebook, and Amazon.  MJB has financial interests in Amazon, Datagen Technologies, and Meshcapade GmbH.
+
+## Contact
+
+For technical questions, please contact [email protected]
+
+For commercial licensing, please contact [email protected] and [email protected]

apps/IFGeo.py

@@ -0,0 +1,186 @@
+# -*- 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.common.seg3d_lossless import Seg3dLossless
+from lib.common.train_util import *
+import torch
+import numpy as np
+import pytorch_lightning as pl
+
+torch.backends.cudnn.benchmark = True
+
+
+class IFGeo(pl.LightningModule):
+
+    def __init__(self, cfg):
+        super(IFGeo, 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.mcube_res = cfg.mcube_res
+        self.clean_mesh_flag = cfg.clean_mesh
+        self.overfit = cfg.overfit
+
+        if cfg.dataset.prior_type == "SMPL":
+            from lib.net.IFGeoNet import IFGeoNet
+            self.netG = IFGeoNet(cfg)
+        else:
+            from lib.net.IFGeoNet_nobody import IFGeoNet
+            self.netG = IFGeoNet(cfg)
+
+
+        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.reconEngine = Seg3dLossless(
+            query_func=query_func_IF,
+            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,
+            visualize=False,
+            debug=False,
+            use_cuda_impl=False,
+            faster=True,
+        )
+
+        self.export_dir = None
+        self.result_eval = {}
+
+    # Training related
+    def configure_optimizers(self):
+
+        # set optimizer
+        weight_decay = self.cfg.weight_decay
+        momentum = self.cfg.momentum
+
+        optim_params_G = [{"params": self.netG.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):
+
+        # cfg log
+        if self.cfg.devices == 1:
+            if not self.cfg.fast_dev and self.global_step == 0:
+                export_cfg(self.logger, osp.join(self.cfg.results_path, self.cfg.name), self.cfg)
+                self.logger.experiment.config.update(convert_to_dict(self.cfg))
+
+        self.netG.train()
+
+        preds_G = self.netG(batch)
+        error_G = self.netG.compute_loss(preds_G, batch["labels_geo"])
+
+        # metrics processing
+        metrics_log = {
+            "loss": error_G,
+        }
+
+        self.log_dict(metrics_log,
+                      prog_bar=True,
+                      logger=True,
+                      on_step=True,
+                      on_epoch=False,
+                      sync_dist=True)
+
+        return metrics_log
+
+    def training_epoch_end(self, outputs):
+
+        # metrics processing
+        metrics_log = {
+            "train/avgloss": batch_mean(outputs, "loss"),
+        }
+
+        self.log_dict(metrics_log,
+                      prog_bar=False,
+                      logger=True,
+                      on_step=False,
+                      on_epoch=True,
+                      rank_zero_only=True)
+
+    def validation_step(self, batch, batch_idx):
+
+        self.netG.eval()
+        self.netG.training = False
+
+        preds_G = self.netG(batch)
+        error_G = self.netG.compute_loss(preds_G, batch["labels_geo"])
+
+        metrics_log = {
+            "val/loss": error_G,
+        }
+
+        self.log_dict(metrics_log,
+                      prog_bar=True,
+                      logger=False,
+                      on_step=True,
+                      on_epoch=False,
+                      sync_dist=True)
+
+        return metrics_log
+
+    def validation_epoch_end(self, outputs):
+
+        # metrics processing
+        metrics_log = {
+            "val/avgloss": batch_mean(outputs, "val/loss"),
+        }
+
+        self.log_dict(metrics_log,
+                      prog_bar=False,
+                      logger=True,
+                      on_step=False,
+                      on_epoch=True,
+                      rank_zero_only=True)

apps/Normal.py

@@ -0,0 +1,220 @@
+from lib.net import NormalNet
+from lib.common.train_util import convert_to_dict, export_cfg, batch_mean
+import torch
+import numpy as np
+import os.path as osp
+from skimage.transform import resize
+import pytorch_lightning as pl
+
+
+class Normal(pl.LightningModule):
+
+    def __init__(self, cfg):
+        super(Normal, self).__init__()
+        self.cfg = cfg
+        self.batch_size = self.cfg.batch_size
+        self.lr_F = self.cfg.lr_netF
+        self.lr_B = self.cfg.lr_netB
+        self.lr_D = self.cfg.lr_netD
+        self.overfit = cfg.overfit
+
+        self.F_losses = [item[0] for item in self.cfg.net.front_losses]
+        self.B_losses = [item[0] for item in self.cfg.net.back_losses]
+        self.ALL_losses = self.F_losses + self.B_losses
+
+        self.automatic_optimization = False
+
+        self.schedulers = []
+
+        self.netG = NormalNet(self.cfg)
+
+        self.in_nml = [item[0] for item in cfg.net.in_nml]
+
+    # Training related
+    def configure_optimizers(self):
+
+        optim_params_N_D = None
+        optimizer_N_D = None
+        scheduler_N_D = None
+
+        # set optimizer
+        optim_params_N_F = [{"params": self.netG.netF.parameters(), "lr": self.lr_F}]
+        optim_params_N_B = [{"params": self.netG.netB.parameters(), "lr": self.lr_B}]
+
+        optimizer_N_F = torch.optim.Adam(optim_params_N_F, lr=self.lr_F, betas=(0.5, 0.999))
+        optimizer_N_B = torch.optim.Adam(optim_params_N_B, lr=self.lr_B, betas=(0.5, 0.999))
+
+        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)
+        if 'gan' in self.ALL_losses:
+            optim_params_N_D = [{"params": self.netG.netD.parameters(), "lr": self.lr_D}]
+            optimizer_N_D = torch.optim.Adam(optim_params_N_D, lr=self.lr_D, betas=(0.5, 0.999))
+            scheduler_N_D = torch.optim.lr_scheduler.MultiStepLR(optimizer_N_D,
+                                                                 milestones=self.cfg.schedule,
+                                                                 gamma=self.cfg.gamma)
+            self.schedulers = [scheduler_N_F, scheduler_N_B, scheduler_N_D]
+            optims = [optimizer_N_F, optimizer_N_B, optimizer_N_D]
+
+        else:
+            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, dataset, idx):
+
+        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,
+                ))
+
+        self.logger.log_image(key=f"Normal/{dataset}/{idx if not self.overfit else 1}",
+                              images=[(np.concatenate(result_list, axis=1) * 255.0).astype(np.uint8)
+                                     ])
+
+    def training_step(self, batch, batch_idx):
+
+        # cfg log
+        if not self.cfg.fast_dev and self.global_step == 0 and self.cfg.devices == 1:
+            export_cfg(self.logger, osp.join(self.cfg.results_path, self.cfg.name), self.cfg)
+            self.logger.experiment.config.update(convert_to_dict(self.cfg))
+
+        self.netG.train()
+
+        # 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"]}
+
+        in_tensor.update(FB_tensor)
+
+        preds_F, preds_B = self.netG(in_tensor)
+        error_dict = self.netG.get_norm_error(preds_F, preds_B, FB_tensor)
+
+        if 'gan' in self.ALL_losses:
+            (opt_F, opt_B, opt_D) = self.optimizers()
+            opt_F.zero_grad()
+            self.manual_backward(error_dict["netF"])
+            opt_B.zero_grad()
+            self.manual_backward(error_dict["netB"], retain_graph=True)
+            opt_D.zero_grad()
+            self.manual_backward(error_dict["netD"])
+            opt_F.step()
+            opt_B.step()
+            opt_D.step()
+        else:
+            (opt_F, opt_B) = self.optimizers()
+            opt_F.zero_grad()
+            self.manual_backward(error_dict["netF"])
+            opt_B.zero_grad()
+            self.manual_backward(error_dict["netB"])
+            opt_F.step()
+            opt_B.step()
+
+        if batch_idx > 0 and batch_idx % int(
+                self.cfg.freq_show_train) == 0 and self.cfg.devices == 1:
+
+            self.netG.eval()
+            with torch.no_grad():
+                nmlF, nmlB = self.netG(in_tensor)
+                in_tensor.update({"nmlF": nmlF, "nmlB": nmlB})
+                self.render_func(in_tensor, "train", self.global_step)
+
+        # metrics processing
+        metrics_log = {"loss": error_dict["netF"] + error_dict["netB"]}
+
+        if "gan" in self.ALL_losses:
+            metrics_log["loss"] += error_dict["netD"]
+
+        for key in error_dict.keys():
+            metrics_log["train/loss_" + key] = error_dict[key].item()
+
+        self.log_dict(metrics_log,
+                      prog_bar=True,
+                      logger=True,
+                      on_step=True,
+                      on_epoch=False,
+                      sync_dist=True)
+
+        return metrics_log
+
+    def training_epoch_end(self, outputs):
+
+        # metrics processing
+        metrics_log = {}
+        for key in outputs[0].keys():
+            if "/" in key:
+                [stage, loss_name] = key.split("/")
+            else:
+                stage = "train"
+                loss_name = key
+            metrics_log[f"{stage}/avg-{loss_name}"] = batch_mean(outputs, key)
+
+        self.log_dict(metrics_log,
+                      prog_bar=False,
+                      logger=True,
+                      on_step=False,
+                      on_epoch=True,
+                      rank_zero_only=True)
+
+    def validation_step(self, batch, batch_idx):
+
+        self.netG.eval()
+        self.netG.training = False
+
+        # 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"]}
+        in_tensor.update(FB_tensor)
+
+        preds_F, preds_B = self.netG(in_tensor)
+        error_dict = self.netG.get_norm_error(preds_F, preds_B, FB_tensor)
+
+        if batch_idx % int(self.cfg.freq_show_train) == 0 and self.cfg.devices == 1:
+
+            with torch.no_grad():
+                nmlF, nmlB = self.netG(in_tensor)
+                in_tensor.update({"nmlF": nmlF, "nmlB": nmlB})
+                self.render_func(in_tensor, "val", batch_idx)
+
+        # metrics processing
+        metrics_log = {"val/loss": error_dict["netF"] + error_dict["netB"]}
+
+        if "gan" in self.ALL_losses:
+            metrics_log["val/loss"] += error_dict["netD"]
+
+        for key in error_dict.keys():
+            metrics_log["val/" + key] = error_dict[key].item()
+
+        return metrics_log
+
+    def validation_epoch_end(self, outputs):
+
+        # metrics processing
+        metrics_log = {}
+        for key in outputs[0].keys():
+            [stage, loss_name] = key.split("/")
+            metrics_log[f"{stage}/avg-{loss_name}"] = batch_mean(outputs, key)
+
+        self.log_dict(metrics_log,
+                      prog_bar=False,
+                      logger=True,
+                      on_step=False,
+                      on_epoch=True,
+                      rank_zero_only=True)

apps/avatarizer.py

@@ -0,0 +1,47 @@
+import numpy as np
+import trimesh
+import torch
+import os.path as osp
+import lib.smplx as smplx
+from lib.dataset.mesh_util import SMPLX
+
+smplx_container = SMPLX()
+
+smpl_npy = "./results/github/econ/obj/304e9c4798a8c3967de7c74c24ef2e38_smpl_00.npy"
+smplx_param = np.load(smpl_npy, allow_pickle=True).item()
+
+for key in smplx_param.keys():
+    smplx_param[key] = smplx_param[key].cpu().view(1, -1)
+    # print(key, smplx_param[key].device, smplx_param[key].shape)
+
+smpl_model = smplx.create(
+    smplx_container.model_dir,
+    model_type="smplx",
+    gender="neutral",
+    age="adult",
+    use_face_contour=False,
+    use_pca=False,
+    num_betas=200,
+    num_expression_coeffs=50,
+    ext='pkl')
+
+smpl_out = smpl_model(
+    body_pose=smplx_param["body_pose"],
+    global_orient=smplx_param["global_orient"],
+    # transl=smplx_param["transl"],
+    betas=smplx_param["betas"],
+    expression=smplx_param["expression"],
+    jaw_pose=smplx_param["jaw_pose"],
+    left_hand_pose=smplx_param["left_hand_pose"],
+    right_hand_pose=smplx_param["right_hand_pose"],
+    return_verts=True,
+    return_joint_transformation=True,
+    return_vertex_transformation=True)
+
+smpl_verts = smpl_out.vertices.detach()[0]
+inv_mat = torch.inverse(smpl_out.vertex_transformation.detach()[0])
+homo_coord = torch.ones_like(smpl_verts)[..., :1]
+smpl_verts = inv_mat @ torch.cat([smpl_verts, homo_coord], dim=1).unsqueeze(-1)
+smpl_verts = smpl_verts[:, :3, 0].cpu()
+
+trimesh.Trimesh(smpl_verts, smpl_model.faces).show()

apps/infer.py

@@ -0,0 +1,528 @@
+# -*- 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 warnings
+import logging
+
+warnings.filterwarnings("ignore")
+logging.getLogger("lightning").setLevel(logging.ERROR)
+logging.getLogger("trimesh").setLevel(logging.ERROR)
+
+import torch, torchvision
+import trimesh
+import numpy as np
+import argparse
+import os
+
+from termcolor import colored
+from tqdm.auto import tqdm
+from apps.Normal import Normal
+from apps.IFGeo import IFGeo
+from lib.common.config import cfg
+from lib.common.train_util import init_loss, load_normal_networks, load_networks
+from lib.common.BNI import BNI
+from lib.common.BNI_utils import save_normal_tensor
+from lib.dataset.TestDataset import TestDataset
+from lib.net.geometry import rot6d_to_rotmat, rotation_matrix_to_angle_axis
+from lib.dataset.mesh_util import *
+from lib.common.voxelize import VoxelGrid
+
+torch.backends.cudnn.benchmark = True
+
+if __name__ == "__main__":
+
+    # loading cfg file
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("-gpu", "--gpu_device", type=int, default=0)
+    parser.add_argument("-loop_smpl", "--loop_smpl", type=int, default=50)
+    parser.add_argument("-patience", "--patience", type=int, default=5)
+    parser.add_argument("-vis_freq", "--vis_freq", type=int, default=1000)
+    parser.add_argument("-multi", action="store_false")
+    parser.add_argument("-in_dir", "--in_dir", type=str, default="./examples")
+    parser.add_argument("-out_dir", "--out_dir", type=str, default="./results")
+    parser.add_argument("-seg_dir", "--seg_dir", type=str, default=None)
+    parser.add_argument("-cfg", "--config", type=str, default="./configs/econ.yaml")
+
+    args = parser.parse_args()
+
+    # cfg read and merge
+    cfg.merge_from_file(args.config)
+    cfg.merge_from_file("./lib/pymafx/configs/pymafx_config.yaml")
+    device = torch.device(f"cuda:{args.gpu_device}")
+
+    # setting for testing on in-the-wild images
+    cfg_show_list = ["test_gpus", [args.gpu_device], "mcube_res", 512, "clean_mesh", True, "test_mode", True, "batch_size", 1]
+
+    cfg.merge_from_list(cfg_show_list)
+    cfg.freeze()
+
+    # load model
+    normal_model = Normal(cfg).to(device)
+    load_normal_networks(normal_model, cfg.normal_path)
+    normal_model.netG.eval()
+
+    # load IFGeo model
+    ifnet_model = IFGeo(cfg).to(device)
+    load_networks(ifnet_model, mlp_path=cfg.ifnet_path)
+    ifnet_model.netG.eval()
+
+    # SMPLX object
+    SMPLX_object = SMPLX()
+
+    dataset_param = {
+        "image_dir": args.in_dir,
+        "seg_dir": args.seg_dir,
+        "use_seg": True,  # w/ or w/o segmentation
+        "hps_type": cfg.bni.hps_type,  # pymafx/pixie
+        "vol_res": cfg.vol_res,
+        "single": args.multi,
+    }
+
+    dataset = TestDataset(dataset_param, device)
+
+    print(colored(f"Dataset Size: {len(dataset)}", "green"))
+
+    pbar = tqdm(dataset)
+
+    for data in pbar:
+
+        losses = init_loss()
+
+        pbar.set_description(f"{data['name']}")
+
+        # final results rendered as image (PNG)
+        # 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)
+        # 4. Blend the cropped image with predicted cloth normal (xxx_crop.png)
+
+        os.makedirs(osp.join(args.out_dir, cfg.name, "png"), exist_ok=True)
+
+        # final reconstruction meshes (OBJ)
+        # 1. SMPL mesh (xxx_smpl_xx.obj)
+        # 2. SMPL params (xxx_smpl.npy)
+        # 3. d-BiNI surfaces (xxx_BNI.obj)
+        # 4. seperate face/hand mesh (xxx_hand/face.obj)
+        # 5. full shape impainted by IF-Nets+, and remeshed shape (xxx_IF_(remesh).obj)
+        # 6. sideded or occluded parts (xxx_side.obj)
+        # 7. final reconstructed clothed human (xxx_full.obj)
+
+        os.makedirs(osp.join(args.out_dir, cfg.name, "obj"), exist_ok=True)
+
+        in_tensor = {
+            "smpl_faces": data["smpl_faces"],
+            "image": data["img_icon"].to(device),
+            "mask": data["img_mask"].to(device)
+        }
+
+        # The optimizer and variables
+        optimed_pose = data["body_pose"].requires_grad_(True)
+        optimed_trans = data["trans"].requires_grad_(True)
+        optimed_betas = data["betas"].requires_grad_(True)
+        optimed_orient = data["global_orient"].requires_grad_(True)
+
+        optimizer_smpl = torch.optim.Adam([optimed_pose, optimed_trans, optimed_betas, optimed_orient], lr=1e-2, amsgrad=True)
+        scheduler_smpl = torch.optim.lr_scheduler.ReduceLROnPlateau(
+            optimizer_smpl,
+            mode="min",
+            factor=0.5,
+            verbose=0,
+            min_lr=1e-5,
+            patience=args.patience,
+        )
+
+        # [result_loop_1, result_loop_2, ...]
+        per_data_lst = []
+
+        N_body, N_pose = optimed_pose.shape[:2]
+
+        smpl_path = osp.join(args.out_dir, "econ", f"png/{data['name']}_smpl.png")
+
+        if osp.exists(smpl_path):
+
+            smpl_verts_lst = []
+            smpl_faces_lst = []
+            for idx in range(N_body):
+
+                smpl_obj = f"{args.out_dir}/{cfg.name}/obj/{data['name']}_smpl_{idx:02d}.obj"
+                smpl_mesh = trimesh.load(smpl_obj)
+                smpl_verts = torch.tensor(smpl_mesh.vertices).to(device).float()
+                smpl_faces = torch.tensor(smpl_mesh.faces).to(device).long()
+                smpl_verts_lst.append(smpl_verts)
+                smpl_faces_lst.append(smpl_faces)
+
+            batch_smpl_verts = torch.stack(smpl_verts_lst)
+            batch_smpl_faces = torch.stack(smpl_faces_lst)
+
+            # render optimized mesh as normal [-1,1]
+            in_tensor["T_normal_F"], in_tensor["T_normal_B"] = dataset.render_normal(batch_smpl_verts, batch_smpl_faces)
+
+            with torch.no_grad():
+                in_tensor["normal_F"], in_tensor["normal_B"] = normal_model.netG(in_tensor)
+
+            in_tensor["smpl_verts"] = batch_smpl_verts * torch.tensor([1., -1., 1.]).to(device)
+            in_tensor["smpl_faces"] = batch_smpl_faces[:, :, [0, 2, 1]]
+        else:
+            # smpl optimization
+            loop_smpl = tqdm(range(args.loop_smpl))
+
+            for i in loop_smpl:
+
+                per_loop_lst = []
+
+                optimizer_smpl.zero_grad()
+
+                N_body, N_pose = optimed_pose.shape[:2]
+
+                # 6d_rot to rot_mat
+                optimed_orient_mat = rot6d_to_rotmat(optimed_orient.view(-1, 6)).view(N_body, 1, 3, 3)
+                optimed_pose_mat = rot6d_to_rotmat(optimed_pose.view(-1, 6)).view(N_body, N_pose, 3, 3)
+
+                smpl_verts, smpl_landmarks, smpl_joints = dataset.smpl_model(
+                    shape_params=optimed_betas,
+                    expression_params=tensor2variable(data["exp"], device),
+                    body_pose=optimed_pose_mat,
+                    global_pose=optimed_orient_mat,
+                    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"]
+                smpl_joints = (smpl_joints + optimed_trans) * data["scale"] * torch.tensor([1.0, 1.0, -1.0]).to(device)
+
+                # landmark errors
+                smpl_joints_3d = (smpl_joints[:, dataset.smpl_data.smpl_joint_ids_45_pixie, :] + 1.0) * 0.5
+                in_tensor["smpl_joint"] = smpl_joints[:, dataset.smpl_data.smpl_joint_ids_24_pixie, :]
+
+                ghum_lmks = data["landmark"][:, SMPLX_object.ghum_smpl_pairs[:, 0], :2].to(device)
+                ghum_conf = data["landmark"][:, SMPLX_object.ghum_smpl_pairs[:, 0], -1].to(device)
+                smpl_lmks = smpl_joints_3d[:, SMPLX_object.ghum_smpl_pairs[:, 1], :2]
+
+                # render optimized mesh as normal [-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_image(type="mask")
+
+                with torch.no_grad():
+                    in_tensor["normal_F"], in_tensor["normal_B"] = normal_model.netG(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"])
+
+                # silhouette loss
+                smpl_arr = torch.cat([T_mask_F, T_mask_B], dim=-1)
+                gt_arr = in_tensor["mask"].repeat(1, 1, 2)
+                diff_S = torch.abs(smpl_arr - gt_arr)
+                losses["silhouette"]["value"] = diff_S.mean()
+
+                # large cloth_overlap --> big difference between body and cloth mask
+                # for loose clothing, reply more on landmarks instead of silhouette+normal loss
+                cloth_overlap = diff_S.sum(dim=[1, 2]) / gt_arr.sum(dim=[1, 2])
+                cloth_overlap_flag = cloth_overlap > cfg.cloth_overlap_thres
+                losses["joint"]["weight"] = [50.0 if flag else 5.0 for flag in cloth_overlap_flag]
+
+                # small body_overlap --> large occlusion or out-of-frame
+                # for highly occluded body, reply only on high-confidence landmarks, no silhouette+normal loss
+
+                # BUG: PyTorch3D silhouette renderer generates dilated mask
+                bg_value = in_tensor["T_normal_F"][0, 0, 0, 0]
+                smpl_arr_fake = torch.cat(
+                    [in_tensor["T_normal_F"][:, 0].ne(bg_value).float(), in_tensor["T_normal_B"][:, 0].ne(bg_value).float()],
+                    dim=-1)
+
+                body_overlap = (gt_arr * smpl_arr_fake.gt(0.0)).sum(dim=[1, 2]) / smpl_arr_fake.gt(0.0).sum(dim=[1, 2])
+                body_overlap_mask = (gt_arr * smpl_arr_fake).unsqueeze(1)
+                body_overlap_flag = body_overlap < cfg.body_overlap_thres
+
+                losses["normal"]["value"] = (diff_F_smpl * body_overlap_mask[..., :512] +
+                                             diff_B_smpl * body_overlap_mask[..., 512:]).mean() / 2.0
+
+                losses["silhouette"]["weight"] = [0 if flag else 1.0 for flag in body_overlap_flag]
+                occluded_idx = torch.where(body_overlap_flag)[0]
+                ghum_conf[occluded_idx] *= ghum_conf[occluded_idx] > 0.95
+                losses["joint"]["value"] = (torch.norm(ghum_lmks - smpl_lmks, dim=2) * ghum_conf).mean(dim=1)
+
+                # Weighted sum of the losses
+                smpl_loss = 0.0
+                pbar_desc = "Body Fitting --- "
+                for k in ["normal", "silhouette", "joint"]:
+                    per_loop_loss = (losses[k]["value"] * torch.tensor(losses[k]["weight"]).to(device)).mean()
+                    pbar_desc += f"{k}: {per_loop_loss:.3f} | "
+                    smpl_loss += per_loop_loss
+                pbar_desc += f"Total: {smpl_loss:.3f}"
+                loose_str = ''.join([str(j) for j in cloth_overlap_flag.int().tolist()])
+                occlude_str = ''.join([str(j) for j in body_overlap_flag.int().tolist()])
+                pbar_desc += colored(f"| loose:{loose_str}, occluded:{occlude_str}", "yellow")
+                loop_smpl.set_description(pbar_desc)
+
+                # save intermediate results / vis_freq and final_step
+                if (i % args.vis_freq == 0) or (i == args.loop_smpl - 1):
+
+                    per_loop_lst.extend([
+                        in_tensor["image"],
+                        in_tensor["T_normal_F"],
+                        in_tensor["normal_F"],
+                        diff_S[:, :, :512].unsqueeze(1).repeat(1, 3, 1, 1),
+                    ])
+                    per_loop_lst.extend([
+                        in_tensor["image"],
+                        in_tensor["T_normal_B"],
+                        in_tensor["normal_B"],
+                        diff_S[:, :, 512:].unsqueeze(1).repeat(1, 3, 1, 1),
+                    ])
+                    per_data_lst.append(get_optim_grid_image(per_loop_lst, None, nrow=N_body * 2, type="smpl"))
+
+                smpl_loss.backward()
+                optimizer_smpl.step()
+                scheduler_smpl.step(smpl_loss)
+
+            in_tensor["smpl_verts"] = smpl_verts * torch.tensor([1.0, 1.0, -1.0]).to(device)
+            in_tensor["smpl_faces"] = in_tensor["smpl_faces"][:, :, [0, 2, 1]]
+            per_data_lst[-1].save(osp.join(args.out_dir, cfg.name, f"png/{data['name']}_smpl.png"))
+
+        img_crop_path = osp.join(args.out_dir, cfg.name, "png", f"{data['name']}_crop.png")
+        torchvision.utils.save_image(
+            torch.cat([
+                data["img_crop"][:, :3], (in_tensor['normal_F'].detach().cpu() + 1.0) * 0.5,
+                (in_tensor['normal_B'].detach().cpu() + 1.0) * 0.5
+            ],
+                      dim=3), img_crop_path)
+
+        rgb_norm_F = blend_rgb_norm(in_tensor["normal_F"], data)
+        rgb_norm_B = blend_rgb_norm(in_tensor["normal_B"], data)
+
+        img_overlap_path = osp.join(args.out_dir, cfg.name, f"png/{data['name']}_overlap.png")
+        torchvision.utils.save_image(
+            torch.Tensor([data["img_raw"], rgb_norm_F, rgb_norm_B]).permute(0, 3, 1, 2) / 255., img_overlap_path)
+
+        smpl_obj_lst = []
+
+        for idx in range(N_body):
+
+            smpl_obj = trimesh.Trimesh(
+                in_tensor["smpl_verts"].detach().cpu()[idx] * torch.tensor([1.0, -1.0, 1.0]),
+                in_tensor["smpl_faces"].detach().cpu()[0][:, [0, 2, 1]],
+                process=False,
+                maintains_order=True,
+            )
+
+            smpl_obj_path = f"{args.out_dir}/{cfg.name}/obj/{data['name']}_smpl_{idx:02d}.obj"
+
+            if not osp.exists(smpl_obj_path):
+                smpl_obj.export(smpl_obj_path)
+                smpl_info = {
+                    "betas": optimed_betas[idx].detach().cpu().unsqueeze(0),
+                    "body_pose": rotation_matrix_to_angle_axis(optimed_pose_mat[idx].detach()).cpu().unsqueeze(0),
+                    "global_orient": rotation_matrix_to_angle_axis(optimed_orient_mat[idx].detach()).cpu().unsqueeze(0),
+                    "transl": optimed_trans[idx].detach().cpu(),
+                    "expression": data["exp"][idx].cpu().unsqueeze(0),
+                    "jaw_pose": rotation_matrix_to_angle_axis(data["jaw_pose"][idx]).cpu().unsqueeze(0),
+                    "left_hand_pose": rotation_matrix_to_angle_axis(data["left_hand_pose"][idx]).cpu().unsqueeze(0),
+                    "right_hand_pose": rotation_matrix_to_angle_axis(data["right_hand_pose"][idx]).cpu().unsqueeze(0),
+                    "scale": data["scale"][idx].cpu(),
+                }
+                np.save(
+                    smpl_obj_path.replace(".obj", ".npy"),
+                    smpl_info,
+                    allow_pickle=True,
+                )
+            smpl_obj_lst.append(smpl_obj)
+
+        del optimizer_smpl
+        del optimed_betas
+        del optimed_orient
+        del optimed_pose
+        del optimed_trans
+
+        torch.cuda.empty_cache()
+
+        # ------------------------------------------------------------------------------------------------------------------
+        # clothing refinement
+
+        per_data_lst = []
+
+        batch_smpl_verts = in_tensor["smpl_verts"].detach() * torch.tensor([1.0, -1.0, 1.0], device=device)
+        batch_smpl_faces = in_tensor["smpl_faces"].detach()[:, :, [0, 2, 1]]
+
+        in_tensor["depth_F"], in_tensor["depth_B"] = dataset.render_depth(batch_smpl_verts, batch_smpl_faces)
+
+        per_loop_lst = []
+
+        in_tensor["BNI_verts"] = []
+        in_tensor["BNI_faces"] = []
+        in_tensor["body_verts"] = []
+        in_tensor["body_faces"] = []
+
+        for idx in range(N_body):
+
+            final_path = f"{args.out_dir}/{cfg.name}/obj/{data['name']}_{idx}_full.obj"
+
+            side_mesh = smpl_obj_lst[idx].copy()
+            face_mesh = smpl_obj_lst[idx].copy()
+            hand_mesh = smpl_obj_lst[idx].copy()
+
+            # save normals, depths and masks
+            BNI_dict = save_normal_tensor(
+                in_tensor,
+                idx,
+                osp.join(args.out_dir, cfg.name, f"BNI/{data['name']}_{idx}"),
+                cfg.bni.thickness,
+            )
+
+            # BNI process
+            BNI_object = BNI(
+                dir_path=osp.join(args.out_dir, cfg.name, "BNI"),
+                name=data["name"],
+                BNI_dict=BNI_dict,
+                cfg=cfg.bni,
+                device=device)
+
+            BNI_object.extract_surface(False)
+
+            in_tensor["body_verts"].append(torch.tensor(smpl_obj_lst[idx].vertices).float())
+            in_tensor["body_faces"].append(torch.tensor(smpl_obj_lst[idx].faces).long())
+
+            # requires shape completion when low overlap
+            # replace SMPL by completed mesh as side_mesh
+
+            if cfg.bni.use_ifnet:
+                print(colored("Use IF-Nets+ for completion\n", "green"))
+
+                side_mesh_path = f"{args.out_dir}/{cfg.name}/obj/{data['name']}_{idx}_IF.obj"
+
+                side_mesh = apply_face_mask(side_mesh, ~SMPLX_object.smplx_eyeball_fid_mask)
+
+                # mesh completion via IF-net
+                in_tensor.update(
+                    dataset.depth_to_voxel({
+                        "depth_F": BNI_object.F_depth.unsqueeze(0),
+                        "depth_B": BNI_object.B_depth.unsqueeze(0)
+                    }))
+
+                occupancies = VoxelGrid.from_mesh(
+                    side_mesh, cfg.vol_res, loc=[
+                        0,
+                    ] * 3, scale=2.0).data.transpose(2, 1, 0)
+                occupancies = np.flip(occupancies, axis=1)
+
+                in_tensor["body_voxels"] = torch.tensor(occupancies.copy()).float().unsqueeze(0).to(device)
+
+                with torch.no_grad():
+                    sdf = ifnet_model.reconEngine(netG=ifnet_model.netG, batch=in_tensor)
+                    verts_IF, faces_IF = ifnet_model.reconEngine.export_mesh(sdf)
+
+                if ifnet_model.clean_mesh_flag:
+                    verts_IF, faces_IF = clean_mesh(verts_IF, faces_IF)
+
+                side_mesh = trimesh.Trimesh(verts_IF, faces_IF)
+                side_mesh = remesh(side_mesh, side_mesh_path)
+
+            else:
+                print(colored("Use SMPL-X body for completion\n", "green"))
+                side_mesh = apply_vertex_mask(
+                    side_mesh,
+                    (SMPLX_object.front_flame_vertex_mask + SMPLX_object.mano_vertex_mask +
+                     SMPLX_object.eyeball_vertex_mask).eq(0).float(),
+                )
+
+            side_verts = torch.tensor(side_mesh.vertices).float().to(device)
+            side_faces = torch.tensor(side_mesh.faces).long().to(device)
+
+            # Possion Fusion between SMPLX and BNI
+            # 1. keep the faces invisible to front+back cameras
+            # 2. keep the front-FLAME+MANO faces
+            # 3. remove eyeball faces
+
+            # export intermediate meshes
+            BNI_object.F_B_trimesh.export(f"{args.out_dir}/{cfg.name}/obj/{data['name']}_{idx}_BNI.obj")
+
+            full_lst = []
+
+            if "face" in cfg.bni.use_smpl:
+
+                # only face
+                face_mesh = apply_vertex_mask(face_mesh, SMPLX_object.front_flame_vertex_mask)
+                face_mesh.vertices = face_mesh.vertices - np.array([0, 0, cfg.bni.thickness])
+
+                # remove face neighbor triangles
+                BNI_object.F_B_trimesh = part_removal(
+                    BNI_object.F_B_trimesh, None, face_mesh, cfg.bni.face_thres, device, camera_ray=True)
+                side_mesh = part_removal(
+                    side_mesh, torch.zeros_like(side_verts[:, 0:1]), face_mesh, cfg.bni.face_thres, device, camera_ray=True)
+                face_mesh.export(f"{args.out_dir}/{cfg.name}/obj/{data['name']}_{idx}_face.obj")
+                full_lst += [face_mesh]
+
+            if "hand" in cfg.bni.use_smpl and (True in data['hands_visibility'][idx]):
+
+                hand_mask = torch.zeros(SMPLX_object.smplx_verts.shape[0],)
+                if data['hands_visibility'][idx][0]:
+                    hand_mask.index_fill_(0, torch.tensor(SMPLX_object.smplx_mano_vid_dict["left_hand"]), 1.0)
+                if data['hands_visibility'][idx][1]:
+                    hand_mask.index_fill_(0, torch.tensor(SMPLX_object.smplx_mano_vid_dict["right_hand"]), 1.0)
+
+                # only hands
+                hand_mesh = apply_vertex_mask(hand_mesh, hand_mask)
+                # remove face neighbor triangles
+                BNI_object.F_B_trimesh = part_removal(
+                    BNI_object.F_B_trimesh, None, hand_mesh, cfg.bni.hand_thres, device, camera_ray=True)
+                side_mesh = part_removal(
+                    side_mesh, torch.zeros_like(side_verts[:, 0:1]), hand_mesh, cfg.bni.hand_thres, device, camera_ray=True)
+                hand_mesh.export(f"{args.out_dir}/{cfg.name}/obj/{data['name']}_{idx}_hand.obj")
+                full_lst += [hand_mesh]
+
+            full_lst += [BNI_object.F_B_trimesh]
+
+            # initial side_mesh could be SMPLX or IF-net
+            side_mesh = part_removal(side_mesh, torch.zeros_like(side_verts[:, 0:1]), sum(full_lst), 2e-2, device, clean=False)
+
+            full_lst += [side_mesh]
+
+            # # export intermediate meshes
+            BNI_object.F_B_trimesh.export(f"{args.out_dir}/{cfg.name}/obj/{data['name']}_{idx}_BNI.obj")
+            side_mesh.export(f"{args.out_dir}/{cfg.name}/obj/{data['name']}_{idx}_side.obj")
+
+            if cfg.bni.use_poisson:
+                final_mesh = poisson(
+                    sum(full_lst),
+                    final_path,
+                    cfg.bni.poisson_depth,
+                )
+            else:
+                final_mesh = sum(full_lst)
+                final_mesh.export(final_path)
+
+            dataset.render.load_meshes(final_mesh.vertices, final_mesh.faces)
+            rotate_recon_lst = dataset.render.get_image(cam_type="four")
+            per_loop_lst.extend([in_tensor['image'][idx:idx + 1]] + rotate_recon_lst)
+
+            # for video rendering
+            in_tensor["BNI_verts"].append(torch.tensor(final_mesh.vertices).float())
+            in_tensor["BNI_faces"].append(torch.tensor(final_mesh.faces).long())
+
+        if len(per_loop_lst) > 0:
+
+            per_data_lst.append(get_optim_grid_image(per_loop_lst, None, nrow=5, type="cloth"))
+            per_data_lst[-1].save(osp.join(args.out_dir, cfg.name, f"png/{data['name']}_cloth.png"))
+
+        os.makedirs(osp.join(args.out_dir, cfg.name, "vid"), exist_ok=True)
+        in_tensor["uncrop_param"] = data["uncrop_param"]
+        in_tensor["img_raw"] = data["img_raw"]
+        torch.save(in_tensor, osp.join(args.out_dir, cfg.name, f"vid/{data['name']}_in_tensor.pt"))

apps/multi_render.py

@@ -0,0 +1,25 @@
+from lib.common.render import Render
+import torch
+import argparse
+
+root = "./results/econ/vid"
+
+# loading cfg file
+parser = argparse.ArgumentParser()
+parser.add_argument("-n", "--name", type=str, default="")
+parser.add_argument("-g", "--gpu", type=int, default=0)
+args = parser.parse_args()
+
+in_tensor = torch.load(f"{root}/{args.name}_in_tensor.pt")
+
+render = Render(size=512, device=torch.device(f"cuda:{args.gpu}"))
+
+# visualize the final results in self-rotation mode
+verts_lst = in_tensor["body_verts"] + in_tensor["BNI_verts"]
+faces_lst = in_tensor["body_faces"] + in_tensor["BNI_faces"]
+
+# self-rotated video
+render.load_meshes(verts_lst, faces_lst)
+render.get_rendered_video_multi(
+    in_tensor,
+    f"{root}/{args.name}_cloth.mp4")

configs/econ.yaml

@@ -0,0 +1,35 @@
+name: econ
+ckpt_dir: "./data/ckpt/"
+normal_path: "./data/ckpt/normal.ckpt"
+ifnet_path: "./data/ckpt/ifnet.ckpt"
+results_path: "./results"
+
+net:
+  in_nml: (('image',3), ('T_normal_F',3), ('T_normal_B',3))
+  in_geo: (('normal_F',3), ('normal_B',3))
+
+test_mode: True
+batch_size: 1
+
+dataset:
+  prior_type: "SMPL"
+
+# user defined
+vol_res: 256      # IF-Net volume resolution
+mcube_res: 256
+clean_mesh: True  # if True, will remove floating pieces
+cloth_overlap_thres: 0.50
+body_overlap_thres: 0.98
+
+bni:
+  k: 2
+  lambda1: 1e-4
+  boundary_consist: 1e-6
+  poisson_depth: 10
+  use_smpl: ["hand", "face"]
+  use_ifnet: True
+  use_poisson: True
+  hand_thres: 4e-2
+  face_thres: 6e-2
+  thickness: 0.02
+  hps_type: "pixie"

docs/installation.md

@@ -0,0 +1,77 @@
+## Getting started
+
+Start by cloning the repo:
+
+```bash
+git clone --depth 1 [email protected]:YuliangXiu/ECON.git
+cd ECON
+```
+
+## Environment
+
+- Ubuntu 20 / 18
+- **CUDA=11.4, GPU Memory > 12GB**
+- Python = 3.8
+- PyTorch >= 1.13.0 (official [Get Started](https://pytorch.org/get-started/locally/))
+- CUPY >= 11.3.0 (offcial [Installation](https://docs.cupy.dev/en/stable/install.html#installing-cupy-from-pypi))
+- PyTorch3D (official [INSTALL.md](https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md), recommend [install-from-local-clone](https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md#2-install-from-a-local-clone))
+
+```bash
+
+sudo apt-get install libeigen3-dev ffmpeg
+
+# install required packages
+cd ECON
+conda env create -f environment.yaml
+conda activate econ
+pip install -r requirements.txt
+
+# install libmesh & libvoxelize
+cd lib/commmon/libmesh
+python setup.py build_ext --inplace
+cd ../libvoxelize
+python setup.py build_ext --inplace
+```
+
+## Register at [ICON's website](https://icon.is.tue.mpg.de/)
+
+![Register](../assets/register.png)
+Required:
+
+- [SMPL](http://smpl.is.tue.mpg.de/): SMPL Model (Male, Female)
+- [SMPL-X](http://smpl-x.is.tue.mpg.de/): SMPL-X Model, used for training
+- [SMPLIFY](http://smplify.is.tue.mpg.de/): SMPL Model (Neutral)
+- [PIXIE](https://icon.is.tue.mpg.de/user.php): PIXIE SMPL-X estimator
+
+:warning: Click **Register now** on all dependencies, then you can download them all with **ONE** account.
+
+## Downloading required models and extra data
+
+```bash
+cd ECON
+bash fetch_data.sh # requires username and password
+```
+
+## Citation
+
+:+1: Please consider citing these awesome HPS approaches: PyMAF-X, PIXIE
+
+
+```
+@article{pymafx2022,
+  title={PyMAF-X: Towards Well-aligned Full-body Model Regression from Monocular Images},
+  author={Zhang, Hongwen and Tian, Yating and Zhang, Yuxiang and Li, Mengcheng and An, Liang and Sun, Zhenan and Liu, Yebin},
+  journal={arXiv preprint arXiv:2207.06400},
+  year={2022}
+}
+
+
+@inproceedings{PIXIE:2021,
+  title={Collaborative Regression of Expressive Bodies using Moderation},
+  author={Yao Feng and Vasileios Choutas and Timo Bolkart and Dimitrios Tzionas and Michael J. Black},
+  booktitle={International Conference on 3D Vision (3DV)},
+  year={2021}
+}
+
+
+```

environment.yaml

@@ -0,0 +1,20 @@
+name: econ
+channels:
+  - pytorch
+  - nvidia
+  - conda-forge
+  - fvcore
+  - iopath
+  - bottler
+  - defaults
+dependencies:
+  - python=3.8
+  - pytorch
+  - torchvision
+  - fvcore
+  - iopath
+  - nvidiacub
+  - pyembree
+  - cupy
+  - cython
+  - pip

fetch_data.sh

@@ -0,0 +1,60 @@
+#!/bin/bash
+urle () { [[ "${1}" ]] || return 1; local LANG=C i x; for (( i = 0; i < ${#1}; i++ )); do x="${1:i:1}"; [[ "${x}" == [a-zA-Z0-9.~-] ]] && echo -n "${x}" || printf '%%%02X' "'${x}"; done; echo; }
+
+mkdir -p data/smpl_related/models
+
+# username and password input
+echo -e "\nYou need to register at https://icon.is.tue.mpg.de/, according to Installation Instruction."
+read -p "Username (ICON):" username
+read -p "Password (ICON):" password
+username=$(urle $username)
+password=$(urle $password)
+
+# SMPL (Male, Female)
+echo -e "\nDownloading SMPL..."
+wget --post-data "username=$username&password=$password" 'https://download.is.tue.mpg.de/download.php?domain=smpl&sfile=SMPL_python_v.1.0.0.zip&resume=1' -O './data/smpl_related/models/SMPL_python_v.1.0.0.zip' --no-check-certificate --continue
+unzip data/smpl_related/models/SMPL_python_v.1.0.0.zip -d data/smpl_related/models
+mv data/smpl_related/models/smpl/models/basicModel_f_lbs_10_207_0_v1.0.0.pkl data/smpl_related/models/smpl/SMPL_FEMALE.pkl
+mv data/smpl_related/models/smpl/models/basicmodel_m_lbs_10_207_0_v1.0.0.pkl data/smpl_related/models/smpl/SMPL_MALE.pkl
+cd data/smpl_related/models
+rm -rf *.zip __MACOSX smpl/models smpl/smpl_webuser
+cd ../../..
+
+# SMPL (Neutral, from SMPLIFY)
+echo -e "\nDownloading SMPLify..."
+wget --post-data "username=$username&password=$password" 'https://download.is.tue.mpg.de/download.php?domain=smplify&sfile=mpips_smplify_public_v2.zip&resume=1' -O './data/smpl_related/models/mpips_smplify_public_v2.zip' --no-check-certificate --continue
+unzip data/smpl_related/models/mpips_smplify_public_v2.zip -d data/smpl_related/models
+mv data/smpl_related/models/smplify_public/code/models/basicModel_neutral_lbs_10_207_0_v1.0.0.pkl data/smpl_related/models/smpl/SMPL_NEUTRAL.pkl
+cd data/smpl_related/models
+rm -rf *.zip smplify_public 
+cd ../../..
+
+# SMPL-X
+echo -e "\nDownloading SMPL-X..."
+wget --post-data "username=$username&password=$password" 'https://download.is.tue.mpg.de/download.php?domain=smplx&sfile=models_smplx_v1_1.zip&resume=1' -O './data/smpl_related/models/models_smplx_v1_1.zip' --no-check-certificate --continue
+unzip data/smpl_related/models/models_smplx_v1_1.zip -d data/smpl_related
+rm -f data/smpl_related/models/models_smplx_v1_1.zip
+
+# ECON
+echo -e "\nDownloading ECON..."
+wget --post-data "username=$username&password=$password" 'https://download.is.tue.mpg.de/download.php?domain=icon&sfile=econ_data.zip&resume=1' -O './data/econ_data.zip' --no-check-certificate --continue
+cd data && unzip econ_data.zip
+mv smpl_data smpl_related/
+rm -f econ_data.zip
+cd ..
+
+mkdir -p data/HPS
+
+# PIXIE
+echo -e "\nDownloading PIXIE..."
+wget --post-data "username=$username&password=$password" 'https://download.is.tue.mpg.de/download.php?domain=icon&sfile=HPS/pixie_data.zip&resume=1' -O './data/HPS/pixie_data.zip' --no-check-certificate --continue
+cd data/HPS && unzip pixie_data.zip
+rm -f pixie_data.zip
+cd ../..
+
+# PyMAF-X
+echo -e "\nDownloading PyMAF-X..."
+wget --post-data "username=$username&password=$password" 'https://download.is.tue.mpg.de/download.php?domain=icon&sfile=HPS/pymafx_data.zip&resume=1' -O './data/HPS/pymafx_data.zip' --no-check-certificate --continue
+cd data/HPS && unzip pymafx_data.zip
+rm -f pymafx_data.zip
+cd ../..