[
{
"path": ".gitignore",
"content": "**/*.old\n**/*.bak\n\n.DS_Store\n# Created by https://www.toptal.com/developers/gitignore/api/vscode,python,jupyternotebooks\n# Edit at https://www.toptal.com/developers/gitignore?templates=vscode,python,jupyternotebooks\n\n### JupyterNotebooks ###\n# gitignore template for Jupyter Notebooks\n# website: http://jupyter.org/\n\n.ipynb_checkpoints\n*/.ipynb_checkpoints/*\n\n# IPython\nprofile_default/\nipython_config.py\n\n# Remove previous ipynb_checkpoints\n# git rm -r .ipynb_checkpoints/\n\n### Python ###\n# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\n#lib/\n#lib64/\nparts/\nsdist/\nvar/\nwheels/\npip-wheel-metadata/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\npytestdebug.log\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\ndoc/_build/\n\n# PyBuilder\ntarget/\n\n# Jupyter Notebook\n\n# IPython\n\n# pyenv\n.python-version\n\n# pipenv\n# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n# However, in case of collaboration, if having platform-specific dependencies or dependencies\n# having no cross-platform support, pipenv may install dependencies that don't work, or not\n# install all needed dependencies.\n#Pipfile.lock\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\npythonenv*\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n\n# pytype static type analyzer\n.pytype/\n\n# profiling data\n.prof\n\n### vscode ###\n.vscode/*\n!.vscode/settings.json\n!.vscode/tasks.json\n!.vscode/launch.json\n!.vscode/extensions.json\n*.code-workspace\n\n# End of https://www.toptal.com/developers/gitignore/api/vscode,python,jupyternotebooks\n"
},
{
"path": "LICENSE.md",
"content": " Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright [yyyy] [name of copyright owner]\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n"
},
{
"path": "README.md",
"content": "# RNNPose: Recurrent 6-DoF Object Pose Refinement with Robust Correspondence Field Estimation and Pose Optimization\n\n[Yan Xu](https://decayale.github.io/), [Kwan-Yee Lin](https://kwanyeelin.github.io/), [Guofeng Zhang](http://www.cad.zju.edu.cn/home/gfzhang/), [Xiaogang Wang](https://www.ee.cuhk.edu.hk/en-gb/people/academic-staff/professors/prof-xiaogang-wang), [Hongsheng Li](http://www.ee.cuhk.edu.hk/~hsli/). \n\n*Conference on Computer Vision and Pattern Recognition (CVPR), 2022.*\n\n[[Paper]](https://scholar.google.com/scholar?hl=zh-CN&as_sdt=0%2C5&q=RNNPose%3A+Recurrent+6-DoF+Object+Pose+Refinement+with+Robust+Correspondence+Field+Estimation+and+Pose+Optimization&btnG=)\n\n\n\n\n## 1. Framework \nThe basic pipeline of our proposed RNNPose. (a) Before refinement, a reference image is rendered according to the object initial pose (shown in a fused view).\n(b) Our RNN-based framework recurrently refines the object pose based on the estimated correspondence field between the reference and target images. The pose is optimized to be consistent with the reliable correspondence estimations highlighted by the similarity score map (built from learned 3D-2D descriptors) via differentiable LM optimization. (c) The output refined pose. \n\n\n\n![\"alt](\"./demo/idea.png\")
\n
\n\n## 2. Pose Estimation with Occlusions and Erroneous Pose Initializations\n\n\n### Estimated Poses and Intermediate System Outputs from Different Recurrent Iterations. \n\n\n ![\"animated\"](\"demo/ape_short_small.gif\")
![\"animated\"](\"demo/driller_short_small.gif\")
\n
\n\n\n### Pose Estimates with Erroneous Pose Initializations\nVisualization of our pose estimations (first row) on Occlusion LINEMOD dataset and the similarity score maps (second row) for downweighting unreliable correspondences during pose optimization. \nFor pose visualization, the white boxes represent the erroneous initial poses, the red boxes are estimated by our algorithm and the ground-truth boxes are in blue. Here, the initial poses for pose refinement are originally from PVNet but added with significant disturbances for robustness testing. \n\n ![](\"./demo/est_vis.png\")
\n\n\n\n## 3. Installation \n### Install the Docker \nA dockerfile is provided to help with the environment setup. \nYou need to install [docker](https://docs.docker.com/get-docker/) and [nvidia-docker2](https://github.com/NVIDIA/nvidia-docker) first and then set up the docker image and start up a container with the following commands: \n\n```\ncd RNNPose/docker\nsudo docker build -t rnnpose . \nsudo docker run -it --runtime=nvidia --ipc=host --volume=\"HOST_VOLUME_YOU_WANT_TO_MAP:DOCKER_VOLUME\" -e DISPLAY=$DISPLAY -e QT_X11_NO_MITSHM=1 rnnpose bash\n\n```\nIf you are not familiar with [docker](https://docs.docker.com/get-docker/), you could also install the dependencies manually following the provided dockerfile. \n\n### Compile the Dependencies\n```\ncd RNNPose/scripts\nbash compile_3rdparty.sh\n```\n\n\n## 4. Data Preparation\nWe follow [DeepIM](https://github.com/liyi14/mx-DeepIM) and [PVNet](https://github.com/zju3dv/pvnet-rendering) to preprocess the training data for LINEMOD. \nYou could follow the steps [here](doc/prepare_data.md) for data preparation. \n\n\n\n## 5. Test with the Pretrained Models\nWe train our model with the mixture of the real data and the synthetic data on LINEMOD dataset. \n\nThe trained models on the LINEMOD dataset have been uploaded to the [OneDrive](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155139432_link_cuhk_edu_hk/ESPTVyUryHdGl65fRAxN51gBBayJJb9NpCqWA-tY2CFKJQ?e=R9bcLW). \nYou can download them \nand put them into the directory *weight/* for testing. \n\n\nAn example bash script is provided below for reference. \n\n```\nexport PYTHONPATH=\"$PROJECT_ROOT_PATH:$PYTHONPATH\"\nexport PYTHONPATH=\"$PROJECT_ROOT_PATH/thirdparty:$PYTHONPATH\"\n\nseq=cat\ngpu=1\nstart_gpu_id=0\nmkdir $model_dir\n\ntrain_file=/home/yxu/Projects/Works/RNNPose_release/tools/eval.py\nconfig_path=/mnt/workspace/Works/RNNPose_release/config/linemod/\"$seq\"_fw0.5.yml\npretrain=$PROJECT_ROOT_PATH/weights/trained_models/\"$seq\".tckpt\n\npython -u $train_file multi_proc_train \\\n --config_path $config_path \\\n --model_dir $model_dir/results \\\n --use_dist True \\\n --dist_port 10000 \\\n --gpus_per_node $gpu \\\n --optim_eval True \\\n --use_apex True \\\n --world_size $gpu \\\n --start_gpu_id $start_gpu_id \\\n --pretrained_path $pretrain \n\n```\n\nNote that you need to specify the PROJECT_ROOT_PATH, i.e. the absolute directory of the project folder *RNNPose* and modify the respective data paths in the configuration files to the locations of downloaded data before executing the commands. You could also refer to the commands below for evaluation with our provide scripts.\n\n### Evaluation on LINEMOD\n```\nbash scripts/eval.sh \n```\n\n### Evaluation on LINEMOD OCCLUSION\n```\nbash scripts/eval_lmocc.sh\n\n```\n\n## Training from Scratch\nAn example training script is provided. \n```\nbash scripts/train.sh \n```\n\n\n\n## 6. Citation\nIf you find our code useful, please cite our paper. \n```\n@inproceedings{xu2022rnnpose,\n title={RNNPose: Recurrent 6-DoF Object Pose Refinement with Robust Correspondence Field Estimation and Pose Optimization},\n author={Xu, Yan and Kwan-Yee Lin and Zhang, Guofeng and Wang, Xiaogang and Li, Hongsheng},\n booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},\n year={2022}\n}\n\n@article{xu2024rnnpose,\n title={Rnnpose: 6-dof object pose estimation via recurrent correspondence field estimation and pose optimization},\n author={Xu, Yan and Lin, Kwan-Yee and Zhang, Guofeng and Wang, Xiaogang and Li, Hongsheng},\n journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},\n year={2024},\n publisher={IEEE}\n```\n\n\n## 7. Acknowledgement\n\nThe skeleton of this code is borrowed from [RSLO](https://github.com/DecaYale/RSLO). We also would like to thank the public codebases [PVNet](https://github.com/zju3dv/pvnet), [RAFT](https://github.com/princeton-vl/RAFT), [SuperGlue](https://github.com/magicleap/SuperGluePretrainedNetwork) and [DeepV2D](https://github.com/princeton-vl/DeepV2D). \n\n\n\n\n\n\n\n"
},
{
"path": "builder/__init__.py",
"content": ""
},
{
"path": "builder/dataset_builder.py",
"content": "\nfrom data.dataset import get_dataset_class\nimport numpy as np\nfrom functools import partial\nfrom data.preprocess import preprocess, preprocess_deepim, patch_crop\n\n\ndef build(input_reader_config,\n training,\n ):\n\n prep_cfg = input_reader_config.preprocess\n dataset_cfg = input_reader_config.dataset\n cfg = input_reader_config\n \n dataset_cls = get_dataset_class(dataset_cfg.dataset_class_name)\n\n if 0:# 'DeepIM' in dataset_cfg.dataset_class_name:\n # patch_cropper = partial(patch_crop, margin_ratio=0.2, output_size=256 )\n patch_cropper = None \n prep_func = partial(preprocess_deepim, \n max_points=dataset_cfg.max_points,\n correspondence_radius=prep_cfg.correspondence_radius_threshold,\n patch_cropper=patch_cropper,\n image_scale=prep_cfg.get('image_scale', 1),\n ) \n\n else:\n prep_func = partial(preprocess, \n max_points=dataset_cfg.max_points,\n correspondence_radius=prep_cfg.correspondence_radius_threshold,\n image_scale=prep_cfg.get('image_scale', 1),\n crop_param=prep_cfg.get('crop_param', None),\n kp_3d_param=prep_cfg.get('kp_3d_param', {\"kp_num\":30} ),\n use_coords_as_3d_feat=prep_cfg.get('use_coords_as_3d_feat', False)\n ) \n \n\n dataset = dataset_cls(\n info_path=dataset_cfg.info_path,\n root_path=dataset_cfg.root_path,\n model_point_dim=dataset_cfg.model_point_dim,\n is_train=training,\n prep_func=prep_func,\n seq_names=dataset_cfg.get('seq_names', None),\n cfg=dataset_cfg\n )\n\n return dataset\n"
},
{
"path": "builder/input_reader_builder.py",
"content": "\nfrom torch.utils.data import Dataset\n\nfrom builder import dataset_builder\n\n\nclass DatasetWrapper(Dataset):\n \"\"\" convert our dataset to Dataset class in pytorch.\n \"\"\"\n\n def __init__(self, dataset):\n self._dataset = dataset\n\n def __len__(self):\n return len(self._dataset)\n\n def __getitem__(self, idx):\n return self._dataset[idx]\n\n @property\n def dataset(self):\n return self._dataset\n\n\ndef build(input_reader_config,\n training,\n ) -> DatasetWrapper:\n\n dataset = dataset_builder.build(\n input_reader_config,\n training,\n )\n dataset = DatasetWrapper(dataset)\n return dataset\n"
},
{
"path": "builder/losses_builder.py",
"content": "\nfrom model import losses\n\ndef build(loss_config):\n\n criterions = {}\n \n criterions[\"metric_loss\"] =losses.MetricLoss(configs=loss_config.metric_loss,)\n\n criterions[\"pose_loss\"] = losses.PointAlignmentLoss(loss_weight=1)\n \n return criterions\n"
},
{
"path": "builder/lr_scheduler_builder.py",
"content": "\nfrom torchplus.train import learning_schedules_fastai as lsf\nimport torch\nimport numpy as np \n\ndef build(optimizer_config, optimizer, total_step):\n\n optimizer_type = list(optimizer_config.keys())[0]\n\n if optimizer_type == 'rms_prop_optimizer':\n config = optimizer_config.rms_prop_optimizer\n lr_scheduler = _create_learning_rate_scheduler(\n config.learning_rate, optimizer, total_step=total_step)\n\n if optimizer_type == 'momentum_optimizer':\n config = optimizer_config.momentum_optimizer\n lr_scheduler = _create_learning_rate_scheduler(\n config.learning_rate, optimizer, total_step=total_step)\n\n if optimizer_type == 'adam_optimizer':\n config = optimizer_config.adam_optimizer\n lr_scheduler = _create_learning_rate_scheduler(\n config.learning_rate, optimizer, total_step=total_step)\n\n return lr_scheduler\n\n\ndef _create_learning_rate_scheduler(learning_rate_config, optimizer, total_step):\n \"\"\"Create optimizer learning rate scheduler based on config.\n\n Args:\n learning_rate_config: A LearningRate proto message.\n\n Returns:\n A learning rate.\n\n Raises:\n ValueError: when using an unsupported input data type.\n \"\"\"\n lr_scheduler = None\n # learning_rate_type = learning_rate_config.WhichOneof('learning_rate')\n learning_rate_type = list(learning_rate_config.keys())[0]\n\n if learning_rate_type == 'multi_phase':\n config = learning_rate_config.multi_phase\n lr_phases = []\n mom_phases = []\n for phase_cfg in config.phases:\n lr_phases.append((phase_cfg.start, phase_cfg.lambda_func))\n mom_phases.append(\n (phase_cfg.start, phase_cfg.momentum_lambda_func))\n lr_scheduler = lsf.LRSchedulerStep(\n optimizer, total_step, lr_phases, mom_phases)\n\n\n\n if learning_rate_type == 'one_cycle':\n config = learning_rate_config.one_cycle\n\n if len(config.lr_maxs)>1:\n assert(len(config.lr_maxs)==4 ) \n lr_max=[]\n # for i in range(len(config.lr_maxs)):\n # lr_max += [config.lr_maxs[i]]*optimizer.param_segs[i] \n\n lr_max = np.array(list(config.lr_maxs) )\n else:\n lr_max = config.lr_max\n\n lr_scheduler = lsf.OneCycle(\n optimizer, total_step, lr_max, list(config.moms), config.div_factor, config.pct_start)\n if learning_rate_type == 'exponential_decay':\n config = learning_rate_config.exponential_decay\n lr_scheduler = lsf.ExponentialDecay(\n optimizer, total_step, config.initial_learning_rate, config.decay_length, config.decay_factor, config.staircase)\n if learning_rate_type == 'exponential_decay_warmup':\n config = learning_rate_config.exponential_decay_warmup\n lr_scheduler = lsf.ExponentialDecayWarmup(\n optimizer, total_step, config.initial_learning_rate, config.decay_length, config.decay_factor, config.div_factor,\n config.pct_start, config.staircase)\n if learning_rate_type == 'manual_stepping':\n config = learning_rate_config.manual_stepping\n lr_scheduler = lsf.ManualStepping(\n optimizer, total_step, list(config.boundaries), list(config.rates))\n\n if lr_scheduler is None:\n raise ValueError('Learning_rate %s not supported.' %\n learning_rate_type)\n\n return lr_scheduler\n"
},
{
"path": "builder/optimizer_builder.py",
"content": "from torchplus.train import learning_schedules\nfrom torchplus.train import optim\nimport torch\nfrom torch import nn\nfrom torchplus.train.fastai_optim import OptimWrapper, FastAIMixedOptim\nfrom functools import partial\n\n\ndef children(m: nn.Module):\n \"Get children of `m`.\"\n return list(m.children())\n\n\ndef num_children(m: nn.Module) -> int:\n \"Get number of children modules in `m`.\"\n return len(children(m))\n\n# return a list of smallest modules dy\n\n\ndef flatten_model(m):\n if m is None:\n return []\n return sum(\n map(flatten_model, m.children()), []) if num_children(m) else [m]\n\n\n# def get_layer_groups(m): return [nn.Sequential(*flatten_model(m))]\ndef get_layer_groups(m): return [nn.ModuleList(flatten_model(m))]\n\ndef get_voxeLO_net_layer_groups(net):\n vfe_grp = get_layer_groups(net)#[0]\n\n other_grp = get_layer_groups(nn.Sequential(net._rotation_loss,\n net._translation_loss,\n net._pyramid_rotation_loss,\n net._pyramid_translation_loss,\n net._consistency_loss, \n ))\n\n return [vfe_grp, mfe_grp, op_grp,other_grp]\n\n\ndef get_voxeLO_net_layer_groups(net):\n vfe_grp = get_layer_groups(net.voxel_feature_extractor)#[0]\n mfe_grp = get_layer_groups(net.middle_feature_extractor)#[0]\n op_grp = get_layer_groups(net.odom_predictor)#[0]\n\n # other_grp = get_layer_groups(net._rotation_loss) + \\\n # get_layer_groups(net._translation_loss) \\\n # + get_layer_groups(net._pyramid_rotation_loss) \\\n # + get_layer_groups(net._pyramid_translation_loss) \\\n # + get_layer_groups(net._consistency_loss)\\\n other_grp = get_layer_groups(nn.Sequential(net._rotation_loss,\n net._translation_loss,\n net._pyramid_rotation_loss,\n net._pyramid_translation_loss,\n net._consistency_loss, \n ))\n\n return [vfe_grp, mfe_grp, op_grp,other_grp]\n\n\ndef build(optimizer_config, net, name=None, mixed=False, loss_scale=512.0):\n\n optimizer_type = list(optimizer_config.keys())[0]\n print(\"Optimizer:\", optimizer_type)\n \n optimizer=None\n\n if optimizer_type == 'rms_prop_optimizer':\n config=optimizer_config.rms_prop_optimizer\n optimizer_func=partial(\n torch.optim.RMSprop,\n alpha=config.decay,\n momentum=config.momentum_optimizer_value,\n eps=config.epsilon)\n\n if optimizer_type == 'momentum_optimizer':\n config=optimizer_config.momentum_optimizer\n optimizer_func=partial(\n torch.optim.SGD,\n momentum=config.momentum_optimizer_value,\n eps=config.epsilon)\n\n if optimizer_type == 'adam_optimizer':\n config=optimizer_config.adam_optimizer\n if optimizer_config.fixed_weight_decay:\n optimizer_func=partial(\n torch.optim.Adam, betas=(0.9, 0.99), amsgrad=config.amsgrad)\n else:\n # regular adam\n optimizer_func=partial(\n torch.optim.Adam, amsgrad=config.amsgrad)\n\n optimizer=OptimWrapper.create(\n optimizer_func,\n 3e-3,\n get_layer_groups(net),\n # get_voxeLO_net_layer_groups(net),\n wd=config.weight_decay,\n true_wd=optimizer_config.fixed_weight_decay,\n bn_wd=True)\n print(hasattr(optimizer, \"_amp_stash\"), '_amp_stash')\n if optimizer is None:\n raise ValueError('Optimizer %s not supported.' % optimizer_type)\n\n if optimizer_config.use_moving_average:\n raise ValueError('torch don\\'t support moving average')\n if name is None:\n # assign a name to optimizer for checkpoint system\n optimizer.name=optimizer_type\n else:\n optimizer.name=name\n return optimizer\n"
},
{
"path": "builder/rnnpose_builder.py",
"content": "from builder import losses_builder\nfrom model.RNNPose import get_posenet_class\nimport model.RNNPose\n\n\ndef build(model_cfg,\n measure_time=False, testing=False):\n \"\"\"build second pytorch instance.\n \"\"\"\n\n criterions=losses_builder.build(model_cfg.loss)\n\n net = get_posenet_class(model_cfg.network_class_name)(\n criterions=criterions,\n opt=model_cfg)\n return net\n"
},
{
"path": "config/default.py",
"content": "from yacs.config import CfgNode as CN\nfrom utils.singleton import Singleton\nimport os\n\ndef _merge_a_into_b(a, b):\n \"\"\"Merge config dictionary a into config dictionary b, clobbering the\n options in b whenever they are also specified in a.\n \"\"\"\n # if type(a) is not dict:\n if not isinstance(a, (dict)):\n return\n\n for k, v in a.items():\n # a must specify keys that are in b\n if not k in b:\n raise KeyError('{} is not a valid config key'.format(k))\n\n # the types must match, too\n old_type = type(b[k])\n if old_type is not type(v):\n if isinstance(b[k], np.ndarray):\n v = np.array(v, dtype=b[k].dtype)\n else:\n raise ValueError(('Type mismatch ({} vs. {}) '\n 'for config key: {}').format(type(b[k]),\n type(v), k))\n\n # recursively merge dicts\n # if type(v) is dict:\n if isinstance(v, (dict)):\n try:\n _merge_a_into_b(a[k], b[k])\n except:\n print('Error under config key: {}'.format(k))\n raise\n else:\n b[k] = v\n\n\n\nclass Config(metaclass=Singleton):\n def __init__(self):\n ############## ↓ Basic ↓ ##############\n self.ROOT_CN = CN()\n self.ROOT_CN.BASIC = CN()\n self.ROOT_CN.BASIC.input_size=[480,640] #h,w\n self.ROOT_CN.BASIC.crop_size=[320,320] #h,w\n self.ROOT_CN.BASIC.zoom_crop_size=[320,320] #h,w\n self.ROOT_CN.BASIC.render_image_size=[320,320]#h,w\n self.ROOT_CN.BASIC.patch_num=64#h,w\n\n ############## ↓ LM OPTIM ↓ ##############\n self.ROOT_CN.LM=CN()\n self.ROOT_CN.LM.LM_LMBDA= 0.0001\n self.ROOT_CN.LM.EP_LMBDA=100\n\n ############## ↓ data ↓ ##############\n self.ROOT_CN.DATA=CN()\n self.ROOT_CN.DATA.OBJ_ROOT=\"\" #h,w\n self.ROOT_CN.DATA.VOC_ROOT=f\"{os.path.dirname(os.path.abspath(__file__)) }/../EXPDATA/\" #h,w\n\n def __get_item__(self, key):\n return self.ROOT_CN.__getitem__(key)\n \n def merge(self, config_dict, sub_key=None):\n\n if sub_key is not None:\n _merge_a_into_b(config_dict, self.ROOT_CN[sub_key])\n else:\n _merge_a_into_b(config_dict, self.ROOT_CN)\n\n############## ↓ Model ↓ ##############\n# _CN.model = CN()\n# _CN.model.input_size=[480,640]\n# _CN.model.crop_size=[320,320] \n# def get_cfg_defaults():\n\ndef get_cfg(Node=None):\n \"\"\"Get a yacs CfgNode object with default values for my_project.\"\"\"\n # Return a clone so that the defaults will not be altered\n # This is for the \"local variable\" use pattern\n # return _CN.clone()\n if Node is None:\n return Config()\n else:\n return Config().__get_item__(Node)\n"
},
{
"path": "config/linemod/copy.sh",
"content": "declare -a arr=(\"glue\" \"ape\" \"cat\" \"phone\" \"eggbox\" \"benchvise\" \"lamp\" \"camera\" \"can\" \"driller\" \"duck\" \"holepuncher\" \"iron\" )\n\n#create training scripts\nfor seq in \"${arr[@]}\"\ndo\n echo \"$seq\"\n cat template_fw0.5.yml > \"$seq\"_fw0.5.yml\n sed -i \"s/SEQ_NAME/$seq/g\" \"$seq\"_fw0.5.yml\ndone\n\narraylength=${#arr[@]}\n"
},
{
"path": "config/linemod/copy_occ.sh",
"content": "declare -a arr=(\"glue\" \"ape\" \"cat\" \"phone\" \"eggbox\" \"benchvise\" \"lamp\" \"camera\" \"can\" \"driller\" \"duck\" \"holepuncher\" \"iron\" )\n\n#create training scripts\nfor seq in \"${arr[@]}\"\ndo\n echo \"$seq\"\n cat template_fw0.5_occ.yml > \"$seq\"_fw0.5_occ.yml\n sed -i \"s/SEQ_NAME/$seq/g\" \"$seq\"_fw0.5_occ.yml\ndone\n\n\n"
},
{
"path": "config/linemod/template_fw0.5.yml",
"content": "vars:\n input_h: &input_h\n 320 \n input_w: &input_w\n 320 \n batch_size: &batch_size\n 1\n descriptor_dim: &descriptor_dim\n 32 \n correspondence_radius_threshold: &correspondence_radius_threshold\n 0.01 #0.04 \n seq_name: &seq_name\n [\"SEQ_NAME\"]\nBASIC:\n zoom_crop_size: [240,240]\nmodel:\n input_h: *input_h\n input_w: *input_w\n batch_size: *batch_size\n seq_len: 2 \n\n\n network_class_name: RNNPose \n descriptor_net:\n module_class_name: HybridFeaNet \n \n keypoints_detector_2d:\n input_dim: 3\n descriptor_dim: *descriptor_dim \n remove_borders: 4\n normalize_output: True \n\n keypoints_detector_3d:\n #KPCONV configurations\n num_layers: 4\n KP_extent: 2.0\n batch_norm_momentum: 0.02\n use_batch_norm: true\n in_points_dim: 3\n fixed_kernel_points: 'center' #['center', 'verticals', 'none']\n KP_influence: 'linear'\n aggregation_mode: 'sum' #['closest', 'sum']\n modulated: false \n first_subsampling_dl: 0.025\n conv_radius: 2.5\n deform_radius: 5\n in_features_dim: 1 #3\n first_feats_dim: 128\n num_kernel_points: 15\n final_feats_dim: *descriptor_dim #256 #32\n normalize_output: True \n gnn_feats_dim: 128 #256\n context_fea_extractor_3d:\n #KPCONV configurations\n num_layers: 4\n KP_extent: 2.0\n batch_norm_momentum: 0.02\n use_batch_norm: true\n in_points_dim: 3\n fixed_kernel_points: 'center' #['center', 'verticals', 'none']\n KP_influence: 'linear'\n aggregation_mode: 'sum' #['closest', 'sum']\n modulated: false \n first_subsampling_dl: 0.025\n conv_radius: 2.5\n deform_radius: 5\n in_features_dim: 1 #3\n first_feats_dim: 128\n num_kernel_points: 15\n final_feats_dim: 256 #*descriptor_dim #256 #32\n normalize_output: False \n gnn_feats_dim: 128 #256\n motion_net:\n IS_CALIBRATED: True\n RESCALE_IMAGES: False\n ITER_COUNT: 4 \n RENDER_ITER_COUNT: 3 #2 #1 #3\n TRAIN_RESIDUAL_WEIGHT: 0 #0.1 \n TRAIN_FLOW_WEIGHT: 0.5 #0.1 #1 \n TRAIN_REPROJ_WEIGHT: 0 \n OPTIM_ITER_COUNT: 1\n FLOW_NET: 'raft' \n SYN_OBSERVED: False\n ONLINE_CROP: True\n raft:\n small: False #True\n fea_net: \"default\"\n mixed_precision: True\n # pretrained_model: \"/mnt/workspace/datasets/weights/models/raft-small.pth\"\n pretrained_model: \"/mnt/workspace/datasets/weights/models/raft-chairs.pth\"\n input_dim: 3 \n iters: 1 \n \n loss:\n metric_loss:\n type: \"normal\" \n pos_radius: *correspondence_radius_threshold # the radius used to find the positive correspondences\n safe_radius: 0.02 #0.13\n pos_margin: 0.1\n neg_margin: 1.4\n max_points: 256\n matchability_radius: 0.06\n weight: 0.001\n saliency_loss:\n loss_weight: 1\n reg_weight: 0.01\n geometric_loss:\n loss_weight: 1\n reg_weight: 0.5 #0.1 \n\n\ntrain_config:\n\n optimizer: \n adam_optimizer: \n learning_rate: \n one_cycle: \n lr_maxs: []\n lr_max: 0.0001 # \n moms: [0.95, 0.85]\n div_factor: 10.0\n pct_start: 0.01 #0.05\n amsgrad: false\n weight_decay: 0.0001 \n fixed_weight_decay: true\n use_moving_average: false\n\n steps: 200000 \n steps_per_eval: 10000\n loss_scale_factor: -1\n clear_metrics_every_epoch: true\n\ntrain_input_reader:\n dataset:\n dataset_class_name: \"LinemodDeepIMSynRealV2\" \n info_path: [\"/home/RNNPose/Projects/Works/RNNPose_release/EXPDATA/data_info/deepim/linemod_orig_deepim.info.train\", \"/home/RNNPose/Projects/Works/RNNPose_release/EXPDATA/data_info/deepim/linemod_syn_deepim.info.train\", \n \"/home/RNNPose/Projects/Works/RNNPose_release/EXPDATA/data_info/linemod_fusesformatted_all10k_deepim.info.train\",\n ] \n root_path: [\"/home/RNNPose/Projects/Works/RNNPose_release/EXPDATA/LM6d_converted/LM6d_refine\", \"/home/RNNPose/Projects/Works/RNNPose_release/EXPDATA/LM6d_converted/LM6d_refine_syn\",\n \"/home/RNNPose/Projects/Works/RNNPose_release/EXPDATA/LINEMOD/fuse_formatted/\"\n ]\n\n model_point_dim: 3\n max_points: 20000\n seq_names: *seq_name \n\n batch_size: *batch_size \n preprocess: \n correspondence_radius_threshold: *correspondence_radius_threshold\n num_workers: 3 \n image_scale: 1\n crop_param:\n rand_crop: false\n margin_ratio: 0.85 \n output_size: *input_h \n crop_with_init_pose: True\n \n\neval_input_reader:\n dataset:\n dataset_class_name: \"LinemodDeepIMSynRealV2\"\n info_path: [\"/home/RNNPose/Projects/Works/RNNPose_release/EXPDATA/data_info/linemod_posecnn.info.eval\" ] \n root_path: [\"/home/RNNPose/Projects/Works/RNNPose_release/EXPDATA/LM6d_converted/LM6d_refine\" ]\n model_point_dim: 3\n max_points: 20000\n seq_names: *seq_name \n batch_size: *batch_size \n preprocess:\n correspondence_radius_threshold: *correspondence_radius_threshold\n num_workers: 3 \n image_scale: 1\n crop_param:\n rand_crop: false\n margin_ratio: 0.85 #0.5 \n output_size: *input_h #\n crop_with_init_pose: True\n \n"
},
{
"path": "config/linemod/template_fw0.5_occ.yml",
"content": "vars:\n input_h: &input_h\n 320 \n input_w: &input_w\n 320 \n batch_size: &batch_size\n 1\n descriptor_dim: &descriptor_dim\n 32 \n correspondence_radius_threshold: &correspondence_radius_threshold\n 0.01 #0.04 \n seq_name: &seq_name\n [\"SEQ_NAME\"]\nBASIC:\n zoom_crop_size: [240,240]\nmodel:\n input_h: *input_h\n input_w: *input_w\n batch_size: *batch_size\n seq_len: 2 \n\n\n network_class_name: RNNPose \n descriptor_net:\n module_class_name: HybridFeaNet \n \n keypoints_detector_2d:\n input_dim: 3\n descriptor_dim: *descriptor_dim \n remove_borders: 4\n normalize_output: True \n\n keypoints_detector_3d:\n #KPCONV configurations\n num_layers: 4\n KP_extent: 2.0\n batch_norm_momentum: 0.02\n use_batch_norm: true\n in_points_dim: 3\n fixed_kernel_points: 'center' #['center', 'verticals', 'none']\n KP_influence: 'linear'\n aggregation_mode: 'sum' #['closest', 'sum']\n modulated: false \n first_subsampling_dl: 0.025\n conv_radius: 2.5\n deform_radius: 5\n in_features_dim: 1 #3\n first_feats_dim: 128\n num_kernel_points: 15\n final_feats_dim: *descriptor_dim #256 #32\n normalize_output: True \n gnn_feats_dim: 128 #256\n context_fea_extractor_3d:\n #KPCONV configurations\n num_layers: 4\n KP_extent: 2.0\n batch_norm_momentum: 0.02\n use_batch_norm: true\n in_points_dim: 3\n fixed_kernel_points: 'center' #['center', 'verticals', 'none']\n KP_influence: 'linear'\n aggregation_mode: 'sum' #['closest', 'sum']\n modulated: false \n first_subsampling_dl: 0.025\n conv_radius: 2.5\n deform_radius: 5\n in_features_dim: 1 #3\n first_feats_dim: 128\n num_kernel_points: 15\n final_feats_dim: 256 #*descriptor_dim #256 #32\n normalize_output: False \n gnn_feats_dim: 128 #256\n motion_net:\n IS_CALIBRATED: True\n RESCALE_IMAGES: False\n ITER_COUNT: 4 \n RENDER_ITER_COUNT: 3 #2 #1 #3\n TRAIN_RESIDUAL_WEIGHT: 0 #0.1 \n TRAIN_FLOW_WEIGHT: 0.5 #0.1 #1 \n TRAIN_REPROJ_WEIGHT: 0 \n OPTIM_ITER_COUNT: 1\n FLOW_NET: 'raft' \n SYN_OBSERVED: False\n ONLINE_CROP: True\n raft:\n small: False #True\n fea_net: \"default\"\n mixed_precision: True\n # pretrained_model: \"/mnt/workspace/datasets/weights/models/raft-small.pth\"\n pretrained_model: \"/mnt/workspace/datasets/weights/models/raft-chairs.pth\"\n input_dim: 3 \n iters: 1 \n \n loss:\n metric_loss:\n type: \"normal\" \n pos_radius: *correspondence_radius_threshold # the radius used to find the positive correspondences\n safe_radius: 0.02 #0.13\n pos_margin: 0.1\n neg_margin: 1.4\n max_points: 256\n matchability_radius: 0.06\n weight: 0.001\n saliency_loss:\n loss_weight: 1\n reg_weight: 0.01\n geometric_loss:\n loss_weight: 1\n reg_weight: 0.5 #0.1 \n\n\ntrain_config:\n\n optimizer: \n adam_optimizer: \n learning_rate: \n one_cycle: \n lr_maxs: []\n lr_max: 0.0001 # \n moms: [0.95, 0.85]\n div_factor: 10.0\n pct_start: 0.01 #0.05\n amsgrad: false\n weight_decay: 0.0001 \n fixed_weight_decay: true\n use_moving_average: false\n\n steps: 200000 \n steps_per_eval: 10000\n loss_scale_factor: -1\n clear_metrics_every_epoch: true\n\ntrain_input_reader:\n dataset:\n dataset_class_name: \"LinemodDeepIMSynRealV2\" \n info_path: [\"/home/RNNPose/Projects/Works/RNNPose_release/EXPDATA/data_info/deepim/linemod_orig_deepim.info.train\", \"/home/RNNPose/Projects/Works/RNNPose_release/EXPDATA/data_info/deepim/linemod_syn_deepim.info.train\", \n \"/home/RNNPose/Projects/Works/RNNPose_release/EXPDATA/data_info/linemod_fusesformatted_all10k_deepim.info.train\",\n ] \n root_path: [\"/home/RNNPose/Projects/Works/RNNPose_release/EXPDATA/LM6d_converted/LM6d_refine\", \"/home/RNNPose/Projects/Works/RNNPose_release/EXPDATA/LM6d_converted/LM6d_refine_syn\",\n \"/home/RNNPose/Projects/Works/RNNPose_release/EXPDATA/LINEMOD/fuse_formatted/\"\n ]\n\n model_point_dim: 3\n max_points: 20000\n seq_names: *seq_name \n\n batch_size: *batch_size \n preprocess: \n correspondence_radius_threshold: *correspondence_radius_threshold\n num_workers: 3 \n image_scale: 1\n crop_param:\n rand_crop: false\n margin_ratio: 0.85 \n output_size: *input_h \n crop_with_init_pose: True\n \n\neval_input_reader:\n dataset:\n dataset_class_name: \"LinemodDeepIMSynRealV2\"\n info_path: [\"/home/RNNPose/Projects/Works/RNNPose_release/EXPDATA/data_info/deepim/linemod_bop_lmocc_pvnetdr.info.eval\"]\n root_path: [\"/home/RNNPose/Projects/Works/RNNPose_release/EXPDATA/lmo\"] \n init_post_type: \"PVNET_LINEMOD_OCC\"\n model_point_dim: 3\n max_points: 20000\n seq_names: *seq_name \n batch_size: *batch_size \n preprocess:\n correspondence_radius_threshold: *correspondence_radius_threshold\n num_workers: 3 \n image_scale: 1\n crop_param:\n rand_crop: false\n margin_ratio: 0.85 #0.5 \n output_size: *input_h #\n crop_with_init_pose: True\n \n \n"
},
{
"path": "data/__init__.py",
"content": "from . import dataset\nfrom . import linemod_dataset"
},
{
"path": "data/dataset.py",
"content": "import pathlib\nimport pickle\nimport time\nfrom functools import partial\n\nimport numpy as np\n\n\nREGISTERED_DATASET_CLASSES = {}\n\n\ndef register_dataset(cls, name=None):\n global REGISTERED_DATASET_CLASSES\n if name is None:\n name = cls.__name__\n assert name not in REGISTERED_DATASET_CLASSES, f\"exist class: {REGISTERED_DATASET_CLASSES}\"\n REGISTERED_DATASET_CLASSES[name] = cls\n return cls\n\n\ndef get_dataset_class(name):\n global REGISTERED_DATASET_CLASSES\n assert name in REGISTERED_DATASET_CLASSES, f\"available class: {REGISTERED_DATASET_CLASSES}\"\n return REGISTERED_DATASET_CLASSES[name]\n\n\nclass Dataset(object):\n NumPointFeatures = -1\n\n def __getitem__(self, index):\n raise NotImplementedError\n\n def __len__(self):\n raise NotImplementedError\n\n def _read_data(self, query):\n\n raise NotImplementedError\n\n def evaluation(self, dt_annos, output_dir):\n \"\"\"Dataset must provide a evaluation function to evaluate model.\"\"\"\n raise NotImplementedError\n"
},
{
"path": "data/linemod/linemod_config.py",
"content": "import numpy as np\ndiameters = {\n 'cat': 15.2633,\n 'ape': 9.74298,\n 'benchvise': 28.6908,\n 'bowl': 17.1185,\n 'cam': 17.1593,\n 'camera': 17.1593,\n 'can': 19.3416,\n 'cup': 12.5961,\n 'driller': 25.9425,\n 'duck': 10.7131,\n 'eggbox': 17.6364,\n 'glue': 16.4857,\n 'holepuncher': 14.8204,\n 'iron': 30.3153,\n 'lamp': 28.5155,\n 'phone': 20.8394\n}\n\nlinemod_cls_names = ['ape', 'cam', 'cat', 'duck', 'glue', 'iron', 'phone', 'benchvise', 'can', 'driller', 'eggbox', 'holepuncher', 'lamp']\n\nlinemod_K = np.array([[572.4114, 0., 325.2611],\n [0., 573.57043, 242.04899],\n [0., 0., 1.]])\n\n\nblender_K = np.array([[700., 0., 320.],\n [0., 700., 240.],\n [0., 0., 1.]])"
},
{
"path": "data/linemod_dataset.py",
"content": "import numpy as np \nimport random\nimport os \nfrom data.dataset import Dataset, register_dataset\nimport pickle\nimport PIL\nimport cv2\nimport torch\nimport time\nimport scipy\n\nfrom utils.geometric import range_to_depth, render_pointcloud\nfrom .transforms import make_transforms \nfrom thirdparty.kpconv.lib.utils import square_distance\nfrom utils.geometric import rotation_angle\n# from utils.visualize import *\nfrom transforms3d.quaternions import mat2quat, quat2mat, qmult\nfrom transforms3d.euler import mat2euler, euler2mat, euler2quat, quat2euler\nimport math\nfrom config.default import get_cfg\n\nCURRENT_DIR=os.path.dirname(os.path.abspath(__file__))\n\ntry:\n from pytorch3d.io import load_obj, load_ply\nexcept:\n print(\"Warning: error occurs when importing pytorch3d \")\n pass\n\n\ndef se3_q2m(se3_q):\n assert se3_q.size == 7\n se3_mx = np.zeros((3, 4))\n # quat = se3_q[0:4] / LA.norm(se3_q[0:4])\n quat = se3_q[:4]\n R = quat2mat(quat)\n se3_mx[:, :3] = R\n se3_mx[:, 3] = se3_q[4:]\n return se3_mx\n\ndef info_convertor(info,):\n \"\"\"\n [Transform the original kitti info file]\n \"\"\"\n\n seqs = info.keys() #['cat']#\n seq_lengths = [len(info[i]) for i in seqs]\n data = []\n for seq in seqs:\n print(seq)\n data.append(info[seq])\n\n new_infos = {\n \"seqs\": list(seqs),\n \"seq_lengths\": seq_lengths,\n \"data\": data\n }\n return new_infos\n\ndef resize(im, target_size, max_size, stride=0, interpolation=cv2.INTER_LINEAR):\n \"\"\"\n only resize input image to target size and return scale\n :param im: BGR image input by opencv\n :param target_size: one dimensional size (the short side)\n :param max_size: one dimensional max size (the long side)\n :param stride: if given, pad the image to designated stride\n :param interpolation: if given, using given interpolation method to resize image\n :return:\n \"\"\"\n im_shape = im.shape\n im_size_min = np.min(im_shape[0:2])\n im_size_max = np.max(im_shape[0:2])\n im_scale = float(target_size) / float(im_size_min)\n # prevent bigger axis from being more than max_size:\n if np.round(im_scale * im_size_max) > max_size:\n im_scale = float(max_size) / float(im_size_max)\n im = cv2.resize(im, None, None, fx=im_scale, fy=im_scale, interpolation=interpolation)\n\n if stride == 0:\n return im, im_scale\n else:\n # pad to product of stride\n im_height = int(np.ceil(im.shape[0] / float(stride)) * stride)\n im_width = int(np.ceil(im.shape[1] / float(stride)) * stride)\n im_channel = im.shape[2]\n padded_im = np.zeros((im_height, im_width, im_channel))\n padded_im[: im.shape[0], : im.shape[1], :] = im\n return padded_im, im_scale\ndef sample_poses(pose_tgt):\n SYN_STD_ROTATION = 15\n SYN_STD_TRANSLATION = 0.01\n ANGLE_MAX=45\n pose_src = pose_tgt.copy()\n num = pose_tgt.shape[0]\n for i in range(num):\n euler = mat2euler(pose_tgt[i, :3, :3])\n euler += SYN_STD_ROTATION * np.random.randn(3) * math.pi / 180.0\n pose_src[i, :3, :3] = euler2mat(euler[0], euler[1], euler[2])\n\n pose_src[i, 0, 3] = pose_tgt[i, 0, 3]+ SYN_STD_TRANSLATION * np.random.randn(1)\n pose_src[i, 1, 3] = pose_tgt[i, 1, 3] + SYN_STD_TRANSLATION * np.random.randn(1)\n pose_src[i, 2, 3] = pose_tgt[i, 2, 3] + 5 * SYN_STD_TRANSLATION * np.random.randn(1)\n\n r_dist = np.arccos((np.trace(pose_src[i, :3,:3].transpose(-1,-2) @ pose_tgt[i, :3,:3]) - 1 )/2)/math.pi*180\n\n while r_dist > ANGLE_MAX:#or not (16 < center_x < (640 - 16) and 16 < center_y < (480 - 16)):\n # print(\"r_dist > ANGLE_MAX, resampling...\")\n print(\"Too large angular differences. Resample the pose...\")\n euler = mat2euler(pose_tgt[i, :3, :3])\n euler += SYN_STD_ROTATION * np.random.randn(3) * math.pi / 180.0\n pose_src[i, :3, :3] = euler2mat(euler[0], euler[1], euler[2])\n\n pose_src[i, 0, 3] = pose_tgt[i, 0, 3]+ SYN_STD_TRANSLATION * np.random.randn(1)\n pose_src[i, 1, 3] = pose_tgt[i, 1, 3] + SYN_STD_TRANSLATION * np.random.randn(1)\n pose_src[i, 2, 3] = pose_tgt[i, 2, 3] + 5 * SYN_STD_TRANSLATION * np.random.randn(1)\n\n r_dist = np.arccos((np.trace(pose_src[i, :3,:3].transpose(-1,-2) @ pose_tgt[i, :3,:3]) - 1 )/2)*math.pi/180\n return pose_src.squeeze()\n\n\n\n\n@register_dataset\nclass LinemodDeepIMSynRealV2(Dataset):\n # use deepim 3d model for geometric feature extraction, mingle the synthetic and real data \n def __init__(self, root_path,\n info_path, model_point_dim,\n is_train,\n prep_func=None,\n seq_names=None, \n cfg={}\n ):\n super().__init__()\n\n assert info_path is not None\n assert isinstance(root_path, (tuple, list)) and isinstance(info_path, (tuple, list))\n assert len(root_path) == len(info_path)\n print(\"Info:\",info_path)\n # assert split in ['train', 'val', 'test']\n self.is_train = is_train\n self.VOC_ROOT = get_cfg('DATA').VOC_ROOT#\"/DATA/yxu/LINEMOD_DEEPIM/\"\n\n infos=[]\n for ipath in info_path:\n with open(ipath, 'rb') as f:\n info = pickle.load(f)\n\n if seq_names is not None:\n for k in list(info.keys()):\n if k not in seq_names:\n del info[k]\n infos.append( info_convertor(info) )\n\n #merge multiple infos \n self.infos = infos[0]\n self.infos['dataset_idx'] = [0]*len(self.infos['seqs'])\n for i, info in enumerate(infos[1:]):\n for k in self.infos:\n if k == 'dataset_idx':\n self.infos[k].extend([i+1]*len(info['seqs']))\n else:\n self.infos[k].extend(info[k])\n\n\n self.root_paths = root_path\n self.model_point_dim = model_point_dim\n # self.max_points=max_points#30000\n self.prep_func=prep_func\n # self.rgb_transformer = None #make_transforms(None, is_train=is_train)\n self.rgb_transformer = make_transforms(None, is_train=is_train)\n print(\"dataset size:\",self.__len__())\n\n self.init_pose_type = cfg.get(\"init_post_type\", \"POSECNN_LINEMOD\" ) \n # self.init_pose_type = cfg.get(\"init_post_type\", \"PVNET_LINEMOD_OCC\" ) \n# self.init_pose_type = cfg.get(\"init_post_type\", \"PVNET_LINEMOD\" ) \n print(\"INIT_POSE_TYPE:\", self.init_pose_type)\n #Load posecnn results\n if not self.is_train:\n with open(f\"{CURRENT_DIR}/../EXPDATA/init_poses/linemod_posecnn_results.pkl\", 'rb') as f:\n self.pose_cnn_results_test_posecnn=pickle.load(f)\n try:\n if self.init_pose_type == \"POSECNN_LINEMOD\":\n #load posecnn results \n self.pose_cnn_results_test=self.pose_cnn_results_test_posecnn\n elif self.init_pose_type ==\"PVNET_LINEMOD\":\n self.pose_cnn_results_test=np.load(f\"{CURRENT_DIR}/../EXPDATA/init_poses/pvnet/pvnet_linemod_test.npy\", allow_pickle=True).flat[0]\n elif self.init_pose_type ==\"PVNET_LINEMOD_OCC\":\n self.pose_cnn_results_test=np.load(f\"{CURRENT_DIR}/../EXPDATA/init_poses/pvnet/pvnet_linemodocc_test.npy\", allow_pickle=True).flat[0]\n else: \n raise NotImplementedError\n except:\n print(\"Loading posecnn results failed!\")\n self.pose_cnn_results_test=None\n try:\n # self.blender_to_bop_pose=np.load(f\"{CURRENT_DIR}/../EXPDATA/init_poses/metricpose/blender2bop_RT.npy\", allow_pickle=True).flat[0]\n self.blender_to_bop_pose=np.load(f\"{CURRENT_DIR}/../EXPDATA/init_poses/pose_conversion/blender2bop_RT.npy\", allow_pickle=True).flat[0]\n except:\n print(\"Loading pose conversion matrix failed!\")\n self.blender_to_bop_pose=None \n \n else:\n self.pose_cnn_results_test=None\n self.blender_to_bop_pose=None\n \n def load_random_background(self, im_observed, mask):\n VOC_root = os.path.join(self.VOC_ROOT, \"VOCdevkit/VOC2012\")\n VOC_image_set_dir = os.path.join(VOC_root, \"ImageSets/Main\")\n VOC_bg_list_path = os.path.join(VOC_image_set_dir, \"diningtable_trainval.txt\")\n with open(VOC_bg_list_path, \"r\") as f:\n VOC_bg_list = [\n line.strip(\"\\r\\n\").split()[0] for line in f.readlines() if line.strip(\"\\r\\n\").split()[1] == \"1\"\n ]\n height, width, channel = im_observed.shape\n target_size = min(height, width)\n max_size = max(height, width)\n observed_hw_ratio = float(height) / float(width)\n\n k = random.randint(0, len(VOC_bg_list) - 1)\n bg_idx = VOC_bg_list[k]\n bg_path = os.path.join(VOC_root, \"JPEGImages/{}.jpg\".format(bg_idx))\n bg_image = cv2.imread(bg_path, cv2.IMREAD_COLOR)[...,::-1] #RGB\n bg_h, bg_w, bg_c = bg_image.shape\n bg_image_resize = np.zeros((height, width, channel), dtype=\"uint8\")\n if (float(height) / float(width) < 1 and float(bg_h) / float(bg_w) < 1) or (\n float(height) / float(width) >= 1 and float(bg_h) / float(bg_w) >= 1\n ):\n if bg_h >= bg_w:\n bg_h_new = int(np.ceil(bg_w * observed_hw_ratio))\n if bg_h_new < bg_h:\n bg_image_crop = bg_image[0:bg_h_new, 0:bg_w, :]\n else:\n bg_image_crop = bg_image\n else:\n bg_w_new = int(np.ceil(bg_h / observed_hw_ratio))\n if bg_w_new < bg_w:\n bg_image_crop = bg_image[0:bg_h, 0:bg_w_new, :]\n else:\n bg_image_crop = bg_image\n else:\n if bg_h >= bg_w:\n bg_h_new = int(np.ceil(bg_w * observed_hw_ratio))\n bg_image_crop = bg_image[0:bg_h_new, 0:bg_w, :]\n else: # bg_h < bg_w\n bg_w_new = int(np.ceil(bg_h / observed_hw_ratio))\n print(bg_w_new)\n bg_image_crop = bg_image[0:bg_h, 0:bg_w_new, :]\n\n bg_image_resize_0, _ = resize(bg_image_crop, target_size, max_size)\n h, w, c = bg_image_resize_0.shape\n bg_image_resize[0:h, 0:w, :] = bg_image_resize_0\n\n # add background to image_observed\n res_image = bg_image_resize.copy()\n res_image[mask>0]=im_observed[mask>0]\n\n # im_observed = res_image\n return res_image\n\n def _read_data(self, idx):\n \"\"\"\n info structure:\n {\n 'cat':[\n {\n \"index\": idx,\n \"model_path\": str,\n \"rgb_path\": str,\n \"depth_path\": str,\n \"RT\": np.array([3,4]),\n \"K\": np.array([3,3]),\n },\n {\n \"index\": idx,\n \"model_path\": str,\n \"rgb_path\": str,\n \"depth_path\": str,\n \"RT\": np.array([3,4]),\n \"K\": np.array([3,3]),\n }\n ...\n ],\n 'dog':[\n\n ]\n ...\n }\n\n \"\"\"\n\n if isinstance(idx, (tuple, list)):\n idx, seed = idx\n else:\n seed = None\n\n seq_lengths = np.array(self.infos['seq_lengths'])\n seq_lengths_cum = np.cumsum(seq_lengths)\n seq_lengths_cum = np.insert(seq_lengths_cum, 0, 0) # insert a dummy 0\n seq_idx = np.nonzero(seq_lengths_cum > idx)[0][0]-1\n\n frame_idx = idx - seq_lengths_cum[seq_idx]\n\n info = self.infos[\"data\"][seq_idx]\n dataset_idx = self.infos[\"dataset_idx\"][seq_idx]\n \n\n model_points_path = os.path.join(f'{os.path.dirname(__file__)}/../EXPDATA/LM6d_converted/models/{self.infos[\"seqs\"][seq_idx]}/textured.obj' ) # TODO: need check\n\n rgb_path = os.path.join(self.root_paths[dataset_idx], info[frame_idx]['rgb_observed_path']) \n depth_path = os.path.join(self.root_paths[dataset_idx], info[frame_idx]['depth_gt_observed_path']) \n\n if info[frame_idx].get('rgb_noisy_rendered', None) is not None:\n rgb_noisy_rendered_path = os.path.join(self.root_paths[dataset_idx], info[frame_idx]['rgb_noisy_rendered']) \n else:\n rgb_noisy_rendered_path = None\n if info[frame_idx].get('depth_noisy_rendered', None) is not None:\n depth_noisy_rendered_path = os.path.join(self.root_paths[dataset_idx], info[frame_idx]['depth_noisy_rendered']) \n else:\n depth_noisy_rendered_path = None\n\n if info[frame_idx].get('pose_noisy_rendered', None) is not None:\n rendered_RT = info[frame_idx]['pose_noisy_rendered'].astype(np.float32)\n # else:\n elif self.is_train:\n rendered_RT = sample_poses( info[frame_idx]['gt_pose'].astype(np.float32)[None] )\n\n K = info[frame_idx]['K'].astype(np.float32)\n RT = info[frame_idx]['gt_pose'].astype(np.float32) #[R,t]\n\n # evaluation \n if not self.is_train:\n if self.pose_cnn_results_test is not None:\n class_name=self.infos[\"seqs\"][seq_idx]\n\n if self.init_pose_type == \"PVNET_LINEMOD\":\n try:\n posecnn_RT = self.pose_cnn_results_test[class_name][frame_idx] # if self.pose_cnn_results_test is not None else np.zeros_like(RT)\n #Transformations are needed as the pvnet has a different coordinate system. \n posecnn_RT[:3,:3] = posecnn_RT[:3,:3]@self.blender_to_bop_pose[class_name][:3,:3].T\n posecnn_RT[:3,3:] = -posecnn_RT[:3,:3] @self.blender_to_bop_pose[class_name][:3,3:] + posecnn_RT[:3,3:] \n except:\n print(\"Warning: frame_idx is out of the range of self.pose_cnn_results_test!\", flush=True)\n posecnn_RT= se3_q2m(self.pose_cnn_results_test_posecnn[class_name][frame_idx]['pose']) #np.zeros_like(RT)\n elif self.init_pose_type ==\"POSECNN_LINEMOD\":\n posecnn_RT= se3_q2m(self.pose_cnn_results_test_posecnn[class_name][frame_idx]['pose']) \n elif self.init_pose_type == \"PVNET_LINEMOD_OCC\":\n try:\n posecnn_RT = self.pose_cnn_results_test[class_name][frame_idx].copy()# if self.pose_cnn_results_test is not None else np.zeros_like(RT)\n #Transformations are needed as the pvnet has a different coordinate system. \n posecnn_RT[:3,:3] = posecnn_RT[:3,:3]@self.blender_to_bop_pose[class_name][:3,:3].T\n posecnn_RT[:3,3:] = -posecnn_RT[:3,:3] @self.blender_to_bop_pose[class_name][:3,3:] + posecnn_RT[:3,3:] \n except:\n # print(frame_idx)\n raise\n else:\n raise NotImplementedError \n \n rendered_RT = posecnn_RT\n else:\n print(\"Warning: fail to load cnn poses!\", flush=True)\n posecnn_RT = np.zeros_like(RT)\n else:\n posecnn_RT = np.zeros_like(RT)\n \n #add noise--for testing purpose only, should always be disabled in normal cases \n# rot_std=0; trans_std=0.04; ang_max=1000;\n# print(f\"Add pose noises rot_std={rot_std}, trans_std={trans_std}\", flush=True)\n# rendered_RT=sample_poses(rendered_RT[None], rot_std=rot_std, trans_std=trans_std, ang_max=ang_max) \n\n # Regularize the matrix to be a valid rotation\n rendered_RT[:3,:3] = rendered_RT[:3,:3]@ np.linalg.inv(scipy.linalg.sqrtm(rendered_RT[:3,:3].T@rendered_RT[:3,:3]))\n \n # model_points = np.fromfile(\n # str(model_points_path), dtype=np.float32, count=-1).reshape([-1, self.model_point_dim]) # N x model_point_dim\n model_points, _,_ = load_obj(str(model_points_path) )\n model_points = model_points.numpy()\n \n visb = model_points[:,-1:] # N x model_point_dim\n\n model_point_features=np.ones_like(model_points[:,:1]).astype(np.float32)\n\n\n rgb = np.asarray(PIL.Image.open(rgb_path))\n\n if depth_path.endswith('.npy'):\n depth = np.load(depth_path) # blender \n else:\n depth = cv2.imread(depth_path, -1) /1000.\n\n if self.is_train and \"LM6d_refine_syn\" in self.root_paths[dataset_idx]: #synthetic data\n rgb = self.load_random_background(rgb, mask=(depth>0)[...,None].repeat(rgb.shape[-1], axis=-1) )\n\n\n \n rgb_rendered = np.asarray(PIL.Image.open(rgb_noisy_rendered_path)) if rgb_noisy_rendered_path is not None else None\n depth_rendered = np.asarray(PIL.Image.open(depth_noisy_rendered_path))/1000 if depth_noisy_rendered_path is not None else None #TODO: need check\n\n ren_mask = render_pointcloud(model_points, rendered_RT[None],K=K[None], render_image_size=rgb.shape[:2] ).squeeze()>0\n # depth = range_to_depth(depth<1, depth*2, K)\n\n return {\n \"class_name\": self.infos[\"seqs\"][seq_idx], \n \"idx\": idx,\n \"model_points\": model_points,\n \"visibility\": visb,\n \"model_point_features\":model_point_features,\n \"image\": rgb,\n \"depth\": depth,\n \"mask\": depth>0,\n \"rendered_image\": rgb_rendered,\n \"rendered_depth\": depth_rendered,\n \"K\": K,\n \"RT\": RT,\n \"rendered_RT\": rendered_RT.astype(np.float32),\n \"ren_mask\":ren_mask,\n \"POSECNN_RT\": posecnn_RT.astype(np.float32), #for test, TODO\n \"scale\": 1 # model_scale * scale = depth_scale\n }\n\n\n\n def __getitem__(self, idx):\n\n data=self._read_data(idx) \n try:\n data_p=self.prep_func(data, rand_rgb_transformer=self.rgb_transformer, find_2d3d_correspondence=self.is_train )\n except Exception as e: \n if e.args[0] in [\"Too few correspondences are found!\"] :\n if isinstance(idx, (tuple, list)):\n # idx, seed = idx\n idx = [(idx[0]+1)%self.__len__(), idx[1]]\n else:\n idx = (idx+1) %self.__len__()\n data_p= self.__getitem__(idx )\n else:\n raise ValueError\n\n return data_p\n\n def __len__(self):\n return np.sum(self.infos['seq_lengths'])\n"
},
{
"path": "data/preprocess.py",
"content": "import open3d as o3d\nimport copy\nimport os\n\nimport pathlib\nimport pickle\nimport time\nfrom collections import defaultdict\nfrom functools import partial\n\nimport cv2\nimport numpy as np\nimport quaternion\nfrom skimage import io as imgio\nfrom utils.timer import simple_timer\n\nimport matplotlib.pyplot as plt\nfrom collections.abc import Iterable\nimport torch\nimport torch.nn.functional as F\nimport quaternion\n\nfrom functools import partial\nimport thirdparty.kpconv.cpp_wrappers.cpp_subsampling.grid_subsampling as cpp_subsampling\nimport thirdparty.kpconv.cpp_wrappers.cpp_neighbors.radius_neighbors as cpp_neighbors\nfrom thirdparty.kpconv.lib.timer import Timer\nfrom utils.geometric import range_to_depth, mask_depth_to_point_cloud\nfrom utils.furthest_point_sample import fragmentation_fps\nfrom utils.rand_utils import truncated_normal\n\n\n\ndef merge_batch(batch_list):\n # [batch][key][seq]->example[key][seq][batch]\n # Or [batch][key]->example[key][batch]\n example_merged = defaultdict(list)\n for example in batch_list: # batch dim\n for k, v in example.items(): # key dim\n # assert isinstance(v, list)\n if isinstance(v, list):\n seq_len = len(v)\n if k not in example_merged:\n example_merged[k] = [[] for i in range(seq_len)]\n for i, vi in enumerate(v): # seq dim\n example_merged[k][i].append(vi)\n\n else:\n example_merged[k].append(v)\n\n ret = {}\n for key, elems in example_merged.items():\n if key in ['model_points', \"original_model_points\", 'visibility']:\n # concat the points of lenghts (N1,N2...) to a longer one with length (N1+N2+...)\n ret[key] = np.concatenate(elems, axis=0)\n # record the point numbers for original batches\n ret['batched_model_point_lengths'] = np.array(\n [len(p) for p in elems], dtype=np.int32)\n elif key in ['rand_model_points', ]:\n # concat the points of lenghts (N1,N2...) to a longer one with length (N1+N2+...)\n ret[key] = np.concatenate(elems, axis=0)\n # record the point numbers for original batches\n ret['batched_rand_model_point_lengths'] = np.array(\n [len(p) for p in elems], dtype=np.int32)\n elif key in ['model_point_features']:\n ret[key] = np.concatenate(elems, axis=0)\n\n # ['odometry/tq','odometry/RT','odometry/invRT' ]:\n elif key in ['image', 'depth', 'K', 'RT', 'original_RT' ,'rand_RT', 'correspondences_2d3d', 'scale', 'POSECNN_RT','rendered_image', 'rendered_depth', 'rendered_RT', '3d_keypoint_inds', '3d_keypoints', 'mask', 'ren_mask']: #'depth_coords2d','lifted_points', \n try:\n ret[key] = np.stack(elems, axis=0)\n except:\n print(key, flush=True)\n raise\n elif key == 'metrics':\n ret[key] = elems\n else:\n ret[key] = []\n for e in elems:\n ret[key].append(e)\n\n return ret\n\n\ndef get_correspondences(src_pcd, tgt_pcd, search_voxel_size, K=None, trans=None):\n if trans is not None:\n src_pcd.transform(trans)\n pcd_tree = o3d.geometry.KDTreeFlann(tgt_pcd)\n\n correspondences = []\n for i, point in enumerate(src_pcd.points):\n [count, idx, _] = pcd_tree.search_radius_vector_3d(\n point, search_voxel_size)\n if K is not None:\n idx = idx[:K]\n for j in idx:\n correspondences.append([i, j])\n\n correspondences = np.array(correspondences)\n # correspondences = torch.from_numpy(correspondences)\n return correspondences\n\n\ndef to_pcd(xyz):\n pcd = o3d.geometry.PointCloud()\n pcd.points = o3d.utility.Vector3dVector(xyz)\n return pcd\n\n\ndef to_tsfm(rot, trans):\n tsfm = np.eye(4)\n tsfm[:3, :3] = rot\n tsfm[:3, 3] = trans.flatten()\n return tsfm\n\n\ndef CameraIntrinsicUpdate(old_K, aug_param):\n '''\n old_K: array of shape (N,3,3), the old camera intrinsic parameters\n aug_pram: dict, the data augmentation parameters\n '''\n aug_type = aug_param['aug_type']\n assert aug_type in ['crop', 'scale', 'flip']\n\n new_K = np.copy(old_K)\n if aug_type == 'crop':\n cx, cy = aug_param['crop/left_top_corner'] # x,y\n new_K[..., 0, 2] = new_K[..., 0, 2] - cx\n new_K[..., 1, 2] = new_K[..., 1, 2] - cy\n elif aug_type == 'scale':\n s_x, s_y = aug_param['scale/scale']\n new_K[..., 0, 0] = s_x * new_K[..., 0, 0]\n new_K[..., 1, 1] = s_y * new_K[..., 1, 1]\n\n new_K[..., 0, 2] = s_x * new_K[..., 0, 2]\n new_K[..., 1, 2] = s_y * new_K[..., 1, 2]\n elif aug_type == 'flip':\n w = aug_param['flip/width']\n # h = aug_param['flip/heigh']\n new_K[..., 0, 2] = w - new_K[..., 0, 2] # px' = w-px\n # new_K[...,1,2] = h- new_K[...,1,2]\n new_K[..., 0, 0] = - new_K[..., 0, 0] # fx' = -fx\n\n return new_K\n\n\ndef crop_transform(images, depths, Ks, crop_param, ):\n assert(len(images) == len(depths) == len(Ks))\n\n crop_type = crop_param[\"crop_type\"]\n assert(crop_type in [\"fixed\", \"center\", \"random\"])\n\n crop_size = crop_param[\"crop_size\"]\n iheight, iwidth = images[0].shape[:2]\n\n if crop_type == \"fixed\":\n lt_corner = crop_param[\"lt_corner\"]\n op = transforms.Crop(\n lt_corner[0], lt_corner[1], crop_size[0], crop_size[1])\n elif crop_type == \"center\":\n op = transforms.CenterCrop(crop_size)\n\n ci, cj, _, _ = op.get_params(images[0], crop_size)\n lt_corner = ci, cj\n\n elif crop_type == \"random\":\n op = transforms.RandomCrop((iheight, iwidth), crop_size)\n\n lt_corner = op.i, op.j\n\n for i, _ in enumerate(images):\n images[i] = op(images[i])\n depths[i] = op(depths[i])\n\n Ks[i] = CameraIntrinsicUpdate(Ks[i],\n {\"aug_type\": \"crop\", \"crop/left_top_corner\": (lt_corner[1], lt_corner[0])})\n\n return images, depths, Ks\n\n\n# def patch_crop(image, depth, mask, K_old, margin_ratio=0.2, output_size=128, offset_ratio=(0,0),bbox=None, mask_depth=True):\ndef patch_crop(image, depth, mask, K_old, margin_ratio=0.2, output_size=128, offset_ratio=(0,0),bbox=None, mask_depth=False):\n '''\n image: HxWx3\n mask: HxW\n K_old: 3x3\n offset: (offset_h, offset_w)\n '''\n\n H, W, _ = image.shape\n \n mask = mask.astype('uint8')*255\n if bbox is None:\n _x, _y, _w, _h = cv2.boundingRect(mask)\n else:\n _x, _y, _w, _h = bbox[1], bbox[0], bbox[3]-bbox[1], bbox[2]-bbox[0]\n\n # center = [_x+_w/2, _y+_h/2]\n center = [_x+_w/2+offset_ratio[1]*_w, _y+_h/2+offset_ratio[0]*_h ]\n\n L = int(max(_w, _h) * (1+2*margin_ratio))\n \n if L<0:\n #TODO\n print(mask.sum(), depth.sum(), '!!!', flush=True)\n L=128\n\n x = max(0, int(center[0] - L/2))\n y = max(0, int(center[1] - L/2))\n\n crop = image[y:y+L, x:x+L]\n # only keep the ROI depth\n\n if mask_depth:\n depth[mask < 1] = 0 # removed by dy at 0810\n depth_crop = depth[y:y+L, x:x+L]\n mask_crop = mask[y:y+L, x:x+L]\n\n \n\n # w=h=int ((1+2*margin_ratio)*L) # actual crop size\n w = h = L # actual crop size\n # automatically handle the \"out of range\" problem\n patch = np.zeros([h, w, 3], dtype=image.dtype)\n # depth_patch = np.ones([h, w], dtype=depth.dtype)\n depth_patch = np.zeros([h, w], dtype=depth.dtype)\n mask_patch = np.zeros([h, w], dtype=depth.dtype)\n\n try:\n xp = 0\n yp = 0\n patch[xp: xp+crop.shape[0], yp:yp+crop.shape[1]] = crop\n depth_patch[xp: xp+crop.shape[0], yp:yp+crop.shape[1]] = depth_crop\n mask_patch[xp: xp+crop.shape[0], yp:yp+crop.shape[1]] = mask_crop\n except:\n import pdb\n pdb.set_trace()\n patch = cv2.resize(patch, (output_size, output_size),\n interpolation=cv2.INTER_LINEAR)\n depth_patch = cv2.resize(\n depth_patch, (output_size, output_size), interpolation=cv2.INTER_NEAREST)\n mask_patch = cv2.resize(\n mask_patch, (output_size, output_size), interpolation=cv2.INTER_NEAREST)\n\n # update the intrinsic parameters\n K_new = np.zeros_like(K_old)\n scale = output_size/L\n K_new[0, 2] = (K_old[0, 2]-x)*scale\n K_new[1, 2] = (K_old[1, 2]-y)*scale\n K_new[0, 0] = K_old[0, 0]*scale\n K_new[1, 1] = K_old[1, 1]*scale\n K_new[2, 2] = 1\n\n # return patch, depth_patch, K_new\n return patch, depth_patch, mask_patch, K_new\n\n\ndef preprocess_deepim(\n input_dict,\n max_points,\n correspondence_radius,\n normalize_model=True,\n rand_transform_model=False, # False,#True,\n rand_rgb_transformer=None,\n image_scale=None,\n patch_cropper=None, # func patch_crop(...)\n \n):\n output_dict = copy.deepcopy(input_dict)\n\n ################################### process 3D point clouds ###################################\n\n if (output_dict['model_points'].shape[0] > max_points):\n # if(output_dict['model_points'].shape[0] > 20000):\n idx = np.random.permutation(\n output_dict['model_points'].shape[0])[:max_points]\n print(idx, output_dict['model_points'].shape, flush=True)\n output_dict['model_points'] = output_dict['model_points'][idx]\n output_dict['model_point_features'] = output_dict['model_point_features'][idx]\n\n output_dict['original_RT'] = copy.deepcopy(output_dict['RT'])\n if normalize_model:\n points = output_dict['model_points']\n mean = points.mean(axis=0)\n scope = points.max(axis=0)-points.min(axis=0)\n points_normalize = (points-mean)/scope.max()\n # points_normalize.tofile(bin_save_path)\n # modify the extrinsic parameters\n output_dict['RT'][:, 3:] = output_dict['RT'][:, :3] @ mean[:,\n None] + output_dict['RT'][:, 3:] # 3x3 @ 3x1 + 3x1\n # input_dict['RT'][:,:3] *=scope.max()\n output_dict['scale'] = scope.max()\n output_dict['original_model_points'] = output_dict['model_points']\n output_dict['model_points'] = points_normalize\n\n\n if rand_transform_model:\n points = output_dict['model_points']\n rand_quat = np.random.randn(1, 4)\n rand_quat = rand_quat/np.linalg.norm(rand_quat, axis=-1)\n rand_rot = quaternion.as_rotation_matrix(\n quaternion.from_float_array(rand_quat)).squeeze() # 3x3\n output_dict['rand_model_points'] = (\n rand_rot@ points.T).T.astype(np.float32)\n output_dict['rand_RT'] = copy.deepcopy(output_dict['RT'])\n # output_dict['RT'][:,:3]@rand_rot.T\n output_dict['rand_RT'][:, :3] = rand_rot\n output_dict['rand_RT'][:, 3] = 0\n\n ################################### process 2D images ###################################\n # carve out image patches\n if patch_cropper is not None:\n ref_mask = output_dict['depth'] > 0\n output_dict['image'], output_dict['depth'], output_dict['K'] = patch_cropper(\n output_dict['image'], output_dict['depth'], ref_mask, output_dict['K'])\n\n output_dict['rendered_image'], output_dict['rendered_depth'], _ = patch_cropper(\n output_dict['rendered_image'], output_dict['rendered_depth'], ref_mask, output_dict['K'].copy() )\n\n # rescale image\n if image_scale is not None:\n output_dict['image'] = cv2.resize(output_dict['image'],\n (output_dict['image'].shape[1]*image_scale,\n output_dict['image'].shape[0]*image_scale),\n interpolation=cv2.INTER_AREA)\n output_dict['depth'] = cv2.resize(output_dict['depth'],\n (output_dict['depth'].shape[1]*image_scale,\n output_dict['depth'].shape[0]*image_scale),\n interpolation=cv2.INTER_NEAREST)\n output_dict['K'][:2] = output_dict['K'][:2]*image_scale\n\n # lift depth\n depth = output_dict['depth'].squeeze() # H,W\n depth_pts, depth_coords2d = mask_depth_to_point_cloud(\n depth != 0, depth, output_dict['K'])\n depth_pts = (output_dict['RT'][:, :3].T@(depth_pts.T - output_dict['RT']\n [:, 3:])).T / output_dict['scale'] # transformed to the model frame\n\n # find 2d-3d correspondences\n tsfm = np.eye(4)\n tsfm[:3] = output_dict['RT']\n model_pcd = output_dict['model_points']\n\n correspondences_2d3d = get_correspondences(\n to_pcd(depth_pts), to_pcd(model_pcd), correspondence_radius, K=5)\n if len(correspondences_2d3d.shape) < 2 or len(correspondences_2d3d) < 10:\n print(depth_pts.shape, model_pcd.shape)\n print(\"correspondences_2d3d.shape:\",\n correspondences_2d3d.shape, flush=True)\n # raise ValueError(\"Too few correspondences are found!\")\n raise Exception(\"Too few correspondences are found!\")\n\n output_dict['depth_coords2d'] = depth_coords2d\n output_dict['lifted_points'] = depth_pts\n # output_dict['correspondences_2d3d'] = np.zeros(1)#correspondences_2d3d\n output_dict['correspondences_2d3d'] = correspondences_2d3d\n\n if rand_rgb_transformer is not None:\n output_dict['image'], _, _ = rand_rgb_transformer(output_dict['image'])\n # TO TENSOR\n output_dict['image'] = (output_dict['image'].astype(\n np.float32)/255.0).transpose([2, 0, 1]) # .mean(axis=0, keepdims=True) # 1,H,W\n output_dict['depth'] = output_dict['depth'].astype(np.float32)[\n None] # 1,H,W\n\n return output_dict\n\ndef preprocess(\n input_dict,\n max_points,\n correspondence_radius,\n normalize_model=True,\n rand_transform_model=False, \n rand_rgb_transformer=None,\n image_scale=None,\n crop_param=None,\n kp_3d_param=None,\n use_coords_as_3d_feat=False,\n find_2d3d_correspondence=True,\n \n):\n output_dict = copy.deepcopy(input_dict)\n\n ################################### process 3D point clouds ###################################\n if use_coords_as_3d_feat:\n output_dict['model_point_features'] = output_dict['model_points'][:,:3]\n\n if (output_dict['model_points'].shape[0] > max_points):\n # if(output_dict['model_points'].shape[0] > 20000):\n idx = np.random.permutation(\n output_dict['model_points'].shape[0])[:max_points]\n print(idx, output_dict['model_points'].shape, flush=True)\n output_dict['model_points'] = output_dict['model_points'][idx]\n output_dict['model_point_features'] = output_dict['model_point_features'][idx]\n\n output_dict['original_RT'] = copy.deepcopy(output_dict['RT'])\n output_dict['original_model_points'] = output_dict['model_points']\n if normalize_model:\n points = output_dict['model_points']\n mean = points.mean(axis=0)\n scope = points.max(axis=0)-points.min(axis=0)\n points_normalize = (points-mean)/scope.max()\n # modify the extrinsic parameters\n output_dict['RT'][:, 3:] = output_dict['RT'][:, :3] @ mean[:,\n None] + output_dict['RT'][:, 3:] # 3x3 @ 3x1 + 3x1\n output_dict['scale'] = scope.max()\n output_dict['model_points'] = points_normalize\n\n\n if rand_transform_model:\n points = output_dict['model_points']\n rand_quat = np.random.randn(1, 4)\n rand_quat = rand_quat/np.linalg.norm(rand_quat, axis=-1)\n rand_rot = quaternion.as_rotation_matrix(\n quaternion.from_float_array(rand_quat)).squeeze() # 3x3\n output_dict['rand_model_points'] = (\n rand_rot@ points.T).T.astype(np.float32)\n output_dict['rand_RT'] = copy.deepcopy(output_dict['RT'])\n output_dict['rand_RT'][:, :3] = rand_rot\n output_dict['rand_RT'][:, 3] = 0\n\n ################################### process 2D images ###################################\n #crop image\n if crop_param is not None:# and output_dict['mask'].sum()>0:\n #without random cropping\n if not crop_param.rand_crop: \n if crop_param.get(\"crop_with_init_pose\", False):\n # bbox= output_dict.get('bbox', None)\n bbox=None\n output_dict['image'], output_dict['depth'], output_dict['mask'], output_dict['K'] = patch_crop(output_dict['image'], output_dict['depth'], mask=output_dict['ren_mask'],\n K_old=output_dict['K'], margin_ratio=crop_param.margin_ratio, output_size=crop_param.output_size, bbox=bbox\n )\n elif crop_param.get(\"crop_with_rand_bbox_shift\", True): \n bbox= output_dict.get('bbox', None)\n # offset_ratio= [truncated_normal(0,0.5,-1,1)*crop_param.max_rand_offset_ratio, truncated_normal(0,0.5,-1,1)*crop_param.max_rand_offset_ratio] \n offset_ratio= [truncated_normal(0,0.33,-1,1)*1, truncated_normal(0,0.33,-1,1)*1] \n output_dict['image'], output_dict['depth'], output_dict['mask'], output_dict['K'] = patch_crop(output_dict['image'], output_dict['depth'], mask=output_dict['mask'],\n K_old=output_dict['K'], margin_ratio=crop_param.margin_ratio, output_size=crop_param.output_size, offset_ratio=offset_ratio, bbox=output_dict.get('bbox', None) \n )\n else:\n bbox= output_dict.get('bbox', None)\n output_dict['image'], output_dict['depth'], output_dict['mask'], output_dict['K'] = patch_crop(output_dict['image'], output_dict['depth'], mask=output_dict['mask'],\n K_old=output_dict['K'], margin_ratio=crop_param.margin_ratio, output_size=crop_param.output_size, bbox=output_dict.get('bbox', None) \n )\n else:\n margin_ratio= truncated_normal(0.5, 0.5, 0, 1) *crop_param.max_rand_margin_ratio\n offset_ratio= [truncated_normal(0,0.5,-1,1)*crop_param.max_rand_offset_ratio, truncated_normal(0,0.5,-1,1)*crop_param.max_rand_offset_ratio] \n output_dict['image'], output_dict['depth'], output_dict['mask'], output_dict['K'] = patch_crop(output_dict['image'], output_dict['depth'], mask=output_dict['mask'],\n K_old=output_dict['K'], margin_ratio=margin_ratio, output_size=crop_param.output_size, offset_ratio=offset_ratio, bbox=output_dict.get('bbox', None) \n )\n \n # rescale image\n if image_scale is not None:\n output_dict['image'] = cv2.resize(output_dict['image'],\n (output_dict['image'].shape[1]*image_scale,\n output_dict['image'].shape[0]*image_scale),\n interpolation=cv2.INTER_AREA)\n output_dict['depth'] = cv2.resize(output_dict['depth'],\n (output_dict['depth'].shape[1]*image_scale,\n output_dict['depth'].shape[0]*image_scale),\n interpolation=cv2.INTER_NEAREST)\n output_dict['K'][:2] = output_dict['K'][:2]*image_scale\n\n # lift depths\n depth = output_dict['depth'].squeeze() # H,W\n depth_pts, depth_coords2d = mask_depth_to_point_cloud(\n depth != 0, depth, output_dict['K'])\n\n depth_pts = (output_dict['RT'][:, :3].T@(depth_pts.T - output_dict['RT']\n [:, 3:])).T / output_dict['scale'] # transformed to the model frame\n\n # find 2d-3d correspondences\n if find_2d3d_correspondence:\n tsfm = np.eye(4)\n tsfm[:3] = output_dict['RT']\n model_pcd = output_dict['model_points']\n correspondences_2d3d = get_correspondences(\n to_pcd(depth_pts), to_pcd(model_pcd), correspondence_radius, K=5)\n if len(correspondences_2d3d.shape) < 2 or len(correspondences_2d3d) < 10:# or ( \"mask\" in output_dict and output_dict['mask'].sum()<10 ) :\n print(depth_pts.shape, model_pcd.shape)\n print(\"correspondences_2d3d.shape:\",\n correspondences_2d3d.shape, flush=True)\n raise Exception(\"Too few correspondences are found!\")\n\n output_dict['depth_coords2d'] = depth_coords2d\n output_dict['lifted_points'] = depth_pts\n output_dict['correspondences_2d3d'] = correspondences_2d3d\n else:\n output_dict['depth_coords2d'] = depth_coords2d\n output_dict['lifted_points'] = depth_pts\n output_dict['correspondences_2d3d'] = np.zeros([10,2], dtype=np.int64) \n\n\n if rand_rgb_transformer is not None:\n output_dict['image'], _, _ = rand_rgb_transformer(output_dict['image'])\n # TO TENSORs\n output_dict['image'] = (output_dict['image'].astype(\n np.float32)/255.0).transpose([2, 0, 1]) # .mean(axis=0, keepdims=True) # 1,H,W\n output_dict['depth'] = output_dict['depth'].astype(np.float32)[\n None] # 1,H,W\n\n return output_dict\n\ndef batch_grid_subsampling_kpconv(points, batches_len, features=None, labels=None, sampleDl=0.1, max_p=0, verbose=0, random_grid_orient=True):\n \"\"\"\n CPP wrapper for a grid subsampling (method = barycenter for points and features)\n \"\"\"\n if (features is None) and (labels is None):\n s_points, s_len = cpp_subsampling.subsample_batch(points,\n batches_len,\n sampleDl=sampleDl,\n max_p=max_p,\n verbose=verbose)\n return torch.from_numpy(s_points), torch.from_numpy(s_len)\n\n elif (labels is None):\n s_points, s_len, s_features = cpp_subsampling.subsample_batch(points,\n batches_len,\n features=features,\n sampleDl=sampleDl,\n max_p=max_p,\n verbose=verbose)\n return torch.from_numpy(s_points), torch.from_numpy(s_len), torch.from_numpy(s_features)\n\n elif (features is None):\n s_points, s_len, s_labels = cpp_subsampling.subsample_batch(points,\n batches_len,\n classes=labels,\n sampleDl=sampleDl,\n max_p=max_p,\n verbose=verbose)\n return torch.from_numpy(s_points), torch.from_numpy(s_len), torch.from_numpy(s_labels)\n\n else:\n s_points, s_len, s_features, s_labels = cpp_subsampling.subsample_batch(points,\n batches_len,\n features=features,\n classes=labels,\n sampleDl=sampleDl,\n max_p=max_p,\n verbose=verbose)\n return torch.from_numpy(s_points), torch.from_numpy(s_len), torch.from_numpy(s_features), torch.from_numpy(s_labels)\n\n\ndef batch_neighbors_kpconv(queries, supports, q_batches, s_batches, radius, max_neighbors):\n \"\"\"\n Computes neighbors for a batch of queries and supports, apply radius search\n :param queries: (N1, 3) the query points\n :param supports: (N2, 3) the support points\n :param q_batches: (B) the list of lengths of batch elements in queries\n :param s_batches: (B)the list of lengths of batch elements in supports\n :param radius: float32\n :return: neighbors indices\n \"\"\"\n\n neighbors = cpp_neighbors.batch_query(\n queries, supports, q_batches, s_batches, radius=radius)\n # print(\"neighbors.shape\" , neighbors.shape, queries.shape,flush=True)\n if max_neighbors > 0:\n return torch.from_numpy(neighbors[:, :max_neighbors])\n else:\n return torch.from_numpy(neighbors)\n\n\ndef collate_fn_descriptor(list_data, config, neighborhood_limits):\n ret = merge_batch(list_data)\n\n batched_points = torch.from_numpy(ret['model_points'])\n batched_lengths = torch.from_numpy(ret['batched_model_point_lengths'])\n batched_features = torch.from_numpy(ret['model_point_features'])\n\n if ret.get('rand_model_points', None) is not None:\n batched_rand_points = torch.from_numpy(ret['rand_model_points'])\n batched_rand_lengths = torch.from_numpy(\n ret['batched_rand_model_point_lengths'])\n\n batched_points = torch.cat(\n [batched_points, batched_rand_points], dim=0)\n batched_lengths = torch.cat(\n [batched_lengths, batched_rand_lengths], dim=0)\n batched_features = torch.cat(\n [batched_features, batched_features], dim=0)\n\n # Starting radius of convolutions\n r_normal = config.first_subsampling_dl * config.conv_radius\n # Starting layer\n layer_blocks = []\n layer = 0\n\n # Lists of inputs\n input_points = []\n input_neighbors = []\n input_pools = []\n input_upsamples = []\n input_batches_len = []\n timer = Timer()\n for block_i, block in enumerate(config.architecture):\n # Stop when meeting a global pooling or upsampling\n if 'global' in block or 'upsample' in block:\n break\n\n # Get all blocks of the layer\n if not ('pool' in block or 'strided' in block):\n layer_blocks += [block]\n if block_i < len(config.architecture) - 1 and not ('upsample' in config.architecture[block_i + 1]):\n continue\n\n # Convolution neighbors indices\n # *****************************\n\n if layer_blocks:\n # Convolutions are done in this layer, compute the neighbors with the good radius\n if np.any(['deformable' in blck for blck in layer_blocks[:-1]]):\n r = r_normal * config.deform_radius / config.conv_radius\n else:\n r = r_normal\n conv_i = batch_neighbors_kpconv(\n batched_points, batched_points, batched_lengths, batched_lengths, r, neighborhood_limits[layer])\n\n else:\n # This layer only perform pooling, no neighbors required\n conv_i = torch.zeros((0, 1), dtype=torch.int64)\n\n # Pooling neighbors indices\n # *************************\n\n if 'pool' in block or 'strided' in block:\n\n # New subsampling length\n dl = 2 * r_normal / config.conv_radius\n\n # Subsampled points\n pool_p, pool_b = batch_grid_subsampling_kpconv(\n batched_points, batched_lengths, sampleDl=dl)\n\n # Radius of pooled neighbors\n if 'deformable' in block:\n r = r_normal * config.deform_radius / config.conv_radius\n else:\n r = r_normal\n\n # Subsample indices\n pool_i = batch_neighbors_kpconv(\n pool_p, batched_points, pool_b, batched_lengths, r, neighborhood_limits[layer])\n\n # Upsample indices (with the radius of the next layer to keep wanted density)\n up_i = batch_neighbors_kpconv(\n batched_points, pool_p, batched_lengths, pool_b, 2 * r, neighborhood_limits[layer])\n\n else:\n # No pooling in the end of this layer, no pooling indices required\n pool_i = torch.zeros((0, 1), dtype=torch.int64)\n pool_p = torch.zeros((0, 3), dtype=torch.float32)\n pool_b = torch.zeros((0,), dtype=torch.int64)\n up_i = torch.zeros((0, 1), dtype=torch.int64)\n\n # Updating input lists\n input_points += [batched_points.float()]\n input_neighbors += [conv_i.long()]\n input_pools += [pool_i.long()]\n input_upsamples += [up_i.long()]\n input_batches_len += [batched_lengths]\n\n # New points for next layer\n batched_points = pool_p\n batched_lengths = pool_b\n\n # Update radius and reset blocks\n r_normal *= 2\n layer += 1\n layer_blocks = []\n\n ###############\n # Return inputs\n ###############\n dict_inputs = {\n \"idx\": ret[\"idx\"],\n 'model_points': input_points,\n 'visibility': torch.from_numpy(ret['visibility']),\n 'neighbors': input_neighbors,\n 'pools': input_pools,\n 'upsamples': input_upsamples,\n 'model_point_features': batched_features.float(),\n 'stack_lengths': input_batches_len,\n 'image': torch.from_numpy(ret['image']),\n 'depth': torch.from_numpy(ret['depth']),\n 'mask': torch.from_numpy(ret['mask']),\n 'ren_mask': torch.from_numpy(ret['ren_mask']),\n 'K': torch.from_numpy(ret['K']),\n 'RT': torch.from_numpy(ret['RT']),\n 'original_RT': torch.from_numpy(ret['original_RT']),\n 'POSECNN_RT': torch.from_numpy(ret.get('POSECNN_RT', np.zeros_like(ret['RT']) ) ),\n 'rand_RT': torch.from_numpy(ret.get('rand_RT', np.zeros_like(ret['RT']))),\n # \"lifted_points\": torch.from_numpy(ret['lifted_points']),\n \"lifted_points\": [torch.from_numpy(d) for d in ret['lifted_points'] ] ,\n # 'depth_coords2d': torch.from_numpy(ret['depth_coords2d']),\n 'depth_coords2d': [torch.from_numpy(d) for d in ret['depth_coords2d']],\n \"correspondences_2d3d\": torch.from_numpy(ret['correspondences_2d3d']),\n \"original_model_points\": torch.from_numpy(ret['original_model_points']),\n \"class_name\": ret['class_name'],\n \"3d_keypoint_inds\": torch.from_numpy(ret['3d_keypoint_inds']),\n \"3d_keypoints\": torch.from_numpy(ret['3d_keypoints'] ) \n }\n\n return dict_inputs\n\n\ndef collate_fn_descriptor_deepim(list_data, config, neighborhood_limits):\n ret = merge_batch(list_data)\n\n batched_points = torch.from_numpy(ret['model_points'])\n batched_lengths = torch.from_numpy(ret['batched_model_point_lengths'])\n batched_features = torch.from_numpy(ret['model_point_features'])\n \n\n if ret.get('rand_model_points', None) is not None:\n # torch.from_numpy(np.concatenate(batched_points_list, axis=0))\n batched_rand_points = torch.from_numpy(ret['rand_model_points'])\n # torch.from_numpy(np.concatenate(batched_points_list, axis=0))\n batched_rand_lengths = torch.from_numpy(\n ret['batched_rand_model_point_lengths'])\n\n batched_points = torch.cat(\n [batched_points, batched_rand_points], dim=0)\n batched_lengths = torch.cat(\n [batched_lengths, batched_rand_lengths], dim=0)\n batched_features = torch.cat(\n [batched_features, batched_features], dim=0)\n\n # Starting radius of convolutions\n r_normal = config.first_subsampling_dl * config.conv_radius\n # Starting layer\n layer_blocks = []\n layer = 0\n\n # Lists of inputs\n input_points = []\n input_neighbors = []\n input_pools = []\n input_upsamples = []\n input_batches_len = []\n timer = Timer()\n for block_i, block in enumerate(config.architecture):\n # timer.tic()\n\n # Stop when meeting a global pooling or upsampling\n if 'global' in block or 'upsample' in block:\n break\n\n # Get all blocks of the layer\n if not ('pool' in block or 'strided' in block):\n layer_blocks += [block]\n if block_i < len(config.architecture) - 1 and not ('upsample' in config.architecture[block_i + 1]):\n continue\n\n # Convolution neighbors indices\n # *****************************\n\n if layer_blocks:\n # Convolutions are done in this layer, compute the neighbors with the good radius\n if np.any(['deformable' in blck for blck in layer_blocks[:-1]]):\n r = r_normal * config.deform_radius / config.conv_radius\n else:\n r = r_normal\n conv_i = batch_neighbors_kpconv(\n batched_points, batched_points, batched_lengths, batched_lengths, r, neighborhood_limits[layer])\n\n else:\n # This layer only perform pooling, no neighbors required\n conv_i = torch.zeros((0, 1), dtype=torch.int64)\n\n # Pooling neighbors indices\n # *************************\n\n # If end of layer is a pooling operation\n if 'pool' in block or 'strided' in block:\n\n # New subsampling length\n dl = 2 * r_normal / config.conv_radius\n\n # Subsampled points\n pool_p, pool_b = batch_grid_subsampling_kpconv(\n batched_points, batched_lengths, sampleDl=dl)\n\n # Radius of pooled neighbors\n if 'deformable' in block:\n r = r_normal * config.deform_radius / config.conv_radius\n else:\n r = r_normal\n\n # Subsample indices\n pool_i = batch_neighbors_kpconv(\n pool_p, batched_points, pool_b, batched_lengths, r, neighborhood_limits[layer])\n\n # Upsample indices (with the radius of the next layer to keep wanted density)\n up_i = batch_neighbors_kpconv(\n batched_points, pool_p, batched_lengths, pool_b, 2 * r, neighborhood_limits[layer])\n\n else:\n # No pooling in the end of this layer, no pooling indices required\n pool_i = torch.zeros((0, 1), dtype=torch.int64)\n pool_p = torch.zeros((0, 3), dtype=torch.float32)\n pool_b = torch.zeros((0,), dtype=torch.int64)\n up_i = torch.zeros((0, 1), dtype=torch.int64)\n\n # Updating input lists\n input_points += [batched_points.float()]\n input_neighbors += [conv_i.long()]\n input_pools += [pool_i.long()]\n input_upsamples += [up_i.long()]\n input_batches_len += [batched_lengths]\n\n # New points for next layer\n batched_points = pool_p\n batched_lengths = pool_b\n\n # Update radius and reset blocks\n r_normal *= 2\n layer += 1\n layer_blocks = []\n\n # timer.toc()\n ###############\n # Return inputs\n ###############\n dict_inputs = {\n \"idx\": ret[\"idx\"],\n 'model_points': input_points,\n 'visibility': torch.from_numpy(ret['visibility']),\n 'neighbors': input_neighbors,\n 'pools': input_pools,\n 'upsamples': input_upsamples,\n 'model_point_features': batched_features.float(),\n 'stack_lengths': input_batches_len,\n 'image': torch.from_numpy(ret['image']),\n 'depth': torch.from_numpy(ret['depth']),\n \"ren_mask\": torch.from_numpy(ret['ren_mask']),\n 'K': torch.from_numpy(ret['K']),\n 'RT': torch.from_numpy(ret['RT']),\n 'original_RT': torch.from_numpy(ret['original_RT']),\n 'rendered_RT': torch.from_numpy(ret['rendered_RT']) if ret.get('rendered_RT', None) is not None else None ,\n 'POSECNN_RT': torch.from_numpy(ret.get('POSECNN_RT', np.zeros_like(ret['RT']) ) ),\n # TODO\n 'rand_RT': torch.from_numpy(ret.get('rand_RT', np.zeros_like(ret['RT']))),\n # \"lifted_points\": torch.from_numpy(ret['lifted_points']),\n \"lifted_points\": [torch.from_numpy(d) for d in ret['lifted_points'] ] ,\n # 'depth_coords2d': torch.from_numpy(ret['depth_coords2d']),\n 'depth_coords2d': [torch.from_numpy(d) for d in ret['depth_coords2d']],\n \"correspondences_2d3d\": torch.from_numpy(ret['correspondences_2d3d']),\n \"original_model_points\": torch.from_numpy(ret['original_model_points']),\n \"class_name\": ret['class_name'],\n }\n\n return dict_inputs\n\n\ndef calibrate_neighbors(dataset, config, collate_fn, keep_ratio=0.8, samples_threshold=2000):\n timer = Timer()\n last_display = timer.total_time\n\n # From config parameter, compute higher bound of neighbors number in a neighborhood\n hist_n = int(np.ceil(4 / 3 * np.pi * (config.deform_radius + 1) ** 3))\n neighb_hists = np.zeros((config.num_layers, hist_n), dtype=np.int32)\n\n # Get histogram of neighborhood sizes i in 1 epoch max.\n for i in range(len(dataset)):\n timer.tic()\n\n batched_input = collate_fn(\n [dataset[i]], config, neighborhood_limits=[hist_n] * 5)\n # update histogram\n counts = [torch.sum(neighb_mat < neighb_mat.shape[0], dim=1).numpy()\n for neighb_mat in batched_input['neighbors']]\n \n hists = [np.bincount(c, minlength=hist_n)[:hist_n] for c in counts]\n neighb_hists += np.vstack(hists)\n timer.toc()\n\n if timer.total_time - last_display > 0.1:\n last_display = timer.total_time\n print(f\"Calib Neighbors {i:08d}: timings {timer.total_time:4.2f}s\")\n\n if np.min(np.sum(neighb_hists, axis=1)) > samples_threshold:\n break\n\n cumsum = np.cumsum(neighb_hists.T, axis=0)\n percentiles = np.sum(cumsum < (keep_ratio * cumsum[hist_n - 1, :]), axis=0)\n\n neighborhood_limits = percentiles\n print('\\n')\n\n return neighborhood_limits\n\n\ndef get_dataloader(dataset, kpconv_config, batch_size=1, num_workers=4, shuffle=True, sampler=None, neighborhood_limits=None):\n if neighborhood_limits is None:\n # neighborhood_limits = calibrate_neighbors(dataset, dataset.config, collate_fn=collate_fn_descriptor)\n neighborhood_limits = calibrate_neighbors(\n dataset, kpconv_config, collate_fn=collate_fn_descriptor)\n print(\"neighborhood:\", neighborhood_limits)\n dataloader = torch.utils.data.DataLoader(\n dataset,\n batch_size=batch_size,\n shuffle=shuffle,\n num_workers=num_workers,\n # https://discuss.pytorch.org/t/supplying-arguments-to-collate-fn/25754/4\n collate_fn=partial(collate_fn_descriptor, config=kpconv_config,\n neighborhood_limits=neighborhood_limits),\n sampler=sampler,\n drop_last=False\n )\n return dataloader, neighborhood_limits\n\ndef get_dataloader_deepim(dataset, kpconv_config, batch_size=1, num_workers=4, shuffle=True, sampler=None, neighborhood_limits=None):\n if neighborhood_limits is None:\n neighborhood_limits = calibrate_neighbors(\n dataset, kpconv_config, collate_fn=collate_fn_descriptor_deepim)\n print(\"Neighborhood:\", neighborhood_limits)\n dataloader = torch.utils.data.DataLoader(\n dataset,\n batch_size=batch_size,\n shuffle=shuffle,\n num_workers=num_workers,\n # https://discuss.pytorch.org/t/supplying-arguments-to-collate-fn/25754/4\n collate_fn=partial(collate_fn_descriptor_deepim, config=kpconv_config,\n neighborhood_limits=neighborhood_limits),\n sampler=sampler,\n drop_last=False\n )\n return dataloader, neighborhood_limits\n\nif __name__ == '__main__':\n pass\n"
},
{
"path": "data/transforms.py",
"content": "import numpy as np\nimport random\nimport torch\nimport torchvision\nfrom torchvision.transforms import functional as F\nimport cv2\nfrom PIL import Image\n\n\nclass Compose(object):\n\n def __init__(self, transforms):\n self.transforms = transforms\n\n def __call__(self, img, kpts=None, mask=None):\n for t in self.transforms:\n img, kpts, mask = t(img, kpts, mask)\n return img, kpts, mask\n\n def __repr__(self):\n format_string = self.__class__.__name__ + \"(\"\n for t in self.transforms:\n format_string += \"\\n\"\n format_string += \" {0}\".format(t)\n format_string += \"\\n)\"\n return format_string\n\n\nclass ToTensor(object):\n\n def __call__(self, img, kpts, mask):\n return np.asarray(img).astype(np.float32) / 255., kpts, mask\n\n\nclass Normalize(object):\n\n def __init__(self, mean, std, to_bgr=True):\n self.mean = mean\n self.std = std\n self.to_bgr = to_bgr\n\n def __call__(self, img, kpts, mask):\n img -= self.mean\n img /= self.std\n if self.to_bgr:\n img = img.transpose(2, 0, 1).astype(np.float32)\n return img, kpts, mask\n\n\nclass ColorJitter(object):\n\n def __init__(self,\n brightness=None,\n contrast=None,\n saturation=None,\n hue=None,\n ):\n self.color_jitter = torchvision.transforms.ColorJitter(\n brightness=brightness,\n contrast=contrast,\n saturation=saturation,\n hue=hue,)\n\n def __call__(self, image, kpts, mask):\n image = np.asarray(self.color_jitter(Image.fromarray(np.ascontiguousarray(image, np.uint8))))\n return image, kpts, mask\n\n\nclass RandomBlur(object):\n\n def __init__(self, prob=0.5):\n self.prob = prob\n\n def __call__(self, image, kpts, mask):\n if random.random() < self.prob:\n sigma = np.random.choice([3, 5, 7, 9])\n image = cv2.GaussianBlur(image, (sigma, sigma), 0)\n return image, kpts, mask\n\n\ndef make_transforms(cfg, is_train):\n if is_train is True:\n transform = Compose(\n [\n RandomBlur(0.5),\n ColorJitter(0.1, 0.1, 0.05, 0.05),\n # ToTensor(),\n # Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n ]\n )\n else:\n transform = Compose(\n [\n # ToTensor(),\n # Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n ]\n )\n\n return transform\n"
},
{
"path": "data/ycb/basic.py",
"content": "import mmcv \nbop_ycb_idx2class={\n 1: '002_master_chef_can', \n 2: '003_cracker_box',\n 3: '004_sugar_box', \n 4: '005_tomato_soup_can',\n 5: '006_mustard_bottle',\n 6: '007_tuna_fish_can',\n 7: '008_pudding_box', \n 8: '009_gelatin_box',\n 9: '010_potted_meat_can', \n 10: '011_banana', \n 11: '019_pitcher_base', \n 12: '021_bleach_cleanser', \n 13: '024_bowl', \n 14: '025_mug', \n 15: '035_power_drill',\n 16: '036_wood_block',\n 17: '037_scissors', \n 18: '040_large_marker',\n 19: '051_large_clamp',\n 20: '052_extra_large_clamp',\n 21: '061_foam_brick', \n }\nbop_ycb_class2idx = dict([[bop_ycb_idx2class[k],k ] for k in bop_ycb_idx2class.keys() ])\n\n\n"
},
{
"path": "doc/prepare_data.md",
"content": "# Data Preparation Tips\nAll the related data for data preparation can be downloaded [here](https://mycuhk-my.sharepoint.com/:f:/g/personal/1155139432_link_cuhk_edu_hk/EoXnZ96Tuy9PpYlZCvDN8vUBPdP1lP-PWQWiZH2KtIQoaQ?e=lpE472). You could download them first and then follow the instructions below for data preparation. \n\n\n\n## Download Datasets \nFirst, the following dataset need to be downloaded and extracted to the folder *EXPDATA/* \n\n[LINEMOD](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155139432_link_cuhk_edu_hk/EYFaYrk0kcdBgC6WMtLJqP0B9Ar0_Nff9qhI2Cs95qDbdA?e=yYxexC)\n\n[LINEMOD_OCC_TEST](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155139432_link_cuhk_edu_hk/EUKcRnwyy9RGu2ASwA3QDXsBnMRrFP-U4X4Eqq-g_MhmIQ?e=hv6H2s)\n\n## Synthetic Data Generation\n\nThe preprocessed data following [DeepIM](https://github.com/liyi14/mx-DeepIM) and [PVNet](https://github.com/zju3dv/pvnet-rendering) can be downloaded from [LM6d_converted](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155139432_link_cuhk_edu_hk/EYFaYrk0kcdBgC6WMtLJqP0B9Ar0_Nff9qhI2Cs95qDbdA) and [raw_data](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155139432_link_cuhk_edu_hk/ESSFXi_7qs1AgNmty7_9y4AB8ffFsGJWOC3ikgD5BIeXHQ?e=qOmvds). \nAfter downloading, you should put the downloaded files into the folder *EXPDATA/* (lying in the repository's root directory). \nTo create occluded objects during training, we follow [PVNet](https://github.com/zju3dv/pvnet-rendering) to randomly create occlusions. \nYou could run the following scripts to transform the data format for our dataloader. \n```\n bash scripts/run_dataformatter.sh\n```\nThe command above will automatically save the formatted data into *EXPDATA/*. \n\n## Download the Object CAD Models\nYou also need to download the [object models](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155139432_link_cuhk_edu_hk/EQScZuLrkPNPmN4eO3kePaUBjOe92EvbKb7kGJk2vKz-bA?e=8McAdh) and put the extracted folder *models* into *./EXPDATA/LM6d_converted/. \n\n## Download Background Images\n[Pascal VOC 2012](http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar) need to be downloaded to folder *EXPDATA/*. These images will be necessary for the random background generation for training. \n\n## Download Initial Poses \nThe initial poses estimated by PoseCNN and PVNet can be downloaded from [here](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155139432_link_cuhk_edu_hk/EQh5y0M_zHVMnbVszjEviCUBNAX_22MFN26Msa48XlJ5MQ?e=rfhT7k). \nThe initial pose folder also should be put into the folder *EXPDATA/*\n\n## Generate the Information Files\nRun the following script to generate the info files, which is put into the folder *EXPDATA/data_info/*\n\n```\nbash scripts/run_datainfo_generation.sh\n```\n\n\nAfter the the data preparation, the expected directory structure should be \n\n\n```\n./EXPDATA\n |──LM6d_converted \n | |──LM6d_refine \n | |──LM6d_refine_syn\n | └──models\n |──LINEMOD\n | └──fuse_formatted\n |──lmo\n |──VOCdevkit\n |──raw_data\n |──init_poses\n └──data_info\n```\n\n"
},
{
"path": "docker/Dockerfile",
"content": "FROM nvidia/cuda:10.2-cudnn7-devel-ubuntu18.04\n\nRUN apt-key del 7fa2af80\nRUN apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub\nRUN rm /etc/apt/sources.list.d/cuda.list\nRUN rm /etc/apt/sources.list.d/nvidia-ml.list\n\n# Dependencies for glvnd and X11.\nRUN apt-get update \nRUN apt-get install -y -qq --no-install-recommends \\\n libglvnd0 \\\n libgl1 \\\n libglx0 \\\n libegl1 \\\n libxext6 \\\n libx11-6 \\\n && rm -rf /var/lib/apt/lists/*\n# Env vars for the nvidia-container-runtime.\nENV NVIDIA_VISIBLE_DEVICES all\nENV NVIDIA_DRIVER_CAPABILITIES graphics,utility,compute\n\n#env vars for cuda\nENV CUDA_HOME /usr/local/cuda\n\n#install miniconda\nRUN apt-get update --fix-missing && \\\n apt-get install -y wget bzip2 ca-certificates curl git && \\\n apt-get clean && \\\n rm -rf /var/lib/apt/lists/*\n\nRUN wget --quiet https://mirrors.tuna.tsinghua.edu.cn/anaconda/miniconda/Miniconda3-py37_4.9.2-Linux-x86_64.sh -O ~/miniconda.sh && \\\n /bin/bash ~/miniconda.sh -b -p /opt/miniconda3 && \\\n rm ~/miniconda.sh && \\\n /opt/miniconda3/bin/conda clean -tipsy && \\\n ln -s /opt/miniconda3/etc/profile.d/conda.sh /etc/profile.d/conda.sh && \\\n echo \". /opt/miniconda3/etc/profile.d/conda.sh\" >> ~/.bashrc && \\\n echo \"conda activate base\" >> ~/.bashrc && \\\n echo \"conda deactivate && conda activate py37\" >> ~/.bashrc\n\n#https://blog.csdn.net/Mao_Jonah/article/details/89502380\nCOPY freeze.yml freeze.yml\nRUN /opt/miniconda3/bin/conda env create -n py37 -f freeze.yml\n\n# WORKDIR /tmp/\n# COPY config.jupyter.tar config.jupyter.tar\n# RUN tar -xvf config.jupyter.tar -C /root/\n\n#install apex\nENV TORCH_CUDA_ARCH_LIST \"6.0 6.2 7.0 7.2\"\n# make sure we don't overwrite some existing directory called \"apex\"\nWORKDIR /tmp/unique_for_apex\n# uninstall Apex if present, twice to make absolutely sure :)\nRUN /opt/miniconda3/envs/py37/bin/pip3 uninstall -y apex || :\nRUN /opt/miniconda3/envs/py37/bin/pip3 uninstall -y apex || :\n# SHA is something the user can touch to force recreation of this Docker layer,\n# and therefore force cloning of the latest version of Apex\nRUN SHA=ToUcHMe git clone https://github.com/NVIDIA/apex.git\nWORKDIR /tmp/unique_for_apex/apex\nRUN /opt/miniconda3/envs/py37/bin/pip3 install -v --no-cache-dir --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext\" .\n#install pytorch3d \n# RUN /opt/miniconda3/envs/py37/bin/pip install pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py37_cu102_pyt171/download.html\n# RUN /opt/miniconda3/envs/py37/bin/pip install \"git+https://github.com/facebookresearch/pytorch3d.git\"\n# RUN /opt/miniconda3/bin/conda install pytorch3d==0.5.0 -c pytorch3d -n py37\n\n\n\n#other pkgs\nRUN apt-get update \\\n && apt-get install -y -qq --no-install-recommends \\\n cmake build-essential vim xvfb unzip tmux psmisc \\\n libx11-dev libassimp-dev \\\n mesa-common-dev freeglut3-dev \\\n rsync \\\n && apt-get clean \\\n && rm -rf /var/lib/apt/lists/*\n\n#create some directories\nRUN mkdir -p /home/RNNPose\n\nEXPOSE 8887 8888 8889 10000 10001 10002 \nWORKDIR /home/RNNPose\n\n"
},
{
"path": "docker/freeze.yml",
"content": "name: py37_tmp\nchannels:\n - pytorch\n - pytorch3d\n - open3d-admin\n - bottler\n - iopath\n - fvcore\n - conda-forge\n - defaults\ndependencies:\n - pytorch3d=0.5.0\n - _libgcc_mutex=0.1=main\n - _openmp_mutex=4.5=1_gnu\n - anyio=2.2.0=py37h06a4308_1\n - argon2-cffi=20.1.0=py37h27cfd23_1\n - async_generator=1.10=py37h28b3542_0\n - attrs=21.2.0=pyhd3eb1b0_0\n - babel=2.9.1=pyhd3eb1b0_0\n - backcall=0.2.0=pyhd3eb1b0_0\n - blas=1.0=mkl\n - bleach=3.3.0=pyhd3eb1b0_0\n - brotlipy=0.7.0=py37h27cfd23_1003\n - ca-certificates=2021.5.30=ha878542_0\n - certifi=2021.5.30=py37h89c1867_0\n - cffi=1.14.5=py37h261ae71_0\n - charset-normalizer=2.0.4=pyhd3eb1b0_0\n - cryptography=3.4.7=py37hd23ed53_0\n - cudatoolkit=10.2.89=h8f6ccaa_8\n - cycler=0.10.0=py37_0\n - dbus=1.13.18=hb2f20db_0\n - defusedxml=0.7.1=pyhd3eb1b0_0\n - entrypoints=0.3=py37_0\n - expat=2.4.1=h2531618_2\n - fontconfig=2.13.1=h6c09931_0\n - freetype=2.10.4=h5ab3b9f_0\n - fvcore=0.1.5.post20210825=py37\n - glib=2.68.2=h36276a3_0\n - gst-plugins-base=1.14.0=h8213a91_2\n - gstreamer=1.14.0=h28cd5cc_2\n - icu=58.2=he6710b0_3\n - idna=3.2=pyhd3eb1b0_0\n - importlib-metadata=3.10.0=py37h06a4308_0\n - importlib_metadata=3.10.0=hd3eb1b0_0\n - intel-openmp=2021.2.0=h06a4308_610\n - iopath=0.1.9=py37\n - ipykernel=5.3.4=py37h5ca1d4c_0\n - ipython=7.22.0=py37hb070fc8_0\n - ipython_genutils=0.2.0=pyhd3eb1b0_1\n - ipywidgets=7.6.3=pyhd3eb1b0_1\n - jedi=0.17.0=py37_0\n - jinja2=3.0.0=pyhd3eb1b0_0\n - joblib=1.0.1=pyhd3eb1b0_0\n - jpeg=9b=h024ee3a_2\n - json5=0.9.6=pyhd3eb1b0_0\n - jsonschema=3.2.0=py_2\n - jupyter=1.0.0=py37h89c1867_6\n - jupyter_client=6.1.12=pyhd3eb1b0_0\n - jupyter_console=6.4.0=pyhd8ed1ab_0\n - jupyter_core=4.7.1=py37h06a4308_0\n - jupyter_server=1.4.1=py37h06a4308_0\n - jupyterlab=3.0.16=pyhd8ed1ab_0\n - jupyterlab_pygments=0.1.2=py_0\n - jupyterlab_server=2.7.1=pyhd3eb1b0_0\n - jupyterlab_widgets=1.0.0=pyhd3eb1b0_1\n - kiwisolver=1.3.1=py37h2531618_0\n - kornia=0.5.3=pyhd8ed1ab_0\n - lcms2=2.12=h3be6417_0\n - ld_impl_linux-64=2.35.1=h7274673_9\n - libffi=3.3=he6710b0_2\n - libgcc-ng=9.3.0=h5101ec6_17\n - libgfortran-ng=7.5.0=ha8ba4b0_17\n - libgfortran4=7.5.0=ha8ba4b0_17\n - libgomp=9.3.0=h5101ec6_17\n - libpng=1.6.37=hbc83047_0\n - libsodium=1.0.18=h7b6447c_0\n - libstdcxx-ng=9.3.0=hd4cf53a_17\n - libtiff=4.2.0=h85742a9_0\n - libuuid=1.0.3=h1bed415_2\n - libuv=1.40.0=h7b6447c_0\n - libwebp-base=1.2.0=h27cfd23_0\n - libxcb=1.14=h7b6447c_0\n - libxml2=2.9.10=hb55368b_3\n - lz4-c=1.9.3=h2531618_0\n - markupsafe=2.0.1=py37h27cfd23_0\n - matplotlib=3.3.4=py37h06a4308_0\n - matplotlib-base=3.3.4=py37h62a2d02_0\n - mistune=0.8.4=py37h14c3975_1001\n - mkl=2021.2.0=h06a4308_296\n - mkl-service=2.3.0=py37h27cfd23_1\n - mkl_fft=1.3.0=py37h42c9631_2\n - mkl_random=1.2.1=py37ha9443f7_2\n - nbclassic=0.2.6=pyhd3eb1b0_0\n - nbclient=0.5.3=pyhd3eb1b0_0\n - nbconvert=6.0.7=py37_0\n - nbformat=5.1.3=pyhd3eb1b0_0\n - ncurses=6.2=he6710b0_1\n - nest-asyncio=1.5.1=pyhd3eb1b0_0\n - ninja=1.10.2=hff7bd54_1\n - notebook=6.4.0=py37h06a4308_0\n - numpy=1.20.2=py37h2d18471_0\n - numpy-base=1.20.2=py37hfae3a4d_0\n - nvidiacub=1.10.0=0\n - olefile=0.46=py37_0\n - open3d=0.13.0=py37_0\n - openssl=1.1.1k=h7f98852_0\n - packaging=20.9=pyhd3eb1b0_0\n - pandas=1.2.4=py37h2531618_0\n - pandoc=2.12=h06a4308_0\n - pandocfilters=1.4.3=py37h06a4308_1\n - parso=0.8.2=pyhd3eb1b0_0\n - pcre=8.44=he6710b0_0\n - pexpect=4.8.0=pyhd3eb1b0_3\n - pickleshare=0.7.5=pyhd3eb1b0_1003\n - pillow=8.2.0=py37he98fc37_0\n - pip=21.1.2=py37h06a4308_0\n - plyfile=0.7.4=pyhd8ed1ab_0\n - portalocker=2.3.0=py37h06a4308_0\n - prometheus_client=0.11.0=pyhd3eb1b0_0\n - prompt-toolkit=3.0.17=pyh06a4308_0\n - prompt_toolkit=3.0.17=hd3eb1b0_0\n - ptyprocess=0.7.0=pyhd3eb1b0_2\n - pycparser=2.20=py_2\n - pygments=2.9.0=pyhd3eb1b0_0\n - pyopenssl=20.0.1=pyhd3eb1b0_1\n - pyparsing=2.4.7=pyhd3eb1b0_0\n - pyqt=5.9.2=py37h05f1152_2\n - pyrsistent=0.17.3=py37h7b6447c_0\n - pysocks=1.7.1=py37_1\n - python=3.7.10=h12debd9_4\n - python-dateutil=2.8.1=pyhd3eb1b0_0\n - python_abi=3.7=1_cp37m\n - pytorch=1.7.1=py3.7_cuda10.2.89_cudnn7.6.5_0\n - pytz=2021.1=pyhd3eb1b0_0\n - pyzmq=20.0.0=py37h2531618_1\n - qt=5.9.7=h5867ecd_1\n - qtconsole=5.1.1=pyhd8ed1ab_0\n - qtpy=1.10.0=pyhd8ed1ab_0\n - readline=8.1=h27cfd23_0\n - requests=2.26.0=pyhd3eb1b0_0\n - scikit-learn=0.24.2=py37ha9443f7_0\n - scipy=1.6.2=py37had2a1c9_1\n - send2trash=1.5.0=pyhd3eb1b0_1\n - setuptools=52.0.0=py37h06a4308_0\n - sip=4.19.8=py37hf484d3e_0\n - six=1.16.0=pyhd3eb1b0_0\n - sniffio=1.2.0=py37h06a4308_1\n - sqlite=3.35.4=hdfb4753_0\n - tabulate=0.8.9=py37h06a4308_0\n - terminado=0.9.4=py37h06a4308_0\n - testpath=0.4.4=pyhd3eb1b0_0\n - threadpoolctl=2.1.0=pyh5ca1d4c_0\n - tk=8.6.10=hbc83047_0\n - torchvision=0.8.2=py37_cu102\n - tornado=6.1=py37h27cfd23_0\n - traitlets=5.0.5=pyhd3eb1b0_0\n - typing_extensions=3.7.4.3=pyha847dfd_0\n - wcwidth=0.2.5=py_0\n - webencodings=0.5.1=py37_1\n - wheel=0.36.2=pyhd3eb1b0_0\n - widgetsnbextension=3.5.1=py37_0\n - xz=5.2.5=h7b6447c_0\n - yacs=0.1.6=py_0\n - yaml=0.2.5=h7b6447c_0\n - zeromq=4.3.4=h2531618_0\n - zipp=3.4.1=pyhd3eb1b0_0\n - zlib=1.2.11=h7b6447c_3\n - zstd=1.4.9=haebb681_0\n - pip:\n - absl-py==0.13.0\n - addict==2.4.0\n - anykeystore==0.2\n - cachetools==4.2.2\n - cryptacular==1.5.5\n - cython==0.29.24\n - decorator==4.4.2\n - easydict==1.9\n - einops==0.3.0\n - fire==0.4.0\n - flow-vis==0.1\n - freetype-py==2.2.0\n - future==0.18.2\n - glumpy==1.2.0\n - google-auth==1.31.0\n - google-auth-oauthlib==0.4.4\n - greenlet==1.1.0\n - grpcio==1.38.0\n - hupper==1.10.3\n - imageio==2.9.0\n - llvmlite==0.36.0\n - loguru==0.5.3\n - markdown==3.3.4\n - networkx==2.5.1\n - numba==0.53.1\n - numpy-quaternion==2021.6.9.13.34.11\n - oauthlib==3.1.1\n - opencv-python==4.5.2.54\n - pastedeploy==2.1.1\n - pbkdf2==1.3\n - plaster==1.0\n - plaster-pastedeploy==0.7\n - protobuf==3.17.3\n - pyasn1==0.4.8\n - pyasn1-modules==0.2.8\n - pyassimp==4.1.3\n - pyglet==1.5.17\n - pyopengl==3.1.5\n - pyopengl-accelerate==3.1.5\n - pyramid==2.0\n - pyramid-mailer==0.15.1\n - python3-openid==3.2.0\n - pywavelets==1.1.1\n - pyyaml==5.4.1\n - repoze-sendmail==4.4.1\n - requests-oauthlib==1.3.0\n - rsa==4.7.2\n - scikit-image==0.18.1\n - sqlalchemy==1.4.18\n - tensorboard==2.5.0\n - tensorboard-data-server==0.6.1\n - tensorboard-plugin-wit==1.8.0\n - tensorboardx==2.2\n - termcolor==1.1.0\n - tifffile==2021.6.14\n - tqdm==4.61.1\n - transaction==3.0.1\n - transforms3d==0.3.1\n - translationstring==1.4\n - triangle==20200424\n - urllib3==1.26.5\n - velruse==1.1.1\n - venusian==3.0.0\n - vispy==0.6.6\n - webob==1.8.7\n - werkzeug==2.0.1\n - zope-deprecation==4.4.0\n - zope-interface==5.4.0\n - mmcv \nprefix: /opt/miniconda3/envs/py37_tmp\n"
},
{
"path": "geometry/__init__.py",
"content": ""
},
{
"path": "geometry/cholesky.py",
"content": "# import tensorflow as tf\nimport torch #as tf\nimport numpy as np\n# from utils.einsum import einsum\nfrom torch import einsum\n\n\n\nclass _cholesky_solve(torch.autograd.Function):\n @staticmethod\n def forward(ctx, H, b):\n chol = torch.cholesky(H)\n xx = torch.cholesky_solve(b, chol)\n ctx.save_for_backward(chol, xx)\n\n return xx\n\n # see OptNet: https://arxiv.org/pdf/1703.00443.pdf\n @staticmethod\n def backward(ctx, dx):\n chol, xx = ctx.saved_tensors\n\n dz = torch.cholesky_solve(dx, chol)\n xs = torch.squeeze(xx, -1)\n zs = torch.squeeze(dz, -1)\n dH = -einsum('...i,...j->...ij', xs, zs)\n\n return dH, dz\ndef cholesky_solve(H, b):\n return _cholesky_solve.apply(H,b)\n\ndef solve(H, b, max_update=1.0):\n \"\"\" Solves the linear system Hx = b, H > 0\"\"\"\n\n # small system, solve on cpu\n H = H.to(dtype=torch.float64) \n b = b.to(dtype=torch.float64) \n\n b = torch.unsqueeze(b, -1)\n x = cholesky_solve(H, b)\n\n # replaces nans and clip large updates\n bad_values = torch.isnan(x) \n x = torch.where(bad_values, torch.zeros_like(x), x)\n x = torch.clamp(x, -max_update, max_update)\n\n x = torch.squeeze(x, -1)\n x = x.to(dtype=torch.float32) \n \n return x\n\n\n\ndef __test__():\n import numpy as np \n np.random.seed(0)\n M=np.random.uniform(size=(3,3))\n H=torch.tensor(M@M.transpose(-1,-2), requires_grad=True )\n\n b=torch.tensor(np.random.uniform(size=(3,) ), requires_grad=True )\n\n x= solve(H,b )\n\n x.backward(torch.ones_like(x) )\n\n\n print(f\"H={H}, b={b}, x={x}, grad={H.grad, b.grad}\")\n\nif __name__==\"__main__\":\n __test__()\n"
},
{
"path": "geometry/diff_render.py",
"content": "import torch\nimport torch.nn as nn \nimport torch.nn.functional as F \n\nimport numpy \n\nfrom pytorch3d.renderer import (\n OpenGLPerspectiveCameras, look_at_view_transform, look_at_rotation,\n RasterizationSettings, MeshRenderer, MeshRasterizer, BlendParams,\n camera_position_from_spherical_angles, HardPhongShader, PointLights,FoVPerspectiveCameras, PerspectiveCameras, SoftPhongShader, Materials\n) \ntry:\n from pytorch3d.structures import Meshes, Textures\n use_textures = True\nexcept:\n from pytorch3d.structures import Meshes\n from pytorch3d.renderer import TexturesVertex\n from pytorch3d.renderer import TexturesVertex as Textures\n\n use_textures = False\n\nimport pytorch3d.renderer.mesh.utils as utils\nfrom pytorch3d.io import load_obj, load_ply, load_objs_as_meshes\nfrom pytorch3d.renderer.mesh.rasterizer import Fragments\n\nfrom plyfile import PlyData\nfrom utils.furthest_point_sample import fragmentation_fps\nimport time\n\n\ndef rasterize(R, T, meshes, rasterizer, blur_radius=0):\n # It will automatically update the camera settings -> R, T in rasterizer.camera\n fragments = rasterizer(meshes, R=R, T=T)\n\n # Copy from pytorch3D source code, try if it is necessary to do gradient decent\n if blur_radius > 0.0:\n clipped_bary_coords = utils._clip_barycentric_coordinates(\n fragments.bary_coords\n )\n clipped_zbuf = utils._interpolate_zbuf(\n fragments.pix_to_face, clipped_bary_coords, meshes\n )\n fragments = Fragments(\n bary_coords=clipped_bary_coords,\n zbuf=clipped_zbuf,\n dists=fragments.dists,\n pix_to_face=fragments.pix_to_face,\n )\n return fragments\n\ndef set_bary_coords_to_nearest(bary_coords_):\n ori_shape = bary_coords_.shape\n exr = bary_coords_ * (bary_coords_ < 0)\n bary_coords_ = bary_coords_.view(-1, bary_coords_.shape[-1])\n arg_max_idx = bary_coords_.argmax(1)\n return torch.zeros_like(bary_coords_).scatter(1, arg_max_idx.unsqueeze(1), 1.0).view(*ori_shape) + exr\n\nclass MeshRendererWithDepth(nn.Module):\n def __init__(self, rasterizer, shader):\n super().__init__()\n self.rasterizer = rasterizer\n self.shader = shader\n\n def to(self, device):\n # Rasterizer and shader have submodules which are not of type nn.Module\n self.rasterizer.to(device)\n self.shader.to(device)\n return self\n\n def forward(self, meshes_world, **kwargs) -> torch.Tensor:\n fragments = self.rasterizer(meshes_world, **kwargs)\n images = self.shader(fragments, meshes_world, **kwargs)\n return images, fragments.zbuf\n\nclass DiffRender(nn.Module):\n def __init__(self, mesh_path, render_texture=False):\n super().__init__()\n\n # self.mesh = mesh\n if mesh_path.endswith('.ply'):\n verts, faces = load_ply(mesh_path)\n self.mesh = Meshes(verts=[verts], faces=[faces])\n elif mesh_path.endswith('.obj'):\n verts, faces,_ = load_obj(mesh_path)\n faces=faces.verts_idx\n self.mesh=load_objs_as_meshes([mesh_path])\n\n self.verts = verts\n self.faces = faces\n self.cam_opencv2pytch3d = torch.tensor(\n [[-1,0,0,0],\n [0,-1,0, 0],\n [0,0, 1, 0],\n [0,0, 0, 1]], dtype=torch.float32\n )\n self.render_texture = render_texture\n\n #get patch infos\n self.pat_centers, self.pat_center_inds, self.vert_frag_ids= fragmentation_fps(verts.detach().cpu().numpy(), 64)\n self.pat_centers = torch.from_numpy(self.pat_centers)\n self.pat_center_inds = torch.from_numpy(self.pat_center_inds)\n self.vert_frag_ids = torch.from_numpy(self.vert_frag_ids)[...,None] #Nx1\n\n\n\n\n def to(self, *args, **kwargs):\n if 'device' in kwargs.keys():\n device = kwargs['device']\n else:\n device = args[0]\n super().to(device)\n self.mesh = self.mesh.to(device)\n self.verts = self.verts.to(device)\n self.faces = self.faces.to(device)\n self.pat_centers = self.pat_centers.to(device)\n self.pat_center_inds = self.pat_center_inds.to(device)\n self.vert_frag_ids = self.vert_frag_ids.to(device)\n \n return self\n\n def get_patch_center_depths(self, T, K):\n #no need to pre-transform, as here we do not use pytorch3d rendering\n # T = self.cam_opencv2pytch3d.to(device=T.device)@T\n\n ## X_cam = X_world R + t\n R = T[...,:3,:3].transpose(-1,-2)\n t = T[...,:3,3]\n\n #render depths\n X_cam= (self.pat_centers@R+t) #BxKx3\n depth= X_cam[...,2:] #BxKx1\n x=X_cam@K.transpose(-1,-2) #BxNx3\n x = x/x[...,-1:]\n img_coords= x[...,:2]\n\n\n return depth, img_coords \n\n # Calculate interpolated maps -> [n, c, h, w]\n # face_memory.shape: [n_face, 3, c]\n @staticmethod\n def forward_interpolate(R, t, meshes, face_memory, rasterizer, blur_radius=0, mode='bilinear', return_depth=True):\n\n fragments = rasterize(R, t, meshes, rasterizer, blur_radius=blur_radius)\n\n # [n, h, w, 1, d]\n if mode == 'nearest':\n out_map = utils.interpolate_face_attributes(fragments.pix_to_face, set_bary_coords_to_nearest(fragments.bary_coords), face_memory)\n else:\n out_map = utils.interpolate_face_attributes(fragments.pix_to_face, fragments.bary_coords, face_memory)\n\n out_map = out_map.squeeze(dim=3)\n out_map = out_map.transpose(3, 2).transpose(2, 1)\n if return_depth:\n return out_map, fragments.zbuf.permute(0,3,1,2) # depth\n else:\n return out_map\n\n def render_mesh(self, T, K, render_image_size, near=0.1, far=6, lights=(1,1,-1) ):\n B=T.shape[0]\n\n device = T.device\n T = self.cam_opencv2pytch3d.to(device=T.device)@T\n\n ## X_cam = X_world R + t\n R = T[...,:3,:3].transpose(-1,-2)\n t = T[...,:3,3]\n\n cameras = PerspectiveCameras(focal_length= torch.stack([K[:,0,0], K[:,1,1] ], dim=-1), \n principal_point=K[:,:2,2], R=R, T=t, image_size=[render_image_size]*B, in_ndc=False, device=device)\n lights = PointLights(device=device, location=[lights])\n\n raster_settings = RasterizationSettings(\n image_size=render_image_size,\n blur_radius=0.0,\n faces_per_pixel=1,\n bin_size=None, #0\n perspective_correct=True\n )\n materials = Materials(\n device=device,\n # specular_color=[[0.0, 1.0, 0.0]],\n shininess=0\n )\n renderer = MeshRendererWithDepth(\n rasterizer=MeshRasterizer(\n cameras=cameras, \n raster_settings=raster_settings\n ),\n shader=SoftPhongShader(\n device=device, \n cameras=cameras,\n lights=lights, \n blend_params=BlendParams(1e-4, 1e-4, (0, 0, 0))\n )\n )\n image,depth =renderer(self.mesh, lights=lights, materials=materials)\n\n return image.permute(0,3,1,2)[:,:3], depth.permute(0,3,1,2) # to BCHW\n\n def render_offset_map(self, T, K, render_image_size, near=0.1, far=6):\n yy, xx = torch.meshgrid(torch.arange(render_image_size[0], device=T.device), torch.arange(render_image_size[1], device=T.device) )\n # xx = xx.to(dtype=torch.float32)\n # yy = yy.to(dtype=torch.float32)\n coords_grid = torch.stack( [ xx.to(dtype=torch.float32), yy.to(dtype=torch.float32)], dim=-1 )\n\n #no need to pre-transform, as here we do not use pytorch3d rendering\n # T = self.cam_opencv2pytch3d.to(device=T.device)@T\n\n ## X_cam = X_world R + t\n R = T[...,:3,:3].transpose(-1,-2)\n t = T[...,:3,3]\n\n #render depths\n X_cam= (self.pat_centers@R+t)\n x=X_cam@K.transpose(-1,-2) #BxNx3\n x = x/x[...,-1:]\n\n offset = x[...,None,None,:2] - coords_grid #BxNx1x1x2-HxWx2\n \n return offset.permute(0,1,4,2,3) #BxNx2xHxW\n\n def forward(self, vert_attribute, T, K, render_image_size, near=0.1, far=6, mode='bilinear') :\n \"\"\"\n Args:\n vert_attribute: (N,C)\n T: (B,3,4) or (B,4,4)\n K: (B,3,3)\n render_image_size (tuple): (h,w)\n near (float, optional): Defaults to 0.1.\n far (int, optional): Defaults to 6.\n \"\"\"\n\n if vert_attribute is None:\n return self.render_mesh(T, K, render_image_size, near=0.1, far=6 )\n if self.render_texture:\n ren_tex=self.render_mesh(T, K, render_image_size, near=0.1, far=6 )\n\n\n B=T.shape[0]\n face_attribute = vert_attribute[self.faces.long()]\n\n device = T.device\n\n T = self.cam_opencv2pytch3d.to(device=T.device)@T\n\n ## X_cam = X_world R + t\n R = T[...,:3,:3].transpose(-1,-2)\n t = T[...,:3,3]\n # t = -(R@T[...,:3,3:]).squeeze(-1)\n\n cameras = PerspectiveCameras(focal_length= torch.stack([K[:,0,0], K[:,1,1] ], dim=-1), \n principal_point=K[:,:2,2], image_size=[render_image_size]*B, in_ndc=False, device=device)\n\n raster_settings = RasterizationSettings(\n image_size=render_image_size,\n blur_radius=0.0,\n faces_per_pixel=1,\n bin_size=None, #0\n perspective_correct=True\n )\n\n rasterizer = MeshRasterizer(\n cameras=cameras,\n raster_settings=raster_settings\n )\n\n out_map, out_depth=self.forward_interpolate(R, t, self.mesh, face_attribute, rasterizer, blur_radius=0, mode=mode)\n \n if not self.render_texture:\n return out_map, out_depth\n else:\n return torch.cat([ren_tex[0], out_map ], dim=1), out_depth\n\n def render_depth(self, T, K, render_image_size, near=0.1, far=6, mode='neareast'):\n \"\"\"\n Args:\n T: (B,3,4) or (B,4,4)\n K: (B,3,3)\n render_image_size (tuple): (h,w)\n near (float, optional): Defaults to 0.1.\n far (int, optional): Defaults to 6.\n mode: 'bilinear' or 'neareast'\n \"\"\"\n\n B=T.shape[0]\n device = T.device\n\n T = self.cam_opencv2pytch3d.to(device=T.device)@T\n\n ## X_cam = X_world R + t\n R = T[...,:3,:3].transpose(-1,-2)\n t = T[...,:3,3]\n cameras = PerspectiveCameras(focal_length= torch.stack([K[:,0,0], K[:,1,1] ], dim=-1), \n principal_point=K[:,:2,2], image_size=[render_image_size]*B, in_ndc=False, device=device)\n\n raster_settings = RasterizationSettings(\n image_size=render_image_size,\n blur_radius=0.0,\n faces_per_pixel=1,\n bin_size=0\n )\n\n rasterizer = MeshRasterizer(\n cameras=cameras,\n raster_settings=raster_settings\n )\n\n\n #render depths\n vert_depths= (self.verts@R+t).squeeze(0)[...,2:]\n face_depths = vert_depths[self.faces.long()]\n out_depth=self.forward_interpolate(R, t, self.mesh, face_depths, rasterizer, blur_radius=0, mode='nearest', return_depth=False)\n\n return out_depth\n\n\nclass DiffRendererWrapper(nn.Module):\n def __init__(self, obj_paths, device=\"cuda\", render_texture=False ):\n super().__init__()\n\n self.renderers = []\n for obj_path in obj_paths:\n self.renderers.append( \n DiffRender(obj_path, render_texture).to(device=device)\n )\n\n self.renderers=nn.ModuleList(self.renderers)\n self.cls2idx=None #could be updated outside\n\n def get_patch_center_depths(self, model_names, T, K):\n \n depths= []\n image_coords= []\n for b,_ in enumerate(model_names):\n model_idx = self.cls2idx[model_names[b]]\n depth, img_coord = self.renderers[model_idx].get_patch_center_depths(T[b:b+1], K )\n depths.append(depth)\n image_coords.append(img_coord)\n \n return torch.cat(depths, dim=0), torch.cat(image_coords, dim=0)\n\n\n def render_offset_map(self, model_names, T, K, render_image_size, near=0.1, far=6):\n offsets= []\n for b,_ in enumerate(model_names):\n model_idx = self.cls2idx[model_names[b]]\n\n offset = self.renderers[model_idx].render_offset_map(T[b:b+1], K[b:b+1], render_image_size, near, far )\n offsets.append(offset)\n \n return torch.cat(offsets, dim=0)\n\n def render_pat_id(self, model_names, T, K, render_image_size, near=0.1, far=6):\n\n pat_ids= []\n for b,_ in enumerate(model_names):\n model_idx = self.cls2idx[model_names[b]]\n \n pat_id,_ = self.renderers[model_idx].forward(self.renderers[model_idx].vert_frag_ids.float()+1,T[b:b+1], K[b:b+1], render_image_size, near, far, 'nearest' )\n pat_ids.append(pat_id-1) #+1 -1, set invalid parts as -1's \n \n return torch.cat(pat_ids, dim=0)\n\n def render_depth(self, model_names, T, K, render_image_size, near=0.1, far=6):\n\n depth_outputs= []\n for b,_ in enumerate(model_names):\n model_idx = self.cls2idx[model_names[b]]\n\n depth = self.renderers[model_idx].render_depth( T[b:b+1], K[b:b+1], render_image_size, near, far, 'nearest' )\n depth_outputs.append(depth)\n \n return torch.cat(depth_outputs, dim=0)\n\n def forward(self, model_names, vert_attribute, T, K, render_image_size, near=0.1, far=6):\n\n map_outputs= []\n depth_outputs= []\n for b,_ in enumerate(model_names):\n model_idx = self.cls2idx[model_names[b]]\n\n feamap, depth= self.renderers[model_idx]( vert_attribute[b], T[b:b+1], K[b:b+1], render_image_size, near, far )\n\n map_outputs.append(feamap)\n depth_outputs.append(depth)\n return torch.cat(map_outputs, dim=0) , torch.cat(depth_outputs, dim=0)\n\n\n"
},
{
"path": "geometry/diff_render_optim.py",
"content": "## Speed optimized: sharing the rasterization among different rendering process\n\nimport torch\nimport torch.nn as nn \nimport torch.nn.functional as F \n\nimport numpy \n\nfrom pytorch3d.renderer import (\n OpenGLPerspectiveCameras, look_at_view_transform, look_at_rotation,\n RasterizationSettings, MeshRenderer, MeshRasterizer, BlendParams,\n camera_position_from_spherical_angles, HardPhongShader, PointLights,FoVPerspectiveCameras, PerspectiveCameras, SoftPhongShader, Materials \n) \ntry:\n from pytorch3d.structures import Meshes, Textures\n use_textures = True\nexcept:\n from pytorch3d.structures import Meshes\n from pytorch3d.renderer import TexturesVertex\n from pytorch3d.renderer import TexturesVertex as Textures\n\n use_textures = False\n\nimport pytorch3d.renderer.mesh.utils as utils\nfrom pytorch3d.io import load_obj, load_ply, load_objs_as_meshes\nfrom pytorch3d.renderer.mesh.rasterizer import Fragments\n\nfrom plyfile import PlyData\nfrom utils.furthest_point_sample import fragmentation_fps\n\n\n\ndef rasterize(R, T, meshes, rasterizer, blur_radius=0):\n # It will automatically update the camera settings -> R, T in rasterizer.camera\n fragments = rasterizer(meshes, R=R, T=T)\n\n # Copy from pytorch3D source code, try if it is necessary to do gradient decent\n if blur_radius > 0.0:\n clipped_bary_coords = utils._clip_barycentric_coordinates(\n fragments.bary_coords\n )\n clipped_zbuf = utils._interpolate_zbuf(\n fragments.pix_to_face, clipped_bary_coords, meshes\n )\n fragments = Fragments(\n bary_coords=clipped_bary_coords,\n zbuf=clipped_zbuf,\n dists=fragments.dists,\n pix_to_face=fragments.pix_to_face,\n )\n return fragments\n\ndef set_bary_coords_to_nearest(bary_coords_):\n ori_shape = bary_coords_.shape\n exr = bary_coords_ * (bary_coords_ < 0)\n bary_coords_ = bary_coords_.view(-1, bary_coords_.shape[-1])\n arg_max_idx = bary_coords_.argmax(1)\n return torch.zeros_like(bary_coords_).scatter(1, arg_max_idx.unsqueeze(1), 1.0).view(*ori_shape) + exr\n\nclass MeshRendererWithDepth(nn.Module):\n def __init__(self, rasterizer, shader):\n super().__init__()\n self.rasterizer = rasterizer\n self.shader = shader\n\n def to(self, device):\n # Rasterizer and shader have submodules which are not of type nn.Module\n self.rasterizer.to(device)\n self.shader.to(device)\n return self\n\n def forward(self, meshes_world, **kwargs) -> torch.Tensor:\n fragments = self.rasterizer(meshes_world, **kwargs)\n images = self.shader(fragments, meshes_world, **kwargs)\n return images, fragments.zbuf\n\nclass MeshRendererWithDepth_v2(nn.Module):\n def __init__(self, rasterizer, shader):\n super().__init__()\n self.rasterizer = rasterizer\n self.shader = shader\n\n # def to(self, device):\n def to(self, *args, **kwargs):\n if 'device' in kwargs.keys():\n device = kwargs['device']\n else:\n device = args[0]\n super().to(device)\n # Rasterizer and shader have submodules which are not of type nn.Module\n self.rasterizer.to(device)\n self.shader.to(device)\n return self\n\n def forward(self, meshes_world, **kwargs) -> torch.Tensor:\n if 'fragments' not in kwargs.keys() or kwargs['fragments'] is None: # sharing fragment results with others for speed, as the rasterizing process occupies most of time\n if 'fragments' in kwargs:\n del kwargs['fragments']\n \n fragments = self.rasterizer(meshes_world, **kwargs)\n else:\n fragments = kwargs['fragments']\n del kwargs['fragments']\n\n images = self.shader(fragments, meshes_world, **kwargs)\n return images, fragments.zbuf\n\nclass DiffRender(nn.Module):\n def __init__(self, mesh_path, render_texture=False):\n super().__init__()\n\n # self.mesh = mesh\n if mesh_path.endswith('.ply'):\n verts, faces = load_ply(mesh_path)\n self.mesh = Meshes(verts=[verts], faces=[faces])\n elif mesh_path.endswith('.obj'):\n verts, faces,_ = load_obj(mesh_path)\n # import pdb; pdb.set_trace()\n faces=faces.verts_idx\n self.mesh=load_objs_as_meshes([mesh_path])\n\n # self.mesh = Meshes(verts=verts, faces=faces, textures=None)\n self.verts = verts\n self.faces = faces\n # self.mesh = Meshes(verts=[verts], faces=[faces])\n # self.feature=feature\n self.cam_opencv2pytch3d = torch.tensor(\n [[-1,0,0,0],\n [0,-1,0, 0],\n [0,0, 1, 0],\n [0,0, 0, 1]], dtype=torch.float32\n )\n self.render_texture = render_texture\n\n #get patch infos\n self.pat_centers, self.pat_center_inds, self.vert_frag_ids= fragmentation_fps(verts.detach().cpu().numpy(), 64)\n self.pat_centers = torch.from_numpy(self.pat_centers)\n self.pat_center_inds = torch.from_numpy(self.pat_center_inds)\n self.vert_frag_ids = torch.from_numpy(self.vert_frag_ids)[...,None] #Nx1\n\n\n\n\n def to(self, *args, **kwargs):\n if 'device' in kwargs.keys():\n device = kwargs['device']\n else:\n device = args[0]\n super().to(device)\n # self.rasterizer.cameras = self.rasterizer.cameras.to(device)\n # self.face_memory = self.face_memory.to(device)\n self.mesh = self.mesh.to(device)\n self.verts = self.verts.to(device)\n self.faces = self.faces.to(device)\n self.pat_centers = self.pat_centers.to(device)\n self.pat_center_inds = self.pat_center_inds.to(device)\n self.vert_frag_ids = self.vert_frag_ids.to(device)\n\n \n # self.cam_opencv2pytch3d = self.cam_opencv2pytch3d.to(device=device)\n return self\n\n def get_patch_center_depths(self, T, K):\n\n #no need to pre-transform, as here we do not use pytorch3d rendering\n # T = self.cam_opencv2pytch3d.to(device=T.device)@T\n\n ## X_cam = X_world R + t\n R = T[...,:3,:3].transpose(-1,-2)\n t = T[...,:3,3]\n\n #render depths\n X_cam= (self.pat_centers@R+t) #BxKx3\n depth= X_cam[...,2:] #BxKx1\n x=X_cam@K.transpose(-1,-2) #BxNx3\n x = x/x[...,-1:]\n img_coords= x[...,:2]\n\n\n return depth, img_coords \n\n # Calculate interpolated maps -> [n, c, h, w]\n # face_memory.shape: [n_face, 3, c]\n @staticmethod\n def forward_interpolate(R, t, meshes, face_memory, rasterizer, blur_radius=0, mode='bilinear', return_depth=True):\n\n fragments = rasterize(R, t, meshes, rasterizer, blur_radius=blur_radius)\n\n # [n, h, w, 1, d]\n if mode == 'nearest':\n out_map = utils.interpolate_face_attributes(fragments.pix_to_face, set_bary_coords_to_nearest(fragments.bary_coords), face_memory)\n else:\n out_map = utils.interpolate_face_attributes(fragments.pix_to_face, fragments.bary_coords, face_memory)\n out_map = out_map.squeeze(dim=3)\n out_map = out_map.transpose(3, 2).transpose(2, 1)\n if return_depth:\n return out_map, fragments.zbuf.permute(0,3,1,2), fragments # depth\n else:\n return out_map, fragments\n\n def render_mesh(self, T, K, render_image_size, near=0.1, far=6, lights=(1,1,-1), fragments=None ):\n B=T.shape[0]\n # face_attribute = vert_attribute[self.faces.long()]\n\n device = T.device\n T = self.cam_opencv2pytch3d.to(device=T.device)@T\n\n ## X_cam = X_world R + t\n R = T[...,:3,:3].transpose(-1,-2)\n t = T[...,:3,3]\n\n cameras = PerspectiveCameras(focal_length= torch.stack([K[:,0,0], K[:,1,1] ], dim=-1), \n principal_point=K[:,:2,2], R=R, T=t, image_size=[render_image_size]*B, in_ndc=False, device=device)\n lights = PointLights(device=device, location=[lights])\n\n raster_settings = RasterizationSettings(\n image_size=render_image_size,\n blur_radius=0.0,\n faces_per_pixel=1, #5,\n bin_size=None, #0\n perspective_correct=True\n )\n materials = Materials(\n device=device,\n # specular_color=[[0.0, 1.0, 0.0]],\n shininess=0\n )\n # renderer = MeshRendererWithDepth(\n renderer = MeshRendererWithDepth_v2(\n rasterizer=MeshRasterizer(\n cameras=cameras, \n raster_settings=raster_settings\n ),\n shader=SoftPhongShader(\n # shader=SoftGouraudShader(\n device=device, \n cameras=cameras,\n lights=lights, \n blend_params=BlendParams(1e-4, 1e-4, (0, 0, 0))\n )\n )\n image,depth =renderer(self.mesh, lights=lights, materials=materials, fragments=fragments)\n\n return image.permute(0,3,1,2)[:,:3], depth.permute(0,3,1,2) # to BCHW\n\n def render_offset_map(self, T, K, render_image_size, near=0.1, far=6):\n yy, xx = torch.meshgrid(torch.arange(render_image_size[0], device=T.device), torch.arange(render_image_size[1], device=T.device) )\n # xx = xx.to(dtype=torch.float32)\n # yy = yy.to(dtype=torch.float32)\n coords_grid = torch.stack( [ xx.to(dtype=torch.float32), yy.to(dtype=torch.float32)], dim=-1 )\n\n #no need to pre-transform, as here we do not use pytorch3d rendering\n # T = self.cam_opencv2pytch3d.to(device=T.device)@T\n\n ## X_cam = X_world R + t\n R = T[...,:3,:3].transpose(-1,-2)\n t = T[...,:3,3]\n\n #render depths\n X_cam= (self.pat_centers@R+t)#.squeeze(0)[...,2:]\n x=X_cam@K.transpose(-1,-2) #BxNx3\n x = x/x[...,-1:]\n\n offset = x[...,None,None,:2] - coords_grid #BxNx1x1x2-HxWx2\n \n return offset.permute(0,1,4,2,3) #BxNx2xHxW\n\n # def forward(self, face_attribute, T, K, render_image_size, near=0.1, far=6):\n def forward(self, vert_attribute, T, K, render_image_size, near=0.1, far=6, render_texture=None, mode='bilinear') :\n \"\"\"\n Args:\n vert_attribute: (N,C)\n T: (B,3,4) or (B,4,4)\n K: (B,3,3)\n render_image_size (tuple): (h,w)\n near (float, optional): Defaults to 0.1.\n far (int, optional): Defaults to 6.\n \"\"\"\n\n # use default rendering settings \n if render_texture is None:\n render_texture= self.render_texture \n \n if vert_attribute is None:\n # only render the rgb image\n return self.render_mesh(T, K, render_image_size, near=0.1, far=6 )\n\n B=T.shape[0]\n face_attribute = vert_attribute[self.faces.long()]\n\n device = T.device\n\n T = self.cam_opencv2pytch3d.to(device=T.device)@T\n\n ## X_cam = X_world R + t\n R = T[...,:3,:3].transpose(-1,-2)\n t = T[...,:3,3]\n # t = -(R@T[...,:3,3:]).squeeze(-1)\n \n cameras = PerspectiveCameras(focal_length= torch.stack([K[:,0,0], K[:,1,1] ], dim=-1), \n principal_point=K[:,:2,2], image_size=[render_image_size]*B, in_ndc=False, device=device)\n\n raster_settings = RasterizationSettings(\n image_size=render_image_size,\n blur_radius=0.0,\n faces_per_pixel=1,\n bin_size=None, #0\n perspective_correct=True\n )\n\n rasterizer = MeshRasterizer(\n cameras=cameras,\n raster_settings=raster_settings\n )\n\n # forward_interpolate(R, T, meshes, face_memory, rasterizer, blur_radius=0, mode='bilinear')\n out_map, out_depth, fragments=self.forward_interpolate(R, t, self.mesh, face_attribute, rasterizer, blur_radius=0, mode=mode)\n \n if not render_texture:\n return out_map, out_depth\n else:\n ren_tex=self.render_mesh(T, K, render_image_size, near=0.1, far=6, fragments=fragments )\n\n #The first 3 channels contain the rendered textures\n return torch.cat([ren_tex[0], out_map ], dim=1), out_depth\n\n def render_depth(self, T, K, render_image_size, near=0.1, far=6, mode='neareast'):\n \"\"\"\n Args:\n T: (B,3,4) or (B,4,4)\n K: (B,3,3)\n render_image_size (tuple): (h,w)\n near (float, optional): Defaults to 0.1.\n far (int, optional): Defaults to 6.\n mode: 'bilinear' or 'neareast'\n \"\"\"\n\n B=T.shape[0]\n device = T.device\n\n T = self.cam_opencv2pytch3d.to(device=T.device)@T\n\n ## X_cam = X_world R + t\n R = T[...,:3,:3].transpose(-1,-2)\n t = T[...,:3,3]\n cameras = PerspectiveCameras(focal_length= torch.stack([K[:,0,0], K[:,1,1] ], dim=-1), \n principal_point=K[:,:2,2], image_size=[render_image_size]*B, in_ndc=False, device=device)\n\n raster_settings = RasterizationSettings(\n image_size=render_image_size,\n blur_radius=0.0,\n faces_per_pixel=1,\n bin_size=0\n )\n\n rasterizer = MeshRasterizer(\n cameras=cameras,\n raster_settings=raster_settings\n )\n\n\n #render depths\n vert_depths= (self.verts@R+t).squeeze(0)[...,2:]\n face_depths = vert_depths[self.faces.long()]\n out_depth, _ =self.forward_interpolate(R, t, self.mesh, face_depths, rasterizer, blur_radius=0, mode='nearest', return_depth=False)\n\n return out_depth\n\n def render_pointcloud(self, T, K, render_image_size, near=0.1, far=6):\n \"\"\"\n Args:\n T: (B,3,4) or (B,4,4)\n K: (B,3,3)\n render_image_size (tuple): (h,w)\n near (float, optional): Defaults to 0.1.\n far (int, optional): Defaults to 6.\n mode: 'bilinear' or 'neareast'\n \"\"\"\n\n B=T.shape[0]\n device = T.device\n\n # T = self.cam_opencv2pytch3d.to(device=T.device)@T\n\n ## X_cam = X_world R + t\n R = T[...,:3,:3].transpose(-1,-2)\n t = T[...,:3,3]\n\n #render depths\n # vert_depths= (self.verts@R+t).squeeze(0)[...,2:]\n X_cam= (self.verts@R+t)#.squeeze(0)\n\n x=X_cam@K.transpose(-1,-2) #BxNx3\n depth = x[...,-1]\n x = x/x[...,-1:]\n\n out = torch.zeros([1,1, *render_image_size], dtype=R.dtype, device=R.device)\n out[:, :, \n torch.round(x[0, :, 1]).long().clamp(0, out.shape[2]-1),\n torch.round(x[0, :, 0]).long().clamp(0, out.shape[3]-1)] = depth \n\n return out #1x1xHxW\n\n\nclass DiffRendererWrapper(nn.Module):\n def __init__(self, obj_paths, device=\"cuda\", render_texture=False ):\n super().__init__()\n\n self.renderers = []\n for obj_path in obj_paths:\n self.renderers.append( \n DiffRender(obj_path, render_texture).to(device=device)\n )\n\n self.renderers=nn.ModuleList(self.renderers)\n self.cls2idx=None #updated outside\n\n def get_patch_center_depths(self, model_names, T, K):\n \n depths= []\n image_coords= []\n for b,_ in enumerate(model_names):\n model_idx = self.cls2idx[model_names[b]]\n depth, img_coord = self.renderers[model_idx].get_patch_center_depths(T[b:b+1], K )\n depths.append(depth)\n image_coords.append(img_coord)\n \n return torch.cat(depths, dim=0), torch.cat(image_coords, dim=0)\n\n\n def render_offset_map(self, model_names, T, K, render_image_size, near=0.1, far=6):\n offsets= []\n for b,_ in enumerate(model_names):\n model_idx = self.cls2idx[model_names[b]]\n\n offset = self.renderers[model_idx].render_offset_map(T[b:b+1], K[b:b+1], render_image_size, near, far )\n offsets.append(offset)\n \n return torch.cat(offsets, dim=0)\n\n def render_pat_id(self, model_names, T, K, render_image_size, near=0.1, far=6):\n\n pat_ids= []\n for b,_ in enumerate(model_names):\n model_idx = self.cls2idx[model_names[b]]\n # face_pat_id = self.renderers[model_idx].vert_frag_ids[self.renderers[model_idx].faces.long()]\n \n pat_id,_ = self.renderers[model_idx].forward(self.renderers[model_idx].vert_frag_ids.float()+1,T[b:b+1], K[b:b+1], render_image_size, near, far, 'nearest' )\n pat_ids.append(pat_id-1) #+1 -1, set invalid parts as -1's \n \n return torch.cat(pat_ids, dim=0)\n\n def render_depth(self, model_names, T, K, render_image_size, near=0.1, far=6):\n \n depth_outputs= []\n for b,_ in enumerate(model_names):\n model_idx = self.cls2idx[model_names[b]]\n\n depth = self.renderers[model_idx].render_depth( T[b:b+1], K[b:b+1], render_image_size, near, far, 'nearest' )\n depth_outputs.append(depth)\n \n return torch.cat(depth_outputs, dim=0)\n def render_mesh(self, model_names, T, K, render_image_size, near=0.1, far=6):\n\n outputs= []\n for b,_ in enumerate(model_names):\n model_idx = self.cls2idx[model_names[b]]\n\n img= self.renderers[model_idx].render_mesh( T[b:b+1], K[b:b+1], render_image_size, near, far, )[0]\n outputs.append(img)\n \n return torch.cat(outputs, dim=0)\n\n def render_pointcloud(self, model_names, T, K, render_image_size, near=0.1, far=6):\n outputs= []\n for b,_ in enumerate(model_names):\n model_idx = self.cls2idx[model_names[b]]\n depth = self.renderers[model_idx].render_pointcloud( T[b:b+1], K[b:b+1], render_image_size, near, far )\n outputs.append(depth)\n \n return torch.cat(outputs, dim=0)\n\n def forward(self, model_names, vert_attribute, T, K, render_image_size, near=0.1, far=6, render_tex=False):\n\n map_outputs= []\n depth_outputs= []\n for b,_ in enumerate(model_names):\n model_idx = self.cls2idx[model_names[b]]\n\n feamap, depth= self.renderers[model_idx]( vert_attribute[b], T[b:b+1], K[b:b+1], render_image_size, near, far, render_texture=render_tex )\n\n map_outputs.append(feamap)\n depth_outputs.append(depth)\n return torch.cat(map_outputs, dim=0) , torch.cat(depth_outputs, dim=0)\n\n\n"
},
{
"path": "geometry/einsum.py",
"content": "# import tensorflow as torch\nimport torch as torch\n\nimport numpy as np\nimport re\nimport string\n\ndef einsum(equation, *inputs):\n\n equation = equation.replace(' ', '')\n # input_shapes = [x.get_shape() for x in list(inputs)]\n input_shapes = [x.shape for x in list(inputs)]\n match = re.match('^([a-zA-Z,.]+)(->[a-zA-Z.]*)?$', equation)\n if not match:\n raise ValueError('Indices have incorrect format: %s' % equation)\n\n input_axis_labels = match.group(1).split(',')\n output_axis_labels = match.group(2)[2:] if match.group(2) else None\n\n if len(input_shapes) != len(input_axis_labels):\n raise ValueError('Got %d arguments for equation \"%s\", expecting %d' %\n (len(input_shapes), equation, len(input_axis_labels)))\n\n # Resolve Ellipsis\n # Assign axes labels for unspecified dimensions in inputs. Labels taken\n # from unused labels. Follow numpy einsum broadcasting conventions for\n # tensors of different length and unlabeled output.\n ellipsis_axes = ''\n if '...' in equation:\n unused = ''.join([c for c in string.ascii_lowercase\n if c not in ''.join(input_axis_labels)])\n for i, ax in enumerate(input_axis_labels):\n if '...' in ax:\n parts = ax.split('...')\n if len(parts) != 2:\n raise ValueError('Unable to resolve ellipsis. Excess number found.')\n\n # n = input_shapes[i].ndims - len(''.join(parts))\n n = len(input_shapes[i]) - len(''.join(parts))\n if n < 0:\n raise ValueError('Ellipses lengths do not match.')\n if len(unused) < n:\n raise ValueError(\n 'Unable to resolve ellipsis, too many distinct labels.')\n replace_axes = unused[-n:] if n > 0 else ''\n input_axis_labels[i] = input_axis_labels[i].replace('...',\n replace_axes)\n if len(replace_axes) > len(ellipsis_axes):\n ellipsis_axes = replace_axes\n \n equation = equation.replace('...', ellipsis_axes)\n out = torch.einsum(equation, *inputs)\n # torch.add_to_collection(\"checkpoints\", out)\n return out\n"
},
{
"path": "geometry/intrinsics.py",
"content": "import torch\nimport numpy as np\n# from utils.einsum import einsum\nfrom .einsum import einsum\n\ndef intrinsics_vec_to_matrix(kvec):\n fx, fy, cx, cy = torch.unbind(kvec, dim=-1)\n z = torch.zeros_like(fx)\n o = torch.ones_like(fx)\n\n K = torch.stack([fx, z, cx, z, fy, cy, z, z, o], dim=-1)\n K = torch.reshape(K, list(kvec.shape)[:-1] + [3,3])\n return K\n\ndef intrinsics_matrix_to_vec(kmat):\n fx = kmat[..., 0, 0]\n fy = kmat[..., 1, 1]\n cx = kmat[..., 0, 2]\n cy = kmat[..., 1, 2]\n return torch.stack([fx, fy, cx, cy], dim=-1)\n\ndef update_intrinsics(intrinsics, delta_focal):\n kvec = intrinsics_matrix_to_vec(intrinsics)\n fx, fy, cx, cy = torch.unstack(kvec, num=4, axis=-1)\n df = torch.squeeze(delta_focal, -1)\n\n # update the focal lengths\n fx = torch.exp(df) * fx\n fy = torch.exp(df) * fy\n\n kvec = torch.stack([fx, fy, cx, cy], axis=-1)\n kmat = intrinsics_vec_to_matrix(kvec)\n return kmat\n\ndef rescale_depth(depth, downscale=4):\n depth = depth[:,None]\n new_shape = [depth.shape[-2]//downscale, depth.shape[-1]//downscale]\n depth = torch.nn.functional.interpolate(depth, new_shape, mode='nearest')\n return torch.squeeze(depth, dim=1)\n\ndef rescale_depth_and_intrinsics(depth, intrinsics, downscale=4):\n sc = torch.tensor([1.0/downscale, 1.0/downscale, 1.0], dtype=torch.float32, device=depth.device)\n intrinsics = einsum('...ij,i->...ij', intrinsics, sc)\n depth = rescale_depth(depth, downscale=downscale)\n return depth, intrinsics\n\ndef rescale_depths_and_intrinsics(depth, intrinsics, downscale=4):\n batch, frames, height, width = [depth.shape[i] for i in range(4)]\n depth = torch.reshape(depth, [batch*frames, height, width])\n depth, intrinsics = rescale_depth_and_intrinsics(depth, intrinsics, downscale)\n depth = torch.reshape(depth,\n [batch, frames]+list(depth.shape)[1:])\n return depth, intrinsics\n"
},
{
"path": "geometry/projective_ops.py",
"content": "import numpy as np\nimport torch \n\n# from utils.einsum import einsum\nfrom torch import einsum\n\n\n# MIN_DEPTH = 0.1\nMIN_DEPTH = 0.01\n\ndef normalize_coords_grid(coords):\n \"\"\" normalize the coordinates to [-1,1]\n\n Args:\n coords: BxKxHxWx2\n \"\"\"\n coords=coords.clone()\n B,K,H,W,_ = coords.shape\n\n coords[...,0] = 2*coords[...,0]/(W-1)-1\n coords[...,1] = 2*coords[...,1]/(H-1)-1\n\n return coords\n\ndef coords_grid(ref, homogeneous=True):\n \"\"\" grid of pixel coordinates \"\"\"\n shape = ref.shape\n\n yy, xx = torch.meshgrid(torch.arange(shape[-2], device=ref.device), torch.arange(shape[-1], device=ref.device) )\n\n xx = xx.to(dtype=torch.float32)\n yy = yy.to(dtype=torch.float32)\n\n if homogeneous:\n coords = torch.stack([xx, yy, torch.ones_like(xx)], dim=-1)\n else:\n coords = torch.stack([xx, yy], dim=-1)\n\n new_shape = [1]*len(shape[:-2]) + list(shape[-2:]) + [-1]\n coords = torch.reshape(coords, new_shape)\n\n tile = list(shape[:-2])+ [1,1,1]\n coords = coords.repeat(tile)\n return coords # BxKxHxWx2\n\n\ndef extract_and_reshape_intrinsics(intrinsics, shape=None):\n \"\"\" Extracts (fx, fy, cx, cy) from intrinsics matrix \"\"\"\n\n fx = intrinsics[:, 0, 0]\n fy = intrinsics[:, 1, 1]\n cx = intrinsics[:, 0, 2]\n cy = intrinsics[:, 1, 2]\n\n if shape is not None:\n batch = list(fx.shape[:1])\n fillr = [1]*len(shape[1:]) \n k_shape = batch+fillr\n\n fx = torch.reshape(fx, k_shape)\n fy = torch.reshape(fy, k_shape)\n cx = torch.reshape(cx, k_shape)\n cy = torch.reshape(cy, k_shape)\n\n return (fx, fy, cx, cy)\n\n\ndef backproject(depth, intrinsics, jacobian=False, depth_coords=None):\n \"\"\" backproject depth map to point cloud \"\"\"\n #depth_coords: (BxKxHxWx2)\n\n if depth_coords is None:\n coords = coords_grid(depth, homogeneous=True)\n x, y, _ = torch.unbind(coords, axis=-1)\n else:\n x, y = torch.unbind(depth_coords, axis=-1)\n\n x_shape = x.shape \n \n fx, fy, cx, cy = extract_and_reshape_intrinsics(intrinsics, x_shape) #Bx1x1x1\n\n Z = depth #BxKxHxW\n X = Z * (x - cx) / fx\n Y = Z * (y - cy) / fy \n points = torch.stack([X, Y, Z], axis=-1)\n\n if jacobian:\n o = torch.zeros_like(Z) # used to fill in zeros\n\n # jacobian w.r.t (fx, fy) , of shape BxKxHxWx4x1\n jacobian_intrinsics = torch.stack([\n torch.stack([-X / fx], dim=-1),\n torch.stack([-Y / fy], dim=-1),\n torch.stack([o], dim=-1),\n torch.stack([o], dim=-1)], axis=-2)\n\n return points, jacobian_intrinsics\n \n return points\n # return points, coords\n\n\ndef project(points, intrinsics, jacobian=False):\n \n \"\"\" project point cloud onto image \"\"\"\n X, Y, Z = torch.unbind(points, axis=-1)\n Z = torch.clamp(Z, min=MIN_DEPTH)\n\n x_shape = X.shape\n fx, fy, cx, cy = extract_and_reshape_intrinsics(intrinsics, x_shape)\n\n x = fx * (X / Z) + cx\n y = fy * (Y / Z) + cy\n coords = torch.stack([x, y], axis=-1)\n\n if jacobian:\n o = torch.zeros_like(x) # used to fill in zeros\n zinv1 = torch.where(Z <= MIN_DEPTH+.01, torch.zeros_like(Z), 1.0 / Z)\n zinv2 = torch.where(Z <= MIN_DEPTH+.01, torch.zeros_like(Z), 1.0 / Z**2)\n\n # jacobian w.r.t (X, Y, Z)\n jacobian_points = torch.stack([\n torch.stack([fx * zinv1, o, -fx * X * zinv2], axis=-1),\n torch.stack([o, fy * zinv1, -fy * Y * zinv2], axis=-1)], axis=-2)\n\n # jacobian w.r.t (fx, fy)\n jacobian_intrinsics = torch.stack([\n torch.stack([X * zinv1], axis=-1),\n torch.stack([Y * zinv1], axis=-1),], axis=-2)\n\n return coords, (jacobian_points, jacobian_intrinsics)\n\n return coords\n"
},
{
"path": "geometry/se3.py",
"content": "\"\"\"\nSO3 and SE3 operations, exponentials and logarithms adapted from Sophus\n\"\"\"\n\nimport numpy as np\nimport torch\nfrom .einsum import einsum\n\n\nMIN_THETA = 1e-4\n\ndef matdotv(A,b):\n return torch.squeeze(torch.matmul(A, torch.expand_dims(b, -1)), -1)\n\ndef hat(a):\n a1, a2, a3 = torch.split(a, [1,1,1], dim=-1)\n zz = torch.zeros_like(a1)\n\n ax = torch.stack([\n torch.cat([zz,-a3,a2], dim=-1),\n torch.cat([a3,zz,-a1], dim=-1),\n torch.cat([-a2,a1,zz], dim=-1)\n ], dim=-2)\n\n return ax\n \n\n### quaternion functions ###\n\ndef quaternion_rotate_point(q, pt, eq=None):\n if eq is None:\n w, vec = torch.split(q, [1, 3], axis=-1)\n uv = 2*matdotv(hat(vec), pt)\n return pt + w*uv + matdotv(hat(vec), uv)\n else:\n w, vec = torch.split(q, [1, 3], axis=-1)\n uv1 = 2*einsum(eq, hat(w*vec), pt)\n uv2 = 2*einsum(eq, hat(vec), pt)\n return pt + uv1 + einsum(eq, hat(vec), uv2)\n\ndef quaternion_rotate_matrix(q, mat, eq=None):\n if eq is None:\n w, vec = torch.split(q, [1, 3], axis=-1)\n uv = 2*torch.matmul(hat(vec), mat)\n return mat + w*uv + torch.matmul(hat(vec), uv)\n else:\n w, vec = torch.split(q, [1, 3], axis=-1)\n uv1 = 2*einsum(eq, hat(w*vec), mat)\n uv2 = 2*einsum(eq, hat(vec), mat)\n return mat + uv1 + einsum(eq, hat(vec), uv2)\n\ndef quaternion_inverse(q):\n return q * [1, -1, -1, -1]\n\ndef quaternion_multiply(a, b):\n aw, ax, ay, az = torch.split(a, [1,1,1,1], axis=-1)\n bw, bx, by, bz = torch.split(b, [1,1,1,1], axis=-1)\n \n q = torch.concat([\n aw * bw - ax * bx - ay * by - az * bz,\n aw * bx + ax * bw + ay * bz - az * by,\n aw * by + ay * bw + az * bx - ax * bz,\n aw * bz + az * bw + ax * by - ay * bx,\n ], axis=-1)\n\n return q\n\ndef quaternion_to_matrix(q):\n w, x, y, z = torch.split(q, [1,1,1,1], axis=-1)\n\n r11 = 1 - 2 * y**2 - 2 * z**2\n r12 = 2 * x * y - 2 * w * z\n r13 = 2 * z * x + 2 * w * y\n\n r21 = 2 * x * y + 2 * w * z\n r22 = 1 - 2 * x**2 - 2 * z**2\n r23 = 2 * y * z - 2 * w * x\n\n r31 = 2 * z * x - 2 * w * y\n r32 = 2 * y * z + 2 * w * x\n r33 = 1 - 2 * x**2 - 2 * y**2\n \n R = torch.stack([\n torch.concat([r11,r12,r13], axis=-1),\n torch.concat([r21,r22,r23], axis=-1),\n torch.concat([r31,r32,r33], axis=-1)\n ], axis=-2)\n\n return R\n\ndef rotation_matrix_to_quaternion(R):\n Rxx, Ryx, Rzx = R[...,0,0], R[...,0,1], R[...,0,2]\n Rxy, Ryy, Rzy = R[...,1,0], R[...,1,1], R[...,1,2]\n Rxz, Ryz, Rzz = R[...,2,0], R[...,2,1], R[...,2,2]\n\n zz = torch.zeros_like(Rxx)\n k1 = torch.stack([Rxx-Ryy-Rzz, zz, zz, zz], axis=-1)\n k2 = torch.stack([Ryx+Rxy, Ryy-Rxx-Rzz, zz, zz], axis=-1)\n k3 = torch.stack([Rzx+Rxz, Rzy+Ryz, Rzz-Rxx-Ryy,zz], axis=-1)\n k4 = torch.stack([Ryz-Rzy, Rzx-Rxz, Rxy-Ryx, Rxx+Ryy+Rzz], axis=-1)\n\n K = torch.stack([k1, k2, k3, k4], axis=-2)\n eigvals, eigvecs = torch.linalg.eigh(K)\n\n x, y, z, w = torch.split(eigvecs[...,-1], [1,1,1,1], axis=-1)\n qvec = torch.concat([w, x, y, z], axis=-1)\n qvec /= torch.sqrt(torch.reduce_sum(qvec**2, axis=-1, keepdims=True))\n\n return qvec * torch.sign(w)\n\ndef so3_expm_and_theta(omega):\n \"\"\" omega in \\so3 \"\"\"\n theta_sq = torch.reduce_sum(omega**2, axis=-1)\n theta = torch.sqrt(theta_sq)\n half_theta = 0.5*theta\n\n ### small ###\n imag_factor = torch.where(theta>MIN_THETA, \n torch.sin(half_theta) / (theta + 1e-12), \n 0.5 - (1.0/48.0)*theta_sq + (1.0/3840.0)*theta_sq*theta_sq)\n\n real_factor = torch.where(theta>MIN_THETA, torch.cos(half_theta),\n 1.0 - (1.0/8.0)*theta_sq + (1.0/384.0)*theta_sq*theta_sq)\n\n qw = real_factor\n qx = imag_factor * omega[...,0]\n qy = imag_factor * omega[...,1]\n qz = imag_factor * omega[...,2]\n\n quat = torch.stack([qw, qx, qy, qz], axis=-1)\n return quat, theta\n \ndef so3_logm_and_theta(so3):\n w, vec = torch.split(so3, [1,3], axis=-1)\n squared_n = torch.reduce_sum(vec**2, axis=-1, keepdims=True)\n n = torch.sqrt(squared_n)\n\n two_atan_nbyw_by_n = torch.where(n SE(3), works for arbitrary batch dimensions\n - Note: gradient is overridden with _se3_matrix_expm_grad, which approximates \n gradient for small upsilon_omega\n \"\"\"\n\n eps=1e-12\n inp_shape = upsilon_omega.shape \n out_shape = list(inp_shape)[:-1]+[4,4] \n\n upsilon_omega = torch.reshape(upsilon_omega, [-1, 6])\n batch = upsilon_omega.shape[0]\n v, w = torch.split(upsilon_omega, [3, 3], dim=-1)\n\n theta_sq = torch.sum(w**2, dim=1 )\n theta_sq = torch.reshape(theta_sq, [-1, 1, 1])\n\n theta = torch.sqrt(theta_sq)\n theta_po4 = theta_sq * theta_sq\n\n wx = hat(w)\n wx_sq = torch.matmul(wx, wx)\n I = torch.eye(3, dtype=upsilon_omega.dtype, device=upsilon_omega.device).repeat([batch,1,1])\n\n ### taylor approximations ###\n R1 = I + (1.0 - (1.0/6.0)*theta_sq + (1.0/120.0)*theta_po4) * wx + \\\n (0.5 - (1.0/12.0)*theta_sq + (1.0/720.0)*theta_po4) * wx_sq\n \n V1 = I + (0.5 - (1.0/24.0)*theta_sq + (1.0/720.0)*theta_po4)*wx + \\\n ((1.0/6.0) - (1.0/120.0)*theta_sq + (1.0/5040.0)*theta_po4)*wx_sq\n\n ### exact values ###\n R2 = I + (torch.sin(theta) / (theta+eps)) * wx +\\\n ((1 - torch.cos(theta)) / (theta_sq+eps)) * wx_sq\n\n V2 = I + ((1 - torch.cos(theta)) / (theta_sq + eps)) * wx + \\\n ((theta - torch.sin(theta))/(theta_sq*theta + eps)) * wx_sq\n\n # print(theta.shape, R1.shape, R2.shape, \">>>\", flush=True)\n # R = torch.where(theta[:, 0, 0] MIN_DEPTH) & (pt_new[..., -1] > MIN_DEPTH)\n # vmask = torch.cast(vmask, torch.float32)[..., torch.newaxis]\n # vmask = vmask.to(dtype=torch.float32)[..., None, :,:] #BxKx1xHxW\n vmask = vmask.to(dtype=torch.float32)[..., :, :, None] # BxKx1xHxW\n return coords, vmask\n return coords\n\n def induced_flow(self, depth, intrinsics, valid_mask=False):\n coords0 = pops.coords_grid(depth, homogeneous=False)\n\n if valid_mask:\n coords1, vmask = self.transform(\n depth, intrinsics, valid_mask=valid_mask)\n return coords1 - coords0, vmask\n coords1 = self.transform(depth, intrinsics, valid_mask=valid_mask)\n return coords1 - coords0\n\n def depth_change(self, depth, intrinsics):\n pt = pops.backproject(depth, intrinsics)\n pt_new = self.__call__(pt)\n return pt_new[..., -1] - pt[..., -1]\n \n def identity(self):\n \"\"\" Push identity transformation to start of collection \"\"\"\n batch, frames = self.shape()\n if self.internal == 'matrix':\n # I = torch.eye(4, batch_shape=[batch, 1])\n I = torch.eye(4, dtype=self.G.dtype, device=self.G.device).repeat(\n [batch, 1, 1, 1])\n # return self.__class__(matrix=I, internal=self.internal, eq=self.eq)\n return self.__class__(matrix=I, internal=self.internal, eq=self.eq)\n else:\n raise NotImplementedError\n\n\n\n\nclass SE3Sequence(SE3):\n \"\"\" Stores collection of SE3 objects \"\"\"\n\n def __init__(self, upsilon=None, matrix=None, so3=None, translation=None, eq= \"aijk,ai...k->ai...j\",internal=DEFAULT_INTERNAL):\n super().__init__(\n upsilon, matrix, so3, translation, internal=internal, eq=eq)\n\n # self.eq = \"aijk,ai...k->ai...j\"\n def __call__(self, pt, inds=None, jacobian=False):\n if self.internal == 'matrix':\n return super().__call__(pt, jacobian=jacobian)\n else:\n raise NotImplementedError\n\n\n def gather(self, inds):\n if self.internal == 'matrix':\n G = torch.index_select(self.G, index=inds, dim=1)\n return SE3Sequence(matrix=G, internal=self.internal)\n else:\n raise NotImplementedError\n\n # def append_identity(self):\n # \"\"\" Push identity transformation to start of collection \"\"\"\n # batch, frames = self.shape()\n # if self.internal == 'matrix':\n # # I = torch.eye(4, batch_shape=[batch, 1])\n # I = torch.eye(4, dtype=self.G.dtype, device=self.G.device).repeat(\n # [batch, 1, 1, 1])\n\n # G = torch.cat([I, self.G], dim=1)\n # return SE3Sequence(matrix=G, internal=self.internal)\n # else:\n # raise NotImplementedError\n\n def reprojction_optim(self,\n target,\n weight,\n depth,\n intrinsics,\n num_iters=2,\n depth_img_coords=None\n ):\n\n target = clip_dangerous_gradients(target).to(dtype=torch.float64)\n weight = clip_dangerous_gradients(weight).to(dtype=torch.float64)\n\n X0 = pops.backproject(depth, intrinsics, depth_coords=depth_img_coords)\n w = weight[..., None] \n\n lm_lmbda = get_cfg(\"LM\").LM_LMBDA\n ep_lmbda = get_cfg(\"LM\").EP_LMBDA\n\n T = self.copy(stop_gradients=False)\n for i in range(num_iters):\n ### compute the jacobians of the transformation ###\n X1, jtran = T(X0, jacobian=True)\n x1, (jproj, jkvec) = pops.project(X1, intrinsics, jacobian=True)\n\n v = (X0[..., -1] > MIN_DEPTH) & (X1[..., -1] > MIN_DEPTH)\n # v = v.to(dtype=torch.float32)[..., None, None]\n v = v.to(dtype=torch.float64)[..., None, None]\n\n ### weighted gauss-newton update ###\n J = einsum('...ij,...jk->...ik', jproj.to(dtype=torch.float64), jtran.to(dtype=torch.float64 )) \n\n H = einsum('ai...j,ai...k->aijk', v*w*J, J)\n b = einsum('ai...j,ai...->aij', v*w*J, target-x1)\n\n ### add dampening and apply increment ###\n H += ep_lmbda*torch.eye(6, dtype=H.dtype, device=H.device) + lm_lmbda*H*torch.eye(6,dtype=H.dtype, device=H.device)\n try:\n delta_upsilon = cholesky_solve(H, b)\n except:\n # print(w.shape,v.shape, w.mean(), v.mean(),H,b, '!!!!')\n raise\n T = T.increment(delta_upsilon)\n\n # update\n if self.internal == 'matrix':\n self.G = T.matrix()\n T = SE3Sequence(\n matrix=T.matrix(), internal=self.internal)\n else:\n raise NotImplementedError\n\n return T\n\n\n def transform(self, depth, intrinsics, valid_mask=False, return3d=False):\n return super().transform(depth, intrinsics, valid_mask, return3d)\n"
},
{
"path": "model/CFNet.py",
"content": "\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom thirdparty.raft.update import BasicUpdateBlock\nfrom thirdparty.raft.extractor import BasicEncoder\nfrom thirdparty.raft.corr import CorrBlock, AlternateCorrBlock\nfrom thirdparty.raft.utils.utils import bilinear_sampler, coords_grid, upflow\n\ntry:\n autocast = torch.cuda.amp.autocast\nexcept:\n # dummy autocast for PyTorch < 1.6\n class autocast:\n def __init__(self, enabled):\n pass\n def __enter__(self):\n pass\n def __exit__(self, *args):\n pass\n\n\nclass ImageFeaEncoder(nn.Module):\n def __init__(self, input_dim=3, output_dim=256):\n super().__init__()\n self.fnet = BasicEncoder(output_dim=output_dim, norm_fn='instance', dropout=False, input_dim=input_dim) \n\n if 1:#self.args.pretrained_model is not None:\n print(\"Loading the weights of RAFT...\")\n import os \n self.load_state_dict(\n # torch.load(self.args.pretrained_model, map_location='cpu'), strict=False\n torch.load( f\"{os.path.dirname(os.path.abspath(__file__)) }/../weights/img_fea_enc.pth\", map_location='cpu'), strict=True\n )\n else:\n print(\"ImageFeaEncoder will be trained from scratch...\")\n\n def forward(self, image1, image2):\n image1 = 2 * (image1 / 255.0) - 1.0\n image2 = 2 * (image2 / 255.0) - 1.0\n\n image1 = image1.contiguous()\n image2 = image2.contiguous()\n with autocast(enabled=True):\n fmap1, fmap2 = self.fnet([image1, image2])\n return fmap1, fmap2\n\n\nclass GRU_CFUpdator(nn.Module):\n def __init__(self, args):\n super().__init__()\n self.args = args\n input_dim = args.get(\"input_dim\", 3)\n\n self.hidden_dim = hdim = 128\n self.context_dim = cdim = 128\n args.corr_levels = 4\n args.corr_radius = 4\n\n if 'alternate_corr' not in self.args:\n self.args.alternate_corr = False\n\n self.update_block = BasicUpdateBlock(self.args, hidden_dim=hdim)\n\n if self.args.pretrained_model is not None:\n print(\"Loading the weights of RAFT...\")\n import os \n self.load_state_dict(\n # torch.load(self.args.pretrained_model, map_location='cpu'), strict=False\n torch.load( f\"{os.path.dirname(os.path.abspath(__file__)) }/../weights/gru_update.pth\", map_location='cpu'), strict=True\n )\n else:\n print(\"GRU_CFUpdator will be trained from scratch...\")\n\n\n \n \n def freeze_bn(self):\n for m in self.modules():\n if isinstance(m, nn.BatchNorm2d):\n m.eval()\n\n def initialize_flow(self, img, downsample_rate=8):\n \"\"\" Flow is represented as difference between two coordinate grids flow = coords1 - coords0\"\"\"\n N, C, H, W = img.shape\n coords0 = coords_grid(N, H//downsample_rate, W//downsample_rate).to(img.device)\n coords1 = coords_grid(N, H//downsample_rate, W//downsample_rate).to(img.device)\n\n # optical flow computed as difference: flow = coords1 - coords0\n return coords0, coords1\n\n def upsample_flow(self, flow, mask, upsample_scale=8):\n \"\"\" Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination \"\"\"\n N, _, H, W = flow.shape\n mask = mask.view(N, 1, 9, upsample_scale, upsample_scale, H, W)\n mask = torch.softmax(mask, dim=2)\n\n up_flow = F.unfold(upsample_scale * flow, [3,3], padding=1)\n up_flow = up_flow.view(N, 2, 9, 1, 1, H, W)\n\n up_flow = torch.sum(mask * up_flow, dim=2)\n up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)\n return up_flow.reshape(N, 2, upsample_scale*H, upsample_scale*W)\n\n\n def forward(self, fmap1, fmap2, iters=1, flow_init=None, upsample=True, test_mode=False, context_fea=None, update_corr_fn=True):\n \"\"\" Estimate optical flow between pair of frames \"\"\"\n\n hdim = self.hidden_dim\n cdim = self.context_dim\n\n if update_corr_fn: # need carful handling outside\n # run the feature network\n self.fmap1 = fmap1.float()\n self.fmap2 = fmap2.float()\n if self.args.alternate_corr:\n self.corr_fn = AlternateCorrBlock(self.fmap1, self.fmap2, radius=self.args.corr_radius)\n else:\n self.corr_fn = CorrBlock(self.fmap1, self.fmap2, radius=self.args.corr_radius)\n\n if update_corr_fn: \n # run the context network\n with autocast(enabled=self.args.mixed_precision):\n assert context_fea is not None\n ds = context_fea.shape[-1]//self.fmap1.shape[-1]\n cnet = F.interpolate(context_fea, scale_factor=1/ds, mode='bilinear', align_corners=True)\n\n self.net, self.inp = torch.split(cnet, [hdim, cdim], dim=1)\n self.net = torch.tanh(self.net)\n self.inp = torch.relu(self.inp)\n\n # coords0, coords1 = self.initialize_flow(image1)\n coords0, coords1 = self.initialize_flow(flow_init)\n\n if flow_init is not None:\n ds = flow_init.shape[-1]//coords0.shape[-1]\n if ds !=1:\n flow_init /=ds\n flow_init = F.interpolate(flow_init, scale_factor=1/ds, mode='bilinear', align_corners=True)\n\n coords1 = coords1 + flow_init\n\n flow_predictions = []\n for itr in range(iters):\n coords1 = coords1.detach()\n corr = self.corr_fn(coords1) # index correlation volume\n\n flow = coords1 - coords0\n with autocast(enabled=self.args.mixed_precision):\n # net, up_mask, delta_flow = self.update_block(net, inp, corr, flow)\n self.net, up_mask, delta_flow = self.update_block(self.net, self.inp, corr, flow)\n\n # F(t+1) = F(t) + \\Delta(t)\n coords1 = coords1 + delta_flow\n\n # upsample predictions\n if up_mask is None:\n flow_up = upflow(coords1 - coords0, scale=image1.shape[2]//coords0.shape[2],)\n else:\n if self.args.fea_net in [\"bigdx4\"]:\n flow_up = self.upsample_flow(coords1 - coords0, up_mask, upsample_scale=4)\n else:\n flow_up = self.upsample_flow(coords1 - coords0, up_mask)\n \n flow_predictions.append(flow_up)\n\n if test_mode:\n return coords1 - coords0, flow_up\n \n return flow_predictions\n\n"
},
{
"path": "model/HybridNet.py",
"content": "\nimport torch \nimport torch.nn as nn \n\nfrom thirdparty.kpconv.kpconv_blocks import *\nimport torch.nn.functional as F\nimport numpy as np\nfrom kpconv.lib.utils import square_distance\nfrom model.descriptor2D import SuperPoint2D\nfrom model.descriptor3D import KPSuperpoint3Dv2\n\n\n\nREGISTERED_HYBRID_NET_CLASSES={}\ndef register_hybrid_net(cls, name=None):\n global REGISTERED_HYBRID_NET_CLASSES\n if name is None:\n name = cls.__name__\n assert name not in REGISTERED_HYBRID_NET_CLASSES, f\"exist class: {REGISTERED_HYBRID_NET_CLASSES}\"\n REGISTERED_HYBRID_NET_CLASSES[name] = cls\n return cls\n\n\ndef get_hybrid_net(name):\n global REGISTERED_HYBRID_NET_CLASSES\n assert name in REGISTERED_HYBRID_NET_CLASSES, f\"available class: {REGISTERED_HYBRID_NET_CLASSES}\"\n return REGISTERED_HYBRID_NET_CLASSES[name]\n\nclass ContextFeatureNet(nn.Module):\n def __init__(self, config):\n super().__init__()\n self.context_fea_extractor_3d= KPSuperpoint3Dv2(config['context_fea_extractor_3d'] )\n \n def forward(self, batch):\n # x = batch['features'].clone().detach()\n # assert len(batch['stack_lengths'][-1])==1, \"Only support bs=1 for now\" \n len_src_c = batch['stack_lengths'][-1][0]\n pcd_c = batch['model_points'][-1]\n pcd_c = pcd_c[:len_src_c]\n\n image=batch['image']\n\n ############### encode 3d and 2d features ###############\n batch3d={\n 'points': batch['model_points'], \n 'neighbors': batch['neighbors'], \n 'pools': batch['pools'], \n 'upsamples': batch['upsamples'],\n 'features': batch['model_point_features'], \n 'stack_lengths': batch['stack_lengths'],\n }\n ctx_descriptors_3d = self.context_fea_extractor_3d(batch3d)\n\n\n return {\n \"ctx_fea_3d\":ctx_descriptors_3d,\n }\n\n\n\n@register_hybrid_net\nclass HybridDescNet(nn.Module):\n #independent 2d and 3d network\n def __init__(self, config):\n super().__init__()\n\n self.corr_fea_extractor_2d= SuperPoint2D(config['keypoints_detector_2d'] )\n self.corr_fea_extractor_3d= KPSuperpoint3Dv2(config['keypoints_detector_3d'] )\n self.descriptors_3d = {}\n\n\n def forward(self, batch):\n assert len(set(batch['class_name']))==1, \"A batch should contain data of the same class.\"\n class_name = batch['class_name'][0]\n\n len_src_c = batch['stack_lengths'][-1][0]\n pcd_c = batch['model_points'][-1]\n pcd_c = pcd_c[:len_src_c]#, pcd_c[len_src_c:]\n\n image=batch['image']\n\n ############### encode 3d and 2d features ###############\n batch3d={\n 'points': batch['model_points'], \n 'neighbors': batch['neighbors'], \n 'pools': batch['pools'], \n 'upsamples': batch['upsamples'],\n 'features': batch['model_point_features'], \n 'stack_lengths': batch['stack_lengths'],\n }\n if self.training:\n self.descriptors_3d[class_name] = self.corr_fea_extractor_3d(batch3d)\n else:\n if class_name not in self.descriptors_3d:\n self.descriptors_3d[class_name] = self.corr_fea_extractor_3d(batch3d)\n\n descriptors_2d = self.corr_fea_extractor_2d(image)['descriptors']\n\n\n return {\n \"descriptors_2d\":descriptors_2d,\n \"descriptors_3d\":self.descriptors_3d[class_name],\n \"scores_saliency_3d\": None, \n \"scores_overlap_3d\":None, \n\n }\n"
},
{
"path": "model/PoseRefiner.py",
"content": "import os \nimport time\nimport cv2\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.distributed as dist\nimport numpy as np\nfrom easydict import EasyDict as edict\nfrom functools import partial\n\nfrom geometry.transformation import *\nfrom geometry.intrinsics import *\nfrom geometry.projective_ops import coords_grid, normalize_coords_grid\nfrom model.CFNet import GRU_CFUpdator , ImageFeaEncoder\nfrom utils.pose_utils import pose_padding\nfrom config.default import get_cfg\n\n\n\nEPS = 1e-5\nMIN_DEPTH = 0.1\nMAX_ERROR = 100.0\n\n# exclude extremly large displacements\nMAX_FLOW = 400\n\n\ndef raft_sequence_flow_loss(flow_preds, flow_gt, valid, gamma=0.8, max_flow=MAX_FLOW):\n \"\"\" Loss function defined over sequence of flow predictions \"\"\"\n\n n_predictions = len(flow_preds) \n flow_loss = 0.0\n \n # exlude invalid pixels and extremely large diplacements\n mag = torch.sum(flow_gt**2, dim=1).sqrt()\n \n valid = (valid >= 0.5) & (mag < max_flow)\n\n for i in range(n_predictions):\n i_weight = gamma**(n_predictions - i - 1)\n i_loss = (flow_preds[i] - flow_gt).abs()\n flow_loss += i_weight * (valid[:, None] * i_loss).mean()\n\n epe = torch.sum((flow_preds[-1] - flow_gt)**2, dim=1).sqrt()\n epe = epe.view(-1)[valid.view(-1)]\n\n metrics = {\n 'epe': epe.mean().item(),\n '1px': (epe < 1).float().mean().item(),\n '3px': (epe < 3).float().mean().item(),\n '5px': (epe < 5).float().mean().item(),\n }\n\n return flow_loss, metrics\n\n\n\n\nclass PoseRefiner(nn.Module):\n def __init__(self, cfg,\n reuse=False,\n schedule=None,\n use_regressor=True,\n is_calibrated=True,\n bn_is_training=False,\n is_training=True,\n renderer=None,\n ):\n\n super().__init__()\n\n self.legacy=True\n self.cfg = cfg\n self.reuse = reuse\n self.sigma=nn.ParameterList( [nn.Parameter(torch.ones(1)*1 )] )\n self.with_corr_weight = self.cfg.get(\"with_corr_weight\", True)\n if not self.with_corr_weight:\n print(\"Warning: the correlation weighting is disabled.\")\n\n self.is_calibrated = cfg.IS_CALIBRATED\n if not is_calibrated:\n self.is_calibrated = is_calibrated\n\n self.is_training = is_training\n self.use_regressor = use_regressor\n\n self.residual_pose_history = []\n self.Ti_history =[]\n self.coords_history = []\n self.residual_history = []\n self.inds_history = []\n # self.weights_history = []\n self.flow_history = []\n self.intrinsics_history = []\n self.summaries = []\n\n if self.cfg.FLOW_NET=='raft':\n self.image_fea_enc= ImageFeaEncoder()\n self.cf_net = GRU_CFUpdator(self.cfg.raft) \n else:\n raise NotImplementedError \n self.renderer = renderer \n\n def _clear(self,):\n self.residual_pose_history = []\n self.Ti_history =[]\n self.coords_history = []\n self.residual_history = []\n self.inds_history = []\n # self.weights_history = []\n self.flow_history = []\n self.intrinsics_history = []\n self.summaries = []\n\n def __len__(self):\n return len(self.residual_pose_history)\n\n def render(self, params, render_tex=False):\n \"\"\" render a batch of images given the intrinsic and extrinsic params \n\n Args:\n params.K: np.array, of shape Bx3x3\n params.camera_extrinsics: np.array, of shape Bx3x4\n\n Returns:\n [type]: [description]\n \"\"\"\n\n bs=params.K.shape[0]\n colors=[]\n depths=[]\n\n color,depth= self.renderer( params.obj_cls, params.vert_attribute, T=params.camera_extrinsics, K=params.K, \n render_image_size=params.render_image_size, near=0.1, far=6, render_tex=render_tex)\n #set the invalid values to zeros\n depth[depth==-1] = 0\n return {\n # 1x3xHxW\n \"syn_img\": color, \n \"syn_depth\": depth.detach(), # 1x1xHxW\n }\n\n\n def get_affine_transformation(self, mask, crop_center, with_intrinsic_transform=False, output_size=None, margin_ratio=0.4):\n B,_,H,W = mask.shape\n ratio = float(H) / float(W)\n affine_matrices = []\n intrinsic_matrices = []\n\n for b in range(B):\n zoom_c_x, zoom_c_y = crop_center[b] #crop_center\n\n ys, xs = np.nonzero(mask[b][0].detach().cpu().numpy() )\n if len(ys)>0 and len(xs)>0:\n obj_imgn_start_x = xs.min() \n obj_imgn_start_y = ys.min()\n obj_imgn_end_x = xs.max()\n obj_imgn_end_y = ys.max()\n else:\n obj_imgn_start_x=0\n obj_imgn_start_y=0\n obj_imgn_end_x=0\n obj_imgn_end_y=0\n\n\n # mask region\n left_dist = zoom_c_x - obj_imgn_start_x\n right_dist = obj_imgn_end_x - zoom_c_x\n up_dist = zoom_c_y - obj_imgn_start_y\n down_dist = obj_imgn_end_y - zoom_c_y\n # crop_height = np.max([ratio * right_dist, ratio * left_dist, up_dist, down_dist]) * 2 * 1.4\n crop_height = np.max([ratio * right_dist, ratio * left_dist, up_dist, down_dist]) * 2 * (1+margin_ratio)\n crop_width = crop_height / ratio\n\n # affine transformation for PyTorch\n x1 = (zoom_c_x - crop_width / 2) * 2 / W - 1;\n x2 = (zoom_c_x + crop_width / 2) * 2 / W - 1;\n y1 = (zoom_c_y - crop_height / 2) * 2 / H - 1;\n y2 = (zoom_c_y + crop_height / 2) * 2 / H - 1;\n\n pts1 = np.float32([[x1, y1], [x1, y2], [x2, y1]])\n pts2 = np.float32([[-1, -1], [-1, 1], [1, -1]])\n affine_matrix = torch.tensor(cv2.getAffineTransform(pts2, pts1), device=mask.device, dtype=torch.float32)\n affine_matrices.append(affine_matrix)\n\n\n if with_intrinsic_transform:\n # affine transformation for PyTorch\n x1 = (zoom_c_x - crop_width / 2)\n x2 = (zoom_c_x + crop_width / 2)\n y1 = (zoom_c_y - crop_height / 2)\n y2 = (zoom_c_y + crop_height / 2)\n\n pts1 = np.float32([[x1, y1], [x1, y2], [x2, y1]])\n # pts2 = np.float32([[0, 0], [0, H-1], [W-1, 0]])\n pts2 = np.float32([[0, 0], [0, output_size[0]-1], [output_size[1]-1, 0]])\n # pts2 = np.float32([[0, 0], [0, 1], [1, 0]])\n intrinsic_matrix = torch.tensor(cv2.getAffineTransform(pts2, pts1), device=mask.device, dtype=torch.float32)\n intrinsic_matrices.append(intrinsic_matrix)\n \n if with_intrinsic_transform:\n return torch.stack(affine_matrices, dim=0), torch.stack(intrinsic_matrices, dim=0)\n else:\n return torch.stack(affine_matrices, dim=0)\n\n def gen_zoom_crop_grids(self, fg_mask, K, T, output_size, model_center=[0,0,0], margin_ratio=0.4):\n ##Get the projected model center in image (assuming the model is zero-centered, which should be reconsidered!) \n crop_center=K@T[:,:3,3:]\n crop_center = crop_center[:,:2]/crop_center[:,2:3]\n\n ##calculate affine transformation parameters\n affine_matrices, crop_intrinsic_transform=self.get_affine_transformation(fg_mask, crop_center=crop_center.detach().cpu().numpy(), with_intrinsic_transform=True, output_size=(output_size[-2], output_size[-1]),margin_ratio=margin_ratio ) \n grids = F.affine_grid(affine_matrices, torch.Size(output_size) )\n ##Get cropped intrinsic_transform\n intrinsics_crop= torch.inverse(pose_padding(crop_intrinsic_transform) )@K\n\n return grids, intrinsics_crop\n\n # def forward(self, image, Ts, intrinsics, fea_3d=None, inds=None, Tj_gt=None, obj_cls=None, geofea_3d=None, geofea_2d=None):\n def forward(self, image, Ts, intrinsics, fea_3d=None, Tj_gt=None, obj_cls=None, geofea_3d=None, geofea_2d=None):\n #clear the history data\n self._clear()\n cfg = self.cfg\n\n RANDER_IMAGE_SIZE = get_cfg(\"BASIC\").render_image_size\n ZOOM_CROP_SIZE=get_cfg(\"BASIC\").zoom_crop_size \n\n if cfg.RESCALE_IMAGES:\n images = 2 * (images / 255.0) - 1.0\n\n Tij_gt=[]\n syn_imgs=[]\n syn_depths=[]\n\n Ti = Ts\n Tij = Ti.copy().identity()\n\n for ren_iter in range(cfg.RENDER_ITER_COUNT):\n # update rendering params\n Ti = Tij*Ti # accumulate Ti \n Tij.identity_() #set Tij to identity matrix at the begining of each ren_iter\n if self.legacy:\n Tij = Ti*Ti.inv() #set Tij to identity matrix at the begining of each ren_iter\n\n render_params = edict({\n \"K\": intrinsics.detach(), \n \"camera_extrinsics\": Ti.matrix().detach().squeeze(1), \n \"obj_cls\": obj_cls,\n \"render_image_size\": RANDER_IMAGE_SIZE, \n })\n\n pc_depth = self.renderer.render_pointcloud(obj_cls, T=render_params.camera_extrinsics, K=render_params.K, \n render_image_size=render_params.render_image_size)\n\n\n if self.cfg.ONLINE_CROP:\n #get the forground mask \n fg_mask = pc_depth>0\n B,C,_,_= pc_depth.size()\n \n ############### Get zoom parameters ###############\n grids, intrinsics_crop = self.gen_zoom_crop_grids(fg_mask, render_params.K, render_params.camera_extrinsics, output_size=[B,C, *ZOOM_CROP_SIZE], model_center=None)\n\n ############### Render reference images ###############\n # Concatentate the 3D ctx feature \"fea_3d\" and 3d descriptor \"geofea_3d\" for feature rendering\n if geofea_3d is not None:\n fea_3d_cat = torch.cat([fea_3d, geofea_3d ], dim=-1) # BxNxC\n else:\n fea_3d_cat = fea_3d\n render_params.vert_attribute = fea_3d_cat #fea_3d\n render_params.K = intrinsics_crop.detach()\n render_params.render_image_size = ZOOM_CROP_SIZE\n ren_res = self.render(render_params, render_tex=True) \n if geofea_3d is not None:\n syn_img, cfea, geofea1 = torch.split(ren_res['syn_img'],[3, fea_3d.shape[-1], geofea_3d.shape[-1] ] ,dim=1)\n syn_depth = ren_res['syn_depth']\n \n else:\n syn_img, cfea = torch.split(ren_res['syn_img'], [3, fea_3d.shape[-1] ] ,dim=1)\n geofea1=None\n syn_depth = ren_res['syn_depth']\n cfea = cfea*0.1 # balance the learning rate with the scale 0.1\n\n ## Crop and zoom images\n syn_image_crop = syn_img \n image_crop= F.grid_sample(image, grids)\n cfea_crop= cfea \n if geofea1 is not None and geofea_2d is not None:\n # geofea1_crop = F.grid_sample(geofea1, grids)\n geofea1_crop = geofea1\n geofea2_crop = F.grid_sample(geofea_2d, grids)\n\n #Render again to get more accurate depth for supervisions in losses \n #TODO: could be merged into the rendering process above. -> Done\n if self.legacy:#self.training:\n depth_render_params=edict({\n \"K\": intrinsics_crop.detach(), \n \"camera_extrinsics\": Ti.matrix().detach().squeeze(1), \n \"obj_cls\": obj_cls,\n \"render_image_size\": ZOOM_CROP_SIZE, \n })\n syn_depth=self.renderer.render_depth(obj_cls, T=depth_render_params.camera_extrinsics, K=depth_render_params.K, \n render_image_size=depth_render_params.render_image_size, near=0.1, far=6)\n\n #for visualization only\n syn_imgs.append(syn_image_crop)\n syn_imgs.append(image_crop)\n\n # encode image features\n feats1, feats2=self.image_fea_enc(syn_image_crop, image_crop)\n # depths = torch.index_select(syn_depth, index=ii, dim=1) + EPS\n depths = syn_depth+EPS\n\n for i in range(cfg.ITER_COUNT):\n # save for loss calculation \n self.intrinsics_history.append(intrinsics_crop)\n syn_depths.append(syn_depth)\n\n Tij = Tij.copy(stop_gradients=True)\n intrinsics_crop = intrinsics_crop.detach()\n \n #Get the projection in frame j of visible model points in frame i with the current relative pose estimation Tij\n reproj_coords, vmask = Tij.transform(\n depths, intrinsics_crop, valid_mask=True)\n\n uniform_grids = coords_grid(depths)\n flow_init = torch.einsum( \"...ijk->...kij\", reproj_coords-uniform_grids[..., :2] ) * (depths>EPS) \n flow = self.cf_net(feats1, feats2, flow_init=flow_init.squeeze(1), context_fea=cfea_crop, update_corr_fn=i==0)\n\n\n self.flow_history.append(flow)\n\n # Get the correspondences in frame j for each point in frame i, based on the current flow estimates\n if isinstance (flow, (list, tuple)): # flow net may return a list of flow maps\n correspondence_target = torch.einsum(\"...ijk->...jki\", flow[-1]) + uniform_grids[..., :2]\n else:\n correspondence_target = torch.einsum(\"...ijk->...jki\", flow) + uniform_grids[..., :2]\n \n # Optimize for the pose by minimizing errors between the constructed correspondence field \n # (with the currently estimated pose) and the estimated correspondence field \n if self.with_corr_weight and geofea1 is not None and geofea_2d is not None:\n geofea2_crop_warpped = F.grid_sample(geofea2_crop, normalize_coords_grid(correspondence_target).squeeze(1) )\n corr_weight= torch.sum(geofea1_crop*geofea2_crop_warpped, dim=1,keepdim=True).permute(0,2,3,1)[:,None] #insert frame axis\n corr_weight = torch.exp(-torch.abs(1-corr_weight)/self.sigma[0]) * (syn_depth>0)[...,None].float()\n else: \n corr_weight = weight\n \n Tij = Tij.reprojction_optim(\n correspondence_target, corr_weight, depths, intrinsics_crop, num_iters=cfg.OPTIM_ITER_COUNT )\n\n reproj_coords, vmask1 = Tij.transform(\n depths, intrinsics_crop, valid_mask=True)\n\n # For the loss calculation later\n self.residual_pose_history.append(Tij)\n\n self.Ti_history.append(Ti.copy(stop_gradients=True) )\n Tij_gt.append( (Tj_gt*Ti.inv()).copy(stop_gradients=True) )\n\n self.residual_history.append(\n vmask*vmask1*(reproj_coords-correspondence_target)) # BxKxHxWx3\n\n # The final update of Ti \n Ti = Tij*Ti\n return {\n \"Tij\": Tij,\n \"Ti_pred\": Ti,\n \"intrinsics\": intrinsics,\n \"flow\": self.flow_history[0],\n \"vmask\": syn_depth>0, \n \"weight\": torch.einsum(\"...ijk->...kij\", corr_weight), \n \"syn_depth\": syn_depths, #ren_res['syn_depth'],\n \"syn_img\": syn_imgs+[image_crop, cfea_crop[:,:3]*10,geofea1[:,:3], geofea2_crop[:,:3]],\n \"Tij_gt\" : Tij_gt\n }\n\n def compute_loss(self, Tij_gts, depths, intrinsics, loss='l1', log_error=True, loss3d=None, ):\n\n total_loss = 0.0\n for i in range(len(self.residual_pose_history)):\n intrinsics = self.intrinsics_history[i]\n\n depth, intrinsics = rescale_depths_and_intrinsics(depths[i], intrinsics, downscale=1)\n\n Tij = self.residual_pose_history[i]\n\n Gij = Tij_gts[i] \n\n # intrinsics_pred = intrinsics\n zstar = depth + EPS\n flow_pred, valid_mask_pred = Tij.induced_flow(\n zstar, intrinsics, valid_mask=True) \n flow_star, valid_mask_star = Gij.induced_flow(\n zstar, intrinsics, valid_mask=True)\n\n valid_mask = valid_mask_pred * valid_mask_star\n\n #3D alignment loss \n loss_3d_proj = 0 \n if loss3d is not None:\n Tj_pred=Tij*self.Ti_history[i]\n Tj_gt=Gij*self.Ti_history[i]\n loss_3d_proj = loss3d(\n R_pred=Tj_pred.G[:, 0, :3, :3], t_pred=Tj_pred.G[:, 0, :3, 3], R_tgt=Tj_gt.G[:, 0, :3, :3], t_tgt=Tj_gt.G[:, 0, :3, 3])\n\n # flow loss\n if isinstance( self.flow_history[0], (list, tuple)):\n #squeeze the frame dimmension\n self.flow_history[i] = [self.flow_history[i][f].squeeze(1) for f in range(len(self.flow_history[i])) ]\n flow_mask = valid_mask.squeeze(1).squeeze(-1)\n loss_flow,_ = raft_sequence_flow_loss(self.flow_history[i], flow_gt=torch.einsum(\"...ijk->...kij\", flow_star.squeeze(1)), valid= flow_mask, gamma=0.8, max_flow=MAX_FLOW)\n else:\n raise NotImplementedError\n \n # reprojection loss \n reproj_diff = valid_mask * \\\n torch.clamp(\n torch.abs(flow_pred - flow_star), -MAX_ERROR, MAX_ERROR)\n reproj_loss = torch.mean(reproj_diff)\n\n\n total_loss +=self.cfg.get(\"TRAIN_PCALIGN_WEIGHT\", 1)*loss_3d_proj+ self.cfg.TRAIN_FLOW_WEIGHT* loss_flow + self.cfg.TRAIN_REPROJ_WEIGHT*reproj_loss\n \n # clear the intermediate values\n self._clear()\n\n return {\n \"total_loss\": total_loss,\n \"reproj_loss\": reproj_loss,\n \"flow_loss\":loss_flow,\n \"loss_3d_proj\": loss_3d_proj,\n \"valid_mask\": valid_mask,\n \"Tij\": Tj_pred.G,\n \"Gij\": Tj_gt.G\n }\n"
},
{
"path": "model/RNNPose.py",
"content": "#\nimport time\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\nimport torch.distributed as dist\nimport apex\nimport numpy as np\nimport os\nfrom easydict import EasyDict as edict\nfrom transforms3d.euler import mat2euler, euler2mat, euler2quat, quat2euler\nfrom functools import partial\n\n\n\nfrom model.HybridNet import HybridDescNet,ContextFeatureNet\nfrom thirdparty.kpconv.lib.utils import square_distance\nfrom model.PoseRefiner import PoseRefiner \n\nfrom utils.pose_utils import pose_padding\nfrom geometry.transformation import SE3Sequence\n\nfrom geometry.diff_render_optim import DiffRendererWrapper\nfrom config.default import get_cfg\nfrom utils.util import dict_recursive_op\n\n# from model.RNNPose import register_posenet\n\nREGISTERED_NETWORK_CLASSES = {}\n\n\ndef register_posenet(cls, name=None):\n global REGISTERED_NETWORK_CLASSES\n if name is None:\n name = cls.__name__\n assert name not in REGISTERED_NETWORK_CLASSES, f\"exist class: {REGISTERED_NETWORK_CLASSES}\"\n REGISTERED_NETWORK_CLASSES[name] = cls\n return cls\n\n\ndef get_posenet_class(name):\n global REGISTERED_NETWORK_CLASSES\n assert name in REGISTERED_NETWORK_CLASSES, f\"available class: {REGISTERED_NETWORK_CLASSES}\"\n return REGISTERED_NETWORK_CLASSES[name]\n\n\n\n\n@register_posenet\nclass RNNPose(nn.Module):\n def __init__(self,\n criterions,\n opt,\n name=\"RNNPose\",\n **kwargs):\n super().__init__()\n\n self.name = name\n self.opt = opt\n\n self.hybrid_desc_net=HybridDescNet(opt.descriptor_net)\n self.ctx_fea_net = ContextFeatureNet(opt.descriptor_net)\n\n self.ctx_fea = {} \n\n self.render_params = edict({\n \"width\": opt.input_w, # 128,\n \"height\": opt.input_h, # 128,\n \"gpu_id\": opt.get('gpu_id', 0),\n \"obj_seqs\": opt.obj_seqs\n })\n\n renderer, diff_renderer= self._render_init(self.render_params)\n self.diff_renderer = diff_renderer\n\n self.motion_net = PoseRefiner(\n opt.motion_net, bn_is_training=self.training, is_training=self.training,\n renderer=diff_renderer \n )\n self.contrastive_loss = criterions.get(\n \"metric_loss\", None)\n self.pose_loss = criterions.get(\"pose_loss\", None)\n\n self.register_buffer(\"global_step\", torch.LongTensor(1).zero_())\n\n\n def update_global_step(self):\n self.global_step += 1\n\n def get_global_step(self):\n return int(self.global_step.cpu().numpy()[0])\n\n def clear_global_step(self):\n self.global_step.zero_()\n \n def sample_poses(self, pose_tgt):\n SYN_STD_ROTATION = 15\n SYN_STD_TRANSLATION = 0.01\n ANGLE_MAX=45\n pose_src = pose_tgt.copy()\n num = pose_tgt.shape[0]\n for i in range(num):\n euler = mat2euler(pose_tgt[i, :3, :3])\n euler += SYN_STD_ROTATION * np.random.randn(3) * np.pi / 180.0\n pose_src[i, :3, :3] = euler2mat(euler[0], euler[1], euler[2])\n\n pose_src[i, 0, 3] = pose_tgt[i, 0, 3]+ SYN_STD_TRANSLATION * np.random.randn(1)\n pose_src[i, 1, 3] = pose_tgt[i, 1, 3] + SYN_STD_TRANSLATION * np.random.randn(1)\n pose_src[i, 2, 3] = pose_tgt[i, 2, 3] + 5 * SYN_STD_TRANSLATION * np.random.randn(1)\n\n r_dist = np.arccos((np.trace(pose_src[i, :3,:3].transpose(-1,-2) @ pose_tgt[i, :3,:3]) - 1 )/2)/np.pi*180\n\n while r_dist > ANGLE_MAX:#or not (16 < center_x < (640 - 16) and 16 < center_y < (480 - 16)):\n print(\"r_dist > ANGLE_MAX, resampling...\")\n euler = mat2euler(pose_tgt[i, :3, :3])\n euler += SYN_STD_ROTATION * np.random.randn(3) * np.pi / 180.0\n pose_src[i, :3, :3] = euler2mat(euler[0], euler[1], euler[2])\n\n pose_src[i, 0, 3] = pose_tgt[i, 0, 3]+ SYN_STD_TRANSLATION * np.random.randn(1)\n pose_src[i, 1, 3] = pose_tgt[i, 1, 3] + SYN_STD_TRANSLATION * np.random.randn(1)\n pose_src[i, 2, 3] = pose_tgt[i, 2, 3] + 5 * SYN_STD_TRANSLATION * np.random.randn(1)\n\n r_dist = np.arccos((np.trace(pose_src[i, :3,:3].transpose(-1,-2) @ pose_tgt[i, :3,:3]) - 1 )/2)*np.pi/180\n return pose_src\n\n def _render_init(self, config):\n # from data.ycb.basic import bop_ycb_class2idx\n print(\"config.gpu_id:\", config.gpu_id)\n\n obj_paths = []\n tex_paths = []\n\n # build cls2idx table for the renderer\n cls2idx = {}\n LM_SEQ=[\"ape\", \"benchvise\", \"camera\",\"cam\", \"can\", \"cat\", \"driller\", \"duck\", \"eggbox\", \"glue\", \"holepuncher\", \"iron\", \"lamp\", \"phone\"]\n # YCB_SEQ=bop_ycb_class2idx.keys()\n for i, seq in enumerate(set(config.obj_seqs)):\n\n if seq in LM_SEQ: \n obj_path = f'{os.path.dirname(__file__)}/../EXPDATA/LM6d_converted/models/{seq}/textured.obj'\n tex_path = f'{os.path.dirname(__file__)}/../EXPDATA/LM6d_converted/models/{seq}/texture_map.png'\n assert os.path.exists(obj_path), f\"'{obj_path}' dose not exist!\" \n assert os.path.exists(tex_path), f\"'{tex_path}' dose not exist!\" \n obj_paths.append(obj_path)\n tex_paths.append(tex_path)\n cls2idx[seq] = i\n else:\n raise NotImplementedError\n renderer=None\n\n diff_renderer = DiffRendererWrapper(obj_paths)\n diff_renderer.cls2idx = cls2idx\n\n return renderer, diff_renderer\n\n\n def forward(self, sample):\n assert len(set(sample['class_name']))==1, \"A batch should contain data of the same class.\"\n class_name = sample['class_name'][0]\n\n #encode 3d-2d descriptors\n preds_dict=self.hybrid_desc_net(sample)\n\n len_src_f = sample['stack_lengths'][0][0]\n geofea_3d = preds_dict.get('descriptors_3d', None) \n geofea_2d = preds_dict.get('descriptors_2d', None) \n\n\n #encode 3D context features \n if self.training:\n self.ctx_fea[class_name]=self.ctx_fea_net(sample)\n else:\n if class_name not in self.ctx_fea:\n self.ctx_fea[class_name]=self.ctx_fea_net(sample)\n\n preds_dict.update(self.ctx_fea[class_name])\n ctx_fea_3d = preds_dict['ctx_fea_3d'][:len_src_f]\n\n if self.training:\n pose=pose_padding(sample['original_RT'])\n\n if sample.get(\"rendered_RT\", None) is not None:\n syn_pose = pose_padding(sample['rendered_RT'])\n else:\n syn_pose = torch.tensor(self.sample_poses(sample['original_RT'].detach().cpu(\n ).numpy()), device=sample['original_RT'].device, dtype=sample['original_RT'].dtype)\n syn_pose = pose_padding(syn_pose)\n else:\n pose = pose_padding(sample['original_RT'])\n syn_pose = pose_padding(sample['rendered_RT'])\n \n\n # calculate the GT relative pose and the initial pose\n \n Ts_pred = SE3Sequence(\n matrix=torch.stack([syn_pose ], dim=1))\n mot_res = self.motion_net(\n Ts=Ts_pred, \n intrinsics=sample['K'],\n image=sample['image'], \n fea_3d=ctx_fea_3d[None], \n Tj_gt=SE3Sequence(matrix=pose[:, None]),\n obj_cls=sample['class_name'],\n geofea_2d = geofea_2d, \n geofea_3d = geofea_3d[None]\n )\n preds_dict.update(mot_res)\n sample['syn_depth'] = mot_res['syn_depth']\n if self.training:\n ret = self.loss(sample, preds_dict)\n\n ret['syn_img'] = mot_res['syn_img']\n ret['syn_depth'] = mot_res['syn_depth']\n ret['flow'] = mot_res['flow']\n ret['weight'] = mot_res['weight']\n\n return ret\n else:\n # ret = self.loss(sample, preds_dict)\n ret={}\n ret.update(preds_dict)\n return ret\n\n \n def loss(self, sample, preds_dict):\n\n len_src_f = sample['stack_lengths'][0][0]\n RT =sample['RT']\n camera_intrinsic=sample['K']\n descriptors_2d_map = preds_dict['descriptors_2d']\n descriptors_3d = preds_dict['descriptors_3d'][:len_src_f]\n rand_descriptors_3d = preds_dict['descriptors_3d'][len_src_f:]\n model_points=sample['model_points'][0][:len_src_f]\n orig_model_points = sample['original_model_points']\n rand_model_points=sample['model_points'][0][len_src_f:]\n lifted_points=sample['lifted_points'][0].squeeze(0)\n correspondence=sample['correspondences_2d3d'].squeeze(0)\n depth = sample['depth']\n \n\n # get the foreground 2d descriptors \n ys_, xs_=torch.nonzero(sample['depth'].squeeze(), as_tuple=True )\n descriptors_2d=descriptors_2d_map[:,:,ys_,xs_].squeeze().permute([1,0])\n\n if self.training: \n device= lifted_points.device\n fg_point_num = len(lifted_points)\n model_point_num = len(model_points)\n \n # append bg features \n ys_bg, xs_bg = torch.nonzero(sample['depth'].squeeze()<=0, as_tuple=True )\n descriptors_2d_bg = descriptors_2d_map[:,:,ys_bg, xs_bg].squeeze().permute([1,0])\n # descriptors_2d_bg = descriptors_2d_bg[np.random.randint(0,len(descriptors_2d_bg), size=len(lifted_points) )]\n descriptors_2d = torch.cat([descriptors_2d, descriptors_2d_bg], dim=0) \n descriptors_3d = torch.cat([descriptors_3d, descriptors_2d_bg], dim=0 )\n\n # append handcrafted coordinates to simplify the code(assign the same coordinates for the bg points far away from the fg points, i.e. 10e6 )\n lifted_points = torch.cat([lifted_points, torch.ones([len(descriptors_2d_bg) ,3],device=device )*10e6 ] ) #append very distant points\n model_points = torch.cat([model_points, torch.ones([len(descriptors_2d_bg) ,3], device=device)*10e6 ] ) #append very distant points\n\n #append one-to-one inds\n if len(descriptors_2d_bg)>0:\n # randomly sample the bg points to balance the learning process\n sample_inds=np.random.randint(0,len(descriptors_2d_bg), size=int(len(correspondence)*0.1) )\n \n bg_corr= torch.stack([\n torch.arange(fg_point_num, fg_point_num+len(descriptors_2d_bg), device=device )[sample_inds], \n torch.arange(model_point_num, model_point_num+len(descriptors_2d_bg), device=device)[sample_inds]\n ], dim=-1) #Nx2\n correspondence = torch.cat([correspondence, bg_corr ], dim=0)\n \n if len(lifted_points)>0:\n contra_loss=self.contrastive_loss(src_pcd=lifted_points, tgt_pcd=model_points,\n src_feats=descriptors_2d, tgt_feats=descriptors_3d,\n correspondence=correspondence,\n scores_overlap=None, scores_saliency=None)\n else:\n print(\"Warning: Contrastive loss is skipped, as no lifted point is found!\", flush=True)\n contra_loss={\n 'circle_loss': torch.zeros([1], device=depth.device),\n 'recall': torch.zeros([1], device=depth.device)\n }\n\n\n loss3d = partial(self.pose_loss.forward,\n points=orig_model_points[None])\n motion_loss = self.motion_net.compute_loss(\n preds_dict['Tij_gt'], sample['syn_depth'], intrinsics=sample['K'], loss='l1', log_error=True, loss3d=loss3d)\n\n \n loss = self.contrastive_loss.weight * contra_loss['circle_loss'] + motion_loss['total_loss'] \n\n res = {\n \"loss\": loss, \n \"circle_loss\": contra_loss['circle_loss'].detach(),\n \"recall\":contra_loss['recall'].detach(),\n # \"geometric_loss\": torch.zeros(1),\n \"reproj_loss\": motion_loss['reproj_loss'].detach(),\n \"loss_3d_proj\": motion_loss['loss_3d_proj'].detach(), \n \"valid_mask\": motion_loss[\"valid_mask\"].detach(),\n }\n return res\n\n"
},
{
"path": "model/descriptor2D.py",
"content": "from easydict import EasyDict as edict\nfrom pathlib import Path\nimport torch\nfrom torch import nn\nfrom torchplus.nn.modules.common import Empty\n\n\n\nclass SuperPoint2D(nn.Module):\n \"\"\"SuperPoint Convolutional Detector and Descriptor\n\n SuperPoint: Self-Supervised Interest Point Detection and\n Description. Daniel DeTone, Tomasz Malisiewicz, and Andrew\n Rabinovich. In CVPRW, 2019. https://arxiv.org/abs/1712.07629\n\n \"\"\"\n default_config = {\n 'descriptor_dim': 256,\n 'nms_radius': 4,\n 'keypoint_threshold': 0.005,\n 'max_keypoints': -1,\n 'remove_borders': 4,\n 'saliency_score_normalization_fuc': 'sigmoid',\n \"use_instance_norm\": True\n }\n\n def __init__(self, config):\n super().__init__()\n self.default_config.update(config) \n self.config=edict(self.default_config)\n\n self.normalize_output=config.get('normalize_output', True)\n\n self.saliency_score_normalization_fuc= self.config.saliency_score_normalization_fuc\n assert self.saliency_score_normalization_fuc in ['sigmoid', 'softmax']\n\n self.relu = nn.ReLU(inplace=True)\n self.pool = nn.MaxPool2d(kernel_size=2, stride=2)\n \n if self.config.use_instance_norm:\n self.Normalization=nn.InstanceNorm2d\n else:\n self.Normalization = Empty\n\n\n c1, c2, c3, c4, c5 = 64, 64, 128, 128, 256\n self.input_dim=config.input_dim\n # self.conv1a = nn.Conv2d(1, c1, kernel_size=3, stride=1, padding=1)\n self.conv1a = nn.Conv2d(config.input_dim, c1, kernel_size=3, stride=1, padding=1)\n self.conv1b = nn.Conv2d(c1, c1, kernel_size=3, stride=1, padding=1)\n self.conv2a = nn.Conv2d(c1, c2, kernel_size=3, stride=1, padding=1)\n self.conv2b = nn.Conv2d(c2, c2, kernel_size=3, stride=1, padding=1)\n self.conv3a = nn.Conv2d(c2, c3, kernel_size=3, stride=1, padding=1)\n self.conv3b = nn.Conv2d(c3, c3, kernel_size=3, stride=1, padding=1)\n self.conv4a = nn.Conv2d(c3, c4, kernel_size=3, stride=1, padding=1)\n self.conv4b = nn.Conv2d(c4, c4, kernel_size=3, stride=1, padding=1)\n\n self.convPa = nn.Sequential(\n nn.Conv2d(c4, c5, kernel_size=3, stride=1, padding=1), \n self.Normalization(c5)\n )\n\n self.convPb = nn.Conv2d(c5, 1, kernel_size=1, stride=1, padding=0)\n\n self.convDa = nn.Conv2d(c4, c5, kernel_size=3, stride=1, padding=1)\n\n self.convDb = nn.Conv2d(\n c5, self.config['descriptor_dim'],\n kernel_size=1, stride=1, padding=0)\n\n self.decode1=nn.Sequential(\n nn.Upsample(scale_factor=2,mode='bilinear'),\n nn.Conv2d(c4,c4, kernel_size=3, stride=1, padding=1),\n self.Normalization(c4),\n nn.ReLU())\n self.decode2=nn.Sequential(\n nn.Upsample(scale_factor=2,mode='bilinear'),\n nn.Conv2d(c4+c3,c4, kernel_size=3, stride=1, padding=1),\n self.Normalization(c4),\n nn.ReLU()\n )\n \n self.decode3=nn.Sequential(\n nn.Upsample(scale_factor=2,mode='bilinear'),\n nn.Conv2d(c4+c2,c4, kernel_size=3, stride=1, padding=1),\n self.Normalization(c4),\n nn.ReLU()\n )\n \n path = Path(__file__).parent.parent/ 'weights/superpoint_v1.pth'\n\n self.load_state_dict(torch.load(str(path), map_location='cpu' ), strict=False)\n\n mk = self.config['max_keypoints']\n if mk == 0 or mk < -1:\n raise ValueError('\\\"max_keypoints\\\" must be positive or \\\"-1\\\"')\n\n print('Loaded SuperPoint model')\n\n def load_state_dict(self,state_dict, strict=True):\n \n if not strict:\n updated_state_dict = {}\n model_dict = self.state_dict()\n for k, v in state_dict.items():\n if k in model_dict and v.shape == model_dict[k].shape:\n updated_state_dict[k] = v\n else:\n updated_state_dict = state_dict\n super().load_state_dict(updated_state_dict, strict)\n\n\n def forward_encoder(self, x):\n if self.input_dim==1:\n x=x.mean(dim=1, keepdims=True) #\n x_skip=[]\n # Shared Encoder\n x = self.relu(self.conv1a(x))\n x = self.relu(self.conv1b(x))\n x_skip.append(x)\n x = self.pool(x)\n x = self.relu(self.conv2a(x))\n x = self.relu(self.conv2b(x))\n x_skip.append(x)\n x = self.pool(x)\n x = self.relu(self.conv3a(x))\n x = self.relu(self.conv3b(x))\n x_skip.append(x)\n x = self.pool(x)\n x = self.relu(self.conv4a(x))\n x = self.relu(self.conv4b(x))\n\n return x, x_skip\n def forward_decoder(self, x, x_skip):\n #upsample first\n x = self.decode1(x)\n x=torch.cat([x, x_skip[-1]], dim=1)\n x = self.decode2(x)\n x=torch.cat([x, x_skip[-2]], dim=1)\n x = self.decode3(x)\n\n\n # Compute the dense keypoint scores\n cPa = self.relu(self.convPa(x))\n scores = self.convPb(cPa)\n\n if self.saliency_score_normalization_fuc == 'sigmoid':\n scores = nn.functional.sigmoid(scores)\n elif self.saliency_score_normalization_fuc == 'softmax':\n scores_shape = scores.shape\n scores = scores.reshape(*scores.shape[0:2], -1)\n scores = nn.functional.softmax(scores/1, dim=-1)\n scores= scores.reshape(*scores_shape)#.clone()\n else:\n raise ValueError\n\n keypoints=None\n\n # Compute the dense descriptors\n cDa = self.relu(self.convDa(x))\n descriptors = self.convDb(cDa)\n if self.normalize_output:\n descriptors = torch.nn.functional.normalize(descriptors, p=2, dim=1)\n\n return keypoints, scores, descriptors\n def forward(self, data):\n \"\"\" Compute keypoints, scores, descriptors for image \"\"\"\n # Shared Encoder\n x, x_skip=self.forward_encoder(data)\n\n # Compute the dense keypoint scores\n keypoints, scores, descriptors = self.forward_decoder(x, x_skip)\n\n return {\n 'keypoints': keypoints,\n 'scores': scores,\n 'descriptors': descriptors,\n }\n\n"
},
{
"path": "model/descriptor3D.py",
"content": "import torch \nimport torch.nn as nn \n\n\nfrom kpconv.kpconv_blocks import *\nimport torch.nn.functional as F\nimport numpy as np\nfrom kpconv.lib.utils import square_distance\n\nclass KPSuperpoint3Dv2(nn.Module):\n #remove useless channels\n def __init__(self, config):\n super().__init__()\n self.normalize_output=config.get('normalize_output', True)\n\n # build the architectures\n config.architecture = [\n 'simple',\n 'resnetb',\n ]\n for i in range(config.num_layers-1):\n config.architecture.append('resnetb_strided')\n config.architecture.append('resnetb')\n config.architecture.append('resnetb')\n for i in range(config.num_layers-2):\n config.architecture.append('nearest_upsample')\n config.architecture.append('unary')\n config.architecture.append('nearest_upsample')\n config.architecture.append('last_unary')\n\n ############\n # Parameters\n ############\n # Current radius of convolution and feature dimension\n layer = 0\n r = config.first_subsampling_dl * config.conv_radius\n in_dim = config.in_features_dim\n out_dim = config.first_feats_dim\n self.K = config.num_kernel_points\n self.epsilon = torch.nn.Parameter(torch.tensor(-5.0))\n self.final_feats_dim = config.final_feats_dim\n\n #####################\n # List Encoder blocks\n #####################\n # Save all block operations in a list of modules\n self.encoder_blocks = nn.ModuleList()\n self.encoder_skip_dims = []\n self.encoder_skips = []\n\n # Loop over consecutive blocks\n for block_i, block in enumerate(config.architecture):\n\n # Check equivariance\n if ('equivariant' in block) and (not out_dim % 3 == 0):\n raise ValueError('Equivariant block but features dimension is not a factor of 3')\n\n # Detect change to next layer for skip connection\n if np.any([tmp in block for tmp in ['pool', 'strided', 'upsample', 'global']]):\n self.encoder_skips.append(block_i)\n self.encoder_skip_dims.append(in_dim)\n\n # Detect upsampling block to stop\n if 'upsample' in block:\n break\n\n # Apply the good block function defining tf ops\n self.encoder_blocks.append(block_decider(block,\n r,\n in_dim,\n out_dim,\n layer,\n config))\n\n # Update dimension of input from output\n if 'simple' in block:\n in_dim = out_dim // 2\n else:\n in_dim = out_dim\n\n # Detect change to a subsampled layer\n if 'pool' in block or 'strided' in block:\n # Update radius and feature dimension for next layer\n layer += 1\n r *= 2\n out_dim *= 2\n\n #####################\n # bottleneck layer \n #####################\n botneck_feats_dim = config.gnn_feats_dim\n self.bottle = nn.Conv1d(in_dim, botneck_feats_dim,kernel_size=1,bias=True)\n # num_head = config.num_head\n self.proj_gnn = nn.Conv1d(botneck_feats_dim,botneck_feats_dim,kernel_size=1, bias=True)\n\n \n #####################\n # List Decoder blocks\n #####################\n out_dim = botneck_feats_dim # + 2\n\n # Save all block operations in a list of modules\n self.decoder_blocks = nn.ModuleList()\n self.decoder_concats = []\n\n # Find first upsampling block\n start_i = 0\n for block_i, block in enumerate(config.architecture):\n if 'upsample' in block:\n start_i = block_i\n break\n \n # Loop over consecutive blocks\n for block_i, block in enumerate(config.architecture[start_i:]):\n\n # Add dimension of skip connection concat\n if block_i > 0 and 'upsample' in config.architecture[start_i + block_i - 1]:\n in_dim += self.encoder_skip_dims[layer]\n self.decoder_concats.append(block_i)\n\n # Apply the good block function defining tf ops\n self.decoder_blocks.append(block_decider(block,\n r,\n in_dim,\n out_dim,\n layer,\n config))\n\n \n\n # Update dimension of input from output\n in_dim = out_dim\n\n # Detect change to a subsampled layer\n if 'upsample' in block:\n # Update radius and feature dimension for next layer\n layer -= 1\n r *= 0.5\n out_dim = out_dim // 2\n return\n\n def regular_score(self,score):\n score = torch.where(torch.isnan(score), torch.zeros_like(score), score)\n score = torch.where(torch.isinf(score), torch.zeros_like(score), score)\n return score\n\n def forward_encoder(self, batch):\n # Get input features\n x = batch['features'].clone().detach()\n len_src_c = batch['stack_lengths'][-1][0]\n len_src_f = batch['stack_lengths'][0][0]\n pcd_c = batch['points'][-1]\n pcd_f = batch['points'][0]\n src_pcd_c, tgt_pcd_c = pcd_c[:len_src_c], pcd_c[len_src_c:]\n\n sigmoid = nn.Sigmoid()\n #################################\n # 1. encoder \n skip_x = []\n for block_i, block_op in enumerate(self.encoder_blocks):\n if block_i in self.encoder_skips:\n skip_x.append(x)\n x = block_op(x, batch)\n\n #################################\n # 2. project the bottleneck features\n feats_c = x.transpose(0,1).unsqueeze(0) #[1, C, N]\n feats_c = self.bottle(feats_c) #[1, C, N]\n \n return feats_c,skip_x\n\n def forward_middle(self, x):\n\n feats_c = self.proj_gnn(x) \n feats_gnn_raw = feats_c.squeeze(0).transpose(0,1)\n\n return feats_gnn_raw\n\n\n def forward_decoder(self, x, skip_x, batch ):\n sigmoid = nn.Sigmoid()\n for block_i, block_op in enumerate(self.decoder_blocks):\n if block_i in self.decoder_concats:\n x = torch.cat([x, skip_x.pop()], dim=1)\n x = block_op(x, batch)\n \n feats_f = x[:,:self.final_feats_dim]\n\n # normalise point-wise features\n if self.normalize_output:\n feats_f = F.normalize(feats_f, p=2, dim=1)\n\n return feats_f \n\n def forward(self, batch):\n # Get input features\n feats_c,skip_x = self.forward_encoder(batch)\n x = self.forward_middle(feats_c)\n feats_f = self.forward_decoder(x, skip_x, batch)\n\n return feats_f\n\n\n\n\n"
},
{
"path": "model/losses.py",
"content": "from sklearn.metrics import precision_recall_fscore_support\nfrom thirdparty.kpconv.lib.utils import square_distance\nfrom abc import ABCMeta, abstractmethod\nimport time\n\nimport numpy as np\nimport torch\nfrom torch import nn\nfrom torch.autograd import Variable\nfrom torch.nn import functional as F\nimport torchplus\n# from utils.pca import pca_tch\n\nimport kornia\n# import self_voxelo.utils.pose_utils as pose_utils\nimport apex.amp as amp\n\n# class Loss(object):\n\n\nclass Loss(nn.Module):\n \"\"\"Abstract base class for loss functions.\"\"\"\n __metaclass__ = ABCMeta\n\n def __init__(self, loss_weight=1):\n super(Loss, self).__init__()\n self._loss_weight = loss_weight\n\n # def __call__(self,\n def forward(self,\n prediction_tensor,\n target_tensor,\n ignore_nan_targets=False,\n scope=None,\n **params):\n \"\"\"Call the loss function.\n\n Args:\n prediction_tensor: an N-d tensor of shape [batch, anchors, ...]\n representing predicted quantities.\n target_tensor: an N-d tensor of shape [batch, anchors, ...] representing\n regression or classification targets.\n ignore_nan_targets: whether to ignore nan targets in the loss computation.\n E.g. can be used if the target tensor is missing groundtruth data that\n shouldn't be factored into the loss.\n scope: Op scope name. Defaults to 'Loss' if None.\n **params: Additional keyword arguments for specific implementations of\n the Loss.\n\n Returns:\n loss: a tensor representing the value of the loss function.\n \"\"\"\n if ignore_nan_targets:\n target_tensor = torch.where(torch.isnan(target_tensor),\n prediction_tensor,\n target_tensor)\n # ret = self._compute_loss(prediction_tensor, target_tensor, **params)\n # if isinstance(ret, (list, tuple)):\n # return [self._loss_weight*ret[0]] + list(ret[1:])\n # else:\n return self._loss_weight*self._compute_loss(prediction_tensor, target_tensor, **params)\n\n @abstractmethod\n @amp.float_function\n def _compute_loss(self, prediction_tensor, target_tensor, **params):\n \"\"\"Method to be overridden by implementations.\n\n Args:\n prediction_tensor: a tensor representing predicted quantities\n target_tensor: a tensor representing regression or classification targets\n **params: Additional keyword arguments for specific implementations of\n the Loss.\n\n Returns:\n loss: an N-d tensor of shape [batch, anchors, ...] containing the loss per\n anchor\n \"\"\"\n\n raise NotImplementedError\n\n\nclass L2Loss(Loss):\n\n def __init__(self, loss_weight=1):\n super(L2Loss, self).__init__(loss_weight)\n\n def _compute_loss(self, prediction_tensor, target_tensor, mask=None):\n \"\"\"Compute loss function.\n\n Args:\n prediction_tensor: A float tensor of shape [batch_size, num_anchors,\n code_size] representing the (encoded) predicted locations of objects.\n target_tensor: A float tensor of shape [batch_size, num_anchors,\n code_size] representing the regression targets\n\n Returns:\n loss: a float tensor of shape [batch_size, num_anchors] tensor\n representing the value of the loss function.\n \"\"\"\n diff = prediction_tensor - target_tensor\n\n if mask is not None:\n mask = mask.expand_as(diff).byte()\n diff = diff[mask]\n # square_diff = 0.5 * weighted_diff * weighted_diff\n square_diff = diff * diff\n return square_diff.mean()\n\n\nclass AdaptiveWeightedL2Loss(Loss):\n\n def __init__(self, init_alpha, learn_alpha=True, loss_weight=1, focal_gamma=0):\n super(AdaptiveWeightedL2Loss, self).__init__(loss_weight)\n self.learn_alpha = learn_alpha\n self.alpha = nn.Parameter(torch.Tensor(\n [init_alpha]), requires_grad=learn_alpha)\n self.focal_gamma = focal_gamma\n # self.alpha_shift = -13 # -10# TODO: temporarily test\n\n def _compute_loss(self, prediction_tensor, target_tensor, mask=None, alpha=None, focal_gamma=None):\n \"\"\"Compute loss function.\n\n Args:\n prediction_tensor: A float tensor of shape [batch_size, num_anchors,\n code_size] representing the (encoded) predicted locations of objects.\n target_tensor: A float tensor of shape [batch_size, num_anchors,\n code_size] representing the regression targets\n\n Returns:\n loss: a float tensor of shape [batch_size, num_anchors] tensor\n representing the value of the loss function.\n \"\"\"\n\n if focal_gamma is None:\n focal_gamma = self.focal_gamma\n _alpha = self.alpha\n if mask is None:\n mask = torch.ones_like(target_tensor)\n else:\n mask = mask.expand_as(target_tensor)\n\n diff = prediction_tensor - target_tensor\n square_diff = (diff * diff) * mask\n\n # loss = square_diff.mean()\n input_shape = prediction_tensor.shape\n loss = torch.sum(square_diff, dim=list(range(1, len(input_shape)))) / \\\n (torch.sum(mask, dim=list(range(1, len(input_shape)))) + 1e-12) # (B,)\n\n focal_weight = (torch.exp(-_alpha) * loss)**focal_gamma\n focal_weight = focal_weight/(torch.sum(focal_weight) + 1e-12)\n\n loss = focal_weight*(torch.exp(-_alpha) * loss)\n loss = loss.sum() + _alpha\n return loss\n\n\nclass MetricLoss(nn.Module):\n \"\"\"\n We evaluate both contrastive loss and circle loss\n \"\"\"\n\n def __init__(self, configs, log_scale=16, pos_optimal=0.1, neg_optimal=1.4, ):\n super(MetricLoss, self).__init__()\n self.log_scale = log_scale\n self.pos_optimal = pos_optimal\n self.neg_optimal = neg_optimal\n\n self.pos_margin = configs.pos_margin\n self.neg_margin = configs.neg_margin\n self.max_points = configs.max_points\n\n self.safe_radius = configs.safe_radius\n self.matchability_radius = configs.matchability_radius\n # just to take care of the numeric precision\n self.pos_radius = configs.pos_radius + 0.001\n self.weight = configs.get('loss_weight', 1)\n\n def get_circle_loss(self, coords_dist, feats_dist):\n \"\"\"\n Modified from: https://github.com/XuyangBai/D3Feat.pytorch\n \"\"\"\n pos_mask = coords_dist < self.pos_radius\n neg_mask = coords_dist > self.safe_radius\n\n # get anchors that have both positive and negative pairs\n row_sel = ((pos_mask.sum(-1) > 0) * (neg_mask.sum(-1) > 0)).detach()\n col_sel = ((pos_mask.sum(-2) > 0) * (neg_mask.sum(-2) > 0)).detach()\n\n # get alpha for both positive and negative pairs\n pos_weight = feats_dist - 1e5 * \\\n (~pos_mask).float() # mask the non-positive\n # mask the uninformative positive\n pos_weight = (pos_weight - self.pos_optimal)\n pos_weight = torch.max(torch.zeros_like(\n pos_weight), pos_weight).detach()\n\n neg_weight = feats_dist + 1e5 * \\\n (~neg_mask).float() # mask the non-negative\n # mask the uninformative negative\n neg_weight = (self.neg_optimal - neg_weight)\n neg_weight = torch.max(torch.zeros_like(\n neg_weight), neg_weight).detach()\n\n lse_pos_row = torch.logsumexp(\n self.log_scale * (feats_dist - self.pos_margin) * pos_weight, dim=-1)\n lse_pos_col = torch.logsumexp(\n self.log_scale * (feats_dist - self.pos_margin) * pos_weight, dim=-2)\n\n lse_neg_row = torch.logsumexp(\n self.log_scale * (self.neg_margin - feats_dist) * neg_weight, dim=-1)\n lse_neg_col = torch.logsumexp(\n self.log_scale * (self.neg_margin - feats_dist) * neg_weight, dim=-2)\n\n loss_row = F.softplus(lse_pos_row + lse_neg_row)/self.log_scale\n loss_col = F.softplus(lse_pos_col + lse_neg_col)/self.log_scale\n\n circle_loss = (loss_row[row_sel].mean() + loss_col[col_sel].mean()) / 2\n\n return circle_loss\n\n def get_recall(self, coords_dist, feats_dist):\n \"\"\"\n Get feature match recall, divided by number of true inliers\n \"\"\"\n pos_mask = coords_dist < self.pos_radius\n n_gt_pos = (pos_mask.sum(-1) > 0).float().sum()+1e-12\n _, sel_idx = torch.min(feats_dist, -1)\n\n sel_dist = torch.gather(coords_dist, dim=-1,\n index=sel_idx[:, None])[pos_mask.sum(-1) > 0]\n n_pred_pos = (sel_dist < self.pos_radius).float().sum()\n recall = n_pred_pos / n_gt_pos\n return recall\n\n def get_weighted_bce_loss(self, prediction, gt):\n loss = nn.BCELoss(reduction='none')\n\n class_loss = loss(prediction, gt)\n\n weights = torch.ones_like(gt)\n w_negative = gt.sum()/gt.size(0)\n w_positive = 1 - w_negative\n\n weights[gt >= 0.5] = w_positive\n weights[gt < 0.5] = w_negative\n w_class_loss = torch.mean(weights * class_loss)\n\n #######################################\n # get classification precision and recall\n predicted_labels = prediction.detach().cpu().round().numpy()\n cls_precision, cls_recall, _, _ = precision_recall_fscore_support(\n gt.cpu().numpy(), predicted_labels, average='binary')\n\n return w_class_loss, cls_precision, cls_recall\n\n def forward(self, src_pcd, tgt_pcd, src_feats, tgt_feats, correspondence, scores_overlap, scores_saliency, rot=None, trans=None):\n \"\"\"\n Circle loss for metric learning, here we feed the positive pairs only\n Input:\n src_pcd: [N, 3], pcd of the 3d model \n tgt_pcd: [M, 3], pcd of the lifted model from 2d depth\n rot: [3, 3], rotation used to rotate the src_pcd to the current frame\n trans: [3, 1], translation used to translate the src_pcd to the current frame\n src_feats: [N, C]\n tgt_feats: [M, C]\n \"\"\"\n\n if rot is not None and trans is not None:\n src_pcd = (torch.matmul(rot, src_pcd.transpose(0, 1)) +\n trans).transpose(0, 1)\n\n stats = dict()\n\n #######################################\n # filter some of correspondence\n if(correspondence.size(0) > self.max_points):\n choice = np.random.permutation(correspondence.size(0))[\n :self.max_points]\n correspondence = correspondence[choice]\n\n src_idx = correspondence[:, 0]\n tgt_idx = correspondence[:, 1]\n src_pcd, tgt_pcd = src_pcd[src_idx], tgt_pcd[tgt_idx]\n src_feats, tgt_feats = src_feats[src_idx], tgt_feats[tgt_idx]\n\n #######################\n # get L2 distance between source / target point cloud\n\n coords_dist = torch.sqrt(square_distance(\n src_pcd[None, :, :], tgt_pcd[None, :, :]).squeeze(0))\n feats_dist = torch.sqrt(square_distance(\n src_feats[None, :, :], tgt_feats[None, :, :], normalised=True)).squeeze(0)\n\n ##############################\n # get FMR and circle loss\n ##############################\n recall = self.get_recall(coords_dist, feats_dist)\n circle_loss = self.get_circle_loss(coords_dist, feats_dist)\n\n stats['circle_loss'] = circle_loss\n stats['recall'] = recall\n\n return stats\n\n\nclass PointAlignmentLoss(nn.Module):\n def __init__(self, loss_weight=1, ):\n super().__init__()\n self._loss_weight = loss_weight\n\n def forward(self, R_pred, t_pred, R_tgt, t_tgt, points):\n return self._loss_weight*self._compute_loss(R_pred, t_pred, R_tgt, t_tgt, points)\n\n def _compute_loss(self, R_pred, t_pred, R_tgt, t_tgt, points, ):\n \"\"\"[summary]\n\n Args:\n R_pred ([type]): [Bx3x3]\n t_pred ([type]): [Bx3]\n R_tgt ([type]): [Bx3x3]\n t_tgt ([type]): [Bx3]\n points ([type]): [BxNx3]\n\n Returns:\n loss [type]: [loss value]\n \"\"\"\n\n loss = 0\n for b in range(len(points)):\n\n diff = points[b]@R_pred[b].transpose(-1, -2) + t_pred[b] - (\n points[b]@R_tgt[b].transpose(-1, -2)+t_tgt[b])\n\n square_diff = diff.abs() \n loss += torch.mean(square_diff)*3\n\n # loss/=len(points)\n\n return loss\n"
},
{
"path": "scripts/compile_3rdparty.sh",
"content": "#!/usr/bin/bash\n\nSCRIPT_DIR=\"$( cd \"$( dirname \"${BASH_SOURCE[0]}\" )\" &> /dev/null && pwd )\"\n\ncd $SCRIPT_DIR/../thirdparty/kpconv/cpp_wrappers\nbash compile_wrappers.sh\n\ncd $SCRIPT_DIR/../thirdparty/nn\npython setup.py build_ext --inplace"
},
{
"path": "scripts/eval.sh",
"content": "export PROJECT_ROOT_PATH=/home/RNNPose/Projects/Works/RNNPose_release\n\nexport PYTHONPATH=\"$PROJECT_ROOT_PATH:$PYTHONPATH\"\nexport PYTHONPATH=\"$PROJECT_ROOT_PATH/thirdparty:$PYTHONPATH\"\nexport model_dir='outputs'\nseq=cat\ngpu=1\nstart_gpu_id=0\nmkdir $model_dir\n\ntrain_file=$PROJECT_ROOT_PATH/tools/eval.py\nconfig_path=/$PROJECT_ROOT_PATH/config/linemod/\"$seq\"_fw0.5.yml\npretrain=$PROJECT_ROOT_PATH/weights/trained_models/\"$seq\".tckpt\n\npython -u $train_file multi_proc_train \\\n --config_path $config_path \\\n --model_dir $model_dir/results \\\n --use_dist True \\\n --dist_port 10000 \\\n --gpus_per_node $gpu \\\n --optim_eval True \\\n --use_apex False \\\n --world_size $gpu \\\n --start_gpu_id $start_gpu_id \\\n --pretrained_path $pretrain \n "
},
{
"path": "scripts/eval_lmocc.sh",
"content": "export PROJECT_ROOT_PATH=/home/RNNPose/Projects/Works/RNNPose_release\n\nexport PYTHONPATH=\"$PROJECT_ROOT_PATH:$PYTHONPATH\"\nexport PYTHONPATH=\"$PROJECT_ROOT_PATH/thirdparty:$PYTHONPATH\"\nexport model_dir='outputs'\nseq=ape\ngpu=1\nstart_gpu_id=0\nmkdir $model_dir\n\ntrain_file=$PROJECT_ROOT_PATH/tools/eval.py\nconfig_path=/$PROJECT_ROOT_PATH/config/linemod/\"$seq\"_fw0.5_occ.yml\npretrain=$PROJECT_ROOT_PATH/weights/trained_models/\"$seq\".tckpt\n\npython -u $train_file multi_proc_train \\\n --config_path $config_path \\\n --model_dir $model_dir/results \\\n --use_dist True \\\n --dist_port 10000 \\\n --gpus_per_node $gpu \\\n --optim_eval True \\\n --use_apex False \\\n --world_size $gpu \\\n --start_gpu_id $start_gpu_id \\\n --pretrained_path $pretrain \n # --use_apex True \\\n "
},
{
"path": "scripts/run_dataformatter.sh",
"content": "SCRIPT_DIR=\"$( cd \"$( dirname \"${BASH_SOURCE[0]}\" )\" &> /dev/null && pwd )\"\nPROJ_ROOT=$SCRIPT_DIR/../\n\npython $PROJ_ROOT/tools/transform_data_format.py run --data_type \"LM_FUSE_PVNET\" --data_info_path $PROJ_ROOT/EXPDATA/\"/data_info/linemod_all_10k_default.info.all\" --image_root $PROJ_ROOT/EXPDATA/raw_data/fuse --depth_root $PROJ_ROOT/EXPDATA/raw_data/orig_renders --save_dir $PROJ_ROOT/EXPDATA/LINEMOD/fuse_formatted && touch 3.finished \n\n\n"
},
{
"path": "scripts/run_datainfo_generation.sh",
"content": "\nSCRIPT_DIR=\"$( cd \"$( dirname \"${BASH_SOURCE[0]}\" )\" &> /dev/null && pwd )\"\nPROJ_ROOT=$SCRIPT_DIR/../\nexport PYTHONPATH=$PROJ_ROOT:PYTHONPATH\n\nmkdir --parent $PROJ_ROOT/EXPDATA/data_info/deepim/\n\npython $PROJ_ROOT/tools/generate_data_info_deepim_0_orig.py create_data_info --data_root $PROJ_ROOT/EXPDATA/LM6d_converted/LM6d_refine --saving_path $PROJ_ROOT/EXPDATA/data_info/deepim/linemod_orig_deepim.info \n\npython $PROJ_ROOT/tools/generate_data_info_deepim_1_syn.py create_data_info --data_root $PROJ_ROOT/EXPDATA/LM6d_converted/LM6d_refine_syn --saving_path $PROJ_ROOT/EXPDATA/data_info/deepim/linemod_syn_deepim.info --with_assertion True\n\npython $PROJ_ROOT/tools/generate_data_info_deepim_2_posecnnval.py create_data_info --data_root $PROJ_ROOT/EXPDATA/LM6d_converted/LM6d_refine --saving_path $PROJ_ROOT/EXPDATA/data_info/linemod_posecnn.info --with_assertion True\n\npython $PROJ_ROOT/tools/generate_data_info_v2_deepim.py create_data_info --data_root $PROJ_ROOT/EXPDATA/LINEMOD/fuse_formatted/ --saving_path $PROJ_ROOT/EXPDATA/data_info/linemod_fusesformatted_all10k_deepim.info --training_data_ratio 1 --shuffle=False\n"
},
{
"path": "scripts/train.sh",
"content": "export PROJECT_ROOT_PATH=/home/RNNPose/Projects/Works/RNNPose_release\n\nexport PYTHONPATH=\"$PROJECT_ROOT_PATH:$PYTHONPATH\"\nexport PYTHONPATH=\"$PROJECT_ROOT_PATH/thirdparty:$PYTHONPATH\"\nexport model_dir='outputs'\nseq=cat\ngpu=1\nstart_gpu_id=0\nmkdir $model_dir\n\ntrain_file=$PROJECT_ROOT_PATH/tools/train.py\nconfig_path=/$PROJECT_ROOT_PATH/config/linemod/\"$seq\"_fw0.5.yml\n# pretrain=$PROJECT_ROOT_PATH/weights/trained_models/\"$seq\".tckpt\n\npython -u $train_file multi_proc_train \\\n --config_path $config_path \\\n --model_dir $model_dir/results \\\n --use_dist True \\\n --dist_port 10000 \\\n --gpus_per_node $gpu \\\n --optim_eval True \\\n --use_apex False \\\n --world_size $gpu \\\n --start_gpu_id $start_gpu_id \\\n # --pretrained_path $pretrain \n "
},
{
"path": "thirdparty/kpconv/__init__.py",
"content": ""
},
{
"path": "thirdparty/kpconv/cpp_wrappers/compile_wrappers.sh",
"content": "#!/bin/bash\n\n# Compile cpp subsampling\ncd cpp_subsampling\npython3 setup.py build_ext --inplace\ncd ..\n\n# Compile cpp neighbors\ncd cpp_neighbors\npython3 setup.py build_ext --inplace\ncd .."
},
{
"path": "thirdparty/kpconv/cpp_wrappers/cpp_neighbors/build.bat",
"content": "@echo off\npy setup.py build_ext --inplace\n\n\npause"
},
{
"path": "thirdparty/kpconv/cpp_wrappers/cpp_neighbors/neighbors/neighbors.cpp",
"content": "\n#include \"neighbors.h\"\n\n\nvoid brute_neighbors(vector& queries, vector& supports, vector& neighbors_indices, float radius, int verbose)\n{\n\n\t// Initialize variables\n\t// ******************\n\n\t// square radius\n\tfloat r2 = radius * radius;\n\n\t// indices\n\tint i0 = 0;\n\n\t// Counting vector\n\tint max_count = 0;\n\tvector> tmp(queries.size());\n\n\t// Search neigbors indices\n\t// ***********************\n\n\tfor (auto& p0 : queries)\n\t{\n\t\tint i = 0;\n\t\tfor (auto& p : supports)\n\t\t{\n\t\t\tif ((p0 - p).sq_norm() < r2)\n\t\t\t{\n\t\t\t\ttmp[i0].push_back(i);\n\t\t\t\tif (tmp[i0].size() > max_count)\n\t\t\t\t\tmax_count = tmp[i0].size();\n\t\t\t}\n\t\t\ti++;\n\t\t}\n\t\ti0++;\n\t}\n\n\t// Reserve the memory\n\tneighbors_indices.resize(queries.size() * max_count);\n\ti0 = 0;\n\tfor (auto& inds : tmp)\n\t{\n\t\tfor (int j = 0; j < max_count; j++)\n\t\t{\n\t\t\tif (j < inds.size())\n\t\t\t\tneighbors_indices[i0 * max_count + j] = inds[j];\n\t\t\telse\n\t\t\t\tneighbors_indices[i0 * max_count + j] = -1;\n\t\t}\n\t\ti0++;\n\t}\n\n\treturn;\n}\n\nvoid ordered_neighbors(vector& queries,\n vector& supports,\n vector& neighbors_indices,\n float radius)\n{\n\n\t// Initialize variables\n\t// ******************\n\n\t// square radius\n\tfloat r2 = radius * radius;\n\n\t// indices\n\tint i0 = 0;\n\n\t// Counting vector\n\tint max_count = 0;\n\tfloat d2;\n\tvector> tmp(queries.size());\n\tvector> dists(queries.size());\n\n\t// Search neigbors indices\n\t// ***********************\n\n\tfor (auto& p0 : queries)\n\t{\n\t\tint i = 0;\n\t\tfor (auto& p : supports)\n\t\t{\n\t\t d2 = (p0 - p).sq_norm();\n\t\t\tif (d2 < r2)\n\t\t\t{\n\t\t\t // Find order of the new point\n\t\t\t auto it = std::upper_bound(dists[i0].begin(), dists[i0].end(), d2);\n\t\t\t int index = std::distance(dists[i0].begin(), it);\n\n\t\t\t // Insert element\n dists[i0].insert(it, d2);\n tmp[i0].insert(tmp[i0].begin() + index, i);\n\n\t\t\t // Update max count\n\t\t\t\tif (tmp[i0].size() > max_count)\n\t\t\t\t\tmax_count = tmp[i0].size();\n\t\t\t}\n\t\t\ti++;\n\t\t}\n\t\ti0++;\n\t}\n\n\t// Reserve the memory\n\tneighbors_indices.resize(queries.size() * max_count);\n\ti0 = 0;\n\tfor (auto& inds : tmp)\n\t{\n\t\tfor (int j = 0; j < max_count; j++)\n\t\t{\n\t\t\tif (j < inds.size())\n\t\t\t\tneighbors_indices[i0 * max_count + j] = inds[j];\n\t\t\telse\n\t\t\t\tneighbors_indices[i0 * max_count + j] = -1;\n\t\t}\n\t\ti0++;\n\t}\n\n\treturn;\n}\n\nvoid batch_ordered_neighbors(vector& queries,\n vector& supports,\n vector& q_batches,\n vector& s_batches,\n vector& neighbors_indices,\n float radius)\n{\n\n\t// Initialize variables\n\t// ******************\n\n\t// square radius\n\tfloat r2 = radius * radius;\n\n\t// indices\n\tint i0 = 0;\n\n\t// Counting vector\n\tint max_count = 0;\n\tfloat d2;\n\tvector> tmp(queries.size());\n\tvector> dists(queries.size());\n\n\t// batch index\n\tint b = 0;\n\tint sum_qb = 0;\n\tint sum_sb = 0;\n\n\n\t// Search neigbors indices\n\t// ***********************\n\n\tfor (auto& p0 : queries)\n\t{\n\t // Check if we changed batch\n\t if (i0 == sum_qb + q_batches[b])\n\t {\n\t sum_qb += q_batches[b];\n\t sum_sb += s_batches[b];\n\t b++;\n\t }\n\n\t // Loop only over the supports of current batch\n\t vector::iterator p_it;\n\t\tint i = 0;\n for(p_it = supports.begin() + sum_sb; p_it < supports.begin() + sum_sb + s_batches[b]; p_it++ )\n {\n\t\t d2 = (p0 - *p_it).sq_norm();\n\t\t\tif (d2 < r2)\n\t\t\t{\n\t\t\t // Find order of the new point\n\t\t\t auto it = std::upper_bound(dists[i0].begin(), dists[i0].end(), d2);\n\t\t\t int index = std::distance(dists[i0].begin(), it);\n\n\t\t\t // Insert element\n dists[i0].insert(it, d2);\n tmp[i0].insert(tmp[i0].begin() + index, sum_sb + i);\n\n\t\t\t // Update max count\n\t\t\t\tif (tmp[i0].size() > max_count)\n\t\t\t\t\tmax_count = tmp[i0].size();\n\t\t\t}\n\t\t\ti++;\n\t\t}\n\t\ti0++;\n\t}\n\n\t// Reserve the memory\n\tneighbors_indices.resize(queries.size() * max_count);\n\ti0 = 0;\n\tfor (auto& inds : tmp)\n\t{\n\t\tfor (int j = 0; j < max_count; j++)\n\t\t{\n\t\t\tif (j < inds.size())\n\t\t\t\tneighbors_indices[i0 * max_count + j] = inds[j];\n\t\t\telse\n\t\t\t\tneighbors_indices[i0 * max_count + j] = supports.size();\n\t\t}\n\t\ti0++;\n\t}\n\n\treturn;\n}\n\n\nvoid batch_nanoflann_neighbors(vector& queries,\n vector& supports,\n vector& q_batches,\n vector& s_batches,\n vector& neighbors_indices,\n float radius)\n{\n\n\t// Initialize variables\n\t// ******************\n\n\t// indices\n\tint i0 = 0;\n\n\t// Square radius\n\tfloat r2 = radius * radius;\n\n\t// Counting vector\n\tint max_count = 0;\n\tfloat d2;\n\tvector>> all_inds_dists(queries.size());\n\n\t// batch index\n\tint b = 0;\n\tint sum_qb = 0;\n\tint sum_sb = 0;\n\n\t// Nanoflann related variables\n\t// ***************************\n\n\t// CLoud variable\n\tPointCloud current_cloud;\n\n\t// Tree parameters\n\tnanoflann::KDTreeSingleIndexAdaptorParams tree_params(10 /* max leaf */);\n\n\t// KDTree type definition\n typedef nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Simple_Adaptor ,\n PointCloud,\n 3 > my_kd_tree_t;\n\n // Pointer to trees\n my_kd_tree_t* index;\n\n // Build KDTree for the first batch element\n current_cloud.pts = vector(supports.begin() + sum_sb, supports.begin() + sum_sb + s_batches[b]);\n index = new my_kd_tree_t(3, current_cloud, tree_params);\n index->buildIndex();\n\n\n\t// Search neigbors indices\n\t// ***********************\n\n // Search params\n nanoflann::SearchParams search_params;\n search_params.sorted = true;\n\n\tfor (auto& p0 : queries)\n\t{\n\n\t // Check if we changed batch\n\t if (i0 == sum_qb + q_batches[b])\n\t {\n\t sum_qb += q_batches[b];\n\t sum_sb += s_batches[b];\n\t b++;\n\n\t // Change the points\n\t current_cloud.pts.clear();\n current_cloud.pts = vector(supports.begin() + sum_sb, supports.begin() + sum_sb + s_batches[b]);\n\n\t // Build KDTree of the current element of the batch\n delete index;\n index = new my_kd_tree_t(3, current_cloud, tree_params);\n index->buildIndex();\n\t }\n\n\t // Initial guess of neighbors size\n all_inds_dists[i0].reserve(max_count);\n\n\t // Find neighbors\n\t float query_pt[3] = { p0.x, p0.y, p0.z};\n\t\tsize_t nMatches = index->radiusSearch(query_pt, r2, all_inds_dists[i0], search_params);\n\n // Update max count\n if (nMatches > max_count)\n max_count = nMatches;\n\n // Increment query idx\n\t\ti0++;\n\t}\n\n\t// Reserve the memory\n\tneighbors_indices.resize(queries.size() * max_count);\n\ti0 = 0;\n\tsum_sb = 0;\n\tsum_qb = 0;\n\tb = 0;\n\tfor (auto& inds_dists : all_inds_dists)\n\t{\n\t // Check if we changed batch\n\t if (i0 == sum_qb + q_batches[b])\n\t {\n\t sum_qb += q_batches[b];\n\t sum_sb += s_batches[b];\n\t b++;\n\t }\n\n\t\tfor (int j = 0; j < max_count; j++)\n\t\t{\n\t\t\tif (j < inds_dists.size())\n\t\t\t\tneighbors_indices[i0 * max_count + j] = inds_dists[j].first + sum_sb;\n\t\t\telse\n\t\t\t\tneighbors_indices[i0 * max_count + j] = supports.size();\n\t\t}\n\t\ti0++;\n\t}\n\n\tdelete index;\n\n\treturn;\n}\n\n"
},
{
"path": "thirdparty/kpconv/cpp_wrappers/cpp_neighbors/neighbors/neighbors.h",
"content": "\n\n#include \"../../cpp_utils/cloud/cloud.h\"\n#include \"../../cpp_utils/nanoflann/nanoflann.hpp\"\n\n#include \n#include \n\nusing namespace std;\n\n\nvoid ordered_neighbors(vector& queries,\n vector& supports,\n vector& neighbors_indices,\n float radius);\n\nvoid batch_ordered_neighbors(vector& queries,\n vector& supports,\n vector& q_batches,\n vector& s_batches,\n vector& neighbors_indices,\n float radius);\n\nvoid batch_nanoflann_neighbors(vector& queries,\n vector& supports,\n vector& q_batches,\n vector& s_batches,\n vector& neighbors_indices,\n float radius);\n"
},
{
"path": "thirdparty/kpconv/cpp_wrappers/cpp_neighbors/setup.py",
"content": "from distutils.core import setup, Extension\nimport numpy.distutils.misc_util\n\n# Adding OpenCV to project\n# ************************\n\n# Adding sources of the project\n# *****************************\n\nSOURCES = [\"../cpp_utils/cloud/cloud.cpp\",\n \"neighbors/neighbors.cpp\",\n \"wrapper.cpp\"]\n\nmodule = Extension(name=\"radius_neighbors\",\n sources=SOURCES,\n extra_compile_args=['-std=c++11',\n '-D_GLIBCXX_USE_CXX11_ABI=0'])\n\n\nsetup(ext_modules=[module], include_dirs=numpy.distutils.misc_util.get_numpy_include_dirs())\n\n\n\n\n\n\n\n\n"
},
{
"path": "thirdparty/kpconv/cpp_wrappers/cpp_neighbors/wrapper.cpp",
"content": "#include \n#include \n#include \"neighbors/neighbors.h\"\n#include \n\n\n\n// docstrings for our module\n// *************************\n\nstatic char module_docstring[] = \"This module provides two methods to compute radius neighbors from pointclouds or batch of pointclouds\";\n\nstatic char batch_query_docstring[] = \"Method to get radius neighbors in a batch of stacked pointclouds\";\n\n\n// Declare the functions\n// *********************\n\nstatic PyObject *batch_neighbors(PyObject *self, PyObject *args, PyObject *keywds);\n\n\n// Specify the members of the module\n// *********************************\n\nstatic PyMethodDef module_methods[] = \n{\n\t{ \"batch_query\", (PyCFunction)batch_neighbors, METH_VARARGS | METH_KEYWORDS, batch_query_docstring },\n\t{NULL, NULL, 0, NULL}\n};\n\n\n// Initialize the module\n// *********************\n\nstatic struct PyModuleDef moduledef = \n{\n PyModuleDef_HEAD_INIT,\n \"radius_neighbors\",\t\t// m_name\n module_docstring, // m_doc\n -1, // m_size\n module_methods, // m_methods\n NULL, // m_reload\n NULL, // m_traverse\n NULL, // m_clear\n NULL, // m_free\n};\n\nPyMODINIT_FUNC PyInit_radius_neighbors(void)\n{\n import_array();\n\treturn PyModule_Create(&moduledef);\n}\n\n\n// Definition of the batch_subsample method\n// **********************************\n\nstatic PyObject* batch_neighbors(PyObject* self, PyObject* args, PyObject* keywds)\n{\n\n\t// Manage inputs\n\t// *************\n\n\t// Args containers\n\tPyObject* queries_obj = NULL;\n\tPyObject* supports_obj = NULL;\n\tPyObject* q_batches_obj = NULL;\n\tPyObject* s_batches_obj = NULL;\n\n\t// Keywords containers\n\tstatic char* kwlist[] = { \"queries\", \"supports\", \"q_batches\", \"s_batches\", \"radius\", NULL };\n\tfloat radius = 0.1;\n\n\t// Parse the input \n\tif (!PyArg_ParseTupleAndKeywords(args, keywds, \"OOOO|$f\", kwlist, &queries_obj, &supports_obj, &q_batches_obj, &s_batches_obj, &radius))\n\t{\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Error parsing arguments\");\n\t\treturn NULL;\n\t}\n\n\n\t// Interpret the input objects as numpy arrays.\n\tPyObject* queries_array = PyArray_FROM_OTF(queries_obj, NPY_FLOAT, NPY_IN_ARRAY);\n\tPyObject* supports_array = PyArray_FROM_OTF(supports_obj, NPY_FLOAT, NPY_IN_ARRAY);\n\tPyObject* q_batches_array = PyArray_FROM_OTF(q_batches_obj, NPY_INT, NPY_IN_ARRAY);\n\tPyObject* s_batches_array = PyArray_FROM_OTF(s_batches_obj, NPY_INT, NPY_IN_ARRAY);\n\n\t// Verify data was load correctly.\n\tif (queries_array == NULL)\n\t{\n\t\tPy_XDECREF(queries_array);\n\t\tPy_XDECREF(supports_array);\n\t\tPy_XDECREF(q_batches_array);\n\t\tPy_XDECREF(s_batches_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Error converting query points to numpy arrays of type float32\");\n\t\treturn NULL;\n\t}\n\tif (supports_array == NULL)\n\t{\n\t\tPy_XDECREF(queries_array);\n\t\tPy_XDECREF(supports_array);\n\t\tPy_XDECREF(q_batches_array);\n\t\tPy_XDECREF(s_batches_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Error converting support points to numpy arrays of type float32\");\n\t\treturn NULL;\n\t}\n\tif (q_batches_array == NULL)\n\t{\n\t\tPy_XDECREF(queries_array);\n\t\tPy_XDECREF(supports_array);\n\t\tPy_XDECREF(q_batches_array);\n\t\tPy_XDECREF(s_batches_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Error converting query batches to numpy arrays of type int32\");\n\t\treturn NULL;\n\t}\n\tif (s_batches_array == NULL)\n\t{\n\t\tPy_XDECREF(queries_array);\n\t\tPy_XDECREF(supports_array);\n\t\tPy_XDECREF(q_batches_array);\n\t\tPy_XDECREF(s_batches_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Error converting support batches to numpy arrays of type int32\");\n\t\treturn NULL;\n\t}\n\n\t// Check that the input array respect the dims\n\tif ((int)PyArray_NDIM(queries_array) != 2 || (int)PyArray_DIM(queries_array, 1) != 3)\n\t{\n\t\tPy_XDECREF(queries_array);\n\t\tPy_XDECREF(supports_array);\n\t\tPy_XDECREF(q_batches_array);\n\t\tPy_XDECREF(s_batches_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Wrong dimensions : query.shape is not (N, 3)\");\n\t\treturn NULL;\n\t}\n\tif ((int)PyArray_NDIM(supports_array) != 2 || (int)PyArray_DIM(supports_array, 1) != 3)\n\t{\n\t\tPy_XDECREF(queries_array);\n\t\tPy_XDECREF(supports_array);\n\t\tPy_XDECREF(q_batches_array);\n\t\tPy_XDECREF(s_batches_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Wrong dimensions : support.shape is not (N, 3)\");\n\t\treturn NULL;\n\t}\n\tif ((int)PyArray_NDIM(q_batches_array) > 1)\n\t{\n\t\tPy_XDECREF(queries_array);\n\t\tPy_XDECREF(supports_array);\n\t\tPy_XDECREF(q_batches_array);\n\t\tPy_XDECREF(s_batches_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Wrong dimensions : queries_batches.shape is not (B,) \");\n\t\treturn NULL;\n\t}\n\tif ((int)PyArray_NDIM(s_batches_array) > 1)\n\t{\n\t\tPy_XDECREF(queries_array);\n\t\tPy_XDECREF(supports_array);\n\t\tPy_XDECREF(q_batches_array);\n\t\tPy_XDECREF(s_batches_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Wrong dimensions : supports_batches.shape is not (B,) \");\n\t\treturn NULL;\n\t}\n\tif ((int)PyArray_DIM(q_batches_array, 0) != (int)PyArray_DIM(s_batches_array, 0))\n\t{\n\t\tPy_XDECREF(queries_array);\n\t\tPy_XDECREF(supports_array);\n\t\tPy_XDECREF(q_batches_array);\n\t\tPy_XDECREF(s_batches_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Wrong number of batch elements: different for queries and supports \");\n\t\treturn NULL;\n\t}\n\n\t// Number of points\n\tint Nq = (int)PyArray_DIM(queries_array, 0);\n\tint Ns= (int)PyArray_DIM(supports_array, 0);\n\n\t// Number of batches\n\tint Nb = (int)PyArray_DIM(q_batches_array, 0);\n\n\t// Call the C++ function\n\t// *********************\n\n\t// Convert PyArray to Cloud C++ class\n\tvector queries;\n\tvector supports;\n\tvector q_batches;\n\tvector s_batches;\n\tqueries = vector((PointXYZ*)PyArray_DATA(queries_array), (PointXYZ*)PyArray_DATA(queries_array) + Nq);\n\tsupports = vector((PointXYZ*)PyArray_DATA(supports_array), (PointXYZ*)PyArray_DATA(supports_array) + Ns);\n\tq_batches = vector((int*)PyArray_DATA(q_batches_array), (int*)PyArray_DATA(q_batches_array) + Nb);\n\ts_batches = vector((int*)PyArray_DATA(s_batches_array), (int*)PyArray_DATA(s_batches_array) + Nb);\n\n\t// Create result containers\n\tvector neighbors_indices;\n\n\t// Compute results\n\t//batch_ordered_neighbors(queries, supports, q_batches, s_batches, neighbors_indices, radius);\n\tbatch_nanoflann_neighbors(queries, supports, q_batches, s_batches, neighbors_indices, radius);\n\n\t// Check result\n\tif (neighbors_indices.size() < 1)\n\t{\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Error\");\n\t\treturn NULL;\n\t}\n\n\t// Manage outputs\n\t// **************\n\n\t// Maximal number of neighbors\n\tint max_neighbors = neighbors_indices.size() / Nq;\n\n\t// Dimension of output containers\n\tnpy_intp* neighbors_dims = new npy_intp[2];\n\tneighbors_dims[0] = Nq;\n\tneighbors_dims[1] = max_neighbors;\n\n\t// Create output array\n\tPyObject* res_obj = PyArray_SimpleNew(2, neighbors_dims, NPY_INT);\n\tPyObject* ret = NULL;\n\n\t// Fill output array with values\n\tsize_t size_in_bytes = Nq * max_neighbors * sizeof(int);\n\tmemcpy(PyArray_DATA(res_obj), neighbors_indices.data(), size_in_bytes);\n\n\t// Merge results\n\tret = Py_BuildValue(\"N\", res_obj);\n\n\t// Clean up\n\t// ********\n\n\tPy_XDECREF(queries_array);\n\tPy_XDECREF(supports_array);\n\tPy_XDECREF(q_batches_array);\n\tPy_XDECREF(s_batches_array);\n\n\treturn ret;\n}\n"
},
{
"path": "thirdparty/kpconv/cpp_wrappers/cpp_subsampling/build.bat",
"content": "@echo off\npy setup.py build_ext --inplace\n\n\npause"
},
{
"path": "thirdparty/kpconv/cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.cpp",
"content": "\n#include \"grid_subsampling.h\"\n\n\nvoid grid_subsampling(vector& original_points,\n vector& subsampled_points,\n vector& original_features,\n vector& subsampled_features,\n vector& original_classes,\n vector& subsampled_classes,\n float sampleDl,\n int verbose) {\n\n\t// Initialize variables\n\t// ******************\n\n\t// Number of points in the cloud\n\tsize_t N = original_points.size();\n\n\t// Dimension of the features\n\tsize_t fdim = original_features.size() / N;\n\tsize_t ldim = original_classes.size() / N;\n\n\t// Limits of the cloud\n\tPointXYZ minCorner = min_point(original_points);\n\tPointXYZ maxCorner = max_point(original_points);\n\tPointXYZ originCorner = floor(minCorner * (1/sampleDl)) * sampleDl;\n\n\t// Dimensions of the grid\n\tsize_t sampleNX = (size_t)floor((maxCorner.x - originCorner.x) / sampleDl) + 1;\n\tsize_t sampleNY = (size_t)floor((maxCorner.y - originCorner.y) / sampleDl) + 1;\n\t//size_t sampleNZ = (size_t)floor((maxCorner.z - originCorner.z) / sampleDl) + 1;\n\n\t// Check if features and classes need to be processed\n\tbool use_feature = original_features.size() > 0;\n\tbool use_classes = original_classes.size() > 0;\n\n\n\t// Create the sampled map\n\t// **********************\n\n\t// Verbose parameters\n\tint i = 0;\n\tint nDisp = N / 100;\n\n\t// Initialize variables\n\tsize_t iX, iY, iZ, mapIdx;\n\tunordered_map data;\n\n\tfor (auto& p : original_points)\n\t{\n\t\t// Position of point in sample map\n\t\tiX = (size_t)floor((p.x - originCorner.x) / sampleDl);\n\t\tiY = (size_t)floor((p.y - originCorner.y) / sampleDl);\n\t\tiZ = (size_t)floor((p.z - originCorner.z) / sampleDl);\n\t\tmapIdx = iX + sampleNX*iY + sampleNX*sampleNY*iZ;\n\n\t\t// If not already created, create key\n\t\tif (data.count(mapIdx) < 1)\n\t\t\tdata.emplace(mapIdx, SampledData(fdim, ldim));\n\n\t\t// Fill the sample map\n\t\tif (use_feature && use_classes)\n\t\t\tdata[mapIdx].update_all(p, original_features.begin() + i * fdim, original_classes.begin() + i * ldim);\n\t\telse if (use_feature)\n\t\t\tdata[mapIdx].update_features(p, original_features.begin() + i * fdim);\n\t\telse if (use_classes)\n\t\t\tdata[mapIdx].update_classes(p, original_classes.begin() + i * ldim);\n\t\telse\n\t\t\tdata[mapIdx].update_points(p);\n\n\t\t// Display\n\t\ti++;\n\t\tif (verbose > 1 && i%nDisp == 0)\n\t\t\tstd::cout << \"\\rSampled Map : \" << std::setw(3) << i / nDisp << \"%\";\n\n\t}\n\n\t// Divide for barycentre and transfer to a vector\n\tsubsampled_points.reserve(data.size());\n\tif (use_feature)\n\t\tsubsampled_features.reserve(data.size() * fdim);\n\tif (use_classes)\n\t\tsubsampled_classes.reserve(data.size() * ldim);\n\tfor (auto& v : data)\n\t{\n\t\tsubsampled_points.push_back(v.second.point * (1.0 / v.second.count));\n\t\tif (use_feature)\n\t\t{\n\t\t float count = (float)v.second.count;\n\t\t transform(v.second.features.begin(),\n v.second.features.end(),\n v.second.features.begin(),\n [count](float f) { return f / count;});\n subsampled_features.insert(subsampled_features.end(),v.second.features.begin(),v.second.features.end());\n\t\t}\n\t\tif (use_classes)\n\t\t{\n\t\t for (int i = 0; i < ldim; i++)\n\t\t subsampled_classes.push_back(max_element(v.second.labels[i].begin(), v.second.labels[i].end(),\n\t\t [](const pair&a, const pair&b){return a.second < b.second;})->first);\n\t\t}\n\t}\n\n\treturn;\n}\n\n\nvoid batch_grid_subsampling(vector& original_points,\n vector& subsampled_points,\n vector& original_features,\n vector& subsampled_features,\n vector& original_classes,\n vector& subsampled_classes,\n vector& original_batches,\n vector& subsampled_batches,\n float sampleDl,\n int max_p)\n{\n\t// Initialize variables\n\t// ******************\n\n\tint b = 0;\n\tint sum_b = 0;\n\n\t// Number of points in the cloud\n\tsize_t N = original_points.size();\n\n\t// Dimension of the features\n\tsize_t fdim = original_features.size() / N;\n\tsize_t ldim = original_classes.size() / N;\n\n\t// Handle max_p = 0\n\tif (max_p < 1)\n\t max_p = N;\n\n\t// Loop over batches\n\t// *****************\n\n\tfor (b = 0; b < original_batches.size(); b++)\n\t{\n\n\t // Extract batch points features and labels\n\t vector b_o_points = vector(original_points.begin () + sum_b,\n\t original_points.begin () + sum_b + original_batches[b]);\n\n vector b_o_features;\n if (original_features.size() > 0)\n {\n b_o_features = vector(original_features.begin () + sum_b * fdim,\n original_features.begin () + (sum_b + original_batches[b]) * fdim);\n\t }\n\n\t vector b_o_classes;\n if (original_classes.size() > 0)\n {\n b_o_classes = vector(original_classes.begin () + sum_b * ldim,\n original_classes.begin () + sum_b + original_batches[b] * ldim);\n\t }\n\n\n // Create result containers\n vector b_s_points;\n vector b_s_features;\n vector b_s_classes;\n\n // Compute subsampling on current batch\n grid_subsampling(b_o_points,\n b_s_points,\n b_o_features,\n b_s_features,\n b_o_classes,\n b_s_classes,\n sampleDl,\n\t\t\t\t\t\t 0);\n\n // Stack batches points features and labels\n // ****************************************\n\n // If too many points remove some\n if (b_s_points.size() <= max_p)\n {\n subsampled_points.insert(subsampled_points.end(), b_s_points.begin(), b_s_points.end());\n\n if (original_features.size() > 0)\n subsampled_features.insert(subsampled_features.end(), b_s_features.begin(), b_s_features.end());\n\n if (original_classes.size() > 0)\n subsampled_classes.insert(subsampled_classes.end(), b_s_classes.begin(), b_s_classes.end());\n\n subsampled_batches.push_back(b_s_points.size());\n }\n else\n {\n subsampled_points.insert(subsampled_points.end(), b_s_points.begin(), b_s_points.begin() + max_p);\n\n if (original_features.size() > 0)\n subsampled_features.insert(subsampled_features.end(), b_s_features.begin(), b_s_features.begin() + max_p * fdim);\n\n if (original_classes.size() > 0)\n subsampled_classes.insert(subsampled_classes.end(), b_s_classes.begin(), b_s_classes.begin() + max_p * ldim);\n\n subsampled_batches.push_back(max_p);\n }\n\n // Stack new batch lengths\n sum_b += original_batches[b];\n\t}\n\n\treturn;\n}\n"
},
{
"path": "thirdparty/kpconv/cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.h",
"content": "\n\n#include \"../../cpp_utils/cloud/cloud.h\"\n\n#include \n#include \n\nusing namespace std;\n\nclass SampledData\n{\npublic:\n\n\t// Elements\n\t// ********\n\n\tint count;\n\tPointXYZ point;\n\tvector features;\n\tvector> labels;\n\n\n\t// Methods\n\t// *******\n\n\t// Constructor\n\tSampledData() \n\t{ \n\t\tcount = 0; \n\t\tpoint = PointXYZ();\n\t}\n\n\tSampledData(const size_t fdim, const size_t ldim)\n\t{\n\t\tcount = 0;\n\t\tpoint = PointXYZ();\n\t features = vector(fdim);\n\t labels = vector>(ldim);\n\t}\n\n\t// Method Update\n\tvoid update_all(const PointXYZ p, vector::iterator f_begin, vector::iterator l_begin)\n\t{\n\t\tcount += 1;\n\t\tpoint += p;\n\t\ttransform (features.begin(), features.end(), f_begin, features.begin(), plus());\n\t\tint i = 0;\n\t\tfor(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it)\n\t\t{\n\t\t labels[i][*it] += 1;\n\t\t i++;\n\t\t}\n\t\treturn;\n\t}\n\tvoid update_features(const PointXYZ p, vector::iterator f_begin)\n\t{\n\t\tcount += 1;\n\t\tpoint += p;\n\t\ttransform (features.begin(), features.end(), f_begin, features.begin(), plus());\n\t\treturn;\n\t}\n\tvoid update_classes(const PointXYZ p, vector::iterator l_begin)\n\t{\n\t\tcount += 1;\n\t\tpoint += p;\n\t\tint i = 0;\n\t\tfor(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it)\n\t\t{\n\t\t labels[i][*it] += 1;\n\t\t i++;\n\t\t}\n\t\treturn;\n\t}\n\tvoid update_points(const PointXYZ p)\n\t{\n\t\tcount += 1;\n\t\tpoint += p;\n\t\treturn;\n\t}\n};\n\nvoid grid_subsampling(vector& original_points,\n vector& subsampled_points,\n vector& original_features,\n vector& subsampled_features,\n vector& original_classes,\n vector& subsampled_classes,\n float sampleDl,\n int verbose);\n\nvoid batch_grid_subsampling(vector& original_points,\n vector& subsampled_points,\n vector& original_features,\n vector& subsampled_features,\n vector& original_classes,\n vector& subsampled_classes,\n vector& original_batches,\n vector& subsampled_batches,\n float sampleDl,\n int max_p);\n\n"
},
{
"path": "thirdparty/kpconv/cpp_wrappers/cpp_subsampling/setup.py",
"content": "from distutils.core import setup, Extension\nimport numpy.distutils.misc_util\n\n# Adding OpenCV to project\n# ************************\n\n# Adding sources of the project\n# *****************************\n\nSOURCES = [\"../cpp_utils/cloud/cloud.cpp\",\n \"grid_subsampling/grid_subsampling.cpp\",\n \"wrapper.cpp\"]\n\nmodule = Extension(name=\"grid_subsampling\",\n sources=SOURCES,\n extra_compile_args=['-std=c++11',\n '-D_GLIBCXX_USE_CXX11_ABI=0'])\n\n\nsetup(ext_modules=[module], include_dirs=numpy.distutils.misc_util.get_numpy_include_dirs())\n\n\n\n\n\n\n\n\n"
},
{
"path": "thirdparty/kpconv/cpp_wrappers/cpp_subsampling/wrapper.cpp",
"content": "#include \n#include \n#include \"grid_subsampling/grid_subsampling.h\"\n#include \n\n\n\n// docstrings for our module\n// *************************\n\nstatic char module_docstring[] = \"This module provides an interface for the subsampling of a batch of stacked pointclouds\";\n\nstatic char subsample_docstring[] = \"function subsampling a pointcloud\";\n\nstatic char subsample_batch_docstring[] = \"function subsampling a batch of stacked pointclouds\";\n\n\n// Declare the functions\n// *********************\n\nstatic PyObject *cloud_subsampling(PyObject* self, PyObject* args, PyObject* keywds);\nstatic PyObject *batch_subsampling(PyObject *self, PyObject *args, PyObject *keywds);\n\n\n// Specify the members of the module\n// *********************************\n\nstatic PyMethodDef module_methods[] = \n{\n\t{ \"subsample\", (PyCFunction)cloud_subsampling, METH_VARARGS | METH_KEYWORDS, subsample_docstring },\n\t{ \"subsample_batch\", (PyCFunction)batch_subsampling, METH_VARARGS | METH_KEYWORDS, subsample_batch_docstring },\n\t{NULL, NULL, 0, NULL}\n};\n\n\n// Initialize the module\n// *********************\n\nstatic struct PyModuleDef moduledef = \n{\n PyModuleDef_HEAD_INIT,\n \"grid_subsampling\", // m_name\n module_docstring, // m_doc\n -1, // m_size\n module_methods, // m_methods\n NULL, // m_reload\n NULL, // m_traverse\n NULL, // m_clear\n NULL, // m_free\n};\n\nPyMODINIT_FUNC PyInit_grid_subsampling(void)\n{\n import_array();\n\treturn PyModule_Create(&moduledef);\n}\n\n\n// Definition of the batch_subsample method\n// **********************************\n\nstatic PyObject* batch_subsampling(PyObject* self, PyObject* args, PyObject* keywds)\n{\n\n\t// Manage inputs\n\t// *************\n\n\t// Args containers\n\tPyObject* points_obj = NULL;\n\tPyObject* features_obj = NULL;\n\tPyObject* classes_obj = NULL;\n\tPyObject* batches_obj = NULL;\n\n\t// Keywords containers\n\tstatic char* kwlist[] = { \"points\", \"batches\", \"features\", \"classes\", \"sampleDl\", \"method\", \"max_p\", \"verbose\", NULL };\n\tfloat sampleDl = 0.1;\n\tconst char* method_buffer = \"barycenters\";\n\tint verbose = 0;\n\tint max_p = 0;\n\n\t// Parse the input \n\tif (!PyArg_ParseTupleAndKeywords(args, keywds, \"OO|$OOfsii\", kwlist, &points_obj, &batches_obj, &features_obj, &classes_obj, &sampleDl, &method_buffer, &max_p, &verbose))\n\t{\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Error parsing arguments\");\n\t\treturn NULL;\n\t}\n\n\t// Get the method argument\n\tstring method(method_buffer);\n\n\t// Interpret method\n\tif (method.compare(\"barycenters\") && method.compare(\"voxelcenters\"))\n\t{\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Error parsing method. Valid method names are \\\"barycenters\\\" and \\\"voxelcenters\\\" \");\n\t\treturn NULL;\n\t}\n\n\t// Check if using features or classes\n\tbool use_feature = true, use_classes = true;\n\tif (features_obj == NULL)\n\t\tuse_feature = false;\n\tif (classes_obj == NULL)\n\t\tuse_classes = false;\n\n\t// Interpret the input objects as numpy arrays.\n\tPyObject* points_array = PyArray_FROM_OTF(points_obj, NPY_FLOAT, NPY_IN_ARRAY);\n\tPyObject* batches_array = PyArray_FROM_OTF(batches_obj, NPY_INT, NPY_IN_ARRAY);\n\tPyObject* features_array = NULL;\n\tPyObject* classes_array = NULL;\n\tif (use_feature)\n\t\tfeatures_array = PyArray_FROM_OTF(features_obj, NPY_FLOAT, NPY_IN_ARRAY);\n\tif (use_classes)\n\t\tclasses_array = PyArray_FROM_OTF(classes_obj, NPY_INT, NPY_IN_ARRAY);\n\n\t// Verify data was load correctly.\n\tif (points_array == NULL)\n\t{\n\t\tPy_XDECREF(points_array);\n\t\tPy_XDECREF(batches_array);\n\t\tPy_XDECREF(classes_array);\n\t\tPy_XDECREF(features_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Error converting input points to numpy arrays of type float32\");\n\t\treturn NULL;\n\t}\n\tif (batches_array == NULL)\n\t{\n\t\tPy_XDECREF(points_array);\n\t\tPy_XDECREF(batches_array);\n\t\tPy_XDECREF(classes_array);\n\t\tPy_XDECREF(features_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Error converting input batches to numpy arrays of type int32\");\n\t\treturn NULL;\n\t}\n\tif (use_feature && features_array == NULL)\n\t{\n\t\tPy_XDECREF(points_array);\n\t\tPy_XDECREF(batches_array);\n\t\tPy_XDECREF(classes_array);\n\t\tPy_XDECREF(features_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Error converting input features to numpy arrays of type float32\");\n\t\treturn NULL;\n\t}\n\tif (use_classes && classes_array == NULL)\n\t{\n\t\tPy_XDECREF(points_array);\n\t\tPy_XDECREF(batches_array);\n\t\tPy_XDECREF(classes_array);\n\t\tPy_XDECREF(features_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Error converting input classes to numpy arrays of type int32\");\n\t\treturn NULL;\n\t}\n\n\t// Check that the input array respect the dims\n\tif ((int)PyArray_NDIM(points_array) != 2 || (int)PyArray_DIM(points_array, 1) != 3)\n\t{\n\t\tPy_XDECREF(points_array);\n\t\tPy_XDECREF(batches_array);\n\t\tPy_XDECREF(classes_array);\n\t\tPy_XDECREF(features_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Wrong dimensions : points.shape is not (N, 3)\");\n\t\treturn NULL;\n\t}\n\tif ((int)PyArray_NDIM(batches_array) > 1)\n\t{\n\t\tPy_XDECREF(points_array);\n\t\tPy_XDECREF(batches_array);\n\t\tPy_XDECREF(classes_array);\n\t\tPy_XDECREF(features_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Wrong dimensions : batches.shape is not (B,) \");\n\t\treturn NULL;\n\t}\n\tif (use_feature && ((int)PyArray_NDIM(features_array) != 2))\n\t{\n\t\tPy_XDECREF(points_array);\n\t\tPy_XDECREF(batches_array);\n\t\tPy_XDECREF(classes_array);\n\t\tPy_XDECREF(features_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Wrong dimensions : features.shape is not (N, d)\");\n\t\treturn NULL;\n\t}\n\n\tif (use_classes && (int)PyArray_NDIM(classes_array) > 2)\n\t{\n\t\tPy_XDECREF(points_array);\n\t\tPy_XDECREF(batches_array);\n\t\tPy_XDECREF(classes_array);\n\t\tPy_XDECREF(features_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Wrong dimensions : classes.shape is not (N,) or (N, d)\");\n\t\treturn NULL;\n\t}\n\n\t// Number of points\n\tint N = (int)PyArray_DIM(points_array, 0);\n\n\t// Number of batches\n\tint Nb = (int)PyArray_DIM(batches_array, 0);\n\n\t// Dimension of the features\n\tint fdim = 0;\n\tif (use_feature)\n\t\tfdim = (int)PyArray_DIM(features_array, 1);\n\n\t//Dimension of labels\n\tint ldim = 1;\n\tif (use_classes && (int)PyArray_NDIM(classes_array) == 2)\n\t\tldim = (int)PyArray_DIM(classes_array, 1);\n\n\t// Check that the input array respect the number of points\n\tif (use_feature && (int)PyArray_DIM(features_array, 0) != N)\n\t{\n\t\tPy_XDECREF(points_array);\n\t\tPy_XDECREF(batches_array);\n\t\tPy_XDECREF(classes_array);\n\t\tPy_XDECREF(features_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Wrong dimensions : features.shape is not (N, d)\");\n\t\treturn NULL;\n\t}\n\tif (use_classes && (int)PyArray_DIM(classes_array, 0) != N)\n\t{\n\t\tPy_XDECREF(points_array);\n\t\tPy_XDECREF(batches_array);\n\t\tPy_XDECREF(classes_array);\n\t\tPy_XDECREF(features_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Wrong dimensions : classes.shape is not (N,) or (N, d)\");\n\t\treturn NULL;\n\t}\n\n\n\t// Call the C++ function\n\t// *********************\n\n\t// Create pyramid\n\tif (verbose > 0)\n\t\tcout << \"Computing cloud pyramid with support points: \" << endl;\n\n\n\t// Convert PyArray to Cloud C++ class\n\tvector original_points;\n\tvector original_batches;\n\tvector original_features;\n\tvector original_classes;\n\toriginal_points = vector((PointXYZ*)PyArray_DATA(points_array), (PointXYZ*)PyArray_DATA(points_array) + N);\n\toriginal_batches = vector((int*)PyArray_DATA(batches_array), (int*)PyArray_DATA(batches_array) + Nb);\n\tif (use_feature)\n\t\toriginal_features = vector((float*)PyArray_DATA(features_array), (float*)PyArray_DATA(features_array) + N * fdim);\n\tif (use_classes)\n\t\toriginal_classes = vector((int*)PyArray_DATA(classes_array), (int*)PyArray_DATA(classes_array) + N * ldim);\n\n\t// Subsample\n\tvector subsampled_points;\n\tvector subsampled_features;\n\tvector subsampled_classes;\n\tvector subsampled_batches;\n\tbatch_grid_subsampling(original_points,\n\t\t\t\t\t\t\tsubsampled_points,\n\t\t\t\t\t\t\toriginal_features,\n\t\t\t\t\t\t\tsubsampled_features,\n\t\t\t\t\t\t\toriginal_classes,\n\t\t\t\t\t\t\tsubsampled_classes,\n\t\t\t\t\t\t\toriginal_batches,\n\t\t\t\t\t\t\tsubsampled_batches,\n\t\t\t\t\t\t\tsampleDl,\n\t\t\t\t\t\t\tmax_p);\n\n\t// Check result\n\tif (subsampled_points.size() < 1)\n\t{\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Error\");\n\t\treturn NULL;\n\t}\n\n\t// Manage outputs\n\t// **************\n\n\t// Dimension of input containers\n\tnpy_intp* point_dims = new npy_intp[2];\n\tpoint_dims[0] = subsampled_points.size();\n\tpoint_dims[1] = 3;\n\tnpy_intp* feature_dims = new npy_intp[2];\n\tfeature_dims[0] = subsampled_points.size();\n\tfeature_dims[1] = fdim;\n\tnpy_intp* classes_dims = new npy_intp[2];\n\tclasses_dims[0] = subsampled_points.size();\n\tclasses_dims[1] = ldim;\n\tnpy_intp* batches_dims = new npy_intp[1];\n\tbatches_dims[0] = Nb;\n\n\t// Create output array\n\tPyObject* res_points_obj = PyArray_SimpleNew(2, point_dims, NPY_FLOAT);\n\tPyObject* res_batches_obj = PyArray_SimpleNew(1, batches_dims, NPY_INT);\n\tPyObject* res_features_obj = NULL;\n\tPyObject* res_classes_obj = NULL;\n\tPyObject* ret = NULL;\n\n\t// Fill output array with values\n\tsize_t size_in_bytes = subsampled_points.size() * 3 * sizeof(float);\n\tmemcpy(PyArray_DATA(res_points_obj), subsampled_points.data(), size_in_bytes);\n\tsize_in_bytes = Nb * sizeof(int);\n\tmemcpy(PyArray_DATA(res_batches_obj), subsampled_batches.data(), size_in_bytes);\n\tif (use_feature)\n\t{\n\t\tsize_in_bytes = subsampled_points.size() * fdim * sizeof(float);\n\t\tres_features_obj = PyArray_SimpleNew(2, feature_dims, NPY_FLOAT);\n\t\tmemcpy(PyArray_DATA(res_features_obj), subsampled_features.data(), size_in_bytes);\n\t}\n\tif (use_classes)\n\t{\n\t\tsize_in_bytes = subsampled_points.size() * ldim * sizeof(int);\n\t\tres_classes_obj = PyArray_SimpleNew(2, classes_dims, NPY_INT);\n\t\tmemcpy(PyArray_DATA(res_classes_obj), subsampled_classes.data(), size_in_bytes);\n\t}\n\n\n\t// Merge results\n\tif (use_feature && use_classes)\n\t\tret = Py_BuildValue(\"NNNN\", res_points_obj, res_batches_obj, res_features_obj, res_classes_obj);\n\telse if (use_feature)\n\t\tret = Py_BuildValue(\"NNN\", res_points_obj, res_batches_obj, res_features_obj);\n\telse if (use_classes)\n\t\tret = Py_BuildValue(\"NNN\", res_points_obj, res_batches_obj, res_classes_obj);\n\telse\n\t\tret = Py_BuildValue(\"NN\", res_points_obj, res_batches_obj);\n\n\t// Clean up\n\t// ********\n\n\tPy_DECREF(points_array);\n\tPy_DECREF(batches_array);\n\tPy_XDECREF(features_array);\n\tPy_XDECREF(classes_array);\n\n\treturn ret;\n}\n\n// Definition of the subsample method\n// ****************************************\n\nstatic PyObject* cloud_subsampling(PyObject* self, PyObject* args, PyObject* keywds)\n{\n\n\t// Manage inputs\n\t// *************\n\n\t// Args containers\n\tPyObject* points_obj = NULL;\n\tPyObject* features_obj = NULL;\n\tPyObject* classes_obj = NULL;\n\n\t// Keywords containers\n\tstatic char* kwlist[] = { \"points\", \"features\", \"classes\", \"sampleDl\", \"method\", \"verbose\", NULL };\n\tfloat sampleDl = 0.1;\n\tconst char* method_buffer = \"barycenters\";\n\tint verbose = 0;\n\n\t// Parse the input \n\tif (!PyArg_ParseTupleAndKeywords(args, keywds, \"O|$OOfsi\", kwlist, &points_obj, &features_obj, &classes_obj, &sampleDl, &method_buffer, &verbose))\n\t{\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Error parsing arguments\");\n\t\treturn NULL;\n\t}\n\n\t// Get the method argument\n\tstring method(method_buffer);\n\n\t// Interpret method\n\tif (method.compare(\"barycenters\") && method.compare(\"voxelcenters\"))\n\t{\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Error parsing method. Valid method names are \\\"barycenters\\\" and \\\"voxelcenters\\\" \");\n\t\treturn NULL;\n\t}\n\n\t// Check if using features or classes\n\tbool use_feature = true, use_classes = true;\n\tif (features_obj == NULL)\n\t\tuse_feature = false;\n\tif (classes_obj == NULL)\n\t\tuse_classes = false;\n\n\t// Interpret the input objects as numpy arrays.\n\tPyObject* points_array = PyArray_FROM_OTF(points_obj, NPY_FLOAT, NPY_IN_ARRAY);\n\tPyObject* features_array = NULL;\n\tPyObject* classes_array = NULL;\n\tif (use_feature)\n\t\tfeatures_array = PyArray_FROM_OTF(features_obj, NPY_FLOAT, NPY_IN_ARRAY);\n\tif (use_classes)\n\t\tclasses_array = PyArray_FROM_OTF(classes_obj, NPY_INT, NPY_IN_ARRAY);\n\n\t// Verify data was load correctly.\n\tif (points_array == NULL)\n\t{\n\t\tPy_XDECREF(points_array);\n\t\tPy_XDECREF(classes_array);\n\t\tPy_XDECREF(features_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Error converting input points to numpy arrays of type float32\");\n\t\treturn NULL;\n\t}\n\tif (use_feature && features_array == NULL)\n\t{\n\t\tPy_XDECREF(points_array);\n\t\tPy_XDECREF(classes_array);\n\t\tPy_XDECREF(features_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Error converting input features to numpy arrays of type float32\");\n\t\treturn NULL;\n\t}\n\tif (use_classes && classes_array == NULL)\n\t{\n\t\tPy_XDECREF(points_array);\n\t\tPy_XDECREF(classes_array);\n\t\tPy_XDECREF(features_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Error converting input classes to numpy arrays of type int32\");\n\t\treturn NULL;\n\t}\n\n\t// Check that the input array respect the dims\n\tif ((int)PyArray_NDIM(points_array) != 2 || (int)PyArray_DIM(points_array, 1) != 3)\n\t{\n\t\tPy_XDECREF(points_array);\n\t\tPy_XDECREF(classes_array);\n\t\tPy_XDECREF(features_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Wrong dimensions : points.shape is not (N, 3)\");\n\t\treturn NULL;\n\t}\n\tif (use_feature && ((int)PyArray_NDIM(features_array) != 2))\n\t{\n\t\tPy_XDECREF(points_array);\n\t\tPy_XDECREF(classes_array);\n\t\tPy_XDECREF(features_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Wrong dimensions : features.shape is not (N, d)\");\n\t\treturn NULL;\n\t}\n\n\tif (use_classes && (int)PyArray_NDIM(classes_array) > 2)\n\t{\n\t\tPy_XDECREF(points_array);\n\t\tPy_XDECREF(classes_array);\n\t\tPy_XDECREF(features_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Wrong dimensions : classes.shape is not (N,) or (N, d)\");\n\t\treturn NULL;\n\t}\n\n\t// Number of points\n\tint N = (int)PyArray_DIM(points_array, 0);\n\n\t// Dimension of the features\n\tint fdim = 0;\n\tif (use_feature)\n\t\tfdim = (int)PyArray_DIM(features_array, 1);\n\n\t//Dimension of labels\n\tint ldim = 1;\n\tif (use_classes && (int)PyArray_NDIM(classes_array) == 2)\n\t\tldim = (int)PyArray_DIM(classes_array, 1);\n\n\t// Check that the input array respect the number of points\n\tif (use_feature && (int)PyArray_DIM(features_array, 0) != N)\n\t{\n\t\tPy_XDECREF(points_array);\n\t\tPy_XDECREF(classes_array);\n\t\tPy_XDECREF(features_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Wrong dimensions : features.shape is not (N, d)\");\n\t\treturn NULL;\n\t}\n\tif (use_classes && (int)PyArray_DIM(classes_array, 0) != N)\n\t{\n\t\tPy_XDECREF(points_array);\n\t\tPy_XDECREF(classes_array);\n\t\tPy_XDECREF(features_array);\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Wrong dimensions : classes.shape is not (N,) or (N, d)\");\n\t\treturn NULL;\n\t}\n\n\n\t// Call the C++ function\n\t// *********************\n\n\t// Create pyramid\n\tif (verbose > 0)\n\t\tcout << \"Computing cloud pyramid with support points: \" << endl;\n\n\n\t// Convert PyArray to Cloud C++ class\n\tvector original_points;\n\tvector original_features;\n\tvector original_classes;\n\toriginal_points = vector((PointXYZ*)PyArray_DATA(points_array), (PointXYZ*)PyArray_DATA(points_array) + N);\n\tif (use_feature)\n\t\toriginal_features = vector((float*)PyArray_DATA(features_array), (float*)PyArray_DATA(features_array) + N * fdim);\n\tif (use_classes)\n\t\toriginal_classes = vector((int*)PyArray_DATA(classes_array), (int*)PyArray_DATA(classes_array) + N * ldim);\n\n\t// Subsample\n\tvector subsampled_points;\n\tvector subsampled_features;\n\tvector subsampled_classes;\n\tgrid_subsampling(original_points,\n\t\tsubsampled_points,\n\t\toriginal_features,\n\t\tsubsampled_features,\n\t\toriginal_classes,\n\t\tsubsampled_classes,\n\t\tsampleDl,\n\t\tverbose);\n\n\t// Check result\n\tif (subsampled_points.size() < 1)\n\t{\n\t\tPyErr_SetString(PyExc_RuntimeError, \"Error\");\n\t\treturn NULL;\n\t}\n\n\t// Manage outputs\n\t// **************\n\n\t// Dimension of input containers\n\tnpy_intp* point_dims = new npy_intp[2];\n\tpoint_dims[0] = subsampled_points.size();\n\tpoint_dims[1] = 3;\n\tnpy_intp* feature_dims = new npy_intp[2];\n\tfeature_dims[0] = subsampled_points.size();\n\tfeature_dims[1] = fdim;\n\tnpy_intp* classes_dims = new npy_intp[2];\n\tclasses_dims[0] = subsampled_points.size();\n\tclasses_dims[1] = ldim;\n\n\t// Create output array\n\tPyObject* res_points_obj = PyArray_SimpleNew(2, point_dims, NPY_FLOAT);\n\tPyObject* res_features_obj = NULL;\n\tPyObject* res_classes_obj = NULL;\n\tPyObject* ret = NULL;\n\n\t// Fill output array with values\n\tsize_t size_in_bytes = subsampled_points.size() * 3 * sizeof(float);\n\tmemcpy(PyArray_DATA(res_points_obj), subsampled_points.data(), size_in_bytes);\n\tif (use_feature)\n\t{\n\t\tsize_in_bytes = subsampled_points.size() * fdim * sizeof(float);\n\t\tres_features_obj = PyArray_SimpleNew(2, feature_dims, NPY_FLOAT);\n\t\tmemcpy(PyArray_DATA(res_features_obj), subsampled_features.data(), size_in_bytes);\n\t}\n\tif (use_classes)\n\t{\n\t\tsize_in_bytes = subsampled_points.size() * ldim * sizeof(int);\n\t\tres_classes_obj = PyArray_SimpleNew(2, classes_dims, NPY_INT);\n\t\tmemcpy(PyArray_DATA(res_classes_obj), subsampled_classes.data(), size_in_bytes);\n\t}\n\n\n\t// Merge results\n\tif (use_feature && use_classes)\n\t\tret = Py_BuildValue(\"NNN\", res_points_obj, res_features_obj, res_classes_obj);\n\telse if (use_feature)\n\t\tret = Py_BuildValue(\"NN\", res_points_obj, res_features_obj);\n\telse if (use_classes)\n\t\tret = Py_BuildValue(\"NN\", res_points_obj, res_classes_obj);\n\telse\n\t\tret = Py_BuildValue(\"N\", res_points_obj);\n\n\t// Clean up\n\t// ********\n\n\tPy_DECREF(points_array);\n\tPy_XDECREF(features_array);\n\tPy_XDECREF(classes_array);\n\n\treturn ret;\n}"
},
{
"path": "thirdparty/kpconv/cpp_wrappers/cpp_utils/cloud/cloud.cpp",
"content": "//\n//\n//\t\t0==========================0\n//\t\t| Local feature test |\n//\t\t0==========================0\n//\n//\t\tversion 1.0 : \n//\t\t\t> \n//\n//---------------------------------------------------\n//\n//\t\tCloud source :\n//\t\tDefine usefull Functions/Methods\n//\n//----------------------------------------------------\n//\n//\t\tHugues THOMAS - 10/02/2017\n//\n\n\n#include \"cloud.h\"\n\n\n// Getters\n// *******\n\nPointXYZ max_point(std::vector points)\n{\n\t// Initialize limits\n\tPointXYZ maxP(points[0]);\n\n\t// Loop over all points\n\tfor (auto p : points)\n\t{\n\t\tif (p.x > maxP.x)\n\t\t\tmaxP.x = p.x;\n\n\t\tif (p.y > maxP.y)\n\t\t\tmaxP.y = p.y;\n\n\t\tif (p.z > maxP.z)\n\t\t\tmaxP.z = p.z;\n\t}\n\n\treturn maxP;\n}\n\nPointXYZ min_point(std::vector points)\n{\n\t// Initialize limits\n\tPointXYZ minP(points[0]);\n\n\t// Loop over all points\n\tfor (auto p : points)\n\t{\n\t\tif (p.x < minP.x)\n\t\t\tminP.x = p.x;\n\n\t\tif (p.y < minP.y)\n\t\t\tminP.y = p.y;\n\n\t\tif (p.z < minP.z)\n\t\t\tminP.z = p.z;\n\t}\n\n\treturn minP;\n}"
},
{
"path": "thirdparty/kpconv/cpp_wrappers/cpp_utils/cloud/cloud.h",
"content": "//\n//\n//\t\t0==========================0\n//\t\t| Local feature test |\n//\t\t0==========================0\n//\n//\t\tversion 1.0 : \n//\t\t\t> \n//\n//---------------------------------------------------\n//\n//\t\tCloud header\n//\n//----------------------------------------------------\n//\n//\t\tHugues THOMAS - 10/02/2017\n//\n\n\n# pragma once\n\n#include \n#include \n#include