├── LICENSE ├── README.md ├── assets ├── smpl_faces.npy ├── template_mesh_uv.obj └── uv_masks │ ├── idx_smpl_posmap32_uniformverts_retrieval.npy │ └── uv_mask32_with_faceid_smpl.npy ├── chamferdist ├── .gitignore ├── LICENSE ├── README.md ├── chamferdist │ ├── ChamferDistance.py │ ├── __init__.py │ ├── chamfer.cu │ └── chamfer_cuda.cpp ├── example.py └── setup.py ├── configs ├── config_demo.yaml └── config_train_demo.yaml ├── data └── README_data.md ├── lib ├── __init__.py ├── config_parser.py ├── dataset.py ├── losses.py ├── modules.py ├── network.py ├── train_eval_funcs.py ├── utils_io.py ├── utils_model.py └── utils_train.py ├── lib_data ├── README.md ├── __init__.py ├── pack_data_example.py └── posmap_generator │ ├── __init__.py │ ├── apps │ ├── __init__.py │ └── example.py │ └── lib │ ├── __init__.py │ └── renderer │ ├── __init__.py │ ├── camera.py │ ├── core.py │ ├── egl │ ├── __init__.py │ ├── data │ │ ├── pos_uv.fs │ │ ├── pos_uv.vs │ │ ├── quad.fs │ │ └── quad.vs │ ├── egl_cam_render.py │ ├── egl_framework.py │ ├── egl_pos_render.py │ ├── egl_render.py │ └── glcontext.py │ ├── gl │ ├── __init__.py │ ├── cam_render.py │ ├── data │ │ ├── pos_uv.fs │ │ ├── pos_uv.vs │ │ ├── quad.fs │ │ └── quad.vs │ ├── framework.py │ ├── pos_render.py │ └── render.py │ ├── glm.py │ ├── mesh.py │ ├── ram_zip.py │ └── tsdf.py ├── main.py ├── render ├── cam_front_extrinsic.npy └── o3d_render_pcl.py ├── requirements.txt └── teasers └── teaser.gif /LICENSE: -------------------------------------------------------------------------------- 1 | License 2 | 3 | Software Copyright License for non-commercial scientific research purposes 4 | Please read carefully the following terms and conditions and any accompanying documentation before you download and/or use the SCALE software, (the "Software"), including 3D meshes, images, videos, textures, software, scripts, and animations. By downloading and/or using the Software (including downloading, cloning, installing, and any other use of the corresponding github repository), you acknowledge that you have read these terms and conditions, understand them, and agree to be bound by them. If you do not agree with these terms and conditions, you must not download and/or use the Software. Any infringement of the terms of this agreement will automatically terminate your rights under this License. 5 | 6 | 7 | Ownership / Licensees 8 | The Software and the associated materials has been developed at the 9 | 10 | Max Planck Institute for Intelligent Systems (hereinafter "MPI"). 11 | 12 | Any copyright or patent right is owned by and proprietary material of the 13 | 14 | Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (hereinafter “MPG”; MPI and MPG hereinafter collectively “Max-Planck”) 15 | 16 | hereinafter the “Licensor”. 17 | 18 | 19 | License Grant 20 | Licensor grants you (Licensee) personally a single-user, non-exclusive, non-transferable, free of charge right: 21 | 22 | To install the Software on computers owned, leased or otherwise controlled by you and/or your organization; 23 | To use the Software for the sole purpose of performing non-commercial scientific research, non-commercial education, or non-commercial artistic projects; 24 | Any other use, in particular any use for commercial, pornographic, military, or surveillance, purposes is prohibited. This includes, without limitation, incorporation in a commercial product, use in a commercial service, or production of other artefacts for commercial purposes. The Software may not be used to create fake, libelous, misleading, or defamatory content of any kind excluding analyses in peer-reviewed scientific research. The Software may not be reproduced, modified and/or made available in any form to any third party without Max-Planck’s prior written permission. 25 | 26 | The Software may not be used for pornographic purposes or to generate pornographic material whether commercial or not. This license also prohibits the use of the Software to train methods/algorithms/neural networks/etc. for commercial, pornographic, military, surveillance, or defamatory use of any kind. By downloading the Software, you agree not to reverse engineer it. 27 | 28 | 29 | No Distribution 30 | The Software and the license herein granted shall not be copied, shared, distributed, re-sold, offered for re-sale, transferred or sub-licensed in whole or in part except that you may make one copy for archive purposes only. 31 | 32 | 33 | Disclaimer of Representations and Warranties 34 | You expressly acknowledge and agree that the Software results from basic research, is provided “AS IS”, may contain errors, and that any use of the Software is at your sole risk. LICENSOR MAKES NO REPRESENTATIONS OR WARRANTIES OF ANY KIND CONCERNING THE SOFTWARE, NEITHER EXPRESS NOR IMPLIED, AND THE ABSENCE OF ANY LEGAL OR ACTUAL DEFECTS, WHETHER DISCOVERABLE OR NOT. Specifically, and not to limit the foregoing, licensor makes no representations or warranties (i) regarding the merchantability or fitness for a particular purpose of the Software, (ii) that the use of the Software will not infringe any patents, copyrights or other intellectual property rights of a third party, and (iii) that the use of the Software will not cause any damage of any kind to you or a third party. 35 | 36 | 37 | Limitation of Liability 38 | Because this Software License Agreement qualifies as a donation, according to Section 521 of the German Civil Code (Bürgerliches Gesetzbuch – BGB) Licensor as a donor is liable for intent and gross negligence only. If the Licensor fraudulently conceals a legal or material defect, they are obliged to compensate the Licensee for the resulting damage. 39 | Licensor shall be liable for loss of data only up to the amount of typical recovery costs which would have arisen had proper and regular data backup measures been taken. For the avoidance of doubt Licensor shall be liable in accordance with the German Product Liability Act in the event of product liability. The foregoing applies also to Licensor’s legal representatives or assistants in performance. Any further liability shall be excluded. 40 | Patent claims generated through the usage of the Software cannot be directed towards the copyright holders. 41 | The Software is provided in the state of development the licensor defines. If modified or extended by Licensee, the Licensor makes no claims about the fitness of the Software and is not responsible for any problems such modifications cause. 42 | 43 | 44 | No Maintenance Services 45 | You understand and agree that Licensor is under no obligation to provide either maintenance services, update services, notices of latent defects, or corrections of defects with regard to the Software. Licensor nevertheless reserves the right to update, modify, or discontinue the Software at any time. 46 | 47 | Defects of the Software must be notified in writing to the Licensor with a comprehensible description of the error symptoms. The notification of the defect should enable the reproduction of the error. The Licensee is encouraged to communicate any use, results, modification or publication. 48 | 49 | 50 | Publications using the Software 51 | You acknowledge that the Software is a valuable scientific resource and agree to appropriately reference the following paper in any publication making use of the Software. 52 | 53 | Citation: 54 | 55 | @inproceedings{Ma:CVPR:2021, 56 | title = {{SCALE}: Modeling Clothed Humans with a Surface Codec of Articulated Local Elements}, 57 | author = {Ma, Qianli and Saito, Shunsuke and Yang, Jinlong and Tang, Siyu and Black, Michael J.}, 58 | booktitle = {Proceedings IEEE/CVF Conf.~on Computer Vision and Pattern Recognition (CVPR)}, 59 | month = jun, 60 | year = {2021}, 61 | month_numeric = {6} 62 | } 63 | 64 | 65 | Commercial licensing opportunities 66 | For commercial uses of the Software, please send email to ps-license@tue.mpg.de 67 | 68 | This Agreement shall be governed by the laws of the Federal Republic of Germany except for the UN Sales Convention. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # SCALE: Modeling Clothed Humans with a Surface Codec of Articulated Local Elements (CVPR 2021) 2 | 3 | [![Paper](https://img.shields.io/badge/arXiv-Paper-b31b1b.svg)](https://arxiv.org/abs/2104.07660) 4 | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1lp6r-A-s1kBorIvg6rLD4Ja3o6JOvu3G?usp=sharing) 5 | 6 | 7 | This repository contains the official PyTorch implementation of the CVPR 2021 paper: 8 | 9 | **SCALE: Modeling Clothed Humans with a Surface Codec of Articulated Local Elements**
10 | Qianli Ma, Shunsuke Saito, Jinlong Yang, Siyu Tang, and Michael J. Black
11 | [Full paper](https://arxiv.org/pdf/2104.07660) | [Video](https://youtu.be/-EvWqFCUb7U) | [Project website](https://qianlim.github.io/SCALE.html) | [Poster](https://ps.is.tuebingen.mpg.de/uploads_file/attachment/attachment/650/SCALE_poster_CVPR_final_compressed.pdf) 12 | 13 | ![](teasers/teaser.gif) 14 | 15 | 16 | ## Installation 17 | - The code has been tested with python 3.6 on both (Ubuntu 18.04 + CUDA 10.0) and (Ubuntu 20.04 + CUDA 11.1). 18 | 19 | - First, in the folder of this SCALE repository, run the following commands to create a new virtual environment and install dependencies: 20 | 21 | ```bash 22 | python3 -m venv $HOME/.virtualenvs/SCALE 23 | source $HOME/.virtualenvs/SCALE/bin/activate 24 | pip install -U pip setuptools 25 | pip install -r requirements.txt 26 | mkdir checkpoints 27 | ``` 28 | 29 | - Install the Chamfer Distance package (MIT license, taken from [this implementation](https://github.com/krrish94/chamferdist/tree/97051583f6fe72d5d4a855696dbfda0ea9b73a6a)). Note: the compilation is verified to be successful under CUDA 10.0, but may not be compatible with later CUDA versions. 30 | 31 | ```bash 32 | cd chamferdist 33 | python setup.py install 34 | cd .. 35 | ``` 36 | 37 | - You are now good to go with the next steps! All the commands below are assumed to be run from the `SCALE` repository folder, within the virtual environment created above. 38 | 39 | 40 | 41 | ## Run SCALE 42 | 43 | - Download our [pre-trained model weights](https://owncloud.tuebingen.mpg.de/index.php/s/pMYCtcpMDjk34Zw), unzip it under the `checkpoints` folder, such that the checkpoints' path is `/checkpoints/SCALE_demo_00000_simuskirt/`. 44 | 45 | - Download the [packed data for demo](https://owncloud.tuebingen.mpg.de/index.php/s/B33dqE5dcwbTbnQ), unzip it under the `data/` folder, such that the data file paths are `/data/packed/00000_simuskirt//`. 46 | 47 | - With the data and pre-trained model ready, the following code will generate a sequence of `.ply` files of the teaser dancing animation in `results/saved_samples/SCALE_demo_00000_simuskirt`: 48 | 49 | ```bash 50 | python main.py --config configs/config_demo.yaml 51 | ``` 52 | 53 | - To render images of the generated point sets, run the following command: 54 | 55 | ```bash 56 | python render/o3d_render_pcl.py --model_name SCALE_demo_00000_simuskirt 57 | ``` 58 | 59 | The images (with both the point normal coloring and patch coloring) will be saved under `results/rendered_imgs/SCALE_demo_00000_simuskirt`. 60 | 61 | 62 | 63 | ## Train SCALE 64 | 65 | ### Training demo with our data examples 66 | 67 | - Assume the demo training data is downloaded from the previous step under `data/packed/`. Now run: 68 | 69 | ```bash 70 | python main.py --config configs/config_train_demo.yaml 71 | ``` 72 | 73 | The training will start! 74 | 75 | - The code will also save the loss curves in the TensorBoard logs under `tb_logs//SCALE_train_demo_00000_simuskirt`. 76 | - Examples from the validation set at every 10 (can be set) epoch will be saved at `results/saved_samples/SCALE_train_demo_00000_simuskirt/val`. 77 | 78 | - Note: the training data provided above are only for demonstration purposes. Due to their very limited number of frames, they will not likely yield a satisfying model. Please refer to the README files in the `data/` and `lib_data/` folders for more information on how to process your customized data. 79 | 80 | ### Training with your own data 81 | 82 | We provide example codes in `lib_data/` to assist you in adapting your own data to the format required by SCALE. Please refer to [`lib_data/README`](./lib_data/README.md) for more details. 83 | 84 | 85 | 86 | ## News 87 | 88 | - [2021/10/29] **We now provide the packed, SCALE-compatible CAPE data on the [CAPE dataset website](https://cape.is.tue.mpg.de/download.php).** Simply register as a user there to access the download links (at the bottom of the Download page). 89 | - [2021/06/24] Code online! 90 | 91 | 92 | 93 | ## License 94 | 95 | Software Copyright License for non-commercial scientific research purposes. Please read carefully the [terms and conditions](./LICENSE) and any accompanying documentation before you download and/or use the SCALE code, including the scripts, animation demos and pre-trained models. By downloading and/or using the Model & Software (including downloading, cloning, installing, and any other use of this GitHub repository), you acknowledge that you have read these terms and conditions, understand them, and agree to be bound by them. If you do not agree with these terms and conditions, you must not download and/or use the Model & Software. Any infringement of the terms of this agreement will automatically terminate your rights under this [License](./LICENSE). 96 | 97 | The SMPL body related files (including `assets/{smpl_faces.npy, template_mesh_uv.obj}` and the UV masks under `assets/uv_masks/`) are subject to the license of the [SMPL model](https://smpl.is.tue.mpg.de/). The provided demo data (including the body pose and the meshes of clothed human bodies) are subject to the license of the [CAPE Dataset](https://cape.is.tue.mpg.de/). The Chamfer Distance implementation is subject to its [original license](./chamferdist/LICENSE). 98 | 99 | 100 | 101 | ## Related Research 102 | [SCANimate: Weakly Supervised Learning of Skinned Clothed Avatar Networks (CVPR 2021)](https://scanimate.is.tue.mpg.de/)
103 | *Shunsuke Saito, Jinlong Yang, Qianli Ma, Michael J. Black* 104 | 105 | Our *implicit* solution to pose-dependent shape modeling: cycle-consistent implicit skinning fields + locally pose-aware implicit function = a fully animatable avatar with implicit surface from raw scans without surface registration! 106 | 107 | [Learning to Dress 3D People in Generative Clothing (CVPR 2020)](https://cape.is.tue.mpg.de/)
108 | *Qianli Ma, Jinlong Yang, Anurag Ranjan, Sergi Pujades, Gerard Pons-Moll, Siyu Tang, Michael J. Black* 109 | 110 | CAPE --- a generative model and a large-scale dataset for 3D clothed human meshes in varied poses and garment types. 111 | We trained SCALE using the [CAPE dataset](https://cape.is.tue.mpg.de/dataset), check it out! 112 | 113 | 114 | 115 | 116 | ## Citations 117 | 118 | ```bibtex 119 | @inproceedings{Ma:CVPR:2021, 120 | title = {{SCALE}: Modeling Clothed Humans with a Surface Codec of Articulated Local Elements}, 121 | author = {Ma, Qianli and Saito, Shunsuke and Yang, Jinlong and Tang, Siyu and Black, Michael J.}, 122 | booktitle = {Proceedings IEEE/CVF Conf.~on Computer Vision and Pattern Recognition (CVPR)}, 123 | month = jun, 124 | year = {2021}, 125 | month_numeric = {6} 126 | } 127 | ``` 128 | 129 | -------------------------------------------------------------------------------- /assets/smpl_faces.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qianlim/SCALE/98f5557e746de3536168c413901ff48c0571c5a4/assets/smpl_faces.npy -------------------------------------------------------------------------------- /assets/uv_masks/idx_smpl_posmap32_uniformverts_retrieval.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qianlim/SCALE/98f5557e746de3536168c413901ff48c0571c5a4/assets/uv_masks/idx_smpl_posmap32_uniformverts_retrieval.npy -------------------------------------------------------------------------------- /assets/uv_masks/uv_mask32_with_faceid_smpl.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qianlim/SCALE/98f5557e746de3536168c413901ff48c0571c5a4/assets/uv_masks/uv_mask32_with_faceid_smpl.npy -------------------------------------------------------------------------------- /chamferdist/.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | 53 | # Translations 54 | *.mo 55 | *.pot 56 | 57 | # Django stuff: 58 | *.log 59 | local_settings.py 60 | db.sqlite3 61 | db.sqlite3-journal 62 | 63 | # Flask stuff: 64 | instance/ 65 | .webassets-cache 66 | 67 | # Scrapy stuff: 68 | .scrapy 69 | 70 | # Sphinx documentation 71 | docs/_build/ 72 | 73 | # PyBuilder 74 | target/ 75 | 76 | # Jupyter Notebook 77 | .ipynb_checkpoints 78 | 79 | # IPython 80 | profile_default/ 81 | ipython_config.py 82 | 83 | # pyenv 84 | .python-version 85 | 86 | # pipenv 87 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 88 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 89 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 90 | # install all needed dependencies. 91 | #Pipfile.lock 92 | 93 | # celery beat schedule file 94 | celerybeat-schedule 95 | 96 | # SageMath parsed files 97 | *.sage.py 98 | 99 | # Environments 100 | .env 101 | .venv 102 | env/ 103 | venv/ 104 | ENV/ 105 | env.bak/ 106 | venv.bak/ 107 | 108 | # Spyder project settings 109 | .spyderproject 110 | .spyproject 111 | 112 | # Rope project settings 113 | .ropeproject 114 | 115 | # mkdocs documentation 116 | /site 117 | 118 | # mypy 119 | .mypy_cache/ 120 | .dmypy.json 121 | dmypy.json 122 | 123 | # Pyre type checker 124 | .pyre/ -------------------------------------------------------------------------------- /chamferdist/LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2019 Krishna Murthy 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining a copy 4 | of this software and associated documentation files (the "Software"), to deal 5 | in the Software without restriction, including without limitation the rights 6 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | copies of the Software, and to permit persons to whom the Software is 8 | furnished to do so, subject to the following conditions: 9 | 10 | The above copyright notice and this permission notice shall be included in all 11 | copies or substantial portions of the Software. 12 | 13 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 19 | SOFTWARE. 20 | 21 | 22 | From AtlasNet: https://github.com/ThibaultGROUEIX/AtlasNet/blob/master/license_MIT 23 | 24 | Copyright (c) 25 | 26 | Permission is hereby granted, free of charge, to any person obtaining a copy 27 | of this software and associated documentation files (the "Software"), to deal 28 | in the Software without restriction, including without limitation the rights 29 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 30 | copies of the Software, and to permit persons to whom the Software is 31 | furnished to do so, subject to the following conditions: 32 | 33 | The above copyright notice and this permission notice shall be included in all 34 | copies or substantial portions of the Software. 35 | 36 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 37 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 38 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 39 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 40 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 41 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 42 | SOFTWARE. -------------------------------------------------------------------------------- /chamferdist/README.md: -------------------------------------------------------------------------------- 1 | This chamfer distance package is taken from [krrish94's implementation](https://github.com/krrish94/chamferdist/tree/97051583f6fe72d5d4a855696dbfda0ea9b73a6a), which is a wrapper on the implementation from [AtlasNet](https://github.com/ThibaultGROUEIX/AtlasNet) (MIT license). Below are the original README from krrish94's repository. 2 | 3 | 4 | 5 | --- 6 | 7 | ## chamferdist: PyTorch Chamfer distance 8 | 9 | > **NOTE**: This is a borrowed implementation from the elegant [AtlasNet](https://github.com/ThibaultGROUEIX/AtlasNet/tree/master/extension) GitHub repo, and all I did was to simply package it. 10 | 11 | A simple example Pytorch module to compute Chamfer distance between two pointclouds. Basically a wrapper around the elegant implementation from [AtlasNet](https://github.com/ThibaultGROUEIX/AtlasNet/tree/master/extension). 12 | 13 | ### Installation 14 | 15 | You can install the package using `pip`. 16 | 17 | ``` 18 | pip install chamferdist 19 | ``` 20 | 21 | ### Building from source 22 | 23 | In your favourite python/conda virtual environment, execute the following commands. 24 | 25 | > **NOTE**: This assumes you have PyTorch installed already (preferably, > 1.1.0; untested for earlier releases). 26 | 27 | ```python 28 | python setup.py install 29 | ``` 30 | 31 | ### Running (example) 32 | 33 | That's it! You're now ready to go. Here's a quick guide to using the package. Fire up a terminal. Import the package. 34 | 35 | ```python 36 | >>> import torch 37 | >>> from chamferdist import ChamferDist 38 | ``` 39 | 40 | Create two random pointclouds. Each pointcloud is a 3D tensor with dimensions `batchsize` x `number of points` x `number of dimensions`. 41 | 42 | ```python 43 | >>> pc1 = torch.randn(1, 100, 3).cuda().contiguous() 44 | >>> pc2 = torch.randn(1, 50, 3).cuda().contiguous() 45 | ``` 46 | 47 | Initialize a `ChamferDist` object. 48 | ```python 49 | >>> chamferDist = ChamferDistance() 50 | ``` 51 | 52 | Now, compute Chamfer distance. 53 | ```python 54 | >>> dist1, dist2, idx1, idx2 = chamferDist(pc1, pc2) 55 | >>> print(dist1.shape, dist2.shape, idx1.shape, idx2.shape) 56 | ``` 57 | 58 | Here, `dist1` is the Chamfer distance between `pc1` and `pc2`. Note that Chamfer distance is not bidirectional (and, in stricter parlance, it is not a _distance metric_). The Chamfer distance in the other direction, i.e., `pc2` to `pc1` is stored in the variable `dist2`. 59 | 60 | For each point in `pc1`, `idx1` stores the index of the closest point in `pc2`. For each point in `pc2`, `idx2` stores the index of the closest point in `pc1`. 61 | 62 | 63 | ### Citing (the original implementation, AtlasNet) 64 | 65 | If you find this work useful, you might want to cite the *original* implementation from which this codebase was borrowed (stolen!) - AtlasNet. 66 | 67 | ``` 68 | @inproceedings{groueix2018, 69 | title={{AtlasNet: A Papier-M\^ach\'e Approach to Learning 3D Surface Generation}}, 70 | author={Groueix, Thibault and Fisher, Matthew and Kim, Vladimir G. and Russell, Bryan and Aubry, Mathieu}, 71 | booktitle={Proceedings IEEE Conf. on Computer Vision and Pattern Recognition (CVPR)}, 72 | year={2018} 73 | } 74 | ``` 75 | -------------------------------------------------------------------------------- /chamferdist/chamferdist/ChamferDistance.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | from torch.autograd import Function 3 | import torch 4 | import chamferdistcuda as chamfer 5 | 6 | 7 | # Chamfer's distance module 8 | # GPU tensors only 9 | class chamferFunction(Function): 10 | @staticmethod 11 | def forward(ctx, xyz1, xyz2): 12 | batchsize, n, _ = xyz1.size() 13 | _, m, _ = xyz2.size() 14 | 15 | dist1 = torch.zeros(batchsize, n) 16 | dist2 = torch.zeros(batchsize, m) 17 | 18 | idx1 = torch.zeros(batchsize, n).type(torch.IntTensor) 19 | idx2 = torch.zeros(batchsize, m).type(torch.IntTensor) 20 | 21 | dist1 = dist1.cuda() 22 | dist2 = dist2.cuda() 23 | idx1 = idx1.cuda() 24 | idx2 = idx2.cuda() 25 | 26 | chamfer.forward(xyz1, xyz2, dist1, dist2, idx1, idx2) 27 | ctx.save_for_backward(xyz1, xyz2, idx1, idx2) 28 | return dist1, dist2, idx1, idx2 29 | 30 | @staticmethod 31 | def backward(ctx, graddist1, graddist2, idx1_, idx2_): 32 | xyz1, xyz2, idx1, idx2 = ctx.saved_tensors 33 | graddist1 = graddist1.contiguous() 34 | graddist2 = graddist2.contiguous() 35 | 36 | gradxyz1 = torch.zeros(xyz1.size()) 37 | gradxyz2 = torch.zeros(xyz2.size()) 38 | 39 | gradxyz1 = gradxyz1.cuda() 40 | gradxyz2 = gradxyz2.cuda() 41 | chamfer.backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, 42 | idx1, idx2) 43 | return gradxyz1, gradxyz2 44 | 45 | 46 | class ChamferDistance(nn.Module): 47 | def __init__(self): 48 | super(ChamferDistance, self).__init__() 49 | 50 | def forward(self, input1, input2): 51 | return chamferFunction.apply(input1, input2) 52 | -------------------------------------------------------------------------------- /chamferdist/chamferdist/__init__.py: -------------------------------------------------------------------------------- 1 | from .ChamferDistance import ChamferDistance 2 | -------------------------------------------------------------------------------- /chamferdist/chamferdist/chamfer.cu: -------------------------------------------------------------------------------- 1 | 2 | #include 3 | #include 4 | 5 | #include 6 | #include 7 | 8 | #include 9 | 10 | 11 | 12 | __global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){ 13 | const int batch=512; 14 | __shared__ float buf[batch*3]; 15 | for (int i=blockIdx.x;ibest){ 127 | result[(i*n+j)]=best; 128 | result_i[(i*n+j)]=best_i; 129 | } 130 | } 131 | __syncthreads(); 132 | } 133 | } 134 | } 135 | // int chamfer_cuda_forward(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i, cudaStream_t stream){ 136 | int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2){ 137 | 138 | const auto batch_size = xyz1.size(0); 139 | const auto n = xyz1.size(1); //num_points point cloud A 140 | const auto m = xyz2.size(1); //num_points point cloud B 141 | 142 | NmDistanceKernel<<>>(batch_size, n, xyz1.data(), m, xyz2.data(), dist1.data(), idx1.data()); 143 | NmDistanceKernel<<>>(batch_size, m, xyz2.data(), n, xyz1.data(), dist2.data(), idx2.data()); 144 | 145 | cudaError_t err = cudaGetLastError(); 146 | if (err != cudaSuccess) { 147 | printf("error in nnd updateOutput: %s\n", cudaGetErrorString(err)); 148 | //THError("aborting"); 149 | return 0; 150 | } 151 | return 1; 152 | 153 | 154 | } 155 | __global__ void NmDistanceGradKernel(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,float * grad_xyz1,float * grad_xyz2){ 156 | for (int i=blockIdx.x;i>>(batch_size,n,xyz1.data(),m,xyz2.data(),graddist1.data(),idx1.data(),gradxyz1.data(),gradxyz2.data()); 185 | NmDistanceGradKernel<<>>(batch_size,m,xyz2.data(),n,xyz1.data(),graddist2.data(),idx2.data(),gradxyz2.data(),gradxyz1.data()); 186 | 187 | cudaError_t err = cudaGetLastError(); 188 | if (err != cudaSuccess) { 189 | printf("error in nnd get grad: %s\n", cudaGetErrorString(err)); 190 | //THError("aborting"); 191 | return 0; 192 | } 193 | return 1; 194 | 195 | } 196 | 197 | -------------------------------------------------------------------------------- /chamferdist/chamferdist/chamfer_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | ///TMP 5 | //#include "common.h" 6 | /// NOT TMP 7 | 8 | 9 | int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2); 10 | 11 | 12 | int chamfer_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2); 13 | 14 | 15 | 16 | 17 | int chamfer_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2) { 18 | return chamfer_cuda_forward(xyz1, xyz2, dist1, dist2, idx1, idx2); 19 | } 20 | 21 | 22 | int chamfer_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, 23 | at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2) { 24 | 25 | return chamfer_cuda_backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2); 26 | } 27 | 28 | 29 | 30 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 31 | m.def("forward", &chamfer_forward, "chamfer forward (CUDA)"); 32 | m.def("backward", &chamfer_backward, "chamfer backward (CUDA)"); 33 | } -------------------------------------------------------------------------------- /chamferdist/example.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example usage 3 | """ 4 | 5 | import torch 6 | from chamferdist import ChamferDistance 7 | 8 | 9 | # Create two random pointclouds 10 | # (Batchsize x Number of points x Number of dims) 11 | pc1 = torch.randn(1, 100, 3).cuda().contiguous() 12 | pc2 = torch.randn(1, 50, 3).cuda().contiguous() 13 | pc1.requires_grad = True 14 | 15 | # Initialize Chamfer distance module 16 | chamferDist = ChamferDistance() 17 | # Compute Chamfer distance, and indices of closest points 18 | # - dist1 is direction pc1 -> pc2 (for each point in pc1, 19 | # gets closest point in pc2) 20 | # - dist 2 is direction pc2 -> pc1 (for each point in pc2, 21 | # gets closest point in pc1) 22 | dist1, dist2, idx1, idx2 = chamferDist(pc1, pc2) 23 | print(dist1.shape, dist2.shape, idx1.shape, idx2.shape) 24 | 25 | # Usually, dist1 is not equal to dist2 (as the associations 26 | # vary, in both directions). To get a consistent measure, 27 | # usually we average both the distances 28 | cdist = 0.5 * (dist1.mean() + dist2.mean()) 29 | print('Chamfer distance:', cdist) 30 | cdist.backward() 31 | -------------------------------------------------------------------------------- /chamferdist/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 3 | 4 | package_name = 'chamferdist' 5 | version = '0.3.0' 6 | requirements = [ 7 | 'Cython', 8 | 'torch>1.1.0', 9 | ] 10 | long_description = 'A pytorch module to compute Chamfer distance \ 11 | between two point sets (pointclouds).' 12 | 13 | setup( 14 | name='chamferdist', 15 | version=version, 16 | description='Pytorch Chamfer distance', 17 | long_description=long_description, 18 | requirements=requirements, 19 | ext_modules=[ 20 | CUDAExtension('chamferdistcuda', [ 21 | 'chamferdist/chamfer_cuda.cpp', 22 | 'chamferdist/chamfer.cu', 23 | ]), 24 | ], 25 | cmdclass={ 26 | 'build_ext': BuildExtension 27 | }) 28 | -------------------------------------------------------------------------------- /configs/config_demo.yaml: -------------------------------------------------------------------------------- 1 | name: SCALE_demo_00000_simuskirt 2 | dataroot: 00000_simuskirt 3 | data_spacing: 1 4 | npoints: 16 5 | save_all_results: 1 6 | mode: test -------------------------------------------------------------------------------- /configs/config_train_demo.yaml: -------------------------------------------------------------------------------- 1 | name: SCALE_train_demo_00000_simuskirt 2 | data_root: 00000_simuskirt 3 | data_spacing: 1 4 | val_every: 10 5 | npoints: 16 # make sure it's a square number. Otherwise the code will use the nearest (smaller) square number for it. 6 | pos_encoding: 0 7 | save_all_results: 1 8 | mode: train -------------------------------------------------------------------------------- /data/README_data.md: -------------------------------------------------------------------------------- 1 | This folder contains data examples that illustrate the data format required by SCALE in the `packed` folder, and the "raw" paired data of {clothed posed body mesh, unclothed posed body mesh} in the `raw` folder, for illustrating how to pack the data. Follow the instructions in the repository's main `README` and `lib_data/README` to download the data. 2 | 3 | Note: the data provided for downloading are only for demonstration purposes, therefore: 4 | 5 | - The data in the `packed` and `raw` folders are not in correspondence. 6 | - There is only a small number of training data frames offered in the `packed` folder; training solely on them will not yield a satisfactory model. Usually training SCALE requires hundreds to thousand of frames to generalize well on unseen poses. For more data, you can download from the [CAPE Dataset website](https://cape.is.tue.mpg.de/dataset) or use your own data, as long as they are processed and packed in the same format as our provided data. 7 | 8 | ### "packed" folder 9 | Each example is stored as a separate `.npz` file that contains the following fields: 10 | - `posmap32`: UV positional map of the minimally-clothed, posed body; 11 | - `body_verts`: vertices of the minimally-clothed, posed body; 12 | - `scan_pc`: a point set sampled from the corresponding clothed, posed body; 13 | - `scan_n`: the points' normals of the point set above. 14 | - `scan_name`: name of the frame in the format of `_.`; 15 | 16 | ### "raw" folder 17 | 18 | This folder contains an example pair of {clothed posed body mesh, unclothed posed body mesh}. Please refer to `lib_data/README` folder and try it out with the data here. 19 | 20 | ### License 21 | 22 | The data provided in this folder, including the minimally-clothed body shapes, body poses and the shape of the clothed body (in the format of meshes and sampled point clouds), are subject to the [license of the CAPE dataset](https://cape.is.tue.mpg.de/license). The UV map design is from the [SMPL](https://smpl.is.tue.mpg.de/) body model and is subject to the license of SMPL. -------------------------------------------------------------------------------- /lib/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qianlim/SCALE/98f5557e746de3536168c413901ff48c0571c5a4/lib/__init__.py -------------------------------------------------------------------------------- /lib/config_parser.py: -------------------------------------------------------------------------------- 1 | def parse_config(argv=None): 2 | import configargparse 3 | arg_formatter = configargparse.ArgumentDefaultsHelpFormatter 4 | cfg_parser = configargparse.DefaultConfigFileParser 5 | description = 'articulated bps project' 6 | parser = configargparse.ArgParser(formatter_class=arg_formatter, 7 | config_file_parser_class=cfg_parser, 8 | description=description, 9 | prog='SCALE') 10 | 11 | # general settings 12 | parser.add_argument('--config', is_config_file=True, help='config file path') 13 | parser.add_argument('--name', type=str, default='debug', help='name of a model/experiment. this name will be used for saving checkpoints and will also appear in saved examples') 14 | parser.add_argument('--mode', type=str, default='test', choices=['train', 'resume', 'test'], help='train, resume or test') 15 | 16 | # architecture related 17 | parser.add_argument('--hsize', type=int, default=256, help='hideen layer size of the ShapeDecoder mlp') 18 | parser.add_argument('--nf', type=int, default=32) 19 | parser.add_argument('--use_dropout', type=int, default=0, help='whether use dropout in the UNet') 20 | parser.add_argument('--up_mode', type=str, default='upconv', choices=['upconv', 'upsample'], help='the method to upsample in the UNet') 21 | parser.add_argument('--latent_size', type=int, default=256, help='the size of a latent vector that conditions the unet, leave it untouched (it is there for historical reason)') 22 | parser.add_argument('--pix_feat_dim', type=int, default=64, help='dim of the pixel-wise latent code output by the UNet') 23 | parser.add_argument('--pos_encoding', type=int, default=0, help='use Positional Encoding (PE) for uv coords instead of plain concat') 24 | parser.add_argument('--posemb_incl_input', type=int, default=0, help='if use PE, then include original coords in the positional encoding') 25 | parser.add_argument('--num_emb_freqs', type=int, default=6, help='if use PE: number of frequencies used in the positional embedding') 26 | 27 | # policy on how to sample each patch (pq coordinates) 28 | parser.add_argument('--npoints', type=int, default=16, help='a square number: number of points (pq coordinates) to sample in a local pixel patch') 29 | parser.add_argument('--pq_include_end', type=int, default=1, help='pq value include 1, i.e. [0,1]; else will be [0,1)') 30 | parser.add_argument('--pqmin', type=float, default=0., help='min val of the pq interval') 31 | parser.add_argument('--pqmax', type=float, default=1., help='max val of the pq interval') 32 | 33 | # data related 34 | parser.add_argument('--data_root', type=str, default='00000_simuskirt', help='the path to the "root" of the packed dataset; under this folder there are train/test/val splits.') 35 | parser.add_argument('--data_spacing', type=int, default=1, help='get every N examples from dataset (set N a large number for fast experimenting)') 36 | parser.add_argument('--img_size', type=int, default=32, help='size of UV positional map') 37 | parser.add_argument('--scan_npoints', type=int, default=-1, help='number of points used in the GT point set. By default -1 will use all points (40000);\ 38 | setting it to another number N will randomly sample N points at each iteration as GT for training.') 39 | 40 | # loss func related 41 | parser.add_argument('--w_rgl', type=float, default=2e3, help='weight for residual length regularization term') 42 | parser.add_argument('--w_normal', type=float, default=1.0, help='weight for the normal loss term') 43 | parser.add_argument('--w_m2s', type=float, default=1e4, help='weight for the Chamfer loss part 1: (m)odel to (s)can, i.e. from the prediction to the GT points') 44 | parser.add_argument('--w_s2m', type=float, default=1e4, help='weight for the Chamfer loss part 2: (s)can to (m)odel, i.e. from the GT points to the predicted points') 45 | 46 | # training / eval related 47 | parser.add_argument('--epochs', type=int, default=500) 48 | parser.add_argument('--decay_start', type=int, default=250, help='start to decay the regularization loss term from the X-th epoch') 49 | parser.add_argument('--rise_start', type=int, default=250, help='start to rise the normal loss term from the X-th epoch') 50 | parser.add_argument('--batch_size', type=int, default=16) 51 | parser.add_argument('--decay_every', type=int, default=400, help='decaly the regularization loss weight every X epochs') 52 | parser.add_argument('--rise_every', type=int, default=400, help='rise the normal loss weight every X epochs') 53 | parser.add_argument('--val_every', type=int, default=20, help='validate every x epochs') 54 | parser.add_argument('--lr', type=float, default=3e-4, help='learning rate') 55 | parser.add_argument('--save_all_results', type=int, default=0, help='save the entire test set results at inference') 56 | 57 | args, _ = parser.parse_known_args() 58 | 59 | return args -------------------------------------------------------------------------------- /lib/dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | from os.path import join 3 | 4 | import torch 5 | from torch.utils.data import Dataset 6 | import numpy as np 7 | from tqdm import tqdm 8 | 9 | SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__)) 10 | 11 | class CloDataSet(Dataset): 12 | def __init__(self, root_dir, 13 | split='train', sample_spacing=1, img_size=32, scan_npoints=-1): 14 | self.datadir = join(root_dir, split) 15 | self.split = split 16 | self.img_size = img_size 17 | self.spacing = sample_spacing 18 | self.scan_npoints = scan_npoints 19 | self.f = np.load(join(SCRIPT_DIR, '..', 'assets', 'smpl_faces.npy')) 20 | 21 | self.posmap, self.scan_n, self.scan_pc = [], [], [] 22 | self.scan_name, self.body_verts, self.pose_params = [], [], [] 23 | self._init_dataset() 24 | self.data_size = int(len(self.posmap)) 25 | 26 | print('Data loaded, in total {} {} examples.\n'.format(self.data_size, self.split)) 27 | 28 | def _init_dataset(self): 29 | print('Loading {} data...'.format(self.split)) 30 | flist = sorted(os.listdir(self.datadir))[::self.spacing] 31 | for fn in tqdm(flist): 32 | dd = np.load(join(self.datadir, fn)) # dd: 'data dict' 33 | 34 | self.posmap.append(torch.tensor(dd['posmap{}'.format(self.img_size)]).float().permute([2,0,1])) 35 | self.scan_name.append(str(dd['scan_name'])) 36 | self.body_verts.append(torch.tensor(dd['body_verts']).float()) 37 | self.scan_n.append(torch.tensor(dd['scan_n']).float()) 38 | self.scan_pc.append(torch.tensor(dd['scan_pc']).float()) # scan_pc: the GT point cloud. 39 | 40 | def __getitem__(self, index): 41 | posmap = self.posmap[index] 42 | scan_name = self.scan_name[index] 43 | body_verts = self.body_verts[index] 44 | 45 | scan_n = self.scan_n[index] 46 | scan_pc = self.scan_pc[index] 47 | 48 | if self.scan_npoints != -1: 49 | selected_idx = torch.randperm(len(scan_n))[:self.scan_npoints] 50 | scan_pc = scan_pc[selected_idx, :] 51 | scan_n = scan_n[selected_idx, :] 52 | 53 | return posmap, scan_n, scan_pc, scan_name, body_verts, torch.tensor(index).long() 54 | 55 | def __len__(self): 56 | return self.data_size 57 | -------------------------------------------------------------------------------- /lib/losses.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | def chamfer_loss_separate(output, target, weight=1e4, phase='train', debug=False): 5 | from chamferdist.chamferdist import ChamferDistance 6 | cdist = ChamferDistance() 7 | model2scan, scan2model, idx1, idx2 = cdist(output, target) 8 | if phase == 'train': 9 | return model2scan, scan2model, idx1, idx2 10 | else: # in test, show both directions, average over points, but keep batch 11 | return torch.mean(model2scan, dim=-1)* weight, torch.mean(scan2model, dim=-1)* weight, 12 | 13 | 14 | def normal_loss(output_normals, target_normals, nearest_idx, weight=1.0, phase='train'): 15 | ''' 16 | Given the set of nearest neighbors found by chamfer distance, calculate the 17 | L1 discrepancy between the predicted and GT normals on each nearest neighbor point pairs. 18 | Note: the input normals are already normalized (length==1). 19 | ''' 20 | nearest_idx = nearest_idx.expand(3, -1, -1).permute([1,2,0]).long() # [batch, N] --> [batch, N, 3], repeat for the last dim 21 | target_normals_chosen = torch.gather(target_normals, dim=1, index=nearest_idx) 22 | 23 | assert output_normals.shape == target_normals_chosen.shape 24 | 25 | if phase == 'train': 26 | lnormal = F.l1_loss(output_normals, target_normals_chosen, reduction='mean') # [batch, 8000, 3]) 27 | return lnormal, target_normals_chosen 28 | else: 29 | lnormal = F.l1_loss(output_normals, target_normals_chosen, reduction='none') 30 | lnormal = lnormal.mean(-1).mean(-1) # avg over all but batch axis 31 | return lnormal, target_normals_chosen 32 | 33 | 34 | def color_loss(output_colors, target_colors, nearest_idx, weight=1.0, phase='train', excl_holes=False): 35 | ''' 36 | Similar to normal loss, used in training a color prediction model. 37 | ''' 38 | nearest_idx = nearest_idx.expand(3, -1, -1).permute([1,2,0]).long() # [batch, N] --> [batch, N, 3], repeat for the last dim 39 | target_colors_chosen = torch.gather(target_colors, dim=1, index=nearest_idx) 40 | 41 | assert output_colors.shape == target_colors_chosen.shape 42 | 43 | if excl_holes: 44 | # scan holes have rgb all=0, exclude these from supervision 45 | colorsum = target_colors_chosen.sum(-1) 46 | mask = (colorsum!=0).float().unsqueeze(-1) 47 | else: 48 | mask = 1. 49 | 50 | if phase == 'train': 51 | lcolor = F.l1_loss(output_colors, target_colors_chosen, reduction='none') # [batch, 8000, 3]) 52 | lcolor = lcolor * mask 53 | lcolor = lcolor.mean() 54 | return lcolor, target_colors_chosen 55 | else: 56 | lcolor = F.l1_loss(output_colors, target_colors_chosen, reduction='none') 57 | lcolor = lcolor * mask 58 | lcolor = lcolor.mean(-1).mean(-1) # avg over all but batch axis 59 | return lcolor, target_colors_chosen -------------------------------------------------------------------------------- /lib/modules.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | class CBatchNorm2d(nn.Module): 5 | ''' Conditional batch normalization layer class. 6 | Borrowed from Occupancy Network repo: https://github.com/autonomousvision/occupancy_networks 7 | Args: 8 | c_dim (int): dimension of latent conditioned code c 9 | f_channels (int): number of channels of the feature maps 10 | norm_method (str): normalization method 11 | ''' 12 | 13 | def __init__(self, c_dim, f_channels, norm_method='batch_norm'): 14 | super().__init__() 15 | self.c_dim = c_dim 16 | self.f_channels = f_channels 17 | self.norm_method = norm_method 18 | # Submodules 19 | self.conv_gamma = nn.Conv1d(c_dim, f_channels, 1) # match the cond dim to num of feature channels 20 | self.conv_beta = nn.Conv1d(c_dim, f_channels, 1) 21 | if norm_method == 'batch_norm': 22 | self.bn = nn.BatchNorm2d(f_channels, affine=False) 23 | elif norm_method == 'instance_norm': 24 | self.bn = nn.InstanceNorm2d(f_channels, affine=False) 25 | elif norm_method == 'group_norm': 26 | self.bn = nn.GroupNorm2d(f_channels, affine=False) 27 | else: 28 | raise ValueError('Invalid normalization method!') 29 | self.reset_parameters() 30 | 31 | def reset_parameters(self): 32 | nn.init.zeros_(self.conv_gamma.weight) 33 | nn.init.zeros_(self.conv_beta.weight) 34 | nn.init.ones_(self.conv_gamma.bias) 35 | nn.init.zeros_(self.conv_beta.bias) 36 | 37 | def forward(self, x, c): 38 | assert(x.size(0) == c.size(0)) 39 | assert(c.size(1) == self.c_dim) 40 | 41 | # c is assumed to be of size batch_size x c_dim x 1 (conv1d needs the 3rd dim) 42 | if len(c.size()) == 2: 43 | c = c.unsqueeze(2) 44 | 45 | # Affine mapping 46 | gamma = self.conv_gamma(c).unsqueeze(-1) # make gamma be of shape [batch, f_dim, 1, 1] 47 | beta = self.conv_beta(c).unsqueeze(-1) 48 | 49 | # Batchnorm 50 | net = self.bn(x) 51 | out = gamma * net + beta 52 | 53 | return out 54 | 55 | 56 | class Conv2DBlock(nn.Module): 57 | def __init__(self, input_nc, output_nc, kernel_size=4, stride=2, padding=1, use_bias=False, use_bn=True, use_relu=True): 58 | super(Conv2DBlock, self).__init__() 59 | self.use_bn = use_bn 60 | self.use_relu = use_relu 61 | self.conv = nn.Conv2d(input_nc, output_nc, kernel_size=kernel_size, stride=stride, padding=padding, bias=use_bias) 62 | if use_bn: 63 | self.bn = nn.BatchNorm2d(output_nc, affine=False) 64 | self.relu = nn.LeakyReLU(0.2, inplace=True) 65 | 66 | def forward(self, x): 67 | if self.use_relu: 68 | x = self.relu(x) 69 | x = self.conv(x) 70 | if self.use_bn: 71 | x = self.bn(x) 72 | return x 73 | 74 | 75 | class UpConv2DBlock(nn.Module): 76 | def __init__(self, input_nc, output_nc, kernel_size=4, stride=2, padding=1, 77 | use_bias=False, use_bn=True, up_mode='upconv', dropout=0.5): 78 | super(UpConv2DBlock, self).__init__() 79 | assert up_mode in ('upconv', 'upsample') 80 | self.use_bn = use_bn 81 | self.relu = nn.ReLU() 82 | if up_mode == 'upconv': 83 | self.up = nn.ConvTranspose2d(input_nc, output_nc, kernel_size=kernel_size, stride=stride, 84 | padding=padding, bias=use_bias) 85 | else: 86 | self.up = nn.Sequential( 87 | nn.Upsample(mode='bilinear', scale_factor=2), 88 | nn.Conv2d(input_nc, output_nc, kernel_size=3, padding=1, stride=1), 89 | ) 90 | if use_bn: 91 | self.bn = nn.BatchNorm2d(output_nc, affine=False) 92 | 93 | def forward(self, x, skip_input=None): 94 | x = self.relu(x) 95 | x = self.up(x) 96 | if self.use_bn: 97 | x = self.bn(x) 98 | 99 | if skip_input is not None: 100 | x = torch.cat([x, skip_input], 1) 101 | return x 102 | 103 | 104 | class UpConv2DBlockCBNCond(nn.Module): 105 | def __init__(self, input_nc, output_nc, 106 | kernel_size=4, stride=2, padding=1, cond_dim=256, 107 | use_bias=False, use_bn=True, up_mode='upconv', use_dropout=False): 108 | super(UpConv2DBlockCBNCond, self).__init__() 109 | assert up_mode in ('upconv', 'upsample') 110 | self.use_bn = use_bn 111 | self.use_dropout = use_dropout 112 | self.relu = nn.ReLU() 113 | if up_mode == 'upconv': 114 | self.up = nn.ConvTranspose2d(input_nc, output_nc, kernel_size=kernel_size, stride=stride, 115 | padding=padding, bias=use_bias) 116 | else: 117 | self.up = nn.Sequential( 118 | nn.Upsample(mode='bilinear', scale_factor=2), 119 | nn.Conv2d(input_nc, output_nc, kernel_size=5, padding=2), 120 | ) 121 | if use_bn: 122 | self.bn = CBatchNorm2d(cond_dim, output_nc) 123 | if use_dropout: 124 | self.drop = nn.Dropout(0.5) 125 | 126 | def forward(self, x, cond, skip_input=None): 127 | x = self.relu(x) 128 | x = self.up(x) 129 | if self.use_bn: 130 | x = self.bn(x, cond) 131 | if self.use_dropout: 132 | x = self.drop(x) 133 | 134 | if skip_input is not None: 135 | x = torch.cat([x, skip_input], 1) 136 | 137 | return x 138 | 139 | 140 | class UnetCond5DS(nn.Module): 141 | ''' 142 | A simple UNet for extracting the pixel-aligned pose features from the input positional maps. 143 | - 5DS: downsample 5 times, for posmap size=32 144 | - For historical reasons the model is conditioned (using Conditional BatchNorm) with a latent vector 145 | - but since the condition vector is the same for all examples, it can essentially be ignored 146 | ''' 147 | def __init__(self, input_nc=3, output_nc=3, nf=64, cond_dim=256, up_mode='upconv', use_dropout=False, return_lowres=False): 148 | super(UnetCond5DS, self).__init__() 149 | assert up_mode in ('upconv', 'upsample') 150 | 151 | self.return_lowres = return_lowres 152 | 153 | self.conv1 = Conv2DBlock(input_nc, nf, 4, 2, 1, use_bias=False, use_bn=False, use_relu=False) 154 | self.conv2 = Conv2DBlock(1 * nf, 2 * nf, 4, 2, 1, use_bias=False, use_bn=True) 155 | self.conv3 = Conv2DBlock(2 * nf, 4 * nf, 4, 2, 1, use_bias=False, use_bn=True) 156 | self.conv4 = Conv2DBlock(4 * nf, 8 * nf, 4, 2, 1, use_bias=False, use_bn=True) 157 | self.conv5 = Conv2DBlock(8 * nf, 8 * nf, 4, 2, 1, use_bias=False, use_bn=False) 158 | 159 | self.upconv1 = UpConv2DBlockCBNCond(8 * nf, 8 * nf, 4, 2, 1, cond_dim=cond_dim, up_mode=up_mode) #2x2, 512 160 | self.upconv2 = UpConv2DBlockCBNCond(8 * nf * 2, 4 * nf, 4, 2, 1, cond_dim=cond_dim, up_mode=up_mode, use_dropout=use_dropout) # 4x4, 512 161 | self.upconv3 = UpConv2DBlockCBNCond(4 * nf * 2, 2 * nf, 4, 2, 1, cond_dim=cond_dim, up_mode=up_mode, use_dropout=use_dropout) # 8x8, 512 162 | self.upconvC4 = UpConv2DBlockCBNCond(2 * nf * 2, 1 * nf, 4, 2, 1, cond_dim=cond_dim, up_mode=up_mode) # 16 163 | self.upconvC5 = UpConv2DBlockCBNCond(1 * nf * 2, output_nc, 4, 2, 1, cond_dim=cond_dim, use_bn=False, use_bias=True, up_mode=up_mode) # 32 164 | 165 | 166 | def forward(self, x, cond): 167 | d1 = self.conv1(x) 168 | d2 = self.conv2(d1) 169 | d3 = self.conv3(d2) 170 | d4 = self.conv4(d3) 171 | d5 = self.conv5(d4) 172 | 173 | u1 = self.upconv1(d5, cond, d4) 174 | u2 = self.upconv2(u1, cond, d3) 175 | u3 = self.upconv3(u2, cond, d2) 176 | uc4 = self.upconvC4(u3, cond, d1) 177 | uc5 = self.upconvC5(uc4, cond) 178 | 179 | return uc5 180 | 181 | class ShapeDecoder(nn.Module): 182 | ''' 183 | Core component of the SCALE pipeline: the "shared MLP" in the SCALE paper Fig. 2 184 | - with skip connection from the input features to the 4th layer's output features (like DeepSDF) 185 | - branches out at the second-to-last layer, one branch for position pred, one for normal pred 186 | ''' 187 | def __init__(self, in_size, hsize = 256, actv_fn='softplus'): 188 | self.hsize = hsize 189 | super(ShapeDecoder, self).__init__() 190 | self.conv1 = torch.nn.Conv1d(in_size, self.hsize, 1) 191 | self.conv2 = torch.nn.Conv1d(self.hsize, self.hsize, 1) 192 | self.conv3 = torch.nn.Conv1d(self.hsize, self.hsize, 1) 193 | self.conv4 = torch.nn.Conv1d(self.hsize, self.hsize, 1) 194 | self.conv5 = torch.nn.Conv1d(self.hsize+in_size, self.hsize, 1) 195 | self.conv6 = torch.nn.Conv1d(self.hsize, self.hsize, 1) 196 | self.conv7 = torch.nn.Conv1d(self.hsize, self.hsize, 1) 197 | self.conv8 = torch.nn.Conv1d(self.hsize, 3, 1) 198 | 199 | self.conv6N = torch.nn.Conv1d(self.hsize, self.hsize, 1) 200 | self.conv7N = torch.nn.Conv1d(self.hsize, self.hsize, 1) 201 | self.conv8N = torch.nn.Conv1d(self.hsize, 3, 1) 202 | 203 | self.bn1 = torch.nn.BatchNorm1d(self.hsize) 204 | self.bn2 = torch.nn.BatchNorm1d(self.hsize) 205 | self.bn3 = torch.nn.BatchNorm1d(self.hsize) 206 | self.bn4 = torch.nn.BatchNorm1d(self.hsize) 207 | 208 | self.bn5 = torch.nn.BatchNorm1d(self.hsize) 209 | self.bn6 = torch.nn.BatchNorm1d(self.hsize) 210 | self.bn7 = torch.nn.BatchNorm1d(self.hsize) 211 | 212 | self.bn6N = torch.nn.BatchNorm1d(self.hsize) 213 | self.bn7N = torch.nn.BatchNorm1d(self.hsize) 214 | 215 | self.actv_fn = nn.ReLU() if actv_fn=='relu' else nn.Softplus() 216 | 217 | def forward(self, x): 218 | x1 = self.actv_fn(self.bn1(self.conv1(x))) 219 | x2 = self.actv_fn(self.bn2(self.conv2(x1))) 220 | x3 = self.actv_fn(self.bn3(self.conv3(x2))) 221 | x4 = self.actv_fn(self.bn4(self.conv4(x3))) 222 | x5 = self.actv_fn(self.bn5(self.conv5(torch.cat([x,x4],dim=1)))) 223 | 224 | # position pred 225 | x6 = self.actv_fn(self.bn6(self.conv6(x5))) 226 | x7 = self.actv_fn(self.bn7(self.conv7(x6))) 227 | x8 = self.conv8(x7) 228 | 229 | # normals pred 230 | xN6 = self.actv_fn(self.bn6N(self.conv6N(x5))) 231 | xN7 = self.actv_fn(self.bn7N(self.conv7N(xN6))) 232 | xN8 = self.conv8N(xN7) 233 | 234 | return x8, xN8 235 | 236 | 237 | class ShapeDecoderTexture(nn.Module): 238 | ''' 239 | ShapeDecoder + another branch to infer texture 240 | ''' 241 | def __init__(self, in_size, hsize = 256, actv_fn='softplus'): 242 | self.hsize = hsize 243 | super(ShapeDecoderTexture, self).__init__() 244 | self.conv1 = torch.nn.Conv1d(in_size, self.hsize, 1) 245 | self.conv2 = torch.nn.Conv1d(self.hsize, self.hsize, 1) 246 | self.conv3 = torch.nn.Conv1d(self.hsize, self.hsize, 1) 247 | self.conv4 = torch.nn.Conv1d(self.hsize, self.hsize, 1) 248 | self.conv5 = torch.nn.Conv1d(self.hsize+in_size, self.hsize, 1) 249 | self.conv6 = torch.nn.Conv1d(self.hsize, self.hsize, 1) 250 | self.conv7 = torch.nn.Conv1d(self.hsize, self.hsize, 1) 251 | self.conv8 = torch.nn.Conv1d(self.hsize, 3, 1) 252 | 253 | self.conv6N = torch.nn.Conv1d(self.hsize, self.hsize, 1) 254 | self.conv7N = torch.nn.Conv1d(self.hsize, self.hsize, 1) 255 | self.conv8N = torch.nn.Conv1d(self.hsize, 3, 1) 256 | 257 | self.conv6T = torch.nn.Conv1d(self.hsize, self.hsize, 1) 258 | self.conv7T = torch.nn.Conv1d(self.hsize, self.hsize, 1) 259 | self.conv8T = torch.nn.Conv1d(self.hsize, 3, 1) 260 | 261 | self.bn1 = torch.nn.BatchNorm1d(self.hsize) 262 | self.bn2 = torch.nn.BatchNorm1d(self.hsize) 263 | self.bn3 = torch.nn.BatchNorm1d(self.hsize) 264 | self.bn4 = torch.nn.BatchNorm1d(self.hsize) 265 | 266 | self.bn5 = torch.nn.BatchNorm1d(self.hsize) 267 | self.bn6 = torch.nn.BatchNorm1d(self.hsize) 268 | self.bn7 = torch.nn.BatchNorm1d(self.hsize) 269 | 270 | self.bn6N = torch.nn.BatchNorm1d(self.hsize) 271 | self.bn7N = torch.nn.BatchNorm1d(self.hsize) 272 | 273 | self.bn6T = torch.nn.BatchNorm1d(self.hsize) 274 | self.bn7T = torch.nn.BatchNorm1d(self.hsize) 275 | 276 | self.actv_fn = nn.ReLU() if actv_fn=='relu' else nn.Softplus() 277 | self.sigm = nn.Sigmoid() 278 | 279 | def forward(self, x): 280 | x1 = self.actv_fn(self.bn1(self.conv1(x))) 281 | x2 = self.actv_fn(self.bn2(self.conv2(x1))) 282 | x3 = self.actv_fn(self.bn3(self.conv3(x2))) 283 | x4 = self.actv_fn(self.bn4(self.conv4(x3))) 284 | x5 = self.actv_fn(self.bn5(self.conv5(torch.cat([x,x4],dim=1)))) 285 | 286 | # position pred 287 | x6 = self.actv_fn(self.bn6(self.conv6(x5))) 288 | x7 = self.actv_fn(self.bn7(self.conv7(x6))) 289 | x8 = self.conv8(x7) 290 | 291 | # normals pred 292 | xN6 = self.actv_fn(self.bn6N(self.conv6N(x5))) 293 | xN7 = self.actv_fn(self.bn7N(self.conv7N(xN6))) 294 | xN8 = self.conv8N(xN7) 295 | 296 | # texture pred 297 | xT6 = self.actv_fn(self.bn6T(self.conv6T(x5))) 298 | xT7 = self.actv_fn(self.bn7T(self.conv7T(xT6))) 299 | xT8 = self.conv8N(xT7) 300 | xT8 = self.sigm(xT8) 301 | 302 | return x8, xN8, xT8 303 | 304 | -------------------------------------------------------------------------------- /lib/network.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from lib.utils_model import PositionalEncoding, normalize_uv 5 | from lib.modules import UnetCond5DS, ShapeDecoder 6 | 7 | class SCALE(nn.Module): 8 | def __init__( 9 | self, 10 | input_nc=3, # num channels of the unet input 11 | output_nc_unet=64, # num channels output by the unet 12 | cond_dim=256, 13 | nf=64, # num filters for the unet 14 | img_size=32, # size of UV positional map 15 | hsize=256, # hidden layer size of the ShapeDecoder MLP 16 | up_mode='upconv', 17 | use_dropout=False, 18 | pos_encoding=False, # whether use Positional Encoding 19 | num_emb_freqs=8, # number of sinusoida frequences if positional encoding is used 20 | posemb_incl_input=False, # wheter include the original coordinate if using Positional Encoding 21 | uv_feat_dim=2, # input dimension of the uv coordinates 22 | pq_feat_dim = 2 # input dimension of the pq coordinates 23 | ): 24 | 25 | super().__init__() 26 | self.cond_dim = cond_dim 27 | self.output_nc_unet = output_nc_unet 28 | self.pos_encoding = pos_encoding 29 | self.num_emb_freqs = num_emb_freqs 30 | self.img_size = img_size 31 | 32 | if self.pos_encoding: 33 | self.embedder = PositionalEncoding(num_freqs=num_emb_freqs, 34 | input_dims=uv_feat_dim, 35 | include_input=posemb_incl_input) 36 | self.embedder.create_embedding_fn() 37 | uv_feat_dim = self.embedder.out_dim 38 | 39 | # U-net: for extracting pixel-aligned pose features from UV positional maps 40 | self.unet = UnetCond5DS(input_nc, output_nc_unet, nf, cond_dim=cond_dim, up_mode=up_mode, 41 | use_dropout=use_dropout) 42 | 43 | # core: maps the contatenated features+uv,pq coords to outputs 44 | self.map2Dto3D = ShapeDecoder(in_size=uv_feat_dim + pq_feat_dim + output_nc_unet, 45 | hsize=hsize, actv_fn='softplus') 46 | 47 | def forward(self, x, clo_code, uv_loc, pq_coords): 48 | ''' 49 | :param x: input posmap, [batch, 3, 256, 256] 50 | :param clo_code: clothing encoding for conditioning [256,], it's here for historical reasons but plays no actual role (as the model is outfit-specific) 51 | :param uv_loc: uv coordinate between 0,1 for each pixel, [B, N_pix, N_subsample, 2]. at each [B, N_pix], the N_subsample rows are the same, 52 | i.e. all subpixel share the same discrete (u,v) value. 53 | :param pq_coords: (p,q) coordinates in subpixel space, range [0,1), shape [B, N_pix, N_subsample, 2] 54 | :returns: 55 | residuals and normals of the points that are grouped into patches, both in shape [B, 3, H, W, N_subsample], 56 | where N_subsample is the number of points sampled per patch. 57 | ''' 58 | pix_feature = self.unet(x, clo_code) 59 | B, C = pix_feature.size()[:2] 60 | H, W = self.img_size, self.img_size 61 | N_subsample = pq_coords.shape[2] 62 | 63 | uv_feat_dim = uv_loc.size()[-1] 64 | pq_coords = pq_coords.reshape(B, -1, 2).transpose(1, 2) # [B, 2, Num of all pq subpixels] 65 | uv_loc = uv_loc.expand(N_subsample, -1, -1, -1).permute([1, 2, 0, 3]) 66 | 67 | # uv and pix feature are shared for all points within each patch 68 | pix_feature = pix_feature.view(B, C, -1).expand(N_subsample, -1,-1,-1).permute([1,2,3,0]) # [B, C, N_pix, N_sample_perpix] 69 | pix_feature = pix_feature.reshape(B, C, -1) 70 | 71 | if self.pos_encoding: 72 | uv_loc = normalize_uv(uv_loc).view(-1,uv_feat_dim) 73 | uv_loc = self.embedder.embed(uv_loc).view(B, -1,self.embedder.out_dim).transpose(1,2) 74 | else: 75 | uv_loc = uv_loc.reshape(B, -1, uv_feat_dim).transpose(1, 2) # [B, N_pix, N_subsample, 2] --> [B, 2, Num of all pq subpixels] 76 | 77 | # Core of this func: concatenated inputs --network--> outputs 78 | residuals, normals = self.map2Dto3D(torch.cat([pix_feature, uv_loc, pq_coords], 1)) # [B, 3, Num of all pq subpixels] 79 | 80 | # shape the output to the shape of 81 | # [batch, height of the positional map, width of positional map, #sampled points per pixel on the positional map (corresponds to a patch)] 82 | residuals = residuals.view(B, 3, H, W, N_subsample) 83 | normals = normals.view(B, 3, H, W, N_subsample) 84 | 85 | return residuals, normals -------------------------------------------------------------------------------- /lib/train_eval_funcs.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from tqdm import tqdm 3 | 4 | from lib.utils_io import PATCH_COLOR_DEF, save_result_examples 5 | from lib.losses import normal_loss, chamfer_loss_separate 6 | from lib.utils_model import gen_transf_mtx_full_uv 7 | 8 | 9 | def train( 10 | model, lat_vecs, device, train_loader, optimizer, 11 | flist_uv, valid_idx, uv_coord_map, 12 | subpixel_sampler=None, 13 | w_s2m=1e4, w_m2s=1e4, w_normal=10, 14 | w_rgl=1e2, w_latent_rgl = 1.0 15 | ): 16 | 17 | n_train_samples = len(train_loader.dataset) 18 | 19 | train_s2m, train_m2s, train_lnormal, train_rgl, train_latent_rgl, train_total = 0, 0, 0, 0, 0, 0 20 | 21 | model.train() 22 | for _, data in enumerate(train_loader): 23 | 24 | # ------------------------------------------------------- 25 | # ------------ load batch data and reshaping ------------ 26 | 27 | [inp_posmap, target_pc_n, target_pc, target_names, body_verts, index] = data 28 | gpu_data = [inp_posmap, target_pc_n, target_pc, body_verts, index] 29 | [inp_posmap, target_pc_n, target_pc, body_verts, index] = list(map(lambda x: x.to(device), gpu_data)) 30 | bs, _, H, W = inp_posmap.size() 31 | 32 | optimizer.zero_grad() 33 | 34 | transf_mtx_map = gen_transf_mtx_full_uv(body_verts, flist_uv) 35 | lat_vec_batch = lat_vecs(torch.tensor(0).cuda()).expand(bs, -1) 36 | 37 | uv_coord_map_batch = uv_coord_map.expand(bs, -1, -1).contiguous() 38 | pq_samples = subpixel_sampler.sample_regular_points() # pq coord grid for one patch 39 | pq_repeated = pq_samples.expand(bs, H * W, -1, -1) # repeat the same pq parameterization for all patches 40 | 41 | N_subsample = pq_samples.shape[1] 42 | 43 | # The same body point is shared by all sampled pq points within each patch 44 | bp_locations = inp_posmap.expand(N_subsample, -1, -1,-1,-1).permute([1, 2, 3, 4, 0]) #[bs, C, H, W, N_sample], 45 | transf_mtx_map = transf_mtx_map.expand(N_subsample, -1, -1, -1, -1, -1).permute([1, 2, 3, 0, 4, 5]) # [bs, H, W, N_subsample, 3, 3] 46 | 47 | # -------------------------------------------------------------------- 48 | # ------------ model pass an coordinate transformation --------------- 49 | 50 | # Core: predict the clothing residual (displacement) from the body, and their normals 51 | pred_res, pred_normals = model(inp_posmap, clo_code=lat_vec_batch, 52 | uv_loc=uv_coord_map_batch, 53 | pq_coords=pq_repeated) 54 | 55 | # local coords --> global coords 56 | pred_res = pred_res.permute([0,2,3,4,1]).unsqueeze(-1) 57 | pred_normals = pred_normals.permute([0,2,3,4,1]).unsqueeze(-1) 58 | 59 | pred_res = torch.matmul(transf_mtx_map, pred_res).squeeze(-1) 60 | pred_normals = torch.matmul(transf_mtx_map, pred_normals).squeeze(-1) 61 | pred_normals = torch.nn.functional.normalize(pred_normals, dim=-1) 62 | 63 | # residual to abosolute locations in space 64 | full_pred = pred_res.permute([0,4,1,2,3]).contiguous() + bp_locations 65 | 66 | # take the selected points and reshape to [Npoints, 3] 67 | full_pred = full_pred.permute([0,2,3,4,1]).reshape(bs, -1, N_subsample, 3)[:, valid_idx, ...] 68 | pred_normals = pred_normals.reshape(bs, -1, N_subsample, 3)[:, valid_idx, ...] 69 | 70 | # reshaping the points that are grouped into patches into a big point set 71 | full_pred = full_pred.reshape(bs, -1, 3).contiguous() 72 | pred_normals = pred_normals.reshape(bs, -1, 3).contiguous() 73 | 74 | # -------------------------------- 75 | # ------------ losses ------------ 76 | 77 | # Chamfer dist from the (s)can to (m)odel: from the GT points to its closest ponit in the predicted point set 78 | m2s, s2m, idx_closest_gt, _ = chamfer_loss_separate(full_pred, target_pc) #idx1: [#pred points] 79 | s2m = torch.mean(s2m) 80 | 81 | # normal loss 82 | lnormal, closest_target_normals = normal_loss(pred_normals, target_pc_n, idx_closest_gt) 83 | 84 | # dist from the predicted points to their respective closest point on the GT, projected by 85 | # the normal of these GT points, to appxoimate the point-to-surface distance 86 | nearest_idx = idx_closest_gt.expand(3, -1, -1).permute([1,2,0]).long() # [batch, N] --> [batch, N, 3], repeat for the last dim 87 | target_points_chosen = torch.gather(target_pc, dim=1, index=nearest_idx) 88 | pc_diff = target_points_chosen - full_pred # vectors from prediction to its closest point in gt pcl 89 | m2s = torch.sum(pc_diff * closest_target_normals, dim=-1) # project on direction of the normal of these gt points 90 | m2s = torch.mean(m2s**2) # the length (squared) is the approx. pred point to scan surface dist. 91 | 92 | rgl_len = torch.mean(pred_res ** 2) 93 | rgl_latent = torch.sum(torch.norm(lat_vec_batch, dim=1)) 94 | 95 | loss = s2m*w_s2m + m2s*w_m2s + rgl_len*w_rgl + lnormal* w_normal + rgl_latent*w_latent_rgl 96 | 97 | loss.backward() 98 | optimizer.step() 99 | 100 | # ------------------------------------------ 101 | # ------------ accumulate stats ------------ 102 | 103 | train_m2s += m2s * bs 104 | train_s2m += s2m * bs 105 | train_lnormal += lnormal * bs 106 | train_rgl += rgl_len * bs 107 | train_latent_rgl += rgl_latent * bs 108 | 109 | train_total += loss * bs 110 | 111 | train_s2m /= n_train_samples 112 | train_m2s /= n_train_samples 113 | train_lnormal /= n_train_samples 114 | train_rgl /= n_train_samples 115 | train_latent_rgl /= n_train_samples 116 | train_total /= n_train_samples 117 | 118 | return train_m2s, train_s2m, train_lnormal, train_rgl, train_latent_rgl, train_total 119 | 120 | 121 | def test(model, lat_vecs, device, test_loader, epoch_idx, samples_dir, 122 | flist_uv, valid_idx, uv_coord_map, 123 | mode='val', subpixel_sampler=None, 124 | model_name=None, save_all_results=False): 125 | 126 | model.eval() 127 | 128 | lat_vecs = lat_vecs(torch.tensor(0).cuda()) # it's here for historical reason, can safely treat it as part of the network weights 129 | 130 | n_test_samples = len(test_loader.dataset) 131 | 132 | test_s2m, test_m2s, test_lnormal, test_rgl, test_latent_rgl = 0, 0, 0, 0, 0 133 | 134 | with torch.no_grad(): 135 | for data in tqdm(test_loader): 136 | 137 | # ------------------------------------------------------- 138 | # ------------ load batch data and reshaping ------------ 139 | 140 | [inp_posmap, target_pc_n, target_pc, target_names, body_verts, index] = data 141 | gpu_data = [inp_posmap, target_pc_n, target_pc, body_verts, index] 142 | [inp_posmap, target_pc_n, target_pc, body_verts, index] = list(map(lambda x: x.to(device, non_blocking=True), gpu_data)) 143 | 144 | bs, C, H, W = inp_posmap.size() 145 | 146 | lat_vec_batch = lat_vecs.expand(bs, -1) 147 | transf_mtx_map = gen_transf_mtx_full_uv(body_verts, flist_uv) 148 | uv_coord_map_batch = uv_coord_map.expand(bs, -1, -1).contiguous() 149 | 150 | pq_samples = subpixel_sampler.sample_regular_points() 151 | pq_repeated = pq_samples.expand(bs, H * W, -1, -1) # [B, H*W, samples_per_pix, 2] 152 | 153 | N_subsample = pq_samples.shape[1] 154 | 155 | # The same body point is shared by all sampled pq points within each patch 156 | bp_locations = inp_posmap.expand(N_subsample, -1, -1,-1,-1).permute([1, 2, 3, 4, 0]) # [B, C, H, W, N_sample] 157 | transf_mtx_map = transf_mtx_map.expand(N_subsample, -1, -1, -1, -1, -1).permute([1, 2, 3, 0, 4, 5]) # [B, H, W, N_subsample, 3, 3] 158 | 159 | # -------------------------------------------------------------------- 160 | # ------------ model pass an coordinate transformation --------------- 161 | 162 | # Core: predict the clothing residual (displacement) from the body, and their normals 163 | pred_res, pred_normals = model(inp_posmap, clo_code=lat_vec_batch, 164 | uv_loc=uv_coord_map_batch, 165 | pq_coords=pq_repeated) 166 | 167 | # local coords --> global coords 168 | pred_res = pred_res.permute([0,2,3,4,1]).unsqueeze(-1) 169 | pred_normals = pred_normals.permute([0, 2, 3, 4, 1]).unsqueeze(-1) 170 | 171 | pred_res = torch.matmul(transf_mtx_map, pred_res).squeeze(-1) 172 | pred_normals = torch.matmul(transf_mtx_map, pred_normals).squeeze(-1) 173 | pred_normals = torch.nn.functional.normalize(pred_normals, dim=-1) 174 | 175 | # residual to abosolute locations in space 176 | full_pred = pred_res.permute([0,4,1,2,3]).contiguous() + bp_locations 177 | 178 | # take the selected points and reshape to [N_valid_points, 3] 179 | full_pred = full_pred.permute([0,2,3,4,1]).reshape(bs, -1, N_subsample, 3)[:, valid_idx, ...] 180 | pred_normals = pred_normals.reshape(bs, -1, N_subsample, 3)[:, valid_idx, ...] 181 | 182 | full_pred = full_pred.reshape(bs, -1, 3).contiguous() 183 | pred_normals = pred_normals.reshape(bs, -1, 3).contiguous() 184 | 185 | # -------------------------------- 186 | # ------------ losses ------------ 187 | 188 | _, s2m, idx_closest_gt, _ = chamfer_loss_separate(full_pred, target_pc) #idx1: [#pred points] 189 | s2m = s2m.mean(1) 190 | lnormal, closest_target_normals = normal_loss(pred_normals, target_pc_n, idx_closest_gt, phase='test') 191 | nearest_idx = idx_closest_gt.expand(3, -1, -1).permute([1,2,0]).long() # [batch, N] --> [batch, N, 3], repeat for the last dim 192 | target_points_chosen = torch.gather(target_pc, dim=1, index=nearest_idx) 193 | pc_diff = target_points_chosen - full_pred # vectors from prediction to its closest point in gt pcl 194 | m2s = torch.sum(pc_diff * closest_target_normals, dim=-1) # project on direction of the normal of these gt points 195 | m2s = torch.mean(m2s**2, 1) 196 | 197 | rgl_len = torch.mean((pred_res ** 2).reshape(bs, -1),1) 198 | rgl_latent = torch.sum(torch.norm(lat_vec_batch, dim=1)) 199 | 200 | # ------------------------------------------ 201 | # ------------ accumulate stats ------------ 202 | 203 | test_m2s += torch.sum(m2s) 204 | test_s2m += torch.sum(s2m) 205 | test_lnormal += torch.sum(lnormal) 206 | test_rgl += torch.sum(rgl_len) 207 | test_latent_rgl += rgl_latent 208 | 209 | patch_colors = PATCH_COLOR_DEF.expand(N_subsample, -1, -1).transpose(0,1).reshape(-1, 3) 210 | 211 | if mode == 'test': 212 | save_spacing = 1 if save_all_results else 10 213 | for i in range(full_pred.shape[0])[::save_spacing]: 214 | save_result_examples(samples_dir, model_name, target_names[i], 215 | points=full_pred[i], normals=pred_normals[i], 216 | patch_color=patch_colors) 217 | 218 | test_m2s /= n_test_samples 219 | test_s2m /= n_test_samples 220 | test_lnormal /= n_test_samples 221 | test_rgl /= n_test_samples 222 | test_latent_rgl /= n_test_samples 223 | 224 | test_s2m, test_m2s, test_lnormal, test_rgl, test_latent_rgl = list(map(lambda x: x.detach().cpu().numpy(), [test_s2m, test_m2s, test_lnormal, test_rgl, test_latent_rgl])) 225 | test_total_loss = test_s2m + test_m2s + test_lnormal + test_rgl + test_latent_rgl 226 | 227 | if mode != 'test': 228 | if epoch_idx == 0 or epoch_idx % 10 == 0: 229 | # only save the first example per batch for quick inspection of validation set results 230 | save_result_examples(samples_dir, model_name, target_names[0], 231 | points=full_pred[0], normals=pred_normals[0], 232 | patch_color=None, epoch=epoch_idx) 233 | 234 | print("model2scan dist: {:.3e}, scan2model dist: {:.3e}, normal loss: {:.3e}" 235 | " rgl term: {:.3e}, latent rgl term:{:.3e},".format(test_m2s, test_s2m, test_lnormal, 236 | test_rgl, test_latent_rgl)) 237 | 238 | return test_m2s, test_s2m, test_lnormal, test_rgl, test_latent_rgl, test_total_loss 239 | -------------------------------------------------------------------------------- /lib/utils_io.py: -------------------------------------------------------------------------------- 1 | import os 2 | from os.path import join 3 | 4 | import torch 5 | import torch.nn.functional as F 6 | import numpy as np 7 | 8 | PATCH_COLOR_DEF = torch.rand([1, 798, 3]) # 32x32 posmap for smpl has 798 valid pixels 9 | 10 | 11 | def getIdxMap_torch(img, offset=False): 12 | # img has shape [channels, H, W] 13 | C, H, W = img.shape 14 | import torch 15 | idx = torch.stack(torch.where(~torch.isnan(img[0]))) 16 | if offset: 17 | idx = idx.float() + 0.5 18 | idx = idx.view(2, H * W).float().contiguous() 19 | idx = idx.transpose(0, 1) 20 | 21 | idx = idx / (H-1) if not offset else idx / H 22 | return idx 23 | 24 | 25 | def load_masks(PROJECT_DIR, posmap_size, body_model='smpl'): 26 | uv_mask_faceid = np.load(join(PROJECT_DIR, 'assets', 'uv_masks', 'uv_mask{}_with_faceid_{}.npy'.format(posmap_size, body_model))).reshape(posmap_size, posmap_size) 27 | uv_mask_faceid = torch.from_numpy(uv_mask_faceid).long().cuda() 28 | 29 | smpl_faces = np.load(join(PROJECT_DIR, 'assets', 'smpl_faces.npy')) # faces = triangle list of the body mesh 30 | flist = torch.tensor(smpl_faces.astype(np.int32)).long() 31 | flist_uv = get_face_per_pixel(uv_mask_faceid, flist).cuda() # Each (valid) pixel on the uv map corresponds to a point on the SMPL body; flist_uv is a list of these triangles 32 | 33 | points_idx_from_posmap = np.load(join(PROJECT_DIR, 'assets', 'uv_masks', 'idx_smpl_posmap{}_uniformverts_retrieval.npy'.format(posmap_size))) 34 | points_idx_from_posmap = torch.from_numpy(points_idx_from_posmap).cuda() 35 | 36 | uv_coord_map = getIdxMap_torch(torch.rand(3, posmap_size, posmap_size)).cuda() 37 | uv_coord_map.requires_grad = True 38 | 39 | return flist_uv, points_idx_from_posmap, uv_coord_map 40 | 41 | 42 | def get_face_per_pixel(mask, flist): 43 | ''' 44 | :param mask: the uv_mask returned from posmap renderer, where -1 stands for background 45 | pixels in the uv map, where other value (int) is the face index that this 46 | pixel point corresponds to. 47 | :param flist: the face list of the body model, 48 | - smpl, it should be an [13776, 3] array 49 | - smplx, it should be an [20908,3] array 50 | :return: 51 | flist_uv: an [img_size, img_size, 3] array, each pixel is the index of the 3 verts that belong to the triangle 52 | Note: we set all background (-1) pixels to be 0 to make it easy to parralelize, but later we 53 | will just mask out these pixels, so it's fine that they are wrong. 54 | ''' 55 | mask2 = mask.clone() 56 | mask2[mask == -1] = 0 #remove the -1 in the mask, so that all mask elements can be seen as meaningful faceid 57 | flist_uv = flist[mask2] 58 | return flist_uv 59 | 60 | 61 | def save_latent_vectors(filepath, latent_vec, epoch): 62 | 63 | all_latents = latent_vec.state_dict() 64 | 65 | torch.save( 66 | {"epoch": epoch, "latent_codes": all_latents}, 67 | os.path.join(filepath), 68 | ) 69 | 70 | 71 | def load_latent_vectors(filepath, lat_vecs): 72 | 73 | full_filename = filepath 74 | 75 | if not os.path.isfile(full_filename): 76 | raise Exception('latent state file "{}" does not exist'.format(full_filename)) 77 | 78 | data = torch.load(full_filename) 79 | 80 | if isinstance(data["latent_codes"], torch.Tensor): 81 | # for backwards compatibility 82 | if not lat_vecs.num_embeddings == data["latent_codes"].size()[0]: 83 | raise Exception( 84 | "num latent codes mismatched: {} vs {}".format( 85 | lat_vecs.num_embeddings, data["latent_codes"].size()[0] 86 | ) 87 | ) 88 | 89 | if not lat_vecs.embedding_dim == data["latent_codes"].size()[2]: 90 | raise Exception("latent code dimensionality mismatch") 91 | 92 | for i, lat_vec in enumerate(data["latent_codes"]): 93 | lat_vecs.weight.data[i, :] = lat_vec 94 | 95 | else: 96 | lat_vecs.load_state_dict(data["latent_codes"]) 97 | 98 | return data["epoch"] 99 | 100 | 101 | def save_model(path, model, epoch, optimizer=None): 102 | model_dict = { 103 | 'epoch': epoch, 104 | 'model_state': model.state_dict() 105 | } 106 | if optimizer is not None: 107 | model_dict['optimizer_state'] = optimizer.state_dict() 108 | 109 | torch.save(model_dict, path) 110 | 111 | 112 | def tensor2numpy(tensor): 113 | if isinstance(tensor, torch.Tensor): 114 | return tensor.detach().cpu().numpy() 115 | 116 | 117 | def customized_export_ply(outfile_name, v, f = None, v_n = None, v_c = None, f_c = None, e = None): 118 | ''' 119 | Author: Jinlong Yang, jyang@tue.mpg.de 120 | 121 | Exports a point cloud / mesh to a .ply file 122 | supports vertex normal and color export 123 | such that the saved file will be correctly displayed in MeshLab 124 | 125 | # v: Vertex position, N_v x 3 float numpy array 126 | # f: Face, N_f x 3 int numpy array 127 | # v_n: Vertex normal, N_v x 3 float numpy array 128 | # v_c: Vertex color, N_v x (3 or 4) uchar numpy array 129 | # f_n: Face normal, N_f x 3 float numpy array 130 | # f_c: Face color, N_f x (3 or 4) uchar numpy array 131 | # e: Edge, N_e x 2 int numpy array 132 | # mode: ascii or binary ply file. Value is {'ascii', 'binary'} 133 | ''' 134 | 135 | v_n_flag=False 136 | v_c_flag=False 137 | f_c_flag=False 138 | 139 | N_v = v.shape[0] 140 | assert(v.shape[1] == 3) 141 | if not type(v_n) == type(None): 142 | assert(v_n.shape[0] == N_v) 143 | if type(v_n) == 'torch.Tensor': 144 | v_n = v_n.detach().cpu().numpy() 145 | v_n_flag = True 146 | if not type(v_c) == type(None): 147 | assert(v_c.shape[0] == N_v) 148 | v_c_flag = True 149 | if v_c.shape[1] == 3: 150 | # warnings.warn("Vertex color does not provide alpha channel, use default alpha = 255") 151 | alpha_channel = np.zeros((N_v, 1), dtype = np.ubyte)+255 152 | v_c = np.hstack((v_c, alpha_channel)) 153 | 154 | N_f = 0 155 | if not type(f) == type(None): 156 | N_f = f.shape[0] 157 | assert(f.shape[1] == 3) 158 | if not type(f_c) == type(None): 159 | assert(f_c.shape[0] == f.shape[0]) 160 | f_c_flag = True 161 | if f_c.shape[1] == 3: 162 | # warnings.warn("Face color does not provide alpha channel, use default alpha = 255") 163 | alpha_channel = np.zeros((N_f, 1), dtype = np.ubyte)+255 164 | f_c = np.hstack((f_c, alpha_channel)) 165 | N_e = 0 166 | if not type(e) == type(None): 167 | N_e = e.shape[0] 168 | 169 | with open(outfile_name, 'w') as file: 170 | # Header 171 | file.write('ply\n') 172 | file.write('format ascii 1.0\n') 173 | file.write('element vertex %d\n'%(N_v)) 174 | file.write('property float x\n') 175 | file.write('property float y\n') 176 | file.write('property float z\n') 177 | 178 | if v_n_flag: 179 | file.write('property float nx\n') 180 | file.write('property float ny\n') 181 | file.write('property float nz\n') 182 | if v_c_flag: 183 | file.write('property uchar red\n') 184 | file.write('property uchar green\n') 185 | file.write('property uchar blue\n') 186 | file.write('property uchar alpha\n') 187 | 188 | file.write('element face %d\n'%(N_f)) 189 | file.write('property list uchar int vertex_indices\n') 190 | if f_c_flag: 191 | file.write('property uchar red\n') 192 | file.write('property uchar green\n') 193 | file.write('property uchar blue\n') 194 | file.write('property uchar alpha\n') 195 | 196 | if not N_e == 0: 197 | file.write('element edge %d\n'%(N_e)) 198 | file.write('property int vertex1\n') 199 | file.write('property int vertex2\n') 200 | 201 | file.write('end_header\n') 202 | 203 | # Main body: 204 | # Vertex 205 | if v_n_flag and v_c_flag: 206 | for i in range(0, N_v): 207 | file.write('%f %f %f %f %f %f %d %d %d %d\n'%\ 208 | (v[i,0], v[i,1], v[i,2],\ 209 | v_n[i,0], v_n[i,1], v_n[i,2], \ 210 | v_c[i,0], v_c[i,1], v_c[i,2], v_c[i,3])) 211 | elif v_n_flag: 212 | for i in range(0, N_v): 213 | file.write('%f %f %f %f %f %f\n'%\ 214 | (v[i,0], v[i,1], v[i,2],\ 215 | v_n[i,0], v_n[i,1], v_n[i,2])) 216 | elif v_c_flag: 217 | for i in range(0, N_v): 218 | file.write('%f %f %f %d %d %d %d\n'%\ 219 | (v[i,0], v[i,1], v[i,2],\ 220 | v_c[i,0], v_c[i,1], v_c[i,2], v_c[i,3])) 221 | else: 222 | for i in range(0, N_v): 223 | file.write('%f %f %f\n'%\ 224 | (v[i,0], v[i,1], v[i,2])) 225 | # Face 226 | if f_c_flag: 227 | for i in range(0, N_f): 228 | file.write('3 %d %d %d %d %d %d %d\n'%\ 229 | (f[i,0], f[i,1], f[i,2],\ 230 | f_c[i,0], f_c[i,1], f_c[i,2], f_c[i,3])) 231 | else: 232 | for i in range(0, N_f): 233 | file.write('3 %d %d %d\n'%\ 234 | (f[i,0], f[i,1], f[i,2])) 235 | 236 | # Edge 237 | if not N_e == 0: 238 | for i in range(0, N_e): 239 | file.write('%d %d\n'%(e[i,0], e[i,1])) 240 | 241 | 242 | def vertex_normal_2_vertex_color(vertex_normal): 243 | # Normalize vertex normal 244 | import torch 245 | if torch.is_tensor(vertex_normal): 246 | vertex_normal = vertex_normal.detach().cpu().numpy() 247 | normal_length = ((vertex_normal**2).sum(1))**0.5 248 | normal_length = normal_length.reshape(-1, 1) 249 | vertex_normal /= normal_length 250 | # Convert normal to color: 251 | color = vertex_normal * 255/2.0 + 128 252 | return color.astype(np.ubyte) 253 | 254 | 255 | def draw_correspondence(pcl_1, pcl_2, output_file): 256 | ''' 257 | Given a pair of (minimal, clothed) point clouds which have same #points, 258 | draw correspondences between each point pair as a line and export to a .ply 259 | file for visualization. 260 | ''' 261 | assert(pcl_1.shape[0] == pcl_2.shape[0]) 262 | N = pcl_2.shape[0] 263 | v = np.vstack((pcl_1, pcl_2)) 264 | arange = np.arange(0, N) 265 | arange = arange.reshape(-1,1) 266 | e = np.hstack((arange, arange+N)) 267 | e = e.astype(np.int32) 268 | customized_export_ply(output_file, v, e = e) 269 | 270 | 271 | def save_result_examples(save_dir, model_name, result_name, 272 | points, normals=None, patch_color=None, 273 | texture=None, coarse_pts=None, 274 | gt=None, epoch=None): 275 | # works on single pcl, i.e. [#num_pts, 3], no batch dimension 276 | from os.path import join 277 | import numpy as np 278 | 279 | if epoch is None: 280 | normal_fn = '{}_{}_pred.ply'.format(model_name,result_name) 281 | else: 282 | normal_fn = '{}_epoch{}_{}_pred.ply'.format(model_name, str(epoch).zfill(4), result_name) 283 | normal_fn = join(save_dir, normal_fn) 284 | points = tensor2numpy(points) 285 | 286 | if normals is not None: 287 | normals = tensor2numpy(normals) 288 | color_normal = vertex_normal_2_vertex_color(normals) 289 | customized_export_ply(normal_fn, v=points, v_n=normals, v_c=color_normal) 290 | 291 | if patch_color is not None: 292 | patch_color = tensor2numpy(patch_color) 293 | if patch_color.max() < 1.1: 294 | patch_color = (patch_color*255.).astype(np.ubyte) 295 | pcolor_fn = normal_fn.replace('pred.ply', 'pred_patchcolor.ply') 296 | customized_export_ply(pcolor_fn, v=points, v_c=patch_color) 297 | 298 | if texture is not None: 299 | texture = tensor2numpy(texture) 300 | if texture.max() < 1.1: 301 | texture = (texture*255.).astype(np.ubyte) 302 | texture_fn = normal_fn.replace('pred.ply', 'pred_texture.ply') 303 | customized_export_ply(texture_fn, v=points, v_c=texture) 304 | 305 | if coarse_pts is not None: 306 | coarse_pts = tensor2numpy(coarse_pts) 307 | coarse_fn = normal_fn.replace('pred.ply', 'interm.ply') 308 | customized_export_ply(coarse_fn, v=coarse_pts) 309 | 310 | if gt is not None: 311 | gt = tensor2numpy(gt) 312 | gt_fn = normal_fn.replace('pred.ply', 'gt.ply') 313 | customized_export_ply(gt_fn, v=gt) -------------------------------------------------------------------------------- /lib/utils_model.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | def gen_transf_mtx_full_uv(verts, faces): 7 | ''' 8 | given a positional uv map, for each of its pixel, get the matrix that transforms the prediction from local to global coordinates 9 | The local coordinates are defined by the posed body mesh (consists of vertcs and faces) 10 | 11 | :param verts: [batch, Nverts, 3] 12 | :param faces: [uv_size, uv_size, 3], uv_size =e.g. 32 13 | 14 | :return: [batch, uv_size, uv_size, 3,3], per example a map of 3x3 rot matrices for local->global transform 15 | 16 | NOTE: local coords are NOT cartesian! uu an vv axis are edges of the triangle, 17 | not perpendicular (more like barycentric coords) 18 | ''' 19 | tris = verts[:, faces] # [batch, uv_size, uv_size, 3, 3] 20 | v1, v2, v3 = tris[:, :, :, 0, :], tris[:, :, :, 1, :], tris[:, :, :, 2, :] 21 | uu = v2 - v1 # u axis of local coords is the first edge, [batch, uv_size, uv_size, 3] 22 | vv = v3 - v1 # v axis, second edge 23 | ww_raw = torch.cross(uu, vv, dim=-1) 24 | ww = F.normalize(ww_raw, p=2, dim=-1) # unit triangle normal as w axis 25 | ww_norm = (torch.norm(uu, dim=-1).mean(-1).mean(-1) + torch.norm(vv, dim=-1).mean(-1).mean(-1)) / 2. 26 | ww = ww*ww_norm.view(len(ww_norm),1,1,1) 27 | 28 | # shape of transf_mtx will be [batch, uv_size, uv_size, 3, 3], where the last two dim is like: 29 | # | | | 30 | #[ uu vv ww] 31 | # | | | 32 | # for local to global, say coord in the local coord system is (r,s,t) 33 | # then the coord in world system should be r*uu + s*vv+ t*ww 34 | # so the uu, vv, ww should be colum vectors of the local->global transf mtx 35 | # so when stack, always stack along dim -1 (i.e. column) 36 | transf_mtx = torch.stack([uu, vv, ww], dim=-1) 37 | 38 | return transf_mtx 39 | 40 | 41 | class SampleSquarePoints(): 42 | def __init__(self, npoints=16, min_val=0, max_val=1, device='cuda', include_end=True): 43 | super(SampleSquarePoints, self).__init__() 44 | self.npoints = npoints 45 | self.device = device 46 | self.min_val = min_val # -1 or 0 47 | self.max_val = max_val # -1 or 0 48 | self.include_end = include_end 49 | 50 | def sample_regular_points(self, N=None): 51 | steps = int(self.npoints ** 0.5) if N is None else int(N ** 0.5) 52 | if self.include_end: 53 | linspace = torch.linspace(self.min_val, self.max_val, steps=steps) # [0,1] 54 | else: 55 | linspace = torch.linspace(self.min_val, self.max_val, steps=steps+1)[: steps] # [0,1) 56 | grid = torch.stack(torch.meshgrid([linspace, linspace]), -1).to(self.device) #[steps, steps, 2] 57 | grid = grid.view(-1,2).unsqueeze(0) #[B, N, 2] 58 | grid.requires_grad = True 59 | 60 | return grid 61 | 62 | def sample_random_points(self, N=None): 63 | npt = self.npoints if N is None else N 64 | shape = torch.Size((1, npt, 2)) 65 | rand_grid = torch.Tensor(shape).float().to(self.device) 66 | rand_grid.data.uniform_(self.min_val, self.max_val) 67 | rand_grid.requires_grad = True #[B, N, 2] 68 | return rand_grid 69 | 70 | 71 | class Embedder(): 72 | ''' 73 | Simple positional encoding, adapted from NeRF: https://github.com/bmild/nerf 74 | ''' 75 | def __init__(self, **kwargs): 76 | 77 | self.kwargs = kwargs 78 | self.create_embedding_fn() 79 | 80 | def create_embedding_fn(self): 81 | 82 | embed_fns = [] 83 | d = self.kwargs['input_dims'] 84 | out_dim = 0 85 | if self.kwargs['include_input']: 86 | embed_fns.append(lambda x: x) 87 | out_dim += d 88 | 89 | max_freq = self.kwargs['max_freq_log2'] 90 | N_freqs = self.kwargs['num_freqs'] 91 | 92 | if self.kwargs['log_sampling']: 93 | freq_bands = 2. ** torch.linspace(0., max_freq, steps=N_freqs) 94 | else: 95 | freq_bands = torch.linspace(2. ** 0., 2. ** max_freq, steps=N_freqs) 96 | 97 | for freq in freq_bands: 98 | for p_fn in self.kwargs['periodic_fns']: 99 | embed_fns.append(lambda x, p_fn=p_fn, freq=freq: p_fn(x * freq)) 100 | out_dim += d 101 | 102 | self.embed_fns = embed_fns 103 | self.out_dim = out_dim 104 | 105 | def embed(self, inputs): 106 | return torch.cat([fn(inputs) for fn in self.embed_fns], -1) 107 | 108 | 109 | def get_embedder(multires, i=0, input_dims=3): 110 | ''' 111 | Helper function for positional encoding, adapted from NeRF: https://github.com/bmild/nerf 112 | ''' 113 | if i == -1: 114 | return nn.Identity(), input_dims 115 | 116 | embed_kwargs = { 117 | 'include_input': True, 118 | 'input_dims': input_dims, 119 | 'max_freq_log2': multires - 1, 120 | 'num_freqs': multires, 121 | 'log_sampling': True, 122 | 'periodic_fns': [torch.sin, torch.cos], 123 | } 124 | 125 | embedder_obj = Embedder(**embed_kwargs) 126 | embed = lambda x, eo=embedder_obj: eo.embed(x) 127 | return embed, embedder_obj.out_dim 128 | 129 | 130 | class PositionalEncoding(): 131 | def __init__(self, input_dims=2, num_freqs=10, include_input=False): 132 | super(PositionalEncoding,self).__init__() 133 | self.include_input = include_input 134 | self.num_freqs = num_freqs 135 | self.input_dims = input_dims 136 | 137 | def create_embedding_fn(self): 138 | embed_fns = [] 139 | out_dim = 0 140 | if self.include_input: 141 | embed_fns.append(lambda x: x) 142 | out_dim += self.input_dims 143 | 144 | freq_bands = 2. ** torch.linspace(0, self.num_freqs-1, self.num_freqs) 145 | periodic_fns = [torch.sin, torch.cos] 146 | 147 | for freq in freq_bands: 148 | for p_fn in periodic_fns: 149 | embed_fns.append(lambda x, p_fn=p_fn, freq=freq:p_fn(math.pi * x * freq)) 150 | # embed_fns.append(lambda x, p_fn=p_fn, freq=freq:p_fn(x * freq)) 151 | out_dim += self.input_dims 152 | 153 | self.embed_fns = embed_fns 154 | self.out_dim = out_dim 155 | 156 | def embed(self,coords): 157 | ''' 158 | use periodic positional encoding to transform cartesian positions to higher dimension 159 | :param coords: [N, 3] 160 | :return: [N, 3*2*num_freqs], where 2 comes from that for each frequency there's a sin() and cos() 161 | ''' 162 | return torch.cat([fn(coords) for fn in self.embed_fns], dim=-1) 163 | 164 | 165 | def normalize_uv(uv): 166 | ''' 167 | normalize uv coords from range [0,1] to range [-1,1] 168 | ''' 169 | return uv * 2. - 1. 170 | 171 | -------------------------------------------------------------------------------- /lib/utils_train.py: -------------------------------------------------------------------------------- 1 | def adjust_loss_weights(init_weight, current_epoch, mode='decay', start=400, every=20): 2 | # decay or rise the loss weights according to the given policy and current epoch 3 | # mode: decay, rise or binary 4 | 5 | if mode != 'binary': 6 | if current_epoch < start: 7 | if mode == 'rise': 8 | weight = init_weight * 1e-6 # use a very small weight for the normal loss in the beginning until the chamfer dist stabalizes 9 | else: 10 | weight = init_weight 11 | else: 12 | if every == 0: 13 | weight = init_weight # don't rise, keep const 14 | else: 15 | if mode == 'rise': 16 | weight = init_weight * (1.05 ** ((current_epoch - start) // every)) 17 | else: 18 | weight = init_weight * (0.85 ** ((current_epoch - start) // every)) 19 | 20 | return weight 21 | 22 | 23 | 24 | -------------------------------------------------------------------------------- /lib_data/README.md: -------------------------------------------------------------------------------- 1 | This folder contains example scripts that pre-processes and packs the data into the format required by SCALE. 2 | 3 | First, download the [example paired data](https://owncloud.tuebingen.mpg.de/index.php/s/3jJCdxEzD4HXW8y), unzip it under the `data/` folder, such that the data file paths are `/data/raw/`. The sample data contain: 4 | - A clothed body mesh (single-layered, as in the case of a real scan), 5 | - A fitted, unclothed SMPL body mesh for this clothed body. 6 | 7 | The following instructions will walk you through the data processing. The commands are to be executed in the `SCALE` repository folder. 8 | 9 | ### Pack the data for the SCALE data loader 10 | 11 | The example code `lib_data/pack_data_example.py` will take a pair of {clothed body mesh, unclothed body mesh}, and pack them into the `.npz` file as required by SCALE: 12 | 13 | ```bash 14 | python lib_data/pack_data_example.py 15 | ``` 16 | 17 | Essentially, the code: 18 | 19 | - saves the vertex locations of the minimally-clothed SMPL body, which is used for calculating local coordinate systems in SCALE; 20 | - renders a UV positional map from the minimally-clothed body (see below for details), which serves as the input to the SCALE network; 21 | - uniformly samples a specified number of points (together with their normals) on the clothed body mesh, which serves as the ground truth clothed body for SCALE training. 22 | 23 | The generated `.npz` file will be saved at `/data/packed/example/`. 24 | 25 | 26 | ### Get UV positional maps given a SMPL body 27 | This section explains the UV positional map rendering step above in more details. The following command won't generate the data needed for training; instead, it runs an example script that demonstrates the rendering of the positional map given the minimally-clothed SMPL template body mesh in `assets/` as input: 28 | 29 | ```bash 30 | cd lib_data/posmap_generator/ 31 | python apps/example.py 32 | ``` 33 | 34 | The script will generate the following outputs under `lib_data/posmap_generator/example_outputs/`: 35 | 36 | - `template_mesh_uv_posmap.png`: the UV positional map, i.e. the input to the SCALE model. Each valid pixel (see below) on this map corresponds to a point on the SMPL body; while their 3D location in R^3 may vary w.r.t. body pose, their relative locations on the body surface (2-manifold) are fixed. The 3-dimensional pixel value of a valid pixel is the (x, y, z) coordinate of its position in R^3. For a clearer demonstration, the default resolution is set to 128x128. In the paper we use the resolution 32x32 for a balance between speed and performance. 37 | - `template_mesh_uv_mask.png`: the binary UV mask of the SMPL model. The white region correspond to the valid pixels on the UV positional map. 38 | - `template_mesh_uv_in3D.obj`: a point cloud consisting the points that correspond to the valid pixels from the UV positional map. 39 | 40 | ### Note: data normalization 41 | 42 | When processing multiple data examples into a dataset, make sure the data are normalized. In the SCALE paper, we normalize the data by: 43 | - setting global orientation to 0 (as they are hardly relevant to the pose-dependent clothing deformation); 44 | - aligning the pelvis location of all data frames. Given the fitted SMPL body, the SMPL model is able to regress the body joint locations. Then simply offset all the meshes by the location of their root joint, i.e. the pelvis. If you need a detailed instruction on this, raise a issue on this repository, or refer to the [SMPL tutorials](https://smpl-made-simple.is.tue.mpg.de/), or contact the [SMPL supporting team](smpl@tuebingen.mpg.de). 45 | 46 | 47 | ### Miscellaneous 48 | Part of the functions under the folder `render_posmap` are adapted from the [PIFu repository](https://github.com/shunsukesaito/PIFu/tree/master/lib/renderer) (MIT license). 49 | -------------------------------------------------------------------------------- /lib_data/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qianlim/SCALE/98f5557e746de3536168c413901ff48c0571c5a4/lib_data/__init__.py -------------------------------------------------------------------------------- /lib_data/pack_data_example.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | from os.path import join, basename, realpath, dirname 4 | 5 | import numpy as np 6 | import trimesh 7 | 8 | PROJECT_DIR = dirname(dirname(realpath(__file__))) 9 | sys.path.append(PROJECT_DIR) 10 | 11 | 12 | def render_posmap(v_minimal, faces, uvs, faces_uvs, img_size=32): 13 | ''' 14 | v_minimal: vertices of the minimally-clothed SMPL body mesh 15 | faces: faces (triangles) of the minimally-clothed SMPL body mesh 16 | uvs: the uv coordinate of vertices of the SMPL body model 17 | faces_uvs: the faces (triangles) on the UV map of the SMPL body model 18 | ''' 19 | from lib_data.posmap_generator.lib.renderer.gl.pos_render import PosRender 20 | 21 | # instantiate renderer 22 | rndr = PosRender(width=img_size, height=img_size) 23 | 24 | # set mesh data on GPU 25 | rndr.set_mesh(v_minimal, faces, uvs, faces_uvs) 26 | 27 | # render 28 | rndr.display() 29 | 30 | # retrieve the rendered buffer 31 | uv_pos = rndr.get_color(0) 32 | uv_mask = uv_pos[:, :, 3] 33 | uv_pos = uv_pos[:, :, :3] 34 | 35 | uv_mask = uv_mask.reshape(-1) 36 | uv_pos = uv_pos.reshape(-1, 3) 37 | 38 | rendered_pos = uv_pos[uv_mask != 0.0] 39 | 40 | uv_pos = uv_pos.reshape(img_size, img_size, 3) 41 | 42 | # get face_id (triangle_id) per pixel 43 | face_id = uv_mask[uv_mask != 0].astype(np.int32) - 1 44 | 45 | assert len(face_id) == len(rendered_pos) 46 | 47 | return uv_pos, uv_mask, face_id 48 | 49 | 50 | class DataProcessor(object): 51 | ''' 52 | Example code for processing the paired data of (clothed body mesh, unclothed SMPL body mesh) into the format required by SCALE. 53 | Try it with the data (to be downloaded) in the data/raw/ folder. 54 | ''' 55 | def __init__(self, dataset_name='03375_blazerlong', 56 | n_sample_scan=40000, posmap_resl=32, 57 | uvs=None, faces_uvs=None): 58 | super().__init__() 59 | 60 | self.uvs = uvs 61 | self.faces_uvs = faces_uvs 62 | 63 | self.n_sample_scan = n_sample_scan # the number of points to sample on the (i.e. for the GT clothed body point cloud) 64 | self.posmap_resl = posmap_resl # resolution of the UV positional map 65 | 66 | self.save_root = join(PROJECT_DIR, 'data', 'packed', dataset_name) 67 | os.makedirs(self.save_root, exist_ok=True) 68 | 69 | 70 | def pack_single_file(self, minimal_fn, clothed_fn): 71 | result = {} 72 | 73 | scan = trimesh.load(clothed_fn, process=False) 74 | 75 | scan_pc, faceid = trimesh.sample.sample_surface_even(scan, self.n_sample_scan+100) # sample_even may cause smaller number of points sampled than wanted 76 | scan_pc = scan_pc[:self.n_sample_scan] 77 | faceid = faceid[:self.n_sample_scan] 78 | scan_n = scan.face_normals[faceid] 79 | result['scan_pc'] = scan_pc 80 | result['scan_n'] = scan_n 81 | result['scan_name'] = basename(clothed_fn).replace('.obj', '') 82 | 83 | minimal_mesh = trimesh.load(minimal_fn, process=False) 84 | posmap, _, _ = render_posmap(minimal_mesh.vertices, minimal_mesh.faces, self.uvs, self.faces_uvs, img_size=self.posmap_resl) 85 | result['posmap{}'.format(self.posmap_resl)] = posmap 86 | result['body_verts'] = minimal_mesh.vertices 87 | 88 | save_fn = join(self.save_root, basename(clothed_fn).replace('_clothed.obj', '_packed.npz')) 89 | np.savez(save_fn, **result) 90 | 91 | return result 92 | 93 | 94 | if __name__ == '__main__': 95 | import argparse 96 | parser = argparse.ArgumentParser() 97 | parser.add_argument('--dataset_name', type=str, default='example', help='name of the dataset to be created') 98 | parser.add_argument('--n_sample_scan', type=int, default=40000, help='number of the poinst to sample from the GT clothed body mesh surface') 99 | parser.add_argument('--posmap_resl', type=int, default=32, help='resolution of the UV positional map to be rendered') 100 | args = parser.parse_args() 101 | 102 | from lib_data.posmap_generator.lib.renderer.mesh import load_obj_mesh 103 | 104 | uv_template_fn = join(PROJECT_DIR, 'assets', 'template_mesh_uv.obj') 105 | verts, faces, uvs, faces_uvs = load_obj_mesh(uv_template_fn, with_texture=True) 106 | 107 | minimal_fn = join(PROJECT_DIR, 'data', 'raw', 'example_body_minimal.obj') 108 | clothed_fn = join(PROJECT_DIR, 'data', 'raw', 'example_body_clothed.obj') 109 | 110 | data = DataProcessor(dataset_name=args.dataset_name, uvs=uvs, faces_uvs=faces_uvs, n_sample_scan=args.n_sample_scan, posmap_resl=args.posmap_resl) 111 | data.pack_single_file(minimal_fn, clothed_fn) 112 | 113 | -------------------------------------------------------------------------------- /lib_data/posmap_generator/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qianlim/SCALE/98f5557e746de3536168c413901ff48c0571c5a4/lib_data/posmap_generator/__init__.py -------------------------------------------------------------------------------- /lib_data/posmap_generator/apps/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qianlim/SCALE/98f5557e746de3536168c413901ff48c0571c5a4/lib_data/posmap_generator/apps/__init__.py -------------------------------------------------------------------------------- /lib_data/posmap_generator/apps/example.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | import numpy as np 4 | import argparse 5 | 6 | from os.path import realpath, join, dirname, basename 7 | 8 | from lib.renderer.gl.pos_render import PosRender 9 | from lib.renderer.mesh import load_obj_mesh, save_obj_mesh 10 | 11 | 12 | def render_posmap(obj_file, H=32, W=32, output_dir='.'): 13 | # load obj file 14 | vertices, faces, uvs, faces_uvs = load_obj_mesh(obj_file, with_texture=True) 15 | 16 | # instantiate renderer 17 | rndr = PosRender(width=W, height=H) 18 | 19 | # set mesh data on GPU 20 | rndr.set_mesh(vertices, faces, uvs, faces_uvs) 21 | 22 | # render 23 | rndr.display() 24 | 25 | # retrieve the rendered buffer 26 | uv_pos = rndr.get_color(0) 27 | uv_mask = uv_pos[:,:,3] 28 | uv_pos = uv_pos[:,:,:3] 29 | 30 | # save mask file 31 | cv2.imwrite(join(output_dir, basename(obj_file).replace('.obj', '_posmap.png')), 255.0*uv_pos) 32 | 33 | # save mask file 34 | cv2.imwrite(join(output_dir, basename(obj_file).replace('.obj', '_mask.png')), 255.0*uv_mask) 35 | 36 | # save the rendered pos map as point cloud 37 | uv_mask = uv_mask.reshape(-1) 38 | uv_pos = uv_pos.reshape(-1,3) 39 | rendered_pos = uv_pos[uv_mask != 0.0] 40 | save_obj_mesh(join(output_dir, basename(obj_file).replace('.obj', '_in3D.obj')), rendered_pos) 41 | 42 | # get face_id per pixel 43 | face_id = uv_mask.astype(np.int32) - 1 44 | 45 | p_list = [] 46 | c_list = [] 47 | 48 | # randomly assign color per face id 49 | for i in range(faces.shape[0]): 50 | pos = uv_pos[face_id == i] 51 | c = np.random.rand(1,3).repeat(pos.shape[0],axis=0) 52 | p_list.append(pos) 53 | c_list.append(c) 54 | 55 | p_list = np.concatenate(p_list, 0) 56 | c_list = np.concatenate(c_list, 0) 57 | 58 | 59 | if __name__ == '__main__': 60 | parser = argparse.ArgumentParser() 61 | parser.add_argument('-i', '--input', type=str) 62 | args = parser.parse_args() 63 | 64 | SCRIPT_DIR = dirname(realpath(__file__)) 65 | smpl_template_pth = join(SCRIPT_DIR, '../../../assets/template_mesh_uv.obj') 66 | 67 | output_dir = join(SCRIPT_DIR, '../example_outputs') 68 | os.makedirs(output_dir, exist_ok=True) 69 | 70 | render_posmap(smpl_template_pth, H=128, W=128, output_dir=output_dir) 71 | -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qianlim/SCALE/98f5557e746de3536168c413901ff48c0571c5a4/lib_data/posmap_generator/lib/__init__.py -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qianlim/SCALE/98f5557e746de3536168c413901ff48c0571c5a4/lib_data/posmap_generator/lib/renderer/__init__.py -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/camera.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | 4 | from .glm import ortho 5 | 6 | 7 | class Camera: 8 | def __init__(self, width=1600, height=1200): 9 | # Focal Length 10 | # equivalent 50mm 11 | focal = np.sqrt(width * width + height * height) 12 | self.focal_x = focal 13 | self.focal_y = focal 14 | # Principal Point Offset 15 | self.principal_x = width / 2 16 | self.principal_y = height / 2 17 | # Axis Skew 18 | self.skew = 0 19 | # Image Size 20 | self.width = width 21 | self.height = height 22 | 23 | self.near = 1 24 | self.far = 10 25 | 26 | # Camera Center 27 | self.center = np.array([0, 0, 1.6]) 28 | self.direction = np.array([0, 0, -1]) 29 | self.right = np.array([1, 0, 0]) 30 | self.up = np.array([0, 1, 0]) 31 | 32 | self.ortho_ratio = None 33 | 34 | def sanity_check(self): 35 | self.center = self.center.reshape([-1]) 36 | self.direction = self.direction.reshape([-1]) 37 | self.right = self.right.reshape([-1]) 38 | self.up = self.up.reshape([-1]) 39 | 40 | assert len(self.center) == 3 41 | assert len(self.direction) == 3 42 | assert len(self.right) == 3 43 | assert len(self.up) == 3 44 | 45 | @staticmethod 46 | def normalize_vector(v): 47 | v_norm = np.linalg.norm(v) 48 | return v if v_norm == 0 else v / v_norm 49 | 50 | def get_real_z_value(self, z): 51 | z_near = self.near 52 | z_far = self.far 53 | z_n = 2.0 * z - 1.0 54 | z_e = 2.0 * z_near * z_far / (z_far + z_near - z_n * (z_far - z_near)) 55 | return z_e 56 | 57 | def get_rotation_matrix(self): 58 | rot_mat = np.eye(3) 59 | s = self.right 60 | s = self.normalize_vector(s) 61 | rot_mat[0, :] = s 62 | u = self.up 63 | u = self.normalize_vector(u) 64 | rot_mat[1, :] = -u 65 | rot_mat[2, :] = self.normalize_vector(self.direction) 66 | 67 | return rot_mat 68 | 69 | def get_translation_vector(self): 70 | rot_mat = self.get_rotation_matrix() 71 | trans = -np.dot(rot_mat, self.center) 72 | return trans 73 | 74 | def get_intrinsic_matrix(self): 75 | int_mat = np.eye(3) 76 | 77 | int_mat[0, 0] = self.focal_x 78 | int_mat[1, 1] = self.focal_y 79 | int_mat[0, 1] = self.skew 80 | int_mat[0, 2] = self.principal_x 81 | int_mat[1, 2] = self.principal_y 82 | 83 | return int_mat 84 | 85 | def get_projection_matrix(self): 86 | ext_mat = self.get_extrinsic_matrix() 87 | int_mat = self.get_intrinsic_matrix() 88 | 89 | return np.matmul(int_mat, ext_mat) 90 | 91 | def get_extrinsic_matrix(self): 92 | rot_mat = self.get_rotation_matrix() 93 | int_mat = self.get_intrinsic_matrix() 94 | trans = self.get_translation_vector() 95 | 96 | extrinsic = np.eye(4) 97 | extrinsic[:3, :3] = rot_mat 98 | extrinsic[:3, 3] = trans 99 | 100 | return extrinsic[:3, :] 101 | 102 | def set_rotation_matrix(self, rot_mat): 103 | self.direction = rot_mat[2, :] 104 | self.up = -rot_mat[1, :] 105 | self.right = rot_mat[0, :] 106 | 107 | def set_intrinsic_matrix(self, int_mat): 108 | self.focal_x = int_mat[0, 0] 109 | self.focal_y = int_mat[1, 1] 110 | self.skew = int_mat[0, 1] 111 | self.principal_x = int_mat[0, 2] 112 | self.principal_y = int_mat[1, 2] 113 | 114 | def set_projection_matrix(self, proj_mat): 115 | res = cv2.decomposeProjectionMatrix(proj_mat) 116 | int_mat, rot_mat, camera_center_homo = res[0], res[1], res[2] 117 | camera_center = camera_center_homo[0:3] / camera_center_homo[3] 118 | camera_center = camera_center.reshape(-1) 119 | int_mat = int_mat / int_mat[2][2] 120 | 121 | self.set_intrinsic_matrix(int_mat) 122 | self.set_rotation_matrix(rot_mat) 123 | self.center = camera_center 124 | 125 | self.sanity_check() 126 | 127 | def get_gl_matrix(self): 128 | z_near = self.near 129 | z_far = self.far 130 | rot_mat = self.get_rotation_matrix() 131 | int_mat = self.get_intrinsic_matrix() 132 | trans = self.get_translation_vector() 133 | 134 | extrinsic = np.eye(4) 135 | extrinsic[:3, :3] = rot_mat 136 | extrinsic[:3, 3] = trans 137 | axis_adj = np.eye(4) 138 | axis_adj[2, 2] = -1 139 | axis_adj[1, 1] = -1 140 | model_view = np.matmul(axis_adj, extrinsic) 141 | 142 | projective = np.zeros([4, 4]) 143 | projective[:2, :2] = int_mat[:2, :2] 144 | projective[:2, 2:3] = -int_mat[:2, 2:3] 145 | projective[3, 2] = -1 146 | projective[2, 2] = (z_near + z_far) 147 | projective[2, 3] = (z_near * z_far) 148 | 149 | if self.ortho_ratio is None: 150 | ndc = ortho(0, self.width, 0, self.height, z_near, z_far) 151 | perspective = np.matmul(ndc, projective) 152 | else: 153 | perspective = ortho(-self.width * self.ortho_ratio / 2, self.width * self.ortho_ratio / 2, 154 | -self.height * self.ortho_ratio / 2, self.height * self.ortho_ratio / 2, 155 | z_near, z_far) 156 | 157 | return perspective, model_view 158 | 159 | 160 | def KRT_from_P(proj_mat, normalize_K=True): 161 | res = cv2.decomposeProjectionMatrix(proj_mat) 162 | K, Rot, camera_center_homog = res[0], res[1], res[2] 163 | camera_center = camera_center_homog[0:3] / camera_center_homog[3] 164 | trans = -Rot.dot(camera_center) 165 | if normalize_K: 166 | K = K / K[2][2] 167 | return K, Rot, trans 168 | 169 | 170 | def MVP_from_P(proj_mat, width, height, near=0.1, far=10000): 171 | ''' 172 | Convert OpenCV camera calibration matrix to OpenGL projection and model view matrix 173 | :param proj_mat: OpenCV camera projeciton matrix 174 | :param width: Image width 175 | :param height: Image height 176 | :param near: Z near value 177 | :param far: Z far value 178 | :return: OpenGL projection matrix and model view matrix 179 | ''' 180 | res = cv2.decomposeProjectionMatrix(proj_mat) 181 | K, Rot, camera_center_homog = res[0], res[1], res[2] 182 | camera_center = camera_center_homog[0:3] / camera_center_homog[3] 183 | trans = -Rot.dot(camera_center) 184 | K = K / K[2][2] 185 | 186 | extrinsic = np.eye(4) 187 | extrinsic[:3, :3] = Rot 188 | extrinsic[:3, 3:4] = trans 189 | axis_adj = np.eye(4) 190 | axis_adj[2, 2] = -1 191 | axis_adj[1, 1] = -1 192 | model_view = np.matmul(axis_adj, extrinsic) 193 | 194 | zFar = far 195 | zNear = near 196 | projective = np.zeros([4, 4]) 197 | projective[:2, :2] = K[:2, :2] 198 | projective[:2, 2:3] = -K[:2, 2:3] 199 | projective[3, 2] = -1 200 | projective[2, 2] = (zNear + zFar) 201 | projective[2, 3] = (zNear * zFar) 202 | 203 | ndc = ortho(0, width, 0, height, zNear, zFar) 204 | 205 | perspective = np.matmul(ndc, projective) 206 | 207 | return perspective, model_view 208 | -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/core.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qianlim/SCALE/98f5557e746de3536168c413901ff48c0571c5a4/lib_data/posmap_generator/lib/renderer/core.py -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/egl/__init__.py: -------------------------------------------------------------------------------- 1 | from .egl_framework import * 2 | from .egl_render import * 3 | from .egl_cam_render import * 4 | -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/egl/data/pos_uv.fs: -------------------------------------------------------------------------------- 1 | #version 410 2 | 3 | in VertexData { 4 | vec3 Position; 5 | } VertexIn; 6 | 7 | in int gl_PrimitiveID; 8 | 9 | layout (location = 0) out vec4 FragPosition; 10 | 11 | void main() 12 | { 13 | FragPosition = vec4(VertexIn.Position,1.0+float(gl_PrimitiveID)); 14 | } -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/egl/data/pos_uv.vs: -------------------------------------------------------------------------------- 1 | #version 330 2 | 3 | layout (location = 0) in vec3 a_Position; 4 | layout (location = 1) in vec2 a_TextureCoord; 5 | 6 | out VertexData { 7 | vec3 Position; 8 | } VertexOut; 9 | 10 | void main() 11 | { 12 | VertexOut.Position = a_Position; 13 | VertexOut.Texcoord = a_TextureCoord; 14 | 15 | gl_Position = vec4(a_TextureCoord, 0.0, 1.0) - vec4(0.5, 0.5, 0, 0); 16 | gl_Position[0] *= 2.0; 17 | gl_Position[1] *= 2.0; 18 | } 19 | -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/egl/data/quad.fs: -------------------------------------------------------------------------------- 1 | #version 330 core 2 | out vec4 FragColor; 3 | 4 | in vec2 TexCoord; 5 | 6 | uniform sampler2D screenTexture; 7 | 8 | void main() 9 | { 10 | FragColor = texture(screenTexture, TexCoord); 11 | } -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/egl/data/quad.vs: -------------------------------------------------------------------------------- 1 | #version 330 core 2 | layout (location = 0) in vec2 aPos; 3 | layout (location = 1) in vec2 aTexCoord; 4 | 5 | out vec2 TexCoord; 6 | 7 | void main() 8 | { 9 | gl_Position = vec4(aPos.x, aPos.y, 0.0, 1.0); 10 | TexCoord = aTexCoord; 11 | } -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/egl/egl_cam_render.py: -------------------------------------------------------------------------------- 1 | from .egl_render import EGLRender 2 | 3 | 4 | class CamRender(EGLRender): 5 | def __init__(self, width=1600, height=1200, name='Cam Renderer', 6 | program_files=['simple.fs', 'simple.vs']): 7 | EGLRender.__init__(self, width, height, name, program_files) 8 | self.camera = None 9 | 10 | def set_camera(self, camera): 11 | self.camera = camera 12 | self.projection_matrix, self.model_view_matrix = camera.get_gl_matrix() 13 | -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/egl/egl_framework.py: -------------------------------------------------------------------------------- 1 | # Mario Rosasco, 2016 2 | # adapted from framework.cpp, Copyright (C) 2010-2012 by Jason L. McKesson 3 | # This file is licensed under the MIT License. 4 | # 5 | # NB: Unlike in the framework.cpp organization, the main loop is contained 6 | # in the tutorial files, not in this framework file. Additionally, a copy of 7 | # this module file must exist in the same directory as the tutorial files 8 | # to be imported properly. 9 | 10 | import os 11 | 12 | 13 | 14 | # Function that creates and compiles shaders according to the given type (a GL enum value) and 15 | # shader program (a file containing a GLSL program). 16 | def loadShader(shaderType, shaderFile): 17 | import OpenGL.GL as gl 18 | 19 | # check if file exists, get full path name 20 | strFilename = findFileOrThrow(shaderFile) 21 | shaderData = None 22 | with open(strFilename, 'r') as f: 23 | shaderData = f.read() 24 | 25 | shader = gl.glCreateShader(shaderType) 26 | 27 | gl.glShaderSource(shader, shaderData) # note that this is a simpler function call than in C 28 | 29 | # This shader compilation is more explicit than the one used in 30 | # framework.cpp, which relies on a glutil wrapper function. 31 | # This is made explicit here mainly to decrease dependence on pyOpenGL 32 | # utilities and wrappers, which docs caution may change in future versions. 33 | gl.glCompileShader(shader) 34 | 35 | status = gl.glGetShaderiv(shader, gl.GL_COMPILE_STATUS) 36 | if status == gl.GL_FALSE: 37 | # Note that getting the error log is much simpler in Python than in C/C++ 38 | # and does not require explicit handling of the string buffer 39 | strInfoLog = gl.glGetShaderInfoLog(shader) 40 | strShaderType = "" 41 | if shaderType is gl.GL_VERTEX_SHADER: 42 | strShaderType = "vertex" 43 | elif shaderType is gl.GL_GEOMETRY_SHADER: 44 | strShaderType = "geometry" 45 | elif shaderType is gl.GL_FRAGMENT_SHADER: 46 | strShaderType = "fragment" 47 | 48 | print("Compilation failure for " + strShaderType + " shader:\n" + str(strInfoLog)) 49 | 50 | return shader 51 | 52 | 53 | # Function that accepts a list of shaders, compiles them, and returns a handle to the compiled program 54 | def createProgram(shaderList): 55 | import OpenGL.GL as gl 56 | 57 | program = gl.glCreateProgram() 58 | 59 | for shader in shaderList: 60 | gl.glAttachShader(program, shader) 61 | 62 | gl.glLinkProgram(program) 63 | 64 | status = gl.glGetProgramiv(program, gl.GL_LINK_STATUS) 65 | if status == gl.GL_FALSE: 66 | # Note that getting the error log is much simpler in Python than in C/C++ 67 | # and does not require explicit handling of the string buffer 68 | strInfoLog = gl.glGetProgramInfoLog(program) 69 | print("Linker failure: \n" + str(strInfoLog)) 70 | 71 | for shader in shaderList: 72 | gl.glDetachShader(program, shader) 73 | 74 | return program 75 | 76 | 77 | # Helper function to locate and open the target file (passed in as a string). 78 | # Returns the full path to the file as a string. 79 | def findFileOrThrow(strBasename): 80 | # Keep constant names in C-style convention, for readability 81 | # when comparing to C(/C++) code. 82 | if os.path.isfile(strBasename): 83 | return strBasename 84 | 85 | LOCAL_FILE_DIR = "data" + os.sep 86 | GLOBAL_FILE_DIR = os.path.dirname(os.path.abspath(__file__)) + os.sep + "../gl/data" + os.sep 87 | 88 | strFilename = LOCAL_FILE_DIR + strBasename 89 | if os.path.isfile(strFilename): 90 | return strFilename 91 | 92 | strFilename = GLOBAL_FILE_DIR + strBasename 93 | if os.path.isfile(strFilename): 94 | return strFilename 95 | 96 | raise IOError('Could not find target file ' + strBasename) 97 | -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/egl/egl_pos_render.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import random 3 | 4 | from .egl_framework import * 5 | from .egl_cam_render import CamRender 6 | 7 | from OpenGL.GL import * 8 | 9 | class PosRender(CamRender): 10 | def __init__(self, width=256, height=256, name='Position Renderer'): 11 | CamRender.__init__(self, width, height, name, program_files=['pos_uv.vs', 'pos_uv.fs']) 12 | 13 | def draw(self): 14 | self.draw_init() 15 | 16 | glEnable(GL_MULTISAMPLE) 17 | 18 | glUseProgram(self.program) 19 | glUniformMatrix4fv(self.model_mat_unif, 1, GL_FALSE, self.model_view_matrix.transpose()) 20 | glUniformMatrix4fv(self.persp_mat_unif, 1, GL_FALSE, self.projection_matrix.transpose()) 21 | 22 | glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer) 23 | 24 | glEnableVertexAttribArray(0) 25 | glVertexAttribPointer(0, self.vertex_dim, GL_DOUBLE, GL_FALSE, 0, None) 26 | 27 | glDrawArrays(GL_TRIANGLES, 0, self.n_vertices) 28 | 29 | glDisableVertexAttribArray(0) 30 | 31 | glBindBuffer(GL_ARRAY_BUFFER, 0) 32 | 33 | glUseProgram(0) 34 | 35 | glDisable(GL_MULTISAMPLE) 36 | 37 | self.draw_end() -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/egl/egl_render.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from .glcontext import create_opengl_context 3 | from .egl_framework import * 4 | 5 | _context_inited = None 6 | import OpenGL.GL as gl 7 | 8 | class EGLRender: 9 | def __init__(self, width=1600, height=1200, name='GL Renderer', 10 | program_files=['simple.fs', 'simple.vs']): 11 | self.width = width 12 | self.height = height 13 | self.name = name 14 | self.use_inverse_depth = False 15 | 16 | self.start = 0 17 | 18 | global _context_inited 19 | if _context_inited is None: 20 | create_opengl_context((width, height)) 21 | _context_inited = True 22 | 23 | gl.glEnable(gl.GL_DEPTH_TEST) 24 | 25 | gl.glClampColor(gl.GL_CLAMP_READ_COLOR, gl.GL_FALSE) 26 | gl.glClampColor(gl.GL_CLAMP_FRAGMENT_COLOR, gl.GL_FALSE) 27 | gl.glClampColor(gl.GL_CLAMP_VERTEX_COLOR, gl.GL_FALSE) 28 | 29 | # init program 30 | shader_list = [] 31 | 32 | for program_file in program_files: 33 | _, ext = os.path.splitext(program_file) 34 | if ext == '.vs': 35 | shader_list.append(loadShader(gl.GL_VERTEX_SHADER, program_file)) 36 | elif ext == '.fs': 37 | shader_list.append(loadShader(gl.GL_FRAGMENT_SHADER, program_file)) 38 | elif ext == '.gs': 39 | shader_list.append(loadShader(gl.GL_GEOMETRY_SHADER, program_file)) 40 | 41 | self.program = createProgram(shader_list) 42 | 43 | for shader in shader_list: 44 | gl.glDeleteShader(shader) 45 | 46 | # Init uniform variables 47 | self.model_mat_unif = gl.glGetUniformLocation(self.program, 'ModelMat') 48 | self.persp_mat_unif = gl.glGetUniformLocation(self.program, 'PerspMat') 49 | 50 | self.vertex_buffer = gl.glGenBuffers(1) 51 | 52 | # Configure frame buffer 53 | self.frame_buffer = gl.glGenFramebuffers(1) 54 | gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, self.frame_buffer) 55 | 56 | # Configure texture buffer to render to 57 | self.color_buffer = gl.glGenTextures(1) 58 | multi_sample_rate = 32 59 | gl.glBindTexture(gl.GL_TEXTURE_2D_MULTISAMPLE, self.color_buffer) 60 | gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_S, gl.GL_CLAMP_TO_EDGE) 61 | gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_T, gl.GL_CLAMP_TO_EDGE) 62 | gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAG_FILTER, gl.GL_LINEAR) 63 | gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MIN_FILTER, gl.GL_LINEAR) 64 | gl.glTexImage2DMultisample(gl.GL_TEXTURE_2D_MULTISAMPLE, multi_sample_rate, gl.GL_RGBA32F, self.width, self.height, gl.GL_TRUE) 65 | gl.glBindTexture(gl.GL_TEXTURE_2D_MULTISAMPLE, 0) 66 | gl.glFramebufferTexture2D(gl.GL_FRAMEBUFFER, gl.GL_COLOR_ATTACHMENT0, gl.GL_TEXTURE_2D_MULTISAMPLE, self.color_buffer, 0) 67 | 68 | self.render_buffer = gl.glGenRenderbuffers(1) 69 | gl.glBindRenderbuffer(gl.GL_RENDERBUFFER, self.render_buffer) 70 | gl.glRenderbufferStorageMultisample(gl.GL_RENDERBUFFER, multi_sample_rate, gl.GL_DEPTH24_STENCIL8, self.width, self.height) 71 | gl.glBindRenderbuffer(gl.GL_RENDERBUFFER, 0) 72 | gl.glFramebufferRenderbuffer(gl.GL_FRAMEBUFFER, gl.GL_DEPTH_STENCIL_ATTACHMENT, gl.GL_RENDERBUFFER, self.render_buffer) 73 | 74 | gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, 0) 75 | 76 | self.intermediate_fbo = gl.glGenFramebuffers(1) 77 | gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, self.intermediate_fbo) 78 | 79 | self.screen_texture = gl.glGenTextures(1) 80 | gl.glBindTexture(gl.GL_TEXTURE_2D, self.screen_texture) 81 | gl.glTexImage2D(gl.GL_TEXTURE_2D, 0, gl.GL_RGBA32F, self.width, self.height, 0, gl.GL_RGBA, gl.GL_FLOAT, None) 82 | gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MIN_FILTER, gl.GL_LINEAR) 83 | gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAG_FILTER, gl.GL_LINEAR) 84 | gl.glFramebufferTexture2D(gl.GL_FRAMEBUFFER, gl.GL_COLOR_ATTACHMENT0, gl.GL_TEXTURE_2D, self.screen_texture, 0) 85 | 86 | gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, 0) 87 | 88 | # Configure texture buffer if needed 89 | self.render_texture = None 90 | 91 | # NOTE: original render_texture only support one input 92 | # this is tentative member of this issue 93 | self.render_texture_v2 = {} 94 | 95 | # Inner storage for buffer data 96 | self.vertex_data = None 97 | self.vertex_dim = None 98 | self.n_vertices = None 99 | 100 | self.model_view_matrix = None 101 | self.projection_matrix = None 102 | 103 | def set_mesh(self, vertices, faces): 104 | self.vertex_data = vertices[faces.reshape([-1])] 105 | self.vertex_dim = self.vertex_data.shape[1] 106 | self.n_vertices = self.vertex_data.shape[0] 107 | 108 | gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.vertex_buffer) 109 | gl.glBufferData(gl.GL_ARRAY_BUFFER, self.vertex_data, gl.GL_STATIC_DRAW) 110 | 111 | gl.glBindBuffer(gl.GL_ARRAY_BUFFER, 0) 112 | 113 | def select_array(self, start, size): 114 | self.start = start 115 | self.size = size 116 | 117 | def set_viewpoint(self, projection, model_view): 118 | self.projection_matrix = projection 119 | self.model_view_matrix = model_view 120 | 121 | def draw_init(self): 122 | gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, self.frame_buffer) 123 | gl.glEnable(gl.GL_DEPTH_TEST) 124 | 125 | gl.glClearColor(0.0, 0.0, 0.0, 0.0) 126 | if self.use_inverse_depth: 127 | gl.glDepthFunc(gl.GL_GREATER) 128 | gl.glClearDepth(0.0) 129 | else: 130 | gl.glDepthFunc(gl.GL_LESS) 131 | gl.glClearDepth(1.0) 132 | gl.glClear(gl.GL_COLOR_BUFFER_BIT | gl.GL_DEPTH_BUFFER_BIT) 133 | 134 | def draw_end(self): 135 | gl.glBindFramebuffer(gl.GL_READ_FRAMEBUFFER, self.frame_buffer) 136 | gl.glBindFramebuffer(gl.GL_DRAW_FRAMEBUFFER, self.intermediate_fbo) 137 | gl.glBlitFramebuffer(0, 0, self.width, self.height, 0, 0, self.width, self.height, gl.GL_COLOR_BUFFER_BIT, gl.GL_NEAREST) 138 | 139 | gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, 0) 140 | gl.glDepthFunc(gl.GL_LESS) 141 | gl.glClearDepth(1.0) 142 | 143 | def draw(self): 144 | self.draw_init() 145 | 146 | gl.glUseProgram(self.program) 147 | gl.glUniformMatrix4fv(self.model_mat_unif, 1, gl.GL_FALSE, self.model_view_matrix.transpose()) 148 | gl.glUniformMatrix4fv(self.persp_mat_unif, 1, gl.GL_FALSE, self.projection_matrix.transpose()) 149 | 150 | gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.vertex_buffer) 151 | 152 | gl.glEnableVertexAttribArray(0) 153 | gl.glVertexAttribPointer(0, self.vertex_dim, gl.GL_DOUBLE, gl.GL_FALSE, 0, None) 154 | 155 | gl.glDrawArrays(gl.GL_TRIANGLES, 0, self.n_vertices) 156 | 157 | gl.glDisableVertexAttribArray(0) 158 | 159 | gl.glBindBuffer(gl.GL_ARRAY_BUFFER, 0) 160 | 161 | gl.glUseProgram(0) 162 | 163 | self.draw_end() 164 | 165 | def get_color(self): 166 | gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, self.intermediate_fbo) 167 | gl.glReadBuffer(gl.GL_COLOR_ATTACHMENT0) 168 | data = gl.glReadPixels(0, 0, self.width, self.height, gl.GL_RGBA, gl.GL_FLOAT, outputType=None) 169 | gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, 0) 170 | rgb = data.reshape(self.height, self.width, -1) 171 | rgb = np.flip(rgb, 0) 172 | return rgb 173 | 174 | def get_z_value(self): 175 | gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, self.intermediate_fbo) 176 | data = gl.glReadPixels(0, 0, self.width, self.height, gl.GL_DEPTH_COMPONENT, gl.GL_FLOAT, outputType=None) 177 | gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, 0) 178 | z = data.reshape(self.height, self.width) 179 | z = np.flip(z, 0) 180 | return z 181 | 182 | -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/egl/glcontext.py: -------------------------------------------------------------------------------- 1 | """Headless GPU-accelerated OpenGL context creation on Google Colaboratory. 2 | 3 | Typical usage: 4 | 5 | # Optional PyOpenGL configuratiopn can be done here. 6 | # import OpenGL 7 | # OpenGL.ERROR_CHECKING = True 8 | 9 | # 'glcontext' must be imported before any OpenGL.* API. 10 | from lucid.misc.gl.glcontext import create_opengl_context 11 | 12 | # Now it's safe to import OpenGL and EGL functions 13 | import OpenGL.GL as gl 14 | 15 | # create_opengl_context() creates a GL context that is attached to an 16 | # offscreen surface of the specified size. Note that rendering to buffers 17 | # of other sizes and formats is still possible with OpenGL Framebuffers. 18 | # 19 | # Users are expected to directly use the EGL API in case more advanced 20 | # context management is required. 21 | width, height = 640, 480 22 | create_opengl_context((width, height)) 23 | 24 | # OpenGL context is available here. 25 | 26 | """ 27 | 28 | from __future__ import print_function 29 | 30 | # pylint: disable=unused-import,g-import-not-at-top,g-statement-before-imports 31 | 32 | try: 33 | import OpenGL 34 | except: 35 | print('This module depends on PyOpenGL.') 36 | print('Please run "\033[1m!pip install -q pyopengl\033[0m" ' 37 | 'prior importing this module.') 38 | raise 39 | 40 | import ctypes 41 | from ctypes import pointer, util 42 | import os 43 | 44 | os.environ['PYOPENGL_PLATFORM'] = 'egl' 45 | 46 | # OpenGL loading workaround. 47 | # 48 | # * PyOpenGL tries to load libGL, but we need libOpenGL, see [1,2]. 49 | # This could have been solved by a symlink libGL->libOpenGL, but: 50 | # 51 | # * Python 2.7 can't find libGL and linEGL due to a bug (see [3]) 52 | # in ctypes.util, that was only wixed in Python 3.6. 53 | # 54 | # So, the only solution I've found is to monkeypatch ctypes.util 55 | # [1] https://devblogs.nvidia.com/egl-eye-opengl-visualization-without-x-server/ 56 | # [2] https://devblogs.nvidia.com/linking-opengl-server-side-rendering/ 57 | # [3] https://bugs.python.org/issue9998 58 | _find_library_old = ctypes.util.find_library 59 | try: 60 | 61 | def _find_library_new(name): 62 | return { 63 | 'GL': 'libOpenGL.so', 64 | 'EGL': 'libEGL.so', 65 | }.get(name, _find_library_old(name)) 66 | util.find_library = _find_library_new 67 | import OpenGL.GL as gl 68 | import OpenGL.EGL as egl 69 | except: 70 | print('Unable to load OpenGL libraries. ' 71 | 'Make sure you use GPU-enabled backend.') 72 | print('Press "Runtime->Change runtime type" and set ' 73 | '"Hardware accelerator" to GPU.') 74 | raise 75 | finally: 76 | util.find_library = _find_library_old 77 | 78 | 79 | def create_opengl_context(surface_size=(640, 480)): 80 | """Create offscreen OpenGL context and make it current. 81 | 82 | Users are expected to directly use EGL API in case more advanced 83 | context management is required. 84 | 85 | Args: 86 | surface_size: (width, height), size of the offscreen rendering surface. 87 | """ 88 | egl_display = egl.eglGetDisplay(egl.EGL_DEFAULT_DISPLAY) 89 | 90 | major, minor = egl.EGLint(), egl.EGLint() 91 | egl.eglInitialize(egl_display, pointer(major), pointer(minor)) 92 | 93 | config_attribs = [ 94 | egl.EGL_SURFACE_TYPE, egl.EGL_PBUFFER_BIT, egl.EGL_BLUE_SIZE, 8, 95 | egl.EGL_GREEN_SIZE, 8, egl.EGL_RED_SIZE, 8, egl.EGL_DEPTH_SIZE, 24, 96 | egl.EGL_RENDERABLE_TYPE, egl.EGL_OPENGL_BIT, egl.EGL_NONE 97 | ] 98 | config_attribs = (egl.EGLint * len(config_attribs))(*config_attribs) 99 | 100 | num_configs = egl.EGLint() 101 | egl_cfg = egl.EGLConfig() 102 | egl.eglChooseConfig(egl_display, config_attribs, pointer(egl_cfg), 1, 103 | pointer(num_configs)) 104 | 105 | width, height = surface_size 106 | pbuffer_attribs = [ 107 | egl.EGL_WIDTH, 108 | width, 109 | egl.EGL_HEIGHT, 110 | height, 111 | egl.EGL_NONE, 112 | ] 113 | pbuffer_attribs = (egl.EGLint * len(pbuffer_attribs))(*pbuffer_attribs) 114 | egl_surf = egl.eglCreatePbufferSurface(egl_display, egl_cfg, pbuffer_attribs) 115 | 116 | egl.eglBindAPI(egl.EGL_OPENGL_API) 117 | 118 | egl_context = egl.eglCreateContext(egl_display, egl_cfg, egl.EGL_NO_CONTEXT, 119 | None) 120 | egl.eglMakeCurrent(egl_display, egl_surf, egl_surf, egl_context) -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/gl/__init__.py: -------------------------------------------------------------------------------- 1 | from .framework import * 2 | from .render import * 3 | from .cam_render import * 4 | -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/gl/cam_render.py: -------------------------------------------------------------------------------- 1 | from OpenGL.GLUT import * 2 | 3 | from .render import Render 4 | 5 | 6 | class CamRender(Render): 7 | def __init__(self, width=1600, height=1200, name='Cam Renderer', 8 | program_files=['simple.fs', 'simple.vs'], color_size=1): 9 | Render.__init__(self, width, height, name, program_files, color_size) 10 | self.camera = None 11 | 12 | glutDisplayFunc(self.display) 13 | glutKeyboardFunc(self.keyboard) 14 | 15 | def set_camera(self, camera): 16 | self.camera = camera 17 | self.projection_matrix, self.model_view_matrix = camera.get_gl_matrix() 18 | 19 | def keyboard(self, key, x, y): 20 | # up 21 | eps = 1 22 | # print(key) 23 | if key == b'w': 24 | self.camera.center += eps * self.camera.direction 25 | elif key == b's': 26 | self.camera.center -= eps * self.camera.direction 27 | if key == b'a': 28 | self.camera.center -= eps * self.camera.right 29 | elif key == b'd': 30 | self.camera.center += eps * self.camera.right 31 | if key == b' ': 32 | self.camera.center += eps * self.camera.up 33 | elif key == b'x': 34 | self.camera.center -= eps * self.camera.up 35 | elif key == b'i': 36 | self.camera.near += 0.1 * eps 37 | self.camera.far += 0.1 * eps 38 | elif key == b'o': 39 | self.camera.near -= 0.1 * eps 40 | self.camera.far -= 0.1 * eps 41 | 42 | self.projection_matrix, self.model_view_matrix = self.camera.get_gl_matrix() 43 | 44 | def show(self): 45 | glutMainLoop() 46 | -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/gl/data/pos_uv.fs: -------------------------------------------------------------------------------- 1 | #version 410 2 | 3 | in VertexData { 4 | vec3 Position; 5 | } VertexIn; 6 | 7 | in int gl_PrimitiveID; 8 | 9 | layout (location = 0) out vec4 FragPosition; 10 | 11 | void main() 12 | { 13 | FragPosition = vec4(VertexIn.Position,1.0+float(gl_PrimitiveID)); 14 | } -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/gl/data/pos_uv.vs: -------------------------------------------------------------------------------- 1 | #version 410 2 | 3 | layout (location = 0) in vec3 a_Position; 4 | layout (location = 1) in vec2 a_TextureCoord; 5 | 6 | out VertexData { 7 | vec3 Position; 8 | } VertexOut; 9 | 10 | void main() 11 | { 12 | VertexOut.Position = a_Position; 13 | 14 | gl_Position = vec4(a_TextureCoord, 0.0, 1.0) - vec4(0.5, 0.5, 0, 0); 15 | gl_Position[0] *= 2.0; 16 | gl_Position[1] *= 2.0; 17 | } 18 | -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/gl/data/quad.fs: -------------------------------------------------------------------------------- 1 | #version 330 core 2 | out vec4 FragColor; 3 | 4 | in vec2 TexCoord; 5 | 6 | uniform sampler2D screenTexture; 7 | 8 | void main() 9 | { 10 | FragColor = texture(screenTexture, TexCoord); 11 | } -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/gl/data/quad.vs: -------------------------------------------------------------------------------- 1 | #version 330 core 2 | layout (location = 0) in vec2 aPos; 3 | layout (location = 1) in vec2 aTexCoord; 4 | 5 | out vec2 TexCoord; 6 | 7 | void main() 8 | { 9 | gl_Position = vec4(aPos.x, aPos.y, 0.0, 1.0); 10 | TexCoord = aTexCoord; 11 | } -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/gl/framework.py: -------------------------------------------------------------------------------- 1 | # Mario Rosasco, 2016 2 | # adapted from framework.cpp, Copyright (C) 2010-2012 by Jason L. McKesson 3 | # This file is licensed under the MIT License. 4 | # 5 | # NB: Unlike in the framework.cpp organization, the main loop is contained 6 | # in the tutorial files, not in this framework file. Additionally, a copy of 7 | # this module file must exist in the same directory as the tutorial files 8 | # to be imported properly. 9 | 10 | import os 11 | 12 | from OpenGL.GL import * 13 | 14 | 15 | # Function that creates and compiles shaders according to the given type (a GL enum value) and 16 | # shader program (a file containing a GLSL program). 17 | def loadShader(shaderType, shaderFile): 18 | # check if file exists, get full path name 19 | strFilename = findFileOrThrow(shaderFile) 20 | shaderData = None 21 | with open(strFilename, 'r') as f: 22 | shaderData = f.read() 23 | 24 | shader = glCreateShader(shaderType) 25 | glShaderSource(shader, shaderData) # note that this is a simpler function call than in C 26 | 27 | # This shader compilation is more explicit than the one used in 28 | # framework.cpp, which relies on a glutil wrapper function. 29 | # This is made explicit here mainly to decrease dependence on pyOpenGL 30 | # utilities and wrappers, which docs caution may change in future versions. 31 | glCompileShader(shader) 32 | 33 | status = glGetShaderiv(shader, GL_COMPILE_STATUS) 34 | if status == GL_FALSE: 35 | # Note that getting the error log is much simpler in Python than in C/C++ 36 | # and does not require explicit handling of the string buffer 37 | strInfoLog = glGetShaderInfoLog(shader) 38 | strShaderType = "" 39 | if shaderType is GL_VERTEX_SHADER: 40 | strShaderType = "vertex" 41 | elif shaderType is GL_GEOMETRY_SHADER: 42 | strShaderType = "geometry" 43 | elif shaderType is GL_FRAGMENT_SHADER: 44 | strShaderType = "fragment" 45 | 46 | print("Compilation failure for " + strShaderType + " shader:\n" + str(strInfoLog)) 47 | 48 | return shader 49 | 50 | 51 | # Function that accepts a list of shaders, compiles them, and returns a handle to the compiled program 52 | def createProgram(shaderList): 53 | program = glCreateProgram() 54 | 55 | for shader in shaderList: 56 | glAttachShader(program, shader) 57 | 58 | glLinkProgram(program) 59 | 60 | status = glGetProgramiv(program, GL_LINK_STATUS) 61 | if status == GL_FALSE: 62 | # Note that getting the error log is much simpler in Python than in C/C++ 63 | # and does not require explicit handling of the string buffer 64 | strInfoLog = glGetProgramInfoLog(program) 65 | print("Linker failure: \n" + str(strInfoLog)) 66 | 67 | for shader in shaderList: 68 | glDetachShader(program, shader) 69 | 70 | return program 71 | 72 | 73 | # Helper function to locate and open the target file (passed in as a string). 74 | # Returns the full path to the file as a string. 75 | def findFileOrThrow(strBasename): 76 | # Keep constant names in C-style convention, for readability 77 | # when comparing to C(/C++) code. 78 | if os.path.isfile(strBasename): 79 | return strBasename 80 | 81 | LOCAL_FILE_DIR = "data" + os.sep 82 | GLOBAL_FILE_DIR = os.path.dirname(os.path.abspath(__file__)) + os.sep + "data" + os.sep 83 | 84 | strFilename = LOCAL_FILE_DIR + strBasename 85 | if os.path.isfile(strFilename): 86 | return strFilename 87 | 88 | strFilename = GLOBAL_FILE_DIR + strBasename 89 | if os.path.isfile(strFilename): 90 | return strFilename 91 | 92 | raise IOError('Could not find target file ' + strBasename) 93 | -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/gl/pos_render.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from .framework import * 4 | from .cam_render import CamRender 5 | 6 | 7 | class PosRender(CamRender): 8 | def __init__(self, width=256, height=256, name='Position Renderer'): 9 | CamRender.__init__(self, width, height, name, program_files=['pos_uv.vs', 'pos_uv.fs']) 10 | 11 | self.uv_buffer = glGenBuffers(1) 12 | self.uv_data = None 13 | 14 | def set_mesh(self, vertices, faces, uvs, faces_uv): 15 | self.vertex_data = vertices[faces.reshape([-1])] 16 | self.vertex_dim = self.vertex_data.shape[1] 17 | self.n_vertices = self.vertex_data.shape[0] 18 | 19 | self.uv_data = uvs[faces_uv.reshape([-1])] 20 | 21 | glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer) 22 | glBufferData(GL_ARRAY_BUFFER, self.vertex_data, GL_STATIC_DRAW) 23 | 24 | glBindBuffer(GL_ARRAY_BUFFER, self.uv_buffer) 25 | glBufferData(GL_ARRAY_BUFFER, self.uv_data, GL_STATIC_DRAW) 26 | 27 | glBindBuffer(GL_ARRAY_BUFFER, 0) 28 | 29 | def draw(self): 30 | self.draw_init() 31 | 32 | glUseProgram(self.program) 33 | 34 | glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer) 35 | glEnableVertexAttribArray(0) 36 | glVertexAttribPointer(0, self.vertex_dim, GL_DOUBLE, GL_FALSE, 0, None) 37 | 38 | glBindBuffer(GL_ARRAY_BUFFER, self.uv_buffer) 39 | glEnableVertexAttribArray(1) 40 | glVertexAttribPointer(1, 2, GL_DOUBLE, GL_FALSE, 0, None) 41 | 42 | glDrawArrays(GL_TRIANGLES, 0, self.n_vertices) 43 | 44 | glDisableVertexAttribArray(1) 45 | glDisableVertexAttribArray(0) 46 | 47 | glBindBuffer(GL_ARRAY_BUFFER, 0) 48 | 49 | glUseProgram(0) 50 | 51 | self.draw_end() -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/gl/render.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from OpenGL.GLUT import * 3 | from .framework import * 4 | 5 | _glut_window = None 6 | 7 | class Render: 8 | def __init__(self, width=1600, height=1200, name='GL Renderer', 9 | program_files=['simple.fs', 'simple.vs'], color_size=1, ms_rate=1): 10 | self.width = width 11 | self.height = height 12 | self.name = name 13 | self.display_mode = GLUT_DOUBLE | GLUT_RGB | GLUT_DEPTH 14 | self.use_inverse_depth = False 15 | 16 | global _glut_window 17 | if _glut_window is None: 18 | glutInit() 19 | glutInitDisplayMode(self.display_mode) 20 | glutInitWindowSize(self.width, self.height) 21 | glutInitWindowPosition(0, 0) 22 | _glut_window = glutCreateWindow("My Render.") 23 | 24 | # glEnable(GL_DEPTH_CLAMP) 25 | glEnable(GL_DEPTH_TEST) 26 | 27 | glClampColor(GL_CLAMP_READ_COLOR, GL_FALSE) 28 | glClampColor(GL_CLAMP_FRAGMENT_COLOR, GL_FALSE) 29 | glClampColor(GL_CLAMP_VERTEX_COLOR, GL_FALSE) 30 | 31 | # init program 32 | shader_list = [] 33 | 34 | for program_file in program_files: 35 | _, ext = os.path.splitext(program_file) 36 | if ext == '.vs': 37 | shader_list.append(loadShader(GL_VERTEX_SHADER, program_file)) 38 | elif ext == '.fs': 39 | shader_list.append(loadShader(GL_FRAGMENT_SHADER, program_file)) 40 | elif ext == '.gs': 41 | shader_list.append(loadShader(GL_GEOMETRY_SHADER, program_file)) 42 | 43 | self.program = createProgram(shader_list) 44 | 45 | for shader in shader_list: 46 | glDeleteShader(shader) 47 | 48 | # Init uniform variables 49 | self.model_mat_unif = glGetUniformLocation(self.program, 'ModelMat') 50 | self.persp_mat_unif = glGetUniformLocation(self.program, 'PerspMat') 51 | 52 | self.vertex_buffer = glGenBuffers(1) 53 | 54 | # Init screen quad program and buffer 55 | self.quad_program, self.quad_buffer = self.init_quad_program() 56 | 57 | # Configure frame buffer 58 | self.frame_buffer = glGenFramebuffers(1) 59 | glBindFramebuffer(GL_FRAMEBUFFER, self.frame_buffer) 60 | 61 | self.intermediate_fbo = None 62 | if ms_rate > 1: 63 | # Configure texture buffer to render to 64 | self.color_buffer = [] 65 | for i in range(color_size): 66 | color_buffer = glGenTextures(1) 67 | multi_sample_rate = ms_rate 68 | glBindTexture(GL_TEXTURE_2D_MULTISAMPLE, color_buffer) 69 | glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE) 70 | glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE) 71 | glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR) 72 | glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR) 73 | glTexImage2DMultisample(GL_TEXTURE_2D_MULTISAMPLE, multi_sample_rate, GL_RGBA32F, self.width, self.height, GL_TRUE) 74 | glBindTexture(GL_TEXTURE_2D_MULTISAMPLE, 0) 75 | glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0 + i, GL_TEXTURE_2D_MULTISAMPLE, color_buffer, 0) 76 | self.color_buffer.append(color_buffer) 77 | 78 | self.render_buffer = glGenRenderbuffers(1) 79 | glBindRenderbuffer(GL_RENDERBUFFER, self.render_buffer) 80 | glRenderbufferStorageMultisample(GL_RENDERBUFFER, multi_sample_rate, GL_DEPTH24_STENCIL8, self.width, self.height) 81 | glBindRenderbuffer(GL_RENDERBUFFER, 0) 82 | glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_DEPTH_STENCIL_ATTACHMENT, GL_RENDERBUFFER, self.render_buffer) 83 | 84 | attachments = [] 85 | for i in range(color_size): 86 | attachments.append(GL_COLOR_ATTACHMENT0 + i) 87 | glDrawBuffers(color_size, attachments) 88 | glBindFramebuffer(GL_FRAMEBUFFER, 0) 89 | 90 | self.intermediate_fbo = glGenFramebuffers(1) 91 | glBindFramebuffer(GL_FRAMEBUFFER, self.intermediate_fbo) 92 | 93 | self.screen_texture = [] 94 | for i in range(color_size): 95 | screen_texture = glGenTextures(1) 96 | glBindTexture(GL_TEXTURE_2D, screen_texture) 97 | glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA32F, self.width, self.height, 0, GL_RGBA, GL_FLOAT, None) 98 | glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR) 99 | glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR) 100 | glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0 + i, GL_TEXTURE_2D, screen_texture, 0) 101 | self.screen_texture.append(screen_texture) 102 | 103 | glDrawBuffers(color_size, attachments) 104 | glBindFramebuffer(GL_FRAMEBUFFER, 0) 105 | else: 106 | self.color_buffer = [] 107 | for i in range(color_size): 108 | color_buffer = glGenTextures(1) 109 | glBindTexture(GL_TEXTURE_2D, color_buffer) 110 | glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE) 111 | glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE) 112 | glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST) 113 | glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST) 114 | glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA32F, self.width, self.height, 0, GL_RGBA, GL_FLOAT, None) 115 | glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0 + i, GL_TEXTURE_2D, color_buffer, 0) 116 | self.color_buffer.append(color_buffer) 117 | 118 | # Configure depth texture map to render to 119 | self.depth_buffer = glGenTextures(1) 120 | glBindTexture(GL_TEXTURE_2D, self.depth_buffer) 121 | glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT) 122 | glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT) 123 | glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST) 124 | glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST) 125 | glTexParameteri(GL_TEXTURE_2D, GL_DEPTH_TEXTURE_MODE, GL_INTENSITY) 126 | glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_COMPARE_MODE, GL_COMPARE_R_TO_TEXTURE) 127 | glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_COMPARE_FUNC, GL_LEQUAL) 128 | glTexImage2D(GL_TEXTURE_2D, 0, GL_DEPTH_COMPONENT, self.width, self.height, 0, GL_DEPTH_COMPONENT, GL_FLOAT, None) 129 | glFramebufferTexture2D(GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, GL_TEXTURE_2D, self.depth_buffer, 0) 130 | 131 | attachments = [] 132 | for i in range(color_size): 133 | attachments.append(GL_COLOR_ATTACHMENT0 + i) 134 | glDrawBuffers(color_size, attachments) 135 | self.screen_texture = self.color_buffer 136 | 137 | glBindFramebuffer(GL_FRAMEBUFFER, 0) 138 | 139 | 140 | # Configure texture buffer if needed 141 | self.render_texture = None 142 | 143 | # NOTE: original render_texture only support one input 144 | # this is tentative member of this issue 145 | self.render_texture_v2 = {} 146 | 147 | # Inner storage for buffer data 148 | self.vertex_data = None 149 | self.vertex_dim = None 150 | self.n_vertices = None 151 | 152 | self.model_view_matrix = None 153 | self.projection_matrix = None 154 | 155 | glutDisplayFunc(self.display) 156 | 157 | 158 | def init_quad_program(self): 159 | shader_list = [] 160 | 161 | shader_list.append(loadShader(GL_VERTEX_SHADER, "quad.vs")) 162 | shader_list.append(loadShader(GL_FRAGMENT_SHADER, "quad.fs")) 163 | 164 | the_program = createProgram(shader_list) 165 | 166 | for shader in shader_list: 167 | glDeleteShader(shader) 168 | 169 | # vertex attributes for a quad that fills the entire screen in Normalized Device Coordinates. 170 | # positions # texCoords 171 | quad_vertices = np.array( 172 | [-1.0, 1.0, 0.0, 1.0, 173 | -1.0, -1.0, 0.0, 0.0, 174 | 1.0, -1.0, 1.0, 0.0, 175 | 176 | -1.0, 1.0, 0.0, 1.0, 177 | 1.0, -1.0, 1.0, 0.0, 178 | 1.0, 1.0, 1.0, 1.0] 179 | ) 180 | 181 | quad_buffer = glGenBuffers(1) 182 | glBindBuffer(GL_ARRAY_BUFFER, quad_buffer) 183 | glBufferData(GL_ARRAY_BUFFER, quad_vertices, GL_STATIC_DRAW) 184 | 185 | glBindBuffer(GL_ARRAY_BUFFER, 0) 186 | 187 | return the_program, quad_buffer 188 | 189 | def set_mesh(self, vertices, faces): 190 | self.vertex_data = vertices[faces.reshape([-1])] 191 | self.vertex_dim = self.vertex_data.shape[1] 192 | self.n_vertices = self.vertex_data.shape[0] 193 | 194 | glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer) 195 | glBufferData(GL_ARRAY_BUFFER, self.vertex_data, GL_STATIC_DRAW) 196 | 197 | glBindBuffer(GL_ARRAY_BUFFER, 0) 198 | 199 | def set_viewpoint(self, projection, model_view): 200 | self.projection_matrix = projection 201 | self.model_view_matrix = model_view 202 | 203 | def draw_init(self): 204 | glBindFramebuffer(GL_FRAMEBUFFER, self.frame_buffer) 205 | glEnable(GL_DEPTH_TEST) 206 | 207 | glClearColor(0.0, 0.0, 0.0, 0.0) 208 | if self.use_inverse_depth: 209 | glDepthFunc(GL_GREATER) 210 | glClearDepth(0.0) 211 | else: 212 | glDepthFunc(GL_LESS) 213 | glClearDepth(1.0) 214 | glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT) 215 | 216 | def draw_end(self): 217 | if self.intermediate_fbo is not None: 218 | for i in range(len(self.color_buffer)): 219 | glBindFramebuffer(GL_READ_FRAMEBUFFER, self.frame_buffer) 220 | glReadBuffer(GL_COLOR_ATTACHMENT0 + i) 221 | glBindFramebuffer(GL_DRAW_FRAMEBUFFER, self.intermediate_fbo) 222 | glDrawBuffer(GL_COLOR_ATTACHMENT0 + i) 223 | glBlitFramebuffer(0, 0, self.width, self.height, 0, 0, self.width, self.height, GL_COLOR_BUFFER_BIT, GL_NEAREST) 224 | 225 | glBindFramebuffer(GL_FRAMEBUFFER, 0) 226 | glDepthFunc(GL_LESS) 227 | glClearDepth(1.0) 228 | 229 | def draw(self): 230 | self.draw_init() 231 | 232 | glUseProgram(self.program) 233 | glUniformMatrix4fv(self.model_mat_unif, 1, GL_FALSE, self.model_view_matrix.transpose()) 234 | glUniformMatrix4fv(self.persp_mat_unif, 1, GL_FALSE, self.projection_matrix.transpose()) 235 | 236 | glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer) 237 | 238 | glEnableVertexAttribArray(0) 239 | glVertexAttribPointer(0, self.vertex_dim, GL_DOUBLE, GL_FALSE, 0, None) 240 | 241 | glDrawArrays(GL_TRIANGLES, 0, self.n_vertices) 242 | 243 | glDisableVertexAttribArray(0) 244 | 245 | glBindBuffer(GL_ARRAY_BUFFER, 0) 246 | 247 | glUseProgram(0) 248 | 249 | self.draw_end() 250 | 251 | def get_color(self, color_id=0): 252 | glBindFramebuffer(GL_FRAMEBUFFER, self.intermediate_fbo if self.intermediate_fbo is not None else self.frame_buffer) 253 | glReadBuffer(GL_COLOR_ATTACHMENT0 + color_id) 254 | data = glReadPixels(0, 0, self.width, self.height, GL_RGBA, GL_FLOAT, outputType=None) 255 | glBindFramebuffer(GL_FRAMEBUFFER, 0) 256 | rgb = data.reshape(self.height, self.width, -1) 257 | rgb = np.flip(rgb, 0) 258 | return rgb 259 | 260 | def get_z_value(self): 261 | glBindFramebuffer(GL_FRAMEBUFFER, self.frame_buffer) 262 | data = glReadPixels(0, 0, self.width, self.height, GL_DEPTH_COMPONENT, GL_FLOAT, outputType=None) 263 | glBindFramebuffer(GL_FRAMEBUFFER, 0) 264 | z = data.reshape(self.height, self.width) 265 | z = np.flip(z, 0) 266 | return z 267 | 268 | def display(self): 269 | # First we draw a scene. 270 | # Notice the result is stored in the texture buffer. 271 | self.draw() 272 | 273 | # Then we return to the default frame buffer since we will display on the screen. 274 | glBindFramebuffer(GL_FRAMEBUFFER, 0) 275 | 276 | # Do the clean-up. 277 | glClearColor(0.0, 0.0, 0.0, 0.0) 278 | glClear(GL_COLOR_BUFFER_BIT) 279 | 280 | # We draw a rectangle which covers the whole screen. 281 | glUseProgram(self.quad_program) 282 | glBindBuffer(GL_ARRAY_BUFFER, self.quad_buffer) 283 | 284 | size_of_double = 8 285 | glEnableVertexAttribArray(0) 286 | glVertexAttribPointer(0, 2, GL_DOUBLE, GL_FALSE, 4 * size_of_double, None) 287 | glEnableVertexAttribArray(1) 288 | glVertexAttribPointer(1, 2, GL_DOUBLE, GL_FALSE, 4 * size_of_double, c_void_p(2 * size_of_double)) 289 | 290 | glDisable(GL_DEPTH_TEST) 291 | 292 | # The stored texture is then mapped to this rectangle. 293 | # properly assing color buffer texture 294 | glActiveTexture(GL_TEXTURE0) 295 | glBindTexture(GL_TEXTURE_2D, self.screen_texture[0]) 296 | glUniform1i(glGetUniformLocation(self.quad_program, 'screenTexture'), 0) 297 | 298 | glDrawArrays(GL_TRIANGLES, 0, 6) 299 | 300 | glDisableVertexAttribArray(1) 301 | glDisableVertexAttribArray(0) 302 | 303 | glEnable(GL_DEPTH_TEST) 304 | glBindBuffer(GL_ARRAY_BUFFER, 0) 305 | glUseProgram(0) 306 | 307 | glutSwapBuffers() 308 | glutPostRedisplay() 309 | 310 | def show(self): 311 | glutMainLoop() 312 | -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/glm.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def vec3(x, y, z): 5 | return np.array([x, y, z], dtype=np.float32) 6 | 7 | 8 | def radians(v): 9 | return np.radians(v) 10 | 11 | 12 | def identity(): 13 | return np.identity(4, dtype=np.float32) 14 | 15 | 16 | def empty(): 17 | return np.zeros([4, 4], dtype=np.float32) 18 | 19 | 20 | def magnitude(v): 21 | return np.linalg.norm(v) 22 | 23 | 24 | def normalize(v): 25 | m = magnitude(v) 26 | return v if m == 0 else v / m 27 | 28 | 29 | def dot(u, v): 30 | return np.sum(u * v) 31 | 32 | 33 | def cross(u, v): 34 | res = vec3(0, 0, 0) 35 | res[0] = u[1] * v[2] - u[2] * v[1] 36 | res[1] = u[2] * v[0] - u[0] * v[2] 37 | res[2] = u[0] * v[1] - u[1] * v[0] 38 | return res 39 | 40 | 41 | # below functions can be optimized 42 | 43 | def translate(m, v): 44 | res = np.copy(m) 45 | res[:, 3] = m[:, 0] * v[0] + m[:, 1] * v[1] + m[:, 2] * v[2] + m[:, 3] 46 | return res 47 | 48 | 49 | def rotate(m, angle, v): 50 | a = angle 51 | c = np.cos(a) 52 | s = np.sin(a) 53 | 54 | axis = normalize(v) 55 | temp = (1 - c) * axis 56 | 57 | rot = empty() 58 | rot[0][0] = c + temp[0] * axis[0] 59 | rot[0][1] = temp[0] * axis[1] + s * axis[2] 60 | rot[0][2] = temp[0] * axis[2] - s * axis[1] 61 | 62 | rot[1][0] = temp[1] * axis[0] - s * axis[2] 63 | rot[1][1] = c + temp[1] * axis[1] 64 | rot[1][2] = temp[1] * axis[2] + s * axis[0] 65 | 66 | rot[2][0] = temp[2] * axis[0] + s * axis[1] 67 | rot[2][1] = temp[2] * axis[1] - s * axis[0] 68 | rot[2][2] = c + temp[2] * axis[2] 69 | 70 | res = empty() 71 | res[:, 0] = m[:, 0] * rot[0][0] + m[:, 1] * rot[0][1] + m[:, 2] * rot[0][2] 72 | res[:, 1] = m[:, 0] * rot[1][0] + m[:, 1] * rot[1][1] + m[:, 2] * rot[1][2] 73 | res[:, 2] = m[:, 0] * rot[2][0] + m[:, 1] * rot[2][1] + m[:, 2] * rot[2][2] 74 | res[:, 3] = m[:, 3] 75 | return res 76 | 77 | 78 | def perspective(fovy, aspect, zNear, zFar): 79 | tanHalfFovy = np.tan(fovy / 2) 80 | 81 | res = empty() 82 | res[0][0] = 1 / (aspect * tanHalfFovy) 83 | res[1][1] = 1 / (tanHalfFovy) 84 | res[2][3] = -1 85 | res[2][2] = - (zFar + zNear) / (zFar - zNear) 86 | res[3][2] = -(2 * zFar * zNear) / (zFar - zNear) 87 | 88 | return res.T 89 | 90 | 91 | def ortho(left, right, bottom, top, zNear, zFar): 92 | # res = np.ones([4, 4], dtype=np.float32) 93 | res = identity() 94 | res[0][0] = 2 / (right - left) 95 | res[1][1] = 2 / (top - bottom) 96 | res[2][2] = - 2 / (zFar - zNear) 97 | res[3][0] = - (right + left) / (right - left) 98 | res[3][1] = - (top + bottom) / (top - bottom) 99 | res[3][2] = - (zFar + zNear) / (zFar - zNear) 100 | return res.T 101 | 102 | 103 | def lookat(eye, center, up): 104 | f = normalize(center - eye) 105 | s = normalize(cross(f, up)) 106 | u = cross(s, f) 107 | 108 | res = identity() 109 | res[0][0] = s[0] 110 | res[1][0] = s[1] 111 | res[2][0] = s[2] 112 | res[0][1] = u[0] 113 | res[1][1] = u[1] 114 | res[2][1] = u[2] 115 | res[0][2] = -f[0] 116 | res[1][2] = -f[1] 117 | res[2][2] = -f[2] 118 | res[3][0] = -dot(s, eye) 119 | res[3][1] = -dot(u, eye) 120 | res[3][2] = -dot(f, eye) 121 | return res.T 122 | 123 | 124 | def transform(d, m): 125 | return np.dot(m, d.T).T 126 | -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/mesh.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def save_obj_mesh(mesh_path, verts, faces=None, color=None): 5 | file = open(mesh_path, 'w') 6 | for i, v in enumerate(verts): 7 | if color is None: 8 | file.write('v %.4f %.4f %.4f\n' % (v[0], v[1], v[2])) 9 | else: 10 | file.write('v %.4f %.4f %.4f %.4f %.4f %.4f\n' % (v[0], v[1], v[2], color[i][0], color[i][1], color[i][2])) 11 | if faces is not None: 12 | for f in faces: 13 | f_plus = f + 1 14 | file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2])) 15 | file.close() 16 | 17 | # https://github.com/ratcave/wavefront_reader 18 | def read_mtlfile(fname): 19 | materials = {} 20 | with open(fname) as f: 21 | lines = f.read().splitlines() 22 | 23 | for line in lines: 24 | if line: 25 | split_line = line.strip().split(' ', 1) 26 | if len(split_line) < 2: 27 | continue 28 | 29 | prefix, data = split_line[0], split_line[1] 30 | if 'newmtl' in prefix: 31 | material = {} 32 | materials[data] = material 33 | elif materials: 34 | if data: 35 | split_data = data.strip().split(' ') 36 | 37 | # assume texture maps are in the same level 38 | # WARNING: do not include space in your filename!! 39 | if 'map' in prefix: 40 | material[prefix] = split_data[-1].split('\\')[-1] 41 | elif len(split_data) > 1: 42 | material[prefix] = tuple(float(d) for d in split_data) 43 | else: 44 | try: 45 | material[prefix] = int(data) 46 | except ValueError: 47 | material[prefix] = float(data) 48 | 49 | return materials 50 | 51 | 52 | def load_obj_mesh_mtl(mesh_file): 53 | vertex_data = [] 54 | norm_data = [] 55 | uv_data = [] 56 | 57 | face_data = [] 58 | face_norm_data = [] 59 | face_uv_data = [] 60 | 61 | # face per material 62 | face_data_mat = {} 63 | face_norm_data_mat = {} 64 | face_uv_data_mat = {} 65 | 66 | # current material name 67 | mtl_data = None 68 | cur_mat = None 69 | 70 | if isinstance(mesh_file, str): 71 | f = open(mesh_file, "r") 72 | else: 73 | f = mesh_file 74 | for line in f: 75 | if isinstance(line, bytes): 76 | line = line.decode("utf-8") 77 | if line.startswith('#'): 78 | continue 79 | values = line.split() 80 | if not values: 81 | continue 82 | 83 | if values[0] == 'v': 84 | v = list(map(float, values[1:4])) 85 | vertex_data.append(v) 86 | elif values[0] == 'vn': 87 | vn = list(map(float, values[1:4])) 88 | norm_data.append(vn) 89 | elif values[0] == 'vt': 90 | vt = list(map(float, values[1:3])) 91 | uv_data.append(vt) 92 | elif values[0] == 'mtllib': 93 | mtl_data = read_mtlfile(mesh_file.replace(mesh_file.split('/')[-1],values[1])) 94 | elif values[0] == 'usemtl': 95 | cur_mat = values[1] 96 | elif values[0] == 'f': 97 | # local triangle data 98 | l_face_data = [] 99 | l_face_uv_data = [] 100 | l_face_norm_data = [] 101 | 102 | # quad mesh 103 | if len(values) > 4: 104 | f = list(map(lambda x: int(x.split('/')[0]) if int(x.split('/')[0]) < 0 else int(x.split('/')[0])-1, values[1:4])) 105 | l_face_data.append(f) 106 | f = list(map(lambda x: int(x.split('/')[0]) if int(x.split('/')[0]) < 0 else int(x.split('/')[0])-1, [values[3], values[4], values[1]])) 107 | l_face_data.append(f) 108 | # tri mesh 109 | else: 110 | f = list(map(lambda x: int(x.split('/')[0]) if int(x.split('/')[0]) < 0 else int(x.split('/')[0])-1, values[1:4])) 111 | l_face_data.append(f) 112 | # deal with texture 113 | if len(values[1].split('/')) >= 2: 114 | # quad mesh 115 | if len(values) > 4: 116 | f = list(map(lambda x: int(x.split('/')[1]) if int(x.split('/')[1]) < 0 else int(x.split('/')[1])-1, values[1:4])) 117 | l_face_uv_data.append(f) 118 | f = list(map(lambda x: int(x.split('/')[1]) if int(x.split('/')[1]) < 0 else int(x.split('/')[1])-1, [values[3], values[4], values[1]])) 119 | l_face_uv_data.append(f) 120 | # tri mesh 121 | elif len(values[1].split('/')[1]) != 0: 122 | f = list(map(lambda x: int(x.split('/')[1]) if int(x.split('/')[1]) < 0 else int(x.split('/')[1])-1, values[1:4])) 123 | l_face_uv_data.append(f) 124 | # deal with normal 125 | if len(values[1].split('/')) == 3: 126 | # quad mesh 127 | if len(values) > 4: 128 | f = list(map(lambda x: int(x.split('/')[2]) if int(x.split('/')[2]) < 0 else int(x.split('/')[2])-1, values[1:4])) 129 | l_face_norm_data.append(f) 130 | f = list(map(lambda x: int(x.split('/')[2]) if int(x.split('/')[2]) < 0 else int(x.split('/')[2])-1, [values[3], values[4], values[1]])) 131 | l_face_norm_data.append(f) 132 | # tri mesh 133 | elif len(values[1].split('/')[2]) != 0: 134 | f = list(map(lambda x: int(x.split('/')[2]) if int(x.split('/')[2]) < 0 else int(x.split('/')[2])-1, values[1:4])) 135 | l_face_norm_data.append(f) 136 | 137 | face_data += l_face_data 138 | face_uv_data += l_face_uv_data 139 | face_norm_data += l_face_norm_data 140 | 141 | if cur_mat is not None: 142 | if cur_mat not in face_data_mat.keys(): 143 | face_data_mat[cur_mat] = [] 144 | if cur_mat not in face_uv_data_mat.keys(): 145 | face_uv_data_mat[cur_mat] = [] 146 | if cur_mat not in face_norm_data_mat.keys(): 147 | face_norm_data_mat[cur_mat] = [] 148 | face_data_mat[cur_mat] += l_face_data 149 | face_uv_data_mat[cur_mat] += l_face_uv_data 150 | face_norm_data_mat[cur_mat] += l_face_norm_data 151 | 152 | vertices = np.array(vertex_data) 153 | faces = np.array(face_data) 154 | 155 | norms = np.array(norm_data) 156 | norms = normalize_v3(norms) 157 | face_normals = np.array(face_norm_data) 158 | 159 | uvs = np.array(uv_data) 160 | face_uvs = np.array(face_uv_data) 161 | 162 | out_tuple = (vertices, faces, norms, face_normals, uvs, face_uvs) 163 | 164 | if cur_mat is not None and mtl_data is not None: 165 | for key in face_data_mat: 166 | face_data_mat[key] = np.array(face_data_mat[key]) 167 | face_uv_data_mat[key] = np.array(face_uv_data_mat[key]) 168 | face_norm_data_mat[key] = np.array(face_norm_data_mat[key]) 169 | 170 | out_tuple += (face_data_mat, face_norm_data_mat, face_uv_data_mat, mtl_data) 171 | 172 | return out_tuple 173 | 174 | 175 | def load_obj_mesh(mesh_file, with_normal=False, with_texture=False): 176 | vertex_data = [] 177 | norm_data = [] 178 | uv_data = [] 179 | 180 | face_data = [] 181 | face_norm_data = [] 182 | face_uv_data = [] 183 | 184 | if isinstance(mesh_file, str): 185 | f = open(mesh_file, "r") 186 | else: 187 | f = mesh_file 188 | for line in f: 189 | if isinstance(line, bytes): 190 | line = line.decode("utf-8") 191 | if line.startswith('#'): 192 | continue 193 | values = line.split() 194 | if not values: 195 | continue 196 | 197 | if values[0] == 'v': 198 | v = list(map(float, values[1:4])) 199 | vertex_data.append(v) 200 | elif values[0] == 'vn': 201 | vn = list(map(float, values[1:4])) 202 | norm_data.append(vn) 203 | elif values[0] == 'vt': 204 | vt = list(map(float, values[1:3])) 205 | uv_data.append(vt) 206 | 207 | elif values[0] == 'f': 208 | # quad mesh 209 | if len(values) > 4: 210 | f = list(map(lambda x: int(x.split('/')[0]), values[1:4])) 211 | face_data.append(f) 212 | f = list(map(lambda x: int(x.split('/')[0]), [values[3], values[4], values[1]])) 213 | face_data.append(f) 214 | # tri mesh 215 | else: 216 | f = list(map(lambda x: int(x.split('/')[0]), values[1:4])) 217 | face_data.append(f) 218 | 219 | # deal with texture 220 | if len(values[1].split('/')) >= 2: 221 | # quad mesh 222 | if len(values) > 4: 223 | f = list(map(lambda x: int(x.split('/')[1]), values[1:4])) 224 | face_uv_data.append(f) 225 | f = list(map(lambda x: int(x.split('/')[1]), [values[3], values[4], values[1]])) 226 | face_uv_data.append(f) 227 | # tri mesh 228 | elif len(values[1].split('/')[1]) != 0: 229 | f = list(map(lambda x: int(x.split('/')[1]), values[1:4])) 230 | face_uv_data.append(f) 231 | # deal with normal 232 | if len(values[1].split('/')) == 3: 233 | # quad mesh 234 | if len(values) > 4: 235 | f = list(map(lambda x: int(x.split('/')[2]), values[1:4])) 236 | face_norm_data.append(f) 237 | f = list(map(lambda x: int(x.split('/')[2]), [values[3], values[4], values[1]])) 238 | face_norm_data.append(f) 239 | # tri mesh 240 | elif len(values[1].split('/')[2]) != 0: 241 | f = list(map(lambda x: int(x.split('/')[2]), values[1:4])) 242 | face_norm_data.append(f) 243 | 244 | vertices = np.array(vertex_data) 245 | faces = np.array(face_data) - 1 246 | 247 | if with_texture and with_normal: 248 | uvs = np.array(uv_data) 249 | face_uvs = np.array(face_uv_data) - 1 250 | norms = np.array(norm_data) 251 | if norms.shape[0] == 0: 252 | norms = compute_normal(vertices, faces) 253 | face_normals = faces 254 | else: 255 | norms = normalize_v3(norms) 256 | face_normals = np.array(face_norm_data) - 1 257 | return vertices, faces, norms, face_normals, uvs, face_uvs 258 | 259 | if with_texture: 260 | uvs = np.array(uv_data) 261 | face_uvs = np.array(face_uv_data) - 1 262 | return vertices, faces, uvs, face_uvs 263 | 264 | if with_normal: 265 | norms = np.array(norm_data) 266 | norms = normalize_v3(norms) 267 | face_normals = np.array(face_norm_data) - 1 268 | return vertices, faces, norms, face_normals 269 | 270 | return vertices, faces 271 | 272 | 273 | def normalize_v3(arr): 274 | ''' Normalize a numpy array of 3 component vectors shape=(n,3) ''' 275 | lens = np.sqrt(arr[:, 0] ** 2 + arr[:, 1] ** 2 + arr[:, 2] ** 2) 276 | eps = 0.00000001 277 | lens[lens < eps] = eps 278 | arr[:, 0] /= lens 279 | arr[:, 1] /= lens 280 | arr[:, 2] /= lens 281 | return arr 282 | 283 | 284 | def compute_normal(vertices, faces): 285 | # Create a zeroed array with the same type and shape as our vertices i.e., per vertex normal 286 | norm = np.zeros(vertices.shape, dtype=vertices.dtype) 287 | # Create an indexed view into the vertex array using the array of three indices for triangles 288 | tris = vertices[faces] 289 | # Calculate the normal for all the triangles, by taking the cross product of the vectors v1-v0, and v2-v0 in each triangle 290 | n = np.cross(tris[::, 1] - tris[::, 0], tris[::, 2] - tris[::, 0]) 291 | # n is now an array of normals per triangle. The length of each normal is dependent the vertices, 292 | # we need to normalize these, so that our next step weights each normal equally. 293 | normalize_v3(n) 294 | # now we have a normalized array of normals, one per triangle, i.e., per triangle normals. 295 | # But instead of one per triangle (i.e., flat shading), we add to each vertex in that triangle, 296 | # the triangles' normal. Multiple triangles would then contribute to every vertex, so we need to normalize again afterwards. 297 | # The cool part, we can actually add the normals through an indexed view of our (zeroed) per vertex normal array 298 | norm[faces[:, 0]] += n 299 | norm[faces[:, 1]] += n 300 | norm[faces[:, 2]] += n 301 | normalize_v3(norm) 302 | 303 | return norm 304 | 305 | # compute tangent and bitangent 306 | def compute_tangent(vertices, faces, normals, uvs, faceuvs): 307 | # NOTE: this could be numerically unstable around [0,0,1] 308 | # but other current solutions are pretty freaky somehow 309 | c1 = np.cross(normals, np.array([0,1,0.0])) 310 | tan = c1 311 | normalize_v3(tan) 312 | btan = np.cross(normals, tan) 313 | 314 | # NOTE: traditional version is below 315 | 316 | # pts_tris = vertices[faces] 317 | # uv_tris = uvs[faceuvs] 318 | 319 | # W = np.stack([pts_tris[::, 1] - pts_tris[::, 0], pts_tris[::, 2] - pts_tris[::, 0]],2) 320 | # UV = np.stack([uv_tris[::, 1] - uv_tris[::, 0], uv_tris[::, 2] - uv_tris[::, 0]], 1) 321 | 322 | # for i in range(W.shape[0]): 323 | # W[i,::] = W[i,::].dot(np.linalg.inv(UV[i,::])) 324 | 325 | # tan = np.zeros(vertices.shape, dtype=vertices.dtype) 326 | # tan[faces[:,0]] += W[:,:,0] 327 | # tan[faces[:,1]] += W[:,:,0] 328 | # tan[faces[:,2]] += W[:,:,0] 329 | 330 | # btan = np.zeros(vertices.shape, dtype=vertices.dtype) 331 | # btan[faces[:,0]] += W[:,:,1] 332 | # btan[faces[:,1]] += W[:,:,1] 333 | # btan[faces[:,2]] += W[:,:,1] 334 | 335 | # normalize_v3(tan) 336 | 337 | # ndott = np.sum(normals*tan, 1, keepdims=True) 338 | # tan = tan - ndott * normals 339 | 340 | # normalize_v3(btan) 341 | # normalize_v3(tan) 342 | 343 | # tan[np.sum(np.cross(normals, tan) * btan, 1) < 0,:] *= -1.0 344 | 345 | return tan, btan 346 | 347 | if __name__ == '__main__': 348 | pts, tri, nml, trin, uvs, triuv = load_obj_mesh('/home/ICT2000/ssaito/Documents/Body/tmp/Baseball_Pitching/0012.obj', True, True) 349 | compute_tangent(pts, tri, uvs, triuv) -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/ram_zip.py: -------------------------------------------------------------------------------- 1 | import zipfile 2 | import io 3 | import os 4 | 5 | 6 | class ZipPool(object): 7 | def __init__(self, append=False): 8 | self.files = dict() 9 | self._append_on_disk = append 10 | 11 | def write_str(self, prefix, path, data): 12 | zip_file = prefix + '.zip' 13 | post_fix = os.path.relpath(path, prefix) 14 | if zip_file not in self.files: 15 | if self._append_on_disk: 16 | self.files[zip_file] = OnDiskZip(zip_file) 17 | else: 18 | self.files[zip_file] = OnDiskWZip(zip_file) 19 | zp = self.files[zip_file] 20 | zp.append_str(post_fix, data) 21 | 22 | def write(self, prefix, path, file_path): 23 | zip_file = prefix + '.zip' 24 | post_fix = os.path.relpath(path, prefix) 25 | if zip_file not in self.files: 26 | if self._append_on_disk: 27 | self.files[zip_file] = OnDiskZip(zip_file) 28 | else: 29 | self.files[zip_file] = OnDiskWZip(zip_file) 30 | zp = self.files[zip_file] 31 | zp.append(post_fix, file_path) 32 | 33 | def flush(self): 34 | for zip_file in self.files: 35 | self.files[zip_file].writetofile() 36 | 37 | 38 | class UnzipPool(object): 39 | def __init__(self, append=False): 40 | self.files = dict() 41 | self._append_on_disk = append 42 | 43 | def open(self, prefix, path): 44 | zip_file = prefix + '.zip' 45 | post_fix = os.path.relpath(path, prefix) 46 | post_fix = post_fix.replace('\\', '/') 47 | if zip_file not in self.files: 48 | self.files[zip_file] = UnZip(zip_file) 49 | zp = self.files[zip_file] 50 | return zp.open(post_fix) 51 | 52 | def read(self, prefix, path): 53 | zip_file = prefix + '.zip' 54 | post_fix = os.path.relpath(path, prefix) 55 | post_fix = post_fix.replace('\\', '/') 56 | if zip_file not in self.files: 57 | self.files[zip_file] = UnZip(zip_file) 58 | zp = self.files[zip_file] 59 | return zp.read(post_fix) 60 | 61 | def listdir(self, prefix, path): 62 | zip_file = prefix + '.zip' 63 | post_fix = os.path.relpath(path, prefix) 64 | post_fix = post_fix.replace('\\', '/') 65 | if zip_file not in self.files: 66 | self.files[zip_file] = UnZip(zip_file) 67 | zp = self.files[zip_file] 68 | file_names = zp.zf.namelist() 69 | 70 | result = set() 71 | for file_name in file_names: 72 | f = os.path.relpath(file_name, post_fix).replace('\\', '/').split('/')[0] 73 | if not f.startswith('.'): 74 | result.add(f) 75 | return result 76 | 77 | def close(self): 78 | for zip in self.files: 79 | self.files[zip].close() 80 | 81 | 82 | class UnZip(object): 83 | def __init__(self, filename): 84 | # Create the in-memory file-like object 85 | self.zf = zipfile.ZipFile(filename, "r") 86 | 87 | def open(self, filename_in_zip): 88 | return self.zf.open(filename_in_zip) 89 | 90 | def read(self, filename_in_zip): 91 | return self.zf.read(filename_in_zip) 92 | 93 | def close(self): 94 | return self.zf.close() 95 | 96 | class InMemoryZip(object): 97 | def __init__(self, filename): 98 | # Create the in-memory file-like object 99 | self.in_memory_zip = io.BytesIO() 100 | self.filename = filename 101 | self.zf = zipfile.ZipFile(self.in_memory_zip, "w", zipfile.ZIP_DEFLATED, True) 102 | 103 | def append_str(self, filename_in_zip, file_contents): 104 | '''Appends a file with name filename_in_zip and contents of 105 | file_contents to the in-memory zip.''' 106 | # Get a handle to the in-memory zip in append mode 107 | zf = self.zf 108 | 109 | # Write the file to the in-memory zip 110 | zf.writestr(filename_in_zip, file_contents) 111 | 112 | # Mark the files as having been created on Windows so that 113 | # Unix permissions are not inferred as 0000 114 | for zfile in zf.filelist: 115 | zfile.create_system = 0 116 | 117 | return self 118 | 119 | def read(self): 120 | '''Returns a string with the contents of the in-memory zip.''' 121 | self.in_memory_zip.seek(0) 122 | return self.in_memory_zip.read() 123 | 124 | def writetofile(self): 125 | '''Writes the in-memory zip to a file.''' 126 | self.zf.close() 127 | f = open(self.filename, "wb") 128 | f.write(self.read()) 129 | f.close() 130 | 131 | 132 | class OnDiskWZip(object): 133 | def __init__(self, filename): 134 | # Create the in-memory file-like object 135 | self.zf = zipfile.ZipFile(filename, "w", zipfile.ZIP_DEFLATED, True) 136 | 137 | def append_str(self, filename_in_zip, file_contents): 138 | '''Appends a file with name filename_in_zip and contents of 139 | file_contents to the in-memory zip.''' 140 | # Get a handle to the in-memory zip in append mode 141 | zf = self.zf 142 | 143 | # Write the file to the in-memory zip 144 | zf.writestr(filename_in_zip, file_contents) 145 | 146 | # Mark the files as having been created on Windows so that 147 | # Unix permissions are not inferred as 0000 148 | for zfile in zf.filelist: 149 | zfile.create_system = 0 150 | 151 | return self 152 | 153 | def append(self, filename_in_zip, file_path): 154 | zf = self.zf 155 | zf.write(file_path, filename_in_zip) 156 | for zfile in zf.filelist: 157 | zfile.create_system = 0 158 | 159 | return self 160 | 161 | 162 | def writetofile(self): 163 | '''Writes the in-memory zip to a file.''' 164 | self.zf.close() 165 | 166 | 167 | class OnDiskZip(object): 168 | def __init__(self, filename): 169 | # Create the in-memory file-like object 170 | self.filename = filename 171 | 172 | def append(self, filename_in_zip, file_contents): 173 | '''Appends a file with name filename_in_zip and contents of 174 | file_contents to the in-memory zip.''' 175 | # Get a handle to the in-memory zip in append mode 176 | zf = zipfile.ZipFile(self.filename, "a", zipfile.ZIP_DEFLATED, True) 177 | 178 | # Write the file to the in-memory zip 179 | zf.writestr(filename_in_zip, file_contents) 180 | 181 | # Mark the files as having been created on Windows so that 182 | # Unix permissions are not inferred as 0000 183 | for zfile in zf.filelist: 184 | zfile.create_system = 0 185 | 186 | return self 187 | 188 | def writetofile(self): 189 | None 190 | -------------------------------------------------------------------------------- /lib_data/posmap_generator/lib/renderer/tsdf.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | ''' 4 | *** Signed distance field file is binary. Format: 5 | - resolutionX,resolutionY,resolutionZ (three signed 4-byte integers (all equal), 12 bytes total) 6 | - bminx,bminy,bminz (coordinates of the lower-left-front corner of the bounding box: (three double precision 8-byte real numbers , 24 bytes total) 7 | - bmaxx,bmaxy,bmaxz (coordinates of the upper-right-back corner of the bounding box: (three double precision 8-byte real numbers , 24 bytes total) 8 | - distance data (in single precision; data alignment: 9 | [0,0,0],...,[resolutionX,0,0], 10 | [0,1,0],...,[resolutionX,resolutionY,0], 11 | [0,0,1],...,[resolutionX,resolutionY,resolutionZ]; 12 | total num bytes: sizeof(float)*(resolutionX+1)*(resolutionY+1)*(resolutionZ+1)) 13 | - closest point for each grid vertex(3 coordinates in single precision) 14 | ''' 15 | 16 | 17 | def create_sdf(b_min, b_max, resX, resY, resZ): 18 | coords = np.mgrid[:resX, :resY, :resZ] 19 | IND = np.eye(4) 20 | length = b_max - b_min 21 | IND[0, 0] = length[0] / resX 22 | IND[1, 1] = length[1] / resY 23 | IND[2, 2] = length[2] / resZ 24 | IND[0:3, 3] = b_min 25 | return coords, IND 26 | 27 | 28 | def load_sdf(file_path, read_closest_points=False, verbose=False): 29 | ''' 30 | :param file_path: file path 31 | :param read_closest_points: whether to read closest points for each grid vertex 32 | :param verbose: verbose flag 33 | :return: 34 | b_min: coordinates of the lower-left-front corner of the bounding box 35 | b_max: coordinates of the upper-right-back corner of the bounding box 36 | volume: distance data in shape (resolutionX+1)*(resolutionY+1)*(resolutionZ+1) 37 | closest_points: closest points in shape (resolutionX+1)*(resolutionY+1)*(resolutionZ+1) 38 | ''' 39 | with open(file_path, 'rb') as fp: 40 | 41 | res_x = int(np.fromfile(fp, dtype=np.int32, 42 | count=1)) # note: the dimension of volume is (1+res_x) x (1+res_y) x (1+res_z) 43 | res_x = - res_x 44 | res_y = -int(np.fromfile(fp, dtype=np.int32, count=1)) 45 | res_z = int(np.fromfile(fp, dtype=np.int32, count=1)) 46 | if verbose: print("resolution: %d %d %d" % (res_x, res_y, res_z)) 47 | 48 | b_min = np.zeros(3, dtype=np.float64) 49 | b_min[0] = np.fromfile(fp, dtype=np.float64, count=1) 50 | b_min[1] = np.fromfile(fp, dtype=np.float64, count=1) 51 | b_min[2] = np.fromfile(fp, dtype=np.float64, count=1) 52 | if verbose: print("b_min: %f %f %f" % (b_min[0], b_min[1], b_min[2])) 53 | 54 | b_max = np.zeros(3, dtype=np.float64) 55 | b_max[0] = np.fromfile(fp, dtype=np.float64, count=1) 56 | b_max[1] = np.fromfile(fp, dtype=np.float64, count=1) 57 | b_max[2] = np.fromfile(fp, dtype=np.float64, count=1) 58 | if verbose: print("b_max: %f %f %f" % (b_max[0], b_max[1], b_max[2])) 59 | 60 | grid_num = (1 + res_x) * (1 + res_y) * (1 + res_z) 61 | volume = np.fromfile(fp, dtype=np.float32, count=grid_num) 62 | volume = volume.reshape(((1 + res_z), (1 + res_y), (1 + res_x))) 63 | volume = np.swapaxes(volume, 0, 2) 64 | if verbose: print("loaded volume from %s" % file_path) 65 | 66 | closest_points = None 67 | if read_closest_points: 68 | closest_points = np.fromfile(fp, dtype=np.float32, count=grid_num * 3) 69 | closest_points = closest_points.reshape(((1 + res_z), (1 + res_y), (1 + res_x), 3)) 70 | closest_points = np.swapaxes(closest_points, 0, 2) 71 | 72 | return b_min, b_max, volume, closest_points 73 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import os 2 | from os.path import join, basename, dirname, realpath 3 | import sys 4 | import time 5 | from datetime import date, datetime 6 | import math 7 | 8 | PROJECT_DIR = dirname(realpath(__file__)) 9 | LOGS_PATH = join(PROJECT_DIR, 'checkpoints') 10 | SAMPLES_PATH = join(PROJECT_DIR, 'results', 'saved_samples') 11 | sys.path.append(PROJECT_DIR) 12 | 13 | import torch 14 | from torch.utils.data import DataLoader 15 | from torch.utils.tensorboard import SummaryWriter 16 | import numpy as np 17 | 18 | from lib.config_parser import parse_config 19 | from lib.dataset import CloDataSet 20 | from lib.network import SCALE 21 | from lib.train_eval_funcs import train, test 22 | from lib.utils_io import load_masks, save_model, save_latent_vectors, load_latent_vectors 23 | from lib.utils_model import SampleSquarePoints 24 | from lib.utils_train import adjust_loss_weights 25 | 26 | torch.manual_seed(12345) 27 | np.random.seed(12345) 28 | 29 | DEVICE = torch.device('cuda') 30 | 31 | def main(): 32 | args = parse_config() 33 | 34 | exp_name = args.name 35 | 36 | # NOTE: when using your custom data, modify the following path to where the packed data is stored. 37 | data_root = join(PROJECT_DIR, 'data', 'packed', '{}'.format(args.data_root)) 38 | 39 | log_dir = join(PROJECT_DIR,'tb_logs/{}/{}'.format(date.today().strftime('%m%d'), exp_name)) 40 | ckpt_dir = join(LOGS_PATH, exp_name) 41 | os.makedirs(ckpt_dir, exist_ok=True) 42 | 43 | samples_dir_val = join(SAMPLES_PATH, exp_name, 'val') 44 | samples_dir_test = join(SAMPLES_PATH, exp_name, 'test') 45 | os.makedirs(samples_dir_test, exist_ok=True) 46 | os.makedirs(samples_dir_val, exist_ok=True) 47 | 48 | flist_uv, valid_idx, uv_coord_map = load_masks(PROJECT_DIR, args.img_size, body_model='smpl') 49 | 50 | # build_model 51 | model = SCALE( 52 | input_nc=3, 53 | output_nc_unet=args.pix_feat_dim, 54 | img_size=args.img_size, 55 | hsize=args.hsize, 56 | nf=args.nf, 57 | up_mode=args.up_mode, 58 | use_dropout=bool(args.use_dropout), 59 | pos_encoding=bool(args.pos_encoding), 60 | num_emb_freqs=args.num_emb_freqs, 61 | posemb_incl_input=bool(args.posemb_incl_input), 62 | uv_feat_dim=2 63 | ) 64 | print(model) 65 | 66 | subpixel_sampler = SampleSquarePoints(npoints=args.npoints, 67 | include_end=bool(args.pq_include_end), 68 | min_val=args.pqmin, 69 | max_val=args.pqmax) 70 | 71 | # Below: for historical reason we have a 256D latent code in the network that is shared for all examples of each garment type. 72 | # Since SCALE is a garment-specific model, this latent code is therefore the same for all examples that the model sees, 73 | # and can be safely seen as part of the network parameters. 74 | lat_vecs = torch.nn.Embedding(1, args.latent_size, max_norm=1.0).cuda() 75 | torch.nn.init.normal_(lat_vecs.weight.data, 0.0, 1e-2 / math.sqrt(args.latent_size)) 76 | 77 | optimizer = torch.optim.Adam( 78 | [ 79 | {"params": model.parameters(), "lr": args.lr,}, 80 | {"params": lat_vecs.parameters(), "lr": args.lr,}, 81 | ]) 82 | 83 | n_epochs = args.epochs 84 | epoch_now = 0 85 | ''' 86 | ------------ Load checkpoints in case of test or resume training ------------ 87 | ''' 88 | if args.mode.lower() in ['test', 'resume']: 89 | checkpoints = sorted([fn for fn in os.listdir(ckpt_dir) if fn.endswith('_model.pt')]) 90 | latest = join(ckpt_dir, checkpoints[-1]) 91 | print('Loading checkpoint {}'.format(basename(latest))) 92 | ckpt_loaded = torch.load(latest) 93 | 94 | model.load_state_dict(ckpt_loaded['model_state']) 95 | 96 | checkpoints = sorted([fn for fn in os.listdir(ckpt_dir) if fn.endswith('_latent_vecs.pt')]) 97 | checkpoint = join(ckpt_dir, checkpoints[-1]) 98 | load_latent_vectors(checkpoint, lat_vecs) 99 | 100 | if args.mode.lower() == 'resume': 101 | optimizer.load_state_dict(ckpt_loaded['optimizer_state']) 102 | for state in optimizer.state.values(): 103 | for k, v in state.items(): 104 | if torch.is_tensor(v): 105 | state[k] = v.to(DEVICE) 106 | epoch_now = ckpt_loaded['epoch'] + 1 107 | print('Resume training from epoch {}'.format(epoch_now)) 108 | 109 | if args.mode.lower() == 'test': 110 | epoch_idx = ckpt_loaded['epoch'] 111 | model.to(DEVICE) 112 | print('Test model with checkpoint at epoch {}'.format(epoch_idx)) 113 | 114 | 115 | ''' 116 | ------------ Training over or resume from saved checkpoints ------------ 117 | ''' 118 | if args.mode.lower() in ['train', 'resume']: 119 | train_set = CloDataSet(root_dir=data_root, split='train', sample_spacing=args.data_spacing, 120 | img_size=args.img_size, scan_npoints=args.scan_npoints) 121 | val_set = CloDataSet(root_dir=data_root, split='val', sample_spacing=args.data_spacing, 122 | img_size=args.img_size, scan_npoints=args.scan_npoints) 123 | 124 | train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=4) 125 | val_loader = DataLoader(val_set, batch_size=args.batch_size, shuffle=False, num_workers=4) 126 | 127 | writer = SummaryWriter(log_dir=log_dir) 128 | 129 | print("Total: {} training examples, {} val examples. Training started..".format(len(train_set), len(val_set))) 130 | 131 | 132 | model.to(DEVICE) 133 | start = time.time() 134 | pbar = range(epoch_now, n_epochs) 135 | for epoch_idx in pbar: 136 | wdecay_rgl = adjust_loss_weights(args.w_rgl, epoch_idx, mode='decay', start=args.decay_start, every=args.decay_every) 137 | wrise_normal = adjust_loss_weights(args.w_normal, epoch_idx, mode='rise', start=args.rise_start, every=args.rise_every) 138 | 139 | train_stats = train( 140 | model, lat_vecs, DEVICE, train_loader, optimizer, 141 | flist_uv, valid_idx, uv_coord_map, 142 | subpixel_sampler=subpixel_sampler, 143 | w_s2m=args.w_s2m, w_m2s=args.w_m2s, w_rgl=wdecay_rgl, 144 | w_latent_rgl=1.0, w_normal=wrise_normal, 145 | ) 146 | 147 | if epoch_idx % 100 == 0 or epoch_idx == n_epochs - 1: 148 | ckpt_path = join(ckpt_dir, '{}_epoch{}_model.pt'.format(exp_name, str(epoch_idx).zfill(5))) 149 | save_model(ckpt_path, model, epoch_idx, optimizer=optimizer) 150 | ckpt_path = join(ckpt_dir, '{}_epoch{}_latent_vecs.pt'.format(exp_name, str(epoch_idx).zfill(5))) 151 | save_latent_vectors(ckpt_path, lat_vecs, epoch_idx) 152 | 153 | # test on val set every N epochs 154 | if epoch_idx % args.val_every == 0: 155 | dur = (time.time() - start) / (60 * (epoch_idx-epoch_now+1)) 156 | now = datetime.now() 157 | dt_string = now.strftime("%d/%m/%Y %H:%M:%S") 158 | print('\n{}, Epoch {}, average {:.2f} min / epoch.'.format(dt_string, epoch_idx, dur)) 159 | print('Weights s2m: {:.1e}, m2s: {:.1e}, normal: {:.1e}, rgl: {:.1e}'.format(args.w_s2m, args.w_m2s, wrise_normal, wdecay_rgl)) 160 | 161 | checkpoints = sorted([fn for fn in os.listdir(ckpt_dir) if fn.endswith('_latent_vecs.pt')]) 162 | checkpoint = join(ckpt_dir, checkpoints[-1]) 163 | load_latent_vectors(checkpoint, lat_vecs) 164 | val_stats = test(model, lat_vecs, 165 | DEVICE, val_loader, epoch_idx, 166 | samples_dir_val, 167 | flist_uv, 168 | valid_idx, uv_coord_map, 169 | subpixel_sampler=subpixel_sampler, 170 | model_name=exp_name, 171 | save_all_results=bool(args.save_all_results), 172 | ) 173 | 174 | tensorboard_tabs = ['model2scan', 'scan2model', 'normal_loss', 'residual_square', 'latent_reg', 'total_loss'] 175 | stats = {'train': train_stats, 'val': val_stats} 176 | 177 | for split in ['train', 'val']: 178 | for (tab, stat) in zip(tensorboard_tabs, stats[split]): 179 | writer.add_scalar('{}/{}'.format(tab, split), stat, epoch_idx) 180 | 181 | 182 | end = time.time() 183 | t_total = (end - start) / 60 184 | print("Training finished, duration: {:.2f} minutes. Now eval on test set..\n".format(t_total)) 185 | writer.close() 186 | 187 | 188 | ''' 189 | ------------ Test model ------------ 190 | ''' 191 | test_set = CloDataSet(root_dir=data_root, split='test', sample_spacing=args.data_spacing, 192 | img_size=args.img_size, scan_npoints=args.scan_npoints) 193 | test_loader = DataLoader(test_set, batch_size=args.batch_size, shuffle=True, num_workers=4) 194 | 195 | print('Eval on test data...') 196 | start = time.time() 197 | test_m2s, test_s2m, test_lnormal, _, _, _ = test(model, lat_vecs, DEVICE, test_loader, epoch_idx, 198 | samples_dir_test, 199 | flist_uv, 200 | valid_idx, uv_coord_map, 201 | mode='test', 202 | subpixel_sampler=subpixel_sampler, 203 | model_name=exp_name, 204 | save_all_results=bool(args.save_all_results), 205 | ) 206 | 207 | print('\nDuration: {}'.format(time.time() - start)) 208 | 209 | testset_result = "Test on test set, {} examples, m2s dist: {:.3e}, s2m dist: {:.3e}, Chamfer total: {:.3e}, normal loss: {:.3e}.\n\n".format(len(test_set), test_m2s, test_s2m, test_m2s+test_s2m, test_lnormal) 210 | print(testset_result) 211 | 212 | with open(join(PROJECT_DIR, 'results', 'results_log.txt'), 'a+') as fn: 213 | fn.write('{}, n_pq={}, epoch={}\n'.format(args.name, args.npoints, epoch_idx)) 214 | fn.write('\t{}'.format(testset_result)) 215 | 216 | 217 | if __name__ == '__main__': 218 | main() -------------------------------------------------------------------------------- /render/cam_front_extrinsic.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qianlim/SCALE/98f5557e746de3536168c413901ff48c0571c5a4/render/cam_front_extrinsic.npy -------------------------------------------------------------------------------- /render/o3d_render_pcl.py: -------------------------------------------------------------------------------- 1 | import os 2 | from os.path import join, basename, dirname, realpath 3 | import glob 4 | 5 | import numpy as np 6 | import open3d as o3d 7 | from tqdm import tqdm 8 | 9 | 10 | def render_pcl_front_view(vis, cam_params=None, fn=None, img_save_fn=None, pt_size=3): 11 | mesh = o3d.io.read_point_cloud(fn) 12 | vis.add_geometry(mesh) 13 | opt = vis.get_render_option() 14 | 15 | opt.point_size = pt_size 16 | 17 | ctr = vis.get_view_control() 18 | ctr.convert_from_pinhole_camera_parameters(cam_params) 19 | 20 | vis.poll_events() 21 | vis.update_renderer() 22 | vis.capture_screen_image(img_save_fn, True) 23 | 24 | vis.clear_geometries() 25 | 26 | 27 | def main(): 28 | import argparse 29 | parser = argparse.ArgumentParser() 30 | parser.add_argument('-c', '--color_opt', default=None, choices=['normal', 'patch'], help='Color the poins by their normals or group them by patches. Leave it the default None will render both colors.') 31 | parser.add_argument('-n', '--model_name', default='SCALE_demo_00000_simuskirt', help='Name of the model (experiment), the same as the --name flag in the main experiment') 32 | parser.add_argument('-r', '--img_res', type=int ,default=1024, help='Resolution of rendered image') 33 | args = parser.parse_args() 34 | 35 | img_res = args.img_res 36 | 37 | # path for saving the rendered images 38 | SCRIPT_DIR = dirname(realpath(__file__)) 39 | target_root = join(SCRIPT_DIR, '..', 'results', 'rendered_imgs') 40 | os.makedirs(target_root, exist_ok=True) 41 | 42 | # set up camera 43 | focal_length = 900 * (args.img_res / 1024.) # 900 is a hand-set focal length when the img resolution=1024. 44 | x0, y0 = (img_res-1)/2, (img_res-1)/2 45 | INTRINSIC = np.array([ 46 | [focal_length, 0., x0], 47 | [0., focal_length, y0], 48 | [0., 0., 1] 49 | ]) 50 | 51 | EXTRINSIC = np.load(join(SCRIPT_DIR, 'cam_front_extrinsic.npy')) 52 | 53 | 54 | cam_intrinsics = o3d.camera.PinholeCameraIntrinsic() 55 | cam_intrinsics.intrinsic_matrix = INTRINSIC 56 | cam_intrinsics.width = img_res 57 | cam_intrinsics.height = img_res 58 | 59 | cam_params_front = o3d.camera.PinholeCameraParameters() 60 | cam_params_front.intrinsic = cam_intrinsics 61 | cam_params_front.extrinsic = EXTRINSIC 62 | 63 | # configs for the rendering 64 | render_opts = { 65 | 'normal': ('normal_colored', 3.0, '_pred.ply'), 66 | 'patch': ('patch_colored', 4.0, '_pred_patchcolor.ply'), 67 | } 68 | 69 | vis = o3d.visualization.Visualizer() 70 | vis.create_window(width=img_res, height=img_res) 71 | 72 | # render 73 | for opt in render_opts.keys(): 74 | if (args.color_opt is not None) and (opt != args.color_opt): 75 | continue 76 | 77 | color_mode, pt_size, ext = render_opts[opt] 78 | 79 | render_savedir = join(target_root, args.model_name, color_mode) 80 | os.makedirs(render_savedir, exist_ok=True) 81 | 82 | ply_folder = join(target_root, '..', 'saved_samples', args.model_name, 'test') 83 | print('parsing pcl files at {}..'.format(ply_folder)) 84 | flist = sorted(glob.glob(join(ply_folder, '*{}'.format(ext)))) 85 | 86 | for fn in tqdm(flist): 87 | bn = basename(fn) 88 | img_save_fn = join(render_savedir, bn.replace('{}'.format(ext), '.png')) 89 | render_pcl_front_view(vis, cam_params_front, fn, img_save_fn, pt_size=pt_size) 90 | 91 | 92 | if __name__ == '__main__': 93 | main() -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | configargparse 2 | trimesh>=3.8 3 | open3d>=0.12 4 | tensorboard==2.4 5 | torch==1.6.0 6 | numpy==1.19.2 7 | 8 | -------------------------------------------------------------------------------- /teasers/teaser.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qianlim/SCALE/98f5557e746de3536168c413901ff48c0571c5a4/teasers/teaser.gif --------------------------------------------------------------------------------