├── .gitattributes ├── LICENSE ├── README.md ├── build ├── lib.linux-x86_64-3.6 │ └── src │ │ └── utils │ │ ├── libkdtree │ │ └── pykdtree │ │ │ └── kdtree.cpython-36m-x86_64-linux-gnu.so │ │ ├── libmcubes │ │ └── mcubes.cpython-36m-x86_64-linux-gnu.so │ │ ├── libmesh │ │ └── triangle_hash.cpython-36m-x86_64-linux-gnu.so │ │ ├── libmise │ │ └── mise.cpython-36m-x86_64-linux-gnu.so │ │ ├── libsimplify │ │ └── simplify_mesh.cpython-36m-x86_64-linux-gnu.so │ │ └── libvoxelize │ │ └── voxelize.cpython-36m-x86_64-linux-gnu.so └── temp.linux-x86_64-3.6 │ └── src │ └── utils │ ├── libkdtree │ └── pykdtree │ │ ├── _kdtree_core.o │ │ └── kdtree.o │ ├── libmcubes │ ├── marchingcubes.o │ ├── mcubes.o │ └── pywrapper.o │ ├── libmesh │ └── triangle_hash.o │ ├── libmise │ └── mise.o │ ├── libsimplify │ └── simplify_mesh.o │ └── libvoxelize │ └── voxelize.o ├── configs ├── .DS_Store ├── default.yaml ├── pointcloud │ ├── room_3plane_attention_sub.yaml │ └── shapenet_3plane_attention_sub.yaml └── voxel │ └── room_grid64_attention_sub_2gpu.yaml ├── environment.yaml ├── eval_meshes.py ├── generate.py ├── media └── pipeline.png ├── modules ├── .DS_Store ├── __init__.py ├── ball_query.py ├── frustum.py ├── functional │ ├── __init__.py │ ├── backend.py │ ├── ball_query.py │ ├── devoxelization.py │ ├── grouping.py │ ├── interpolatation.py │ ├── loss.py │ ├── sampling.py │ ├── src │ │ ├── ball_query │ │ │ ├── ball_query.cpp │ │ │ ├── ball_query.cppZone.Identifier │ │ │ ├── ball_query.cu │ │ │ ├── ball_query.cuh │ │ │ ├── ball_query.cuhZone.Identifier │ │ │ ├── ball_query.cuZone.Identifier │ │ │ ├── ball_query.hpp │ │ │ └── ball_query.hppZone.Identifier │ │ ├── bindings.cpp │ │ ├── cuda_utils.cuh │ │ ├── grouping │ │ │ ├── grouping.cpp │ │ │ ├── grouping.cppZone.Identifier │ │ │ ├── grouping.cu │ │ │ ├── grouping.cuh │ │ │ ├── grouping.cuhZone.Identifier │ │ │ ├── grouping.cuZone.Identifier │ │ │ ├── grouping.hpp │ │ │ └── grouping.hppZone.Identifier │ │ ├── interpolate │ │ │ ├── neighbor_interpolate.cpp │ │ │ ├── neighbor_interpolate.cppZone.Identifier │ │ │ ├── neighbor_interpolate.cu │ │ │ ├── neighbor_interpolate.cuh │ │ │ ├── neighbor_interpolate.cuhZone.Identifier │ │ │ ├── neighbor_interpolate.cuZone.Identifier │ │ │ ├── neighbor_interpolate.hpp │ │ │ ├── neighbor_interpolate.hppZone.Identifier │ │ │ ├── trilinear_devox.cpp │ │ │ ├── trilinear_devox.cppZone.Identifier │ │ │ ├── trilinear_devox.cu │ │ │ ├── trilinear_devox.cuh │ │ │ ├── trilinear_devox.cuhZone.Identifier │ │ │ ├── trilinear_devox.cuZone.Identifier │ │ │ ├── trilinear_devox.hpp │ │ │ └── trilinear_devox.hppZone.Identifier │ │ ├── sampling │ │ │ ├── sampling.cpp │ │ │ ├── sampling.cppZone.Identifier │ │ │ ├── sampling.cu │ │ │ ├── sampling.cuh │ │ │ ├── sampling.cuhZone.Identifier │ │ │ ├── sampling.cuZone.Identifier │ │ │ ├── sampling.hpp │ │ │ └── sampling.hppZone.Identifier │ │ ├── utils.hpp │ │ └── voxelization │ │ │ ├── vox.cpp │ │ │ ├── vox.cppZone.Identifier │ │ │ ├── vox.cu │ │ │ ├── vox.cuh │ │ │ ├── vox.cuhZone.Identifier │ │ │ ├── vox.cuZone.Identifier │ │ │ ├── vox.hpp │ │ │ └── vox.hppZone.Identifier │ └── voxelization.py ├── loss.py ├── pointnet.py ├── se.py ├── shared_mlp.py ├── utils.py └── voxelization.py ├── scripts ├── dataset_matterport │ └── build_dataset.py ├── dataset_scannet │ ├── SensorData.py │ ├── build_dataset.py │ └── scannet_generation.log ├── dataset_synthetic_room │ └── build_dataset.py ├── download_data.sh ├── download_demo_data.sh └── download_shapenet_data.sh ├── setup.py ├── src ├── __init__.py ├── attention.py ├── checkpoints.py ├── common.py ├── config.py ├── conv_onet │ ├── __init__.py │ ├── config.py │ ├── generation.py │ ├── models │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-36.pyc │ │ │ └── decoder.cpython-36.pyc │ │ └── decoder.py │ └── training.py ├── data │ ├── __init__.py │ ├── core.py │ ├── fields.py │ └── transforms.py ├── encoder │ ├── __init__.py │ ├── pointnet.py │ ├── pointnetpp.py │ ├── unet.py │ ├── unet3d.py │ └── voxels.py ├── eval.py ├── layers.py ├── training.py ├── transformer.py └── utils │ ├── .DS_Store │ ├── __init__.py │ ├── binvox_rw.py │ ├── icp.py │ ├── io.py │ ├── libkdtree │ ├── .DS_Store │ ├── .gitignore │ ├── MANIFEST.in │ ├── README │ ├── README.rst │ ├── __init__.py │ ├── pykdtree │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ └── __init__.cpython-36.pyc │ │ ├── _kdtree_core.c │ │ ├── _kdtree_core.c.mako │ │ ├── kdtree.c │ │ ├── kdtree.cpython-36m-x86_64-linux-gnu.so │ │ ├── kdtree.pyx │ │ ├── render_template.py │ │ └── test_tree.py │ └── setup.cfg │ ├── libmcubes │ ├── .DS_Store │ ├── .gitignore │ ├── LICENSE │ ├── README.rst │ ├── __init__.py │ ├── exporter.py │ ├── marchingcubes.cpp │ ├── marchingcubes.h │ ├── mcubes.cpp │ ├── mcubes.cpython-36m-x86_64-linux-gnu.so │ ├── mcubes.pyx │ ├── pyarray_symbol.h │ ├── pyarraymodule.h │ ├── pywrapper.cpp │ └── pywrapper.h │ ├── libmesh │ ├── .DS_Store │ ├── .gitignore │ ├── __init__.py │ ├── inside_mesh.py │ ├── triangle_hash.cpython-36m-x86_64-linux-gnu.so │ └── triangle_hash.pyx │ ├── libmise │ ├── .DS_Store │ ├── .gitignore │ ├── __init__.py │ ├── mise.cpython-36m-x86_64-linux-gnu.so │ ├── mise.pyx │ └── test.py │ ├── libsimplify │ ├── .DS_Store │ ├── Simplify.h │ ├── __init__.py │ ├── simplify_mesh.cpp │ ├── simplify_mesh.cpython-36m-x86_64-linux-gnu.so │ ├── simplify_mesh.pyx │ └── test.py │ ├── libvoxelize │ ├── .DS_Store │ ├── .gitignore │ ├── __init__.py │ ├── tribox2.h │ ├── voxelize.cpython-36m-x86_64-linux-gnu.so │ └── voxelize.pyx │ ├── mesh.py │ ├── visualize.py │ └── voxels.py └── train.py /.gitattributes: -------------------------------------------------------------------------------- 1 | # Auto detect text files and perform LF normalization 2 | * text=auto 3 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 Songyou Peng, Michael Niemeyer, Lars Mescheder, Marc Pollefeys, Andreas Geiger 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ALTO: Alternating Latent Topologies for Implicit 3D Reconstruction 2 | [**Paper**](https://arxiv.org/abs/2212.04096) | [**Project Page**](https://visual.ee.ucla.edu/alto.htm/)
3 | 4 |
5 | 6 |
7 | 8 | This repository contains the implementation of the paper: 9 | 10 | ALTO: Alternating Latent Topologies for Implicit 3D Reconstruction 11 | 12 | If you find our code or paper useful, please consider citing 13 | ```bibtex 14 | @inproceedings{Wang2023CVPR, 15 | title = {ALTO: Alternating Latent Topologies for Implicit 3D Reconstruction}, 16 | author = {Wang, Zhen and Zhou, Shijie and Park, Jeong Joon and Paschalidou, Despoina and You, Suya and Wetzstein, Gordon and Guibas, Leonidas and Kadambi, Achuta}, 17 | booktitle = {Proceedings IEEE Conf. on Computer Vision and Pattern Recognition (CVPR)}, 18 | year = {2023} 19 | } 20 | ``` 21 | 22 | ## Installation 23 | You can create an anaconda environment called `alto` using 24 | ``` 25 | conda env create -f environment.yaml 26 | conda activate alto 27 | ``` 28 | **Note**: you might need to install **torch-scatter** mannually following [the official instruction](https://github.com/rusty1s/pytorch_scatter#pytorch-140): 29 | ``` 30 | pip install torch-scatter==2.0.4 -f https://pytorch-geometric.com/whl/torch-1.4.0+cu101.html 31 | ``` 32 | 33 | Next, compile the extension modules. 34 | You can do this via 35 | ``` 36 | python setup.py build_ext --inplace 37 | ``` 38 | 39 | ## Dataset 40 | In this paper, we consider 3 different datasets: 41 | ### Synthetic Indoor Scene Dataset 42 | You can download the preprocessed data (144 GB) using 43 | 44 | ``` 45 | bash scripts/download_data.sh 46 | ``` 47 | 48 | This script should download and unpack the data automatically into the `data/synthetic_room_dataset` folder. 49 | 50 | ### ShapeNet 51 | You can download the dataset (73.4 GB) by running the [script](https://github.com/autonomousvision/occupancy_networks#preprocessed-data) from Occupancy Networks. After, you should have the dataset in `data/ShapeNet` folder. 52 | 53 | ### ScanNet 54 | Download ScanNet v2 data from the [official ScanNet website](https://github.com/ScanNet/ScanNet). 55 | Then, you can preprocess data with: 56 | `scripts/dataset_scannet/build_dataset.py` and put into `data/ScanNet` folder. 57 | 58 | ## Experiments 59 | ### Training 60 | To train a network, run: 61 | ``` 62 | python train.py CONFIG.yaml 63 | ``` 64 | For available training options, please take a look at `configs/default.yaml`. 65 | 66 | **Note**: We implement the code in a multiple-GPU version. Please make sure to call the right version of our encoder at `Line 99` for feature triplane or `Line 100` for feature volume in `train.py`. 67 | 68 | ### Mesh Generation 69 | To generate meshes using a trained model, use 70 | ``` 71 | python generate.py CONFIG.yaml 72 | ``` 73 | where you replace `CONFIG.yaml` with the correct config file. 74 | 75 | 76 | ### Evaluation 77 | For evaluation of the models, we provide the script `eval_meshes.py`. You can run it using: 78 | ``` 79 | python eval_meshes.py CONFIG.yaml 80 | ``` 81 | The script takes the meshes generated in the previous step and evaluates them using a standardized protocol. The output will be written to `.pkl/.csv` files in the corresponding generation folder which can be processed using [pandas](https://pandas.pydata.org/). 82 | 83 | ### Acknowledgement 84 | The code is largely based on [ConvONet](https://github.com/autonomousvision/convolutional_occupancy_networks). Many thanks to the authors for opensourcing the codebase. 85 | 86 | --- 87 | ## Pretrained models 88 | 89 | [ShapeNet 3k](https://drive.google.com/file/d/17AHuISu1f8xWQFevQ2K1lWE1A4LlmsI2/view?usp=share_link) 90 | 91 | [Synthetic Room 10k](https://drive.google.com/file/d/1cffLRxa6mGZlMuUwTEKrMJIxyfC7MDJu/view?usp=share_link) 92 | -------------------------------------------------------------------------------- /build/lib.linux-x86_64-3.6/src/utils/libkdtree/pykdtree/kdtree.cpython-36m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/build/lib.linux-x86_64-3.6/src/utils/libkdtree/pykdtree/kdtree.cpython-36m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /build/lib.linux-x86_64-3.6/src/utils/libmcubes/mcubes.cpython-36m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/build/lib.linux-x86_64-3.6/src/utils/libmcubes/mcubes.cpython-36m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /build/lib.linux-x86_64-3.6/src/utils/libmesh/triangle_hash.cpython-36m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/build/lib.linux-x86_64-3.6/src/utils/libmesh/triangle_hash.cpython-36m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /build/lib.linux-x86_64-3.6/src/utils/libmise/mise.cpython-36m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/build/lib.linux-x86_64-3.6/src/utils/libmise/mise.cpython-36m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /build/lib.linux-x86_64-3.6/src/utils/libsimplify/simplify_mesh.cpython-36m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/build/lib.linux-x86_64-3.6/src/utils/libsimplify/simplify_mesh.cpython-36m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /build/lib.linux-x86_64-3.6/src/utils/libvoxelize/voxelize.cpython-36m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/build/lib.linux-x86_64-3.6/src/utils/libvoxelize/voxelize.cpython-36m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /build/temp.linux-x86_64-3.6/src/utils/libkdtree/pykdtree/_kdtree_core.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/build/temp.linux-x86_64-3.6/src/utils/libkdtree/pykdtree/_kdtree_core.o -------------------------------------------------------------------------------- /build/temp.linux-x86_64-3.6/src/utils/libkdtree/pykdtree/kdtree.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/build/temp.linux-x86_64-3.6/src/utils/libkdtree/pykdtree/kdtree.o -------------------------------------------------------------------------------- /build/temp.linux-x86_64-3.6/src/utils/libmcubes/marchingcubes.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/build/temp.linux-x86_64-3.6/src/utils/libmcubes/marchingcubes.o -------------------------------------------------------------------------------- /build/temp.linux-x86_64-3.6/src/utils/libmcubes/mcubes.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/build/temp.linux-x86_64-3.6/src/utils/libmcubes/mcubes.o -------------------------------------------------------------------------------- /build/temp.linux-x86_64-3.6/src/utils/libmcubes/pywrapper.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/build/temp.linux-x86_64-3.6/src/utils/libmcubes/pywrapper.o -------------------------------------------------------------------------------- /build/temp.linux-x86_64-3.6/src/utils/libmesh/triangle_hash.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/build/temp.linux-x86_64-3.6/src/utils/libmesh/triangle_hash.o -------------------------------------------------------------------------------- /build/temp.linux-x86_64-3.6/src/utils/libmise/mise.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/build/temp.linux-x86_64-3.6/src/utils/libmise/mise.o -------------------------------------------------------------------------------- /build/temp.linux-x86_64-3.6/src/utils/libsimplify/simplify_mesh.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/build/temp.linux-x86_64-3.6/src/utils/libsimplify/simplify_mesh.o -------------------------------------------------------------------------------- /build/temp.linux-x86_64-3.6/src/utils/libvoxelize/voxelize.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/build/temp.linux-x86_64-3.6/src/utils/libvoxelize/voxelize.o -------------------------------------------------------------------------------- /configs/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/configs/.DS_Store -------------------------------------------------------------------------------- /configs/default.yaml: -------------------------------------------------------------------------------- 1 | method: conv_onet 2 | data: 3 | dataset: Shapes3D 4 | path: data/ShapeNet 5 | watertight_path: data/watertight 6 | classes: null 7 | input_type: img 8 | train_split: train 9 | val_split: val 10 | test_split: test 11 | dim: 3 12 | points_file: points.npz 13 | points_iou_file: points.npz 14 | multi_files: null 15 | points_subsample: 1024 16 | points_unpackbits: true 17 | model_file: model.off 18 | watertight_file: model_watertight.off 19 | img_folder: img 20 | img_size: 224 21 | img_with_camera: false 22 | img_augment: false 23 | n_views: 24 24 | pointcloud_file: pointcloud.npz 25 | pointcloud_chamfer_file: pointcloud.npz 26 | pointcloud_n: 256 27 | pointcloud_target_n: 1024 28 | pointcloud_noise: 0.05 29 | voxels_file: 'model.binvox' 30 | padding: 0.1 31 | model: 32 | decoder: simple 33 | encoder: resnet18 34 | decoder_kwargs: {} 35 | encoder_kwargs: {} 36 | multi_gpu: false 37 | c_dim: 512 38 | training: 39 | out_dir: out/default 40 | batch_size: 64 41 | print_every: 200 42 | visualize_every: 1000 43 | checkpoint_every: 1000 44 | validate_every: 2000 45 | backup_every: 100000 46 | eval_sample: false 47 | model_selection_metric: loss 48 | model_selection_mode: minimize 49 | n_workers: 4 50 | n_workers_val: 4 51 | test: 52 | threshold: 0.5 53 | eval_mesh: true 54 | eval_pointcloud: true 55 | remove_wall: false 56 | model_file: model_best.pt 57 | generation: 58 | batch_size: 100000 59 | refinement_step: 0 60 | vis_n_outputs: 30 61 | generate_mesh: true 62 | generate_pointcloud: true 63 | generation_dir: generation 64 | use_sampling: false 65 | resolution_0: 32 66 | upsampling_steps: 3 # 2 67 | simplify_nfaces: null 68 | copy_groundtruth: false 69 | copy_input: true 70 | latent_number: 4 71 | latent_H: 8 72 | latent_W: 8 73 | latent_ny: 2 74 | latent_nx: 2 75 | latent_repeat: true 76 | sliding_window: False # added for crop generation -------------------------------------------------------------------------------- /configs/pointcloud/room_3plane_attention_sub.yaml: -------------------------------------------------------------------------------- 1 | method: conv_onet 2 | data: 3 | input_type: pointcloud 4 | classes: ['rooms_04', 'rooms_05', 'rooms_06', 'rooms_07', 'rooms_08'] 5 | path: ../alter_conv_onet/data/synthetic_room_dataset 6 | pointcloud_n: 10000 7 | pointcloud_noise: 0.005 8 | points_subsample: 2048 9 | points_file: points_iou 10 | points_iou_file: points_iou 11 | pointcloud_file: pointcloud 12 | pointcloud_chamfer_file: pointcloud 13 | multi_files: 10 14 | voxels_file: null 15 | model: 16 | encoder: pointnet_local_pool 17 | encoder_kwargs: 18 | hidden_dim: 32 19 | plane_type: ['xz', 'xy', 'yz'] 20 | plane_resolution: 128 21 | unet: True 22 | unet_kwargs: 23 | depth: 4 #5 24 | merge_mode: concat 25 | start_filts: 32 26 | decoder: simple_local_attention_sub #simple_local_transformer_knn10 #simple_local 27 | decoder_kwargs: 28 | sample_mode: bilinear # bilinear / nearest 29 | hidden_size: 32 30 | K_neighbors: 10 # for transformer only 31 | plane_type: ['xz', 'xy', 'yz'] 32 | c_dim: 32 33 | training: 34 | out_dir: out/pointcloud/room_3plane_attention_sub 35 | batch_size: 24 #32 #9 #24 #10 #32 36 | model_selection_metric: iou 37 | model_selection_mode: maximize 38 | print_every: 100 39 | visualize_every: 100000 #10000 #5000 #10000 40 | validate_every: 5000 #5000 #5000 #5000 #10000 41 | checkpoint_every: 2000 42 | backup_every: 10000 43 | n_workers: 8 44 | n_workers_val: 4 45 | test: 46 | threshold: 0.2 47 | eval_mesh: true 48 | eval_pointcloud: false 49 | remove_wall: true 50 | model_file: model_best.pt 51 | generation: 52 | vis_n_outputs: 2 53 | refine: false 54 | n_x: 128 55 | n_z: 1 56 | -------------------------------------------------------------------------------- /configs/pointcloud/shapenet_3plane_attention_sub.yaml: -------------------------------------------------------------------------------- 1 | method: conv_onet 2 | data: 3 | input_type: pointcloud 4 | classes: null 5 | path: ../occupancy_networks-master/data/ShapeNet 6 | pointcloud_n: 3000 # 300 7 | pointcloud_noise: 0.005 8 | points_subsample: 2048 9 | points_file: points.npz 10 | points_iou_file: points.npz 11 | voxels_file: null 12 | model: 13 | encoder: pointnet_local_pool 14 | encoder_kwargs: 15 | hidden_dim: 32 16 | plane_type: ['xz', 'xy', 'yz'] 17 | plane_resolution: 64 18 | unet: True 19 | unet_kwargs: 20 | depth: 4 21 | merge_mode: concat 22 | start_filts: 32 23 | decoder: simple_local_attention_sub #simple_local_transformer_knn10 #simple_local 24 | decoder_kwargs: 25 | sample_mode: bilinear # bilinear / nearest 26 | hidden_size: 32 27 | plane_type: ['xz', 'xy', 'yz'] 28 | c_dim: 32 29 | training: 30 | out_dir: out/pointcloud/shapenet_3plane_attention 31 | batch_size: 32 32 | model_selection_metric: iou 33 | model_selection_mode: maximize 34 | print_every: 100 35 | visualize_every: 10000 36 | validate_every: 10000 37 | checkpoint_every: 2000 38 | backup_every: 10000 39 | n_workers: 8 40 | n_workers_val: 4 41 | test: 42 | threshold: 0.2 43 | eval_mesh: true 44 | eval_pointcloud: false 45 | model_file: model_best.pt 46 | generation: 47 | generation_dir: generation 48 | vis_n_outputs: 2 49 | refine: false 50 | n_x: 128 51 | n_z: 1 52 | -------------------------------------------------------------------------------- /configs/voxel/room_grid64_attention_sub_2gpu.yaml: -------------------------------------------------------------------------------- 1 | method: conv_onet 2 | data: 3 | input_type: pointcloud 4 | classes: ['rooms_04', 'rooms_05', 'rooms_06', 'rooms_07', 'rooms_08'] 5 | path: ../alter_conv_onet/data/synthetic_room_dataset #data/synthetic_room_dataset 6 | pointcloud_n: 10000 7 | pointcloud_noise: 0.005 8 | points_subsample: 2048 9 | points_file: points_iou 10 | points_iou_file: points_iou 11 | pointcloud_file: pointcloud 12 | pointcloud_chamfer_file: pointcloud 13 | multi_files: 10 14 | voxels_file: null 15 | model: 16 | encoder: pointnet_local_pool 17 | encoder_kwargs: 18 | hidden_dim: 32 19 | plane_type: 'grid' 20 | grid_resolution: 64 21 | unet3d: True 22 | unet3d_kwargs: 23 | num_levels: 4 24 | f_maps: 32 25 | in_channels: 32 26 | out_channels: 32 27 | is_unet: True 28 | decoder: simple_local_attention_sub 29 | decoder_kwargs: 30 | sample_mode: bilinear # bilinear / nearest 31 | hidden_size: 32 32 | plane_type: 'grid' 33 | num_heads: 4 34 | c_dim: 32 35 | training: 36 | out_dir: out/pointcloud/room_grid64_attention 37 | batch_size: 10 #12 38 | model_selection_metric: iou 39 | model_selection_mode: maximize 40 | print_every: 100 41 | visualize_every: 20000 42 | validate_every: 5000 43 | checkpoint_every: 2000 44 | backup_every: 10000 45 | n_workers: 8 46 | n_workers_val: 4 47 | test: 48 | threshold: 0.2 49 | eval_mesh: true 50 | eval_pointcloud: false 51 | remove_wall: true 52 | model_file: model_best.pt 53 | generation: 54 | generation_dir: mc 55 | vis_n_outputs: 2 56 | refine: false 57 | n_x: 128 58 | n_z: 1 -------------------------------------------------------------------------------- /environment.yaml: -------------------------------------------------------------------------------- 1 | name: alto 2 | channels: 3 | - conda-forge 4 | - pytorch 5 | - defaults 6 | dependencies: 7 | - cython=0.29.2 8 | - imageio=2.4.1 9 | - numpy=1.15.4 10 | - numpy-base=1.15.4 11 | - matplotlib=3.0.3 12 | - matplotlib-base=3.0.3 13 | - pandas=0.23.4 14 | - pillow=5.3.0 15 | #- pyembree=0.1.4 16 | - pytest=4.0.2 17 | - python=3.6.7 18 | - pytorch=1.4.0 19 | - pyyaml=3.13 20 | - scikit-image=0.14.1 21 | - scipy=1.1.0 22 | - tensorboardx=1.4 23 | - torchvision=0.2.1 24 | - tqdm=4.28.1 25 | - trimesh=2.37.7 26 | - pip: 27 | - h5py==2.9.0 28 | - plyfile==0.7 29 | - torch_scatter==2.0.4 30 | 31 | -------------------------------------------------------------------------------- /eval_meshes.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | from tqdm import tqdm 4 | import pandas as pd 5 | import trimesh 6 | import torch 7 | from src import config, data 8 | from src.eval import MeshEvaluator 9 | from src.utils.io import load_pointcloud 10 | 11 | 12 | parser = argparse.ArgumentParser( 13 | description='Evaluate mesh algorithms.' 14 | ) 15 | parser.add_argument('config', type=str, help='Path to config file.') 16 | parser.add_argument('--no-cuda', action='store_true', help='Do not use cuda.') 17 | parser.add_argument('--eval_input', action='store_true', 18 | help='Evaluate inputs instead.') 19 | 20 | args = parser.parse_args() 21 | cfg = config.load_config(args.config, 'configs/default.yaml') 22 | is_cuda = (torch.cuda.is_available() and not args.no_cuda) 23 | device = torch.device("cuda" if is_cuda else "cpu") 24 | 25 | # Shorthands 26 | out_dir = cfg['training']['out_dir'] 27 | generation_dir = os.path.join(out_dir, cfg['generation']['generation_dir']) 28 | if not args.eval_input: 29 | out_file = os.path.join(generation_dir, 'eval_meshes_full.pkl') 30 | out_file_class = os.path.join(generation_dir, 'eval_meshes.csv') 31 | else: 32 | out_file = os.path.join(generation_dir, 'eval_input_full.pkl') 33 | out_file_class = os.path.join(generation_dir, 'eval_input.csv') 34 | 35 | # Dataset 36 | points_field = data.PointsField( 37 | cfg['data']['points_iou_file'], 38 | unpackbits=cfg['data']['points_unpackbits'], 39 | multi_files=cfg['data']['multi_files'] 40 | ) 41 | pointcloud_field = data.PointCloudField( 42 | cfg['data']['pointcloud_chamfer_file'], 43 | multi_files=cfg['data']['multi_files'] 44 | ) 45 | fields = { 46 | 'points_iou': points_field, 47 | 'pointcloud_chamfer': pointcloud_field, 48 | 'idx': data.IndexField(), 49 | } 50 | 51 | print('Test split: ', cfg['data']['test_split']) 52 | 53 | dataset_folder = cfg['data']['path'] 54 | dataset = data.Shapes3dDataset( 55 | dataset_folder, fields, 56 | cfg['data']['test_split'], 57 | categories=cfg['data']['classes'], 58 | cfg=cfg 59 | ) 60 | 61 | # Evaluator 62 | evaluator = MeshEvaluator(n_points=100000) 63 | 64 | # Loader 65 | test_loader = torch.utils.data.DataLoader( 66 | dataset, batch_size=1, num_workers=0, shuffle=False) 67 | 68 | # Evaluate all classes 69 | eval_dicts = [] 70 | print('Evaluating meshes...') 71 | for it, data in enumerate(tqdm(test_loader)): 72 | if data is None: 73 | print('Invalid data.') 74 | continue 75 | 76 | # Output folders 77 | if not args.eval_input: 78 | mesh_dir = os.path.join(generation_dir, 'meshes') 79 | pointcloud_dir = os.path.join(generation_dir, 'pointcloud') 80 | else: 81 | mesh_dir = os.path.join(generation_dir, 'input') 82 | pointcloud_dir = os.path.join(generation_dir, 'input') 83 | 84 | # Get index etc. 85 | idx = data['idx'].item() 86 | 87 | try: 88 | model_dict = dataset.get_model_dict(idx) 89 | except AttributeError: 90 | model_dict = {'model': str(idx), 'category': 'n/a'} 91 | 92 | modelname = model_dict['model'] 93 | category_id = model_dict['category'] 94 | 95 | try: 96 | category_name = dataset.metadata[category_id].get('name', 'n/a') 97 | # for room dataset 98 | if category_name == 'n/a': 99 | category_name = category_id 100 | except AttributeError: 101 | category_name = 'n/a' 102 | 103 | if category_id != 'n/a': 104 | mesh_dir = os.path.join(mesh_dir, category_id) 105 | pointcloud_dir = os.path.join(pointcloud_dir, category_id) 106 | 107 | # Evaluate 108 | pointcloud_tgt = data['pointcloud_chamfer'].squeeze(0).numpy() 109 | normals_tgt = data['pointcloud_chamfer.normals'].squeeze(0).numpy() 110 | points_tgt = data['points_iou'].squeeze(0).numpy() 111 | occ_tgt = data['points_iou.occ'].squeeze(0).numpy() 112 | 113 | # Evaluating mesh and pointcloud 114 | # Start row and put basic informatin inside 115 | eval_dict = { 116 | 'idx': idx, 117 | 'class id': category_id, 118 | 'class name': category_name, 119 | 'modelname': modelname, 120 | } 121 | eval_dicts.append(eval_dict) 122 | 123 | # Evaluate mesh 124 | if cfg['test']['eval_mesh']: 125 | #mesh_file = os.path.join(mesh_dir, '%s.off' % modelname) # mc 126 | mesh_file = os.path.join(mesh_dir, '%s.ply' % modelname) # poco mc 127 | 128 | if os.path.exists(mesh_file): 129 | try: 130 | mesh = trimesh.load(mesh_file, process=False) 131 | eval_dict_mesh = evaluator.eval_mesh( 132 | mesh, pointcloud_tgt, normals_tgt, points_tgt, occ_tgt, remove_wall=cfg['test']['remove_wall']) 133 | for k, v in eval_dict_mesh.items(): 134 | eval_dict[k + ' (mesh)'] = v 135 | except Exception as e: 136 | print("Error: Could not evaluate mesh: %s" % mesh_file) 137 | else: 138 | print('Warning: mesh does not exist: %s' % mesh_file) 139 | 140 | # Evaluate point cloud 141 | if cfg['test']['eval_pointcloud']: 142 | pointcloud_file = os.path.join( 143 | pointcloud_dir, '%s.ply' % modelname) 144 | 145 | if os.path.exists(pointcloud_file): 146 | pointcloud = load_pointcloud(pointcloud_file) 147 | eval_dict_pcl = evaluator.eval_pointcloud( 148 | pointcloud, pointcloud_tgt) 149 | for k, v in eval_dict_pcl.items(): 150 | eval_dict[k + ' (pcl)'] = v 151 | else: 152 | print('Warning: pointcloud does not exist: %s' 153 | % pointcloud_file) 154 | 155 | # Create pandas dataframe and save 156 | eval_df = pd.DataFrame(eval_dicts) 157 | eval_df.set_index(['idx'], inplace=True) 158 | eval_df.to_pickle(out_file) 159 | 160 | # Create CSV file with main statistics 161 | eval_df_class = eval_df.groupby(by=['class name']).mean() 162 | eval_df_class.to_csv(out_file_class) 163 | 164 | # Print results 165 | eval_df_class.loc['mean'] = eval_df_class.mean() 166 | print(eval_df_class) 167 | -------------------------------------------------------------------------------- /media/pipeline.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/media/pipeline.png -------------------------------------------------------------------------------- /modules/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/modules/.DS_Store -------------------------------------------------------------------------------- /modules/__init__.py: -------------------------------------------------------------------------------- 1 | from modules.ball_query import BallQuery 2 | from modules.frustum import FrustumPointNetLoss 3 | from modules.loss import KLLoss 4 | from modules.pointnet import PointNetAModule, PointNetSAModule, PointNetFPModule 5 | from modules.pvconv import PVConv 6 | from modules.se import SE3d 7 | from modules.shared_mlp import SharedMLP 8 | from modules.voxelization import Voxelization 9 | -------------------------------------------------------------------------------- /modules/ball_query.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | import modules.functional as F 5 | 6 | __all__ = ['BallQuery'] 7 | 8 | 9 | class BallQuery(nn.Module): 10 | def __init__(self, radius, num_neighbors, include_coordinates=True): 11 | super().__init__() 12 | self.radius = radius 13 | self.num_neighbors = num_neighbors 14 | self.include_coordinates = include_coordinates 15 | 16 | def forward(self, points_coords, centers_coords, points_features=None): 17 | points_coords = points_coords.contiguous() 18 | centers_coords = centers_coords.contiguous() 19 | neighbor_indices = F.ball_query(centers_coords, points_coords, self.radius, self.num_neighbors) 20 | neighbor_coordinates = F.grouping(points_coords, neighbor_indices) 21 | neighbor_coordinates = neighbor_coordinates - centers_coords.unsqueeze(-1) 22 | 23 | if points_features is None: 24 | assert self.include_coordinates, 'No Features For Grouping' 25 | neighbor_features = neighbor_coordinates 26 | else: 27 | neighbor_features = F.grouping(points_features, neighbor_indices) 28 | if self.include_coordinates: 29 | neighbor_features = torch.cat([neighbor_coordinates, neighbor_features], dim=1) 30 | return neighbor_features 31 | 32 | def extra_repr(self): 33 | return 'radius={}, num_neighbors={}{}'.format( 34 | self.radius, self.num_neighbors, ', include coordinates' if self.include_coordinates else '') 35 | -------------------------------------------------------------------------------- /modules/functional/__init__.py: -------------------------------------------------------------------------------- 1 | from modules.functional.ball_query import ball_query 2 | from modules.functional.devoxelization import trilinear_devoxelize 3 | from modules.functional.grouping import grouping 4 | from modules.functional.interpolatation import nearest_neighbor_interpolate 5 | from modules.functional.loss import kl_loss, huber_loss 6 | from modules.functional.sampling import gather, furthest_point_sample, logits_mask 7 | from modules.functional.voxelization import avg_voxelize 8 | -------------------------------------------------------------------------------- /modules/functional/backend.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | from torch.utils.cpp_extension import load 4 | 5 | _src_path = os.path.dirname(os.path.abspath(__file__)) 6 | _backend = load(name='_pvcnn_backend', 7 | extra_cflags=['-O3', '-std=c++17'], 8 | sources=[os.path.join(_src_path,'src', f) for f in [ 9 | 'ball_query/ball_query.cpp', 10 | 'ball_query/ball_query.cu', 11 | 'grouping/grouping.cpp', 12 | 'grouping/grouping.cu', 13 | 'interpolate/neighbor_interpolate.cpp', 14 | 'interpolate/neighbor_interpolate.cu', 15 | 'interpolate/trilinear_devox.cpp', 16 | 'interpolate/trilinear_devox.cu', 17 | 'sampling/sampling.cpp', 18 | 'sampling/sampling.cu', 19 | 'voxelization/vox.cpp', 20 | 'voxelization/vox.cu', 21 | 'bindings.cpp', 22 | ]] 23 | ) 24 | 25 | __all__ = ['_backend'] 26 | -------------------------------------------------------------------------------- /modules/functional/ball_query.py: -------------------------------------------------------------------------------- 1 | from torch.autograd import Function 2 | 3 | from modules.functional.backend import _backend 4 | 5 | __all__ = ['ball_query'] 6 | 7 | 8 | def ball_query(centers_coords, points_coords, radius, num_neighbors): 9 | """ 10 | :param centers_coords: coordinates of centers, FloatTensor[B, 3, M] 11 | :param points_coords: coordinates of points, FloatTensor[B, 3, N] 12 | :param radius: float, radius of ball query 13 | :param num_neighbors: int, maximum number of neighbors 14 | :return: 15 | neighbor_indices: indices of neighbors, IntTensor[B, M, U] 16 | """ 17 | centers_coords = centers_coords.contiguous() 18 | points_coords = points_coords.contiguous() 19 | return _backend.ball_query(centers_coords, points_coords, radius, num_neighbors) 20 | -------------------------------------------------------------------------------- /modules/functional/devoxelization.py: -------------------------------------------------------------------------------- 1 | from torch.autograd import Function 2 | 3 | from modules.functional.backend import _backend 4 | 5 | __all__ = ['trilinear_devoxelize'] 6 | 7 | 8 | class TrilinearDevoxelization(Function): 9 | @staticmethod 10 | def forward(ctx, features, coords, resolution, is_training=True): 11 | """ 12 | :param ctx: 13 | :param coords: the coordinates of points, FloatTensor[B, 3, N] 14 | :param features: FloatTensor[B, C, R, R, R] 15 | :param resolution: int, the voxel resolution 16 | :param is_training: bool, training mode 17 | :return: 18 | FloatTensor[B, C, N] 19 | """ 20 | B, C = features.shape[:2] 21 | features = features.contiguous().view(B, C, -1) 22 | coords = coords.contiguous() 23 | outs, inds, wgts = _backend.trilinear_devoxelize_forward(resolution, is_training, coords, features) 24 | if is_training: 25 | ctx.save_for_backward(inds, wgts) 26 | ctx.r = resolution 27 | return outs 28 | 29 | @staticmethod 30 | def backward(ctx, grad_output): 31 | """ 32 | :param ctx: 33 | :param grad_output: gradient of outputs, FloatTensor[B, C, N] 34 | :return: 35 | gradient of inputs, FloatTensor[B, C, R, R, R] 36 | """ 37 | inds, wgts = ctx.saved_tensors 38 | grad_inputs = _backend.trilinear_devoxelize_backward(grad_output.contiguous(), inds, wgts, ctx.r) 39 | return grad_inputs.view(grad_output.size(0), grad_output.size(1), ctx.r, ctx.r, ctx.r), None, None, None 40 | 41 | 42 | trilinear_devoxelize = TrilinearDevoxelization.apply 43 | -------------------------------------------------------------------------------- /modules/functional/grouping.py: -------------------------------------------------------------------------------- 1 | from torch.autograd import Function 2 | 3 | from modules.functional.backend import _backend 4 | 5 | __all__ = ['grouping'] 6 | 7 | 8 | class Grouping(Function): 9 | @staticmethod 10 | def forward(ctx, features, indices): 11 | """ 12 | :param ctx: 13 | :param features: features of points, FloatTensor[B, C, N] 14 | :param indices: neighbor indices of centers, IntTensor[B, M, U], M is #centers, U is #neighbors 15 | :return: 16 | grouped_features: grouped features, FloatTensor[B, C, M, U] 17 | """ 18 | features = features.contiguous() 19 | indices = indices.contiguous() 20 | ctx.save_for_backward(indices) 21 | ctx.num_points = features.size(-1) 22 | return _backend.grouping_forward(features, indices) 23 | 24 | @staticmethod 25 | def backward(ctx, grad_output): 26 | indices, = ctx.saved_tensors 27 | grad_features = _backend.grouping_backward(grad_output.contiguous(), indices, ctx.num_points) 28 | return grad_features, None 29 | 30 | 31 | grouping = Grouping.apply 32 | -------------------------------------------------------------------------------- /modules/functional/interpolatation.py: -------------------------------------------------------------------------------- 1 | from torch.autograd import Function 2 | 3 | from modules.functional.backend import _backend 4 | 5 | __all__ = ['nearest_neighbor_interpolate'] 6 | 7 | 8 | class NeighborInterpolation(Function): 9 | @staticmethod 10 | def forward(ctx, points_coords, centers_coords, centers_features): 11 | """ 12 | :param ctx: 13 | :param points_coords: coordinates of points, FloatTensor[B, 3, N] 14 | :param centers_coords: coordinates of centers, FloatTensor[B, 3, M] 15 | :param centers_features: features of centers, FloatTensor[B, C, M] 16 | :return: 17 | points_features: features of points, FloatTensor[B, C, N] 18 | """ 19 | centers_coords = centers_coords.contiguous() 20 | points_coords = points_coords.contiguous() 21 | centers_features = centers_features.contiguous() 22 | points_features, indices, weights = _backend.three_nearest_neighbors_interpolate_forward( 23 | points_coords, centers_coords, centers_features 24 | ) 25 | ctx.save_for_backward(indices, weights) 26 | ctx.num_centers = centers_coords.size(-1) 27 | return points_features 28 | 29 | @staticmethod 30 | def backward(ctx, grad_output): 31 | indices, weights = ctx.saved_tensors 32 | grad_centers_features = _backend.three_nearest_neighbors_interpolate_backward( 33 | grad_output.contiguous(), indices, weights, ctx.num_centers 34 | ) 35 | return None, None, grad_centers_features 36 | 37 | 38 | nearest_neighbor_interpolate = NeighborInterpolation.apply 39 | -------------------------------------------------------------------------------- /modules/functional/loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | __all__ = ['kl_loss', 'huber_loss'] 5 | 6 | 7 | def kl_loss(x, y): 8 | x = F.softmax(x.detach(), dim=1) 9 | y = F.log_softmax(y, dim=1) 10 | return torch.mean(torch.sum(x * (torch.log(x) - y), dim=1)) 11 | 12 | 13 | def huber_loss(error, delta): 14 | abs_error = torch.abs(error) 15 | quadratic = torch.min(abs_error, torch.full_like(abs_error, fill_value=delta)) 16 | losses = 0.5 * (quadratic ** 2) + delta * (abs_error - quadratic) 17 | return torch.mean(losses) 18 | -------------------------------------------------------------------------------- /modules/functional/sampling.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from torch.autograd import Function 4 | 5 | from modules.functional.backend import _backend 6 | 7 | __all__ = ['gather', 'furthest_point_sample', 'logits_mask'] 8 | 9 | 10 | class Gather(Function): 11 | @staticmethod 12 | def forward(ctx, features, indices): 13 | """ 14 | Gather 15 | :param ctx: 16 | :param features: features of points, FloatTensor[B, C, N] 17 | :param indices: centers' indices in points, IntTensor[b, m] 18 | :return: 19 | centers_coords: coordinates of sampled centers, FloatTensor[B, C, M] 20 | """ 21 | features = features.contiguous() 22 | indices = indices.int().contiguous() 23 | ctx.save_for_backward(indices) 24 | ctx.num_points = features.size(-1) 25 | return _backend.gather_features_forward(features, indices) 26 | 27 | @staticmethod 28 | def backward(ctx, grad_output): 29 | indices, = ctx.saved_tensors 30 | grad_features = _backend.gather_features_backward(grad_output.contiguous(), indices, ctx.num_points) 31 | return grad_features, None 32 | 33 | 34 | gather = Gather.apply 35 | 36 | 37 | def furthest_point_sample(coords, num_samples): 38 | """ 39 | Uses iterative furthest point sampling to select a set of npoint features that have the largest 40 | minimum distance to the sampled point set 41 | :param coords: coordinates of points, FloatTensor[B, 3, N] 42 | :param num_samples: int, M 43 | :return: 44 | centers_coords: coordinates of sampled centers, FloatTensor[B, 3, M] 45 | """ 46 | coords = coords.contiguous() 47 | indices = _backend.furthest_point_sampling(coords, num_samples) 48 | return gather(coords, indices) 49 | 50 | 51 | def logits_mask(coords, logits, num_points_per_object): 52 | """ 53 | Use logits to sample points 54 | :param coords: coords of points, FloatTensor[B, 3, N] 55 | :param logits: binary classification logits, FloatTensor[B, 2, N] 56 | :param num_points_per_object: M, #points per object after masking, int 57 | :return: 58 | selected_coords: FloatTensor[B, 3, M] 59 | masked_coords_mean: mean coords of selected points, FloatTensor[B, 3] 60 | mask: mask to select points, BoolTensor[B, N] 61 | """ 62 | batch_size, _, num_points = coords.shape 63 | mask = torch.lt(logits[:, 0, :], logits[:, 1, :]) # [B, N] 64 | num_candidates = torch.sum(mask, dim=-1, keepdim=True) # [B, 1] 65 | masked_coords = coords * mask.view(batch_size, 1, num_points) # [B, C, N] 66 | masked_coords_mean = torch.sum(masked_coords, dim=-1) / torch.max(num_candidates, 67 | torch.ones_like(num_candidates)).float() # [B, C] 68 | selected_indices = torch.zeros((batch_size, num_points_per_object), device=coords.device, dtype=torch.int32) 69 | for i in range(batch_size): 70 | current_mask = mask[i] # [N] 71 | current_candidates = current_mask.nonzero().view(-1) 72 | current_num_candidates = current_candidates.numel() 73 | if current_num_candidates >= num_points_per_object: 74 | choices = np.random.choice(current_num_candidates, num_points_per_object, replace=False) 75 | selected_indices[i] = current_candidates[choices] 76 | elif current_num_candidates > 0: 77 | choices = np.concatenate([ 78 | np.arange(current_num_candidates).repeat(num_points_per_object // current_num_candidates), 79 | np.random.choice(current_num_candidates, num_points_per_object % current_num_candidates, replace=False) 80 | ]) 81 | np.random.shuffle(choices) 82 | selected_indices[i] = current_candidates[choices] 83 | selected_coords = gather(masked_coords - masked_coords_mean.view(batch_size, -1, 1), selected_indices) 84 | return selected_coords, masked_coords_mean, mask 85 | -------------------------------------------------------------------------------- /modules/functional/src/ball_query/ball_query.cpp: -------------------------------------------------------------------------------- 1 | #include "ball_query.hpp" 2 | #include "ball_query.cuh" 3 | 4 | #include "../utils.hpp" 5 | 6 | at::Tensor ball_query_forward(at::Tensor centers_coords, 7 | at::Tensor points_coords, const float radius, 8 | const int num_neighbors) { 9 | CHECK_CUDA(centers_coords); 10 | CHECK_CUDA(points_coords); 11 | CHECK_CONTIGUOUS(centers_coords); 12 | CHECK_CONTIGUOUS(points_coords); 13 | CHECK_IS_FLOAT(centers_coords); 14 | CHECK_IS_FLOAT(points_coords); 15 | 16 | int b = centers_coords.size(0); 17 | int m = centers_coords.size(2); 18 | int n = points_coords.size(2); 19 | 20 | at::Tensor neighbors_indices = torch::zeros( 21 | {b, m, num_neighbors}, 22 | at::device(centers_coords.device()).dtype(at::ScalarType::Int)); 23 | 24 | ball_query(b, n, m, radius * radius, num_neighbors, 25 | centers_coords.data_ptr(), 26 | points_coords.data_ptr(), 27 | neighbors_indices.data_ptr()); 28 | 29 | return neighbors_indices; 30 | } 31 | -------------------------------------------------------------------------------- /modules/functional/src/ball_query/ball_query.cppZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/ball_query/ball_query.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include "../cuda_utils.cuh" 6 | 7 | /* 8 | Function: ball query 9 | Args: 10 | b : batch size 11 | n : number of points in point clouds 12 | m : number of query centers 13 | r2 : ball query radius ** 2 14 | u : maximum number of neighbors 15 | centers_coords: coordinates of centers, FloatTensor[b, 3, m] 16 | points_coords : coordinates of points, FloatTensor[b, 3, n] 17 | neighbors_indices : neighbor indices in points, IntTensor[b, m, u] 18 | */ 19 | __global__ void ball_query_kernel(int b, int n, int m, float r2, int u, 20 | const float *__restrict__ centers_coords, 21 | const float *__restrict__ points_coords, 22 | int *__restrict__ neighbors_indices) { 23 | int batch_index = blockIdx.x; 24 | int index = threadIdx.x; 25 | int stride = blockDim.x; 26 | points_coords += batch_index * n * 3; 27 | centers_coords += batch_index * m * 3; 28 | neighbors_indices += batch_index * m * u; 29 | 30 | for (int j = index; j < m; j += stride) { 31 | float center_x = centers_coords[j]; 32 | float center_y = centers_coords[j + m]; 33 | float center_z = centers_coords[j + m + m]; 34 | for (int k = 0, cnt = 0; k < n && cnt < u; ++k) { 35 | float dx = center_x - points_coords[k]; 36 | float dy = center_y - points_coords[k + n]; 37 | float dz = center_z - points_coords[k + n + n]; 38 | float d2 = dx * dx + dy * dy + dz * dz; 39 | if (d2 < r2) { 40 | if (cnt == 0) { 41 | for (int v = 0; v < u; ++v) { 42 | neighbors_indices[j * u + v] = k; 43 | } 44 | } 45 | neighbors_indices[j * u + cnt] = k; 46 | ++cnt; 47 | } 48 | } 49 | } 50 | } 51 | 52 | void ball_query(int b, int n, int m, float r2, int u, 53 | const float *centers_coords, const float *points_coords, 54 | int *neighbors_indices) { 55 | ball_query_kernel<<>>( 57 | b, n, m, r2, u, centers_coords, points_coords, neighbors_indices); 58 | CUDA_CHECK_ERRORS(); 59 | } 60 | -------------------------------------------------------------------------------- /modules/functional/src/ball_query/ball_query.cuh: -------------------------------------------------------------------------------- 1 | #ifndef _BALL_QUERY_CUH 2 | #define _BALL_QUERY_CUH 3 | 4 | void ball_query(int b, int n, int m, float r2, int u, 5 | const float *centers_coords, const float *points_coords, 6 | int *neighbors_indices); 7 | 8 | #endif 9 | -------------------------------------------------------------------------------- /modules/functional/src/ball_query/ball_query.cuhZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/ball_query/ball_query.cuZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/ball_query/ball_query.hpp: -------------------------------------------------------------------------------- 1 | #ifndef _BALL_QUERY_HPP 2 | #define _BALL_QUERY_HPP 3 | 4 | #include 5 | 6 | at::Tensor ball_query_forward(at::Tensor centers_coords, 7 | at::Tensor points_coords, const float radius, 8 | const int num_neighbors); 9 | 10 | #endif 11 | -------------------------------------------------------------------------------- /modules/functional/src/ball_query/ball_query.hppZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/bindings.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include "ball_query/ball_query.hpp" 4 | #include "grouping/grouping.hpp" 5 | #include "interpolate/neighbor_interpolate.hpp" 6 | #include "interpolate/trilinear_devox.hpp" 7 | #include "sampling/sampling.hpp" 8 | #include "voxelization/vox.hpp" 9 | 10 | PYBIND11_MODULE(_pvcnn_backend, m) { 11 | m.def("gather_features_forward", &gather_features_forward, 12 | "Gather Centers' Features forward (CUDA)"); 13 | m.def("gather_features_backward", &gather_features_backward, 14 | "Gather Centers' Features backward (CUDA)"); 15 | m.def("furthest_point_sampling", &furthest_point_sampling_forward, 16 | "Furthest Point Sampling (CUDA)"); 17 | m.def("ball_query", &ball_query_forward, "Ball Query (CUDA)"); 18 | m.def("grouping_forward", &grouping_forward, 19 | "Grouping Features forward (CUDA)"); 20 | m.def("grouping_backward", &grouping_backward, 21 | "Grouping Features backward (CUDA)"); 22 | m.def("three_nearest_neighbors_interpolate_forward", 23 | &three_nearest_neighbors_interpolate_forward, 24 | "3 Nearest Neighbors Interpolate forward (CUDA)"); 25 | m.def("three_nearest_neighbors_interpolate_backward", 26 | &three_nearest_neighbors_interpolate_backward, 27 | "3 Nearest Neighbors Interpolate backward (CUDA)"); 28 | 29 | m.def("trilinear_devoxelize_forward", &trilinear_devoxelize_forward, 30 | "Trilinear Devoxelization forward (CUDA)"); 31 | m.def("trilinear_devoxelize_backward", &trilinear_devoxelize_backward, 32 | "Trilinear Devoxelization backward (CUDA)"); 33 | m.def("avg_voxelize_forward", &avg_voxelize_forward, 34 | "Voxelization forward with average pooling (CUDA)"); 35 | m.def("avg_voxelize_backward", &avg_voxelize_backward, 36 | "Voxelization backward (CUDA)"); 37 | } 38 | -------------------------------------------------------------------------------- /modules/functional/src/cuda_utils.cuh: -------------------------------------------------------------------------------- 1 | #ifndef _CUDA_UTILS_H 2 | #define _CUDA_UTILS_H 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | #include 9 | #include 10 | 11 | #include 12 | 13 | #define MAXIMUM_THREADS 512 14 | 15 | inline int optimal_num_threads(int work_size) { 16 | const int pow_2 = std::log2(static_cast(work_size)); 17 | return max(min(1 << pow_2, MAXIMUM_THREADS), 1); 18 | } 19 | 20 | inline dim3 optimal_block_config(int x, int y) { 21 | const int x_threads = optimal_num_threads(x); 22 | const int y_threads = 23 | max(min(optimal_num_threads(y), MAXIMUM_THREADS / x_threads), 1); 24 | dim3 block_config(x_threads, y_threads, 1); 25 | return block_config; 26 | } 27 | 28 | #define CUDA_CHECK_ERRORS() \ 29 | { \ 30 | cudaError_t err = cudaGetLastError(); \ 31 | if (cudaSuccess != err) { \ 32 | fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n", \ 33 | cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, \ 34 | __FILE__); \ 35 | exit(-1); \ 36 | } \ 37 | } 38 | 39 | #endif 40 | -------------------------------------------------------------------------------- /modules/functional/src/grouping/grouping.cpp: -------------------------------------------------------------------------------- 1 | #include "grouping.hpp" 2 | #include "grouping.cuh" 3 | 4 | #include "../utils.hpp" 5 | 6 | at::Tensor grouping_forward(at::Tensor features, at::Tensor indices) { 7 | CHECK_CUDA(features); 8 | CHECK_CUDA(indices); 9 | CHECK_CONTIGUOUS(features); 10 | CHECK_CONTIGUOUS(indices); 11 | CHECK_IS_FLOAT(features); 12 | CHECK_IS_INT(indices); 13 | 14 | int b = features.size(0); 15 | int c = features.size(1); 16 | int n = features.size(2); 17 | int m = indices.size(1); 18 | int u = indices.size(2); 19 | at::Tensor output = torch::zeros( 20 | {b, c, m, u}, at::device(features.device()).dtype(at::ScalarType::Float)); 21 | grouping(b, c, n, m, u, features.data_ptr(), indices.data_ptr(), 22 | output.data_ptr()); 23 | return output; 24 | } 25 | 26 | at::Tensor grouping_backward(at::Tensor grad_y, at::Tensor indices, 27 | const int n) { 28 | CHECK_CUDA(grad_y); 29 | CHECK_CUDA(indices); 30 | CHECK_CONTIGUOUS(grad_y); 31 | CHECK_CONTIGUOUS(indices); 32 | CHECK_IS_FLOAT(grad_y); 33 | CHECK_IS_INT(indices); 34 | 35 | int b = grad_y.size(0); 36 | int c = grad_y.size(1); 37 | int m = indices.size(1); 38 | int u = indices.size(2); 39 | at::Tensor grad_x = torch::zeros( 40 | {b, c, n}, at::device(grad_y.device()).dtype(at::ScalarType::Float)); 41 | grouping_grad(b, c, n, m, u, grad_y.data_ptr(), 42 | indices.data_ptr(), grad_x.data_ptr()); 43 | return grad_x; 44 | } 45 | -------------------------------------------------------------------------------- /modules/functional/src/grouping/grouping.cppZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/grouping/grouping.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "../cuda_utils.cuh" 5 | 6 | /* 7 | Function: grouping features of neighbors (forward) 8 | Args: 9 | b : batch size 10 | c : #channles of features 11 | n : number of points in point clouds 12 | m : number of query centers 13 | u : maximum number of neighbors 14 | features: points' features, FloatTensor[b, c, n] 15 | indices : neighbor indices in points, IntTensor[b, m, u] 16 | out : gathered features, FloatTensor[b, c, m, u] 17 | */ 18 | __global__ void grouping_kernel(int b, int c, int n, int m, int u, 19 | const float *__restrict__ features, 20 | const int *__restrict__ indices, 21 | float *__restrict__ out) { 22 | int batch_index = blockIdx.x; 23 | features += batch_index * n * c; 24 | indices += batch_index * m * u; 25 | out += batch_index * m * u * c; 26 | 27 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 28 | const int stride = blockDim.y * blockDim.x; 29 | for (int i = index; i < c * m; i += stride) { 30 | const int l = i / m; 31 | const int j = i % m; 32 | for (int k = 0; k < u; ++k) { 33 | out[(l * m + j) * u + k] = features[l * n + indices[j * u + k]]; 34 | } 35 | } 36 | } 37 | 38 | void grouping(int b, int c, int n, int m, int u, const float *features, 39 | const int *indices, float *out) { 40 | grouping_kernel<<>>(b, c, n, m, u, features, 42 | indices, out); 43 | CUDA_CHECK_ERRORS(); 44 | } 45 | 46 | /* 47 | Function: grouping features of neighbors (backward) 48 | Args: 49 | b : batch size 50 | c : #channles of features 51 | n : number of points in point clouds 52 | m : number of query centers 53 | u : maximum number of neighbors 54 | grad_y : grad of gathered features, FloatTensor[b, c, m, u] 55 | indices : neighbor indices in points, IntTensor[b, m, u] 56 | grad_x: grad of points' features, FloatTensor[b, c, n] 57 | */ 58 | __global__ void grouping_grad_kernel(int b, int c, int n, int m, int u, 59 | const float *__restrict__ grad_y, 60 | const int *__restrict__ indices, 61 | float *__restrict__ grad_x) { 62 | int batch_index = blockIdx.x; 63 | grad_y += batch_index * m * u * c; 64 | indices += batch_index * m * u; 65 | grad_x += batch_index * n * c; 66 | 67 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 68 | const int stride = blockDim.y * blockDim.x; 69 | for (int i = index; i < c * m; i += stride) { 70 | const int l = i / m; 71 | const int j = i % m; 72 | for (int k = 0; k < u; ++k) { 73 | atomicAdd(grad_x + l * n + indices[j * u + k], 74 | grad_y[(l * m + j) * u + k]); 75 | } 76 | } 77 | } 78 | 79 | void grouping_grad(int b, int c, int n, int m, int u, const float *grad_y, 80 | const int *indices, float *grad_x) { 81 | grouping_grad_kernel<<>>( 83 | b, c, n, m, u, grad_y, indices, grad_x); 84 | CUDA_CHECK_ERRORS(); 85 | } 86 | -------------------------------------------------------------------------------- /modules/functional/src/grouping/grouping.cuh: -------------------------------------------------------------------------------- 1 | #ifndef _GROUPING_CUH 2 | #define _GROUPING_CUH 3 | 4 | void grouping(int b, int c, int n, int m, int u, const float *features, 5 | const int *indices, float *out); 6 | void grouping_grad(int b, int c, int n, int m, int u, const float *grad_y, 7 | const int *indices, float *grad_x); 8 | 9 | #endif -------------------------------------------------------------------------------- /modules/functional/src/grouping/grouping.cuhZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/grouping/grouping.cuZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/grouping/grouping.hpp: -------------------------------------------------------------------------------- 1 | #ifndef _GROUPING_HPP 2 | #define _GROUPING_HPP 3 | 4 | #include 5 | 6 | at::Tensor grouping_forward(at::Tensor features, at::Tensor indices); 7 | at::Tensor grouping_backward(at::Tensor grad_y, at::Tensor indices, 8 | const int n); 9 | 10 | #endif 11 | -------------------------------------------------------------------------------- /modules/functional/src/grouping/grouping.hppZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/interpolate/neighbor_interpolate.cpp: -------------------------------------------------------------------------------- 1 | #include "neighbor_interpolate.hpp" 2 | #include "neighbor_interpolate.cuh" 3 | 4 | #include "../utils.hpp" 5 | 6 | std::vector 7 | three_nearest_neighbors_interpolate_forward(at::Tensor points_coords, 8 | at::Tensor centers_coords, 9 | at::Tensor centers_features) { 10 | CHECK_CUDA(points_coords); 11 | CHECK_CUDA(centers_coords); 12 | CHECK_CUDA(centers_features); 13 | CHECK_CONTIGUOUS(points_coords); 14 | CHECK_CONTIGUOUS(centers_coords); 15 | CHECK_CONTIGUOUS(centers_features); 16 | CHECK_IS_FLOAT(points_coords); 17 | CHECK_IS_FLOAT(centers_coords); 18 | CHECK_IS_FLOAT(centers_features); 19 | 20 | int b = centers_features.size(0); 21 | int c = centers_features.size(1); 22 | int m = centers_features.size(2); 23 | int n = points_coords.size(2); 24 | 25 | at::Tensor indices = torch::zeros( 26 | {b, 3, n}, at::device(points_coords.device()).dtype(at::ScalarType::Int)); 27 | at::Tensor weights = torch::zeros( 28 | {b, 3, n}, 29 | at::device(points_coords.device()).dtype(at::ScalarType::Float)); 30 | at::Tensor output = torch::zeros( 31 | {b, c, n}, 32 | at::device(centers_features.device()).dtype(at::ScalarType::Float)); 33 | 34 | three_nearest_neighbors_interpolate( 35 | b, c, m, n, points_coords.data_ptr(), 36 | centers_coords.data_ptr(), centers_features.data_ptr(), 37 | indices.data_ptr(), weights.data_ptr(), 38 | output.data_ptr()); 39 | return {output, indices, weights}; 40 | } 41 | 42 | at::Tensor three_nearest_neighbors_interpolate_backward(at::Tensor grad_y, 43 | at::Tensor indices, 44 | at::Tensor weights, 45 | const int m) { 46 | CHECK_CUDA(grad_y); 47 | CHECK_CUDA(indices); 48 | CHECK_CUDA(weights); 49 | CHECK_CONTIGUOUS(grad_y); 50 | CHECK_CONTIGUOUS(indices); 51 | CHECK_CONTIGUOUS(weights); 52 | CHECK_IS_FLOAT(grad_y); 53 | CHECK_IS_INT(indices); 54 | CHECK_IS_FLOAT(weights); 55 | 56 | int b = grad_y.size(0); 57 | int c = grad_y.size(1); 58 | int n = grad_y.size(2); 59 | at::Tensor grad_x = torch::zeros( 60 | {b, c, m}, at::device(grad_y.device()).dtype(at::ScalarType::Float)); 61 | three_nearest_neighbors_interpolate_grad( 62 | b, c, n, m, grad_y.data_ptr(), indices.data_ptr(), 63 | weights.data_ptr(), grad_x.data_ptr()); 64 | return grad_x; 65 | } 66 | -------------------------------------------------------------------------------- /modules/functional/src/interpolate/neighbor_interpolate.cppZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/interpolate/neighbor_interpolate.cuh: -------------------------------------------------------------------------------- 1 | #ifndef _NEIGHBOR_INTERPOLATE_CUH 2 | #define _NEIGHBOR_INTERPOLATE_CUH 3 | 4 | void three_nearest_neighbors_interpolate(int b, int c, int m, int n, 5 | const float *points_coords, 6 | const float *centers_coords, 7 | const float *centers_features, 8 | int *indices, float *weights, 9 | float *out); 10 | void three_nearest_neighbors_interpolate_grad(int b, int c, int n, int m, 11 | const float *grad_y, 12 | const int *indices, 13 | const float *weights, 14 | float *grad_x); 15 | 16 | #endif 17 | -------------------------------------------------------------------------------- /modules/functional/src/interpolate/neighbor_interpolate.cuhZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/interpolate/neighbor_interpolate.cuZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/interpolate/neighbor_interpolate.hpp: -------------------------------------------------------------------------------- 1 | #ifndef _NEIGHBOR_INTERPOLATE_HPP 2 | #define _NEIGHBOR_INTERPOLATE_HPP 3 | 4 | #include 5 | #include 6 | 7 | std::vector 8 | three_nearest_neighbors_interpolate_forward(at::Tensor points_coords, 9 | at::Tensor centers_coords, 10 | at::Tensor centers_features); 11 | at::Tensor three_nearest_neighbors_interpolate_backward(at::Tensor grad_y, 12 | at::Tensor indices, 13 | at::Tensor weights, 14 | const int m); 15 | 16 | #endif 17 | -------------------------------------------------------------------------------- /modules/functional/src/interpolate/neighbor_interpolate.hppZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/interpolate/trilinear_devox.cpp: -------------------------------------------------------------------------------- 1 | #include "trilinear_devox.hpp" 2 | #include "trilinear_devox.cuh" 3 | 4 | #include "../utils.hpp" 5 | 6 | /* 7 | Function: trilinear devoxelization (forward) 8 | Args: 9 | r : voxel resolution 10 | trainig : whether is training mode 11 | coords : the coordinates of points, FloatTensor[b, 3, n] 12 | features : features, FloatTensor[b, c, s], s = r ** 3 13 | Return: 14 | outs : outputs, FloatTensor[b, c, n] 15 | inds : the voxel coordinates of point cube, IntTensor[b, 8, n] 16 | wgts : weight for trilinear interpolation, FloatTensor[b, 8, n] 17 | */ 18 | std::vector 19 | trilinear_devoxelize_forward(const int r, const bool is_training, 20 | const at::Tensor coords, 21 | const at::Tensor features) { 22 | CHECK_CUDA(features); 23 | CHECK_CUDA(coords); 24 | CHECK_CONTIGUOUS(features); 25 | CHECK_CONTIGUOUS(coords); 26 | CHECK_IS_FLOAT(features); 27 | CHECK_IS_FLOAT(coords); 28 | 29 | int b = features.size(0); 30 | int c = features.size(1); 31 | int n = coords.size(2); 32 | int r2 = r * r; 33 | int r3 = r2 * r; 34 | at::Tensor outs = torch::zeros( 35 | {b, c, n}, at::device(features.device()).dtype(at::ScalarType::Float)); 36 | if (is_training) { 37 | at::Tensor inds = torch::zeros( 38 | {b, 8, n}, at::device(features.device()).dtype(at::ScalarType::Int)); 39 | at::Tensor wgts = torch::zeros( 40 | {b, 8, n}, at::device(features.device()).dtype(at::ScalarType::Float)); 41 | trilinear_devoxelize(b, c, n, r, r2, r3, true, coords.data_ptr(), 42 | features.data_ptr(), inds.data_ptr(), 43 | wgts.data_ptr(), outs.data_ptr()); 44 | return {outs, inds, wgts}; 45 | } else { 46 | at::Tensor inds = torch::zeros( 47 | {1}, at::device(features.device()).dtype(at::ScalarType::Int)); 48 | at::Tensor wgts = torch::zeros( 49 | {1}, at::device(features.device()).dtype(at::ScalarType::Float)); 50 | trilinear_devoxelize(b, c, n, r, r2, r3, false, coords.data_ptr(), 51 | features.data_ptr(), inds.data_ptr(), 52 | wgts.data_ptr(), outs.data_ptr()); 53 | return {outs, inds, wgts}; 54 | } 55 | } 56 | 57 | /* 58 | Function: trilinear devoxelization (backward) 59 | Args: 60 | grad_y : grad outputs, FloatTensor[b, c, n] 61 | indices : the voxel coordinates of point cube, IntTensor[b, 8, n] 62 | weights : weight for trilinear interpolation, FloatTensor[b, 8, n] 63 | r : voxel resolution 64 | Return: 65 | grad_x : grad inputs, FloatTensor[b, c, s], s = r ** 3 66 | */ 67 | at::Tensor trilinear_devoxelize_backward(const at::Tensor grad_y, 68 | const at::Tensor indices, 69 | const at::Tensor weights, 70 | const int r) { 71 | CHECK_CUDA(grad_y); 72 | CHECK_CUDA(weights); 73 | CHECK_CUDA(indices); 74 | CHECK_CONTIGUOUS(grad_y); 75 | CHECK_CONTIGUOUS(weights); 76 | CHECK_CONTIGUOUS(indices); 77 | CHECK_IS_FLOAT(grad_y); 78 | CHECK_IS_FLOAT(weights); 79 | CHECK_IS_INT(indices); 80 | 81 | int b = grad_y.size(0); 82 | int c = grad_y.size(1); 83 | int n = grad_y.size(2); 84 | int r3 = r * r * r; 85 | at::Tensor grad_x = torch::zeros( 86 | {b, c, r3}, at::device(grad_y.device()).dtype(at::ScalarType::Float)); 87 | trilinear_devoxelize_grad(b, c, n, r3, indices.data_ptr(), 88 | weights.data_ptr(), grad_y.data_ptr(), 89 | grad_x.data_ptr()); 90 | return grad_x; 91 | } 92 | -------------------------------------------------------------------------------- /modules/functional/src/interpolate/trilinear_devox.cppZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/interpolate/trilinear_devox.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "../cuda_utils.cuh" 5 | 6 | /* 7 | Function: trilinear devoxlization (forward) 8 | Args: 9 | b : batch size 10 | c : #channels 11 | n : number of points 12 | r : voxel resolution 13 | r2 : r ** 2 14 | r3 : r ** 3 15 | coords : the coordinates of points, FloatTensor[b, 3, n] 16 | feat : features, FloatTensor[b, c, r3] 17 | inds : the voxel indices of point cube, IntTensor[b, 8, n] 18 | wgts : weight for trilinear interpolation, FloatTensor[b, 8, n] 19 | outs : outputs, FloatTensor[b, c, n] 20 | */ 21 | __global__ void trilinear_devoxelize_kernel(int b, int c, int n, int r, int r2, 22 | int r3, bool is_training, 23 | const float *__restrict__ coords, 24 | const float *__restrict__ feat, 25 | int *__restrict__ inds, 26 | float *__restrict__ wgts, 27 | float *__restrict__ outs) { 28 | int batch_index = blockIdx.x; 29 | int stride = blockDim.x; 30 | int index = threadIdx.x; 31 | coords += batch_index * n * 3; 32 | inds += batch_index * n * 8; 33 | wgts += batch_index * n * 8; 34 | feat += batch_index * c * r3; 35 | outs += batch_index * c * n; 36 | 37 | for (int i = index; i < n; i += stride) { 38 | float x = coords[i]; 39 | float y = coords[i + n]; 40 | float z = coords[i + n + n]; 41 | float x_lo_f = floorf(x); 42 | float y_lo_f = floorf(y); 43 | float z_lo_f = floorf(z); 44 | 45 | float x_d_1 = x - x_lo_f; // / (x_hi_f - x_lo_f + 1e-8f) 46 | float y_d_1 = y - y_lo_f; 47 | float z_d_1 = z - z_lo_f; 48 | float x_d_0 = 1.0f - x_d_1; 49 | float y_d_0 = 1.0f - y_d_1; 50 | float z_d_0 = 1.0f - z_d_1; 51 | 52 | float wgt000 = x_d_0 * y_d_0 * z_d_0; 53 | float wgt001 = x_d_0 * y_d_0 * z_d_1; 54 | float wgt010 = x_d_0 * y_d_1 * z_d_0; 55 | float wgt011 = x_d_0 * y_d_1 * z_d_1; 56 | float wgt100 = x_d_1 * y_d_0 * z_d_0; 57 | float wgt101 = x_d_1 * y_d_0 * z_d_1; 58 | float wgt110 = x_d_1 * y_d_1 * z_d_0; 59 | float wgt111 = x_d_1 * y_d_1 * z_d_1; 60 | 61 | int x_lo = static_cast(x_lo_f); 62 | int y_lo = static_cast(y_lo_f); 63 | int z_lo = static_cast(z_lo_f); 64 | int x_hi = (x_d_1 > 0) ? -1 : 0; 65 | int y_hi = (y_d_1 > 0) ? -1 : 0; 66 | int z_hi = (z_d_1 > 0) ? 1 : 0; 67 | 68 | int idx000 = x_lo * r2 + y_lo * r + z_lo; 69 | int idx001 = idx000 + z_hi; // x_lo * r2 + y_lo * r + z_hi; 70 | int idx010 = idx000 + (y_hi & r); // x_lo * r2 + y_hi * r + z_lo; 71 | int idx011 = idx010 + z_hi; // x_lo * r2 + y_hi * r + z_hi; 72 | int idx100 = idx000 + (x_hi & r2); // x_hi * r2 + y_lo * r + z_lo; 73 | int idx101 = idx100 + z_hi; // x_hi * r2 + y_lo * r + z_hi; 74 | int idx110 = idx100 + (y_hi & r); // x_hi * r2 + y_hi * r + z_lo; 75 | int idx111 = idx110 + z_hi; // x_hi * r2 + y_hi * r + z_hi; 76 | 77 | if (is_training) { 78 | wgts[i] = wgt000; 79 | wgts[i + n] = wgt001; 80 | wgts[i + n * 2] = wgt010; 81 | wgts[i + n * 3] = wgt011; 82 | wgts[i + n * 4] = wgt100; 83 | wgts[i + n * 5] = wgt101; 84 | wgts[i + n * 6] = wgt110; 85 | wgts[i + n * 7] = wgt111; 86 | inds[i] = idx000; 87 | inds[i + n] = idx001; 88 | inds[i + n * 2] = idx010; 89 | inds[i + n * 3] = idx011; 90 | inds[i + n * 4] = idx100; 91 | inds[i + n * 5] = idx101; 92 | inds[i + n * 6] = idx110; 93 | inds[i + n * 7] = idx111; 94 | } 95 | 96 | for (int j = 0; j < c; j++) { 97 | int jr3 = j * r3; 98 | outs[j * n + i] = 99 | wgt000 * feat[jr3 + idx000] + wgt001 * feat[jr3 + idx001] + 100 | wgt010 * feat[jr3 + idx010] + wgt011 * feat[jr3 + idx011] + 101 | wgt100 * feat[jr3 + idx100] + wgt101 * feat[jr3 + idx101] + 102 | wgt110 * feat[jr3 + idx110] + wgt111 * feat[jr3 + idx111]; 103 | } 104 | } 105 | } 106 | 107 | /* 108 | Function: trilinear devoxlization (backward) 109 | Args: 110 | b : batch size 111 | c : #channels 112 | n : number of points 113 | r3 : voxel cube size = voxel resolution ** 3 114 | inds : the voxel indices of point cube, IntTensor[b, 8, n] 115 | wgts : weight for trilinear interpolation, FloatTensor[b, 8, n] 116 | grad_y : grad outputs, FloatTensor[b, c, n] 117 | grad_x : grad inputs, FloatTensor[b, c, r3] 118 | */ 119 | __global__ void trilinear_devoxelize_grad_kernel( 120 | int b, int c, int n, int r3, const int *__restrict__ inds, 121 | const float *__restrict__ wgts, const float *__restrict__ grad_y, 122 | float *__restrict__ grad_x) { 123 | int batch_index = blockIdx.x; 124 | int stride = blockDim.x; 125 | int index = threadIdx.x; 126 | inds += batch_index * n * 8; 127 | wgts += batch_index * n * 8; 128 | grad_x += batch_index * c * r3; 129 | grad_y += batch_index * c * n; 130 | 131 | for (int i = index; i < n; i += stride) { 132 | int idx000 = inds[i]; 133 | int idx001 = inds[i + n]; 134 | int idx010 = inds[i + n * 2]; 135 | int idx011 = inds[i + n * 3]; 136 | int idx100 = inds[i + n * 4]; 137 | int idx101 = inds[i + n * 5]; 138 | int idx110 = inds[i + n * 6]; 139 | int idx111 = inds[i + n * 7]; 140 | float wgt000 = wgts[i]; 141 | float wgt001 = wgts[i + n]; 142 | float wgt010 = wgts[i + n * 2]; 143 | float wgt011 = wgts[i + n * 3]; 144 | float wgt100 = wgts[i + n * 4]; 145 | float wgt101 = wgts[i + n * 5]; 146 | float wgt110 = wgts[i + n * 6]; 147 | float wgt111 = wgts[i + n * 7]; 148 | 149 | for (int j = 0; j < c; j++) { 150 | int jr3 = j * r3; 151 | float g = grad_y[j * n + i]; 152 | atomicAdd(grad_x + jr3 + idx000, wgt000 * g); 153 | atomicAdd(grad_x + jr3 + idx001, wgt001 * g); 154 | atomicAdd(grad_x + jr3 + idx010, wgt010 * g); 155 | atomicAdd(grad_x + jr3 + idx011, wgt011 * g); 156 | atomicAdd(grad_x + jr3 + idx100, wgt100 * g); 157 | atomicAdd(grad_x + jr3 + idx101, wgt101 * g); 158 | atomicAdd(grad_x + jr3 + idx110, wgt110 * g); 159 | atomicAdd(grad_x + jr3 + idx111, wgt111 * g); 160 | } 161 | } 162 | } 163 | 164 | void trilinear_devoxelize(int b, int c, int n, int r, int r2, int r3, 165 | bool training, const float *coords, const float *feat, 166 | int *inds, float *wgts, float *outs) { 167 | trilinear_devoxelize_kernel<<>>( 168 | b, c, n, r, r2, r3, training, coords, feat, inds, wgts, outs); 169 | CUDA_CHECK_ERRORS(); 170 | } 171 | 172 | void trilinear_devoxelize_grad(int b, int c, int n, int r3, const int *inds, 173 | const float *wgts, const float *grad_y, 174 | float *grad_x) { 175 | trilinear_devoxelize_grad_kernel<<>>( 176 | b, c, n, r3, inds, wgts, grad_y, grad_x); 177 | CUDA_CHECK_ERRORS(); 178 | } 179 | -------------------------------------------------------------------------------- /modules/functional/src/interpolate/trilinear_devox.cuh: -------------------------------------------------------------------------------- 1 | #ifndef _TRILINEAR_DEVOX_CUH 2 | #define _TRILINEAR_DEVOX_CUH 3 | 4 | // CUDA function declarations 5 | void trilinear_devoxelize(int b, int c, int n, int r, int r2, int r3, 6 | bool is_training, const float *coords, 7 | const float *feat, int *inds, float *wgts, 8 | float *outs); 9 | void trilinear_devoxelize_grad(int b, int c, int n, int r3, const int *inds, 10 | const float *wgts, const float *grad_y, 11 | float *grad_x); 12 | 13 | #endif 14 | -------------------------------------------------------------------------------- /modules/functional/src/interpolate/trilinear_devox.cuhZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/interpolate/trilinear_devox.cuZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/interpolate/trilinear_devox.hpp: -------------------------------------------------------------------------------- 1 | #ifndef _TRILINEAR_DEVOX_HPP 2 | #define _TRILINEAR_DEVOX_HPP 3 | 4 | #include 5 | #include 6 | 7 | std::vector trilinear_devoxelize_forward(const int r, 8 | const bool is_training, 9 | const at::Tensor coords, 10 | const at::Tensor features); 11 | 12 | at::Tensor trilinear_devoxelize_backward(const at::Tensor grad_y, 13 | const at::Tensor indices, 14 | const at::Tensor weights, const int r); 15 | 16 | #endif 17 | -------------------------------------------------------------------------------- /modules/functional/src/interpolate/trilinear_devox.hppZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/sampling/sampling.cpp: -------------------------------------------------------------------------------- 1 | #include "sampling.hpp" 2 | #include "sampling.cuh" 3 | 4 | #include "../utils.hpp" 5 | 6 | at::Tensor gather_features_forward(at::Tensor features, at::Tensor indices) { 7 | CHECK_CUDA(features); 8 | CHECK_CUDA(indices); 9 | CHECK_CONTIGUOUS(features); 10 | CHECK_CONTIGUOUS(indices); 11 | CHECK_IS_FLOAT(features); 12 | CHECK_IS_INT(indices); 13 | 14 | int b = features.size(0); 15 | int c = features.size(1); 16 | int n = features.size(2); 17 | int m = indices.size(1); 18 | at::Tensor output = torch::zeros( 19 | {b, c, m}, at::device(features.device()).dtype(at::ScalarType::Float)); 20 | gather_features(b, c, n, m, features.data_ptr(), 21 | indices.data_ptr(), output.data_ptr()); 22 | return output; 23 | } 24 | 25 | at::Tensor gather_features_backward(at::Tensor grad_y, at::Tensor indices, 26 | const int n) { 27 | CHECK_CUDA(grad_y); 28 | CHECK_CUDA(indices); 29 | CHECK_CONTIGUOUS(grad_y); 30 | CHECK_CONTIGUOUS(indices); 31 | CHECK_IS_FLOAT(grad_y); 32 | CHECK_IS_INT(indices); 33 | 34 | int b = grad_y.size(0); 35 | int c = grad_y.size(1); 36 | at::Tensor grad_x = torch::zeros( 37 | {b, c, n}, at::device(grad_y.device()).dtype(at::ScalarType::Float)); 38 | gather_features_grad(b, c, n, indices.size(1), grad_y.data_ptr(), 39 | indices.data_ptr(), grad_x.data_ptr()); 40 | return grad_x; 41 | } 42 | 43 | at::Tensor furthest_point_sampling_forward(at::Tensor coords, 44 | const int num_samples) { 45 | CHECK_CUDA(coords); 46 | CHECK_CONTIGUOUS(coords); 47 | CHECK_IS_FLOAT(coords); 48 | 49 | int b = coords.size(0); 50 | int n = coords.size(2); 51 | at::Tensor indices = torch::zeros( 52 | {b, num_samples}, at::device(coords.device()).dtype(at::ScalarType::Int)); 53 | at::Tensor distances = torch::full( 54 | {b, n}, 1e38f, at::device(coords.device()).dtype(at::ScalarType::Float)); 55 | furthest_point_sampling(b, n, num_samples, coords.data_ptr(), 56 | distances.data_ptr(), indices.data_ptr()); 57 | return indices; 58 | } 59 | -------------------------------------------------------------------------------- /modules/functional/src/sampling/sampling.cppZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/sampling/sampling.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "../cuda_utils.cuh" 5 | 6 | /* 7 | Function: gather centers' features (forward) 8 | Args: 9 | b : batch size 10 | c : #channles of features 11 | n : number of points in point clouds 12 | m : number of query/sampled centers 13 | features: points' features, FloatTensor[b, c, n] 14 | indices : centers' indices in points, IntTensor[b, m] 15 | out : gathered features, FloatTensor[b, c, m] 16 | */ 17 | __global__ void gather_features_kernel(int b, int c, int n, int m, 18 | const float *__restrict__ features, 19 | const int *__restrict__ indices, 20 | float *__restrict__ out) { 21 | int batch_index = blockIdx.x; 22 | int channel_index = blockIdx.y; 23 | int temp_index = batch_index * c + channel_index; 24 | features += temp_index * n; 25 | indices += batch_index * m; 26 | out += temp_index * m; 27 | 28 | for (int j = threadIdx.x; j < m; j += blockDim.x) { 29 | out[j] = features[indices[j]]; 30 | } 31 | } 32 | 33 | void gather_features(int b, int c, int n, int m, const float *features, 34 | const int *indices, float *out) { 35 | gather_features_kernel<<>>( 37 | b, c, n, m, features, indices, out); 38 | CUDA_CHECK_ERRORS(); 39 | } 40 | 41 | /* 42 | Function: gather centers' features (backward) 43 | Args: 44 | b : batch size 45 | c : #channles of features 46 | n : number of points in point clouds 47 | m : number of query/sampled centers 48 | grad_y : grad of gathered features, FloatTensor[b, c, m] 49 | indices : centers' indices in points, IntTensor[b, m] 50 | grad_x : grad of points' features, FloatTensor[b, c, n] 51 | */ 52 | __global__ void gather_features_grad_kernel(int b, int c, int n, int m, 53 | const float *__restrict__ grad_y, 54 | const int *__restrict__ indices, 55 | float *__restrict__ grad_x) { 56 | int batch_index = blockIdx.x; 57 | int channel_index = blockIdx.y; 58 | int temp_index = batch_index * c + channel_index; 59 | grad_y += temp_index * m; 60 | indices += batch_index * m; 61 | grad_x += temp_index * n; 62 | 63 | for (int j = threadIdx.x; j < m; j += blockDim.x) { 64 | atomicAdd(grad_x + indices[j], grad_y[j]); 65 | } 66 | } 67 | 68 | void gather_features_grad(int b, int c, int n, int m, const float *grad_y, 69 | const int *indices, float *grad_x) { 70 | gather_features_grad_kernel<<>>( 72 | b, c, n, m, grad_y, indices, grad_x); 73 | CUDA_CHECK_ERRORS(); 74 | } 75 | 76 | /* 77 | Function: furthest point sampling 78 | Args: 79 | b : batch size 80 | n : number of points in point clouds 81 | m : number of query/sampled centers 82 | coords : points' coords, FloatTensor[b, 3, n] 83 | distances : minimum distance of a point to the set, IntTensor[b, n] 84 | indices : sampled centers' indices in points, IntTensor[b, m] 85 | */ 86 | __global__ void furthest_point_sampling_kernel(int b, int n, int m, 87 | const float *__restrict__ coords, 88 | float *__restrict__ distances, 89 | int *__restrict__ indices) { 90 | if (m <= 0) 91 | return; 92 | int batch_index = blockIdx.x; 93 | coords += batch_index * n * 3; 94 | distances += batch_index * n; 95 | indices += batch_index * m; 96 | 97 | const int BlockSize = 512; 98 | __shared__ float dists[BlockSize]; 99 | __shared__ int dists_i[BlockSize]; 100 | const int BufferSize = 3072; 101 | __shared__ float buf[BufferSize * 3]; 102 | 103 | int old = 0; 104 | if (threadIdx.x == 0) 105 | indices[0] = old; 106 | 107 | for (int j = threadIdx.x; j < min(BufferSize, n); j += blockDim.x) { 108 | buf[j] = coords[j]; 109 | buf[j + BufferSize] = coords[j + n]; 110 | buf[j + BufferSize + BufferSize] = coords[j + n + n]; 111 | } 112 | __syncthreads(); 113 | 114 | for (int j = 1; j < m; j++) { 115 | int besti = 0; // best index 116 | float best = -1; // farthest distance 117 | // calculating the distance with the latest sampled point 118 | float x1 = coords[old]; 119 | float y1 = coords[old + n]; 120 | float z1 = coords[old + n + n]; 121 | for (int k = threadIdx.x; k < n; k += blockDim.x) { 122 | // fetch distance at block n, thread k 123 | float td = distances[k]; 124 | float x2, y2, z2; 125 | if (k < BufferSize) { 126 | x2 = buf[k]; 127 | y2 = buf[k + BufferSize]; 128 | z2 = buf[k + BufferSize + BufferSize]; 129 | } else { 130 | x2 = coords[k]; 131 | y2 = coords[k + n]; 132 | z2 = coords[k + n + n]; 133 | } 134 | float d = 135 | (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + (z2 - z1) * (z2 - z1); 136 | float d2 = min(d, td); 137 | // update "point-to-set" distance 138 | if (d2 != td) 139 | distances[k] = d2; 140 | // update the farthest distance at sample step j 141 | if (d2 > best) { 142 | best = d2; 143 | besti = k; 144 | } 145 | } 146 | 147 | dists[threadIdx.x] = best; 148 | dists_i[threadIdx.x] = besti; 149 | for (int u = 0; (1 << u) < blockDim.x; u++) { 150 | __syncthreads(); 151 | if (threadIdx.x < (blockDim.x >> (u + 1))) { 152 | int i1 = (threadIdx.x * 2) << u; 153 | int i2 = (threadIdx.x * 2 + 1) << u; 154 | if (dists[i1] < dists[i2]) { 155 | dists[i1] = dists[i2]; 156 | dists_i[i1] = dists_i[i2]; 157 | } 158 | } 159 | } 160 | __syncthreads(); 161 | 162 | // finish sample step j; old is the sampled index 163 | old = dists_i[0]; 164 | if (threadIdx.x == 0) 165 | indices[j] = old; 166 | } 167 | } 168 | 169 | void furthest_point_sampling(int b, int n, int m, const float *coords, 170 | float *distances, int *indices) { 171 | furthest_point_sampling_kernel<<>>(b, n, m, coords, distances, 172 | indices); 173 | CUDA_CHECK_ERRORS(); 174 | } 175 | -------------------------------------------------------------------------------- /modules/functional/src/sampling/sampling.cuh: -------------------------------------------------------------------------------- 1 | #ifndef _SAMPLING_CUH 2 | #define _SAMPLING_CUH 3 | 4 | void gather_features(int b, int c, int n, int m, const float *features, 5 | const int *indices, float *out); 6 | void gather_features_grad(int b, int c, int n, int m, const float *grad_y, 7 | const int *indices, float *grad_x); 8 | void furthest_point_sampling(int b, int n, int m, const float *coords, 9 | float *distances, int *indices); 10 | 11 | #endif 12 | -------------------------------------------------------------------------------- /modules/functional/src/sampling/sampling.cuhZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/sampling/sampling.cuZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/sampling/sampling.hpp: -------------------------------------------------------------------------------- 1 | #ifndef _SAMPLING_HPP 2 | #define _SAMPLING_HPP 3 | 4 | #include 5 | 6 | at::Tensor gather_features_forward(at::Tensor features, at::Tensor indices); 7 | at::Tensor gather_features_backward(at::Tensor grad_y, at::Tensor indices, 8 | const int n); 9 | at::Tensor furthest_point_sampling_forward(at::Tensor coords, 10 | const int num_samples); 11 | 12 | #endif 13 | -------------------------------------------------------------------------------- /modules/functional/src/sampling/sampling.hppZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/utils.hpp: -------------------------------------------------------------------------------- 1 | #ifndef _UTILS_HPP 2 | #define _UTILS_HPP 3 | 4 | #include 5 | #include 6 | 7 | #define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor") 8 | 9 | #define CHECK_CONTIGUOUS(x) \ 10 | TORCH_CHECK(x.is_contiguous(), #x " must be a contiguous tensor") 11 | 12 | #define CHECK_IS_INT(x) \ 13 | TORCH_CHECK(x.scalar_type() == at::ScalarType::Int, \ 14 | #x " must be an int tensor") 15 | 16 | #define CHECK_IS_FLOAT(x) \ 17 | TORCH_CHECK(x.scalar_type() == at::ScalarType::Float, \ 18 | #x " must be a float tensor") 19 | 20 | #endif 21 | -------------------------------------------------------------------------------- /modules/functional/src/voxelization/vox.cpp: -------------------------------------------------------------------------------- 1 | #include "vox.hpp" 2 | #include "vox.cuh" 3 | 4 | #include "../utils.hpp" 5 | 6 | /* 7 | Function: average pool voxelization (forward) 8 | Args: 9 | features: features, FloatTensor[b, c, n] 10 | coords : coords of each point, IntTensor[b, 3, n] 11 | resolution : voxel resolution 12 | Return: 13 | out : outputs, FloatTensor[b, c, s], s = r ** 3 14 | ind : voxel index of each point, IntTensor[b, n] 15 | cnt : #points in each voxel index, IntTensor[b, s] 16 | */ 17 | std::vector avg_voxelize_forward(const at::Tensor features, 18 | const at::Tensor coords, 19 | const int resolution) { 20 | CHECK_CUDA(features); 21 | CHECK_CUDA(coords); 22 | CHECK_CONTIGUOUS(features); 23 | CHECK_CONTIGUOUS(coords); 24 | CHECK_IS_FLOAT(features); 25 | CHECK_IS_INT(coords); 26 | 27 | int b = features.size(0); 28 | int c = features.size(1); 29 | int n = features.size(2); 30 | int r = resolution; 31 | int r2 = r * r; 32 | int r3 = r2 * r; 33 | at::Tensor ind = torch::zeros( 34 | {b, n}, at::device(features.device()).dtype(at::ScalarType::Int)); 35 | at::Tensor out = torch::zeros( 36 | {b, c, r3}, at::device(features.device()).dtype(at::ScalarType::Float)); 37 | at::Tensor cnt = torch::zeros( 38 | {b, r3}, at::device(features.device()).dtype(at::ScalarType::Int)); 39 | avg_voxelize(b, c, n, r, r2, r3, coords.data_ptr(), 40 | features.data_ptr(), ind.data_ptr(), 41 | cnt.data_ptr(), out.data_ptr()); 42 | return {out, ind, cnt}; 43 | } 44 | 45 | /* 46 | Function: average pool voxelization (backward) 47 | Args: 48 | grad_y : grad outputs, FloatTensor[b, c, s] 49 | indices: voxel index of each point, IntTensor[b, n] 50 | cnt : #points in each voxel index, IntTensor[b, s] 51 | Return: 52 | grad_x : grad inputs, FloatTensor[b, c, n] 53 | */ 54 | at::Tensor avg_voxelize_backward(const at::Tensor grad_y, 55 | const at::Tensor indices, 56 | const at::Tensor cnt) { 57 | CHECK_CUDA(grad_y); 58 | CHECK_CUDA(indices); 59 | CHECK_CUDA(cnt); 60 | CHECK_CONTIGUOUS(grad_y); 61 | CHECK_CONTIGUOUS(indices); 62 | CHECK_CONTIGUOUS(cnt); 63 | CHECK_IS_FLOAT(grad_y); 64 | CHECK_IS_INT(indices); 65 | CHECK_IS_INT(cnt); 66 | 67 | int b = grad_y.size(0); 68 | int c = grad_y.size(1); 69 | int s = grad_y.size(2); 70 | int n = indices.size(1); 71 | at::Tensor grad_x = torch::zeros( 72 | {b, c, n}, at::device(grad_y.device()).dtype(at::ScalarType::Float)); 73 | avg_voxelize_grad(b, c, n, s, indices.data_ptr(), cnt.data_ptr(), 74 | grad_y.data_ptr(), grad_x.data_ptr()); 75 | return grad_x; 76 | } 77 | -------------------------------------------------------------------------------- /modules/functional/src/voxelization/vox.cppZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/voxelization/vox.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "../cuda_utils.cuh" 5 | 6 | /* 7 | Function: get how many points in each voxel grid 8 | Args: 9 | b : batch size 10 | n : number of points 11 | r : voxel resolution 12 | r2 : = r * r 13 | r3 : s, voxel cube size = r ** 3 14 | coords : coords of each point, IntTensor[b, 3, n] 15 | ind : voxel index of each point, IntTensor[b, n] 16 | cnt : #points in each voxel index, IntTensor[b, s] 17 | */ 18 | __global__ void grid_stats_kernel(int b, int n, int r, int r2, int r3, 19 | const int *__restrict__ coords, 20 | int *__restrict__ ind, int *cnt) { 21 | int batch_index = blockIdx.x; 22 | int stride = blockDim.x; 23 | int index = threadIdx.x; 24 | coords += batch_index * n * 3; 25 | ind += batch_index * n; 26 | cnt += batch_index * r3; 27 | 28 | for (int i = index; i < n; i += stride) { 29 | // if (ind[i] == -1) 30 | // continue; 31 | ind[i] = coords[i] * r2 + coords[i + n] * r + coords[i + n + n]; 32 | atomicAdd(cnt + ind[i], 1); 33 | } 34 | } 35 | 36 | /* 37 | Function: average pool voxelization (forward) 38 | Args: 39 | b : batch size 40 | c : #channels 41 | n : number of points 42 | s : voxel cube size = voxel resolution ** 3 43 | ind : voxel index of each point, IntTensor[b, n] 44 | cnt : #points in each voxel index, IntTensor[b, s] 45 | feat: features, FloatTensor[b, c, n] 46 | out : outputs, FloatTensor[b, c, s] 47 | */ 48 | __global__ void avg_voxelize_kernel(int b, int c, int n, int s, 49 | const int *__restrict__ ind, 50 | const int *__restrict__ cnt, 51 | const float *__restrict__ feat, 52 | float *__restrict__ out) { 53 | int batch_index = blockIdx.x; 54 | int stride = blockDim.x; 55 | int index = threadIdx.x; 56 | ind += batch_index * n; 57 | feat += batch_index * c * n; 58 | out += batch_index * c * s; 59 | cnt += batch_index * s; 60 | for (int i = index; i < n; i += stride) { 61 | int pos = ind[i]; 62 | // if (pos == -1) 63 | // continue; 64 | int cur_cnt = cnt[pos]; 65 | if (cur_cnt > 0) { 66 | float div_cur_cnt = 1.0 / static_cast(cur_cnt); 67 | for (int j = 0; j < c; j++) { 68 | atomicAdd(out + j * s + pos, feat[j * n + i] * div_cur_cnt); 69 | } 70 | } 71 | } 72 | } 73 | 74 | /* 75 | Function: average pool voxelization (backward) 76 | Args: 77 | b : batch size 78 | c : #channels 79 | n : number of points 80 | r3 : voxel cube size = voxel resolution ** 3 81 | ind : voxel index of each point, IntTensor[b, n] 82 | cnt : #points in each voxel index, IntTensor[b, s] 83 | grad_y : grad outputs, FloatTensor[b, c, s] 84 | grad_x : grad inputs, FloatTensor[b, c, n] 85 | */ 86 | __global__ void avg_voxelize_grad_kernel(int b, int c, int n, int r3, 87 | const int *__restrict__ ind, 88 | const int *__restrict__ cnt, 89 | const float *__restrict__ grad_y, 90 | float *__restrict__ grad_x) { 91 | int batch_index = blockIdx.x; 92 | int stride = blockDim.x; 93 | int index = threadIdx.x; 94 | ind += batch_index * n; 95 | grad_x += batch_index * c * n; 96 | grad_y += batch_index * c * r3; 97 | cnt += batch_index * r3; 98 | for (int i = index; i < n; i += stride) { 99 | int pos = ind[i]; 100 | // if (pos == -1) 101 | // continue; 102 | int cur_cnt = cnt[pos]; 103 | if (cur_cnt > 0) { 104 | float div_cur_cnt = 1.0 / static_cast(cur_cnt); 105 | for (int j = 0; j < c; j++) { 106 | atomicAdd(grad_x + j * n + i, grad_y[j * r3 + pos] * div_cur_cnt); 107 | } 108 | } 109 | } 110 | } 111 | 112 | void avg_voxelize(int b, int c, int n, int r, int r2, int r3, const int *coords, 113 | const float *feat, int *ind, int *cnt, float *out) { 114 | grid_stats_kernel<<>>(b, n, r, r2, r3, coords, ind, 115 | cnt); 116 | avg_voxelize_kernel<<>>(b, c, n, r3, ind, cnt, 117 | feat, out); 118 | CUDA_CHECK_ERRORS(); 119 | } 120 | 121 | void avg_voxelize_grad(int b, int c, int n, int s, const int *ind, 122 | const int *cnt, const float *grad_y, float *grad_x) { 123 | avg_voxelize_grad_kernel<<>>(b, c, n, s, ind, cnt, 124 | grad_y, grad_x); 125 | CUDA_CHECK_ERRORS(); 126 | } 127 | -------------------------------------------------------------------------------- /modules/functional/src/voxelization/vox.cuh: -------------------------------------------------------------------------------- 1 | #ifndef _VOX_CUH 2 | #define _VOX_CUH 3 | 4 | // CUDA function declarations 5 | void avg_voxelize(int b, int c, int n, int r, int r2, int r3, const int *coords, 6 | const float *feat, int *ind, int *cnt, float *out); 7 | void avg_voxelize_grad(int b, int c, int n, int s, const int *idx, 8 | const int *cnt, const float *grad_y, float *grad_x); 9 | 10 | #endif 11 | -------------------------------------------------------------------------------- /modules/functional/src/voxelization/vox.cuhZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/voxelization/vox.cuZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/src/voxelization/vox.hpp: -------------------------------------------------------------------------------- 1 | #ifndef _VOX_HPP 2 | #define _VOX_HPP 3 | 4 | #include 5 | #include 6 | 7 | std::vector avg_voxelize_forward(const at::Tensor features, 8 | const at::Tensor coords, 9 | const int resolution); 10 | 11 | at::Tensor avg_voxelize_backward(const at::Tensor grad_y, 12 | const at::Tensor indices, 13 | const at::Tensor cnt); 14 | 15 | #endif 16 | -------------------------------------------------------------------------------- /modules/functional/src/voxelization/vox.hppZone.Identifier: -------------------------------------------------------------------------------- 1 | [ZoneTransfer] 2 | ZoneId=3 3 | ReferrerUrl=C:\Users\13476\Downloads\pvcnn-master.zip 4 | -------------------------------------------------------------------------------- /modules/functional/voxelization.py: -------------------------------------------------------------------------------- 1 | from torch.autograd import Function 2 | 3 | from modules.functional.backend import _backend 4 | 5 | __all__ = ['avg_voxelize'] 6 | 7 | 8 | class AvgVoxelization(Function): 9 | @staticmethod 10 | def forward(ctx, features, coords, resolution): 11 | """ 12 | :param ctx: 13 | :param features: Features of the point cloud, FloatTensor[B, C, N] 14 | :param coords: Voxelized Coordinates of each point, IntTensor[B, 3, N] 15 | :param resolution: Voxel resolution 16 | :return: 17 | Voxelized Features, FloatTensor[B, C, R, R, R] 18 | """ 19 | features = features.contiguous() 20 | coords = coords.int().contiguous() 21 | b, c, _ = features.shape 22 | out, indices, counts = _backend.avg_voxelize_forward(features, coords, resolution) 23 | ctx.save_for_backward(indices, counts) 24 | return out.view(b, c, resolution, resolution, resolution) 25 | 26 | @staticmethod 27 | def backward(ctx, grad_output): 28 | """ 29 | :param ctx: 30 | :param grad_output: gradient of output, FloatTensor[B, C, R, R, R] 31 | :return: 32 | gradient of inputs, FloatTensor[B, C, N] 33 | """ 34 | b, c = grad_output.shape[:2] 35 | indices, counts = ctx.saved_tensors 36 | grad_features = _backend.avg_voxelize_backward(grad_output.contiguous().view(b, c, -1), indices, counts) 37 | return grad_features, None, None 38 | 39 | 40 | avg_voxelize = AvgVoxelization.apply 41 | -------------------------------------------------------------------------------- /modules/loss.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | import modules.functional as F 4 | 5 | __all__ = ['KLLoss'] 6 | 7 | 8 | class KLLoss(nn.Module): 9 | def forward(self, x, y): 10 | return F.kl_loss(x, y) 11 | -------------------------------------------------------------------------------- /modules/pointnet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | import modules.functional as F 5 | from modules.ball_query import BallQuery 6 | from modules.shared_mlp import SharedMLP 7 | 8 | __all__ = ['PointNetAModule', 'PointNetSAModule', 'PointNetFPModule'] 9 | 10 | 11 | class PointNetAModule(nn.Module): 12 | def __init__(self, in_channels, out_channels, include_coordinates=True): 13 | super().__init__() 14 | if not isinstance(out_channels, (list, tuple)): 15 | out_channels = [[out_channels]] 16 | elif not isinstance(out_channels[0], (list, tuple)): 17 | out_channels = [out_channels] 18 | 19 | mlps = [] 20 | total_out_channels = 0 21 | for _out_channels in out_channels: 22 | mlps.append( 23 | SharedMLP(in_channels=in_channels + (3 if include_coordinates else 0), 24 | out_channels=_out_channels, dim=1) 25 | ) 26 | total_out_channels += _out_channels[-1] 27 | 28 | self.include_coordinates = include_coordinates 29 | self.out_channels = total_out_channels 30 | self.mlps = nn.ModuleList(mlps) 31 | 32 | def forward(self, inputs): 33 | features, coords = inputs 34 | if self.include_coordinates: 35 | features = torch.cat([features, coords], dim=1) 36 | coords = torch.zeros((coords.size(0), 3, 1), device=coords.device) 37 | if len(self.mlps) > 1: 38 | features_list = [] 39 | for mlp in self.mlps: 40 | features_list.append(mlp(features).max(dim=-1, keepdim=True).values) 41 | return torch.cat(features_list, dim=1), coords 42 | else: 43 | return self.mlps[0](features).max(dim=-1, keepdim=True).values, coords 44 | 45 | def extra_repr(self): 46 | return f'out_channels={self.out_channels}, include_coordinates={self.include_coordinates}' 47 | 48 | 49 | class PointNetSAModule(nn.Module): 50 | def __init__(self, num_centers, radius, num_neighbors, in_channels, out_channels, include_coordinates=True): 51 | super().__init__() 52 | if not isinstance(radius, (list, tuple)): 53 | radius = [radius] 54 | if not isinstance(num_neighbors, (list, tuple)): 55 | num_neighbors = [num_neighbors] * len(radius) 56 | assert len(radius) == len(num_neighbors) 57 | if not isinstance(out_channels, (list, tuple)): 58 | out_channels = [[out_channels]] * len(radius) 59 | elif not isinstance(out_channels[0], (list, tuple)): 60 | out_channels = [out_channels] * len(radius) 61 | assert len(radius) == len(out_channels) 62 | 63 | groupers, mlps = [], [] 64 | total_out_channels = 0 65 | for _radius, _out_channels, _num_neighbors in zip(radius, out_channels, num_neighbors): 66 | groupers.append( 67 | BallQuery(radius=_radius, num_neighbors=_num_neighbors, include_coordinates=include_coordinates) 68 | ) 69 | mlps.append( 70 | SharedMLP(in_channels=in_channels + (3 if include_coordinates else 0), 71 | out_channels=_out_channels, dim=2) 72 | ) 73 | total_out_channels += _out_channels[-1] 74 | 75 | self.num_centers = num_centers 76 | self.out_channels = total_out_channels 77 | self.groupers = nn.ModuleList(groupers) 78 | self.mlps = nn.ModuleList(mlps) 79 | 80 | def forward(self, inputs): 81 | features, coords = inputs 82 | centers_coords = F.furthest_point_sample(coords, self.num_centers) 83 | features_list = [] 84 | for grouper, mlp in zip(self.groupers, self.mlps): 85 | features_list.append(mlp(grouper(coords, centers_coords, features)).max(dim=-1).values) 86 | if len(features_list) > 1: 87 | return torch.cat(features_list, dim=1), centers_coords 88 | else: 89 | return features_list[0], centers_coords 90 | 91 | def extra_repr(self): 92 | return f'num_centers={self.num_centers}, out_channels={self.out_channels}' 93 | 94 | 95 | class PointNetFPModule(nn.Module): 96 | def __init__(self, in_channels, out_channels): 97 | super().__init__() 98 | self.mlp = SharedMLP(in_channels=in_channels, out_channels=out_channels, dim=1) 99 | 100 | def forward(self, inputs): 101 | if len(inputs) == 3: 102 | points_coords, centers_coords, centers_features = inputs 103 | points_features = None 104 | else: 105 | points_coords, centers_coords, centers_features, points_features = inputs 106 | interpolated_features = F.nearest_neighbor_interpolate(points_coords, centers_coords, centers_features) 107 | if points_features is not None: 108 | interpolated_features = torch.cat( 109 | [interpolated_features, points_features], dim=1 110 | ) 111 | return self.mlp(interpolated_features), points_coords 112 | -------------------------------------------------------------------------------- /modules/se.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | __all__ = ['SE3d'] 4 | 5 | 6 | class SE3d(nn.Module): 7 | def __init__(self, channel, reduction=8): 8 | super().__init__() 9 | self.fc = nn.Sequential( 10 | nn.Linear(channel, channel // reduction, bias=False), 11 | nn.ReLU(inplace=True), 12 | nn.Linear(channel // reduction, channel, bias=False), 13 | nn.Sigmoid() 14 | ) 15 | 16 | def forward(self, inputs): 17 | return inputs * self.fc(inputs.mean(-1).mean(-1).mean(-1)).view(inputs.shape[0], inputs.shape[1], 1, 1, 1) 18 | -------------------------------------------------------------------------------- /modules/shared_mlp.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | __all__ = ['SharedMLP'] 4 | 5 | 6 | class SharedMLP(nn.Module): 7 | def __init__(self, in_channels, out_channels, dim=1): 8 | super().__init__() 9 | if dim == 1: 10 | conv = nn.Conv1d 11 | bn = nn.BatchNorm1d 12 | elif dim == 2: 13 | conv = nn.Conv2d 14 | bn = nn.BatchNorm2d 15 | else: 16 | raise ValueError 17 | if not isinstance(out_channels, (list, tuple)): 18 | out_channels = [out_channels] 19 | layers = [] 20 | for oc in out_channels: 21 | layers.extend([ 22 | conv(in_channels, oc, 1), 23 | bn(oc), 24 | nn.ReLU(True), 25 | ]) 26 | in_channels = oc 27 | self.layers = nn.Sequential(*layers) 28 | 29 | def forward(self, inputs): 30 | if isinstance(inputs, (list, tuple)): 31 | return (self.layers(inputs[0]), *inputs[1:]) 32 | else: 33 | return self.layers(inputs) 34 | -------------------------------------------------------------------------------- /modules/utils.py: -------------------------------------------------------------------------------- 1 | import functools 2 | 3 | import torch.nn as nn 4 | 5 | from modules import SharedMLP, PVConv, PointNetSAModule, PointNetAModule, PointNetFPModule 6 | 7 | __all__ = ['create_mlp_components', 'create_pointnet_components', 8 | 'create_pointnet2_sa_components', 'create_pointnet2_fp_modules'] 9 | 10 | 11 | def _linear_bn_relu(in_channels, out_channels): 12 | return nn.Sequential(nn.Linear(in_channels, out_channels), nn.BatchNorm1d(out_channels), nn.ReLU(True)) 13 | 14 | 15 | def create_mlp_components(in_channels, out_channels, classifier=False, dim=2, width_multiplier=1): 16 | r = width_multiplier 17 | 18 | if dim == 1: 19 | block = _linear_bn_relu 20 | else: 21 | block = SharedMLP 22 | if not isinstance(out_channels, (list, tuple)): 23 | out_channels = [out_channels] 24 | if len(out_channels) == 0 or (len(out_channels) == 1 and out_channels[0] is None): 25 | return nn.Sequential(), in_channels, in_channels 26 | 27 | layers = [] 28 | for oc in out_channels[:-1]: 29 | if oc < 1: 30 | layers.append(nn.Dropout(oc)) 31 | else: 32 | oc = int(r * oc) 33 | layers.append(block(in_channels, oc)) 34 | in_channels = oc 35 | if dim == 1: 36 | if classifier: 37 | layers.append(nn.Linear(in_channels, out_channels[-1])) 38 | else: 39 | layers.append(_linear_bn_relu(in_channels, int(r * out_channels[-1]))) 40 | else: 41 | if classifier: 42 | layers.append(nn.Conv1d(in_channels, out_channels[-1], 1)) 43 | else: 44 | layers.append(SharedMLP(in_channels, int(r * out_channels[-1]))) 45 | return layers, out_channels[-1] if classifier else int(r * out_channels[-1]) 46 | 47 | 48 | def create_pointnet_components(blocks, in_channels, with_se=False, normalize=True, eps=0, 49 | width_multiplier=1, voxel_resolution_multiplier=1): 50 | r, vr = width_multiplier, voxel_resolution_multiplier 51 | 52 | layers, concat_channels = [], 0 53 | for out_channels, num_blocks, voxel_resolution in blocks: 54 | out_channels = int(r * out_channels) 55 | if voxel_resolution is None: 56 | block = SharedMLP 57 | else: 58 | block = functools.partial(PVConv, kernel_size=3, resolution=int(vr * voxel_resolution), 59 | with_se=with_se, normalize=normalize, eps=eps) 60 | for _ in range(num_blocks): 61 | layers.append(block(in_channels, out_channels)) 62 | in_channels = out_channels 63 | concat_channels += out_channels 64 | return layers, in_channels, concat_channels 65 | 66 | 67 | def create_pointnet2_sa_components(sa_blocks, extra_feature_channels, with_se=False, normalize=True, eps=0, 68 | width_multiplier=1, voxel_resolution_multiplier=1): 69 | r, vr = width_multiplier, voxel_resolution_multiplier 70 | in_channels = extra_feature_channels + 3 71 | 72 | sa_layers, sa_in_channels = [], [] 73 | for conv_configs, sa_configs in sa_blocks: 74 | sa_in_channels.append(in_channels) 75 | sa_blocks = [] 76 | if conv_configs is not None: 77 | out_channels, num_blocks, voxel_resolution = conv_configs 78 | out_channels = int(r * out_channels) 79 | if voxel_resolution is None: 80 | block = SharedMLP 81 | else: 82 | block = functools.partial(PVConv, kernel_size=3, resolution=int(vr * voxel_resolution), 83 | with_se=with_se, normalize=normalize, eps=eps) 84 | for _ in range(num_blocks): 85 | sa_blocks.append(block(in_channels, out_channels)) 86 | in_channels = out_channels 87 | extra_feature_channels = in_channels 88 | num_centers, radius, num_neighbors, out_channels = sa_configs 89 | _out_channels = [] 90 | for oc in out_channels: 91 | if isinstance(oc, (list, tuple)): 92 | _out_channels.append([int(r * _oc) for _oc in oc]) 93 | else: 94 | _out_channels.append(int(r * oc)) 95 | out_channels = _out_channels 96 | if num_centers is None: 97 | block = PointNetAModule 98 | else: 99 | block = functools.partial(PointNetSAModule, num_centers=num_centers, radius=radius, 100 | num_neighbors=num_neighbors) 101 | sa_blocks.append(block(in_channels=extra_feature_channels, out_channels=out_channels, 102 | include_coordinates=True)) 103 | in_channels = extra_feature_channels = sa_blocks[-1].out_channels 104 | if len(sa_blocks) == 1: 105 | sa_layers.append(sa_blocks[0]) 106 | else: 107 | sa_layers.append(nn.Sequential(*sa_blocks)) 108 | 109 | return sa_layers, sa_in_channels, in_channels, 1 if num_centers is None else num_centers 110 | 111 | 112 | def create_pointnet2_fp_modules(fp_blocks, in_channels, sa_in_channels, with_se=False, normalize=True, eps=0, 113 | width_multiplier=1, voxel_resolution_multiplier=1): 114 | r, vr = width_multiplier, voxel_resolution_multiplier 115 | 116 | fp_layers = [] 117 | for fp_idx, (fp_configs, conv_configs) in enumerate(fp_blocks): 118 | fp_blocks = [] 119 | out_channels = tuple(int(r * oc) for oc in fp_configs) 120 | fp_blocks.append( 121 | PointNetFPModule(in_channels=in_channels + sa_in_channels[-1 - fp_idx], out_channels=out_channels) 122 | ) 123 | in_channels = out_channels[-1] 124 | if conv_configs is not None: 125 | out_channels, num_blocks, voxel_resolution = conv_configs 126 | out_channels = int(r * out_channels) 127 | if voxel_resolution is None: 128 | block = SharedMLP 129 | else: 130 | block = functools.partial(PVConv, kernel_size=3, resolution=int(vr * voxel_resolution), 131 | with_se=with_se, normalize=normalize, eps=eps) 132 | for _ in range(num_blocks): 133 | fp_blocks.append(block(in_channels, out_channels)) 134 | in_channels = out_channels 135 | if len(fp_blocks) == 1: 136 | fp_layers.append(fp_blocks[0]) 137 | else: 138 | fp_layers.append(nn.Sequential(*fp_blocks)) 139 | 140 | return fp_layers, in_channels 141 | -------------------------------------------------------------------------------- /modules/voxelization.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | import modules.functional as F 5 | 6 | __all__ = ['Voxelization'] 7 | 8 | 9 | class Voxelization(nn.Module): 10 | def __init__(self, resolution, normalize=True, eps=0): 11 | super().__init__() 12 | self.r = int(resolution) 13 | self.normalize = normalize 14 | self.eps = eps 15 | 16 | def forward(self, features, coords): 17 | coords = coords.detach() 18 | norm_coords = coords - coords.mean(2, keepdim=True) 19 | if self.normalize: 20 | norm_coords = norm_coords / (norm_coords.norm(dim=1, keepdim=True).max(dim=2, keepdim=True).values * 2.0 + self.eps) + 0.5 21 | else: 22 | norm_coords = (norm_coords + 1) / 2.0 23 | norm_coords = torch.clamp(norm_coords * self.r, 0, self.r - 1) 24 | vox_coords = torch.round(norm_coords).to(torch.int32) 25 | return F.avg_voxelize(features, vox_coords, self.r), norm_coords 26 | 27 | def extra_repr(self): 28 | return 'resolution={}{}'.format(self.r, ', normalized eps = {}'.format(self.eps) if self.normalize else '') 29 | -------------------------------------------------------------------------------- /scripts/dataset_matterport/build_dataset.py: -------------------------------------------------------------------------------- 1 | from os import listdir, makedirs, getcwd 2 | from os.path import join, exists, isdir, exists 3 | import json 4 | import trimesh 5 | import numpy as np 6 | from copy import deepcopy 7 | import shutil 8 | import zipfile 9 | from tqdm import tqdm 10 | from src.utils.io import export_pointcloud 11 | 12 | def create_dir(dir_in): 13 | if not exists(dir_in): 14 | makedirs(dir_in) 15 | 16 | base_path = 'data/Matterport3D/v1/scans' 17 | scene_name = 'JmbYfDe2QKZ' 18 | out_path = 'data/Matterport3D_processed' 19 | scene_path = join(base_path, scene_name, 'region_segmentations') 20 | regions = [join(scene_path, 'region'+str(m)+'.ply') for m in range(100) if exists(join(scene_path, 'region'+str(m)+'.ply'))] 21 | outfile = join(out_path, scene_name) 22 | create_dir(outfile) 23 | 24 | n_pointcloud_points = 500000 25 | dtype = np.float16 26 | cut_mesh =True 27 | save_part_mesh = False 28 | 29 | mat_permute = np.array([ 30 | [1, 0, 0, 0], 31 | [0, 0, 1, 0], 32 | [0, 1, 0, 0], 33 | [0, 0, 0, 1]]) 34 | for idx, r_path in tqdm(enumerate(regions)): 35 | mesh = trimesh.load(r_path) 36 | z_max = max(mesh.vertices[:, 2]) 37 | z_range = max(mesh.vertices[:, 2]) - min(mesh.vertices[:, 2]) 38 | x_min = min(mesh.vertices[:, 0]) 39 | y_min = min(mesh.vertices[:, 1]) 40 | # For better visualization, cut the ceilings and parts of walls 41 | if cut_mesh: 42 | mesh = trimesh.intersections.slice_mesh_plane(mesh, np.array([0, 0, -1]), np.array([0, 0, z_max - 0.5*z_range])) 43 | # mesh = trimesh.intersections.slice_mesh_plane(mesh, np.array([0, 1, 0]), np.array([0, y_min + 0.5, 0])) 44 | mesh = trimesh.intersections.slice_mesh_plane(mesh, np.array([1, 0, 0]), np.array([x_min + 0.2, 0, 0])) 45 | mesh = deepcopy(mesh) 46 | mesh.apply_transform(mat_permute) 47 | if save_part_mesh == True: 48 | out_file = join(outfile, 'mesh_fused%d.ply'%idx) 49 | mesh.export(out_file) 50 | 51 | if idx == 0: 52 | faces = mesh.faces 53 | vertices = mesh.vertices 54 | else: 55 | faces = np.concatenate([faces, mesh.faces + vertices.shape[0]]) 56 | vertices = np.concatenate([vertices, mesh.vertices]) 57 | 58 | mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False) 59 | out_file = join(outfile, 'mesh_fused.ply') 60 | mesh.export(out_file) 61 | 62 | # Sample surface points 63 | pcl, face_idx = mesh.sample(n_pointcloud_points, return_index=True) 64 | normals = mesh.face_normals[face_idx] 65 | 66 | # save surface points 67 | out_file = join(outfile, 'pointcloud.npz') 68 | np.savez(out_file, points=pcl.astype(dtype), normals=normals.astype(dtype)) 69 | export_pointcloud(pcl, join(outfile, 'pointcloud.ply')) 70 | 71 | # create test.lst file 72 | with open(join(out_path, 'test.lst'), "w") as file: 73 | file.write(scene_name) -------------------------------------------------------------------------------- /scripts/dataset_scannet/build_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import numpy as np 4 | import argparse 5 | import trimesh 6 | from SensorData import SensorData 7 | import tqdm 8 | from os.path import join 9 | from os import listdir 10 | import numpy as np 11 | import multiprocessing 12 | 13 | path_in = '/mnt/e/scans' # '/dir/to/scannet_v2' 14 | path_out = '/home/continuum_linux/alter_conv_onet/data' # '/dir/to/scannet_out' 15 | 16 | if not os.path.exists(path_out): 17 | os.makedirs(path_out) 18 | 19 | path_out = join(path_out, 'scenes') 20 | if not os.path.exists(path_out): 21 | os.makedirs(path_out) 22 | 23 | 24 | def align_axis(file_name, mesh): 25 | rotation_matrix = np.array([ 26 | [1., 0., 0., 0.], 27 | [0., 0., 1., 0.], 28 | [0., 1., 0., 0.], 29 | [0., 0., 0., 1.], 30 | ]) 31 | lines = open(file_name).readlines() 32 | for line in lines: 33 | if 'axisAlignment' in line: 34 | axis_align_matrix = [float(x) for x in line.rstrip().strip('axisAlignment = ').split(' ')] 35 | break 36 | axis_align_matrix = np.array(axis_align_matrix).reshape((4,4)) 37 | axis_align_matrix = rotation_matrix @ axis_align_matrix 38 | mesh.apply_transform(axis_align_matrix) 39 | return mesh, axis_align_matrix 40 | 41 | def sample_points(mesh, n_points=100000, p_type=np.float16): 42 | pcl, idx = mesh.sample(n_points, return_index=True) 43 | normals = mesh.face_normals[idx] 44 | out_dict = { 45 | 'points': pcl.astype(p_type), 46 | 'normals': normals.astype(p_type), 47 | 48 | } 49 | return out_dict 50 | 51 | def scale_to_unit_cube(mesh, y_level=-0.5): 52 | bbox = mesh.bounds 53 | loc = (bbox[0] + bbox[1]) / 2 54 | scale = 1. / (bbox[1] - bbox[0]).max() 55 | vertices_t = (mesh.vertices - loc.reshape(-1, 3)) * scale 56 | y_min = min(vertices_t[:, 1]) 57 | 58 | # create_transform_matrix 59 | S_loc = np.eye(4) 60 | S_loc[:-1, -1] = -loc 61 | # create scale mat 62 | S_scale = np.eye(4) * scale 63 | S_scale[-1, -1] = 1 64 | # create last translate matrix 65 | S_loc2 = np.eye(4) 66 | S_loc2[1, -1] = -y_min + y_level 67 | 68 | S = S_loc2 @ S_scale @ S_loc 69 | mesh.apply_transform(S) 70 | 71 | return mesh, S 72 | 73 | 74 | def process(scene_name): 75 | out_path_cur = os.path.join(path_out, scene_name) 76 | if not os.path.exists(out_path_cur): 77 | os.makedirs(out_path_cur) 78 | 79 | # load mesh 80 | mesh = trimesh.load(os.path.join(path_in, scene_name, scene_name+'_vh_clean_2.ply'), process=False) 81 | # mesh = trimesh.load(trimesh.util.wrap_as_stream(os.path.join(path_in, scene_name, scene_name+'_vh_clean.ply')), file_type='ply') 82 | txt_file = os.path.join(path_in, scene_name, '%s.txt' % scene_name) 83 | mesh, align_mat = align_axis(txt_file, mesh) 84 | mesh, scale_mat = scale_to_unit_cube(mesh) 85 | scale_matrix = np.linalg.inv(scale_mat @ align_mat) 86 | 87 | # file_cur = os.path.join(path_in, scene_name, scene_name+'.sens') 88 | # sd = SensorData(file_cur) 89 | # sd.export_depth_images(os.path.join(path_out, scene_name, 'depth'), frame_skip=1) 90 | # sd.process_camera_dict(join(path_out, scene_name), scale_matrix) 91 | pcl = sample_points(mesh) 92 | out_file = join(path_out, scene_name, 'pointcloud.npz') 93 | np.savez(out_file, **pcl) 94 | 95 | file_list = listdir(path_in) 96 | file_list.sort() 97 | pbar = tqdm.tqdm() 98 | # pool = multiprocessing.Pool(processes=1) 99 | for f in file_list: 100 | print(f) 101 | process(f) 102 | # pool.apply_async(process, args=(f,), callback=lambda _: pbar.update()) 103 | # pool.close() 104 | # pool.join() 105 | -------------------------------------------------------------------------------- /scripts/download_data.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | mkdir -p data 3 | cd data 4 | echo "Start downloading ..." 5 | wget https://s3.eu-central-1.amazonaws.com/avg-projects/convolutional_occupancy_networks/data/synthetic_room_dataset.zip 6 | unzip synthetic_room_dataset.zip 7 | echo "Done!" -------------------------------------------------------------------------------- /scripts/download_demo_data.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | mkdir -p data 3 | cd data 4 | echo "Downloading demo data..." 5 | wget https://s3.eu-central-1.amazonaws.com/avg-projects/convolutional_occupancy_networks/data/demo_data.zip 6 | unzip demo_data.zip 7 | echo "Done!" -------------------------------------------------------------------------------- /scripts/download_shapenet_data.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | mkdir data 3 | cd data 4 | wget https://s3.eu-central-1.amazonaws.com/avg-projects/occupancy_networks/data/dataset_small_v1.1.zip 5 | unzip dataset_small_v1.1.zip 6 | 7 | if [ ! -f "ShapeNet/metadata.yaml" ]; then 8 | cp metadata.yaml ShapeNet/metadata.yaml 9 | fi -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | try: 2 | from setuptools import setup 3 | except ImportError: 4 | from distutils.core import setup 5 | from distutils.extension import Extension 6 | from Cython.Build import cythonize 7 | # from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension 8 | import numpy 9 | 10 | 11 | # Get the numpy include directory. 12 | numpy_include_dir = numpy.get_include() 13 | 14 | # Extensions 15 | # pykdtree (kd tree) 16 | pykdtree = Extension( 17 | 'src.utils.libkdtree.pykdtree.kdtree', 18 | sources=[ 19 | 'src/utils/libkdtree/pykdtree/kdtree.c', 20 | 'src/utils/libkdtree/pykdtree/_kdtree_core.c' 21 | ], 22 | language='c', 23 | extra_compile_args=['-std=c99', '-O3', '-fopenmp'], 24 | extra_link_args=['-lgomp'], 25 | include_dirs=[numpy_include_dir] 26 | ) 27 | 28 | # mcubes (marching cubes algorithm) 29 | mcubes_module = Extension( 30 | 'src.utils.libmcubes.mcubes', 31 | sources=[ 32 | 'src/utils/libmcubes/mcubes.pyx', 33 | 'src/utils/libmcubes/pywrapper.cpp', 34 | 'src/utils/libmcubes/marchingcubes.cpp' 35 | ], 36 | language='c++', 37 | extra_compile_args=['-std=c++11'], 38 | include_dirs=[numpy_include_dir] 39 | ) 40 | 41 | # triangle hash (efficient mesh intersection) 42 | triangle_hash_module = Extension( 43 | 'src.utils.libmesh.triangle_hash', 44 | sources=[ 45 | 'src/utils/libmesh/triangle_hash.pyx' 46 | ], 47 | libraries=['m'], # Unix-like specific 48 | include_dirs=[numpy_include_dir] 49 | ) 50 | 51 | # mise (efficient mesh extraction) 52 | mise_module = Extension( 53 | 'src.utils.libmise.mise', 54 | sources=[ 55 | 'src/utils/libmise/mise.pyx' 56 | ], 57 | ) 58 | 59 | # simplify (efficient mesh simplification) 60 | simplify_mesh_module = Extension( 61 | 'src.utils.libsimplify.simplify_mesh', 62 | sources=[ 63 | 'src/utils/libsimplify/simplify_mesh.pyx' 64 | ], 65 | include_dirs=[numpy_include_dir] 66 | ) 67 | 68 | # voxelization (efficient mesh voxelization) 69 | voxelize_module = Extension( 70 | 'src.utils.libvoxelize.voxelize', 71 | sources=[ 72 | 'src/utils/libvoxelize/voxelize.pyx' 73 | ], 74 | libraries=['m'] # Unix-like specific 75 | ) 76 | 77 | # Gather all extension modules 78 | ext_modules = [ 79 | pykdtree, 80 | mcubes_module, 81 | triangle_hash_module, 82 | mise_module, 83 | simplify_mesh_module, 84 | voxelize_module, 85 | ] 86 | 87 | setup( 88 | ext_modules=cythonize(ext_modules), 89 | # cmdclass={ 90 | # 'build_ext': BuildExtension 91 | # } 92 | ) 93 | -------------------------------------------------------------------------------- /src/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/__init__.py -------------------------------------------------------------------------------- /src/attention.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | import math 4 | #from torch_cluster import knn as tc_knn 5 | 6 | class DotProductAttention(nn.Module): 7 | def __init__(self): 8 | super().__init__() 9 | ''' 10 | q: (bs, num_p, 1, fea_dim) 11 | k: (bs, num_p, 9, fea_dim) 12 | v: (bs, num_p, 9, fea_dim) 13 | ''' 14 | def forward(self, q, k, v): 15 | ''' 16 | q: (bs, num_p, 1, fea_dim) 17 | k: (bs, num_p, 9, fea_dim) 18 | v: (bs, num_p, 9, fea_dim) 19 | ''' 20 | d = q.shape[-1] 21 | scores = torch.einsum('ijkl,ijlm->ijkm', q, k.transpose(2, 3)) / math.sqrt(d) 22 | attention_weights = nn.functional.softmax(scores, dim=-1) 23 | output = torch.einsum('ijkm,ijmn->ijkn', attention_weights, v) 24 | return output 25 | 26 | class SubAttention(nn.Module): 27 | def __init__(self, hidden_size, num_heads): 28 | super().__init__() 29 | self.attention_mlp = nn.Sequential( 30 | nn.Linear(hidden_size, hidden_size * 4), 31 | nn.ReLU(), 32 | nn.Linear(hidden_size * 4, hidden_size), 33 | ) 34 | self.num_heads = num_heads 35 | 36 | def transpose_qkv(self, X): 37 | # input X (batch_size, no. ,num_neighbor, num_hiddens) 38 | # output (batch_size*num_heads, no. , num_nei, num_hiddens/num_heads) 39 | bs = X.shape[0] 40 | num_p = X.shape[1] 41 | num_neighbor = X.shape[2] 42 | 43 | X = X.reshape(bs, num_p, num_neighbor, self.num_heads, -1) # (batch_size, no. , num_nei, num_heads, num_hiddens / num_heads) 44 | X = X.permute(0, 3, 1, 2, 4) # (bs, num_heads, num_p, num_nei, num_hiddens/num_heads) 45 | output = X.reshape(-1, X.shape[2], X.shape[3], X.shape[4]) # (batch_size*num_heads, no. , num_nei, num_hiddens/num_heads) 46 | return output 47 | 48 | def transpose_output(self, X): 49 | ''' 50 | input: (batch_size * num_heads, no. of queries, num_hiddens / num_heads) 51 | output: (batch_size, no. of queries, num_hiddens) 52 | ''' 53 | X = X.reshape(-1, self.num_heads, X.shape[1], X.shape[2]) 54 | X = X.permute(0, 2, 1, 3) 55 | return X.reshape(X.shape[0], X.shape[1], -1) 56 | 57 | def forward(self, q, k, v, pos_encoding): 58 | ''' 59 | q: (bs, num_p, 1, fea_dim) 60 | k: (bs, num_p, 9, fea_dim) 61 | v: (bs, num_p, 9, fea_dim) 62 | ''' 63 | 64 | q = self.transpose_qkv(q) 65 | k = self.transpose_qkv(k) 66 | v = self.transpose_qkv(v) 67 | pos_encoding = self.transpose_qkv(pos_encoding) 68 | 69 | scores = self.attention_mlp( q - k + pos_encoding) # (bs,256,9,32) 70 | attention_weights = nn.functional.softmax(scores, dim=-2) 71 | v = v + pos_encoding # (bs, 256, 9, 32) 72 | output = torch.einsum('bijd,bijd->bid', attention_weights, v) # (bs, 256, 32) 73 | 74 | output = self.transpose_output(output) 75 | 76 | return output 77 | 78 | -------------------------------------------------------------------------------- /src/checkpoints.py: -------------------------------------------------------------------------------- 1 | import os 2 | import urllib 3 | import torch 4 | from torch.utils import model_zoo 5 | 6 | 7 | class CheckpointIO(object): 8 | ''' CheckpointIO class. 9 | 10 | It handles saving and loading checkpoints. 11 | 12 | Args: 13 | checkpoint_dir (str): path where checkpoints are saved 14 | ''' 15 | def __init__(self, checkpoint_dir='./chkpts', **kwargs): 16 | self.module_dict = kwargs 17 | self.checkpoint_dir = checkpoint_dir 18 | if not os.path.exists(checkpoint_dir): 19 | os.makedirs(checkpoint_dir) 20 | 21 | def register_modules(self, **kwargs): 22 | ''' Registers modules in current module dictionary. 23 | ''' 24 | self.module_dict.update(kwargs) 25 | 26 | def save(self, filename, **kwargs): 27 | ''' Saves the current module dictionary. 28 | 29 | Args: 30 | filename (str): name of output file 31 | ''' 32 | if not os.path.isabs(filename): 33 | filename = os.path.join(self.checkpoint_dir, filename) 34 | 35 | outdict = kwargs 36 | for k, v in self.module_dict.items(): 37 | outdict[k] = v.state_dict() 38 | torch.save(outdict, filename) 39 | 40 | def load(self, filename): 41 | '''Loads a module dictionary from local file or url. 42 | 43 | Args: 44 | filename (str): name of saved module dictionary 45 | ''' 46 | if is_url(filename): 47 | return self.load_url(filename) 48 | else: 49 | return self.load_file(filename) 50 | 51 | def load_file(self, filename): 52 | '''Loads a module dictionary from file. 53 | 54 | Args: 55 | filename (str): name of saved module dictionary 56 | ''' 57 | 58 | if not os.path.isabs(filename): 59 | filename = os.path.join(self.checkpoint_dir, filename) 60 | 61 | if os.path.exists(filename): 62 | print(filename) 63 | print('=> Loading checkpoint from local file...') 64 | state_dict = torch.load(filename, map_location='cuda:0') ### for generate.py 65 | scalars = self.parse_state_dict(state_dict) 66 | return scalars 67 | else: 68 | raise FileExistsError 69 | 70 | def load_url(self, url): 71 | '''Load a module dictionary from url. 72 | 73 | Args: 74 | url (str): url to saved model 75 | ''' 76 | print(url) 77 | print('=> Loading checkpoint from url...') 78 | state_dict = model_zoo.load_url(url, progress=True) 79 | scalars = self.parse_state_dict(state_dict) 80 | return scalars 81 | 82 | def parse_state_dict(self, state_dict): 83 | '''Parse state_dict of model and return scalars. 84 | 85 | Args: 86 | state_dict (dict): State dict of model 87 | ''' 88 | 89 | for k, v in self.module_dict.items(): 90 | if k in state_dict: 91 | v.load_state_dict(state_dict[k]) 92 | else: 93 | print('Warning: Could not find %s in checkpoint!' % k) 94 | scalars = {k: v for k, v in state_dict.items() 95 | if k not in self.module_dict} 96 | return scalars 97 | 98 | def is_url(url): 99 | scheme = urllib.parse.urlparse(url).scheme 100 | return scheme in ('http', 'https') -------------------------------------------------------------------------------- /src/config.py: -------------------------------------------------------------------------------- 1 | import yaml 2 | from torchvision import transforms 3 | from src import data 4 | from src import conv_onet 5 | 6 | 7 | method_dict = { 8 | 'conv_onet': conv_onet 9 | } 10 | 11 | 12 | # General config 13 | def load_config(path, default_path=None): 14 | ''' Loads config file. 15 | 16 | Args: 17 | path (str): path to config file 18 | default_path (bool): whether to use default path 19 | ''' 20 | # Load configuration from file itself 21 | with open(path, 'r') as f: 22 | #cfg_special = yaml.load(f) ### 23 | cfg_special = yaml.safe_load(f) 24 | 25 | # Check if we should inherit from a config 26 | inherit_from = cfg_special.get('inherit_from') 27 | 28 | # If yes, load this config first as default 29 | # If no, use the default_path 30 | if inherit_from is not None: 31 | cfg = load_config(inherit_from, default_path) 32 | elif default_path is not None: 33 | with open(default_path, 'r') as f: 34 | #cfg = yaml.load(f) 35 | cfg = yaml.safe_load(f) ### 36 | else: 37 | cfg = dict() 38 | 39 | # Include main configuration 40 | update_recursive(cfg, cfg_special) 41 | 42 | return cfg 43 | 44 | 45 | def update_recursive(dict1, dict2): 46 | ''' Update two config dictionaries recursively. 47 | 48 | Args: 49 | dict1 (dict): first dictionary to be updated 50 | dict2 (dict): second dictionary which entries should be used 51 | 52 | ''' 53 | for k, v in dict2.items(): 54 | if k not in dict1: 55 | dict1[k] = dict() 56 | if isinstance(v, dict): 57 | update_recursive(dict1[k], v) 58 | else: 59 | dict1[k] = v 60 | 61 | 62 | # Models 63 | def get_model(cfg, device=None, dataset=None): 64 | ''' Returns the model instance. 65 | 66 | Args: 67 | cfg (dict): config dictionary 68 | device (device): pytorch device 69 | dataset (dataset): dataset 70 | ''' 71 | method = cfg['method'] 72 | model = method_dict[method].config.get_model( 73 | cfg, device=device, dataset=dataset) 74 | return model 75 | 76 | 77 | # Trainer 78 | def get_trainer(model, optimizer, scheduler, cfg, device): 79 | ''' Returns a trainer instance. 80 | 81 | Args: 82 | model (nn.Module): the model which is used 83 | optimizer (optimizer): pytorch optimizer 84 | cfg (dict): config dictionary 85 | device (device): pytorch device 86 | ''' 87 | method = cfg['method'] 88 | trainer = method_dict[method].config.get_trainer( 89 | model, optimizer, scheduler, cfg, device) 90 | return trainer 91 | 92 | 93 | # Generator for final mesh extraction 94 | def get_generator(model, cfg, device): 95 | ''' Returns a generator instance. 96 | 97 | Args: 98 | model (nn.Module): the model which is used 99 | cfg (dict): config dictionary 100 | device (device): pytorch device 101 | ''' 102 | method = cfg['method'] 103 | generator = method_dict[method].config.get_generator(model, cfg, device) 104 | return generator 105 | 106 | 107 | # Datasets 108 | def get_dataset(mode, cfg, return_idx=False): 109 | ''' Returns the dataset. 110 | 111 | Args: 112 | model (nn.Module): the model which is used 113 | cfg (dict): config dictionary 114 | return_idx (bool): whether to include an ID field 115 | ''' 116 | method = cfg['method'] 117 | dataset_type = cfg['data']['dataset'] 118 | dataset_folder = cfg['data']['path'] 119 | categories = cfg['data']['classes'] 120 | 121 | # Get split 122 | splits = { 123 | 'train': cfg['data']['train_split'], 124 | 'val': cfg['data']['val_split'], 125 | 'test': cfg['data']['test_split'], 126 | } 127 | 128 | split = splits[mode] 129 | 130 | # Create dataset 131 | if dataset_type == 'Shapes3D': 132 | # Dataset fields 133 | # Method specific fields (usually correspond to output) 134 | fields = method_dict[method].config.get_data_fields(mode, cfg) 135 | # Input fields 136 | inputs_field = get_inputs_field(mode, cfg) 137 | if inputs_field is not None: 138 | fields['inputs'] = inputs_field 139 | 140 | if return_idx: 141 | fields['idx'] = data.IndexField() 142 | 143 | dataset = data.Shapes3dDataset( 144 | dataset_folder, fields, 145 | split=split, 146 | categories=categories, 147 | cfg = cfg 148 | ) 149 | else: 150 | raise ValueError('Invalid dataset "%s"' % cfg['data']['dataset']) 151 | 152 | return dataset 153 | 154 | 155 | def get_inputs_field(mode, cfg): 156 | ''' Returns the inputs fields. 157 | 158 | Args: 159 | mode (str): the mode which is used 160 | cfg (dict): config dictionary 161 | ''' 162 | input_type = cfg['data']['input_type'] 163 | 164 | if input_type is None: 165 | inputs_field = None 166 | elif input_type == 'pointcloud': 167 | transform = transforms.Compose([ 168 | data.SubsamplePointcloud(cfg['data']['pointcloud_n']), 169 | data.PointcloudNoise(cfg['data']['pointcloud_noise']) 170 | ]) 171 | inputs_field = data.PointCloudField( 172 | cfg['data']['pointcloud_file'], transform, 173 | multi_files= cfg['data']['multi_files'] 174 | ) 175 | elif input_type == 'partial_pointcloud': 176 | transform = transforms.Compose([ 177 | data.SubsamplePointcloud(cfg['data']['pointcloud_n']), 178 | data.PointcloudNoise(cfg['data']['pointcloud_noise']) 179 | ]) 180 | inputs_field = data.PartialPointCloudField( 181 | cfg['data']['pointcloud_file'], transform, 182 | multi_files= cfg['data']['multi_files'] 183 | ) 184 | elif input_type == 'pointcloud_crop': 185 | transform = transforms.Compose([ 186 | data.SubsamplePointcloud(cfg['data']['pointcloud_n']), 187 | data.PointcloudNoise(cfg['data']['pointcloud_noise']) 188 | ]) 189 | 190 | inputs_field = data.PatchPointCloudField( 191 | cfg['data']['pointcloud_file'], 192 | transform, 193 | multi_files= cfg['data']['multi_files'], 194 | ) 195 | 196 | elif input_type == 'voxels': 197 | inputs_field = data.VoxelsField( 198 | cfg['data']['voxels_file'] 199 | ) 200 | elif input_type == 'idx': 201 | inputs_field = data.IndexField() 202 | else: 203 | raise ValueError( 204 | 'Invalid input type (%s)' % input_type) 205 | return inputs_field -------------------------------------------------------------------------------- /src/conv_onet/__init__.py: -------------------------------------------------------------------------------- 1 | from src.conv_onet import ( 2 | config, generation, training, models 3 | ) 4 | 5 | __all__ = [ 6 | config, generation, training, models 7 | ] 8 | -------------------------------------------------------------------------------- /src/conv_onet/models/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch import distributions as dist 4 | from src.conv_onet.models import decoder 5 | 6 | 7 | # Decoder dictionary 8 | decoder_dict = { 9 | 'simple_local': decoder.LocalDecoder, 10 | 'simple_local_attention_sub': decoder.LocalDecoder_attention_sub 11 | } 12 | 13 | class ConvolutionalOccupancyNetwork(nn.Module): 14 | ''' Occupancy Network class. 15 | 16 | Args: 17 | decoder (nn.Module): decoder network 18 | encoder (nn.Module): encoder network 19 | device (device): torch device 20 | ''' 21 | 22 | def __init__(self, decoder, encoder=None, device=None): 23 | super().__init__() 24 | 25 | self.decoder = decoder.to(device) 26 | 27 | if encoder is not None: 28 | self.encoder = encoder.to(device) 29 | else: 30 | self.encoder = None 31 | 32 | self._device = device 33 | 34 | def forward(self, p, inputs, sample=True, **kwargs): 35 | ''' Performs a forward pass through the network. 36 | 37 | Args: 38 | p (tensor): sampled points 39 | inputs (tensor): conditioning input 40 | sample (bool): whether to sample for z 41 | ''' 42 | if isinstance(p, dict): 43 | batch_size = p['p'].size(0) 44 | else: 45 | batch_size = p.size(0) 46 | 47 | c_plane = self.encoder(inputs) 48 | p_r = self.decode(p, c_plane) 49 | 50 | 51 | return p_r 52 | 53 | def encode_inputs(self, inputs): 54 | ''' Encodes the input. 55 | 56 | Args: 57 | input (tensor): the input 58 | ''' 59 | 60 | if self.encoder is not None: 61 | c = self.encoder(inputs) 62 | else: 63 | # Return inputs? 64 | c = torch.empty(inputs.size(0), 0) 65 | 66 | return c 67 | 68 | def decode(self, p, c, loop_num=None, point_feature=None): 69 | ''' Returns occupancy probabilities for the sampled points. 70 | 71 | Args: 72 | p (tensor): points 73 | c (tensor): latent conditioned code c 74 | ''' 75 | 76 | # out = self.decoder(p, c, loop_num, point_feature) 77 | out = self.decoder(p, c) 78 | 79 | return out 80 | 81 | def to(self, device): 82 | ''' Puts the model to the device. 83 | 84 | Args: 85 | device (device): pytorch device 86 | ''' 87 | model = super().to(device) 88 | model._device = device 89 | return model -------------------------------------------------------------------------------- /src/conv_onet/models/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/conv_onet/models/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /src/conv_onet/models/__pycache__/decoder.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/conv_onet/models/__pycache__/decoder.cpython-36.pyc -------------------------------------------------------------------------------- /src/conv_onet/training.py: -------------------------------------------------------------------------------- 1 | import os 2 | from tqdm import trange 3 | import torch 4 | from torch.nn import functional as F 5 | from torch import distributions as dist 6 | from src.common import ( 7 | compute_iou, make_3d_grid, add_key, 8 | ) 9 | from src.utils import visualize as vis 10 | from src.training import BaseTrainer 11 | import numpy as np 12 | 13 | class Trainer(BaseTrainer): 14 | ''' Trainer object for the Occupancy Network. 15 | 16 | Args: 17 | model (nn.Module): Occupancy Network model 18 | optimizer (optimizer): pytorch optimizer object 19 | device (device): pytorch device 20 | input_type (str): input type 21 | vis_dir (str): visualization directory 22 | threshold (float): threshold value 23 | eval_sample (bool): whether to evaluate samples 24 | 25 | ''' 26 | 27 | def __init__(self, model, optimizer, scheduler, device=None, input_type='pointcloud', 28 | vis_dir=None, threshold=0.5, eval_sample=False): 29 | self.model = model 30 | self.optimizer = optimizer 31 | self.scheduler = scheduler 32 | self.device = device 33 | self.input_type = input_type 34 | self.vis_dir = vis_dir 35 | self.threshold = threshold 36 | self.eval_sample = eval_sample 37 | 38 | if vis_dir is not None and not os.path.exists(vis_dir): 39 | os.makedirs(vis_dir) 40 | 41 | def train_step(self, data): 42 | ''' Performs a training step. 43 | 44 | Args: 45 | data (dict): data dictionary 46 | ''' 47 | self.model.train() 48 | self.optimizer.zero_grad() 49 | loss = self.compute_loss(data) 50 | loss.backward() 51 | self.optimizer.step() 52 | 53 | return loss.item() 54 | 55 | def schedule(self): 56 | self.scheduler.step() 57 | 58 | def eval_step(self, data): 59 | ''' Performs an evaluation step. 60 | 61 | Args: 62 | data (dict): data dictionary 63 | ''' 64 | self.model.eval() 65 | 66 | device = self.device 67 | threshold = self.threshold 68 | eval_dict = {} 69 | 70 | points = data.get('points').to(device) 71 | occ = data.get('points.occ').to(device) 72 | 73 | inputs = data.get('inputs', torch.empty(points.size(0), 0)).to(device) 74 | voxels_occ = data.get('voxels') 75 | 76 | points_iou = data.get('points_iou').to(device) 77 | occ_iou = data.get('points_iou.occ').to(device) 78 | 79 | batch_size = points.size(0) 80 | 81 | kwargs = {} 82 | 83 | # add pre-computed index 84 | inputs = add_key(inputs, data.get('inputs.ind'), 'points', 'index', device=device) 85 | # add pre-computed normalized coordinates 86 | points = add_key(points, data.get('points.normalized'), 'p', 'p_n', device=device) 87 | points_iou = add_key(points_iou, data.get('points_iou.normalized'), 'p', 'p_n', device=device) 88 | 89 | # Compute iou 90 | with torch.no_grad(): 91 | p_out = self.model.module(points_iou, inputs, sample=self.eval_sample, **kwargs) 92 | 93 | occ_iou_np = (occ_iou >= 0.5).cpu().numpy() 94 | occ_iou_hat_np = (p_out.probs >= threshold).cpu().numpy() 95 | 96 | iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean() 97 | eval_dict['iou'] = iou 98 | 99 | # Estimate voxel iou 100 | if voxels_occ is not None: 101 | voxels_occ = voxels_occ.to(device) 102 | points_voxels = make_3d_grid( 103 | (-0.5 + 1/64,) * 3, (0.5 - 1/64,) * 3, voxels_occ.shape[1:]) 104 | points_voxels = points_voxels.expand( 105 | batch_size, *points_voxels.size()) 106 | points_voxels = points_voxels.to(device) 107 | with torch.no_grad(): 108 | p_out = self.model(points_voxels, inputs, 109 | sample=self.eval_sample, **kwargs) 110 | 111 | voxels_occ_np = (voxels_occ >= 0.5).cpu().numpy() 112 | occ_hat_np = (p_out.probs >= threshold).cpu().numpy() 113 | iou_voxels = compute_iou(voxels_occ_np, occ_hat_np).mean() 114 | 115 | eval_dict['iou_voxels'] = iou_voxels 116 | 117 | return eval_dict 118 | 119 | def compute_loss(self, data): 120 | ''' Computes the loss. 121 | 122 | Args: 123 | data (dict): data dictionary 124 | ''' 125 | device = self.device 126 | p = data.get('points').to(device) 127 | occ = data.get('points.occ').to(device) 128 | inputs = data.get('inputs', torch.empty(p.size(0), 0)).to(device) 129 | 130 | if 'pointcloud_crop' in data.keys(): 131 | # add pre-computed index 132 | inputs = add_key(inputs, data.get('inputs.ind'), 'points', 'index', device=device) 133 | inputs['mask'] = data.get('inputs.mask').to(device) 134 | # add pre-computed normalized coordinates 135 | p = add_key(p, data.get('points.normalized'), 'p', 'p_n', device=device) 136 | 137 | c_plane = self.model.module.encoder(inputs) 138 | p_r = self.model.module.decode(p, c_plane) 139 | 140 | 141 | kwargs = {} 142 | # General points 143 | # logits = self.model.decode(p_r, c, **kwargs).logits 144 | logits = p_r.logits 145 | loss_i = F.binary_cross_entropy_with_logits( 146 | logits, occ, reduction='none') 147 | loss = loss_i.sum(-1).mean() 148 | 149 | return loss 150 | -------------------------------------------------------------------------------- /src/data/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | from src.data.core import ( 3 | Shapes3dDataset, collate_remove_none, worker_init_fn 4 | ) 5 | from src.data.fields import ( 6 | IndexField, PointsField, 7 | VoxelsField, PatchPointsField, PointCloudField, PatchPointCloudField, PartialPointCloudField, 8 | ) 9 | from src.data.transforms import ( 10 | PointcloudNoise, SubsamplePointcloud, 11 | SubsamplePoints, 12 | ) 13 | __all__ = [ 14 | # Core 15 | Shapes3dDataset, 16 | collate_remove_none, 17 | worker_init_fn, 18 | # Fields 19 | IndexField, 20 | PointsField, 21 | VoxelsField, 22 | PointCloudField, 23 | PartialPointCloudField, 24 | PatchPointCloudField, 25 | PatchPointsField, 26 | # Transforms 27 | PointcloudNoise, 28 | SubsamplePointcloud, 29 | SubsamplePoints, 30 | ] 31 | -------------------------------------------------------------------------------- /src/data/transforms.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | # Transforms 5 | class PointcloudNoise(object): 6 | ''' Point cloud noise transformation class. 7 | 8 | It adds noise to point cloud data. 9 | 10 | Args: 11 | stddev (int): standard deviation 12 | ''' 13 | 14 | def __init__(self, stddev): 15 | self.stddev = stddev 16 | 17 | def __call__(self, data): 18 | ''' Calls the transformation. 19 | 20 | Args: 21 | data (dictionary): data dictionary 22 | ''' 23 | data_out = data.copy() 24 | points = data[None] 25 | noise = self.stddev * np.random.randn(*points.shape) 26 | noise = noise.astype(np.float32) 27 | data_out[None] = points + noise 28 | return data_out 29 | 30 | class SubsamplePointcloud(object): 31 | ''' Point cloud subsampling transformation class. 32 | 33 | It subsamples the point cloud data. 34 | 35 | Args: 36 | N (int): number of points to be subsampled 37 | ''' 38 | def __init__(self, N): 39 | self.N = N 40 | # self.N_max = N_max 41 | 42 | def __call__(self, data): 43 | ''' Calls the transformation. 44 | 45 | Args: 46 | data (dict): data dictionary 47 | ''' 48 | data_out = data.copy() 49 | points = data[None] 50 | normals = data['normals'] 51 | 52 | # num_samples = np.random.randint(self.N, self.N_max) 53 | # print('points.shape', points.shape) 54 | 55 | indices = np.random.randint(points.shape[0], size=self.N) 56 | 57 | # indices = np.random.randint(points.shape[0], size=num_samples) 58 | data_out[None] = points[indices, :] 59 | # print('=======', data_out[None].shape) 60 | data_out['normals'] = normals[indices, :] 61 | 62 | # zw 63 | # data_out['sub_points_idx'] = indices 64 | 65 | return data_out 66 | 67 | 68 | class SubsamplePoints(object): 69 | ''' Points subsampling transformation class. 70 | 71 | It subsamples the points data. 72 | 73 | Args: 74 | N (int): number of points to be subsampled 75 | ''' 76 | def __init__(self, N): 77 | self.N = N # 2048 78 | 79 | def __call__(self, data): 80 | ''' Calls the transformation. 81 | 82 | Args: 83 | data (dictionary): data dictionary 84 | ''' 85 | points = data[None] 86 | occ = data['occ'] 87 | # print('points.shape', points.shape) 88 | data_out = data.copy() 89 | if isinstance(self.N, int): 90 | idx = np.random.randint(points.shape[0], size=self.N) 91 | # print(idx) 92 | # print('idx.shape', idx.shape) 93 | data_out.update({ 94 | None: points[idx, :], 95 | 'occ': occ[idx], 96 | 'sub_points_idx': idx, 97 | }) 98 | else: 99 | Nt_out, Nt_in = self.N 100 | occ_binary = (occ >= 0.5) 101 | points0 = points[~occ_binary] 102 | points1 = points[occ_binary] 103 | 104 | idx0 = np.random.randint(points0.shape[0], size=Nt_out) 105 | idx1 = np.random.randint(points1.shape[0], size=Nt_in) 106 | 107 | points0 = points0[idx0, :] 108 | points1 = points1[idx1, :] 109 | points = np.concatenate([points0, points1], axis=0) 110 | 111 | occ0 = np.zeros(Nt_out, dtype=np.float32) 112 | occ1 = np.ones(Nt_in, dtype=np.float32) 113 | occ = np.concatenate([occ0, occ1], axis=0) 114 | 115 | volume = occ_binary.sum() / len(occ_binary) 116 | volume = volume.astype(np.float32) 117 | 118 | data_out.update({ 119 | None: points, 120 | 'occ': occ, 121 | 'volume': volume, 122 | }) 123 | return data_out 124 | -------------------------------------------------------------------------------- /src/encoder/__init__.py: -------------------------------------------------------------------------------- 1 | from src.encoder import ( 2 | pointnet, voxels, pointnetpp 3 | ) 4 | 5 | 6 | encoder_dict = { 7 | 'pointnet_local_pool': pointnet.LocalPoolPointnet 8 | } 9 | -------------------------------------------------------------------------------- /src/encoder/voxels.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch_scatter import scatter_mean 5 | from src.encoder.unet import UNet 6 | from src.encoder.unet3d import UNet3D 7 | from src.common import coordinate2index, normalize_coordinate, normalize_3d_coordinate 8 | 9 | 10 | class LocalVoxelEncoder(nn.Module): 11 | ''' 3D-convolutional encoder network for voxel input. 12 | 13 | Args: 14 | dim (int): input dimension 15 | c_dim (int): dimension of latent code c 16 | hidden_dim (int): hidden dimension of the network 17 | unet (bool): weather to use U-Net 18 | unet_kwargs (str): U-Net parameters 19 | unet3d (bool): weather to use 3D U-Net 20 | unet3d_kwargs (str): 3D U-Net parameters 21 | plane_resolution (int): defined resolution for plane feature 22 | grid_resolution (int): defined resolution for grid feature 23 | plane_type (str): 'xz' - 1-plane, ['xz', 'xy', 'yz'] - 3-plane, ['grid'] - 3D grid volume 24 | kernel_size (int): kernel size for the first layer of CNN 25 | padding (float): conventional padding paramter of ONet for unit cube, so [-0.5, 0.5] -> [-0.55, 0.55] 26 | 27 | ''' 28 | 29 | def __init__(self, dim=3, c_dim=128, unet=False, unet_kwargs=None, unet3d=False, unet3d_kwargs=None, 30 | plane_resolution=512, grid_resolution=None, plane_type='xz', kernel_size=3, padding=0.1): 31 | super().__init__() 32 | self.actvn = F.relu 33 | if kernel_size == 1: 34 | self.conv_in = nn.Conv3d(1, c_dim, 1) 35 | else: 36 | self.conv_in = nn.Conv3d(1, c_dim, kernel_size, padding=1) 37 | 38 | if unet: 39 | self.unet = UNet(c_dim, in_channels=c_dim, **unet_kwargs) 40 | else: 41 | self.unet = None 42 | 43 | if unet3d: 44 | self.unet3d = UNet3D(**unet3d_kwargs) 45 | else: 46 | self.unet3d = None 47 | 48 | self.c_dim = c_dim 49 | 50 | self.reso_plane = plane_resolution 51 | self.reso_grid = grid_resolution 52 | 53 | self.plane_type = plane_type 54 | self.padding = padding 55 | 56 | def generate_plane_features(self, p, c, plane='xz'): 57 | # acquire indices of features in plane 58 | xy = normalize_coordinate(p.clone(), plane=plane, padding=self.padding) 59 | index = coordinate2index(xy, self.reso_plane) 60 | 61 | # scatter plane features from points 62 | fea_plane = c.new_zeros(p.size(0), self.c_dim, self.reso_plane**2) 63 | c = c.permute(0, 2, 1) 64 | fea_plane = scatter_mean(c, index, out=fea_plane) 65 | fea_plane = fea_plane.reshape(p.size(0), self.c_dim, self.reso_plane, self.reso_plane) 66 | 67 | # process the plane features with UNet 68 | if self.unet is not None: 69 | fea_plane = self.unet(fea_plane) 70 | 71 | return fea_plane 72 | 73 | def generate_grid_features(self, p, c): 74 | p_nor = normalize_3d_coordinate(p.clone(), padding=self.padding) 75 | index = coordinate2index(p_nor, self.reso_grid, coord_type='3d') 76 | # scatter grid features from points 77 | fea_grid = c.new_zeros(p.size(0), self.c_dim, self.reso_grid**3) 78 | c = c.permute(0, 2, 1) 79 | fea_grid = scatter_mean(c, index, out=fea_grid) 80 | fea_grid = fea_grid.reshape(p.size(0), self.c_dim, self.reso_grid, self.reso_grid, self.reso_grid) 81 | 82 | if self.unet3d is not None: 83 | fea_grid = self.unet3d(fea_grid) 84 | 85 | return fea_grid 86 | 87 | 88 | def forward(self, x): 89 | batch_size = x.size(0) 90 | device = x.device 91 | n_voxel = x.size(1) * x.size(2) * x.size(3) 92 | 93 | # voxel 3D coordintates 94 | coord1 = torch.linspace(-0.5, 0.5, x.size(1)).to(device) 95 | coord2 = torch.linspace(-0.5, 0.5, x.size(2)).to(device) 96 | coord3 = torch.linspace(-0.5, 0.5, x.size(3)).to(device) 97 | 98 | coord1 = coord1.view(1, -1, 1, 1).expand_as(x) 99 | coord2 = coord2.view(1, 1, -1, 1).expand_as(x) 100 | coord3 = coord3.view(1, 1, 1, -1).expand_as(x) 101 | p = torch.stack([coord1, coord2, coord3], dim=4) 102 | p = p.view(batch_size, n_voxel, -1) 103 | 104 | # Acquire voxel-wise feature 105 | x = x.unsqueeze(1) 106 | c = self.actvn(self.conv_in(x)).view(batch_size, self.c_dim, -1) 107 | c = c.permute(0, 2, 1) 108 | 109 | fea = {} 110 | if 'grid' in self.plane_type: 111 | fea['grid'] = self.generate_grid_features(p, c) 112 | else: 113 | if 'xz' in self.plane_type: 114 | fea['xz'] = self.generate_plane_features(p, c, plane='xz') 115 | if 'xy' in self.plane_type: 116 | fea['xy'] = self.generate_plane_features(p, c, plane='xy') 117 | if 'yz' in self.plane_type: 118 | fea['yz'] = self.generate_plane_features(p, c, plane='yz') 119 | return fea 120 | 121 | class VoxelEncoder(nn.Module): 122 | ''' 3D-convolutional encoder network for voxel input. 123 | 124 | Args: 125 | dim (int): input dimension 126 | c_dim (int): output dimension 127 | ''' 128 | 129 | def __init__(self, dim=3, c_dim=128): 130 | super().__init__() 131 | self.actvn = F.relu 132 | 133 | self.conv_in = nn.Conv3d(1, 32, 3, padding=1) 134 | 135 | self.conv_0 = nn.Conv3d(32, 64, 3, padding=1, stride=2) 136 | self.conv_1 = nn.Conv3d(64, 128, 3, padding=1, stride=2) 137 | self.conv_2 = nn.Conv3d(128, 256, 3, padding=1, stride=2) 138 | self.conv_3 = nn.Conv3d(256, 512, 3, padding=1, stride=2) 139 | self.fc = nn.Linear(512 * 2 * 2 * 2, c_dim) 140 | 141 | def forward(self, x): 142 | batch_size = x.size(0) 143 | 144 | x = x.unsqueeze(1) 145 | net = self.conv_in(x) 146 | net = self.conv_0(self.actvn(net)) 147 | net = self.conv_1(self.actvn(net)) 148 | net = self.conv_2(self.actvn(net)) 149 | net = self.conv_3(self.actvn(net)) 150 | 151 | hidden = net.view(batch_size, 512 * 2 * 2 * 2) 152 | c = self.fc(self.actvn(hidden)) 153 | 154 | return c -------------------------------------------------------------------------------- /src/layers.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | # Resnet Blocks 6 | class ResnetBlockFC(nn.Module): 7 | ''' Fully connected ResNet Block class. 8 | 9 | Args: 10 | size_in (int): input dimension 11 | size_out (int): output dimension 12 | size_h (int): hidden dimension 13 | ''' 14 | 15 | def __init__(self, size_in, size_out=None, size_h=None): 16 | super().__init__() 17 | # Attributes 18 | if size_out is None: 19 | size_out = size_in 20 | 21 | if size_h is None: 22 | size_h = min(size_in, size_out) 23 | 24 | self.size_in = size_in 25 | self.size_h = size_h 26 | self.size_out = size_out 27 | # Submodules 28 | self.fc_0 = nn.Linear(size_in, size_h) 29 | self.fc_1 = nn.Linear(size_h, size_out) 30 | self.actvn = nn.ReLU() 31 | 32 | if size_in == size_out: 33 | self.shortcut = None 34 | else: 35 | self.shortcut = nn.Linear(size_in, size_out, bias=False) 36 | # Initialization 37 | nn.init.zeros_(self.fc_1.weight) 38 | 39 | def forward(self, x): 40 | net = self.fc_0(self.actvn(x)) 41 | dx = self.fc_1(self.actvn(net)) 42 | 43 | if self.shortcut is not None: 44 | x_s = self.shortcut(x) 45 | else: 46 | x_s = x 47 | 48 | return x_s + dx -------------------------------------------------------------------------------- /src/training.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from collections import defaultdict 3 | from tqdm import tqdm 4 | 5 | 6 | class BaseTrainer(object): 7 | ''' Base trainer class. 8 | ''' 9 | 10 | def evaluate(self, val_loader): 11 | ''' Performs an evaluation. 12 | Args: 13 | val_loader (dataloader): pytorch dataloader 14 | ''' 15 | eval_list = defaultdict(list) 16 | 17 | for data in tqdm(val_loader): 18 | eval_step_dict = self.eval_step(data) 19 | 20 | for k, v in eval_step_dict.items(): 21 | eval_list[k].append(v) 22 | 23 | eval_dict = {k: np.mean(v) for k, v in eval_list.items()} 24 | return eval_dict 25 | 26 | def train_step(self, *args, **kwargs): 27 | ''' Performs a training step. 28 | ''' 29 | raise NotImplementedError 30 | 31 | def eval_step(self, *args, **kwargs): 32 | ''' Performs an evaluation step. 33 | ''' 34 | raise NotImplementedError 35 | 36 | def visualize(self, *args, **kwargs): 37 | ''' Performs visualization. 38 | ''' 39 | raise NotImplementedError 40 | -------------------------------------------------------------------------------- /src/utils/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/utils/.DS_Store -------------------------------------------------------------------------------- /src/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/utils/__init__.py -------------------------------------------------------------------------------- /src/utils/icp.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from sklearn.neighbors import NearestNeighbors 3 | 4 | 5 | def best_fit_transform(A, B): 6 | ''' 7 | Calculates the least-squares best-fit transform that maps corresponding 8 | points A to B in m spatial dimensions 9 | Input: 10 | A: Nxm numpy array of corresponding points 11 | B: Nxm numpy array of corresponding points 12 | Returns: 13 | T: (m+1)x(m+1) homogeneous transformation matrix that maps A on to B 14 | R: mxm rotation matrix 15 | t: mx1 translation vector 16 | ''' 17 | 18 | assert A.shape == B.shape 19 | 20 | # get number of dimensions 21 | m = A.shape[1] 22 | 23 | # translate points to their centroids 24 | centroid_A = np.mean(A, axis=0) 25 | centroid_B = np.mean(B, axis=0) 26 | AA = A - centroid_A 27 | BB = B - centroid_B 28 | 29 | # rotation matrix 30 | H = np.dot(AA.T, BB) 31 | U, S, Vt = np.linalg.svd(H) 32 | R = np.dot(Vt.T, U.T) 33 | 34 | # special reflection case 35 | if np.linalg.det(R) < 0: 36 | Vt[m-1,:] *= -1 37 | R = np.dot(Vt.T, U.T) 38 | 39 | # translation 40 | t = centroid_B.T - np.dot(R,centroid_A.T) 41 | 42 | # homogeneous transformation 43 | T = np.identity(m+1) 44 | T[:m, :m] = R 45 | T[:m, m] = t 46 | 47 | return T, R, t 48 | 49 | 50 | def nearest_neighbor(src, dst): 51 | ''' 52 | Find the nearest (Euclidean) neighbor in dst for each point in src 53 | Input: 54 | src: Nxm array of points 55 | dst: Nxm array of points 56 | Output: 57 | distances: Euclidean distances of the nearest neighbor 58 | indices: dst indices of the nearest neighbor 59 | ''' 60 | 61 | assert src.shape == dst.shape 62 | 63 | neigh = NearestNeighbors(n_neighbors=1) 64 | neigh.fit(dst) 65 | distances, indices = neigh.kneighbors(src, return_distance=True) 66 | return distances.ravel(), indices.ravel() 67 | 68 | 69 | def icp(A, B, init_pose=None, max_iterations=20, tolerance=0.001): 70 | ''' 71 | The Iterative Closest Point method: finds best-fit transform that maps 72 | points A on to points B 73 | Input: 74 | A: Nxm numpy array of source mD points 75 | B: Nxm numpy array of destination mD point 76 | init_pose: (m+1)x(m+1) homogeneous transformation 77 | max_iterations: exit algorithm after max_iterations 78 | tolerance: convergence criteria 79 | Output: 80 | T: final homogeneous transformation that maps A on to B 81 | distances: Euclidean distances (errors) of the nearest neighbor 82 | i: number of iterations to converge 83 | ''' 84 | 85 | assert A.shape == B.shape 86 | 87 | # get number of dimensions 88 | m = A.shape[1] 89 | 90 | # make points homogeneous, copy them to maintain the originals 91 | src = np.ones((m+1,A.shape[0])) 92 | dst = np.ones((m+1,B.shape[0])) 93 | src[:m,:] = np.copy(A.T) 94 | dst[:m,:] = np.copy(B.T) 95 | 96 | # apply the initial pose estimation 97 | if init_pose is not None: 98 | src = np.dot(init_pose, src) 99 | 100 | prev_error = 0 101 | 102 | for i in range(max_iterations): 103 | # find the nearest neighbors between the current source and destination points 104 | distances, indices = nearest_neighbor(src[:m,:].T, dst[:m,:].T) 105 | 106 | # compute the transformation between the current source and nearest destination points 107 | T,_,_ = best_fit_transform(src[:m,:].T, dst[:m,indices].T) 108 | 109 | # update the current source 110 | src = np.dot(T, src) 111 | 112 | # check error 113 | mean_error = np.mean(distances) 114 | if np.abs(prev_error - mean_error) < tolerance: 115 | break 116 | prev_error = mean_error 117 | 118 | # calculate final transformation 119 | T,_,_ = best_fit_transform(A, src[:m,:].T) 120 | 121 | return T, distances, i 122 | -------------------------------------------------------------------------------- /src/utils/io.py: -------------------------------------------------------------------------------- 1 | import os 2 | from plyfile import PlyElement, PlyData 3 | import numpy as np 4 | 5 | 6 | def export_pointcloud(vertices, out_file, as_text=True): 7 | assert(vertices.shape[1] == 3) 8 | vertices = vertices.astype(np.float32) 9 | vertices = np.ascontiguousarray(vertices) 10 | vector_dtype = [('x', 'f4'), ('y', 'f4'), ('z', 'f4')] 11 | vertices = vertices.view(dtype=vector_dtype).flatten() 12 | plyel = PlyElement.describe(vertices, 'vertex') 13 | plydata = PlyData([plyel], text=as_text) 14 | plydata.write(out_file) 15 | 16 | 17 | def load_pointcloud(in_file): 18 | plydata = PlyData.read(in_file) 19 | vertices = np.stack([ 20 | plydata['vertex']['x'], 21 | plydata['vertex']['y'], 22 | plydata['vertex']['z'] 23 | ], axis=1) 24 | return vertices 25 | 26 | 27 | def read_off(file): 28 | """ 29 | Reads vertices and faces from an off file. 30 | 31 | :param file: path to file to read 32 | :type file: str 33 | :return: vertices and faces as lists of tuples 34 | :rtype: [(float)], [(int)] 35 | """ 36 | 37 | assert os.path.exists(file), 'file %s not found' % file 38 | 39 | with open(file, 'r') as fp: 40 | lines = fp.readlines() 41 | lines = [line.strip() for line in lines] 42 | 43 | # Fix for ModelNet bug were 'OFF' and the number of vertices and faces 44 | # are all in the first line. 45 | if len(lines[0]) > 3: 46 | assert lines[0][:3] == 'OFF' or lines[0][:3] == 'off', \ 47 | 'invalid OFF file %s' % file 48 | 49 | parts = lines[0][3:].split(' ') 50 | assert len(parts) == 3 51 | 52 | num_vertices = int(parts[0]) 53 | assert num_vertices > 0 54 | 55 | num_faces = int(parts[1]) 56 | assert num_faces > 0 57 | 58 | start_index = 1 59 | # This is the regular case! 60 | else: 61 | assert lines[0] == 'OFF' or lines[0] == 'off', \ 62 | 'invalid OFF file %s' % file 63 | 64 | parts = lines[1].split(' ') 65 | assert len(parts) == 3 66 | 67 | num_vertices = int(parts[0]) 68 | assert num_vertices > 0 69 | 70 | num_faces = int(parts[1]) 71 | assert num_faces > 0 72 | 73 | start_index = 2 74 | 75 | vertices = [] 76 | for i in range(num_vertices): 77 | vertex = lines[start_index + i].split(' ') 78 | vertex = [float(point.strip()) for point in vertex if point != ''] 79 | assert len(vertex) == 3 80 | 81 | vertices.append(vertex) 82 | 83 | faces = [] 84 | for i in range(num_faces): 85 | face = lines[start_index + num_vertices + i].split(' ') 86 | face = [index.strip() for index in face if index != ''] 87 | 88 | # check to be sure 89 | for index in face: 90 | assert index != '', \ 91 | 'found empty vertex index: %s (%s)' \ 92 | % (lines[start_index + num_vertices + i], file) 93 | 94 | face = [int(index) for index in face] 95 | 96 | assert face[0] == len(face) - 1, \ 97 | 'face should have %d vertices but as %d (%s)' \ 98 | % (face[0], len(face) - 1, file) 99 | assert face[0] == 3, \ 100 | 'only triangular meshes supported (%s)' % file 101 | for index in face: 102 | assert index >= 0 and index < num_vertices, \ 103 | 'vertex %d (of %d vertices) does not exist (%s)' \ 104 | % (index, num_vertices, file) 105 | 106 | assert len(face) > 1 107 | 108 | faces.append(face) 109 | 110 | return vertices, faces 111 | 112 | assert False, 'could not open %s' % file 113 | -------------------------------------------------------------------------------- /src/utils/libkdtree/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/utils/libkdtree/.DS_Store -------------------------------------------------------------------------------- /src/utils/libkdtree/.gitignore: -------------------------------------------------------------------------------- 1 | build 2 | -------------------------------------------------------------------------------- /src/utils/libkdtree/MANIFEST.in: -------------------------------------------------------------------------------- 1 | exclude pykdtree/render_template.py 2 | include LICENSE.txt 3 | -------------------------------------------------------------------------------- /src/utils/libkdtree/README: -------------------------------------------------------------------------------- 1 | README.rst -------------------------------------------------------------------------------- /src/utils/libkdtree/README.rst: -------------------------------------------------------------------------------- 1 | .. image:: https://travis-ci.org/storpipfugl/pykdtree.svg?branch=master 2 | :target: https://travis-ci.org/storpipfugl/pykdtree 3 | .. image:: https://ci.appveyor.com/api/projects/status/ubo92368ktt2d25g/branch/master 4 | :target: https://ci.appveyor.com/project/storpipfugl/pykdtree 5 | 6 | ======== 7 | pykdtree 8 | ======== 9 | 10 | Objective 11 | --------- 12 | pykdtree is a kd-tree implementation for fast nearest neighbour search in Python. 13 | The aim is to be the fastest implementation around for common use cases (low dimensions and low number of neighbours) for both tree construction and queries. 14 | 15 | The implementation is based on scipy.spatial.cKDTree and libANN by combining the best features from both and focus on implementation efficiency. 16 | 17 | The interface is similar to that of scipy.spatial.cKDTree except only Euclidean distance measure is supported. 18 | 19 | Queries are optionally multithreaded using OpenMP. 20 | 21 | Installation 22 | ------------ 23 | Default build of pykdtree with OpenMP enabled queries using libgomp 24 | 25 | .. code-block:: bash 26 | 27 | $ cd 28 | $ python setup.py install 29 | 30 | If it fails with undefined compiler flags or you want to use another OpenMP implementation please modify setup.py at the indicated point to match your system. 31 | 32 | Building without OpenMP support is controlled by the USE_OMP environment variable 33 | 34 | .. code-block:: bash 35 | 36 | $ cd 37 | $ export USE_OMP=0 38 | $ python setup.py install 39 | 40 | Note evironment variables are by default not exported when using sudo so in this case do 41 | 42 | .. code-block:: bash 43 | 44 | $ USE_OMP=0 sudo -E python setup.py install 45 | 46 | Usage 47 | ----- 48 | The usage of pykdtree is similar to scipy.spatial.cKDTree so for now refer to its documentation 49 | 50 | >>> from pykdtree.kdtree import KDTree 51 | >>> kd_tree = KDTree(data_pts) 52 | >>> dist, idx = kd_tree.query(query_pts, k=8) 53 | 54 | The number of threads to be used in OpenMP enabled queries can be controlled with the standard OpenMP environment variable OMP_NUM_THREADS. 55 | 56 | The **leafsize** argument (number of data points per leaf) for the tree creation can be used to control the memory overhead of the kd-tree. pykdtree uses a default **leafsize=16**. 57 | Increasing **leafsize** will reduce the memory overhead and construction time but increase query time. 58 | 59 | pykdtree accepts data in double precision (numpy.float64) or single precision (numpy.float32) floating point. If data of another type is used an internal copy in double precision is made resulting in a memory overhead. If the kd-tree is constructed on single precision data the query points must be single precision as well. 60 | 61 | Benchmarks 62 | ---------- 63 | Comparison with scipy.spatial.cKDTree and libANN. This benchmark is on geospatial 3D data with 10053632 data points and 4276224 query points. The results are indexed relative to the construction time of scipy.spatial.cKDTree. A leafsize of 10 (scipy.spatial.cKDTree default) is used. 64 | 65 | Note: libANN is *not* thread safe. In this benchmark libANN is compiled with "-O3 -funroll-loops -ffast-math -fprefetch-loop-arrays" in order to achieve optimum performance. 66 | 67 | ================== ===================== ====== ======== ================== 68 | Operation scipy.spatial.cKDTree libANN pykdtree pykdtree 4 threads 69 | ------------------ --------------------- ------ -------- ------------------ 70 | 71 | Construction 100 304 96 96 72 | 73 | query 1 neighbour 1267 294 223 70 74 | 75 | Total 1 neighbour 1367 598 319 166 76 | 77 | query 8 neighbours 2193 625 449 143 78 | 79 | Total 8 neighbours 2293 929 545 293 80 | ================== ===================== ====== ======== ================== 81 | 82 | Looking at the combined construction and query this gives the following performance improvement relative to scipy.spatial.cKDTree 83 | 84 | ========== ====== ======== ================== 85 | Neighbours libANN pykdtree pykdtree 4 threads 86 | ---------- ------ -------- ------------------ 87 | 1 129% 329% 723% 88 | 89 | 8 147% 320% 682% 90 | ========== ====== ======== ================== 91 | 92 | Note: mileage will vary with the dataset at hand and computer architecture. 93 | 94 | Test 95 | ---- 96 | Run the unit tests using nosetest 97 | 98 | .. code-block:: bash 99 | 100 | $ cd 101 | $ python setup.py nosetests 102 | 103 | Installing on AppVeyor 104 | ---------------------- 105 | 106 | Pykdtree requires the "stdint.h" header file which is not available on certain 107 | versions of Windows or certain Windows compilers including those on the 108 | continuous integration platform AppVeyor. To get around this the header file(s) 109 | can be downloaded and placed in the correct "include" directory. This can 110 | be done by adding the `anaconda/missing-headers.ps1` script to your repository 111 | and running it the install step of `appveyor.yml`: 112 | 113 | # install missing headers that aren't included with MSVC 2008 114 | # https://github.com/omnia-md/conda-recipes/pull/524 115 | - "powershell ./appveyor/missing-headers.ps1" 116 | 117 | In addition to this, AppVeyor does not support OpenMP so this feature must be 118 | turned off by adding the following to `appveyor.yml` in the 119 | `environment` section: 120 | 121 | environment: 122 | global: 123 | # Don't build with openmp because it isn't supported in appveyor's compilers 124 | USE_OMP: "0" 125 | 126 | Changelog 127 | --------- 128 | v1.3.1 : Fix masking in the "query" method introduced in 1.3.0 129 | 130 | v1.3.0 : Keyword argument "mask" added to "query" method. OpenMP compilation now works for MS Visual Studio compiler 131 | 132 | v1.2.2 : Build process fixes 133 | 134 | v1.2.1 : Fixed OpenMP thread safety issue introduced in v1.2.0 135 | 136 | v1.2.0 : 64 and 32 bit MSVC Windows support added 137 | 138 | v1.1.1 : Same as v1.1 release due to incorrect pypi release 139 | 140 | v1.1 : Build process improvements. Add data attribute to kdtree class for scipy interface compatibility 141 | 142 | v1.0 : Switched license from GPLv3 to LGPLv3 143 | 144 | v0.3 : Avoid zipping of installed egg 145 | 146 | v0.2 : Reduced memory footprint. Can now handle single precision data internally avoiding copy conversion to double precision. Default leafsize changed from 10 to 16 as this reduces the memory footprint and makes it a cache line multiplum (negligible if any query performance observed in benchmarks). Reduced memory allocation for leaf nodes. Applied patch for building on OS X. 147 | 148 | v0.1 : Initial version. 149 | -------------------------------------------------------------------------------- /src/utils/libkdtree/__init__.py: -------------------------------------------------------------------------------- 1 | from .pykdtree.kdtree import KDTree 2 | 3 | 4 | __all__ = [ 5 | KDTree 6 | ] 7 | -------------------------------------------------------------------------------- /src/utils/libkdtree/pykdtree/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/utils/libkdtree/pykdtree/__init__.py -------------------------------------------------------------------------------- /src/utils/libkdtree/pykdtree/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/utils/libkdtree/pykdtree/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /src/utils/libkdtree/pykdtree/kdtree.cpython-36m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/utils/libkdtree/pykdtree/kdtree.cpython-36m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /src/utils/libkdtree/pykdtree/render_template.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from mako.template import Template 4 | 5 | mytemplate = Template(filename='_kdtree_core.c.mako') 6 | with open('_kdtree_core.c', 'w') as fp: 7 | fp.write(mytemplate.render()) 8 | -------------------------------------------------------------------------------- /src/utils/libkdtree/setup.cfg: -------------------------------------------------------------------------------- 1 | [bdist_rpm] 2 | requires=numpy 3 | release=1 4 | 5 | 6 | -------------------------------------------------------------------------------- /src/utils/libmcubes/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/utils/libmcubes/.DS_Store -------------------------------------------------------------------------------- /src/utils/libmcubes/.gitignore: -------------------------------------------------------------------------------- 1 | PyMCubes.egg-info 2 | build 3 | -------------------------------------------------------------------------------- /src/utils/libmcubes/LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2012-2015, P. M. Neila 2 | All rights reserved. 3 | 4 | Redistribution and use in source and binary forms, with or without 5 | modification, are permitted provided that the following conditions are met: 6 | 7 | * Redistributions of source code must retain the above copyright notice, this 8 | list of conditions and the following disclaimer. 9 | 10 | * Redistributions in binary form must reproduce the above copyright notice, 11 | this list of conditions and the following disclaimer in the documentation 12 | and/or other materials provided with the distribution. 13 | 14 | * Neither the name of the copyright holder nor the names of its 15 | contributors may be used to endorse or promote products derived from 16 | this software without specific prior written permission. 17 | 18 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | -------------------------------------------------------------------------------- /src/utils/libmcubes/README.rst: -------------------------------------------------------------------------------- 1 | ======== 2 | PyMCubes 3 | ======== 4 | 5 | PyMCubes is an implementation of the marching cubes algorithm to extract 6 | isosurfaces from volumetric data. The volumetric data can be given as a 7 | three-dimensional NumPy array or as a Python function ``f(x, y, z)``. The first 8 | option is much faster, but it requires more memory and becomes unfeasible for 9 | very large volumes. 10 | 11 | PyMCubes also provides a function to export the results of the marching cubes as 12 | COLLADA ``(.dae)`` files. This requires the 13 | `PyCollada `_ library. 14 | 15 | Installation 16 | ============ 17 | 18 | Just as any standard Python package, clone or download the project 19 | and run:: 20 | 21 | $ cd path/to/PyMCubes 22 | $ python setup.py build 23 | $ python setup.py install 24 | 25 | If you do not have write permission on the directory of Python packages, 26 | install with the ``--user`` option:: 27 | 28 | $ python setup.py install --user 29 | 30 | Example 31 | ======= 32 | 33 | The following example creates a data volume with spherical isosurfaces and 34 | extracts one of them (i.e., a sphere) with PyMCubes. The result is exported as 35 | ``sphere.dae``:: 36 | 37 | >>> import numpy as np 38 | >>> import mcubes 39 | 40 | # Create a data volume (30 x 30 x 30) 41 | >>> X, Y, Z = np.mgrid[:30, :30, :30] 42 | >>> u = (X-15)**2 + (Y-15)**2 + (Z-15)**2 - 8**2 43 | 44 | # Extract the 0-isosurface 45 | >>> vertices, triangles = mcubes.marching_cubes(u, 0) 46 | 47 | # Export the result to sphere.dae 48 | >>> mcubes.export_mesh(vertices, triangles, "sphere.dae", "MySphere") 49 | 50 | The second example is very similar to the first one, but it uses a function 51 | to represent the volume instead of a NumPy array:: 52 | 53 | >>> import numpy as np 54 | >>> import mcubes 55 | 56 | # Create the volume 57 | >>> f = lambda x, y, z: x**2 + y**2 + z**2 58 | 59 | # Extract the 16-isosurface 60 | >>> vertices, triangles = mcubes.marching_cubes_func((-10,-10,-10), (10,10,10), 61 | ... 100, 100, 100, f, 16) 62 | 63 | # Export the result to sphere2.dae 64 | >>> mcubes.export_mesh(vertices, triangles, "sphere2.dae", "MySphere") 65 | -------------------------------------------------------------------------------- /src/utils/libmcubes/__init__.py: -------------------------------------------------------------------------------- 1 | from src.utils.libmcubes.mcubes import ( 2 | marching_cubes, marching_cubes_func 3 | ) 4 | from src.utils.libmcubes.exporter import ( 5 | export_mesh, export_obj, export_off 6 | ) 7 | 8 | 9 | __all__ = [ 10 | marching_cubes, marching_cubes_func, 11 | export_mesh, export_obj, export_off 12 | ] 13 | -------------------------------------------------------------------------------- /src/utils/libmcubes/exporter.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | 4 | 5 | def export_obj(vertices, triangles, filename): 6 | """ 7 | Exports a mesh in the (.obj) format. 8 | """ 9 | 10 | with open(filename, 'w') as fh: 11 | 12 | for v in vertices: 13 | fh.write("v {} {} {}\n".format(*v)) 14 | 15 | for f in triangles: 16 | fh.write("f {} {} {}\n".format(*(f + 1))) 17 | 18 | 19 | def export_off(vertices, triangles, filename): 20 | """ 21 | Exports a mesh in the (.off) format. 22 | """ 23 | 24 | with open(filename, 'w') as fh: 25 | fh.write('OFF\n') 26 | fh.write('{} {} 0\n'.format(len(vertices), len(triangles))) 27 | 28 | for v in vertices: 29 | fh.write("{} {} {}\n".format(*v)) 30 | 31 | for f in triangles: 32 | fh.write("3 {} {} {}\n".format(*f)) 33 | 34 | 35 | def export_mesh(vertices, triangles, filename, mesh_name="mcubes_mesh"): 36 | """ 37 | Exports a mesh in the COLLADA (.dae) format. 38 | 39 | Needs PyCollada (https://github.com/pycollada/pycollada). 40 | """ 41 | 42 | import collada 43 | 44 | mesh = collada.Collada() 45 | 46 | vert_src = collada.source.FloatSource("verts-array", vertices, ('X','Y','Z')) 47 | geom = collada.geometry.Geometry(mesh, "geometry0", mesh_name, [vert_src]) 48 | 49 | input_list = collada.source.InputList() 50 | input_list.addInput(0, 'VERTEX', "#verts-array") 51 | 52 | triset = geom.createTriangleSet(np.copy(triangles), input_list, "") 53 | geom.primitives.append(triset) 54 | mesh.geometries.append(geom) 55 | 56 | geomnode = collada.scene.GeometryNode(geom, []) 57 | node = collada.scene.Node(mesh_name, children=[geomnode]) 58 | 59 | myscene = collada.scene.Scene("mcubes_scene", [node]) 60 | mesh.scenes.append(myscene) 61 | mesh.scene = myscene 62 | 63 | mesh.write(filename) 64 | -------------------------------------------------------------------------------- /src/utils/libmcubes/mcubes.cpython-36m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/utils/libmcubes/mcubes.cpython-36m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /src/utils/libmcubes/mcubes.pyx: -------------------------------------------------------------------------------- 1 | 2 | # distutils: language = c++ 3 | # cython: embedsignature = True 4 | 5 | # from libcpp.vector cimport vector 6 | import numpy as np 7 | 8 | # Define PY_ARRAY_UNIQUE_SYMBOL 9 | cdef extern from "pyarray_symbol.h": 10 | pass 11 | 12 | cimport numpy as np 13 | 14 | np.import_array() 15 | 16 | cdef extern from "pywrapper.h": 17 | cdef object c_marching_cubes "marching_cubes"(np.ndarray, double) except + 18 | cdef object c_marching_cubes2 "marching_cubes2"(np.ndarray, double) except + 19 | cdef object c_marching_cubes3 "marching_cubes3"(np.ndarray, double) except + 20 | cdef object c_marching_cubes_func "marching_cubes_func"(tuple, tuple, int, int, int, object, double) except + 21 | 22 | def marching_cubes(np.ndarray volume, float isovalue): 23 | 24 | verts, faces = c_marching_cubes(volume, isovalue) 25 | verts.shape = (-1, 3) 26 | faces.shape = (-1, 3) 27 | return verts, faces 28 | 29 | def marching_cubes2(np.ndarray volume, float isovalue): 30 | 31 | verts, faces = c_marching_cubes2(volume, isovalue) 32 | verts.shape = (-1, 3) 33 | faces.shape = (-1, 3) 34 | return verts, faces 35 | 36 | def marching_cubes3(np.ndarray volume, float isovalue): 37 | 38 | verts, faces = c_marching_cubes3(volume, isovalue) 39 | verts.shape = (-1, 3) 40 | faces.shape = (-1, 3) 41 | return verts, faces 42 | 43 | def marching_cubes_func(tuple lower, tuple upper, int numx, int numy, int numz, object f, double isovalue): 44 | 45 | verts, faces = c_marching_cubes_func(lower, upper, numx, numy, numz, f, isovalue) 46 | verts.shape = (-1, 3) 47 | faces.shape = (-1, 3) 48 | return verts, faces 49 | -------------------------------------------------------------------------------- /src/utils/libmcubes/pyarray_symbol.h: -------------------------------------------------------------------------------- 1 | 2 | #define PY_ARRAY_UNIQUE_SYMBOL mcubes_PyArray_API 3 | -------------------------------------------------------------------------------- /src/utils/libmcubes/pyarraymodule.h: -------------------------------------------------------------------------------- 1 | 2 | #ifndef _EXTMODULE_H 3 | #define _EXTMODULE_H 4 | 5 | #include 6 | #include 7 | 8 | // #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION 9 | #define PY_ARRAY_UNIQUE_SYMBOL mcubes_PyArray_API 10 | #define NO_IMPORT_ARRAY 11 | #include "numpy/arrayobject.h" 12 | 13 | #include 14 | 15 | template 16 | struct numpy_typemap; 17 | 18 | #define define_numpy_type(ctype, dtype) \ 19 | template<> \ 20 | struct numpy_typemap \ 21 | {static const int type = dtype;}; 22 | 23 | define_numpy_type(bool, NPY_BOOL); 24 | define_numpy_type(char, NPY_BYTE); 25 | define_numpy_type(short, NPY_SHORT); 26 | define_numpy_type(int, NPY_INT); 27 | define_numpy_type(long, NPY_LONG); 28 | define_numpy_type(long long, NPY_LONGLONG); 29 | define_numpy_type(unsigned char, NPY_UBYTE); 30 | define_numpy_type(unsigned short, NPY_USHORT); 31 | define_numpy_type(unsigned int, NPY_UINT); 32 | define_numpy_type(unsigned long, NPY_ULONG); 33 | define_numpy_type(unsigned long long, NPY_ULONGLONG); 34 | define_numpy_type(float, NPY_FLOAT); 35 | define_numpy_type(double, NPY_DOUBLE); 36 | define_numpy_type(long double, NPY_LONGDOUBLE); 37 | define_numpy_type(std::complex, NPY_CFLOAT); 38 | define_numpy_type(std::complex, NPY_CDOUBLE); 39 | define_numpy_type(std::complex, NPY_CLONGDOUBLE); 40 | 41 | template 42 | T PyArray_SafeGet(const PyArrayObject* aobj, const npy_intp* indaux) 43 | { 44 | // HORROR. 45 | npy_intp* ind = const_cast(indaux); 46 | void* ptr = PyArray_GetPtr(const_cast(aobj), ind); 47 | switch(PyArray_TYPE(aobj)) 48 | { 49 | case NPY_BOOL: 50 | return static_cast(*reinterpret_cast(ptr)); 51 | case NPY_BYTE: 52 | return static_cast(*reinterpret_cast(ptr)); 53 | case NPY_SHORT: 54 | return static_cast(*reinterpret_cast(ptr)); 55 | case NPY_INT: 56 | return static_cast(*reinterpret_cast(ptr)); 57 | case NPY_LONG: 58 | return static_cast(*reinterpret_cast(ptr)); 59 | case NPY_LONGLONG: 60 | return static_cast(*reinterpret_cast(ptr)); 61 | case NPY_UBYTE: 62 | return static_cast(*reinterpret_cast(ptr)); 63 | case NPY_USHORT: 64 | return static_cast(*reinterpret_cast(ptr)); 65 | case NPY_UINT: 66 | return static_cast(*reinterpret_cast(ptr)); 67 | case NPY_ULONG: 68 | return static_cast(*reinterpret_cast(ptr)); 69 | case NPY_ULONGLONG: 70 | return static_cast(*reinterpret_cast(ptr)); 71 | case NPY_FLOAT: 72 | return static_cast(*reinterpret_cast(ptr)); 73 | case NPY_DOUBLE: 74 | return static_cast(*reinterpret_cast(ptr)); 75 | case NPY_LONGDOUBLE: 76 | return static_cast(*reinterpret_cast(ptr)); 77 | default: 78 | throw std::runtime_error("data type not supported"); 79 | } 80 | } 81 | 82 | template 83 | T PyArray_SafeSet(PyArrayObject* aobj, const npy_intp* indaux, const T& value) 84 | { 85 | // HORROR. 86 | npy_intp* ind = const_cast(indaux); 87 | void* ptr = PyArray_GetPtr(aobj, ind); 88 | switch(PyArray_TYPE(aobj)) 89 | { 90 | case NPY_BOOL: 91 | *reinterpret_cast(ptr) = static_cast(value); 92 | break; 93 | case NPY_BYTE: 94 | *reinterpret_cast(ptr) = static_cast(value); 95 | break; 96 | case NPY_SHORT: 97 | *reinterpret_cast(ptr) = static_cast(value); 98 | break; 99 | case NPY_INT: 100 | *reinterpret_cast(ptr) = static_cast(value); 101 | break; 102 | case NPY_LONG: 103 | *reinterpret_cast(ptr) = static_cast(value); 104 | break; 105 | case NPY_LONGLONG: 106 | *reinterpret_cast(ptr) = static_cast(value); 107 | break; 108 | case NPY_UBYTE: 109 | *reinterpret_cast(ptr) = static_cast(value); 110 | break; 111 | case NPY_USHORT: 112 | *reinterpret_cast(ptr) = static_cast(value); 113 | break; 114 | case NPY_UINT: 115 | *reinterpret_cast(ptr) = static_cast(value); 116 | break; 117 | case NPY_ULONG: 118 | *reinterpret_cast(ptr) = static_cast(value); 119 | break; 120 | case NPY_ULONGLONG: 121 | *reinterpret_cast(ptr) = static_cast(value); 122 | break; 123 | case NPY_FLOAT: 124 | *reinterpret_cast(ptr) = static_cast(value); 125 | break; 126 | case NPY_DOUBLE: 127 | *reinterpret_cast(ptr) = static_cast(value); 128 | break; 129 | case NPY_LONGDOUBLE: 130 | *reinterpret_cast(ptr) = static_cast(value); 131 | break; 132 | default: 133 | throw std::runtime_error("data type not supported"); 134 | } 135 | } 136 | 137 | #endif 138 | -------------------------------------------------------------------------------- /src/utils/libmcubes/pywrapper.h: -------------------------------------------------------------------------------- 1 | 2 | #ifndef _PYWRAPPER_H 3 | #define _PYWRAPPER_H 4 | 5 | #include 6 | #include "pyarraymodule.h" 7 | 8 | #include 9 | 10 | PyObject* marching_cubes(PyArrayObject* arr, double isovalue); 11 | PyObject* marching_cubes2(PyArrayObject* arr, double isovalue); 12 | PyObject* marching_cubes3(PyArrayObject* arr, double isovalue); 13 | PyObject* marching_cubes_func(PyObject* lower, PyObject* upper, 14 | int numx, int numy, int numz, PyObject* f, double isovalue); 15 | 16 | #endif // _PYWRAPPER_H 17 | -------------------------------------------------------------------------------- /src/utils/libmesh/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/utils/libmesh/.DS_Store -------------------------------------------------------------------------------- /src/utils/libmesh/.gitignore: -------------------------------------------------------------------------------- 1 | triangle_hash.cpp 2 | build 3 | -------------------------------------------------------------------------------- /src/utils/libmesh/__init__.py: -------------------------------------------------------------------------------- 1 | from .inside_mesh import ( 2 | check_mesh_contains, MeshIntersector, TriangleIntersector2d 3 | ) 4 | 5 | 6 | __all__ = [ 7 | check_mesh_contains, MeshIntersector, TriangleIntersector2d 8 | ] 9 | -------------------------------------------------------------------------------- /src/utils/libmesh/inside_mesh.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from .triangle_hash import TriangleHash as _TriangleHash 3 | 4 | 5 | def check_mesh_contains(mesh, points, hash_resolution=512): 6 | intersector = MeshIntersector(mesh, hash_resolution) 7 | contains = intersector.query(points) 8 | return contains 9 | 10 | 11 | class MeshIntersector: 12 | def __init__(self, mesh, resolution=512): 13 | triangles = mesh.vertices[mesh.faces].astype(np.float64) 14 | n_tri = triangles.shape[0] 15 | 16 | self.resolution = resolution 17 | self.bbox_min = triangles.reshape(3 * n_tri, 3).min(axis=0) 18 | self.bbox_max = triangles.reshape(3 * n_tri, 3).max(axis=0) 19 | # Tranlate and scale it to [0.5, self.resolution - 0.5]^3 20 | self.scale = (resolution - 1) / (self.bbox_max - self.bbox_min) 21 | self.translate = 0.5 - self.scale * self.bbox_min 22 | 23 | self._triangles = triangles = self.rescale(triangles) 24 | # assert(np.allclose(triangles.reshape(-1, 3).min(0), 0.5)) 25 | # assert(np.allclose(triangles.reshape(-1, 3).max(0), resolution - 0.5)) 26 | 27 | triangles2d = triangles[:, :, :2] 28 | self._tri_intersector2d = TriangleIntersector2d( 29 | triangles2d, resolution) 30 | 31 | def query(self, points): 32 | # Rescale points 33 | points = self.rescale(points) 34 | 35 | # placeholder result with no hits we'll fill in later 36 | contains = np.zeros(len(points), dtype=np.bool) 37 | 38 | # cull points outside of the axis aligned bounding box 39 | # this avoids running ray tests unless points are close 40 | inside_aabb = np.all( 41 | (0 <= points) & (points <= self.resolution), axis=1) 42 | if not inside_aabb.any(): 43 | return contains 44 | 45 | # Only consider points inside bounding box 46 | mask = inside_aabb 47 | points = points[mask] 48 | 49 | # Compute intersection depth and check order 50 | points_indices, tri_indices = self._tri_intersector2d.query(points[:, :2]) 51 | 52 | triangles_intersect = self._triangles[tri_indices] 53 | points_intersect = points[points_indices] 54 | 55 | depth_intersect, abs_n_2 = self.compute_intersection_depth( 56 | points_intersect, triangles_intersect) 57 | 58 | # Count number of intersections in both directions 59 | smaller_depth = depth_intersect >= points_intersect[:, 2] * abs_n_2 60 | bigger_depth = depth_intersect < points_intersect[:, 2] * abs_n_2 61 | points_indices_0 = points_indices[smaller_depth] 62 | points_indices_1 = points_indices[bigger_depth] 63 | 64 | nintersect0 = np.bincount(points_indices_0, minlength=points.shape[0]) 65 | nintersect1 = np.bincount(points_indices_1, minlength=points.shape[0]) 66 | 67 | # Check if point contained in mesh 68 | contains1 = (np.mod(nintersect0, 2) == 1) 69 | contains2 = (np.mod(nintersect1, 2) == 1) 70 | if (contains1 != contains2).any(): 71 | print('Warning: contains1 != contains2 for some points.') 72 | contains[mask] = (contains1 & contains2) 73 | return contains 74 | 75 | def compute_intersection_depth(self, points, triangles): 76 | t1 = triangles[:, 0, :] 77 | t2 = triangles[:, 1, :] 78 | t3 = triangles[:, 2, :] 79 | 80 | v1 = t3 - t1 81 | v2 = t2 - t1 82 | # v1 = v1 / np.linalg.norm(v1, axis=-1, keepdims=True) 83 | # v2 = v2 / np.linalg.norm(v2, axis=-1, keepdims=True) 84 | 85 | normals = np.cross(v1, v2) 86 | alpha = np.sum(normals[:, :2] * (t1[:, :2] - points[:, :2]), axis=1) 87 | 88 | n_2 = normals[:, 2] 89 | t1_2 = t1[:, 2] 90 | s_n_2 = np.sign(n_2) 91 | abs_n_2 = np.abs(n_2) 92 | 93 | mask = (abs_n_2 != 0) 94 | 95 | depth_intersect = np.full(points.shape[0], np.nan) 96 | depth_intersect[mask] = \ 97 | t1_2[mask] * abs_n_2[mask] + alpha[mask] * s_n_2[mask] 98 | 99 | # Test the depth: 100 | # TODO: remove and put into tests 101 | # points_new = np.concatenate([points[:, :2], depth_intersect[:, None]], axis=1) 102 | # alpha = (normals * t1).sum(-1) 103 | # mask = (depth_intersect == depth_intersect) 104 | # assert(np.allclose((points_new[mask] * normals[mask]).sum(-1), 105 | # alpha[mask])) 106 | return depth_intersect, abs_n_2 107 | 108 | def rescale(self, array): 109 | array = self.scale * array + self.translate 110 | return array 111 | 112 | 113 | class TriangleIntersector2d: 114 | def __init__(self, triangles, resolution=128): 115 | self.triangles = triangles 116 | self.tri_hash = _TriangleHash(triangles, resolution) 117 | 118 | def query(self, points): 119 | point_indices, tri_indices = self.tri_hash.query(points) 120 | point_indices = np.array(point_indices, dtype=np.int64) 121 | tri_indices = np.array(tri_indices, dtype=np.int64) 122 | points = points[point_indices] 123 | triangles = self.triangles[tri_indices] 124 | mask = self.check_triangles(points, triangles) 125 | point_indices = point_indices[mask] 126 | tri_indices = tri_indices[mask] 127 | return point_indices, tri_indices 128 | 129 | def check_triangles(self, points, triangles): 130 | contains = np.zeros(points.shape[0], dtype=np.bool) 131 | A = triangles[:, :2] - triangles[:, 2:] 132 | A = A.transpose([0, 2, 1]) 133 | y = points - triangles[:, 2] 134 | 135 | detA = A[:, 0, 0] * A[:, 1, 1] - A[:, 0, 1] * A[:, 1, 0] 136 | 137 | mask = (np.abs(detA) != 0.) 138 | A = A[mask] 139 | y = y[mask] 140 | detA = detA[mask] 141 | 142 | s_detA = np.sign(detA) 143 | abs_detA = np.abs(detA) 144 | 145 | u = (A[:, 1, 1] * y[:, 0] - A[:, 0, 1] * y[:, 1]) * s_detA 146 | v = (-A[:, 1, 0] * y[:, 0] + A[:, 0, 0] * y[:, 1]) * s_detA 147 | 148 | sum_uv = u + v 149 | contains[mask] = ( 150 | (0 < u) & (u < abs_detA) & (0 < v) & (v < abs_detA) 151 | & (0 < sum_uv) & (sum_uv < abs_detA) 152 | ) 153 | return contains 154 | 155 | -------------------------------------------------------------------------------- /src/utils/libmesh/triangle_hash.cpython-36m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/utils/libmesh/triangle_hash.cpython-36m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /src/utils/libmesh/triangle_hash.pyx: -------------------------------------------------------------------------------- 1 | 2 | # distutils: language=c++ 3 | import numpy as np 4 | cimport numpy as np 5 | cimport cython 6 | from libcpp.vector cimport vector 7 | from libc.math cimport floor, ceil 8 | 9 | cdef class TriangleHash: 10 | cdef vector[vector[int]] spatial_hash 11 | cdef int resolution 12 | 13 | def __cinit__(self, double[:, :, :] triangles, int resolution): 14 | self.spatial_hash.resize(resolution * resolution) 15 | self.resolution = resolution 16 | self._build_hash(triangles) 17 | 18 | @cython.boundscheck(False) # Deactivate bounds checking 19 | @cython.wraparound(False) # Deactivate negative indexing. 20 | cdef int _build_hash(self, double[:, :, :] triangles): 21 | assert(triangles.shape[1] == 3) 22 | assert(triangles.shape[2] == 2) 23 | 24 | cdef int n_tri = triangles.shape[0] 25 | cdef int bbox_min[2] 26 | cdef int bbox_max[2] 27 | 28 | cdef int i_tri, j, x, y 29 | cdef int spatial_idx 30 | 31 | for i_tri in range(n_tri): 32 | # Compute bounding box 33 | for j in range(2): 34 | bbox_min[j] = min( 35 | triangles[i_tri, 0, j], triangles[i_tri, 1, j], triangles[i_tri, 2, j] 36 | ) 37 | bbox_max[j] = max( 38 | triangles[i_tri, 0, j], triangles[i_tri, 1, j], triangles[i_tri, 2, j] 39 | ) 40 | bbox_min[j] = min(max(bbox_min[j], 0), self.resolution - 1) 41 | bbox_max[j] = min(max(bbox_max[j], 0), self.resolution - 1) 42 | 43 | # Find all voxels where bounding box intersects 44 | for x in range(bbox_min[0], bbox_max[0] + 1): 45 | for y in range(bbox_min[1], bbox_max[1] + 1): 46 | spatial_idx = self.resolution * x + y 47 | self.spatial_hash[spatial_idx].push_back(i_tri) 48 | 49 | @cython.boundscheck(False) # Deactivate bounds checking 50 | @cython.wraparound(False) # Deactivate negative indexing. 51 | cpdef query(self, double[:, :] points): 52 | assert(points.shape[1] == 2) 53 | cdef int n_points = points.shape[0] 54 | 55 | cdef vector[int] points_indices 56 | cdef vector[int] tri_indices 57 | # cdef int[:] points_indices_np 58 | # cdef int[:] tri_indices_np 59 | 60 | cdef int i_point, k, x, y 61 | cdef int spatial_idx 62 | 63 | for i_point in range(n_points): 64 | x = int(points[i_point, 0]) 65 | y = int(points[i_point, 1]) 66 | if not (0 <= x < self.resolution and 0 <= y < self.resolution): 67 | continue 68 | 69 | spatial_idx = self.resolution * x + y 70 | for i_tri in self.spatial_hash[spatial_idx]: 71 | points_indices.push_back(i_point) 72 | tri_indices.push_back(i_tri) 73 | 74 | points_indices_np = np.zeros(points_indices.size(), dtype=np.int32) 75 | tri_indices_np = np.zeros(tri_indices.size(), dtype=np.int32) 76 | 77 | cdef int[:] points_indices_view = points_indices_np 78 | cdef int[:] tri_indices_view = tri_indices_np 79 | 80 | for k in range(points_indices.size()): 81 | points_indices_view[k] = points_indices[k] 82 | 83 | for k in range(tri_indices.size()): 84 | tri_indices_view[k] = tri_indices[k] 85 | 86 | return points_indices_np, tri_indices_np 87 | -------------------------------------------------------------------------------- /src/utils/libmise/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/utils/libmise/.DS_Store -------------------------------------------------------------------------------- /src/utils/libmise/.gitignore: -------------------------------------------------------------------------------- 1 | mise.c 2 | mise.cpp 3 | mise.html 4 | -------------------------------------------------------------------------------- /src/utils/libmise/__init__.py: -------------------------------------------------------------------------------- 1 | from .mise import MISE 2 | 3 | __all__ = [ 4 | MISE 5 | ] 6 | -------------------------------------------------------------------------------- /src/utils/libmise/mise.cpython-36m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/utils/libmise/mise.cpython-36m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /src/utils/libmise/test.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from mise import MISE 3 | import time 4 | 5 | t0 = time.time() 6 | extractor = MISE(1, 2, 0.) 7 | 8 | p = extractor.query() 9 | i = 0 10 | 11 | while p.shape[0] != 0: 12 | print(i) 13 | print(p) 14 | v = 2 * (p.sum(axis=-1) > 2).astype(np.float64) - 1 15 | extractor.update(p, v) 16 | p = extractor.query() 17 | i += 1 18 | if (i >= 8): 19 | break 20 | 21 | print(extractor.to_dense()) 22 | # p, v = extractor.get_points() 23 | # print(p) 24 | # print(v) 25 | print('Total time: %f' % (time.time() - t0)) 26 | -------------------------------------------------------------------------------- /src/utils/libsimplify/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/utils/libsimplify/.DS_Store -------------------------------------------------------------------------------- /src/utils/libsimplify/__init__.py: -------------------------------------------------------------------------------- 1 | from .simplify_mesh import ( 2 | mesh_simplify 3 | ) 4 | import trimesh 5 | 6 | 7 | def simplify_mesh(mesh, f_target=10000, agressiveness=7.): 8 | vertices = mesh.vertices 9 | faces = mesh.faces 10 | 11 | vertices, faces = mesh_simplify(vertices, faces, f_target, agressiveness) 12 | 13 | mesh_simplified = trimesh.Trimesh(vertices, faces, process=False) 14 | 15 | return mesh_simplified 16 | -------------------------------------------------------------------------------- /src/utils/libsimplify/simplify_mesh.cpython-36m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/utils/libsimplify/simplify_mesh.cpython-36m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /src/utils/libsimplify/simplify_mesh.pyx: -------------------------------------------------------------------------------- 1 | # distutils: language = c++ 2 | from libcpp.vector cimport vector 3 | import numpy as np 4 | cimport numpy as np 5 | 6 | 7 | cdef extern from "Simplify.h": 8 | cdef struct vec3f: 9 | double x, y, z 10 | 11 | cdef cppclass SymetricMatrix: 12 | SymetricMatrix() except + 13 | 14 | 15 | cdef extern from "Simplify.h" namespace "Simplify": 16 | cdef struct Triangle: 17 | int v[3] 18 | double err[4] 19 | int deleted, dirty, attr 20 | vec3f uvs[3] 21 | int material 22 | 23 | cdef struct Vertex: 24 | vec3f p 25 | int tstart, tcount 26 | SymetricMatrix q 27 | int border 28 | 29 | cdef vector[Triangle] triangles 30 | cdef vector[Vertex] vertices 31 | cdef void simplify_mesh(int, double) 32 | 33 | 34 | cpdef mesh_simplify(double[:, ::1] vertices_in, long[:, ::1] triangles_in, 35 | int f_target, double agressiveness=7.) except +: 36 | vertices.clear() 37 | triangles.clear() 38 | 39 | # Read in vertices and triangles 40 | cdef Vertex v 41 | for iv in range(vertices_in.shape[0]): 42 | v = Vertex() 43 | v.p.x = vertices_in[iv, 0] 44 | v.p.y = vertices_in[iv, 1] 45 | v.p.z = vertices_in[iv, 2] 46 | vertices.push_back(v) 47 | 48 | cdef Triangle t 49 | for it in range(triangles_in.shape[0]): 50 | t = Triangle() 51 | t.v[0] = triangles_in[it, 0] 52 | t.v[1] = triangles_in[it, 1] 53 | t.v[2] = triangles_in[it, 2] 54 | triangles.push_back(t) 55 | 56 | # Simplify 57 | # print('Simplify...') 58 | simplify_mesh(f_target, agressiveness) 59 | 60 | # Only use triangles that are not deleted 61 | cdef vector[Triangle] triangles_notdel 62 | triangles_notdel.reserve(triangles.size()) 63 | 64 | for t in triangles: 65 | if not t.deleted: 66 | triangles_notdel.push_back(t) 67 | 68 | # Read out triangles 69 | vertices_out = np.empty((vertices.size(), 3), dtype=np.float64) 70 | triangles_out = np.empty((triangles_notdel.size(), 3), dtype=np.int64) 71 | 72 | cdef double[:, :] vertices_out_view = vertices_out 73 | cdef long[:, :] triangles_out_view = triangles_out 74 | 75 | for iv in range(vertices.size()): 76 | vertices_out_view[iv, 0] = vertices[iv].p.x 77 | vertices_out_view[iv, 1] = vertices[iv].p.y 78 | vertices_out_view[iv, 2] = vertices[iv].p.z 79 | 80 | for it in range(triangles_notdel.size()): 81 | triangles_out_view[it, 0] = triangles_notdel[it].v[0] 82 | triangles_out_view[it, 1] = triangles_notdel[it].v[1] 83 | triangles_out_view[it, 2] = triangles_notdel[it].v[2] 84 | 85 | # Clear vertices and triangles 86 | vertices.clear() 87 | triangles.clear() 88 | 89 | return vertices_out, triangles_out -------------------------------------------------------------------------------- /src/utils/libsimplify/test.py: -------------------------------------------------------------------------------- 1 | from simplify_mesh import mesh_simplify 2 | import numpy as np 3 | 4 | v = np.random.rand(100, 3) 5 | f = np.random.choice(range(100), (50, 3)) 6 | 7 | mesh_simplify(v, f, 50) -------------------------------------------------------------------------------- /src/utils/libvoxelize/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/utils/libvoxelize/.DS_Store -------------------------------------------------------------------------------- /src/utils/libvoxelize/.gitignore: -------------------------------------------------------------------------------- 1 | voxelize.c 2 | voxelize.html 3 | build 4 | -------------------------------------------------------------------------------- /src/utils/libvoxelize/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/utils/libvoxelize/__init__.py -------------------------------------------------------------------------------- /src/utils/libvoxelize/tribox2.h: -------------------------------------------------------------------------------- 1 | /********************************************************/ 2 | /* AABB-triangle overlap test code */ 3 | /* by Tomas Akenine-M�ller */ 4 | /* Function: int triBoxOverlap(float boxcenter[3], */ 5 | /* float boxhalfsize[3],float triverts[3][3]); */ 6 | /* History: */ 7 | /* 2001-03-05: released the code in its first version */ 8 | /* 2001-06-18: changed the order of the tests, faster */ 9 | /* */ 10 | /* Acknowledgement: Many thanks to Pierre Terdiman for */ 11 | /* suggestions and discussions on how to optimize code. */ 12 | /* Thanks to David Hunt for finding a ">="-bug! */ 13 | /********************************************************/ 14 | #include 15 | #include 16 | 17 | #define X 0 18 | #define Y 1 19 | #define Z 2 20 | 21 | #define CROSS(dest,v1,v2) \ 22 | dest[0]=v1[1]*v2[2]-v1[2]*v2[1]; \ 23 | dest[1]=v1[2]*v2[0]-v1[0]*v2[2]; \ 24 | dest[2]=v1[0]*v2[1]-v1[1]*v2[0]; 25 | 26 | #define DOT(v1,v2) (v1[0]*v2[0]+v1[1]*v2[1]+v1[2]*v2[2]) 27 | 28 | #define SUB(dest,v1,v2) \ 29 | dest[0]=v1[0]-v2[0]; \ 30 | dest[1]=v1[1]-v2[1]; \ 31 | dest[2]=v1[2]-v2[2]; 32 | 33 | #define FINDMINMAX(x0,x1,x2,min,max) \ 34 | min = max = x0; \ 35 | if(x1max) max=x1;\ 37 | if(x2max) max=x2; 39 | 40 | int planeBoxOverlap(float normal[3],float d, float maxbox[3]) 41 | { 42 | int q; 43 | float vmin[3],vmax[3]; 44 | for(q=X;q<=Z;q++) 45 | { 46 | if(normal[q]>0.0f) 47 | { 48 | vmin[q]=-maxbox[q]; 49 | vmax[q]=maxbox[q]; 50 | } 51 | else 52 | { 53 | vmin[q]=maxbox[q]; 54 | vmax[q]=-maxbox[q]; 55 | } 56 | } 57 | if(DOT(normal,vmin)+d>0.0f) return 0; 58 | if(DOT(normal,vmax)+d>=0.0f) return 1; 59 | 60 | return 0; 61 | } 62 | 63 | 64 | /*======================== X-tests ========================*/ 65 | #define AXISTEST_X01(a, b, fa, fb) \ 66 | p0 = a*v0[Y] - b*v0[Z]; \ 67 | p2 = a*v2[Y] - b*v2[Z]; \ 68 | if(p0rad || max<-rad) return 0; 71 | 72 | #define AXISTEST_X2(a, b, fa, fb) \ 73 | p0 = a*v0[Y] - b*v0[Z]; \ 74 | p1 = a*v1[Y] - b*v1[Z]; \ 75 | if(p0rad || max<-rad) return 0; 78 | 79 | /*======================== Y-tests ========================*/ 80 | #define AXISTEST_Y02(a, b, fa, fb) \ 81 | p0 = -a*v0[X] + b*v0[Z]; \ 82 | p2 = -a*v2[X] + b*v2[Z]; \ 83 | if(p0rad || max<-rad) return 0; 86 | 87 | #define AXISTEST_Y1(a, b, fa, fb) \ 88 | p0 = -a*v0[X] + b*v0[Z]; \ 89 | p1 = -a*v1[X] + b*v1[Z]; \ 90 | if(p0rad || max<-rad) return 0; 93 | 94 | /*======================== Z-tests ========================*/ 95 | 96 | #define AXISTEST_Z12(a, b, fa, fb) \ 97 | p1 = a*v1[X] - b*v1[Y]; \ 98 | p2 = a*v2[X] - b*v2[Y]; \ 99 | if(p2rad || max<-rad) return 0; 102 | 103 | #define AXISTEST_Z0(a, b, fa, fb) \ 104 | p0 = a*v0[X] - b*v0[Y]; \ 105 | p1 = a*v1[X] - b*v1[Y]; \ 106 | if(p0rad || max<-rad) return 0; 109 | 110 | int triBoxOverlap(float boxcenter[3],float boxhalfsize[3],float tri0[3], float tri1[3], float tri2[3]) 111 | { 112 | 113 | /* use separating axis theorem to test overlap between triangle and box */ 114 | /* need to test for overlap in these directions: */ 115 | /* 1) the {x,y,z}-directions (actually, since we use the AABB of the triangle */ 116 | /* we do not even need to test these) */ 117 | /* 2) normal of the triangle */ 118 | /* 3) crossproduct(edge from tri, {x,y,z}-directin) */ 119 | /* this gives 3x3=9 more tests */ 120 | float v0[3],v1[3],v2[3]; 121 | float min,max,d,p0,p1,p2,rad,fex,fey,fez; 122 | float normal[3],e0[3],e1[3],e2[3]; 123 | 124 | /* This is the fastest branch on Sun */ 125 | /* move everything so that the boxcenter is in (0,0,0) */ 126 | SUB(v0, tri0, boxcenter); 127 | SUB(v1, tri1, boxcenter); 128 | SUB(v2, tri2, boxcenter); 129 | 130 | /* compute triangle edges */ 131 | SUB(e0,v1,v0); /* tri edge 0 */ 132 | SUB(e1,v2,v1); /* tri edge 1 */ 133 | SUB(e2,v0,v2); /* tri edge 2 */ 134 | 135 | /* Bullet 3: */ 136 | /* test the 9 tests first (this was faster) */ 137 | fex = fabs(e0[X]); 138 | fey = fabs(e0[Y]); 139 | fez = fabs(e0[Z]); 140 | AXISTEST_X01(e0[Z], e0[Y], fez, fey); 141 | AXISTEST_Y02(e0[Z], e0[X], fez, fex); 142 | AXISTEST_Z12(e0[Y], e0[X], fey, fex); 143 | 144 | fex = fabs(e1[X]); 145 | fey = fabs(e1[Y]); 146 | fez = fabs(e1[Z]); 147 | AXISTEST_X01(e1[Z], e1[Y], fez, fey); 148 | AXISTEST_Y02(e1[Z], e1[X], fez, fex); 149 | AXISTEST_Z0(e1[Y], e1[X], fey, fex); 150 | 151 | fex = fabs(e2[X]); 152 | fey = fabs(e2[Y]); 153 | fez = fabs(e2[Z]); 154 | AXISTEST_X2(e2[Z], e2[Y], fez, fey); 155 | AXISTEST_Y1(e2[Z], e2[X], fez, fex); 156 | AXISTEST_Z12(e2[Y], e2[X], fey, fex); 157 | 158 | /* Bullet 1: */ 159 | /* first test overlap in the {x,y,z}-directions */ 160 | /* find min, max of the triangle each direction, and test for overlap in */ 161 | /* that direction -- this is equivalent to testing a minimal AABB around */ 162 | /* the triangle against the AABB */ 163 | 164 | /* test in X-direction */ 165 | FINDMINMAX(v0[X],v1[X],v2[X],min,max); 166 | if(min>boxhalfsize[X] || max<-boxhalfsize[X]) return 0; 167 | 168 | /* test in Y-direction */ 169 | FINDMINMAX(v0[Y],v1[Y],v2[Y],min,max); 170 | if(min>boxhalfsize[Y] || max<-boxhalfsize[Y]) return 0; 171 | 172 | /* test in Z-direction */ 173 | FINDMINMAX(v0[Z],v1[Z],v2[Z],min,max); 174 | if(min>boxhalfsize[Z] || max<-boxhalfsize[Z]) return 0; 175 | 176 | /* Bullet 2: */ 177 | /* test if the box intersects the plane of the triangle */ 178 | /* compute plane equation of triangle: normal*x+d=0 */ 179 | CROSS(normal,e0,e1); 180 | d=-DOT(normal,v0); /* plane eq: normal.x+d=0 */ 181 | if(!planeBoxOverlap(normal,d,boxhalfsize)) return 0; 182 | 183 | return 1; /* box and triangle overlaps */ 184 | } 185 | -------------------------------------------------------------------------------- /src/utils/libvoxelize/voxelize.cpython-36m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wzhen1/ALTO/2b89d81a89d65b72f0538a5318f6b490e2aa33fa/src/utils/libvoxelize/voxelize.cpython-36m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /src/utils/libvoxelize/voxelize.pyx: -------------------------------------------------------------------------------- 1 | cimport cython 2 | from libc.math cimport floor, ceil 3 | from cython.view cimport array as cvarray 4 | 5 | cdef extern from "tribox2.h": 6 | int triBoxOverlap(float boxcenter[3], float boxhalfsize[3], 7 | float tri0[3], float tri1[3], float tri2[3]) 8 | 9 | 10 | @cython.boundscheck(False) # Deactivate bounds checking 11 | @cython.wraparound(False) # Deactivate negative indexing. 12 | cpdef int voxelize_mesh_(bint[:, :, :] occ, float[:, :, ::1] faces): 13 | assert(faces.shape[1] == 3) 14 | assert(faces.shape[2] == 3) 15 | 16 | n_faces = faces.shape[0] 17 | cdef int i 18 | for i in range(n_faces): 19 | voxelize_triangle_(occ, faces[i]) 20 | 21 | 22 | @cython.boundscheck(False) # Deactivate bounds checking 23 | @cython.wraparound(False) # Deactivate negative indexing. 24 | cpdef int voxelize_triangle_(bint[:, :, :] occupancies, float[:, ::1] triverts): 25 | cdef int bbox_min[3] 26 | cdef int bbox_max[3] 27 | cdef int i, j, k 28 | cdef float boxhalfsize[3] 29 | cdef float boxcenter[3] 30 | cdef bint intersection 31 | 32 | boxhalfsize[:] = (0.5, 0.5, 0.5) 33 | 34 | for i in range(3): 35 | bbox_min[i] = ( 36 | min(triverts[0, i], triverts[1, i], triverts[2, i]) 37 | ) 38 | bbox_min[i] = min(max(bbox_min[i], 0), occupancies.shape[i] - 1) 39 | 40 | for i in range(3): 41 | bbox_max[i] = ( 42 | max(triverts[0, i], triverts[1, i], triverts[2, i]) 43 | ) 44 | bbox_max[i] = min(max(bbox_max[i], 0), occupancies.shape[i] - 1) 45 | 46 | for i in range(bbox_min[0], bbox_max[0] + 1): 47 | for j in range(bbox_min[1], bbox_max[1] + 1): 48 | for k in range(bbox_min[2], bbox_max[2] + 1): 49 | boxcenter[:] = (i + 0.5, j + 0.5, k + 0.5) 50 | intersection = triBoxOverlap(&boxcenter[0], &boxhalfsize[0], 51 | &triverts[0, 0], &triverts[1, 0], &triverts[2, 0]) 52 | occupancies[i, j, k] |= intersection 53 | 54 | 55 | @cython.boundscheck(False) # Deactivate bounds checking 56 | @cython.wraparound(False) # Deactivate negative indexing. 57 | cdef int test_triangle_aabb(float[::1] boxcenter, float[::1] boxhalfsize, float[:, ::1] triverts): 58 | assert(boxcenter.shape[0] == 3) 59 | assert(boxhalfsize.shape[0] == 3) 60 | assert(triverts.shape[0] == triverts.shape[1] == 3) 61 | 62 | # print(triverts) 63 | # Call functions 64 | cdef int result = triBoxOverlap(&boxcenter[0], &boxhalfsize[0], 65 | &triverts[0, 0], &triverts[1, 0], &triverts[2, 0]) 66 | return result 67 | -------------------------------------------------------------------------------- /src/utils/visualize.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from matplotlib import pyplot as plt 3 | from mpl_toolkits.mplot3d import Axes3D 4 | import src.common as common 5 | 6 | 7 | def visualize_data(data, data_type, out_file): 8 | r''' Visualizes the data with regard to its type. 9 | 10 | Args: 11 | data (tensor): batch of data 12 | data_type (string): data type (img, voxels or pointcloud) 13 | out_file (string): output file 14 | ''' 15 | if data_type == 'voxels': 16 | visualize_voxels(data, out_file=out_file) 17 | elif data_type == 'pointcloud': 18 | visualize_pointcloud(data, out_file=out_file) 19 | elif data_type is None or data_type == 'idx': 20 | pass 21 | else: 22 | raise ValueError('Invalid data_type "%s"' % data_type) 23 | 24 | 25 | def visualize_voxels(voxels, out_file=None, show=False): 26 | r''' Visualizes voxel data. 27 | 28 | Args: 29 | voxels (tensor): voxel data 30 | out_file (string): output file 31 | show (bool): whether the plot should be shown 32 | ''' 33 | # Use numpy 34 | voxels = np.asarray(voxels) 35 | # Create plot 36 | fig = plt.figure() 37 | ax = fig.gca(projection=Axes3D.name) 38 | voxels = voxels.transpose(2, 0, 1) 39 | ax.voxels(voxels, edgecolor='k') 40 | ax.set_xlabel('Z') 41 | ax.set_ylabel('X') 42 | ax.set_zlabel('Y') 43 | ax.view_init(elev=30, azim=45) 44 | if out_file is not None: 45 | plt.savefig(out_file) 46 | if show: 47 | plt.show() 48 | plt.close(fig) 49 | 50 | 51 | def visualize_pointcloud(points, normals=None, 52 | out_file=None, show=False): 53 | r''' Visualizes point cloud data. 54 | 55 | Args: 56 | points (tensor): point data 57 | normals (tensor): normal data (if existing) 58 | out_file (string): output file 59 | show (bool): whether the plot should be shown 60 | ''' 61 | # Use numpy 62 | points = np.asarray(points) 63 | # Create plot 64 | fig = plt.figure() 65 | ax = fig.gca(projection=Axes3D.name) 66 | ax.scatter(points[:, 2], points[:, 0], points[:, 1]) 67 | if normals is not None: 68 | ax.quiver( 69 | points[:, 2], points[:, 0], points[:, 1], 70 | normals[:, 2], normals[:, 0], normals[:, 1], 71 | length=0.1, color='k' 72 | ) 73 | ax.set_xlabel('Z') 74 | ax.set_ylabel('X') 75 | ax.set_zlabel('Y') 76 | ax.set_xlim(-0.5, 0.5) 77 | ax.set_ylim(-0.5, 0.5) 78 | ax.set_zlim(-0.5, 0.5) 79 | ax.view_init(elev=30, azim=45) 80 | if out_file is not None: 81 | plt.savefig(out_file) 82 | if show: 83 | plt.show() 84 | plt.close(fig) 85 | 86 | --------------------------------------------------------------------------------