├── .github └── workflows │ └── formatter.yml ├── .gitignore ├── LICENSE ├── README.md ├── assets ├── pipeline.png └── radar.png ├── configs ├── _base_ │ └── default_runtime.py ├── nuscenes │ ├── pretrain-ponder-spunet-v1m1-0-base-color-amp.py │ ├── pretrain-ponder-spunet-v1m1-0-base-color.py │ ├── pretrain-ponder-spunet-v1m1-0-base-semantic.py │ └── pretrain-ponder-spunet-v1m1-0-base.py ├── s3dis │ ├── pretrain-ponder-spunet-v1m1-0-base.py │ ├── semseg-ppt-v1m1-0-s3-sc-st-spunet-lovasz-ft.py │ └── semseg-spunet-v1m1-0-base.py ├── scannet │ ├── insseg-ppt-v1m1-0-pointgroup-spunet-ft.py │ ├── pretrain-ponder-ppt-v1m1-0-sc-s3-st-spunet.py │ ├── pretrain-ponder-spunet-v1m1-0-base.py │ ├── semseg-ppt-v1m1-0-sc-s3-st-spunet-lovasz-ft.py │ └── semseg-spunet-v1m1-0-base.py ├── scannet200 │ ├── insseg-ppt-v1m1-0-pointgroup-spunet-ft.py │ └── semseg-ppt-v1m1-0-spunet-lovasz-ft.py └── structured3d │ ├── pretrain-ponder-spunet-v1m1-0-base.py │ └── semseg-spunet-v1m1-0-base.py ├── data └── scannet │ └── skip.lst ├── docs ├── data_preparation.md ├── getting_started.md └── model_zoo.md ├── libs ├── pointgroup_ops │ ├── functions │ │ ├── __init__.py │ │ └── functions.py │ ├── setup.py │ └── src │ │ ├── bfs_cluster.cpp │ │ └── bfs_cluster_kernel.cu ├── pointops │ ├── __init__.py │ ├── functions │ │ ├── __init__.py │ │ ├── aggregation.py │ │ ├── attention.py │ │ ├── grouping.py │ │ ├── interpolation.py │ │ ├── query.py │ │ ├── sampling.py │ │ ├── subtraction.py │ │ └── utils.py │ ├── setup.py │ └── src │ │ ├── __init__.py │ │ ├── aggregation │ │ ├── aggregation_cuda.cpp │ │ ├── aggregation_cuda_kernel.cu │ │ └── aggregation_cuda_kernel.h │ │ ├── attention │ │ ├── attention_cuda.cpp │ │ ├── attention_cuda_kernel.cu │ │ └── attention_cuda_kernel.h │ │ ├── ball_query │ │ ├── ball_query_cuda.cpp │ │ ├── ball_query_cuda_kernel.cu │ │ └── ball_query_cuda_kernel.h │ │ ├── cuda_utils.h │ │ ├── grouping │ │ ├── grouping_cuda.cpp │ │ ├── grouping_cuda_kernel.cu │ │ └── grouping_cuda_kernel.h │ │ ├── interpolation │ │ ├── interpolation_cuda.cpp │ │ ├── interpolation_cuda_kernel.cu │ │ └── interpolation_cuda_kernel.h │ │ ├── knn_query │ │ ├── knn_query_cuda.cpp │ │ ├── knn_query_cuda_kernel.cu │ │ └── knn_query_cuda_kernel.h │ │ ├── pointops_api.cpp │ │ ├── random_ball_query │ │ ├── random_ball_query_cuda.cpp │ │ ├── random_ball_query_cuda_kernel.cu │ │ └── random_ball_query_cuda_kernel.h │ │ ├── sampling │ │ ├── sampling_cuda.cpp │ │ ├── sampling_cuda_kernel.cu │ │ └── sampling_cuda_kernel.h │ │ └── subtraction │ │ ├── subtraction_cuda.cpp │ │ ├── subtraction_cuda_kernel.cu │ │ └── subtraction_cuda_kernel.h └── smooth-sampler │ ├── .gitignore │ ├── LICENSE │ ├── README.md │ ├── setup.py │ └── smooth_sampler │ ├── __init__.py │ ├── csrc │ ├── smooth_sampler.cpp │ └── smooth_sampler_kernel.cu │ └── modules.py ├── ponder ├── __init__.py ├── datasets │ ├── __init__.py │ ├── builder.py │ ├── dataloader.py │ ├── defaults.py │ ├── nuscenes.py │ ├── preprocessing │ │ ├── nuscenes │ │ │ └── preprocess_nuscenes_info.py │ │ ├── s3dis │ │ │ ├── preprocess_s3dis.py │ │ │ └── preprocess_s3dis_voxelized.py │ │ ├── scannet │ │ │ ├── SensorData.py │ │ │ ├── meta_data │ │ │ │ ├── classes_ObjClassification-ShapeNetCore55.txt │ │ │ │ ├── classes_SemVoxLabel-nyu40id.txt │ │ │ │ ├── scannet200_constants.py │ │ │ │ ├── scannet200_splits.py │ │ │ │ ├── scannet_means.npz │ │ │ │ ├── scannetv1_test.txt │ │ │ │ ├── scannetv1_train.txt │ │ │ │ ├── scannetv1_val.txt │ │ │ │ ├── scannetv2-labels-old.combined.tsv │ │ │ │ ├── scannetv2-labels.combined.tsv │ │ │ │ ├── scannetv2_test.txt │ │ │ │ ├── scannetv2_train.txt │ │ │ │ └── scannetv2_val.txt │ │ │ ├── preprocess_scannet.py │ │ │ └── reader.py │ │ └── structured3d │ │ │ └── preprocess_structured3d.py │ ├── s3dis.py │ ├── scannet.py │ ├── structure3d.py │ ├── transform.py │ └── utils.py ├── engines │ ├── __init__.py │ ├── defaults.py │ ├── hooks │ │ ├── __init__.py │ │ ├── builder.py │ │ ├── default.py │ │ ├── evaluator.py │ │ └── misc.py │ ├── launch.py │ ├── test.py │ └── train.py ├── models │ ├── __init__.py │ ├── builder.py │ ├── default.py │ ├── losses │ │ ├── __init__.py │ │ ├── builder.py │ │ ├── lovasz.py │ │ └── misc.py │ ├── point_group │ │ ├── __init__.py │ │ ├── point_group_v1m1_base.py │ │ └── utils.py │ ├── point_prompt_training │ │ ├── __init__.py │ │ ├── point_prompt_training_v1m1_language_guided.py │ │ └── point_prompt_training_v1m2_decoupled.py │ ├── ponder │ │ ├── __init__.py │ │ ├── ponder_indoor_base.py │ │ ├── ponder_outdoor_base.py │ │ ├── render_utils │ │ │ ├── __init__.py │ │ │ ├── builder.py │ │ │ ├── decoders.py │ │ │ ├── fields │ │ │ │ ├── __init__.py │ │ │ │ └── sdf_field.py │ │ │ ├── models │ │ │ │ ├── __init__.py │ │ │ │ ├── base_surface_model.py │ │ │ │ ├── neus.py │ │ │ │ └── volsdf.py │ │ │ ├── ray_samplers.py │ │ │ ├── rays.py │ │ │ ├── renderers.py │ │ │ └── scene_colliders.py │ │ └── unet3d.py │ ├── sparse_unet │ │ ├── __init__.py │ │ ├── mink_unet.py │ │ ├── spconv_unet_v1m1_base.py │ │ ├── spconv_unet_v1m2_bn_momentum.py │ │ └── spconv_unet_v1m3_pdnorm.py │ └── utils.py └── utils │ ├── __init__.py │ ├── cache.py │ ├── comm.py │ ├── config.py │ ├── env.py │ ├── events.py │ ├── logger.py │ ├── misc.py │ ├── optimizer.py │ ├── path.py │ ├── registry.py │ ├── scheduler.py │ ├── timer.py │ └── visualization.py ├── scripts ├── test.sh └── train.sh └── tools ├── test.py ├── test_s3dis_6fold.py └── train.py /.github/workflows/formatter.yml: -------------------------------------------------------------------------------- 1 | name: Formatter 2 | 3 | on: 4 | workflow_dispatch: 5 | push: 6 | branches: 7 | - main 8 | pull_request: 9 | types: [opened, reopened, synchronize] 10 | 11 | concurrency: 12 | group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} 13 | cancel-in-progress: true 14 | 15 | jobs: 16 | formatter: 17 | runs-on: ubuntu-latest 18 | steps: 19 | - uses: actions/checkout@v3 20 | - uses: psf/black@stable 21 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | env/ 106 | .venv 107 | venv/ 108 | ENV/ 109 | env.bak/ 110 | venv.bak/ 111 | 112 | # Spyder project settings 113 | .spyderproject 114 | .spyproject 115 | 116 | # Rope project settings 117 | .ropeproject 118 | 119 | # mkdocs documentation 120 | /site 121 | 122 | # mypy 123 | .mypy_cache/ 124 | .dmypy.json 125 | dmypy.json 126 | 127 | # Pyre type checker 128 | .pyre/ 129 | 130 | ### VisualStudioCode 131 | .vscode/* 132 | !.vscode/settings.json 133 | !.vscode/tasks.json 134 | !.vscode/launch.json 135 | !.vscode/extensions.json 136 | *.code-workspace 137 | **/.vscode 138 | 139 | # JetBrains 140 | .idea/ 141 | 142 | # Data & Models 143 | *.h5 144 | *.tar 145 | *.tar.gz 146 | 147 | # Lightning-Hydra-Template 148 | configs/local/default.yaml 149 | /data/ 150 | /logs/ 151 | .env 152 | 153 | # Aim logging 154 | .aim 155 | 156 | # local 157 | exp/ 158 | data/ 159 | .cache/ 160 | data 161 | *.pth 162 | *.out -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 PonderV2 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 | -------------------------------------------------------------------------------- /assets/pipeline.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/OpenGVLab/PonderV2/c152af1df9d04a98bc19696e8d69d047834a80db/assets/pipeline.png -------------------------------------------------------------------------------- /assets/radar.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/OpenGVLab/PonderV2/c152af1df9d04a98bc19696e8d69d047834a80db/assets/radar.png -------------------------------------------------------------------------------- /configs/_base_/default_runtime.py: -------------------------------------------------------------------------------- 1 | weight = None # path to model weight 2 | resume = False # whether to resume training process 3 | evaluate = True # evaluate after each epoch training process 4 | test_only = False # test process 5 | 6 | seed = None # train process will init a random seed and record 7 | save_path = "exp/default" 8 | num_worker = 16 # total worker in all gpu 9 | batch_size = 16 # total batch size in all gpu 10 | batch_size_val = None # auto adapt to bs 1 for each gpu 11 | batch_size_test = None # auto adapt to bs 1 for each gpu 12 | epoch = 100 # total epoch, data loop = epoch // eval_epoch 13 | eval_epoch = 100 # sche total eval & checkpoint epoch 14 | 15 | sync_bn = False 16 | enable_amp = False 17 | empty_cache = False 18 | find_unused_parameters = False 19 | 20 | mix_prob = 0 21 | param_dicts = None # example: param_dicts = [dict(keyword="block", lr_scale=0.1)] 22 | 23 | # hook 24 | hooks = [ 25 | dict(type="CheckpointLoader"), 26 | dict(type="IterationTimer", warmup_iter=2), 27 | dict(type="InformationWriter"), 28 | dict(type="SemSegEvaluator"), 29 | dict(type="CheckpointSaver", save_freq=None), 30 | dict(type="PreciseEvaluator", test_last=False), 31 | ] 32 | 33 | # Trainer 34 | train = dict(type="DefaultTrainer") 35 | 36 | # Tester 37 | test = dict(type="SemSegTester", verbose=True) 38 | -------------------------------------------------------------------------------- /docs/data_preparation.md: -------------------------------------------------------------------------------- 1 | ## Data Preparation 2 | 3 | ### ScanNet v2 4 | 5 | The preprocessing support semantic and instance segmentation for both `ScanNet20`, `ScanNet200` and `ScanNet Data Efficient`. 6 | 7 | - Download the [ScanNet](http://www.scan-net.org/) v2 dataset. 8 | - Run preprocessing code for raw ScanNet as follows: 9 | ```bash 10 | # RAW_SCANNET_DIR: the directory of downloaded ScanNet v2 raw dataset. 11 | # PROCESSED_SCANNET_DIR: the directory of processed ScanNet dataset (output dir). 12 | python ponder/datasets/preprocessing/scannet/preprocess_scannet.py --dataset_root ${RAW_SCANNET_DIR} --output_root ${PROCESSED_SCANNET_DIR} 13 | # extract RGB-D iamges and 2D semantic labels (ScanNet20 only): 14 | python ponder/datasets/preprocessing/scannet/reader.py --scans_path ${RAW_SCANNET_DIR}/scans --output_path ${PROCESSED_SCANNET_DIR}/rgbd --export_depth_images --export_color_images --export_poses --export_intrinsics --export_label 15 | ``` 16 | 17 | - (Optional) Download ScanNet Data Efficient files: 18 | ```bash 19 | # download-scannet.py is the official download script 20 | # or follow instruction here: https://kaldir.vc.in.tum.de/scannet_benchmark/data_efficient/documentation#download 21 | python download-scannet.py --data_efficient -o ${RAW_SCANNET_DIR} 22 | # unzip downloads 23 | cd ${RAW_SCANNET_DIR}/tasks 24 | unzip limited-annotation-points.zip 25 | unzip limited-bboxes.zip 26 | unzip limited-reconstruction-scenes.zip 27 | # copy files to processed dataset folder 28 | cp -r ${RAW_SCANNET_DIR}/tasks ${PROCESSED_SCANNET_DIR} 29 | ``` 30 | 31 | - Link processed dataset to codebase: 32 | ```bash 33 | # PROCESSED_SCANNET_DIR: the directory of processed ScanNet dataset. 34 | mkdir data 35 | ln -s ${PROCESSED_SCANNET_DIR} ${CODEBASE_DIR}/data/scannet 36 | ``` 37 | 38 | ## S3DIS 39 | - Download S3DIS data by filling this [Google form](https://docs.google.com/forms/d/e/1FAIpQLScDimvNMCGhy_rmBA2gHfDu3naktRm6A8BPwAWWDv-Uhm6Shw/viewform?c=0&w=1). Download the `Stanford3dDataset_v1.2.zip` file and unzip it. 40 | - Run preprocessing code for S3DIS as follows: 41 | ```bash 42 | # S3DIS_DIR: the directory of downloaded Stanford3dDataset_v1.2 dataset. 43 | # RAW_S3DIS_DIR: the directory of Stanford2d3dDataset_noXYZ dataset. (optional, for parsing normal) 44 | # PROCESSED_S3DIS_DIR: the directory of processed S3DIS dataset (output dir). 45 | 46 | # S3DIS with normal vector, RGB-D images and 2D semantic labels 47 | python ponder/datasets/preprocessing/s3dis/preprocess_s3dis.py --dataset_root ${S3DIS_DIR} --output_root ${PROCESSED_S3DIS_DIR} --raw_root ${RAW_S3DIS_DIR} --parse_normal --parse_rgbd 48 | # if you want S3DIS with aligned angle: 49 | python ponder/datasets/preprocessing/s3dis/preprocess_s3dis.py --dataset_root ${S3DIS_DIR} --output_root ${PROCESSED_S3DIS_DIR} --raw_root ${RAW_S3DIS_DIR} --align_angle --parse_normal --parse_rgbd 50 | ``` 51 | - Link processed dataset to codebase. 52 | ```bash 53 | # PROCESSED_S3DIS_DIR: the directory of processed S3DIS dataset. 54 | mkdir data 55 | ln -s ${PROCESSED_S3DIS_DIR} ${CODEBASE_DIR}/data/s3dis 56 | ``` 57 | 58 | ## Structured3D 59 | - Download Structured3D panorama related and perspective (full) related zip files by filling this [Google form](https://docs.google.com/forms/d/e/1FAIpQLSc0qtvh4vHSoZaW6UvlXYy79MbcGdZfICjh4_t4bYofQIVIdw/viewform?pli=1) (no need to unzip them). 60 | - Organize all downloaded zip file in one folder (`${STRUCT3D_DIR}`). 61 | - Run preprocessing code for Structured3D as follows: 62 | ```bash 63 | # STRUCT3D_DIR: the directory of downloaded Structured3D dataset. 64 | # PROCESSED_STRUCT3D_DIR: the directory of processed Structured3D dataset (output dir). 65 | # NUM_WORKERS: Number for workers for preprocessing, default same as cpu count (might OOM). 66 | export PYTHONPATH=./ 67 | python ponder/datasets/preprocessing/structured3d/preprocess_structured3d.py --dataset_root ${STRUCT3D_DIR} --output_root ${PROCESSED_STRUCT3D_DIR} --num_workers ${NUM_WORKERS} --grid_size 0.01 --fuse_prsp --fuse_pano --parse_rgbd 68 | ``` 69 | 70 | Following the instruction of [Swin3D](https://arxiv.org/abs/2304.06906), we keep 25 categories with frequencies of more than 0.001, out of the original 40 categories. 71 | 72 | - Link processed dataset to codebase. 73 | ```bash 74 | # PROCESSED_STRUCT3D_DIR: the directory of processed Structured3D dataset (output dir). 75 | mkdir data 76 | ln -s ${PROCESSED_STRUCT3D_DIR} ${CODEBASE_DIR}/data/structured3d 77 | ``` 78 | 79 | ## nuScenes 80 | - Download the official [NuScene](https://www.nuscenes.org/nuscenes#download) dataset (with Lidar Segmentation) and organize the downloaded files as follows: 81 | ```bash 82 | NUSCENES_DIR 83 | │── samples 84 | │── sweeps 85 | │── lidarseg 86 | ... 87 | │── v1.0-trainval 88 | │── v1.0-test 89 | ``` 90 | 91 | - Run information preprocessing code (modified from OpenPCDet) for nuScenes as follows: 92 | ```bash 93 | # NUSCENES_DIR: the directory of downloaded nuScenes dataset. 94 | # PROCESSED_NUSCENES_DIR: the directory of processed nuScenes dataset (output dir). 95 | # MAX_SWEEPS: Max number of sweeps. Default: 10. 96 | pip install nuscenes-devkit pyquaternion 97 | python ponder/datasets/preprocessing/nuscenes/preprocess_nuscenes_info.py --dataset_root ${NUSCENES_DIR} --output_root ${PROCESSED_NUSCENES_DIR} --max_sweeps ${MAX_SWEEPS} --with_camera 98 | ``` 99 | 100 | - Link raw dataset to processed NuScene dataset folder: 101 | ```bash 102 | # NUSCENES_DIR: the directory of downloaded nuScenes dataset. 103 | # PROCESSED_NUSCENES_DIR: the directory of processed nuScenes dataset (output dir). 104 | ln -s ${NUSCENES_DIR} {PROCESSED_NUSCENES_DIR}/raw 105 | ``` 106 | 107 | then the processed nuscenes folder is organized as follows: 108 | ```bash 109 | nuscene 110 | |── raw 111 | │── samples 112 | │── sweeps 113 | │── lidarseg 114 | ... 115 | │── v1.0-trainval 116 | │── v1.0-test 117 | |── info 118 | ``` 119 | 120 | - Link processed dataset to codebase. 121 | ```bash 122 | # PROCESSED_NUSCENES_DIR: the directory of processed nuScenes dataset (output dir). 123 | mkdir data 124 | ln -s ${PROCESSED_NUSCENES_DIR} ${CODEBASE_DIR}/data/nuscenes 125 | ``` -------------------------------------------------------------------------------- /docs/getting_started.md: -------------------------------------------------------------------------------- 1 | ## Getting Started 2 | 3 | ### (Pre)Training 4 | 5 | **Train from scratch.** The training processing is based on configs in `configs` folder. 6 | The training script will generate an experiment folder in `exp` folder and backup essential code in the experiment folder. 7 | Training config, log, tensorboard and checkpoints will also be saved into the experiment folder during the training process. 8 | ```bash 9 | export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} 10 | # Script (Recommended) 11 | sh scripts/train.sh -p ${INTERPRETER_PATH} -g ${NUM_GPU} -d ${DATASET_NAME} -c ${CONFIG_NAME} -n ${EXP_NAME} 12 | # Direct 13 | export PYTHONPATH=./ 14 | python tools/train.py --config-file ${CONFIG_PATH} --num-gpus ${NUM_GPU} --options save_path=${SAVE_PATH} 15 | ``` 16 | 17 | For example: 18 | ```bash 19 | # By script (Recommended) 20 | # -p is default set as python and can be ignored 21 | sh scripts/train.sh -p python -d scannet -c semseg-spunet-v1m1-0-base -n semseg-spunet-v1m1-0-base 22 | # Direct 23 | export PYTHONPATH=./ 24 | python tools/train.py --config-file configs/scannet/semseg-spunet-v1m1-0-base.py --options save_path=exp/scannet/semseg-spunet-v1m1-0-base 25 | ``` 26 | **Resume training from checkpoint.** If the training process is interrupted by accident, the following script can resume training from a given checkpoint. 27 | ```bash 28 | export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} 29 | # Script (Recommended) 30 | # simply add "-r true" 31 | sh scripts/train.sh -p ${INTERPRETER_PATH} -g ${NUM_GPU} -d ${DATASET_NAME} -c ${CONFIG_NAME} -n ${EXP_NAME} -r true 32 | # Direct 33 | export PYTHONPATH=./ 34 | python tools/train.py --config-file ${CONFIG_PATH} --num-gpus ${NUM_GPU} --options save_path=${SAVE_PATH} resume=True weight=${CHECKPOINT_PATH} 35 | ``` 36 | 37 | ### Testing 38 | During training, model evaluation is performed on point clouds after grid sampling (voxelization), providing an initial assessment of model performance. However, to obtain precise evaluation results, testing is **essential**. The testing process involves subsampling a dense point cloud into a sequence of voxelized point clouds, ensuring comprehensive coverage of all points. These sub-results are then predicted and collected to form a complete prediction of the entire point cloud. This approach yields higher evaluation results compared to simply mapping/interpolating the prediction. In addition, our testing code supports TTA (test time augmentation) testing, which further enhances the stability of evaluation performance. 39 | 40 | ```bash 41 | # By script (Based on experiment folder created by training script) 42 | sh scripts/test.sh -p ${INTERPRETER_PATH} -g ${NUM_GPU} -d ${DATASET_NAME} -n ${EXP_NAME} -w ${CHECKPOINT_NAME} 43 | # Direct 44 | export PYTHONPATH=./ 45 | python tools/test.py --config-file ${CONFIG_PATH} --num-gpus ${NUM_GPU} --options save_path=${SAVE_PATH} weight=${CHECKPOINT_PATH} 46 | ``` 47 | For example: 48 | ```bash 49 | # By script (Based on experiment folder created by training script) 50 | # -p is default set as python and can be ignored 51 | # -w is default set as model_best and can be ignored 52 | sh scripts/test.sh -p python -d scannet -n semseg-spunet-v1m1-0-base -w model_best 53 | # Direct 54 | export PYTHONPATH=./ 55 | python tools/test.py --config-file configs/scannet/semseg-spunet-v1m1-0-base.py --options save_path=exp/scannet/semseg-spunet-v1m1-0-base weight=exp/scannet/semseg-spunet-v1m1-0-base/model/model_best.pth 56 | ``` 57 | 58 | The TTA can be disabled by replace `data.test.test_cfg.aug_transform = [...]` with: 59 | 60 | ```python 61 | data = dict( 62 | train = dict(...), 63 | val = dict(...), 64 | test = dict( 65 | ..., 66 | test_cfg = dict( 67 | ..., 68 | aug_transform = [ 69 | [dict(type="RandomRotateTargetAngle", angle=[0], axis="z", center=[0, 0, 0], p=1)] 70 | ] 71 | ) 72 | ) 73 | ) 74 | ``` -------------------------------------------------------------------------------- /docs/model_zoo.md: -------------------------------------------------------------------------------- 1 | ## Model Zoo 2 | 3 | *For outdoor models, please check out [UniPAD](https://github.com/Nightmare-n/UniPAD)!* 4 | 5 | ### Pretraining 6 | 7 | Model | Backbone | Google Drive | Baidu Pan | Config 8 | ----- | ----- | -----| ----- | ----- 9 | [PonderIndoor-v2](../ponder/models/ponder/ponder_indoor_base.py) | [SpUNet-v1m3](../ponder/models/sparse_unet/spconv_unet_v1m3_pdnorm.py) | [ckpt](https://drive.google.com/file/d/1oFHhr6YPZwgUCtn0y9M6zAHL7XNkmpVM/view?usp=sharing) | [ckpt](https://pan.baidu.com/s/1mF5BdjvS2DDjFrqntm-kYA?pwd=soin) | [cfg](../configs/scannet/pretrain-ponder-ppt-v1m1-0-sc-s3-st-spunet.py) 10 | 11 | ### Finetuning 12 | #### Indoor Semantic Segmentation 13 | 14 | Benchmark | Model | Backbone | Val mIoU | Test mIoU | Google Drive | Baidu Pan | Config 15 | ----- | ----- | ----- | ----- | ----- | ----- | ----- | ----- 16 | ScanNet | [PPT-v1m1](../ponder/models/point_prompt_training/point_prompt_training_v1m1_language_guided.py) | [SpUNet-v1m3](../ponder/models/sparse_unet/spconv_unet_v1m3_pdnorm.py) | 77.0 | 78.5 | [ckpt](https://drive.google.com/file/d/16RhUSJDxtsS7Z_FeRJwk0L3q7Dt-LN1A/view?usp=sharing) | [ckpt](https://pan.baidu.com/s/1l0_k0h9fvnI38By6bSdijQ?pwd=wks7) | [cfg](../configs/scannet/semseg-ppt-v1m1-0-sc-s3-st-spunet-lovasz-ft.py) 17 | ScanNet200 | [PPT-v1m1](../ponder/models/point_prompt_training/point_prompt_training_v1m1_language_guided.py) | [SpUNet-v1m3](../ponder/models/sparse_unet/spconv_unet_v1m3_pdnorm.py) | 32.3 | 34.6 | [ckpt](https://drive.google.com/file/d/1d_we6SsNJLDeRc1LepZwcOCgOGJy11rr/view?usp=sharing) | [ckpt](https://pan.baidu.com/s/1fvH5wA60wl2In0BaUMuFOw?pwd=3ron) | [cfg](../configs/scannet200/semseg-ppt-v1m1-0-spunet-lovasz-ft.py) 18 | S3DIS (Area5) | [PPT-v1m1](../ponder/models/point_prompt_training/point_prompt_training_v1m1_language_guided.py) | [SpUNet-v1m3](../ponder/models/sparse_unet/spconv_unet_v1m3_pdnorm.py) | 73.2 | 79.1 | [ckpt](https://drive.google.com/file/d/1GZgfxWJC9hNEHKV30t3PuIZMmTXngHXg/view?usp=sharing) | [ckpt](https://pan.baidu.com/s/1KJ-PwvROofcGeTzkKdsXCQ?pwd=bbaa) | [cfg](../configs/s3dis/semseg-ppt-v1m1-0-s3-sc-st-spunet-lovasz-ft.py) 19 | 20 | #### Indoor Instance Segmentation 21 | Benchmark | Model | Backbone | mAP@25 | mAP@50 | mAP | Google Drive | Baidu Pan | Config 22 | ----- | ----- | ----- | ----- | ----- | ----- | ----- | ----- | ----- 23 | ScanNet | [PPT-v1m1](../ponder/models/point_prompt_training/point_prompt_training_v1m1_language_guided.py) | [SpUNet-v1m3](../ponder/models/sparse_unet/spconv_unet_v1m3_pdnorm.py) | 77.0 | 62.6 | 40.9 | [ckpt](https://drive.google.com/file/d/15tjsGY6bgZiQSXJel7yywzdYBDrpystk/view?usp=sharing) | [ckpt](https://pan.baidu.com/s/10BifGbWQ6CW_FcAnaw-XQg?pwd=jmd9) | [cfg](../configs/scannet/insseg-ppt-v1m1-0-pointgroup-spunet-ft.py) 24 | ScanNet200 | [PPT-v1m1](../ponder/models/point_prompt_training/point_prompt_training_v1m1_language_guided.py) | [SpUNet-v1m3](../ponder/models/sparse_unet/spconv_unet_v1m3_pdnorm.py) | 37.6| 30.5 | 20.1 | [ckpt](https://drive.google.com/file/d/1MVC2xSgXqbFDzIlni1KyPT28thbAIwm6/view?usp=sharing) | [ckpt](https://pan.baidu.com/s/1MbCPJEWbgOOmoB3riCEfcg?pwd=6pm0) | [cfg](../configs/scannet200/insseg-ppt-v1m1-0-pointgroup-spunet-ft.py) 25 | 26 | -------------------------------------------------------------------------------- /libs/pointgroup_ops/functions/__init__.py: -------------------------------------------------------------------------------- 1 | from .functions import Clustering, ballquery_batch_p, bfs_cluster 2 | -------------------------------------------------------------------------------- /libs/pointgroup_ops/setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | from distutils.sysconfig import get_config_vars 3 | from sys import argv 4 | 5 | from setuptools import setup 6 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 7 | 8 | (opt,) = get_config_vars("OPT") 9 | os.environ["OPT"] = " ".join( 10 | flag for flag in opt.split() if flag != "-Wstrict-prototypes" 11 | ) 12 | 13 | 14 | def _argparse(pattern, argv, is_flag=True, is_list=False): 15 | if is_flag: 16 | found = pattern in argv 17 | if found: 18 | argv.remove(pattern) 19 | return found, argv 20 | else: 21 | arr = [arg for arg in argv if pattern == arg.split("=")[0]] 22 | if is_list: 23 | if len(arr) == 0: # not found 24 | return False, argv 25 | else: 26 | assert "=" in arr[0], f"{arr[0]} requires a value." 27 | argv.remove(arr[0]) 28 | val = arr[0].split("=")[1] 29 | if "," in val: 30 | return val.split(","), argv 31 | else: 32 | return [val], argv 33 | else: 34 | if len(arr) == 0: # not found 35 | return False, argv 36 | else: 37 | assert "=" in arr[0], f"{arr[0]} requires a value." 38 | argv.remove(arr[0]) 39 | return arr[0].split("=")[1], argv 40 | 41 | 42 | INCLUDE_DIRS, argv = _argparse("--include_dirs", argv, False, is_list=True) 43 | include_dirs = [] 44 | if not (INCLUDE_DIRS is False): 45 | include_dirs += INCLUDE_DIRS 46 | 47 | setup( 48 | name="pointgroup_ops", 49 | packages=["pointgroup_ops"], 50 | package_dir={"pointgroup_ops": "functions"}, 51 | ext_modules=[ 52 | CUDAExtension( 53 | name="pointgroup_ops_cuda", 54 | sources=["src/bfs_cluster.cpp", "src/bfs_cluster_kernel.cu"], 55 | extra_compile_args={"cxx": ["-g"], "nvcc": ["-O2"]}, 56 | ) 57 | ], 58 | include_dirs=[*include_dirs], 59 | cmdclass={"build_ext": BuildExtension}, 60 | ) 61 | -------------------------------------------------------------------------------- /libs/pointgroup_ops/src/bfs_cluster.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | Ball Query with BatchIdx & Clustering Algorithm 3 | Written by Li Jiang 4 | All Rights Reserved 2020. 5 | */ 6 | 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | #include 16 | 17 | int ballquery_batch_p_cuda(int n, int meanActive, float radius, const float *xyz, const int *batch_idxs, const int *batch_offsets, int *idx, int *start_len, cudaStream_t stream); 18 | 19 | 20 | using Int = int32_t; 21 | class ConnectedComponent{ 22 | public: 23 | std::vector pt_idxs {}; 24 | 25 | ConnectedComponent(){}; 26 | void addPoint(Int pt_idx) 27 | { 28 | pt_idxs.push_back(pt_idx); 29 | 30 | } 31 | }; 32 | using ConnectedComponents = std::vector; 33 | 34 | /* ================================== ballquery_batch_p ================================== */ 35 | // input xyz: (n, 3) float 36 | // input batch_idxs: (n) int 37 | // input batch_offsets: (B+1) int, batch_offsets[-1] 38 | // output idx: (n * meanActive) dim 0 for number of points in the ball, idx in n 39 | // output start_len: (n, 2), int 40 | int ballquery_batch_p(at::Tensor xyz_tensor, at::Tensor batch_idxs_tensor, at::Tensor batch_offsets_tensor, at::Tensor idx_tensor, at::Tensor start_len_tensor, int n, int meanActive, float radius){ 41 | const float *xyz = xyz_tensor.data(); 42 | const int *batch_idxs = batch_idxs_tensor.data(); 43 | const int *batch_offsets = batch_offsets_tensor.data(); 44 | int *idx = idx_tensor.data(); 45 | int *start_len = start_len_tensor.data(); 46 | 47 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 48 | int cumsum = ballquery_batch_p_cuda(n, meanActive, radius, xyz, batch_idxs, batch_offsets, idx, start_len, stream); 49 | return cumsum; 50 | } 51 | 52 | /* ================================== bfs_cluster ================================== */ 53 | ConnectedComponent find_cc(Int idx, int *semantic_label, Int *ball_query_idxs, int *start_len, int *visited){ 54 | ConnectedComponent cc; 55 | cc.addPoint(idx); 56 | visited[idx] = 1; 57 | 58 | std::queue Q; 59 | assert(Q.empty()); 60 | Q.push(idx); 61 | 62 | while(!Q.empty()){ 63 | Int cur = Q.front(); Q.pop(); 64 | int start = start_len[cur * 2]; 65 | int len = start_len[cur * 2 + 1]; 66 | int label_cur = semantic_label[cur]; 67 | for(Int i = start; i < start + len; i++){ 68 | Int idx_i = ball_query_idxs[i]; 69 | if(semantic_label[idx_i] != label_cur) continue; 70 | if(visited[idx_i] == 1) continue; 71 | 72 | cc.addPoint(idx_i); 73 | visited[idx_i] = 1; 74 | 75 | Q.push(idx_i); 76 | } 77 | } 78 | return cc; 79 | } 80 | 81 | //input: semantic_label, int, N 82 | //input: ball_query_idxs, Int, (nActive) 83 | //input: start_len, int, (N, 2) 84 | //output: clusters, CCs 85 | int get_clusters(int *semantic_label, Int *ball_query_idxs, int *start_len, const Int nPoint, int threshold, ConnectedComponents &clusters){ 86 | int visited[nPoint] = {0}; 87 | 88 | int sumNPoint = 0; 89 | for(Int i = 0; i < nPoint; i++){ 90 | if(visited[i] == 0){ 91 | ConnectedComponent CC = find_cc(i, semantic_label, ball_query_idxs, start_len, visited); 92 | if((int)CC.pt_idxs.size() >= threshold){ 93 | clusters.push_back(CC); 94 | sumNPoint += (int)CC.pt_idxs.size(); 95 | } 96 | } 97 | } 98 | 99 | return sumNPoint; 100 | } 101 | 102 | void fill_cluster_idxs_(ConnectedComponents &CCs, int *cluster_idxs, int *cluster_offsets){ 103 | for(int i = 0; i < (int)CCs.size(); i++){ 104 | cluster_offsets[i + 1] = cluster_offsets[i] + (int)CCs[i].pt_idxs.size(); 105 | for(int j = 0; j < (int)CCs[i].pt_idxs.size(); j++){ 106 | int idx = CCs[i].pt_idxs[j]; 107 | cluster_idxs[(cluster_offsets[i] + j) * 2 + 0] = i; 108 | cluster_idxs[(cluster_offsets[i] + j) * 2 + 1] = idx; 109 | } 110 | } 111 | } 112 | 113 | //input: semantic_label, int, N 114 | //input: ball_query_idxs, int, (nActive) 115 | //input: start_len, int, (N, 2) 116 | //output: cluster_idxs, int (sumNPoint, 2), dim 0 for cluster_id, dim 1 for corresponding point idxs in N 117 | //output: cluster_offsets, int (nCluster + 1) 118 | void bfs_cluster(at::Tensor semantic_label_tensor, at::Tensor ball_query_idxs_tensor, at::Tensor start_len_tensor, 119 | at::Tensor cluster_idxs_tensor, at::Tensor cluster_offsets_tensor, const int N, int threshold){ 120 | int *semantic_label = semantic_label_tensor.data(); 121 | Int *ball_query_idxs = ball_query_idxs_tensor.data(); 122 | int *start_len = start_len_tensor.data(); 123 | 124 | ConnectedComponents CCs; 125 | int sumNPoint = get_clusters(semantic_label, ball_query_idxs, start_len, N, threshold, CCs); 126 | 127 | int nCluster = (int)CCs.size(); 128 | cluster_idxs_tensor.resize_({sumNPoint, 2}); 129 | cluster_offsets_tensor.resize_({nCluster + 1}); 130 | cluster_idxs_tensor.zero_(); 131 | cluster_offsets_tensor.zero_(); 132 | 133 | int *cluster_idxs = cluster_idxs_tensor.data(); 134 | int *cluster_offsets = cluster_offsets_tensor.data(); 135 | 136 | fill_cluster_idxs_(CCs, cluster_idxs, cluster_offsets); 137 | } 138 | 139 | //------------------------------------API------------------------------------------ 140 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){ 141 | 142 | m.def("ballquery_batch_p", &ballquery_batch_p, "ballquery_batch_p"); 143 | m.def("bfs_cluster", &bfs_cluster, "bfs_cluster"); 144 | 145 | } 146 | -------------------------------------------------------------------------------- /libs/pointgroup_ops/src/bfs_cluster_kernel.cu: -------------------------------------------------------------------------------- 1 | /* 2 | Ball Query with BatchIdx 3 | Written by Li Jiang 4 | All Rights Reserved 2020. 5 | */ 6 | #include 7 | #include 8 | #include 9 | 10 | #define TOTAL_THREADS 1024 11 | #define THREADS_PER_BLOCK 512 12 | #define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0)) 13 | 14 | 15 | /* ================================== ballquery_batch_p ================================== */ 16 | __global__ void ballquery_batch_p_cuda_(int n, int meanActive, float radius, const float *xyz, const int *batch_idxs, const int *batch_offsets, int *idx, int *start_len, int *cumsum) { 17 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 18 | if (pt_idx >= n) return; 19 | 20 | start_len += (pt_idx * 2); 21 | int idx_temp[1000]; 22 | 23 | float radius2 = radius * radius; 24 | float o_x = xyz[pt_idx * 3 + 0]; 25 | float o_y = xyz[pt_idx * 3 + 1]; 26 | float o_z = xyz[pt_idx * 3 + 2]; 27 | 28 | int batch_idx = batch_idxs[pt_idx]; 29 | int start = batch_offsets[batch_idx]; 30 | int end = batch_offsets[batch_idx + 1]; 31 | 32 | int cnt = 0; 33 | for(int k = start; k < end; k++){ 34 | float x = xyz[k * 3 + 0]; 35 | float y = xyz[k * 3 + 1]; 36 | float z = xyz[k * 3 + 2]; 37 | float d2 = (o_x - x) * (o_x - x) + (o_y - y) * (o_y - y) + (o_z - z) * (o_z - z); 38 | if(d2 < radius2){ 39 | if(cnt < 1000){ 40 | idx_temp[cnt] = k; 41 | } 42 | else{ 43 | break; 44 | } 45 | ++cnt; 46 | } 47 | } 48 | 49 | start_len[0] = atomicAdd(cumsum, cnt); 50 | start_len[1] = cnt; 51 | 52 | int thre = n * meanActive; 53 | if(start_len[0] >= thre) return; 54 | 55 | idx += start_len[0]; 56 | if(start_len[0] + cnt >= thre) cnt = thre - start_len[0]; 57 | 58 | for(int k = 0; k < cnt; k++){ 59 | idx[k] = idx_temp[k]; 60 | } 61 | } 62 | 63 | 64 | int ballquery_batch_p_cuda(int n, int meanActive, float radius, const float *xyz, const int *batch_idxs, const int *batch_offsets, int *idx, int *start_len, cudaStream_t stream) { 65 | // param xyz: (n, 3) 66 | // param batch_idxs: (n) 67 | // param batch_offsets: (B + 1) 68 | // output idx: (n * meanActive) dim 0 for number of points in the ball, idx in n 69 | // output start_len: (n, 2), int 70 | 71 | cudaError_t err; 72 | 73 | dim3 blocks(DIVUP(n, THREADS_PER_BLOCK)); 74 | dim3 threads(THREADS_PER_BLOCK); 75 | 76 | int cumsum = 0; 77 | int* p_cumsum; 78 | cudaMalloc((void**)&p_cumsum, sizeof(int)); 79 | cudaMemcpy(p_cumsum, &cumsum, sizeof(int), cudaMemcpyHostToDevice); 80 | 81 | ballquery_batch_p_cuda_<<>>(n, meanActive, radius, xyz, batch_idxs, batch_offsets, idx, start_len, p_cumsum); 82 | 83 | err = cudaGetLastError(); 84 | if (cudaSuccess != err) { 85 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 86 | exit(-1); 87 | } 88 | 89 | cudaMemcpy(&cumsum, p_cumsum, sizeof(int), cudaMemcpyDeviceToHost); 90 | return cumsum; 91 | } 92 | -------------------------------------------------------------------------------- /libs/pointops/__init__.py: -------------------------------------------------------------------------------- 1 | from .functions import * 2 | -------------------------------------------------------------------------------- /libs/pointops/functions/__init__.py: -------------------------------------------------------------------------------- 1 | from .aggregation import aggregation 2 | from .attention import attention_fusion_step, attention_relation_step 3 | from .grouping import grouping, grouping2 4 | from .interpolation import interpolation, interpolation2 5 | from .query import ball_query, knn_query, random_ball_query 6 | from .sampling import farthest_point_sampling 7 | from .subtraction import subtraction 8 | from .utils import ( 9 | ball_query_and_group, 10 | batch2offset, 11 | knn_query_and_group, 12 | offset2batch, 13 | query_and_group, 14 | ) 15 | -------------------------------------------------------------------------------- /libs/pointops/functions/aggregation.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from pointops._C import aggregation_backward_cuda, aggregation_forward_cuda 3 | from torch.autograd import Function 4 | 5 | 6 | class Aggregation(Function): 7 | @staticmethod 8 | def forward(ctx, input, position, weight, idx): 9 | """ 10 | input: input: (n, c), position: (n, nsample, c), weight : (n, nsample, c'), idx: (n, nsample) 11 | output: (n, c) 12 | """ 13 | assert ( 14 | input.is_contiguous() 15 | and position.is_contiguous() 16 | and weight.is_contiguous() 17 | ) 18 | n, nsample, c = position.shape 19 | w_c = weight.shape[-1] 20 | output = torch.cuda.FloatTensor(n, c).zero_() 21 | aggregation_forward_cuda( 22 | n, nsample, c, w_c, input, position, weight, idx, output 23 | ) 24 | ctx.save_for_backward(input, position, weight, idx) 25 | return output 26 | 27 | @staticmethod 28 | def backward(ctx, grad_output): 29 | """ 30 | input: grad_out: (n, c) 31 | output: grad_input: (n, c), grad_position: (n, nsample, c), grad_weight : (n, nsample, c') 32 | """ 33 | input, position, weight, idx = ctx.saved_tensors 34 | n, nsample, c = position.shape 35 | w_c = weight.shape[-1] 36 | grad_input = torch.cuda.FloatTensor(n, c).zero_() 37 | grad_position = torch.cuda.FloatTensor(n, nsample, c).zero_() 38 | grad_weight = torch.cuda.FloatTensor(n, nsample, w_c).zero_() 39 | aggregation_backward_cuda( 40 | n, 41 | nsample, 42 | c, 43 | w_c, 44 | input, 45 | position, 46 | weight, 47 | idx, 48 | grad_output, 49 | grad_input, 50 | grad_position, 51 | grad_weight, 52 | ) 53 | return grad_input, grad_position, grad_weight, None 54 | 55 | 56 | aggregation = Aggregation.apply 57 | -------------------------------------------------------------------------------- /libs/pointops/functions/attention.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from pointops._C import ( 3 | attention_fusion_step_backward_cuda, 4 | attention_fusion_step_forward_cuda, 5 | attention_relation_step_backward_cuda, 6 | attention_relation_step_forward_cuda, 7 | ) 8 | from torch.autograd import Function 9 | 10 | 11 | class AttentionRelationStep(Function): 12 | @staticmethod 13 | def forward(ctx, query, key, weight, index_target, index_refer): 14 | """ 15 | input - query: (n, g, c), key: (n, g, c), weight: (c) 1_c for scatter attention, 16 | index_target: (m), index_refer: (m) 17 | output - relation: (M, g) 18 | """ 19 | 20 | assert ( 21 | query.is_contiguous() 22 | and key.is_contiguous() 23 | and index_target.is_contiguous() 24 | and index_refer.is_contiguous() 25 | and weight.is_contiguous() 26 | ) 27 | 28 | assert index_target.shape[0] == index_refer.shape[0] 29 | 30 | _, g, c = query.shape 31 | m = index_target.shape[0] 32 | output = torch.cuda.FloatTensor(m, g).zero_() 33 | attention_relation_step_forward_cuda( 34 | m, g, c, query, key, weight, index_target.int(), index_refer.int(), output 35 | ) 36 | ctx.save_for_backward(query, key, weight, index_target, index_refer) 37 | return output 38 | 39 | @staticmethod 40 | def backward(ctx, grad_output): 41 | query, key, weight, index_target, index_refer = ctx.saved_tensors 42 | n, g, c = query.shape 43 | m = index_target.shape[0] 44 | grad_query = torch.cuda.FloatTensor(n, g, c).zero_() 45 | grad_key = torch.cuda.FloatTensor(n, g, c).zero_() 46 | grad_weight = torch.cuda.FloatTensor(c).zero_() 47 | attention_relation_step_backward_cuda( 48 | m, 49 | g, 50 | c, 51 | query, 52 | grad_query, 53 | key, 54 | grad_key, 55 | weight, 56 | grad_weight, 57 | index_target.int(), 58 | index_refer.int(), 59 | grad_output, 60 | ) 61 | return grad_query, grad_key, None, None, None 62 | 63 | 64 | class AttentionFusionStep(Function): 65 | @staticmethod 66 | def forward(ctx, weight, value, index_target, index_refer): 67 | """ 68 | input - weight: (m, g), value: (n, g, c) 69 | index_target: (m), index_value: (m) 70 | output - output: (n, g, c) 71 | """ 72 | 73 | assert ( 74 | weight.is_contiguous() 75 | and value.is_contiguous() 76 | and index_target.is_contiguous() 77 | and index_refer.is_contiguous() 78 | and weight.is_contiguous() 79 | ) 80 | 81 | assert index_target.shape[0] == index_refer.shape[0] 82 | 83 | n, g, c = value.shape 84 | m = index_refer.shape[0] 85 | output = torch.cuda.FloatTensor(n, g, c).zero_() 86 | attention_fusion_step_forward_cuda( 87 | m, g, c, weight, value, index_target.int(), index_refer.int(), output 88 | ) 89 | ctx.save_for_backward(weight, value, index_target, index_refer) 90 | return output 91 | 92 | @staticmethod 93 | def backward(ctx, grad_output): 94 | """ 95 | input: grad_output: (n, g, c) 96 | output: grad_weight: (m, g), grad_value: (n, g, c), none, none 97 | """ 98 | weight, value, index_target, index_refer = ctx.saved_tensors 99 | n, g, c = value.shape 100 | m = index_target.shape[0] 101 | grad_weight = torch.cuda.FloatTensor(m, g).zero_() 102 | grad_value = torch.cuda.FloatTensor(n, g, c).zero_() 103 | attention_fusion_step_backward_cuda( 104 | m, 105 | g, 106 | c, 107 | weight, 108 | grad_weight, 109 | value, 110 | grad_value, 111 | index_target.int(), 112 | index_refer.int(), 113 | grad_output, 114 | ) 115 | return grad_weight, grad_value, None, None 116 | 117 | 118 | attention_relation_step = AttentionRelationStep.apply 119 | attention_fusion_step = AttentionFusionStep.apply 120 | -------------------------------------------------------------------------------- /libs/pointops/functions/grouping.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from pointops._C import grouping_backward_cuda, grouping_forward_cuda 3 | from torch.autograd import Function 4 | 5 | 6 | class Grouping(Function): 7 | @staticmethod 8 | def forward(ctx, input, idx): 9 | """ 10 | input: input: (n, c), idx : (m, nsample) 11 | output: (m, nsample, c) 12 | """ 13 | assert input.is_contiguous() and idx.is_contiguous() 14 | m, nsample, n, c = idx.shape[0], idx.shape[1], input.shape[0], input.shape[1] 15 | output = torch.cuda.FloatTensor(m, nsample, c) 16 | grouping_forward_cuda(m, nsample, c, input, idx, output) 17 | ctx.n = n 18 | ctx.save_for_backward(idx) 19 | return output 20 | 21 | @staticmethod 22 | def backward(ctx, grad_output): 23 | """ 24 | input: grad_out: (m, c, nsample) 25 | output: (n, c), None 26 | """ 27 | n = ctx.n 28 | (idx,) = ctx.saved_tensors 29 | m, nsample, c = grad_output.shape 30 | grad_input = torch.cuda.FloatTensor(n, c).zero_() 31 | grouping_backward_cuda(m, nsample, c, grad_output, idx, grad_input) 32 | return grad_input, None 33 | 34 | 35 | def grouping(idx, feat, xyz, new_xyz=None, with_xyz=False): 36 | if new_xyz is None: 37 | new_xyz = xyz 38 | assert xyz.is_contiguous() and feat.is_contiguous() 39 | m, nsample, c = idx.shape[0], idx.shape[1], feat.shape[1] 40 | xyz = torch.cat([xyz, torch.zeros([1, 3]).to(xyz.device)], dim=0) 41 | feat = torch.cat([feat, torch.zeros([1, c]).to(feat.device)], dim=0) 42 | grouped_feat = feat[idx.view(-1).long(), :].view( 43 | m, nsample, c 44 | ) # (m, num_sample, c) 45 | 46 | if with_xyz: 47 | assert new_xyz.is_contiguous() 48 | mask = torch.sign(idx + 1) 49 | grouped_xyz = xyz[idx.view(-1).long(), :].view( 50 | m, nsample, 3 51 | ) - new_xyz.unsqueeze( 52 | 1 53 | ) # (m, num_sample, 3) 54 | grouped_xyz = torch.einsum( 55 | "n s c, n s -> n s c", grouped_xyz, mask 56 | ) # (m, num_sample, 3) 57 | return torch.cat((grouped_xyz, grouped_feat), -1) 58 | else: 59 | return grouped_feat 60 | 61 | 62 | grouping2 = Grouping.apply 63 | -------------------------------------------------------------------------------- /libs/pointops/functions/interpolation.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from pointops._C import interpolation_backward_cuda, interpolation_forward_cuda 3 | from torch.autograd import Function 4 | 5 | from .query import knn_query 6 | 7 | 8 | def interpolation(xyz, new_xyz, feat, offset, new_offset, k=3): 9 | """ 10 | input: coords: (m, 3), new_xyz: (n, 3), color: (m, c), offset: (b), new_offset: (b) 11 | output: (n, c) 12 | """ 13 | assert xyz.is_contiguous() and new_xyz.is_contiguous() and feat.is_contiguous() 14 | idx, dist = knn_query(k, xyz, offset, new_xyz, new_offset) # (n, 3), (n, 3) 15 | dist_recip = 1.0 / (dist + 1e-8) # (n, 3) 16 | norm = torch.sum(dist_recip, dim=1, keepdim=True) 17 | weight = dist_recip / norm # (n, 3) 18 | 19 | new_feat = torch.cuda.FloatTensor(new_xyz.shape[0], feat.shape[1]).zero_() 20 | for i in range(k): 21 | new_feat += feat[idx[:, i].long(), :] * weight[:, i].unsqueeze(-1) 22 | return new_feat 23 | 24 | 25 | class Interpolation(Function): 26 | @staticmethod 27 | def forward(ctx, xyz, new_xyz, input, offset, new_offset, k=3): 28 | """ 29 | input: coords: (m, 3), new_xyz: (n, 3), input: (m, c), offset: (b), new_offset: (b) 30 | output: (n, c) 31 | """ 32 | assert xyz.is_contiguous() and new_xyz.is_contiguous() and input.is_contiguous() 33 | idx, dist = knn_query(k, xyz, offset, new_xyz, new_offset) # (n, k), (n, k) 34 | dist_recip = 1.0 / (dist + 1e-8) # (n, k) 35 | norm = torch.sum(dist_recip, dim=1, keepdim=True) 36 | weight = dist_recip / norm # (n, k) 37 | 38 | n, c, m = new_xyz.shape[0], input.shape[1], input.shape[0] 39 | output = torch.cuda.FloatTensor(n, c).zero_() 40 | interpolation_forward_cuda(n, c, k, input, idx, weight, output) 41 | ctx.m, ctx.k = m, k 42 | ctx.save_for_backward(idx, weight) 43 | return output 44 | 45 | @staticmethod 46 | def backward(ctx, grad_output): 47 | """ 48 | input: coords: (m, 3), new_xyz: (n, 3), input: (m, c), offset: (b), new_offset: (b) 49 | output: (n, c) 50 | """ 51 | m, k = ctx.m, ctx.k 52 | idx, weight = ctx.saved_tensors 53 | n, c = grad_output.shape 54 | grad_input = torch.cuda.FloatTensor(m, c).zero_() 55 | interpolation_backward_cuda(n, c, k, grad_output, idx, weight, grad_input) 56 | return None, None, grad_input, None, None, None 57 | 58 | 59 | interpolation2 = Interpolation.apply 60 | -------------------------------------------------------------------------------- /libs/pointops/functions/query.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from pointops._C import ball_query_cuda, knn_query_cuda, random_ball_query_cuda 3 | from torch.autograd import Function 4 | 5 | 6 | class KNNQuery(Function): 7 | @staticmethod 8 | def forward(ctx, nsample, xyz, offset, new_xyz=None, new_offset=None): 9 | """ 10 | input: coords: (n, 3), new_xyz: (m, 3), offset: (b), new_offset: (b) 11 | output: idx: (m, nsample) -1 is placeholder, dist2: (m, nsample) 12 | """ 13 | if new_xyz is None or new_offset is None: 14 | new_xyz = xyz 15 | new_offset = offset 16 | assert xyz.is_contiguous() and new_xyz.is_contiguous() 17 | m = new_xyz.shape[0] 18 | idx = torch.cuda.IntTensor(m, nsample).zero_() 19 | dist2 = torch.cuda.FloatTensor(m, nsample).zero_() 20 | knn_query_cuda( 21 | m, nsample, xyz, new_xyz, offset.int(), new_offset.int(), idx, dist2 22 | ) 23 | return idx, torch.sqrt(dist2) 24 | 25 | 26 | class RandomBallQuery(Function): 27 | """Random Ball Query. 28 | 29 | Find nearby points in spherical space. 30 | """ 31 | 32 | @staticmethod 33 | def forward( 34 | ctx, nsample, max_radius, min_radius, xyz, offset, new_xyz=None, new_offset=None 35 | ): 36 | """ 37 | input: coords: (n, 3), new_xyz: (m, 3), offset: (b), new_offset: (b) 38 | output: idx: (m, nsample), dist2: (m, nsample) 39 | """ 40 | if new_xyz is None or new_offset is None: 41 | new_xyz = xyz 42 | new_offset = offset 43 | assert xyz.is_contiguous() and new_xyz.is_contiguous() 44 | assert min_radius < max_radius 45 | 46 | m = new_xyz.shape[0] 47 | order = [] 48 | for k in range(offset.shape[0]): 49 | s_k, e_k = (0, offset[0]) if k == 0 else (offset[k - 1], offset[k]) 50 | order.append( 51 | torch.randperm(e_k - s_k, dtype=torch.int32, device=offset.device) + s_k 52 | ) 53 | order = torch.cat(order, dim=0) 54 | idx = torch.cuda.IntTensor(m, nsample).zero_() 55 | dist2 = torch.cuda.FloatTensor(m, nsample).zero_() 56 | random_ball_query_cuda( 57 | m, 58 | nsample, 59 | min_radius, 60 | max_radius, 61 | order, 62 | xyz, 63 | new_xyz, 64 | offset.int(), 65 | new_offset.int(), 66 | idx, 67 | dist2, 68 | ) 69 | return idx, torch.sqrt(dist2) 70 | 71 | 72 | class BallQuery(Function): 73 | """Ball Query. 74 | 75 | Find nearby points in spherical space. 76 | """ 77 | 78 | @staticmethod 79 | def forward( 80 | ctx, nsample, max_radius, min_radius, xyz, offset, new_xyz=None, new_offset=None 81 | ): 82 | """ 83 | input: coords: (n, 3), new_xyz: (m, 3), offset: (b), new_offset: (b) 84 | output: idx: (m, nsample), dist2: (m, nsample) 85 | """ 86 | if new_xyz is None or new_offset is None: 87 | new_xyz = xyz 88 | new_offset = offset 89 | assert xyz.is_contiguous() and new_xyz.is_contiguous() 90 | assert min_radius < max_radius 91 | 92 | m = new_xyz.shape[0] 93 | idx = torch.cuda.IntTensor(m, nsample).zero_() 94 | dist2 = torch.cuda.FloatTensor(m, nsample).zero_() 95 | ball_query_cuda( 96 | m, 97 | nsample, 98 | min_radius, 99 | max_radius, 100 | xyz, 101 | new_xyz, 102 | offset.int(), 103 | new_offset.int(), 104 | idx, 105 | dist2, 106 | ) 107 | return idx, torch.sqrt(dist2) 108 | 109 | 110 | knn_query = KNNQuery.apply 111 | ball_query = BallQuery.apply 112 | random_ball_query = RandomBallQuery.apply 113 | -------------------------------------------------------------------------------- /libs/pointops/functions/sampling.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from pointops._C import farthest_point_sampling_cuda 3 | from torch.autograd import Function 4 | 5 | 6 | class FarthestPointSampling(Function): 7 | @staticmethod 8 | def forward(ctx, xyz, offset, new_offset): 9 | """ 10 | input: coords: (n, 3), offset: (b), new_offset: (b) 11 | output: idx: (m) 12 | """ 13 | assert xyz.is_contiguous() 14 | n, b, n_max = xyz.shape[0], offset.shape[0], offset[0] 15 | for i in range(1, b): 16 | n_max = max(offset[i] - offset[i - 1], n_max) 17 | idx = torch.cuda.IntTensor(new_offset[b - 1].item()).zero_() 18 | tmp = torch.cuda.FloatTensor(n).fill_(1e10) 19 | farthest_point_sampling_cuda( 20 | b, n_max, xyz, offset.int(), new_offset.int(), tmp, idx 21 | ) 22 | del tmp 23 | return idx 24 | 25 | 26 | farthest_point_sampling = FarthestPointSampling.apply 27 | -------------------------------------------------------------------------------- /libs/pointops/functions/subtraction.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from pointops._C import subtraction_backward_cuda, subtraction_forward_cuda 3 | from torch.autograd import Function 4 | 5 | 6 | class Subtraction(Function): 7 | @staticmethod 8 | def forward(ctx, input1, input2, idx): 9 | """ 10 | input: input1: (n, c), input2: (n, c), idx: (n, nsample) 11 | output: (n, nsample, c) 12 | """ 13 | assert input1.is_contiguous() and input2.is_contiguous() 14 | n, c = input1.shape 15 | nsample = idx.shape[-1] 16 | output = torch.cuda.FloatTensor(n, nsample, c).zero_() 17 | subtraction_forward_cuda(n, nsample, c, input1, input2, idx, output) 18 | ctx.save_for_backward(idx) 19 | return output 20 | 21 | @staticmethod 22 | def backward(ctx, grad_output): 23 | """ 24 | input: grad_out: (n, nsample, c) 25 | output: grad_input1: (n, c), grad_input2: (n, c) 26 | """ 27 | (idx,) = ctx.saved_tensors 28 | n, nsample, c = grad_output.shape 29 | grad_input1 = torch.cuda.FloatTensor(n, c).zero_() 30 | grad_input2 = torch.cuda.FloatTensor(n, c).zero_() 31 | subtraction_backward_cuda( 32 | n, nsample, c, idx, grad_output, grad_input1, grad_input2 33 | ) 34 | return grad_input1, grad_input2, None 35 | 36 | 37 | subtraction = Subtraction.apply 38 | -------------------------------------------------------------------------------- /libs/pointops/functions/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from pointops import ball_query, grouping, knn_query 3 | 4 | 5 | def knn_query_and_group( 6 | feat, 7 | xyz, 8 | offset=None, 9 | new_xyz=None, 10 | new_offset=None, 11 | idx=None, 12 | nsample=None, 13 | with_xyz=False, 14 | ): 15 | if idx is None: 16 | assert nsample is not None 17 | idx, _ = knn_query(nsample, xyz, offset, new_xyz, new_offset) 18 | return grouping(idx, feat, xyz, new_xyz, with_xyz), idx 19 | 20 | 21 | def ball_query_and_group( 22 | feat, 23 | xyz, 24 | offset=None, 25 | new_xyz=None, 26 | new_offset=None, 27 | idx=None, 28 | max_radio=None, 29 | min_radio=0, 30 | nsample=None, 31 | with_xyz=False, 32 | ): 33 | if idx is None: 34 | assert nsample is not None and offset is not None 35 | assert max_radio is not None and min_radio is not None 36 | idx, _ = ball_query( 37 | nsample, max_radio, min_radio, xyz, offset, new_xyz, new_offset 38 | ) 39 | return grouping(idx, feat, xyz, new_xyz, with_xyz), idx 40 | 41 | 42 | def query_and_group( 43 | nsample, 44 | xyz, 45 | new_xyz, 46 | feat, 47 | idx, 48 | offset, 49 | new_offset, 50 | dilation=0, 51 | with_feat=True, 52 | with_xyz=True, 53 | ): 54 | """ 55 | input: coords: (n, 3), new_xyz: (m, 3), color: (n, c), idx: (m, nsample), offset: (b), new_offset: (b) 56 | output: new_feat: (m, nsample, c+3), grouped_idx: (m, nsample) 57 | """ 58 | assert xyz.is_contiguous() and new_xyz.is_contiguous() and feat.is_contiguous() 59 | if new_xyz is None: 60 | new_xyz = xyz 61 | 62 | if idx is None: 63 | num_samples_total = 1 + (nsample - 1) * (dilation + 1) 64 | # num points in a batch might < num_samples_total => [n1, n2, ..., nk, ns, ns, ns, ...] 65 | idx_no_dilation, _ = knn_query( 66 | num_samples_total, xyz, offset, new_xyz, new_offset 67 | ) # (m, nsample * (d + 1)) 68 | idx = [] 69 | batch_end = offset.tolist() 70 | batch_start = [0] + batch_end[:-1] 71 | new_batch_end = new_offset.tolist() 72 | new_batch_start = [0] + new_batch_end[:-1] 73 | for i in range(offset.shape[0]): 74 | if batch_end[i] - batch_start[i] < num_samples_total: 75 | soft_dilation = (batch_end[i] - batch_start[i] - 1) / (nsample - 1) - 1 76 | else: 77 | soft_dilation = dilation 78 | idx.append( 79 | idx_no_dilation[ 80 | new_batch_start[i] : new_batch_end[i], 81 | [int((soft_dilation + 1) * i) for i in range(nsample)], 82 | ] 83 | ) 84 | idx = torch.cat(idx, dim=0) 85 | 86 | if not with_feat: 87 | return idx 88 | 89 | n, m, c = xyz.shape[0], new_xyz.shape[0], feat.shape[1] 90 | grouped_xyz = xyz[idx.view(-1).long(), :].view(m, nsample, 3) # (m, nsample, 3) 91 | # grouped_xyz = grouping(coords, idx) # (m, nsample, 3) 92 | grouped_xyz -= new_xyz.unsqueeze(1) # (m, nsample, 3) 93 | grouped_feat = feat[idx.view(-1).long(), :].view(m, nsample, c) # (m, nsample, c) 94 | # grouped_feat = grouping(color, idx) # (m, nsample, c) 95 | 96 | if with_xyz: 97 | return torch.cat((grouped_xyz, grouped_feat), -1), idx # (m, nsample, 3+c) 98 | else: 99 | return grouped_feat, idx 100 | 101 | 102 | def offset2batch(offset): 103 | return ( 104 | torch.cat( 105 | [ 106 | ( 107 | torch.tensor([i] * (o - offset[i - 1])) 108 | if i > 0 109 | else torch.tensor([i] * o) 110 | ) 111 | for i, o in enumerate(offset) 112 | ], 113 | dim=0, 114 | ) 115 | .long() 116 | .to(offset.device) 117 | ) 118 | 119 | 120 | def batch2offset(batch): 121 | return torch.cumsum(batch.bincount(), dim=0).int() 122 | -------------------------------------------------------------------------------- /libs/pointops/setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | from distutils.sysconfig import get_config_vars 3 | 4 | from setuptools import setup 5 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 6 | 7 | (opt,) = get_config_vars("OPT") 8 | os.environ["OPT"] = " ".join( 9 | flag for flag in opt.split() if flag != "-Wstrict-prototypes" 10 | ) 11 | 12 | src = "src" 13 | sources = [ 14 | os.path.join(root, file) 15 | for root, dirs, files in os.walk(src) 16 | for file in files 17 | if file.endswith(".cpp") or file.endswith(".cu") 18 | ] 19 | 20 | setup( 21 | name="pointops", 22 | version="1.0", 23 | install_requires=["torch", "numpy"], 24 | packages=["pointops"], 25 | package_dir={"pointops": "functions"}, 26 | ext_modules=[ 27 | CUDAExtension( 28 | name="pointops._C", 29 | sources=sources, 30 | extra_compile_args={"cxx": ["-g"], "nvcc": ["-O2"]}, 31 | ) 32 | ], 33 | cmdclass={"build_ext": BuildExtension}, 34 | ) 35 | -------------------------------------------------------------------------------- /libs/pointops/src/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/OpenGVLab/PonderV2/c152af1df9d04a98bc19696e8d69d047834a80db/libs/pointops/src/__init__.py -------------------------------------------------------------------------------- /libs/pointops/src/aggregation/aggregation_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "aggregation_cuda_kernel.h" 5 | 6 | 7 | void aggregation_forward_cuda(int n, int nsample, int c, int w_c, at::Tensor input_tensor, at::Tensor position_tensor, at::Tensor weight_tensor, at::Tensor idx_tensor, at::Tensor output_tensor) 8 | { 9 | const float *input = input_tensor.data_ptr(); 10 | const float *position = position_tensor.data_ptr(); 11 | const float *weight = weight_tensor.data_ptr(); 12 | const int *idx = idx_tensor.data_ptr(); 13 | float *output = output_tensor.data_ptr(); 14 | aggregation_forward_cuda_launcher(n, nsample, c, w_c, input, position, weight, idx, output); 15 | } 16 | 17 | void aggregation_backward_cuda(int n, int nsample, int c, int w_c, at::Tensor input_tensor, at::Tensor position_tensor, at::Tensor weight_tensor, at::Tensor idx_tensor, at::Tensor grad_output_tensor, at::Tensor grad_input_tensor, at::Tensor grad_position_tensor, at::Tensor grad_weight_tensor) 18 | { 19 | const float *input = input_tensor.data_ptr(); 20 | const float *position = position_tensor.data_ptr(); 21 | const float *weight = weight_tensor.data_ptr(); 22 | const int *idx = idx_tensor.data_ptr(); 23 | const float *grad_output = grad_output_tensor.data_ptr(); 24 | float *grad_input = grad_input_tensor.data_ptr(); 25 | float *grad_position = grad_position_tensor.data_ptr(); 26 | float *grad_weight = grad_weight_tensor.data_ptr(); 27 | aggregation_backward_cuda_launcher(n, nsample, c, w_c, input, position, weight, idx, grad_output, grad_input, grad_position, grad_weight); 28 | } 29 | -------------------------------------------------------------------------------- /libs/pointops/src/aggregation/aggregation_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "aggregation_cuda_kernel.h" 3 | 4 | 5 | __global__ void aggregation_forward_cuda_kernel(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, float *output) { 6 | // input: input: (n, c), position: (n, nsample, c), weight: (n, nsample, w_c), idx: (n, nsample), output: (n, c) 7 | int index = blockIdx.x * blockDim.x + threadIdx.x; 8 | if (index >= n * c) return; 9 | const int c_idx = index % c; 10 | const int n_idx = index / c; 11 | const int w_c_idx = c_idx % w_c; 12 | for (int nsample_idx = 0; nsample_idx < nsample; nsample_idx++) 13 | { 14 | int idx_idx = n_idx * nsample + nsample_idx; 15 | int input_idx = idx[idx_idx] * c + c_idx; 16 | int position_idx = n_idx * nsample * c + nsample_idx * c + c_idx; 17 | int weight_idx = n_idx * nsample * w_c + nsample_idx * w_c + w_c_idx; 18 | output[index] += (input[input_idx] + position[position_idx]) * weight[weight_idx]; 19 | } 20 | } 21 | 22 | __global__ void aggregation_backward_cuda_kernel(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, const float *grad_output, float *grad_input, float *grad_position, float *grad_weight) { 23 | // input: grad_output: (n, c), output: grad_input: (n, c), grad_position: (n, nsample, c), grad_weight: (n, nsample, w_c) 24 | int index = blockIdx.x * blockDim.x + threadIdx.x; 25 | if (index >= n * c) return; 26 | const int c_idx = index % c; 27 | const int n_idx = index / c; 28 | const int w_c_idx = c_idx % w_c; 29 | for (int nsample_idx = 0; nsample_idx < nsample; nsample_idx++) 30 | { 31 | int idx_idx = n_idx * nsample + nsample_idx; 32 | int input_idx = idx[idx_idx] * c + c_idx; 33 | int position_idx = n_idx * nsample * c + nsample_idx * c + c_idx; 34 | int weight_idx = n_idx * nsample * w_c + nsample_idx * w_c + w_c_idx; 35 | atomicAdd(grad_input + input_idx, grad_output[index] * weight[weight_idx]); 36 | grad_position[position_idx] = grad_output[index] * weight[weight_idx]; 37 | atomicAdd(grad_weight + weight_idx, grad_output[index] * (input[input_idx] + position[position_idx])); 38 | } 39 | } 40 | 41 | void aggregation_forward_cuda_launcher(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, float *output) { 42 | // input: input: (n, c), position: (n, nsample, c), weight: (n, nsample, w_c), idx: (n, nsample), output: (n, c) 43 | dim3 blocks(DIVUP(n * c, THREADS_PER_BLOCK)); 44 | dim3 threads(THREADS_PER_BLOCK); 45 | aggregation_forward_cuda_kernel<<>>(n, nsample, c, w_c, input, position, weight, idx, output); 46 | } 47 | 48 | void aggregation_backward_cuda_launcher(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, const float *grad_output, float *grad_input, float *grad_position, float *grad_weight) { 49 | // input: grad_output: (n, c), output: grad_input: (n, c), grad_position: (n, nsample, c), grad_weight: (n, nsample, w_c) 50 | dim3 blocks(DIVUP(n * c, THREADS_PER_BLOCK)); 51 | dim3 threads(THREADS_PER_BLOCK); 52 | aggregation_backward_cuda_kernel<<>>(n, nsample, c, w_c, input, position, weight, idx, grad_output, grad_input, grad_position, grad_weight); 53 | } 54 | -------------------------------------------------------------------------------- /libs/pointops/src/aggregation/aggregation_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _AGGREGATION_CUDA_KERNEL 2 | #define _AGGREGATION_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void aggregation_forward_cuda(int n, int nsample, int c, int w_c, at::Tensor input_tensor, at::Tensor position_tensor, at::Tensor weight_tensor, at::Tensor idx_tensor, at::Tensor output_tensor); 8 | void aggregation_backward_cuda(int n, int nsample, int c, int w_c, at::Tensor input_tensor, at::Tensor position_tensor, at::Tensor weight_tensor, at::Tensor idx_tensor, at::Tensor grad_output_tensor, at::Tensor grad_input_tensor, at::Tensor grad_position_tensor, at::Tensor grad_weight_tensor); 9 | 10 | #ifdef __cplusplus 11 | extern "C" { 12 | #endif 13 | 14 | void aggregation_forward_cuda_launcher(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, float *output); 15 | void aggregation_backward_cuda_launcher(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, const float *grad_output, float *grad_input, float *grad_position, float *grad_weight); 16 | 17 | #ifdef __cplusplus 18 | } 19 | #endif 20 | #endif 21 | -------------------------------------------------------------------------------- /libs/pointops/src/attention/attention_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "attention_cuda_kernel.h" 5 | 6 | 7 | void attention_relation_step_forward_cuda(int m, int g, int c, 8 | at::Tensor query_tensor, at::Tensor key_tensor, at::Tensor weight_tensor, 9 | at::Tensor index_target_tensor, at::Tensor index_refer_tensor, 10 | at::Tensor output_tensor) 11 | { 12 | const float *query = query_tensor.data_ptr(); 13 | const float *key = key_tensor.data_ptr(); 14 | const float *weight = weight_tensor.data_ptr(); 15 | const int *index_target = index_target_tensor.data_ptr(); 16 | const int *index_refer = index_refer_tensor.data_ptr(); 17 | float *output = output_tensor.data_ptr(); 18 | attention_relation_step_forward_cuda_launcher(m, g, c, query, key, weight, index_target, index_refer, output); 19 | } 20 | 21 | void attention_relation_step_backward_cuda(int m, int g, int c, 22 | at::Tensor query_tensor, at::Tensor grad_query_tensor, 23 | at::Tensor key_tensor, at::Tensor grad_key_tensor, 24 | at::Tensor weight_tensor, at::Tensor grad_weight_tensor, 25 | at::Tensor index_target_tensor, at::Tensor index_refer_tensor, 26 | at::Tensor grad_output_tensor) 27 | { 28 | const float *query = query_tensor.data_ptr(); 29 | float *grad_query = grad_query_tensor.data_ptr(); 30 | const float *key = key_tensor.data_ptr(); 31 | float *grad_key = grad_key_tensor.data_ptr(); 32 | const float *weight = weight_tensor.data_ptr(); 33 | float *grad_weight = grad_weight_tensor.data_ptr(); 34 | const int *index_target = index_target_tensor.data_ptr(); 35 | const int *index_refer = index_refer_tensor.data_ptr(); 36 | const float *grad_output = grad_output_tensor.data_ptr(); 37 | attention_relation_step_backward_cuda_launcher(m, g, c, 38 | query, grad_query, 39 | key, grad_key, 40 | weight, grad_weight, 41 | index_target, index_refer, grad_output); 42 | } 43 | 44 | 45 | void attention_fusion_step_forward_cuda(int m, int g, int c, 46 | at::Tensor weight_tensor, at::Tensor value_tensor, 47 | at::Tensor index_target_tensor, at::Tensor index_refer_tensor, 48 | at::Tensor output_tensor) 49 | { 50 | const float *weight = weight_tensor.data_ptr(); 51 | const float *value = value_tensor.data_ptr(); 52 | const int *index_target = index_target_tensor.data_ptr(); 53 | const int *index_refer = index_refer_tensor.data_ptr(); 54 | float *output = output_tensor.data_ptr(); 55 | attention_fusion_step_forward_cuda_launcher(m, g, c, weight, value, index_target, index_refer, output); 56 | } 57 | 58 | 59 | void attention_fusion_step_backward_cuda(int m, int g, int c, 60 | at::Tensor weight_tensor, at::Tensor grad_weight_tensor, 61 | at::Tensor value_tensor, at::Tensor grad_value_tensor, 62 | at::Tensor index_target_tensor, at::Tensor index_refer_tensor, 63 | at::Tensor grad_output_tensor) 64 | { 65 | const float *weight = weight_tensor.data_ptr(); 66 | float *grad_weight = grad_weight_tensor.data_ptr(); 67 | const float *value = value_tensor.data_ptr(); 68 | float *grad_value = grad_value_tensor.data_ptr(); 69 | const int *index_target = index_target_tensor.data_ptr(); 70 | const int *index_refer = index_refer_tensor.data_ptr(); 71 | const float *grad_output = grad_output_tensor.data_ptr(); 72 | attention_fusion_step_backward_cuda_launcher(m, g, c, 73 | weight, grad_weight, 74 | value, grad_value, 75 | index_target, index_refer, grad_output); 76 | } 77 | -------------------------------------------------------------------------------- /libs/pointops/src/attention/attention_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _ATTENTION_CUDA_KERNEL 2 | #define _ATTENTION_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void attention_relation_step_forward_cuda(int m, int g, int c, 8 | at::Tensor query_tensor, at::Tensor key_tensor, at::Tensor weight_tensor, 9 | at::Tensor index_target_tensor, at::Tensor index_refer_tensor, 10 | at::Tensor output_tensor); 11 | void attention_relation_step_backward_cuda(int m, int g, int c, 12 | at::Tensor query_tensor, at::Tensor grad_query_tensor, 13 | at::Tensor key_tensor, at::Tensor grad_key_tensor, 14 | at::Tensor weight_tensor, at::Tensor grad_weight_tensor, 15 | at::Tensor index_target_tensor, at::Tensor index_refer_tensor, 16 | at::Tensor grad_output_tensor); 17 | void attention_fusion_step_forward_cuda(int m, int g, int c, 18 | at::Tensor weight_tensor, at::Tensor value_tensor, 19 | at::Tensor index_target_tensor, at::Tensor index_refer_tensor, 20 | at::Tensor output_tensor); 21 | void attention_fusion_step_backward_cuda(int m, int g, int c, 22 | at::Tensor weight_tensor, at::Tensor grad_weight_tensor, 23 | at::Tensor value_tensor, at::Tensor grad_value_tensor, 24 | at::Tensor index_target_tensor, at::Tensor index_refer_tensor, 25 | at::Tensor grad_output_tensor); 26 | 27 | #ifdef __cplusplus 28 | extern "C" { 29 | #endif 30 | 31 | void attention_relation_step_forward_cuda_launcher(int m, int g, int c, 32 | const float *query, const float *key, const float *weight, 33 | const int *index_target, const int *index_refer, 34 | float *output); 35 | void attention_relation_step_backward_cuda_launcher(int m, int g, int c, 36 | const float *query, float *grad_query, 37 | const float *key, float *grad_key, 38 | const float *weight, float *grad_weight, 39 | const int *index_target, const int *index_refer, 40 | const float *grad_output); 41 | void attention_fusion_step_forward_cuda_launcher(int m, int g, int c, 42 | const float *weight, const float *value, 43 | const int *index_target, const int *index_refer, 44 | float *output); 45 | void attention_fusion_step_backward_cuda_launcher(int m, int g, int c, 46 | const float *weight, float *grad_weight, 47 | const float *value, float *grad_value, 48 | const int *index_target, const int *index_refer, 49 | const float *grad_output); 50 | 51 | #ifdef __cplusplus 52 | } 53 | #endif 54 | #endif 55 | -------------------------------------------------------------------------------- /libs/pointops/src/ball_query/ball_query_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "ball_query_cuda_kernel.h" 5 | 6 | 7 | void ball_query_cuda(int m, int nsample, 8 | float min_radius, float max_radius, 9 | at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, 10 | at::Tensor offset_tensor, at::Tensor new_offset_tensor, 11 | at::Tensor idx_tensor, at::Tensor dist2_tensor) 12 | { 13 | const float *xyz = xyz_tensor.data_ptr(); 14 | const float *new_xyz = new_xyz_tensor.data_ptr(); 15 | const int *offset = offset_tensor.data_ptr(); 16 | const int *new_offset = new_offset_tensor.data_ptr(); 17 | int *idx = idx_tensor.data_ptr(); 18 | float *dist2 = dist2_tensor.data_ptr(); 19 | ball_query_cuda_launcher(m, nsample, min_radius, max_radius, xyz, new_xyz, offset, new_offset, idx, dist2); 20 | } 21 | -------------------------------------------------------------------------------- /libs/pointops/src/ball_query/ball_query_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _BALL_QUERY_CUDA_KERNEL 2 | #define _BALL_QUERY_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void ball_query_cuda(int m, int nsample, 8 | float min_radius, float max_radius, 9 | at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, 10 | at::Tensor offset_tensor, at::Tensor new_offset_tensor, 11 | at::Tensor idx_tensor, at::Tensor dist2_tensor); 12 | 13 | #ifdef __cplusplus 14 | extern "C" { 15 | #endif 16 | 17 | void ball_query_cuda_launcher(int m, int nsample, 18 | float min_radius, float max_radius, 19 | const float *xyz, const float *new_xyz, 20 | const int *offset, const int *new_offset, 21 | int *idx, float *dist2); 22 | 23 | #ifdef __cplusplus 24 | } 25 | #endif 26 | #endif 27 | -------------------------------------------------------------------------------- /libs/pointops/src/cuda_utils.h: -------------------------------------------------------------------------------- 1 | #ifndef _CUDA_UTILS_H 2 | #define _CUDA_UTILS_H 3 | 4 | #include 5 | #include 6 | 7 | #define TOTAL_THREADS 1024 8 | #define THREADS_PER_BLOCK 512 9 | #define DIVUP(m, n) ((m) / (n) + ((m) % (n) > 0)) 10 | 11 | inline int opt_n_threads(int work_size) { 12 | const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0); 13 | return std::max(std::min(1 << pow_2, TOTAL_THREADS), 1); 14 | } 15 | 16 | inline dim3 opt_block_config(int x, int y) { 17 | const int x_threads = opt_n_threads(x); 18 | const int y_threads = std::max(std::min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1); 19 | dim3 block_config(x_threads, y_threads, 1); 20 | return block_config; 21 | } 22 | 23 | #endif 24 | -------------------------------------------------------------------------------- /libs/pointops/src/grouping/grouping_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "grouping_cuda_kernel.h" 5 | 6 | 7 | void grouping_forward_cuda(int m, int nsample, int c, at::Tensor input_tensor, at::Tensor idx_tensor, at::Tensor output_tensor) 8 | { 9 | const float *input = input_tensor.data_ptr(); 10 | const int *idx = idx_tensor.data_ptr(); 11 | float *output = output_tensor.data_ptr(); 12 | grouping_forward_cuda_launcher(m, nsample, c, input, idx, output); 13 | } 14 | 15 | void grouping_backward_cuda(int m, int nsample, int c, at::Tensor grad_output_tensor, at::Tensor idx_tensor, at::Tensor grad_input_tensor) 16 | { 17 | const float *grad_output = grad_output_tensor.data_ptr(); 18 | const int *idx = idx_tensor.data_ptr(); 19 | float *grad_input = grad_input_tensor.data_ptr(); 20 | grouping_backward_cuda_launcher(m, nsample, c, grad_output, idx, grad_input); 21 | } 22 | -------------------------------------------------------------------------------- /libs/pointops/src/grouping/grouping_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "grouping_cuda_kernel.h" 3 | 4 | 5 | __global__ void grouping_forward_cuda_kernel(int m, int nsample, int c, const float *__restrict__ input, const int *__restrict__ idx, float *__restrict__ output) { 6 | // input: input: (n, c), idx: (m, nsample), output: (m, nsample, c) 7 | int index = blockIdx.x * blockDim.x + threadIdx.x; 8 | if (index >= m * nsample * c) return; 9 | const int c_idx = index % c; 10 | const int nsample_idx = (index / c) % nsample; 11 | const int m_idx = index / nsample / c; 12 | const int input_idx = idx[m_idx * nsample + nsample_idx] * c + c_idx; 13 | output[index] = input[input_idx]; 14 | } 15 | 16 | __global__ void grouping_backward_cuda_kernel(int m, int nsample, int c, const float *__restrict__ grad_output, const int *__restrict__ idx, float *__restrict__ grad_input) { 17 | // input: grad_output: (m, nsample, c), idx: (m, nsample), output: grad_input: (n, c) 18 | int index = blockIdx.x * blockDim.x + threadIdx.x; 19 | if (index >= m * nsample * c) return; 20 | const int c_idx = index % c; 21 | const int nsample_idx = (index / c) % nsample; 22 | const int m_idx = index / nsample / c; 23 | const int input_idx = idx[m_idx * nsample + nsample_idx] * c + c_idx; 24 | atomicAdd(grad_input + input_idx, grad_output[index]); 25 | } 26 | 27 | void grouping_forward_cuda_launcher(int m, int nsample, int c, const float *input, const int *idx, float *output) { 28 | // input: input: (n, c), idx: (m, nsample), output: (m, nsample, c) 29 | dim3 blocks(DIVUP(m * nsample * c, THREADS_PER_BLOCK)); 30 | dim3 threads(THREADS_PER_BLOCK); 31 | grouping_forward_cuda_kernel<<>>(m, nsample, c, input, idx, output); 32 | } 33 | 34 | void grouping_backward_cuda_launcher(int m, int nsample, int c, const float *grad_output, const int *idx, float *grad_input) 35 | { 36 | // input: grad_output: (m, nsample, c), idx: (m, nsample), output: grad_input: (n, c) 37 | dim3 blocks(DIVUP(m * nsample * c, THREADS_PER_BLOCK)); 38 | dim3 threads(THREADS_PER_BLOCK); 39 | grouping_backward_cuda_kernel<<>>(m, nsample, c, grad_output, idx, grad_input); 40 | } 41 | -------------------------------------------------------------------------------- /libs/pointops/src/grouping/grouping_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _GROUPING_CUDA_KERNEL 2 | #define _GROUPING_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void grouping_forward_cuda(int m, int nsample, int c, at::Tensor input_tensor, at::Tensor idx_tensor, at::Tensor output_tensor); 8 | void grouping_backward_cuda(int m, int nsample, int c, at::Tensor grad_output_tensor, at::Tensor idx_tensor, at::Tensor grad_input_tensor); 9 | 10 | #ifdef __cplusplus 11 | extern "C" { 12 | #endif 13 | 14 | void grouping_forward_cuda_launcher(int m, int nsample, int c, const float *input, const int *idx, float *output); 15 | void grouping_backward_cuda_launcher(int m, int nsample, int c, const float *grad_output, const int *idx, float *grad_input); 16 | 17 | #ifdef __cplusplus 18 | } 19 | #endif 20 | #endif 21 | -------------------------------------------------------------------------------- /libs/pointops/src/interpolation/interpolation_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "interpolation_cuda_kernel.h" 5 | 6 | 7 | void interpolation_forward_cuda(int n, int c, int k, at::Tensor input_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor output_tensor) 8 | { 9 | const float *input = input_tensor.data_ptr(); 10 | const int *idx = idx_tensor.data_ptr(); 11 | const float *weight = weight_tensor.data_ptr(); 12 | float *output = output_tensor.data_ptr(); 13 | interpolation_forward_cuda_launcher(n, c, k, input, idx, weight, output); 14 | } 15 | 16 | void interpolation_backward_cuda(int n, int c, int k, at::Tensor grad_output_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor grad_input_tensor) 17 | { 18 | const float *grad_output = grad_output_tensor.data_ptr(); 19 | const int *idx = idx_tensor.data_ptr(); 20 | const float *weight = weight_tensor.data_ptr(); 21 | float *grad_input = grad_input_tensor.data_ptr(); 22 | interpolation_backward_cuda_launcher(n, c, k, grad_output, idx, weight, grad_input); 23 | } 24 | -------------------------------------------------------------------------------- /libs/pointops/src/interpolation/interpolation_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "interpolation_cuda_kernel.h" 3 | 4 | 5 | __global__ void interpolation_forward_cuda_kernel(int n, int c, int k, const float *input, const int *idx, const float *weight, float *output) 6 | { 7 | // input: input: (m, c), idx: (n, k), weight: (n, k), output: output (n, c) 8 | int index = blockIdx.x * blockDim.x + threadIdx.x; 9 | if (index >= n * c) return; 10 | int c_idx = index % c; 11 | int n_idx = index / c; 12 | for (int i = 0; i < k; i++) 13 | { 14 | int idx_idx = n_idx * k + i; 15 | int input_idx = idx[idx_idx] * c + c_idx; 16 | output[index] += input[input_idx] * weight[idx_idx]; 17 | } 18 | } 19 | 20 | __global__ void interpolation_backward_cuda_kernel(int n, int c, int k, const float *grad_output, const int *idx, const float *weight, float *grad_input) 21 | { 22 | // input: grad_output: (n, c), idx: (n, k), weight: (n, k), output: grad_input (m, c) 23 | int index = blockIdx.x * blockDim.x + threadIdx.x; 24 | if (index >= n * c) return; 25 | int c_idx = index % c; 26 | int n_idx = index / c; 27 | for (int i = 0; i < k; i++) 28 | { 29 | int idx_idx = n_idx * k + i; 30 | int input_idx = idx[idx_idx] * c + c_idx; 31 | atomicAdd(grad_input + input_idx, grad_output[index] * weight[idx_idx]); 32 | } 33 | } 34 | 35 | void interpolation_forward_cuda_launcher(int n, int c, int k, const float *input, const int *idx, const float *weight, float *output) { 36 | // input: input: (m, c), idx: (n, k), weight: (n, k), output: output (n, c) 37 | dim3 blocks(DIVUP(n * c, THREADS_PER_BLOCK)); 38 | dim3 threads(THREADS_PER_BLOCK); 39 | interpolation_forward_cuda_kernel<<>>(n, c, k, input, idx, weight, output); 40 | } 41 | 42 | void interpolation_backward_cuda_launcher(int n, int c, int k, const float *grad_output, const int *idx, const float *weight, float *grad_input) { 43 | // input: grad_output: (n, c), idx: (n, k), weight: (n, k), output: grad_input (m, c) 44 | dim3 blocks(DIVUP(n * c, THREADS_PER_BLOCK)); 45 | dim3 threads(THREADS_PER_BLOCK); 46 | interpolation_backward_cuda_kernel<<>>(n, c, k, grad_output, idx, weight, grad_input); 47 | } 48 | -------------------------------------------------------------------------------- /libs/pointops/src/interpolation/interpolation_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _INTERPOLATION_CUDA_KERNEL 2 | #define _INTERPOLATION_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void interpolation_forward_cuda(int n, int c, int k, at::Tensor input_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor output_tensor); 8 | void interpolation_backward_cuda(int n, int c, int k, at::Tensor grad_output_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor grad_input_tensor); 9 | 10 | #ifdef __cplusplus 11 | extern "C" { 12 | #endif 13 | 14 | void interpolation_forward_cuda_launcher(int n, int c, int k, const float *input, const int *idx, const float *weight, float *output); 15 | void interpolation_backward_cuda_launcher(int n, int c, int k, const float *grad_output, const int *idx, const float *weight, float *grad_input); 16 | 17 | #ifdef __cplusplus 18 | } 19 | #endif 20 | #endif 21 | -------------------------------------------------------------------------------- /libs/pointops/src/knn_query/knn_query_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "knn_query_cuda_kernel.h" 5 | 6 | 7 | void knn_query_cuda(int m, int nsample, at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor idx_tensor, at::Tensor dist2_tensor) 8 | { 9 | const float *xyz = xyz_tensor.data_ptr(); 10 | const float *new_xyz = new_xyz_tensor.data_ptr(); 11 | const int *offset = offset_tensor.data_ptr(); 12 | const int *new_offset = new_offset_tensor.data_ptr(); 13 | int *idx = idx_tensor.data_ptr(); 14 | float *dist2 = dist2_tensor.data_ptr(); 15 | knn_query_cuda_launcher(m, nsample, xyz, new_xyz, offset, new_offset, idx, dist2); 16 | } 17 | -------------------------------------------------------------------------------- /libs/pointops/src/knn_query/knn_query_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "knn_query_cuda_kernel.h" 3 | 4 | 5 | namespace knn_query_utils{ 6 | 7 | template 8 | __device__ void swap(DType *x, DType *y) 9 | { 10 | DType tmp = *x; 11 | *x = *y; 12 | *y = tmp; 13 | } 14 | 15 | __device__ void reheap(float *dist, int *idx, int k) 16 | { 17 | int root = 0; 18 | int child = root * 2 + 1; 19 | while (child < k) 20 | { 21 | if(child + 1 < k && dist[child+1] > dist[child]) 22 | child++; 23 | if(dist[root] > dist[child]) 24 | return; 25 | swap(&dist[root], &dist[child]); 26 | swap(&idx[root], &idx[child]); 27 | root = child; 28 | child = root * 2 + 1; 29 | } 30 | } 31 | 32 | 33 | __device__ void heap_sort(float *dist, int *idx, int k) 34 | { 35 | int i; 36 | for (i = k - 1; i > 0; i--) 37 | { 38 | swap(&dist[0], &dist[i]); 39 | swap(&idx[0], &idx[i]); 40 | reheap(dist, idx, i); 41 | } 42 | } 43 | 44 | 45 | __device__ int get_bt_idx(int idx, const int *offset) 46 | { 47 | int i = 0; 48 | while (1) 49 | { 50 | if (idx < offset[i]) 51 | break; 52 | else 53 | i++; 54 | } 55 | return i; 56 | } 57 | } // namespace knn_query_utils 58 | 59 | 60 | __global__ void knn_query_cuda_kernel(int m, int nsample, const float *__restrict__ xyz, const float *__restrict__ new_xyz, const int *__restrict__ offset, const int *__restrict__ new_offset, int *__restrict__ idx, float *__restrict__ dist2) { 61 | // input: xyz (n, 3) new_xyz (m, 3) 62 | // output: idx (m, nsample) dist2 (m, nsample) 63 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 64 | if (pt_idx >= m) return; 65 | 66 | new_xyz += pt_idx * 3; 67 | idx += pt_idx * nsample; 68 | dist2 += pt_idx * nsample; 69 | 70 | int bt_idx = knn_query_utils::get_bt_idx(pt_idx, new_offset); 71 | int start; 72 | if (bt_idx == 0) 73 | start = 0; 74 | else 75 | start = offset[bt_idx - 1]; 76 | int end = offset[bt_idx]; 77 | 78 | float new_x = new_xyz[0]; 79 | float new_y = new_xyz[1]; 80 | float new_z = new_xyz[2]; 81 | 82 | float best_dist[128]; 83 | int best_idx[128]; 84 | for(int i = 0; i < nsample; i++){ 85 | best_dist[i] = 1e10; 86 | best_idx[i] = -1; 87 | } 88 | for(int i = start; i < end; i++){ 89 | float x = xyz[i * 3 + 0]; 90 | float y = xyz[i * 3 + 1]; 91 | float z = xyz[i * 3 + 2]; 92 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + (new_z - z) * (new_z - z); 93 | if (d2 < best_dist[0]){ 94 | best_dist[0] = d2; 95 | best_idx[0] = i; 96 | knn_query_utils::reheap(best_dist, best_idx, nsample); 97 | } 98 | } 99 | knn_query_utils::heap_sort(best_dist, best_idx, nsample); 100 | for(int i = 0; i < nsample; i++){ 101 | idx[i] = best_idx[i]; 102 | dist2[i] = best_dist[i]; 103 | } 104 | } 105 | 106 | 107 | void knn_query_cuda_launcher(int m, int nsample, const float *xyz, const float *new_xyz, const int *offset, const int *new_offset, int *idx, float *dist2) { 108 | // input: new_xyz: (m, 3), xyz: (n, 3), idx: (m, nsample) 109 | dim3 blocks(DIVUP(m, THREADS_PER_BLOCK)); 110 | dim3 threads(THREADS_PER_BLOCK); 111 | knn_query_cuda_kernel<<>>(m, nsample, xyz, new_xyz, offset, new_offset, idx, dist2); 112 | } 113 | -------------------------------------------------------------------------------- /libs/pointops/src/knn_query/knn_query_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _KNN_QUERY_CUDA_KERNEL 2 | #define _KNN_QUERY_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void knn_query_cuda(int m, int nsample, at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor idx_tensor, at::Tensor dist2_tensor); 8 | 9 | #ifdef __cplusplus 10 | extern "C" { 11 | #endif 12 | 13 | void knn_query_cuda_launcher(int m, int nsample, const float *xyz, const float *new_xyz, const int *offset, const int *new_offset, int *idx, float *dist2); 14 | 15 | #ifdef __cplusplus 16 | } 17 | #endif 18 | #endif 19 | -------------------------------------------------------------------------------- /libs/pointops/src/pointops_api.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "knn_query/knn_query_cuda_kernel.h" 5 | #include "ball_query/ball_query_cuda_kernel.h" 6 | #include "random_ball_query/random_ball_query_cuda_kernel.h" 7 | #include "sampling/sampling_cuda_kernel.h" 8 | #include "grouping/grouping_cuda_kernel.h" 9 | #include "interpolation/interpolation_cuda_kernel.h" 10 | #include "aggregation/aggregation_cuda_kernel.h" 11 | #include "subtraction/subtraction_cuda_kernel.h" 12 | #include "attention/attention_cuda_kernel.h" 13 | 14 | 15 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 16 | m.def("knn_query_cuda", &knn_query_cuda, "knn_query_cuda"); 17 | m.def("ball_query_cuda", &ball_query_cuda, "ball_query_cuda"); 18 | m.def("random_ball_query_cuda", &random_ball_query_cuda, "random_ball_query_cuda"); 19 | m.def("farthest_point_sampling_cuda", &farthest_point_sampling_cuda, "farthest_point_sampling_cuda"); 20 | m.def("grouping_forward_cuda", &grouping_forward_cuda, "grouping_forward_cuda"); 21 | m.def("grouping_backward_cuda", &grouping_backward_cuda, "grouping_backward_cuda"); 22 | m.def("interpolation_forward_cuda", &interpolation_forward_cuda, "interpolation_forward_cuda"); 23 | m.def("interpolation_backward_cuda", &interpolation_backward_cuda, "interpolation_backward_cuda"); 24 | m.def("subtraction_forward_cuda", &subtraction_forward_cuda, "subtraction_forward_cuda"); 25 | m.def("subtraction_backward_cuda", &subtraction_backward_cuda, "subtraction_backward_cuda"); 26 | m.def("aggregation_forward_cuda", &aggregation_forward_cuda, "aggregation_forward_cuda"); 27 | m.def("aggregation_backward_cuda", &aggregation_backward_cuda, "aggregation_backward_cuda"); 28 | m.def("attention_relation_step_forward_cuda", &attention_relation_step_forward_cuda, "attention_relation_step_forward_cuda"); 29 | m.def("attention_relation_step_backward_cuda", &attention_relation_step_backward_cuda, "attention_relation_step_backward_cuda"); 30 | m.def("attention_fusion_step_forward_cuda", &attention_fusion_step_forward_cuda, "attention_fusion_step_forward_cuda"); 31 | m.def("attention_fusion_step_backward_cuda", &attention_fusion_step_backward_cuda, "attention_fusion_step_backward_cuda"); 32 | } 33 | -------------------------------------------------------------------------------- /libs/pointops/src/random_ball_query/random_ball_query_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "random_ball_query_cuda_kernel.h" 5 | 6 | 7 | void random_ball_query_cuda(int m, int nsample, 8 | float min_radius, float max_radius, at::Tensor order_tensor, 9 | at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, 10 | at::Tensor offset_tensor, at::Tensor new_offset_tensor, 11 | at::Tensor idx_tensor, at::Tensor dist2_tensor) 12 | { 13 | const int *order = order_tensor.data_ptr(); 14 | const float *xyz = xyz_tensor.data_ptr(); 15 | const float *new_xyz = new_xyz_tensor.data_ptr(); 16 | const int *offset = offset_tensor.data_ptr(); 17 | const int *new_offset = new_offset_tensor.data_ptr(); 18 | int *idx = idx_tensor.data_ptr(); 19 | float *dist2 = dist2_tensor.data_ptr(); 20 | random_ball_query_cuda_launcher(m, nsample, min_radius, max_radius, order, xyz, new_xyz, offset, new_offset, idx, dist2); 21 | } 22 | -------------------------------------------------------------------------------- /libs/pointops/src/random_ball_query/random_ball_query_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "random_ball_query_cuda_kernel.h" 3 | 4 | 5 | namespace random_ball_query_utils{ 6 | 7 | template 8 | __device__ void swap(DType *x, DType *y) 9 | { 10 | DType tmp = *x; 11 | *x = *y; 12 | *y = tmp; 13 | } 14 | 15 | __device__ void reheap(float *dist, int *idx, int k) 16 | { 17 | int root = 0; 18 | int child = root * 2 + 1; 19 | while (child < k) 20 | { 21 | if(child + 1 < k && dist[child+1] > dist[child]) 22 | child++; 23 | if(dist[root] > dist[child]) 24 | return; 25 | swap(&dist[root], &dist[child]); 26 | swap(&idx[root], &idx[child]); 27 | root = child; 28 | child = root * 2 + 1; 29 | } 30 | } 31 | 32 | 33 | __device__ void heap_sort(float *dist, int *idx, int k) 34 | { 35 | int i; 36 | for (i = k - 1; i > 0; i--) 37 | { 38 | swap(&dist[0], &dist[i]); 39 | swap(&idx[0], &idx[i]); 40 | reheap(dist, idx, i); 41 | } 42 | } 43 | 44 | __device__ int get_bt_idx(int idx, const int *offset) 45 | { 46 | int i = 0; 47 | while (1) 48 | { 49 | if (idx < offset[i]) 50 | break; 51 | else 52 | i++; 53 | } 54 | return i; 55 | } 56 | } // namespace ball_query_utils 57 | 58 | __global__ void random_ball_query_cuda_kernel(int m, int nsample, 59 | float min_radius, float max_radius, const int *__restrict__ order, 60 | const float *__restrict__ xyz, const float *__restrict__ new_xyz, 61 | const int *__restrict__ offset, const int *__restrict__ new_offset, 62 | int *__restrict__ idx, float *__restrict__ dist2) { 63 | // input: xyz (n, 3) new_xyz (m, 3) 64 | // output: idx (m, nsample) dist (m, nsample) 65 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 66 | if (pt_idx >= m) return; 67 | 68 | new_xyz += pt_idx * 3; 69 | idx += pt_idx * nsample; 70 | dist2 += pt_idx * nsample; 71 | 72 | int bt_idx = random_ball_query_utils::get_bt_idx(pt_idx, new_offset); 73 | int start; 74 | if (bt_idx == 0) 75 | start = 0; 76 | else 77 | start = offset[bt_idx - 1]; 78 | int end = offset[bt_idx]; 79 | 80 | float max_radius2 = max_radius * max_radius; 81 | float min_radius2 = min_radius * min_radius; 82 | float new_x = new_xyz[0]; 83 | float new_y = new_xyz[1]; 84 | float new_z = new_xyz[2]; 85 | 86 | int cnt = 0; 87 | 88 | for(int i = start; i < end; i++){ 89 | float x = xyz[order[i] * 3 + 0]; 90 | float y = xyz[order[i] * 3 + 1]; 91 | float z = xyz[order[i] * 3 + 2]; 92 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + (new_z - z) * (new_z - z); 93 | 94 | if (d2 <= 1e-5 || (d2 >= min_radius2 && d2 < max_radius2)){ 95 | dist2[cnt] = d2; 96 | idx[cnt] = order[i]; 97 | cnt += 1; 98 | if (cnt >= nsample) break; 99 | } 100 | } 101 | 102 | if (cnt < nsample) { 103 | for (int i = cnt; i < nsample; i++){ 104 | idx[i] = -1; 105 | dist2[i] = 1e10; 106 | } 107 | } 108 | } 109 | 110 | void random_ball_query_cuda_launcher(int m, int nsample, 111 | float min_radius, float max_radius, const int *order, 112 | const float *xyz, const float *new_xyz, 113 | const int *offset, const int *new_offset, 114 | int *idx, float *dist2) { 115 | // input: new_xyz: (m, 3), xyz: (n, 3), idx: (m, nsample) 116 | dim3 blocks(DIVUP(m, THREADS_PER_BLOCK)); 117 | dim3 threads(THREADS_PER_BLOCK); 118 | random_ball_query_cuda_kernel<<>>(m, nsample, 119 | min_radius, max_radius, order, 120 | xyz, new_xyz, 121 | offset, new_offset, 122 | idx, dist2); 123 | } 124 | -------------------------------------------------------------------------------- /libs/pointops/src/random_ball_query/random_ball_query_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _RANDOM_BALL_QUERY_CUDA_KERNEL 2 | #define _RANDOM_BALL_QUERY_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void random_ball_query_cuda(int m, int nsample, 8 | float min_radius, float max_radius, at::Tensor order_tensor, 9 | at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, 10 | at::Tensor offset_tensor, at::Tensor new_offset_tensor, 11 | at::Tensor idx_tensor, at::Tensor dist2_tensor); 12 | 13 | #ifdef __cplusplus 14 | extern "C" { 15 | #endif 16 | 17 | void random_ball_query_cuda_launcher(int m, int nsample, 18 | float min_radius, float max_radius, const int *order, 19 | const float *xyz, const float *new_xyz, 20 | const int *offset, const int *new_offset, 21 | int *idx, float *dist2); 22 | 23 | #ifdef __cplusplus 24 | } 25 | #endif 26 | #endif 27 | -------------------------------------------------------------------------------- /libs/pointops/src/sampling/sampling_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "sampling_cuda_kernel.h" 5 | 6 | 7 | void farthest_point_sampling_cuda(int b, int n, at::Tensor xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor tmp_tensor, at::Tensor idx_tensor) 8 | { 9 | const float *xyz = xyz_tensor.data_ptr(); 10 | const int *offset = offset_tensor.data_ptr(); 11 | const int *new_offset = new_offset_tensor.data_ptr(); 12 | float *tmp = tmp_tensor.data_ptr(); 13 | int *idx = idx_tensor.data_ptr(); 14 | farthest_point_sampling_cuda_launcher(b, n, xyz, offset, new_offset, tmp, idx); 15 | } 16 | -------------------------------------------------------------------------------- /libs/pointops/src/sampling/sampling_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _SAMPLING_CUDA_KERNEL 2 | #define _SAMPLING_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void farthest_point_sampling_cuda(int b, int n, at::Tensor xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor tmp_tensor, at::Tensor idx_tensor); 8 | 9 | #ifdef __cplusplus 10 | extern "C" { 11 | #endif 12 | 13 | void farthest_point_sampling_cuda_launcher(int b, int n, const float *xyz, const int *offset, const int *new_offset, float *tmp, int *idx); 14 | 15 | #ifdef __cplusplus 16 | } 17 | #endif 18 | #endif 19 | -------------------------------------------------------------------------------- /libs/pointops/src/subtraction/subtraction_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "subtraction_cuda_kernel.h" 5 | 6 | 7 | void subtraction_forward_cuda(int n, int nsample, int c, at::Tensor input1_tensor, at::Tensor input2_tensor, at::Tensor idx_tensor, at::Tensor output_tensor) 8 | { 9 | const float *input1 = input1_tensor.data_ptr(); 10 | const float *input2 = input2_tensor.data_ptr(); 11 | const int *idx = idx_tensor.data_ptr(); 12 | float *output = output_tensor.data_ptr(); 13 | subtraction_forward_cuda_launcher(n, nsample, c, input1, input2, idx, output); 14 | } 15 | 16 | void subtraction_backward_cuda(int n, int nsample, int c, at::Tensor idx_tensor, at::Tensor grad_output_tensor, at::Tensor grad_input1_tensor, at::Tensor grad_input2_tensor) 17 | { 18 | const int *idx = idx_tensor.data_ptr(); 19 | const float *grad_output = grad_output_tensor.data_ptr(); 20 | float *grad_input1 = grad_input1_tensor.data_ptr(); 21 | float *grad_input2 = grad_input2_tensor.data_ptr(); 22 | subtraction_backward_cuda_launcher(n, nsample, c, idx, grad_output, grad_input1, grad_input2); 23 | } 24 | -------------------------------------------------------------------------------- /libs/pointops/src/subtraction/subtraction_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "subtraction_cuda_kernel.h" 3 | 4 | 5 | __global__ void subtraction_forward_cuda_kernel(int n, int nsample, int c, const float *input1, const float *input2, const int *idx, float *output) { 6 | // input: input1: (n, c), input2: (n, c), idx: (n, nsample), output: (n, nsample, c) 7 | int index = blockIdx.x * blockDim.x + threadIdx.x; 8 | if (index >= n * nsample * c) return; 9 | const int c_idx = index % c; 10 | const int nsample_idx = (index / c) % nsample; 11 | const int n_idx = index / nsample / c; 12 | const int idx_idx = n_idx * nsample + nsample_idx; 13 | const int input1_idx = n_idx * c + c_idx; 14 | const int input2_idx = idx[idx_idx] * c + c_idx; 15 | output[index] = input1[input1_idx] - input2[input2_idx]; 16 | } 17 | 18 | __global__ void subtraction_backward_cuda_kernel(int n, int nsample, int c, const int *idx, const float *grad_output, float *grad_input1, float *grad_input2) { 19 | // input: grad_output: (n, nsample, c), output: grad_input1: (n, c), grad_input2: (n, c) 20 | int index = blockIdx.x * blockDim.x + threadIdx.x; 21 | if (index >= n * nsample * c) return; 22 | const int c_idx = index % c; 23 | const int nsample_idx = (index / c) % nsample; 24 | const int n_idx = index / nsample / c; 25 | const int idx_idx = n_idx * nsample + nsample_idx; 26 | const int input1_idx = n_idx * c + c_idx; 27 | const int input2_idx = idx[idx_idx] * c + c_idx; 28 | atomicAdd(grad_input1 + input1_idx, grad_output[index]); 29 | atomicAdd(grad_input2 + input2_idx, -grad_output[index]); 30 | } 31 | 32 | void subtraction_forward_cuda_launcher(int n, int nsample, int c, const float *input1, const float *input2, const int *idx, float *output) { 33 | // input: input1: (n, c), input2: (n, c), idx: (n, nsample), output: (n, nsample, c) 34 | dim3 blocks(DIVUP(n * nsample * c, THREADS_PER_BLOCK)); 35 | dim3 threads(THREADS_PER_BLOCK); 36 | subtraction_forward_cuda_kernel<<>>(n, nsample, c, input1, input2, idx, output); 37 | } 38 | 39 | void subtraction_backward_cuda_launcher(int n, int nsample, int c, const int *idx, const float *grad_output, float *grad_input1, float *grad_input2) { 40 | // input: grad_output: (n, nsample, c), output: grad_input1: (n, c), grad_input2: (n, c) 41 | dim3 blocks(DIVUP(n * nsample * c, THREADS_PER_BLOCK)); 42 | dim3 threads(THREADS_PER_BLOCK); 43 | subtraction_backward_cuda_kernel<<>>(n, nsample, c, idx, grad_output, grad_input1, grad_input2); 44 | } 45 | -------------------------------------------------------------------------------- /libs/pointops/src/subtraction/subtraction_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _SUBTRACTION_CUDA_KERNEL 2 | #define _SUBTRACTION_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void subtraction_forward_cuda(int n, int nsample, int c, at::Tensor input1_tensor, at::Tensor input2_tensor, at::Tensor idx_tensor, at::Tensor output_tensor); 8 | void subtraction_backward_cuda(int n, int nsample, int c, at::Tensor idx_tensor, at::Tensor grad_output_tensor, at::Tensor grad_input1_tensor, at::Tensor grad_input2_tensor); 9 | 10 | #ifdef __cplusplus 11 | extern "C" { 12 | #endif 13 | 14 | void subtraction_forward_cuda_launcher(int n, int nsample, int c, const float *input1, const float *input2, const int *idx, float *output); 15 | void subtraction_backward_cuda_launcher(int n, int nsample, int c, const int *idx, const float *grad_output, float *grad_input1, float *grad_input2); 16 | 17 | #ifdef __cplusplus 18 | } 19 | #endif 20 | #endif 21 | -------------------------------------------------------------------------------- /libs/smooth-sampler/.gitignore: -------------------------------------------------------------------------------- 1 | build 2 | dist 3 | *.egg-info 4 | __pycache__ -------------------------------------------------------------------------------- /libs/smooth-sampler/LICENSE: -------------------------------------------------------------------------------- 1 | 2 | All contributions by Tymoteusz Bleja: 3 | Copyright (c) 2022 Tymoteusz Bleja 4 | 5 | From PyTorch: 6 | 7 | Copyright (c) 2016- Facebook, Inc (Adam Paszke) 8 | Copyright (c) 2014- Facebook, Inc (Soumith Chintala) 9 | Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert) 10 | Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu) 11 | Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu) 12 | Copyright (c) 2011-2013 NYU (Clement Farabet) 13 | Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston) 14 | Copyright (c) 2006 Idiap Research Institute (Samy Bengio) 15 | Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz) 16 | 17 | From Caffe2: 18 | 19 | Copyright (c) 2016-present, Facebook Inc. All rights reserved. 20 | 21 | All contributions by Facebook: 22 | Copyright (c) 2016 Facebook Inc. 23 | 24 | All contributions by Google: 25 | Copyright (c) 2015 Google Inc. 26 | All rights reserved. 27 | 28 | All contributions by Yangqing Jia: 29 | Copyright (c) 2015 Yangqing Jia 30 | All rights reserved. 31 | 32 | All contributions by Kakao Brain: 33 | Copyright 2019-2020 Kakao Brain 34 | 35 | All contributions by Cruise LLC: 36 | Copyright (c) 2022 Cruise LLC. 37 | All rights reserved. 38 | 39 | All contributions from Caffe: 40 | Copyright(c) 2013, 2014, 2015, the respective contributors 41 | All rights reserved. 42 | 43 | All other contributions: 44 | Copyright(c) 2015, 2016 the respective contributors 45 | All rights reserved. 46 | 47 | Caffe2 uses a copyright model similar to Caffe: each contributor holds 48 | copyright over their contributions to Caffe2. The project versioning records 49 | all such contribution and copyright details. If a contributor wants to further 50 | mark their specific copyright on a particular contribution, they should 51 | indicate their copyright solely in the commit message of the change when it is 52 | committed. 53 | 54 | All rights reserved. 55 | 56 | Redistribution and use in source and binary forms, with or without 57 | modification, are permitted provided that the following conditions are met: 58 | 59 | 1. Redistributions of source code must retain the above copyright 60 | notice, this list of conditions and the following disclaimer. 61 | 62 | 2. Redistributions in binary form must reproduce the above copyright 63 | notice, this list of conditions and the following disclaimer in the 64 | documentation and/or other materials provided with the distribution. 65 | 66 | 3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America 67 | and IDIAP Research Institute nor the names of its contributors may be 68 | used to endorse or promote products derived from this software without 69 | specific prior written permission. 70 | 71 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 72 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 73 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 74 | ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE 75 | LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 76 | CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 77 | SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 78 | INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 79 | CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 80 | ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 81 | POSSIBILITY OF SUCH DAMAGE. -------------------------------------------------------------------------------- /libs/smooth-sampler/README.md: -------------------------------------------------------------------------------- 1 | # Smooth sampler 2 | 3 | A drop-in replacement for Pytorch's [grid_sample](https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html) supporting smoothstep activation for interpolation weights (proposed in [Instant NGP](https://nvlabs.github.io/instant-ngp/) by Müller et al.) as well as double backpropagation. Currently supports 3D inputs and trilinear interpolation mode. Based on Pytorch's native grid_sampler code. Used in [GO-Surf](https://jingwenwang95.github.io/go_surf/) by Wang et al. 4 | 5 | # Installation 6 | 7 | On Python3 environment with Pytorch CUDA installation run: 8 | 9 | ```bash 10 | pip install git+https://github.com/tymoteuszb/smooth-sampler 11 | ``` 12 | 13 | # Usage 14 | 15 | The API is consistent with Pytorch's grid_sample: 16 | 17 | ```python 18 | import torch 19 | from smooth_sampler import SmoothSampler 20 | 21 | align_corners = True 22 | padding_mode = "zeros" 23 | 24 | input = (torch.rand([2,2,2,3,11], device="cuda")).requires_grad_(True) 25 | grid = (torch.rand([2,2,1,5,3], device="cuda") * 2. - 1.).requires_grad_(True) 26 | 27 | out1 = SmoothSampler.apply(input, grid, padding_mode, align_corners, False) 28 | out2 = torch.nn.functional.grid_sample(input, grid, padding_mode=padding_mode, align_corners=align_corners) 29 | assert torch.allclose(out1, out2) 30 | 31 | grad1_input, grad1_grid = torch.autograd.grad(out1, [input, grid], torch.ones_like(out1), create_graph=True) 32 | grad2_input, grad2_grid = torch.autograd.grad(out2, [input, grid], torch.ones_like(out2), create_graph=True) 33 | assert torch.allclose(grad1_input, grad2_input) 34 | assert torch.allclose(grad1_grid, grad2_grid) 35 | 36 | loss1 = out1.sum() + grad1_input.sum() + grad1_grid.sum() 37 | loss1.backward() # Works! 38 | 39 | loss2 = out2.sum() + grad2_input.sum() + grad2_grid.sum() 40 | loss2.backward() # RuntimeException: derivative for aten::grid_sampler_3d_backward is not implemented 41 | ``` 42 | 43 | # Citation 44 | 45 | If you use this code in your project, please consider adding a citation: 46 | 47 | ``` 48 | @article{wang2022go-surf, 49 | title={GO-Surf: Neural Feature Grid Optimization for Fast, High-Fidelity RGB-D Surface Reconstruction}, 50 | author={Wang, Jingwen and Bleja, Tymoteusz and Agapito, Lourdes}, 51 | journal={arXiv preprint}, 52 | year={2022} 53 | } 54 | ``` 55 | -------------------------------------------------------------------------------- /libs/smooth-sampler/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from torch.utils import cpp_extension 3 | 4 | setup( 5 | name="smooth_sampler", 6 | version="0.0.1", 7 | author="Tymoteusz Bleja, Jingwen Wang", 8 | author_email="tymoteusz.bleja@gmail.com, jingwen.wang.17@ucl.ac.uk", 9 | description=( 10 | "Trilinear sampler with smoothstep and double backpropagation - Pytorch extension." 11 | ), 12 | ext_modules=[ 13 | cpp_extension.CUDAExtension( 14 | "smooth_sampler._C", 15 | [ 16 | "smooth_sampler/csrc/smooth_sampler.cpp", 17 | "smooth_sampler/csrc/smooth_sampler_kernel.cu", 18 | ], 19 | ) 20 | ], 21 | packages=["smooth_sampler"], 22 | cmdclass={"build_ext": cpp_extension.BuildExtension}, 23 | classifiers=[ 24 | "Development Status :: 3 - Alpha", 25 | "Environment :: GPU :: NVIDIA CUDA", 26 | "Topic :: Utilities", 27 | "License :: OSI Approved :: BSD License", 28 | "Programming Language :: C++", 29 | "Programming Language :: CUDA", 30 | "Programming Language :: Python :: 3 :: Only", 31 | "Topic :: Scientific/Engineering :: Artificial Intelligence", 32 | ], 33 | ) 34 | -------------------------------------------------------------------------------- /libs/smooth-sampler/smooth_sampler/__init__.py: -------------------------------------------------------------------------------- 1 | from smooth_sampler.modules import SmoothSampler 2 | -------------------------------------------------------------------------------- /libs/smooth-sampler/smooth_sampler/csrc/smooth_sampler.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | Based on https://github.com/pytorch/pytorch/blob/v1.12.0/aten/src/ATen/native/cuda/GridSampler.cpp 3 | */ 4 | 5 | #include 6 | #include 7 | 8 | #define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor") 9 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 10 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 11 | 12 | void launch_smooth_sampler_forward_kernel( 13 | const torch::TensorBase &output, const torch::TensorBase &input, const torch::TensorBase &grid, 14 | int64_t padding_mode, bool align_corners, bool apply_smoothstep); 15 | 16 | void launch_smooth_sampler_backward_kernel( 17 | const torch::TensorBase& grad_input, const torch::TensorBase &grad_grid, 18 | const torch::TensorBase& grad_output, const torch::TensorBase& input, 19 | const torch::TensorBase& grid, int64_t padding_mode, bool align_corners, 20 | bool apply_smoothstep, bool input_requires_grad); 21 | 22 | void launch_smooth_sampler_backward_backward_kernel( 23 | const torch::TensorBase& grad_input, 24 | const torch::TensorBase& grad_grid, 25 | const torch::TensorBase& grad_grad_out, 26 | const torch::TensorBase& input, 27 | const torch::TensorBase& grid, 28 | const torch::TensorBase& grad_out_input, 29 | const torch::TensorBase& grad_out_grid, 30 | const torch::TensorBase& grad_output, 31 | int64_t padding_mode, 32 | const bool align_corners, 33 | const bool apply_smoothstep, 34 | const bool input_requires_grad); 35 | 36 | torch::Tensor smooth_sampler_forward(torch::Tensor input, torch::Tensor grid, 37 | int64_t padding_mode, bool align_corners, bool apply_smoothstep) { 38 | CHECK_INPUT(input) 39 | CHECK_INPUT(grid) 40 | const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); 41 | 42 | auto in_size = input.sizes(); 43 | auto grid_size = grid.sizes(); 44 | auto output = torch::empty( 45 | {in_size[0], in_size[1], grid_size[1], grid_size[2], grid_size[3]}, 46 | input.options()); 47 | launch_smooth_sampler_forward_kernel( 48 | output, input, grid, padding_mode, align_corners, apply_smoothstep); 49 | return output; 50 | } 51 | 52 | std::tuple smooth_sampler_backward(torch::Tensor grad_output, torch::Tensor input, 53 | torch::Tensor grid, int64_t padding_mode, bool align_corners, 54 | bool apply_smoothstep, bool input_requires_grad) { 55 | CHECK_INPUT(grad_output) 56 | CHECK_INPUT(input) 57 | CHECK_INPUT(grid) 58 | const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); 59 | 60 | torch::Tensor grad_input = ([&]() { 61 | if (input_requires_grad) { 62 | return torch::zeros_like(input, LEGACY_CONTIGUOUS_MEMORY_FORMAT); 63 | } else { 64 | return torch::Tensor(); 65 | } 66 | })(); 67 | auto grad_grid = torch::empty_like(grid, LEGACY_CONTIGUOUS_MEMORY_FORMAT); 68 | launch_smooth_sampler_backward_kernel( 69 | grad_input, grad_grid, grad_output, input, 70 | grid, padding_mode, align_corners, apply_smoothstep, input_requires_grad); 71 | return std::make_tuple(grad_input, grad_grid); 72 | } 73 | 74 | std::tuple smooth_sampler_backward_backward(torch::Tensor grad_out_input, torch::Tensor grad_out_grid, 75 | torch::Tensor input, torch::Tensor grid, torch::Tensor grad_output, 76 | int64_t padding_mode, bool align_corners, 77 | bool apply_smoothstep, bool input_requires_grad) { 78 | CHECK_INPUT(grad_out_input) 79 | CHECK_INPUT(grad_out_grid) 80 | CHECK_INPUT(input) 81 | CHECK_INPUT(grid) 82 | CHECK_INPUT(grad_output) 83 | const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); 84 | 85 | auto grad_input = torch::zeros_like(input, LEGACY_CONTIGUOUS_MEMORY_FORMAT); 86 | auto grad_grid = torch::empty_like(grid, LEGACY_CONTIGUOUS_MEMORY_FORMAT); 87 | auto grad_grad_out = torch::zeros_like(grad_output, LEGACY_CONTIGUOUS_MEMORY_FORMAT); 88 | launch_smooth_sampler_backward_backward_kernel(grad_input, grad_grid, grad_grad_out, input, grid, 89 | grad_out_input, grad_out_grid, grad_output, 90 | padding_mode, align_corners, apply_smoothstep, input_requires_grad); 91 | return std::make_tuple(grad_input, grad_grid, grad_grad_out); 92 | } 93 | 94 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 95 | m.def("forward", &smooth_sampler_forward, "Smooth sampler forward (CUDA)"); 96 | m.def("backward", &smooth_sampler_backward, "Smooth sampler backward (CUDA)"); 97 | m.def("backward_backward", &smooth_sampler_backward_backward, "Smooth sampler backward backward (CUDA)"); 98 | } -------------------------------------------------------------------------------- /libs/smooth-sampler/smooth_sampler/modules.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from smooth_sampler import _C 3 | 4 | 5 | def padding_mode_enum(padding_mode): 6 | if padding_mode == "zeros": 7 | return 0 8 | elif padding_mode == "border": 9 | return 1 10 | else: # padding_mode == 'reflection' 11 | return 2 12 | 13 | 14 | class SmoothSamplerBackward(torch.autograd.Function): 15 | @staticmethod 16 | def forward( 17 | ctx, 18 | input, 19 | grid, 20 | grad_out, 21 | padding_mode="zeros", 22 | align_corners=True, 23 | apply_smoothstep=False, 24 | ): 25 | ctx.align_corners = align_corners 26 | ctx.apply_smoothstep = apply_smoothstep 27 | ctx.padding_mode = padding_mode 28 | grad_input, grad_grid = _C.backward( 29 | grad_out, 30 | input, 31 | grid, 32 | padding_mode_enum(padding_mode), 33 | ctx.align_corners, 34 | apply_smoothstep, 35 | input.requires_grad, 36 | ) 37 | ctx.save_for_backward(input, grid, grad_out) 38 | 39 | return grad_input, grad_grid 40 | 41 | @staticmethod 42 | def backward(ctx, grad_out_input, grad_out_grid): 43 | input, grid, grad_out = ctx.saved_tensors 44 | 45 | input_requires_grad = ( 46 | grad_out_input is not None and (grad_out_input != 0.0).any().item() 47 | ) 48 | grad_input, grad_grid, grad_grad_out = _C.backward_backward( 49 | grad_out_input.contiguous(), 50 | grad_out_grid.contiguous(), 51 | input, 52 | grid, 53 | grad_out, 54 | padding_mode_enum(ctx.padding_mode), 55 | ctx.align_corners, 56 | ctx.apply_smoothstep, 57 | input_requires_grad, 58 | ) 59 | 60 | return grad_input, grad_grid, grad_grad_out, None, None, None 61 | 62 | 63 | class SmoothSampler(torch.autograd.Function): 64 | @staticmethod 65 | def forward( 66 | ctx, 67 | input, 68 | grid, 69 | padding_mode="zeros", 70 | align_corners=True, 71 | apply_smoothstep=False, 72 | ): 73 | output = _C.forward( 74 | input, 75 | grid, 76 | padding_mode_enum(padding_mode), 77 | align_corners, 78 | apply_smoothstep, 79 | ) 80 | ctx.save_for_backward(input, grid) 81 | ctx.align_corners = align_corners 82 | ctx.apply_smoothstep = apply_smoothstep 83 | ctx.padding_mode = padding_mode 84 | return output 85 | 86 | @staticmethod 87 | def backward(ctx, grad_out): 88 | input, grid = ctx.saved_tensors 89 | 90 | if (grad_out == 0.0).all().item(): 91 | return torch.zeros_like(input), torch.zeros_like(grid), None, None, None 92 | 93 | d_input, d_grid = SmoothSamplerBackward.apply( 94 | input, 95 | grid, 96 | grad_out.contiguous(), 97 | ctx.padding_mode, 98 | ctx.align_corners, 99 | ctx.apply_smoothstep, 100 | ) 101 | return d_input, d_grid, None, None, None 102 | 103 | 104 | if __name__ == "__main__": 105 | torch.manual_seed(3) 106 | torch.cuda.manual_seed(3) 107 | 108 | for padding_mode in ["zeros", "border", "reflection"]: 109 | for align_corners in [True, False]: 110 | input = (torch.rand([2, 2, 2, 3, 11], device="cuda")).requires_grad_(True) 111 | grid = ( 112 | torch.rand([2, 2, 1, 5, 3], device="cuda") * 2.0 - 1.0 113 | ).requires_grad_(True) 114 | 115 | # SmoothSampler forward vs native forward 116 | out1 = SmoothSampler.apply(input, grid, padding_mode, align_corners, False) 117 | out2 = torch.nn.functional.grid_sample( 118 | input, grid, padding_mode=padding_mode, align_corners=align_corners 119 | ) 120 | assert torch.allclose(out1, out2) 121 | 122 | # SmoothSampler backward vs native backward 123 | grad1_input, grad1_grid = torch.autograd.grad( 124 | out1, [input, grid], torch.ones_like(out1), create_graph=True 125 | ) 126 | grad2_input, grad2_grid = torch.autograd.grad( 127 | out2, [input, grid], torch.ones_like(out2), create_graph=True 128 | ) 129 | assert torch.allclose(grad1_input, grad2_input) 130 | assert torch.allclose(grad1_grid, grad2_grid) 131 | 132 | for apply_smoothstep in [True, False]: 133 | input = ( 134 | torch.rand([2, 2, 2, 3, 11], device="cuda").double() 135 | ).requires_grad_(True) 136 | grid = ( 137 | (torch.rand([2, 2, 1, 5, 3], device="cuda") * 2.0 - 1.0) 138 | .double() 139 | .requires_grad_(True) 140 | ) 141 | 142 | # Analytic gradients vs finite differences gradients 143 | torch.autograd.gradcheck( 144 | SmoothSampler.apply, 145 | [input, grid, padding_mode, align_corners, apply_smoothstep], 146 | eps=1e-4, 147 | atol=1e-3, 148 | rtol=1e-2, 149 | ) 150 | torch.autograd.gradgradcheck( 151 | SmoothSampler.apply, 152 | [input, grid, padding_mode, align_corners, apply_smoothstep], 153 | eps=1e-4, 154 | atol=1e-3, 155 | rtol=1e-2, 156 | ) 157 | -------------------------------------------------------------------------------- /ponder/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/OpenGVLab/PonderV2/c152af1df9d04a98bc19696e8d69d047834a80db/ponder/__init__.py -------------------------------------------------------------------------------- /ponder/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | from .builder import build_dataset 2 | from .dataloader import MultiDatasetDataloader 3 | from .defaults import DefaultDataset 4 | from .nuscenes import NuScenesDataset 5 | from .s3dis import S3DISDataset, S3DISRGBDDataset 6 | from .scannet import ScanNet200Dataset, ScanNetDataset, ScanNetRGBDDataset 7 | from .structure3d import Structured3DDataset, Structured3DRGBDDataset 8 | from .utils import collate_fn, point_collate_fn 9 | -------------------------------------------------------------------------------- /ponder/datasets/builder.py: -------------------------------------------------------------------------------- 1 | """ 2 | Dataset Builder 3 | 4 | Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) 5 | Please cite our work if the code is helpful to you. 6 | """ 7 | 8 | from ponder.utils.registry import Registry 9 | 10 | DATASETS = Registry("datasets") 11 | 12 | 13 | def build_dataset(cfg): 14 | """Build datasets.""" 15 | return DATASETS.build(cfg) 16 | -------------------------------------------------------------------------------- /ponder/datasets/dataloader.py: -------------------------------------------------------------------------------- 1 | import weakref 2 | from functools import partial 3 | 4 | import torch 5 | import torch.utils.data 6 | 7 | import ponder.utils.comm as comm 8 | from ponder.utils.env import set_seed 9 | 10 | from .defaults import ConcatDataset 11 | from .utils import point_collate_fn 12 | 13 | 14 | class MultiDatasetDummySampler: 15 | def __init__(self): 16 | self.dataloader = None 17 | 18 | def set_epoch(self, epoch): 19 | if comm.get_world_size() > 1: 20 | for dataloader in self.dataloader.dataloaders: 21 | dataloader.sampler.set_epoch(epoch) 22 | return 23 | 24 | 25 | class MultiDatasetDataloader: 26 | """ 27 | Multiple Datasets Dataloader, batch data from a same dataset and mix up ratio determined by loop of each sub dataset. 28 | The overall length is determined by the main dataset (first) and loop of concat dataset. 29 | """ 30 | 31 | def __init__( 32 | self, 33 | concat_dataset: ConcatDataset, 34 | batch_size_per_gpu: int, 35 | num_worker_per_gpu: int, 36 | mix_prob=0, 37 | seed=None, 38 | max_point=-1, 39 | ): 40 | self.datasets = concat_dataset.datasets 41 | self.ratios = [dataset.loop for dataset in self.datasets] 42 | # reset data loop, original loop serve as ratios 43 | for dataset in self.datasets: 44 | dataset.loop = 1 45 | # determine union training epoch by main dataset 46 | self.datasets[0].loop = concat_dataset.loop 47 | # build sub-dataloaders 48 | num_workers = num_worker_per_gpu // len(self.datasets) 49 | self.dataloaders = [] 50 | for dataset_id, dataset in enumerate(self.datasets): 51 | if comm.get_world_size() > 1: 52 | sampler = torch.utils.data.distributed.DistributedSampler(dataset) 53 | else: 54 | sampler = None 55 | 56 | init_fn = ( 57 | partial( 58 | self._worker_init_fn, 59 | dataset_id=dataset_id, 60 | num_workers=num_workers, 61 | num_datasets=len(self.datasets), 62 | rank=comm.get_rank(), 63 | seed=seed, 64 | ) 65 | if seed is not None 66 | else None 67 | ) 68 | self.dataloaders.append( 69 | torch.utils.data.DataLoader( 70 | dataset, 71 | batch_size=batch_size_per_gpu, 72 | shuffle=(sampler is None), 73 | num_workers=num_worker_per_gpu, 74 | sampler=sampler, 75 | collate_fn=partial( 76 | point_collate_fn, mix_prob=mix_prob, max_point=max_point 77 | ), 78 | pin_memory=True, 79 | worker_init_fn=init_fn, 80 | drop_last=True, 81 | persistent_workers=True, 82 | ) 83 | ) 84 | self.sampler = MultiDatasetDummySampler() 85 | self.sampler.dataloader = weakref.proxy(self) 86 | 87 | def __iter__(self): 88 | iterator = [iter(dataloader) for dataloader in self.dataloaders] 89 | while True: 90 | for i in range(len(self.ratios)): 91 | for _ in range(self.ratios[i]): 92 | try: 93 | batch = next(iterator[i]) 94 | except StopIteration: 95 | if i == 0: 96 | return 97 | else: 98 | iterator[i] = iter(self.dataloaders[i]) 99 | batch = next(iterator[i]) 100 | yield batch 101 | 102 | def __len__(self): 103 | main_data_loader_length = len(self.dataloaders[0]) 104 | return ( 105 | main_data_loader_length // self.ratios[0] * sum(self.ratios) 106 | + main_data_loader_length % self.ratios[0] 107 | ) 108 | 109 | @staticmethod 110 | def _worker_init_fn(worker_id, num_workers, dataset_id, num_datasets, rank, seed): 111 | worker_seed = ( 112 | num_workers * num_datasets * rank 113 | + num_workers * dataset_id 114 | + worker_id 115 | + seed 116 | ) 117 | set_seed(worker_seed) 118 | -------------------------------------------------------------------------------- /ponder/datasets/preprocessing/s3dis/preprocess_s3dis_voxelized.py: -------------------------------------------------------------------------------- 1 | """ 2 | Preprocessing Script for S3DIS 3 | Parsing normal vectors has a large consumption of memory. Please reduce max_workers if memory is limited. 4 | 5 | Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) 6 | Please cite our work if the code is helpful to you. 7 | """ 8 | 9 | import argparse 10 | import glob 11 | import multiprocessing as mp 12 | import os 13 | from concurrent.futures import ProcessPoolExecutor 14 | from itertools import repeat 15 | 16 | import numpy as np 17 | import torch 18 | 19 | from ponder.datasets.transform import GridSample 20 | 21 | 22 | def voxelize_parser(data_path, dataset_root, output_root, voxel_size): 23 | print(f"Parsing data: {data_path}") 24 | out_path = data_path.replace(dataset_root, output_root) 25 | os.makedirs(os.path.dirname(out_path), exist_ok=True) 26 | data = torch.load(data_path) 27 | data = GridSample( 28 | grid_size=voxel_size, hash_type="fnv", mode="train", keys=data.keys() 29 | )(data) 30 | torch.save(data, out_path) 31 | 32 | 33 | def main_process(): 34 | parser = argparse.ArgumentParser() 35 | parser.add_argument( 36 | "--dataset_root", required=True, help="Path to processed S3DIS dataset" 37 | ) 38 | parser.add_argument( 39 | "--output_root", 40 | required=True, 41 | help="Output path where area folders will be located", 42 | ) 43 | parser.add_argument( 44 | "--voxel_size", default=0.01, type=float, help="Voxel size for voxelization" 45 | ) 46 | args = parser.parse_args() 47 | 48 | data_list = glob.glob(os.path.join(args.dataset_root, "*/*.pth")) 49 | # Preprocess data. 50 | print("Processing scenes...") 51 | pool = ProcessPoolExecutor(max_workers=mp.cpu_count()) 52 | # pool = ProcessPoolExecutor(max_workers=1) 53 | _ = list( 54 | pool.map( 55 | voxelize_parser, 56 | data_list, 57 | repeat(args.dataset_root), 58 | repeat(args.output_root), 59 | repeat(args.voxel_size), 60 | ) 61 | ) 62 | 63 | 64 | if __name__ == "__main__": 65 | main_process() 66 | -------------------------------------------------------------------------------- /ponder/datasets/preprocessing/scannet/meta_data/classes_ObjClassification-ShapeNetCore55.txt: -------------------------------------------------------------------------------- 1 | 1 trash 2 | 3 basket 3 | 4 bathtub 4 | 5 bed 5 | 9 shelf 6 | 13 cabinet 7 | 18 chair 8 | 20 keyboard 9 | 22 tv 10 | 30 lamp 11 | 31 laptop 12 | 35 microwave 13 | 39 pillow 14 | 42 printer 15 | 47 sofa 16 | 48 stove 17 | 49 table 18 | -------------------------------------------------------------------------------- /ponder/datasets/preprocessing/scannet/meta_data/classes_SemVoxLabel-nyu40id.txt: -------------------------------------------------------------------------------- 1 | 1 wall 2 | 2 floor 3 | 3 cabinet 4 | 4 bed 5 | 5 chair 6 | 6 sofa 7 | 7 table 8 | 8 door 9 | 9 window 10 | 10 bookshelf 11 | 11 picture 12 | 12 counter 13 | 14 desk 14 | 16 curtain 15 | 24 refridgerator 16 | 28 shower curtain 17 | 33 toilet 18 | 34 sink 19 | 36 bathtub 20 | 39 otherfurniture -------------------------------------------------------------------------------- /ponder/datasets/preprocessing/scannet/meta_data/scannet_means.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/OpenGVLab/PonderV2/c152af1df9d04a98bc19696e8d69d047834a80db/ponder/datasets/preprocessing/scannet/meta_data/scannet_means.npz -------------------------------------------------------------------------------- /ponder/datasets/preprocessing/scannet/meta_data/scannetv1_test.txt: -------------------------------------------------------------------------------- 1 | scene0568_00 2 | scene0568_01 3 | scene0568_02 4 | scene0304_00 5 | scene0488_00 6 | scene0488_01 7 | scene0412_00 8 | scene0412_01 9 | scene0217_00 10 | scene0019_00 11 | scene0019_01 12 | scene0414_00 13 | scene0575_00 14 | scene0575_01 15 | scene0575_02 16 | scene0426_00 17 | scene0426_01 18 | scene0426_02 19 | scene0426_03 20 | scene0549_00 21 | scene0549_01 22 | scene0578_00 23 | scene0578_01 24 | scene0578_02 25 | scene0665_00 26 | scene0665_01 27 | scene0050_00 28 | scene0050_01 29 | scene0050_02 30 | scene0257_00 31 | scene0025_00 32 | scene0025_01 33 | scene0025_02 34 | scene0583_00 35 | scene0583_01 36 | scene0583_02 37 | scene0701_00 38 | scene0701_01 39 | scene0701_02 40 | scene0580_00 41 | scene0580_01 42 | scene0565_00 43 | scene0169_00 44 | scene0169_01 45 | scene0655_00 46 | scene0655_01 47 | scene0655_02 48 | scene0063_00 49 | scene0221_00 50 | scene0221_01 51 | scene0591_00 52 | scene0591_01 53 | scene0591_02 54 | scene0678_00 55 | scene0678_01 56 | scene0678_02 57 | scene0462_00 58 | scene0427_00 59 | scene0595_00 60 | scene0193_00 61 | scene0193_01 62 | scene0164_00 63 | scene0164_01 64 | scene0164_02 65 | scene0164_03 66 | scene0598_00 67 | scene0598_01 68 | scene0598_02 69 | scene0599_00 70 | scene0599_01 71 | scene0599_02 72 | scene0328_00 73 | scene0300_00 74 | scene0300_01 75 | scene0354_00 76 | scene0458_00 77 | scene0458_01 78 | scene0423_00 79 | scene0423_01 80 | scene0423_02 81 | scene0307_00 82 | scene0307_01 83 | scene0307_02 84 | scene0606_00 85 | scene0606_01 86 | scene0606_02 87 | scene0432_00 88 | scene0432_01 89 | scene0608_00 90 | scene0608_01 91 | scene0608_02 92 | scene0651_00 93 | scene0651_01 94 | scene0651_02 95 | scene0430_00 96 | scene0430_01 97 | scene0689_00 98 | scene0357_00 99 | scene0357_01 100 | scene0574_00 101 | scene0574_01 102 | scene0574_02 103 | scene0329_00 104 | scene0329_01 105 | scene0329_02 106 | scene0153_00 107 | scene0153_01 108 | scene0616_00 109 | scene0616_01 110 | scene0671_00 111 | scene0671_01 112 | scene0618_00 113 | scene0382_00 114 | scene0382_01 115 | scene0490_00 116 | scene0621_00 117 | scene0607_00 118 | scene0607_01 119 | scene0149_00 120 | scene0695_00 121 | scene0695_01 122 | scene0695_02 123 | scene0695_03 124 | scene0389_00 125 | scene0377_00 126 | scene0377_01 127 | scene0377_02 128 | scene0342_00 129 | scene0139_00 130 | scene0629_00 131 | scene0629_01 132 | scene0629_02 133 | scene0496_00 134 | scene0633_00 135 | scene0633_01 136 | scene0518_00 137 | scene0652_00 138 | scene0406_00 139 | scene0406_01 140 | scene0406_02 141 | scene0144_00 142 | scene0144_01 143 | scene0494_00 144 | scene0278_00 145 | scene0278_01 146 | scene0316_00 147 | scene0609_00 148 | scene0609_01 149 | scene0609_02 150 | scene0609_03 151 | scene0084_00 152 | scene0084_01 153 | scene0084_02 154 | scene0696_00 155 | scene0696_01 156 | scene0696_02 157 | scene0351_00 158 | scene0351_01 159 | scene0643_00 160 | scene0644_00 161 | scene0645_00 162 | scene0645_01 163 | scene0645_02 164 | scene0081_00 165 | scene0081_01 166 | scene0081_02 167 | scene0647_00 168 | scene0647_01 169 | scene0535_00 170 | scene0353_00 171 | scene0353_01 172 | scene0353_02 173 | scene0559_00 174 | scene0559_01 175 | scene0559_02 176 | scene0593_00 177 | scene0593_01 178 | scene0246_00 179 | scene0653_00 180 | scene0653_01 181 | scene0064_00 182 | scene0064_01 183 | scene0356_00 184 | scene0356_01 185 | scene0356_02 186 | scene0030_00 187 | scene0030_01 188 | scene0030_02 189 | scene0222_00 190 | scene0222_01 191 | scene0338_00 192 | scene0338_01 193 | scene0338_02 194 | scene0378_00 195 | scene0378_01 196 | scene0378_02 197 | scene0660_00 198 | scene0553_00 199 | scene0553_01 200 | scene0553_02 201 | scene0527_00 202 | scene0663_00 203 | scene0663_01 204 | scene0663_02 205 | scene0664_00 206 | scene0664_01 207 | scene0664_02 208 | scene0334_00 209 | scene0334_01 210 | scene0334_02 211 | scene0046_00 212 | scene0046_01 213 | scene0046_02 214 | scene0203_00 215 | scene0203_01 216 | scene0203_02 217 | scene0088_00 218 | scene0088_01 219 | scene0088_02 220 | scene0088_03 221 | scene0086_00 222 | scene0086_01 223 | scene0086_02 224 | scene0670_00 225 | scene0670_01 226 | scene0256_00 227 | scene0256_01 228 | scene0256_02 229 | scene0249_00 230 | scene0441_00 231 | scene0658_00 232 | scene0704_00 233 | scene0704_01 234 | scene0187_00 235 | scene0187_01 236 | scene0131_00 237 | scene0131_01 238 | scene0131_02 239 | scene0207_00 240 | scene0207_01 241 | scene0207_02 242 | scene0461_00 243 | scene0011_00 244 | scene0011_01 245 | scene0343_00 246 | scene0251_00 247 | scene0077_00 248 | scene0077_01 249 | scene0684_00 250 | scene0684_01 251 | scene0550_00 252 | scene0686_00 253 | scene0686_01 254 | scene0686_02 255 | scene0208_00 256 | scene0500_00 257 | scene0500_01 258 | scene0552_00 259 | scene0552_01 260 | scene0648_00 261 | scene0648_01 262 | scene0435_00 263 | scene0435_01 264 | scene0435_02 265 | scene0435_03 266 | scene0690_00 267 | scene0690_01 268 | scene0693_00 269 | scene0693_01 270 | scene0693_02 271 | scene0700_00 272 | scene0700_01 273 | scene0700_02 274 | scene0699_00 275 | scene0231_00 276 | scene0231_01 277 | scene0231_02 278 | scene0697_00 279 | scene0697_01 280 | scene0697_02 281 | scene0697_03 282 | scene0474_00 283 | scene0474_01 284 | scene0474_02 285 | scene0474_03 286 | scene0474_04 287 | scene0474_05 288 | scene0355_00 289 | scene0355_01 290 | scene0146_00 291 | scene0146_01 292 | scene0146_02 293 | scene0196_00 294 | scene0702_00 295 | scene0702_01 296 | scene0702_02 297 | scene0314_00 298 | scene0277_00 299 | scene0277_01 300 | scene0277_02 301 | scene0095_00 302 | scene0095_01 303 | scene0015_00 304 | scene0100_00 305 | scene0100_01 306 | scene0100_02 307 | scene0558_00 308 | scene0558_01 309 | scene0558_02 310 | scene0685_00 311 | scene0685_01 312 | scene0685_02 313 | -------------------------------------------------------------------------------- /ponder/datasets/preprocessing/scannet/meta_data/scannetv1_val.txt: -------------------------------------------------------------------------------- 1 | scene0534_00 2 | scene0534_01 3 | scene0319_00 4 | scene0273_00 5 | scene0273_01 6 | scene0225_00 7 | scene0198_00 8 | scene0003_00 9 | scene0003_01 10 | scene0003_02 11 | scene0409_00 12 | scene0409_01 13 | scene0331_00 14 | scene0331_01 15 | scene0505_00 16 | scene0505_01 17 | scene0505_02 18 | scene0505_03 19 | scene0505_04 20 | scene0506_00 21 | scene0057_00 22 | scene0057_01 23 | scene0074_00 24 | scene0074_01 25 | scene0074_02 26 | scene0091_00 27 | scene0112_00 28 | scene0112_01 29 | scene0112_02 30 | scene0240_00 31 | scene0102_00 32 | scene0102_01 33 | scene0513_00 34 | scene0514_00 35 | scene0514_01 36 | scene0537_00 37 | scene0516_00 38 | scene0516_01 39 | scene0495_00 40 | scene0617_00 41 | scene0133_00 42 | scene0520_00 43 | scene0520_01 44 | scene0635_00 45 | scene0635_01 46 | scene0054_00 47 | scene0473_00 48 | scene0473_01 49 | scene0524_00 50 | scene0524_01 51 | scene0379_00 52 | scene0471_00 53 | scene0471_01 54 | scene0471_02 55 | scene0566_00 56 | scene0248_00 57 | scene0248_01 58 | scene0248_02 59 | scene0529_00 60 | scene0529_01 61 | scene0529_02 62 | scene0391_00 63 | scene0264_00 64 | scene0264_01 65 | scene0264_02 66 | scene0675_00 67 | scene0675_01 68 | scene0350_00 69 | scene0350_01 70 | scene0350_02 71 | scene0450_00 72 | scene0068_00 73 | scene0068_01 74 | scene0237_00 75 | scene0237_01 76 | scene0365_00 77 | scene0365_01 78 | scene0365_02 79 | scene0605_00 80 | scene0605_01 81 | scene0539_00 82 | scene0539_01 83 | scene0539_02 84 | scene0540_00 85 | scene0540_01 86 | scene0540_02 87 | scene0170_00 88 | scene0170_01 89 | scene0170_02 90 | scene0433_00 91 | scene0340_00 92 | scene0340_01 93 | scene0340_02 94 | scene0160_00 95 | scene0160_01 96 | scene0160_02 97 | scene0160_03 98 | scene0160_04 99 | scene0059_00 100 | scene0059_01 101 | scene0059_02 102 | scene0056_00 103 | scene0056_01 104 | scene0478_00 105 | scene0478_01 106 | scene0548_00 107 | scene0548_01 108 | scene0548_02 109 | scene0204_00 110 | scene0204_01 111 | scene0204_02 112 | scene0033_00 113 | scene0145_00 114 | scene0483_00 115 | scene0508_00 116 | scene0508_01 117 | scene0508_02 118 | scene0180_00 119 | scene0148_00 120 | scene0556_00 121 | scene0556_01 122 | scene0416_00 123 | scene0416_01 124 | scene0416_02 125 | scene0416_03 126 | scene0416_04 127 | scene0073_00 128 | scene0073_01 129 | scene0073_02 130 | scene0073_03 131 | scene0034_00 132 | scene0034_01 133 | scene0034_02 134 | scene0639_00 135 | scene0561_00 136 | scene0561_01 137 | scene0298_00 138 | scene0692_00 139 | scene0692_01 140 | scene0692_02 141 | scene0692_03 142 | scene0692_04 143 | scene0642_00 144 | scene0642_01 145 | scene0642_02 146 | scene0642_03 147 | scene0630_00 148 | scene0630_01 149 | scene0630_02 150 | scene0630_03 151 | scene0630_04 152 | scene0630_05 153 | scene0630_06 154 | scene0706_00 155 | scene0567_00 156 | scene0567_01 157 | -------------------------------------------------------------------------------- /ponder/datasets/preprocessing/scannet/meta_data/scannetv2_test.txt: -------------------------------------------------------------------------------- 1 | scene0707_00 2 | scene0708_00 3 | scene0709_00 4 | scene0710_00 5 | scene0711_00 6 | scene0712_00 7 | scene0713_00 8 | scene0714_00 9 | scene0715_00 10 | scene0716_00 11 | scene0717_00 12 | scene0718_00 13 | scene0719_00 14 | scene0720_00 15 | scene0721_00 16 | scene0722_00 17 | scene0723_00 18 | scene0724_00 19 | scene0725_00 20 | scene0726_00 21 | scene0727_00 22 | scene0728_00 23 | scene0729_00 24 | scene0730_00 25 | scene0731_00 26 | scene0732_00 27 | scene0733_00 28 | scene0734_00 29 | scene0735_00 30 | scene0736_00 31 | scene0737_00 32 | scene0738_00 33 | scene0739_00 34 | scene0740_00 35 | scene0741_00 36 | scene0742_00 37 | scene0743_00 38 | scene0744_00 39 | scene0745_00 40 | scene0746_00 41 | scene0747_00 42 | scene0748_00 43 | scene0749_00 44 | scene0750_00 45 | scene0751_00 46 | scene0752_00 47 | scene0753_00 48 | scene0754_00 49 | scene0755_00 50 | scene0756_00 51 | scene0757_00 52 | scene0758_00 53 | scene0759_00 54 | scene0760_00 55 | scene0761_00 56 | scene0762_00 57 | scene0763_00 58 | scene0764_00 59 | scene0765_00 60 | scene0766_00 61 | scene0767_00 62 | scene0768_00 63 | scene0769_00 64 | scene0770_00 65 | scene0771_00 66 | scene0772_00 67 | scene0773_00 68 | scene0774_00 69 | scene0775_00 70 | scene0776_00 71 | scene0777_00 72 | scene0778_00 73 | scene0779_00 74 | scene0780_00 75 | scene0781_00 76 | scene0782_00 77 | scene0783_00 78 | scene0784_00 79 | scene0785_00 80 | scene0786_00 81 | scene0787_00 82 | scene0788_00 83 | scene0789_00 84 | scene0790_00 85 | scene0791_00 86 | scene0792_00 87 | scene0793_00 88 | scene0794_00 89 | scene0795_00 90 | scene0796_00 91 | scene0797_00 92 | scene0798_00 93 | scene0799_00 94 | scene0800_00 95 | scene0801_00 96 | scene0802_00 97 | scene0803_00 98 | scene0804_00 99 | scene0805_00 100 | scene0806_00 101 | -------------------------------------------------------------------------------- /ponder/datasets/preprocessing/scannet/meta_data/scannetv2_val.txt: -------------------------------------------------------------------------------- 1 | scene0568_00 2 | scene0568_01 3 | scene0568_02 4 | scene0304_00 5 | scene0488_00 6 | scene0488_01 7 | scene0412_00 8 | scene0412_01 9 | scene0217_00 10 | scene0019_00 11 | scene0019_01 12 | scene0414_00 13 | scene0575_00 14 | scene0575_01 15 | scene0575_02 16 | scene0426_00 17 | scene0426_01 18 | scene0426_02 19 | scene0426_03 20 | scene0549_00 21 | scene0549_01 22 | scene0578_00 23 | scene0578_01 24 | scene0578_02 25 | scene0665_00 26 | scene0665_01 27 | scene0050_00 28 | scene0050_01 29 | scene0050_02 30 | scene0257_00 31 | scene0025_00 32 | scene0025_01 33 | scene0025_02 34 | scene0583_00 35 | scene0583_01 36 | scene0583_02 37 | scene0701_00 38 | scene0701_01 39 | scene0701_02 40 | scene0580_00 41 | scene0580_01 42 | scene0565_00 43 | scene0169_00 44 | scene0169_01 45 | scene0655_00 46 | scene0655_01 47 | scene0655_02 48 | scene0063_00 49 | scene0221_00 50 | scene0221_01 51 | scene0591_00 52 | scene0591_01 53 | scene0591_02 54 | scene0678_00 55 | scene0678_01 56 | scene0678_02 57 | scene0462_00 58 | scene0427_00 59 | scene0595_00 60 | scene0193_00 61 | scene0193_01 62 | scene0164_00 63 | scene0164_01 64 | scene0164_02 65 | scene0164_03 66 | scene0598_00 67 | scene0598_01 68 | scene0598_02 69 | scene0599_00 70 | scene0599_01 71 | scene0599_02 72 | scene0328_00 73 | scene0300_00 74 | scene0300_01 75 | scene0354_00 76 | scene0458_00 77 | scene0458_01 78 | scene0423_00 79 | scene0423_01 80 | scene0423_02 81 | scene0307_00 82 | scene0307_01 83 | scene0307_02 84 | scene0606_00 85 | scene0606_01 86 | scene0606_02 87 | scene0432_00 88 | scene0432_01 89 | scene0608_00 90 | scene0608_01 91 | scene0608_02 92 | scene0651_00 93 | scene0651_01 94 | scene0651_02 95 | scene0430_00 96 | scene0430_01 97 | scene0689_00 98 | scene0357_00 99 | scene0357_01 100 | scene0574_00 101 | scene0574_01 102 | scene0574_02 103 | scene0329_00 104 | scene0329_01 105 | scene0329_02 106 | scene0153_00 107 | scene0153_01 108 | scene0616_00 109 | scene0616_01 110 | scene0671_00 111 | scene0671_01 112 | scene0618_00 113 | scene0382_00 114 | scene0382_01 115 | scene0490_00 116 | scene0621_00 117 | scene0607_00 118 | scene0607_01 119 | scene0149_00 120 | scene0695_00 121 | scene0695_01 122 | scene0695_02 123 | scene0695_03 124 | scene0389_00 125 | scene0377_00 126 | scene0377_01 127 | scene0377_02 128 | scene0342_00 129 | scene0139_00 130 | scene0629_00 131 | scene0629_01 132 | scene0629_02 133 | scene0496_00 134 | scene0633_00 135 | scene0633_01 136 | scene0518_00 137 | scene0652_00 138 | scene0406_00 139 | scene0406_01 140 | scene0406_02 141 | scene0144_00 142 | scene0144_01 143 | scene0494_00 144 | scene0278_00 145 | scene0278_01 146 | scene0316_00 147 | scene0609_00 148 | scene0609_01 149 | scene0609_02 150 | scene0609_03 151 | scene0084_00 152 | scene0084_01 153 | scene0084_02 154 | scene0696_00 155 | scene0696_01 156 | scene0696_02 157 | scene0351_00 158 | scene0351_01 159 | scene0643_00 160 | scene0644_00 161 | scene0645_00 162 | scene0645_01 163 | scene0645_02 164 | scene0081_00 165 | scene0081_01 166 | scene0081_02 167 | scene0647_00 168 | scene0647_01 169 | scene0535_00 170 | scene0353_00 171 | scene0353_01 172 | scene0353_02 173 | scene0559_00 174 | scene0559_01 175 | scene0559_02 176 | scene0593_00 177 | scene0593_01 178 | scene0246_00 179 | scene0653_00 180 | scene0653_01 181 | scene0064_00 182 | scene0064_01 183 | scene0356_00 184 | scene0356_01 185 | scene0356_02 186 | scene0030_00 187 | scene0030_01 188 | scene0030_02 189 | scene0222_00 190 | scene0222_01 191 | scene0338_00 192 | scene0338_01 193 | scene0338_02 194 | scene0378_00 195 | scene0378_01 196 | scene0378_02 197 | scene0660_00 198 | scene0553_00 199 | scene0553_01 200 | scene0553_02 201 | scene0527_00 202 | scene0663_00 203 | scene0663_01 204 | scene0663_02 205 | scene0664_00 206 | scene0664_01 207 | scene0664_02 208 | scene0334_00 209 | scene0334_01 210 | scene0334_02 211 | scene0046_00 212 | scene0046_01 213 | scene0046_02 214 | scene0203_00 215 | scene0203_01 216 | scene0203_02 217 | scene0088_00 218 | scene0088_01 219 | scene0088_02 220 | scene0088_03 221 | scene0086_00 222 | scene0086_01 223 | scene0086_02 224 | scene0670_00 225 | scene0670_01 226 | scene0256_00 227 | scene0256_01 228 | scene0256_02 229 | scene0249_00 230 | scene0441_00 231 | scene0658_00 232 | scene0704_00 233 | scene0704_01 234 | scene0187_00 235 | scene0187_01 236 | scene0131_00 237 | scene0131_01 238 | scene0131_02 239 | scene0207_00 240 | scene0207_01 241 | scene0207_02 242 | scene0461_00 243 | scene0011_00 244 | scene0011_01 245 | scene0343_00 246 | scene0251_00 247 | scene0077_00 248 | scene0077_01 249 | scene0684_00 250 | scene0684_01 251 | scene0550_00 252 | scene0686_00 253 | scene0686_01 254 | scene0686_02 255 | scene0208_00 256 | scene0500_00 257 | scene0500_01 258 | scene0552_00 259 | scene0552_01 260 | scene0648_00 261 | scene0648_01 262 | scene0435_00 263 | scene0435_01 264 | scene0435_02 265 | scene0435_03 266 | scene0690_00 267 | scene0690_01 268 | scene0693_00 269 | scene0693_01 270 | scene0693_02 271 | scene0700_00 272 | scene0700_01 273 | scene0700_02 274 | scene0699_00 275 | scene0231_00 276 | scene0231_01 277 | scene0231_02 278 | scene0697_00 279 | scene0697_01 280 | scene0697_02 281 | scene0697_03 282 | scene0474_00 283 | scene0474_01 284 | scene0474_02 285 | scene0474_03 286 | scene0474_04 287 | scene0474_05 288 | scene0355_00 289 | scene0355_01 290 | scene0146_00 291 | scene0146_01 292 | scene0146_02 293 | scene0196_00 294 | scene0702_00 295 | scene0702_01 296 | scene0702_02 297 | scene0314_00 298 | scene0277_00 299 | scene0277_01 300 | scene0277_02 301 | scene0095_00 302 | scene0095_01 303 | scene0015_00 304 | scene0100_00 305 | scene0100_01 306 | scene0100_02 307 | scene0558_00 308 | scene0558_01 309 | scene0558_02 310 | scene0685_00 311 | scene0685_01 312 | scene0685_02 313 | -------------------------------------------------------------------------------- /ponder/datasets/preprocessing/scannet/reader.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | # All rights reserved. 3 | # 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import argparse 8 | import csv 9 | import os 10 | import sys 11 | import zipfile 12 | from glob import glob 13 | 14 | import cv2 15 | import imageio.v2 as imageio 16 | import numpy as np 17 | 18 | from ponder.datasets.preprocessing.scannet.SensorData import SensorData 19 | 20 | # params 21 | parser = argparse.ArgumentParser() 22 | # data paths 23 | parser.add_argument("--scans_path", required=True, help="path to scans folder") 24 | parser.add_argument("--output_path", required=True, help="path to output folder") 25 | parser.add_argument( 26 | "--export_depth_images", dest="export_depth_images", action="store_true" 27 | ) 28 | parser.add_argument( 29 | "--export_color_images", dest="export_color_images", action="store_true" 30 | ) 31 | parser.add_argument("--export_poses", dest="export_poses", action="store_true") 32 | parser.add_argument( 33 | "--export_intrinsics", dest="export_intrinsics", action="store_true" 34 | ) 35 | parser.add_argument("--export_label", dest="export_label", action="store_true") 36 | parser.set_defaults( 37 | export_depth_images=False, 38 | export_color_images=False, 39 | export_poses=False, 40 | export_intrinsics=False, 41 | export_label=False, 42 | ) 43 | 44 | opt = parser.parse_args() 45 | print(opt) 46 | 47 | 48 | def represents_int(s): 49 | try: 50 | int(s) 51 | return True 52 | except ValueError: 53 | return False 54 | 55 | 56 | def read_label_mapping(filename, label_from="raw_category", label_to="nyu40id"): 57 | # assert os.path.isfile(filename) 58 | mapping = dict() 59 | # print(filename) 60 | with open(filename, "r") as csvfile: 61 | reader = csv.DictReader(csvfile, delimiter="\t") 62 | for row in reader: 63 | mapping[row[label_from]] = int(row[label_to]) 64 | # if ints convert 65 | if represents_int(list(mapping.keys())[0]): 66 | mapping = {int(k): v for k, v in mapping.items()} 67 | return mapping 68 | 69 | 70 | def main(): 71 | scans = glob(opt.scans_path + "/*") 72 | scans.sort() 73 | 74 | label_mapping = None 75 | if opt.export_label: 76 | root = os.path.dirname(opt.scans_path) 77 | label_map = read_label_mapping( 78 | filename=os.path.join(root, "scannetv2-labels.combined.tsv"), 79 | label_from="id", 80 | label_to="nyu40id", 81 | ) 82 | 83 | for scan in scans: 84 | scenename = scan.split("/")[-1] 85 | filename = os.path.join(scan, scenename + ".sens") 86 | if not os.path.exists(opt.output_path): 87 | os.makedirs(opt.output_path) 88 | # os.makedirs(os.path.join(opt.output_path, 'depth')) 89 | # os.makedirs(os.path.join(opt.output_path, 'color')) 90 | # os.makedirs(os.path.join(opt.output_path, 'pose')) 91 | # os.makedirs(os.path.join(opt.output_path, 'intrinsic')) 92 | os.makedirs(os.path.join(opt.output_path, scenename)) 93 | # load the data 94 | print("loading %s..." % filename) 95 | sd = SensorData(filename) 96 | print("loaded!\n") 97 | if opt.export_depth_images: 98 | # sd.export_depth_images(os.path.join(opt.output_path, 'depth', scenename)) 99 | sd.export_depth_images(os.path.join(opt.output_path, scenename, "depth")) 100 | if opt.export_color_images: 101 | # sd.export_color_images(os.path.join(opt.output_path, 'color', scenename)) 102 | sd.export_color_images(os.path.join(opt.output_path, scenename, "color")) 103 | if opt.export_poses: 104 | # sd.export_poses(os.path.join(opt.output_path, 'pose', scenename)) 105 | sd.export_poses(os.path.join(opt.output_path, scenename, "pose")) 106 | if opt.export_intrinsics: 107 | # sd.export_intrinsics(os.path.join(opt.output_path, 'intrinsic', scenename)) 108 | sd.export_intrinsics(os.path.join(opt.output_path, scenename, "intrinsic")) 109 | 110 | os.system(f"cp {scan}/scene*.txt {opt.output_path}/{scenename}/") 111 | 112 | if opt.export_label: 113 | 114 | def map_label_image(image, label_mapping): 115 | mapped = np.copy(image) 116 | for k, v in label_mapping.items(): 117 | mapped[image == k] = v 118 | return mapped.astype(np.uint8) 119 | 120 | label_zip_path = os.path.join( 121 | opt.scans_path, scenename, f"{scenename}_2d-label-filt.zip" 122 | ) 123 | print("process labels") 124 | with open(label_zip_path, "rb") as f: 125 | zip_file = zipfile.ZipFile(f) 126 | for frame in range(0, len(sd.frames)): 127 | label_file = f"label-filt/{frame}.png" 128 | with zip_file.open(label_file) as lf: 129 | image = np.array(imageio.imread(lf)) 130 | 131 | mapped_image = map_label_image(image, label_map) 132 | output_path = os.path.join(opt.output_path, scenename, "label") 133 | os.makedirs(output_path, exist_ok=True) 134 | print("output:", output_path) 135 | cv2.imwrite(os.path.join(output_path, f"{frame}.png"), mapped_image) 136 | 137 | 138 | if __name__ == "__main__": 139 | main() 140 | -------------------------------------------------------------------------------- /ponder/datasets/structure3d.py: -------------------------------------------------------------------------------- 1 | """ 2 | Structured3D Datasets 3 | 4 | Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) 5 | Please cite our work if the code is helpful to you. 6 | """ 7 | 8 | import glob 9 | import os 10 | from collections.abc import Sequence 11 | 12 | import numpy as np 13 | import torch 14 | 15 | from .builder import DATASETS 16 | from .defaults import DefaultDataset 17 | 18 | 19 | @DATASETS.register_module() 20 | class Structured3DDataset(DefaultDataset): 21 | def get_data_list(self): 22 | if isinstance(self.split, str): 23 | data_list = glob.glob(os.path.join(self.data_root, self.split, "*/*.pth")) 24 | elif isinstance(self.split, Sequence): 25 | data_list = [] 26 | for split in self.split: 27 | data_list += glob.glob(os.path.join(self.data_root, split, "*/*.pth")) 28 | else: 29 | raise NotImplementedError 30 | return data_list 31 | 32 | def get_data_name(self, idx): 33 | file_path = self.data_list[idx % len(self.data_list)] 34 | dir_path, file_name = os.path.split(file_path) 35 | scene_name = os.path.basename(dir_path) 36 | room_name = os.path.splitext(file_name)[0] 37 | data_name = f"{scene_name}_{room_name}" 38 | return data_name 39 | 40 | 41 | @DATASETS.register_module() 42 | class Structured3DRGBDDataset(Structured3DDataset): 43 | def __init__( 44 | self, 45 | split="train", 46 | data_root="data/dataset", 47 | transform=None, 48 | test_mode=False, 49 | test_cfg=None, 50 | num_cameras=5, 51 | render_semantic=True, 52 | loop=1, 53 | ): 54 | super(Structured3DRGBDDataset, self).__init__( 55 | split=split, 56 | data_root=data_root, 57 | transform=transform, 58 | test_mode=test_mode, 59 | test_cfg=test_cfg, 60 | loop=loop, 61 | ) 62 | self.num_cameras = num_cameras 63 | self.render_semantic = render_semantic 64 | 65 | def get_data_list(self): 66 | if isinstance(self.split, str): 67 | data_list = glob.glob(os.path.join(self.data_root, self.split, "*/*.pth")) 68 | elif isinstance(self.split, Sequence): 69 | data_list = [] 70 | for split in self.split: 71 | data_list += glob.glob(os.path.join(self.data_root, split, "*/*.pth")) 72 | else: 73 | raise NotImplementedError 74 | 75 | print("Filtering Structured3D RGBD dataset...") 76 | filtered_data_list = [] 77 | for data_path in data_list: 78 | rgbd_paths = glob.glob( 79 | os.path.join(data_path.split(".pth")[0] + "_rgbd", "*.pth") 80 | ) 81 | if len(rgbd_paths) <= 0: 82 | # print(f"{data_path} has no rgbd data.") 83 | continue 84 | filtered_data_list.append(data_path) 85 | print( 86 | f"Finish filtering! Totally {len(filtered_data_list)} from {len(data_list)} data." 87 | ) 88 | return filtered_data_list 89 | 90 | def get_data(self, idx): 91 | data_path = self.data_list[idx % len(self.data_list)] 92 | data = torch.load(data_path) 93 | coord = data["coord"] 94 | color = data["color"] 95 | normal = data["normal"] 96 | if "semantic_gt" in data.keys(): 97 | segment = data["semantic_gt"].reshape([-1]) 98 | else: 99 | segment = np.ones(coord.shape[0]) * -1 100 | 101 | rgbd_paths = glob.glob( 102 | os.path.join(data_path.split(".pth")[0] + "_rgbd", "*.pth") 103 | ) 104 | 105 | if len(rgbd_paths) <= 0: 106 | print(f"{data_path} has no rgbd data.") 107 | return self.get_data(np.random.randint(0, self.__len__())) 108 | 109 | rgbd_paths = np.random.choice( 110 | rgbd_paths, self.num_cameras, replace=self.num_cameras > len(rgbd_paths) 111 | ) 112 | rgbd_dicts = [torch.load(p) for p in rgbd_paths] 113 | 114 | for i in range(len(rgbd_dicts)): 115 | if (rgbd_dicts[i]["depth_mask"]).mean() < 0.25: 116 | os.rename(rgbd_paths[i], rgbd_paths[i] + ".bad") 117 | return self.get_data(idx) 118 | 119 | data_dict = dict( 120 | coord=coord, 121 | normal=normal, 122 | color=color, 123 | segment=segment, 124 | intrinsic=np.stack([d["intrinsic"] for d in rgbd_dicts], axis=0).astype( 125 | np.float32 126 | ), 127 | extrinsic=np.stack( 128 | [np.linalg.inv(d["extrinsic"]) for d in rgbd_dicts], axis=0 129 | ).astype(np.float32), 130 | rgb=np.stack([d["rgb"].astype(np.float32) for d in rgbd_dicts], axis=0), 131 | depth=np.stack( 132 | [ 133 | d["depth"].astype(np.float32) 134 | * d["depth_mask"].astype(np.float32) 135 | * (d["depth"] < 65535).astype(np.float32) 136 | for d in rgbd_dicts 137 | ], 138 | axis=0, 139 | ), 140 | depth_scale=1.0 / 1000.0, 141 | ) 142 | if self.render_semantic: 143 | for d in rgbd_dicts: 144 | d["semantic_map"][d["semantic_map"] <= 0] = -1 145 | d["semantic_map"][d["semantic_map"] > 40] = -1 146 | d["semantic_map"] = d["semantic_map"].astype(np.int16) 147 | data_dict.update( 148 | dict(semantic=np.stack([d["semantic_map"] for d in rgbd_dicts], axis=0)) 149 | ) 150 | return data_dict 151 | -------------------------------------------------------------------------------- /ponder/datasets/utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Utils for Datasets 3 | 4 | Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) 5 | Please cite our work if the code is helpful to you. 6 | """ 7 | 8 | import random 9 | from collections.abc import Mapping, Sequence 10 | 11 | import numpy as np 12 | import torch 13 | from torch.utils.data.dataloader import default_collate 14 | 15 | 16 | def collate_fn(batch, max_point=-1): 17 | """ 18 | collate function for point cloud which support dict and list, 19 | 'coord' is necessary to determine 'offset' 20 | """ 21 | if not isinstance(batch, Sequence): 22 | raise TypeError(f"{batch.dtype} is not supported.") 23 | 24 | # we drop a large data if it exceeds max_point 25 | # note that directly drop the last one may cause problem 26 | if max_point > 0: 27 | accum_num_points = 0 28 | ret_batches = [] 29 | for batch_id, data in enumerate(batch): 30 | num_coords = data["coord"].shape[0] 31 | if accum_num_points + num_coords > max_point: 32 | print( 33 | "SKIP: accum_num_points", accum_num_points, "num_coords", num_coords 34 | ) 35 | continue 36 | accum_num_points += num_coords 37 | ret_batches.append(data) 38 | return collate_fn(ret_batches) 39 | 40 | if isinstance(batch[0], torch.Tensor): 41 | return torch.cat(list(batch)) 42 | elif isinstance(batch[0], str): 43 | # str is also a kind of Sequence, judgement should before Sequence 44 | return list(batch) 45 | elif isinstance(batch[0], Sequence): 46 | for data in batch: 47 | data.append(torch.tensor([data[0].shape[0]])) 48 | batch = [collate_fn(samples) for samples in zip(*batch)] 49 | batch[-1] = torch.cumsum(batch[-1], dim=0).int() 50 | return batch 51 | elif isinstance(batch[0], Mapping): 52 | batch = {key: collate_fn([d[key] for d in batch]) for key in batch[0]} 53 | for key in batch.keys(): 54 | if "offset" in key: 55 | batch[key] = torch.cumsum(batch[key], dim=0) 56 | return batch 57 | else: 58 | return default_collate(batch) 59 | 60 | 61 | def point_collate_fn(batch, mix_prob=0, max_point=-1): 62 | assert isinstance( 63 | batch[0], Mapping 64 | ) # currently, only support input_dict, rather than input_list 65 | batch = collate_fn(batch, max_point=max_point) 66 | if "offset" in batch.keys(): 67 | # Mix3d (https://arxiv.org/pdf/2110.02210.pdf) 68 | if random.random() < mix_prob: 69 | batch["offset"] = torch.cat( 70 | [batch["offset"][1:-1:2], batch["offset"][-1].unsqueeze(0)], dim=0 71 | ) 72 | return batch 73 | 74 | 75 | def gaussian_kernel(dist2: np.array, a: float = 1, c: float = 5): 76 | return a * np.exp(-dist2 / (2 * c**2)) 77 | -------------------------------------------------------------------------------- /ponder/engines/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/OpenGVLab/PonderV2/c152af1df9d04a98bc19696e8d69d047834a80db/ponder/engines/__init__.py -------------------------------------------------------------------------------- /ponder/engines/hooks/__init__.py: -------------------------------------------------------------------------------- 1 | from .builder import build_hooks 2 | from .default import HookBase 3 | from .evaluator import * 4 | from .misc import * 5 | -------------------------------------------------------------------------------- /ponder/engines/hooks/builder.py: -------------------------------------------------------------------------------- 1 | """ 2 | Hook Builder 3 | 4 | Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) 5 | Please cite our work if the code is helpful to you. 6 | """ 7 | 8 | from ponder.utils.registry import Registry 9 | 10 | HOOKS = Registry("hooks") 11 | 12 | 13 | def build_hooks(cfg): 14 | hooks = [] 15 | for hook_cfg in cfg: 16 | hooks.append(HOOKS.build(hook_cfg)) 17 | return hooks 18 | -------------------------------------------------------------------------------- /ponder/engines/hooks/default.py: -------------------------------------------------------------------------------- 1 | """ 2 | Default Hook 3 | 4 | Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) 5 | Please cite our work if the code is helpful to you. 6 | """ 7 | 8 | 9 | class HookBase: 10 | """ 11 | Base class for hooks that can be registered with :class:`TrainerBase`. 12 | """ 13 | 14 | trainer = None # A weak reference to the trainer object. 15 | 16 | def before_train(self): 17 | pass 18 | 19 | def before_epoch(self): 20 | pass 21 | 22 | def before_step(self): 23 | pass 24 | 25 | def after_step(self): 26 | pass 27 | 28 | def after_epoch(self): 29 | pass 30 | 31 | def after_train(self): 32 | pass 33 | -------------------------------------------------------------------------------- /ponder/models/__init__.py: -------------------------------------------------------------------------------- 1 | from .builder import build_model 2 | 3 | # Semantic Segmentation 4 | from .default import DefaultClassifier, DefaultSegmentor 5 | 6 | # PPT 7 | from .point_prompt_training import * 8 | 9 | # Pretraining 10 | from .ponder import * 11 | 12 | # Backbones 13 | from .sparse_unet import * 14 | 15 | # from .point_group import * 16 | -------------------------------------------------------------------------------- /ponder/models/builder.py: -------------------------------------------------------------------------------- 1 | """ 2 | Model Builder 3 | 4 | Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) 5 | Please cite our work if the code is helpful to you. 6 | """ 7 | 8 | from ponder.utils.registry import Registry 9 | 10 | MODELS = Registry("models") 11 | MODULES = Registry("modules") 12 | 13 | 14 | def build_model(cfg): 15 | """Build models.""" 16 | return MODELS.build(cfg) 17 | -------------------------------------------------------------------------------- /ponder/models/default.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | from ponder.models.losses import build_criteria 4 | 5 | from .builder import MODELS, build_model 6 | 7 | 8 | @MODELS.register_module() 9 | class DefaultSegmentor(nn.Module): 10 | def __init__(self, backbone=None, criteria=None): 11 | super().__init__() 12 | self.backbone = build_model(backbone) 13 | self.criteria = build_criteria(criteria) 14 | 15 | def forward(self, input_dict): 16 | if "condition" in input_dict.keys(): 17 | # PPT (https://arxiv.org/abs/2308.09718) 18 | # currently, only support one batch one condition 19 | input_dict["condition"] = input_dict["condition"][0] 20 | seg_logits = self.backbone(input_dict) 21 | # train 22 | if self.training: 23 | loss = self.criteria(seg_logits, input_dict["segment"]) 24 | return dict(loss=loss) 25 | # eval 26 | elif "segment" in input_dict.keys(): 27 | loss = self.criteria(seg_logits, input_dict["segment"]) 28 | return dict(loss=loss, seg_logits=seg_logits) 29 | # test 30 | else: 31 | return dict(seg_logits=seg_logits) 32 | 33 | 34 | @MODELS.register_module() 35 | class DefaultClassifier(nn.Module): 36 | def __init__( 37 | self, 38 | backbone=None, 39 | criteria=None, 40 | num_classes=40, 41 | backbone_embed_dim=256, 42 | ): 43 | super().__init__() 44 | self.backbone = build_model(backbone) 45 | self.criteria = build_criteria(criteria) 46 | self.num_classes = num_classes 47 | self.backbone_embed_dim = backbone_embed_dim 48 | self.cls_head = nn.Sequential( 49 | nn.Linear(backbone_embed_dim, 256), 50 | nn.BatchNorm1d(256), 51 | nn.ReLU(inplace=True), 52 | nn.Dropout(p=0.5), 53 | nn.Linear(256, 128), 54 | nn.BatchNorm1d(128), 55 | nn.ReLU(inplace=True), 56 | nn.Dropout(p=0.5), 57 | nn.Linear(128, num_classes), 58 | ) 59 | 60 | def forward(self, input_dict): 61 | feat = self.backbone(input_dict) 62 | cls_logits = self.cls_head(feat) 63 | if self.training: 64 | loss = self.criteria(cls_logits, input_dict["category"]) 65 | return dict(loss=loss) 66 | elif "category" in input_dict.keys(): 67 | loss = self.criteria(cls_logits, input_dict["category"]) 68 | return dict(loss=loss, cls_logits=cls_logits) 69 | else: 70 | return dict(cls_logits=cls_logits) 71 | -------------------------------------------------------------------------------- /ponder/models/losses/__init__.py: -------------------------------------------------------------------------------- 1 | from .builder import build_criteria 2 | from .lovasz import LovaszLoss 3 | from .misc import BinaryFocalLoss, CrossEntropyLoss, DiceLoss, FocalLoss, SmoothCELoss 4 | -------------------------------------------------------------------------------- /ponder/models/losses/builder.py: -------------------------------------------------------------------------------- 1 | """ 2 | Criteria Builder 3 | 4 | Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) 5 | Please cite our work if the code is helpful to you. 6 | """ 7 | 8 | from ponder.utils.registry import Registry 9 | 10 | LOSSES = Registry("losses") 11 | 12 | 13 | class Criteria(object): 14 | def __init__(self, cfg=None): 15 | self.cfg = cfg if cfg is not None else [] 16 | self.criteria = [] 17 | for loss_cfg in self.cfg: 18 | self.criteria.append(LOSSES.build(cfg=loss_cfg)) 19 | 20 | def __call__(self, pred, target): 21 | if len(self.criteria) == 0: 22 | # loss computation occur in model 23 | return pred 24 | loss = 0 25 | for c in self.criteria: 26 | loss += c(pred, target) 27 | return loss 28 | 29 | 30 | def build_criteria(cfg): 31 | return Criteria(cfg) 32 | -------------------------------------------------------------------------------- /ponder/models/point_group/__init__.py: -------------------------------------------------------------------------------- 1 | from .point_group_v1m1_base import PointGroup 2 | -------------------------------------------------------------------------------- /ponder/models/point_prompt_training/__init__.py: -------------------------------------------------------------------------------- 1 | from .point_prompt_training_v1m1_language_guided import * 2 | from .point_prompt_training_v1m2_decoupled import * 3 | -------------------------------------------------------------------------------- /ponder/models/point_prompt_training/point_prompt_training_v1m2_decoupled.py: -------------------------------------------------------------------------------- 1 | """ 2 | Point Prompt Training with decoupled segmentation head 3 | 4 | Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) 5 | Please cite our work if the code is helpful to you. 6 | """ 7 | 8 | from collections import OrderedDict 9 | from functools import partial 10 | 11 | import torch 12 | import torch.nn as nn 13 | 14 | from ponder.models.builder import MODELS 15 | from ponder.models.losses import build_criteria 16 | 17 | 18 | @MODELS.register_module("PPT-v1m2") 19 | class PointPromptTraining(nn.Module): 20 | """ 21 | PointPromptTraining v1m2 provides Data-driven Context and enables multi-dataset training with 22 | Decoupled Segmentation Head. PDNorm is supported by SpUNet-v1m3 to adapt the 23 | backbone to a specific dataset with a given dataset condition and context. 24 | """ 25 | 26 | def __init__( 27 | self, 28 | backbone=None, 29 | criteria=None, 30 | backbone_out_channels=96, 31 | context_channels=256, 32 | conditions=("Structured3D", "ScanNet", "S3DIS"), 33 | num_classes=(25, 20, 13), 34 | ): 35 | super().__init__() 36 | assert len(conditions) == len(num_classes) 37 | assert backbone.type in ["SpUNet-v1m3"] # SpUNet v1m3: Sparse UNet with PDNorm 38 | self.backbone = MODELS.build(backbone) 39 | self.criteria = build_criteria(criteria) 40 | self.conditions = conditions 41 | self.embedding_table = nn.Embedding(len(conditions), context_channels) 42 | self.seg_heads = nn.ModuleList( 43 | [nn.Linear(backbone_out_channels, num_cls) for num_cls in num_classes] 44 | ) 45 | 46 | def forward(self, data_dict): 47 | condition = data_dict["condition"][0] 48 | assert condition in self.conditions 49 | context = self.embedding_table( 50 | torch.tensor( 51 | [self.conditions.index(condition)], device=data_dict["coord"].device 52 | ) 53 | ) 54 | data_dict["context"] = context 55 | feat = self.backbone(data_dict) 56 | seg_head = self.seg_heads[self.conditions.index(condition)] 57 | seg_logits = seg_head(feat) 58 | # train 59 | if self.training: 60 | loss = self.criteria(seg_logits, data_dict["segment"]) 61 | return dict(loss=loss) 62 | # eval 63 | elif "segment" in data_dict.keys(): 64 | loss = self.criteria(seg_logits, data_dict["segment"]) 65 | return dict(loss=loss, seg_logits=seg_logits) 66 | # test 67 | else: 68 | return dict(seg_logits=seg_logits) 69 | -------------------------------------------------------------------------------- /ponder/models/ponder/__init__.py: -------------------------------------------------------------------------------- 1 | from .ponder_indoor_base import PonderIndoor 2 | from .ponder_outdoor_base import PonderOutdoor 3 | from .unet3d import SimpleConv3D, UNet3D, UNet3Dv1m2 4 | -------------------------------------------------------------------------------- /ponder/models/ponder/render_utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .builder import build_collider, build_field, build_renderer, build_sampler 2 | from .fields import * 3 | from .models import * 4 | from .ray_samplers import * 5 | from .rays import RayBundle 6 | from .scene_colliders import * 7 | -------------------------------------------------------------------------------- /ponder/models/ponder/render_utils/builder.py: -------------------------------------------------------------------------------- 1 | from ponder.utils.registry import Registry 2 | 3 | RENDERERS = Registry("renderers") 4 | FIELDS = Registry("fields") 5 | COLLIDERS = Registry("colliders") 6 | SAMPLERS = Registry("samplers") 7 | 8 | 9 | def build_renderer(cfg, **kwargs): 10 | """Build renderers.""" 11 | return RENDERERS.build(cfg, default_args=kwargs) 12 | 13 | 14 | def build_field(cfg, **kwargs): 15 | """Build fields.""" 16 | return FIELDS.build(cfg, default_args=kwargs) 17 | 18 | 19 | def build_collider(cfg, **kwargs): 20 | """Build colliders.""" 21 | return COLLIDERS.build(cfg, default_args=kwargs) 22 | 23 | 24 | def build_sampler(cfg, **kwargs): 25 | """Build samplers.""" 26 | return SAMPLERS.build(cfg, default_args=kwargs) 27 | -------------------------------------------------------------------------------- /ponder/models/ponder/render_utils/decoders.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | 5 | 6 | class SDFDecoder(nn.Module): 7 | def __init__( 8 | self, in_dim, out_dim, hidden_size=256, n_blocks=5, points_factor=1.0, **kwargs 9 | ): 10 | super().__init__() 11 | 12 | dims = [hidden_size] + [hidden_size for _ in range(n_blocks)] + [out_dim] 13 | self.num_layers = len(dims) 14 | 15 | for l in range(self.num_layers - 1): 16 | lin = nn.Linear(dims[l], dims[l + 1]) 17 | setattr(self, "lin" + str(l), lin) 18 | 19 | self.fc_c = nn.ModuleList( 20 | [nn.Linear(in_dim, hidden_size) for i in range(self.num_layers - 1)] 21 | ) 22 | self.fc_p = nn.Linear(3, hidden_size) 23 | 24 | self.activation = nn.Softplus(beta=100) 25 | 26 | self.points_factor = points_factor 27 | 28 | def forward(self, points, point_feats): 29 | x = self.fc_p(points) * self.points_factor 30 | for l in range(self.num_layers - 1): 31 | x = x + self.fc_c[l](point_feats) 32 | lin = getattr(self, "lin" + str(l)) 33 | x = lin(x) 34 | if l < self.num_layers - 2: 35 | x = self.activation(x) 36 | return x 37 | 38 | 39 | class RGBDecoder(nn.Module): 40 | def __init__( 41 | self, 42 | in_dim, 43 | out_dim=3, 44 | hidden_size=256, 45 | n_blocks=5, 46 | points_factor=1.0, 47 | **kwargs 48 | ): 49 | super().__init__() 50 | 51 | dims = [hidden_size] + [hidden_size for _ in range(n_blocks)] + [out_dim] 52 | self.num_layers = len(dims) 53 | 54 | for l in range(self.num_layers - 1): 55 | lin = nn.Linear(dims[l], dims[l + 1]) 56 | setattr(self, "lin" + str(l), lin) 57 | 58 | self.fc_p = nn.Linear(3, hidden_size) 59 | 60 | self.fc_c = nn.ModuleList( 61 | [nn.Linear(in_dim, hidden_size) for i in range(self.num_layers - 1)] 62 | ) 63 | self.activation = nn.ReLU() 64 | 65 | self.points_factor = points_factor 66 | 67 | def forward(self, points, point_feats): 68 | x = self.fc_p(points) * self.points_factor 69 | for l in range(self.num_layers - 1): 70 | x = x + self.fc_c[l](point_feats) 71 | lin = getattr(self, "lin" + str(l)) 72 | x = lin(x) 73 | if l < self.num_layers - 2: 74 | x = self.activation(x) 75 | x = torch.sigmoid(x) 76 | return x 77 | 78 | 79 | class SemanticDecoder(nn.Module): 80 | def __init__( 81 | self, in_dim, out_dim, hidden_size=256, n_blocks=5, points_factor=1.0, **kwargs 82 | ): 83 | super().__init__() 84 | 85 | dims = [hidden_size] + [hidden_size for _ in range(n_blocks)] + [out_dim] 86 | self.num_layers = len(dims) 87 | 88 | for l in range(self.num_layers - 1): 89 | lin = nn.Linear(dims[l], dims[l + 1]) 90 | setattr(self, "lin" + str(l), lin) 91 | 92 | self.fc_p = nn.Linear(3, hidden_size) 93 | 94 | self.fc_c = nn.ModuleList( 95 | [nn.Linear(in_dim, hidden_size) for i in range(self.num_layers - 1)] 96 | ) 97 | self.activation = nn.ReLU() 98 | 99 | self.points_factor = points_factor 100 | 101 | def forward(self, points, point_feats): 102 | x = self.fc_p(points) * self.points_factor 103 | for l in range(self.num_layers - 1): 104 | x = x + self.fc_c[l](point_feats) 105 | lin = getattr(self, "lin" + str(l)) 106 | x = lin(x) 107 | if l < self.num_layers - 2: 108 | x = self.activation(x) 109 | return x 110 | -------------------------------------------------------------------------------- /ponder/models/ponder/render_utils/fields/__init__.py: -------------------------------------------------------------------------------- 1 | from .sdf_field import SDFField 2 | 3 | __all__ = ["SDFField"] 4 | -------------------------------------------------------------------------------- /ponder/models/ponder/render_utils/models/__init__.py: -------------------------------------------------------------------------------- 1 | from .neus import NeuSModel 2 | from .volsdf import VolSDFModel 3 | 4 | __all__ = ["NeuSModel", "VolSDFModel"] 5 | -------------------------------------------------------------------------------- /ponder/models/ponder/render_utils/models/neus.py: -------------------------------------------------------------------------------- 1 | from functools import partial 2 | 3 | from ..builder import RENDERERS 4 | from .base_surface_model import SurfaceModel 5 | 6 | 7 | @RENDERERS.register_module() 8 | class NeuSModel(SurfaceModel): 9 | def __init__(self, field, collider, sampler, loss, **kwargs): 10 | super().__init__(field=field, collider=collider, sampler=sampler, loss=loss) 11 | self.anneal_end = 50000 12 | 13 | def get_training_callbacks(self): 14 | raise NotImplementedError 15 | 16 | def sample_and_forward_field(self, ray_bundle, volume_feature): 17 | sampler_out_dict = self.sampler( 18 | ray_bundle, 19 | occupancy_fn=self.field.get_occupancy, 20 | sdf_fn=partial(self.field.get_sdf, volume_feature=volume_feature), 21 | ) 22 | ray_samples = sampler_out_dict.pop("ray_samples") 23 | field_outputs = self.field(ray_samples, volume_feature, return_alphas=True) 24 | weights, _ = ray_samples.get_weights_and_transmittance_from_alphas( 25 | field_outputs["alphas"] 26 | ) 27 | 28 | samples_and_field_outputs = { 29 | "ray_samples": ray_samples, 30 | "field_outputs": field_outputs, 31 | "weights": weights, # (num_rays, num_smaples+num_importance, 1) 32 | "sampled_points": ray_samples.frustums.get_start_positions(), # (num_rays, num_smaples+num_importance, 3) 33 | **sampler_out_dict, 34 | } 35 | 36 | return samples_and_field_outputs 37 | -------------------------------------------------------------------------------- /ponder/models/ponder/render_utils/models/volsdf.py: -------------------------------------------------------------------------------- 1 | from functools import partial 2 | 3 | from ..builder import RENDERERS 4 | from .base_surface_model import SurfaceModel 5 | 6 | 7 | @RENDERERS.register_module() 8 | class VolSDFModel(SurfaceModel): 9 | def __init__(self, field, collider, sampler, loss, **kwargs): 10 | super().__init__(field=field, collider=collider, sampler=sampler, loss=loss) 11 | 12 | def sample_and_forward_field(self, ray_bundle, volume_feature): 13 | sampler_out_dict = self.sampler( 14 | ray_bundle, 15 | density_fn=self.field.laplace_density, 16 | sdf_fn=partial(self.field.get_sdf, volume_feature=volume_feature), 17 | ) 18 | ray_samples = sampler_out_dict.pop("ray_samples") 19 | field_outputs = self.field(ray_samples, volume_feature) 20 | weights, _ = ray_samples.get_weights_and_transmittance(field_outputs["density"]) 21 | 22 | samples_and_field_outputs = { 23 | "ray_samples": ray_samples, 24 | "field_outputs": field_outputs, 25 | "weights": weights, 26 | "sampled_points": ray_samples.frustums.get_start_positions(), # (num_rays, num_smaples+num_importance, 3) 27 | **sampler_out_dict, 28 | } 29 | return samples_and_field_outputs 30 | -------------------------------------------------------------------------------- /ponder/models/ponder/render_utils/renderers.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | 4 | 5 | class RGBRenderer(nn.Module): 6 | """Standard volumetic rendering.""" 7 | 8 | def __init__(self, background_color=(0.0, 0.0, 0.0)): 9 | super().__init__() 10 | self.background_color = background_color 11 | 12 | def forward(self, rgb, weights): 13 | """Composite samples along ray and render color image 14 | 15 | Args: 16 | rgb: RGB for each sample, (num_rays, num_samples, 3) 17 | weights: Weights for each sample, (num_rays, num_samples, 1) 18 | Returns: 19 | Outputs of rgb values. 20 | """ 21 | comp_rgb = torch.sum(weights * rgb, dim=-2) # (num_rays, 3) 22 | accumulated_weight = torch.sum(weights, dim=-2) 23 | comp_rgb = comp_rgb + comp_rgb.new_tensor(self.background_color) * ( 24 | 1.0 - accumulated_weight 25 | ) 26 | if not self.training: 27 | torch.clamp_(comp_rgb, min=0.0, max=1.0) 28 | return comp_rgb 29 | 30 | 31 | class DepthRenderer(nn.Module): 32 | """Calculate depth along ray.""" 33 | 34 | def __init__(self, **kwargs): 35 | super().__init__() 36 | 37 | def forward(self, ray_samples, weights): 38 | """Composite samples along ray and calculate depths. 39 | 40 | Args: 41 | weights: Weights for each sample. 42 | ray_samples: Set of ray samples. 43 | Returns: 44 | Outputs of depth values. 45 | """ 46 | eps = 1e-10 47 | # steps = (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2 48 | steps = ray_samples.frustums.starts 49 | depth = torch.sum(weights * steps, dim=-2) / (torch.sum(weights, -2) + eps) 50 | depth = torch.clip(depth, steps.min(), steps.max()) 51 | return depth 52 | 53 | 54 | class NormalRenderer(nn.Module): 55 | """Calculate normals along the ray.""" 56 | 57 | def __init__(self, **kwargs): 58 | super().__init__() 59 | 60 | def forward(self, normals, weights): 61 | """Calculate normals along the ray.""" 62 | n = torch.sum(weights * normals, dim=-2) 63 | return n 64 | 65 | 66 | class SemanticRenderer(nn.Module): 67 | """Calculate semantic features along the ray.""" 68 | 69 | def __init__(self, **kwargs): 70 | super().__init__() 71 | 72 | def forward(self, semantic, weights): 73 | """Calculate semantic features along the ray.""" 74 | f = torch.sum(weights * semantic, dim=-2) 75 | return f 76 | -------------------------------------------------------------------------------- /ponder/models/ponder/render_utils/scene_colliders.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from .builder import COLLIDERS 5 | 6 | 7 | class SceneCollider(nn.Module): 8 | """Module for setting near and far values for rays.""" 9 | 10 | def __init__(self, **kwargs): 11 | self.kwargs = kwargs 12 | super().__init__() 13 | 14 | def set_nears_and_fars(self, ray_bundle): 15 | """To be implemented.""" 16 | raise NotImplementedError 17 | 18 | def forward(self, ray_bundle): 19 | """Sets the nears and fars if they are not set already.""" 20 | if ray_bundle.nears is not None and ray_bundle.fars is not None: 21 | return ray_bundle 22 | return self.set_nears_and_fars(ray_bundle) 23 | 24 | 25 | @COLLIDERS.register_module() 26 | class AABBBoxCollider(SceneCollider): 27 | """Module for colliding rays with the scene box to compute near and far values. 28 | 29 | Args: 30 | scene_box: scene box to apply to dataset 31 | """ 32 | 33 | def __init__(self, bbox, near_plane, **kwargs): 34 | super().__init__(**kwargs) 35 | self.bbox = bbox 36 | self.near_plane = near_plane 37 | 38 | def _intersect_with_aabb(self, rays_o, rays_d, aabb): 39 | """Returns collection of valid rays within a specified near/far bounding box along with a mask 40 | specifying which rays are valid 41 | 42 | Args: 43 | rays_o: (num_rays, 3) ray origins, scaled 44 | rays_d: (num_rays, 3) ray directions 45 | aabb: (6, ) This is [min point (x,y,z), max point (x,y,z)], scaled 46 | """ 47 | # avoid divide by zero 48 | dir_fraction = 1.0 / (rays_d + 1e-6) 49 | 50 | # x 51 | t1 = (aabb[0] - rays_o[:, 0:1]) * dir_fraction[:, 0:1] 52 | t2 = (aabb[3] - rays_o[:, 0:1]) * dir_fraction[:, 0:1] 53 | # y 54 | t3 = (aabb[1] - rays_o[:, 1:2]) * dir_fraction[:, 1:2] 55 | t4 = (aabb[4] - rays_o[:, 1:2]) * dir_fraction[:, 1:2] 56 | # z 57 | t5 = (aabb[2] - rays_o[:, 2:3]) * dir_fraction[:, 2:3] 58 | t6 = (aabb[5] - rays_o[:, 2:3]) * dir_fraction[:, 2:3] 59 | 60 | nears = torch.max( 61 | torch.cat( 62 | [torch.minimum(t1, t2), torch.minimum(t3, t4), torch.minimum(t5, t6)], 63 | dim=1, 64 | ), 65 | dim=1, 66 | ).values 67 | fars = torch.min( 68 | torch.cat( 69 | [torch.maximum(t1, t2), torch.maximum(t3, t4), torch.maximum(t5, t6)], 70 | dim=1, 71 | ), 72 | dim=1, 73 | ).values 74 | 75 | # clamp to near plane 76 | nears = torch.clamp(nears, min=self.near_plane) 77 | # fars = torch.maximum(fars, nears + 1e-6) 78 | # if self.training: 79 | # assert (nears < fars).all() 80 | # else: 81 | mask_at_box = nears < fars 82 | nears[~mask_at_box] = 0.0 83 | fars[~mask_at_box] = 0.0 84 | 85 | return nears, fars 86 | 87 | def set_nears_and_fars(self, ray_bundle): 88 | """Intersects the rays with the scene box and updates the near and far values. 89 | Populates nears and fars fields and returns the ray_bundle. 90 | Returns: 91 | nears: (num_rays, 1) 92 | fars: (num_rays, 1) 93 | """ 94 | nears, fars = self._intersect_with_aabb( 95 | ray_bundle.origins, ray_bundle.directions, self.bbox 96 | ) 97 | ray_bundle.nears = nears[..., None] 98 | ray_bundle.fars = fars[..., None] 99 | return ray_bundle 100 | 101 | 102 | @COLLIDERS.register_module() 103 | class NearFarCollider(SceneCollider): 104 | """Sets the nears and fars with fixed values. 105 | 106 | Args: 107 | near_plane: distance to near plane 108 | far_plane: distance to far plane 109 | """ 110 | 111 | def __init__(self, near_plane, far_plane, **kwargs): 112 | super().__init__(**kwargs) 113 | self.near_plane = near_plane 114 | self.far_plane = far_plane 115 | 116 | def set_nears_and_fars(self, ray_bundle): 117 | ones = torch.ones_like(ray_bundle.origins[..., 0:1]) 118 | ray_bundle.nears = ones * self.near_plane 119 | ray_bundle.fars = ones * self.far_plane 120 | return ray_bundle 121 | -------------------------------------------------------------------------------- /ponder/models/sparse_unet/__init__.py: -------------------------------------------------------------------------------- 1 | # from .mink_unet import * 2 | from .spconv_unet_v1m1_base import * 3 | from .spconv_unet_v1m2_bn_momentum import * 4 | from .spconv_unet_v1m3_pdnorm import * 5 | -------------------------------------------------------------------------------- /ponder/models/utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | General Utils for Models 3 | 4 | Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) 5 | Please cite our work if the code is helpful to you. 6 | """ 7 | 8 | import torch 9 | 10 | 11 | def offset2batch(offset): 12 | return ( 13 | torch.cat( 14 | [ 15 | ( 16 | torch.tensor([i] * (o - offset[i - 1])) 17 | if i > 0 18 | else torch.tensor([i] * o) 19 | ) 20 | for i, o in enumerate(offset) 21 | ], 22 | dim=0, 23 | ) 24 | .long() 25 | .to(offset.device) 26 | ) 27 | 28 | 29 | def batch2offset(batch): 30 | return torch.cumsum(batch.bincount(), dim=0).long() 31 | 32 | 33 | def off_diagonal(x): 34 | # return a flattened view of the off-diagonal elements of a square matrix 35 | n, m = x.shape 36 | assert n == m 37 | return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten() 38 | 39 | 40 | def checkpoint(func, inputs, params, flag): 41 | """ 42 | Evaluate a function without caching intermediate activations, allowing for 43 | reduced memory at the expense of extra compute in the backward pass. 44 | :param func: the function to evaluate. 45 | :param inputs: the argument sequence to pass to `func`. 46 | :param params: a sequence of parameters `func` depends on but does not 47 | explicitly take as arguments. 48 | :param flag: if False, disable gradient checkpointing. 49 | """ 50 | if flag: 51 | args = tuple(inputs) + tuple(params) 52 | return CheckpointFunction.apply(func, len(inputs), *args) 53 | else: 54 | return func(*inputs) 55 | 56 | 57 | class CheckpointFunction(torch.autograd.Function): 58 | @staticmethod 59 | def forward(ctx, run_function, length, *args): 60 | ctx.run_function = run_function 61 | ctx.input_tensors = list(args[:length]) 62 | ctx.input_params = list(args[length:]) 63 | 64 | with torch.no_grad(): 65 | output_tensors = ctx.run_function(*ctx.input_tensors) 66 | return output_tensors 67 | 68 | @staticmethod 69 | def backward(ctx, *output_grads): 70 | ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors] 71 | with torch.enable_grad(): 72 | # Fixes a bug where the first op in run_function modifies the 73 | # Tensor storage in place, which is not allowed for detach()'d 74 | # Tensors. 75 | shallow_copies = [x.view_as(x) for x in ctx.input_tensors] 76 | output_tensors = ctx.run_function(*shallow_copies) 77 | input_grads = torch.autograd.grad( 78 | output_tensors, 79 | ctx.input_tensors + ctx.input_params, 80 | output_grads, 81 | allow_unused=True, 82 | ) 83 | del ctx.input_tensors 84 | del ctx.input_params 85 | del output_tensors 86 | return (None, None) + input_grads 87 | -------------------------------------------------------------------------------- /ponder/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/OpenGVLab/PonderV2/c152af1df9d04a98bc19696e8d69d047834a80db/ponder/utils/__init__.py -------------------------------------------------------------------------------- /ponder/utils/cache.py: -------------------------------------------------------------------------------- 1 | """ 2 | Data Cache Utils 3 | 4 | Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) 5 | Please cite our work if the code is helpful to you. 6 | """ 7 | 8 | import os 9 | 10 | import SharedArray 11 | 12 | try: 13 | from multiprocessing.shared_memory import ShareableList 14 | except ImportError: 15 | import warnings 16 | 17 | warnings.warn("Please update python version >= 3.8 to enable shared_memory") 18 | import numpy as np 19 | 20 | 21 | def shared_array(name, var=None): 22 | if var is not None: 23 | # check exist 24 | if os.path.exists(f"/dev/shm/{name}"): 25 | return SharedArray.attach(f"shm://{name}") 26 | # create shared_array 27 | data = SharedArray.create(f"shm://{name}", var.shape, dtype=var.dtype) 28 | data[...] = var[...] 29 | data.flags.writeable = False 30 | else: 31 | data = SharedArray.attach(f"shm://{name}").copy() 32 | return data 33 | 34 | 35 | def shared_dict(name, var=None): 36 | name = str(name) 37 | assert "." not in name # '.' is used as sep flag 38 | data = {} 39 | if var is not None: 40 | assert isinstance(var, dict) 41 | keys = var.keys() 42 | # current version only cache np.array 43 | keys_valid = [] 44 | for key in keys: 45 | if isinstance(var[key], np.ndarray): 46 | keys_valid.append(key) 47 | keys = keys_valid 48 | 49 | ShareableList(sequence=keys, name=name + ".keys") 50 | for key in keys: 51 | if isinstance(var[key], np.ndarray): 52 | data[key] = shared_array(name=f"{name}.{key}", var=var[key]) 53 | else: 54 | keys = list(ShareableList(name=name + ".keys")) 55 | for key in keys: 56 | data[key] = shared_array(name=f"{name}.{key}") 57 | return data 58 | -------------------------------------------------------------------------------- /ponder/utils/env.py: -------------------------------------------------------------------------------- 1 | """ 2 | Environment Utils 3 | 4 | Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) 5 | Please cite our work if the code is helpful to you. 6 | """ 7 | 8 | import os 9 | import random 10 | from datetime import datetime 11 | 12 | import numpy as np 13 | import torch 14 | import torch.backends.cudnn as cudnn 15 | 16 | 17 | def get_random_seed(): 18 | seed = ( 19 | os.getpid() 20 | + int(datetime.now().strftime("%S%f")) 21 | + int.from_bytes(os.urandom(2), "big") 22 | ) 23 | return seed 24 | 25 | 26 | def set_seed(seed=None): 27 | if seed is None: 28 | seed = get_random_seed() 29 | random.seed(seed) 30 | np.random.seed(seed) 31 | torch.manual_seed(seed) 32 | torch.cuda.manual_seed(seed) 33 | torch.cuda.manual_seed_all(seed) 34 | cudnn.benchmark = False 35 | cudnn.deterministic = True 36 | os.environ["PYTHONHASHSEED"] = str(seed) 37 | -------------------------------------------------------------------------------- /ponder/utils/misc.py: -------------------------------------------------------------------------------- 1 | """ 2 | Misc 3 | 4 | Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) 5 | Please cite our work if the code is helpful to you. 6 | """ 7 | 8 | import os 9 | import warnings 10 | from collections import abc 11 | from importlib import import_module 12 | 13 | import numpy as np 14 | import torch 15 | 16 | 17 | class AverageMeter(object): 18 | """Computes and stores the average and current value""" 19 | 20 | def __init__(self): 21 | self.val = 0 22 | self.avg = 0 23 | self.sum = 0 24 | self.count = 0 25 | 26 | def reset(self): 27 | self.val = 0 28 | self.avg = 0 29 | self.sum = 0 30 | self.count = 0 31 | 32 | def update(self, val, n=1): 33 | self.val = val 34 | self.sum += val * n 35 | self.count += n 36 | self.avg = self.sum / self.count 37 | 38 | 39 | def intersection_and_union(output, target, K, ignore_index=-1): 40 | # 'K' classes, output and target sizes are N or N * L or N * H * W, each value in range 0 to K - 1. 41 | assert output.ndim in [1, 2, 3] 42 | assert output.shape == target.shape 43 | output = output.reshape(output.size).copy() 44 | target = target.reshape(target.size) 45 | output[np.where(target == ignore_index)[0]] = ignore_index 46 | intersection = output[np.where(output == target)[0]] 47 | area_intersection, _ = np.histogram(intersection, bins=np.arange(K + 1)) 48 | area_output, _ = np.histogram(output, bins=np.arange(K + 1)) 49 | area_target, _ = np.histogram(target, bins=np.arange(K + 1)) 50 | area_union = area_output + area_target - area_intersection 51 | return area_intersection, area_union, area_target 52 | 53 | 54 | def intersection_and_union_gpu(output, target, k, ignore_index=-1): 55 | # 'K' classes, output and target sizes are N or N * L or N * H * W, each value in range 0 to K - 1. 56 | assert output.dim() in [1, 2, 3] 57 | assert output.shape == target.shape 58 | output = output.view(-1) 59 | target = target.view(-1) 60 | output[target == ignore_index] = ignore_index 61 | intersection = output[output == target] 62 | area_intersection = torch.histc(intersection, bins=k, min=0, max=k - 1) 63 | area_output = torch.histc(output, bins=k, min=0, max=k - 1) 64 | area_target = torch.histc(target, bins=k, min=0, max=k - 1) 65 | area_union = area_output + area_target - area_intersection 66 | return area_intersection, area_union, area_target 67 | 68 | 69 | def make_dirs(dir_name): 70 | if not os.path.exists(dir_name): 71 | os.makedirs(dir_name, exist_ok=True) 72 | 73 | 74 | def find_free_port(): 75 | import socket 76 | 77 | sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 78 | # Binding to port 0 will cause the OS to find an available port for us 79 | sock.bind(("", 0)) 80 | port = sock.getsockname()[1] 81 | sock.close() 82 | # NOTE: there is still a chance the port could be taken by other processes. 83 | return port 84 | 85 | 86 | def is_seq_of(seq, expected_type, seq_type=None): 87 | """Check whether it is a sequence of some type. 88 | 89 | Args: 90 | seq (Sequence): The sequence to be checked. 91 | expected_type (type): Expected type of sequence items. 92 | seq_type (type, optional): Expected sequence type. 93 | 94 | Returns: 95 | bool: Whether the sequence is valid. 96 | """ 97 | if seq_type is None: 98 | exp_seq_type = abc.Sequence 99 | else: 100 | assert isinstance(seq_type, type) 101 | exp_seq_type = seq_type 102 | if not isinstance(seq, exp_seq_type): 103 | return False 104 | for item in seq: 105 | if not isinstance(item, expected_type): 106 | return False 107 | return True 108 | 109 | 110 | def is_str(x): 111 | """Whether the input is an string instance. 112 | 113 | Note: This method is deprecated since python 2 is no longer supported. 114 | """ 115 | return isinstance(x, str) 116 | 117 | 118 | def import_modules_from_strings(imports, allow_failed_imports=False): 119 | """Import modules from the given list of strings. 120 | 121 | Args: 122 | imports (list | str | None): The given module names to be imported. 123 | allow_failed_imports (bool): If True, the failed imports will return 124 | None. Otherwise, an ImportError is raise. Default: False. 125 | 126 | Returns: 127 | list[module] | module | None: The imported modules. 128 | 129 | Examples: 130 | >>> osp, sys = import_modules_from_strings( 131 | ... ['os.path', 'sys']) 132 | >>> import os.path as osp_ 133 | >>> import sys as sys_ 134 | >>> assert osp == osp_ 135 | >>> assert sys == sys_ 136 | """ 137 | if not imports: 138 | return 139 | single_import = False 140 | if isinstance(imports, str): 141 | single_import = True 142 | imports = [imports] 143 | if not isinstance(imports, list): 144 | raise TypeError(f"custom_imports must be a list but got type {type(imports)}") 145 | imported = [] 146 | for imp in imports: 147 | if not isinstance(imp, str): 148 | raise TypeError(f"{imp} is of type {type(imp)} and cannot be imported.") 149 | try: 150 | imported_tmp = import_module(imp) 151 | except ImportError: 152 | if allow_failed_imports: 153 | warnings.warn(f"{imp} failed to import and is ignored.", UserWarning) 154 | imported_tmp = None 155 | else: 156 | raise ImportError 157 | imported.append(imported_tmp) 158 | if single_import: 159 | imported = imported[0] 160 | return imported 161 | -------------------------------------------------------------------------------- /ponder/utils/optimizer.py: -------------------------------------------------------------------------------- 1 | """ 2 | Optimizer 3 | 4 | Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) 5 | Please cite our work if the code is helpful to you. 6 | """ 7 | 8 | import torch 9 | 10 | from ponder.utils.logger import get_root_logger 11 | from ponder.utils.registry import Registry 12 | 13 | OPTIMIZERS = Registry("optimizers") 14 | 15 | 16 | OPTIMIZERS.register_module(module=torch.optim.SGD, name="SGD") 17 | OPTIMIZERS.register_module(module=torch.optim.Adam, name="Adam") 18 | OPTIMIZERS.register_module(module=torch.optim.AdamW, name="AdamW") 19 | 20 | 21 | def build_optimizer(cfg, model, param_dicts=None): 22 | if param_dicts is None: 23 | cfg.params = model.parameters() 24 | else: 25 | cfg.params = [dict(names=[], params=[], lr=cfg.lr)] 26 | for i in range(len(param_dicts)): 27 | param_group = dict(names=[], params=[]) 28 | if "lr" in param_dicts[i].keys(): 29 | param_group["lr"] = param_dicts[i].lr 30 | if "momentum" in param_dicts[i].keys(): 31 | param_group["momentum"] = param_dicts[i].momentum 32 | if "weight_decay" in param_dicts[i].keys(): 33 | param_group["weight_decay"] = param_dicts[i].weight_decay 34 | cfg.params.append(param_group) 35 | 36 | for n, p in model.named_parameters(): 37 | flag = False 38 | for i in range(len(param_dicts)): 39 | if param_dicts[i].keyword in n: 40 | cfg.params[i + 1]["names"].append(n) 41 | cfg.params[i + 1]["params"].append(p) 42 | flag = True 43 | break 44 | if not flag: 45 | cfg.params[0]["names"].append(n) 46 | cfg.params[0]["params"].append(p) 47 | 48 | logger = get_root_logger() 49 | for i in range(len(cfg.params)): 50 | param_names = cfg.params[i].pop("names") 51 | message = "" 52 | for key in cfg.params[i].keys(): 53 | if key != "params": 54 | message += f" {key}: {cfg.params[i][key]};" 55 | logger.info(f"Params Group {i+1} -{message} Params: {param_names}.") 56 | return OPTIMIZERS.build(cfg=cfg) 57 | -------------------------------------------------------------------------------- /ponder/utils/path.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | import os 3 | import os.path as osp 4 | from pathlib import Path 5 | 6 | from .misc import is_str 7 | 8 | 9 | def is_filepath(x): 10 | return is_str(x) or isinstance(x, Path) 11 | 12 | 13 | def fopen(filepath, *args, **kwargs): 14 | if is_str(filepath): 15 | return open(filepath, *args, **kwargs) 16 | elif isinstance(filepath, Path): 17 | return filepath.open(*args, **kwargs) 18 | raise ValueError("`filepath` should be a string or a Path") 19 | 20 | 21 | def check_file_exist(filename, msg_tmpl='file "{}" does not exist'): 22 | if not osp.isfile(filename): 23 | raise FileNotFoundError(msg_tmpl.format(filename)) 24 | 25 | 26 | def mkdir_or_exist(dir_name, mode=0o777): 27 | if dir_name == "": 28 | return 29 | dir_name = osp.expanduser(dir_name) 30 | os.makedirs(dir_name, mode=mode, exist_ok=True) 31 | 32 | 33 | def symlink(src, dst, overwrite=True, **kwargs): 34 | if os.path.lexists(dst) and overwrite: 35 | os.remove(dst) 36 | os.symlink(src, dst, **kwargs) 37 | 38 | 39 | def scandir(dir_path, suffix=None, recursive=False, case_sensitive=True): 40 | """Scan a directory to find the interested files. 41 | 42 | Args: 43 | dir_path (str | obj:`Path`): Path of the directory. 44 | suffix (str | tuple(str), optional): File suffix that we are 45 | interested in. Default: None. 46 | recursive (bool, optional): If set to True, recursively scan the 47 | directory. Default: False. 48 | case_sensitive (bool, optional) : If set to False, ignore the case of 49 | suffix. Default: True. 50 | 51 | Returns: 52 | A generator for all the interested files with relative paths. 53 | """ 54 | if isinstance(dir_path, (str, Path)): 55 | dir_path = str(dir_path) 56 | else: 57 | raise TypeError('"dir_path" must be a string or Path object') 58 | 59 | if (suffix is not None) and not isinstance(suffix, (str, tuple)): 60 | raise TypeError('"suffix" must be a string or tuple of strings') 61 | 62 | if suffix is not None and not case_sensitive: 63 | suffix = ( 64 | suffix.lower() 65 | if isinstance(suffix, str) 66 | else tuple(item.lower() for item in suffix) 67 | ) 68 | 69 | root = dir_path 70 | 71 | def _scandir(dir_path, suffix, recursive, case_sensitive): 72 | for entry in os.scandir(dir_path): 73 | if not entry.name.startswith(".") and entry.is_file(): 74 | rel_path = osp.relpath(entry.path, root) 75 | _rel_path = rel_path if case_sensitive else rel_path.lower() 76 | if suffix is None or _rel_path.endswith(suffix): 77 | yield rel_path 78 | elif recursive and os.path.isdir(entry.path): 79 | # scan recursively if entry.path is a directory 80 | yield from _scandir(entry.path, suffix, recursive, case_sensitive) 81 | 82 | return _scandir(dir_path, suffix, recursive, case_sensitive) 83 | 84 | 85 | def find_vcs_root(path, markers=(".git",)): 86 | """Finds the root directory (including itself) of specified markers. 87 | 88 | Args: 89 | path (str): Path of directory or file. 90 | markers (list[str], optional): List of file or directory names. 91 | 92 | Returns: 93 | The directory contained one of the markers or None if not found. 94 | """ 95 | if osp.isfile(path): 96 | path = osp.dirname(path) 97 | 98 | prev, cur = None, osp.abspath(osp.expanduser(path)) 99 | while cur != prev: 100 | if any(osp.exists(osp.join(cur, marker)) for marker in markers): 101 | return cur 102 | prev, cur = cur, osp.split(cur)[0] 103 | return None 104 | -------------------------------------------------------------------------------- /ponder/utils/scheduler.py: -------------------------------------------------------------------------------- 1 | """ 2 | Scheduler 3 | 4 | Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) 5 | Please cite our work if the code is helpful to you. 6 | """ 7 | 8 | import torch.optim.lr_scheduler as lr_scheduler 9 | 10 | from .registry import Registry 11 | 12 | SCHEDULERS = Registry("schedulers") 13 | 14 | 15 | @SCHEDULERS.register_module() 16 | class MultiStepLR(lr_scheduler.MultiStepLR): 17 | def __init__( 18 | self, 19 | optimizer, 20 | milestones, 21 | total_steps, 22 | gamma=0.1, 23 | last_epoch=-1, 24 | verbose=False, 25 | ): 26 | super().__init__( 27 | optimizer=optimizer, 28 | milestones=[rate * total_steps for rate in milestones], 29 | gamma=gamma, 30 | last_epoch=last_epoch, 31 | verbose=verbose, 32 | ) 33 | 34 | 35 | @SCHEDULERS.register_module() 36 | class MultiStepWithWarmupLR(lr_scheduler.LambdaLR): 37 | def __init__( 38 | self, 39 | optimizer, 40 | milestones, 41 | total_steps, 42 | gamma=0.1, 43 | warmup_rate=0.05, 44 | warmup_scale=1e-6, 45 | last_epoch=-1, 46 | verbose=False, 47 | ): 48 | milestones = [rate * total_steps for rate in milestones] 49 | 50 | def multi_step_with_warmup(s): 51 | factor = 1.0 52 | for i in range(len(milestones)): 53 | if s < milestones[i]: 54 | break 55 | factor *= gamma 56 | 57 | if s <= warmup_rate * total_steps: 58 | warmup_coefficient = 1 - (1 - s / warmup_rate / total_steps) * ( 59 | 1 - warmup_scale 60 | ) 61 | else: 62 | warmup_coefficient = 1.0 63 | return warmup_coefficient * factor 64 | 65 | super().__init__( 66 | optimizer=optimizer, 67 | lr_lambda=multi_step_with_warmup, 68 | last_epoch=last_epoch, 69 | verbose=verbose, 70 | ) 71 | 72 | 73 | @SCHEDULERS.register_module() 74 | class PolyLR(lr_scheduler.LambdaLR): 75 | def __init__(self, optimizer, total_steps, power=0.9, last_epoch=-1, verbose=False): 76 | super().__init__( 77 | optimizer=optimizer, 78 | lr_lambda=lambda s: (1 - s / (total_steps + 1)) ** power, 79 | last_epoch=last_epoch, 80 | verbose=verbose, 81 | ) 82 | 83 | 84 | @SCHEDULERS.register_module() 85 | class ExpLR(lr_scheduler.LambdaLR): 86 | def __init__(self, optimizer, total_steps, gamma=0.9, last_epoch=-1, verbose=False): 87 | super().__init__( 88 | optimizer=optimizer, 89 | lr_lambda=lambda s: gamma ** (s / total_steps), 90 | last_epoch=last_epoch, 91 | verbose=verbose, 92 | ) 93 | 94 | 95 | @SCHEDULERS.register_module() 96 | class CosineAnnealingLR(lr_scheduler.CosineAnnealingLR): 97 | def __init__(self, optimizer, total_steps, eta_min=0, last_epoch=-1, verbose=False): 98 | super().__init__( 99 | optimizer=optimizer, 100 | T_max=total_steps, 101 | eta_min=eta_min, 102 | last_epoch=last_epoch, 103 | verbose=verbose, 104 | ) 105 | 106 | 107 | @SCHEDULERS.register_module() 108 | class OneCycleLR(lr_scheduler.OneCycleLR): 109 | r""" 110 | torch.optim.lr_scheduler.OneCycleLR, Block total_steps 111 | """ 112 | 113 | def __init__( 114 | self, 115 | optimizer, 116 | max_lr, 117 | total_steps=None, 118 | pct_start=0.3, 119 | anneal_strategy="cos", 120 | cycle_momentum=True, 121 | base_momentum=0.85, 122 | max_momentum=0.95, 123 | div_factor=25.0, 124 | final_div_factor=1e4, 125 | three_phase=False, 126 | last_epoch=-1, 127 | verbose=False, 128 | ): 129 | super().__init__( 130 | optimizer=optimizer, 131 | max_lr=max_lr, 132 | total_steps=total_steps, 133 | pct_start=pct_start, 134 | anneal_strategy=anneal_strategy, 135 | cycle_momentum=cycle_momentum, 136 | base_momentum=base_momentum, 137 | max_momentum=max_momentum, 138 | div_factor=div_factor, 139 | final_div_factor=final_div_factor, 140 | three_phase=three_phase, 141 | last_epoch=last_epoch, 142 | verbose=verbose, 143 | ) 144 | 145 | 146 | def build_scheduler(cfg, optimizer): 147 | cfg.optimizer = optimizer 148 | return SCHEDULERS.build(cfg=cfg) 149 | -------------------------------------------------------------------------------- /ponder/utils/timer.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | # -*- coding: utf-8 -*- 3 | 4 | from time import perf_counter 5 | from typing import Optional 6 | 7 | 8 | class Timer: 9 | """ 10 | A timer which computes the time elapsed since the start/reset of the timer. 11 | """ 12 | 13 | def __init__(self) -> None: 14 | self.reset() 15 | 16 | def reset(self) -> None: 17 | """ 18 | Reset the timer. 19 | """ 20 | self._start = perf_counter() 21 | self._paused: Optional[float] = None 22 | self._total_paused = 0 23 | self._count_start = 1 24 | 25 | def pause(self) -> None: 26 | """ 27 | Pause the timer. 28 | """ 29 | if self._paused is not None: 30 | raise ValueError("Trying to pause a Timer that is already paused!") 31 | self._paused = perf_counter() 32 | 33 | def is_paused(self) -> bool: 34 | """ 35 | Returns: 36 | bool: whether the timer is currently paused 37 | """ 38 | return self._paused is not None 39 | 40 | def resume(self) -> None: 41 | """ 42 | Resume the timer. 43 | """ 44 | if self._paused is None: 45 | raise ValueError("Trying to resume a Timer that is not paused!") 46 | # pyre-fixme[58]: `-` is not supported for operand types `float` and 47 | # `Optional[float]`. 48 | self._total_paused += perf_counter() - self._paused 49 | self._paused = None 50 | self._count_start += 1 51 | 52 | def seconds(self) -> float: 53 | """ 54 | Returns: 55 | (float): the total number of seconds since the start/reset of the 56 | timer, excluding the time when the timer is paused. 57 | """ 58 | if self._paused is not None: 59 | end_time: float = self._paused # type: ignore 60 | else: 61 | end_time = perf_counter() 62 | return end_time - self._start - self._total_paused 63 | 64 | def avg_seconds(self) -> float: 65 | """ 66 | Returns: 67 | (float): the average number of seconds between every start/reset and 68 | pause. 69 | """ 70 | return self.seconds() / self._count_start 71 | -------------------------------------------------------------------------------- /ponder/utils/visualization.py: -------------------------------------------------------------------------------- 1 | """ 2 | Visualization Utils 3 | 4 | Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) 5 | Please cite our work if the code is helpful to you. 6 | """ 7 | 8 | import os 9 | 10 | import numpy as np 11 | import open3d as o3d 12 | import torch 13 | 14 | 15 | def to_numpy(x): 16 | if isinstance(x, torch.Tensor): 17 | x = x.clone().detach().cpu().numpy() 18 | assert isinstance(x, np.ndarray) 19 | return x 20 | 21 | 22 | def save_point_cloud(coord, color=None, file_path="pc.ply", logger=None): 23 | os.makedirs(os.path.dirname(file_path), exist_ok=True) 24 | coord = to_numpy(coord) 25 | if color is not None: 26 | color = to_numpy(color) 27 | pcd = o3d.geometry.PointCloud() 28 | pcd.points = o3d.utility.Vector3dVector(coord) 29 | pcd.colors = o3d.utility.Vector3dVector( 30 | np.ones_like(coord) if color is None else color 31 | ) 32 | o3d.io.write_point_cloud(file_path, pcd) 33 | if logger is not None: 34 | logger.info(f"Save Point Cloud to: {file_path}") 35 | 36 | 37 | def save_bounding_boxes( 38 | bboxes_corners, color=(1.0, 0.0, 0.0), file_path="bbox.ply", logger=None 39 | ): 40 | bboxes_corners = to_numpy(bboxes_corners) 41 | # point list 42 | points = bboxes_corners.reshape(-1, 3) 43 | # line list 44 | box_lines = np.array( 45 | [ 46 | [0, 1], 47 | [1, 2], 48 | [2, 3], 49 | [3, 0], 50 | [4, 5], 51 | [5, 6], 52 | [6, 7], 53 | [7, 0], 54 | [0, 4], 55 | [1, 5], 56 | [2, 6], 57 | [3, 7], 58 | ] 59 | ) 60 | lines = [] 61 | for i, _ in enumerate(bboxes_corners): 62 | lines.append(box_lines + i * 8) 63 | lines = np.concatenate(lines) 64 | # color list 65 | color = np.array([color for _ in range(len(lines))]) 66 | # generate line set 67 | line_set = o3d.geometry.LineSet() 68 | line_set.points = o3d.utility.Vector3dVector(points) 69 | line_set.lines = o3d.utility.Vector2iVector(lines) 70 | line_set.colors = o3d.utility.Vector3dVector(color) 71 | o3d.io.write_line_set(file_path, line_set) 72 | 73 | if logger is not None: 74 | logger.info(f"Save Boxes to: {file_path}") 75 | 76 | 77 | def save_lines( 78 | points, lines, color=(1.0, 0.0, 0.0), file_path="lines.ply", logger=None 79 | ): 80 | points = to_numpy(points) 81 | lines = to_numpy(lines) 82 | colors = np.array([color for _ in range(len(lines))]) 83 | line_set = o3d.geometry.LineSet() 84 | line_set.points = o3d.utility.Vector3dVector(points) 85 | line_set.lines = o3d.utility.Vector2iVector(lines) 86 | line_set.colors = o3d.utility.Vector3dVector(colors) 87 | o3d.io.write_line_set(file_path, line_set) 88 | 89 | if logger is not None: 90 | logger.info(f"Save Lines to: {file_path}") 91 | -------------------------------------------------------------------------------- /scripts/test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | 3 | cd $(dirname $(dirname "$0")) || exit 4 | PYTHON=python 5 | 6 | TEST_CODE=test.py 7 | 8 | DATASET=scannet 9 | CONFIG="None" 10 | EXP_NAME=debug 11 | WEIGHT=model_best 12 | GPU=None 13 | 14 | while getopts "p:d:c:n:w:g:" opt; do 15 | case $opt in 16 | p) 17 | PYTHON=$OPTARG 18 | ;; 19 | d) 20 | DATASET=$OPTARG 21 | ;; 22 | c) 23 | CONFIG=$OPTARG 24 | ;; 25 | n) 26 | EXP_NAME=$OPTARG 27 | ;; 28 | w) 29 | WEIGHT=$OPTARG 30 | ;; 31 | g) 32 | GPU=$OPTARG 33 | ;; 34 | \?) 35 | echo "Invalid option: -$OPTARG" 36 | ;; 37 | esac 38 | done 39 | 40 | if [ "${NUM_GPU}" = 'None' ] 41 | then 42 | NUM_GPU=`$PYTHON -c 'import torch; print(torch.cuda.device_count())'` 43 | fi 44 | 45 | echo "Experiment name: $EXP_NAME" 46 | echo "Python interpreter dir: $PYTHON" 47 | echo "Dataset: $DATASET" 48 | echo "GPU Num: $GPU" 49 | 50 | EXP_DIR=exp/${DATASET}/${EXP_NAME} 51 | MODEL_DIR=${EXP_DIR}/model 52 | CODE_DIR=${EXP_DIR}/code 53 | CONFIG_DIR=${EXP_DIR}/config.py 54 | 55 | if [ "${CONFIG}" = "None" ] 56 | then 57 | CONFIG_DIR=${EXP_DIR}/config.py 58 | else 59 | CONFIG_DIR=configs/${DATASET}/${CONFIG}.py 60 | fi 61 | 62 | echo "Loading config in:" $CONFIG_DIR 63 | #export PYTHONPATH=./$CODE_DIR 64 | export PYTHONPATH=./ 65 | echo "Running code in: $CODE_DIR" 66 | 67 | 68 | echo " =========> RUN TASK <=========" 69 | 70 | #$PYTHON -u "$CODE_DIR"/tools/$TEST_CODE \ 71 | $PYTHON -u tools/$TEST_CODE \ 72 | --config-file "$CONFIG_DIR" \ 73 | --num-gpus "$GPU" \ 74 | --options save_path="$EXP_DIR" weight="${MODEL_DIR}"/"${WEIGHT}".pth 75 | -------------------------------------------------------------------------------- /scripts/train.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | 3 | cd $(dirname $(dirname "$0")) || exit 4 | ROOT_DIR=$(pwd) 5 | PYTHON=python 6 | 7 | TRAIN_CODE=train.py 8 | 9 | DATASET=scannet 10 | CONFIG="None" 11 | EXP_NAME=debug 12 | WEIGHT="None" 13 | RESUME=false 14 | GPU=None 15 | 16 | 17 | while getopts "p:d:c:n:w:g:r:" opt; do 18 | case $opt in 19 | p) 20 | PYTHON=$OPTARG 21 | ;; 22 | d) 23 | DATASET=$OPTARG 24 | ;; 25 | c) 26 | CONFIG=$OPTARG 27 | ;; 28 | n) 29 | EXP_NAME=$OPTARG 30 | ;; 31 | w) 32 | WEIGHT=$OPTARG 33 | ;; 34 | r) 35 | RESUME=$OPTARG 36 | ;; 37 | g) 38 | GPU=$OPTARG 39 | ;; 40 | \?) 41 | echo "Invalid option: -$OPTARG" 42 | ;; 43 | esac 44 | done 45 | 46 | if [ "${NUM_GPU}" = 'None' ] 47 | then 48 | NUM_GPU=`$PYTHON -c 'import torch; print(torch.cuda.device_count())'` 49 | fi 50 | 51 | echo "Experiment name: $EXP_NAME" 52 | echo "Python interpreter dir: $PYTHON" 53 | echo "Dataset: $DATASET" 54 | echo "Config: $CONFIG" 55 | echo "GPU Num: $GPU" 56 | 57 | EXP_DIR=exp/${DATASET}/${EXP_NAME} 58 | MODEL_DIR=${EXP_DIR}/model 59 | CODE_DIR=${EXP_DIR}/code 60 | CONFIG_DIR=configs/${DATASET}/${CONFIG}.py 61 | 62 | 63 | echo " =========> CREATE EXP DIR <=========" 64 | echo "Experiment dir: $ROOT_DIR/$EXP_DIR" 65 | if ${RESUME} 66 | then 67 | CONFIG_DIR=${EXP_DIR}/config.py 68 | WEIGHT=$MODEL_DIR/model_last.pth 69 | else 70 | mkdir -p "$MODEL_DIR" "$CODE_DIR" 71 | cp -r scripts tools ponder "$CODE_DIR" 72 | fi 73 | 74 | echo "Loading config in:" $CONFIG_DIR 75 | export PYTHONPATH=./$CODE_DIR 76 | echo "Running code in: $CODE_DIR" 77 | 78 | 79 | echo " =========> RUN TASK <=========" 80 | 81 | if [ "${WEIGHT}" = "None" ] 82 | then 83 | $PYTHON "$CODE_DIR"/tools/$TRAIN_CODE \ 84 | --config-file "$CONFIG_DIR" \ 85 | --num-gpus "$GPU" \ 86 | --options save_path="$EXP_DIR" 87 | else 88 | $PYTHON "$CODE_DIR"/tools/$TRAIN_CODE \ 89 | --config-file "$CONFIG_DIR" \ 90 | --num-gpus "$GPU" \ 91 | --options save_path="$EXP_DIR" resume="$RESUME" weight="$WEIGHT" 92 | fi -------------------------------------------------------------------------------- /tools/test.py: -------------------------------------------------------------------------------- 1 | """ 2 | Main Testing Script 3 | 4 | Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) 5 | Please cite our work if the code is helpful to you. 6 | """ 7 | 8 | from ponder.engines.defaults import ( 9 | default_argument_parser, 10 | default_config_parser, 11 | default_setup, 12 | ) 13 | from ponder.engines.launch import launch 14 | from ponder.engines.test import TESTERS 15 | 16 | 17 | def main_worker(cfg): 18 | cfg = default_setup(cfg) 19 | tester = TESTERS.build(dict(type=cfg.test.type, cfg=cfg)) 20 | tester.test() 21 | 22 | 23 | def main(): 24 | args = default_argument_parser().parse_args() 25 | cfg = default_config_parser(args.config_file, args.options) 26 | 27 | launch( 28 | main_worker, 29 | num_gpus_per_machine=args.num_gpus, 30 | num_machines=args.num_machines, 31 | machine_rank=args.machine_rank, 32 | dist_url=args.dist_url, 33 | cfg=(cfg,), 34 | ) 35 | 36 | 37 | if __name__ == "__main__": 38 | main() 39 | -------------------------------------------------------------------------------- /tools/test_s3dis_6fold.py: -------------------------------------------------------------------------------- 1 | """ 2 | Test script for S3DIS 6-fold cross validation 3 | 4 | Gathering Area_X.pth from result folder of experiment record of each area as follows: 5 | |- RECORDS_PATH 6 | |- Area_1.pth 7 | |- Area_2.pth 8 | |- Area_3.pth 9 | |- Area_4.pth 10 | |- Area_5.pth 11 | |- Area_6.pth 12 | 13 | Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) 14 | Please cite our work if the code is helpful to you. 15 | """ 16 | 17 | import argparse 18 | import glob 19 | import os 20 | 21 | import numpy as np 22 | import torch 23 | 24 | from ponder.utils.logger import get_root_logger 25 | 26 | CLASS_NAMES = [ 27 | "ceiling", 28 | "floor", 29 | "wall", 30 | "beam", 31 | "column", 32 | "window", 33 | "door", 34 | "table", 35 | "chair", 36 | "sofa", 37 | "bookcase", 38 | "board", 39 | "clutter", 40 | ] 41 | 42 | 43 | def evaluation(intersection, union, target, logger=None): 44 | iou_class = intersection / (union + 1e-10) 45 | accuracy_class = intersection / (target + 1e-10) 46 | mIoU = np.mean(iou_class) 47 | mAcc = np.mean(accuracy_class) 48 | allAcc = sum(intersection) / (sum(target) + 1e-10) 49 | 50 | if logger is not None: 51 | logger.info( 52 | "Val result: mIoU/mAcc/allAcc {:.4f}/{:.4f}/{:.4f}".format( 53 | mIoU, mAcc, allAcc 54 | ) 55 | ) 56 | for i in range(len(CLASS_NAMES)): 57 | logger.info( 58 | "Class_{idx} - {name} Result: iou/accuracy {iou:.4f}/{accuracy:.4f}".format( 59 | idx=i, 60 | name=CLASS_NAMES[i], 61 | iou=iou_class[i], 62 | accuracy=accuracy_class[i], 63 | ) 64 | ) 65 | 66 | 67 | def main(): 68 | parser = argparse.ArgumentParser() 69 | parser.add_argument( 70 | "--record_root", 71 | required=True, 72 | help="Path to the S3DIS record of each split", 73 | ) 74 | config = parser.parse_args() 75 | logger = get_root_logger( 76 | log_file=os.path.join(config.record_root, "6-fold.log"), 77 | file_mode="w", 78 | ) 79 | 80 | records = sorted(glob.glob(os.path.join(config.record_root, "Area_*.pth"))) 81 | assert len(records) == 6 82 | intersection_ = np.zeros(len(CLASS_NAMES), dtype=int) 83 | union_ = np.zeros(len(CLASS_NAMES), dtype=int) 84 | target_ = np.zeros(len(CLASS_NAMES), dtype=int) 85 | 86 | for record in records: 87 | area = os.path.basename(record).split(".")[0] 88 | info = torch.load(record) 89 | logger.info(f"<<<<<<<<<<<<<<<<< Parsing {area} <<<<<<<<<<<<<<<<<") 90 | intersection = info["intersection"] 91 | union = info["union"] 92 | target = info["target"] 93 | evaluation(intersection, union, target, logger=logger) 94 | intersection_ += intersection 95 | union_ += union 96 | target_ += target 97 | 98 | logger.info(f"<<<<<<<<<<<<<<<<< Parsing 6-fold <<<<<<<<<<<<<<<<<") 99 | evaluation(intersection_, union_, target_, logger=logger) 100 | 101 | 102 | if __name__ == "__main__": 103 | main() 104 | -------------------------------------------------------------------------------- /tools/train.py: -------------------------------------------------------------------------------- 1 | """ 2 | Main Training Script 3 | 4 | Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) 5 | Please cite our work if the code is helpful to you. 6 | """ 7 | 8 | from ponder.engines.defaults import ( 9 | default_argument_parser, 10 | default_config_parser, 11 | default_setup, 12 | ) 13 | from ponder.engines.launch import launch, slurm_launch 14 | from ponder.engines.train import TRAINERS 15 | 16 | 17 | def main_worker(cfg): 18 | cfg = default_setup(cfg) 19 | trainer = TRAINERS.build(dict(type=cfg.train.type, cfg=cfg)) 20 | trainer.train() 21 | 22 | 23 | def main(): 24 | args = default_argument_parser().parse_args() 25 | cfg = default_config_parser(args.config_file, args.options) 26 | 27 | if args.launcher == "pytorch": 28 | launcher = launch 29 | elif args.launcher == "slurm": 30 | launcher = slurm_launch 31 | 32 | launcher( 33 | main_worker, 34 | num_gpus_per_machine=args.num_gpus, 35 | num_machines=args.num_machines, 36 | machine_rank=args.machine_rank, 37 | dist_url=args.dist_url, 38 | port=args.master_port, 39 | cfg=(cfg,), 40 | ) 41 | 42 | 43 | if __name__ == "__main__": 44 | main() 45 | --------------------------------------------------------------------------------