├── scannet ├── scans │ └── .gitignore ├── meta_data │ ├── scannet_means.npz │ ├── scannetv2_test.txt │ └── scannetv2_train_0.1.txt ├── README.md ├── data_viz.py ├── scannet_utils.py ├── batch_load_scannet_data.py └── model_util_scannet.py ├── sunrgbd ├── OFFICIAL_SUNRGBD │ └── .gitignore ├── README.md ├── matlab │ ├── extract_split.m │ ├── extract_rgbd_data_v1.m │ └── extract_rgbd_data_v2.m └── sunrgbd_trainval │ └── sunrgbd_v1_train_0.05.txt ├── OpenPCDet ├── pcdet │ ├── version.py │ ├── models │ │ ├── backbones_2d │ │ │ ├── __init__.py │ │ │ ├── map_to_bev │ │ │ │ ├── __init__.py │ │ │ │ ├── height_compression.py │ │ │ │ └── pointpillar_scatter.py │ │ │ └── base_bev_backbone.py │ │ ├── backbones_3d │ │ │ ├── pfe │ │ │ │ └── __init__.py │ │ │ ├── vfe │ │ │ │ ├── __init__.py │ │ │ │ ├── vfe_template.py │ │ │ │ └── mean_vfe.py │ │ │ └── __init__.py │ │ ├── roi_heads │ │ │ └── __init__.py │ │ ├── dense_heads │ │ │ ├── __init__.py │ │ │ ├── anchor_head_single.py │ │ │ ├── point_head_simple.py │ │ │ └── target_assigner │ │ │ │ └── anchor_generator.py │ │ ├── detectors │ │ │ ├── __init__.py │ │ │ ├── point_rcnn.py │ │ │ ├── second_net.py │ │ │ ├── pointpillar.py │ │ │ ├── pv_rcnn.py │ │ │ └── PartA2_net.py │ │ ├── __init__.py │ │ └── model_utils │ │ │ └── model_nms_utils.py │ ├── ops │ │ ├── pointnet2 │ │ │ ├── pointnet2_stack │ │ │ │ └── src │ │ │ │ │ ├── cuda_utils.h │ │ │ │ │ ├── sampling_gpu.h │ │ │ │ │ ├── ball_query_gpu.h │ │ │ │ │ ├── pointnet2_api.cpp │ │ │ │ │ ├── sampling.cpp │ │ │ │ │ ├── group_points_gpu.h │ │ │ │ │ ├── interpolate_gpu.h │ │ │ │ │ ├── ball_query.cpp │ │ │ │ │ ├── group_points.cpp │ │ │ │ │ ├── ball_query_gpu.cu │ │ │ │ │ └── interpolate.cpp │ │ │ └── pointnet2_batch │ │ │ │ └── src │ │ │ │ ├── cuda_utils.h │ │ │ │ ├── ball_query_gpu.h │ │ │ │ ├── group_points_gpu.h │ │ │ │ ├── sampling_gpu.h │ │ │ │ ├── interpolate_gpu.h │ │ │ │ ├── pointnet2_api.cpp │ │ │ │ ├── group_points.cpp │ │ │ │ ├── ball_query.cpp │ │ │ │ ├── sampling.cpp │ │ │ │ ├── interpolate.cpp │ │ │ │ ├── ball_query_gpu.cu │ │ │ │ └── group_points_gpu.cu │ │ ├── iou3d_nms │ │ │ ├── src │ │ │ │ ├── iou3d_cpu.h │ │ │ │ ├── iou3d_nms.h │ │ │ │ └── iou3d_nms_api.cpp │ │ │ └── iou3d_nms_utils.py │ │ ├── roipoint_pool3d │ │ │ ├── src │ │ │ │ └── roipoint_pool3d.cpp │ │ │ └── roipoint_pool3d_utils.py │ │ └── roiaware_pool3d │ │ │ └── roiaware_pool3d_utils.py │ ├── __init__.py │ ├── datasets │ │ ├── kitti │ │ │ └── kitti_object_eval_python │ │ │ │ ├── evaluate.py │ │ │ │ ├── LICENSE │ │ │ │ └── README.md │ │ ├── processor │ │ │ └── point_feature_encoder.py │ │ ├── __init__.py │ │ └── augmentor │ │ │ ├── augmentor_utils.py │ │ │ └── data_augmentor.py │ ├── config.py │ └── utils │ │ └── object3d_kitti.py ├── tools │ ├── scripts │ │ ├── dist_test.sh │ │ ├── dist_train.sh │ │ ├── slurm_test_single.sh │ │ ├── slurm_test_mgpu.sh │ │ └── slurm_train.sh │ ├── cfgs │ │ ├── dataset_configs │ │ │ ├── kitti_dataset.yaml │ │ │ └── nuscenes_dataset.yaml │ │ └── kitti_models │ │ │ ├── second.yaml │ │ │ ├── PartA2_free.yaml │ │ │ └── second_multihead.yaml │ ├── train_utils │ │ └── optimization │ │ │ ├── __init__.py │ │ │ └── learning_schedules_fastai.py │ └── demo.py └── setup.py ├── imgs └── teaser.png ├── requirements.txt ├── run_pretrain.sh ├── run_eval.sh ├── run_train.sh ├── run_eval_opt.sh ├── pointnet2 ├── _ext_src │ ├── src │ │ ├── ball_query.h │ │ ├── group_points.h │ │ ├── sampling.h │ │ ├── interpolate.h │ │ ├── bindings.cpp │ │ ├── utils.h │ │ ├── ball_query.cpp │ │ ├── cuda_utils.h │ │ ├── ball_query_gpu.cu │ │ ├── group_points.cpp │ │ ├── group_points_gpu.cu │ │ ├── sampling.cpp │ │ └── interpolate.cpp │ └── include │ │ ├── ball_query.h │ │ ├── group_points.h │ │ ├── sampling.h │ │ ├── interpolate.h │ │ ├── utils.h │ │ └── cuda_utils.h └── setup.py ├── utils ├── tf_visualizer.py └── tf_logger.py ├── .gitignore ├── models └── voting_module.py └── generate_random_split.py /scannet/scans/.gitignore: -------------------------------------------------------------------------------- 1 | * 2 | !.gitignore 3 | -------------------------------------------------------------------------------- /sunrgbd/OFFICIAL_SUNRGBD/.gitignore: -------------------------------------------------------------------------------- 1 | * 2 | !.gitignore 3 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/version.py: -------------------------------------------------------------------------------- 1 | __version__ = "0.3.0+0000000" 2 | -------------------------------------------------------------------------------- /imgs/teaser.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yezhen17/3DIoUMatch/HEAD/imgs/teaser.png -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | matplotlib 2 | opencv-python 3 | plyfile 4 | trimesh==2.35.39 5 | mayavi 6 | vtk_visualizer -------------------------------------------------------------------------------- /scannet/meta_data/scannet_means.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yezhen17/3DIoUMatch/HEAD/scannet/meta_data/scannet_means.npz -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/backbones_2d/__init__.py: -------------------------------------------------------------------------------- 1 | from .base_bev_backbone import BaseBEVBackbone 2 | 3 | __all__ = { 4 | 'BaseBEVBackbone': BaseBEVBackbone 5 | } 6 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/backbones_3d/pfe/__init__.py: -------------------------------------------------------------------------------- 1 | from .voxel_set_abstraction import VoxelSetAbstraction 2 | 3 | __all__ = { 4 | 'VoxelSetAbstraction': VoxelSetAbstraction 5 | } 6 | -------------------------------------------------------------------------------- /OpenPCDet/tools/scripts/dist_test.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -x 4 | NGPUS=$1 5 | PY_ARGS=${@:2} 6 | 7 | python -m torch.distributed.launch --nproc_per_node=${NGPUS} test.py --launcher pytorch ${PY_ARGS} 8 | 9 | -------------------------------------------------------------------------------- /OpenPCDet/tools/scripts/dist_train.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -x 4 | NGPUS=$1 5 | PY_ARGS=${@:2} 6 | 7 | python -m torch.distributed.launch --nproc_per_node=${NGPUS} train.py --launcher pytorch ${PY_ARGS} 8 | 9 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_stack/src/cuda_utils.h: -------------------------------------------------------------------------------- 1 | #ifndef _STACK_CUDA_UTILS_H 2 | #define _STACK_CUDA_UTILS_H 3 | 4 | #include 5 | 6 | #define THREADS_PER_BLOCK 256 7 | #define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0)) 8 | 9 | #endif 10 | -------------------------------------------------------------------------------- /run_pretrain.sh: -------------------------------------------------------------------------------- 1 | export CUDA_VISIBLE_DEVICES=$1 2 | LOG_DIR=$2 3 | DATASET=$3 4 | LABELED_LIST=$4 5 | mkdir -p "${LOG_DIR}" 6 | python -u pretrain.py --log_dir="${LOG_DIR}" --dataset="${DATASET}" \ 7 | --labeled_sample_list="${LABELED_LIST}" 2>&1|tee "${LOG_DIR}"/LOG_ALL.log & -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/backbones_2d/map_to_bev/__init__.py: -------------------------------------------------------------------------------- 1 | from .height_compression import HeightCompression 2 | from .pointpillar_scatter import PointPillarScatter 3 | 4 | __all__ = { 5 | 'HeightCompression': HeightCompression, 6 | 'PointPillarScatter': PointPillarScatter 7 | } 8 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/backbones_3d/vfe/__init__.py: -------------------------------------------------------------------------------- 1 | from .mean_vfe import MeanVFE 2 | from .pillar_vfe import PillarVFE 3 | from .vfe_template import VFETemplate 4 | 5 | __all__ = { 6 | 'VFETemplate': VFETemplate, 7 | 'MeanVFE': MeanVFE, 8 | 'PillarVFE': PillarVFE 9 | } 10 | -------------------------------------------------------------------------------- /run_eval.sh: -------------------------------------------------------------------------------- 1 | export CUDA_VISIBLE_DEVICES=$1 2 | LOG_DIR=$2 3 | DATASET=$3 4 | LABELED_LIST=$4 5 | CKPT=$5 6 | mkdir -p "${LOG_DIR}"; 7 | python -u train.py --log_dir="${LOG_DIR}" --dataset="${DATASET}" --detector_checkpoint="${CKPT}" \ 8 | --labeled_sample_list="${LABELED_LIST}" --use_iou_for_nms --eval 9 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/iou3d_nms/src/iou3d_cpu.h: -------------------------------------------------------------------------------- 1 | #ifndef IOU3D_CPU_H 2 | #define IOU3D_CPU_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | 9 | int boxes_iou_bev_cpu(at::Tensor boxes_a_tensor, at::Tensor boxes_b_tensor, at::Tensor ans_iou_tensor); 10 | 11 | #endif 12 | -------------------------------------------------------------------------------- /run_train.sh: -------------------------------------------------------------------------------- 1 | export CUDA_VISIBLE_DEVICES=$1 2 | LOG_DIR=$2 3 | DATASET=$3 4 | LABELED_LIST=$4 5 | PRETRAIN_CKPT=$5 6 | mkdir -p "${LOG_DIR}"; 7 | python -u train.py --log_dir="${LOG_DIR}" --dataset="${DATASET}" \ 8 | --labeled_sample_list="${LABELED_LIST}" --detector_checkpoint="${PRETRAIN_CKPT}" --view_stats \ 9 | 2>&1|tee "${LOG_DIR}"/LOG.log & 10 | 11 | -------------------------------------------------------------------------------- /run_eval_opt.sh: -------------------------------------------------------------------------------- 1 | export CUDA_VISIBLE_DEVICES=$1 2 | LOG_DIR=$2 3 | DATASET=$3 4 | LABELED_LIST=$4 5 | CKPT=$5 6 | OPT_RATE=$6 7 | mkdir -p "${LOG_DIR}"; 8 | python -u train.py --log_dir="${LOG_DIR}" --dataset="${DATASET}" --detector_checkpoint="${CKPT}" \ 9 | --labeled_sample_list="${LABELED_LIST}" --use_iou_for_nms --eval --opt_step=10 --opt_rate="${OPT_RATE}" 10 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/roi_heads/__init__.py: -------------------------------------------------------------------------------- 1 | from .partA2_head import PartA2FCHead 2 | from .pointrcnn_head import PointRCNNHead 3 | from .pvrcnn_head import PVRCNNHead 4 | from .roi_head_template import RoIHeadTemplate 5 | 6 | __all__ = { 7 | 'RoIHeadTemplate': RoIHeadTemplate, 8 | 'PartA2FCHead': PartA2FCHead, 9 | 'PVRCNNHead': PVRCNNHead, 10 | 'PointRCNNHead': PointRCNNHead 11 | } 12 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/ball_query.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | #include 8 | 9 | at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, 10 | const int nsample); 11 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/include/ball_query.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | #include 8 | 9 | at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, 10 | const int nsample); 11 | -------------------------------------------------------------------------------- /scannet/README.md: -------------------------------------------------------------------------------- 1 | ### Prepare ScanNet Data 2 | 3 | 1. Download ScanNet v2 data [HERE](https://github.com/ScanNet/ScanNet). Move/link the `scans` folder such that under `scans` there should be folders with names such as `scene0001_01`. 4 | 5 | 2. Extract point clouds and annotations (semantic seg, instance seg etc.) by running `python batch_load_scannet_data.py`, which will create a folder named `scannet_train_detection_data` here. 6 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/include/group_points.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | #include 8 | 9 | at::Tensor group_points(at::Tensor points, at::Tensor idx); 10 | at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); 11 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/group_points.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | #include 8 | 9 | at::Tensor group_points(at::Tensor points, at::Tensor idx); 10 | at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); 11 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/backbones_3d/__init__.py: -------------------------------------------------------------------------------- 1 | from .pointnet2_backbone import PointNet2Backbone, PointNet2MSG 2 | from .spconv_backbone import VoxelBackBone8x, VoxelResBackBone8x 3 | from .spconv_unet import UNetV2 4 | 5 | __all__ = { 6 | 'VoxelBackBone8x': VoxelBackBone8x, 7 | 'UNetV2': UNetV2, 8 | 'PointNet2Backbone': PointNet2Backbone, 9 | 'PointNet2MSG': PointNet2MSG, 10 | 'VoxelResBackBone8x': VoxelResBackBone8x, 11 | } 12 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_batch/src/cuda_utils.h: -------------------------------------------------------------------------------- 1 | #ifndef _CUDA_UTILS_H 2 | #define _CUDA_UTILS_H 3 | 4 | #include 5 | 6 | #define TOTAL_THREADS 1024 7 | #define THREADS_PER_BLOCK 256 8 | #define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0)) 9 | 10 | inline int opt_n_threads(int work_size) { 11 | const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0); 12 | 13 | return max(min(1 << pow_2, TOTAL_THREADS), 1); 14 | } 15 | #endif 16 | -------------------------------------------------------------------------------- /OpenPCDet/tools/scripts/slurm_test_single.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -x 4 | 5 | PARTITION=$1 6 | GPUS=1 7 | GPUS_PER_NODE=1 8 | PY_ARGS=${@:2} 9 | JOB_NAME=eval 10 | SRUN_ARGS=${SRUN_ARGS:-""} 11 | 12 | srun -p ${PARTITION} \ 13 | --job-name=${JOB_NAME} \ 14 | --gres=gpu:${GPUS_PER_NODE} \ 15 | --ntasks=${GPUS} \ 16 | --ntasks-per-node=${GPUS_PER_NODE} \ 17 | --kill-on-bad-exit=1 \ 18 | ${SRUN_ARGS} \ 19 | python -u test.py ${PY_ARGS} 20 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/sampling.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | #include 8 | 9 | at::Tensor gather_points(at::Tensor points, at::Tensor idx); 10 | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); 11 | at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples); 12 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/include/sampling.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | #include 8 | 9 | at::Tensor gather_points(at::Tensor points, at::Tensor idx); 10 | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); 11 | at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples); 12 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_stack/src/sampling_gpu.h: -------------------------------------------------------------------------------- 1 | #ifndef _SAMPLING_GPU_H 2 | #define _SAMPLING_GPU_H 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | 9 | int furthest_point_sampling_wrapper(int b, int n, int m, 10 | at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor idx_tensor); 11 | 12 | void furthest_point_sampling_kernel_launcher(int b, int n, int m, 13 | const float *dataset, float *temp, int *idxs); 14 | 15 | #endif 16 | -------------------------------------------------------------------------------- /OpenPCDet/tools/scripts/slurm_test_mgpu.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -x 4 | 5 | PARTITION=$1 6 | GPUS=$2 7 | GPUS_PER_NODE=$GPUS 8 | PY_ARGS=${@:3} 9 | JOB_NAME=eval 10 | SRUN_ARGS=${SRUN_ARGS:-""} 11 | 12 | PORT=$(( ( RANDOM % 10000 ) + 10000 )) 13 | 14 | srun -p ${PARTITION} \ 15 | --job-name=${JOB_NAME} \ 16 | --gres=gpu:${GPUS_PER_NODE} \ 17 | --ntasks=${GPUS} \ 18 | --ntasks-per-node=${GPUS_PER_NODE} \ 19 | --kill-on-bad-exit=1 \ 20 | ${SRUN_ARGS} \ 21 | python -u test.py --launcher slurm ${PY_ARGS} --tcp_port $PORT 22 | 23 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/iou3d_nms/src/iou3d_nms.h: -------------------------------------------------------------------------------- 1 | #ifndef IOU3D_NMS_H 2 | #define IOU3D_NMS_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | 9 | int boxes_overlap_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b, at::Tensor ans_overlap); 10 | int boxes_iou_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b, at::Tensor ans_iou); 11 | int nms_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh); 12 | int nms_normal_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh); 13 | 14 | #endif 15 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_batch/src/ball_query_gpu.h: -------------------------------------------------------------------------------- 1 | #ifndef _BALL_QUERY_GPU_H 2 | #define _BALL_QUERY_GPU_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | 9 | int ball_query_wrapper_fast(int b, int n, int m, float radius, int nsample, 10 | at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor idx_tensor); 11 | 12 | void ball_query_kernel_launcher_fast(int b, int n, int m, float radius, int nsample, 13 | const float *xyz, const float *new_xyz, int *idx); 14 | 15 | #endif 16 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/backbones_3d/vfe/vfe_template.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | 4 | class VFETemplate(nn.Module): 5 | def __init__(self, model_cfg, **kwargs): 6 | super().__init__() 7 | self.model_cfg = model_cfg 8 | 9 | def get_output_feature_dim(self): 10 | raise NotImplementedError 11 | 12 | def forward(self, **kwargs): 13 | """ 14 | Args: 15 | **kwargs: 16 | 17 | Returns: 18 | batch_dict: 19 | ... 20 | vfe_features: (num_voxels, C) 21 | """ 22 | raise NotImplementedError 23 | -------------------------------------------------------------------------------- /OpenPCDet/tools/scripts/slurm_train.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -x 4 | 5 | PARTITION=$1 6 | JOB_NAME=$2 7 | GPUS=$3 8 | PY_ARGS=${@:4} 9 | 10 | GPUS_PER_NODE=${GPUS_PER_NODE:-8} 11 | CPUS_PER_TASK=${CPUS_PER_TASK:-5} 12 | SRUN_ARGS=${SRUN_ARGS:-""} 13 | 14 | PORT=$(( ( RANDOM % 10000 ) + 10000 )) 15 | 16 | srun -p ${PARTITION} \ 17 | --job-name=${JOB_NAME} \ 18 | --gres=gpu:${GPUS_PER_NODE} \ 19 | --ntasks=${GPUS} \ 20 | --ntasks-per-node=${GPUS_PER_NODE} \ 21 | --cpus-per-task=${CPUS_PER_TASK} \ 22 | --kill-on-bad-exit=1 \ 23 | ${SRUN_ARGS} \ 24 | python -u train.py --launcher slurm ${PY_ARGS} --tcp_port $PORT 25 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/__init__.py: -------------------------------------------------------------------------------- 1 | import subprocess 2 | from pathlib import Path 3 | 4 | from .version import __version__ 5 | 6 | __all__ = [ 7 | '__version__' 8 | ] 9 | 10 | 11 | def get_git_commit_number(): 12 | if not (Path(__file__).parent / '../.git').exists(): 13 | return '0000000' 14 | 15 | cmd_out = subprocess.run(['git', 'rev-parse', 'HEAD'], stdout=subprocess.PIPE) 16 | git_commit_number = cmd_out.stdout.decode('utf-8')[:7] 17 | return git_commit_number 18 | 19 | 20 | script_version = get_git_commit_number() 21 | 22 | 23 | if script_version not in __version__: 24 | __version__ = __version__ + '+py%s' % script_version 25 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/interpolate.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | 8 | #include 9 | #include 10 | 11 | std::vector three_nn(at::Tensor unknowns, at::Tensor knows); 12 | at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, 13 | at::Tensor weight); 14 | at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, 15 | at::Tensor weight, const int m); 16 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/include/interpolate.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | 8 | #include 9 | #include 10 | 11 | std::vector three_nn(at::Tensor unknowns, at::Tensor knows); 12 | at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, 13 | at::Tensor weight); 14 | at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, 15 | at::Tensor weight, const int m); 16 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/iou3d_nms/src/iou3d_nms_api.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include 6 | 7 | #include "iou3d_cpu.h" 8 | #include "iou3d_nms.h" 9 | 10 | 11 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 12 | m.def("boxes_overlap_bev_gpu", &boxes_overlap_bev_gpu, "oriented boxes overlap"); 13 | m.def("boxes_iou_bev_gpu", &boxes_iou_bev_gpu, "oriented boxes iou"); 14 | m.def("nms_gpu", &nms_gpu, "oriented nms gpu"); 15 | m.def("nms_normal_gpu", &nms_normal_gpu, "nms gpu"); 16 | m.def("boxes_iou_bev_cpu", &boxes_iou_bev_cpu, "oriented boxes iou"); 17 | } 18 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/dense_heads/__init__.py: -------------------------------------------------------------------------------- 1 | from .anchor_head_multi import AnchorHeadMulti 2 | from .anchor_head_single import AnchorHeadSingle 3 | from .anchor_head_template import AnchorHeadTemplate 4 | from .point_head_box import PointHeadBox 5 | from .point_head_simple import PointHeadSimple 6 | from .point_intra_part_head import PointIntraPartOffsetHead 7 | 8 | __all__ = { 9 | 'AnchorHeadTemplate': AnchorHeadTemplate, 10 | 'AnchorHeadSingle': AnchorHeadSingle, 11 | 'PointIntraPartOffsetHead': PointIntraPartOffsetHead, 12 | 'PointHeadSimple': PointHeadSimple, 13 | 'PointHeadBox': PointHeadBox, 14 | 'AnchorHeadMulti': AnchorHeadMulti, 15 | } 16 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/detectors/__init__.py: -------------------------------------------------------------------------------- 1 | from .detector3d_template import Detector3DTemplate 2 | from .PartA2_net import PartA2Net 3 | from .point_rcnn import PointRCNN 4 | from .pointpillar import PointPillar 5 | from .pv_rcnn import PVRCNN 6 | from .second_net import SECONDNet 7 | 8 | __all__ = { 9 | 'Detector3DTemplate': Detector3DTemplate, 10 | 'SECONDNet': SECONDNet, 11 | 'PartA2Net': PartA2Net, 12 | 'PVRCNN': PVRCNN, 13 | 'PointPillar': PointPillar, 14 | 'PointRCNN': PointRCNN 15 | } 16 | 17 | 18 | def build_detector(model_cfg, num_class, dataset): 19 | model = __all__[model_cfg.NAME]( 20 | model_cfg=model_cfg, num_class=num_class, dataset=dataset 21 | ) 22 | 23 | return model 24 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/bindings.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include "ball_query.h" 7 | #include "group_points.h" 8 | #include "interpolate.h" 9 | #include "sampling.h" 10 | 11 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 12 | m.def("gather_points", &gather_points); 13 | m.def("gather_points_grad", &gather_points_grad); 14 | m.def("furthest_point_sampling", &furthest_point_sampling); 15 | 16 | m.def("three_nn", &three_nn); 17 | m.def("three_interpolate", &three_interpolate); 18 | m.def("three_interpolate_grad", &three_interpolate_grad); 19 | 20 | m.def("ball_query", &ball_query); 21 | 22 | m.def("group_points", &group_points); 23 | m.def("group_points_grad", &group_points_grad); 24 | } 25 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_stack/src/ball_query_gpu.h: -------------------------------------------------------------------------------- 1 | /* 2 | Stacked-batch-data version of ball query, modified from the original implementation of official PointNet++ codes. 3 | Written by Shaoshuai Shi 4 | All Rights Reserved 2019-2020. 5 | */ 6 | 7 | 8 | #ifndef _STACK_BALL_QUERY_GPU_H 9 | #define _STACK_BALL_QUERY_GPU_H 10 | 11 | #include 12 | #include 13 | #include 14 | #include 15 | 16 | int ball_query_wrapper_stack(int B, int M, float radius, int nsample, 17 | at::Tensor new_xyz_tensor, at::Tensor new_xyz_batch_cnt_tensor, 18 | at::Tensor xyz_tensor, at::Tensor xyz_batch_cnt_tensor, at::Tensor idx_tensor); 19 | 20 | 21 | void ball_query_kernel_launcher_stack(int B, int M, float radius, int nsample, 22 | const float *new_xyz, const int *new_xyz_batch_cnt, const float *xyz, const int *xyz_batch_cnt, int *idx); 23 | 24 | 25 | #endif 26 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_batch/src/group_points_gpu.h: -------------------------------------------------------------------------------- 1 | #ifndef _GROUP_POINTS_GPU_H 2 | #define _GROUP_POINTS_GPU_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | 9 | 10 | int group_points_wrapper_fast(int b, int c, int n, int npoints, int nsample, 11 | at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor); 12 | 13 | void group_points_kernel_launcher_fast(int b, int c, int n, int npoints, int nsample, 14 | const float *points, const int *idx, float *out); 15 | 16 | int group_points_grad_wrapper_fast(int b, int c, int n, int npoints, int nsample, 17 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor); 18 | 19 | void group_points_grad_kernel_launcher_fast(int b, int c, int n, int npoints, int nsample, 20 | const float *grad_out, const int *idx, float *grad_points); 21 | 22 | #endif 23 | -------------------------------------------------------------------------------- /sunrgbd/README.md: -------------------------------------------------------------------------------- 1 | ### Prepare SUN RGB-D Data 2 | 3 | 1. Download SUNRGBD v2 data [HERE](http://rgbd.cs.princeton.edu/data/) (SUNRGBD.zip, SUNRGBDMeta2DBB_v2.mat, SUNRGBDMeta3DBB_v2.mat) and the toolkits (SUNRGBDtoolbox.zip). Move all the downloaded files under OFFICIAL_SUNRGBD. Unzip the zip files. 4 | 5 | 2. Extract point clouds and annotations (class, v2 2D -- xmin,ymin,xmax,ymax, and 3D bounding boxes -- centroids, size, 2D heading) by running `extract_split.m`, `extract_rgbd_data_v2.m` and `extract_rgbd_data_v1.m` under the `matlab` folder. 6 | 7 | 3. Prepare data by running `python sunrgbd_data.py --gen_v1_data` 8 | 9 | You can also examine and visualize the data with `python sunrgbd_data.py --viz` and use MeshLab to view the generated PLY files at `data_viz_dump`. 10 | 11 | NOTE: SUNRGBDtoolbox.zip should have MD5 hash `18d22e1761d36352f37232cba102f91f` (you can check the hash with `md5 SUNRGBDtoolbox.zip` on Mac OS or `md5sum SUNRGBDtoolbox.zip` on Linux) 12 | -------------------------------------------------------------------------------- /pointnet2/setup.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | from setuptools import setup 7 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 8 | import glob 9 | 10 | _ext_src_root = "_ext_src" 11 | _ext_sources = glob.glob("{}/src/*.cpp".format(_ext_src_root)) + glob.glob( 12 | "{}/src/*.cu".format(_ext_src_root) 13 | ) 14 | # _ext_headers = glob.glob("{}/include/*".format(_ext_src_root)) 15 | 16 | setup( 17 | name='pointnet2', 18 | ext_modules=[ 19 | CUDAExtension( 20 | name='pointnet2._ext', 21 | sources=_ext_sources, 22 | extra_compile_args={ 23 | "cxx": ["-O2"], 24 | "nvcc": ["-O2"], 25 | }, 26 | ) 27 | ], 28 | cmdclass={ 29 | 'build_ext': BuildExtension 30 | } 31 | ) 32 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/backbones_2d/map_to_bev/height_compression.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | 4 | class HeightCompression(nn.Module): 5 | def __init__(self, model_cfg, **kwargs): 6 | super().__init__() 7 | self.model_cfg = model_cfg 8 | self.num_bev_features = self.model_cfg.NUM_BEV_FEATURES 9 | 10 | def forward(self, batch_dict): 11 | """ 12 | Args: 13 | batch_dict: 14 | encoded_spconv_tensor: sparse tensor 15 | Returns: 16 | batch_dict: 17 | spatial_features: 18 | 19 | """ 20 | encoded_spconv_tensor = batch_dict['encoded_spconv_tensor'] 21 | spatial_features = encoded_spconv_tensor.dense() 22 | N, C, D, H, W = spatial_features.shape 23 | spatial_features = spatial_features.view(N, C * D, H, W) 24 | batch_dict['spatial_features'] = spatial_features 25 | batch_dict['spatial_features_stride'] = batch_dict['encoded_spconv_tensor_stride'] 26 | return batch_dict 27 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_stack/src/pointnet2_api.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "ball_query_gpu.h" 5 | #include "group_points_gpu.h" 6 | #include "sampling_gpu.h" 7 | #include "interpolate_gpu.h" 8 | 9 | 10 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 11 | m.def("ball_query_wrapper", &ball_query_wrapper_stack, "ball_query_wrapper_stack"); 12 | 13 | m.def("furthest_point_sampling_wrapper", &furthest_point_sampling_wrapper, "furthest_point_sampling_wrapper"); 14 | 15 | m.def("group_points_wrapper", &group_points_wrapper_stack, "group_points_wrapper_stack"); 16 | m.def("group_points_grad_wrapper", &group_points_grad_wrapper_stack, "group_points_grad_wrapper_stack"); 17 | 18 | m.def("three_nn_wrapper", &three_nn_wrapper_stack, "three_nn_wrapper_stack"); 19 | m.def("three_interpolate_wrapper", &three_interpolate_wrapper_stack, "three_interpolate_wrapper_stack"); 20 | m.def("three_interpolate_grad_wrapper", &three_interpolate_grad_wrapper_stack, "three_interpolate_grad_wrapper_stack"); 21 | } 22 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/datasets/kitti/kitti_object_eval_python/evaluate.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | import fire 4 | 5 | import .kitti_common as kitti 6 | from .eval import get_coco_eval_result, get_official_eval_result 7 | 8 | 9 | def _read_imageset_file(path): 10 | with open(path, 'r') as f: 11 | lines = f.readlines() 12 | return [int(line) for line in lines] 13 | 14 | 15 | def evaluate(label_path, 16 | result_path, 17 | label_split_file, 18 | current_class=0, 19 | coco=False, 20 | score_thresh=-1): 21 | dt_annos = kitti.get_label_annos(result_path) 22 | if score_thresh > 0: 23 | dt_annos = kitti.filter_annos_low_score(dt_annos, score_thresh) 24 | val_image_ids = _read_imageset_file(label_split_file) 25 | gt_annos = kitti.get_label_annos(label_path, val_image_ids) 26 | if coco: 27 | return get_coco_eval_result(gt_annos, dt_annos, current_class) 28 | else: 29 | return get_official_eval_result(gt_annos, dt_annos, current_class) 30 | 31 | 32 | if __name__ == '__main__': 33 | fire.Fire() 34 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_batch/src/sampling_gpu.h: -------------------------------------------------------------------------------- 1 | #ifndef _SAMPLING_GPU_H 2 | #define _SAMPLING_GPU_H 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | 9 | int gather_points_wrapper_fast(int b, int c, int n, int npoints, 10 | at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor); 11 | 12 | void gather_points_kernel_launcher_fast(int b, int c, int n, int npoints, 13 | const float *points, const int *idx, float *out); 14 | 15 | 16 | int gather_points_grad_wrapper_fast(int b, int c, int n, int npoints, 17 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor); 18 | 19 | void gather_points_grad_kernel_launcher_fast(int b, int c, int n, int npoints, 20 | const float *grad_out, const int *idx, float *grad_points); 21 | 22 | 23 | int furthest_point_sampling_wrapper(int b, int n, int m, 24 | at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor idx_tensor); 25 | 26 | void furthest_point_sampling_kernel_launcher(int b, int n, int m, 27 | const float *dataset, float *temp, int *idxs); 28 | 29 | #endif 30 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/detectors/point_rcnn.py: -------------------------------------------------------------------------------- 1 | from .detector3d_template import Detector3DTemplate 2 | 3 | 4 | class PointRCNN(Detector3DTemplate): 5 | def __init__(self, model_cfg, num_class, dataset): 6 | super().__init__(model_cfg=model_cfg, num_class=num_class, dataset=dataset) 7 | self.module_list = self.build_networks() 8 | 9 | def forward(self, batch_dict): 10 | for cur_module in self.module_list: 11 | batch_dict = cur_module(batch_dict) 12 | 13 | if self.training: 14 | loss, tb_dict, disp_dict = self.get_training_loss() 15 | 16 | ret_dict = { 17 | 'loss': loss 18 | } 19 | return ret_dict, tb_dict, disp_dict 20 | else: 21 | pred_dicts, recall_dicts = self.post_processing(batch_dict) 22 | return pred_dicts, recall_dicts 23 | 24 | def get_training_loss(self): 25 | disp_dict = {} 26 | loss_point, tb_dict = self.point_head.get_loss() 27 | loss_rcnn, tb_dict = self.roi_head.get_loss(tb_dict) 28 | 29 | loss = loss_point + loss_rcnn 30 | return loss, tb_dict, disp_dict 31 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/datasets/kitti/kitti_object_eval_python/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 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 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/detectors/second_net.py: -------------------------------------------------------------------------------- 1 | from .detector3d_template import Detector3DTemplate 2 | 3 | 4 | class SECONDNet(Detector3DTemplate): 5 | def __init__(self, model_cfg, num_class, dataset): 6 | super().__init__(model_cfg=model_cfg, num_class=num_class, dataset=dataset) 7 | self.module_list = self.build_networks() 8 | 9 | def forward(self, batch_dict): 10 | for cur_module in self.module_list: 11 | batch_dict = cur_module(batch_dict) 12 | 13 | if self.training: 14 | loss, tb_dict, disp_dict = self.get_training_loss() 15 | 16 | ret_dict = { 17 | 'loss': loss 18 | } 19 | return ret_dict, tb_dict, disp_dict 20 | else: 21 | pred_dicts, recall_dicts = self.post_processing(batch_dict) 22 | return pred_dicts, recall_dicts 23 | 24 | def get_training_loss(self): 25 | disp_dict = {} 26 | 27 | loss_rpn, tb_dict = self.dense_head.get_loss() 28 | tb_dict = { 29 | 'loss_rpn': loss_rpn.item(), 30 | **tb_dict 31 | } 32 | 33 | loss = loss_rpn 34 | return loss, tb_dict, disp_dict 35 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/detectors/pointpillar.py: -------------------------------------------------------------------------------- 1 | from .detector3d_template import Detector3DTemplate 2 | 3 | 4 | class PointPillar(Detector3DTemplate): 5 | def __init__(self, model_cfg, num_class, dataset): 6 | super().__init__(model_cfg=model_cfg, num_class=num_class, dataset=dataset) 7 | self.module_list = self.build_networks() 8 | 9 | def forward(self, batch_dict): 10 | for cur_module in self.module_list: 11 | batch_dict = cur_module(batch_dict) 12 | 13 | if self.training: 14 | loss, tb_dict, disp_dict = self.get_training_loss() 15 | 16 | ret_dict = { 17 | 'loss': loss 18 | } 19 | return ret_dict, tb_dict, disp_dict 20 | else: 21 | pred_dicts, recall_dicts = self.post_processing(batch_dict) 22 | return pred_dicts, recall_dicts 23 | 24 | def get_training_loss(self): 25 | disp_dict = {} 26 | 27 | loss_rpn, tb_dict = self.dense_head.get_loss() 28 | tb_dict = { 29 | 'loss_rpn': loss_rpn.item(), 30 | **tb_dict 31 | } 32 | 33 | loss = loss_rpn 34 | return loss, tb_dict, disp_dict 35 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/backbones_3d/vfe/mean_vfe.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from .vfe_template import VFETemplate 4 | 5 | 6 | class MeanVFE(VFETemplate): 7 | def __init__(self, model_cfg, num_point_features, **kwargs): 8 | super().__init__(model_cfg=model_cfg) 9 | self.num_point_features = num_point_features 10 | 11 | def get_output_feature_dim(self): 12 | return self.num_point_features 13 | 14 | def forward(self, batch_dict, **kwargs): 15 | """ 16 | Args: 17 | batch_dict: 18 | voxels: (num_voxels, max_points_per_voxel, C) 19 | voxel_num_points: optional (num_voxels) 20 | **kwargs: 21 | 22 | Returns: 23 | vfe_features: (num_voxels, C) 24 | """ 25 | voxel_features, voxel_num_points = batch_dict['voxels'], batch_dict['voxel_num_points'] 26 | points_mean = voxel_features[:, :, :].sum(dim=1, keepdim=False) 27 | normalizer = torch.clamp_min(voxel_num_points.view(-1, 1), min=1.0).type_as(voxel_features) 28 | points_mean = points_mean / normalizer 29 | batch_dict['voxel_features'] = points_mean.contiguous() 30 | 31 | return batch_dict 32 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/detectors/pv_rcnn.py: -------------------------------------------------------------------------------- 1 | from .detector3d_template import Detector3DTemplate 2 | 3 | 4 | class PVRCNN(Detector3DTemplate): 5 | def __init__(self, model_cfg, num_class, dataset): 6 | super().__init__(model_cfg=model_cfg, num_class=num_class, dataset=dataset) 7 | self.module_list = self.build_networks() 8 | 9 | def forward(self, batch_dict): 10 | for cur_module in self.module_list: 11 | batch_dict = cur_module(batch_dict) 12 | 13 | if self.training: 14 | loss, tb_dict, disp_dict = self.get_training_loss() 15 | 16 | ret_dict = { 17 | 'loss': loss 18 | } 19 | return ret_dict, tb_dict, disp_dict 20 | else: 21 | pred_dicts, recall_dicts = self.post_processing(batch_dict) 22 | return pred_dicts, recall_dicts 23 | 24 | def get_training_loss(self): 25 | disp_dict = {} 26 | loss_rpn, tb_dict = self.dense_head.get_loss() 27 | loss_point, tb_dict = self.point_head.get_loss(tb_dict) 28 | loss_rcnn, tb_dict = self.roi_head.get_loss(tb_dict) 29 | 30 | loss = loss_rpn + loss_point + loss_rcnn 31 | return loss, tb_dict, disp_dict 32 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/detectors/PartA2_net.py: -------------------------------------------------------------------------------- 1 | from .detector3d_template import Detector3DTemplate 2 | 3 | 4 | class PartA2Net(Detector3DTemplate): 5 | def __init__(self, model_cfg, num_class, dataset): 6 | super().__init__(model_cfg=model_cfg, num_class=num_class, dataset=dataset) 7 | self.module_list = self.build_networks() 8 | 9 | def forward(self, batch_dict): 10 | for cur_module in self.module_list: 11 | batch_dict = cur_module(batch_dict) 12 | 13 | if self.training: 14 | loss, tb_dict, disp_dict = self.get_training_loss() 15 | 16 | ret_dict = { 17 | 'loss': loss 18 | } 19 | return ret_dict, tb_dict, disp_dict 20 | else: 21 | pred_dicts, recall_dicts = self.post_processing(batch_dict) 22 | return pred_dicts, recall_dicts 23 | 24 | def get_training_loss(self): 25 | disp_dict = {} 26 | loss_rpn, tb_dict = self.dense_head.get_loss() 27 | loss_point, tb_dict = self.point_head.get_loss(tb_dict) 28 | loss_rcnn, tb_dict = self.roi_head.get_loss(tb_dict) 29 | 30 | loss = loss_rpn + loss_point + loss_rcnn 31 | return loss, tb_dict, disp_dict 32 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_stack/src/sampling.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | 6 | #include "sampling_gpu.h" 7 | 8 | extern THCState *state; 9 | #define CHECK_CUDA(x) do { \ 10 | if (!x.type().is_cuda()) { \ 11 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 12 | exit(-1); \ 13 | } \ 14 | } while (0) 15 | #define CHECK_CONTIGUOUS(x) do { \ 16 | if (!x.is_contiguous()) { \ 17 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 18 | exit(-1); \ 19 | } \ 20 | } while (0) 21 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x) 22 | 23 | 24 | int furthest_point_sampling_wrapper(int b, int n, int m, 25 | at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor idx_tensor) { 26 | 27 | CHECK_INPUT(points_tensor); 28 | CHECK_INPUT(temp_tensor); 29 | CHECK_INPUT(idx_tensor); 30 | 31 | const float *points = points_tensor.data(); 32 | float *temp = temp_tensor.data(); 33 | int *idx = idx_tensor.data(); 34 | 35 | furthest_point_sampling_kernel_launcher(b, n, m, points, temp, idx); 36 | return 1; 37 | } 38 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/__init__.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | 3 | import numpy as np 4 | import torch 5 | 6 | from .detectors import build_detector 7 | 8 | 9 | def build_network(model_cfg, num_class, dataset): 10 | model = build_detector( 11 | model_cfg=model_cfg, num_class=num_class, dataset=dataset 12 | ) 13 | return model 14 | 15 | 16 | def load_data_to_gpu(batch_dict): 17 | for key, val in batch_dict.items(): 18 | if not isinstance(val, np.ndarray): 19 | continue 20 | if key in ['frame_id', 'metadata', 'calib', 'image_shape']: 21 | continue 22 | batch_dict[key] = torch.from_numpy(val).float().cuda() 23 | 24 | 25 | def model_fn_decorator(): 26 | ModelReturn = namedtuple('ModelReturn', ['loss', 'tb_dict', 'disp_dict']) 27 | 28 | def model_func(model, batch_dict): 29 | load_data_to_gpu(batch_dict) 30 | ret_dict, tb_dict, disp_dict = model(batch_dict) 31 | 32 | loss = ret_dict['loss'].mean() 33 | if hasattr(model, 'update_global_step'): 34 | model.update_global_step() 35 | else: 36 | model.module.update_global_step() 37 | 38 | return ModelReturn(loss, tb_dict, disp_dict) 39 | 40 | return model_func 41 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_batch/src/interpolate_gpu.h: -------------------------------------------------------------------------------- 1 | #ifndef _INTERPOLATE_GPU_H 2 | #define _INTERPOLATE_GPU_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | 9 | 10 | void three_nn_wrapper_fast(int b, int n, int m, at::Tensor unknown_tensor, 11 | at::Tensor known_tensor, at::Tensor dist2_tensor, at::Tensor idx_tensor); 12 | 13 | void three_nn_kernel_launcher_fast(int b, int n, int m, const float *unknown, 14 | const float *known, float *dist2, int *idx); 15 | 16 | 17 | void three_interpolate_wrapper_fast(int b, int c, int m, int n, at::Tensor points_tensor, 18 | at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor out_tensor); 19 | 20 | void three_interpolate_kernel_launcher_fast(int b, int c, int m, int n, 21 | const float *points, const int *idx, const float *weight, float *out); 22 | 23 | 24 | void three_interpolate_grad_wrapper_fast(int b, int c, int n, int m, at::Tensor grad_out_tensor, 25 | at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor grad_points_tensor); 26 | 27 | void three_interpolate_grad_kernel_launcher_fast(int b, int c, int n, int m, const float *grad_out, 28 | const int *idx, const float *weight, float *grad_points); 29 | 30 | #endif 31 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_batch/src/pointnet2_api.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "ball_query_gpu.h" 5 | #include "group_points_gpu.h" 6 | #include "sampling_gpu.h" 7 | #include "interpolate_gpu.h" 8 | 9 | 10 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 11 | m.def("ball_query_wrapper", &ball_query_wrapper_fast, "ball_query_wrapper_fast"); 12 | 13 | m.def("group_points_wrapper", &group_points_wrapper_fast, "group_points_wrapper_fast"); 14 | m.def("group_points_grad_wrapper", &group_points_grad_wrapper_fast, "group_points_grad_wrapper_fast"); 15 | 16 | m.def("gather_points_wrapper", &gather_points_wrapper_fast, "gather_points_wrapper_fast"); 17 | m.def("gather_points_grad_wrapper", &gather_points_grad_wrapper_fast, "gather_points_grad_wrapper_fast"); 18 | 19 | m.def("furthest_point_sampling_wrapper", &furthest_point_sampling_wrapper, "furthest_point_sampling_wrapper"); 20 | 21 | m.def("three_nn_wrapper", &three_nn_wrapper_fast, "three_nn_wrapper_fast"); 22 | m.def("three_interpolate_wrapper", &three_interpolate_wrapper_fast, "three_interpolate_wrapper_fast"); 23 | m.def("three_interpolate_grad_wrapper", &three_interpolate_grad_wrapper_fast, "three_interpolate_grad_wrapper_fast"); 24 | } 25 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/utils.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | #include 8 | #include 9 | 10 | #define CHECK_CUDA(x) \ 11 | do { \ 12 | TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor"); \ 13 | } while (0) 14 | 15 | #define CHECK_CONTIGUOUS(x) \ 16 | do { \ 17 | TORCH_CHECK(x.is_contiguous(), #x " must be a contiguous tensor"); \ 18 | } while (0) 19 | 20 | #define CHECK_IS_INT(x) \ 21 | do { \ 22 | TORCH_CHECK(x.scalar_type() == at::ScalarType::Int, \ 23 | #x " must be an int tensor"); \ 24 | } while (0) 25 | 26 | #define CHECK_IS_FLOAT(x) \ 27 | do { \ 28 | TORCH_CHECK(x.scalar_type() == at::ScalarType::Float, \ 29 | #x " must be a float tensor"); \ 30 | } while (0) 31 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/include/utils.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | #include 8 | #include 9 | 10 | #define CHECK_CUDA(x) \ 11 | do { \ 12 | TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor"); \ 13 | } while (0) 14 | 15 | #define CHECK_CONTIGUOUS(x) \ 16 | do { \ 17 | TORCH_CHECK(x.is_contiguous(), #x " must be a contiguous tensor"); \ 18 | } while (0) 19 | 20 | #define CHECK_IS_INT(x) \ 21 | do { \ 22 | TORCH_CHECK(x.scalar_type() == at::ScalarType::Int, \ 23 | #x " must be an int tensor"); \ 24 | } while (0) 25 | 26 | #define CHECK_IS_FLOAT(x) \ 27 | do { \ 28 | TORCH_CHECK(x.scalar_type() == at::ScalarType::Float, \ 29 | #x " must be a float tensor"); \ 30 | } while (0) 31 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_stack/src/group_points_gpu.h: -------------------------------------------------------------------------------- 1 | /* 2 | Stacked-batch-data version of point grouping, modified from the original implementation of official PointNet++ codes. 3 | Written by Shaoshuai Shi 4 | All Rights Reserved 2019-2020. 5 | */ 6 | 7 | 8 | #ifndef _STACK_GROUP_POINTS_GPU_H 9 | #define _STACK_GROUP_POINTS_GPU_H 10 | 11 | #include 12 | #include 13 | #include 14 | #include 15 | 16 | 17 | int group_points_wrapper_stack(int B, int M, int C, int nsample, 18 | at::Tensor features_tensor, at::Tensor features_batch_cnt_tensor, 19 | at::Tensor idx_tensor, at::Tensor idx_batch_cnt_tensor, at::Tensor out_tensor); 20 | 21 | void group_points_kernel_launcher_stack(int B, int M, int C, int nsample, 22 | const float *features, const int *features_batch_cnt, const int *idx, const int *idx_batch_cnt, float *out); 23 | 24 | int group_points_grad_wrapper_stack(int B, int M, int C, int N, int nsample, 25 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor idx_batch_cnt_tensor, 26 | at::Tensor features_batch_cnt_tensor, at::Tensor grad_features_tensor); 27 | 28 | void group_points_grad_kernel_launcher_stack(int B, int M, int C, int N, int nsample, 29 | const float *grad_out, const int *idx, const int *idx_batch_cnt, const int *features_batch_cnt, float *grad_features); 30 | 31 | #endif 32 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/ball_query.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include "ball_query.h" 7 | #include "utils.h" 8 | 9 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 10 | int nsample, const float *new_xyz, 11 | const float *xyz, int *idx); 12 | 13 | at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, 14 | const int nsample) { 15 | CHECK_CONTIGUOUS(new_xyz); 16 | CHECK_CONTIGUOUS(xyz); 17 | CHECK_IS_FLOAT(new_xyz); 18 | CHECK_IS_FLOAT(xyz); 19 | 20 | if (new_xyz.type().is_cuda()) { 21 | CHECK_CUDA(xyz); 22 | } 23 | 24 | at::Tensor idx = 25 | torch::zeros({new_xyz.size(0), new_xyz.size(1), nsample}, 26 | at::device(new_xyz.device()).dtype(at::ScalarType::Int)); 27 | 28 | if (new_xyz.type().is_cuda()) { 29 | query_ball_point_kernel_wrapper(xyz.size(0), xyz.size(1), new_xyz.size(1), 30 | radius, nsample, new_xyz.data(), 31 | xyz.data(), idx.data()); 32 | } else { 33 | TORCH_CHECK(false, "CPU not supported"); 34 | } 35 | 36 | return idx; 37 | } 38 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_stack/src/interpolate_gpu.h: -------------------------------------------------------------------------------- 1 | #ifndef _INTERPOLATE_GPU_H 2 | #define _INTERPOLATE_GPU_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | 9 | 10 | void three_nn_wrapper_stack(at::Tensor unknown_tensor, 11 | at::Tensor unknown_batch_cnt_tensor, at::Tensor known_tensor, 12 | at::Tensor known_batch_cnt_tensor, at::Tensor dist2_tensor, at::Tensor idx_tensor); 13 | 14 | 15 | void three_interpolate_wrapper_stack(at::Tensor features_tensor, 16 | at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor out_tensor); 17 | 18 | 19 | 20 | void three_interpolate_grad_wrapper_stack(at::Tensor grad_out_tensor, at::Tensor idx_tensor, 21 | at::Tensor weight_tensor, at::Tensor grad_features_tensor); 22 | 23 | 24 | void three_nn_kernel_launcher_stack(int batch_size, int N, int M, const float *unknown, 25 | const int *unknown_batch_cnt, const float *known, const int *known_batch_cnt, 26 | float *dist2, int *idx); 27 | 28 | 29 | void three_interpolate_kernel_launcher_stack(int N, int channels, 30 | const float *features, const int *idx, const float *weight, float *out); 31 | 32 | 33 | 34 | void three_interpolate_grad_kernel_launcher_stack(int N, int channels, const float *grad_out, 35 | const int *idx, const float *weight, float *grad_features); 36 | 37 | 38 | 39 | #endif -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_batch/src/group_points.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | batch version of point grouping, modified from the original implementation of official PointNet++ codes. 3 | Written by Shaoshuai Shi 4 | All Rights Reserved 2018. 5 | */ 6 | 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include "group_points_gpu.h" 14 | 15 | extern THCState *state; 16 | 17 | 18 | int group_points_grad_wrapper_fast(int b, int c, int n, int npoints, int nsample, 19 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor) { 20 | 21 | float *grad_points = grad_points_tensor.data(); 22 | const int *idx = idx_tensor.data(); 23 | const float *grad_out = grad_out_tensor.data(); 24 | 25 | group_points_grad_kernel_launcher_fast(b, c, n, npoints, nsample, grad_out, idx, grad_points); 26 | return 1; 27 | } 28 | 29 | 30 | int group_points_wrapper_fast(int b, int c, int n, int npoints, int nsample, 31 | at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor) { 32 | 33 | const float *points = points_tensor.data(); 34 | const int *idx = idx_tensor.data(); 35 | float *out = out_tensor.data(); 36 | 37 | group_points_kernel_launcher_fast(b, c, n, npoints, nsample, points, idx, out); 38 | return 1; 39 | } 40 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_batch/src/ball_query.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | batch version of ball query, modified from the original implementation of official PointNet++ codes. 3 | Written by Shaoshuai Shi 4 | All Rights Reserved 2018. 5 | */ 6 | 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include "ball_query_gpu.h" 14 | 15 | extern THCState *state; 16 | 17 | #define CHECK_CUDA(x) do { \ 18 | if (!x.type().is_cuda()) { \ 19 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 20 | exit(-1); \ 21 | } \ 22 | } while (0) 23 | #define CHECK_CONTIGUOUS(x) do { \ 24 | if (!x.is_contiguous()) { \ 25 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 26 | exit(-1); \ 27 | } \ 28 | } while (0) 29 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x) 30 | 31 | 32 | int ball_query_wrapper_fast(int b, int n, int m, float radius, int nsample, 33 | at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor idx_tensor) { 34 | CHECK_INPUT(new_xyz_tensor); 35 | CHECK_INPUT(xyz_tensor); 36 | const float *new_xyz = new_xyz_tensor.data(); 37 | const float *xyz = xyz_tensor.data(); 38 | int *idx = idx_tensor.data(); 39 | 40 | ball_query_kernel_launcher_fast(b, n, m, radius, nsample, new_xyz, xyz, idx); 41 | return 1; 42 | } 43 | -------------------------------------------------------------------------------- /scannet/data_viz.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import sys 7 | import os 8 | 9 | BASE_DIR = os.path.dirname(__file__) 10 | sys.path.append(BASE_DIR) 11 | 12 | import numpy as np 13 | import pc_util 14 | 15 | scene_name = 'scannet_train_detection_data/scene0002_00' 16 | output_folder = 'data_viz_dump' 17 | 18 | data = np.load(scene_name+'_vert.npy') 19 | scene_points = data[:,0:3] 20 | colors = data[:,3:] 21 | instance_labels = np.load(scene_name+'_ins_label.npy') 22 | semantic_labels = np.load(scene_name+'_sem_label.npy') 23 | instance_bboxes = np.load(scene_name+'_bbox.npy') 24 | 25 | print(np.unique(instance_labels)) 26 | print(np.unique(semantic_labels)) 27 | input() 28 | if not os.path.exists(output_folder): 29 | os.mkdir(output_folder) 30 | 31 | # Write scene as OBJ file for visualization 32 | pc_util.write_ply_rgb(scene_points, colors, os.path.join(output_folder, 'scene.obj')) 33 | pc_util.write_ply_color(scene_points, instance_labels, os.path.join(output_folder, 'scene_instance.obj')) 34 | pc_util.write_ply_color(scene_points, semantic_labels, os.path.join(output_folder, 'scene_semantic.obj')) 35 | 36 | from model_util_scannet import ScannetDatasetConfig 37 | DC = ScannetDatasetConfig() 38 | print(instance_bboxes.shape) 39 | bbox, bbox_color = DC.param2bbox(instance_bboxes, True) 40 | pc_util.write_ply_color(bbox, bbox_color, os.path.join(output_folder, 'scene_bbox.obj')) 41 | -------------------------------------------------------------------------------- /sunrgbd/matlab/extract_split.m: -------------------------------------------------------------------------------- 1 | % Copyright (c) Facebook, Inc. and its affiliates. 2 | % 3 | % This source code is licensed under the MIT license found in the 4 | % LICENSE file in the root directory of this source tree. 5 | 6 | %% Dump train/val split. 7 | % Author: Charles R. Qi 8 | 9 | addpath('../OFFICIAL_SUNRGBD/SUNRGBDtoolbox') 10 | 11 | %% Construct Hash Map 12 | hash_train = java.util.Hashtable; 13 | hash_val = java.util.Hashtable; 14 | 15 | split = load('../OFFICIAL_SUNRGBD/SUNRGBDtoolbox/traintestSUNRGBD/allsplit.mat'); 16 | 17 | N_train = length(split.alltrain); 18 | N_val = length(split.alltest); 19 | 20 | for i = 1:N_train 21 | folder_path = split.alltrain{i}; 22 | folder_path(1:16) = ''; 23 | hash_train.put(folder_path,0); 24 | end 25 | for i = 1:N_val 26 | folder_path = split.alltest{i}; 27 | folder_path(1:16) = ''; 28 | hash_val.put(folder_path,0); 29 | end 30 | 31 | %% Map data to train or val set. 32 | load('../OFFICIAL_SUNRGBD/SUNRGBDMeta3DBB_v2.mat'); 33 | 34 | fid_train = fopen('../sunrgbd_trainval/train_data_idx.txt', 'w'); 35 | fid_val = fopen('../sunrgbd_trainval/val_data_idx.txt', 'w'); 36 | 37 | for imageId = 1:10335 38 | data = SUNRGBDMeta(imageId); 39 | depthpath = data.depthpath; 40 | depthpath(1:16) = ''; 41 | [filepath,name,ext] = fileparts(depthpath); 42 | [filepath,name,ext] = fileparts(filepath); 43 | if hash_train.containsKey(filepath) 44 | fprintf(fid_train, '%d\n', imageId); 45 | elseif hash_val.containsKey(filepath) 46 | fprintf(fid_val, '%d\n', imageId); 47 | else 48 | a = 1; 49 | end 50 | end 51 | fclose(fid_train); 52 | fclose(fid_val); 53 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/cuda_utils.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #ifndef _CUDA_UTILS_H 7 | #define _CUDA_UTILS_H 8 | 9 | #include 10 | #include 11 | #include 12 | 13 | #include 14 | #include 15 | 16 | #include 17 | 18 | #define TOTAL_THREADS 512 19 | 20 | inline int opt_n_threads(int work_size) { 21 | const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0); 22 | 23 | return max(min(1 << pow_2, TOTAL_THREADS), 1); 24 | } 25 | 26 | inline dim3 opt_block_config(int x, int y) { 27 | const int x_threads = opt_n_threads(x); 28 | const int y_threads = 29 | max(min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1); 30 | dim3 block_config(x_threads, y_threads, 1); 31 | 32 | return block_config; 33 | } 34 | 35 | #define CUDA_CHECK_ERRORS() \ 36 | do { \ 37 | cudaError_t err = cudaGetLastError(); \ 38 | if (cudaSuccess != err) { \ 39 | fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n", \ 40 | cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, \ 41 | __FILE__); \ 42 | exit(-1); \ 43 | } \ 44 | } while (0) 45 | 46 | #endif 47 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/include/cuda_utils.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #ifndef _CUDA_UTILS_H 7 | #define _CUDA_UTILS_H 8 | 9 | #include 10 | #include 11 | #include 12 | 13 | #include 14 | #include 15 | 16 | #include 17 | 18 | #define TOTAL_THREADS 512 19 | 20 | inline int opt_n_threads(int work_size) { 21 | const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0); 22 | 23 | return max(min(1 << pow_2, TOTAL_THREADS), 1); 24 | } 25 | 26 | inline dim3 opt_block_config(int x, int y) { 27 | const int x_threads = opt_n_threads(x); 28 | const int y_threads = 29 | max(min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1); 30 | dim3 block_config(x_threads, y_threads, 1); 31 | 32 | return block_config; 33 | } 34 | 35 | #define CUDA_CHECK_ERRORS() \ 36 | do { \ 37 | cudaError_t err = cudaGetLastError(); \ 38 | if (cudaSuccess != err) { \ 39 | fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n", \ 40 | cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, \ 41 | __FILE__); \ 42 | exit(-1); \ 43 | } \ 44 | } while (0) 45 | 46 | #endif 47 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/backbones_2d/map_to_bev/pointpillar_scatter.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | class PointPillarScatter(nn.Module): 6 | def __init__(self, model_cfg, grid_size, **kwargs): 7 | super().__init__() 8 | 9 | self.model_cfg = model_cfg 10 | self.num_bev_features = self.model_cfg.NUM_BEV_FEATURES 11 | self.nx, self.ny, self.nz = grid_size 12 | assert self.nz == 1 13 | 14 | def forward(self, batch_dict, **kwargs): 15 | pillar_features, coords = batch_dict['pillar_features'], batch_dict['voxel_coords'] 16 | batch_spatial_features = [] 17 | batch_size = coords[:, 0].max().int().item() + 1 18 | for batch_idx in range(batch_size): 19 | spatial_feature = torch.zeros( 20 | self.num_bev_features, 21 | self.nz * self.nx * self.ny, 22 | dtype=pillar_features.dtype, 23 | device=pillar_features.device) 24 | 25 | batch_mask = coords[:, 0] == batch_idx 26 | this_coords = coords[batch_mask, :] 27 | indices = this_coords[:, 1] + this_coords[:, 2] * self.nx + this_coords[:, 3] 28 | indices = indices.type(torch.long) 29 | pillars = pillar_features[batch_mask, :] 30 | pillars = pillars.t() 31 | spatial_feature[:, indices] = pillars 32 | batch_spatial_features.append(spatial_feature) 33 | 34 | batch_spatial_features = torch.stack(batch_spatial_features, 0) 35 | batch_spatial_features = batch_spatial_features.view(batch_size, self.num_bev_features * self.nz, self.ny, self.nx) 36 | batch_dict['spatial_features'] = batch_spatial_features 37 | return batch_dict 38 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_batch/src/sampling.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | batch version of point sampling and gathering, modified from the original implementation of official PointNet++ codes. 3 | Written by Shaoshuai Shi 4 | All Rights Reserved 2018. 5 | */ 6 | 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | 13 | #include "sampling_gpu.h" 14 | 15 | extern THCState *state; 16 | 17 | 18 | int gather_points_wrapper_fast(int b, int c, int n, int npoints, 19 | at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor){ 20 | const float *points = points_tensor.data(); 21 | const int *idx = idx_tensor.data(); 22 | float *out = out_tensor.data(); 23 | 24 | gather_points_kernel_launcher_fast(b, c, n, npoints, points, idx, out); 25 | return 1; 26 | } 27 | 28 | 29 | int gather_points_grad_wrapper_fast(int b, int c, int n, int npoints, 30 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor) { 31 | 32 | const float *grad_out = grad_out_tensor.data(); 33 | const int *idx = idx_tensor.data(); 34 | float *grad_points = grad_points_tensor.data(); 35 | 36 | gather_points_grad_kernel_launcher_fast(b, c, n, npoints, grad_out, idx, grad_points); 37 | return 1; 38 | } 39 | 40 | 41 | int furthest_point_sampling_wrapper(int b, int n, int m, 42 | at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor idx_tensor) { 43 | 44 | const float *points = points_tensor.data(); 45 | float *temp = temp_tensor.data(); 46 | int *idx = idx_tensor.data(); 47 | 48 | furthest_point_sampling_kernel_launcher(b, n, m, points, temp, idx); 49 | return 1; 50 | } 51 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_stack/src/ball_query.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | Stacked-batch-data version of ball query, modified from the original implementation of official PointNet++ codes. 3 | Written by Shaoshuai Shi 4 | All Rights Reserved 2019-2020. 5 | */ 6 | 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include "ball_query_gpu.h" 14 | 15 | extern THCState *state; 16 | 17 | #define CHECK_CUDA(x) do { \ 18 | if (!x.type().is_cuda()) { \ 19 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 20 | exit(-1); \ 21 | } \ 22 | } while (0) 23 | #define CHECK_CONTIGUOUS(x) do { \ 24 | if (!x.is_contiguous()) { \ 25 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 26 | exit(-1); \ 27 | } \ 28 | } while (0) 29 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x) 30 | 31 | int ball_query_wrapper_stack(int B, int M, float radius, int nsample, 32 | at::Tensor new_xyz_tensor, at::Tensor new_xyz_batch_cnt_tensor, 33 | at::Tensor xyz_tensor, at::Tensor xyz_batch_cnt_tensor, at::Tensor idx_tensor) { 34 | CHECK_INPUT(new_xyz_tensor); 35 | CHECK_INPUT(xyz_tensor); 36 | CHECK_INPUT(new_xyz_batch_cnt_tensor); 37 | CHECK_INPUT(xyz_batch_cnt_tensor); 38 | 39 | const float *new_xyz = new_xyz_tensor.data(); 40 | const float *xyz = xyz_tensor.data(); 41 | const int *new_xyz_batch_cnt = new_xyz_batch_cnt_tensor.data(); 42 | const int *xyz_batch_cnt = xyz_batch_cnt_tensor.data(); 43 | int *idx = idx_tensor.data(); 44 | 45 | ball_query_kernel_launcher_stack(B, M, radius, nsample, new_xyz, new_xyz_batch_cnt, xyz, xyz_batch_cnt, idx); 46 | return 1; 47 | } 48 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/datasets/kitti/kitti_object_eval_python/README.md: -------------------------------------------------------------------------------- 1 | # kitti-object-eval-python 2 | **Note**: This is borrowed from [traveller59/kitti-object-eval-python](https://github.com/traveller59/kitti-object-eval-python) 3 | 4 | Fast kitti object detection eval in python(finish eval in less than 10 second), support 2d/bev/3d/aos. , support coco-style AP. If you use command line interface, numba need some time to compile jit functions. 5 | ## Dependencies 6 | Only support python 3.6+, need `numpy`, `skimage`, `numba`, `fire`. If you have Anaconda, just install `cudatoolkit` in anaconda. Otherwise, please reference to this [page](https://github.com/numba/numba#custom-python-environments) to set up llvm and cuda for numba. 7 | * Install by conda: 8 | ``` 9 | conda install -c numba cudatoolkit=x.x (8.0, 9.0, 9.1, depend on your environment) 10 | ``` 11 | ## Usage 12 | * commandline interface: 13 | ``` 14 | python evaluate.py evaluate --label_path=/path/to/your_gt_label_folder --result_path=/path/to/your_result_folder --label_split_file=/path/to/val.txt --current_class=0 --coco=False 15 | ``` 16 | * python interface: 17 | ```Python 18 | import kitti_common as kitti 19 | from eval import get_official_eval_result, get_coco_eval_result 20 | def _read_imageset_file(path): 21 | with open(path, 'r') as f: 22 | lines = f.readlines() 23 | return [int(line) for line in lines] 24 | det_path = "/path/to/your_result_folder" 25 | dt_annos = kitti.get_label_annos(det_path) 26 | gt_path = "/path/to/your_gt_label_folder" 27 | gt_split_file = "/path/to/val.txt" # from https://xiaozhichen.github.io/files/mv3d/imagesets.tar.gz 28 | val_image_ids = _read_imageset_file(gt_split_file) 29 | gt_annos = kitti.get_label_annos(gt_path, val_image_ids) 30 | print(get_official_eval_result(gt_annos, dt_annos, 0)) # 6s in my computer 31 | print(get_coco_eval_result(gt_annos, dt_annos, 0)) # 18s in my computer 32 | ``` 33 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/datasets/processor/point_feature_encoder.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | class PointFeatureEncoder(object): 5 | def __init__(self, config, point_cloud_range=None): 6 | super().__init__() 7 | self.point_encoding_config = config 8 | assert list(self.point_encoding_config.src_feature_list[0:3]) == ['x', 'y', 'z'] 9 | self.used_feature_list = self.point_encoding_config.used_feature_list 10 | self.src_feature_list = self.point_encoding_config.src_feature_list 11 | self.point_cloud_range = point_cloud_range 12 | 13 | @property 14 | def num_point_features(self): 15 | return getattr(self, self.point_encoding_config.encoding_type)(points=None) 16 | 17 | def forward(self, data_dict): 18 | """ 19 | Args: 20 | data_dict: 21 | points: (N, 3 + C_in) 22 | ... 23 | Returns: 24 | data_dict: 25 | points: (N, 3 + C_out), 26 | use_lead_xyz: whether to use xyz as point-wise features 27 | ... 28 | """ 29 | data_dict['points'], use_lead_xyz = getattr(self, self.point_encoding_config.encoding_type)( 30 | data_dict['points'] 31 | ) 32 | data_dict['use_lead_xyz'] = use_lead_xyz 33 | return data_dict 34 | 35 | def absolute_coordinates_encoding(self, points=None): 36 | if points is None: 37 | num_output_features = len(self.used_feature_list) 38 | return num_output_features 39 | 40 | point_feature_list = [points[:, 0:3]] 41 | for x in self.used_feature_list: 42 | if x in ['x', 'y', 'z']: 43 | continue 44 | idx = self.src_feature_list.index(x) 45 | point_feature_list.append(points[:, idx:idx+1]) 46 | point_features = np.concatenate(point_feature_list, axis=1) 47 | return point_features, True 48 | -------------------------------------------------------------------------------- /OpenPCDet/tools/cfgs/dataset_configs/kitti_dataset.yaml: -------------------------------------------------------------------------------- 1 | DATASET: 'KittiDataset' 2 | DATA_PATH: '../data/kitti' 3 | 4 | POINT_CLOUD_RANGE: [0, -40, -3, 70.4, 40, 1] 5 | 6 | DATA_SPLIT: { 7 | 'train': train, 8 | 'test': val 9 | } 10 | 11 | INFO_PATH: { 12 | 'train': [kitti_infos_train.pkl], 13 | 'test': [kitti_infos_val.pkl], 14 | } 15 | 16 | FOV_POINTS_ONLY: True 17 | 18 | 19 | DATA_AUGMENTOR: 20 | DISABLE_AUG_LIST: ['placeholder'] 21 | AUG_CONFIG_LIST: 22 | - NAME: gt_sampling 23 | USE_ROAD_PLANE: True 24 | DB_INFO_PATH: 25 | - kitti_dbinfos_train.pkl 26 | PREPARE: { 27 | filter_by_min_points: ['Car:5', 'Pedestrian:5', 'Cyclist:5'], 28 | filter_by_difficulty: [-1], 29 | } 30 | 31 | SAMPLE_GROUPS: ['Car:20','Pedestrian:15', 'Cyclist:15'] 32 | NUM_POINT_FEATURES: 4 33 | DATABASE_WITH_FAKELIDAR: False 34 | REMOVE_EXTRA_WIDTH: [0.0, 0.0, 0.0] 35 | LIMIT_WHOLE_SCENE: True 36 | 37 | - NAME: random_world_flip 38 | ALONG_AXIS_LIST: ['x'] 39 | 40 | - NAME: random_world_rotation 41 | WORLD_ROT_ANGLE: [-0.78539816, 0.78539816] 42 | 43 | - NAME: random_world_scaling 44 | WORLD_SCALE_RANGE: [0.95, 1.05] 45 | 46 | 47 | POINT_FEATURE_ENCODING: { 48 | encoding_type: absolute_coordinates_encoding, 49 | used_feature_list: ['x', 'y', 'z', 'intensity'], 50 | src_feature_list: ['x', 'y', 'z', 'intensity'], 51 | } 52 | 53 | 54 | DATA_PROCESSOR: 55 | - NAME: mask_points_and_boxes_outside_range 56 | REMOVE_OUTSIDE_BOXES: True 57 | 58 | - NAME: shuffle_points 59 | SHUFFLE_ENABLED: { 60 | 'train': True, 61 | 'test': False 62 | } 63 | 64 | - NAME: transform_points_to_voxels 65 | VOXEL_SIZE: [0.05, 0.05, 0.1] 66 | MAX_POINTS_PER_VOXEL: 5 67 | MAX_NUMBER_OF_VOXELS: { 68 | 'train': 16000, 69 | 'test': 40000 70 | } 71 | -------------------------------------------------------------------------------- /utils/tf_visualizer.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | '''Code adapted from https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix''' 7 | import os 8 | import time 9 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 10 | import sys 11 | sys.path.append(BASE_DIR) 12 | import tf_logger 13 | 14 | 15 | class Visualizer(): 16 | def __init__(self, log_dir, name='train'): 17 | # self.opt = opt 18 | #self.logger = tf_logger.Logger(os.path.join(opt.logging_dir, opt.name)) 19 | #self.log_name = os.path.join(opt.checkpoint_dir, opt.name, 'loss_log.txt') 20 | self.logger = tf_logger.Logger(os.path.join(log_dir, name)) 21 | self.log_name = os.path.join(log_dir, 'tf_visualizer_log.txt') 22 | with open(self.log_name, "a") as log_file: 23 | now = time.strftime("%c") 24 | log_file.write('================ Training Loss (%s) ================\n' % now) 25 | 26 | # |visuals|: dictionary of images to save 27 | def log_images(self, visuals, step): 28 | for label, image_numpy in visuals.items(): 29 | self.logger.image_summary( 30 | label, [image_numpy], step) 31 | 32 | # scalars: dictionary of scalar labels and values 33 | def log_scalars(self, scalars, step): 34 | for label, val in scalars.items(): 35 | self.logger.scalar_summary(label, val, step) 36 | 37 | # scatter plots 38 | def plot_current_points(self, points, disp_offset=10): 39 | pass 40 | 41 | # scalars: same format as |scalars| of plot_current_scalars 42 | def print_current_scalars(self, epoch, i, scalars): 43 | message = '(epoch: %d, iters: %d) ' % (epoch, i) 44 | for k, v in scalars.items(): 45 | message += '%s: %.3f ' % (k, v) 46 | 47 | print(message) 48 | with open(self.log_name, "a") as log_file: 49 | log_file.write('%s\n' % message) 50 | -------------------------------------------------------------------------------- /sunrgbd/matlab/extract_rgbd_data_v1.m: -------------------------------------------------------------------------------- 1 | % Copyright (c) Facebook, Inc. and its affiliates. 2 | % 3 | % This source code is licensed under the MIT license found in the 4 | % LICENSE file in the root directory of this source tree. 5 | 6 | %% Dump SUNRGBD data to our format 7 | % for each sample, we have RGB image, 2d boxes. 8 | % point cloud (in camera coordinate), calibration and 3d boxes. 9 | % 10 | % Extract using V1 labels. 11 | % 12 | % Author: Charles R. Qi 13 | % 14 | clear; close all; clc; 15 | addpath(genpath('.')) 16 | addpath('../OFFICIAL_SUNRGBD/SUNRGBDtoolbox') 17 | %% V1 2D&3D BB and Seg masks 18 | load('../OFFICIAL_SUNRGBD/SUNRGBDtoolbox/Metadata/SUNRGBDMeta.mat') 19 | % load('./Metadata/SUNRGBD2Dseg.mat') 20 | 21 | %% Create folders 22 | det_label_folder = '../sunrgbd_trainval/label_v1/'; 23 | mkdir(det_label_folder); 24 | %% Read 25 | for imageId = 1:10335 26 | imageId 27 | try 28 | data = SUNRGBDMeta(imageId); 29 | data.depthpath(1:16) = ''; 30 | data.depthpath = strcat('../OFFICIAL_SUNRGBD/SUNRGBD', data.depthpath); 31 | data.rgbpath(1:16) = ''; 32 | data.rgbpath = strcat('../OFFICIAL_SUNRGBD/SUNRGBD', data.rgbpath); 33 | 34 | % MAT files are 3x smaller than TXT files. In Python we can use 35 | % scipy.io.loadmat('xxx.mat')['points3d_rgb'] to load the data. 36 | mat_filename = strcat(num2str(imageId,'%06d'), '.mat'); 37 | txt_filename = strcat(num2str(imageId,'%06d'), '.txt'); 38 | 39 | % Write 2D and 3D box label 40 | data2d = data; 41 | fid = fopen(strcat(det_label_folder, txt_filename), 'w'); 42 | for j = 1:length(data.groundtruth3DBB) 43 | centroid = data.groundtruth3DBB(j).centroid; 44 | classname = data.groundtruth3DBB(j).classname; 45 | orientation = data.groundtruth3DBB(j).orientation; 46 | coeffs = abs(data.groundtruth3DBB(j).coeffs); 47 | box2d = data2d.groundtruth2DBB(j).gtBb2D; 48 | fprintf(fid, '%s %d %d %d %d %f %f %f %f %f %f %f %f\n', classname, box2d(1), box2d(2), box2d(3), box2d(4), centroid(1), centroid(2), centroid(3), coeffs(1), coeffs(2), coeffs(3), orientation(1), orientation(2)); 49 | end 50 | fclose(fid); 51 | 52 | catch 53 | end 54 | 55 | end 56 | -------------------------------------------------------------------------------- /OpenPCDet/setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | import subprocess 3 | 4 | from setuptools import find_packages, setup 5 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 6 | 7 | 8 | def get_git_commit_number(): 9 | if not os.path.exists('.git'): 10 | return '0000000' 11 | 12 | cmd_out = subprocess.run(['git', 'rev-parse', 'HEAD'], stdout=subprocess.PIPE) 13 | git_commit_number = cmd_out.stdout.decode('utf-8')[:7] 14 | return git_commit_number 15 | 16 | 17 | def make_cuda_ext(name, module, sources): 18 | cuda_ext = CUDAExtension( 19 | name='%s.%s' % (module, name), 20 | sources=[os.path.join(*module.split('.'), src) for src in sources] 21 | ) 22 | return cuda_ext 23 | 24 | 25 | def write_version_to_file(version, target_file): 26 | with open(target_file, 'w') as f: 27 | print('__version__ = "%s"' % version, file=f) 28 | 29 | # Deleted most of the ext_modules for faster installation 30 | if __name__ == '__main__': 31 | version = '0.3.0+%s' % get_git_commit_number() 32 | write_version_to_file(version, 'pcdet/version.py') 33 | 34 | setup( 35 | name='pcdet', 36 | version=version, 37 | description='OpenPCDet is a general codebase for 3D object detection from point cloud', 38 | install_requires=[ 39 | 'numpy', 40 | 'torch>=1.1', 41 | 'tensorboardX', 42 | 'easydict', 43 | 'pyyaml' 44 | ], 45 | author='Shaoshuai Shi', 46 | author_email='shaoshuaics@gmail.com', 47 | license='Apache License 2.0', 48 | packages=find_packages(exclude=['tools', 'data', 'output']), 49 | cmdclass={'build_ext': BuildExtension}, 50 | ext_modules=[ 51 | make_cuda_ext( 52 | name='iou3d_nms_cuda', 53 | module='pcdet.ops.iou3d_nms', 54 | sources=[ 55 | 'src/iou3d_cpu.cpp', 56 | 'src/iou3d_nms_api.cpp', 57 | 'src/iou3d_nms.cpp', 58 | 'src/iou3d_nms_kernel.cu', 59 | ] 60 | ), 61 | ], 62 | ) 63 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/ball_query_gpu.cu: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include 7 | #include 8 | #include 9 | 10 | #include "cuda_utils.h" 11 | 12 | // input: new_xyz(b, m, 3) xyz(b, n, 3) 13 | // output: idx(b, m, nsample) 14 | __global__ void query_ball_point_kernel(int b, int n, int m, float radius, 15 | int nsample, 16 | const float *__restrict__ new_xyz, 17 | const float *__restrict__ xyz, 18 | int *__restrict__ idx) { 19 | int batch_index = blockIdx.x; 20 | xyz += batch_index * n * 3; 21 | new_xyz += batch_index * m * 3; 22 | idx += m * nsample * batch_index; 23 | 24 | int index = threadIdx.x; 25 | int stride = blockDim.x; 26 | 27 | float radius2 = radius * radius; 28 | for (int j = index; j < m; j += stride) { 29 | float new_x = new_xyz[j * 3 + 0]; 30 | float new_y = new_xyz[j * 3 + 1]; 31 | float new_z = new_xyz[j * 3 + 2]; 32 | for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) { 33 | float x = xyz[k * 3 + 0]; 34 | float y = xyz[k * 3 + 1]; 35 | float z = xyz[k * 3 + 2]; 36 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + 37 | (new_z - z) * (new_z - z); 38 | if (d2 < radius2) { 39 | if (cnt == 0) { 40 | for (int l = 0; l < nsample; ++l) { 41 | idx[j * nsample + l] = k; 42 | } 43 | } 44 | idx[j * nsample + cnt] = k; 45 | ++cnt; 46 | } 47 | } 48 | } 49 | } 50 | 51 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 52 | int nsample, const float *new_xyz, 53 | const float *xyz, int *idx) { 54 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 55 | query_ball_point_kernel<<>>( 56 | b, n, m, radius, nsample, new_xyz, xyz, idx); 57 | 58 | CUDA_CHECK_ERRORS(); 59 | } 60 | -------------------------------------------------------------------------------- /scannet/meta_data/scannetv2_train_0.1.txt: -------------------------------------------------------------------------------- 1 | scene0601_00 2 | scene0443_00 3 | scene0610_01 4 | scene0440_02 5 | scene0324_01 6 | scene0226_01 7 | scene0091_00 8 | scene0214_02 9 | scene0469_02 10 | scene0401_00 11 | scene0292_00 12 | scene0332_01 13 | scene0654_00 14 | scene0096_00 15 | scene0698_00 16 | scene0588_01 17 | scene0070_00 18 | scene0524_01 19 | scene0093_01 20 | scene0589_00 21 | scene0348_01 22 | scene0290_00 23 | scene0373_00 24 | scene0255_02 25 | scene0291_02 26 | scene0315_00 27 | scene0294_02 28 | scene0600_01 29 | scene0132_01 30 | scene0533_01 31 | scene0234_00 32 | scene0248_02 33 | scene0463_01 34 | scene0281_00 35 | scene0541_02 36 | scene0057_01 37 | scene0383_00 38 | scene0078_01 39 | scene0361_00 40 | scene0212_00 41 | scene0073_03 42 | scene0469_01 43 | scene0594_00 44 | scene0407_00 45 | scene0177_01 46 | scene0073_01 47 | scene0136_01 48 | scene0694_01 49 | scene0677_02 50 | scene0478_01 51 | scene0185_00 52 | scene0452_01 53 | scene0561_00 54 | scene0265_00 55 | scene0119_00 56 | scene0291_01 57 | scene0418_02 58 | scene0451_03 59 | scene0439_00 60 | scene0097_00 61 | scene0666_02 62 | scene0358_01 63 | scene0112_02 64 | scene0557_01 65 | scene0399_00 66 | scene0067_02 67 | scene0501_02 68 | scene0210_00 69 | scene0265_02 70 | scene0369_02 71 | scene0122_01 72 | scene0116_00 73 | scene0012_02 74 | scene0630_00 75 | scene0044_01 76 | scene0392_02 77 | scene0587_01 78 | scene0101_00 79 | scene0166_01 80 | scene0418_00 81 | scene0625_01 82 | scene0181_00 83 | scene0301_01 84 | scene0279_02 85 | scene0668_00 86 | scene0082_00 87 | scene0615_00 88 | scene0120_01 89 | scene0225_00 90 | scene0399_01 91 | scene0078_00 92 | scene0269_02 93 | scene0267_00 94 | scene0456_00 95 | scene0440_01 96 | scene0197_02 97 | scene0259_00 98 | scene0451_05 99 | scene0588_03 100 | scene0223_02 101 | scene0656_01 102 | scene0376_01 103 | scene0416_04 104 | scene0472_02 105 | scene0495_00 106 | scene0611_01 107 | scene0451_01 108 | scene0052_01 109 | scene0347_00 110 | scene0065_01 111 | scene0628_00 112 | scene0642_01 113 | scene0038_00 114 | scene0411_02 115 | scene0195_00 116 | scene0009_02 117 | scene0250_02 118 | scene0352_02 119 | scene0313_02 120 | scene0545_00 121 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_batch/src/interpolate.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | batch version of point interpolation, modified from the original implementation of official PointNet++ codes. 3 | Written by Shaoshuai Shi 4 | All Rights Reserved 2018. 5 | */ 6 | 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | #include 16 | #include "interpolate_gpu.h" 17 | 18 | extern THCState *state; 19 | 20 | 21 | void three_nn_wrapper_fast(int b, int n, int m, at::Tensor unknown_tensor, 22 | at::Tensor known_tensor, at::Tensor dist2_tensor, at::Tensor idx_tensor) { 23 | const float *unknown = unknown_tensor.data(); 24 | const float *known = known_tensor.data(); 25 | float *dist2 = dist2_tensor.data(); 26 | int *idx = idx_tensor.data(); 27 | 28 | three_nn_kernel_launcher_fast(b, n, m, unknown, known, dist2, idx); 29 | } 30 | 31 | 32 | void three_interpolate_wrapper_fast(int b, int c, int m, int n, 33 | at::Tensor points_tensor, 34 | at::Tensor idx_tensor, 35 | at::Tensor weight_tensor, 36 | at::Tensor out_tensor) { 37 | 38 | const float *points = points_tensor.data(); 39 | const float *weight = weight_tensor.data(); 40 | float *out = out_tensor.data(); 41 | const int *idx = idx_tensor.data(); 42 | 43 | three_interpolate_kernel_launcher_fast(b, c, m, n, points, idx, weight, out); 44 | } 45 | 46 | void three_interpolate_grad_wrapper_fast(int b, int c, int n, int m, 47 | at::Tensor grad_out_tensor, 48 | at::Tensor idx_tensor, 49 | at::Tensor weight_tensor, 50 | at::Tensor grad_points_tensor) { 51 | 52 | const float *grad_out = grad_out_tensor.data(); 53 | const float *weight = weight_tensor.data(); 54 | float *grad_points = grad_points_tensor.data(); 55 | const int *idx = idx_tensor.data(); 56 | 57 | three_interpolate_grad_kernel_launcher_fast(b, c, n, m, grad_out, idx, weight, grad_points); 58 | } 59 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/roipoint_pool3d/src/roipoint_pool3d.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #define CHECK_CUDA(x) do { \ 5 | if (!x.type().is_cuda()) { \ 6 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 7 | exit(-1); \ 8 | } \ 9 | } while (0) 10 | #define CHECK_CONTIGUOUS(x) do { \ 11 | if (!x.is_contiguous()) { \ 12 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 13 | exit(-1); \ 14 | } \ 15 | } while (0) 16 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x) 17 | 18 | 19 | void roipool3dLauncher(int batch_size, int pts_num, int boxes_num, int feature_in_len, int sampled_pts_num, 20 | const float *xyz, const float *boxes3d, const float *pts_feature, float *pooled_features, int *pooled_empty_flag); 21 | 22 | 23 | int roipool3d_gpu(at::Tensor xyz, at::Tensor boxes3d, at::Tensor pts_feature, at::Tensor pooled_features, at::Tensor pooled_empty_flag){ 24 | // params xyz: (B, N, 3) 25 | // params boxes3d: (B, M, 7) 26 | // params pts_feature: (B, N, C) 27 | // params pooled_features: (B, M, 512, 3+C) 28 | // params pooled_empty_flag: (B, M) 29 | CHECK_INPUT(xyz); 30 | CHECK_INPUT(boxes3d); 31 | CHECK_INPUT(pts_feature); 32 | CHECK_INPUT(pooled_features); 33 | CHECK_INPUT(pooled_empty_flag); 34 | 35 | int batch_size = xyz.size(0); 36 | int pts_num = xyz.size(1); 37 | int boxes_num = boxes3d.size(1); 38 | int feature_in_len = pts_feature.size(2); 39 | int sampled_pts_num = pooled_features.size(2); 40 | 41 | 42 | const float * xyz_data = xyz.data(); 43 | const float * boxes3d_data = boxes3d.data(); 44 | const float * pts_feature_data = pts_feature.data(); 45 | float * pooled_features_data = pooled_features.data(); 46 | int * pooled_empty_flag_data = pooled_empty_flag.data(); 47 | 48 | roipool3dLauncher(batch_size, pts_num, boxes_num, feature_in_len, sampled_pts_num, 49 | xyz_data, boxes3d_data, pts_feature_data, pooled_features_data, pooled_empty_flag_data); 50 | 51 | 52 | 53 | return 1; 54 | } 55 | 56 | 57 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 58 | m.def("forward", &roipool3d_gpu, "roipool3d forward (CUDA)"); 59 | } 60 | 61 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/group_points.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include "group_points.h" 7 | #include "utils.h" 8 | 9 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 10 | const float *points, const int *idx, 11 | float *out); 12 | 13 | void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 14 | int nsample, const float *grad_out, 15 | const int *idx, float *grad_points); 16 | 17 | at::Tensor group_points(at::Tensor points, at::Tensor idx) { 18 | CHECK_CONTIGUOUS(points); 19 | CHECK_CONTIGUOUS(idx); 20 | CHECK_IS_FLOAT(points); 21 | CHECK_IS_INT(idx); 22 | 23 | if (points.type().is_cuda()) { 24 | CHECK_CUDA(idx); 25 | } 26 | 27 | at::Tensor output = 28 | torch::zeros({points.size(0), points.size(1), idx.size(1), idx.size(2)}, 29 | at::device(points.device()).dtype(at::ScalarType::Float)); 30 | 31 | if (points.type().is_cuda()) { 32 | group_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), 33 | idx.size(1), idx.size(2), points.data(), 34 | idx.data(), output.data()); 35 | } else { 36 | TORCH_CHECK(false, "CPU not supported"); 37 | } 38 | 39 | return output; 40 | } 41 | 42 | at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n) { 43 | CHECK_CONTIGUOUS(grad_out); 44 | CHECK_CONTIGUOUS(idx); 45 | CHECK_IS_FLOAT(grad_out); 46 | CHECK_IS_INT(idx); 47 | 48 | if (grad_out.type().is_cuda()) { 49 | CHECK_CUDA(idx); 50 | } 51 | 52 | at::Tensor output = 53 | torch::zeros({grad_out.size(0), grad_out.size(1), n}, 54 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 55 | 56 | if (grad_out.type().is_cuda()) { 57 | group_points_grad_kernel_wrapper( 58 | grad_out.size(0), grad_out.size(1), n, idx.size(1), idx.size(2), 59 | grad_out.data(), idx.data(), output.data()); 60 | } else { 61 | TORCH_CHECK(false, "CPU not supported"); 62 | } 63 | 64 | return output; 65 | } 66 | -------------------------------------------------------------------------------- /OpenPCDet/tools/cfgs/dataset_configs/nuscenes_dataset.yaml: -------------------------------------------------------------------------------- 1 | DATASET: 'NuScenesDataset' 2 | DATA_PATH: '../data/nuscenes' 3 | 4 | VERSION: 'v1.0-trainval' 5 | MAX_SWEEPS: 10 6 | PRED_VELOCITY: True 7 | SET_NAN_VELOCITY_TO_ZEROS: True 8 | FILTER_MIN_POINTS_IN_GT: 1 9 | 10 | DATA_SPLIT: { 11 | 'train': train, 12 | 'test': val 13 | } 14 | 15 | INFO_PATH: { 16 | 'train': [nuscenes_infos_10sweeps_train.pkl], 17 | 'test': [nuscenes_infos_10sweeps_val.pkl], 18 | } 19 | 20 | POINT_CLOUD_RANGE: [-51.2, -51.2, -5.0, 51.2, 51.2, 3.0] 21 | 22 | BALANCED_RESAMPLING: True 23 | 24 | DATA_AUGMENTOR: 25 | DISABLE_AUG_LIST: ['placeholder'] 26 | AUG_CONFIG_LIST: 27 | - NAME: gt_sampling 28 | DB_INFO_PATH: 29 | - nuscenes_dbinfos_10sweeps_withvelo.pkl 30 | PREPARE: { 31 | filter_by_min_points: [ 32 | 'car:5','truck:5', 'construction_vehicle:5', 'bus:5', 'trailer:5', 33 | 'barrier:5', 'motorcycle:5', 'bicycle:5', 'pedestrian:5', 'traffic_cone:5' 34 | ], 35 | } 36 | 37 | SAMPLE_GROUPS: [ 38 | 'car:2','truck:3', 'construction_vehicle:7', 'bus:4', 'trailer:6', 39 | 'barrier:2', 'motorcycle:6', 'bicycle:6', 'pedestrian:2', 'traffic_cone:2' 40 | ] 41 | 42 | NUM_POINT_FEATURES: 5 43 | DATABASE_WITH_FAKELIDAR: False 44 | REMOVE_EXTRA_WIDTH: [0.0, 0.0, 0.0] 45 | LIMIT_WHOLE_SCENE: True 46 | 47 | - NAME: random_world_flip 48 | ALONG_AXIS_LIST: ['x', 'y'] 49 | 50 | - NAME: random_world_rotation 51 | WORLD_ROT_ANGLE: [-0.3925, 0.3925] 52 | 53 | - NAME: random_world_scaling 54 | WORLD_SCALE_RANGE: [0.95, 1.05] 55 | 56 | 57 | POINT_FEATURE_ENCODING: { 58 | encoding_type: absolute_coordinates_encoding, 59 | used_feature_list: ['x', 'y', 'z', 'intensity', 'timestamp'], 60 | src_feature_list: ['x', 'y', 'z', 'intensity', 'timestamp'], 61 | } 62 | 63 | 64 | DATA_PROCESSOR: 65 | - NAME: mask_points_and_boxes_outside_range 66 | REMOVE_OUTSIDE_BOXES: True 67 | 68 | - NAME: shuffle_points 69 | SHUFFLE_ENABLED: { 70 | 'train': True, 71 | 'test': True 72 | } 73 | 74 | - NAME: transform_points_to_voxels 75 | VOXEL_SIZE: [0.1, 0.1, 0.2] 76 | MAX_POINTS_PER_VOXEL: 10 77 | MAX_NUMBER_OF_VOXELS: { 78 | 'train': 60000, 79 | 'test': 60000 80 | } 81 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_batch/src/ball_query_gpu.cu: -------------------------------------------------------------------------------- 1 | /* 2 | batch version of ball query, modified from the original implementation of official PointNet++ codes. 3 | Written by Shaoshuai Shi 4 | All Rights Reserved 2018. 5 | */ 6 | 7 | #include 8 | #include 9 | #include 10 | 11 | #include "ball_query_gpu.h" 12 | #include "cuda_utils.h" 13 | 14 | 15 | __global__ void ball_query_kernel_fast(int b, int n, int m, float radius, int nsample, 16 | const float *__restrict__ new_xyz, const float *__restrict__ xyz, int *__restrict__ idx) { 17 | // new_xyz: (B, M, 3) 18 | // xyz: (B, N, 3) 19 | // output: 20 | // idx: (B, M, nsample) 21 | int bs_idx = blockIdx.y; 22 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 23 | if (bs_idx >= b || pt_idx >= m) return; 24 | 25 | new_xyz += bs_idx * m * 3 + pt_idx * 3; 26 | xyz += bs_idx * n * 3; 27 | idx += bs_idx * m * nsample + pt_idx * nsample; 28 | 29 | float radius2 = radius * radius; 30 | float new_x = new_xyz[0]; 31 | float new_y = new_xyz[1]; 32 | float new_z = new_xyz[2]; 33 | 34 | int cnt = 0; 35 | for (int k = 0; k < n; ++k) { 36 | float x = xyz[k * 3 + 0]; 37 | float y = xyz[k * 3 + 1]; 38 | float z = xyz[k * 3 + 2]; 39 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + (new_z - z) * (new_z - z); 40 | if (d2 < radius2){ 41 | if (cnt == 0){ 42 | for (int l = 0; l < nsample; ++l) { 43 | idx[l] = k; 44 | } 45 | } 46 | idx[cnt] = k; 47 | ++cnt; 48 | if (cnt >= nsample) break; 49 | } 50 | } 51 | } 52 | 53 | 54 | void ball_query_kernel_launcher_fast(int b, int n, int m, float radius, int nsample, \ 55 | const float *new_xyz, const float *xyz, int *idx) { 56 | // new_xyz: (B, M, 3) 57 | // xyz: (B, N, 3) 58 | // output: 59 | // idx: (B, M, nsample) 60 | 61 | cudaError_t err; 62 | 63 | dim3 blocks(DIVUP(m, THREADS_PER_BLOCK), b); // blockIdx.x(col), blockIdx.y(row) 64 | dim3 threads(THREADS_PER_BLOCK); 65 | 66 | ball_query_kernel_fast<<>>(b, n, m, radius, nsample, new_xyz, xyz, idx); 67 | // cudaDeviceSynchronize(); // for using printf in kernel function 68 | err = cudaGetLastError(); 69 | if (cudaSuccess != err) { 70 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 71 | exit(-1); 72 | } 73 | } 74 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/roipoint_pool3d/roipoint_pool3d_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.autograd import Function 4 | 5 | from ...utils import box_utils 6 | from . import roipoint_pool3d_cuda 7 | 8 | 9 | class RoIPointPool3d(nn.Module): 10 | def __init__(self, num_sampled_points=512, pool_extra_width=1.0): 11 | super().__init__() 12 | self.num_sampled_points = num_sampled_points 13 | self.pool_extra_width = pool_extra_width 14 | 15 | def forward(self, points, point_features, boxes3d): 16 | """ 17 | Args: 18 | points: (B, N, 3) 19 | point_features: (B, N, C) 20 | boxes3d: (B, M, 7), [x, y, z, dx, dy, dz, heading] 21 | 22 | Returns: 23 | pooled_features: (B, M, 512, 3 + C) 24 | pooled_empty_flag: (B, M) 25 | """ 26 | return RoIPointPool3dFunction.apply( 27 | points, point_features, boxes3d, self.pool_extra_width, self.num_sampled_points 28 | ) 29 | 30 | 31 | class RoIPointPool3dFunction(Function): 32 | @staticmethod 33 | def forward(ctx, points, point_features, boxes3d, pool_extra_width, num_sampled_points=512): 34 | """ 35 | Args: 36 | ctx: 37 | points: (B, N, 3) 38 | point_features: (B, N, C) 39 | boxes3d: (B, num_boxes, 7), [x, y, z, dx, dy, dz, heading] 40 | pool_extra_width: 41 | num_sampled_points: 42 | 43 | Returns: 44 | pooled_features: (B, num_boxes, 512, 3 + C) 45 | pooled_empty_flag: (B, num_boxes) 46 | """ 47 | assert points.shape.__len__() == 3 and points.shape[2] == 3 48 | batch_size, boxes_num, feature_len = points.shape[0], boxes3d.shape[1], point_features.shape[2] 49 | pooled_boxes3d = box_utils.enlarge_box3d(boxes3d.view(-1, 7), pool_extra_width).view(batch_size, -1, 7) 50 | 51 | pooled_features = point_features.new_zeros((batch_size, boxes_num, num_sampled_points, 3 + feature_len)) 52 | pooled_empty_flag = point_features.new_zeros((batch_size, boxes_num)).int() 53 | 54 | roipoint_pool3d_cuda.forward( 55 | points.contiguous(), pooled_boxes3d.contiguous(), 56 | point_features.contiguous(), pooled_features, pooled_empty_flag 57 | ) 58 | 59 | return pooled_features, pooled_empty_flag 60 | 61 | @staticmethod 62 | def backward(ctx, grad_out): 63 | raise NotImplementedError 64 | 65 | 66 | if __name__ == '__main__': 67 | pass 68 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | /ckpts/ 131 | scannet/scannet_train_detection_data/* 132 | .idea/ 133 | -------------------------------------------------------------------------------- /scannet/scannet_utils.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | ''' Ref: https://github.com/ScanNet/ScanNet/blob/master/BenchmarkScripts ''' 7 | import os 8 | import sys 9 | import json 10 | import csv 11 | 12 | try: 13 | import numpy as np 14 | except: 15 | print("Failed to import numpy package.") 16 | sys.exit(-1) 17 | 18 | try: 19 | from plyfile import PlyData, PlyElement 20 | except: 21 | print("Please install the module 'plyfile' for PLY i/o, e.g.") 22 | print("pip install plyfile") 23 | sys.exit(-1) 24 | 25 | def represents_int(s): 26 | ''' if string s represents an int. ''' 27 | try: 28 | int(s) 29 | return True 30 | except ValueError: 31 | return False 32 | 33 | 34 | def read_label_mapping(filename, label_from='raw_category', label_to='nyu40id'): 35 | assert os.path.isfile(filename) 36 | mapping = dict() 37 | with open(filename) as csvfile: 38 | reader = csv.DictReader(csvfile, delimiter='\t') 39 | for row in reader: 40 | mapping[row[label_from]] = int(row[label_to]) 41 | if represents_int(list(mapping.keys())[0]): 42 | mapping = {int(k):v for k,v in mapping.items()} 43 | return mapping 44 | 45 | def read_mesh_vertices(filename): 46 | """ read XYZ for each vertex. 47 | """ 48 | assert os.path.isfile(filename) 49 | with open(filename, 'rb') as f: 50 | plydata = PlyData.read(f) 51 | num_verts = plydata['vertex'].count 52 | vertices = np.zeros(shape=[num_verts, 3], dtype=np.float32) 53 | vertices[:,0] = plydata['vertex'].data['x'] 54 | vertices[:,1] = plydata['vertex'].data['y'] 55 | vertices[:,2] = plydata['vertex'].data['z'] 56 | return vertices 57 | 58 | def read_mesh_vertices_rgb(filename): 59 | """ read XYZ RGB for each vertex. 60 | Note: RGB values are in 0-255 61 | """ 62 | assert os.path.isfile(filename) 63 | with open(filename, 'rb') as f: 64 | plydata = PlyData.read(f) 65 | num_verts = plydata['vertex'].count 66 | vertices = np.zeros(shape=[num_verts, 6], dtype=np.float32) 67 | vertices[:,0] = plydata['vertex'].data['x'] 68 | vertices[:,1] = plydata['vertex'].data['y'] 69 | vertices[:,2] = plydata['vertex'].data['z'] 70 | vertices[:,3] = plydata['vertex'].data['red'] 71 | vertices[:,4] = plydata['vertex'].data['green'] 72 | vertices[:,5] = plydata['vertex'].data['blue'] 73 | return vertices 74 | 75 | 76 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.utils.data import DataLoader 3 | from torch.utils.data import DistributedSampler as _DistributedSampler 4 | 5 | from pcdet.utils import common_utils 6 | 7 | from .dataset import DatasetTemplate 8 | from .kitti.kitti_dataset import KittiDataset 9 | from .nuscenes.nuscenes_dataset import NuScenesDataset 10 | 11 | __all__ = { 12 | 'DatasetTemplate': DatasetTemplate, 13 | 'KittiDataset': KittiDataset, 14 | 'NuScenesDataset': NuScenesDataset 15 | } 16 | 17 | 18 | class DistributedSampler(_DistributedSampler): 19 | 20 | def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True): 21 | super().__init__(dataset, num_replicas=num_replicas, rank=rank) 22 | self.shuffle = shuffle 23 | 24 | def __iter__(self): 25 | if self.shuffle: 26 | g = torch.Generator() 27 | g.manual_seed(self.epoch) 28 | indices = torch.randperm(len(self.dataset), generator=g).tolist() 29 | else: 30 | indices = torch.arange(len(self.dataset)).tolist() 31 | 32 | indices += indices[:(self.total_size - len(indices))] 33 | assert len(indices) == self.total_size 34 | 35 | indices = indices[self.rank:self.total_size:self.num_replicas] 36 | assert len(indices) == self.num_samples 37 | 38 | return iter(indices) 39 | 40 | 41 | def build_dataloader(dataset_cfg, class_names, batch_size, dist, root_path=None, workers=4, 42 | logger=None, training=True, merge_all_iters_to_one_epoch=False, total_epochs=0): 43 | 44 | dataset = __all__[dataset_cfg.DATASET]( 45 | dataset_cfg=dataset_cfg, 46 | class_names=class_names, 47 | root_path=root_path, 48 | training=training, 49 | logger=logger, 50 | ) 51 | 52 | if merge_all_iters_to_one_epoch: 53 | assert hasattr(dataset, 'merge_all_iters_to_one_epoch') 54 | dataset.merge_all_iters_to_one_epoch(merge=True, epochs=total_epochs) 55 | 56 | if dist: 57 | if training: 58 | sampler = torch.utils.data.distributed.DistributedSampler(dataset) 59 | else: 60 | rank, world_size = common_utils.get_dist_info() 61 | sampler = DistributedSampler(dataset, world_size, rank, shuffle=False) 62 | else: 63 | sampler = None 64 | dataloader = DataLoader( 65 | dataset, batch_size=batch_size, pin_memory=True, num_workers=workers, 66 | shuffle=(sampler is None) and training, collate_fn=dataset.collate_batch, 67 | drop_last=False, sampler=sampler, timeout=0 68 | ) 69 | 70 | return dataset, dataloader, sampler 71 | -------------------------------------------------------------------------------- /OpenPCDet/tools/train_utils/optimization/__init__.py: -------------------------------------------------------------------------------- 1 | from functools import partial 2 | 3 | import torch.nn as nn 4 | import torch.optim as optim 5 | import torch.optim.lr_scheduler as lr_sched 6 | 7 | from .fastai_optim import OptimWrapper 8 | from .learning_schedules_fastai import CosineWarmupLR, OneCycle 9 | 10 | 11 | def build_optimizer(model, optim_cfg): 12 | if optim_cfg.OPTIMIZER == 'adam': 13 | optimizer = optim.Adam(model.parameters(), lr=optim_cfg.LR, weight_decay=optim_cfg.WEIGHT_DECAY) 14 | elif optim_cfg.OPTIMIZER == 'sgd': 15 | optimizer = optim.SGD( 16 | model.parameters(), lr=optim_cfg.LR, weight_decay=optim_cfg.WEIGHT_DECAY, 17 | momentum=optim_cfg.MOMENTUM 18 | ) 19 | elif optim_cfg.OPTIMIZER == 'adam_onecycle': 20 | def children(m: nn.Module): 21 | return list(m.children()) 22 | 23 | def num_children(m: nn.Module) -> int: 24 | return len(children(m)) 25 | 26 | flatten_model = lambda m: sum(map(flatten_model, m.children()), []) if num_children(m) else [m] 27 | get_layer_groups = lambda m: [nn.Sequential(*flatten_model(m))] 28 | 29 | optimizer_func = partial(optim.Adam, betas=(0.9, 0.99)) 30 | optimizer = OptimWrapper.create( 31 | optimizer_func, 3e-3, get_layer_groups(model), wd=optim_cfg.WEIGHT_DECAY, true_wd=True, bn_wd=True 32 | ) 33 | else: 34 | raise NotImplementedError 35 | 36 | return optimizer 37 | 38 | 39 | def build_scheduler(optimizer, total_iters_each_epoch, total_epochs, last_epoch, optim_cfg): 40 | decay_steps = [x * total_iters_each_epoch for x in optim_cfg.DECAY_STEP_LIST] 41 | def lr_lbmd(cur_epoch): 42 | cur_decay = 1 43 | for decay_step in decay_steps: 44 | if cur_epoch >= decay_step: 45 | cur_decay = cur_decay * optim_cfg.LR_DECAY 46 | return max(cur_decay, optim_cfg.LR_CLIP / optim_cfg.LR) 47 | 48 | lr_warmup_scheduler = None 49 | total_steps = total_iters_each_epoch * total_epochs 50 | if optim_cfg.OPTIMIZER == 'adam_onecycle': 51 | lr_scheduler = OneCycle( 52 | optimizer, total_steps, optim_cfg.LR, list(optim_cfg.MOMS), optim_cfg.DIV_FACTOR, optim_cfg.PCT_START 53 | ) 54 | else: 55 | lr_scheduler = lr_sched.LambdaLR(optimizer, lr_lbmd, last_epoch=last_epoch) 56 | 57 | if optim_cfg.LR_WARMUP: 58 | lr_warmup_scheduler = CosineWarmupLR( 59 | optimizer, T_max=optim_cfg.WARMUP_EPOCH * len(total_iters_each_epoch), 60 | eta_min=optim_cfg.LR / optim_cfg.DIV_FACTOR 61 | ) 62 | 63 | return lr_scheduler, lr_warmup_scheduler 64 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/datasets/augmentor/augmentor_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from ...utils import common_utils 4 | 5 | 6 | def random_flip_along_x(gt_boxes, points): 7 | """ 8 | Args: 9 | gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]] 10 | points: (M, 3 + C) 11 | Returns: 12 | """ 13 | enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5]) 14 | if enable: 15 | gt_boxes[:, 1] = -gt_boxes[:, 1] 16 | gt_boxes[:, 6] = -gt_boxes[:, 6] 17 | points[:, 1] = -points[:, 1] 18 | 19 | if gt_boxes.shape[1] > 7: 20 | gt_boxes[:, 8] = -gt_boxes[:, 8] 21 | 22 | return gt_boxes, points 23 | 24 | 25 | def random_flip_along_y(gt_boxes, points): 26 | """ 27 | Args: 28 | gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]] 29 | points: (M, 3 + C) 30 | Returns: 31 | """ 32 | enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5]) 33 | if enable: 34 | gt_boxes[:, 0] = -gt_boxes[:, 0] 35 | gt_boxes[:, 6] = -(gt_boxes[:, 6] + np.pi) 36 | points[:, 0] = -points[:, 0] 37 | 38 | if gt_boxes.shape[1] > 7: 39 | gt_boxes[:, 7] = -gt_boxes[:, 7] 40 | 41 | return gt_boxes, points 42 | 43 | 44 | def global_rotation(gt_boxes, points, rot_range): 45 | """ 46 | Args: 47 | gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]] 48 | points: (M, 3 + C), 49 | rot_range: [min, max] 50 | Returns: 51 | """ 52 | noise_rotation = np.random.uniform(rot_range[0], rot_range[1]) 53 | points = common_utils.rotate_points_along_z(points[np.newaxis, :, :], np.array([noise_rotation]))[0] 54 | gt_boxes[:, 0:3] = common_utils.rotate_points_along_z(gt_boxes[np.newaxis, :, 0:3], np.array([noise_rotation]))[0] 55 | gt_boxes[:, 6] += noise_rotation 56 | if gt_boxes.shape[1] > 7: 57 | gt_boxes[:, 7:9] = common_utils.rotate_points_along_z( 58 | np.hstack((gt_boxes[:, 7:9], np.zeros((gt_boxes.shape[0], 1))))[np.newaxis, :, :], 59 | np.array([noise_rotation]) 60 | )[0][:, 0:2] 61 | 62 | return gt_boxes, points 63 | 64 | 65 | def global_scaling(gt_boxes, points, scale_range): 66 | """ 67 | Args: 68 | gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading] 69 | points: (M, 3 + C), 70 | scale_range: [min, max] 71 | Returns: 72 | """ 73 | if scale_range[1] - scale_range[0] < 1e-3: 74 | return gt_boxes, points 75 | noise_scale = np.random.uniform(scale_range[0], scale_range[1]) 76 | points[:, :3] *= noise_scale 77 | gt_boxes[:, :6] *= noise_scale 78 | return gt_boxes, points 79 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/model_utils/model_nms_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from ...ops.iou3d_nms import iou3d_nms_utils 4 | 5 | 6 | def class_agnostic_nms(box_scores, box_preds, nms_config, score_thresh=None): 7 | src_box_scores = box_scores 8 | if score_thresh is not None: 9 | scores_mask = (box_scores >= score_thresh) 10 | box_scores = box_scores[scores_mask] 11 | box_preds = box_preds[scores_mask] 12 | 13 | selected = [] 14 | if box_scores.shape[0] > 0: 15 | box_scores_nms, indices = torch.topk(box_scores, k=min(nms_config.NMS_PRE_MAXSIZE, box_scores.shape[0])) 16 | boxes_for_nms = box_preds[indices] 17 | keep_idx, selected_scores = getattr(iou3d_nms_utils, nms_config.NMS_TYPE)( 18 | boxes_for_nms[:, 0:7], box_scores_nms, nms_config.NMS_THRESH, **nms_config 19 | ) 20 | selected = indices[keep_idx[:nms_config.NMS_POST_MAXSIZE]] 21 | 22 | if score_thresh is not None: 23 | original_idxs = scores_mask.nonzero().view(-1) 24 | selected = original_idxs[selected] 25 | return selected, src_box_scores[selected] 26 | 27 | 28 | def multi_classes_nms(cls_scores, box_preds, nms_config, score_thresh=None): 29 | """ 30 | Args: 31 | cls_scores: (N, num_class) 32 | box_preds: (N, 7 + C) 33 | nms_config: 34 | score_thresh: 35 | 36 | Returns: 37 | 38 | """ 39 | pred_scores, pred_labels, pred_boxes = [], [], [] 40 | for k in range(cls_scores.shape[1]): 41 | if score_thresh is not None: 42 | scores_mask = (cls_scores[:, k] >= score_thresh) 43 | box_scores = cls_scores[scores_mask, k] 44 | cur_box_preds = box_preds[scores_mask] 45 | else: 46 | box_scores = cls_scores[:, k] 47 | 48 | selected = [] 49 | if box_scores.shape[0] > 0: 50 | box_scores_nms, indices = torch.topk(box_scores, k=min(nms_config.NMS_PRE_MAXSIZE, box_scores.shape[0])) 51 | boxes_for_nms = cur_box_preds[indices] 52 | keep_idx, selected_scores = getattr(iou3d_nms_utils, nms_config.NMS_TYPE)( 53 | boxes_for_nms[:, 0:7], box_scores_nms, nms_config.NMS_THRESH, **nms_config 54 | ) 55 | selected = indices[keep_idx[:nms_config.NMS_POST_MAXSIZE]] 56 | 57 | pred_scores.append(box_scores[selected]) 58 | pred_labels.append(box_scores.new_ones(len(selected)).long() * k) 59 | pred_boxes.append(cur_box_preds[selected]) 60 | 61 | pred_scores = torch.cat(pred_scores, dim=0) 62 | pred_labels = torch.cat(pred_labels, dim=0) 63 | pred_boxes = torch.cat(pred_boxes, dim=0) 64 | 65 | return pred_scores, pred_labels, pred_boxes 66 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_stack/src/group_points.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | Stacked-batch-data version of point grouping, modified from the original implementation of official PointNet++ codes. 3 | Written by Shaoshuai Shi 4 | All Rights Reserved 2019-2020. 5 | */ 6 | 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include "group_points_gpu.h" 14 | 15 | extern THCState *state; 16 | #define CHECK_CUDA(x) do { \ 17 | if (!x.type().is_cuda()) { \ 18 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 19 | exit(-1); \ 20 | } \ 21 | } while (0) 22 | #define CHECK_CONTIGUOUS(x) do { \ 23 | if (!x.is_contiguous()) { \ 24 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 25 | exit(-1); \ 26 | } \ 27 | } while (0) 28 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x) 29 | 30 | 31 | int group_points_grad_wrapper_stack(int B, int M, int C, int N, int nsample, 32 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor idx_batch_cnt_tensor, 33 | at::Tensor features_batch_cnt_tensor, at::Tensor grad_features_tensor) { 34 | 35 | CHECK_INPUT(grad_out_tensor); 36 | CHECK_INPUT(idx_tensor); 37 | CHECK_INPUT(idx_batch_cnt_tensor); 38 | CHECK_INPUT(features_batch_cnt_tensor); 39 | CHECK_INPUT(grad_features_tensor); 40 | 41 | const float *grad_out = grad_out_tensor.data(); 42 | const int *idx = idx_tensor.data(); 43 | const int *idx_batch_cnt = idx_batch_cnt_tensor.data(); 44 | const int *features_batch_cnt = features_batch_cnt_tensor.data(); 45 | float *grad_features = grad_features_tensor.data(); 46 | 47 | group_points_grad_kernel_launcher_stack(B, M, C, N, nsample, grad_out, idx, idx_batch_cnt, features_batch_cnt, grad_features); 48 | return 1; 49 | } 50 | 51 | 52 | int group_points_wrapper_stack(int B, int M, int C, int nsample, 53 | at::Tensor features_tensor, at::Tensor features_batch_cnt_tensor, 54 | at::Tensor idx_tensor, at::Tensor idx_batch_cnt_tensor, at::Tensor out_tensor) { 55 | 56 | CHECK_INPUT(features_tensor); 57 | CHECK_INPUT(features_batch_cnt_tensor); 58 | CHECK_INPUT(idx_tensor); 59 | CHECK_INPUT(idx_batch_cnt_tensor); 60 | CHECK_INPUT(out_tensor); 61 | 62 | const float *features = features_tensor.data(); 63 | const int *idx = idx_tensor.data(); 64 | const int *features_batch_cnt = features_batch_cnt_tensor.data(); 65 | const int *idx_batch_cnt = idx_batch_cnt_tensor.data(); 66 | float *out = out_tensor.data(); 67 | 68 | group_points_kernel_launcher_stack(B, M, C, nsample, features, features_batch_cnt, idx, idx_batch_cnt, out); 69 | return 1; 70 | } -------------------------------------------------------------------------------- /utils/tf_logger.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import tensorflow as tf 7 | import numpy as np 8 | import scipy.misc 9 | import tensorflow.compat.v1 as tf 10 | tf.disable_v2_behavior() 11 | try: 12 | from StringIO import StringIO # Python 2.7 13 | except ImportError: 14 | from io import BytesIO # Python 3.x 15 | 16 | 17 | class Logger(object): 18 | 19 | def __init__(self, log_dir): 20 | """Create a summary writer logging to log_dir.""" 21 | self.writer = tf.summary.FileWriter(log_dir) 22 | 23 | def scalar_summary(self, tag, value, step): 24 | """Log a scalar variable.""" 25 | summary = tf.Summary(value=[tf.Summary.Value(tag=tag, simple_value=value)]) 26 | self.writer.add_summary(summary, step) 27 | 28 | def image_summary(self, tag, images, step): 29 | """Log a list of images.""" 30 | 31 | img_summaries = [] 32 | for i, img in enumerate(images): 33 | # Write the image to a string 34 | try: 35 | s = StringIO() 36 | except: 37 | s = BytesIO() 38 | scipy.misc.toimage(img).save(s, format="png") 39 | 40 | # Create an Image object 41 | img_sum = tf.Summary.Image(encoded_image_string=s.getvalue(), 42 | height=img.shape[0], 43 | width=img.shape[1]) 44 | # Create a Summary value 45 | img_summaries.append(tf.Summary.Value(tag='%s/%d' % (tag, i), image=img_sum)) 46 | 47 | # Create and write Summary 48 | summary = tf.Summary(value=img_summaries) 49 | self.writer.add_summary(summary, step) 50 | 51 | def histo_summary(self, tag, values, step, bins=1000): 52 | """Log a histogram of the tensor of values.""" 53 | 54 | # Create a histogram using numpy 55 | counts, bin_edges = np.histogram(values, bins=bins) 56 | 57 | # Fill the fields of the histogram proto 58 | hist = tf.HistogramProto() 59 | hist.min = float(np.min(values)) 60 | hist.max = float(np.max(values)) 61 | hist.num = int(np.prod(values.shape)) 62 | hist.sum = float(np.sum(values)) 63 | hist.sum_squares = float(np.sum(values**2)) 64 | 65 | # Drop the start of the first bin 66 | bin_edges = bin_edges[1:] 67 | 68 | # Add bin edges and counts 69 | for edge in bin_edges: 70 | hist.bucket_limit.append(edge) 71 | for c in counts: 72 | hist.bucket.append(c) 73 | 74 | # Create and write Summary 75 | summary = tf.Summary(value=[tf.Summary.Value(tag=tag, histo=hist)]) 76 | self.writer.add_summary(summary, step) 77 | self.writer.flush() 78 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/config.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | 3 | import yaml 4 | from easydict import EasyDict 5 | 6 | 7 | def log_config_to_file(cfg, pre='cfg', logger=None): 8 | for key, val in cfg.items(): 9 | if isinstance(cfg[key], EasyDict): 10 | logger.info('\n%s.%s = edict()' % (pre, key)) 11 | log_config_to_file(cfg[key], pre=pre + '.' + key, logger=logger) 12 | continue 13 | logger.info('%s.%s: %s' % (pre, key, val)) 14 | 15 | 16 | def cfg_from_list(cfg_list, config): 17 | """Set config keys via list (e.g., from command line).""" 18 | from ast import literal_eval 19 | assert len(cfg_list) % 2 == 0 20 | for k, v in zip(cfg_list[0::2], cfg_list[1::2]): 21 | key_list = k.split('.') 22 | d = config 23 | for subkey in key_list[:-1]: 24 | assert subkey in d, 'NotFoundKey: %s' % subkey 25 | d = d[subkey] 26 | subkey = key_list[-1] 27 | assert subkey in d, 'NotFoundKey: %s' % subkey 28 | try: 29 | value = literal_eval(v) 30 | except: 31 | value = v 32 | 33 | if type(value) != type(d[subkey]) and isinstance(d[subkey], EasyDict): 34 | key_val_list = value.split(',') 35 | for src in key_val_list: 36 | cur_key, cur_val = src.split(':') 37 | val_type = type(d[subkey][cur_key]) 38 | cur_val = val_type(cur_val) 39 | d[subkey][cur_key] = cur_val 40 | elif type(value) != type(d[subkey]) and isinstance(d[subkey], list): 41 | val_list = value.split(',') 42 | for k, x in enumerate(val_list): 43 | val_list[k] = type(d[subkey][0])(x) 44 | d[subkey] = val_list 45 | else: 46 | assert type(value) == type(d[subkey]), \ 47 | 'type {} does not match original type {}'.format(type(value), type(d[subkey])) 48 | d[subkey] = value 49 | 50 | 51 | def merge_new_config(config, new_config): 52 | if '_BASE_CONFIG_' in new_config: 53 | with open(new_config['_BASE_CONFIG_'], 'r') as f: 54 | try: 55 | yaml_config = yaml.load(f, Loader=yaml.FullLoader) 56 | except: 57 | yaml_config = yaml.load(f) 58 | config.update(EasyDict(yaml_config)) 59 | 60 | for key, val in new_config.items(): 61 | if not isinstance(val, dict): 62 | config[key] = val 63 | continue 64 | if key not in config: 65 | config[key] = EasyDict() 66 | merge_new_config(config[key], val) 67 | 68 | return config 69 | 70 | 71 | def cfg_from_yaml_file(cfg_file, config): 72 | with open(cfg_file, 'r') as f: 73 | try: 74 | new_config = yaml.load(f, Loader=yaml.FullLoader) 75 | except: 76 | new_config = yaml.load(f) 77 | 78 | merge_new_config(config=config, new_config=new_config) 79 | 80 | return config 81 | 82 | 83 | cfg = EasyDict() 84 | cfg.ROOT_DIR = (Path(__file__).resolve().parent / '../').resolve() 85 | cfg.LOCAL_RANK = 0 86 | -------------------------------------------------------------------------------- /models/voting_module.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | ''' Voting module: generate votes from XYZ and features of seed points. 7 | 8 | Date: July, 2019 9 | Author: Charles R. Qi and Or Litany 10 | ''' 11 | 12 | import torch 13 | import torch.nn as nn 14 | import torch.nn.functional as F 15 | 16 | class VotingModule(nn.Module): 17 | def __init__(self, vote_factor, seed_feature_dim): 18 | """ Votes generation from seed point features. 19 | 20 | Args: 21 | vote_factor: int 22 | number of votes generated from each seed point 23 | seed_feature_dim: int 24 | number of channels of seed point features 25 | vote_feature_dim: int 26 | number of channels of vote features 27 | """ 28 | super().__init__() 29 | self.vote_factor = vote_factor 30 | self.in_dim = seed_feature_dim 31 | self.out_dim = self.in_dim # due to residual feature, in_dim has to be == out_dim 32 | self.conv1 = torch.nn.Conv1d(self.in_dim, self.in_dim, 1) 33 | self.conv2 = torch.nn.Conv1d(self.in_dim, self.in_dim, 1) 34 | self.conv3 = torch.nn.Conv1d(self.in_dim, (3+self.out_dim) * self.vote_factor, 1) 35 | self.bn1 = torch.nn.BatchNorm1d(self.in_dim) 36 | self.bn2 = torch.nn.BatchNorm1d(self.in_dim) 37 | 38 | def forward(self, seed_xyz, seed_features): 39 | """ Forward pass. 40 | 41 | Arguments: 42 | seed_xyz: (batch_size, num_seed, 3) Pytorch tensor 43 | seed_features: (batch_size, feature_dim, num_seed) Pytorch tensor 44 | Returns: 45 | vote_xyz: (batch_size, num_seed*vote_factor, 3) 46 | vote_features: (batch_size, vote_feature_dim, num_seed*vote_factor) 47 | """ 48 | batch_size = seed_xyz.shape[0] 49 | num_seed = seed_xyz.shape[1] 50 | num_vote = num_seed*self.vote_factor 51 | net = F.relu(self.bn1(self.conv1(seed_features))) 52 | net = F.relu(self.bn2(self.conv2(net))) 53 | net = self.conv3(net) # (batch_size, (3+out_dim)*vote_factor, num_seed) 54 | 55 | net = net.transpose(2,1).view(batch_size, num_seed, self.vote_factor, 3+self.out_dim) 56 | offset = net[:,:,:,0:3] 57 | vote_xyz = seed_xyz.unsqueeze(2) + offset.contiguous() 58 | vote_xyz = vote_xyz.contiguous().view(batch_size, num_vote, 3) 59 | 60 | residual_features = net[:,:,:,3:] # (batch_size, num_seed, vote_factor, out_dim) 61 | vote_features = seed_features.transpose(2,1).unsqueeze(2) + residual_features 62 | vote_features = vote_features.contiguous().view(batch_size, num_vote, self.out_dim) 63 | vote_features = vote_features.transpose(2,1).contiguous() 64 | 65 | return vote_xyz, vote_features 66 | 67 | if __name__=='__main__': 68 | net = VotingModule(2, 256).cuda() 69 | xyz, features = net(torch.rand(8,1024,3).cuda(), torch.rand(8,256,1024).cuda()) 70 | print('xyz', xyz.shape) 71 | print('features', features.shape) 72 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/dense_heads/anchor_head_single.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch.nn as nn 3 | 4 | from .anchor_head_template import AnchorHeadTemplate 5 | 6 | 7 | class AnchorHeadSingle(AnchorHeadTemplate): 8 | def __init__(self, model_cfg, input_channels, num_class, class_names, grid_size, point_cloud_range, 9 | predict_boxes_when_training=True): 10 | super().__init__( 11 | model_cfg=model_cfg, num_class=num_class, class_names=class_names, grid_size=grid_size, point_cloud_range=point_cloud_range, 12 | predict_boxes_when_training=predict_boxes_when_training 13 | ) 14 | 15 | self.num_anchors_per_location = sum(self.num_anchors_per_location) 16 | 17 | self.conv_cls = nn.Conv2d( 18 | input_channels, self.num_anchors_per_location * self.num_class, 19 | kernel_size=1 20 | ) 21 | self.conv_box = nn.Conv2d( 22 | input_channels, self.num_anchors_per_location * self.box_coder.code_size, 23 | kernel_size=1 24 | ) 25 | 26 | if self.model_cfg.get('USE_DIRECTION_CLASSIFIER', None) is not None: 27 | self.conv_dir_cls = nn.Conv2d( 28 | input_channels, 29 | self.num_anchors_per_location * self.model_cfg.NUM_DIR_BINS, 30 | kernel_size=1 31 | ) 32 | else: 33 | self.conv_dir_cls = None 34 | self.init_weights() 35 | 36 | def init_weights(self): 37 | pi = 0.01 38 | nn.init.constant_(self.conv_cls.bias, -np.log((1 - pi) / pi)) 39 | nn.init.normal_(self.conv_box.weight, mean=0, std=0.001) 40 | 41 | def forward(self, data_dict): 42 | spatial_features_2d = data_dict['spatial_features_2d'] 43 | 44 | cls_preds = self.conv_cls(spatial_features_2d) 45 | box_preds = self.conv_box(spatial_features_2d) 46 | 47 | cls_preds = cls_preds.permute(0, 2, 3, 1).contiguous() # [N, H, W, C] 48 | box_preds = box_preds.permute(0, 2, 3, 1).contiguous() # [N, H, W, C] 49 | 50 | self.forward_ret_dict['cls_preds'] = cls_preds 51 | self.forward_ret_dict['box_preds'] = box_preds 52 | 53 | if self.conv_dir_cls is not None: 54 | dir_cls_preds = self.conv_dir_cls(spatial_features_2d) 55 | dir_cls_preds = dir_cls_preds.permute(0, 2, 3, 1).contiguous() 56 | self.forward_ret_dict['dir_cls_preds'] = dir_cls_preds 57 | else: 58 | dir_cls_preds = None 59 | 60 | if self.training: 61 | targets_dict = self.assign_targets( 62 | gt_boxes=data_dict['gt_boxes'] 63 | ) 64 | self.forward_ret_dict.update(targets_dict) 65 | 66 | if not self.training or self.predict_boxes_when_training: 67 | batch_cls_preds, batch_box_preds = self.generate_predicted_boxes( 68 | batch_size=data_dict['batch_size'], 69 | cls_preds=cls_preds, box_preds=box_preds, dir_cls_preds=dir_cls_preds 70 | ) 71 | data_dict['batch_cls_preds'] = batch_cls_preds 72 | data_dict['batch_box_preds'] = batch_box_preds 73 | data_dict['cls_preds_normalized'] = False 74 | 75 | return data_dict 76 | -------------------------------------------------------------------------------- /sunrgbd/matlab/extract_rgbd_data_v2.m: -------------------------------------------------------------------------------- 1 | % Copyright (c) Facebook, Inc. and its affiliates. 2 | % 3 | % This source code is licensed under the MIT license found in the 4 | % LICENSE file in the root directory of this source tree. 5 | 6 | %% Dump SUNRGBD data to our format 7 | % for each sample, we have RGB image, 2d boxes. 8 | % point cloud (in camera coordinate), calibration and 3d boxes. 9 | % 10 | % Compared to extract_rgbd_data.m in frustum_pointents, use v2 2D and 3D 11 | % bboxes. 12 | % 13 | % Author: Charles R. Qi 14 | % 15 | clear; close all; clc; 16 | addpath(genpath('.')) 17 | addpath('../OFFICIAL_SUNRGBD/SUNRGBDtoolbox/readData') 18 | %% V1 2D&3D BB and Seg masks 19 | % load('./Metadata/SUNRGBDMeta.mat') 20 | % load('./Metadata/SUNRGBD2Dseg.mat') 21 | 22 | %% V2 3DBB annotations (overwrites SUNRGBDMeta) 23 | load('../OFFICIAL_SUNRGBD/SUNRGBDMeta3DBB_v2.mat'); 24 | load('../OFFICIAL_SUNRGBD/SUNRGBDMeta2DBB_v2.mat'); 25 | %% Create folders 26 | depth_folder = '../sunrgbd_trainval/depth/'; 27 | image_folder = '../sunrgbd_trainval/image/'; 28 | calib_folder = '../sunrgbd_trainval/calib/'; 29 | det_label_folder = '../sunrgbd_trainval/label/'; 30 | seg_label_folder = '../sunrgbd_trainval/seg_label/'; 31 | mkdir(depth_folder); 32 | mkdir(image_folder); 33 | mkdir(calib_folder); 34 | mkdir(det_label_folder); 35 | mkdir(seg_label_folder); 36 | %% Read 37 | parfor imageId = 1:10335 38 | imageId 39 | try 40 | data = SUNRGBDMeta(imageId); 41 | data.depthpath(1:16) = ''; 42 | data.depthpath = strcat('../OFFICIAL_SUNRGBD', data.depthpath); 43 | data.rgbpath(1:16) = ''; 44 | data.rgbpath = strcat('../OFFICIAL_SUNRGBD', data.rgbpath); 45 | 46 | % Write point cloud in depth map 47 | [rgb,points3d,depthInpaint,imsize]=read3dPoints(data); 48 | rgb(isnan(points3d(:,1)),:) = []; 49 | points3d(isnan(points3d(:,1)),:) = []; 50 | points3d_rgb = [points3d, rgb]; 51 | 52 | % MAT files are 3x smaller than TXT files. In Python we can use 53 | % scipy.io.loadmat('xxx.mat')['points3d_rgb'] to load the data. 54 | mat_filename = strcat(num2str(imageId,'%06d'), '.mat'); 55 | txt_filename = strcat(num2str(imageId,'%06d'), '.txt'); 56 | parsave(strcat(depth_folder, mat_filename), points3d_rgb); 57 | 58 | % Write images 59 | copyfile(data.rgbpath, sprintf('%s/%06d.jpg', image_folder, imageId)); 60 | 61 | % Write calibration 62 | dlmwrite(strcat(calib_folder, txt_filename), data.Rtilt(:)', 'delimiter', ' '); 63 | dlmwrite(strcat(calib_folder, txt_filename), data.K(:)', 'delimiter', ' ', '-append'); 64 | 65 | % Write 2D and 3D box label 66 | data2d = SUNRGBDMeta2DBB(imageId); 67 | fid = fopen(strcat(det_label_folder, txt_filename), 'w'); 68 | for j = 1:length(data.groundtruth3DBB) 69 | centroid = data.groundtruth3DBB(j).centroid; 70 | classname = data.groundtruth3DBB(j).classname; 71 | orientation = data.groundtruth3DBB(j).orientation; 72 | coeffs = abs(data.groundtruth3DBB(j).coeffs); 73 | box2d = data2d.groundtruth2DBB(j).gtBb2D; 74 | assert(strcmp(data2d.groundtruth2DBB(j).classname, classname)); 75 | fprintf(fid, '%s %d %d %d %d %f %f %f %f %f %f %f %f\n', classname, box2d(1), box2d(2), box2d(3), box2d(4), centroid(1), centroid(2), centroid(3), coeffs(1), coeffs(2), coeffs(3), orientation(1), orientation(2)); 76 | end 77 | fclose(fid); 78 | 79 | catch 80 | end 81 | 82 | end 83 | 84 | function parsave(filename, instance) 85 | save(filename, 'instance'); 86 | end 87 | -------------------------------------------------------------------------------- /sunrgbd/sunrgbd_trainval/sunrgbd_v1_train_0.05.txt: -------------------------------------------------------------------------------- 1 | 007173 2 | 010130 3 | 007381 4 | 008860 5 | 007155 6 | 010266 7 | 007610 8 | 005986 9 | 006838 10 | 009871 11 | 010234 12 | 008765 13 | 006882 14 | 006474 15 | 006619 16 | 005533 17 | 009789 18 | 005976 19 | 008494 20 | 005131 21 | 006403 22 | 009054 23 | 009751 24 | 007602 25 | 005471 26 | 006770 27 | 006537 28 | 005592 29 | 005794 30 | 009961 31 | 007857 32 | 006286 33 | 009972 34 | 007863 35 | 010327 36 | 009197 37 | 006320 38 | 009412 39 | 006441 40 | 009527 41 | 009302 42 | 008227 43 | 008189 44 | 008436 45 | 007619 46 | 007594 47 | 007821 48 | 007228 49 | 008820 50 | 010291 51 | 009343 52 | 005604 53 | 007659 54 | 007176 55 | 005364 56 | 005439 57 | 010123 58 | 008815 59 | 010020 60 | 008623 61 | 008135 62 | 006833 63 | 009847 64 | 007423 65 | 007765 66 | 006177 67 | 007898 68 | 007425 69 | 009950 70 | 007845 71 | 005322 72 | 009536 73 | 009822 74 | 009446 75 | 006647 76 | 007084 77 | 008017 78 | 010191 79 | 006239 80 | 009979 81 | 008385 82 | 009291 83 | 005886 84 | 007532 85 | 006855 86 | 006113 87 | 009068 88 | 008869 89 | 008584 90 | 009002 91 | 009543 92 | 009243 93 | 008172 94 | 010076 95 | 006763 96 | 007728 97 | 006752 98 | 007960 99 | 009336 100 | 008176 101 | 005628 102 | 005458 103 | 010204 104 | 009525 105 | 005423 106 | 007622 107 | 008517 108 | 007380 109 | 006011 110 | 005982 111 | 005815 112 | 006374 113 | 007941 114 | 007059 115 | 008424 116 | 006115 117 | 007485 118 | 010031 119 | 006190 120 | 008456 121 | 009133 122 | 006167 123 | 007301 124 | 006482 125 | 008307 126 | 007803 127 | 005531 128 | 009862 129 | 009044 130 | 006808 131 | 005200 132 | 008854 133 | 007125 134 | 006305 135 | 009638 136 | 010140 137 | 009863 138 | 006937 139 | 008953 140 | 007653 141 | 008578 142 | 006429 143 | 007487 144 | 005273 145 | 005464 146 | 008173 147 | 009378 148 | 007151 149 | 010265 150 | 005637 151 | 009581 152 | 007704 153 | 005474 154 | 006304 155 | 009324 156 | 008893 157 | 008395 158 | 006798 159 | 005574 160 | 010046 161 | 008344 162 | 005225 163 | 006865 164 | 006258 165 | 008466 166 | 006234 167 | 007937 168 | 008442 169 | 007159 170 | 006388 171 | 007479 172 | 005737 173 | 007761 174 | 009199 175 | 007010 176 | 005746 177 | 007130 178 | 009176 179 | 007191 180 | 009948 181 | 006262 182 | 009819 183 | 008386 184 | 007133 185 | 008670 186 | 008449 187 | 006483 188 | 009102 189 | 005559 190 | 007562 191 | 008074 192 | 010295 193 | 007407 194 | 009533 195 | 010080 196 | 007497 197 | 007214 198 | 005241 199 | 008443 200 | 006625 201 | 006863 202 | 005824 203 | 010041 204 | 006252 205 | 007278 206 | 006532 207 | 008944 208 | 005182 209 | 007662 210 | 006994 211 | 005829 212 | 008421 213 | 007107 214 | 006668 215 | 007643 216 | 008736 217 | 010268 218 | 005491 219 | 010036 220 | 006650 221 | 008907 222 | 009108 223 | 008019 224 | 009218 225 | 008198 226 | 007242 227 | 005252 228 | 006053 229 | 006221 230 | 005467 231 | 009522 232 | 008009 233 | 008862 234 | 005621 235 | 006913 236 | 006567 237 | 005759 238 | 007131 239 | 005958 240 | 006246 241 | 006826 242 | 005386 243 | 008890 244 | 006661 245 | 006324 246 | 005979 247 | 009187 248 | 005749 249 | 007746 250 | 009655 251 | 010116 252 | 005514 253 | 009634 254 | 009868 255 | 005570 256 | 008430 257 | 008617 258 | 005168 259 | 008016 260 | 007088 261 | 005424 262 | 005220 263 | 006316 264 | 005378 265 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/group_points_gpu.cu: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include 7 | #include 8 | 9 | #include "cuda_utils.h" 10 | 11 | // input: points(b, c, n) idx(b, npoints, nsample) 12 | // output: out(b, c, npoints, nsample) 13 | __global__ void group_points_kernel(int b, int c, int n, int npoints, 14 | int nsample, 15 | const float *__restrict__ points, 16 | const int *__restrict__ idx, 17 | float *__restrict__ out) { 18 | int batch_index = blockIdx.x; 19 | points += batch_index * n * c; 20 | idx += batch_index * npoints * nsample; 21 | out += batch_index * npoints * nsample * c; 22 | 23 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 24 | const int stride = blockDim.y * blockDim.x; 25 | for (int i = index; i < c * npoints; i += stride) { 26 | const int l = i / npoints; 27 | const int j = i % npoints; 28 | for (int k = 0; k < nsample; ++k) { 29 | int ii = idx[j * nsample + k]; 30 | out[(l * npoints + j) * nsample + k] = points[l * n + ii]; 31 | } 32 | } 33 | } 34 | 35 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 36 | const float *points, const int *idx, 37 | float *out) { 38 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 39 | 40 | group_points_kernel<<>>( 41 | b, c, n, npoints, nsample, points, idx, out); 42 | 43 | CUDA_CHECK_ERRORS(); 44 | } 45 | 46 | // input: grad_out(b, c, npoints, nsample), idx(b, npoints, nsample) 47 | // output: grad_points(b, c, n) 48 | __global__ void group_points_grad_kernel(int b, int c, int n, int npoints, 49 | int nsample, 50 | const float *__restrict__ grad_out, 51 | const int *__restrict__ idx, 52 | float *__restrict__ grad_points) { 53 | int batch_index = blockIdx.x; 54 | grad_out += batch_index * npoints * nsample * c; 55 | idx += batch_index * npoints * nsample; 56 | grad_points += batch_index * n * c; 57 | 58 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 59 | const int stride = blockDim.y * blockDim.x; 60 | for (int i = index; i < c * npoints; i += stride) { 61 | const int l = i / npoints; 62 | const int j = i % npoints; 63 | for (int k = 0; k < nsample; ++k) { 64 | int ii = idx[j * nsample + k]; 65 | atomicAdd(grad_points + l * n + ii, 66 | grad_out[(l * npoints + j) * nsample + k]); 67 | } 68 | } 69 | } 70 | 71 | void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 72 | int nsample, const float *grad_out, 73 | const int *idx, float *grad_points) { 74 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 75 | 76 | group_points_grad_kernel<<>>( 77 | b, c, n, npoints, nsample, grad_out, idx, grad_points); 78 | 79 | CUDA_CHECK_ERRORS(); 80 | } 81 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/sampling.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include "sampling.h" 7 | #include "utils.h" 8 | 9 | void gather_points_kernel_wrapper(int b, int c, int n, int npoints, 10 | const float *points, const int *idx, 11 | float *out); 12 | void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 13 | const float *grad_out, const int *idx, 14 | float *grad_points); 15 | 16 | void furthest_point_sampling_kernel_wrapper(int b, int n, int m, 17 | const float *dataset, float *temp, 18 | int *idxs); 19 | 20 | at::Tensor gather_points(at::Tensor points, at::Tensor idx) { 21 | CHECK_CONTIGUOUS(points); 22 | CHECK_CONTIGUOUS(idx); 23 | CHECK_IS_FLOAT(points); 24 | CHECK_IS_INT(idx); 25 | 26 | if (points.type().is_cuda()) { 27 | CHECK_CUDA(idx); 28 | } 29 | 30 | at::Tensor output = 31 | torch::zeros({points.size(0), points.size(1), idx.size(1)}, 32 | at::device(points.device()).dtype(at::ScalarType::Float)); 33 | 34 | if (points.type().is_cuda()) { 35 | gather_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), 36 | idx.size(1), points.data(), 37 | idx.data(), output.data()); 38 | } else { 39 | TORCH_CHECK(false, "CPU not supported"); 40 | } 41 | 42 | return output; 43 | } 44 | 45 | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, 46 | const int n) { 47 | CHECK_CONTIGUOUS(grad_out); 48 | CHECK_CONTIGUOUS(idx); 49 | CHECK_IS_FLOAT(grad_out); 50 | CHECK_IS_INT(idx); 51 | 52 | if (grad_out.type().is_cuda()) { 53 | CHECK_CUDA(idx); 54 | } 55 | 56 | at::Tensor output = 57 | torch::zeros({grad_out.size(0), grad_out.size(1), n}, 58 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 59 | 60 | if (grad_out.type().is_cuda()) { 61 | gather_points_grad_kernel_wrapper(grad_out.size(0), grad_out.size(1), n, 62 | idx.size(1), grad_out.data(), 63 | idx.data(), output.data()); 64 | } else { 65 | TORCH_CHECK(false, "CPU not supported"); 66 | } 67 | 68 | return output; 69 | } 70 | at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples) { 71 | CHECK_CONTIGUOUS(points); 72 | CHECK_IS_FLOAT(points); 73 | 74 | at::Tensor output = 75 | torch::zeros({points.size(0), nsamples}, 76 | at::device(points.device()).dtype(at::ScalarType::Int)); 77 | 78 | at::Tensor tmp = 79 | torch::full({points.size(0), points.size(1)}, 1e10, 80 | at::device(points.device()).dtype(at::ScalarType::Float)); 81 | 82 | if (points.type().is_cuda()) { 83 | furthest_point_sampling_kernel_wrapper( 84 | points.size(0), points.size(1), nsamples, points.data(), 85 | tmp.data(), output.data()); 86 | } else { 87 | TORCH_CHECK(false, "CPU not supported"); 88 | } 89 | 90 | return output; 91 | } 92 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_stack/src/ball_query_gpu.cu: -------------------------------------------------------------------------------- 1 | /* 2 | Stacked-batch-data version of ball query, modified from the original implementation of official PointNet++ codes. 3 | Written by Shaoshuai Shi 4 | All Rights Reserved 2019-2020. 5 | */ 6 | 7 | 8 | #include 9 | #include 10 | #include 11 | 12 | #include "ball_query_gpu.h" 13 | #include "cuda_utils.h" 14 | 15 | 16 | __global__ void ball_query_kernel_stack(int B, int M, float radius, int nsample, \ 17 | const float *new_xyz, const int *new_xyz_batch_cnt, const float *xyz, const int *xyz_batch_cnt, int *idx) { 18 | // :param xyz: (N1 + N2 ..., 3) xyz coordinates of the features 19 | // :param xyz_batch_cnt: (batch_size), [N1, N2, ...] 20 | // :param new_xyz: (M1 + M2 ..., 3) centers of the ball query 21 | // :param new_xyz_batch_cnt: (batch_size), [M1, M2, ...] 22 | // output: 23 | // idx: (M, nsample) 24 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 25 | if (pt_idx >= M) return; 26 | 27 | int bs_idx = 0, pt_cnt = new_xyz_batch_cnt[0]; 28 | for (int k = 1; k < B; k++){ 29 | if (pt_idx < pt_cnt) break; 30 | pt_cnt += new_xyz_batch_cnt[k]; 31 | bs_idx = k; 32 | } 33 | 34 | int xyz_batch_start_idx = 0; 35 | for (int k = 0; k < bs_idx; k++) xyz_batch_start_idx += xyz_batch_cnt[k]; 36 | // for (int k = 0; k < bs_idx; k++) new_xyz_batch_start_idx += new_xyz_batch_cnt[k]; 37 | 38 | new_xyz += pt_idx * 3; 39 | xyz += xyz_batch_start_idx * 3; 40 | idx += pt_idx * nsample; 41 | 42 | float radius2 = radius * radius; 43 | float new_x = new_xyz[0]; 44 | float new_y = new_xyz[1]; 45 | float new_z = new_xyz[2]; 46 | int n = xyz_batch_cnt[bs_idx]; 47 | 48 | int cnt = 0; 49 | for (int k = 0; k < n; ++k) { 50 | float x = xyz[k * 3 + 0]; 51 | float y = xyz[k * 3 + 1]; 52 | float z = xyz[k * 3 + 2]; 53 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + (new_z - z) * (new_z - z); 54 | if (d2 < radius2){ 55 | if (cnt == 0){ 56 | for (int l = 0; l < nsample; ++l) { 57 | idx[l] = k; 58 | } 59 | } 60 | idx[cnt] = k; 61 | ++cnt; 62 | if (cnt >= nsample) break; 63 | } 64 | } 65 | if (cnt == 0) idx[0] = -1; 66 | } 67 | 68 | 69 | void ball_query_kernel_launcher_stack(int B, int M, float radius, int nsample, 70 | const float *new_xyz, const int *new_xyz_batch_cnt, const float *xyz, const int *xyz_batch_cnt, int *idx){ 71 | // :param xyz: (N1 + N2 ..., 3) xyz coordinates of the features 72 | // :param xyz_batch_cnt: (batch_size), [N1, N2, ...] 73 | // :param new_xyz: (M1 + M2 ..., 3) centers of the ball query 74 | // :param new_xyz_batch_cnt: (batch_size), [M1, M2, ...] 75 | // output: 76 | // idx: (M, nsample) 77 | 78 | cudaError_t err; 79 | 80 | dim3 blocks(DIVUP(M, THREADS_PER_BLOCK)); // blockIdx.x(col), blockIdx.y(row) 81 | dim3 threads(THREADS_PER_BLOCK); 82 | 83 | ball_query_kernel_stack<<>>(B, M, radius, nsample, new_xyz, new_xyz_batch_cnt, xyz, xyz_batch_cnt, idx); 84 | // cudaDeviceSynchronize(); // for using printf in kernel function 85 | err = cudaGetLastError(); 86 | if (cudaSuccess != err) { 87 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 88 | exit(-1); 89 | } 90 | } 91 | -------------------------------------------------------------------------------- /OpenPCDet/tools/cfgs/kitti_models/second.yaml: -------------------------------------------------------------------------------- 1 | CLASS_NAMES: ['Car', 'Pedestrian', 'Cyclist'] 2 | 3 | DATA_CONFIG: 4 | _BASE_CONFIG_: cfgs/dataset_configs/kitti_dataset.yaml 5 | 6 | 7 | MODEL: 8 | NAME: SECONDNet 9 | 10 | VFE: 11 | NAME: MeanVFE 12 | 13 | BACKBONE_3D: 14 | NAME: VoxelBackBone8x 15 | 16 | MAP_TO_BEV: 17 | NAME: HeightCompression 18 | NUM_BEV_FEATURES: 256 19 | 20 | BACKBONE_2D: 21 | NAME: BaseBEVBackbone 22 | 23 | LAYER_NUMS: [5, 5] 24 | LAYER_STRIDES: [1, 2] 25 | NUM_FILTERS: [128, 256] 26 | UPSAMPLE_STRIDES: [1, 2] 27 | NUM_UPSAMPLE_FILTERS: [256, 256] 28 | 29 | DENSE_HEAD: 30 | NAME: AnchorHeadSingle 31 | CLASS_AGNOSTIC: False 32 | 33 | USE_DIRECTION_CLASSIFIER: True 34 | DIR_OFFSET: 0.78539 35 | DIR_LIMIT_OFFSET: 0.0 36 | NUM_DIR_BINS: 2 37 | 38 | ANCHOR_GENERATOR_CONFIG: [ 39 | { 40 | 'class_name': 'Car', 41 | 'anchor_sizes': [[3.9, 1.6, 1.56]], 42 | 'anchor_rotations': [0, 1.57], 43 | 'anchor_bottom_heights': [-1.78], 44 | 'align_center': False, 45 | 'feature_map_stride': 8, 46 | 'matched_threshold': 0.6, 47 | 'unmatched_threshold': 0.45 48 | }, 49 | { 50 | 'class_name': 'Pedestrian', 51 | 'anchor_sizes': [[0.8, 0.6, 1.73]], 52 | 'anchor_rotations': [0, 1.57], 53 | 'anchor_bottom_heights': [-0.6], 54 | 'align_center': False, 55 | 'feature_map_stride': 8, 56 | 'matched_threshold': 0.5, 57 | 'unmatched_threshold': 0.35 58 | }, 59 | { 60 | 'class_name': 'Cyclist', 61 | 'anchor_sizes': [[1.76, 0.6, 1.73]], 62 | 'anchor_rotations': [0, 1.57], 63 | 'anchor_bottom_heights': [-0.6], 64 | 'align_center': False, 65 | 'feature_map_stride': 8, 66 | 'matched_threshold': 0.5, 67 | 'unmatched_threshold': 0.35 68 | } 69 | ] 70 | 71 | TARGET_ASSIGNER_CONFIG: 72 | NAME: AxisAlignedTargetAssigner 73 | POS_FRACTION: -1.0 74 | SAMPLE_SIZE: 512 75 | NORM_BY_NUM_EXAMPLES: False 76 | MATCH_HEIGHT: False 77 | BOX_CODER: ResidualCoder 78 | 79 | LOSS_CONFIG: 80 | LOSS_WEIGHTS: { 81 | 'cls_weight': 1.0, 82 | 'loc_weight': 2.0, 83 | 'dir_weight': 0.2, 84 | 'code_weights': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0] 85 | } 86 | 87 | POST_PROCESSING: 88 | RECALL_THRESH_LIST: [0.3, 0.5, 0.7] 89 | SCORE_THRESH: 0.1 90 | OUTPUT_RAW_SCORE: False 91 | 92 | EVAL_METRIC: kitti 93 | 94 | NMS_CONFIG: 95 | MULTI_CLASSES_NMS: False 96 | NMS_TYPE: nms_gpu 97 | NMS_THRESH: 0.01 98 | NMS_PRE_MAXSIZE: 4096 99 | NMS_POST_MAXSIZE: 500 100 | 101 | 102 | OPTIMIZATION: 103 | BATCH_SIZE_PER_GPU: 4 104 | NUM_EPOCHS: 80 105 | 106 | OPTIMIZER: adam_onecycle 107 | LR: 0.003 108 | WEIGHT_DECAY: 0.01 109 | MOMENTUM: 0.9 110 | 111 | MOMS: [0.95, 0.85] 112 | PCT_START: 0.4 113 | DIV_FACTOR: 10 114 | DECAY_STEP_LIST: [35, 45] 115 | LR_DECAY: 0.1 116 | LR_CLIP: 0.0000001 117 | 118 | LR_WARMUP: False 119 | WARMUP_EPOCH: 1 120 | 121 | GRAD_NORM_CLIP: 10 122 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/utils/object3d_kitti.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def get_objects_from_label(label_file): 5 | with open(label_file, 'r') as f: 6 | lines = f.readlines() 7 | objects = [Object3d(line) for line in lines] 8 | return objects 9 | 10 | 11 | def cls_type_to_id(cls_type): 12 | type_to_id = {'Car': 1, 'Pedestrian': 2, 'Cyclist': 3, 'Van': 4} 13 | if cls_type not in type_to_id.keys(): 14 | return -1 15 | return type_to_id[cls_type] 16 | 17 | 18 | class Object3d(object): 19 | def __init__(self, line): 20 | label = line.strip().split(' ') 21 | self.src = line 22 | self.cls_type = label[0] 23 | self.cls_id = cls_type_to_id(self.cls_type) 24 | self.truncation = float(label[1]) 25 | self.occlusion = float(label[2]) # 0:fully visible 1:partly occluded 2:largely occluded 3:unknown 26 | self.alpha = float(label[3]) 27 | self.box2d = np.array((float(label[4]), float(label[5]), float(label[6]), float(label[7])), dtype=np.float32) 28 | self.h = float(label[8]) 29 | self.w = float(label[9]) 30 | self.l = float(label[10]) 31 | self.loc = np.array((float(label[11]), float(label[12]), float(label[13])), dtype=np.float32) 32 | self.dis_to_cam = np.linalg.norm(self.loc) 33 | self.ry = float(label[14]) 34 | self.score = float(label[15]) if label.__len__() == 16 else -1.0 35 | self.level_str = None 36 | self.level = self.get_kitti_obj_level() 37 | 38 | def get_kitti_obj_level(self): 39 | height = float(self.box2d[3]) - float(self.box2d[1]) + 1 40 | 41 | if height >= 40 and self.truncation <= 0.15 and self.occlusion <= 0: 42 | self.level_str = 'Easy' 43 | return 0 # Easy 44 | elif height >= 25 and self.truncation <= 0.3 and self.occlusion <= 1: 45 | self.level_str = 'Moderate' 46 | return 1 # Moderate 47 | elif height >= 25 and self.truncation <= 0.5 and self.occlusion <= 2: 48 | self.level_str = 'Hard' 49 | return 2 # Hard 50 | else: 51 | self.level_str = 'UnKnown' 52 | return -1 53 | 54 | def generate_corners3d(self): 55 | """ 56 | generate corners3d representation for this object 57 | :return corners_3d: (8, 3) corners of box3d in camera coord 58 | """ 59 | l, h, w = self.l, self.h, self.w 60 | x_corners = [l / 2, l / 2, -l / 2, -l / 2, l / 2, l / 2, -l / 2, -l / 2] 61 | y_corners = [0, 0, 0, 0, -h, -h, -h, -h] 62 | z_corners = [w / 2, -w / 2, -w / 2, w / 2, w / 2, -w / 2, -w / 2, w / 2] 63 | 64 | R = np.array([[np.cos(self.ry), 0, np.sin(self.ry)], 65 | [0, 1, 0], 66 | [-np.sin(self.ry), 0, np.cos(self.ry)]]) 67 | corners3d = np.vstack([x_corners, y_corners, z_corners]) # (3, 8) 68 | corners3d = np.dot(R, corners3d).T 69 | corners3d = corners3d + self.loc 70 | return corners3d 71 | 72 | def to_str(self): 73 | print_str = '%s %.3f %.3f %.3f box2d: %s hwl: [%.3f %.3f %.3f] pos: %s ry: %.3f' \ 74 | % (self.cls_type, self.truncation, self.occlusion, self.alpha, self.box2d, self.h, self.w, self.l, 75 | self.loc, self.ry) 76 | return print_str 77 | 78 | def to_kitti_format(self): 79 | kitti_str = '%s %.2f %d %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f' \ 80 | % (self.cls_type, self.truncation, int(self.occlusion), self.alpha, self.box2d[0], self.box2d[1], 81 | self.box2d[2], self.box2d[3], self.h, self.w, self.l, self.loc[0], self.loc[1], self.loc[2], 82 | self.ry) 83 | return kitti_str 84 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_batch/src/group_points_gpu.cu: -------------------------------------------------------------------------------- 1 | /* 2 | batch version of point grouping, modified from the original implementation of official PointNet++ codes. 3 | Written by Shaoshuai Shi 4 | All Rights Reserved 2018. 5 | */ 6 | 7 | #include 8 | #include 9 | 10 | #include "cuda_utils.h" 11 | #include "group_points_gpu.h" 12 | 13 | 14 | __global__ void group_points_grad_kernel_fast(int b, int c, int n, int npoints, int nsample, 15 | const float *__restrict__ grad_out, const int *__restrict__ idx, float *__restrict__ grad_points) { 16 | // grad_out: (B, C, npoints, nsample) 17 | // idx: (B, npoints, nsample) 18 | // output: 19 | // grad_points: (B, C, N) 20 | int bs_idx = blockIdx.z; 21 | int c_idx = blockIdx.y; 22 | int index = blockIdx.x * blockDim.x + threadIdx.x; 23 | int pt_idx = index / nsample; 24 | if (bs_idx >= b || c_idx >= c || pt_idx >= npoints) return; 25 | 26 | int sample_idx = index % nsample; 27 | grad_out += bs_idx * c * npoints * nsample + c_idx * npoints * nsample + pt_idx * nsample + sample_idx; 28 | idx += bs_idx * npoints * nsample + pt_idx * nsample + sample_idx; 29 | 30 | atomicAdd(grad_points + bs_idx * c * n + c_idx * n + idx[0] , grad_out[0]); 31 | } 32 | 33 | void group_points_grad_kernel_launcher_fast(int b, int c, int n, int npoints, int nsample, 34 | const float *grad_out, const int *idx, float *grad_points) { 35 | // grad_out: (B, C, npoints, nsample) 36 | // idx: (B, npoints, nsample) 37 | // output: 38 | // grad_points: (B, C, N) 39 | cudaError_t err; 40 | dim3 blocks(DIVUP(npoints * nsample, THREADS_PER_BLOCK), c, b); // blockIdx.x(col), blockIdx.y(row) 41 | dim3 threads(THREADS_PER_BLOCK); 42 | 43 | group_points_grad_kernel_fast<<>>(b, c, n, npoints, nsample, grad_out, idx, grad_points); 44 | 45 | err = cudaGetLastError(); 46 | if (cudaSuccess != err) { 47 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 48 | exit(-1); 49 | } 50 | } 51 | 52 | 53 | __global__ void group_points_kernel_fast(int b, int c, int n, int npoints, int nsample, 54 | const float *__restrict__ points, const int *__restrict__ idx, float *__restrict__ out) { 55 | // points: (B, C, N) 56 | // idx: (B, npoints, nsample) 57 | // output: 58 | // out: (B, C, npoints, nsample) 59 | int bs_idx = blockIdx.z; 60 | int c_idx = blockIdx.y; 61 | int index = blockIdx.x * blockDim.x + threadIdx.x; 62 | int pt_idx = index / nsample; 63 | if (bs_idx >= b || c_idx >= c || pt_idx >= npoints) return; 64 | 65 | int sample_idx = index % nsample; 66 | 67 | idx += bs_idx * npoints * nsample + pt_idx * nsample + sample_idx; 68 | int in_idx = bs_idx * c * n + c_idx * n + idx[0]; 69 | int out_idx = bs_idx * c * npoints * nsample + c_idx * npoints * nsample + pt_idx * nsample + sample_idx; 70 | 71 | out[out_idx] = points[in_idx]; 72 | } 73 | 74 | 75 | void group_points_kernel_launcher_fast(int b, int c, int n, int npoints, int nsample, 76 | const float *points, const int *idx, float *out) { 77 | // points: (B, C, N) 78 | // idx: (B, npoints, nsample) 79 | // output: 80 | // out: (B, C, npoints, nsample) 81 | cudaError_t err; 82 | dim3 blocks(DIVUP(npoints * nsample, THREADS_PER_BLOCK), c, b); // blockIdx.x(col), blockIdx.y(row) 83 | dim3 threads(THREADS_PER_BLOCK); 84 | 85 | group_points_kernel_fast<<>>(b, c, n, npoints, nsample, points, idx, out); 86 | // cudaDeviceSynchronize(); // for using printf in kernel function 87 | err = cudaGetLastError(); 88 | if (cudaSuccess != err) { 89 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 90 | exit(-1); 91 | } 92 | } 93 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/dense_heads/point_head_simple.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from ...utils import box_utils 4 | from .point_head_template import PointHeadTemplate 5 | 6 | 7 | class PointHeadSimple(PointHeadTemplate): 8 | """ 9 | A simple point-based segmentation head, which are used for PV-RCNN keypoint segmentaion. 10 | Reference Paper: https://arxiv.org/abs/1912.13192 11 | PV-RCNN: Point-Voxel Feature Set Abstraction for 3D Object Detection 12 | """ 13 | def __init__(self, num_class, input_channels, model_cfg, **kwargs): 14 | super().__init__(model_cfg=model_cfg, num_class=num_class) 15 | self.cls_layers = self.make_fc_layers( 16 | fc_cfg=self.model_cfg.CLS_FC, 17 | input_channels=input_channels, 18 | output_channels=num_class 19 | ) 20 | 21 | def assign_targets(self, input_dict): 22 | """ 23 | Args: 24 | input_dict: 25 | point_features: (N1 + N2 + N3 + ..., C) 26 | batch_size: 27 | point_coords: (N1 + N2 + N3 + ..., 4) [bs_idx, x, y, z] 28 | gt_boxes (optional): (B, M, 8) 29 | Returns: 30 | point_cls_labels: (N1 + N2 + N3 + ...), long type, 0:background, -1:ignored 31 | point_part_labels: (N1 + N2 + N3 + ..., 3) 32 | """ 33 | point_coords = input_dict['point_coords'] 34 | gt_boxes = input_dict['gt_boxes'] 35 | assert gt_boxes.shape.__len__() == 3, 'gt_boxes.shape=%s' % str(gt_boxes.shape) 36 | assert point_coords.shape.__len__() in [2], 'points.shape=%s' % str(point_coords.shape) 37 | 38 | batch_size = gt_boxes.shape[0] 39 | extend_gt_boxes = box_utils.enlarge_box3d( 40 | gt_boxes.view(-1, gt_boxes.shape[-1]), extra_width=self.model_cfg.TARGET_CONFIG.GT_EXTRA_WIDTH 41 | ).view(batch_size, -1, gt_boxes.shape[-1]) 42 | targets_dict = self.assign_stack_targets( 43 | points=point_coords, gt_boxes=gt_boxes, extend_gt_boxes=extend_gt_boxes, 44 | set_ignore_flag=True, use_ball_constraint=False, 45 | ret_part_labels=False 46 | ) 47 | 48 | return targets_dict 49 | 50 | def get_loss(self, tb_dict=None): 51 | tb_dict = {} if tb_dict is None else tb_dict 52 | point_loss_cls, tb_dict_1 = self.get_cls_layer_loss() 53 | 54 | point_loss = point_loss_cls 55 | tb_dict.update(tb_dict_1) 56 | return point_loss, tb_dict 57 | 58 | def forward(self, batch_dict): 59 | """ 60 | Args: 61 | batch_dict: 62 | batch_size: 63 | point_features: (N1 + N2 + N3 + ..., C) or (B, N, C) 64 | point_features_before_fusion: (N1 + N2 + N3 + ..., C) 65 | point_coords: (N1 + N2 + N3 + ..., 4) [bs_idx, x, y, z] 66 | point_labels (optional): (N1 + N2 + N3 + ...) 67 | gt_boxes (optional): (B, M, 8) 68 | Returns: 69 | batch_dict: 70 | point_cls_scores: (N1 + N2 + N3 + ..., 1) 71 | point_part_offset: (N1 + N2 + N3 + ..., 3) 72 | """ 73 | if self.model_cfg.get('USE_POINT_FEATURES_BEFORE_FUSION', False): 74 | point_features = batch_dict['point_features_before_fusion'] 75 | else: 76 | point_features = batch_dict['point_features'] 77 | point_cls_preds = self.cls_layers(point_features) # (total_points, num_class) 78 | 79 | ret_dict = { 80 | 'point_cls_preds': point_cls_preds, 81 | } 82 | 83 | point_cls_scores = torch.sigmoid(point_cls_preds) 84 | batch_dict['point_cls_scores'], _ = point_cls_scores.max(dim=-1) 85 | 86 | if self.training: 87 | targets_dict = self.assign_targets(batch_dict) 88 | ret_dict['point_cls_labels'] = targets_dict['point_cls_labels'] 89 | self.forward_ret_dict = ret_dict 90 | 91 | return batch_dict 92 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/interpolate.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include "interpolate.h" 7 | #include "utils.h" 8 | 9 | void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown, 10 | const float *known, float *dist2, int *idx); 11 | void three_interpolate_kernel_wrapper(int b, int c, int m, int n, 12 | const float *points, const int *idx, 13 | const float *weight, float *out); 14 | void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m, 15 | const float *grad_out, 16 | const int *idx, const float *weight, 17 | float *grad_points); 18 | 19 | std::vector three_nn(at::Tensor unknowns, at::Tensor knows) { 20 | CHECK_CONTIGUOUS(unknowns); 21 | CHECK_CONTIGUOUS(knows); 22 | CHECK_IS_FLOAT(unknowns); 23 | CHECK_IS_FLOAT(knows); 24 | 25 | if (unknowns.type().is_cuda()) { 26 | CHECK_CUDA(knows); 27 | } 28 | 29 | at::Tensor idx = 30 | torch::zeros({unknowns.size(0), unknowns.size(1), 3}, 31 | at::device(unknowns.device()).dtype(at::ScalarType::Int)); 32 | at::Tensor dist2 = 33 | torch::zeros({unknowns.size(0), unknowns.size(1), 3}, 34 | at::device(unknowns.device()).dtype(at::ScalarType::Float)); 35 | 36 | if (unknowns.type().is_cuda()) { 37 | three_nn_kernel_wrapper(unknowns.size(0), unknowns.size(1), knows.size(1), 38 | unknowns.data(), knows.data(), 39 | dist2.data(), idx.data()); 40 | } else { 41 | TORCH_CHECK(false, "CPU not supported"); 42 | } 43 | 44 | return {dist2, idx}; 45 | } 46 | 47 | at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, 48 | at::Tensor weight) { 49 | CHECK_CONTIGUOUS(points); 50 | CHECK_CONTIGUOUS(idx); 51 | CHECK_CONTIGUOUS(weight); 52 | CHECK_IS_FLOAT(points); 53 | CHECK_IS_INT(idx); 54 | CHECK_IS_FLOAT(weight); 55 | 56 | if (points.type().is_cuda()) { 57 | CHECK_CUDA(idx); 58 | CHECK_CUDA(weight); 59 | } 60 | 61 | at::Tensor output = 62 | torch::zeros({points.size(0), points.size(1), idx.size(1)}, 63 | at::device(points.device()).dtype(at::ScalarType::Float)); 64 | 65 | if (points.type().is_cuda()) { 66 | three_interpolate_kernel_wrapper( 67 | points.size(0), points.size(1), points.size(2), idx.size(1), 68 | points.data(), idx.data(), weight.data(), 69 | output.data()); 70 | } else { 71 | TORCH_CHECK(false, "CPU not supported"); 72 | } 73 | 74 | return output; 75 | } 76 | at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, 77 | at::Tensor weight, const int m) { 78 | CHECK_CONTIGUOUS(grad_out); 79 | CHECK_CONTIGUOUS(idx); 80 | CHECK_CONTIGUOUS(weight); 81 | CHECK_IS_FLOAT(grad_out); 82 | CHECK_IS_INT(idx); 83 | CHECK_IS_FLOAT(weight); 84 | 85 | if (grad_out.type().is_cuda()) { 86 | CHECK_CUDA(idx); 87 | CHECK_CUDA(weight); 88 | } 89 | 90 | at::Tensor output = 91 | torch::zeros({grad_out.size(0), grad_out.size(1), m}, 92 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 93 | 94 | if (grad_out.type().is_cuda()) { 95 | three_interpolate_kernel_wrapper( 96 | grad_out.size(0), grad_out.size(1), grad_out.size(2), m, 97 | grad_out.data(), idx.data(), weight.data(), 98 | output.data()); 99 | } else { 100 | TORCH_CHECK(false, "CPU not supported"); 101 | } 102 | 103 | return output; 104 | } 105 | -------------------------------------------------------------------------------- /OpenPCDet/tools/cfgs/kitti_models/PartA2_free.yaml: -------------------------------------------------------------------------------- 1 | CLASS_NAMES: ['Car', 'Pedestrian', 'Cyclist'] 2 | 3 | DATA_CONFIG: 4 | _BASE_CONFIG_: cfgs/dataset_configs/kitti_dataset.yaml 5 | 6 | 7 | MODEL: 8 | NAME: PointRCNN 9 | 10 | VFE: 11 | NAME: MeanVFE 12 | 13 | BACKBONE_3D: 14 | NAME: UNetV2 15 | RETURN_ENCODED_TENSOR: False 16 | 17 | POINT_HEAD: 18 | NAME: PointIntraPartOffsetHead 19 | CLS_FC: [128, 128] 20 | PART_FC: [128, 128] 21 | REG_FC: [128, 128] 22 | CLASS_AGNOSTIC: False 23 | USE_POINT_FEATURES_BEFORE_FUSION: False 24 | TARGET_CONFIG: 25 | GT_EXTRA_WIDTH: [0.2, 0.2, 0.2] 26 | BOX_CODER: PointResidualCoder 27 | BOX_CODER_CONFIG: { 28 | 'use_mean_size': True, 29 | 'mean_size': [ 30 | [3.9, 1.6, 1.56], 31 | [0.8, 0.6, 1.73], 32 | [1.76, 0.6, 1.73] 33 | ] 34 | } 35 | 36 | LOSS_CONFIG: 37 | LOSS_REG: WeightedSmoothL1Loss 38 | LOSS_WEIGHTS: { 39 | 'point_cls_weight': 1.0, 40 | 'point_box_weight': 1.0, 41 | 'point_part_weight': 1.0, 42 | 'code_weights': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0] 43 | } 44 | 45 | ROI_HEAD: 46 | NAME: PartA2FCHead 47 | CLASS_AGNOSTIC: True 48 | 49 | SHARED_FC: [256, 256, 256] 50 | CLS_FC: [256, 256] 51 | REG_FC: [256, 256] 52 | DP_RATIO: 0.3 53 | DISABLE_PART: True 54 | SEG_MASK_SCORE_THRESH: 0.0 55 | 56 | NMS_CONFIG: 57 | TRAIN: 58 | NMS_TYPE: nms_gpu 59 | MULTI_CLASSES_NMS: False 60 | NMS_PRE_MAXSIZE: 9000 61 | NMS_POST_MAXSIZE: 512 62 | NMS_THRESH: 0.8 63 | TEST: 64 | NMS_TYPE: nms_gpu 65 | MULTI_CLASSES_NMS: False 66 | NMS_PRE_MAXSIZE: 9000 67 | NMS_POST_MAXSIZE: 100 68 | NMS_THRESH: 0.85 69 | 70 | ROI_AWARE_POOL: 71 | POOL_SIZE: 12 72 | NUM_FEATURES: 128 73 | MAX_POINTS_PER_VOXEL: 128 74 | 75 | TARGET_CONFIG: 76 | BOX_CODER: ResidualCoder 77 | ROI_PER_IMAGE: 128 78 | FG_RATIO: 0.5 79 | 80 | SAMPLE_ROI_BY_EACH_CLASS: True 81 | CLS_SCORE_TYPE: roi_iou 82 | 83 | CLS_FG_THRESH: 0.75 84 | CLS_BG_THRESH: 0.25 85 | CLS_BG_THRESH_LO: 0.1 86 | HARD_BG_RATIO: 0.8 87 | 88 | REG_FG_THRESH: 0.65 89 | 90 | LOSS_CONFIG: 91 | CLS_LOSS: BinaryCrossEntropy 92 | REG_LOSS: smooth-l1 93 | CORNER_LOSS_REGULARIZATION: True 94 | LOSS_WEIGHTS: { 95 | 'rcnn_cls_weight': 1.0, 96 | 'rcnn_reg_weight': 1.0, 97 | 'rcnn_corner_weight': 1.0, 98 | 'code_weights': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0] 99 | } 100 | 101 | POST_PROCESSING: 102 | RECALL_THRESH_LIST: [0.3, 0.5, 0.7] 103 | SCORE_THRESH: 0.1 104 | OUTPUT_RAW_SCORE: False 105 | 106 | EVAL_METRIC: kitti 107 | 108 | NMS_CONFIG: 109 | MULTI_CLASSES_NMS: False 110 | NMS_TYPE: nms_gpu 111 | NMS_THRESH: 0.1 112 | NMS_PRE_MAXSIZE: 4096 113 | NMS_POST_MAXSIZE: 500 114 | 115 | 116 | OPTIMIZATION: 117 | BATCH_SIZE_PER_GPU: 4 118 | NUM_EPOCHS: 80 119 | 120 | OPTIMIZER: adam_onecycle 121 | LR: 0.003 122 | WEIGHT_DECAY: 0.01 123 | MOMENTUM: 0.9 124 | 125 | MOMS: [0.95, 0.85] 126 | PCT_START: 0.4 127 | DIV_FACTOR: 10 128 | DECAY_STEP_LIST: [35, 45] 129 | LR_DECAY: 0.1 130 | LR_CLIP: 0.0000001 131 | 132 | LR_WARMUP: False 133 | WARMUP_EPOCH: 1 134 | 135 | GRAD_NORM_CLIP: 10 136 | -------------------------------------------------------------------------------- /OpenPCDet/tools/demo.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import glob 3 | from pathlib import Path 4 | 5 | import mayavi.mlab as mlab 6 | import numpy as np 7 | import torch 8 | 9 | from pcdet.config import cfg, cfg_from_yaml_file 10 | from pcdet.datasets import DatasetTemplate 11 | from pcdet.models import build_network, load_data_to_gpu 12 | from pcdet.utils import common_utils 13 | from visual_utils import visualize_utils as V 14 | 15 | 16 | class DemoDataset(DatasetTemplate): 17 | def __init__(self, dataset_cfg, class_names, training=True, root_path=None, logger=None, ext='.bin'): 18 | """ 19 | Args: 20 | root_path: 21 | dataset_cfg: 22 | class_names: 23 | training: 24 | logger: 25 | """ 26 | super().__init__( 27 | dataset_cfg=dataset_cfg, class_names=class_names, training=training, root_path=root_path, logger=logger 28 | ) 29 | self.root_path = root_path 30 | self.ext = ext 31 | data_file_list = glob.glob(str(root_path / f'*{self.ext}')) if self.root_path.is_dir() else [self.root_path] 32 | 33 | data_file_list.sort() 34 | self.sample_file_list = data_file_list 35 | 36 | def __len__(self): 37 | return len(self.sample_file_list) 38 | 39 | def __getitem__(self, index): 40 | if self.ext == '.bin': 41 | points = np.fromfile(self.sample_file_list[index], dtype=np.float32).reshape(-1, 4) 42 | elif self.ext == '.npy': 43 | points = np.load(self.sample_file_list[index]) 44 | else: 45 | raise NotImplementedError 46 | 47 | input_dict = { 48 | 'points': points, 49 | 'frame_id': index, 50 | } 51 | 52 | data_dict = self.prepare_data(data_dict=input_dict) 53 | return data_dict 54 | 55 | 56 | def parse_config(): 57 | parser = argparse.ArgumentParser(description='arg parser') 58 | parser.add_argument('--cfg_file', type=str, default='cfgs/kitti_models/second.yaml', 59 | help='specify the config for demo') 60 | parser.add_argument('--data_path', type=str, default='demo_data', 61 | help='specify the point cloud data file or directory') 62 | parser.add_argument('--ckpt', type=str, default=None, help='specify the pretrained model') 63 | parser.add_argument('--ext', type=str, default='.bin', help='specify the extension of your point cloud data file') 64 | 65 | args = parser.parse_args() 66 | 67 | cfg_from_yaml_file(args.cfg_file, cfg) 68 | 69 | return args, cfg 70 | 71 | 72 | def main(): 73 | args, cfg = parse_config() 74 | logger = common_utils.create_logger() 75 | logger.info('-----------------Quick Demo of OpenPCDet-------------------------') 76 | demo_dataset = DemoDataset( 77 | dataset_cfg=cfg.DATA_CONFIG, class_names=cfg.CLASS_NAMES, training=False, 78 | root_path=Path(args.data_path), ext=args.ext, logger=logger 79 | ) 80 | logger.info(f'Total number of samples: \t{len(demo_dataset)}') 81 | 82 | model = build_network(model_cfg=cfg.MODEL, num_class=len(cfg.CLASS_NAMES), dataset=demo_dataset) 83 | model.load_params_from_file(filename=args.ckpt, logger=logger, to_cpu=True) 84 | model.cuda() 85 | model.eval() 86 | with torch.no_grad(): 87 | for idx, data_dict in enumerate(demo_dataset): 88 | logger.info(f'Visualized sample index: \t{idx + 1}') 89 | data_dict = demo_dataset.collate_batch([data_dict]) 90 | load_data_to_gpu(data_dict) 91 | pred_dicts, _ = model.forward(data_dict) 92 | 93 | V.draw_scenes( 94 | points=data_dict['points'][:, 1:], ref_boxes=pred_dicts[0]['pred_boxes'], 95 | ref_scores=pred_dicts[0]['pred_scores'], ref_labels=pred_dicts[0]['pred_labels'] 96 | ) 97 | mlab.show(stop=True) 98 | 99 | logger.info('Demo done.') 100 | 101 | 102 | if __name__ == '__main__': 103 | main() 104 | -------------------------------------------------------------------------------- /generate_random_split.py: -------------------------------------------------------------------------------- 1 | """ Generate random data splits 2 | 3 | Written by: Yezhen Cong, 2020 4 | """ 5 | 6 | import os 7 | import sys 8 | import numpy as np 9 | 10 | from scannet.model_util_scannet import ScannetDatasetConfig 11 | from sunrgbd.model_util_sunrgbd import SunrgbdDatasetConfig 12 | 13 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 14 | ROOT_DIR = BASE_DIR 15 | 16 | 17 | def gen_scannet_split(labeled_ratio, count): 18 | DC = ScannetDatasetConfig() 19 | split_set = 'train' 20 | split_filenames = os.path.join(ROOT_DIR, 'scannet/meta_data', 21 | 'scannetv2_{}.txt'.format(split_set)) 22 | with open(split_filenames, 'r') as f: 23 | scan_names = f.read().splitlines() 24 | # remove unavailiable scans 25 | num_scans = len(scan_names) 26 | scan2label = np.zeros((num_scans, DC.num_class)) 27 | num_labeled_scans = int(labeled_ratio * num_scans) 28 | data_path = os.path.join(BASE_DIR, 'scannet/scannet_train_detection_data') 29 | for i, scan_name in enumerate(scan_names): 30 | instance_bboxes = np.load(os.path.join(data_path, scan_name) + '_bbox.npy') 31 | class_ind = [DC.nyu40id2class[x] for x in instance_bboxes[:, -1]] 32 | if class_ind != []: 33 | unique_class_ind = list(set(class_ind)) 34 | else: 35 | continue 36 | for j in unique_class_ind: 37 | scan2label[i, j] = 1 38 | 39 | while True: 40 | choices = np.random.choice(num_scans, num_labeled_scans, replace=False) 41 | class_distr = np.sum(scan2label[choices], axis=0) 42 | class_mask = np.where(class_distr > 0, 1, 0) 43 | if np.sum(class_mask) == DC.num_class: 44 | labeled_scan_names = list(np.array(scan_names)[choices]) 45 | with open(os.path.join(ROOT_DIR, 'scannet/meta_data/scannetv2_train_{}_{}.txt'.format(labeled_ratio, count)), 46 | 'w') as f: 47 | for scan_name in labeled_scan_names: 48 | f.write(scan_name + '\n') 49 | break 50 | 51 | unlabeled_scan_names = list(set(scan_names) - set(labeled_scan_names)) 52 | print('\tSelected {} labeled scans, remained {} unlabeled scans'.format(len(labeled_scan_names), len(unlabeled_scan_names))) 53 | 54 | 55 | def gen_sunrgbd_split(labeled_ratio, count): 56 | DC = SunrgbdDatasetConfig() 57 | data_path = os.path.join(ROOT_DIR, 'sunrgbd/sunrgbd_pc_bbox_votes_50k_v1_%s' % ('train')) 58 | raw_data_path = os.path.join(ROOT_DIR, 'sunrgbd/sunrgbd_trainval') 59 | scan_names = sorted(list(set([os.path.basename(x)[0:6] for x in os.listdir(data_path)]))) 60 | num_scans = len(scan_names) 61 | num_labeled_scans = int(labeled_ratio * num_scans) 62 | scan2label = np.zeros((num_scans, DC.num_class)) 63 | for i, scan_name in enumerate(scan_names): 64 | bboxes = np.load(os.path.join(data_path, scan_name) + '_bbox.npy') # K,8 65 | class_ind = bboxes[:, -1] 66 | if len(class_ind) != 0: 67 | unique_class_ind = np.unique(class_ind) 68 | else: 69 | continue 70 | for j in unique_class_ind: 71 | scan2label[i, int(j)] = 1 72 | 73 | while True: 74 | choices = np.random.choice(num_scans, num_labeled_scans, replace=False) 75 | class_distr = np.sum(scan2label[choices], axis=0) 76 | class_mask = np.where(class_distr > 0, 1, 0) 77 | if np.sum(class_mask) == DC.num_class: 78 | labeled_scan_names = list(np.array(scan_names)[choices]) 79 | with open(os.path.join(raw_data_path, 'sunrgbd_v1_train_{}_{}.txt'.format(labeled_ratio, count)), 'w') as f: 80 | for scan_name in labeled_scan_names: 81 | f.write(scan_name + '\n') 82 | break 83 | 84 | unlabeled_scan_names = list(set(scan_names) - set(labeled_scan_names)) 85 | print('Selected {} labeled scans, remained {} unlabeled scans'.format(len(labeled_scan_names), len(unlabeled_scan_names))) 86 | 87 | ratio = float(sys.argv[1]) 88 | dataset = sys.argv[2] 89 | count = int(sys.argv[3]) 90 | 91 | if dataset == 'scannet': 92 | gen_scannet_split(ratio, count) 93 | else: 94 | gen_sunrgbd_split(ratio, count) 95 | -------------------------------------------------------------------------------- /OpenPCDet/tools/cfgs/kitti_models/second_multihead.yaml: -------------------------------------------------------------------------------- 1 | CLASS_NAMES: ['Car', 'Pedestrian', 'Cyclist'] 2 | DATA_CONFIG: 3 | _BASE_CONFIG_: cfgs/dataset_configs/kitti_dataset.yaml 4 | 5 | 6 | MODEL: 7 | NAME: SECONDNet 8 | 9 | VFE: 10 | NAME: MeanVFE 11 | 12 | BACKBONE_3D: 13 | NAME: VoxelBackBone8x 14 | 15 | MAP_TO_BEV: 16 | NAME: HeightCompression 17 | NUM_BEV_FEATURES: 256 18 | 19 | BACKBONE_2D: 20 | NAME: BaseBEVBackbone 21 | 22 | LAYER_NUMS: [5, 5] 23 | LAYER_STRIDES: [1, 2] 24 | NUM_FILTERS: [128, 256] 25 | UPSAMPLE_STRIDES: [1, 2] 26 | NUM_UPSAMPLE_FILTERS: [256, 256] 27 | 28 | DENSE_HEAD: 29 | NAME: AnchorHeadMulti 30 | CLASS_AGNOSTIC: False 31 | 32 | USE_DIRECTION_CLASSIFIER: True 33 | DIR_OFFSET: 0.78539 34 | DIR_LIMIT_OFFSET: 0.0 35 | NUM_DIR_BINS: 2 36 | 37 | USE_MULTIHEAD: True 38 | SEPARATE_MULTIHEAD: True 39 | ANCHOR_GENERATOR_CONFIG: [ 40 | { 41 | 'class_name': 'Car', 42 | 'anchor_sizes': [[3.9, 1.6, 1.56]], 43 | 'anchor_rotations': [0, 1.57], 44 | 'anchor_bottom_heights': [-1.6], 45 | 'align_center': False, 46 | 'feature_map_stride': 8, 47 | 'matched_threshold': 0.6, 48 | 'unmatched_threshold': 0.45 49 | }, 50 | { 51 | 'class_name': 'Pedestrian', 52 | 'anchor_sizes': [[0.8, 0.6, 1.73]], 53 | 'anchor_rotations': [0, 1.57], 54 | 'anchor_bottom_heights': [-1.6], 55 | 'align_center': False, 56 | 'feature_map_stride': 8, 57 | 'matched_threshold': 0.5, 58 | 'unmatched_threshold': 0.35 59 | }, 60 | { 61 | 'class_name': 'Cyclist', 62 | 'anchor_sizes': [[1.76, 0.6, 1.73]], 63 | 'anchor_rotations': [0, 1.57], 64 | 'anchor_bottom_heights': [-1.6], 65 | 'align_center': False, 66 | 'feature_map_stride': 8, 67 | 'matched_threshold': 0.5, 68 | 'unmatched_threshold': 0.35 69 | } 70 | ] 71 | 72 | SHARED_CONV_NUM_FILTER: 64 73 | 74 | RPN_HEAD_CFGS: [ 75 | { 76 | 'HEAD_CLS_NAME': ['Car'], 77 | }, 78 | { 79 | 'HEAD_CLS_NAME': ['Pedestrian'], 80 | }, 81 | { 82 | 'HEAD_CLS_NAME': ['Cyclist'], 83 | } 84 | ] 85 | 86 | TARGET_ASSIGNER_CONFIG: 87 | NAME: AxisAlignedTargetAssigner 88 | POS_FRACTION: -1.0 89 | SAMPLE_SIZE: 512 90 | NORM_BY_NUM_EXAMPLES: False 91 | MATCH_HEIGHT: False 92 | BOX_CODER: ResidualCoder 93 | 94 | LOSS_CONFIG: 95 | LOSS_WEIGHTS: { 96 | 'cls_weight': 1.0, 97 | 'loc_weight': 2.0, 98 | 'dir_weight': 0.2, 99 | 'code_weights': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0] 100 | } 101 | 102 | POST_PROCESSING: 103 | RECALL_THRESH_LIST: [0.3, 0.5, 0.7] 104 | MULTI_CLASSES_NMS: False 105 | SCORE_THRESH: 0.1 106 | OUTPUT_RAW_SCORE: False 107 | 108 | EVAL_METRIC: kitti 109 | 110 | NMS_CONFIG: 111 | MULTI_CLASSES_NMS: False 112 | NMS_TYPE: nms_gpu 113 | NMS_THRESH: 0.1 114 | NMS_PRE_MAXSIZE: 4096 115 | NMS_POST_MAXSIZE: 500 116 | 117 | 118 | OPTIMIZATION: 119 | BATCH_SIZE_PER_GPU: 4 120 | NUM_EPOCHS: 80 121 | 122 | OPTIMIZER: adam_onecycle 123 | LR: 0.003 124 | WEIGHT_DECAY: 0.01 125 | MOMENTUM: 0.9 126 | 127 | MOMS: [0.95, 0.85] 128 | PCT_START: 0.4 129 | DIV_FACTOR: 10 130 | DECAY_STEP_LIST: [35, 45] 131 | LR_DECAY: 0.1 132 | LR_CLIP: 0.0000001 133 | 134 | LR_WARMUP: False 135 | WARMUP_EPOCH: 1 136 | 137 | GRAD_NORM_CLIP: 10 138 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/iou3d_nms/iou3d_nms_utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3D IoU Calculation and Rotated NMS 3 | Written by Shaoshuai Shi 4 | All Rights Reserved 2019-2020. 5 | """ 6 | import torch 7 | 8 | from ...utils import common_utils 9 | from . import iou3d_nms_cuda 10 | 11 | 12 | def boxes_bev_iou_cpu(boxes_a, boxes_b): 13 | """ 14 | Args: 15 | boxes_a: (N, 7) [x, y, z, dx, dy, dz, heading] 16 | boxes_b: (N, 7) [x, y, z, dx, dy, dz, heading] 17 | 18 | Returns: 19 | 20 | """ 21 | boxes_a, is_numpy = common_utils.check_numpy_to_torch(boxes_a) 22 | boxes_b, is_numpy = common_utils.check_numpy_to_torch(boxes_b) 23 | assert not (boxes_a.is_cuda or boxes_b.is_cuda), 'Only support CPU tensors' 24 | assert boxes_a.shape[1] == 7 and boxes_b.shape[1] == 7 25 | ans_iou = boxes_a.new_zeros(torch.Size((boxes_a.shape[0], boxes_b.shape[0]))) 26 | iou3d_nms_cuda.boxes_iou_bev_cpu(boxes_a.contiguous(), boxes_b.contiguous(), ans_iou) 27 | 28 | return ans_iou.numpy() if is_numpy else ans_iou 29 | 30 | 31 | def boxes_iou_bev(boxes_a, boxes_b): 32 | """ 33 | Args: 34 | boxes_a: (N, 7) [x, y, z, dx, dy, dz, heading] 35 | boxes_b: (N, 7) [x, y, z, dx, dy, dz, heading] 36 | 37 | Returns: 38 | ans_iou: (N, M) 39 | """ 40 | assert boxes_a.shape[1] == boxes_b.shape[1] == 7 41 | ans_iou = torch.cuda.FloatTensor(torch.Size((boxes_a.shape[0], boxes_b.shape[0]))).zero_() 42 | 43 | iou3d_nms_cuda.boxes_iou_bev_gpu(boxes_a.contiguous(), boxes_b.contiguous(), ans_iou) 44 | 45 | return ans_iou 46 | 47 | 48 | def boxes_iou3d_gpu(boxes_a, boxes_b): 49 | """ 50 | Args: 51 | boxes_a: (N, 7) [x, y, z, dx, dy, dz, heading] 52 | boxes_b: (N, 7) [x, y, z, dx, dy, dz, heading] 53 | 54 | Returns: 55 | ans_iou: (N, M) 56 | """ 57 | assert boxes_a.shape[1] == boxes_b.shape[1] == 7 58 | 59 | # height overlap 60 | boxes_a_height_max = (boxes_a[:, 2] + boxes_a[:, 5] / 2).view(-1, 1) 61 | boxes_a_height_min = (boxes_a[:, 2] - boxes_a[:, 5] / 2).view(-1, 1) 62 | boxes_b_height_max = (boxes_b[:, 2] + boxes_b[:, 5] / 2).view(1, -1) 63 | boxes_b_height_min = (boxes_b[:, 2] - boxes_b[:, 5] / 2).view(1, -1) 64 | 65 | # bev overlap 66 | overlaps_bev = torch.cuda.FloatTensor(torch.Size((boxes_a.shape[0], boxes_b.shape[0]))).zero_() # (N, M) 67 | iou3d_nms_cuda.boxes_overlap_bev_gpu(boxes_a.contiguous(), boxes_b.contiguous(), overlaps_bev) 68 | 69 | max_of_min = torch.max(boxes_a_height_min, boxes_b_height_min) 70 | min_of_max = torch.min(boxes_a_height_max, boxes_b_height_max) 71 | overlaps_h = torch.clamp(min_of_max - max_of_min, min=0) 72 | 73 | # 3d iou 74 | overlaps_3d = overlaps_bev * overlaps_h 75 | 76 | vol_a = (boxes_a[:, 3] * boxes_a[:, 4] * boxes_a[:, 5]).view(-1, 1) 77 | vol_b = (boxes_b[:, 3] * boxes_b[:, 4] * boxes_b[:, 5]).view(1, -1) 78 | 79 | iou3d = overlaps_3d / torch.clamp(vol_a + vol_b - overlaps_3d, min=1e-6) 80 | 81 | return iou3d 82 | 83 | 84 | def nms_gpu(boxes, scores, thresh, pre_maxsize=None, **kwargs): 85 | """ 86 | :param boxes: (N, 7) [x, y, z, dx, dy, dz, heading] 87 | :param scores: (N) 88 | :param thresh: 89 | :return: 90 | """ 91 | assert boxes.shape[1] == 7 92 | order = scores.sort(0, descending=True)[1] 93 | if pre_maxsize is not None: 94 | order = order[:pre_maxsize] 95 | 96 | boxes = boxes[order].contiguous() 97 | keep = torch.LongTensor(boxes.size(0)) 98 | num_out = iou3d_nms_cuda.nms_gpu(boxes, keep, thresh) 99 | return order[keep[:num_out].cuda()].contiguous(), None 100 | 101 | 102 | def nms_normal_gpu(boxes, scores, thresh, **kwargs): 103 | """ 104 | :param boxes: (N, 7) [x, y, z, dx, dy, dz, heading] 105 | :param scores: (N) 106 | :param thresh: 107 | :return: 108 | """ 109 | assert boxes.shape[1] == 7 110 | order = scores.sort(0, descending=True)[1] 111 | 112 | boxes = boxes[order].contiguous() 113 | 114 | keep = torch.LongTensor(boxes.size(0)) 115 | num_out = iou3d_nms_cuda.nms_normal_gpu(boxes, keep, thresh) 116 | return order[keep[:num_out].cuda()].contiguous(), None 117 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/dense_heads/target_assigner/anchor_generator.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | class AnchorGenerator(object): 5 | def __init__(self, anchor_range, anchor_generator_config): 6 | super().__init__() 7 | self.anchor_generator_cfg = anchor_generator_config 8 | self.anchor_range = anchor_range 9 | self.anchor_sizes = [config['anchor_sizes'] for config in anchor_generator_config] 10 | self.anchor_rotations = [config['anchor_rotations'] for config in anchor_generator_config] 11 | self.anchor_heights = [config['anchor_bottom_heights'] for config in anchor_generator_config] 12 | self.align_center = [config.get('align_center', False) for config in anchor_generator_config] 13 | 14 | assert len(self.anchor_sizes) == len(self.anchor_rotations) == len(self.anchor_heights) 15 | self.num_of_anchor_sets = len(self.anchor_sizes) 16 | 17 | def generate_anchors(self, grid_sizes): 18 | assert len(grid_sizes) == self.num_of_anchor_sets 19 | all_anchors = [] 20 | num_anchors_per_location = [] 21 | for grid_size, anchor_size, anchor_rotation, anchor_height, align_center in zip( 22 | grid_sizes, self.anchor_sizes, self.anchor_rotations, self.anchor_heights, self.align_center): 23 | 24 | num_anchors_per_location.append(len(anchor_rotation) * len(anchor_size) * len(anchor_height)) 25 | if align_center: 26 | x_stride = (self.anchor_range[3] - self.anchor_range[0]) / grid_size[0] 27 | y_stride = (self.anchor_range[4] - self.anchor_range[1]) / grid_size[1] 28 | x_offset, y_offset = x_stride / 2, y_stride / 2 29 | else: 30 | x_stride = (self.anchor_range[3] - self.anchor_range[0]) / (grid_size[0] - 1) 31 | y_stride = (self.anchor_range[4] - self.anchor_range[1]) / (grid_size[1] - 1) 32 | x_offset, y_offset = 0, 0 33 | 34 | x_shifts = torch.arange( 35 | self.anchor_range[0] + x_offset, self.anchor_range[3] + 1e-5, step=x_stride, dtype=torch.float32, 36 | ).cuda() 37 | y_shifts = torch.arange( 38 | self.anchor_range[1] + y_offset, self.anchor_range[4] + 1e-5, step=y_stride, dtype=torch.float32, 39 | ).cuda() 40 | z_shifts = x_shifts.new_tensor(anchor_height) 41 | 42 | num_anchor_size, num_anchor_rotation = anchor_size.__len__(), anchor_rotation.__len__() 43 | anchor_rotation = x_shifts.new_tensor(anchor_rotation) 44 | anchor_size = x_shifts.new_tensor(anchor_size) 45 | x_shifts, y_shifts, z_shifts = torch.meshgrid([ 46 | x_shifts, y_shifts, z_shifts 47 | ]) # [x_grid, y_grid, z_grid] 48 | anchors = torch.stack((x_shifts, y_shifts, z_shifts), dim=-1) # [x, y, z, 3] 49 | anchors = anchors[:, :, :, None, :].repeat(1, 1, 1, anchor_size.shape[0], 1) 50 | anchor_size = anchor_size.view(1, 1, 1, -1, 3).repeat([*anchors.shape[0:3], 1, 1]) 51 | anchors = torch.cat((anchors, anchor_size), dim=-1) 52 | anchors = anchors[:, :, :, :, None, :].repeat(1, 1, 1, 1, num_anchor_rotation, 1) 53 | anchor_rotation = anchor_rotation.view(1, 1, 1, 1, -1, 1).repeat([*anchors.shape[0:3], num_anchor_size, 1, 1]) 54 | anchors = torch.cat((anchors, anchor_rotation), dim=-1) # [x, y, z, num_size, num_rot, 7] 55 | 56 | anchors = anchors.permute(2, 1, 0, 3, 4, 5).contiguous() 57 | #anchors = anchors.view(-1, anchors.shape[-1]) 58 | anchors[..., 2] += anchors[..., 5] / 2 # shift to box centers 59 | all_anchors.append(anchors) 60 | return all_anchors, num_anchors_per_location 61 | 62 | 63 | if __name__ == '__main__': 64 | from easydict import EasyDict 65 | config = [ 66 | EasyDict({ 67 | 'anchor_sizes': [[2.1, 4.7, 1.7], [0.86, 0.91, 1.73], [0.84, 1.78, 1.78]], 68 | 'anchor_rotations': [0, 1.57], 69 | 'anchor_heights': [0, 0.5] 70 | }) 71 | ] 72 | 73 | A = AnchorGenerator( 74 | anchor_range=[-75.2, -75.2, -2, 75.2, 75.2, 4], 75 | anchor_generator_config=config 76 | ) 77 | import pdb 78 | pdb.set_trace() 79 | A.generate_anchors([[188, 188]]) 80 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/datasets/augmentor/data_augmentor.py: -------------------------------------------------------------------------------- 1 | from functools import partial 2 | 3 | import numpy as np 4 | 5 | from ...utils import common_utils 6 | from . import augmentor_utils, database_sampler 7 | 8 | 9 | class DataAugmentor(object): 10 | def __init__(self, root_path, augmentor_configs, class_names, logger=None): 11 | self.root_path = root_path 12 | self.class_names = class_names 13 | self.logger = logger 14 | 15 | self.data_augmentor_queue = [] 16 | aug_config_list = augmentor_configs if isinstance(augmentor_configs, list) \ 17 | else augmentor_configs.AUG_CONFIG_LIST 18 | 19 | for cur_cfg in aug_config_list: 20 | if not isinstance(augmentor_configs, list): 21 | if cur_cfg.NAME in augmentor_configs.DISABLE_AUG_LIST: 22 | continue 23 | cur_augmentor = getattr(self, cur_cfg.NAME)(config=cur_cfg) 24 | self.data_augmentor_queue.append(cur_augmentor) 25 | 26 | def gt_sampling(self, config=None): 27 | db_sampler = database_sampler.DataBaseSampler( 28 | root_path=self.root_path, 29 | sampler_cfg=config, 30 | class_names=self.class_names, 31 | logger=self.logger 32 | ) 33 | return db_sampler 34 | 35 | def __getstate__(self): 36 | d = dict(self.__dict__) 37 | del d['logger'] 38 | return d 39 | 40 | def __setstate__(self, d): 41 | self.__dict__.update(d) 42 | 43 | def random_world_flip(self, data_dict=None, config=None): 44 | if data_dict is None: 45 | return partial(self.random_world_flip, config=config) 46 | gt_boxes, points = data_dict['gt_boxes'], data_dict['points'] 47 | for cur_axis in config['ALONG_AXIS_LIST']: 48 | assert cur_axis in ['x', 'y'] 49 | gt_boxes, points = getattr(augmentor_utils, 'random_flip_along_%s' % cur_axis)( 50 | gt_boxes, points, 51 | ) 52 | 53 | data_dict['gt_boxes'] = gt_boxes 54 | data_dict['points'] = points 55 | return data_dict 56 | 57 | def random_world_rotation(self, data_dict=None, config=None): 58 | if data_dict is None: 59 | return partial(self.random_world_rotation, config=config) 60 | rot_range = config['WORLD_ROT_ANGLE'] 61 | if not isinstance(rot_range, list): 62 | rot_range = [-rot_range, rot_range] 63 | gt_boxes, points = augmentor_utils.global_rotation( 64 | data_dict['gt_boxes'], data_dict['points'], rot_range=rot_range 65 | ) 66 | 67 | data_dict['gt_boxes'] = gt_boxes 68 | data_dict['points'] = points 69 | return data_dict 70 | 71 | def random_world_scaling(self, data_dict=None, config=None): 72 | if data_dict is None: 73 | return partial(self.random_world_scaling, config=config) 74 | gt_boxes, points = augmentor_utils.global_scaling( 75 | data_dict['gt_boxes'], data_dict['points'], config['WORLD_SCALE_RANGE'] 76 | ) 77 | data_dict['gt_boxes'] = gt_boxes 78 | data_dict['points'] = points 79 | return data_dict 80 | 81 | def forward(self, data_dict): 82 | """ 83 | Args: 84 | data_dict: 85 | points: (N, 3 + C_in) 86 | gt_boxes: optional, (N, 7) [x, y, z, dx, dy, dz, heading] 87 | gt_names: optional, (N), string 88 | ... 89 | 90 | Returns: 91 | """ 92 | for cur_augmentor in self.data_augmentor_queue: 93 | data_dict = cur_augmentor(data_dict=data_dict) 94 | 95 | data_dict['gt_boxes'][:, 6] = common_utils.limit_period( 96 | data_dict['gt_boxes'][:, 6], offset=0.5, period=2 * np.pi 97 | ) 98 | if 'calib' in data_dict: 99 | data_dict.pop('calib') 100 | if 'road_plane' in data_dict: 101 | data_dict.pop('road_plane') 102 | if 'gt_boxes_mask' in data_dict: 103 | gt_boxes_mask = data_dict['gt_boxes_mask'] 104 | data_dict['gt_boxes'] = data_dict['gt_boxes'][gt_boxes_mask] 105 | data_dict['gt_names'] = data_dict['gt_names'][gt_boxes_mask] 106 | data_dict.pop('gt_boxes_mask') 107 | return data_dict 108 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/pointnet2/pointnet2_stack/src/interpolate.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | Stacked-batch-data version of point interpolation, modified from the original implementation of official PointNet++ codes. 3 | Written by Shaoshuai Shi 4 | All Rights Reserved 2019-2020. 5 | */ 6 | 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | #include 16 | #include "interpolate_gpu.h" 17 | 18 | extern THCState *state; 19 | 20 | #define CHECK_CUDA(x) do { \ 21 | if (!x.type().is_cuda()) { \ 22 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 23 | exit(-1); \ 24 | } \ 25 | } while (0) 26 | #define CHECK_CONTIGUOUS(x) do { \ 27 | if (!x.is_contiguous()) { \ 28 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 29 | exit(-1); \ 30 | } \ 31 | } while (0) 32 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x) 33 | 34 | 35 | void three_nn_wrapper_stack(at::Tensor unknown_tensor, 36 | at::Tensor unknown_batch_cnt_tensor, at::Tensor known_tensor, 37 | at::Tensor known_batch_cnt_tensor, at::Tensor dist2_tensor, at::Tensor idx_tensor){ 38 | // unknown: (N1 + N2 ..., 3) 39 | // unknown_batch_cnt: (batch_size), [N1, N2, ...] 40 | // known: (M1 + M2 ..., 3) 41 | // known_batch_cnt: (batch_size), [M1, M2, ...] 42 | // Return: 43 | // dist: (N1 + N2 ..., 3) l2 distance to the three nearest neighbors 44 | // idx: (N1 + N2 ..., 3) index of the three nearest neighbors 45 | CHECK_INPUT(unknown_tensor); 46 | CHECK_INPUT(unknown_batch_cnt_tensor); 47 | CHECK_INPUT(known_tensor); 48 | CHECK_INPUT(known_batch_cnt_tensor); 49 | CHECK_INPUT(dist2_tensor); 50 | CHECK_INPUT(idx_tensor); 51 | 52 | int batch_size = unknown_batch_cnt_tensor.size(0); 53 | int N = unknown_tensor.size(0); 54 | int M = known_tensor.size(0); 55 | const float *unknown = unknown_tensor.data(); 56 | const int *unknown_batch_cnt = unknown_batch_cnt_tensor.data(); 57 | const float *known = known_tensor.data(); 58 | const int *known_batch_cnt = known_batch_cnt_tensor.data(); 59 | float *dist2 = dist2_tensor.data(); 60 | int *idx = idx_tensor.data(); 61 | 62 | three_nn_kernel_launcher_stack(batch_size, N, M, unknown, unknown_batch_cnt, known, known_batch_cnt, dist2, idx); 63 | } 64 | 65 | 66 | void three_interpolate_wrapper_stack(at::Tensor features_tensor, 67 | at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor out_tensor) { 68 | // features_tensor: (M1 + M2 ..., C) 69 | // idx_tensor: [N1 + N2 ..., 3] 70 | // weight_tensor: [N1 + N2 ..., 3] 71 | // Return: 72 | // out_tensor: (N1 + N2 ..., C) 73 | CHECK_INPUT(features_tensor); 74 | CHECK_INPUT(idx_tensor); 75 | CHECK_INPUT(weight_tensor); 76 | CHECK_INPUT(out_tensor); 77 | 78 | int N = out_tensor.size(0); 79 | int channels = features_tensor.size(1); 80 | const float *features = features_tensor.data(); 81 | const float *weight = weight_tensor.data(); 82 | const int *idx = idx_tensor.data(); 83 | float *out = out_tensor.data(); 84 | 85 | three_interpolate_kernel_launcher_stack(N, channels, features, idx, weight, out); 86 | } 87 | 88 | 89 | void three_interpolate_grad_wrapper_stack(at::Tensor grad_out_tensor, at::Tensor idx_tensor, 90 | at::Tensor weight_tensor, at::Tensor grad_features_tensor) { 91 | // grad_out_tensor: (N1 + N2 ..., C) 92 | // idx_tensor: [N1 + N2 ..., 3] 93 | // weight_tensor: [N1 + N2 ..., 3] 94 | // Return: 95 | // grad_features_tensor: (M1 + M2 ..., C) 96 | CHECK_INPUT(grad_out_tensor); 97 | CHECK_INPUT(idx_tensor); 98 | CHECK_INPUT(weight_tensor); 99 | CHECK_INPUT(grad_features_tensor); 100 | 101 | int N = grad_out_tensor.size(0); 102 | int channels = grad_out_tensor.size(1); 103 | const float *grad_out = grad_out_tensor.data(); 104 | const float *weight = weight_tensor.data(); 105 | const int *idx = idx_tensor.data(); 106 | float *grad_features = grad_features_tensor.data(); 107 | 108 | // printf("N=%d, channels=%d\n", N, channels); 109 | three_interpolate_grad_kernel_launcher_stack(N, channels, grad_out, idx, weight, grad_features); 110 | } -------------------------------------------------------------------------------- /OpenPCDet/pcdet/ops/roiaware_pool3d/roiaware_pool3d_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.autograd import Function 4 | 5 | from ...utils import common_utils 6 | from . import roiaware_pool3d_cuda 7 | 8 | 9 | def points_in_boxes_cpu(points, boxes): 10 | """ 11 | Args: 12 | points: (num_points, 3) 13 | boxes: [x, y, z, dx, dy, dz, heading], (x, y, z) is the box center, each box DO NOT overlaps 14 | Returns: 15 | point_indices: (N, num_points) 16 | """ 17 | assert boxes.shape[1] == 7 18 | assert points.shape[1] == 3 19 | points, is_numpy = common_utils.check_numpy_to_torch(points) 20 | boxes, is_numpy = common_utils.check_numpy_to_torch(boxes) 21 | 22 | point_indices = points.new_zeros((boxes.shape[0], points.shape[0]), dtype=torch.int) 23 | roiaware_pool3d_cuda.points_in_boxes_cpu(boxes.float().contiguous(), points.float().contiguous(), point_indices) 24 | 25 | return point_indices.numpy() if is_numpy else point_indices 26 | 27 | 28 | def points_in_boxes_gpu(points, boxes): 29 | """ 30 | :param points: (B, M, 3) 31 | :param boxes: (B, T, 7), num_valid_boxes <= T 32 | :return box_idxs_of_pts: (B, M), default background = -1 33 | """ 34 | assert boxes.shape[0] == points.shape[0] 35 | assert boxes.shape[2] == 7 and points.shape[2] == 3 36 | batch_size, num_points, _ = points.shape 37 | 38 | box_idxs_of_pts = points.new_zeros((batch_size, num_points), dtype=torch.int).fill_(-1) 39 | roiaware_pool3d_cuda.points_in_boxes_gpu(boxes.contiguous(), points.contiguous(), box_idxs_of_pts) 40 | 41 | return box_idxs_of_pts 42 | 43 | 44 | class RoIAwarePool3d(nn.Module): 45 | def __init__(self, out_size, max_pts_each_voxel=128): 46 | super().__init__() 47 | self.out_size = out_size 48 | self.max_pts_each_voxel = max_pts_each_voxel 49 | 50 | def forward(self, rois, pts, pts_feature, pool_method='max'): 51 | assert pool_method in ['max', 'avg'] 52 | return RoIAwarePool3dFunction.apply(rois, pts, pts_feature, self.out_size, self.max_pts_each_voxel, pool_method) 53 | 54 | 55 | class RoIAwarePool3dFunction(Function): 56 | @staticmethod 57 | def forward(ctx, rois, pts, pts_feature, out_size, max_pts_each_voxel, pool_method): 58 | """ 59 | Args: 60 | ctx: 61 | rois: (N, 7) [x, y, z, dx, dy, dz, heading] (x, y, z) is the box center 62 | pts: (npoints, 3) 63 | pts_feature: (npoints, C) 64 | out_size: int or tuple, like 7 or (7, 7, 7) 65 | max_pts_each_voxel: 66 | pool_method: 'max' or 'avg' 67 | 68 | Returns: 69 | pooled_features: (N, out_x, out_y, out_z, C) 70 | """ 71 | assert rois.shape[1] == 7 and pts.shape[1] == 3 72 | if isinstance(out_size, int): 73 | out_x = out_y = out_z = out_size 74 | else: 75 | assert len(out_size) == 3 76 | for k in range(3): 77 | assert isinstance(out_size[k], int) 78 | out_x, out_y, out_z = out_size 79 | 80 | num_rois = rois.shape[0] 81 | num_channels = pts_feature.shape[-1] 82 | num_pts = pts.shape[0] 83 | 84 | pooled_features = pts_feature.new_zeros((num_rois, out_x, out_y, out_z, num_channels)) 85 | argmax = pts_feature.new_zeros((num_rois, out_x, out_y, out_z, num_channels), dtype=torch.int) 86 | pts_idx_of_voxels = pts_feature.new_zeros((num_rois, out_x, out_y, out_z, max_pts_each_voxel), dtype=torch.int) 87 | 88 | pool_method_map = {'max': 0, 'avg': 1} 89 | pool_method = pool_method_map[pool_method] 90 | roiaware_pool3d_cuda.forward(rois, pts, pts_feature, argmax, pts_idx_of_voxels, pooled_features, pool_method) 91 | 92 | ctx.roiaware_pool3d_for_backward = (pts_idx_of_voxels, argmax, pool_method, num_pts, num_channels) 93 | return pooled_features 94 | 95 | @staticmethod 96 | def backward(ctx, grad_out): 97 | """ 98 | :param grad_out: (N, out_x, out_y, out_z, C) 99 | :return: 100 | grad_in: (npoints, C) 101 | """ 102 | pts_idx_of_voxels, argmax, pool_method, num_pts, num_channels = ctx.roiaware_pool3d_for_backward 103 | 104 | grad_in = grad_out.new_zeros((num_pts, num_channels)) 105 | roiaware_pool3d_cuda.backward(pts_idx_of_voxels, argmax, grad_out.contiguous(), grad_in, pool_method) 106 | 107 | return None, None, grad_in, None, None, None 108 | 109 | 110 | if __name__ == '__main__': 111 | pass 112 | -------------------------------------------------------------------------------- /scannet/batch_load_scannet_data.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | # Modified by Zhao Na, 2019 7 | 8 | """ Batch mode in loading Scannet scenes with vertices and ground truth labels 9 | for semantic and instance segmentations 10 | 11 | Usage example: python ./batch_load_scannet_data.py 12 | """ 13 | import os 14 | import sys 15 | import datetime 16 | import numpy as np 17 | from load_scannet_data import export 18 | # import pdb 19 | # from vtk_visualizer.plot3d import * 20 | # import cv2 21 | 22 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 23 | ROOT_DIR = os.path.dirname(BASE_DIR) 24 | sys.path.append(ROOT_DIR) 25 | sys.path.append(os.path.join(ROOT_DIR, 'utils')) 26 | from model_util_scannet import ScannetDatasetConfig 27 | DC = ScannetDatasetConfig() 28 | 29 | 30 | # SCANNET_DIR = './scannet/scans' 31 | SCANNET_DIR = 'scans' 32 | TRAIN_SCAN_NAMES = [line.rstrip() for line in open('meta_data/scannet_train.txt')] 33 | LABEL_MAP_FILE = 'meta_data/scannetv2-labels.combined.tsv' 34 | DONOTCARE_CLASS_IDS = np.array([]) 35 | OBJ_CLASS_IDS = np.array([3,4,5,6,7,8,9,10,11,12,14,16,24,28,33,34,36,39]) 36 | MAX_NUM_POINT = 50000 37 | # OUTPUT_FOLDER = './scannet/scannet_train_detection_data' 38 | OUTPUT_FOLDER = 'scannet_train_detection_data' 39 | 40 | X_RANGE = [-6., 6.] # (-left, right) in metres 41 | Y_RANGE = [-9., 9.] # (-behind, front) in metres 42 | RESOLUTION = 0.01 #Desired resolution in metres to use. Each output pixel will represent an square region res x res in size 43 | # INITIALIZE EMPTY ARRAY - of the dimensions we want 44 | BOUND_X = [int(X_RANGE[0]/RESOLUTION), int(X_RANGE[1]/RESOLUTION)] 45 | BOUND_Y = [int(Y_RANGE[0]/RESOLUTION), int(Y_RANGE[1]/RESOLUTION)] 46 | GIRD_SIZE = [BOUND_X[1]-BOUND_X[0], BOUND_Y[1]-BOUND_Y[0]] 47 | 48 | 49 | def export_one_scan(scan_name, output_filename_prefix): 50 | mesh_file = os.path.join(SCANNET_DIR, scan_name, scan_name + '_vh_clean_2.ply') 51 | agg_file = os.path.join(SCANNET_DIR, scan_name, scan_name + '.aggregation.json') 52 | seg_file = os.path.join(SCANNET_DIR, scan_name, scan_name + '_vh_clean_2.0.010000.segs.json') 53 | meta_file = os.path.join(SCANNET_DIR, scan_name, scan_name + '.txt') # includes axisAlignment info for the train set scans. 54 | mesh_vertices, semantic_labels, instance_labels, instance_bboxes, instance2semantic = \ 55 | export(mesh_file, agg_file, seg_file, meta_file, LABEL_MAP_FILE, None) 56 | 57 | mask = np.logical_not(np.in1d(semantic_labels, DONOTCARE_CLASS_IDS)) 58 | mesh_vertices = mesh_vertices[mask,:] 59 | semantic_labels = semantic_labels[mask] 60 | instance_labels = instance_labels[mask] 61 | 62 | num_instances = len(np.unique(instance_labels)) 63 | print('Num of instances: ', num_instances) 64 | 65 | bbox_mask = np.in1d(instance_bboxes[:,-1], OBJ_CLASS_IDS) 66 | instance_bboxes = instance_bboxes[bbox_mask,:] 67 | print('Num of care instances: ', instance_bboxes.shape[0]) 68 | 69 | N = mesh_vertices.shape[0] 70 | if N > MAX_NUM_POINT: 71 | choices = np.random.choice(N, MAX_NUM_POINT, replace=False) 72 | mesh_vertices = mesh_vertices[choices, :] 73 | semantic_labels = semantic_labels[choices] 74 | instance_labels = instance_labels[choices] 75 | 76 | np.save(output_filename_prefix+'_vert.npy', mesh_vertices) 77 | np.save(output_filename_prefix+'_sem_label.npy', semantic_labels) 78 | np.save(output_filename_prefix+'_ins_label.npy', instance_labels) 79 | np.save(output_filename_prefix+'_bbox.npy', instance_bboxes) 80 | 81 | def batch_export(): 82 | if not os.path.exists(OUTPUT_FOLDER): 83 | print('Creating new data folder: {}'.format(OUTPUT_FOLDER)) 84 | os.mkdir(OUTPUT_FOLDER) 85 | 86 | for scan_name in TRAIN_SCAN_NAMES: 87 | print('-'*20+'begin') 88 | print(datetime.datetime.now()) 89 | print(scan_name) 90 | output_filename_prefix = os.path.join(OUTPUT_FOLDER, scan_name) 91 | if os.path.isfile(output_filename_prefix+'_vert.npy'): 92 | print('File already exists. skipping.') 93 | print('-'*20+'done') 94 | continue 95 | try: 96 | export_one_scan(scan_name, output_filename_prefix) 97 | except: 98 | print('Failed export scan: %s'%(scan_name)) 99 | print('-'*20+'done') 100 | 101 | if __name__=='__main__': 102 | batch_export() 103 | -------------------------------------------------------------------------------- /OpenPCDet/tools/train_utils/optimization/learning_schedules_fastai.py: -------------------------------------------------------------------------------- 1 | # This file is modified from https://github.com/traveller59/second.pytorch 2 | 3 | import math 4 | from functools import partial 5 | 6 | import numpy as np 7 | import torch.optim.lr_scheduler as lr_sched 8 | 9 | from .fastai_optim import OptimWrapper 10 | 11 | 12 | class LRSchedulerStep(object): 13 | def __init__(self, fai_optimizer: OptimWrapper, total_step, lr_phases, 14 | mom_phases): 15 | # if not isinstance(fai_optimizer, OptimWrapper): 16 | # raise TypeError('{} is not a fastai OptimWrapper'.format( 17 | # type(fai_optimizer).__name__)) 18 | self.optimizer = fai_optimizer 19 | self.total_step = total_step 20 | self.lr_phases = [] 21 | 22 | for i, (start, lambda_func) in enumerate(lr_phases): 23 | if len(self.lr_phases) != 0: 24 | assert self.lr_phases[-1][0] < start 25 | if isinstance(lambda_func, str): 26 | lambda_func = eval(lambda_func) 27 | if i < len(lr_phases) - 1: 28 | self.lr_phases.append((int(start * total_step), int(lr_phases[i + 1][0] * total_step), lambda_func)) 29 | else: 30 | self.lr_phases.append((int(start * total_step), total_step, lambda_func)) 31 | assert self.lr_phases[0][0] == 0 32 | self.mom_phases = [] 33 | for i, (start, lambda_func) in enumerate(mom_phases): 34 | if len(self.mom_phases) != 0: 35 | assert self.mom_phases[-1][0] < start 36 | if isinstance(lambda_func, str): 37 | lambda_func = eval(lambda_func) 38 | if i < len(mom_phases) - 1: 39 | self.mom_phases.append((int(start * total_step), int(mom_phases[i + 1][0] * total_step), lambda_func)) 40 | else: 41 | self.mom_phases.append((int(start * total_step), total_step, lambda_func)) 42 | assert self.mom_phases[0][0] == 0 43 | 44 | def step(self, step): 45 | for start, end, func in self.lr_phases: 46 | if step >= start: 47 | self.optimizer.lr = func((step - start) / (end - start)) 48 | for start, end, func in self.mom_phases: 49 | if step >= start: 50 | self.optimizer.mom = func((step - start) / (end - start)) 51 | 52 | 53 | def annealing_cos(start, end, pct): 54 | # print(pct, start, end) 55 | "Cosine anneal from `start` to `end` as pct goes from 0.0 to 1.0." 56 | cos_out = np.cos(np.pi * pct) + 1 57 | return end + (start - end) / 2 * cos_out 58 | 59 | 60 | class OneCycle(LRSchedulerStep): 61 | def __init__(self, fai_optimizer, total_step, lr_max, moms, div_factor, 62 | pct_start): 63 | self.lr_max = lr_max 64 | self.moms = moms 65 | self.div_factor = div_factor 66 | self.pct_start = pct_start 67 | a1 = int(total_step * self.pct_start) 68 | a2 = total_step - a1 69 | low_lr = self.lr_max / self.div_factor 70 | lr_phases = ((0, partial(annealing_cos, low_lr, self.lr_max)), 71 | (self.pct_start, 72 | partial(annealing_cos, self.lr_max, low_lr / 1e4))) 73 | mom_phases = ((0, partial(annealing_cos, *self.moms)), 74 | (self.pct_start, partial(annealing_cos, 75 | *self.moms[::-1]))) 76 | fai_optimizer.lr, fai_optimizer.mom = low_lr, self.moms[0] 77 | super().__init__(fai_optimizer, total_step, lr_phases, mom_phases) 78 | 79 | 80 | class CosineWarmupLR(lr_sched._LRScheduler): 81 | def __init__(self, optimizer, T_max, eta_min=0, last_epoch=-1): 82 | self.T_max = T_max 83 | self.eta_min = eta_min 84 | super(CosineWarmupLR, self).__init__(optimizer, last_epoch) 85 | 86 | def get_lr(self): 87 | return [self.eta_min + (base_lr - self.eta_min) * 88 | (1 - math.cos(math.pi * self.last_epoch / self.T_max)) / 2 89 | for base_lr in self.base_lrs] 90 | 91 | 92 | class FakeOptim: 93 | def __init__(self): 94 | self.lr = 0 95 | self.mom = 0 96 | 97 | 98 | if __name__ == "__main__": 99 | import matplotlib.pyplot as plt 100 | 101 | opt = FakeOptim() # 3e-3, wd=0.4, div_factor=10 102 | schd = OneCycle(opt, 100, 3e-3, (0.95, 0.85), 10.0, 0.1) 103 | 104 | lrs = [] 105 | moms = [] 106 | for i in range(100): 107 | schd.step(i) 108 | lrs.append(opt.lr) 109 | moms.append(opt.mom) 110 | plt.plot(lrs) 111 | # plt.plot(moms) 112 | plt.show() 113 | plt.plot(moms) 114 | plt.show() 115 | -------------------------------------------------------------------------------- /scannet/model_util_scannet.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | # Modified by Yezhen Cong, 2020 7 | 8 | import numpy as np 9 | import sys 10 | import os 11 | 12 | import torch 13 | 14 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 15 | sys.path.append(BASE_DIR) 16 | ROOT_DIR = os.path.dirname(BASE_DIR) 17 | sys.path.append(os.path.join(ROOT_DIR, 'utils')) 18 | 19 | class ScannetDatasetConfig(object): 20 | def __init__(self): 21 | self.num_class = 18 22 | self.num_heading_bin = 1 23 | self.num_size_cluster = 18 24 | 25 | self.type2class = {'cabinet':0, 'bed':1, 'chair':2, 'sofa':3, 'table':4, 'door':5, 26 | 'window':6,'bookshelf':7,'picture':8, 'counter':9, 'desk':10, 'curtain':11, 27 | 'refrigerator':12, 'showercurtrain':13, 'toilet':14, 'sink':15, 'bathtub':16, 'garbagebin':17} 28 | self.class2type = {self.type2class[t]:t for t in self.type2class} 29 | self.nyu40ids = np.array([3,4,5,6,7,8,9,10,11,12,14,16,24,28,33,34,36,39]) 30 | self.nyu40id2class = {nyu40id: i for i,nyu40id in enumerate(list(self.nyu40ids))} 31 | self.mean_size_arr = np.load(os.path.join(ROOT_DIR,'scannet/meta_data/scannet_means.npz'))['arr_0'] 32 | self.type_mean_size = {} 33 | for i in range(self.num_size_cluster): 34 | self.type_mean_size[self.class2type[i]] = self.mean_size_arr[i,:] 35 | self.mean_size_arr_gpu = torch.from_numpy(self.mean_size_arr).cuda().float() 36 | 37 | def angle2class(self, angle): 38 | ''' Convert continuous angle to discrete class 39 | [optinal] also small regression number from 40 | class center angle to current angle. 41 | 42 | angle is from 0-2pi (or -pi~pi), class center at 0, 1*(2pi/N), 2*(2pi/N) ... (N-1)*(2pi/N) 43 | return is class of int32 of 0,1,...,N-1 and a number such that 44 | class*(2pi/N) + number = angle 45 | 46 | NOT USED. 47 | ''' 48 | assert(False) 49 | 50 | def class2angle_gpu(self, pred_cls, residual, to_label_format=True): 51 | ''' Inverse function to angle2class. 52 | 53 | As ScanNet only has axis-alined boxes so angles are always 0. ''' 54 | return torch.zeros(pred_cls.shape).cuda() # 0 55 | 56 | def class2size_gpu(self, pred_cls, residual): 57 | ''' Inverse function to size2class ''' 58 | return self.mean_size_arr_gpu[pred_cls, :] + residual 59 | 60 | def class2angle(self, pred_cls, residual, to_label_format=True): 61 | ''' Inverse function to angle2class. 62 | 63 | As ScanNet only has axis-alined boxes so angles are always 0. ''' 64 | return np.zeros(pred_cls.shape) # 0 65 | 66 | def size2class(self, size, type_name): 67 | ''' Convert 3D box size (l,w,h) to size class and size residual ''' 68 | size_class = self.type2class[type_name] 69 | size_residual = size - self.type_mean_size[type_name] 70 | return size_class, size_residual 71 | 72 | def class2size(self, pred_cls, residual): 73 | ''' Inverse function to size2class ''' 74 | return self.mean_size_arr[pred_cls, :] + residual 75 | 76 | def param2obb(self, center, heading_class, heading_residual, size_class, size_residual): 77 | heading_angle = self.class2angle(heading_class, heading_residual) 78 | box_size = self.class2size(int(size_class), size_residual) 79 | obb = np.zeros((7,)) 80 | obb[0:3] = center 81 | obb[3:6] = box_size 82 | obb[6] = heading_angle*-1 83 | return obb 84 | 85 | def rotate_aligned_boxes(input_boxes, rot_mat): 86 | centers, lengths = input_boxes[:,0:3], input_boxes[:,3:6] 87 | new_centers = np.dot(centers, np.transpose(rot_mat)) 88 | 89 | dx, dy = lengths[:,0]/2.0, lengths[:,1]/2.0 90 | new_x = np.zeros((dx.shape[0], 4)) 91 | new_y = np.zeros((dx.shape[0], 4)) 92 | 93 | for i, crnr in enumerate([(-1,-1), (1, -1), (1, 1), (-1, 1)]): 94 | crnrs = np.zeros((dx.shape[0], 3)) 95 | crnrs[:,0] = crnr[0]*dx 96 | crnrs[:,1] = crnr[1]*dy 97 | crnrs = np.dot(crnrs, np.transpose(rot_mat)) 98 | new_x[:,i] = crnrs[:,0] 99 | new_y[:,i] = crnrs[:,1] 100 | 101 | 102 | new_dx = 2.0*np.max(new_x, 1) 103 | new_dy = 2.0*np.max(new_y, 1) 104 | new_lengths = np.stack((new_dx, new_dy, lengths[:,2]), axis=1) 105 | 106 | return np.concatenate([new_centers, new_lengths], axis=1) 107 | -------------------------------------------------------------------------------- /OpenPCDet/pcdet/models/backbones_2d/base_bev_backbone.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | 5 | 6 | class BaseBEVBackbone(nn.Module): 7 | def __init__(self, model_cfg, input_channels): 8 | super().__init__() 9 | self.model_cfg = model_cfg 10 | 11 | if self.model_cfg.get('LAYER_NUMS', None) is not None: 12 | assert len(self.model_cfg.LAYER_NUMS) == len(self.model_cfg.LAYER_STRIDES) == len(self.model_cfg.NUM_FILTERS) 13 | layer_nums = self.model_cfg.LAYER_NUMS 14 | layer_strides = self.model_cfg.LAYER_STRIDES 15 | num_filters = self.model_cfg.NUM_FILTERS 16 | else: 17 | layer_nums = layer_strides = num_filters = [] 18 | 19 | if self.model_cfg.get('UPSAMPLE_STRIDES', None) is not None: 20 | assert len(self.model_cfg.UPSAMPLE_STRIDES) == len(self.model_cfg.NUM_UPSAMPLE_FILTERS) 21 | num_upsample_filters = self.model_cfg.NUM_UPSAMPLE_FILTERS 22 | upsample_strides = self.model_cfg.UPSAMPLE_STRIDES 23 | else: 24 | upsample_strides = num_upsample_filters = [] 25 | 26 | num_levels = len(layer_nums) 27 | c_in_list = [input_channels, *num_filters[:-1]] 28 | self.blocks = nn.ModuleList() 29 | self.deblocks = nn.ModuleList() 30 | for idx in range(num_levels): 31 | cur_layers = [ 32 | nn.ZeroPad2d(1), 33 | nn.Conv2d( 34 | c_in_list[idx], num_filters[idx], kernel_size=3, 35 | stride=layer_strides[idx], padding=0, bias=False 36 | ), 37 | nn.BatchNorm2d(num_filters[idx], eps=1e-3, momentum=0.01), 38 | nn.ReLU() 39 | ] 40 | for k in range(layer_nums[idx]): 41 | cur_layers.extend([ 42 | nn.Conv2d(num_filters[idx], num_filters[idx], kernel_size=3, padding=1, bias=False), 43 | nn.BatchNorm2d(num_filters[idx], eps=1e-3, momentum=0.01), 44 | nn.ReLU() 45 | ]) 46 | self.blocks.append(nn.Sequential(*cur_layers)) 47 | if len(upsample_strides) > 0: 48 | stride = upsample_strides[idx] 49 | if stride >= 1: 50 | self.deblocks.append(nn.Sequential( 51 | nn.ConvTranspose2d( 52 | num_filters[idx], num_upsample_filters[idx], 53 | upsample_strides[idx], 54 | stride=upsample_strides[idx], bias=False 55 | ), 56 | nn.BatchNorm2d(num_upsample_filters[idx], eps=1e-3, momentum=0.01), 57 | nn.ReLU() 58 | )) 59 | else: 60 | stride = np.round(1 / stride).astype(np.int) 61 | self.deblocks.append(nn.Sequential( 62 | nn.Conv2d( 63 | num_filters[idx], num_upsample_filters[idx], 64 | stride, 65 | stride=stride, bias=False 66 | ), 67 | nn.BatchNorm2d(num_upsample_filters[idx], eps=1e-3, momentum=0.01), 68 | nn.ReLU() 69 | )) 70 | 71 | c_in = sum(num_upsample_filters) 72 | if len(upsample_strides) > num_levels: 73 | self.deblocks.append(nn.Sequential( 74 | nn.ConvTranspose2d(c_in, c_in, upsample_strides[-1], stride=upsample_strides[-1], bias=False), 75 | nn.BatchNorm2d(c_in, eps=1e-3, momentum=0.01), 76 | nn.ReLU(), 77 | )) 78 | 79 | self.num_bev_features = c_in 80 | 81 | def forward(self, data_dict): 82 | """ 83 | Args: 84 | data_dict: 85 | spatial_features 86 | Returns: 87 | """ 88 | spatial_features = data_dict['spatial_features'] 89 | ups = [] 90 | ret_dict = {} 91 | x = spatial_features 92 | for i in range(len(self.blocks)): 93 | x = self.blocks[i](x) 94 | 95 | stride = int(spatial_features.shape[2] / x.shape[2]) 96 | ret_dict['spatial_features_%dx' % stride] = x 97 | if len(self.deblocks) > 0: 98 | ups.append(self.deblocks[i](x)) 99 | else: 100 | ups.append(x) 101 | 102 | if len(ups) > 1: 103 | x = torch.cat(ups, dim=1) 104 | elif len(ups) == 1: 105 | x = ups[0] 106 | 107 | if len(self.deblocks) > len(self.blocks): 108 | x = self.deblocks[-1](x) 109 | 110 | data_dict['spatial_features_2d'] = x 111 | 112 | return data_dict 113 | --------------------------------------------------------------------------------