support PaMIR

YuliangXiu · Jun 15, 2022 · 23b3d13 · 23b3d13
1 parent 55a74e0
commit 23b3d13
Show file tree

Hide file tree

Showing 10 changed files with 246 additions and 116 deletions.
diff --git a/README.md b/README.md
@@ -24,7 +24,7 @@
   <br>
     <a href="https://pytorch.org/get-started/locally/"><img alt="PyTorch" src="https://img.shields.io/badge/PyTorch-ee4c2c?logo=pytorch&logoColor=white"></a>
     <a href="https://pytorchlightning.ai/"><img alt="Lightning" src="https://img.shields.io/badge/-Lightning-792ee5?logo=pytorchlightning&logoColor=white"></a>
-    <a href="https://paperswithcode.com/sota/3d-human-reconstruction-on-cape?p=icon-implicit-clothed-humans-obtained-from"><img src="https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/icon-implicit-clothed-humans-obtained-from/3d-human-reconstruction-on-cape"></a><br></br>
+    <a href='https://colab.research.google.com/drive/1-AWeWhPvCTBX0KfMtgtMk10uPU05ihoA?usp=sharing' style='padding-left: 0.5rem;'><img src='https://colab.research.google.com/assets/colab-badge.svg' alt='Google Colab'></a><br></br>
     <a href='https://arxiv.org/abs/2112.09127'>
       <img src='https://img.shields.io/badge/Paper-PDF-green?style=for-the-badge&logo=arXiv&logoColor=green' alt='Paper PDF'>
     </a>
@@ -56,9 +56,6 @@
     <li>
       <a href="#who-needs-ICON">Who needs ICON</a>
     </li>
-    <li>
-      <a href="#todo">TODO</a>
-    </li>
     <li>
       <a href="#installation">Installation</a>
     </li>
@@ -99,10 +96,16 @@
 
 ## Who needs ICON?
 
-- Given an RGB image, you could get:
-  - image (png): segmentation, normal images (body + cloth), overlap result (rgb + normal)
-  - mesh (obj): SMPL-(X) body, reconstructed clothed human
-  - video (mp4): self-rotated clothed human
+- Given a raw RGB image, you could get:
+  - image (png): 
+    - segmented human RGB
+    - normal maps of body and cloth
+    - pixel-aligned normal-RGB overlap
+  - mesh (obj):
+    - SMPL-(X) body from *PyMAF, PIXIE, PARE, HybrIK, BEV*
+    - 3D clothed human reconstruction
+  - video (mp4): 
+    - self-rotated clothed human
 
 |          ![Intermediate Results](assets/intermediate_results.png)          |
 | :------------------------------------------------------------------------: |
@@ -121,39 +124,24 @@
 | :----------------------------------------------------------------------------------------: |
 | _3D Clothed Avatar, created from 400+ images using **ICON+SCANimate**, animated by AIST++_ |
 
-## TODO
-
-- [x] testing code and pretrained models (\*self-implemented version)
-  - [x] ICON (w/ & w/o global encoder, w/ PyMAF/HybrIK/BEV/PIXIE/PARE as HPS)
-  - [x] PIFu\* (RGB image + predicted normal map as input)
-  - [x] PaMIR\* (RGB image + predicted normal map as input, w/ PyMAF/HybrIK/BEV/PIXIE/PARE as HPS)
-- [x] colab notebook <a href='https://colab.research.google.com/drive/1-AWeWhPvCTBX0KfMtgtMk10uPU05ihoA?usp=sharing' style='padding-left: 0.5rem;'>
-      <img src='https://colab.research.google.com/assets/colab-badge.svg' alt='Google Colab'>
-      </a>
-- [x] dataset processing
-  - [x] RGB/Normal Rendering
-  - [x] Visibility computing
-- [ ] training code
-  - [x] ICON (w/ & w/o global encoder, w/ SMPL-X)
-  - [x] PIFu\* (RGB image + predicted normal map as input)
-  - [ ] PaMIR\* (RGB image + predicted normal map as input, w/ voxelized SMPL)
-- [x] evaluation code
+- If you want to **Train & Evaluate** on **PIFu/PaMIR/ICON** using your own data, please check [dataset.md](./docs/dataset.md) to prepare dataset, [training.md](./docs/training.md) for training, and [evaluation.md](./docs/evaluation.md) for benchmark evaluation.
 
+<br><br>
 ## Installation
 
 Please follow the [Installation Instruction](docs/installation.md) to setup all the required packages, extra data, and models.
 
 ## Dataset
 
-Please follow the [Dataset Instruction](docs/dataset.md) to generate the train/val/test dataset from raw scans (THuman2.0).
+Please follow the [Dataset Instruction](docs/dataset.md) to generate the train/val/test dataset from THuman2.0
 
 ## Training
 
-Please follow the [Training Instruction](docs/training.md) to train your own model using THuman2.0.
+Please follow the [Training Instruction](docs/training.md) to train your own model using THuman2.0
 
 ## Evaluation
 
-Please follow the [Evaluation Instruction](docs/evaluation.md) to benchmark models on THuman2.0.
+Please follow the [Evaluation Instruction](docs/evaluation.md) to benchmark trained models on THuman2.0
 
 ## Add-on
 1. [Garment Extraction from Fashion Images](docs/garment-extraction.md), supported by ETH Zürich students as 3DV course project.
@@ -221,6 +209,12 @@ Some images used in the qualitative examples come from [pinterest.com](https://w
 
 This project has received funding from the European Union’s Horizon 2020 research and innovation programme under the Marie Skłodowska-Curie grant agreement No.860768 ([CLIPE Project](https://www.clipe-itn.eu)).
 
+<br>
+
+--------------
+
+<br>
+
 ## License
 
 This code and model are available for non-commercial scientific research purposes as defined in the [LICENSE](LICENSE) file. By downloading and using the code and model you agree to the terms in the [LICENSE](LICENSE).

diff --git a/apps/ICON.py b/apps/ICON.py
@@ -536,13 +536,10 @@ def test_step(self, batch, batch_idx):
         # export paths
         mesh_name = batch["subject"][0]
         mesh_rot = batch["rotation"][0].item()
-        ckpt_dir = self.cfg.resume_path.split("/")[-2]
+        ckpt_dir = self.cfg.name
 
-        if (self.cfg.dataset.online_smpl) or (self.cfg.optim_body):
-            for kid, key in enumerate(self.cfg.dataset.noise_type):
-                ckpt_dir += f"_{key}_{self.cfg.dataset.noise_scale[kid]}"
-        else:
-            ckpt_dir += "_perfect_smpl"
+        for kid, key in enumerate(self.cfg.dataset.noise_type):
+            ckpt_dir += f"_{key}_{self.cfg.dataset.noise_scale[kid]}"
 
         if self.cfg.optim_cloth:
             ckpt_dir += "_optim_cloth"

diff --git a/configs/train/pamir.yaml b/configs/train/pamir.yaml
@@ -21,7 +21,7 @@ dataset:
   zray_type: False
   semantic_p: False
   remove_outlier: False
-
+  
   noise_type: ['pose', 'beta']
   noise_scale: [0.05, 0.5]
 

diff --git a/docs/dataset.md b/docs/dataset.md
@@ -11,7 +11,15 @@
 
 Please refer to [THuman2.0-Dataset](https://github.com/ytrock/THuman2.0-Dataset) to download the original scans into `data/thuman2/scans`, and its SMPL-X(male) fits into `data/thuman2/fits`. Then generate `all.txt` by `ls > ../all.txt` under `data/thuman2/scans`, which contains all the subject names (0000~0525).
 
-:eyes: `./sample_data` contains one example of THuman2.0 and shows the data folder structure.
+PaMIR only support SMPL instead of SMPL-X, please download THuman2.0's SMPL fits (.obj, .pkl) and put them at `./data/thuman2.0/smpl`
+
+```bash
+wget https://download.is.tue.mpg.de/icon/smpl.zip --no-check-certificate -O ./data/thuman2/smpl.zip
+unzip ./data/thuman2/smpl.zip -d ./data/thuman2/
+rm ./data/thuman2/smpl.zip
+```
+
+:eyes: `./sample_data` contains one example of THuman2.0 which shows the data folder structure.
 
 ## Debug Mode
 

diff --git a/docs/evaluation.md b/docs/evaluation.md
@@ -4,8 +4,7 @@
 
 Make sure you have already generated all the required synthetic data (refer to [Dataset Instruction](dataset.md)) under `./data/thuman2_{num_views}views`, which includes the rendered RGB (`render/`), normal images(`normal_B/`, `normal_F/`, `T_normal_B/`, `T_normal_F/`), corresponding calibration matrix (`calib/`) and pre-computed visibility arrays (`vis/`).
 
-
-:warning: Don't support headless mode currently, `unset PYOPENGL_PLATFORM` before training. (will fix it later...)
+:warning: Don't support headless mode currently, `unset PYOPENGL_PLATFORM` before training.
 
 ## Command
 
@@ -18,21 +17,11 @@ CUDA_VISIBLE_DEVICES=0 python -m apps.train -cfg ./configs/train/icon-filter.yam
 # ICON w/o filter
 CUDA_VISIBLE_DEVICES=0 python -m apps.train -cfg ./configs/train/icon-nofilter.yaml -test
 
-# PIFu
+# PIFu* (*: re-implementation)
 CUDA_VISIBLE_DEVICES=0 python -m apps.train -cfg ./configs/train/pifu.yaml -test
 
-# PaMIR
+# PaMIR* (*: re-implementation)
 CUDA_VISIBLE_DEVICES=0 python -m apps.train -cfg ./configs/train/pamir.yaml -test
 ```
 
-## Intermediate Results
-
-All the intermediate results are exported at `./results/ckpt_perfect_smpl`
-
-## Benchmark on THuman2.0 (testset: 0500-0504, 3 views)
-
-|Metrics|ICON w/ filter|ICON w/o filter|PIFu|PaMIR|
-|---|---|---|---|---|
-|Chamfer|1.068|1.218|1.711|-|
-|P2S|1.068|1.210|1.159|-|
-|NC|0.061|0.073|0.075|-|
+The qualitative results are located at `./results/`
diff --git a/docs/training.md b/docs/training.md
@@ -5,15 +5,21 @@
 Make sure you have already generated all the required synthetic data (refer to [Dataset Instruction](dataset.md)) under `./data/thuman2_{num_views}views`, which includes the rendered RGB (`render/`), normal images(`normal_B/`, `normal_F/`, `T_normal_B/`, `T_normal_F/`), corresponding calibration matrix (`calib/`) and pre-computed visibility arrays (`vis/`).
 
 :eyes: Test your dataloader with [vedo](https://vedo.embl.es/)
+
 ```bash
-python -m lib.dataloader_demo -v
+
+# visualization for SMPL-X mesh
+python -m lib.dataloader_demo -v -c ./configs/train/icon-filter.yaml
+
+# visualization for voxelized SMPL
+python -m lib.dataloader_demo -v -c ./configs/train/pamir.yaml
 ```
 
 <p align="center">
     <img src="../assets/vedo.gif" width=50%>
 </p>
 
-:warning: Don't support headless mode currently, `unset PYOPENGL_PLATFORM` before training. (will fix it later...)
+:warning: Don't support headless mode currently, `unset PYOPENGL_PLATFORM` before training.
 ## Command
 
 ```bash
@@ -25,12 +31,15 @@ CUDA_VISIBLE_DEVICES=0 python -m apps.train -cfg ./configs/train/icon-filter.yam
 # ICON w/o filter (name: icon-nofilter)
 CUDA_VISIBLE_DEVICES=0 python -m apps.train -cfg ./configs/train/icon-nofilter.yaml
 
-# ICON-MVP (name: icon-mvp), mvp = minimal viable product, simple (used features) yet efficient (GPU)
+# ICON-MVP (name: icon-mvp), mvp = minimal viable product, simple (used features) yet efficient (GPU usage)
 # https://en.wikipedia.org/wiki/Minimum_viable_product
 CUDA_VISIBLE_DEVICES=0 python -m apps.train -cfg ./configs/train/icon-mvp.yaml
 
 # PIFu (name: pifu)
 CUDA_VISIBLE_DEVICES=0 python -m apps.train -cfg ./configs/train/pifu.yaml
+
+# PaMIR (name: pamir)
+CUDA_VISIBLE_DEVICES=0 python -m apps.train -cfg ./configs/train/pamir.yaml
 ```
 
 ## Tensorboard

diff --git a/lib/dataloader_demo.py b/lib/dataloader_demo.py
@@ -2,52 +2,57 @@
 from lib.common.config import get_cfg_defaults
 from lib.dataset.PIFuDataset import PIFuDataset
 
-args = get_cfg_defaults()
-args.merge_from_file("./configs/train/icon-filter.yaml")
+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()
 
-# loading cfg file
-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')
-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]
+    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.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.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("--------------------")
+    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("--------------------")