Fully specified Python version

.gitattributes

@@ -29,9 +29,3 @@ saved_model/**/* 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

@@ -4,15 +4,14 @@ data/thuman*
 __pycache__
 debug/
 log/
+results/*
 .vscode
 !.gitignore
 force_push.sh
 .idea
+smplx/
 human_det/
 kaolin/
 neural_voxelization_layer/
 pytorch3d/
 force_push.sh
-results/
-gradio_cached_examples/
-gradio_queue.db

Dockerfile

@@ -1,107 +0,0 @@
-FROM nvidia/cuda:11.3.1-cudnn8-devel-ubuntu20.04
-
-ARG DEBIAN_FRONTEND=noninteractive
-
-RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
-    curl \
-    git \
-    wget \
-    freeglut3-dev \
-    unzip \
-    ffmpeg \
-    libsm6 \
-    libxext6 \
-    libgomp1 \
-    libfontconfig1 \
-    libgl1-mesa-glx \
-    libgl1-mesa-dev \
-    libglfw3 \
-    libglfw3-dev \
-    libglew2.1 \
-    libglew-dev \
-    mesa-utils \
-    libc6 \
-    libxdamage1 \
-    libxfixes3 \
-    libxcb-glx0 \
-    libxcb-dri2-0 \
-    libxcb-dri3-0 \
-    libxcb-present0 \
-    libxcb-sync1 \
-    libxshmfence1 \
-    libxxf86vm1 \
-    libxrender1 \
-    libgbm1 \
-    build-essential \
-    libeigen3-dev \
-    python3.8 \
-    python3-pip \
-    python-is-python3 \
-    nvidia-cuda-toolkit \
-    && rm -rf /var/lib/apt/lists/*
-
-
-# Set up a new user named "user" with user ID 1000
-RUN useradd -m -u 1000 user
-
-# Switch to the "user" user
-USER user
-
-FROM python:3.8
-
-ENV PYTHONUNBUFFERED=1
-
-# ENV TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6"
-# ENV TCNN_CUDA_ARCHITECTURES=86;80;75;70;61;60
-ENV FORCE_CUDA=1
-
-# Set the environment variable to specify the GPU device
-ENV CUDA_HOME=/usr/local/cuda
-# ENV CUDA_DEVICE_ORDER=PCI_BUS_ID
-# ENV CUDA_VISIBLE_DEVICES=0
-
-ENV PATH=${CUDA_HOME}/bin:/home/${USER_NAME}/.local/bin:/usr/bin:${PATH}
-ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:/usr/lib:/usr/lib64:${LD_LIBRARY_PATH}
-
-# Set home to the user's home directory
-ENV HOME=/home/user \
-    PATH=/home/user/.local/bin:$PATH \
-    PYTHONPATH=$HOME/app:$PYTHONPATH \
-	PYTHONUNBUFFERED=1 \
-	GRADIO_ALLOW_FLAGGING=never \
-	GRADIO_NUM_PORTS=1 \
-	GRADIO_SERVER_NAME=0.0.0.0 \
-	GRADIO_THEME=huggingface \
-	SYSTEM=spaces
-
-RUN pip install --upgrade pip ninja
-RUN pip install setuptools==69.5.1
-RUN pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
-
-RUN python -c "import torch, os; print(torch.version.cuda); print(os.environ.get('CUDA_HOME'))"
-COPY requirements.txt /tmp
-RUN cd /tmp && pip install -r requirements.txt
-
-RUN pip install https://download.is.tue.mpg.de/icon/HF/kaolin-0.11.0-cp38-cp38-linux_x86_64.whl
-RUN pip install https://download.is.tue.mpg.de/icon/HF/pytorch3d-0.7.0-cp38-cp38-linux_x86_64.whl
-RUN pip install pyembree
-
-ENV TRANSFORMERS_CACHE=/tmp
-ENV MPLCONFIGDIR=/tmp
-
-# cannot cache function '_make_tree': no locator available for file '/usr/local/lib/python3.8/site-packages/pymatting/util/kdtree.py'
-ENV NUMBA_CACHE_DIR=/tmp/numba_cache
-ENV HF_HOME=$HOME/.cache/huggingface
-RUN mkdir -p $HF_HOME
-
-RUN chmod 777 -R $HOME
-
-# Copy the current directory contents into the container at $HOME/app setting the owner to the user
-COPY --chown=user . $HOME/app
-
-# Set the working directory to the user's home directory
-WORKDIR $HOME/app
-
-ENV DISPLAY=:0
-
-CMD ["python", "app.py"]

README.md

@@ -1,42 +1,12 @@
 ---
-title: ICON - Clothed Human Digitization
-metaTitle: ICON-Avatarify from Photo
+title: ICON
+metaTitle: "Image2Human by Yuliang Xiu"
 emoji: 🤼
 colorFrom: indigo
 colorTo: yellow
-sdk: docker
+sdk: gradio
+sdk_version: 3.1.1
+app_file: ./apps/app.py
 pinned: true
----
-
-# ICON Clothed Human Digitization 
-### ICON: Implicit Clothed humans Obtained from Normals (CVPR 2022)
-
-<table>
-<th>
-<ul>
-<li><strong>Homepage</strong> <a href="http://icon.is.tue.mpg.de">icon.is.tue.mpg.de</a></li>
-<li><strong>Code</strong> <a href="https://github.com/YuliangXiu/ICON">YuliangXiu/ICON</a></li>
-<li><strong>Paper</strong> <a href="https://arxiv.org/abs/2112.09127">arXiv</a>, <a href="https://readpaper.com/paper/4569785684533977089">ReadPaper</a></li>
-<li><strong>Chatroom</strong> <a href="https://discord.gg/Vqa7KBGRyk">Discord</a></li>
-<li><strong>Colab Notebook</strong> <a href="https://colab.research.google.com/drive/1-AWeWhPvCTBX0KfMtgtMk10uPU05ihoA?usp=sharing">Google Colab</a></li>
-</ul>
-<a href="https://twitter.com/yuliangxiu"><img alt="Twitter Follow" src="https://img.shields.io/twitter/follow/yuliangxiu?style=social"></a>
-<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>
-<a href="https://youtu.be/hZd6AYin2DE"><img alt="YouTube Video Views" src="https://img.shields.io/youtube/views/hZd6AYin2DE?style=social"></a>
-</th>
-<th>
-<iframe width="560" height="315" src="https://www.youtube.com/embed/hZd6AYin2DE" title="YouTube video player" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture" allowfullscreen></iframe>
-</th>
-</table>
-  
-#### 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 = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
-  month     = {June},
-  year      = {2022},
-  pages     = {13296-13306}
-} 
-```
+python_version: 3.8.13
+---

app.py

@@ -1,148 +0,0 @@
-# install
-
-
-import glob
-import gradio as gr
-import numpy as np
-import os
-import subprocess
-
-from apps.infer import generate_model
-
-# running
-
-description = '''
-# ICON Clothed Human Digitization 
-### ICON: Implicit Clothed humans Obtained from Normals (CVPR 2022)
-
-<table>
-<th>
-<ul>
-<li><strong>Homepage</strong> <a href="http://icon.is.tue.mpg.de">icon.is.tue.mpg.de</a></li>
-<li><strong>Code</strong> <a href="https://github.com/YuliangXiu/ICON">YuliangXiu/ICON</a></li>
-<li><strong>Paper</strong> <a href="https://arxiv.org/abs/2112.09127">arXiv</a>, <a href="https://readpaper.com/paper/4569785684533977089">ReadPaper</a></li>
-<li><strong>Chatroom</strong> <a href="https://discord.gg/Vqa7KBGRyk">Discord</a></li>
-<li><strong>Colab Notebook</strong> <a href="https://colab.research.google.com/drive/1-AWeWhPvCTBX0KfMtgtMk10uPU05ihoA?usp=sharing">Google Colab</a></li>
-</ul>
-<a href="https://twitter.com/yuliangxiu"><img alt="Twitter Follow" src="https://img.shields.io/twitter/follow/yuliangxiu?style=social"></a>
-<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>
-<a href="https://youtu.be/hZd6AYin2DE"><img alt="YouTube Video Views" src="https://img.shields.io/youtube/views/hZd6AYin2DE?style=social"></a>
-</th>
-<th>
-<iframe width="560" height="315" src="https://www.youtube.com/embed/hZd6AYin2DE" title="YouTube video player" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture" allowfullscreen></iframe>
-</th>
-</table>
-
-<center>
-<a href="https://huggingface.co/spaces/Yuliang/ICON?duplicate=true"><img src="https://huggingface.co/datasets/huggingface/badges/raw/main/duplicate-this-space-lg-dark.svg"/></a>
-<h2> The reconstruction + refinement + video takes about 200 seconds for a single image. </h2> 
-<h2><span style="color:red"> ICON is only suitable for humanoid images and will not work well on cartoons with non-human shapes.</span></h2>
-</center>
-
-<details>
-
-<summary>More</summary>
-
-#### 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 = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
-  month     = {June},
-  year      = {2022},
-  pages     = {13296-13306}
-} 
-```
-
-#### 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)
-
-#### Image Credits
-
-* [Pinterest](https://www.pinterest.com/search/pins/?q=parkour&rs=sitelinks_searchbox)
-
-#### 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
-
-
-model_types = ['ICON', 'PIFu']
-examples_names = glob.glob('examples/*.png')
-examples_types = np.random.choice(
-    model_types, len(examples_names), p=[0.6, 0.2, 0.2])
-
-examples = [list(item) for item in zip(examples_names, examples_types)]
-
-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, 1000, 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("Generate Image")
-                btn_submit = gr.Button("Submit Image")
-
-            gr.Examples(examples=examples,
-                        inputs=[inp, radio_choice],
-                        cache_examples=False,
-                        fn=generate_model,
-                        outputs=out_lst)
-
-            out_vid = gr.Video(
-                label="Image + Normal + SMPL Body + Clothed Human")
-            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_final = gr.Model3D(
-                clear_color=[0.0, 0.0, 0.0, 0.0],  label="Clothed human")
-            out_final_download = gr.File(
-                label="Download clothed human mesh")
-            out_smpl = gr.Model3D(
-                clear_color=[0.0, 0.0, 0.0, 0.0],  label="SMPL body")
-            out_smpl_download = gr.File(label="Download SMPL body mesh")
-            out_smpl_npy_download = gr.File(label="Download SMPL params")
-
-    out_lst = [out_smpl, out_smpl_download, out_smpl_npy_download,
-               out_final, out_final_download, out_vid, 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=(os.environ['USER'], os.environ['PASSWORD']),
-    #             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

@@ -14,26 +14,21 @@
 #
 # 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.renderer.gl.init_gl import initialize_GL_context
 from lib.common.render import Render
 from lib.dataset.mesh_util import SMPLX, update_mesh_shape_prior_losses, get_visibility
 import warnings
 import logging
 import torch
-import lib.smplx as smplx
+import 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
 
@@ -102,8 +97,10 @@ class ICON(pl.LightningModule):
 
         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']),
+            kid_template_path=osp.join(
+                osp.realpath(self.smpl_data.model_dir),
+                f"{smpl_type}/{smpl_type}_kid_template.npy",
+            ),
             model_type=smpl_type,
             gender=gender,
             age=age,

apps/app.py

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

apps/infer.py

@@ -14,33 +14,34 @@
 #
 # Contact: [email protected]
 
-import os
-import gc
-
 import logging
+from lib.common.render import query_color, image2vid
 from lib.common.config import cfg
+from lib.common.cloth_extraction import extract_cloth
 from lib.dataset.mesh_util import (
     load_checkpoint,
     update_mesh_shape_prior_losses,
+    get_optim_grid_image,
     blend_rgb_norm,
     unwrap,
     remesh,
     tensor2variable,
-    rot6d_to_rotmat
+    normal_loss
 )
 
 from lib.dataset.TestDataset import TestDataset
-from lib.common.render import query_color
 from lib.net.local_affine import LocalAffine
 from pytorch3d.structures import Meshes
 from apps.ICON import ICON
 
+import os
 from termcolor import colored
+import argparse
 import numpy as np
 from PIL import Image
 import trimesh
+import pickle
 import numpy as np
-from tqdm import tqdm
 
 import torch
 torch.backends.cudnn.benchmark = True
@@ -48,31 +49,36 @@ torch.backends.cudnn.benchmark = True
 logging.getLogger("trimesh").setLevel(logging.ERROR)
 
 
-def generate_model(in_path, model_type):
+if __name__ == "__main__":
 
-    torch.cuda.empty_cache()
-    
-    if model_type == 'ICON':
-        model_type = 'icon-filter'
-    else:
-        model_type = model_type.lower()
+    # loading cfg file
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("-gpu", "--gpu_device", type=int, default=0)
+    parser.add_argument("-colab", action="store_true")
+    parser.add_argument("-loop_smpl", "--loop_smpl", type=int, default=100)
+    parser.add_argument("-patience", "--patience", type=int, default=5)
+    parser.add_argument("-vis_freq", "--vis_freq", type=int, default=10)
+    parser.add_argument("-loop_cloth", "--loop_cloth", type=int, default=200)
+    parser.add_argument("-hps_type", "--hps_type", type=str, default="pymaf")
+    parser.add_argument("-export_video", action="store_true")
+    parser.add_argument("-in_dir", "--in_dir", type=str, default="./examples")
+    parser.add_argument("-out_dir", "--out_dir",
+                        type=str, default="./results")
+    parser.add_argument('-seg_dir', '--seg_dir', type=str, default=None)
+    parser.add_argument(
+        "-cfg", "--config", type=str, default="./configs/icon-filter.yaml"
+    )
 
-    config_dict = {'loop_smpl': 100,
-                   'loop_cloth': 200,
-                   'patience': 5,
-                   'out_dir': './results',
-                   'hps_type': 'pymaf',
-                   'config': f"./configs/{model_type}.yaml"}
+    args = parser.parse_args()
 
     # cfg read and merge
-    cfg.merge_from_file(config_dict['config'])
+    cfg.merge_from_file(args.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],
+        [args.gpu_device],
         "mcube_res",
         256,
         "clean_mesh",
@@ -82,21 +88,28 @@ def generate_model(in_path, model_type):
     cfg.merge_from_list(cfg_show_list)
     cfg.freeze()
 
-    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
-    device = torch.device(f"cuda:0")
+    os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
+    device = torch.device(f"cuda:{args.gpu_device}")
+
+    if args.colab:
+        print(colored("colab environment...", "red"))
+        from tqdm.notebook import tqdm
+    else:
+        print(colored("normal environment...", "red"))
+        from tqdm import tqdm
 
     # load model and dataloader
     model = ICON(cfg)
     model = load_checkpoint(model, cfg)
 
     dataset_param = {
-        'image_path': in_path,
-        'seg_dir': None,
+        'image_dir': args.in_dir,
+        'seg_dir': args.seg_dir,
         'has_det': True,            # w/ or w/o detection
-        'hps_type': 'pymaf'   # pymaf/pare/pixie
+        'hps_type': args.hps_type   # pymaf/pare/pixie
     }
 
-    if config_dict['hps_type'] == "pixie" and "pamir" in config_dict['config']:
+    if args.hps_type == "pixie" and "pamir" in args.config:
         print(colored("PIXIE isn't compatible with PaMIR, thus switch to PyMAF", "red"))
         dataset_param["hps_type"] = "pymaf"
 
@@ -126,10 +139,10 @@ def generate_model(in_path, model_type):
             data["global_orient"], device=device, requires_grad=True
         )  # [1,1,3,3]
 
-        optimizer_smpl = torch.optim.Adam(
+        optimizer_smpl = torch.optim.SGD(
             [optimed_pose, optimed_trans, optimed_betas, optimed_orient],
             lr=1e-3,
-            amsgrad=True,
+            momentum=0.9,
         )
         scheduler_smpl = torch.optim.lr_scheduler.ReduceLROnPlateau(
             optimizer_smpl,
@@ -137,45 +150,38 @@ def generate_model(in_path, model_type):
             factor=0.5,
             verbose=0,
             min_lr=1e-5,
-            patience=config_dict['patience'],
+            patience=args.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},
+            "cloth": {"weight": 1e1, "value": 0.0},             # Cloth: Normal_recon - Normal_pred
+            "stiffness": {"weight": 1e5, "value": 0.0},         # Cloth: [RT]_v1 - [RT]_v2 (v1-edge-v2)
+            "rigid": {"weight": 1e5, "value": 0.0},             # Cloth: det(R) = 1
+            "edge": {"weight": 0, "value": 0.0},                # Cloth: edge length
+            "nc": {"weight": 0, "value": 0.0},                  # Cloth: normal consistency
+            "laplacian": {"weight": 1e2, "value": 0.0},         # Cloth: laplacian smoonth
+            "normal": {"weight": 1e0, "value": 0.0},            # Body: Normal_pred - Normal_smpl
+            "silhouette": {"weight": 1e1, "value": 0.0},        # Body: Silhouette_pred - Silhouette_smpl
         }
 
         # smpl optimization
 
-        loop_smpl = tqdm(range(config_dict['loop_smpl']))
+        loop_smpl = tqdm(
+            range(args.loop_smpl if cfg.net.prior_type != "pifu" else 1))
+
+        per_data_lst = []
 
-        for _ in loop_smpl:
+        for i in loop_smpl:
+
+            per_loop_lst = []
 
             optimizer_smpl.zero_grad()
-            
-            # 6d_rot to rot_mat
-            optimed_orient_mat = rot6d_to_rotmat(optimed_orient.view(-1,6)).unsqueeze(0)
-            optimed_pose_mat = rot6d_to_rotmat(optimed_pose.view(-1,6)).unsqueeze(0)
 
             if dataset_param["hps_type"] != "pixie":
                 smpl_out = dataset.smpl_model(
                     betas=optimed_betas,
-                    body_pose=optimed_pose_mat,
-                    global_orient=optimed_orient_mat,
+                    body_pose=optimed_pose,
+                    global_orient=optimed_orient,
                     pose2rot=False,
                 )
 
@@ -185,8 +191,8 @@ def generate_model(in_path, model_type):
                 smpl_verts, _, _ = dataset.smpl_model(
                     shape_params=optimed_betas,
                     expression_params=tensor2variable(data["exp"], device),
-                    body_pose=optimed_pose_mat,
-                    global_pose=optimed_orient_mat,
+                    body_pose=optimed_pose,
+                    global_pose=optimed_orient,
                     jaw_pose=tensor2variable(data["jaw_pose"], device),
                     left_hand_pose=tensor2variable(
                         data["left_hand_pose"], device),
@@ -214,7 +220,12 @@ def generate_model(in_path, model_type):
             diff_B_smpl = torch.abs(
                 in_tensor["T_normal_B"] - in_tensor["normal_B"])
 
-            losses["normal"]["value"] = (diff_F_smpl + diff_B_smpl).mean()
+            loss_F_smpl = normal_loss(
+                in_tensor["T_normal_F"], in_tensor["normal_F"])
+            loss_B_smpl = normal_loss(
+                in_tensor["T_normal_B"], in_tensor["normal_B"])
+
+            losses["normal"]["value"] = (loss_F_smpl + loss_B_smpl).mean()
 
             # silhouette loss
             smpl_arr = torch.cat([T_mask_F, T_mask_B], dim=-1)[0]
@@ -239,6 +250,33 @@ def generate_model(in_path, model_type):
             pbar_desc += f"Total: {smpl_loss:.3f}"
             loop_smpl.set_description(pbar_desc)
 
+            if i % args.vis_freq == 0:
+
+                per_loop_lst.extend(
+                    [
+                        in_tensor["image"],
+                        in_tensor["T_normal_F"],
+                        in_tensor["normal_F"],
+                        diff_F_smpl / 2.0,
+                        diff_S[:, :512].unsqueeze(
+                            0).unsqueeze(0).repeat(1, 3, 1, 1),
+                    ]
+                )
+                per_loop_lst.extend(
+                    [
+                        in_tensor["image"],
+                        in_tensor["T_normal_B"],
+                        in_tensor["normal_B"],
+                        diff_B_smpl / 2.0,
+                        diff_S[:, 512:].unsqueeze(
+                            0).unsqueeze(0).repeat(1, 3, 1, 1),
+                    ]
+                )
+                per_data_lst.append(
+                    get_optim_grid_image(
+                        per_loop_lst, None, nrow=5, type="smpl")
+                )
+
             smpl_loss.backward()
             optimizer_smpl.step()
             scheduler_smpl.step(smpl_loss)
@@ -249,20 +287,18 @@ def generate_model(in_path, model_type):
         # 1. SMPL Fitting
         # 2. Clothes Refinement
 
-        os.makedirs(os.path.join(config_dict['out_dir'], cfg.name,
+        os.makedirs(os.path.join(args.out_dir, cfg.name,
                     "refinement"), exist_ok=True)
 
         # visualize the final results in self-rotation mode
-        os.makedirs(os.path.join(config_dict['out_dir'],
-                    cfg.name, "vid"), exist_ok=True)
+        os.makedirs(os.path.join(args.out_dir, cfg.name, "vid"), exist_ok=True)
 
         # final results rendered as image
         # 1. Render the final fitted SMPL (xxx_smpl.png)
         # 2. Render the final reconstructed clothed human (xxx_cloth.png)
         # 3. Blend the original image with predicted cloth normal (xxx_overlap.png)
 
-        os.makedirs(os.path.join(config_dict['out_dir'],
-                    cfg.name, "png"), exist_ok=True)
+        os.makedirs(os.path.join(args.out_dir, cfg.name, "png"), exist_ok=True)
 
         # final reconstruction meshes
         # 1. SMPL mesh (xxx_smpl.obj)
@@ -271,41 +307,54 @@ def generate_model(in_path, model_type):
         # 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)
+        os.makedirs(os.path.join(args.out_dir, cfg.name, "obj"), exist_ok=True)
+
+        if cfg.net.prior_type != "pifu":
+
+            per_data_lst[0].save(
+                os.path.join(
+                    args.out_dir, cfg.name, f"refinement/{data['name']}_smpl.gif"
+                ),
+                save_all=True,
+                append_images=per_data_lst[1:],
+                duration=500,
+                loop=0,
+            )
+
+            if args.vis_freq == 1:
+                image2vid(
+                    per_data_lst,
+                    os.path.join(
+                        args.out_dir, cfg.name, f"refinement/{data['name']}_smpl.avi"
+                    ),
+                )
+
+            per_data_lst[-1].save(
+                os.path.join(args.out_dir, cfg.name,
+                             f"png/{data['name']}_smpl.png")
+            )
 
-        norm_pred_F = (
+        norm_pred = (
             ((in_tensor["normal_F"][0].permute(1, 2, 0) + 1.0) * 255.0 / 2.0)
             .detach()
             .cpu()
             .numpy()
             .astype(np.uint8)
         )
-        
-        norm_pred_B = (
-            ((in_tensor["normal_B"][0].permute(1, 2, 0) + 1.0) * 255.0 / 2.0)
-            .detach()
-            .cpu()
-            .numpy()
-            .astype(np.uint8)
-        )
 
-        norm_orig_F = unwrap(norm_pred_F, data)
-        norm_orig_B = unwrap(norm_pred_B, data)
-        
+        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_F = blend_rgb_norm(data["ori_image"], norm_orig_F, mask_orig)
-        rgb_norm_B = blend_rgb_norm(data["ori_image"], norm_orig_B, mask_orig)
+        rgb_norm = blend_rgb_norm(data["ori_image"], norm_orig, mask_orig)
 
         Image.fromarray(
             np.concatenate(
-                [data["ori_image"].astype(np.uint8), rgb_norm_F, rgb_norm_B], axis=1)
-        ).save(os.path.join(config_dict['out_dir'], cfg.name, f"png/{data['name']}_overlap.png"))
+                [data["ori_image"].astype(np.uint8), rgb_norm], axis=1)
+        ).save(os.path.join(args.out_dir, cfg.name, f"png/{data['name']}_overlap.png"))
 
         smpl_obj = trimesh.Trimesh(
             in_tensor["smpl_verts"].detach().cpu()[0] *
@@ -314,24 +363,23 @@ def generate_model(in_path, model_type):
             process=False,
             maintains_order=True
         )
-        smpl_obj.visual.vertex_colors = (smpl_obj.vertex_normals+1.0)*255.0*0.5
-        smpl_obj.export(
-            f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_smpl.obj")
         smpl_obj.export(
-            f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_smpl.glb")
+            f"{args.out_dir}/{cfg.name}/obj/{data['name']}_smpl.obj")
 
         smpl_info = {'betas': optimed_betas,
-                     'pose': optimed_pose_mat,
-                     'orient': optimed_orient_mat,
+                     '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)
+            f"{args.out_dir}/{cfg.name}/obj/{data['name']}_smpl.npy", smpl_info, allow_pickle=True)
 
         # ------------------------------------------------------------------------------------------------------------------
 
         # cloth optimization
 
+        per_data_lst = []
+
         # cloth recon
         in_tensor.update(
             dataset.compute_vis_cmap(
@@ -356,15 +404,13 @@ def generate_model(in_path, model_type):
         recon_obj = trimesh.Trimesh(
             verts_pr, faces_pr, process=False, maintains_order=True
         )
-        recon_obj.visual.vertex_colors = (
-            recon_obj.vertex_normals+1.0)*255.0*0.5
         recon_obj.export(
-            os.path.join(config_dict['out_dir'], cfg.name,
+            os.path.join(args.out_dir, cfg.name,
                          f"obj/{data['name']}_recon.obj")
         )
 
         # Isotropic Explicit Remeshing for better geometry topology
-        verts_refine, faces_refine = remesh(os.path.join(config_dict['out_dir'], cfg.name,
+        verts_refine, faces_refine = remesh(os.path.join(args.out_dir, cfg.name,
                                                          f"obj/{data['name']}_recon.obj"), 0.5, device)
 
         # define local_affine deform verts
@@ -380,16 +426,26 @@ def generate_model(in_path, model_type):
             factor=0.1,
             verbose=0,
             min_lr=1e-5,
-            patience=config_dict['patience'],
+            patience=args.patience,
         )
 
+        with torch.no_grad():
+            per_loop_lst = []
+            rotate_recon_lst = dataset.render.get_rgb_image(cam_ids=[
+                0, 1, 2, 3])
+            per_loop_lst.extend(rotate_recon_lst)
+            per_data_lst.append(get_optim_grid_image(
+                per_loop_lst, None, type="cloth"))
+
         final = None
 
-        if config_dict['loop_cloth'] > 0:
+        if args.loop_cloth > 0:
 
-            loop_cloth = tqdm(range(config_dict['loop_cloth']))
+            loop_cloth = tqdm(range(args.loop_cloth))
 
-            for _ in loop_cloth:
+            for i in loop_cloth:
+
+                per_loop_lst = []
 
                 optimizer_cloth.zero_grad()
 
@@ -426,67 +482,135 @@ def generate_model(in_path, model_type):
                 loop_cloth.set_description(pbar_desc)
 
                 # update params
-                cloth_loss.backward()
+                cloth_loss.backward(retain_graph=True)
                 optimizer_cloth.step()
                 scheduler_cloth.step(cloth_loss)
 
+                # for vis
+                with torch.no_grad():
+                    if i % args.vis_freq == 0:
+
+                        rotate_recon_lst = dataset.render.get_rgb_image(cam_ids=[
+                            0, 1, 2, 3])
+
+                        per_loop_lst.extend(
+                            [
+                                in_tensor["image"],
+                                in_tensor["P_normal_F"],
+                                in_tensor["normal_F"],
+                                diff_F_cloth / 2.0,
+                            ]
+                        )
+                        per_loop_lst.extend(
+                            [
+                                in_tensor["image"],
+                                in_tensor["P_normal_B"],
+                                in_tensor["normal_B"],
+                                diff_B_cloth / 2.0,
+                            ]
+                        )
+                        per_loop_lst.extend(rotate_recon_lst)
+                        per_data_lst.append(
+                            get_optim_grid_image(
+                                per_loop_lst, None, type="cloth")
+                        )
+
+            # gif for optimization
+            per_data_lst[1].save(
+                os.path.join(
+                    args.out_dir, cfg.name, f"refinement/{data['name']}_cloth.gif"
+                ),
+                save_all=True,
+                append_images=per_data_lst[2:],
+                duration=500,
+                loop=0,
+            )
+
+            if args.vis_freq == 1:
+                image2vid(
+                    per_data_lst,
+                    os.path.join(
+                        args.out_dir, cfg.name, f"refinement/{data['name']}_cloth.avi"
+                    ),
+                )
+
             final = trimesh.Trimesh(
                 mesh_pr.verts_packed().detach().squeeze(0).cpu(),
                 mesh_pr.faces_packed().detach().squeeze(0).cpu(),
                 process=False, maintains_order=True
             )
-            
-            # only with front texture
-            tex_colors = query_color(
+            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,
             )
-
-            # full normal textures
-            norm_colors = (mesh_pr.verts_normals_padded().squeeze(
-                0).detach().cpu() + 1.0) * 0.5 * 255.0
-                
-            final.visual.vertex_colors = tex_colors
-            final.export(
-                f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_refine.obj")
-                
-            final.visual.vertex_colors = norm_colors
+            final.visual.vertex_colors = final_colors
             final.export(
-                f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_refine.glb")
+                f"{args.out_dir}/{cfg.name}/obj/{data['name']}_refine.obj")
 
-        # always export visualized video regardless of the cloth refinment
-        verts_lst = [smpl_obj.vertices, final.vertices]
-        faces_lst = [smpl_obj.faces, final.faces]
-
-        # self-rotated video
-        dataset.render.load_meshes(
-            verts_lst, faces_lst)
-        dataset.render.get_rendered_video(
-            [data["ori_image"], rgb_norm_F, rgb_norm_B],
-            os.path.join(config_dict['out_dir'], cfg.name,
-                         f"vid/{data['name']}_cloth.mp4"),
+        # always export visualized png regardless of the cloth refinment
+        per_data_lst[-1].save(
+            os.path.join(args.out_dir, cfg.name,
+                         f"png/{data['name']}_cloth.png")
         )
 
-    smpl_obj_path = f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_smpl.obj"
-    smpl_glb_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"
-    refine_obj_path = f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_refine.obj"
-    refine_glb_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")
-
-    # clean all the variables
-    for element in dir():
-        if 'path' not in element:
-            del locals()[element]
-    gc.collect()
-    torch.cuda.empty_cache()
-    
-    return [smpl_glb_path, smpl_obj_path,smpl_npy_path,
-            refine_glb_path, refine_obj_path,
-            video_path, video_path, overlap_path]
+        # always export visualized video regardless of the cloth refinment
+        if args.export_video:
+            if final is not None:
+                verts_lst = [verts_pr, final.vertices]
+                faces_lst = [faces_pr, final.faces]
+            else:
+                verts_lst = [verts_pr]
+                faces_lst = [faces_pr]
+
+            # self-rotated video
+            dataset.render.load_meshes(
+                verts_lst, faces_lst)
+            dataset.render.get_rendered_video(
+                [data["ori_image"], rgb_norm],
+                os.path.join(args.out_dir, cfg.name,
+                             f"vid/{data['name']}_cloth.mp4"),
+            )
+
+        # garment extraction from deepfashion images
+        if not (args.seg_dir is None):
+            if final is not None:
+                recon_obj = final.copy()
+
+            os.makedirs(os.path.join(
+                args.out_dir, cfg.name, "clothes"), exist_ok=True)
+            os.makedirs(os.path.join(args.out_dir, cfg.name,
+                        "clothes", "info"), exist_ok=True)
+            for seg in data['segmentations']:
+                # These matrices work for PyMaf, not sure about the other hps type
+                K = np.array([[1.0000,  0.0000,  0.0000,  0.0000],
+                              [0.0000,  1.0000,  0.0000,  0.0000],
+                              [0.0000,  0.0000, -0.5000,  0.0000],
+                              [-0.0000, -0.0000,  0.5000,  1.0000]]).T
+
+                R = np.array([[-1.,  0.,  0.],
+                              [0.,  1.,  0.],
+                              [0.,  0., -1.]])
+
+                t = np.array([[-0.,  -0., 100.]])
+                clothing_obj = extract_cloth(recon_obj, seg, K, R, t, smpl_obj)
+                if clothing_obj is not None:
+                    cloth_type = seg['type'].replace(' ', '_')
+                    cloth_info = {
+                        'betas': optimed_betas,
+                        'body_pose': optimed_pose,
+                        'global_orient': optimed_orient,
+                        'pose2rot': False,
+                        'clothing_type': cloth_type,
+                    }
+
+                    file_id = f"{data['name']}_{cloth_type}"
+                    with open(os.path.join(args.out_dir, cfg.name, "clothes", "info", f"{file_id}_info.pkl"), 'wb') as fp:
+                        pickle.dump(cloth_info, fp)
+
+                    clothing_obj.export(os.path.join(
+                        args.out_dir, cfg.name, "clothes", f"{file_id}.obj"))
+                else:
+                    print(
+                        f"Unable to extract clothing of type {seg['type']} from image {data['name']}")

configs/icon-filter.yaml

@@ -1,7 +1,7 @@
 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"
+resume_path: "./data/ckpt/icon-filter.ckpt"
+normal_path: "./data/ckpt/normal.ckpt"
 
 test_mode: True
 batch_size: 1

configs/icon-nofilter.yaml

@@ -1,7 +1,7 @@
 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"
+resume_path: "./data/ckpt/icon-nofilter.ckpt"
+normal_path: "./data/ckpt/normal.ckpt"
 
 test_mode: True
 batch_size: 1

configs/pamir.yaml

@@ -1,7 +1,7 @@
 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"
+resume_path: "./data/ckpt/pamir.ckpt"
+normal_path: "./data/ckpt/normal.ckpt"
 
 test_mode: True
 batch_size: 1

configs/pifu.yaml

@@ -1,7 +1,7 @@
 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"
+resume_path: "./data/ckpt/pifu.ckpt"
+normal_path: "./data/ckpt/normal.ckpt"
 
 test_mode: True
 batch_size: 1

environment.yaml

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

examples/segmentation/003883.json

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

examples/segmentation/028009.json

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

fetch_data.sh

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

install.sh

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

lib/common/render.py

@@ -33,15 +33,14 @@ from pytorch3d.renderer import (
 )
 from pytorch3d.renderer.mesh import TexturesVertex
 from pytorch3d.structures import Meshes
-
-import os, subprocess
-
 from lib.dataset.mesh_util import SMPLX, get_visibility
+
 import lib.common.render_utils as util
 import torch
 import numpy as np
 from PIL import Image
 from tqdm import tqdm
+import os
 import cv2
 import math
 from termcolor import colored
@@ -327,10 +326,8 @@ class Render:
 
     def get_rendered_video(self, images, save_path):
         
-        tmp_path = save_path.replace('cloth', 'tmp')
-        
         self.cam_pos = []
-        for angle in range(0, 360, 3):
+        for angle in range(360):
             self.cam_pos.append(
                 (
                     100.0 * math.cos(np.pi / 180 * angle),
@@ -345,7 +342,7 @@ class Render:
 
         fourcc = cv2.VideoWriter_fourcc(*"mp4v")
         video = cv2.VideoWriter(
-            tmp_path, fourcc, 30, (self.size * len(self.meshes) +
+            save_path, fourcc, 30, (self.size * len(self.meshes) +
                                     new_shape[1] * len(images), self.size)
         )
 
@@ -375,8 +372,6 @@ class Render:
             video.write(final_img)
 
         video.release()
-        
-        os.system(f'ffmpeg -y -loglevel quiet -stats -i {tmp_path} -c:v libx264 {save_path}')
 
     def get_silhouette_image(self, cam_ids=[0, 2]):
 

lib/common/train_util.py

@@ -32,8 +32,6 @@ import os
 from termcolor import colored
 
 
-
-
 def reshape_sample_tensor(sample_tensor, num_views):
     if num_views == 1:
         return sample_tensor

lib/dataloader_demo.py

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

lib/dataset/Evaluator.py

@@ -15,11 +15,11 @@
 #
 # 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 os
 import numpy as np
 import torch
 from torch import nn

lib/dataset/PIFuDataset.py

@@ -9,12 +9,12 @@ import os.path as osp
 import numpy as np
 from PIL import Image
 import random
-import os
 import trimesh
 import torch
+import vedo
 from kaolin.ops.mesh import check_sign
 import torchvision.transforms as transforms
-from huggingface_hub import hf_hub_download, cached_download
+from ipdb import set_trace
 
 
 class PIFuDataset():
@@ -343,9 +343,9 @@ class PIFuDataset():
             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_path = osp.join(self.smplx.model_dir, "smpl/SMPL_MALE.pkl")
+        tetra_path = osp.join(self.smplx.tedra_dir,
+                              "tetra_male_adult_smpl.npz")
 
         smpl_model = TetraSMPLModel(smpl_path, tetra_path, 'adult')
 
@@ -365,7 +365,7 @@ class PIFuDataset():
         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']),
+        faces = np.loadtxt(osp.join(self.smplx.tedra_dir, "tetrahedrons_male_adult.txt"),
                            dtype=np.int32) - 1
 
         pad_v_num = int(8000 - verts.shape[0])
@@ -586,4 +586,77 @@ class PIFuDataset():
         labels = torch.from_numpy(labels).float()
         normals = torch.from_numpy(normals).float()
 
-        return {'samples_geo': samples, 'labels_geo': labels}
+        return {'samples_geo': samples, 'labels_geo': labels}
+
+    def visualize_sampling3D(self, data_dict, mode='vis'):
+
+        # create plot
+        vp = vedo.Plotter(title="", size=(1500, 1500), axes=0, bg='white')
+        vis_list = []
+
+        assert mode in ['vis', 'sdf', 'normal', 'cmap', 'occ']
+
+        # sdf-1 cmap-3 norm-3 vis-1
+        if mode == 'vis':
+            labels = data_dict[f'smpl_feat'][:, [-1]]  # visibility
+            colors = np.concatenate([labels, labels, labels], axis=1)
+        elif mode == 'occ':
+            labels = data_dict[f'labels_geo'][..., None]  # occupancy
+            colors = np.concatenate([labels, labels, labels], axis=1)
+        elif mode == 'sdf':
+            labels = data_dict[f'smpl_feat'][:, [0]]  # sdf
+            labels -= labels.min()
+            labels /= labels.max()
+            colors = np.concatenate([labels, labels, labels], axis=1)
+        elif mode == 'normal':
+            labels = data_dict[f'smpl_feat'][:, -4:-1]  # normal
+            colors = (labels + 1.0) * 0.5
+        elif mode == 'cmap':
+            labels = data_dict[f'smpl_feat'][:, -7:-4]  # colormap
+            colors = np.array(labels)
+
+        points = projection(data_dict['samples_geo'], data_dict['calib'])
+        verts = projection(data_dict['verts'], data_dict['calib'])
+        points[:, 1] *= -1
+        verts[:, 1] *= -1
+
+        # create a mesh
+        mesh = trimesh.Trimesh(verts, data_dict['faces'], process=True)
+        mesh.visual.vertex_colors = [128.0, 128.0, 128.0, 255.0]
+        vis_list.append(mesh)
+
+        if 'voxel_verts' in data_dict.keys():
+            print(colored("voxel verts", "green"))
+            voxel_verts = data_dict['voxel_verts'] * 2.0
+            voxel_faces = data_dict['voxel_faces']
+            voxel_verts[:, 1] *= -1
+            voxel = trimesh.Trimesh(
+                voxel_verts, voxel_faces[:, [0, 2, 1]], process=False, maintain_order=True)
+            voxel.visual.vertex_colors = [0.0, 128.0, 0.0, 255.0]
+            vis_list.append(voxel)
+
+        if 'smpl_verts' in data_dict.keys():
+            print(colored("smpl verts", "green"))
+            smplx_verts = data_dict['smpl_verts']
+            smplx_faces = data_dict['smpl_faces']
+            smplx_verts[:, 1] *= -1
+            smplx = trimesh.Trimesh(
+                smplx_verts, smplx_faces[:, [0, 2, 1]], process=False, maintain_order=True)
+            smplx.visual.vertex_colors = [128.0, 128.0, 0.0, 255.0]
+            vis_list.append(smplx)
+
+        # create a picure
+        img_pos = [1.0, 0.0, -1.0]
+        for img_id, img_key in enumerate(['normal_F', 'image', 'T_normal_B']):
+            image_arr = (data_dict[img_key].detach().cpu().permute(
+                1, 2, 0).numpy() + 1.0) * 0.5 * 255.0
+            image_dim = image_arr.shape[0]
+            image = vedo.Picture(image_arr).scale(
+                2.0 / image_dim).pos(-1.0, -1.0, img_pos[img_id])
+            vis_list.append(image)
+
+        # create a pointcloud
+        pc = vedo.Points(points, r=15, c=np.float32(colors))
+        vis_list.append(pc)
+
+        vp.show(*vis_list, bg="white", axes=1.0, interactive=True)

lib/dataset/TestDataset.py

@@ -15,9 +15,7 @@
 #
 # Contact: [email protected]
 
-import os
-
-import lib.smplx as smplx
+import smplx
 from lib.pymaf.utils.geometry import rotation_matrix_to_angle_axis, batch_rodrigues
 from lib.pymaf.utils.imutils import process_image
 from lib.pymaf.core import path_config
@@ -27,12 +25,14 @@ from lib.common.render import Render
 from lib.dataset.body_model import TetraSMPLModel
 from lib.dataset.mesh_util import get_visibility, SMPLX
 import os.path as osp
+import os
 import torch
+import glob
 import numpy as np
 import random
+import human_det
 from termcolor import colored
 from PIL import ImageFile
-from huggingface_hub import cached_download
 
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 
@@ -42,7 +42,7 @@ class TestDataset():
 
         random.seed(1993)
 
-        self.image_path = cfg['image_path']
+        self.image_dir = cfg['image_dir']
         self.seg_dir = cfg['seg_dir']
         self.has_det = cfg['has_det']
         self.hps_type = cfg['hps_type']
@@ -51,7 +51,19 @@ class TestDataset():
 
         self.device = device
 
-        self.subject_list = [self.image_path]
+        if self.has_det:
+            self.det = human_det.Detection()
+        else:
+            self.det = None
+
+        keep_lst = sorted(glob.glob(f"{self.image_dir}/*"))
+        img_fmts = ['jpg', 'png', 'jpeg', "JPG", 'bmp']
+        keep_lst = [
+            item for item in keep_lst if item.split(".")[-1] in img_fmts
+        ]
+
+        self.subject_list = sorted(
+            [item for item in keep_lst if item.split(".")[-1] in img_fmts])
 
         # smpl related
         self.smpl_data = SMPLX()
@@ -100,9 +112,9 @@ class TestDataset():
     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_path = osp.join(self.smpl_data.model_dir, "smpl/SMPL_NEUTRAL.pkl")
+        tetra_path = osp.join(self.smpl_data.tedra_dir,
+                              'tetra_neutral_adult_smpl.npz')
         smpl_model = TetraSMPLModel(smpl_path, tetra_path, 'adult')
 
         pose = torch.cat([global_orient[0], body_pose[0]], dim=0)
@@ -112,8 +124,8 @@ class TestDataset():
         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']),
+        faces = np.loadtxt(osp.join(self.smpl_data.tedra_dir,
+                                    'tetrahedrons_neutral_adult.txt'),
                            dtype=np.int32) - 1
 
         pad_v_num = int(8000 - verts.shape[0])
@@ -148,7 +160,7 @@ class TestDataset():
 
         if self.seg_dir is None:
             img_icon, img_hps, img_ori, img_mask, uncrop_param = process_image(
-                img_path, self.hps_type, 512, self.device)
+                img_path, self.det, self.hps_type, 512, self.device)
 
             data_dict = {
                 'name': img_name,
@@ -160,7 +172,7 @@ class TestDataset():
 
         else:
             img_icon, img_hps, img_ori, img_mask, uncrop_param, segmentations = process_image(
-                img_path, self.hps_type, 512, self.device,
+                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,
@@ -233,11 +245,6 @@ class TestDataset():
         # body_pose - [1, 23, 3, 3] / [1, 21, 3, 3]
         # global_orient - [1, 1, 3, 3]
         # smpl_verts - [1, 6890, 3] / [1, 10475, 3]
-        
-        # from rot_mat to rot_6d for better optimization
-        N_body = data_dict["body_pose"].shape[1]
-        data_dict["body_pose"] =  data_dict["body_pose"][:, :, :, :2].reshape(1, N_body,-1)
-        data_dict["global_orient"] =  data_dict["global_orient"][:, :, :, :2].reshape(1, 1,-1)
 
         return data_dict
 
@@ -252,3 +259,84 @@ class TestDataset():
         # render optimized mesh (normal, T_normal, image [-1,1])
         self.render.load_meshes(verts, faces)
         return self.render.get_depth_map(cam_ids=[0, 2])
+
+    def visualize_alignment(self, data):
+
+        import vedo
+        import trimesh
+
+        if self.hps_type != 'pixie':
+            smpl_out = self.smpl_model(betas=data['betas'],
+                                       body_pose=data['body_pose'],
+                                       global_orient=data['global_orient'],
+                                       pose2rot=False)
+            smpl_verts = (
+                (smpl_out.vertices + data['trans']) * data['scale']).detach().cpu().numpy()[0]
+        else:
+            smpl_verts, _, _ = self.smpl_model(shape_params=data['betas'],
+                                               expression_params=data['exp'],
+                                               body_pose=data['body_pose'],
+                                               global_pose=data['global_orient'],
+                                               jaw_pose=data['jaw_pose'],
+                                               left_hand_pose=data['left_hand_pose'],
+                                               right_hand_pose=data['right_hand_pose'])
+
+            smpl_verts = (
+                (smpl_verts + data['trans']) * data['scale']).detach().cpu().numpy()[0]
+
+        smpl_verts *= np.array([1.0, -1.0, -1.0])
+        faces = data['smpl_faces'][0].detach().cpu().numpy()
+
+        image_P = data['image']
+        image_F, image_B = self.render_normal(smpl_verts, faces)
+
+        # create plot
+        vp = vedo.Plotter(title="", size=(1500, 1500))
+        vis_list = []
+
+        image_F = (
+            0.5 * (1.0 + image_F[0].permute(1, 2, 0).detach().cpu().numpy()) * 255.0)
+        image_B = (
+            0.5 * (1.0 + image_B[0].permute(1, 2, 0).detach().cpu().numpy()) * 255.0)
+        image_P = (
+            0.5 * (1.0 + image_P[0].permute(1, 2, 0).detach().cpu().numpy()) * 255.0)
+
+        vis_list.append(vedo.Picture(image_P*0.5+image_F *
+                        0.5).scale(2.0/image_P.shape[0]).pos(-1.0, -1.0, 1.0))
+        vis_list.append(vedo.Picture(image_F).scale(
+            2.0/image_F.shape[0]).pos(-1.0, -1.0, -0.5))
+        vis_list.append(vedo.Picture(image_B).scale(
+            2.0/image_B.shape[0]).pos(-1.0, -1.0, -1.0))
+
+        # create a mesh
+        mesh = trimesh.Trimesh(smpl_verts, faces, process=False)
+        mesh.visual.vertex_colors = [200, 200, 0]
+        vis_list.append(mesh)
+
+        vp.show(*vis_list, bg="white", axes=1, interactive=True)
+
+
+if __name__ == '__main__':
+
+    cfg.merge_from_file("./configs/icon-filter.yaml")
+    cfg.merge_from_file('./lib/pymaf/configs/pymaf_config.yaml')
+
+    cfg_show_list = [
+        'test_gpus', ['0'], 'mcube_res', 512, 'clean_mesh', False
+    ]
+
+    cfg.merge_from_list(cfg_show_list)
+    cfg.freeze()
+
+    os.environ['CUDA_VISIBLE_DEVICES'] = "0"
+    device = torch.device('cuda:0')
+
+    dataset = TestDataset(
+        {
+            'image_dir': "./examples",
+            'has_det': True,    # w/ or w/o detection
+            'hps_type': 'bev'  # pymaf/pare/pixie/hybrik/bev
+        }, device)
+
+    for i in range(len(dataset)):
+        dataset.visualize_alignment(dataset[i])

lib/dataset/mesh_util.py

@@ -17,17 +17,18 @@
 
 import numpy as np
 import cv2
+import pymeshlab
 import torch
 import torchvision
 import trimesh
 from pytorch3d.io import load_obj
+import os
 from termcolor import colored
+import os.path as osp
 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
 
@@ -41,24 +42,6 @@ from pytorch3d.loss import (
     mesh_normal_consistency
 )
 
-from huggingface_hub import hf_hub_download, hf_hub_url, cached_download
-
-def rot6d_to_rotmat(x):
-    """Convert 6D rotation representation to 3x3 rotation matrix.
-    Based on Zhou et al., "On the Continuity of Rotation Representations in Neural Networks", CVPR 2019
-    Input:
-        (B,6) Batch of 6-D rotation representations
-    Output:
-        (B,3,3) Batch of corresponding rotation matrices
-    """
-    x = x.view(-1, 3, 2)
-    a1 = x[:, :, 0]
-    a2 = x[:, :, 1]
-    b1 = F.normalize(a1)
-    b2 = F.normalize(a2 - torch.einsum("bi,bi->b", b1, a2).unsqueeze(-1) * b1)
-    b3 = torch.cross(b1, b2)
-    return torch.stack((b1, b2, b3), dim=-1)
-
 
 def tensor2variable(tensor, device):
     # [1,23,3,3]
@@ -137,18 +120,32 @@ def mesh_edge_loss(meshes, target_length: float = 0.0):
 
 
 def remesh(obj_path, perc, device):
-    mesh = trimesh.load(obj_path)
-    mesh = mesh.simplify_quadratic_decimation(50000)
-    mesh = trimesh.smoothing.filter_humphrey(
-        mesh, alpha=0.1, beta=0.5, iterations=10, laplacian_operator=None
-    )
-    mesh.export(obj_path.replace("recon", "remesh"))
-    verts_pr = torch.tensor(mesh.vertices).float().unsqueeze(0).to(device)
-    faces_pr = torch.tensor(mesh.faces).long().unsqueeze(0).to(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
@@ -208,31 +205,33 @@ def load_checkpoint(model, cfg):
 
     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'))
+    if os.path.exists(cfg.resume_path) and cfg.resume_path.endswith("ckpt"):
+        main_dict = torch.load(cfg.resume_path,
+                               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'))
+
+    if os.path.exists(cfg.normal_path) and cfg.normal_path.endswith("ckpt"):
+        normal_dict = torch.load(cfg.normal_path,
+                                 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)
@@ -253,7 +252,7 @@ def load_checkpoint(model, cfg):
 
 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']))
+    smpl_vtx_std = np.loadtxt(os.path.join(folder, 'vertices.txt'))
     min_x = np.min(smpl_vtx_std[:, 0])
     max_x = np.max(smpl_vtx_std[:, 0])
     min_y = np.min(smpl_vtx_std[:, 1])
@@ -266,12 +265,12 @@ def read_smpl_constants(folder):
     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']),
+    smpl_faces = np.loadtxt(os.path.join(folder, 'faces.txt'),
                             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']),
+    smpl_tetras = np.loadtxt(os.path.join(folder, 'tetrahedrons.txt'),
                              dtype=np.int32) - 1
 
     return smpl_vertex_code, smpl_face_code, smpl_faces, smpl_tetras
@@ -397,12 +396,11 @@ def cal_sdf_batch(verts, faces, cmaps, vis, points):
     bary_weights = barycentric_coordinates_of_projection(
         points.view(-1, 3), closest_triangles)
 
-    pts_cmap = (closest_cmaps*bary_weights[:, :, None]).sum(1).unsqueeze(0).clamp_(min=0.0, max=1.0)
+    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_norm = F.normalize(pts_norm, dim=2)
     pts_dist = torch.sqrt(residues) / torch.sqrt(torch.tensor(3))
 
     pts_signs = 2.0 * (check_sign(verts, faces[0], points).float() - 0.5)
@@ -844,21 +842,26 @@ def mesh_move(mesh_lst, step, scale=1.0):
 
 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.current_dir = osp.join(osp.dirname(__file__),
+                                    "../../data/smpl_related")
+
+        self.smpl_verts_path = osp.join(self.current_dir,
+                                        "smpl_data/smpl_verts.npy")
+        self.smplx_verts_path = osp.join(self.current_dir,
+                                         "smpl_data/smplx_verts.npy")
+        self.faces_path = osp.join(self.current_dir,
+                                   "smpl_data/smplx_faces.npy")
+        self.cmap_vert_path = osp.join(self.current_dir,
+                                       "smpl_data/smplx_cmap.npy")
 
         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')
-        
+        self.model_dir = osp.join(self.current_dir, "models")
+        self.tedra_dir = osp.join(self.current_dir, "../tedra_data")
+
     def get_smpl_mat(self, vert_ids):
 
         mat = torch.as_tensor(np.load(self.cmap_vert_path)).float()

lib/net/FBNet.py

@@ -81,8 +81,7 @@ def define_G(input_nc,
     # print(netG)
     if len(gpu_ids) > 0:
         assert (torch.cuda.is_available())
-        device=torch.device(f"cuda:{gpu_ids[0]}")
-        netG = netG.to(device)
+        netG.cuda(gpu_ids[0])
     netG.apply(weights_init)
     return netG
 

lib/net/HGPIFuNet.py

@@ -15,7 +15,7 @@
 #
 # Contact: [email protected]
 
-# from lib.net.voxelize import Voxelization
+from lib.net.voxelize import Voxelization
 from lib.dataset.mesh_util import cal_sdf_batch, feat_select, read_smpl_constants
 from lib.net.NormalNet import NormalNet
 from lib.net.MLP import MLP
@@ -26,6 +26,7 @@ from termcolor import colored
 from lib.net.BasePIFuNet import BasePIFuNet
 import torch.nn as nn
 import torch
+import os
 
 
 maskout = False
@@ -293,8 +294,14 @@ class HGPIFuNet(BasePIFuNet):
             # smpl_cmap [B, N, 3]
             # smpl_vis [B, N, 1]
 
+            # set ourlier point features as uniform values
+            smpl_outlier = torch.abs(smpl_sdf).ge(self.sdf_clip)
+            smpl_sdf[smpl_outlier] = torch.sign(smpl_sdf[smpl_outlier])
+
             feat_lst = [smpl_sdf]
             if 'cmap' in self.smpl_feats:
+                smpl_cmap[smpl_outlier.repeat(
+                    1, 1, 3)] = smpl_sdf[smpl_outlier].repeat(1, 1, 3)
                 feat_lst.append(smpl_cmap)
             if 'norm' in self.smpl_feats:
                 feat_lst.append(smpl_norm)

lib/net/net_util.py

@@ -316,7 +316,7 @@ class Vgg19(torch.nn.Module):
 class VGGLoss(nn.Module):
     def __init__(self):
         super(VGGLoss, self).__init__()
-        self.vgg = Vgg19()
+        self.vgg = Vgg19().cuda()
         self.criterion = nn.L1Loss()
         self.weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0]
 

lib/pymaf/core/path_config.py

@@ -6,18 +6,33 @@ for the datasets and data files necessary to run the code.
 Things you need to change: *_ROOT that indicate the path to each dataset
 """
 import os
-from huggingface_hub import hf_hub_url, cached_download
 
 # pymaf
-pymaf_data_dir = hf_hub_url('Yuliang/PyMAF', '')
-smpl_data_dir = hf_hub_url('Yuliang/SMPL', '')
-SMPL_MODEL_DIR = os.path.join(smpl_data_dir, 'models/smpl')
+pymaf_data_dir = os.path.join(os.path.dirname(__file__),
+                              "../../../data/pymaf_data")
 
-SMPL_MEAN_PARAMS = cached_download(os.path.join(pymaf_data_dir, 'smpl_mean_params.npz'), use_auth_token=os.environ['ICON'])
-MESH_DOWNSAMPLEING = cached_download(os.path.join(pymaf_data_dir, 'mesh_downsampling.npz'), use_auth_token=os.environ['ICON'])
-CUBE_PARTS_FILE = cached_download(os.path.join(pymaf_data_dir, 'cube_parts.npy'), use_auth_token=os.environ['ICON'])
-JOINT_REGRESSOR_TRAIN_EXTRA = cached_download(os.path.join(pymaf_data_dir, 'J_regressor_extra.npy'), use_auth_token=os.environ['ICON'])
-JOINT_REGRESSOR_H36M = cached_download(os.path.join(pymaf_data_dir, 'J_regressor_h36m.npy'), use_auth_token=os.environ['ICON'])
-VERTEX_TEXTURE_FILE = cached_download(os.path.join(pymaf_data_dir, 'vertex_texture.npy'), use_auth_token=os.environ['ICON'])
-SMPL_MEAN_PARAMS = cached_download(os.path.join(pymaf_data_dir, 'smpl_mean_params.npz'), use_auth_token=os.environ['ICON'])
-CHECKPOINT_FILE = cached_download(os.path.join(pymaf_data_dir, 'pretrained_model/PyMAF_model_checkpoint.pt'), use_auth_token=os.environ['ICON'])
+SMPL_MEAN_PARAMS = os.path.join(pymaf_data_dir, 'smpl_mean_params.npz')
+SMPL_MODEL_DIR = os.path.join(pymaf_data_dir, '../smpl_related/models/smpl')
+
+CUBE_PARTS_FILE = os.path.join(pymaf_data_dir, 'cube_parts.npy')
+JOINT_REGRESSOR_TRAIN_EXTRA = os.path.join(pymaf_data_dir,
+                                           'J_regressor_extra.npy')
+JOINT_REGRESSOR_H36M = os.path.join(pymaf_data_dir, 'J_regressor_h36m.npy')
+VERTEX_TEXTURE_FILE = os.path.join(pymaf_data_dir, 'vertex_texture.npy')
+SMPL_MEAN_PARAMS = os.path.join(pymaf_data_dir, 'smpl_mean_params.npz')
+SMPL_MODEL_DIR = os.path.join(pymaf_data_dir, '../smpl_related/models/smpl')
+CHECKPOINT_FILE = os.path.join(pymaf_data_dir,
+                               'pretrained_model/PyMAF_model_checkpoint.pt')
+
+# pare
+pare_data_dir = os.path.join(os.path.dirname(__file__),
+                             "../../../data/pare_data")
+CFG = os.path.join(pare_data_dir, 'pare/checkpoints/pare_w_3dpw_config.yaml')
+CKPT = os.path.join(pare_data_dir,
+                    'pare/checkpoints/pare_w_3dpw_checkpoint.ckpt')
+
+# hybrik
+hybrik_data_dir = os.path.join(os.path.dirname(__file__),
+                               "../../../data/hybrik_data")
+HYBRIK_CFG = os.path.join(hybrik_data_dir, 'hybrik_config.yaml')
+HYBRIK_CKPT = os.path.join(hybrik_data_dir, 'pretrained_w_cam.pth')

lib/pymaf/models/maf_extractor.py

@@ -3,13 +3,13 @@
 from packaging import version
 import torch
 import scipy
+import os
 import numpy as np
 import torch.nn as nn
 import torch.nn.functional as F
 
 from lib.common.config import cfg
 from lib.pymaf.utils.geometry import projection
-from lib.pymaf.core.path_config import MESH_DOWNSAMPLEING
 
 import logging
 
@@ -48,7 +48,10 @@ class MAF_Extractor(nn.Module):
 
         # downsample SMPL mesh and assign part labels
         # from https://github.com/nkolot/GraphCMR/blob/master/data/mesh_downsampling.npz
-        smpl_mesh_graph = np.load(MESH_DOWNSAMPLEING,
+        mesh_downsampling_path = os.path.join(
+            os.path.dirname(__file__),
+            "../../../data/pymaf_data/mesh_downsampling.npz")
+        smpl_mesh_graph = np.load(mesh_downsampling_path,
                                   allow_pickle=True,
                                   encoding='latin1')
 

lib/pymaf/models/res_module.py

@@ -4,11 +4,11 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import os
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from collections import OrderedDict
-import os
 from lib.pymaf.core.cfgs import cfg
 
 import logging

lib/pymaf/models/smpl.py

@@ -2,9 +2,9 @@
 
 import torch
 import numpy as np
-from lib.smplx import SMPL as _SMPL
-from lib.smplx.body_models import ModelOutput
-from lib.smplx.lbs import vertices2joints
+from smplx import SMPL as _SMPL
+from smplx.body_models import ModelOutput
+from smplx.lbs import vertices2joints
 from collections import namedtuple
 
 from lib.pymaf.core import path_config, constants

lib/pymaf/utils/imutils.py

@@ -1,14 +1,15 @@
 """
 This file contains functions that are used to perform data augmentation.
 """
+from turtle import reset
 import cv2
 import io
 import torch
 import numpy as np
+import scipy.misc
 from PIL import Image
-from rembg import remove
-from rembg.session_factory import new_session
-from torchvision.models import detection
+from rembg.bg import remove
+import human_det
 
 from lib.pymaf.core import constants
 from lib.pymaf.utils.streamer import aug_matrix
@@ -44,7 +45,6 @@ def get_bbox(img, det):
     bbox = bboxes[0, 0, 0].cpu().numpy()
 
     return bbox
-# Michael Black is
 
 
 def get_transformer(input_res):
@@ -86,7 +86,7 @@ def get_transformer(input_res):
     return [image_to_tensor, mask_to_tensor, image_to_pymaf_tensor, image_to_pixie_tensor, image_to_hybrik_tensor]
 
 
-def process_image(img_file, hps_type, input_res=512, device=None, seg_path=None):
+def process_image(img_file, det, hps_type, input_res=512, device=None, seg_path=None):
     """Read image, do preprocessing and possibly crop it according to the bounding box.
     If there are bounding box annotations, use them to crop the image.
     If no bounding box is specified but openpose detections are available, use them to get the bounding box.
@@ -104,19 +104,21 @@ def process_image(img_file, hps_type, input_res=512, device=None, seg_path=None)
     img_for_crop = cv2.warpAffine(img_ori, M[0:2, :],
                                   (input_res*2, input_res*2), flags=cv2.INTER_CUBIC)
 
-    # detection for bbox
-    detector = detection.maskrcnn_resnet50_fpn(pretrained=True)
-    detector.eval()
-    predictions = detector(
-        [torch.from_numpy(img_for_crop).permute(2, 0, 1) / 255.])[0]
-    human_ids = torch.where(
-        predictions["scores"] == predictions["scores"][predictions['labels'] == 1].max())
-    bbox = predictions["boxes"][human_ids, :].flatten().detach().cpu().numpy()
-    
-    width = bbox[2] - bbox[0]
-    height = bbox[3] - bbox[1]
-    center = np.array([(bbox[0] + bbox[2]) / 2.0,
-                        (bbox[1] + bbox[3]) / 2.0])
+    if det is not None:
+
+        # detection for bbox
+        bbox = get_bbox(img_for_crop, det)
+
+        width = bbox[2] - bbox[0]
+        height = bbox[3] - bbox[1]
+        center = np.array([(bbox[0] + bbox[2]) / 2.0,
+                           (bbox[1] + bbox[3]) / 2.0])
+
+    else:
+        # Assume that the person is centerered in the image
+        height = img_for_crop.shape[0]
+        width = img_for_crop.shape[1]
+        center = np.array([width // 2, height // 2])
 
     scale = max(height, width) / 180
 
@@ -127,8 +129,12 @@ def process_image(img_file, hps_type, input_res=512, device=None, seg_path=None)
         img_np, cropping_parameters = crop(
             img_for_crop, center, scale, (input_res, input_res))
 
-    img_pil = Image.fromarray(remove(img_np, post_process_mask=True, session=new_session("u2net")))
-    
+    with torch.no_grad():
+        buf = io.BytesIO()
+        Image.fromarray(img_np).save(buf, format='png')
+        img_pil = Image.open(
+            io.BytesIO(remove(buf.getvalue()))).convert("RGBA")
+
     # for icon
     img_rgb = image_to_tensor(img_pil.convert("RGB"))
     img_mask = torch.tensor(1.0) - (mask_to_tensor(img_pil.split()[-1]) <

lib/renderer/mesh.py

@@ -18,7 +18,7 @@
 from lib.dataset.mesh_util import SMPLX
 from lib.common.render_utils import face_vertices
 import numpy as np
-import lib.smplx as smplx
+import smplx
 import trimesh
 import torch
 import torch.nn.functional as F