├── .gitignore ├── LICENSE ├── README.md ├── config ├── __init__.py ├── cmdline.py ├── config_files │ ├── coco │ │ ├── eval_pix2surf_mv_car.yaml │ │ ├── eval_pix2surf_mv_chair.yaml │ │ ├── eval_pix2surf_mv_plane.yaml │ │ ├── eval_pix2surf_sv_car.yaml │ │ ├── eval_pix2surf_sv_chair.yaml │ │ ├── eval_pix2surf_sv_plane.yaml │ │ ├── pix2surf_mv_car.yaml │ │ ├── pix2surf_mv_chair.yaml │ │ ├── pix2surf_mv_plane.yaml │ │ ├── pix2surf_sv_car.yaml │ │ ├── pix2surf_sv_chair.yaml │ │ ├── pix2surf_sv_plane.yaml │ │ ├── render_pix2surf_mv_car.yaml │ │ ├── render_pix2surf_mv_chair.yaml │ │ ├── render_pix2surf_mv_plane.yaml │ │ ├── render_pix2surf_sv_car.yaml │ │ ├── render_pix2surf_sv_chair.yaml │ │ ├── render_pix2surf_sv_plane.yaml │ │ ├── xnocsinit_sv_car.yaml │ │ ├── xnocsinit_sv_chair.yaml │ │ └── xnocsinit_sv_plane.yaml │ ├── plain │ │ ├── eval_pix2surf_mv_car.yaml │ │ ├── eval_pix2surf_mv_chair.yaml │ │ ├── eval_pix2surf_mv_plane.yaml │ │ ├── eval_pix2surf_sv_car.yaml │ │ ├── eval_pix2surf_sv_chair.yaml │ │ ├── eval_pix2surf_sv_plane.yaml │ │ ├── pix2surf_mv_car.yaml │ │ ├── pix2surf_mv_chair.yaml │ │ ├── pix2surf_mv_plane.yaml │ │ ├── pix2surf_sv_car.yaml │ │ ├── pix2surf_sv_chair.yaml │ │ ├── pix2surf_sv_plane.yaml │ │ ├── render_pix2surf_mv_car.yaml │ │ ├── render_pix2surf_mv_chair.yaml │ │ ├── render_pix2surf_mv_plane.yaml │ │ ├── render_pix2surf_sv_car.yaml │ │ ├── render_pix2surf_sv_chair.yaml │ │ ├── render_pix2surf_sv_plane.yaml │ │ ├── xnocsinit_sv_car.yaml │ │ ├── xnocsinit_sv_chair.yaml │ │ └── xnocsinit_sv_plane.yaml │ └── viz │ │ ├── render_pix2surf_mv_car.yaml │ │ ├── render_pix2surf_mv_chair.yaml │ │ └── render_pix2surf_mv_plane.yaml ├── default.py └── startup.py ├── core ├── __init__.py ├── evaluation.py ├── models │ ├── __init__.py │ ├── modelbase_v2.py │ ├── neurips_baseline.py │ ├── pix2surf_mv.py │ ├── pix2surf_mv_mveval.py │ ├── pix2surf_mv_render.py │ ├── pix2surf_sv.py │ ├── pix2surf_sv_mveval.py │ ├── pix2surf_sv_render.py │ └── utils │ │ ├── __init__.py │ │ ├── model_utils.py │ │ └── render_utils.py ├── net_bank │ ├── __init__.py │ ├── loss.py │ ├── mlp.py │ ├── modules │ │ ├── __init__.py │ │ └── neurips_nox_network_util.py │ ├── pix2surf_cnn.py │ ├── xnocs_segnet.py │ └── xnocs_setsegnet.py └── trainer.py ├── dataset ├── __init__.py ├── mv5_correspondence_dataset.py └── neurips_dataset.py ├── logger ├── __init__.py ├── logger_meta │ ├── __init__.py │ ├── base_logger.py │ ├── image_logger.py │ ├── metric_logger.py │ ├── model_logger.py │ ├── obj_logger.py │ └── xls_logger.py └── logger_s1.py ├── requirements.txt ├── resource ├── index │ ├── pix2surf_viz_car_test.json │ ├── pix2surf_viz_chair_test.json │ ├── pix2surf_viz_plane_test.json │ ├── shapenet_coco_car_test.json │ ├── shapenet_coco_car_train.json │ ├── shapenet_coco_car_vali.json │ ├── shapenet_coco_chair_test.json │ ├── shapenet_coco_chair_train.json │ ├── shapenet_coco_chair_vali.json │ ├── shapenet_coco_plane_test.json │ ├── shapenet_coco_plane_train.json │ ├── shapenet_coco_plane_vali.json │ ├── shapenet_plain_car_mini_test.json │ ├── shapenet_plain_car_test.json │ ├── shapenet_plain_car_train.json │ ├── shapenet_plain_chair_mini_test.json │ ├── shapenet_plain_chair_test.json │ ├── shapenet_plain_chair_train.json │ ├── shapenet_plain_plane_mini_test.json │ ├── shapenet_plain_plane_test.json │ └── shapenet_plain_plane_train.json └── utils │ ├── car_vali.npy │ ├── chair_vali.npy │ ├── make_coco_validation.py │ └── plane_vali.npy ├── run.py ├── run_batch.py └── scripts ├── evaluate_pix2surf.sh ├── evaluate_pix2surf_coco.sh ├── pix2surf_car.sh ├── pix2surf_car_coco.sh ├── pix2surf_chair.sh ├── pix2surf_chair_coco.sh ├── pix2surf_plane.sh ├── pix2surf_plane_coco.sh ├── render_pix2surf_coco.sh ├── render_pix2surf_plain.sh └── render_pix2surf_viz.sh /.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 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settin 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | .idea 131 | resource/data/shapenet_plain 132 | log 133 | resource/data 134 | resource/*weight* 135 | config/config_files/test.yaml -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 Ray 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Pix2Surf 2 | One implementation of our ECCV2020 paper **[Pix2Surf: Learning Parametric 3D Surface Models of Objects from Images](https://geometry.stanford.edu/projects/pix2surf/)** 3 | 4 | If you use the code please cite our paper. 5 | 6 | ```latex 7 | @inproceedings{pix2surf_2020, 8 | author = {Lei, Jiahui and Sridhar, Srinath and Guerrero, Paul and Sung, Minhyuk and Mitra, Niloy and Guibas, Leonidas J.}, 9 | title = {Pix2Surf: Learning Parametric 3D Surface Models of Objects from Images}, 10 | booktitle = {Proceedings of European Conference on Computer Vision ({ECCV})}, 11 | url = {https://geometry.stanford.edu/projects/pix2surf}, 12 | month = August, 13 | year = {2020} 14 | } 15 | ``` 16 | 17 | 18 | 19 | ## Installation 20 | 21 | We use [pytorch 1.1](https://pytorch.org/get-started/previous-versions/) and currently only support linux with GPUs. We also depend on [tk3dv](https://github.com/drsrinathsridhar/tk3dv). To install the code: 22 | ```shell script 23 | # create virtue env 24 | conda create -n pix2surf python=3.6 25 | source activate pix2surf 26 | 27 | # install pytorch with corresponding cuda version, for example we use cuda 10.0 here 28 | conda install pytorch==1.1.0 torchvision==0.3.0 cudatoolkit=10.0 -c pytorch 29 | 30 | # install other requirements 31 | pip install -r ./requirements.txt 32 | 33 | # install tk3dv 34 | pip install git+https://github.com/drsrinathsridhar/tk3dv.git 35 | ``` 36 | 37 | 38 | 39 | ## Dataset 40 | 41 | We use the dataset from X-NOCS project. Current code supports both ShapeNet-COCO (~172GB) and ShapeNet-Plain (~5GB) Dataset. Please download any or both of them from [here](https://github.com/drsrinathsridhar/xnocs/blob/master/dataset/README.md) and unzip. Then, link the dataset to the code and pre-process the files by: 42 | 43 | ```shell 44 | # prepare dirs 45 | cd [...]/ProjectRoot 46 | mkdir ./log # or link to somewhere 47 | mkdir ./resource/weight 48 | mkdir ./resource/data 49 | 50 | # link COCO or Plain 51 | ln -s [...]/PathToShapeNetCOCO ./resource/data/shapenet_coco 52 | ln -s [...]/PathToShapeNetPlain ./resource/data/shapenet_plain 53 | 54 | # do some minor modification to the dataset structure 55 | cd ./resource/utils 56 | python make_coco_validation.py 57 | ``` 58 | 59 | Now, you will see the `./resource/data` folder structure like this: 60 | 61 | ```shell 62 | ├── shapenet_coco 63 | │   ├── test 64 | │   ├── train 65 | │   └── vali 66 | └── shapenet_plain 67 | ├── test 68 | └── train 69 | ``` 70 | 71 | 72 | 73 | ## Run the code 74 | 75 | Please check the `ProjectRoot/scripts` and run the `.sh` scripts under ProjectRoot dir. Here we give some examples. 76 | 77 | ### Evaluation 78 | 79 | Here is an evaluation example for pre-trained model on ShapeNet-Plain Dataset. 80 | 81 | First, download pre-trained models on ShapeNet-Plain Dataset from **[here](http://download.cs.stanford.edu/orion/pix2surf/plain-weight.zip)**. Unzip and move all `*.model` to `ProjectRoot/resource/weight` and then just run: 82 | 83 | ``` 84 | cd [...]/ProjectRoot 85 | bash ./scripts/evaluate_pix2surf.sh 86 | ``` 87 | 88 | The log will be stored at `./log/eval-pix2surf-XXXXX` and you can find the metric reports under the `xls` sub-folder. This will reproduce the numbers of both single view and multi view version of Pix2Surf in table 1 in our paper. 89 | 90 | Pre-trained weights on ShapeNet-COCO dataset can be found [here](http://download.cs.stanford.edu/orion/pix2surf/coco-weight.zip). 91 | 92 | ### Train your own model 93 | 94 | Here is a training example for car category on ShapeNet-COCO Dataset. 95 | 96 | ```shell 97 | cd [...]/ProjectRoot 98 | bash ./scripts/pix2surf_car_coco.sh 99 | ``` 100 | 101 | All the trainings currently are configured with 2 GPUs, but 4 or more GPUs with larger batch size will probably lead to higher performance. 102 | 103 | ### Visualization 104 | 105 | We provide a naive post-processing code to generate meshes and high resolution Geometry Images with Texture Maps. Here is an example for generating the visualization in the paper. 106 | 107 | Please download a small dataset with consistent lighting condition across multiple views **[here](http://download.cs.stanford.edu/orion/pix2surf/pix2surf_viz_dataset.zip)**. As our method has unsupervised learned component, to reproduce the same learned chart pattern as shown in the paper, please download our weight **[here](http://download.cs.stanford.edu/orion/pix2surf/viz-weight.zip)**. (It's interesting to see that different patterns are discovered in different trainings) 108 | 109 | ```shell 110 | # link visualziation dataset 111 | cd [...]/ProjectRoot 112 | ln -s [...]/pix2surf_viz ./resource/data 113 | 114 | # move weights to resource/weight 115 | mv [...]/PathToDownloadedWeight/*.model ./resource/weight 116 | 117 | # run postprocessing 118 | bash ./scripts/render_pix2surf_viz.sh 119 | ``` 120 | 121 | Under the `log\render-pix2surf-XXXXXX`, you will find `obj` sub-folder that contains the mesh results. In the `image` sub-folder, you will find Geometry Image: `*GIM.png`or `*GIM-uni.png` and Texture Map `*TEX-uni.png` or `*TEX.png` . Geometry Images along with Texture Maps have similar data structure as NOCS Maps along with RGB colors (different arrangement), so you can use [tk3dv tools](https://github.com/drsrinathsridhar/tk3dv/blob/master/examples/visualizeNOCSMap.py) to visualize it. 122 | 123 | 124 | 125 | ## Check more about NOCS series 126 | 127 | **[[NOCS]](https://geometry.stanford.edu/projects/NOCS_CVPR2019/)** NOCS for Pose Estimation [CVPR2019] 128 | 129 | **[[X-NOCS]](https://geometry.stanford.edu/projects/xnocs/)** Two-intersection NOCS for shape reconstruction [NeurIPS2019] 130 | 131 | **[[ANCSH]](https://articulated-pose.github.io/)** Articulated Pose Estimation [CVPR2020] 132 | 133 | **[[S-NOCS]](https://geometry.stanford.edu/projects/pix2surf/)** Shape reconstruction in NOCS with Surfaces (This work) [ECCV2020] 134 | 135 | **[[T-NOCS]](https://geometry.stanford.edu/projects/caspr/)** NOCS along Time Axis [arXiv2020 Pre-print] -------------------------------------------------------------------------------- /config/__init__.py: -------------------------------------------------------------------------------- 1 | from .default import get_default_cfg 2 | from .cmdline import merge_from_cmdline 3 | from .startup import startup 4 | from os.path import join 5 | 6 | 7 | def get_cfg(filename=None, interactive=True): 8 | """ 9 | config priority: cmd>cfg-file>default 10 | :param filename: filename XXX in $project/config/config_files/XXX.yaml 11 | :return: a frozen configuration 12 | """ 13 | # get default 14 | cfg = get_default_cfg() 15 | # get config file parameter, will do another time later to overwrite local file config 16 | cfg = merge_from_cmdline(cfg) 17 | # merge local file, cmd>file 18 | if cfg.CONFIG_FILE is not 'None': 19 | filepath = join(cfg.ROOT_DIR, 'config', 'config_files', cfg.CONFIG_FILE) 20 | cfg.merge_from_file(filepath) 21 | elif filename is not None: 22 | filepath = join(cfg.ROOT_DIR, 'config', 'config_files', filename) 23 | cfg.merge_from_file(filepath) 24 | # merge cmd line 25 | cfg = merge_from_cmdline(cfg) 26 | # startup 27 | startup(cfg, interactive) 28 | cfg.freeze() 29 | return cfg 30 | -------------------------------------------------------------------------------- /config/cmdline.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | 4 | def merge_from_cmdline(cfg): 5 | """ 6 | Merge some usually changed settings from cmd 7 | """ 8 | parser = argparse.ArgumentParser() 9 | parser.add_argument('--config', type=str, default='None', help="Choose config file") 10 | parser.add_argument('--gpu', type=str, default=None, help="Set GPU device, type=str, e.g: --gpu=0,1,2 ") 11 | parser.add_argument('--logdir', type=str, default=None, help='log dir name in $project/log/XXXX/.... e.g. exp') 12 | cmd = vars(parser.parse_args()) 13 | if cmd['config'] is not 'None': 14 | cfg.CONFIG_FILE = cmd['config'] 15 | if cmd['gpu'] is not None: 16 | cfg.GPU = [int(id) for id in cmd['gpu'].split(",")] 17 | if cmd['logdir'] is not None: 18 | cfg.LOG_DIR = cmd['logdir'] 19 | return cfg 20 | -------------------------------------------------------------------------------- /config/config_files/coco/eval_pix2surf_mv_car.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_coco' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_coco_car_train.json', 5 | 'resource/index/shapenet_coco_car_test.json'] 6 | MODES: ['test'] 7 | MODEL: 'pix2surf_mv_mveval' 8 | 9 | GPU: [0] 10 | BATCHSIZE: 2 11 | DATALOADER_WORKERS: 4 12 | EPOCH_TOTAL: 1 13 | LR: 0.0 14 | 15 | LOG_DIR: 'eval-pix2surf-mv-car-coco' 16 | LOGGER: 'logger_s1' 17 | LOGGER_SELECT: ['metric','image','model','xls'] 18 | MODEL_SAVE_PER_N_EPOCH: 1 19 | VIS_PER_N_EPOCH: 1 20 | VIS_PER_N_BATCH: 10 21 | VIS_ONE_PER_BATCH: True 22 | 23 | RESUME: False 24 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_mv_car_coco_epoch_30.model'] 25 | 26 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 27 | 'core/net_bank/mlp.py', 28 | 'core/net_bank/loss.py', 29 | 'core/evaluation.py', 30 | 'core/models/pix2surf_mv.py', 31 | 'core/models/utils/model_utils.py', 32 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/coco/eval_pix2surf_mv_chair.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_coco' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_coco_chair_train.json', 5 | 'resource/index/shapenet_coco_chair_test.json'] 6 | MODES: ['test'] 7 | MODEL: 'pix2surf_mv_mveval' 8 | 9 | GPU: [0] 10 | BATCHSIZE: 2 11 | DATALOADER_WORKERS: 4 12 | EPOCH_TOTAL: 1 13 | LR: 0.0 14 | 15 | LOG_DIR: 'eval-pix2surf-mv-chair-coco' 16 | LOGGER: 'logger_s1' 17 | LOGGER_SELECT: ['metric','image','model','xls'] 18 | MODEL_SAVE_PER_N_EPOCH: 1 19 | VIS_PER_N_EPOCH: 1 20 | VIS_PER_N_BATCH: 10 21 | VIS_ONE_PER_BATCH: True 22 | 23 | RESUME: False 24 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_mv_chair_coco_epoch_30.model'] 25 | 26 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 27 | 'core/net_bank/mlp.py', 28 | 'core/net_bank/loss.py', 29 | 'core/evaluation.py', 30 | 'core/models/pix2surf_mv.py', 31 | 'core/models/utils/model_utils.py', 32 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/coco/eval_pix2surf_mv_plane.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_coco' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_coco_plane_train.json', 5 | 'resource/index/shapenet_coco_plane_test.json'] 6 | MODES: ['test'] 7 | MODEL: 'pix2surf_mv_mveval' 8 | 9 | GPU: [0] 10 | BATCHSIZE: 2 11 | DATALOADER_WORKERS: 4 12 | EPOCH_TOTAL: 1 13 | LR: 0.0 14 | 15 | LOG_DIR: 'eval-pix2surf-mv-plane-coco' 16 | LOGGER: 'logger_s1' 17 | LOGGER_SELECT: ['metric','image','model','xls'] 18 | MODEL_SAVE_PER_N_EPOCH: 1 19 | VIS_PER_N_EPOCH: 1 20 | VIS_PER_N_BATCH: 10 21 | VIS_ONE_PER_BATCH: True 22 | 23 | RESUME: False 24 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_mv_plane_coco_epoch_30.model'] 25 | 26 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 27 | 'core/net_bank/mlp.py', 28 | 'core/net_bank/loss.py', 29 | 'core/evaluation.py', 30 | 'core/models/pix2surf_mv.py', 31 | 'core/models/utils/model_utils.py', 32 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/coco/eval_pix2surf_sv_car.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_coco' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_coco_car_train.json', 5 | 'resource/index/shapenet_coco_car_test.json'] 6 | MODES: ['test'] 7 | MODEL: 'pix2surf_sv_mveval' 8 | 9 | GPU: [0] 10 | BATCHSIZE: 2 11 | DATALOADER_WORKERS: 4 12 | EPOCH_TOTAL: 1 13 | LR: 0.0 14 | 15 | LOG_DIR: 'eval-pix2surf-sv-car-coco' 16 | LOGGER: 'logger_s1' 17 | LOGGER_SELECT: ['metric','image','model','xls'] 18 | MODEL_SAVE_PER_N_EPOCH: 1 19 | VIS_PER_N_EPOCH: 1 20 | VIS_PER_N_BATCH: 10 21 | VIS_ONE_PER_BATCH: True 22 | 23 | RESUME: False 24 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_sv_car_coco_epoch_30.model'] 25 | 26 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 27 | 'core/net_bank/mlp.py', 28 | 'core/net_bank/loss.py', 29 | 'core/evaluation.py', 30 | 'core/models/pix2surf_sv.py', 31 | 'core/models/utils/model_utils.py', 32 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/coco/eval_pix2surf_sv_chair.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_coco' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_coco_chair_train.json', 5 | 'resource/index/shapenet_coco_chair_test.json'] 6 | MODES: ['test'] 7 | MODEL: 'pix2surf_sv_mveval' 8 | 9 | GPU: [0] 10 | BATCHSIZE: 2 11 | DATALOADER_WORKERS: 4 12 | EPOCH_TOTAL: 1 13 | LR: 0.0 14 | 15 | LOG_DIR: 'eval-pix2surf-sv-chair-coco' 16 | LOGGER: 'logger_s1' 17 | LOGGER_SELECT: ['metric','image','model','xls'] 18 | MODEL_SAVE_PER_N_EPOCH: 1 19 | VIS_PER_N_EPOCH: 1 20 | VIS_PER_N_BATCH: 10 21 | VIS_ONE_PER_BATCH: True 22 | 23 | RESUME: False 24 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_sv_chair_coco_epoch_30.model'] 25 | 26 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 27 | 'core/net_bank/mlp.py', 28 | 'core/net_bank/loss.py', 29 | 'core/evaluation.py', 30 | 'core/models/pix2surf_sv.py', 31 | 'core/models/utils/model_utils.py', 32 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/coco/eval_pix2surf_sv_plane.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_coco' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_coco_plane_train.json', 5 | 'resource/index/shapenet_coco_plane_test.json'] 6 | MODES: ['test'] 7 | MODEL: 'pix2surf_sv_mveval' 8 | 9 | GPU: [0] 10 | BATCHSIZE: 2 11 | DATALOADER_WORKERS: 4 12 | EPOCH_TOTAL: 1 13 | LR: 0.0 14 | 15 | LOG_DIR: 'eval-pix2surf-sv-plane-coco' 16 | LOGGER: 'logger_s1' 17 | LOGGER_SELECT: ['metric','image','model','xls'] 18 | MODEL_SAVE_PER_N_EPOCH: 1 19 | VIS_PER_N_EPOCH: 1 20 | VIS_PER_N_BATCH: 10 21 | VIS_ONE_PER_BATCH: True 22 | 23 | RESUME: False 24 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_sv_plane_coco_epoch_30.model'] 25 | 26 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 27 | 'core/net_bank/mlp.py', 28 | 'core/net_bank/loss.py', 29 | 'core/evaluation.py', 30 | 'core/models/pix2surf_sv.py', 31 | 'core/models/utils/model_utils.py', 32 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/coco/pix2surf_mv_car.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_coco' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_coco_car_train.json', 5 | 'resource/index/shapenet_coco_car_vali.json', 6 | 'resource/index/shapenet_coco_car_test.json'] 7 | MODES: ['train','vali'] 8 | MODEL: 'pix2surf_mv' 9 | 10 | GPU: [0,1] 11 | BATCHSIZE: 4 12 | DATALOADER_WORKERS: 6 13 | EPOCH_TOTAL: 50 14 | LR: 1e-4 15 | 16 | LOG_DIR: 'pix2surf-mv-car-coco' 17 | LOGGER: 'logger_s1' 18 | LOGGER_SELECT: ['metric','image','model','xls'] 19 | MODEL_SAVE_PER_N_EPOCH: 5 20 | VIS_PER_N_EPOCH: 1 21 | VIS_PER_N_BATCH: 50 22 | VIS_ONE_PER_BATCH: True 23 | VIS_TRAIN_PER_BATCH: 100 24 | 25 | RESUME: False 26 | MODEL_INIT_PATH: ['log/pix2surf-sv-car-coco/model/epoch_30.model'] 27 | 28 | 29 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 30 | 'core/net_bank/mlp.py', 31 | 'core/net_bank/loss.py', 32 | 'core/models/utils/model_utils.py', 33 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/coco/pix2surf_mv_chair.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_coco' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_coco_chair_train.json', 5 | 'resource/index/shapenet_coco_chair_vali.json', 6 | 'resource/index/shapenet_coco_chair_test.json'] 7 | MODES: ['train','vali'] 8 | MODEL: 'pix2surf_mv' 9 | 10 | GPU: [0,1] 11 | BATCHSIZE: 4 12 | DATALOADER_WORKERS: 6 13 | EPOCH_TOTAL: 50 14 | LR: 3e-5 15 | 16 | LOG_DIR: 'pix2surf-mv-chair-coco' 17 | LOGGER: 'logger_s1' 18 | LOGGER_SELECT: ['metric','image','model','xls'] 19 | MODEL_SAVE_PER_N_EPOCH: 5 20 | VIS_PER_N_EPOCH: 1 21 | VIS_PER_N_BATCH: 50 22 | VIS_ONE_PER_BATCH: True 23 | VIS_TRAIN_PER_BATCH: 100 24 | 25 | RESUME: False 26 | MODEL_INIT_PATH: ['log/pix2surf-sv-chair-coco/model/epoch_30.model'] 27 | 28 | 29 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 30 | 'core/net_bank/mlp.py', 31 | 'core/net_bank/loss.py', 32 | 'core/models/utils/model_utils.py', 33 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/coco/pix2surf_mv_plane.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_coco' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_coco_plane_train.json', 5 | 'resource/index/shapenet_coco_plane_vali.json', 6 | 'resource/index/shapenet_coco_plane_test.json'] 7 | MODES: ['train','vali'] 8 | MODEL: 'pix2surf_mv' 9 | 10 | GPU: [0,1] 11 | BATCHSIZE: 4 12 | DATALOADER_WORKERS: 6 13 | EPOCH_TOTAL: 50 14 | LR: 3e-5 15 | 16 | LOG_DIR: 'pix2surf-mv-plane-coco' 17 | LOGGER: 'logger_s1' 18 | LOGGER_SELECT: ['metric','image','model','xls'] 19 | MODEL_SAVE_PER_N_EPOCH: 5 20 | VIS_PER_N_EPOCH: 1 21 | VIS_PER_N_BATCH: 50 22 | VIS_ONE_PER_BATCH: True 23 | VIS_TRAIN_PER_BATCH: 100 24 | 25 | RESUME: False 26 | MODEL_INIT_PATH: ['log/pix2surf-sv-plane-coco/model/epoch_30.model'] 27 | 28 | 29 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 30 | 'core/net_bank/mlp.py', 31 | 'core/net_bank/loss.py', 32 | 'core/models/utils/model_utils.py', 33 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/coco/pix2surf_sv_car.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_coco' 2 | DATASET: 'neurips_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_coco_car_train.json', 5 | 'resource/index/shapenet_coco_car_vali.json', 6 | 'resource/index/shapenet_coco_car_test.json'] 7 | MODES: ['train','vali'] 8 | MODEL: 'pix2surf_sv' 9 | 10 | GPU: [0,1] 11 | BATCHSIZE: 16 12 | DATALOADER_WORKERS: 6 13 | EPOCH_TOTAL: 30 14 | LR: 1e-4 15 | 16 | LOG_DIR: 'pix2surf-sv-car-coco' 17 | LOGGER: 'logger_s1' 18 | LOGGER_SELECT: ['metric','image','model','xls'] 19 | MODEL_SAVE_PER_N_EPOCH: 5 20 | VIS_PER_N_EPOCH: 1 21 | VIS_PER_N_BATCH: 50 22 | VIS_ONE_PER_BATCH: True 23 | VIS_TRAIN_PER_BATCH: 100 24 | 25 | RESUME: False 26 | MODEL_INIT_PATH: ['log/xnocs-init-sv-car-coco/model/epoch_30.model'] 27 | 28 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 29 | 'core/net_bank/mlp.py', 30 | 'core/net_bank/loss.py', 31 | 'core/models/utils/model_utils.py', 32 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/coco/pix2surf_sv_chair.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_coco' 2 | DATASET: 'neurips_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_coco_chair_train.json', 5 | 'resource/index/shapenet_coco_chair_vali.json', 6 | 'resource/index/shapenet_coco_chair_test.json'] 7 | MODES: ['train','vali'] 8 | MODEL: 'pix2surf_sv' 9 | 10 | GPU: [0,1] 11 | BATCHSIZE: 16 12 | DATALOADER_WORKERS: 6 13 | EPOCH_TOTAL: 30 14 | LR: 3e-5 15 | 16 | LOG_DIR: 'pix2surf-sv-chair-coco' 17 | LOGGER: 'logger_s1' 18 | LOGGER_SELECT: ['metric','image','model','xls'] 19 | MODEL_SAVE_PER_N_EPOCH: 5 20 | VIS_PER_N_EPOCH: 1 21 | VIS_PER_N_BATCH: 50 22 | VIS_ONE_PER_BATCH: True 23 | VIS_TRAIN_PER_BATCH: 100 24 | 25 | RESUME: False 26 | MODEL_INIT_PATH: ['log/xnocs-init-sv-chair-coco/model/epoch_30.model'] 27 | 28 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 29 | 'core/net_bank/mlp.py', 30 | 'core/net_bank/loss.py', 31 | 'core/models/utils/model_utils.py', 32 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/coco/pix2surf_sv_plane.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_coco' 2 | DATASET: 'neurips_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_coco_plane_train.json', 5 | 'resource/index/shapenet_coco_plane_vali.json', 6 | 'resource/index/shapenet_coco_plane_test.json'] 7 | MODES: ['train','vali'] 8 | MODEL: 'pix2surf_sv' 9 | 10 | GPU: [0,1] 11 | BATCHSIZE: 16 12 | DATALOADER_WORKERS: 6 13 | EPOCH_TOTAL: 30 14 | LR: 3e-5 15 | 16 | LOG_DIR: 'pix2surf-sv-plane-coco' 17 | LOGGER: 'logger_s1' 18 | LOGGER_SELECT: ['metric','image','model','xls'] 19 | MODEL_SAVE_PER_N_EPOCH: 5 20 | VIS_PER_N_EPOCH: 1 21 | VIS_PER_N_BATCH: 50 22 | VIS_ONE_PER_BATCH: True 23 | VIS_TRAIN_PER_BATCH: 100 24 | 25 | RESUME: False 26 | MODEL_INIT_PATH: ['log/xnocs-init-sv-plane-coco/model/epoch_30.model'] 27 | 28 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 29 | 'core/net_bank/mlp.py', 30 | 'core/net_bank/loss.py', 31 | 'core/models/utils/model_utils.py', 32 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/coco/render_pix2surf_mv_car.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'pix2surf_viz' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/pix2surf_viz_car_test.json'] 5 | MODES: ['test'] 6 | MODEL: 'pix2surf_mv_render' 7 | 8 | GPU: [0] 9 | BATCHSIZE: 2 10 | DATALOADER_WORKERS: 4 11 | EPOCH_TOTAL: 1 12 | LR: 0.0 13 | 14 | LOG_DIR: 'render-pix2surf-mv-car-coco' 15 | LOGGER: 'logger_s1' 16 | LOGGER_SELECT: ['metric','image','model','xls','obj'] 17 | MODEL_SAVE_PER_N_EPOCH: 1 18 | VIS_PER_N_EPOCH: 1 19 | VIS_PER_N_BATCH: 1 20 | VIS_ONE_PER_BATCH: False 21 | 22 | RESUME: False 23 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_mv_car_coco_epoch_30.model'] 24 | 25 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 26 | 'core/net_bank/mlp.py', 27 | 'core/net_bank/loss.py', 28 | 'core/models/pix2surf_mv.py', 29 | 'core/models/utils/model_utils.py', 30 | 'core/models/utils/render_utils.py'] 31 | RENDER_MASK_TH: 0.02 32 | RENDER_IMAGE_MASK_NK: 6 -------------------------------------------------------------------------------- /config/config_files/coco/render_pix2surf_mv_chair.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'pix2surf_viz' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/pix2surf_viz_chair_test.json'] 5 | MODES: ['test'] 6 | MODEL: 'pix2surf_mv_render' 7 | 8 | GPU: [0] 9 | BATCHSIZE: 2 10 | DATALOADER_WORKERS: 4 11 | EPOCH_TOTAL: 1 12 | LR: 0.0 13 | 14 | LOG_DIR: 'render-pix2surf-mv-chair-coco' 15 | LOGGER: 'logger_s1' 16 | LOGGER_SELECT: ['metric','image','model','xls','obj'] 17 | MODEL_SAVE_PER_N_EPOCH: 1 18 | VIS_PER_N_EPOCH: 1 19 | VIS_PER_N_BATCH: 1 20 | VIS_ONE_PER_BATCH: False 21 | 22 | RESUME: False 23 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_mv_chair_coco_epoch_30.model'] 24 | 25 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 26 | 'core/net_bank/mlp.py', 27 | 'core/net_bank/loss.py', 28 | 'core/models/pix2surf_mv.py', 29 | 'core/models/utils/model_utils.py', 30 | 'core/models/utils/render_utils.py'] 31 | 32 | RENDER_MASK_TH: 0.03 33 | RENDER_IMAGE_MASK_NK: 2 -------------------------------------------------------------------------------- /config/config_files/coco/render_pix2surf_mv_plane.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'pix2surf_viz' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/pix2surf_viz_plane_test.json'] 5 | MODES: ['test'] 6 | MODEL: 'pix2surf_mv_render' 7 | 8 | GPU: [0] 9 | BATCHSIZE: 2 10 | DATALOADER_WORKERS: 4 11 | EPOCH_TOTAL: 1 12 | LR: 0.0 13 | 14 | LOG_DIR: 'render-pix2surf-mv-plane-coco' 15 | LOGGER: 'logger_s1' 16 | LOGGER_SELECT: ['metric','image','model','xls','obj'] 17 | MODEL_SAVE_PER_N_EPOCH: 1 18 | VIS_PER_N_EPOCH: 1 19 | VIS_PER_N_BATCH: 1 20 | VIS_ONE_PER_BATCH: False 21 | 22 | RESUME: False 23 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_mv_plane_coco_epoch_30.model'] 24 | 25 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 26 | 'core/net_bank/mlp.py', 27 | 'core/net_bank/loss.py', 28 | 'core/models/pix2surf_mv.py', 29 | 'core/models/utils/model_utils.py', 30 | 'core/models/utils/render_utils.py'] 31 | 32 | RENDER_MASK_TH: 0.02 33 | RENDER_IMAGE_MASK_NK: 1 34 | -------------------------------------------------------------------------------- /config/config_files/coco/render_pix2surf_sv_car.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'pix2surf_viz' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/pix2surf_viz_car_test.json'] 5 | MODES: ['test'] 6 | MODEL: 'pix2surf_sv_render' 7 | 8 | GPU: [0] 9 | BATCHSIZE: 2 10 | DATALOADER_WORKERS: 4 11 | EPOCH_TOTAL: 1 12 | LR: 0.0 13 | 14 | LOG_DIR: 'render-pix2surf-sv-car-coco' 15 | LOGGER: 'logger_s1' 16 | LOGGER_SELECT: ['metric','image','model','xls','obj'] 17 | MODEL_SAVE_PER_N_EPOCH: 1 18 | VIS_PER_N_EPOCH: 1 19 | VIS_PER_N_BATCH: 1 20 | VIS_ONE_PER_BATCH: False 21 | 22 | RESUME: False 23 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_sv_car_coco_epoch_30.model'] 24 | 25 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 26 | 'core/net_bank/mlp.py', 27 | 'core/net_bank/loss.py', 28 | 'core/models/pix2surf_sv.py', 29 | 'core/models/utils/model_utils.py', 30 | 'core/models/utils/render_utils.py'] 31 | 32 | RENDER_MASK_TH: 0.02 33 | RENDER_IMAGE_MASK_NK: 6 -------------------------------------------------------------------------------- /config/config_files/coco/render_pix2surf_sv_chair.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'pix2surf_viz' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/pix2surf_viz_chair_test.json'] 5 | MODES: ['test'] 6 | MODEL: 'pix2surf_sv_render' 7 | 8 | GPU: [0] 9 | BATCHSIZE: 2 10 | DATALOADER_WORKERS: 4 11 | EPOCH_TOTAL: 1 12 | LR: 0.0 13 | 14 | LOG_DIR: 'render-pix2surf-sv-chair-coco' 15 | LOGGER: 'logger_s1' 16 | LOGGER_SELECT: ['metric','image','model','xls','obj'] 17 | MODEL_SAVE_PER_N_EPOCH: 1 18 | VIS_PER_N_EPOCH: 1 19 | VIS_PER_N_BATCH: 1 20 | VIS_ONE_PER_BATCH: False 21 | 22 | RESUME: False 23 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_sv_chair_coco_epoch_30.model'] 24 | 25 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 26 | 'core/net_bank/mlp.py', 27 | 'core/net_bank/loss.py', 28 | 'core/models/pix2surf_sv.py', 29 | 'core/models/utils/model_utils.py', 30 | 'core/models/utils/render_utils.py'] 31 | 32 | RENDER_MASK_TH: 0.03 33 | RENDER_IMAGE_MASK_NK: 2 -------------------------------------------------------------------------------- /config/config_files/coco/render_pix2surf_sv_plane.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'pix2surf_viz' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/pix2surf_viz_plane_test.json'] 5 | MODES: ['test'] 6 | MODEL: 'pix2surf_sv_render' 7 | 8 | GPU: [0] 9 | BATCHSIZE: 2 10 | DATALOADER_WORKERS: 4 11 | EPOCH_TOTAL: 1 12 | LR: 0.0 13 | 14 | LOG_DIR: 'render-pix2surf-sv-plane-coco' 15 | LOGGER: 'logger_s1' 16 | LOGGER_SELECT: ['metric','image','model','xls','obj'] 17 | MODEL_SAVE_PER_N_EPOCH: 1 18 | VIS_PER_N_EPOCH: 1 19 | VIS_PER_N_BATCH: 1 20 | VIS_ONE_PER_BATCH: False 21 | 22 | RESUME: False 23 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_sv_plane_coco_epoch_30.model'] 24 | 25 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 26 | 'core/net_bank/mlp.py', 27 | 'core/net_bank/loss.py', 28 | 'core/models/pix2surf_sv.py', 29 | 'core/models/utils/model_utils.py', 30 | 'core/models/utils/render_utils.py'] 31 | 32 | RENDER_MASK_TH: 0.02 33 | RENDER_IMAGE_MASK_NK: 1 34 | -------------------------------------------------------------------------------- /config/config_files/coco/xnocsinit_sv_car.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_coco' 2 | DATASET: 'neurips_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_coco_car_train.json', 5 | 'resource/index/shapenet_coco_car_vali.json', 6 | 'resource/index/shapenet_coco_car_test.json'] 7 | 8 | MODES: ['train','vali'] 9 | MODEL: 'neurips_baseline' 10 | 11 | GPU: [0,1] 12 | BATCHSIZE: 16 13 | DATALOADER_WORKERS: 6 14 | EPOCH_TOTAL: 30 15 | LR: 1e-4 16 | 17 | LOG_DIR: 'xnocs-init-sv-car-coco' 18 | LOGGER: 'logger_s1' 19 | LOGGER_SELECT: ['metric','image','model'] 20 | MODEL_SAVE_PER_N_EPOCH: 5 21 | VIS_PER_N_BATCH: 50 22 | VIS_ONE_PER_BATCH: True 23 | VIS_TRAIN_PER_BATCH: 100 24 | 25 | BACKUP_FILES: ['core/net_bank/xnocs_segnet.py', 26 | 'core/net_bank/loss.py'] -------------------------------------------------------------------------------- /config/config_files/coco/xnocsinit_sv_chair.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_coco' 2 | DATASET: 'neurips_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_coco_chair_train.json', 5 | 'resource/index/shapenet_coco_chair_vali.json', 6 | 'resource/index/shapenet_coco_chair_test.json'] 7 | 8 | MODES: ['train','vali'] 9 | MODEL: 'neurips_baseline' 10 | 11 | GPU: [0,1] 12 | BATCHSIZE: 16 13 | DATALOADER_WORKERS: 6 14 | EPOCH_TOTAL: 30 15 | LR: 1e-4 16 | 17 | LOG_DIR: 'xnocs-init-sv-chair-coco' 18 | LOGGER: 'logger_s1' 19 | LOGGER_SELECT: ['metric','image','model'] 20 | MODEL_SAVE_PER_N_EPOCH: 5 21 | VIS_PER_N_BATCH: 50 22 | VIS_ONE_PER_BATCH: True 23 | VIS_TRAIN_PER_BATCH: 100 24 | 25 | BACKUP_FILES: ['core/net_bank/xnocs_segnet.py', 26 | 'core/net_bank/loss.py'] -------------------------------------------------------------------------------- /config/config_files/coco/xnocsinit_sv_plane.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_coco' 2 | DATASET: 'neurips_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_coco_plane_train.json', 5 | 'resource/index/shapenet_coco_plane_vali.json', 6 | 'resource/index/shapenet_coco_plane_test.json'] 7 | 8 | MODES: ['train','vali'] 9 | MODEL: 'neurips_baseline' 10 | 11 | GPU: [0,1] 12 | BATCHSIZE: 16 13 | DATALOADER_WORKERS: 6 14 | EPOCH_TOTAL: 30 15 | LR: 1e-4 16 | 17 | LOG_DIR: 'xnocs-init-sv-plane-coco' 18 | LOGGER: 'logger_s1' 19 | LOGGER_SELECT: ['metric','image','model'] 20 | MODEL_SAVE_PER_N_EPOCH: 5 21 | VIS_PER_N_BATCH: 50 22 | VIS_ONE_PER_BATCH: True 23 | VIS_TRAIN_PER_BATCH: 100 24 | 25 | BACKUP_FILES: ['core/net_bank/xnocs_segnet.py', 26 | 'core/net_bank/loss.py'] -------------------------------------------------------------------------------- /config/config_files/plain/eval_pix2surf_mv_car.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_plain_car_train.json', 5 | 'resource/index/shapenet_plain_car_test.json'] 6 | MODES: ['test'] 7 | MODEL: 'pix2surf_mv_mveval' 8 | 9 | GPU: [0] 10 | BATCHSIZE: 2 11 | DATALOADER_WORKERS: 4 12 | EPOCH_TOTAL: 1 13 | LR: 0.0 14 | 15 | LOG_DIR: 'eval-pix2surf-mv-car-plain' 16 | LOGGER: 'logger_s1' 17 | LOGGER_SELECT: ['metric','image','model','xls'] 18 | MODEL_SAVE_PER_N_EPOCH: 1 19 | VIS_PER_N_EPOCH: 1 20 | VIS_PER_N_BATCH: 10 21 | VIS_ONE_PER_BATCH: True 22 | 23 | RESUME: False 24 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_mv_car_plain_epoch_21.model'] 25 | 26 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 27 | 'core/net_bank/mlp.py', 28 | 'core/net_bank/loss.py', 29 | 'core/evaluation.py', 30 | 'core/models/pix2surf_mv.py', 31 | 'core/models/utils/model_utils.py', 32 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/plain/eval_pix2surf_mv_chair.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_plain_chair_train.json', 5 | 'resource/index/shapenet_plain_chair_test.json'] 6 | MODES: ['test'] 7 | MODEL: 'pix2surf_mv_mveval' 8 | 9 | GPU: [0] 10 | BATCHSIZE: 2 11 | DATALOADER_WORKERS: 4 12 | EPOCH_TOTAL: 1 13 | LR: 0.0 14 | 15 | LOG_DIR: 'eval-pix2surf-mv-chair-plain' 16 | LOGGER: 'logger_s1' 17 | LOGGER_SELECT: ['metric','image','model','xls'] 18 | MODEL_SAVE_PER_N_EPOCH: 1 19 | VIS_PER_N_EPOCH: 1 20 | VIS_PER_N_BATCH: 10 21 | VIS_ONE_PER_BATCH: True 22 | 23 | RESUME: False 24 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_mv_chair_plain_epoch_21.model'] 25 | 26 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 27 | 'core/net_bank/mlp.py', 28 | 'core/net_bank/loss.py', 29 | 'core/evaluation.py', 30 | 'core/models/pix2surf_mv.py', 31 | 'core/models/utils/model_utils.py', 32 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/plain/eval_pix2surf_mv_plane.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_plain_plane_train.json', 5 | 'resource/index/shapenet_plain_plane_test.json'] 6 | MODES: ['test'] 7 | MODEL: 'pix2surf_mv_mveval' 8 | 9 | GPU: [0] 10 | BATCHSIZE: 2 11 | DATALOADER_WORKERS: 4 12 | EPOCH_TOTAL: 1 13 | LR: 0.0 14 | 15 | LOG_DIR: 'eval-pix2surf-mv-plane-plain' 16 | LOGGER: 'logger_s1' 17 | LOGGER_SELECT: ['metric','image','model','xls'] 18 | MODEL_SAVE_PER_N_EPOCH: 1 19 | VIS_PER_N_EPOCH: 1 20 | VIS_PER_N_BATCH: 10 21 | VIS_ONE_PER_BATCH: True 22 | 23 | RESUME: False 24 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_mv_plane_plain_epoch_21.model'] 25 | 26 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 27 | 'core/net_bank/mlp.py', 28 | 'core/net_bank/loss.py', 29 | 'core/evaluation.py', 30 | 'core/models/pix2surf_mv.py', 31 | 'core/models/utils/model_utils.py', 32 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/plain/eval_pix2surf_sv_car.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_plain_car_train.json', 5 | 'resource/index/shapenet_plain_car_test.json'] 6 | MODES: ['test'] 7 | MODEL: 'pix2surf_sv_mveval' 8 | 9 | GPU: [0] 10 | BATCHSIZE: 2 11 | DATALOADER_WORKERS: 4 12 | EPOCH_TOTAL: 1 13 | LR: 0.0 14 | 15 | LOG_DIR: 'eval-pix2surf-sv-car-plain' 16 | LOGGER: 'logger_s1' 17 | LOGGER_SELECT: ['metric','image','model','xls'] 18 | MODEL_SAVE_PER_N_EPOCH: 1 19 | VIS_PER_N_EPOCH: 1 20 | VIS_PER_N_BATCH: 10 21 | VIS_ONE_PER_BATCH: True 22 | 23 | RESUME: False 24 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_sv_car_plain_epoch_30.model'] 25 | 26 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 27 | 'core/net_bank/mlp.py', 28 | 'core/net_bank/loss.py', 29 | 'core/evaluation.py', 30 | 'core/models/pix2surf_sv.py', 31 | 'core/models/utils/model_utils.py', 32 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/plain/eval_pix2surf_sv_chair.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_plain_chair_train.json', 5 | 'resource/index/shapenet_plain_chair_test.json'] 6 | MODES: ['test'] 7 | MODEL: 'pix2surf_sv_mveval' 8 | 9 | GPU: [0] 10 | BATCHSIZE: 2 11 | DATALOADER_WORKERS: 4 12 | EPOCH_TOTAL: 1 13 | LR: 0.0 14 | 15 | LOG_DIR: 'eval-pix2surf-sv-chair-plain' 16 | LOGGER: 'logger_s1' 17 | LOGGER_SELECT: ['metric','image','model','xls'] 18 | MODEL_SAVE_PER_N_EPOCH: 1 19 | VIS_PER_N_EPOCH: 1 20 | VIS_PER_N_BATCH: 10 21 | VIS_ONE_PER_BATCH: True 22 | 23 | RESUME: False 24 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_sv_chair_plain_epoch_30.model'] 25 | 26 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 27 | 'core/net_bank/mlp.py', 28 | 'core/net_bank/loss.py', 29 | 'core/evaluation.py', 30 | 'core/models/pix2surf_sv.py', 31 | 'core/models/utils/model_utils.py', 32 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/plain/eval_pix2surf_sv_plane.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_plain_plane_train.json', 5 | 'resource/index/shapenet_plain_plane_test.json'] 6 | MODES: ['test'] 7 | MODEL: 'pix2surf_sv_mveval' 8 | 9 | GPU: [0] 10 | BATCHSIZE: 2 11 | DATALOADER_WORKERS: 4 12 | EPOCH_TOTAL: 1 13 | LR: 0.0 14 | 15 | LOG_DIR: 'eval-pix2surf-sv-plane-plain' 16 | LOGGER: 'logger_s1' 17 | LOGGER_SELECT: ['metric','image','model','xls'] 18 | MODEL_SAVE_PER_N_EPOCH: 1 19 | VIS_PER_N_EPOCH: 1 20 | VIS_PER_N_BATCH: 10 21 | VIS_ONE_PER_BATCH: True 22 | 23 | RESUME: False 24 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_sv_plane_plain_epoch_30.model'] 25 | 26 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 27 | 'core/net_bank/mlp.py', 28 | 'core/net_bank/loss.py', 29 | 'core/evaluation.py', 30 | 'core/models/pix2surf_sv.py', 31 | 'core/models/utils/model_utils.py', 32 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/plain/pix2surf_mv_car.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'mv5_correspondence_dataset' 3 | ADD_BACKGROUND_NOISE: True 4 | PROPORTION: 1.0 5 | DATASET_INDEX: ['resource/index/shapenet_plain_car_train.json', 6 | 'resource/index/shapenet_plain_car_test.json'] 7 | MODES: ['train'] 8 | MODEL: 'pix2surf_mv' 9 | 10 | GPU: [0,1] 11 | BATCHSIZE: 4 12 | DATALOADER_WORKERS: 6 13 | EPOCH_TOTAL: 30 14 | LR: 1e-4 15 | 16 | LOG_DIR: 'pix2surf-mv-car-plain' 17 | LOGGER: 'logger_s1' 18 | LOGGER_SELECT: ['metric','image','model','xls'] 19 | MODEL_SAVE_PER_N_EPOCH: 3 20 | VIS_PER_N_EPOCH: 1 21 | VIS_PER_N_BATCH: 50 22 | VIS_ONE_PER_BATCH: True 23 | VIS_TRAIN_PER_BATCH: 100 24 | 25 | RESUME: False 26 | MODEL_INIT_PATH: ['log/pix2surf-sv-car-plain/model/epoch_30.model'] 27 | 28 | 29 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 30 | 'core/net_bank/mlp.py', 31 | 'core/net_bank/loss.py', 32 | 'core/models/utils/model_utils.py', 33 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/plain/pix2surf_mv_chair.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'mv5_correspondence_dataset' 3 | ADD_BACKGROUND_NOISE: True 4 | PROPORTION: 1.0 5 | DATASET_INDEX: ['resource/index/shapenet_plain_chair_train.json', 6 | 'resource/index/shapenet_plain_chair_test.json'] 7 | MODES: ['train'] 8 | MODEL: 'pix2surf_mv' 9 | 10 | GPU: [0,1] 11 | BATCHSIZE: 4 12 | DATALOADER_WORKERS: 6 13 | EPOCH_TOTAL: 30 14 | LR: 3e-5 15 | 16 | LOG_DIR: 'pix2surf-mv-chair-plain' 17 | LOGGER: 'logger_s1' 18 | LOGGER_SELECT: ['metric','image','model','xls'] 19 | MODEL_SAVE_PER_N_EPOCH: 3 20 | VIS_PER_N_EPOCH: 1 21 | VIS_PER_N_BATCH: 50 22 | VIS_ONE_PER_BATCH: True 23 | VIS_TRAIN_PER_BATCH: 100 24 | 25 | RESUME: False 26 | MODEL_INIT_PATH: ['log/pix2surf-sv-chair-plain/model/epoch_30.model'] 27 | 28 | 29 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 30 | 'core/net_bank/mlp.py', 31 | 'core/net_bank/loss.py', 32 | 'core/models/utils/model_utils.py', 33 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/plain/pix2surf_mv_plane.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'mv5_correspondence_dataset' 3 | ADD_BACKGROUND_NOISE: True 4 | PROPORTION: 1.0 5 | DATASET_INDEX: ['resource/index/shapenet_plain_plane_train.json', 6 | 'resource/index/shapenet_plain_plane_test.json'] 7 | MODES: ['train'] 8 | MODEL: 'pix2surf_mv' 9 | 10 | GPU: [0,1] 11 | BATCHSIZE: 4 12 | DATALOADER_WORKERS: 6 13 | EPOCH_TOTAL: 30 14 | LR: 3e-5 15 | 16 | LOG_DIR: 'pix2surf-mv-plane-plain' 17 | LOGGER: 'logger_s1' 18 | LOGGER_SELECT: ['metric','image','model','xls'] 19 | MODEL_SAVE_PER_N_EPOCH: 3 20 | VIS_PER_N_EPOCH: 1 21 | VIS_PER_N_BATCH: 50 22 | VIS_ONE_PER_BATCH: True 23 | VIS_TRAIN_PER_BATCH: 100 24 | 25 | RESUME: False 26 | MODEL_INIT_PATH: ['log/pix2surf-sv-plane-plain/model/epoch_30.model'] 27 | 28 | 29 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 30 | 'core/net_bank/mlp.py', 31 | 'core/net_bank/loss.py', 32 | 'core/models/utils/model_utils.py', 33 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/plain/pix2surf_sv_car.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'neurips_dataset' 3 | ADD_BACKGROUND_NOISE: True 4 | PROPORTION: 1.0 5 | DATASET_INDEX: ['resource/index/shapenet_plain_car_train.json', 6 | 'resource/index/shapenet_plain_car_test.json'] 7 | MODES: ['train'] 8 | MODEL: 'pix2surf_sv' 9 | 10 | GPU: [0,1] 11 | BATCHSIZE: 16 12 | DATALOADER_WORKERS: 6 13 | EPOCH_TOTAL: 30 14 | LR: 1e-4 15 | 16 | LOG_DIR: 'pix2surf-sv-car-plain' 17 | LOGGER: 'logger_s1' 18 | LOGGER_SELECT: ['metric','image','model','xls'] 19 | MODEL_SAVE_PER_N_EPOCH: 5 20 | VIS_PER_N_EPOCH: 1 21 | VIS_PER_N_BATCH: 50 22 | VIS_ONE_PER_BATCH: True 23 | VIS_TRAIN_PER_BATCH: 100 24 | 25 | RESUME: False 26 | MODEL_INIT_PATH: ['log/xnocs-init-sv-car-plain/model/epoch_30.model'] 27 | 28 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 29 | 'core/net_bank/mlp.py', 30 | 'core/net_bank/loss.py', 31 | 'core/models/utils/model_utils.py', 32 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/plain/pix2surf_sv_chair.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'neurips_dataset' 3 | ADD_BACKGROUND_NOISE: True 4 | PROPORTION: 1.0 5 | DATASET_INDEX: ['resource/index/shapenet_plain_chair_train.json', 6 | 'resource/index/shapenet_plain_chair_test.json'] 7 | MODES: ['train'] 8 | MODEL: 'pix2surf_sv' 9 | 10 | GPU: [0,1] 11 | BATCHSIZE: 16 12 | DATALOADER_WORKERS: 6 13 | EPOCH_TOTAL: 30 14 | LR: 3e-5 15 | 16 | LOG_DIR: 'pix2surf-sv-chair-plain' 17 | LOGGER: 'logger_s1' 18 | LOGGER_SELECT: ['metric','image','model','xls'] 19 | MODEL_SAVE_PER_N_EPOCH: 5 20 | VIS_PER_N_EPOCH: 1 21 | VIS_PER_N_BATCH: 50 22 | VIS_ONE_PER_BATCH: True 23 | VIS_TRAIN_PER_BATCH: 100 24 | 25 | RESUME: False 26 | MODEL_INIT_PATH: ['log/xnocs-init-sv-chair-plain/model/epoch_30.model'] 27 | 28 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 29 | 'core/net_bank/mlp.py', 30 | 'core/net_bank/loss.py', 31 | 'core/models/utils/model_utils.py', 32 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/plain/pix2surf_sv_plane.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'neurips_dataset' 3 | ADD_BACKGROUND_NOISE: True 4 | PROPORTION: 1.0 5 | DATASET_INDEX: ['resource/index/shapenet_plain_plane_train.json', 6 | 'resource/index/shapenet_plain_plane_test.json'] 7 | MODES: ['train'] 8 | MODEL: 'pix2surf_sv' 9 | 10 | GPU: [0,1] 11 | BATCHSIZE: 16 12 | DATALOADER_WORKERS: 6 13 | EPOCH_TOTAL: 30 14 | LR: 3e-5 15 | 16 | LOG_DIR: 'pix2surf-sv-plane-plain' 17 | LOGGER: 'logger_s1' 18 | LOGGER_SELECT: ['metric','image','model','xls'] 19 | MODEL_SAVE_PER_N_EPOCH: 5 20 | VIS_PER_N_EPOCH: 1 21 | VIS_PER_N_BATCH: 50 22 | VIS_ONE_PER_BATCH: True 23 | VIS_TRAIN_PER_BATCH: 100 24 | 25 | RESUME: False 26 | MODEL_INIT_PATH: ['log/xnocs-init-sv-plane-plain/model/epoch_30.model'] 27 | 28 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 29 | 'core/net_bank/mlp.py', 30 | 'core/net_bank/loss.py', 31 | 'core/models/utils/model_utils.py', 32 | 'core/models/utils/render_utils.py'] -------------------------------------------------------------------------------- /config/config_files/plain/render_pix2surf_mv_car.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_plain_car_mini_test.json'] 5 | MODES: ['test'] 6 | MODEL: 'pix2surf_mv_render' 7 | 8 | GPU: [0] 9 | BATCHSIZE: 2 10 | DATALOADER_WORKERS: 4 11 | EPOCH_TOTAL: 1 12 | LR: 0.0 13 | 14 | LOG_DIR: 'render-pix2surf-mv-car-plain' 15 | LOGGER: 'logger_s1' 16 | LOGGER_SELECT: ['metric','image','model','xls','obj'] 17 | MODEL_SAVE_PER_N_EPOCH: 1 18 | VIS_PER_N_EPOCH: 1 19 | VIS_PER_N_BATCH: 1 20 | VIS_ONE_PER_BATCH: False 21 | 22 | RESUME: False 23 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_mv_car_plain_epoch_21.model'] 24 | 25 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 26 | 'core/net_bank/mlp.py', 27 | 'core/net_bank/loss.py', 28 | 'core/models/pix2surf_mv.py', 29 | 'core/models/utils/model_utils.py', 30 | 'core/models/utils/render_utils.py'] 31 | RENDER_MASK_TH: 0.02 32 | RENDER_IMAGE_MASK_NK: 6 -------------------------------------------------------------------------------- /config/config_files/plain/render_pix2surf_mv_chair.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_plain_chair_mini_test.json'] 5 | MODES: ['test'] 6 | MODEL: 'pix2surf_mv_render' 7 | 8 | GPU: [0] 9 | BATCHSIZE: 2 10 | DATALOADER_WORKERS: 4 11 | EPOCH_TOTAL: 1 12 | LR: 0.0 13 | 14 | LOG_DIR: 'render-pix2surf-mv-chair-plain' 15 | LOGGER: 'logger_s1' 16 | LOGGER_SELECT: ['metric','image','model','xls','obj'] 17 | MODEL_SAVE_PER_N_EPOCH: 1 18 | VIS_PER_N_EPOCH: 1 19 | VIS_PER_N_BATCH: 1 20 | VIS_ONE_PER_BATCH: False 21 | 22 | RESUME: False 23 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_mv_chair_plain_epoch_21.model'] 24 | 25 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 26 | 'core/net_bank/mlp.py', 27 | 'core/net_bank/loss.py', 28 | 'core/models/pix2surf_mv.py', 29 | 'core/models/utils/model_utils.py', 30 | 'core/models/utils/render_utils.py'] 31 | 32 | RENDER_MASK_TH: 0.03 33 | RENDER_IMAGE_MASK_NK: 2 -------------------------------------------------------------------------------- /config/config_files/plain/render_pix2surf_mv_plane.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_plain_plane_mini_test.json'] 5 | MODES: ['test'] 6 | MODEL: 'pix2surf_mv_render' 7 | 8 | GPU: [0] 9 | BATCHSIZE: 2 10 | DATALOADER_WORKERS: 4 11 | EPOCH_TOTAL: 1 12 | LR: 0.0 13 | 14 | LOG_DIR: 'render-pix2surf-mv-plane-plain' 15 | LOGGER: 'logger_s1' 16 | LOGGER_SELECT: ['metric','image','model','xls','obj'] 17 | MODEL_SAVE_PER_N_EPOCH: 1 18 | VIS_PER_N_EPOCH: 1 19 | VIS_PER_N_BATCH: 1 20 | VIS_ONE_PER_BATCH: False 21 | 22 | RESUME: False 23 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_mv_plane_plain_epoch_21.model'] 24 | 25 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 26 | 'core/net_bank/mlp.py', 27 | 'core/net_bank/loss.py', 28 | 'core/models/pix2surf_mv.py', 29 | 'core/models/utils/model_utils.py', 30 | 'core/models/utils/render_utils.py'] 31 | 32 | RENDER_MASK_TH: 0.02 33 | RENDER_IMAGE_MASK_NK: 4 -------------------------------------------------------------------------------- /config/config_files/plain/render_pix2surf_sv_car.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_plain_car_mini_test.json'] 5 | MODES: ['test'] 6 | MODEL: 'pix2surf_sv_render' 7 | 8 | GPU: [0] 9 | BATCHSIZE: 2 10 | DATALOADER_WORKERS: 4 11 | EPOCH_TOTAL: 1 12 | LR: 0.0 13 | 14 | LOG_DIR: 'render-pix2surf-sv-car-plain' 15 | LOGGER: 'logger_s1' 16 | LOGGER_SELECT: ['metric','image','model','xls','obj'] 17 | MODEL_SAVE_PER_N_EPOCH: 1 18 | VIS_PER_N_EPOCH: 1 19 | VIS_PER_N_BATCH: 1 20 | VIS_ONE_PER_BATCH: False 21 | 22 | RESUME: False 23 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_sv_car_plain_epoch_30.model'] 24 | 25 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 26 | 'core/net_bank/mlp.py', 27 | 'core/net_bank/loss.py', 28 | 'core/models/pix2surf_sv.py', 29 | 'core/models/utils/model_utils.py', 30 | 'core/models/utils/render_utils.py'] 31 | 32 | RENDER_MASK_TH: 0.02 33 | RENDER_IMAGE_MASK_NK: 6 -------------------------------------------------------------------------------- /config/config_files/plain/render_pix2surf_sv_chair.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_plain_chair_mini_test.json'] 5 | MODES: ['test'] 6 | MODEL: 'pix2surf_sv_render' 7 | 8 | GPU: [0] 9 | BATCHSIZE: 2 10 | DATALOADER_WORKERS: 4 11 | EPOCH_TOTAL: 1 12 | LR: 0.0 13 | 14 | LOG_DIR: 'render-pix2surf-sv-chair-plain' 15 | LOGGER: 'logger_s1' 16 | LOGGER_SELECT: ['metric','image','model','xls','obj'] 17 | MODEL_SAVE_PER_N_EPOCH: 1 18 | VIS_PER_N_EPOCH: 1 19 | VIS_PER_N_BATCH: 1 20 | VIS_ONE_PER_BATCH: False 21 | 22 | RESUME: False 23 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_sv_chair_plain_epoch_30.model'] 24 | 25 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 26 | 'core/net_bank/mlp.py', 27 | 'core/net_bank/loss.py', 28 | 'core/models/pix2surf_sv.py', 29 | 'core/models/utils/model_utils.py', 30 | 'core/models/utils/render_utils.py'] 31 | 32 | RENDER_MASK_TH: 0.03 33 | RENDER_IMAGE_MASK_NK: 2 -------------------------------------------------------------------------------- /config/config_files/plain/render_pix2surf_sv_plane.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_plain_plane_mini_test.json'] 5 | MODES: ['test'] 6 | MODEL: 'pix2surf_sv_render' 7 | 8 | GPU: [0] 9 | BATCHSIZE: 2 10 | DATALOADER_WORKERS: 4 11 | EPOCH_TOTAL: 1 12 | LR: 0.0 13 | 14 | LOG_DIR: 'render-pix2surf-sv-plane-plain' 15 | LOGGER: 'logger_s1' 16 | LOGGER_SELECT: ['metric','image','model','xls','obj'] 17 | MODEL_SAVE_PER_N_EPOCH: 1 18 | VIS_PER_N_EPOCH: 1 19 | VIS_PER_N_BATCH: 1 20 | VIS_ONE_PER_BATCH: False 21 | 22 | RESUME: False 23 | MODEL_INIT_PATH: ['resource/weight/eccv_pix2surf_sv_plane_plain_epoch_30.model'] 24 | 25 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 26 | 'core/net_bank/mlp.py', 27 | 'core/net_bank/loss.py', 28 | 'core/models/pix2surf_sv.py', 29 | 'core/models/utils/model_utils.py', 30 | 'core/models/utils/render_utils.py'] 31 | 32 | RENDER_MASK_TH: 0.02 33 | RENDER_IMAGE_MASK_NK: 4 -------------------------------------------------------------------------------- /config/config_files/plain/xnocsinit_sv_car.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'neurips_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_plain_car_train.json', 5 | 'resource/index/shapenet_plain_car_test.json'] 6 | 7 | MODES: ['train'] 8 | MODEL: 'neurips_baseline' 9 | 10 | GPU: [0,1] 11 | BATCHSIZE: 16 12 | DATALOADER_WORKERS: 6 13 | EPOCH_TOTAL: 30 14 | LR: 1e-4 15 | 16 | LOG_DIR: 'xnocs-init-sv-car-plain' 17 | LOGGER: 'logger_s1' 18 | LOGGER_SELECT: ['metric','image','model'] 19 | MODEL_SAVE_PER_N_EPOCH: 5 20 | VIS_PER_N_BATCH: 50 21 | VIS_ONE_PER_BATCH: True 22 | VIS_TRAIN_PER_BATCH: 100 23 | 24 | BACKUP_FILES: ['core/net_bank/xnocs_segnet.py', 25 | 'core/net_bank/loss.py'] -------------------------------------------------------------------------------- /config/config_files/plain/xnocsinit_sv_chair.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'neurips_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_plain_chair_train.json', 5 | 'resource/index/shapenet_plain_chair_test.json'] 6 | 7 | MODES: ['train'] 8 | MODEL: 'neurips_baseline' 9 | 10 | GPU: [0,1] 11 | BATCHSIZE: 16 12 | DATALOADER_WORKERS: 6 13 | EPOCH_TOTAL: 30 14 | LR: 1e-4 15 | 16 | LOG_DIR: 'xnocs-init-sv-chair-plain' 17 | LOGGER: 'logger_s1' 18 | LOGGER_SELECT: ['metric','image','model'] 19 | MODEL_SAVE_PER_N_EPOCH: 5 20 | VIS_PER_N_BATCH: 50 21 | VIS_ONE_PER_BATCH: True 22 | VIS_TRAIN_PER_BATCH: 100 23 | 24 | BACKUP_FILES: ['core/net_bank/xnocs_segnet.py', 25 | 'core/net_bank/loss.py'] -------------------------------------------------------------------------------- /config/config_files/plain/xnocsinit_sv_plane.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'shapenet_plain' 2 | DATASET: 'neurips_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/shapenet_plain_plane_train.json', 5 | 'resource/index/shapenet_plain_plane_test.json'] 6 | 7 | MODES: ['train'] 8 | MODEL: 'neurips_baseline' 9 | 10 | GPU: [0,1] 11 | BATCHSIZE: 16 12 | DATALOADER_WORKERS: 6 13 | EPOCH_TOTAL: 30 14 | LR: 1e-4 15 | 16 | LOG_DIR: 'xnocs-init-sv-plane-plain' 17 | LOGGER: 'logger_s1' 18 | LOGGER_SELECT: ['metric','image','model'] 19 | MODEL_SAVE_PER_N_EPOCH: 5 20 | VIS_PER_N_BATCH: 50 21 | VIS_ONE_PER_BATCH: True 22 | VIS_TRAIN_PER_BATCH: 100 23 | 24 | BACKUP_FILES: ['core/net_bank/xnocs_segnet.py', 25 | 'core/net_bank/loss.py'] -------------------------------------------------------------------------------- /config/config_files/viz/render_pix2surf_mv_car.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'pix2surf_viz' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/pix2surf_viz_car_test.json'] 5 | MODES: ['test'] 6 | MODEL: 'pix2surf_mv_render' 7 | 8 | GPU: [0] 9 | BATCHSIZE: 2 10 | DATALOADER_WORKERS: 4 11 | EPOCH_TOTAL: 1 12 | LR: 0.0 13 | 14 | LOG_DIR: 'render-pix2surf-mv-car-viz' 15 | LOGGER: 'logger_s1' 16 | LOGGER_SELECT: ['metric','image','model','xls','obj'] 17 | MODEL_SAVE_PER_N_EPOCH: 1 18 | VIS_PER_N_EPOCH: 1 19 | VIS_PER_N_BATCH: 1 20 | VIS_ONE_PER_BATCH: False 21 | 22 | RESUME: False 23 | MODEL_INIT_PATH: ['resource/weight/viz_car.model'] 24 | 25 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 26 | 'core/net_bank/mlp.py', 27 | 'core/net_bank/loss.py', 28 | 'core/models/pix2surf_mv.py', 29 | 'core/models/utils/model_utils.py', 30 | 'core/models/utils/render_utils.py'] 31 | RENDER_MASK_TH: 0.02 32 | RENDER_IMAGE_MASK_NK: 6 -------------------------------------------------------------------------------- /config/config_files/viz/render_pix2surf_mv_chair.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'pix2surf_viz' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/pix2surf_viz_chair_test.json'] 5 | MODES: ['test'] 6 | MODEL: 'pix2surf_mv_render' 7 | 8 | GPU: [0] 9 | BATCHSIZE: 2 10 | DATALOADER_WORKERS: 4 11 | EPOCH_TOTAL: 1 12 | LR: 0.0 13 | 14 | LOG_DIR: 'render-pix2surf-mv-chair-viz' 15 | LOGGER: 'logger_s1' 16 | LOGGER_SELECT: ['metric','image','model','xls','obj'] 17 | MODEL_SAVE_PER_N_EPOCH: 1 18 | VIS_PER_N_EPOCH: 1 19 | VIS_PER_N_BATCH: 1 20 | VIS_ONE_PER_BATCH: False 21 | 22 | RESUME: False 23 | MODEL_INIT_PATH: ['resource/weight/viz_chair.model'] 24 | 25 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 26 | 'core/net_bank/mlp.py', 27 | 'core/net_bank/loss.py', 28 | 'core/models/pix2surf_mv.py', 29 | 'core/models/utils/model_utils.py', 30 | 'core/models/utils/render_utils.py'] 31 | 32 | RENDER_MASK_TH: 0.03 33 | RENDER_IMAGE_MASK_NK: 2 -------------------------------------------------------------------------------- /config/config_files/viz/render_pix2surf_mv_plane.yaml: -------------------------------------------------------------------------------- 1 | DATASET_ROOT: 'pix2surf_viz' 2 | DATASET: 'mv5_correspondence_dataset' 3 | PROPORTION: 1.0 4 | DATASET_INDEX: ['resource/index/pix2surf_viz_plane_test.json'] 5 | MODES: ['test'] 6 | MODEL: 'pix2surf_mv_render' 7 | 8 | GPU: [0] 9 | BATCHSIZE: 2 10 | DATALOADER_WORKERS: 4 11 | EPOCH_TOTAL: 1 12 | LR: 0.0 13 | 14 | LOG_DIR: 'render-pix2surf-mv-plane-viz' 15 | LOGGER: 'logger_s1' 16 | LOGGER_SELECT: ['metric','image','model','xls','obj'] 17 | MODEL_SAVE_PER_N_EPOCH: 1 18 | VIS_PER_N_EPOCH: 1 19 | VIS_PER_N_BATCH: 1 20 | VIS_ONE_PER_BATCH: False 21 | 22 | RESUME: False 23 | MODEL_INIT_PATH: ['resource/weight/viz_plane.model'] 24 | 25 | BACKUP_FILES: ['core/net_bank/pix2surf_cnn.py', 26 | 'core/net_bank/mlp.py', 27 | 'core/net_bank/loss.py', 28 | 'core/models/pix2surf_mv.py', 29 | 'core/models/utils/model_utils.py', 30 | 'core/models/utils/render_utils.py'] 31 | 32 | RENDER_MASK_TH: 0.02 33 | RENDER_IMAGE_MASK_NK: 1 34 | -------------------------------------------------------------------------------- /config/default.py: -------------------------------------------------------------------------------- 1 | """ 2 | Default configuration 3 | """ 4 | from yacs.config import CfgNode as CN 5 | import os 6 | 7 | 8 | def get_default_cfg(): 9 | cfg = CN() 10 | # dataset 11 | cfg.DATASET = '' 12 | cfg.DATASET_ROOT = '' 13 | cfg.DATASET_CATES = [] 14 | cfg.DATASET_INDEX = [] 15 | cfg.PROPORTION = 1.0 # Ratio of first K samples in the dataset list 16 | cfg.ADD_BACKGROUND_NOISE = False 17 | cfg.ROOT_DIR = os.getcwd() 18 | cfg.CONFIG_FILE = 'None' 19 | cfg.MODES = ['train', 'test'] 20 | cfg.GPU = [0] 21 | 22 | cfg.BATCHSIZE = 2 23 | cfg.DATALOADER_WORKERS = 2 24 | 25 | cfg.MODEL = 'None' 26 | cfg.MODEL_INIT_PATH = ['None'] 27 | cfg.LR = 0.001 28 | cfg.EPOCH_TOTAL = 1 29 | 30 | cfg.ADAM_BETA1 = 0.5 31 | cfg.ADAM_BETA2 = 0.9 32 | cfg.OPTI_DECAY_RATE = 1.0 33 | cfg.OPTI_DECAY_INTERVAL = 1000 34 | cfg.OPTI_DECAY_MIN = 0.00001 35 | 36 | cfg.LOG_DIR = 'debug' # this param is the name under $ProjectRoot/log/ 37 | 38 | cfg.RESUME = False # If true, check whether the log dir exists, if not, just start a new one 39 | cfg.RESUME_EPOCH_ID = 0 40 | 41 | cfg.LOGGER = 'logger_s1' 42 | cfg.LOGGER_SELECT = ['metric'] 43 | 44 | cfg.MODEL_SAVE_PER_N_EPOCH = 5 45 | 46 | cfg.VIS_PER_N_EPOCH = 1 47 | cfg.VIS_ONE_PER_BATCH = True 48 | # visualization config for non-training phase 49 | cfg.VIS_PER_N_BATCH = 1 50 | # visualization config for trianing phase 51 | cfg.VIS_TRAIN_PER_BATCH = 20 52 | 53 | cfg.BACKUP_FILES = [] 54 | 55 | # for postprocessing (render) 56 | cfg.RENDER_MASK_TH = 0.02 # distance th for valid neighborhood point in 3D 57 | cfg.RENDER_IMAGE_MASK_NK = 4 # image mask optimization n nearest neighbor 58 | 59 | return cfg 60 | -------------------------------------------------------------------------------- /config/startup.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import shutil 4 | import platform 5 | 6 | 7 | def startup(cfg, interactive=True): 8 | """ 9 | preparation before start 10 | """ 11 | # check required logger are selected 12 | assert 'model' in cfg.LOGGER_SELECT 13 | assert 'metric' in cfg.LOGGER_SELECT 14 | # check config 15 | os_cate = platform.system().lower() 16 | assert os_cate in ['linux'] 17 | print("=" * shutil.get_terminal_size()[0]) 18 | print("Please check the configuration") 19 | print('-' * shutil.get_terminal_size()[0]) 20 | print(cfg) 21 | print('-' * shutil.get_terminal_size()[0]) 22 | print('y/n?', end='') 23 | if interactive: 24 | need_input = True 25 | while need_input: 26 | response = input().lower() 27 | if response in ('y', 'n'): 28 | need_input = False 29 | if response == 'n': 30 | sys.exit() 31 | else: 32 | print("Warning, NO INTERACTIVE CONFIRMATION!") 33 | print("=" * shutil.get_terminal_size()[0]) 34 | 35 | # Set visible GPUs 36 | os.environ["CUDA_VISIBLE_DEVICES"] = str(cfg.GPU)[1:-1] 37 | print("Set GPU: " + str(cfg.GPU)[1:-1] + '...') 38 | 39 | # check log dir 40 | abs_log_dir = os.path.join(cfg.ROOT_DIR, 'log', cfg.LOG_DIR) 41 | if cfg.RESUME: 42 | # if resume, don't remove the old log 43 | if not os.path.exists(abs_log_dir): 44 | print("Warning: Need resume but the log dir: " + abs_log_dir + " doesn't exists, create new one", end='') 45 | os.makedirs(abs_log_dir) 46 | else: 47 | if os.path.exists(abs_log_dir): 48 | print("Warning! Not resume but log dir: " + abs_log_dir + ' already existed. Remove the old log dir? y/n', 49 | end='') 50 | if interactive: 51 | need_input = True 52 | while need_input: 53 | response = input().lower() 54 | if response in ('y', 'n'): 55 | need_input = False 56 | if response == 'n': 57 | sys.exit() 58 | else: 59 | print("Warning, NO INTERACTIVE CONFIRMATION, REMOVE OLD LOG DIR!") 60 | os.system("rm " + abs_log_dir + " -r") 61 | os.makedirs(abs_log_dir) 62 | print('Log dir confirmed...') 63 | 64 | # save configuration in log dir 65 | with open(os.path.join(abs_log_dir, 'configuration_%d_ep.txt' % cfg.RESUME_EPOCH_ID), 'w+') as f: 66 | print(cfg, file=f) 67 | f.close() 68 | print('Save configuration to local file...') 69 | 70 | # backup model, dataset, config 71 | file_backup_path = os.path.join(abs_log_dir, 'files_backup') 72 | os.makedirs(file_backup_path) 73 | os.system('cp ' + os.path.join(cfg.ROOT_DIR, 'core', 'models', cfg.MODEL + '.py') + ' ' + file_backup_path) 74 | os.system('cp ' + os.path.join(cfg.ROOT_DIR, 'dataset', cfg.DATASET + '.py') + ' ' + file_backup_path) 75 | os.system('cp ' + os.path.join(cfg.ROOT_DIR, 'config', 'config_files', cfg.CONFIG_FILE) + ' ' + file_backup_path) 76 | for filename in cfg.BACKUP_FILES: 77 | os.system('cp ' + os.path.join(cfg.ROOT_DIR, filename) + ' ' + file_backup_path) 78 | -------------------------------------------------------------------------------- /core/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JiahuiLei/Pix2Surf/ac56ce5a20adab8cb563d96c915e1cc253631bf8/core/__init__.py -------------------------------------------------------------------------------- /core/evaluation.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from sklearn.neighbors import NearestNeighbors 3 | from multiprocessing.dummy import Pool as ThreadPool 4 | import cv2 as cv 5 | from copy import deepcopy 6 | 7 | 8 | def eval_warp(batch, method_name, nox_gt_key, nox_pred_key): 9 | """ 10 | Parameters 11 | ---------- 12 | batch: the batch in post-processing, must be a multi-view like batch 13 | method_name: string of name e.g 'pix2surf-sv' that will be written to the xls 14 | nox_gt_key: the name of nox-gt; WARNING! the nocs-map should have white background 15 | nox_pred_key: see above, for prediction 16 | 17 | Returns: the batch that added the 'metric-report' the report xls 18 | ------- 19 | """ 20 | 21 | # multi thread eval (Initial motivation of multi-thread is for some slow computation like normal) 22 | n_batch = batch[nox_gt_key][0].shape[0] 23 | n_view = len(batch[nox_gt_key]) 24 | 25 | # make the parameter tuple list for multi-thread 26 | TASKS, RESULTS = [], [] 27 | id = 0 28 | for bdx in range(n_batch): 29 | for vdx in range(n_view): 30 | arg = [id] 31 | id += 1 32 | _nocs_v_gt = batch[nox_gt_key][vdx][bdx].detach().cpu().numpy().transpose(1, 2, 0) 33 | arg.append(_nocs_v_gt) 34 | _nox_v_pred = batch[nox_pred_key][vdx][bdx].detach().cpu().numpy().transpose(1, 2, 0) 35 | arg.append(_nox_v_pred) 36 | TASKS.append(tuple(arg)) 37 | assert id == n_batch * n_view 38 | 39 | with ThreadPool(max(id, 16)) as pool: 40 | _results = [pool.apply_async(eval_thread, t) for t in TASKS] 41 | RESULTS = [r.get() for r in _results] 42 | 43 | ordered_results = [] 44 | for idx in range(id): 45 | for r in RESULTS: 46 | if r[0] == idx: 47 | ordered_results.append(r) 48 | assert len(ordered_results) == len(RESULTS) 49 | 50 | accuracy_list, correspondence_error_list, discontinuity_score_list = [], [], [] 51 | id = 0 52 | for bdx in range(n_batch): 53 | _cd, _corr_l2, _disconti_score = 0, 0, 0 54 | for vdx in range(n_view): 55 | r = ordered_results[id] 56 | id += 1 57 | # for different viewpoint of each object, average across views 58 | _cd += r[1] / n_view 59 | _corr_l2 += r[2] / n_view 60 | _disconti_score += r[3] / n_view 61 | accuracy_list.append(_cd) 62 | correspondence_error_list.append(_corr_l2) 63 | discontinuity_score_list.append(_disconti_score) 64 | 65 | # make xls 66 | report_xls = dict() 67 | if 'metric-report' in batch.keys(): 68 | report_xls = batch['metric-report'] 69 | report_xls[method_name + '-accuracy'] = accuracy_list 70 | report_xls[method_name + '-2D3D-correspondence'] = correspondence_error_list 71 | report_xls[method_name + '-discontinuity-score'] = discontinuity_score_list 72 | batch['metric-report'] = report_xls 73 | 74 | return batch 75 | 76 | 77 | def eval_thread(bid, nocs_gt, nocs_pred): 78 | """ 79 | Parameters 80 | ---------- 81 | bid: batch idx, used to reconstruct a batch in right order 82 | nocs_gt: numpy, [H,W,3] 83 | nocs_pred: numpy, [H,W,3] 84 | Returns 85 | ------- 86 | a tuple of returned values elements 87 | """ 88 | nocs_gt_pts = nocs_gt.reshape(-1, 3) 89 | mask_gt_pts = np.sum(nocs_gt_pts, axis=1) < 3.0 90 | nocs_gt_pts = nocs_gt_pts[mask_gt_pts, :] 91 | 92 | nocs_pred_pts = nocs_pred.reshape(-1, 3) 93 | mask_pred_pts = np.sum(nocs_pred_pts, axis=1) < 3.0 94 | nocs_pred_pts = nocs_pred_pts[mask_pred_pts, :] 95 | 96 | # based on xyz, find nearest neighbor in each other 97 | neigh = NearestNeighbors(n_neighbors=1) 98 | neigh.fit(nocs_pred_pts) 99 | gt_nn_distance, gt_nn_index_of_pred = neigh.kneighbors(nocs_gt_pts, return_distance=True) 100 | neigh = NearestNeighbors(n_neighbors=1) 101 | neigh.fit(nocs_gt_pts) 102 | pred_nn_distance, pred_nn_index_of_gt = neigh.kneighbors(nocs_pred_pts, return_distance=True) 103 | gt_nn_index_of_pred = gt_nn_index_of_pred.squeeze(1) 104 | pred_nn_index_of_gt = pred_nn_index_of_gt.squeeze(1) 105 | 106 | # Compute 2 way Chamfer distance 107 | cd_dist_gt2pred = np.sum((nocs_gt_pts - nocs_pred_pts[gt_nn_index_of_pred, :]) ** 2, axis=1) 108 | cd_dist_pred2gt = np.sum((nocs_pred_pts - nocs_gt_pts[pred_nn_index_of_gt, :]) ** 2, axis=1) 109 | visible_2way_chamfer_distance = cd_dist_gt2pred.mean() + cd_dist_pred2gt.mean() 110 | 111 | # Compute Correspondence error 112 | mask_gt = (nocs_gt.sum(2) < 3.0).astype(np.float) 113 | mask_pred = (nocs_pred.sum(2) < 3.0).astype(np.float) 114 | mask_intersection = mask_gt * mask_pred # H,W,1 115 | xyz_dif = np.sum((deepcopy(nocs_gt) - deepcopy(nocs_pred)) ** 2, axis=2) 116 | xyz_correspondence_distance = (xyz_dif * mask_intersection).sum() / (mask_intersection.sum() + 1e-5) 117 | 118 | # Compute Discontinuity score 119 | pair_dist_gt = compute_pixel_neighbor_diff(deepcopy(nocs_gt)) 120 | pair_dist_pred = compute_pixel_neighbor_diff(deepcopy(nocs_pred)) 121 | k1, k2 = 30, 20 122 | th = 0.05 123 | gt_hist_normalized, gt_count = pair_dist2hist(pair_dist_gt, k1, k2, th) 124 | pred_hist_normalized, pred_count = pair_dist2hist(pair_dist_pred, k1, k2, th) 125 | large_dist_pairs_conv = np.sum((gt_count[k1:] / (gt_count[k1:].sum() + 1e-5)) * \ 126 | (pred_count[k1:] / (pred_count[k1:].sum() + 1e-5))) 127 | 128 | # Cross-View Consistency Error is computed outside here, in the network 129 | return ( 130 | bid, 131 | visible_2way_chamfer_distance, # accuracy 132 | xyz_correspondence_distance, # correspondence 133 | large_dist_pairs_conv, # continuity 134 | ) 135 | 136 | 137 | def compute_pixel_neighbor_diff(nocs_map): 138 | mask = (np.sum(nocs_map, axis=2) < 3.0).astype(np.float) 139 | kernel = cv.getStructuringElement(cv.MORPH_RECT, (3, 3)) 140 | mask_smaller = cv.erode(mask, kernel, iterations=1) 141 | 142 | d_r = deepcopy(nocs_map) 143 | d_r[:, :-1, :] -= deepcopy(nocs_map)[:, 1:, :] 144 | d_r = np.sqrt(np.sum(d_r ** 2, axis=2)) * mask_smaller 145 | 146 | d_l = deepcopy(nocs_map) 147 | d_l[:, 1:, :] -= deepcopy(nocs_map)[:, :-1, :] 148 | d_l = np.sqrt(np.sum(d_l ** 2, axis=2)) * mask_smaller 149 | 150 | d_d = deepcopy(nocs_map) 151 | d_d[:-1, :, :] -= deepcopy(nocs_map)[1:, :, :] 152 | d_d = np.sqrt(np.sum(d_d ** 2, axis=2)) * mask_smaller 153 | 154 | d_u = deepcopy(nocs_map) 155 | d_u[1:, :, :] -= deepcopy(nocs_map)[:-1, :, :] 156 | d_u = np.sqrt(np.sum(d_u ** 2, axis=2)) * mask_smaller 157 | 158 | select_mask = mask_smaller.reshape(-1) > 0.5 159 | dr = d_r.reshape(-1)[select_mask] 160 | dl = d_l.reshape(-1)[select_mask] 161 | du = d_u.reshape(-1)[select_mask] 162 | dd = d_d.reshape(-1)[select_mask] 163 | distance = np.concatenate((dr, dl, du, dd), 0) 164 | 165 | return distance 166 | 167 | 168 | def pair_dist2hist(pair_dist, k1=20, k2=20, th=0.04): 169 | bin1 = np.linspace(0, th, k1) 170 | bin2 = np.linspace(th, np.sqrt(3), k2) 171 | bin = np.concatenate((bin1, bin2, np.array([np.sqrt(3)])), 0) 172 | conunt_list = [] 173 | for idx in range(k1 + k2): 174 | mask = (pair_dist >= bin[idx]).astype(np.float) * (pair_dist < bin[idx + 1]).astype(np.float) 175 | conunt_list.append(mask.sum()) 176 | count = np.array(conunt_list) 177 | hist = count / (len(pair_dist) + 1e-5) 178 | return hist, count 179 | -------------------------------------------------------------------------------- /core/models/__init__.py: -------------------------------------------------------------------------------- 1 | import importlib 2 | 3 | 4 | def get_model(alias): 5 | module = importlib.import_module('core.models.' + alias) 6 | return module.Model 7 | -------------------------------------------------------------------------------- /core/models/modelbase_v2.py: -------------------------------------------------------------------------------- 1 | """ 2 | Model base Version 2 3 | """ 4 | import torch.nn as nn 5 | from torch.nn import init 6 | import torch 7 | 8 | 9 | class ModelBase(object): 10 | def __init__(self, cfg): 11 | """ 12 | Key components: 13 | gpu_list; gpu_list from cfg 14 | optimizer 15 | network: one network in on model, compute loss inside each network on different gpus 16 | Key functions need to implement: 17 | _preprocess 18 | _predict 19 | _postprocess 20 | train/vali_batch: composed of the three methods above 21 | """ 22 | self.model_name = 'Model Base' 23 | # init device 24 | self.gpu_list = cfg.GPU 25 | self.__dataparallel_flag__ = False # flag: whether use DataParallel 26 | # init optimizer 27 | self.lr = cfg.LR 28 | self.optimizer = None 29 | # init network 30 | self.network = None 31 | # clarify output meanings 32 | self.output_info_dict = { 33 | 'metric': list(), 34 | } 35 | 36 | # models core methods ######################################################### 37 | def _preprocess(self, batch): 38 | """ 39 | get a batch from dataset.__getitem__ following a collate_fn 40 | Pass batch dict from CPU to GPU 41 | """ 42 | device = torch.device("cuda") 43 | for k in batch.keys(): 44 | batch[k] = batch[k].to(device) 45 | return batch 46 | 47 | def _predict(self, batch, is_train=True): 48 | """ 49 | forward through the network, 50 | """ 51 | nn_out = self.network(batch['to-nn'], is_train=is_train) 52 | for k, v in nn_out.items(): 53 | batch[k] = v 54 | return batch 55 | 56 | def _postprocess(self, batch): 57 | """ 58 | Post process, get multi GPU batch dicts, process on one gpu or cpu 59 | :return: a dictionary 60 | """ 61 | return batch 62 | 63 | def train_batch(self, batch): 64 | batch = self._preprocess(batch) 65 | self.set_train() 66 | self.zero_grad() 67 | batch = self._predict(batch) 68 | batch = self._postprocess(batch) 69 | if self.__dataparallel_flag__: 70 | for k in batch.keys(): 71 | if k.endswith('loss') or k in self.output_info_dict['metric']: 72 | if isinstance(batch[k], list): 73 | for idx in len(batch[k]): 74 | batch[k][idx] = batch[k][idx].mean() 75 | else: 76 | batch[k] = batch[k].mean() 77 | batch['batch-loss'].backward() 78 | self.optimizer.step() 79 | return batch 80 | 81 | def vali_batch(self, batch): 82 | batch = self._preprocess(batch) 83 | self.set_eval() 84 | with torch.no_grad(): 85 | batch = self._predict(batch, is_train=False) 86 | batch = self._postprocess(batch) 87 | if self.__dataparallel_flag__: 88 | for k in batch.keys(): 89 | if k.endswith('loss') or k in self.output_info_dict['metric']: 90 | if isinstance(batch[k], list): 91 | for idx in len(batch[k]): 92 | batch[k][idx] = batch[k][idx].mean() 93 | else: 94 | batch[k] = batch[k].mean() 95 | return batch 96 | 97 | def test_batch(self, batch): 98 | raise NotImplementedError 99 | 100 | def zero_grad(self): 101 | self.optimizer.zero_grad() 102 | 103 | # models I/O methods ########################################################## 104 | # load, save, init refers to https://github.com/xiumingzhang/GenRe-ShapeHD 105 | def load_model(self, loadpath, current_model_state='cpu', load_optimizer=False, strict=True): 106 | assert current_model_state in ['cpu'], "Model Loading Error!" 107 | print("Load Model from: ") 108 | print(loadpath) 109 | device = torch.device(current_model_state) 110 | if isinstance(loadpath, list): 111 | for path in loadpath: 112 | checkpoint = torch.load(path, map_location=device) 113 | if self.__dataparallel_flag__: 114 | for k in self.network.module.network_dict.keys(): 115 | if k in checkpoint.keys(): 116 | self.network.module.network_dict[k].load_state_dict(checkpoint[k], strict=strict) 117 | if load_optimizer: 118 | self.optimizer.module.load_state_dict(checkpoint['optimizer'], strict=strict) 119 | else: 120 | for k in self.network.network_dict.keys(): 121 | if k in checkpoint.keys(): 122 | self.network.network_dict[k].load_state_dict(checkpoint[k], strict=strict) 123 | if load_optimizer: 124 | self.optimizer.load_state_dict(checkpoint['optimizer'], strict=strict) 125 | else: 126 | path = loadpath 127 | checkpoint = torch.load(path, map_location=device) 128 | if self.__dataparallel_flag__: 129 | for k in self.network.module.network_dict.keys(): 130 | if k in checkpoint.keys(): 131 | self.network.module.network_dict[k].load_state_dict(checkpoint[k]) 132 | if load_optimizer: 133 | self.optimizer.module.load_state_dict(checkpoint['optimizer']) 134 | else: 135 | for k in self.network.network_dict.keys(): 136 | if k in checkpoint.keys(): 137 | self.network.network_dict[k].load_state_dict(checkpoint[k]) 138 | if load_optimizer: 139 | self.optimizer.load_state_dict(checkpoint['optimizer']) 140 | return checkpoint 141 | 142 | def save_model(self, filepath, additional_dict=None): 143 | save_dict = {} 144 | # add additional info 145 | if additional_dict is not None: 146 | for k, v in additional_dict.items(): 147 | save_dict[k] = v 148 | # save all self.__networks_dict param 149 | if self.__dataparallel_flag__: 150 | for net_name in self.network.module.network_dict.keys(): 151 | save_dict[net_name] = self.network.module.network_dict[net_name].state_dict() 152 | else: 153 | for net_name in self.network.network_dict.keys(): 154 | save_dict[net_name] = self.network.network_dict[net_name].state_dict() 155 | save_dict['optimizer'] = self.optimizer.state_dict() 156 | # save 157 | torch.save(save_dict, filepath) 158 | 159 | def init_weight(self, net=None, init_type='kaiming', init_param=0.02): 160 | """ 161 | Use nn.Module.apply(fn) to recursively apply initialization to the Model 162 | If pass in a network, init the passed in net 163 | else, init self.net 164 | """ 165 | 166 | def init_func(m, init_type=init_type): 167 | classname = m.__class__.__name__ 168 | if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1): 169 | if init_type == 'normal': 170 | init.normal_(m.weight.data, 0.0, init_param) 171 | elif init_type == 'xavier': 172 | init.xavier_normal_(m.weight.data, gain=init_param) 173 | elif init_type == 'kaiming': 174 | init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') 175 | elif init_type == 'orth': 176 | init.orthogonal_(m.weight.data, gain=init_param) 177 | else: 178 | raise NotImplementedError('initialization method [%s] is not implemented' % init_type) 179 | if hasattr(m, 'bias') and m.bias is not None: 180 | init.constant_(m.bias.data, 0.0) 181 | elif classname.find('BatchNorm') != -1: 182 | try: 183 | init.normal_(m.weight.data, 1.0, init_param) 184 | init.constant_(m.bias.data, 0.0) 185 | except: 186 | try: 187 | init.normal_(m.bn.weight.data, 1.0, init_param) 188 | init.constant_(m.bn.bias.data, 0.0) 189 | except: 190 | raise ValueError("Can't Initialize BN") 191 | 192 | if net is not None: 193 | net.apply(init_func) 194 | else: 195 | self.network.apply(init_func) 196 | 197 | # models state methods ######################################################## 198 | 199 | def to_gpus(self): 200 | # set device (if no gpu is available or gpu=[-1]) 201 | device = torch.device("cuda") 202 | if len(self.gpu_list) > 1: 203 | self.network = nn.DataParallel(self.network) 204 | self.__dataparallel_flag__ = True 205 | self.network.to(device) 206 | 207 | def set_train(self): 208 | """ 209 | Set models's network to train mode 210 | """ 211 | self.network.train() 212 | 213 | def set_eval(self): 214 | """ 215 | Set models's network to eval mode 216 | WARNING! THIS METHOD IS VERY IMPORTANT DURING VALIDATION BECAUSE OF PYTORCH BN MECHANISM 217 | """ 218 | self.network.eval() 219 | 220 | 221 | class Network(nn.Module): 222 | def __init__(self): 223 | super(Network, self).__init__() 224 | self.network_dict = None 225 | 226 | def forward(self, *input): 227 | raise NotImplementedError 228 | -------------------------------------------------------------------------------- /core/models/neurips_baseline.py: -------------------------------------------------------------------------------- 1 | """ 2 | ECCV Pix2Surf 3 | Initialization with Neurips NOX training 4 | """ 5 | 6 | from .modelbase_v2 import ModelBase 7 | from .modelbase_v2 import Network as NetBase 8 | from core.net_bank.xnocs_segnet import SegNet 9 | from core.net_bank.loss import MaskL2Loss 10 | import os 11 | import torch 12 | from torch import nn 13 | 14 | 15 | class Model(ModelBase): 16 | def __init__(self, cfg): 17 | super(Model, self).__init__(cfg) 18 | self.name = 'neurips-nox' 19 | self.cfg = cfg 20 | self.network = Network() 21 | self.optimizer = torch.optim.Adam(params=self.network.parameters(), lr=self.lr, 22 | betas=(self.cfg.ADAM_BETA1, self.cfg.ADAM_BETA2)) 23 | self.resume = cfg.RESUME 24 | if self.resume: 25 | self.resume_id = cfg.RESUME_EPOCH_ID 26 | load_path = os.path.join(cfg.ROOT_DIR, 'log', cfg.LOG_DIR, 'model', 27 | 'epoch_%d' % cfg.RESUME_EPOCH_ID + '.model') 28 | self.load_model(loadpath=load_path, current_model_state='cpu') 29 | elif cfg.MODEL_INIT_PATH != ['None']: 30 | self.load_model(loadpath=cfg.MODEL_INIT_PATH) 31 | self.to_gpus() 32 | # config output meaning 33 | self.output_info_dict = { 34 | 'metric': ['batch-loss', 'reg-v-loss', 'reg-x-loss', 'mask-v-loss', 'mask-x-loss'], 35 | 'image': ['rgb-v', 'nox-v-gt', 'nox-x-gt', 'nox-v-pred', 'nox-x-pred', 'mask-v', 'mask-x'], 36 | } 37 | 38 | def _preprocess(self, in_batch): 39 | device = torch.device("cuda") 40 | nox_v = in_batch['nox-v'].float().permute(0, 3, 1, 2).to(device) / 255.0 # [0,1] 41 | nox_x = in_batch['nox-x'].float().permute(0, 3, 1, 2).to(device) / 255.0 # [0,1] 42 | rgb_v = in_batch['rgb-v'].float().permute(0, 3, 1, 2).to(device) / 255.0 # [0,1] 43 | rgb_x = in_batch['rgb-x'].float().permute(0, 3, 1, 2).to(device) / 255.0 # [0,1] 44 | mask_v = in_batch['mask-v'].float().permute(0, 3, 1, 2).to(device) # 0,1 45 | mask_x = in_batch['mask-x'].float().permute(0, 3, 1, 2).to(device) # 0,1 46 | pack = {'rgb-v': rgb_v, 'rgb-x': rgb_x, 'nox-v': nox_v, 'nox-x': nox_x, 47 | 'mask-v': mask_v, 'mask-x': mask_x} 48 | return {'to-nn': pack, 'meta-info': in_batch['info']} 49 | 50 | 51 | class Network(NetBase): 52 | def __init__(self): 53 | super(Network, self).__init__() 54 | net_dict = { 55 | 'seg-net': SegNet(out_channels=10) 56 | } 57 | self.network_dict = nn.ModuleDict(net_dict) 58 | # loss 59 | self.cls_criterion = nn.CrossEntropyLoss() # not masked, for all pixel 60 | self.ml2_criterion = MaskL2Loss() # masked for nocs regression 61 | 62 | def forward(self, pack, is_train=True): 63 | batch = dict() 64 | 65 | # make cnn prediction 66 | pred = self.network_dict['seg-net'](pack['rgb-v']) 67 | pred_nox_v = pred[:, :3, :, :] 68 | pred_nox_x = pred[:, 3:6, :, :] 69 | pred_score_v = pred[:, 6:8, :, :] 70 | pred_score_x = pred[:, 8:10, :, :] 71 | 72 | mask1c_v = pack['mask-v'][:, 0, :, :].unsqueeze(1).detach() 73 | mask_v_loss = self.cls_criterion(pred_score_v, mask1c_v.squeeze(1).long().detach()) 74 | pred_mask_v = torch.argmax(pred_score_v, dim=1, keepdim=True).float() 75 | 76 | mask1c_x = pack['mask-x'][:, 0, :, :].unsqueeze(1).detach() 77 | mask_x_loss = self.cls_criterion(pred_score_x, mask1c_x.squeeze(1).long().detach()) 78 | pred_mask_x = torch.argmax(pred_score_x, dim=1, keepdim=True).float() 79 | 80 | reg_v_loss = self.ml2_criterion(pred_nox_v, pack['nox-v'], mask1c_v, True) 81 | reg_x_loss = self.ml2_criterion(pred_nox_x, pack['nox-x'], mask1c_x, True) 82 | 83 | # summary 84 | batch['batch-loss'] = ((reg_v_loss + reg_x_loss) * 0.3 + (mask_v_loss + mask_x_loss) * 0.7).unsqueeze(0) 85 | 86 | batch['reg-v-loss'] = reg_v_loss.detach().unsqueeze(0) 87 | batch['reg-x-loss'] = reg_x_loss.detach().unsqueeze(0) 88 | batch['mask-v-loss'] = mask_v_loss.detach().unsqueeze(0) 89 | batch['mask-x-loss'] = mask_x_loss.detach().unsqueeze(0) 90 | 91 | batch['nox-v-gt'] = pack['nox-v'] * pack['mask-v'] + (1.0 - pack['mask-v']) 92 | batch['nox-x-gt'] = pack['nox-x'] * pack['mask-x'] + (1.0 - pack['mask-x']) 93 | batch['nox-v-pred'] = pred_nox_v * pred_mask_v + (1.0 - pred_mask_v) 94 | batch['nox-x-pred'] = pred_nox_x * pred_mask_x + (1.0 - pred_mask_x) 95 | batch['mask-v'] = pred_mask_v 96 | batch['mask-x'] = pred_mask_x 97 | batch['rgb-v'] = pack['rgb-v'] 98 | return batch 99 | -------------------------------------------------------------------------------- /core/models/pix2surf_mv.py: -------------------------------------------------------------------------------- 1 | """ 2 | Pix2Surf Multi-View Version 3 | """ 4 | 5 | from .modelbase_v2 import ModelBase 6 | from .modelbase_v2 import Network as NetBase 7 | from core.models.utils import * 8 | from core.net_bank.pix2surf_cnn import SegNetGroup 9 | from core.net_bank.mlp import NOCS_AMP_MLP 10 | from core.net_bank.loss import MaskL2Loss 11 | import os 12 | import torch 13 | from torch import nn 14 | 15 | 16 | class Model(ModelBase): 17 | def __init__(self, cfg): 18 | super(Model, self).__init__(cfg) 19 | self.name = 'pix2surf-mv' 20 | self.cfg = cfg 21 | # register key component 22 | self.network = Network() 23 | self.optimizer = torch.optim.Adam(params=self.network.parameters(), lr=self.lr, 24 | betas=(self.cfg.ADAM_BETA1, self.cfg.ADAM_BETA2)) 25 | # initialize models 26 | self.resume = cfg.RESUME 27 | if self.resume: 28 | self.resume_id = cfg.RESUME_EPOCH_ID 29 | load_path = os.path.join(cfg.ROOT_DIR, 'log', cfg.LOG_DIR, 'model', 30 | 'epoch_%d' % cfg.RESUME_EPOCH_ID + '.model') 31 | self.load_model(loadpath=load_path, current_model_state='cpu', strict=False) 32 | elif cfg.MODEL_INIT_PATH != ['None']: 33 | self.load_model(loadpath=cfg.MODEL_INIT_PATH, strict=False) 34 | self.to_gpus() 35 | # config output meaning 36 | self.output_info_dict = { 37 | 'metric': ['batch-loss', 'reg-v-loss', 'reg-x-loss', 'mask-v-loss', 'mask-x-loss', 38 | 'sp-loss', 'crr-xyz-loss'], 39 | 'image': ['rgb-v', 'nox-v-gt', 'mask-v', 'sp-image'] + # sp-image is a nocs map 40 | ['unwrapped-chart', 'unwrapped-chart-uni', 'learned-chart'], 41 | # learned-chart is color coded uv in image space, unwrapped-chart is in uv space visualization 42 | } 43 | 44 | def _preprocess(self, in_batch): 45 | return load_multiview_batch(in_batch) 46 | 47 | 48 | class Network(NetBase): 49 | def __init__(self): 50 | super(Network, self).__init__() 51 | net_dict = { 52 | 'seg-net': SegNetGroup(out_channels=10, additional=2), 53 | 'global-code': nn.Sequential( 54 | nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=0, stride=1), 55 | nn.BatchNorm2d(512), 56 | nn.ELU(), 57 | nn.Conv2d(in_channels=512, out_channels=1024, kernel_size=3, padding=0, stride=1), 58 | nn.MaxPool2d(kernel_size=(3, 6)) 59 | ), 60 | 'mlp': NOCS_AMP_MLP(latent_dim=1024, amp_dim=256, p_in=2, c_out=3) 61 | } 62 | self.network_dict = nn.ModuleDict(net_dict) 63 | # loss 64 | self.cls_criterion = nn.CrossEntropyLoss() # not masked, for all pixel 65 | self.ml2_criterion = MaskL2Loss() 66 | self.sgmd = nn.Sigmoid() 67 | # visualization resolution 68 | self.vis_chart_res = 256 69 | self.vis_chart_container = torch.zeros(1, 3, self.vis_chart_res, self.vis_chart_res) 70 | 71 | def forward(self, pack, is_train=True): 72 | batch = dict() 73 | n_batch = pack['nox-v'][0].shape[0] 74 | n_view = len(pack['rgb-v']) 75 | 76 | code_list = list() 77 | pred_list, featuremap_list = self.network_dict['seg-net'](pack['rgb-v'], return_code=True) 78 | for fm in featuremap_list: # do for each view 79 | code_list.append(self.network_dict['global-code'](fm).reshape(n_batch, -1, 1).contiguous()) 80 | global_z = torch.max(torch.cat(code_list, 2), dim=2).values.contiguous() 81 | 82 | # prepare gather container 83 | pred_nox_v_list, pred_nox_x_list, pred_mask_v_list, pred_mask_x_list = [], [], [], [] 84 | pred_xyz_list, pred_uv_list = [], [] 85 | learned_chart, unwrapped_chart_list, sp_image = [], [], [] 86 | reg_v_loss, reg_x_loss, mask_v_loss, mask_x_loss, sp_loss = 0, 0, 0, 0, 0 87 | 88 | for ii in range(n_view): 89 | mask_v = pack['mask-v'][ii] 90 | mask_x = pack['mask-x'][ii] 91 | 92 | # make cnn prediction 93 | pred = pred_list[ii] 94 | pred_nox_v = pred[:, :3, :, :] 95 | pred_nox_x = pred[:, 3:6, :, :] 96 | pred_score_v = pred[:, 6:8, :, :] 97 | pred_score_x = pred[:, 8:10, :, :] 98 | learned_uv = self.sgmd(pred[:, 10:12, :, :]) 99 | 100 | # make NOCS-regression branch 101 | mask1c_v = mask_v[:, 0, :, :].unsqueeze(1).detach() 102 | mask_v_loss = mask_v_loss + self.cls_criterion(pred_score_v, mask1c_v.squeeze(1).long().detach()) / n_view 103 | pred_mask_v = torch.argmax(pred_score_v, dim=1, keepdim=True).float() 104 | mask1c_x = mask_x[:, 0, :, :].unsqueeze(1).detach() 105 | mask_x_loss = mask_x_loss + self.cls_criterion(pred_score_x, mask1c_x.squeeze(1).long().detach()) / n_view 106 | pred_mask_x = torch.argmax(pred_score_x, dim=1, keepdim=True).float() 107 | reg_v_loss = reg_v_loss + self.ml2_criterion(pred_nox_v, pack['nox-v'][ii], mask1c_v, True) / n_view 108 | reg_x_loss = reg_x_loss + self.ml2_criterion(pred_nox_x, pack['nox-x'][ii], mask1c_x, True) / n_view 109 | 110 | # make mlp prediction 111 | eachview_z = code_list[ii].squeeze(2) 112 | latent_dim = eachview_z.shape[1] 113 | c = torch.cat((eachview_z[:, :latent_dim // 2], global_z[:, latent_dim // 2:]), dim=1) 114 | queried_uv = query_feature(learned_uv, pack['uv-v'][ii]) 115 | pred_xyz = self.network_dict['mlp'](c, queried_uv, unique_code=True) 116 | pred_xyz = self.sgmd(pred_xyz) 117 | sp_loss = sp_loss + self.ml2_criterion(pred_xyz, pack['uv-xyz-v'][ii], pack['uv-mask-v'][ii]) / n_view 118 | 119 | # vis 120 | unwrapped_chart = self.vis_chart_container.repeat(n_batch, 1, 1, 1).cuda() 121 | unwrapped_chart = spread_feature(unwrapped_chart, learned_uv, pack['rgb-v'][ii], pack['mask-v'][ii]) 122 | vis_sampled_xyz = torch.ones_like(pack['rgb-v'][ii]).float() 123 | uv = pack['uv-v'][ii] 124 | uv[:, 0, :, :] = uv[:, 0, :, :] * mask1c_v.shape[2] 125 | uv[:, 1, :, :] = uv[:, 1, :, :] * mask1c_v.shape[3] 126 | uv = uv.long() 127 | idx = uv[:, 0, :, :] * mask1c_v.shape[3] + uv[:, 1, :, :] # B,N,1 128 | idx = idx.permute(0, 2, 1) # B,1,N 129 | vis_sampled_xyz = vis_sampled_xyz.reshape(n_batch, 3, -1) # B,3,R*R 130 | vis_sampled_xyz = vis_sampled_xyz.scatter(dim=2, index=idx.repeat(1, 3, 1), src=pred_xyz.squeeze(3)) 131 | vis_sampled_xyz = vis_sampled_xyz.reshape(n_batch, 3, mask1c_v.shape[2], mask1c_v.shape[3]) 132 | 133 | # gather 134 | pred_nox_v_list.append(pred_nox_v) 135 | pred_nox_x_list.append(pred_nox_x) 136 | pred_mask_v_list.append(pred_mask_v) 137 | pred_mask_x_list.append(pred_mask_x) 138 | 139 | pred_xyz_list.append(pred_xyz) 140 | pred_uv_list.append(queried_uv) 141 | unwrapped_chart_list.append(unwrapped_chart) 142 | learned_chart.append(learned_uv.repeat(1, 2, 1, 1)[:, :3, :, :] * mask1c_v + (1.0 - mask1c_v)) 143 | sp_image.append(vis_sampled_xyz) 144 | 145 | # make naive multi-view constrain: 146 | _p1_list, _p2_list, _m_list = [], [], [] 147 | _uv1_list, _uv2_list = [], [] 148 | for base_view_id in range(len(pack['crr-idx-mtx'])): 149 | for query_view_id in range(len(pack['crr-idx-mtx'][base_view_id])): 150 | base_pc = pred_xyz_list[base_view_id] 151 | query_pc = pred_xyz_list[base_view_id + query_view_id + 1] 152 | base_uv = pred_uv_list[base_view_id] 153 | query_uv = pred_uv_list[base_view_id + query_view_id + 1] 154 | pair_idx = pack['crr-idx-mtx'][base_view_id][query_view_id].squeeze(3) 155 | paired_pc_from_base_to_query = torch.gather(base_pc.squeeze(3), dim=2, 156 | index=pair_idx.repeat(1, 3, 1)).unsqueeze(3) 157 | paired_uv_from_base_to_query = torch.gather(base_uv.squeeze(3), dim=2, 158 | index=pair_idx.repeat(1, 2, 1)).unsqueeze(3) 159 | _p1_list.append(paired_pc_from_base_to_query) 160 | _p2_list.append(query_pc) 161 | _uv1_list.append(paired_uv_from_base_to_query) 162 | _uv2_list.append(query_uv) 163 | _m_list.append(pack['crr-mask-mtx'][base_view_id][query_view_id]) 164 | 165 | crr_xyz_loss = self.ml2_criterion(torch.cat(_p1_list, dim=2).contiguous(), 166 | torch.cat(_p2_list, dim=2).contiguous(), 167 | torch.cat(_m_list, dim=2).contiguous(), detach=False) 168 | 169 | crr_uv_loss = self.ml2_criterion(torch.cat(_uv1_list, dim=2).contiguous(), 170 | torch.cat(_uv2_list, dim=2).contiguous(), 171 | torch.cat(_m_list, dim=2).contiguous(), detach=False) # not used 172 | 173 | # summary 174 | batch['batch-loss'] = (((reg_v_loss + reg_x_loss) * 0.1 + (mask_v_loss + mask_x_loss) * 0.1) * 0.1 + \ 175 | sp_loss * 0.9 + crr_xyz_loss * 0.9).unsqueeze(0) # + crr_uv_loss * 0.1 176 | 177 | batch['reg-v-loss'] = reg_v_loss.detach().unsqueeze(0) 178 | batch['reg-x-loss'] = reg_x_loss.detach().unsqueeze(0) 179 | batch['mask-v-loss'] = mask_v_loss.detach().unsqueeze(0) 180 | batch['mask-x-loss'] = mask_x_loss.detach().unsqueeze(0) 181 | batch['sp-loss'] = sp_loss.detach().unsqueeze(0) 182 | batch['crr-xyz-loss'] = crr_xyz_loss.detach().unsqueeze(0) 183 | 184 | batch['nox-v-gt'] = torch.cat([p * m + (1.0 - m) for p, m in zip(pack['nox-v'], pack['mask-v'])], 3) 185 | batch['nox-x-gt'] = torch.cat([p * m + (1.0 - m) for p, m in zip(pack['nox-x'], pack['mask-x'])], 3) 186 | batch['mask-v'] = torch.cat(pred_mask_v_list, 3) 187 | batch['mask-x'] = torch.cat(pred_mask_x_list, 3) 188 | batch['rgb-v'] = torch.cat(pack['rgb-v'], 3) 189 | batch['sp-image'] = torch.cat(sp_image, 3) 190 | batch['unwrapped-chart'] = torch.cat(unwrapped_chart_list, 3) 191 | # vis all learned chart in one unified uv space, curious to see what happens 192 | vis_nsc_uni = unwrapped_chart_list[0] 193 | for new_scatter in unwrapped_chart_list: 194 | vis_nsc_uni = torch.max(new_scatter, vis_nsc_uni) 195 | batch['unwrapped-chart-uni'] = vis_nsc_uni 196 | batch['learned-chart'] = torch.cat(learned_chart, 3) 197 | 198 | return batch 199 | -------------------------------------------------------------------------------- /core/models/pix2surf_mv_mveval.py: -------------------------------------------------------------------------------- 1 | """ 2 | Pix2Surf Multi-View Version Evaluation in multi view protocol 3 | """ 4 | 5 | from .modelbase_v2 import ModelBase 6 | from core.models.utils import * 7 | from .pix2surf_mv import Network as MV_Net 8 | 9 | import os 10 | import torch 11 | from core.evaluation import eval_warp 12 | 13 | 14 | class Model(ModelBase): 15 | def __init__(self, cfg): 16 | super(Model, self).__init__(cfg) 17 | self.name = 'pix2surf-mv' 18 | self.cfg = cfg 19 | # register key component 20 | self.network = Network() 21 | self.optimizer = torch.optim.Adam(params=self.network.parameters(), lr=self.lr, 22 | betas=(self.cfg.ADAM_BETA1, self.cfg.ADAM_BETA2)) 23 | # initialize models 24 | self.resume = cfg.RESUME 25 | if self.resume: 26 | self.resume_id = cfg.RESUME_EPOCH_ID 27 | load_path = os.path.join(cfg.ROOT_DIR, 'log', cfg.LOG_DIR, 'model', 28 | 'epoch_%d' % cfg.RESUME_EPOCH_ID + '.model') 29 | self.load_model(loadpath=load_path, current_model_state='cpu', strict=False) 30 | elif cfg.MODEL_INIT_PATH != ['None']: 31 | self.load_model(loadpath=cfg.MODEL_INIT_PATH, strict=False) 32 | self.to_gpus() 33 | # config output meaning 34 | self.output_info_dict = { 35 | 'metric': ['batch-loss', 'reg-v-loss', 'reg-x-loss', 'mask-v-loss', 'mask-x-loss', 36 | 'sp-loss', 'crr-xyz-loss'], 37 | 'image': ['uni-rgb-v', 'nox-v-gt-uni', 'mask-v'] + 38 | ['unwrapped-chart', 'unwrapped-chart-uni', 'learned-chart', 'sp-image-uni'], 39 | 'xls': ['metric-report'] 40 | } 41 | 42 | def _preprocess(self, in_batch): 43 | return load_multiview_batch(in_batch) 44 | 45 | def _postprocess(self, batch): 46 | # compute metric in multi thread 47 | batch = eval_warp(batch, method_name='pix2surf-mv', nox_gt_key='nox-v-gt', nox_pred_key='sp-image') 48 | # add crr_loss to xls report 49 | batch['metric-report']['consistency-error'] = [float(i) for i in 50 | batch['crr-xyz-loss-xls'].detach().cpu().numpy()] 51 | return batch 52 | 53 | 54 | class Network(MV_Net): 55 | def __init__(self): 56 | super(Network, self).__init__() 57 | # make eval config 58 | self.eval_image_res = (240, 320) 59 | self.eval_image_grid = make_grid(self.eval_image_res) 60 | 61 | def forward(self, pack, is_train=True): 62 | batch = dict() 63 | n_batch = pack['nox-v'][0].shape[0] 64 | n_view = len(pack['rgb-v']) 65 | 66 | code_list = list() 67 | pred_list, featuremap_list = self.network_dict['seg-net'](pack['rgb-v'], return_code=True) 68 | for fm in featuremap_list: # do for each view 69 | code_list.append(self.network_dict['global-code'](fm).reshape(n_batch, -1, 1).contiguous()) 70 | global_z = torch.max(torch.cat(code_list, 2), dim=2).values.contiguous() 71 | 72 | # prepare gather container 73 | pred_nox_v_list, pred_nox_x_list, pred_mask_v_list, pred_mask_x_list = [], [], [], [] 74 | pred_xyz_list, pred_uv_list = [], [] 75 | learned_chart_list, unwrapped_chart_list = [], [] 76 | reg_v_loss, reg_x_loss, mask_v_loss, mask_x_loss, sp_loss = 0, 0, 0, 0, 0 77 | eval_rendered_list = [] 78 | 79 | for ii in range(n_view): 80 | mask_v = pack['mask-v'][ii] 81 | mask_x = pack['mask-x'][ii] 82 | 83 | # make cnn prediction 84 | pred = pred_list[ii] 85 | pred_nox_v = pred[:, :3, :, :] 86 | pred_nox_x = pred[:, 3:6, :, :] 87 | pred_score_v = pred[:, 6:8, :, :] 88 | pred_score_x = pred[:, 8:10, :, :] 89 | learned_uv = self.sgmd(pred[:, 10:12, :, :]) 90 | 91 | # make NOCS-regression branch 92 | mask1c_v = mask_v[:, 0, :, :].unsqueeze(1).detach() 93 | mask_v_loss = mask_v_loss + self.cls_criterion(pred_score_v, mask1c_v.squeeze(1).long().detach()) / n_view 94 | pred_mask_v = torch.argmax(pred_score_v, dim=1, keepdim=True).float() 95 | mask1c_x = mask_x[:, 0, :, :].unsqueeze(1).detach() 96 | mask_x_loss = mask_x_loss + self.cls_criterion(pred_score_x, mask1c_x.squeeze(1).long().detach()) / n_view 97 | pred_mask_x = torch.argmax(pred_score_x, dim=1, keepdim=True).float() 98 | reg_v_loss = reg_v_loss + self.ml2_criterion(pred_nox_v, pack['nox-v'][ii], mask1c_v, True) / n_view 99 | reg_x_loss = reg_x_loss + self.ml2_criterion(pred_nox_x, pack['nox-x'][ii], mask1c_x, True) / n_view 100 | 101 | # make mlp prediction 102 | eachview_z = code_list[ii].squeeze(2) 103 | latent_dim = eachview_z.shape[1] 104 | c = torch.cat((eachview_z[:, :latent_dim // 2], global_z[:, latent_dim // 2:]), dim=1) 105 | queried_uv = query_feature(learned_uv, pack['uv-v'][ii]) 106 | pred_xyz = self.network_dict['mlp'](c, queried_uv, unique_code=True) 107 | pred_xyz = self.sgmd(pred_xyz) 108 | sp_loss = sp_loss + self.ml2_criterion(pred_xyz, pack['uv-xyz-v'][ii], pack['uv-mask-v'][ii]) / n_view 109 | 110 | # Do SP evaluation 111 | _eval_rendered_list = list() 112 | for bid in range(n_batch): 113 | # select mask 114 | _mask = pred_mask_v[bid, ...].reshape(-1) # H*W 115 | _learned_uv = learned_uv[bid, ...].reshape(1, 2, -1) # 1,2,H*W 116 | _learned_uv = _learned_uv[:, :, _mask > 0] # 1,2,S 117 | uv = self.eval_image_grid.cuda().reshape(1, 2, -1)[:, :, _mask > 0].unsqueeze(3) # 1,2,S,1 118 | # do Surface Parametrization 119 | eval_xyz_v = self.network_dict['mlp'](c[bid, ...].unsqueeze(0), _learned_uv.unsqueeze(3), 120 | unique_code=True) 121 | eval_xyz_v = self.sgmd(eval_xyz_v) # 1,3,S,1 122 | uv[:, 0, :, :] = uv[:, 0, :, :] * mask1c_v.shape[2] 123 | uv[:, 1, :, :] = uv[:, 1, :, :] * mask1c_v.shape[3] 124 | uv = uv.long() 125 | idx = uv[:, 0, :, :] * mask1c_v.shape[3] + uv[:, 1, :, :] # B,N,1 126 | idx = idx.permute(0, 2, 1) # B,1,N 127 | vis_eval = torch.ones_like(pack['rgb-v'][ii]).float()[bid, ...].unsqueeze(0) 128 | vis_eval = vis_eval.reshape(1, 3, -1) # B,3,R*R 129 | vis_eval = vis_eval.scatter(dim=2, index=idx.repeat(1, 3, 1), src=eval_xyz_v.squeeze(3)) 130 | vis_eval = vis_eval.reshape(1, 3, mask1c_v.shape[2], mask1c_v.shape[3]) 131 | _eval_rendered_list.append(vis_eval) 132 | eval_rendered = torch.cat(_eval_rendered_list, 0) 133 | eval_rendered_list.append(eval_rendered) 134 | 135 | # vis unwrapped chart 136 | unwrapped_chart = self.vis_chart_container.repeat(n_batch, 1, 1, 1).cuda() 137 | unwrapped_chart = spread_feature(unwrapped_chart, learned_uv, pack['rgb-v'][ii], pack['mask-v'][ii]) 138 | 139 | # gather 140 | pred_nox_v_list.append(pred_nox_v) 141 | pred_nox_x_list.append(pred_nox_x) 142 | pred_mask_v_list.append(pred_mask_v) 143 | pred_mask_x_list.append(pred_mask_x) 144 | 145 | pred_xyz_list.append(pred_xyz) 146 | pred_uv_list.append(queried_uv) 147 | unwrapped_chart_list.append(unwrapped_chart) 148 | learned_chart_list.append(learned_uv.repeat(1, 2, 1, 1)[:, :3, :, :] * pred_mask_x + (1.0 - pred_mask_v)) 149 | 150 | # make naive multi-view constrain: 151 | _p1_list, _p2_list, _m_list = [], [], [] 152 | _uv1_list, _uv2_list = [], [] 153 | for base_view_id in range(len(pack['crr-idx-mtx'])): 154 | for query_view_id in range(len(pack['crr-idx-mtx'][base_view_id])): 155 | base_pc = pred_xyz_list[base_view_id] 156 | query_pc = pred_xyz_list[base_view_id + query_view_id + 1] 157 | base_uv = pred_uv_list[base_view_id] 158 | query_uv = pred_uv_list[base_view_id + query_view_id + 1] 159 | pair_idx = pack['crr-idx-mtx'][base_view_id][query_view_id].squeeze(3) 160 | paired_pc_from_base_to_query = torch.gather(base_pc.squeeze(3), dim=2, 161 | index=pair_idx.repeat(1, 3, 1)).unsqueeze(3) 162 | paired_uv_from_base_to_query = torch.gather(base_uv.squeeze(3), dim=2, 163 | index=pair_idx.repeat(1, 2, 1)).unsqueeze(3) 164 | _p1_list.append(paired_pc_from_base_to_query) 165 | _p2_list.append(query_pc) 166 | _uv1_list.append(paired_uv_from_base_to_query) 167 | _uv2_list.append(query_uv) 168 | _m_list.append(pack['crr-mask-mtx'][base_view_id][query_view_id]) 169 | 170 | crr_xyz_loss_each = self.ml2_criterion(torch.cat(_p1_list, dim=2).contiguous(), 171 | torch.cat(_p2_list, dim=2).contiguous(), 172 | torch.cat(_m_list, dim=2).contiguous(), 173 | detach=False, reduce_batch=False) 174 | crr_xyz_loss = crr_xyz_loss_each.mean() 175 | 176 | crr_uv_loss_each = self.ml2_criterion(torch.cat(_uv1_list, dim=2).contiguous(), 177 | torch.cat(_uv2_list, dim=2).contiguous(), 178 | torch.cat(_m_list, dim=2).contiguous(), 179 | detach=False, reduce_batch=False) 180 | crr_uv_loss = crr_uv_loss_each.mean() 181 | 182 | # summary 183 | batch['batch-loss'] = (((reg_v_loss + reg_x_loss) * 0.1 + (mask_v_loss + mask_x_loss) * 0.1) * 0.1 + \ 184 | sp_loss * 0.9 + crr_xyz_loss * 0.9).unsqueeze(0) # + crr_uv_loss * 0.1 185 | 186 | batch['reg-v-loss'] = reg_v_loss.detach().unsqueeze(0) 187 | batch['reg-x-loss'] = reg_x_loss.detach().unsqueeze(0) 188 | batch['mask-v-loss'] = mask_v_loss.detach().unsqueeze(0) 189 | batch['mask-x-loss'] = mask_x_loss.detach().unsqueeze(0) 190 | batch['sp-loss'] = sp_loss.detach().unsqueeze(0) 191 | batch['crr-xyz-loss'] = crr_xyz_loss.detach().unsqueeze(0) 192 | batch['crr-xyz-loss-xls'] = crr_xyz_loss_each.detach() 193 | 194 | batch['mask-v'] = torch.cat(pred_mask_v_list, 3) 195 | batch['mask-x'] = torch.cat(pred_mask_x_list, 3) 196 | batch['rgb-v'] = pack['rgb-v'] 197 | batch['uni-rgb-v'] = torch.cat(pack['rgb-v'], 3) 198 | 199 | batch['nox-v-gt'] = [p * m + (1.0 - m) for p, m in zip(pack['nox-v'], pack['mask-v'])] 200 | batch['nox-x-gt'] = [p * m + (1.0 - m) for p, m in zip(pack['nox-x'], pack['mask-x'])] 201 | batch['nox-v-gt-uni'] = torch.cat([p * m + (1.0 - m) for p, m in zip(pack['nox-v'], pack['mask-v'])], 3) 202 | batch['nox-x-gt-uni'] = torch.cat([p * m + (1.0 - m) for p, m in zip(pack['nox-x'], pack['mask-x'])], 3) 203 | 204 | batch['sp-image'] = eval_rendered_list 205 | batch['sp-image-uni'] = torch.cat(eval_rendered_list, 3) 206 | 207 | batch['learned-chart'] = torch.cat(learned_chart_list, 3) 208 | batch['unwrapped-chart'] = torch.cat(unwrapped_chart_list, 3) 209 | vis_nsc_uni = unwrapped_chart_list[0] 210 | for new_scatter in unwrapped_chart_list: 211 | vis_nsc_uni = torch.max(new_scatter, vis_nsc_uni) 212 | batch['unwrapped-chart-uni'] = vis_nsc_uni 213 | 214 | return batch 215 | -------------------------------------------------------------------------------- /core/models/pix2surf_sv.py: -------------------------------------------------------------------------------- 1 | """ 2 | Pix2Surf Single View Single Chart Version 3 | """ 4 | 5 | from .modelbase_v2 import ModelBase 6 | from .modelbase_v2 import Network as NetBase 7 | from core.models.utils import * 8 | from core.net_bank.pix2surf_cnn import SegNet 9 | from core.net_bank.mlp import NOCS_AMP_MLP 10 | from core.net_bank.loss import MaskL2Loss 11 | import os 12 | import torch 13 | from torch import nn 14 | 15 | 16 | class Model(ModelBase): 17 | def __init__(self, cfg): 18 | super(Model, self).__init__(cfg) 19 | self.name = 'pix2surf-sv' 20 | self.cfg = cfg 21 | # register key component 22 | self.network = Network() 23 | self.optimizer = torch.optim.Adam(params=self.network.parameters(), lr=self.lr, 24 | betas=(self.cfg.ADAM_BETA1, self.cfg.ADAM_BETA2)) 25 | # initialize models 26 | self.resume = cfg.RESUME 27 | if self.resume: 28 | self.resume_id = cfg.RESUME_EPOCH_ID 29 | load_path = os.path.join(cfg.ROOT_DIR, 'log', cfg.LOG_DIR, 'model', 30 | 'epoch_%d' % cfg.RESUME_EPOCH_ID + '.model') 31 | self.load_model(loadpath=load_path, current_model_state='cpu') 32 | elif cfg.MODEL_INIT_PATH != ['None']: 33 | self.load_model(loadpath=cfg.MODEL_INIT_PATH, strict=False) 34 | self.to_gpus() 35 | # config output meaning 36 | self.output_info_dict = { 37 | 'metric': ['batch-loss', 'reg-v-loss', 'reg-x-loss', 'mask-v-loss', 'mask-x-loss', 'mlp-v-loss'], 38 | 'image': ['rgb-v', 'nox-v-gt', 'mask-v', 'tex', 'gim', 'learned-chart', 'sp-image'], 39 | # tex(ture) is unwrapped chart with color visualization 40 | # gim (Geometry Image) is unwrapped chart with NOCS XYZ visualization 41 | # learned-chart is the 2 channel output color coded visualization in image space 42 | } 43 | 44 | def _preprocess(self, in_batch): 45 | return load_singleview_batch(in_batch) 46 | 47 | 48 | class Network(NetBase): 49 | def __init__(self): 50 | super(Network, self).__init__() 51 | net_dict = { 52 | 'seg-net': SegNet(out_channels=10, additional=2), 53 | 'global-code': nn.Sequential( 54 | nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=0, stride=1), 55 | nn.BatchNorm2d(512), 56 | nn.ELU(), 57 | nn.Conv2d(in_channels=512, out_channels=1024, kernel_size=3, padding=0, stride=1), 58 | nn.MaxPool2d(kernel_size=(3, 6)) 59 | ), 60 | 'mlp': NOCS_AMP_MLP(latent_dim=1024, amp_dim=256, p_in=2, c_out=3) 61 | } 62 | self.network_dict = nn.ModuleDict(net_dict) 63 | self.sgmd = nn.Sigmoid() 64 | # loss 65 | self.cls_criterion = nn.CrossEntropyLoss() # not masked, for all pixel 66 | self.ml2_criterion = MaskL2Loss() # masked for nocs regression 67 | # visualization 68 | self.vis_chart_res = 128 69 | self.vis_chart_container = torch.zeros(1, 3, self.vis_chart_res, self.vis_chart_res) 70 | 71 | def forward(self, pack, is_train=True): 72 | batch = dict() 73 | n_batch = pack['nox-v'].shape[0] 74 | 75 | # make cnn prediction 76 | pred, fm = self.network_dict['seg-net'](pack['rgb-v'], return_code=True) 77 | pred_nox_v = pred[:, :3, :, :] 78 | pred_nox_x = pred[:, 3:6, :, :] 79 | pred_score_v = pred[:, 6:8, :, :] 80 | pred_score_x = pred[:, 8:10, :, :] 81 | learned_uv = self.sgmd(pred[:, 10:12, :, :]) 82 | 83 | # make NOCS-regression branch 84 | mask1c_v = pack['mask-v'][:, 0, :, :].unsqueeze(1).detach() 85 | mask_v_loss = self.cls_criterion(pred_score_v, mask1c_v.squeeze(1).long().detach()) 86 | pred_mask_v = torch.argmax(pred_score_v, dim=1, keepdim=True).float() 87 | mask1c_x = pack['mask-x'][:, 0, :, :].unsqueeze(1).detach() 88 | mask_x_loss = self.cls_criterion(pred_score_x, mask1c_x.squeeze(1).long().detach()) 89 | pred_mask_x = torch.argmax(pred_score_x, dim=1, keepdim=True).float() 90 | reg_v_loss = self.ml2_criterion(pred_nox_v, pack['nox-v'], mask1c_v, True) 91 | reg_x_loss = self.ml2_criterion(pred_nox_x, pack['nox-x'], mask1c_x, True) 92 | 93 | # make mlp prediction 94 | z = self.network_dict['global-code'](fm).reshape(n_batch, -1).contiguous() 95 | queried_uv_v = query_feature(learned_uv, pack['uv-v']) 96 | pred_xyz_v = self.network_dict['mlp'](z, queried_uv_v, unique_code=True) 97 | pred_xyz_v = self.sgmd(pred_xyz_v) 98 | sp_v_loss = self.ml2_criterion(pred_xyz_v, pack['uv-xyz-v'], pack['uv-mask-v']) 99 | 100 | # vis 101 | tex = self.vis_chart_container.repeat(n_batch, 1, 1, 1).cuda() 102 | tex = spread_feature(tex, learned_uv, pack['rgb-v'], pack['mask-v']) 103 | gim = self.vis_chart_container.repeat(n_batch, 1, 1, 1).cuda() 104 | gim = spread_feature(gim, queried_uv_v, pred_xyz_v, pack['uv-mask-v']) 105 | 106 | vis_sampled_xyz = torch.ones_like(pack['rgb-v']).float() 107 | uv = pack['uv-v'] 108 | uv[:, 0, :, :] = uv[:, 0, :, :] * mask1c_v.shape[2] 109 | uv[:, 1, :, :] = uv[:, 1, :, :] * mask1c_v.shape[3] 110 | uv = uv.long() 111 | idx = uv[:, 0, :, :] * mask1c_v.shape[3] + uv[:, 1, :, :] # B,N,1 112 | idx = idx.permute(0, 2, 1) # B,1,N 113 | vis_sampled_xyz = vis_sampled_xyz.reshape(n_batch, 3, -1) # B,3,R*R 114 | vis_sampled_xyz = vis_sampled_xyz.scatter(dim=2, index=idx.repeat(1, 3, 1), src=pred_xyz_v.squeeze(3)) 115 | vis_sampled_xyz = vis_sampled_xyz.reshape(n_batch, 3, mask1c_v.shape[2], mask1c_v.shape[3]) 116 | 117 | # summary 118 | batch['batch-loss'] = (((reg_v_loss + reg_x_loss) * 0.3 + (mask_v_loss + mask_x_loss) * 0.7) * 0.1 + \ 119 | sp_v_loss * 0.9).unsqueeze(0) 120 | batch['reg-v-loss'] = reg_v_loss.detach().unsqueeze(0) 121 | batch['reg-x-loss'] = reg_x_loss.detach().unsqueeze(0) 122 | batch['mask-v-loss'] = mask_v_loss.detach().unsqueeze(0) 123 | batch['mask-x-loss'] = mask_x_loss.detach().unsqueeze(0) 124 | batch['mlp-v-loss'] = sp_v_loss.detach().unsqueeze(0) 125 | 126 | batch['nox-v-gt'] = pack['nox-v'] * pack['mask-v'] + (1.0 - pack['mask-v']) 127 | batch['nox-x-gt'] = pack['nox-x'] * pack['mask-x'] + (1.0 - pack['mask-x']) 128 | batch['mask-v'] = pred_mask_v 129 | batch['mask-x'] = pred_mask_x 130 | batch['rgb-v'] = pack['rgb-v'] 131 | 132 | batch['learned-chart'] = learned_uv.repeat(1, 2, 1, 1)[:, :3, :, :] * pred_mask_v + (1.0 - pred_mask_v) 133 | batch['tex'] = tex.detach() 134 | batch['gim'] = gim.detach() 135 | batch['sp-image'] = vis_sampled_xyz 136 | 137 | return batch 138 | -------------------------------------------------------------------------------- /core/models/pix2surf_sv_mveval.py: -------------------------------------------------------------------------------- 1 | """ 2 | Pix2Surf Single-View Version Evaluation in multi view protocol 3 | """ 4 | 5 | from .modelbase_v2 import ModelBase 6 | from .pix2surf_sv import Network as SV_Net 7 | from core.models.utils import * 8 | 9 | import os 10 | import torch 11 | from core.evaluation import eval_warp 12 | 13 | 14 | class Model(ModelBase): 15 | def __init__(self, cfg): 16 | super(Model, self).__init__(cfg) 17 | self.name = 'pix2surf-sv' 18 | self.cfg = cfg 19 | # register key component 20 | self.network = Network() 21 | self.optimizer = torch.optim.Adam(params=self.network.parameters(), lr=self.lr, 22 | betas=(self.cfg.ADAM_BETA1, self.cfg.ADAM_BETA2)) 23 | # initialize models 24 | self.resume = cfg.RESUME 25 | if self.resume: 26 | self.resume_id = cfg.RESUME_EPOCH_ID 27 | load_path = os.path.join(cfg.ROOT_DIR, 'log', cfg.LOG_DIR, 'model', 28 | 'epoch_%d' % cfg.RESUME_EPOCH_ID + '.model') 29 | self.load_model(loadpath=load_path, current_model_state='cpu', strict=False) 30 | elif cfg.MODEL_INIT_PATH != ['None']: 31 | self.load_model(loadpath=cfg.MODEL_INIT_PATH, strict=False) 32 | self.to_gpus() 33 | # config output meaning 34 | self.output_info_dict = { 35 | 'metric': ['batch-loss', 'reg-v-loss', 'reg-x-loss', 'mask-v-loss', 'mask-x-loss', 36 | 'sp-loss', 'crr-xyz-loss'], 37 | 'image': ['uni-rgb-v', 'nox-v-gt-uni', 'mask-v'] + 38 | ['unwrapped-chart', 'unwrapped-chart-uni', 'learned-chart', 'sp-image-uni'], 39 | 'xls': ['metric-report'] 40 | } 41 | 42 | def _preprocess(self, in_batch): 43 | return load_multiview_batch(in_batch) 44 | 45 | def _postprocess(self, batch): 46 | # compute metric in multi thread 47 | batch = eval_warp(batch, method_name='pix2surf-sv', nox_gt_key='nox-v-gt', nox_pred_key='sp-image') 48 | # add crr_loss to xls report 49 | batch['metric-report']['consistency-error'] = [float(i) for i in 50 | batch['crr-xyz-loss-xls'].detach().cpu().numpy()] 51 | return batch 52 | 53 | 54 | class Network(SV_Net): 55 | def __init__(self): 56 | super(Network, self).__init__() 57 | # visualization resolution 58 | self.vis_chart_res = 256 59 | self.vis_chart_container = torch.zeros(1, 3, self.vis_chart_res, self.vis_chart_res) 60 | # make eval config 61 | self.eval_image_res = (240, 320) 62 | self.eval_image_grid = make_grid(self.eval_image_res) 63 | 64 | def forward(self, pack, is_train=True): 65 | batch = dict() 66 | n_batch = pack['nox-v'][0].shape[0] 67 | n_view = len(pack['rgb-v']) 68 | 69 | # pred list is a list for each view 70 | pred_list, code_list = list(), list() 71 | for ii in range(n_view): # do for each view 72 | pred, fm = self.network_dict['seg-net'](pack['rgb-v'][ii], return_code=True) 73 | pred_list.append(pred) 74 | code_list.append(self.network_dict['global-code'](fm).reshape(n_batch, -1, 1).contiguous()) 75 | 76 | # prepare gather container 77 | pred_nox_v_list, pred_nox_x_list, pred_mask_v_list, pred_mask_x_list = [], [], [], [] 78 | pred_xyz_list, pred_uv_list = [], [] 79 | learned_chart_list, unwrapped_chart_list = [], [] 80 | reg_v_loss, reg_x_loss, mask_v_loss, mask_x_loss, sp_loss = 0, 0, 0, 0, 0 81 | eval_rendered_list = [] 82 | 83 | for ii in range(n_view): 84 | mask_v = pack['mask-v'][ii] 85 | mask_x = pack['mask-x'][ii] 86 | 87 | # make cnn prediction 88 | pred = pred_list[ii] 89 | pred_nox_v = pred[:, :3, :, :] 90 | pred_nox_x = pred[:, 3:6, :, :] 91 | pred_score_v = pred[:, 6:8, :, :] 92 | pred_score_x = pred[:, 8:10, :, :] 93 | learned_uv = self.sgmd(pred[:, 10:12, :, :]) 94 | 95 | # make NOCS-regression branch 96 | mask1c_v = mask_v[:, 0, :, :].unsqueeze(1).detach() 97 | mask_v_loss = mask_v_loss + self.cls_criterion(pred_score_v, mask1c_v.squeeze(1).long().detach()) / n_view 98 | pred_mask_v = torch.argmax(pred_score_v, dim=1, keepdim=True).float() 99 | mask1c_x = mask_x[:, 0, :, :].unsqueeze(1).detach() 100 | mask_x_loss = mask_x_loss + self.cls_criterion(pred_score_x, mask1c_x.squeeze(1).long().detach()) / n_view 101 | pred_mask_x = torch.argmax(pred_score_x, dim=1, keepdim=True).float() 102 | reg_v_loss = reg_v_loss + self.ml2_criterion(pred_nox_v, pack['nox-v'][ii], mask1c_v, True) / n_view 103 | reg_x_loss = reg_x_loss + self.ml2_criterion(pred_nox_x, pack['nox-x'][ii], mask1c_x, True) / n_view 104 | 105 | # make mlp prediction 106 | eachview_z = code_list[ii].squeeze(2) 107 | queried_uv = query_feature(learned_uv, pack['uv-v'][ii]) 108 | pred_xyz = self.network_dict['mlp'](eachview_z, queried_uv, unique_code=True) 109 | pred_xyz = self.sgmd(pred_xyz) 110 | sp_loss = sp_loss + self.ml2_criterion(pred_xyz, pack['uv-xyz-v'][ii], pack['uv-mask-v'][ii]) / n_view 111 | 112 | # Do SP evaluation 113 | _eval_rendered_list = list() 114 | for bid in range(n_batch): 115 | # select mask 116 | _mask = pred_mask_v[bid, ...].reshape(-1) # H*W 117 | _learned_uv = learned_uv[bid, ...].reshape(1, 2, -1) # 1,2,H*W 118 | _learned_uv = _learned_uv[:, :, _mask > 0] # 1,2,S 119 | uv = self.eval_image_grid.cuda().reshape(1, 2, -1)[:, :, _mask > 0].unsqueeze(3) # 1,2,S,1 120 | # do Surface Parametrization 121 | eval_xyz_v = self.network_dict['mlp'](eachview_z[bid, ...].unsqueeze(0), _learned_uv.unsqueeze(3), 122 | unique_code=True) 123 | eval_xyz_v = self.sgmd(eval_xyz_v) # 1,3,S,1 124 | uv[:, 0, :, :] = uv[:, 0, :, :] * mask1c_v.shape[2] 125 | uv[:, 1, :, :] = uv[:, 1, :, :] * mask1c_v.shape[3] 126 | uv = uv.long() 127 | idx = uv[:, 0, :, :] * mask1c_v.shape[3] + uv[:, 1, :, :] # B,N,1 128 | idx = idx.permute(0, 2, 1) # B,1,N 129 | vis_eval = torch.ones_like(pack['rgb-v'][ii]).float()[bid, ...].unsqueeze(0) 130 | vis_eval = vis_eval.reshape(1, 3, -1) # B,3,R*R 131 | vis_eval = vis_eval.scatter(dim=2, index=idx.repeat(1, 3, 1), src=eval_xyz_v.squeeze(3)) 132 | vis_eval = vis_eval.reshape(1, 3, mask1c_v.shape[2], mask1c_v.shape[3]) 133 | _eval_rendered_list.append(vis_eval) 134 | eval_rendered = torch.cat(_eval_rendered_list, 0) 135 | eval_rendered_list.append(eval_rendered) 136 | 137 | # vis unwrapped chart 138 | unwrapped_chart = self.vis_chart_container.repeat(n_batch, 1, 1, 1).cuda() 139 | unwrapped_chart = spread_feature(unwrapped_chart, learned_uv, pack['rgb-v'][ii], pack['mask-v'][ii]) 140 | 141 | # gather 142 | pred_nox_v_list.append(pred_nox_v) 143 | pred_nox_x_list.append(pred_nox_x) 144 | pred_mask_v_list.append(pred_mask_v) 145 | pred_mask_x_list.append(pred_mask_x) 146 | 147 | pred_xyz_list.append(pred_xyz) 148 | pred_uv_list.append(queried_uv) 149 | unwrapped_chart_list.append(unwrapped_chart) 150 | learned_chart_list.append(learned_uv.repeat(1, 2, 1, 1)[:, :3, :, :] * pred_mask_x + (1.0 - pred_mask_v)) 151 | 152 | # make naive multi-view constrain: 153 | _p1_list, _p2_list, _m_list = [], [], [] 154 | _uv1_list, _uv2_list = [], [] 155 | for base_view_id in range(len(pack['crr-idx-mtx'])): 156 | for query_view_id in range(len(pack['crr-idx-mtx'][base_view_id])): 157 | base_pc = pred_xyz_list[base_view_id] 158 | query_pc = pred_xyz_list[base_view_id + query_view_id + 1] 159 | base_uv = pred_uv_list[base_view_id] 160 | query_uv = pred_uv_list[base_view_id + query_view_id + 1] 161 | pair_idx = pack['crr-idx-mtx'][base_view_id][query_view_id].squeeze(3) 162 | paired_pc_from_base_to_query = torch.gather(base_pc.squeeze(3), dim=2, 163 | index=pair_idx.repeat(1, 3, 1)).unsqueeze(3) 164 | paired_uv_from_base_to_query = torch.gather(base_uv.squeeze(3), dim=2, 165 | index=pair_idx.repeat(1, 2, 1)).unsqueeze(3) 166 | _p1_list.append(paired_pc_from_base_to_query) 167 | _p2_list.append(query_pc) 168 | _uv1_list.append(paired_uv_from_base_to_query) 169 | _uv2_list.append(query_uv) 170 | _m_list.append(pack['crr-mask-mtx'][base_view_id][query_view_id]) 171 | 172 | crr_xyz_loss_each = self.ml2_criterion(torch.cat(_p1_list, dim=2).contiguous(), 173 | torch.cat(_p2_list, dim=2).contiguous(), 174 | torch.cat(_m_list, dim=2).contiguous(), 175 | detach=False, reduce_batch=False) 176 | crr_xyz_loss = crr_xyz_loss_each.mean() 177 | 178 | crr_uv_loss_each = self.ml2_criterion(torch.cat(_uv1_list, dim=2).contiguous(), 179 | torch.cat(_uv2_list, dim=2).contiguous(), 180 | torch.cat(_m_list, dim=2).contiguous(), 181 | detach=False, reduce_batch=False) 182 | crr_uv_loss = crr_uv_loss_each.mean() 183 | 184 | # summary 185 | batch['batch-loss'] = (((reg_v_loss + reg_x_loss) * 0.3 + (mask_v_loss + mask_x_loss) * 0.7) * 0.1 + \ 186 | sp_loss * 0.9).unsqueeze(0) 187 | 188 | batch['reg-v-loss'] = reg_v_loss.detach().unsqueeze(0) 189 | batch['reg-x-loss'] = reg_x_loss.detach().unsqueeze(0) 190 | batch['mask-v-loss'] = mask_v_loss.detach().unsqueeze(0) 191 | batch['mask-x-loss'] = mask_x_loss.detach().unsqueeze(0) 192 | batch['sp-loss'] = sp_loss.detach().unsqueeze(0) 193 | batch['crr-xyz-loss'] = crr_xyz_loss.detach().unsqueeze(0) 194 | batch['crr-xyz-loss-xls'] = crr_xyz_loss_each.detach() 195 | 196 | batch['mask-v'] = torch.cat(pred_mask_v_list, 3) 197 | batch['mask-x'] = torch.cat(pred_mask_x_list, 3) 198 | batch['rgb-v'] = pack['rgb-v'] 199 | batch['uni-rgb-v'] = torch.cat(pack['rgb-v'], 3) 200 | 201 | batch['nox-v-gt'] = [p * m + (1.0 - m) for p, m in zip(pack['nox-v'], pack['mask-v'])] 202 | batch['nox-x-gt'] = [p * m + (1.0 - m) for p, m in zip(pack['nox-x'], pack['mask-x'])] 203 | batch['nox-v-gt-uni'] = torch.cat([p * m + (1.0 - m) for p, m in zip(pack['nox-v'], pack['mask-v'])], 3) 204 | batch['nox-x-gt-uni'] = torch.cat([p * m + (1.0 - m) for p, m in zip(pack['nox-x'], pack['mask-x'])], 3) 205 | 206 | batch['sp-image'] = eval_rendered_list 207 | batch['sp-image-uni'] = torch.cat(eval_rendered_list, 3) 208 | 209 | batch['learned-chart'] = torch.cat(learned_chart_list, 3) 210 | batch['unwrapped-chart'] = torch.cat(unwrapped_chart_list, 3) 211 | vis_nsc_uni = unwrapped_chart_list[0] 212 | for new_scatter in unwrapped_chart_list: 213 | vis_nsc_uni = torch.max(new_scatter, vis_nsc_uni) 214 | batch['unwrapped-chart-uni'] = vis_nsc_uni 215 | 216 | return batch 217 | -------------------------------------------------------------------------------- /core/models/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .model_utils import * 2 | from .render_utils import * 3 | -------------------------------------------------------------------------------- /core/models/utils/model_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | 4 | 5 | def load_multiview_batch(in_batch): 6 | device = torch.device("cuda") 7 | nox_v = [item.float().permute(0, 3, 1, 2).to(device) for item in in_batch['nox-v']] # [0,1] 8 | nox_x = [item.float().permute(0, 3, 1, 2).to(device) for item in in_batch['nox-x']] # [0,1] 9 | rgb_v = [item.float().permute(0, 3, 1, 2).to(device) for item in in_batch['rgb-v']] # [0,1] 10 | rgb_x = [item.float().permute(0, 3, 1, 2).to(device) for item in in_batch['rgb-x']] # [0,1] 11 | mask_v = [item.float().permute(0, 3, 1, 2).to(device) for item in in_batch['mask-v']] # 0,1 12 | mask_x = [item.float().permute(0, 3, 1, 2).to(device) for item in in_batch['mask-x']] # 0,1 13 | 14 | uv_v = [item.float().permute(0, 2, 1).unsqueeze(3).to(device) for item in in_batch['uv-v']] 15 | uv_x = [item.float().permute(0, 2, 1).unsqueeze(3).to(device) for item in in_batch['uv-x']] 16 | uv_mask_v = [item.float().permute(0, 2, 1).unsqueeze(3).to(device) for item in in_batch['uv-mask-v']] 17 | uv_mask_x = [item.float().permute(0, 2, 1).unsqueeze(3).to(device) for item in in_batch['uv-mask-x']] 18 | uv_xyz_v = [item.float().permute(0, 2, 1).unsqueeze(3).to(device) for item in in_batch['uv-xyz-v']] 19 | uv_xyz_x = [item.float().permute(0, 2, 1).unsqueeze(3).to(device) for item in in_batch['uv-xyz-x']] 20 | 21 | crr_idx_mtx = [[ii.long().permute(0, 2, 1).unsqueeze(3).to(device) for ii in item] 22 | for item in in_batch['crr-idx-mtx']] 23 | crr_mask_mtx = [[ii.float().permute(0, 2, 1).unsqueeze(3).to(device) for ii in item] 24 | for item in in_batch['crr-mask-mtx']] 25 | 26 | pack = {'rgb-v': rgb_v, 'rgb-x': rgb_x, 'nox-v': nox_v, 'nox-x': nox_x, 27 | 'mask-v': mask_v, 'mask-x': mask_x, 28 | 'uv-v': uv_v, 'uv-x': uv_x, 29 | 'uv-mask-v': uv_mask_v, 'uv-mask-x': uv_mask_x, 30 | 'uv-xyz-v': uv_xyz_v, 'uv-xyz-x': uv_xyz_x, 31 | 'crr-idx-mtx': crr_idx_mtx, 'crr-mask-mtx': crr_mask_mtx} 32 | return {'to-nn': pack, 'meta-info': in_batch['info']} 33 | 34 | 35 | def load_singleview_batch(in_batch): 36 | device = torch.device("cuda") 37 | nox_v = in_batch['nox-v'].float().permute(0, 3, 1, 2).to(device) / 255.0 # [0,1] 38 | nox_x = in_batch['nox-x'].float().permute(0, 3, 1, 2).to(device) / 255.0 # [0,1] 39 | rgb_v = in_batch['rgb-v'].float().permute(0, 3, 1, 2).to(device) / 255.0 # [0,1] 40 | rgb_x = in_batch['rgb-x'].float().permute(0, 3, 1, 2).to(device) / 255.0 # [0,1] 41 | mask_v = in_batch['mask-v'].float().permute(0, 3, 1, 2).to(device) # 0,1 42 | mask_x = in_batch['mask-x'].float().permute(0, 3, 1, 2).to(device) # 0,1 43 | 44 | uv_v = in_batch['uv-v'].float().permute(0, 2, 1).unsqueeze(3).to(device) 45 | uv_x = in_batch['uv-x'].float().permute(0, 2, 1).unsqueeze(3).to(device) 46 | uv_mask_v = in_batch['uv-mask-v'].float().permute(0, 2, 1).unsqueeze(3).to(device) 47 | uv_mask_x = in_batch['uv-mask-x'].float().permute(0, 2, 1).unsqueeze(3).to(device) 48 | uv_xyz_v = in_batch['uv-xyz-v'].float().permute(0, 2, 1).unsqueeze(3).to(device) / 255 49 | uv_xyz_x = in_batch['uv-xyz-x'].float().permute(0, 2, 1).unsqueeze(3).to(device) / 255 50 | pack = {'rgb-v': rgb_v, 'rgb-x': rgb_x, 'nox-v': nox_v, 'nox-x': nox_x, 51 | 'mask-v': mask_v, 'mask-x': mask_x, 52 | 'uv-v': uv_v, 'uv-x': uv_x, 53 | 'uv-mask-v': uv_mask_v, 'uv-mask-x': uv_mask_x, 54 | 'uv-xyz-v': uv_xyz_v, 'uv-xyz-x': uv_xyz_x} 55 | return {'to-nn': pack, 'meta-info': in_batch['info']} 56 | 57 | 58 | def spread_feature(container, learned_uv, feature, mask1c): 59 | """ 60 | :param container: B,C,R,R 61 | :param learned_uv: B,2,H,W 62 | :param feature: B,C,H,W aligned with latent uv map 63 | :param mask1c: B,1,H,W used to mask latent uv and feature 64 | :return: container 65 | """ 66 | assert float(mask1c.max()) < (1.0 + 1e-9) 67 | assert container.shape[1] == feature.shape[1] 68 | c = container.shape[1] 69 | res = container.shape[2] 70 | _learned_uv = learned_uv * mask1c.repeat(1, 2, 1, 1) 71 | _feature = feature * mask1c.repeat(1, c, 1, 1) 72 | learned_uv = torch.clamp((_learned_uv * res).long(), 0, res - 1) 73 | learned_uv = learned_uv.reshape(learned_uv.shape[0], 2, -1) 74 | learned_uv = learned_uv[:, 0, :] * res + learned_uv[:, 1, :] # B, R*R 75 | learned_uv = learned_uv.unsqueeze(1).repeat(1, c, 1) # B,C,R*R 76 | container = container.reshape(container.shape[0], container.shape[1], -1) 77 | container = container.scatter(2, learned_uv, _feature.reshape(feature.shape[0], c, -1)) 78 | container = container.reshape(container.shape[0], container.shape[1], res, res) 79 | return container 80 | 81 | 82 | def query_feature(feature_map, query_uv): 83 | """ 84 | query features from feature map 85 | :param feature_map: B,C,res1,res2 86 | :param query_uv: B,2,K,1 in [0,1] 87 | :return B,C,K,1 88 | """ 89 | assert float(query_uv.max()) < 1 + 1e-9 90 | assert query_uv.shape[1] == 2 91 | res1 = feature_map.shape[2] 92 | res2 = feature_map.shape[3] 93 | query_index = query_uv.clone() 94 | query_index[:, 0, ...] = torch.clamp((query_uv[:, 0, ...] * res1).long(), 0, res1 - 1) 95 | query_index[:, 1, ...] = torch.clamp((query_uv[:, 1, ...] * res2).long(), 0, res2 - 1) 96 | if query_index.ndimension() > 3: 97 | index = query_index.squeeze(3) # B,2,K 98 | else: 99 | index = query_index # B*2*K 100 | # cvt to 1D index 101 | index = index[:, 0, :] * feature_map.shape[3] + index[:, 1, :] # B,K 102 | index = index.unsqueeze(2).repeat(1, 1, feature_map.shape[1]) 103 | flatten_feature_map = feature_map.reshape(feature_map.shape[0], 104 | feature_map.shape[1], -1).permute(0, 2, 1) # B,N,C 105 | query = torch.gather(flatten_feature_map, 1, index.long()).contiguous() # B,K,C 106 | query = query.permute(0, 2, 1).unsqueeze(3) 107 | return query 108 | 109 | 110 | def make_grid(res, return_np=False): 111 | dim0 = np.arange(0, res[0]) + 0.5 112 | dim0 = dim0 / len(dim0) 113 | dim1 = np.arange(0, res[1]) + 0.5 114 | dim1 = dim1 / len(dim1) 115 | col_uv, row_uv = np.meshgrid(dim1, dim0) 116 | super_uv = np.concatenate((row_uv[..., np.newaxis], col_uv[..., np.newaxis]), 2) # R,R,2 117 | super_uv_tensor = torch.from_numpy(super_uv.astype(np.float32)) 118 | if return_np: 119 | return super_uv_tensor.permute(2, 0, 1).unsqueeze(0), super_uv 120 | else: 121 | return super_uv_tensor.permute(2, 0, 1).unsqueeze(0) 122 | -------------------------------------------------------------------------------- /core/models/utils/render_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import cv2 as cv 3 | import numpy as np 4 | from sklearn.neighbors import NearestNeighbors 5 | from .model_utils import spread_feature 6 | 7 | def optimize_image_mask(image_mask, sp_image, nK=4, th=1e-2): 8 | mask_pts = image_mask.reshape(-1) 9 | xyz_pts = sp_image.reshape(-1, 3) 10 | xyz_pts = xyz_pts[mask_pts > 0.5, :] 11 | Neighbors = NearestNeighbors(n_neighbors=nK + 1, algorithm='kd_tree').fit(xyz_pts) 12 | nn_dist, nn_idx = Neighbors.kneighbors(xyz_pts) # N,nK 13 | nn_dist = nn_dist[:, 1:] 14 | valid = (np.sum((nn_dist < th).astype(np.float), axis=1) == nK).astype(np.float) 15 | optimized_mask = image_mask.reshape(-1) 16 | optimized_mask[mask_pts > 0.5] = valid 17 | optimized_mask = optimized_mask.reshape(image_mask.shape[0], image_mask.shape[1]) 18 | return optimized_mask 19 | 20 | 21 | def generate_final_mask(image_learned_uv, image_mask, 22 | image_resize_factor, mask_container_low_res, final_gim): 23 | """ 24 | Post Process Algorithm to generate mask of the unwrapped chart 25 | Parameters 26 | ---------- 27 | image_learned_uv: [H,W,2] 28 | image_mask: [H,W] 29 | image_resize_factor: float 30 | mask_container_low_res: a predefined tensor with intermediate low resolution 31 | final_gim: a predefined tensor with target high resolution 32 | """ 33 | # resize (larger) rgb and uv with Bi-linear up-sampling 34 | resized_uv = cv.resize(image_learned_uv, dsize=(image_resize_factor * image_learned_uv.shape[0], 35 | image_resize_factor * image_learned_uv.shape[1]), 36 | interpolation=cv.INTER_LINEAR) 37 | resized_mask = cv.resize(image_mask, dsize=(image_resize_factor * image_learned_uv.shape[0], 38 | image_resize_factor * image_learned_uv.shape[1]), 39 | interpolation=cv.INTER_LINEAR) 40 | resized_mask = (resized_mask > 0.5).astype(np.float) 41 | # use gradient to remove the edge 42 | discontinuous_mask_u = cv.Laplacian(image_learned_uv[..., 0], ddepth=cv.CV_32F) # small gradient map 43 | discontinuous_mask_v = cv.Laplacian(image_learned_uv[..., 1], ddepth=cv.CV_32F) # small gradient map 44 | # use the max and min in latent u and v to find the threshhold 45 | u_max = (image_learned_uv[..., 0] * image_mask).max() 46 | v_max = (image_learned_uv[..., 1] * image_mask).max() 47 | u_min = (image_learned_uv[..., 0] * image_mask + (1.0 - image_mask)).min() 48 | v_min = (image_learned_uv[..., 1] * image_mask + (1.0 - image_mask)).min() 49 | u_th = (u_max - u_min) / 30 50 | v_th = (v_max - v_min) / 30 51 | discontinuous_mask_u = (discontinuous_mask_u > u_th).astype(np.float) * image_mask 52 | discontinuous_mask_v = (discontinuous_mask_v > v_th).astype(np.float) * image_mask 53 | discontinuous_mask = ((discontinuous_mask_u + discontinuous_mask_v) > 0).astype(np.float) 54 | # use the mask to remove the boundary 55 | boundary_recovery_mask = (cv.Laplacian(image_mask, ddepth=cv.CV_32F) > 0.01).astype(np.float) 56 | discontinuous_mask = discontinuous_mask * (1.0 - boundary_recovery_mask) 57 | resized_discontinuous_mask = cv.resize(discontinuous_mask, 58 | dsize=(image_resize_factor * image_learned_uv.shape[0], 59 | image_resize_factor * image_learned_uv.shape[1]), 60 | interpolation=cv.INTER_NEAREST) 61 | # make the small mask & texture 62 | high_res_mask = torch.from_numpy(resized_mask * (1.0 - resized_discontinuous_mask)) \ 63 | .unsqueeze(0).unsqueeze(0).cuda().float() # 1,1,R,R 64 | high_res_uv = torch.from_numpy(resized_uv).permute(2, 0, 1).unsqueeze(0).cuda().float() 65 | low_res_mask = mask_container_low_res.cuda() 66 | low_res_mask = spread_feature(low_res_mask, high_res_uv, high_res_mask, high_res_mask) 67 | # use close to remove the holes in small mask and then resize 68 | low_res_mask_closed = low_res_mask.detach().cpu().squeeze(0).squeeze(0).numpy() # R,R 69 | close_k_size = int(final_gim.shape[2] / 100) 70 | close_kernel = cv.getStructuringElement(cv.MORPH_ELLIPSE, (close_k_size, close_k_size)) 71 | final_mask_np = cv.resize(low_res_mask_closed, dsize=(final_gim.shape[2], 72 | final_gim.shape[2]), 73 | interpolation=cv.INTER_NEAREST) # R,R,3 74 | final_mask_np = (final_mask_np > 0).astype(np.float) 75 | final_mask_np = cv.morphologyEx(final_mask_np, cv.MORPH_OPEN, close_kernel) 76 | return final_mask_np 77 | 78 | 79 | def generate_texture(sp_image, full_gim, image_rgb, image_mask, final_mask_np, final_res, nK=4, th=1e-2): 80 | # prepare root and query points form the image and from the high-res chart 81 | root_xyz_np = sp_image.reshape(-1, 3) # H*W,3 82 | root_rgb_np = image_rgb.reshape(-1, 3) # H*W,3 83 | _image_mask = image_mask.reshape(-1) # H*W 84 | root_xyz_np = root_xyz_np[_image_mask > 0.5, :] # M,2 [0,1] 85 | root_rgb_np = root_rgb_np[_image_mask > 0.5, :] # M,3 [0,1] 86 | query_xyz_np = full_gim.reshape(-1, 3) # R*R,3 87 | _final_mask_np = final_mask_np.reshape(-1) # R*R 88 | query_xyz_np = query_xyz_np[_final_mask_np > 0.5, :] # N,3 [0,1] 89 | # finding nearest root pixel points 90 | Neighbors = NearestNeighbors(n_neighbors=nK, algorithm='kd_tree').fit(root_xyz_np) 91 | nn_dist, nn_idx = Neighbors.kneighbors(query_xyz_np) # N,nK 92 | # optimize the gim mask 93 | valid = (nn_dist[:, 0] < th).astype(np.float) 94 | optimized_final_mask_np = final_mask_np.reshape(-1).copy() 95 | optimized_final_mask_np[_final_mask_np > 0.5] = valid 96 | optimized_final_mask_np = optimized_final_mask_np.reshape(final_mask_np.shape[0], final_mask_np.shape[1]) 97 | # do interpolation based on chart distance 98 | interpolation_weight = nn_dist.copy() 99 | interpolation_weight = 1 - interpolation_weight / np.sum(interpolation_weight, 1, keepdims=True) 100 | interpolation_weight = interpolation_weight / np.sum(interpolation_weight, 1, keepdims=True) 101 | query_rgb_np = np.zeros((query_xyz_np.shape[0], 3)) 102 | for kdx in range(nK): 103 | nn_color = root_rgb_np[nn_idx[:, kdx], :] 104 | query_rgb_np += nn_color * interpolation_weight[:, kdx][..., np.newaxis] 105 | final_texture_np = np.ones((final_res ** 2, 3)) 106 | final_texture_np[_final_mask_np > 0.5, :] = query_rgb_np 107 | final_texture_np = final_texture_np.reshape(final_res, final_res, 3) 108 | return final_texture_np, optimized_final_mask_np -------------------------------------------------------------------------------- /core/net_bank/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JiahuiLei/Pix2Surf/ac56ce5a20adab8cb563d96c915e1cc253631bf8/core/net_bank/__init__.py -------------------------------------------------------------------------------- /core/net_bank/loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | 4 | 5 | class MaskL2Loss(nn.Module): 6 | def __init__(self): 7 | super(MaskL2Loss, self).__init__() 8 | 9 | def forward(self, pred, target, mask, detach=True, reduce_batch=True): 10 | assert mask.max() <= 1 + 1e-6 11 | if detach: 12 | target = target.detach() 13 | mask = mask.detach() 14 | if mask.shape[1] > 1: 15 | mask = mask[:, 0, :, :].unsqueeze(1) 16 | assert pred.shape == target.shape 17 | dif = (pred - target) ** 2 * (mask.float()) 18 | loss = torch.sum(dif.reshape(mask.shape[0], -1).contiguous(), 1) 19 | count = torch.sum(mask.reshape(mask.shape[0], -1).contiguous(), 1).detach() 20 | loss[count == 0] = loss[count == 0] * 0 21 | loss = loss / (count + 1) 22 | if reduce_batch: 23 | non_zero_count = torch.sum((count > 0).float()) 24 | if non_zero_count == 0: 25 | loss = torch.sum(loss) * 0 26 | else: 27 | loss = torch.sum(loss) / non_zero_count 28 | return loss 29 | else: 30 | return loss 31 | 32 | 33 | if __name__ == "__main__": 34 | l = MaskL2Loss().cuda() 35 | # l2 = MaskMSELoss().cuda() 36 | x = torch.rand(2, 3, 128, 128).cuda() 37 | y = torch.rand(2, 3, 128, 128).cuda() 38 | m = torch.cat((torch.zeros(1, 3, 128, 128), torch.zeros(1, 3, 128, 128)), 0).cuda() 39 | m = (m > 0.5).float() 40 | loss = l(x, y, m) 41 | # loss2 = l2(x, y, m) 42 | -------------------------------------------------------------------------------- /core/net_bank/mlp.py: -------------------------------------------------------------------------------- 1 | """ 2 | 2019.8.10 3 | """ 4 | import torch 5 | from torch import nn 6 | 7 | 8 | class BasicMLPBlock(nn.Module): 9 | def __init__(self, f_in, p_in, f_out, p_out): 10 | super().__init__() 11 | self.f_in = f_in 12 | self.f_out = f_out 13 | self.p_in = p_in 14 | self.p_out = p_out 15 | self.f = nn.Sequential( 16 | nn.Conv2d(in_channels=f_in + p_in, out_channels=f_out, kernel_size=1, stride=1, padding=0), 17 | nn.ReLU(inplace=True), 18 | nn.Conv2d(in_channels=f_out, out_channels=f_out + p_out, kernel_size=1, stride=1, padding=0), 19 | ) 20 | self.relu = nn.ReLU(inplace=True) 21 | 22 | def forward(self, f_and_p): 23 | assert f_and_p.shape[1] == self.f_in + self.p_in 24 | raw_feature = self.f(f_and_p) 25 | f = self.relu(raw_feature[:, self.p_out:, :, :]) 26 | xyz = raw_feature[:, :self.p_out, :, :] 27 | act_feature = torch.cat((xyz, f), 1) 28 | return act_feature 29 | 30 | 31 | class NOCS_MLP(nn.Module): 32 | def __init__(self, c_in=512, p_in=2, p_out=3, bw=None, compress=True): 33 | super().__init__() 34 | self.bw = c_in if bw is None else bw 35 | self.p_out = p_out 36 | self.block1 = BasicMLPBlock(f_in=c_in, p_in=p_in, f_out=self.bw, p_out=p_out) 37 | self.block2 = BasicMLPBlock(f_in=self.bw, p_in=p_out, f_out=self.bw, p_out=p_out) 38 | self.block3v = BasicMLPBlock(f_in=self.bw, p_in=p_out, f_out=self.bw, p_out=p_out) 39 | self.block4v = BasicMLPBlock(f_in=self.bw, p_in=p_out, f_out=self.bw, p_out=p_out) 40 | self.layer5v = nn.Conv2d(in_channels=self.bw + p_out, out_channels=p_out, kernel_size=1, stride=1, padding=0) 41 | self.compress_fn = nn.Sigmoid() 42 | self.compress = compress 43 | 44 | def forward(self, code, uv, unique_code=True): 45 | if unique_code: 46 | input = torch.cat((code.unsqueeze(2).unsqueeze(2).repeat(1, 1, uv.shape[2], uv.shape[3]), uv), 47 | 1).contiguous() 48 | else: 49 | input = torch.cat((code, uv), 1).contiguous() 50 | 51 | f1 = self.block1(input) 52 | xyz1 = f1[:, :self.p_out, :, :] 53 | 54 | f2 = self.block2(f1) 55 | xyz2 = f2[:, :self.p_out, :, :] 56 | 57 | f3_v = self.block3v(f2) 58 | xyz3_v = f3_v[:, :self.p_out, :, :] 59 | 60 | f4_v = self.block4v(f3_v) 61 | xyz4_v = f4_v[:, :self.p_out, :, :] 62 | 63 | xyz5_v = self.layer5v(f4_v) 64 | 65 | if self.compress: 66 | xyz_v = self.compress_fn(xyz1 + xyz2 + xyz3_v + xyz4_v + xyz5_v) 67 | else: 68 | xyz_v = xyz1 + xyz2 + xyz3_v + xyz4_v + xyz5_v 69 | return xyz_v 70 | 71 | class NOCS_AMP_MLP(nn.Module): 72 | def __init__(self, latent_dim=1024, amp_dim=256, p_in=2, c_out=3): 73 | super().__init__() 74 | self.main_mlp = NOCS_MLP(c_in=latent_dim, p_in=amp_dim, p_out=c_out, bw=512, compress=False) 75 | self.amp1 = torch.nn.Conv2d(p_in, amp_dim // 4, 1) 76 | self.amp2 = torch.nn.Conv2d(amp_dim // 4, amp_dim // 2, 1) 77 | self.amp3 = torch.nn.Conv2d(amp_dim // 2, amp_dim, 1) 78 | self.activation = nn.ReLU() 79 | 80 | def forward(self, code, uv, unique_code=True): 81 | am = self.activation(self.amp1(uv)) 82 | am = self.activation(self.amp2(am)) 83 | am = self.activation(self.amp3(am)) 84 | return self.main_mlp(code, am, unique_code) 85 | 86 | if __name__ == "__main__": 87 | uv = torch.rand(5, 2, 100, 1).cuda() 88 | c = torch.rand(5, 512).cuda() 89 | net = NOCS_MLP().cuda() 90 | y = net(c, uv) 91 | print(y.shape) 92 | 93 | net = NOCS_AMP_MLP(1024, 256, 2, 3).cuda() 94 | c = torch.rand(3, 1024).cuda() 95 | uv = torch.rand(3, 2, 100, 1).cuda() 96 | xyz = net(c, uv, True) 97 | print(xyz.shape) 98 | 99 | -------------------------------------------------------------------------------- /core/net_bank/modules/__init__.py: -------------------------------------------------------------------------------- 1 | from .neurips_nox_network_util import * -------------------------------------------------------------------------------- /core/net_bank/pix2surf_cnn.py: -------------------------------------------------------------------------------- 1 | # Borrowed from https://github.com/meetshah1995/pytorch-semseg 2 | # refer to https://github.com/drsrinathsridhar/xnocs 3 | # minor modified 4 | 5 | import torch.nn as nn 6 | import torchvision.models as models 7 | from core.net_bank.modules import segnetDown2, segnetDown3, segnetUp2, segnetUp3 8 | import torch 9 | 10 | class SegNet(nn.Module): 11 | def __init__(self, out_channels=8, in_channels=3, 12 | pretrained=True, withSkipConnections=True, new_version=False, additional=None): 13 | """ 14 | :param out_channels: 15 | :param in_channels: 16 | :param pretrained: 17 | :param withSkipConnections: 18 | :param new_version: 19 | :param additional: all additional output layer are new version 20 | """ 21 | super().__init__() 22 | 23 | self.in_channels = in_channels 24 | self.withSkipConnections = withSkipConnections 25 | 26 | self.down1 = segnetDown2(self.in_channels, 64, withFeatureMap=self.withSkipConnections) 27 | self.down2 = segnetDown2(64, 128, withFeatureMap=self.withSkipConnections) 28 | self.down3 = segnetDown3(128, 256, withFeatureMap=self.withSkipConnections) 29 | self.down4 = segnetDown3(256, 512, withFeatureMap=self.withSkipConnections) 30 | self.down5 = segnetDown3(512, 512, withFeatureMap=self.withSkipConnections) 31 | 32 | self.up5 = segnetUp3(512, 512, withSkipConnections=self.withSkipConnections) 33 | self.up4 = segnetUp3(512, 256, withSkipConnections=self.withSkipConnections) 34 | self.up3 = segnetUp3(256, 128, withSkipConnections=self.withSkipConnections) 35 | self.up2 = segnetUp2(128, 64, withSkipConnections=self.withSkipConnections) 36 | self.up1 = segnetUp2(64, out_channels, last_layer=True if new_version else False, 37 | withSkipConnections=self.withSkipConnections) 38 | if additional is not None: 39 | self.additional_last_layer = segnetUp2(64, additional, last_layer=True, 40 | withSkipConnections=self.withSkipConnections) 41 | self.additional = True 42 | else: 43 | self.additional = False 44 | 45 | if pretrained: 46 | vgg16 = models.vgg16(pretrained=True) 47 | Arch = 'SegNet' 48 | if self.withSkipConnections: 49 | Arch = 'SegNetSkip' 50 | print('[ INFO ]: Using pre-trained weights from VGG16 with {}.'.format(Arch)) 51 | self.init_vgg16_params(vgg16) 52 | 53 | def forward(self, inputs, return_code=False): 54 | down1, indices_1, unpool_shape1, FM1 = self.down1(inputs) 55 | down2, indices_2, unpool_shape2, FM2 = self.down2(down1) 56 | down3, indices_3, unpool_shape3, FM3 = self.down3(down2) 57 | down4, indices_4, unpool_shape4, FM4 = self.down4(down3) 58 | down5, indices_5, unpool_shape5, FM5 = self.down5(down4) 59 | 60 | up5 = self.up5(down5, indices_5, unpool_shape5, SkipFeatureMap=FM5) 61 | up4 = self.up4(up5, indices_4, unpool_shape4, SkipFeatureMap=FM4) 62 | up3 = self.up3(up4, indices_3, unpool_shape3, SkipFeatureMap=FM3) 63 | up2 = self.up2(up3, indices_2, unpool_shape2, SkipFeatureMap=FM2) 64 | up1 = self.up1(up2, indices_1, unpool_shape1, SkipFeatureMap=FM1) 65 | 66 | # # DEBUG: print sizes 67 | # print('down1:', down1.size()) 68 | # print('down2:', down2.size()) 69 | # print('down3:', down3.size()) 70 | # print('down4:', down4.size()) 71 | # print('down5:', down5.size()) 72 | # 73 | # print('up5:', up5.size()) 74 | # print('up4:', up4.size()) 75 | # print('up3:', up3.size()) 76 | # print('up2:', up2.size()) 77 | # print('up1:', up1.size()) 78 | if self.additional: 79 | add = self.additional_last_layer(up2, indices_1, unpool_shape1, SkipFeatureMap=FM1) 80 | up1 = torch.cat((up1, add), dim=1) 81 | if return_code: 82 | return up1, down5 83 | else: 84 | return up1 85 | 86 | def init_vgg16_params(self, vgg16): 87 | blocks = [self.down1, self.down2, self.down3, self.down4, self.down5] 88 | 89 | features = list(vgg16.features.children()) 90 | 91 | vgg_layers = [] 92 | for _layer in features: 93 | if isinstance(_layer, nn.Conv2d): 94 | vgg_layers.append(_layer) 95 | 96 | merged_layers = [] 97 | for idx, conv_block in enumerate(blocks): 98 | if idx < 2: 99 | units = [conv_block.conv1.cbr_unit, conv_block.conv2.cbr_unit] 100 | else: 101 | units = [ 102 | conv_block.conv1.cbr_unit, 103 | conv_block.conv2.cbr_unit, 104 | conv_block.conv3.cbr_unit, 105 | ] 106 | for _unit in units: 107 | for _layer in _unit: 108 | if isinstance(_layer, nn.Conv2d): 109 | merged_layers.append(_layer) 110 | 111 | assert len(vgg_layers) == len(merged_layers) 112 | 113 | for l1, l2 in zip(vgg_layers, merged_layers): 114 | if isinstance(l1, nn.Conv2d) and isinstance(l2, nn.Conv2d): 115 | assert l1.weight.size() == l2.weight.size() 116 | assert l1.bias.size() == l2.bias.size() 117 | l2.weight.data = l1.weight.data 118 | l2.bias.data = l1.bias.data 119 | 120 | 121 | class SegNetGroup(nn.Module): 122 | def __init__(self, out_channels=8, in_channels=3, 123 | pretrained=True, withSkipConnections=True, new_version=False, additional=None): 124 | """ 125 | :param out_channels: 126 | :param in_channels: 127 | :param pretrained: 128 | :param withSkipConnections: 129 | :param new_version: 130 | :param additional: all additional output layer are new version 131 | """ 132 | super().__init__() 133 | 134 | self.in_channels = in_channels 135 | self.withSkipConnections = withSkipConnections 136 | 137 | self.down1 = segnetDown2(self.in_channels, 64, withFeatureMap=self.withSkipConnections) 138 | self.down2 = segnetDown2(64, 128, withFeatureMap=self.withSkipConnections) 139 | self.down3 = segnetDown3(128, 256, withFeatureMap=self.withSkipConnections) 140 | self.down4 = segnetDown3(256, 512, withFeatureMap=self.withSkipConnections) 141 | self.down5 = segnetDown3(512, 512, withFeatureMap=self.withSkipConnections) 142 | 143 | self.up5 = segnetUp3(512, 512, withSkipConnections=self.withSkipConnections) 144 | self.up4 = segnetUp3(512, 256, withSkipConnections=self.withSkipConnections) 145 | self.up3 = segnetUp3(256, 128, withSkipConnections=self.withSkipConnections) 146 | self.up2 = segnetUp2(128, 64, withSkipConnections=self.withSkipConnections) 147 | self.up1 = segnetUp2(64, out_channels, last_layer=True if new_version else False, 148 | withSkipConnections=self.withSkipConnections) 149 | if additional is not None: 150 | self.additional_last_layer = segnetUp2(64, additional, last_layer=True, 151 | withSkipConnections=self.withSkipConnections) 152 | self.additional = True 153 | else: 154 | self.additional = False 155 | 156 | if pretrained: 157 | vgg16 = models.vgg16(pretrained=True) 158 | Arch = 'SegNet' 159 | if self.withSkipConnections: 160 | Arch = 'SegNetSkip' 161 | print('[ INFO ]: Using pre-trained weights from VGG16 with {}.'.format(Arch)) 162 | self.init_vgg16_params(vgg16) 163 | 164 | def forward(self, inputs, return_code=False): 165 | """ 166 | :param inputs: is a list of rgb 167 | :param return_code: before the max pooling 168 | :return: 169 | """ 170 | l1, l2, l3, l4, l5 = [], [], [], [], [] 171 | l5_feature = [] 172 | for rgb in inputs: 173 | down1, indices_1, unpool_shape1, FM1 = self.down1(rgb) 174 | down2, indices_2, unpool_shape2, FM2 = self.down2(down1) 175 | down3, indices_3, unpool_shape3, FM3 = self.down3(down2) 176 | down4, indices_4, unpool_shape4, FM4 = self.down4(down3) 177 | down5, indices_5, unpool_shape5, FM5 = self.down5(down4) 178 | l1.append([indices_1, unpool_shape1, FM1]) 179 | l2.append([indices_2, unpool_shape2, FM2]) 180 | l3.append([indices_3, unpool_shape3, FM3]) 181 | l4.append([indices_4, unpool_shape4, FM4]) 182 | l5.append([indices_5, unpool_shape5, FM5]) 183 | l5_feature.append(down5.unsqueeze(0)) 184 | max_pooled_feature = torch.max(torch.cat(l5_feature, 0), dim=0).values 185 | f_dim = max_pooled_feature.shape[1] 186 | 187 | pred_list = [] 188 | for i in range(len(inputs)): 189 | down5 = torch.cat((max_pooled_feature[:, :f_dim // 2, :, :], 190 | l5_feature[i].squeeze(0)[:, f_dim // 2:, :, :]), dim=1) 191 | up5 = self.up5(down5, *l5[i]) 192 | up4 = self.up4(up5, *l4[i]) 193 | up3 = self.up3(up4, *l3[i]) 194 | up2 = self.up2(up3, *l2[i]) 195 | up1 = self.up1(up2, *l1[i]) 196 | 197 | if self.additional: 198 | add = self.additional_last_layer(up2, *l1[i]) 199 | up1 = torch.cat((up1, add), dim=1) 200 | pred_list.append(up1) 201 | feature_list = [item.squeeze(0) for item in l5_feature] 202 | 203 | if return_code: 204 | return pred_list, feature_list 205 | else: 206 | return pred_list 207 | 208 | def init_vgg16_params(self, vgg16): 209 | blocks = [self.down1, self.down2, self.down3, self.down4, self.down5] 210 | 211 | features = list(vgg16.features.children()) 212 | 213 | vgg_layers = [] 214 | for _layer in features: 215 | if isinstance(_layer, nn.Conv2d): 216 | vgg_layers.append(_layer) 217 | 218 | merged_layers = [] 219 | for idx, conv_block in enumerate(blocks): 220 | if idx < 2: 221 | units = [conv_block.conv1.cbr_unit, conv_block.conv2.cbr_unit] 222 | else: 223 | units = [ 224 | conv_block.conv1.cbr_unit, 225 | conv_block.conv2.cbr_unit, 226 | conv_block.conv3.cbr_unit, 227 | ] 228 | for _unit in units: 229 | for _layer in _unit: 230 | if isinstance(_layer, nn.Conv2d): 231 | merged_layers.append(_layer) 232 | 233 | assert len(vgg_layers) == len(merged_layers) 234 | 235 | for l1, l2 in zip(vgg_layers, merged_layers): 236 | if isinstance(l1, nn.Conv2d) and isinstance(l2, nn.Conv2d): 237 | assert l1.weight.size() == l2.weight.size() 238 | assert l1.bias.size() == l2.bias.size() 239 | l2.weight.data = l1.weight.data 240 | l2.bias.data = l1.bias.data 241 | 242 | 243 | class SegNetEncoder(nn.Module): 244 | def __init__(self, in_channels=3, withSkipConnections=True): 245 | """ 246 | :param in_channels: 247 | :param pretrained: 248 | :param withSkipConnections: 249 | """ 250 | super().__init__() 251 | 252 | self.in_channels = in_channels 253 | self.withSkipConnections = withSkipConnections 254 | 255 | self.down1 = segnetDown2(self.in_channels, 64, withFeatureMap=self.withSkipConnections) 256 | self.down2 = segnetDown2(64, 128, withFeatureMap=self.withSkipConnections) 257 | self.down3 = segnetDown3(128, 256, withFeatureMap=self.withSkipConnections) 258 | self.down4 = segnetDown3(256, 512, withFeatureMap=self.withSkipConnections) 259 | self.down5 = segnetDown3(512, 512, withFeatureMap=self.withSkipConnections) 260 | 261 | def forward(self, inputs, return_code=False): 262 | down1, indices_1, unpool_shape1, FM1 = self.down1(inputs) 263 | down2, indices_2, unpool_shape2, FM2 = self.down2(down1) 264 | down3, indices_3, unpool_shape3, FM3 = self.down3(down2) 265 | down4, indices_4, unpool_shape4, FM4 = self.down4(down3) 266 | down5, indices_5, unpool_shape5, FM5 = self.down5(down4) 267 | return down5 268 | 269 | 270 | if __name__ == '__main__': 271 | import torch 272 | 273 | net = SegNetGroup(withSkipConnections=True, out_channels=8).cuda() 274 | x = torch.rand(2, 3, 320, 240).cuda() 275 | y, f = net([x, x], return_code=True) 276 | print(f.shape) 277 | print(y.shape) 278 | -------------------------------------------------------------------------------- /core/net_bank/xnocs_segnet.py: -------------------------------------------------------------------------------- 1 | # Borrowed from https://github.com/meetshah1995/pytorch-semseg 2 | # refer to https://github.com/drsrinathsridhar/xnocs 3 | 4 | import torch.nn as nn 5 | import torchvision.models as models 6 | from core.net_bank.modules import segnetDown2, segnetDown3, segnetUp2, segnetUp3 7 | 8 | 9 | # import os, sys 10 | # FileDirPath = os.path.dirname(os.path.realpath(__file__)) 11 | # sys.path.append(os.path.join(FileDirPath, '..')) 12 | # from tk3dv.ptTools import ptUtils 13 | # from tk3dv.ptTools import ptNets 14 | # sys.path.append(os.path.join(FileDirPath, '.')) 15 | 16 | class SegNet(nn.Module): 17 | def __init__(self, out_channels=8, in_channels=3, 18 | pretrained=True, withSkipConnections=True): 19 | super().__init__() 20 | 21 | self.in_channels = in_channels 22 | self.withSkipConnections = withSkipConnections 23 | 24 | self.down1 = segnetDown2(self.in_channels, 64, withFeatureMap=self.withSkipConnections) 25 | self.down2 = segnetDown2(64, 128, withFeatureMap=self.withSkipConnections) 26 | self.down3 = segnetDown3(128, 256, withFeatureMap=self.withSkipConnections) 27 | self.down4 = segnetDown3(256, 512, withFeatureMap=self.withSkipConnections) 28 | self.down5 = segnetDown3(512, 512, withFeatureMap=self.withSkipConnections) 29 | 30 | self.up5 = segnetUp3(512, 512, withSkipConnections=self.withSkipConnections) 31 | self.up4 = segnetUp3(512, 256, withSkipConnections=self.withSkipConnections) 32 | self.up3 = segnetUp3(256, 128, withSkipConnections=self.withSkipConnections) 33 | self.up2 = segnetUp2(128, 64, withSkipConnections=self.withSkipConnections) 34 | self.up1 = segnetUp2(64, out_channels, withSkipConnections=self.withSkipConnections) 35 | 36 | if pretrained: 37 | vgg16 = models.vgg16(pretrained=True) 38 | Arch = 'SegNet' 39 | if self.withSkipConnections: 40 | Arch = 'SegNetSkip' 41 | print('[ INFO ]: Using pre-trained weights from VGG16 with {}.'.format(Arch)) 42 | self.init_vgg16_params(vgg16) 43 | 44 | def forward(self, inputs): 45 | down1, indices_1, unpool_shape1, FM1 = self.down1(inputs) 46 | down2, indices_2, unpool_shape2, FM2 = self.down2(down1) 47 | down3, indices_3, unpool_shape3, FM3 = self.down3(down2) 48 | down4, indices_4, unpool_shape4, FM4 = self.down4(down3) 49 | down5, indices_5, unpool_shape5, FM5 = self.down5(down4) 50 | 51 | up5 = self.up5(down5, indices_5, unpool_shape5, SkipFeatureMap=FM5) 52 | up4 = self.up4(up5, indices_4, unpool_shape4, SkipFeatureMap=FM4) 53 | up3 = self.up3(up4, indices_3, unpool_shape3, SkipFeatureMap=FM3) 54 | up2 = self.up2(up3, indices_2, unpool_shape2, SkipFeatureMap=FM2) 55 | up1 = self.up1(up2, indices_1, unpool_shape1, SkipFeatureMap=FM1) 56 | 57 | # # DEBUG: print sizes 58 | # print('down1:', down1.size()) 59 | # print('down2:', down2.size()) 60 | # print('down3:', down3.size()) 61 | # print('down4:', down4.size()) 62 | # print('down5:', down5.size()) 63 | # 64 | # print('up5:', up5.size()) 65 | # print('up4:', up4.size()) 66 | # print('up3:', up3.size()) 67 | # print('up2:', up2.size()) 68 | # print('up1:', up1.size()) 69 | 70 | return up1 71 | 72 | def init_vgg16_params(self, vgg16): 73 | blocks = [self.down1, self.down2, self.down3, self.down4, self.down5] 74 | 75 | features = list(vgg16.features.children()) 76 | 77 | vgg_layers = [] 78 | for _layer in features: 79 | if isinstance(_layer, nn.Conv2d): 80 | vgg_layers.append(_layer) 81 | 82 | merged_layers = [] 83 | for idx, conv_block in enumerate(blocks): 84 | if idx < 2: 85 | units = [conv_block.conv1.cbr_unit, conv_block.conv2.cbr_unit] 86 | else: 87 | units = [ 88 | conv_block.conv1.cbr_unit, 89 | conv_block.conv2.cbr_unit, 90 | conv_block.conv3.cbr_unit, 91 | ] 92 | for _unit in units: 93 | for _layer in _unit: 94 | if isinstance(_layer, nn.Conv2d): 95 | merged_layers.append(_layer) 96 | 97 | assert len(vgg_layers) == len(merged_layers) 98 | 99 | for l1, l2 in zip(vgg_layers, merged_layers): 100 | if isinstance(l1, nn.Conv2d) and isinstance(l2, nn.Conv2d): 101 | assert l1.weight.size() == l2.weight.size() 102 | assert l1.bias.size() == l2.bias.size() 103 | l2.weight.data = l1.weight.data 104 | l2.bias.data = l1.bias.data 105 | 106 | 107 | if __name__ == '__main__': 108 | import torch 109 | 110 | net = SegNet(withSkipConnections=True, out_channels=8).cuda() 111 | x = torch.rand(2, 3, 640, 480).cuda() 112 | y = net(x) 113 | print(y.shape) 114 | -------------------------------------------------------------------------------- /core/net_bank/xnocs_setsegnet.py: -------------------------------------------------------------------------------- 1 | # Borrowed and extensively modified from https://github.com/meetshah1995/pytorch-semseg 2 | # refer to https://github.com/drsrinathsridhar/xnocs 3 | 4 | import torch.nn.functional as F 5 | import os, sys, torch 6 | import torch.nn as nn 7 | import torchvision.models as models 8 | from core.net_bank.modules import segnetDown2, segnetDown3, segnetUp2, segnetUp3 9 | 10 | 11 | # Permutation equivariant set version of SegNet (with or without skip connections) 12 | class SetSegNet(nn.Module): 13 | def __init__(self, n_classes=21, in_channels=3, is_unpooling=True, pretrained=True, withSkipConnections=False, 14 | enablePermEq=True): 15 | super().__init__() 16 | 17 | self.in_channels = in_channels 18 | self.is_unpooling = is_unpooling 19 | self.withSkipConnections = withSkipConnections 20 | self.enablePermEq = enablePermEq 21 | 22 | self.down1 = segnetDown2(self.in_channels, 64, withFeatureMap=self.withSkipConnections) 23 | self.down2 = segnetDown2(64, 128, withFeatureMap=self.withSkipConnections) 24 | self.down3 = segnetDown3(128, 256, withFeatureMap=self.withSkipConnections) 25 | self.down4 = segnetDown3(256, 512, withFeatureMap=self.withSkipConnections) 26 | self.down5 = segnetDown3(512, 512, withFeatureMap=self.withSkipConnections) 27 | 28 | self.up5 = segnetUp3(512, 512, withSkipConnections=self.withSkipConnections) 29 | self.up4 = segnetUp3(512, 256, withSkipConnections=self.withSkipConnections) 30 | self.up3 = segnetUp3(256, 128, withSkipConnections=self.withSkipConnections) 31 | self.up2 = segnetUp2(128, 64, withSkipConnections=self.withSkipConnections) 32 | self.up1 = segnetUp2(64, n_classes, withSkipConnections=self.withSkipConnections) 33 | 34 | if pretrained: 35 | vgg16 = models.vgg16(pretrained=True) 36 | Arch = 'SetSegNet' 37 | if self.withSkipConnections: 38 | Arch = 'SetSegNetSkip' 39 | print( 40 | '[ INFO ]: Using pre-trained weights from VGG16 with {}. Permutation equivariant layers are {}.'.format( 41 | Arch, 'ENABLED' if self.enablePermEq else 'DISABLED')) 42 | self.init_vgg16_params(vgg16) 43 | 44 | def avgSubtract(self, FeatureMap): 45 | # print('-'*50) 46 | # print('FeatureMap:', FeatureMap.size()) 47 | B, S, C, W, H = FeatureMap.size() 48 | FeatureMap_p = FeatureMap.view(B, S, -1) 49 | # print('View:', FeatureMap_p.size()) 50 | FeatureMap_p = FeatureMap_p.permute(0, 2, 1) # Set size should be the last dimension 51 | # print('Permuted:', FeatureMap_p.size()) 52 | # print('S:', S) 53 | # MP = FeatureMap_p[:, :, 0].unsqueeze(2) # TEMP TESTING TODO 54 | MP = F.avg_pool1d(FeatureMap_p, S) 55 | # MP = self.MaxPool1D(FeatureMap_p) 56 | # print('MaxPooled:', MP.size()) 57 | MP = MP.permute(0, 2, 1) # Undo previous permute 58 | # print('Permuted:', MP.size()) 59 | MP = MP.view(B, 1, C, W, H) 60 | # print('Final:', MP.size()) 61 | 62 | MS = [] 63 | for i in range(S): 64 | MS.append(FeatureMap[:, i, :, :, :].unsqueeze(1) - MP) 65 | 66 | MS = torch.cat(MS, dim=1) 67 | # print('MaxSubtracted:', MS.size()) 68 | # print('-' * 50) 69 | 70 | return MS 71 | 72 | def maxSubtract(self, FeatureMap): 73 | # TODO: There is a CUDA bug in max_pool1d, so using avgSubtract 74 | return self.avgSubtract(FeatureMap) 75 | 76 | # print('-'*50) 77 | # print('FeatureMap:', FeatureMap.size()) 78 | B, S, C, W, H = FeatureMap.size() 79 | FeatureMap_p = FeatureMap.view(B, S, -1) 80 | # print('View:', FeatureMap_p.size()) 81 | FeatureMap_p = FeatureMap_p.permute(0, 2, 1) # Set size should be the last dimension 82 | # print('Permuted:', FeatureMap_p.size()) 83 | # print('S:', S) 84 | # MP = FeatureMap_p[:, :, 0].unsqueeze(2) # TEMP TESTING TODO 85 | MP = F.max_pool1d(FeatureMap_p, S) 86 | # MP = self.MaxPool1D(FeatureMap_p) 87 | # print('MaxPooled:', MP.size()) 88 | MP = MP.permute(0, 2, 1) # Undo previous permute 89 | # print('Permuted:', MP.size()) 90 | MP = MP.view(B, 1, C, W, H) 91 | # print('Final:', MP.size()) 92 | 93 | MS = [] 94 | for i in range(S): 95 | MS.append(FeatureMap[:, i, :, :, :].unsqueeze(1) - MP) 96 | 97 | MS = torch.cat(MS, dim=1) 98 | # print('MaxSubtracted:', MS.size()) 99 | # print('-' * 50) 100 | 101 | return MS 102 | 103 | def forward(self, inputs): 104 | B, S, C, W, H = inputs.size() 105 | # print('B, S, C, W, H:', B, S, C, W, H) 106 | if self.withSkipConnections: 107 | downs1 = [] 108 | all_indices_1 = [] 109 | all_unpool_shape1 = [] 110 | all_FM1 = [] 111 | for s in range(S): 112 | down1, indices_1, unpool_shape1, FM1 = self.down1(torch.squeeze(inputs[:, s, :, :, :], dim=1)) 113 | downs1.append(down1.unsqueeze(dim=1)) 114 | all_indices_1.append(indices_1) 115 | all_unpool_shape1.append(unpool_shape1) 116 | all_FM1.append(FM1) 117 | downs1 = torch.cat(downs1, dim=1) 118 | if self.enablePermEq: 119 | downs1 = self.maxSubtract(downs1) 120 | 121 | downs2 = [] 122 | all_indices_2 = [] 123 | all_unpool_shape2 = [] 124 | all_FM2 = [] 125 | for s in range(S): 126 | down2, indices_2, unpool_shape2, FM2 = self.down2(torch.squeeze(downs1[:, s, :, :, :], dim=1)) 127 | downs2.append(down2.unsqueeze(dim=1)) 128 | all_indices_2.append(indices_2) 129 | all_unpool_shape2.append(unpool_shape2) 130 | all_FM2.append(FM2) 131 | downs2 = torch.cat(downs2, dim=1) 132 | if self.enablePermEq: 133 | downs2 = self.maxSubtract(downs2) 134 | 135 | downs3 = [] 136 | all_indices_3 = [] 137 | all_unpool_shape3 = [] 138 | all_FM3 = [] 139 | for s in range(S): 140 | down3, indices_3, unpool_shape3, FM3 = self.down3(torch.squeeze(downs2[:, s, :, :, :], dim=1)) 141 | downs3.append(down3.unsqueeze(dim=1)) 142 | all_indices_3.append(indices_3) 143 | all_unpool_shape3.append(unpool_shape3) 144 | all_FM3.append(FM3) 145 | downs3 = torch.cat(downs3, dim=1) 146 | if self.enablePermEq: 147 | downs3 = self.maxSubtract(downs3) 148 | 149 | downs4 = [] 150 | all_indices_4 = [] 151 | all_unpool_shape4 = [] 152 | all_FM4 = [] 153 | for s in range(S): 154 | down4, indices_4, unpool_shape4, FM4 = self.down4(torch.squeeze(downs3[:, s, :, :, :], dim=1)) 155 | downs4.append(down4.unsqueeze(dim=1)) 156 | all_indices_4.append(indices_4) 157 | all_unpool_shape4.append(unpool_shape4) 158 | all_FM4.append(FM4) 159 | downs4 = torch.cat(downs4, dim=1) 160 | if self.enablePermEq: 161 | downs4 = self.maxSubtract(downs4) 162 | 163 | downs5 = [] 164 | all_indices_5 = [] 165 | all_unpool_shape5 = [] 166 | all_FM5 = [] 167 | for s in range(S): 168 | down5, indices_5, unpool_shape5, FM5 = self.down5(torch.squeeze(downs4[:, s, :, :, :], dim=1)) 169 | downs5.append(down5.unsqueeze(dim=1)) 170 | all_indices_5.append(indices_5) 171 | all_unpool_shape5.append(unpool_shape5) 172 | all_FM5.append(FM5) 173 | downs5 = torch.cat(downs5, dim=1) 174 | if self.enablePermEq: 175 | downs5 = self.maxSubtract(downs5) 176 | 177 | ups5 = [] 178 | ups4 = [] 179 | ups3 = [] 180 | ups2 = [] 181 | ups1 = [] 182 | for s in range(S): 183 | up5 = self.up5(torch.squeeze(downs5[:, s, :, :, :], dim=1), all_indices_5[s], all_unpool_shape5[s], 184 | SkipFeatureMap=all_FM5[s]) 185 | ups5.append(up5.unsqueeze(dim=1)) 186 | ups5 = torch.cat(ups5, dim=1) 187 | if self.enablePermEq: 188 | ups5 = self.maxSubtract(ups5) 189 | 190 | for s in range(S): 191 | up4 = self.up4(torch.squeeze(ups5[:, s, :, :, :], dim=1), all_indices_4[s], all_unpool_shape4[s], 192 | SkipFeatureMap=all_FM4[s]) 193 | ups4.append(up4.unsqueeze(dim=1)) 194 | ups4 = torch.cat(ups4, dim=1) 195 | if self.enablePermEq: 196 | ups4 = self.maxSubtract(ups4) 197 | 198 | for s in range(S): 199 | up3 = self.up3(torch.squeeze(ups4[:, s, :, :, :], dim=1), all_indices_3[s], all_unpool_shape3[s], 200 | SkipFeatureMap=all_FM3[s]) 201 | ups3.append(up3.unsqueeze(dim=1)) 202 | ups3 = torch.cat(ups3, dim=1) 203 | if self.enablePermEq: 204 | ups3 = self.maxSubtract(ups3) 205 | 206 | for s in range(S): 207 | up2 = self.up2(torch.squeeze(ups3[:, s, :, :, :], dim=1), all_indices_2[s], all_unpool_shape2[s], 208 | SkipFeatureMap=all_FM2[s]) 209 | ups2.append(up2.unsqueeze(dim=1)) 210 | ups2 = torch.cat(ups2, dim=1) 211 | if self.enablePermEq: 212 | ups2 = self.maxSubtract(ups2) 213 | 214 | for s in range(S): 215 | up1 = self.up1(torch.squeeze(ups2[:, s, :, :, :], dim=1), all_indices_1[s], all_unpool_shape1[s], 216 | SkipFeatureMap=all_FM1[s]) 217 | ups1.append(up1.unsqueeze(dim=1)) 218 | 219 | SetUps1 = torch.cat(ups1, dim=1) 220 | # print(SetUps1.size()) 221 | 222 | return SetUps1 223 | 224 | def init_vgg16_params(self, vgg16): 225 | blocks = [self.down1, self.down2, self.down3, self.down4, self.down5] 226 | 227 | features = list(vgg16.features.children()) 228 | 229 | vgg_layers = [] 230 | for _layer in features: 231 | if isinstance(_layer, nn.Conv2d): 232 | vgg_layers.append(_layer) 233 | 234 | merged_layers = [] 235 | for idx, conv_block in enumerate(blocks): 236 | if idx < 2: 237 | units = [conv_block.conv1.cbr_unit, conv_block.conv2.cbr_unit] 238 | else: 239 | units = [ 240 | conv_block.conv1.cbr_unit, 241 | conv_block.conv2.cbr_unit, 242 | conv_block.conv3.cbr_unit, 243 | ] 244 | for _unit in units: 245 | for _layer in _unit: 246 | if isinstance(_layer, nn.Conv2d): 247 | merged_layers.append(_layer) 248 | 249 | assert len(vgg_layers) == len(merged_layers) 250 | 251 | for l1, l2 in zip(vgg_layers, merged_layers): 252 | if isinstance(l1, nn.Conv2d) and isinstance(l2, nn.Conv2d): 253 | assert l1.weight.size() == l2.weight.size() 254 | assert l1.bias.size() == l2.bias.size() 255 | l2.weight.data = l1.weight.data 256 | l2.bias.data = l1.bias.data 257 | 258 | 259 | if __name__ == '__main__': 260 | with torch.no_grad(): 261 | model = SetSegNet(n_classes=21, in_channels=3, withSkipConnections=True, enablePermEq=True).cuda() 262 | x = torch.rand(1, 5, 3, 240, 320).cuda() 263 | y = model(x) 264 | print(y.shape) 265 | -------------------------------------------------------------------------------- /core/trainer.py: -------------------------------------------------------------------------------- 1 | """ 2 | Wrapper 3 | """ 4 | 5 | 6 | class Trainer(object): 7 | def __init__(self, cfg, model, dataloader_dict, logger): 8 | self.cfg = cfg 9 | self.modes = self.cfg.MODES 10 | self.model = model 11 | self.dataloader_dict = dataloader_dict 12 | self.logger = logger 13 | self.current_epoch = cfg.RESUME_EPOCH_ID if cfg.RESUME else 1 14 | self.batch_count = 0 15 | 16 | def do_epoch(self): 17 | # for multi phase of one epoch 18 | for mode in self.modes: 19 | if mode in self.dataloader_dict.keys(): 20 | # start this phase, pass forward all samples in data-set 21 | self.batch_count = 0 22 | batch_total_num = len(self.dataloader_dict[mode]) 23 | for batch in iter(self.dataloader_dict[mode]): 24 | self.batch_count += 1 25 | if mode == 'train': 26 | batch = self.model.train_batch(batch) # actual training line 27 | else: 28 | batch = self.model.vali_batch(batch) # actual validation/test line 29 | batch = self.wrap_batchwise_log(batch, batch_total_num, mode=mode, ) 30 | self.logger.log(batch) 31 | self.current_epoch += 1 32 | self.adjust_learning_rate(optimizer=self.model.optimizer, epoch=self.current_epoch, 33 | lr_decay_rate=self.cfg.OPTI_DECAY_RATE, 34 | lr_decay_epoch=self.cfg.OPTI_DECAY_INTERVAL, min_lr=self.cfg.OPTI_DECAY_MIN) 35 | if self.current_epoch == self.cfg.EPOCH_TOTAL + 1: # if the last epoch 36 | self.logger.end_log() 37 | return self.current_epoch 38 | 39 | def wrap_batchwise_log(self, batch, batch_total, mode='train'): 40 | try: 41 | assert 'meta-info' in batch.keys() # Valid batch must have a key called meta-info 42 | except: 43 | raise ValueError("A valid batch passing to logger must contain a member called 'meta-info'!") 44 | wrapped = dict() 45 | # add logger head 46 | wrapped['batch-id'] = self.batch_count 47 | wrapped['batch-total'] = batch_total 48 | wrapped['epoch-id'] = self.current_epoch 49 | wrapped['phase'] = mode.lower() 50 | # add parser 51 | wrapped['parser'] = self.model.output_info_dict 52 | # add save method 53 | wrapped['save-method'] = self.model.save_model 54 | # add meta info 55 | wrapped['meta-info'] = batch['meta-info'] 56 | # main data ################################################# 57 | wrapped['data'] = batch 58 | return wrapped 59 | 60 | @staticmethod 61 | def adjust_learning_rate(optimizer, epoch, lr_decay_rate, lr_decay_epoch, min_lr=1e-5): 62 | if ((epoch + 1) % lr_decay_epoch) != 0: 63 | return 64 | 65 | for param_group in optimizer.param_groups: 66 | # print(param_group) 67 | lr_before = param_group['lr'] 68 | param_group['lr'] = param_group['lr'] * lr_decay_rate 69 | param_group['lr'] = max(param_group['lr'], min_lr) 70 | print('changing learning rate {:5f} to {:.5f}'.format(lr_before, max(param_group['lr'], min_lr))) 71 | 72 | @staticmethod 73 | def reset_learning_rate(optimizer, lr): 74 | for param_group in optimizer.param_groups: 75 | # print(param_group) 76 | lr_before = param_group['lr'] 77 | param_group['lr'] = lr 78 | print('changing learning rate {:5f} to {:.5f}'.format(lr_before, lr)) 79 | -------------------------------------------------------------------------------- /dataset/__init__.py: -------------------------------------------------------------------------------- 1 | import importlib 2 | 3 | 4 | def get_dataset(alias): 5 | module = importlib.import_module('dataset.' + alias.lower()) 6 | return module.Dataset 7 | -------------------------------------------------------------------------------- /dataset/neurips_dataset.py: -------------------------------------------------------------------------------- 1 | """ 2 | NOX dataset for neurips nocs dataset 3 | """ 4 | 5 | import torch.utils.data as data 6 | import os 7 | import numpy as np 8 | import cv2 as cv 9 | import json 10 | 11 | 12 | class Dataset(data.Dataset): 13 | 14 | def __init__(self, cfg, mode): 15 | 16 | super().__init__() 17 | self.mode = mode 18 | self.dataset_root = os.path.join(cfg.ROOT_DIR, 'resource', 'data', cfg.DATASET_ROOT) 19 | # each phase has separate dir ends with phase name: e.g. xxxx_train xxxx_test 20 | modes_name = os.listdir(self.dataset_root) 21 | self.split_dir_name = None 22 | for name in modes_name: 23 | if name.lower().endswith(mode): 24 | self.split_dir_name = name 25 | if self.split_dir_name is None: 26 | raise ValueError('No corresponding phase data') 27 | self.cates_list = cfg.DATASET_CATES 28 | 29 | self.shape = (320, 240) 30 | self.supervision_cap = 4096 # how many pixel in nocs map foreground will be sampled to trian SP branch 31 | 32 | # if use noise background augmentation 33 | self.add_noise_flag = cfg.ADD_BACKGROUND_NOISE 34 | 35 | # build or directly index 36 | index_fn = None 37 | for f_index in cfg.DATASET_INDEX: 38 | if f_index.endswith(self.mode + '.json'): 39 | index_fn = f_index 40 | if index_fn is None: 41 | self.index = self.build_index(check=True) 42 | else: 43 | with open(index_fn, 'r') as filehandle: 44 | self.index = json.load(filehandle) 45 | # self.index = self.build_index(check=True) 46 | 47 | # use the proportion configuration to reduce the dataset for development 48 | if cfg.PROPORTION < 1.0: 49 | p = cfg.PROPORTION 50 | all_num = len(self.index) 51 | self.index = self.index[:int(all_num * p - 1)] 52 | return 53 | 54 | def __len__(self): 55 | return len(self.index) 56 | 57 | def __getitem__(self, idx): 58 | meta_info = self.index[idx] 59 | return self.grab(meta_info) 60 | 61 | def grab(self, meta_info): 62 | # get direct rgb, nox, mask from NOCS dataset 63 | v_rgb_w_bg = cv.imread(os.path.join(self.dataset_root, meta_info['view0'])) 64 | x_rgb_x_bg = cv.imread(os.path.join(self.dataset_root, meta_info['view1'])) 65 | v_nocs_x_bg = cv.imread(os.path.join(self.dataset_root, meta_info['nox0'])) 66 | x_nocs_x_bg = cv.imread(os.path.join(self.dataset_root, meta_info['nox1'])) 67 | 68 | v_rgb_w_bg = cv.resize(v_rgb_w_bg, self.shape, interpolation=cv.INTER_NEAREST) 69 | x_rgb_x_bg = cv.resize(x_rgb_x_bg, self.shape, interpolation=cv.INTER_NEAREST) 70 | v_nocs_x_bg = cv.resize(v_nocs_x_bg, self.shape, interpolation=cv.INTER_NEAREST) 71 | x_nocs_x_bg = cv.resize(x_nocs_x_bg, self.shape, interpolation=cv.INTER_NEAREST) 72 | 73 | mask_v = (np.sum(v_nocs_x_bg.astype(float), axis=2, keepdims=True) < 255 * 3).astype(np.float) 74 | mask_x = (np.sum(x_nocs_x_bg.astype(float), axis=2, keepdims=True) < 255 * 3).astype(np.float) 75 | 76 | if self.add_noise_flag: 77 | # add noise to background 78 | if self.mode == 'train': 79 | # have some possibility to have the noise 80 | seed = np.random.rand() 81 | if seed > 0.5: 82 | noise = np.random.rand(self.shape[1], self.shape[0], 3).astype(np.float32) 83 | v_rgb_w_bg = v_rgb_w_bg * mask_v + noise * (1.0 - mask_v) * 255 84 | v_rgb_w_bg = v_rgb_w_bg.astype(np.float32) 85 | 86 | sample_uv_v, sample_mask_v, sample_xyz_v = self.get_uv_mask_xyz(v_nocs_x_bg) 87 | sample_uv_x, sample_mask_x, sample_xyz_x = self.get_uv_mask_xyz(x_nocs_x_bg) 88 | 89 | meta_info_reduced = { 90 | 'cate': meta_info['cate'], 91 | 'object': meta_info['object'], 92 | 'id': meta_info['id'], 93 | 'phase': meta_info['phase'] 94 | } 95 | return { 96 | 'rgb-v': v_rgb_w_bg, 97 | 'rgb-x': x_rgb_x_bg, 98 | 'nox-v': v_nocs_x_bg, 99 | 'nox-x': x_nocs_x_bg, 100 | 'mask-v': mask_v, 101 | 'mask-x': mask_x, 102 | 'uv-v': sample_uv_v, 103 | 'uv-x': sample_uv_x, 104 | 'uv-mask-v': sample_mask_v, 105 | 'uv-mask-x': sample_mask_x, 106 | 'uv-xyz-v': sample_xyz_v, 107 | 'uv-xyz-x': sample_xyz_x, 108 | 'info': meta_info_reduced, # meta info 109 | 'pose': meta_info['pose'], 110 | } 111 | 112 | def get_uv_mask_xyz(self, nocs_map): 113 | # get image space uv grid 114 | crop_param = ((0, nocs_map.shape[0]), (0, nocs_map.shape[1])) 115 | u = np.arange(0, crop_param[0][1] - crop_param[0][0]) + 0.5 116 | u = u / len(u) 117 | v = np.arange(0, crop_param[1][1] - crop_param[1][0]) + 0.5 118 | v = v / len(v) 119 | col_uv, row_uv = np.meshgrid(v, u) 120 | uv_map = np.concatenate((row_uv[..., np.newaxis], col_uv[..., np.newaxis]), 2) 121 | 122 | nocs_croped_v = nocs_map[crop_param[0][0]:crop_param[0][1], crop_param[1][0]: crop_param[1][1], :] # not used 123 | uv_list = uv_map.reshape(-1, 2) # [K,2] 124 | nocs_list_v = nocs_croped_v.reshape(-1, 3) # [K,3] 125 | # mask the background 126 | mask = np.sum(nocs_list_v, axis=1) < 255 * 3 127 | uv_list = uv_list[mask, :] 128 | nocs_list_v = nocs_list_v[mask, :] 129 | # sample 130 | if uv_list.shape[0] >= self.supervision_cap: 131 | sample_index = np.random.randint(low=0, high=max(uv_list.shape[0] - 1, 1), size=self.supervision_cap) 132 | uv_list = uv_list[sample_index, :] 133 | nocs_list_v = nocs_list_v[sample_index, :] 134 | # randomly shuffle 135 | arr = np.array(range(0, uv_list.shape[0])) 136 | np.random.shuffle(arr) 137 | uv_list = uv_list[arr, :] 138 | nocs_list_v = nocs_list_v[arr, :] 139 | # now two lists are both shorter/equal to the supervision size 140 | # safely padding 141 | uv = np.zeros((self.supervision_cap, 2)) 142 | mask = np.zeros((self.supervision_cap, 1)) 143 | xyz_v = np.zeros((self.supervision_cap, 3)) 144 | # fill value 145 | uv[0:uv_list.shape[0], :] = uv_list 146 | mask[0:uv_list.shape[0], :] = 1 147 | xyz_v[0:uv_list.shape[0], :] = nocs_list_v 148 | return uv.astype(np.float), mask.astype(np.float), xyz_v.astype(np.float) 149 | 150 | def build_index(self, check=False): 151 | """ 152 | make inner index for each data, to make sure each data sample is usable 153 | """ 154 | index_list = list() 155 | root = os.path.join(self.dataset_root, self.split_dir_name) 156 | for cate in self.cates_list: 157 | cate_root = os.path.join(root, cate) 158 | models_list = os.listdir(cate_root) 159 | for obj in models_list: 160 | obj_root = os.path.join(cate_root, obj) 161 | files_list = os.listdir(obj_root) 162 | id_set = set(id.split('_')[1] for id in files_list) 163 | for id in id_set: 164 | relative_obj_root = os.path.join(self.split_dir_name, cate, obj) 165 | sample = { 166 | 'cate': cate, 'object': obj, 'id': id, 'phase': self.mode, 167 | 'view0': os.path.join(relative_obj_root, 'frame_' + id + '_Color_00.png'), 168 | 'view1': os.path.join(relative_obj_root, 'frame_' + id + '_Color_01.png'), 169 | 'nox0': os.path.join(relative_obj_root, 'frame_' + id + '_NOXRayTL_00.png'), 170 | 'nox1': os.path.join(relative_obj_root, 'frame_' + id + '_NOXRayTL_01.png'), 171 | 'pose': os.path.join(relative_obj_root, 'frame_' + id + '_CameraPose.json'), 172 | } 173 | # index_list.append(sample) # used to debug 174 | if check: # check this sample 175 | try: 176 | self.grab(sample) 177 | index_list.append(sample) 178 | print("\r Check %d meta data" % len(index_list), end='') 179 | except: 180 | print("Cate " + cate + " Objcet " + obj + " ID " + id + " get error during check") 181 | if check: 182 | # save 183 | with open(self.mode + '.json', 'w') as filehandle: 184 | json.dump(index_list, filehandle) 185 | return index_list 186 | 187 | 188 | if __name__ == '__main__': 189 | from yacs.config import CfgNode as CN 190 | from matplotlib import pyplot as plt 191 | 192 | plt.ioff() 193 | cfg = CN() 194 | 195 | ROOT_DIR = os.getcwd().split('/') 196 | ROOT_DIR.pop(-1) 197 | ROOT_DIR = ['/'] + ROOT_DIR 198 | cfg.ROOT_DIR = os.path.join(*ROOT_DIR) 199 | cfg.DATASET_ROOT = "pix2surf_viz" 200 | cfg.DATASET_CATES = ['airplanes'] 201 | cfg.PROPORTION = 1.0 202 | cfg.ADD_BACKGROUND_NOISE = False 203 | # cfg.DATASET_INDEX = ['../resource/index/shapenet_plain_car_train.json', 204 | # '../resource/index/shapenet_plain_car_test.json'] 205 | cfg.DATASET_INDEX = [] 206 | # dataset = Dataset(cfg, 'train') 207 | dataset = Dataset(cfg, 'test') 208 | for dbg in dataset: 209 | plt.subplot(2, 2, 1) 210 | plt.imshow(dbg['rgb-v'].squeeze()[..., [2, 1, 0]] / 255) 211 | plt.subplot(2, 2, 2) 212 | plt.imshow(dbg['nox-v'].squeeze()) 213 | plt.subplot(2, 2, 3) 214 | plt.imshow(dbg['nox-x'].squeeze()) 215 | plt.subplot(2, 2, 4) 216 | plt.imshow(dbg['mask-v'].squeeze()) 217 | plt.show() 218 | plt.pause(1) 219 | -------------------------------------------------------------------------------- /logger/__init__.py: -------------------------------------------------------------------------------- 1 | import importlib 2 | 3 | 4 | def get_logger(alias): 5 | module = importlib.import_module('logger.' + alias.lower()) 6 | return module.Logger 7 | -------------------------------------------------------------------------------- /logger/logger_meta/__init__.py: -------------------------------------------------------------------------------- 1 | from .metric_logger import MetricLogger 2 | from .image_logger import ImageLogger 3 | from .model_logger import ModelLogger 4 | from .xls_logger import XLSLogger 5 | from .obj_logger import ObjectLogger 6 | 7 | LOGGER_REGISTED = { 8 | 'metric': MetricLogger, 9 | 'image': ImageLogger, 10 | 'model': ModelLogger, 11 | 'xls': XLSLogger, 12 | 'obj': ObjectLogger, 13 | } 14 | -------------------------------------------------------------------------------- /logger/logger_meta/base_logger.py: -------------------------------------------------------------------------------- 1 | class BaseLogger(object): 2 | def __init__(self, tb_logger, log_path, cfg): 3 | super().__init__() 4 | self.cfg = cfg 5 | self.NAME = 'base' 6 | self.tb = tb_logger 7 | self.log_path = log_path 8 | # make dir 9 | 10 | def log_phase(self): 11 | pass 12 | 13 | def log_batch(self, batch): 14 | pass 15 | -------------------------------------------------------------------------------- /logger/logger_meta/image_logger.py: -------------------------------------------------------------------------------- 1 | from .base_logger import BaseLogger 2 | import matplotlib 3 | import os 4 | import numpy as np 5 | from cv2 import imwrite 6 | import cv2 7 | import torch 8 | from copy import deepcopy 9 | 10 | matplotlib.use('Agg') 11 | 12 | 13 | class ImageLogger(BaseLogger): 14 | def __init__(self, tb_logger, log_path, cfg): 15 | super().__init__(tb_logger, log_path, cfg) 16 | self.NAME = 'image' 17 | os.makedirs(self.log_path, exist_ok=True) 18 | self.visual_interval_epoch = cfg.VIS_PER_N_EPOCH 19 | self.visual_interval_batch = cfg.VIS_PER_N_BATCH 20 | self.visual_one = cfg.VIS_ONE_PER_BATCH 21 | self.visual_train_interval_batch = cfg.VIS_TRAIN_PER_BATCH 22 | 23 | def log_batch(self, batch): 24 | # get data 25 | if not self.NAME in batch['parser'].keys(): 26 | return 27 | keys_list = batch['parser'][self.NAME] 28 | if len(keys_list) == 0: 29 | return 30 | data = batch['data'] 31 | phase = batch['phase'] 32 | current_batch = batch['batch-id'] 33 | current_epoch = batch['epoch-id'] 34 | meta_info = batch['meta-info'] # {k:[info0,info1,info2,....]} 35 | # check dir 36 | os.makedirs(os.path.join(self.log_path, 'epoch_%d' % current_epoch), exist_ok=True) 37 | # check whether need log 38 | if phase.lower() == 'train': 39 | if current_batch % self.visual_train_interval_batch != 0 or current_batch % self.visual_interval_batch != 0: 40 | return 41 | else: 42 | if current_epoch % self.visual_interval_epoch != 0 or current_batch % self.visual_interval_batch != 0: 43 | return 44 | # for each key 45 | for img_key in keys_list: # for each key 46 | kdata = data[img_key] 47 | if isinstance(kdata, list): 48 | assert len(kdata[0].shape) == 4 49 | nbatch = kdata[0].shape[0] 50 | kdata = deepcopy(kdata) 51 | else: 52 | assert len(kdata.shape) == 4 53 | nbatch = kdata.shape[0] 54 | if isinstance(kdata, torch.Tensor): 55 | kdata = [deepcopy(kdata.detach().cpu().numpy())] 56 | else: 57 | kdata = [deepcopy(kdata)] 58 | # convert to ndarray 59 | if isinstance(kdata[0], torch.Tensor): 60 | for i, tensor in enumerate(kdata): 61 | kdata[i] = tensor.detach().cpu().numpy() 62 | # for each sample in batch 63 | for batch_id in range(nbatch): 64 | # get meta postfix 65 | prefix = "" 66 | for k, v in meta_info.items(): 67 | prefix += k + "_" + str(v[batch_id]) + "_" 68 | # now all cases are converted to list of image 69 | nview = len(kdata) 70 | for view_id in range(nview): 71 | img = kdata[view_id][batch_id] # 3*W*H / 1*W*H 72 | assert img.ndim == 3 73 | # first process image 74 | color_flag = False 75 | if img.shape[0] == 1: 76 | color_flag = True 77 | cm = matplotlib.cm.get_cmap('viridis') 78 | img = cm(img.squeeze(0))[..., :3] 79 | img = img.transpose(2, 0, 1) 80 | img *= 255 81 | else: 82 | img *= 255.0 if img.max() < 200 else 1 83 | img = np.clip(img, a_min=0, a_max=255) 84 | img = img.astype(np.uint8) 85 | self.tb.add_image(phase + '/' + img_key, img if color_flag else img[[2, 1, 0], ...], current_epoch) 86 | # save to file 87 | img = img.transpose(1, 2, 0) 88 | filename = os.path.join( 89 | self.log_path, 'epoch_%d' % current_epoch, 90 | prefix + 'b%d_v%d_' % (batch_id, view_id) + img_key + '.png' 91 | ) 92 | if color_flag: 93 | img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 94 | imwrite(filename, img) 95 | if self.visual_one: 96 | break 97 | 98 | def log_phase(self): 99 | pass 100 | -------------------------------------------------------------------------------- /logger/logger_meta/metric_logger.py: -------------------------------------------------------------------------------- 1 | import matplotlib 2 | import os 3 | import shutil 4 | import csv 5 | import numpy as np 6 | import matplotlib.pyplot as plt 7 | from .base_logger import BaseLogger 8 | import time 9 | 10 | matplotlib.use('Agg') 11 | 12 | 13 | class MetricLogger(BaseLogger): 14 | def __init__(self, tb_logger, log_path, cfg): 15 | super().__init__(tb_logger, log_path, cfg) 16 | self.NAME = 'metric' 17 | os.makedirs(self.log_path, exist_ok=True) 18 | os.makedirs(os.path.join(self.log_path, 'batchwise'), exist_ok=True) 19 | self.phase = None 20 | self.current_epoch = -1 21 | self.current_batch = -1 22 | self.metric_container = dict() 23 | self.batch_size = self.cfg.BATCHSIZE 24 | self.epoch_total = self.cfg.EPOCH_TOTAL 25 | self.phase_time_start = time.time() 26 | 27 | def log_batch(self, batch): 28 | """ 29 | - add each metric to tensorboard 30 | - record each metric for epoch save 31 | - display in terminal, displayed metric is averaged 32 | """ 33 | if not self.NAME in batch['parser'].keys(): 34 | return 35 | keys_list = batch['parser'][self.NAME] 36 | if len(keys_list) == 0: 37 | return 38 | data = batch['data'] 39 | total_batch = batch['batch-total'] 40 | self.phase = batch['phase'] 41 | self.current_batch = batch['batch-id'] 42 | self.current_epoch = batch['epoch-id'] 43 | for k in keys_list: 44 | if not k in self.metric_container.keys(): 45 | self.metric_container[k] = [data[k]] 46 | else: 47 | self.metric_container[k].append(data[k]) 48 | # add to tb 49 | self.tb.add_scalars('Metric-BatchWise/' + k, {self.phase: float(data[k])}, 50 | self.current_batch + self.current_epoch * total_batch) 51 | self.update_progbar(self.phase, self.current_epoch, self.current_batch, total_batch, 52 | sum(self.metric_container[keys_list[0]]) / len(self.metric_container[keys_list[0]]), 53 | (time.time() - self.phase_time_start) / 60) 54 | 55 | def log_phase(self): 56 | """ 57 | - save the batch-wise scalar in file 58 | - save corresponding figure for each scalar 59 | For epoch wise metric 60 | - add to tb 61 | """ 62 | phase = self.phase 63 | batchwise_root = os.path.join(self.log_path, 'batchwise', 'epoch_%d' % self.current_epoch) 64 | os.makedirs(batchwise_root, exist_ok=True) 65 | for k, v in self.metric_container.items(): 66 | # first log this phase's batch wise csv 67 | with open(os.path.join(batchwise_root, '%d_ep_' % self.current_epoch + '_' + k + '_' + phase + '.csv'), 68 | 'w', newline='') as myfile: 69 | wr = csv.writer(myfile, quoting=csv.QUOTE_ALL) 70 | wr.writerow(v) 71 | # log figure 72 | plt.plot(np.array(v)) 73 | plt.title('%d_ep_' % self.current_epoch + k + '_' + phase) 74 | plt.savefig(os.path.join(batchwise_root, '%d_ep_' % self.current_epoch + '_' + k + '_' + phase + '.png')) 75 | plt.close() 76 | # add average to tb epoch wise 77 | mean = sum(v) / len(v) 78 | self.tb.add_scalars('Metric-EpochWise/' + k, {phase: float(mean)}, int(self.current_epoch)) 79 | print('\n' + '=' * shutil.get_terminal_size()[0]) 80 | self.metric_container = dict() 81 | self.phase_time_start = time.time() 82 | 83 | def update_progbar(self, phase, epoch_current, current, total, first_metric, time): 84 | prefix = '|EP %4d/%4d|SP %4d/%4d|M %.5f|%.1f min|%.1f min est.|' % ( 85 | epoch_current, self.epoch_total, current * self.batch_size, total * self.batch_size, 86 | first_metric, time, time * (total / current - 1)) + phase + '|' 87 | termWidth = shutil.get_terminal_size()[0] 88 | barLength = termWidth - len(prefix) - 14 89 | print(self.ProgressBar(current / total, prefix, barLength, True), end=' ') 90 | 91 | @classmethod 92 | def ProgressBar(cls, percent, prefix=None, notches=50, numericalpercent=True, unicode=False): 93 | """ 94 | https://github.com/paulojraposo/ProgBar 95 | """ 96 | outString = u"" # Unicode string. 97 | if prefix: 98 | prefix = "{} ".format(prefix) 99 | outString = outString + prefix 100 | x_of_notches = int(round(percent * notches)) 101 | startCap = "[" 102 | endCap = "]" 103 | fullSegment = '>' # '▣' # "✔" # ">" # "#" 104 | blankSegment = '-' # '▢' # "." 105 | if unicode: 106 | fullSegment = "\u25AE" # Full block in Unicode 107 | blankSegment = "\u25AF" # Empty block in Unicode 108 | outString = outString + startCap 109 | for i in range(x_of_notches): 110 | outString = outString + fullSegment # Full block 111 | for i in range(notches - x_of_notches): 112 | outString = outString + blankSegment 113 | outString = outString + endCap 114 | if numericalpercent: 115 | outString = outString + " [{}%]".format(str(round(percent * 100, 2))) 116 | return '\r' + outString 117 | -------------------------------------------------------------------------------- /logger/logger_meta/model_logger.py: -------------------------------------------------------------------------------- 1 | import os 2 | from .base_logger import BaseLogger 3 | 4 | 5 | class ModelLogger(BaseLogger): 6 | def __init__(self, tb_logger, log_path, cfg): 7 | super().__init__(tb_logger, log_path, cfg) 8 | self.NAME = 'model' 9 | os.makedirs(self.log_path, exist_ok=True) 10 | self.phase = 'train' 11 | self.current_epoch = -1 12 | self.save_interval = cfg.MODEL_SAVE_PER_N_EPOCH 13 | self.save_method = None 14 | 15 | def log_batch(self, batch): 16 | self.phase = batch['phase'] 17 | self.current_epoch = batch['epoch-id'] 18 | self.save_method = batch['save-method'] 19 | 20 | def log_phase(self): 21 | if self.phase == 'train' and (self.current_epoch % self.save_interval == 0): 22 | self.save_method(os.path.join(self.log_path, "epoch_%d.model" % self.current_epoch)) 23 | if self.phase == 'train': 24 | if any([True if fn.endswith('latest.model') else False for fn in os.listdir(self.log_path)]): 25 | os.system('rm ' + os.path.join(self.log_path, "epoch_*_latest.model")) 26 | self.save_method(os.path.join(self.log_path, "epoch_%d_latest.model" % self.current_epoch)) 27 | -------------------------------------------------------------------------------- /logger/logger_meta/obj_logger.py: -------------------------------------------------------------------------------- 1 | from .base_logger import BaseLogger 2 | import matplotlib 3 | import os 4 | import torch 5 | from tk3dv.nocstools import datastructures as ds 6 | import cv2 as cv 7 | from copy import deepcopy 8 | 9 | matplotlib.use('Agg') 10 | 11 | 12 | class ObjectLogger(BaseLogger): 13 | def __init__(self, tb_logger, log_path, cfg): 14 | super().__init__(tb_logger, log_path, cfg) 15 | self.NAME = 'obj' 16 | os.makedirs(self.log_path, exist_ok=True) 17 | self.visual_interval_epoch = cfg.VIS_PER_N_EPOCH 18 | self.visual_interval_batch = cfg.VIS_PER_N_BATCH 19 | self.visual_one = cfg.VIS_ONE_PER_BATCH 20 | self.visual_train_interval_batch = cfg.VIS_TRAIN_PER_BATCH 21 | 22 | def log_batch(self, batch): 23 | # get data 24 | if not self.NAME in batch['parser'].keys(): 25 | return 26 | keys_list = batch['parser'][self.NAME] 27 | if len(keys_list) == 0: 28 | return 29 | data = batch['data'] 30 | phase = batch['phase'] 31 | current_batch = batch['batch-id'] 32 | current_epoch = batch['epoch-id'] 33 | meta_info = batch['meta-info'] # {k:[info0,info1,info2,....]} 34 | # check dir 35 | os.makedirs(os.path.join(self.log_path, 'epoch_%d' % current_epoch), exist_ok=True) 36 | # check whether need log 37 | if phase.lower() == 'train': 38 | if current_batch % self.visual_train_interval_batch != 0 or current_batch % self.visual_interval_batch != 0: 39 | return 40 | else: 41 | if current_epoch % self.visual_interval_epoch != 0 or current_batch % self.visual_interval_batch != 0: 42 | return 43 | # for each key 44 | for obj_key in keys_list: # for each key 45 | nocs = data[obj_key[0]] 46 | color = data[obj_key[1]] 47 | if isinstance(nocs, list): 48 | assert isinstance(color, list) 49 | assert len(nocs[0].shape) == 4 50 | assert len(color[0].shape) == 4 51 | assert nocs[0].shape[0] == color[0].shape[0] 52 | nbatch = nocs[0].shape[0] 53 | else: 54 | assert len(nocs.shape) == 4 55 | assert len(color.shape) == 4 56 | assert nocs.shape[0] == color.shape[0] 57 | nbatch = nocs.shape[0] 58 | nocs = [nocs] 59 | color = [color] 60 | 61 | # convert to ndarray 62 | if isinstance(nocs[0], torch.Tensor): 63 | for i, tensor in enumerate(nocs): 64 | nocs[i] = tensor.detach().cpu().numpy() 65 | if isinstance(color[0], torch.Tensor): 66 | for i, tensor in enumerate(color): 67 | color[i] = tensor.detach().cpu().numpy() 68 | nocs = deepcopy(nocs) 69 | color = deepcopy(color) 70 | 71 | # for each sample in batch 72 | for batch_id in range(nbatch): 73 | # get meta postfix 74 | prefix = "" 75 | for k, v in meta_info.items(): 76 | prefix += k + "_" + str(v[batch_id]) + "_" 77 | # now all case convert to list of image 78 | nview = len(nocs) 79 | for view_id in range(nview): 80 | nocs_map = nocs[view_id][batch_id] # 3*H*W [0,1] 81 | color_map = color[view_id][batch_id] # 3*H*W [0,1] 82 | 83 | assert nocs_map.ndim == 3 84 | assert nocs_map.shape[0] == 3 85 | assert color_map.ndim == 3 86 | assert color_map.shape[0] == 3 87 | nocs_map = nocs_map.transpose(1, 2, 0) # H*W*3 88 | color_map = color_map.transpose(1, 2, 0) # H*W*3 89 | 90 | # smaller meshes 91 | nocs_map = cv.resize(nocs_map, dsize=(int(nocs_map.shape[1] / 4), int(nocs_map.shape[0] / 4)), 92 | interpolation=cv.INTER_NEAREST) 93 | color_map = cv.resize(color_map, dsize=(int(color_map.shape[1] / 4), int(color_map.shape[0] / 4)), 94 | interpolation=cv.INTER_NEAREST) 95 | 96 | # save to file 97 | filename = os.path.join( 98 | self.log_path, 'epoch_%d' % current_epoch, 99 | prefix + 'b%d_v%d_' % (batch_id, view_id) + 100 | obj_key[0] + '-' + obj_key[1] + '.obj' 101 | ) 102 | # save here 103 | # need to convert color, need to *255 104 | if nocs_map.max() < 1 + 1e-4: 105 | nocs_map = nocs_map[:, :, [2, 1, 0]] * 255.0 106 | if color_map.max() < 1 + 1e-4: 107 | color_map = color_map[:, :, [2, 1, 0]] * 255.0 108 | tk3dv_nocs_mp = ds.NOCSMap(nocs_map, RGB=color_map) 109 | tk3dv_nocs_mp.serialize(filename) 110 | if self.visual_one: 111 | break 112 | 113 | def log_phase(self): 114 | pass 115 | -------------------------------------------------------------------------------- /logger/logger_meta/xls_logger.py: -------------------------------------------------------------------------------- 1 | """ 2 | excel logger 3 | data structure: 4 | - one head-key is one file 5 | - each passed in data is a dict {col-name:list of values}, each value will be recorded into one row 6 | - there is some basic meta info for each row 7 | """ 8 | 9 | import pandas as pd 10 | from .base_logger import BaseLogger 11 | import os 12 | 13 | 14 | class XLSLogger(BaseLogger): 15 | def __init__(self, tb_logger, log_path, cfg): 16 | super().__init__(tb_logger, log_path, cfg) 17 | self.NAME = 'xls' 18 | os.makedirs(self.log_path, exist_ok=True) 19 | self.visual_interval_epoch = cfg.VIS_PER_N_EPOCH 20 | self.record_interval_batch = cfg.VIS_PER_N_BATCH 21 | self.visual_one = cfg.VIS_ONE_PER_BATCH 22 | self.visual_train_interval_batch = cfg.VIS_TRAIN_PER_BATCH 23 | self.pd_container = dict() 24 | 25 | self.current_epoch = 1 26 | self.current_phase = '' 27 | 28 | def log_batch(self, batch): 29 | # get data 30 | if not self.NAME in batch['parser'].keys(): 31 | return 32 | keys_list = batch['parser'][self.NAME] 33 | if len(keys_list) == 0: 34 | return 35 | data = batch['data'] 36 | phase = batch['phase'] 37 | current_epoch = batch['epoch-id'] 38 | self.current_epoch = current_epoch 39 | self.current_phase = phase 40 | meta_info = batch['meta-info'] 41 | # for each key (file) 42 | for sheet_key in keys_list: 43 | kdata = data[sheet_key] 44 | assert isinstance(kdata, dict) 45 | if sheet_key not in self.pd_container.keys(): 46 | self.pd_container[sheet_key] = pd.DataFrame() 47 | add_list = list() 48 | count = len(meta_info['object']) 49 | for ii in range(count): 50 | data = dict() 51 | for k, v in meta_info.items(): 52 | data[k] = v[ii] 53 | for k, v in kdata.items(): 54 | data[k] = v[ii] 55 | prefix = "" 56 | for k, v in meta_info.items(): 57 | prefix += k + "_" + str(v[ii]) + "_" 58 | data['prefix'] = prefix 59 | add_list.append(data) 60 | self.pd_container[sheet_key] = self.pd_container[sheet_key].append(add_list, ignore_index=True) 61 | 62 | def log_phase(self): 63 | for k in self.pd_container.keys(): 64 | self.pd_container[k].to_excel( 65 | os.path.join(self.log_path, k + '_' + str(self.current_epoch) + '_' + self.current_phase + '.xls')) 66 | self.pd_container[k] = pd.DataFrame() 67 | -------------------------------------------------------------------------------- /logger/logger_s1.py: -------------------------------------------------------------------------------- 1 | """ 2 | 06/31/2019 3 | Each passed in batch must contain info: 4 | 1. method of model to save itself 5 | 2. the way to parsing the batch keys 6 | 3. meta information of each sample 7 | 4. batch head (phase) 8 | Logger can be selected from config file 9 | """ 10 | import os 11 | from tensorboardX import SummaryWriter as writer 12 | from .logger_meta import LOGGER_REGISTED 13 | 14 | 15 | class Logger(object): 16 | 17 | def __init__(self, cfg): 18 | self.cfg = cfg 19 | self.phase = 'INIT' 20 | self.epoch = -1 21 | tb_path = os.path.join(cfg.ROOT_DIR, 'log', cfg.LOG_DIR, 'tensorboardx') 22 | self.tb_writer = writer(tb_path) 23 | self.lg_list = self.compose(self.cfg.LOGGER_SELECT) 24 | 25 | def compose(self, names): 26 | loggers_list = list() 27 | mapping = LOGGER_REGISTED 28 | for name in names: 29 | if name in mapping.keys(): 30 | loggers_list.append(mapping[name](self.tb_writer, os.path.join( 31 | os.path.join(self.cfg.ROOT_DIR, 'log', self.cfg.LOG_DIR, name)), self.cfg)) 32 | else: 33 | raise Warning('Required logger ' + name + ' not found!') 34 | return loggers_list 35 | 36 | def log(self, batch): 37 | # if phase changes, first log phase then log batch 38 | phase = batch['phase'] 39 | epoch = batch['epoch-id'] 40 | if self.phase != 'INIT' and phase != self.phase: 41 | for lgr in self.lg_list: 42 | lgr.log_phase() 43 | elif self.epoch != -1 and epoch != self.epoch: 44 | for lgr in self.lg_list: 45 | lgr.log_phase() 46 | self.phase = phase 47 | self.epoch = epoch 48 | for lgr in self.lg_list: 49 | lgr.log_batch(batch) 50 | 51 | def end_log(self): 52 | for lgr in self.lg_list: 53 | lgr.log_phase() 54 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | # pytorch==1.1.0 2 | # torchvision==0.3.0 3 | # cudatoolkit==10.0.130 4 | 5 | tensorboard==1.15.0 6 | tensorboardx==1.6 7 | matplotlib==3.3.0 8 | opencv-python==4.2.0.34 9 | pandas==1.1.0 10 | transforms3d==0.3.1 11 | xlwt==1.3.0 12 | yacs==0.1.7 13 | sklearn==0.0 14 | pyqt5==5.14.0 -------------------------------------------------------------------------------- /resource/utils/car_vali.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JiahuiLei/Pix2Surf/ac56ce5a20adab8cb563d96c915e1cc253631bf8/resource/utils/car_vali.npy -------------------------------------------------------------------------------- /resource/utils/chair_vali.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JiahuiLei/Pix2Surf/ac56ce5a20adab8cb563d96c915e1cc253631bf8/resource/utils/chair_vali.npy -------------------------------------------------------------------------------- /resource/utils/make_coco_validation.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | 4 | # modify new version of shapenet plain to old version 5 | if os.path.exists('../data/shapenet_plain'): 6 | if os.path.exists('../data/shapenet_plain/val'): 7 | os.system('mv ../data/shapenet_plain/val ../data/shapenet_plain/test') 8 | 9 | # modify new version of shapenet coco to old version 10 | # split a small validation set from training set 11 | if os.path.exists('../data/shapenet_coco'): 12 | if os.path.exists('../data/shapenet_coco/val'): 13 | os.system('mv ../data/shapenet_coco/val ../data/shapenet_coco/test') 14 | source_root = '../data/shapenet_coco/train' 15 | target_root = '../data/shapenet_coco/vali' 16 | os.makedirs(target_root, exist_ok=False) 17 | 18 | for cate, vali_list_fn in zip(['03001627', '02958343', '02691156'], 19 | ['../utils/chair_vali.npy', 20 | '../utils/car_vali.npy', 21 | '../utils/plane_vali.npy']): 22 | fn_list = np.load(vali_list_fn) 23 | os.makedirs(os.path.join(target_root, cate)) 24 | for fn in fn_list: 25 | src = os.path.join(source_root, cate, fn) 26 | tgt = os.path.join(target_root, cate, fn) 27 | cmd = 'mv ' + src + ' ' + tgt 28 | print(cmd) 29 | os.system(cmd) 30 | -------------------------------------------------------------------------------- /resource/utils/plane_vali.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JiahuiLei/Pix2Surf/ac56ce5a20adab8cb563d96c915e1cc253631bf8/resource/utils/plane_vali.npy -------------------------------------------------------------------------------- /run.py: -------------------------------------------------------------------------------- 1 | """ 2 | Main program 3 | May the Force be with you. 4 | 2019.3 5 | """ 6 | from torch.utils.data import DataLoader 7 | 8 | from dataset import get_dataset 9 | from logger import get_logger 10 | from core.models import get_model 11 | from core.trainer import Trainer 12 | from config import get_cfg 13 | 14 | # preparer configuration 15 | cfg = get_cfg() 16 | 17 | # prepare dataset 18 | DatasetClass = get_dataset(cfg.DATASET) 19 | dataloader_dict = dict() 20 | for mode in cfg.MODES: 21 | phase_dataset = DatasetClass(cfg, mode=mode) 22 | dataloader_dict[mode] = DataLoader(phase_dataset, batch_size=cfg.BATCHSIZE, 23 | shuffle=True if mode in ['train'] else False, 24 | num_workers=cfg.DATALOADER_WORKERS, pin_memory=True, 25 | drop_last=True) 26 | 27 | # prepare models 28 | ModelClass = get_model(cfg.MODEL) 29 | model = ModelClass(cfg) 30 | 31 | # prepare logger 32 | LoggerClass = get_logger(cfg.LOGGER) 33 | logger = LoggerClass(cfg) 34 | 35 | # register dataset, models, logger to trainer 36 | trainer = Trainer(cfg, model, dataloader_dict, logger) 37 | 38 | # start training 39 | epoch_total = cfg.EPOCH_TOTAL + (cfg.RESUME_EPOCH_ID if cfg.RESUME else 0) 40 | while trainer.do_epoch() <= cfg.EPOCH_TOTAL: 41 | pass 42 | -------------------------------------------------------------------------------- /run_batch.py: -------------------------------------------------------------------------------- 1 | """ 2 | Main program 3 | May the Force be with you. 4 | 5 | This main file is used on slurm server without interactive check of config 6 | """ 7 | from torch.utils.data import DataLoader 8 | 9 | from dataset import get_dataset 10 | from logger import get_logger 11 | from core.models import get_model 12 | from core.trainer import Trainer 13 | from config import get_cfg 14 | 15 | # preparer configuration 16 | cfg = get_cfg(interactive=False) 17 | 18 | # prepare dataset 19 | DatasetClass = get_dataset(cfg.DATASET) 20 | dataloader_dict = dict() 21 | for mode in cfg.MODES: 22 | phase_dataset = DatasetClass(cfg, mode=mode) 23 | dataloader_dict[mode] = DataLoader(phase_dataset, batch_size=cfg.BATCHSIZE, 24 | shuffle=True if mode in ['train'] else False, 25 | num_workers=cfg.DATALOADER_WORKERS, pin_memory=True, 26 | drop_last=True) 27 | 28 | # prepare models 29 | ModelClass = get_model(cfg.MODEL) 30 | model = ModelClass(cfg) 31 | 32 | # prepare logger 33 | LoggerClass = get_logger(cfg.LOGGER) 34 | logger = LoggerClass(cfg) 35 | 36 | # register dataset, models, logger to trainer 37 | trainer = Trainer(cfg, model, dataloader_dict, logger) 38 | 39 | # start training 40 | epoch_total = cfg.EPOCH_TOTAL + (cfg.RESUME_EPOCH_ID if cfg.RESUME else 0) 41 | while trainer.do_epoch() <= cfg.EPOCH_TOTAL: 42 | pass 43 | -------------------------------------------------------------------------------- /scripts/evaluate_pix2surf.sh: -------------------------------------------------------------------------------- 1 | source activate pix2surf 2 | 3 | python run_batch.py --config=plain/eval_pix2surf_sv_car.yaml 4 | python run_batch.py --config=plain/eval_pix2surf_mv_car.yaml 5 | 6 | python run_batch.py --config=plain/eval_pix2surf_sv_chair.yaml 7 | python run_batch.py --config=plain/eval_pix2surf_mv_chair.yaml 8 | 9 | python run_batch.py --config=plain/eval_pix2surf_sv_plane.yaml 10 | python run_batch.py --config=plain/eval_pix2surf_mv_plane.yaml -------------------------------------------------------------------------------- /scripts/evaluate_pix2surf_coco.sh: -------------------------------------------------------------------------------- 1 | source activate pix2surf 2 | 3 | python run_batch.py --config=coco/eval_pix2surf_sv_car.yaml 4 | python run_batch.py --config=coco/eval_pix2surf_mv_car.yaml 5 | 6 | python run_batch.py --config=coco/eval_pix2surf_sv_chair.yaml 7 | python run_batch.py --config=coco/eval_pix2surf_mv_chair.yaml 8 | 9 | python run_batch.py --config=coco/eval_pix2surf_sv_plane.yaml 10 | python run_batch.py --config=coco/eval_pix2surf_mv_plane.yaml -------------------------------------------------------------------------------- /scripts/pix2surf_car.sh: -------------------------------------------------------------------------------- 1 | source activate pix2surf 2 | 3 | python run_batch.py --config=plain/xnocsinit_sv_car.yaml 4 | python run_batch.py --config=plain/pix2surf_sv_car.yaml 5 | python run_batch.py --config=plain/pix2surf_mv_car.yaml -------------------------------------------------------------------------------- /scripts/pix2surf_car_coco.sh: -------------------------------------------------------------------------------- 1 | source activate pix2surf 2 | 3 | python run_batch.py --config=coco/xnocsinit_sv_car.yaml 4 | python run_batch.py --config=coco/pix2surf_sv_car.yaml 5 | python run_batch.py --config=coco/pix2surf_mv_car.yaml -------------------------------------------------------------------------------- /scripts/pix2surf_chair.sh: -------------------------------------------------------------------------------- 1 | source activate pix2surf 2 | 3 | python run_batch.py --config=plain/xnocsinit_sv_chair.yaml 4 | python run_batch.py --config=plain/pix2surf_sv_chair.yaml 5 | python run_batch.py --config=plain/pix2surf_mv_chair.yaml -------------------------------------------------------------------------------- /scripts/pix2surf_chair_coco.sh: -------------------------------------------------------------------------------- 1 | source activate pix2surf 2 | 3 | python run_batch.py --config=coco/xnocsinit_sv_chair.yaml 4 | python run_batch.py --config=coco/pix2surf_sv_chair.yaml 5 | python run_batch.py --config=coco/pix2surf_mv_chair.yaml -------------------------------------------------------------------------------- /scripts/pix2surf_plane.sh: -------------------------------------------------------------------------------- 1 | source activate pix2surf 2 | 3 | python run_batch.py --config=plain/xnocsinit_sv_plane.yaml 4 | python run_batch.py --config=plain/pix2surf_sv_plane.yaml 5 | python run_batch.py --config=plain/pix2surf_mv_plane.yaml -------------------------------------------------------------------------------- /scripts/pix2surf_plane_coco.sh: -------------------------------------------------------------------------------- 1 | source activate pix2surf 2 | 3 | python run_batch.py --config=coco/xnocsinit_sv_plane.yaml 4 | python run_batch.py --config=coco/pix2surf_sv_plane.yaml 5 | python run_batch.py --config=coco/pix2surf_mv_plane.yaml -------------------------------------------------------------------------------- /scripts/render_pix2surf_coco.sh: -------------------------------------------------------------------------------- 1 | source activate pix2surf 2 | 3 | python run_batch.py --config=coco/render_pix2surf_sv_car.yaml 4 | python run_batch.py --config=coco/render_pix2surf_mv_car.yaml 5 | 6 | python run_batch.py --config=coco/render_pix2surf_sv_chair.yaml 7 | python run_batch.py --config=coco/render_pix2surf_mv_chair.yaml 8 | 9 | python run_batch.py --config=coco/render_pix2surf_sv_plane.yaml 10 | python run_batch.py --config=coco/render_pix2surf_mv_plane.yaml -------------------------------------------------------------------------------- /scripts/render_pix2surf_plain.sh: -------------------------------------------------------------------------------- 1 | source activate pix2surf 2 | 3 | python run_batch.py --config=plain/render_pix2surf_sv_car.yaml 4 | python run_batch.py --config=plain/render_pix2surf_mv_car.yaml 5 | 6 | python run_batch.py --config=plain/render_pix2surf_sv_chair.yaml 7 | python run_batch.py --config=plain/render_pix2surf_mv_chair.yaml 8 | 9 | python run_batch.py --config=plain/render_pix2surf_sv_plane.yaml 10 | python run_batch.py --config=plain/render_pix2surf_mv_plane.yaml -------------------------------------------------------------------------------- /scripts/render_pix2surf_viz.sh: -------------------------------------------------------------------------------- 1 | source activate pix2surf 2 | 3 | python run_batch.py --config=viz/render_pix2surf_mv_car.yaml 4 | python run_batch.py --config=viz/render_pix2surf_mv_chair.yaml 5 | python run_batch.py --config=viz/render_pix2surf_mv_plane.yaml --------------------------------------------------------------------------------