├── pcdet ├── ops │ ├── __init__.py │ ├── iou3d_nms │ │ ├── __init__.py │ │ └── src │ │ │ ├── iou3d_cpu.h │ │ │ ├── iou3d_nms.h │ │ │ └── iou3d_nms_api.cpp │ ├── pointnet2 │ │ ├── __init__.py │ │ ├── pointnet2_batch │ │ │ ├── __init__.py │ │ │ └── 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 │ │ └── pointnet2_stack │ │ │ ├── __init__.py │ │ │ └── src │ │ │ ├── cuda_utils.h │ │ │ ├── voxel_query_gpu.h │ │ │ ├── sampling_gpu.h │ │ │ ├── ball_query_gpu.h │ │ │ ├── group_points_gpu.h │ │ │ ├── interpolate_gpu.h │ │ │ ├── voxel_query.cpp │ │ │ ├── ball_query.cpp │ │ │ ├── pointnet2_api.cpp │ │ │ ├── sampling.cpp │ │ │ ├── group_points.cpp │ │ │ ├── ball_query_gpu.cu │ │ │ └── vector_pool_gpu.h │ ├── roiaware_pool3d │ │ └── __init__.py │ ├── roipoint_pool3d │ │ ├── __init__.py │ │ ├── src │ │ │ └── roipoint_pool3d.cpp │ │ └── roipoint_pool3d_utils.py │ ├── bev_pool │ │ ├── __init__.py │ │ ├── bev_pool.py │ │ └── src │ │ │ └── bev_pool.cpp │ └── ingroup_inds │ │ ├── ingroup_inds_op.py │ │ └── src │ │ ├── error.cuh │ │ ├── ingroup_inds.cpp │ │ └── ingroup_inds_kernel.cu ├── utils │ ├── __init__.py │ ├── transform_utils.py │ ├── object3d_kitti.py │ └── object3d_custom.py ├── datasets │ ├── custom │ │ └── __init__.py │ ├── kitti │ │ ├── __init__.py │ │ ├── kitti_object_eval_python │ │ │ ├── __init__.py │ │ │ ├── evaluate.py │ │ │ ├── LICENSE │ │ │ └── README.md │ │ └── kitti_utils.py │ ├── lyft │ │ ├── __init__.py │ │ └── lyft_mAP_eval │ │ │ └── __init__.py │ ├── once │ │ ├── __init__.py │ │ └── once_eval │ │ │ └── eval_utils.py │ ├── waymo │ │ └── __init__.py │ ├── augmentor │ │ └── __init__.py │ ├── nuscenes │ │ └── __init__.py │ ├── pandaset │ │ └── __init__.py │ ├── processor │ │ ├── __init__.py │ │ └── point_feature_encoder.py │ ├── argo2 │ │ ├── __init__.py │ │ └── argo2_utils │ │ │ └── constants.py │ └── __init__.py ├── models │ ├── model_utils │ │ ├── __init__.py │ │ ├── basic_block_2d.py │ │ └── transfusion_utils.py │ ├── dense_heads │ │ ├── target_assigner │ │ │ ├── __init__.py │ │ │ └── anchor_generator.py │ │ ├── __init__.py │ │ ├── anchor_head_single.py │ │ └── point_head_simple.py │ ├── roi_heads │ │ ├── target_assigner │ │ │ └── __init__.py │ │ └── __init__.py │ ├── backbones_3d │ │ ├── vfe │ │ │ ├── image_vfe_modules │ │ │ │ ├── __init__.py │ │ │ │ ├── ffn │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── ddn_loss │ │ │ │ │ │ ├── __init__.py │ │ │ │ │ │ ├── balancer.py │ │ │ │ │ │ └── ddn_loss.py │ │ │ │ │ ├── ddn │ │ │ │ │ │ ├── __init__.py │ │ │ │ │ │ └── ddn_deeplabv3.py │ │ │ │ │ └── depth_ffn.py │ │ │ │ └── f2v │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── sampler.py │ │ │ │ │ └── frustum_to_voxel.py │ │ │ ├── vfe_template.py │ │ │ ├── __init__.py │ │ │ ├── mean_vfe.py │ │ │ ├── image_vfe.py │ │ │ └── dynamic_mean_vfe.py │ │ ├── pfe │ │ │ └── __init__.py │ │ ├── __init__.py │ │ └── focal_sparse_conv │ │ │ └── SemanticSeg │ │ │ ├── basic_blocks.py │ │ │ └── pyramid_ffn.py │ ├── backbones_2d │ │ ├── fuser │ │ │ ├── __init__.py │ │ │ └── convfuser.py │ │ ├── __init__.py │ │ └── map_to_bev │ │ │ ├── __init__.py │ │ │ ├── height_compression.py │ │ │ ├── conv2d_collapse.py │ │ │ └── pointpillar_scatter.py │ ├── backbones_image │ │ ├── __init__.py │ │ └── img_neck │ │ │ ├── __init__.py │ │ │ └── generalized_lss.py │ ├── view_transforms │ │ └── __init__.py │ ├── detectors │ │ ├── point_rcnn.py │ │ ├── second_net.py │ │ ├── pointpillar.py │ │ ├── PartA2_net.py │ │ ├── caddn.py │ │ ├── voxel_rcnn.py │ │ ├── pv_rcnn.py │ │ ├── voxelnext.py │ │ ├── voxelnext_kp.py │ │ ├── pillarnet.py │ │ ├── centerpoint.py │ │ ├── __init__.py │ │ ├── transfusion.py │ │ ├── pv_rcnn_plusplus.py │ │ └── bevfusion.py │ └── __init__.py ├── __init__.py └── config.py ├── docs ├── abl.png ├── arch.png ├── vis.png ├── results.png └── DEMO.md ├── .gitignore ├── tools ├── visual_utils │ ├── camera_pose.json │ └── camera_pose_2.json ├── cfgs │ ├── dataset_configs │ │ ├── waymo_dataset_multiframe.yaml │ │ ├── waymo_dataset.yaml │ │ ├── waymo_dataset_kp.yaml │ │ └── waymo_dataset_kp_augv2.yaml │ └── waymo_models │ │ └── kp_effv2next4_voxelnext_iou_aug_bev_channel.yaml ├── train_utils │ └── optimization │ │ └── __init__.py ├── test.py └── process_tools │ └── create_integrated_database.py └── README.md /pcdet/ops/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/datasets/custom/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/datasets/kitti/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/datasets/lyft/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/datasets/once/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/datasets/waymo/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/ops/iou3d_nms/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/ops/pointnet2/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/datasets/augmentor/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/datasets/nuscenes/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/datasets/pandaset/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/datasets/processor/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/models/model_utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/ops/roiaware_pool3d/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/ops/roipoint_pool3d/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/datasets/argo2/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | -------------------------------------------------------------------------------- /pcdet/datasets/lyft/lyft_mAP_eval/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/ops/pointnet2/pointnet2_batch/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/ops/pointnet2/pointnet2_stack/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/models/dense_heads/target_assigner/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/models/roi_heads/target_assigner/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/datasets/kitti/kitti_object_eval_python/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/models/backbones_3d/vfe/image_vfe_modules/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pcdet/ops/bev_pool/__init__.py: -------------------------------------------------------------------------------- 1 | from .bev_pool import bev_pool -------------------------------------------------------------------------------- /docs/abl.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shijianjian/VoxelKP/HEAD/docs/abl.png -------------------------------------------------------------------------------- /docs/arch.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shijianjian/VoxelKP/HEAD/docs/arch.png -------------------------------------------------------------------------------- /docs/vis.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shijianjian/VoxelKP/HEAD/docs/vis.png -------------------------------------------------------------------------------- /docs/results.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shijianjian/VoxelKP/HEAD/docs/results.png -------------------------------------------------------------------------------- /pcdet/models/backbones_2d/fuser/__init__.py: -------------------------------------------------------------------------------- 1 | from .convfuser import ConvFuser 2 | __all__ = { 3 | 'ConvFuser':ConvFuser 4 | } -------------------------------------------------------------------------------- /pcdet/models/backbones_image/__init__.py: -------------------------------------------------------------------------------- 1 | from .swin import SwinTransformer 2 | __all__ = { 3 | 'SwinTransformer':SwinTransformer, 4 | } -------------------------------------------------------------------------------- /pcdet/models/view_transforms/__init__.py: -------------------------------------------------------------------------------- 1 | from .depth_lss import DepthLSSTransform 2 | __all__ = { 3 | 'DepthLSSTransform': DepthLSSTransform, 4 | } -------------------------------------------------------------------------------- /pcdet/models/backbones_3d/vfe/image_vfe_modules/ffn/__init__.py: -------------------------------------------------------------------------------- 1 | from .depth_ffn import DepthFFN 2 | 3 | __all__ = { 4 | 'DepthFFN': DepthFFN 5 | } 6 | -------------------------------------------------------------------------------- /pcdet/models/backbones_3d/vfe/image_vfe_modules/ffn/ddn_loss/__init__.py: -------------------------------------------------------------------------------- 1 | from .ddn_loss import DDNLoss 2 | 3 | __all__ = { 4 | "DDNLoss": DDNLoss 5 | } 6 | -------------------------------------------------------------------------------- /pcdet/models/backbones_image/img_neck/__init__.py: -------------------------------------------------------------------------------- 1 | from .generalized_lss import GeneralizedLSSFPN 2 | __all__ = { 3 | 'GeneralizedLSSFPN':GeneralizedLSSFPN, 4 | } -------------------------------------------------------------------------------- /pcdet/models/backbones_3d/pfe/__init__.py: -------------------------------------------------------------------------------- 1 | from .voxel_set_abstraction import VoxelSetAbstraction 2 | 3 | __all__ = { 4 | 'VoxelSetAbstraction': VoxelSetAbstraction 5 | } 6 | -------------------------------------------------------------------------------- /pcdet/models/backbones_3d/vfe/image_vfe_modules/ffn/ddn/__init__.py: -------------------------------------------------------------------------------- 1 | from .ddn_deeplabv3 import DDNDeepLabV3 2 | 3 | __all__ = { 4 | 'DDNDeepLabV3': DDNDeepLabV3 5 | } 6 | -------------------------------------------------------------------------------- /pcdet/models/backbones_3d/vfe/image_vfe_modules/f2v/__init__.py: -------------------------------------------------------------------------------- 1 | from .frustum_to_voxel import FrustumToVoxel 2 | 3 | __all__ = { 4 | 'FrustumToVoxel': FrustumToVoxel 5 | } 6 | -------------------------------------------------------------------------------- /pcdet/datasets/argo2/argo2_utils/constants.py: -------------------------------------------------------------------------------- 1 | LABEL_ATTR = ( 2 | "tx_m", 3 | "ty_m", 4 | "tz_m", 5 | "length_m", 6 | "width_m", 7 | "height_m", 8 | "qw", 9 | "qx", 10 | "qy", 11 | "qz", 12 | ) -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | **__pycache__** 3 | **build** 4 | **egg-info** 5 | **dist** 6 | data/ 7 | *.pyc 8 | venv/ 9 | *.idea/ 10 | *.so 11 | *.sh 12 | *.pth 13 | *.pkl 14 | *.zip 15 | *.bin 16 | output 17 | version.py 18 | *.out 19 | *.err 20 | checkpoints -------------------------------------------------------------------------------- /pcdet/models/backbones_2d/__init__.py: -------------------------------------------------------------------------------- 1 | from .base_bev_backbone import BaseBEVBackbone, BaseBEVBackboneV1, BaseBEVResBackbone 2 | 3 | __all__ = { 4 | 'BaseBEVBackbone': BaseBEVBackbone, 5 | 'BaseBEVBackboneV1': BaseBEVBackboneV1, 6 | 'BaseBEVResBackbone': BaseBEVResBackbone, 7 | } 8 | -------------------------------------------------------------------------------- /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 | int boxes_aligned_iou_bev_cpu(at::Tensor boxes_a_tensor, at::Tensor boxes_b_tensor, at::Tensor ans_iou_tensor); 11 | #endif 12 | -------------------------------------------------------------------------------- /pcdet/models/backbones_2d/map_to_bev/__init__.py: -------------------------------------------------------------------------------- 1 | from .height_compression import HeightCompression 2 | from .pointpillar_scatter import PointPillarScatter, PointPillarScatter3d 3 | from .conv2d_collapse import Conv2DCollapse 4 | 5 | __all__ = { 6 | 'HeightCompression': HeightCompression, 7 | 'PointPillarScatter': PointPillarScatter, 8 | 'Conv2DCollapse': Conv2DCollapse, 9 | 'PointPillarScatter3d': PointPillarScatter3d, 10 | } 11 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /pcdet/models/backbones_3d/vfe/__init__.py: -------------------------------------------------------------------------------- 1 | from .mean_vfe import MeanVFE 2 | from .pillar_vfe import PillarVFE 3 | from .dynamic_mean_vfe import DynamicMeanVFE 4 | from .dynamic_pillar_vfe import DynamicPillarVFE, DynamicPillarVFESimple2D 5 | from .dynamic_voxel_vfe import DynamicVoxelVFE 6 | from .image_vfe import ImageVFE 7 | from .vfe_template import VFETemplate 8 | 9 | __all__ = { 10 | 'VFETemplate': VFETemplate, 11 | 'MeanVFE': MeanVFE, 12 | 'PillarVFE': PillarVFE, 13 | 'ImageVFE': ImageVFE, 14 | 'DynMeanVFE': DynamicMeanVFE, 15 | 'DynPillarVFE': DynamicPillarVFE, 16 | 'DynamicPillarVFESimple2D': DynamicPillarVFESimple2D, 17 | 'DynamicVoxelVFE': DynamicVoxelVFE, 18 | } 19 | -------------------------------------------------------------------------------- /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 | #include 9 | 10 | int boxes_aligned_overlap_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b, at::Tensor ans_overlap); 11 | int boxes_overlap_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b, at::Tensor ans_overlap); 12 | int paired_boxes_overlap_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b, at::Tensor ans_overlap); 13 | int boxes_iou_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b, at::Tensor ans_iou); 14 | int nms_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh); 15 | int nms_normal_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh); 16 | 17 | #endif 18 | -------------------------------------------------------------------------------- /pcdet/ops/ingroup_inds/ingroup_inds_op.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | try: 4 | from . import ingroup_inds_cuda 5 | # import ingroup_indices 6 | except ImportError: 7 | ingroup_indices = None 8 | print('Can not import ingroup indices') 9 | 10 | ingroup_indices = ingroup_inds_cuda 11 | 12 | from torch.autograd import Function 13 | class IngroupIndicesFunction(Function): 14 | 15 | @staticmethod 16 | def forward(ctx, group_inds): 17 | 18 | out_inds = torch.zeros_like(group_inds) - 1 19 | 20 | ingroup_indices.forward(group_inds, out_inds) 21 | 22 | ctx.mark_non_differentiable(out_inds) 23 | 24 | return out_inds 25 | 26 | @staticmethod 27 | def backward(ctx, g): 28 | 29 | return None 30 | 31 | ingroup_inds = IngroupIndicesFunction.apply -------------------------------------------------------------------------------- /pcdet/ops/pointnet2/pointnet2_stack/src/voxel_query_gpu.h: -------------------------------------------------------------------------------- 1 | #ifndef _STACK_VOXEL_QUERY_GPU_H 2 | #define _STACK_VOXEL_QUERY_GPU_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | 9 | int voxel_query_wrapper_stack(int M, int R1, int R2, int R3, int nsample, float radius, 10 | int z_range, int y_range, int x_range, at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, 11 | at::Tensor new_coords_tensor, at::Tensor point_indices_tensor, at::Tensor idx_tensor); 12 | 13 | 14 | void voxel_query_kernel_launcher_stack(int M, int R1, int R2, int R3, int nsample, 15 | float radius, int z_range, int y_range, int x_range, const float *new_xyz, 16 | const float *xyz, const int *new_coords, const int *point_indices, int *idx); 17 | 18 | 19 | #endif 20 | -------------------------------------------------------------------------------- /tools/visual_utils/camera_pose.json: -------------------------------------------------------------------------------- 1 | { 2 | "class_name" : "PinholeCameraParameters", 3 | "extrinsic" : 4 | [ 5 | -0.77394908176636501, 6 | 0.058152235847115118, 7 | -0.6305720706620157, 8 | 0.0, 9 | 0.63324003321685673, 10 | 0.076014705935911045, 11 | -0.77021349300891329, 12 | 0.0, 13 | 0.0031431138246869675, 14 | -0.99540950464994993, 15 | -0.095655835619300245, 16 | 0.0, 17 | -3.3618450725840994, 18 | 2.9274178750202857, 19 | 0.27275870653822759, 20 | 1.0 21 | ], 22 | "intrinsic" : 23 | { 24 | "height" : 1061, 25 | "intrinsic_matrix" : 26 | [ 27 | 918.85295341528945, 28 | 0.0, 29 | 0.0, 30 | 0.0, 31 | 918.85295341528945, 32 | 0.0, 33 | 959.5, 34 | 530.0, 35 | 1.0 36 | ], 37 | "width" : 1920 38 | }, 39 | "version_major" : 1, 40 | "version_minor" : 0 41 | } -------------------------------------------------------------------------------- /tools/visual_utils/camera_pose_2.json: -------------------------------------------------------------------------------- 1 | { 2 | "class_name" : "PinholeCameraParameters", 3 | "extrinsic" : 4 | [ 5 | -0.13096528263703272, 6 | -0.31409150324157692, 7 | 0.94031623528217823, 8 | 0.0, 9 | -0.99114204947364004, 10 | 0.020400942735956197, 11 | -0.13122971957859597, 12 | 0.0, 13 | 0.022034802222730174, 14 | -0.94917349790590488, 15 | -0.3139811114765893, 16 | 0.0, 17 | -0.77890921423199455, 18 | 6.7462791826161936, 19 | 7.6009766263140097, 20 | 1.0 21 | ], 22 | "intrinsic" : 23 | { 24 | "height" : 1061, 25 | "intrinsic_matrix" : 26 | [ 27 | 918.85295341528945, 28 | 0.0, 29 | 0.0, 30 | 0.0, 31 | 918.85295341528945, 32 | 0.0, 33 | 959.5, 34 | 530.0, 35 | 1.0 36 | ], 37 | "width" : 1920 38 | }, 39 | "version_major" : 1, 40 | "version_minor" : 0 41 | } -------------------------------------------------------------------------------- /pcdet/models/backbones_3d/vfe/image_vfe_modules/ffn/ddn/ddn_deeplabv3.py: -------------------------------------------------------------------------------- 1 | from .ddn_template import DDNTemplate 2 | 3 | try: 4 | import torchvision 5 | except: 6 | pass 7 | 8 | 9 | class DDNDeepLabV3(DDNTemplate): 10 | 11 | def __init__(self, backbone_name, **kwargs): 12 | """ 13 | Initializes DDNDeepLabV3 model 14 | Args: 15 | backbone_name: string, ResNet Backbone Name [ResNet50/ResNet101] 16 | """ 17 | if backbone_name == "ResNet50": 18 | constructor = torchvision.models.segmentation.deeplabv3_resnet50 19 | elif backbone_name == "ResNet101": 20 | constructor = torchvision.models.segmentation.deeplabv3_resnet101 21 | else: 22 | raise NotImplementedError 23 | 24 | super().__init__(constructor=constructor, **kwargs) 25 | -------------------------------------------------------------------------------- /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 farthest_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 farthest_point_sampling_kernel_launcher(int b, int n, int m, 13 | const float *dataset, float *temp, int *idxs); 14 | 15 | int stack_farthest_point_sampling_wrapper( 16 | at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor xyz_batch_cnt_tensor, 17 | at::Tensor idx_tensor, at::Tensor num_sampled_points_tensor); 18 | 19 | 20 | void stack_farthest_point_sampling_kernel_launcher(int N, int batch_size, 21 | const float *dataset, float *temp, int *xyz_batch_cnt, int *idxs, int *num_sampled_points); 22 | 23 | #endif 24 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /pcdet/ops/ingroup_inds/src/error.cuh: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | 4 | #define CHECK_CALL(call) \ 5 | do \ 6 | { \ 7 | const cudaError_t error_code = call; \ 8 | if (error_code != cudaSuccess) \ 9 | { \ 10 | printf("CUDA Error:\n"); \ 11 | printf(" File: %s\n", __FILE__); \ 12 | printf(" Line: %d\n", __LINE__); \ 13 | printf(" Error code: %d\n", error_code); \ 14 | printf(" Error text: %s\n", \ 15 | cudaGetErrorString(error_code)); \ 16 | exit(1); \ 17 | } \ 18 | } while (0) -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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_aligned_overlap_bev_gpu", &boxes_aligned_overlap_bev_gpu, "aligned oriented boxes overlap"); 13 | m.def("boxes_overlap_bev_gpu", &boxes_overlap_bev_gpu, "oriented boxes overlap"); 14 | m.def("paired_boxes_overlap_bev_gpu", &paired_boxes_overlap_bev_gpu, "oriented boxes overlap"); 15 | m.def("boxes_iou_bev_gpu", &boxes_iou_bev_gpu, "oriented boxes iou"); 16 | m.def("nms_gpu", &nms_gpu, "oriented nms gpu"); 17 | m.def("nms_normal_gpu", &nms_normal_gpu, "nms gpu"); 18 | m.def("boxes_aligned_iou_bev_cpu", &boxes_aligned_iou_bev_cpu, "aligned oriented boxes iou"); 19 | m.def("boxes_iou_bev_cpu", &boxes_iou_bev_cpu, "oriented boxes iou"); 20 | } 21 | -------------------------------------------------------------------------------- /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 .pvrcnn_head_kp import PVRCNNHeadKP 5 | from .pvrcnn_head_kp_v2 import PVRCNNHeadKPV2 6 | from .pvrcnn_head_kp_multihead import PVRCNNHeadKPMultihead 7 | from .second_head import SECONDHead 8 | from .voxelrcnn_head import VoxelRCNNHead 9 | from .roi_head_template import RoIHeadTemplate 10 | from .mppnet_head import MPPNetHead 11 | from .mppnet_memory_bank_e2e import MPPNetHeadE2E 12 | 13 | __all__ = { 14 | 'RoIHeadTemplate': RoIHeadTemplate, 15 | 'PartA2FCHead': PartA2FCHead, 16 | 'PVRCNNHead': PVRCNNHead, 17 | 'PVRCNNHeadKP': PVRCNNHeadKP, 18 | 'PVRCNNHeadKPV2': PVRCNNHeadKPV2, 19 | 'PVRCNNHeadKPMultihead': PVRCNNHeadKPMultihead, 20 | 'SECONDHead': SECONDHead, 21 | 'PointRCNNHead': PointRCNNHead, 22 | 'VoxelRCNNHead': VoxelRCNNHead, 23 | 'MPPNetHead': MPPNetHead, 24 | 'MPPNetHeadE2E': MPPNetHeadE2E, 25 | } 26 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 farthest_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 farthest_point_sampling_kernel_launcher(int b, int n, int m, 27 | const float *dataset, float *temp, int *idxs); 28 | 29 | #endif 30 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /pcdet/models/model_utils/basic_block_2d.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | 4 | class BasicBlock2D(nn.Module): 5 | 6 | def __init__(self, in_channels, out_channels, **kwargs): 7 | """ 8 | Initializes convolutional block 9 | Args: 10 | in_channels: int, Number of input channels 11 | out_channels: int, Number of output channels 12 | **kwargs: Dict, Extra arguments for nn.Conv2d 13 | """ 14 | super().__init__() 15 | self.in_channels = in_channels 16 | self.out_channels = out_channels 17 | self.conv = nn.Conv2d(in_channels=in_channels, 18 | out_channels=out_channels, 19 | **kwargs) 20 | self.bn = nn.BatchNorm2d(out_channels) 21 | self.relu = nn.ReLU(inplace=True) 22 | 23 | def forward(self, features): 24 | """ 25 | Applies convolutional block 26 | Args: 27 | features: (B, C_in, H, W), Input features 28 | Returns: 29 | x: (B, C_out, H, W), Output features 30 | """ 31 | x = self.conv(features) 32 | x = self.bn(x) 33 | x = self.relu(x) 34 | return x 35 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /pcdet/models/backbones_3d/__init__.py: -------------------------------------------------------------------------------- 1 | from .pointnet2_backbone import PointNet2Backbone, PointNet2MSG 2 | from .spconv_backbone import VoxelBackBone8x, VoxelResBackBone8x 3 | from .spconv_backbone_2d import PillarBackBone8x, PillarRes18BackBone8x 4 | from .spconv_backbone_focal import VoxelBackBone8xFocal 5 | from .spconv_backbone_voxelnext import VoxelResBackBone8xVoxelNeXt 6 | from .spconv_backbone_voxelnext_deBEV_effv2next4 import VoxelResBackBone8xVoxelNeXtEffv2Next4 7 | from .spconv_backbone_voxelnext2d import VoxelResBackBone8xVoxelNeXt2D 8 | from .spconv_unet import UNetV2 9 | from .dsvt import DSVT 10 | 11 | __all__ = { 12 | 'VoxelBackBone8x': VoxelBackBone8x, 13 | 'UNetV2': UNetV2, 14 | 'PointNet2Backbone': PointNet2Backbone, 15 | 'PointNet2MSG': PointNet2MSG, 16 | 'VoxelResBackBone8x': VoxelResBackBone8x, 17 | 'VoxelBackBone8xFocal': VoxelBackBone8xFocal, 18 | 'VoxelResBackBone8xVoxelNeXt': VoxelResBackBone8xVoxelNeXt, 19 | 'VoxelResBackBone8xVoxelNeXtEffv2Next4': VoxelResBackBone8xVoxelNeXtEffv2Next4, 20 | 'VoxelResBackBone8xVoxelNeXt2D': VoxelResBackBone8xVoxelNeXt2D, 21 | 'PillarBackBone8x': PillarBackBone8x, 22 | 'PillarRes18BackBone8x': PillarRes18BackBone8x, 23 | 'DSVT': DSVT, 24 | } 25 | -------------------------------------------------------------------------------- /pcdet/models/backbones_2d/fuser/convfuser.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | 4 | 5 | class ConvFuser(nn.Module): 6 | def __init__(self,model_cfg) -> None: 7 | super().__init__() 8 | self.model_cfg = model_cfg 9 | in_channel = self.model_cfg.IN_CHANNEL 10 | out_channel = self.model_cfg.OUT_CHANNEL 11 | self.conv = nn.Sequential( 12 | nn.Conv2d(in_channel, out_channel, 3, padding=1, bias=False), 13 | nn.BatchNorm2d(out_channel), 14 | nn.ReLU(True) 15 | ) 16 | 17 | def forward(self,batch_dict): 18 | """ 19 | Args: 20 | batch_dict: 21 | spatial_features_img (tensor): Bev features from image modality 22 | spatial_features (tensor): Bev features from lidar modality 23 | 24 | Returns: 25 | batch_dict: 26 | spatial_features (tensor): Bev features after muli-modal fusion 27 | """ 28 | img_bev = batch_dict['spatial_features_img'] 29 | lidar_bev = batch_dict['spatial_features'] 30 | cat_bev = torch.cat([img_bev,lidar_bev],dim=1) 31 | mm_bev = self.conv(cat_bev) 32 | batch_dict['spatial_features'] = mm_bev 33 | return batch_dict -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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("farthest_point_sampling_wrapper", &farthest_point_sampling_wrapper, "farthest_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 | -------------------------------------------------------------------------------- /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 | from .center_head import CenterHead 8 | from .center_head_kp import CenterHeadKP 9 | from .center_head_kp_naive import CenterHeadKPNaive 10 | from .center_head_kp_direct import CenterHeadKPDirect 11 | from .voxelnext_head import VoxelNeXtHead 12 | from .voxelnext_head_kp_merge import VoxelNeXtHeadKPMerge 13 | from .transfusion_head import TransFusionHead 14 | 15 | __all__ = { 16 | 'AnchorHeadTemplate': AnchorHeadTemplate, 17 | 'AnchorHeadSingle': AnchorHeadSingle, 18 | 'PointIntraPartOffsetHead': PointIntraPartOffsetHead, 19 | 'PointHeadSimple': PointHeadSimple, 20 | 'PointHeadBox': PointHeadBox, 21 | 'AnchorHeadMulti': AnchorHeadMulti, 22 | 'CenterHead': CenterHead, 23 | 'CenterHeadKP': CenterHeadKP, 24 | 'CenterHeadKPNaive': CenterHeadKPNaive, 25 | 'CenterHeadKPDirect': CenterHeadKPDirect, 26 | 'VoxelNeXtHead': VoxelNeXtHead, 27 | 'VoxelNeXtHeadKPMerge': VoxelNeXtHeadKPMerge, 28 | 'TransFusionHead': TransFusionHead, 29 | } 30 | -------------------------------------------------------------------------------- /pcdet/models/detectors/caddn.py: -------------------------------------------------------------------------------- 1 | from .detector3d_template import Detector3DTemplate 2 | 3 | 4 | class CaDDN(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_rpn = self.dense_head.get_loss() 28 | loss_depth, tb_dict_depth = self.vfe.get_loss() 29 | 30 | tb_dict = { 31 | 'loss_rpn': loss_rpn.item(), 32 | 'loss_depth': loss_depth.item(), 33 | **tb_dict_rpn, 34 | **tb_dict_depth 35 | } 36 | 37 | loss = loss_rpn + loss_depth 38 | return loss, tb_dict, disp_dict 39 | -------------------------------------------------------------------------------- /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 "group_points_gpu.h" 13 | 14 | 15 | int group_points_grad_wrapper_fast(int b, int c, int n, int npoints, int nsample, 16 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor) { 17 | 18 | float *grad_points = grad_points_tensor.data(); 19 | const int *idx = idx_tensor.data(); 20 | const float *grad_out = grad_out_tensor.data(); 21 | 22 | group_points_grad_kernel_launcher_fast(b, c, n, npoints, nsample, grad_out, idx, grad_points); 23 | return 1; 24 | } 25 | 26 | 27 | int group_points_wrapper_fast(int b, int c, int n, int npoints, int nsample, 28 | at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor) { 29 | 30 | const float *points = points_tensor.data(); 31 | const int *idx = idx_tensor.data(); 32 | float *out = out_tensor.data(); 33 | 34 | group_points_kernel_launcher_fast(b, c, n, npoints, nsample, points, idx, out); 35 | return 1; 36 | } 37 | -------------------------------------------------------------------------------- /pcdet/ops/ingroup_inds/src/ingroup_inds.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | 6 | #define CHECK_CUDA(x) \ 7 | TORCH_CHECK(x.device().is_cuda(), #x, " must be a CUDAtensor ") 8 | #define CHECK_CONTIGUOUS(x) \ 9 | TORCH_CHECK(x.is_contiguous(), #x, " must be contiguous ") 10 | #define CHECK_INPUT(x) \ 11 | CHECK_CUDA(x); \ 12 | CHECK_CONTIGUOUS(x) 13 | 14 | 15 | void ingroup_inds_launcher( 16 | const long *group_inds_data, 17 | long *out_inds_data, 18 | int N, 19 | int max_group_id 20 | ); 21 | 22 | 23 | void ingroup_inds_gpu( 24 | at::Tensor group_inds, 25 | at::Tensor out_inds 26 | ); 27 | 28 | void ingroup_inds_gpu( 29 | at::Tensor group_inds, 30 | at::Tensor out_inds 31 | ) { 32 | 33 | CHECK_INPUT(group_inds); 34 | CHECK_INPUT(out_inds); 35 | int N = group_inds.size(0); 36 | int max_group_id = group_inds.max().item().toLong(); 37 | 38 | 39 | long *group_inds_data = group_inds.data_ptr(); 40 | long *out_inds_data = out_inds.data_ptr(); 41 | 42 | ingroup_inds_launcher( 43 | group_inds_data, 44 | out_inds_data, 45 | N, 46 | max_group_id 47 | ); 48 | 49 | } 50 | 51 | 52 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 53 | m.def("forward", &ingroup_inds_gpu, "cuda version of get_inner_win_inds of SST"); 54 | } -------------------------------------------------------------------------------- /pcdet/models/detectors/voxel_rcnn.py: -------------------------------------------------------------------------------- 1 | from .detector3d_template import Detector3DTemplate 2 | 3 | 4 | class VoxelRCNN(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 = 0 27 | 28 | loss_rpn, tb_dict = self.dense_head.get_loss() 29 | loss_rcnn, tb_dict = self.roi_head.get_loss(tb_dict) 30 | 31 | loss = loss + loss_rpn + loss_rcnn 32 | 33 | if hasattr(self.backbone_3d, 'get_loss'): 34 | loss_backbone3d, tb_dict = self.backbone_3d.get_loss(tb_dict) 35 | loss += loss_backbone3d 36 | 37 | return loss, tb_dict, disp_dict 38 | -------------------------------------------------------------------------------- /pcdet/models/backbones_3d/vfe/image_vfe_modules/f2v/sampler.py: -------------------------------------------------------------------------------- 1 | from functools import partial 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | 7 | 8 | class Sampler(nn.Module): 9 | 10 | def __init__(self, mode="bilinear", padding_mode="zeros"): 11 | """ 12 | Initializes module 13 | Args: 14 | mode: string, Sampling mode [bilinear/nearest] 15 | padding_mode: string, Padding mode for outside grid values [zeros/border/reflection] 16 | """ 17 | super().__init__() 18 | self.mode = mode 19 | self.padding_mode = padding_mode 20 | 21 | if torch.__version__ >= '1.3': 22 | self.grid_sample = partial(F.grid_sample, align_corners=True) 23 | else: 24 | self.grid_sample = F.grid_sample 25 | 26 | def forward(self, input_features, grid): 27 | """ 28 | Samples input using sampling grid 29 | Args: 30 | input_features: (B, C, D, H, W), Input frustum features 31 | grid: (B, X, Y, Z, 3), Sampling grids for input features 32 | Returns 33 | output_features: (B, C, X, Y, Z) Output voxel features 34 | """ 35 | # Sample from grid 36 | output = self.grid_sample(input=input_features, grid=grid, mode=self.mode, padding_mode=self.padding_mode) 37 | return output 38 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 -------------------------------------------------------------------------------- /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 "ball_query_gpu.h" 13 | 14 | #define CHECK_CUDA(x) do { \ 15 | if (!x.type().is_cuda()) { \ 16 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 17 | exit(-1); \ 18 | } \ 19 | } while (0) 20 | #define CHECK_CONTIGUOUS(x) do { \ 21 | if (!x.is_contiguous()) { \ 22 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 23 | exit(-1); \ 24 | } \ 25 | } while (0) 26 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x) 27 | 28 | 29 | int ball_query_wrapper_fast(int b, int n, int m, float radius, int nsample, 30 | at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor idx_tensor) { 31 | CHECK_INPUT(new_xyz_tensor); 32 | CHECK_INPUT(xyz_tensor); 33 | const float *new_xyz = new_xyz_tensor.data(); 34 | const float *xyz = xyz_tensor.data(); 35 | int *idx = idx_tensor.data(); 36 | 37 | ball_query_kernel_launcher_fast(b, n, m, radius, nsample, new_xyz, xyz, idx); 38 | return 1; 39 | } 40 | -------------------------------------------------------------------------------- /pcdet/models/backbones_2d/map_to_bev/conv2d_collapse.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from pcdet.models.model_utils.basic_block_2d import BasicBlock2D 5 | 6 | 7 | class Conv2DCollapse(nn.Module): 8 | 9 | def __init__(self, model_cfg, grid_size): 10 | """ 11 | Initializes 2D convolution collapse module 12 | Args: 13 | model_cfg: EasyDict, Model configuration 14 | grid_size: (X, Y, Z) Voxel grid size 15 | """ 16 | super().__init__() 17 | self.model_cfg = model_cfg 18 | self.num_heights = grid_size[-1] 19 | self.num_bev_features = self.model_cfg.NUM_BEV_FEATURES 20 | self.block = BasicBlock2D(in_channels=self.num_bev_features * self.num_heights, 21 | out_channels=self.num_bev_features, 22 | **self.model_cfg.ARGS) 23 | 24 | def forward(self, batch_dict): 25 | """ 26 | Collapses voxel features to BEV via concatenation and channel reduction 27 | Args: 28 | batch_dict: 29 | voxel_features: (B, C, Z, Y, X), Voxel feature representation 30 | Returns: 31 | batch_dict: 32 | spatial_features: (B, C, Y, X), BEV feature representation 33 | """ 34 | voxel_features = batch_dict["voxel_features"] 35 | bev_features = voxel_features.flatten(start_dim=1, end_dim=2) # (B, C, Z, Y, X) -> (B, C*Z, Y, X) 36 | bev_features = self.block(bev_features) # (B, C*Z, Y, X) -> (B, C, Y, X) 37 | batch_dict["spatial_features"] = bev_features 38 | return batch_dict 39 | -------------------------------------------------------------------------------- /pcdet/ops/pointnet2/pointnet2_stack/src/voxel_query.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include "voxel_query_gpu.h" 9 | 10 | #define CHECK_CUDA(x) do { \ 11 | if (!x.type().is_cuda()) { \ 12 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 13 | exit(-1); \ 14 | } \ 15 | } while (0) 16 | #define CHECK_CONTIGUOUS(x) do { \ 17 | if (!x.is_contiguous()) { \ 18 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 19 | exit(-1); \ 20 | } \ 21 | } while (0) 22 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x) 23 | 24 | 25 | int voxel_query_wrapper_stack(int M, int R1, int R2, int R3, int nsample, float radius, 26 | int z_range, int y_range, int x_range, at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, 27 | at::Tensor new_coords_tensor, at::Tensor point_indices_tensor, at::Tensor idx_tensor) { 28 | CHECK_INPUT(new_coords_tensor); 29 | CHECK_INPUT(point_indices_tensor); 30 | CHECK_INPUT(new_xyz_tensor); 31 | CHECK_INPUT(xyz_tensor); 32 | 33 | const float *new_xyz = new_xyz_tensor.data(); 34 | const float *xyz = xyz_tensor.data(); 35 | const int *new_coords = new_coords_tensor.data(); 36 | const int *point_indices = point_indices_tensor.data(); 37 | int *idx = idx_tensor.data(); 38 | 39 | voxel_query_kernel_launcher_stack(M, R1, R2, R3, nsample, radius, z_range, y_range, x_range, new_xyz, xyz, new_coords, point_indices, idx); 40 | return 1; 41 | } 42 | -------------------------------------------------------------------------------- /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 | print("=========================") 12 | print("before", cur_module.__class__.__name__, batch_dict.keys()) 13 | batch_dict = cur_module(batch_dict) 14 | print("after", cur_module.__class__.__name__, batch_dict.keys()) 15 | print("=========================") 16 | 17 | if self.training: 18 | loss, tb_dict, disp_dict = self.get_training_loss() 19 | 20 | ret_dict = { 21 | 'loss': loss 22 | } 23 | return ret_dict, tb_dict, disp_dict 24 | else: 25 | pred_dicts, recall_dicts = self.post_processing(batch_dict) 26 | return pred_dicts, recall_dicts 27 | 28 | def get_training_loss(self): 29 | disp_dict = {} 30 | loss_rpn, tb_dict = self.dense_head.get_loss() 31 | loss_point, tb_dict = self.point_head.get_loss(tb_dict) 32 | loss_rcnn, tb_dict = self.roi_head.get_loss(tb_dict) 33 | 34 | loss = loss_rpn + loss_point + loss_rcnn 35 | 36 | if hasattr(self.backbone_3d, 'get_loss'): 37 | loss_backbone3d, tb_dict = self.backbone_3d.get_loss(tb_dict) 38 | loss += loss_backbone3d 39 | 40 | return loss, tb_dict, disp_dict 41 | -------------------------------------------------------------------------------- /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 "sampling_gpu.h" 12 | 13 | 14 | int gather_points_wrapper_fast(int b, int c, int n, int npoints, 15 | at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor){ 16 | const float *points = points_tensor.data(); 17 | const int *idx = idx_tensor.data(); 18 | float *out = out_tensor.data(); 19 | 20 | gather_points_kernel_launcher_fast(b, c, n, npoints, points, idx, out); 21 | return 1; 22 | } 23 | 24 | 25 | int gather_points_grad_wrapper_fast(int b, int c, int n, int npoints, 26 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor) { 27 | 28 | const float *grad_out = grad_out_tensor.data(); 29 | const int *idx = idx_tensor.data(); 30 | float *grad_points = grad_points_tensor.data(); 31 | 32 | gather_points_grad_kernel_launcher_fast(b, c, n, npoints, grad_out, idx, grad_points); 33 | return 1; 34 | } 35 | 36 | 37 | int farthest_point_sampling_wrapper(int b, int n, int m, 38 | at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor idx_tensor) { 39 | 40 | const float *points = points_tensor.data(); 41 | float *temp = temp_tensor.data(); 42 | int *idx = idx_tensor.data(); 43 | 44 | farthest_point_sampling_kernel_launcher(b, n, m, points, temp, idx); 45 | return 1; 46 | } 47 | -------------------------------------------------------------------------------- /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 "ball_query_gpu.h" 13 | 14 | #define CHECK_CUDA(x) do { \ 15 | if (!x.type().is_cuda()) { \ 16 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 17 | exit(-1); \ 18 | } \ 19 | } while (0) 20 | #define CHECK_CONTIGUOUS(x) do { \ 21 | if (!x.is_contiguous()) { \ 22 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 23 | exit(-1); \ 24 | } \ 25 | } while (0) 26 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x) 27 | 28 | 29 | int ball_query_wrapper_stack(int B, int M, float radius, int nsample, 30 | at::Tensor new_xyz_tensor, at::Tensor new_xyz_batch_cnt_tensor, 31 | at::Tensor xyz_tensor, at::Tensor xyz_batch_cnt_tensor, at::Tensor idx_tensor) { 32 | CHECK_INPUT(new_xyz_tensor); 33 | CHECK_INPUT(xyz_tensor); 34 | CHECK_INPUT(new_xyz_batch_cnt_tensor); 35 | CHECK_INPUT(xyz_batch_cnt_tensor); 36 | 37 | const float *new_xyz = new_xyz_tensor.data(); 38 | const float *xyz = xyz_tensor.data(); 39 | const int *new_xyz_batch_cnt = new_xyz_batch_cnt_tensor.data(); 40 | const int *xyz_batch_cnt = xyz_batch_cnt_tensor.data(); 41 | int *idx = idx_tensor.data(); 42 | 43 | ball_query_kernel_launcher_stack(B, M, radius, nsample, new_xyz, new_xyz_batch_cnt, xyz, xyz_batch_cnt, idx); 44 | return 1; 45 | } 46 | -------------------------------------------------------------------------------- /pcdet/models/detectors/voxelnext.py: -------------------------------------------------------------------------------- 1 | from .detector3d_template import Detector3DTemplate 2 | 3 | 4 | class VoxelNeXt(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 | 11 | for cur_module in self.module_list: 12 | batch_dict = cur_module(batch_dict) 13 | 14 | if self.training: 15 | loss, tb_dict, disp_dict = self.get_training_loss() 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 | 26 | disp_dict = {} 27 | loss, tb_dict = self.dense_head.get_loss() 28 | 29 | return loss, tb_dict, disp_dict 30 | 31 | def post_processing(self, batch_dict): 32 | post_process_cfg = self.model_cfg.POST_PROCESSING 33 | batch_size = batch_dict['batch_size'] 34 | final_pred_dict = batch_dict['final_box_dicts'] 35 | recall_dict = {} 36 | for index in range(batch_size): 37 | pred_boxes = final_pred_dict[index]['pred_boxes'] 38 | 39 | recall_dict = self.generate_recall_record( 40 | box_preds=pred_boxes, 41 | recall_dict=recall_dict, batch_index=index, data_dict=batch_dict, 42 | thresh_list=post_process_cfg.RECALL_THRESH_LIST 43 | ) 44 | 45 | return final_pred_dict, recall_dict 46 | -------------------------------------------------------------------------------- /pcdet/models/detectors/voxelnext_kp.py: -------------------------------------------------------------------------------- 1 | from .detector3d_template import Detector3DTemplate 2 | 3 | 4 | class VoxelNeXt_KP(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 | 11 | for cur_module in self.module_list: 12 | batch_dict = cur_module(batch_dict) 13 | 14 | if self.training: 15 | loss, tb_dict, disp_dict = self.get_training_loss() 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 | 26 | disp_dict = {} 27 | loss, tb_dict = self.dense_head.get_loss() 28 | 29 | return loss, tb_dict, disp_dict 30 | 31 | def post_processing(self, batch_dict): 32 | post_process_cfg = self.model_cfg.POST_PROCESSING 33 | batch_size = batch_dict['batch_size'] 34 | final_pred_dict = batch_dict['final_box_dicts'] 35 | recall_dict = {} 36 | for index in range(batch_size): 37 | pred_boxes = final_pred_dict[index]['pred_boxes'] 38 | 39 | recall_dict = self.generate_recall_record( 40 | box_preds=pred_boxes, 41 | recall_dict=recall_dict, batch_index=index, data_dict=batch_dict, 42 | thresh_list=post_process_cfg.RECALL_THRESH_LIST 43 | ) 44 | 45 | return final_pred_dict, recall_dict 46 | -------------------------------------------------------------------------------- /pcdet/datasets/once/once_eval/eval_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def compute_split_parts(num_samples, num_parts): 5 | part_samples = num_samples // num_parts 6 | remain_samples = num_samples % num_parts 7 | if part_samples == 0: 8 | return [num_samples] 9 | if remain_samples == 0: 10 | return [part_samples] * num_parts 11 | else: 12 | return [part_samples] * num_parts + [remain_samples] 13 | 14 | 15 | def overall_filter(boxes): 16 | ignore = np.zeros(boxes.shape[0], dtype=bool) # all false 17 | return ignore 18 | 19 | 20 | def distance_filter(boxes, level): 21 | ignore = np.ones(boxes.shape[0], dtype=bool) # all true 22 | dist = np.sqrt(np.sum(boxes[:, 0:3] * boxes[:, 0:3], axis=1)) 23 | 24 | if level == 0: # 0-30m 25 | flag = dist < 30 26 | elif level == 1: # 30-50m 27 | flag = (dist >= 30) & (dist < 50) 28 | elif level == 2: # 50m-inf 29 | flag = dist >= 50 30 | else: 31 | assert False, 'level < 3 for distance metric, found level %s' % (str(level)) 32 | 33 | ignore[flag] = False 34 | return ignore 35 | 36 | 37 | def overall_distance_filter(boxes, level): 38 | ignore = np.ones(boxes.shape[0], dtype=bool) # all true 39 | dist = np.sqrt(np.sum(boxes[:, 0:3] * boxes[:, 0:3], axis=1)) 40 | 41 | if level == 0: 42 | flag = np.ones(boxes.shape[0], dtype=bool) 43 | elif level == 1: # 0-30m 44 | flag = dist < 30 45 | elif level == 2: # 30-50m 46 | flag = (dist >= 30) & (dist < 50) 47 | elif level == 3: # 50m-inf 48 | flag = dist >= 50 49 | else: 50 | assert False, 'level < 4 for overall & distance metric, found level %s' % (str(level)) 51 | 52 | ignore[flag] = False 53 | return ignore -------------------------------------------------------------------------------- /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 | try: 9 | import kornia 10 | except: 11 | pass 12 | # print('Warning: kornia is not installed. This package is only required by CaDDN') 13 | 14 | 15 | 16 | def build_network(model_cfg, num_class, dataset): 17 | model = build_detector( 18 | model_cfg=model_cfg, num_class=num_class, dataset=dataset 19 | ) 20 | return model 21 | 22 | 23 | def load_data_to_gpu(batch_dict): 24 | for key, val in batch_dict.items(): 25 | if key == 'camera_imgs': 26 | batch_dict[key] = val.cuda() 27 | elif not isinstance(val, np.ndarray): 28 | continue 29 | elif key in ['frame_id', 'metadata', 'calib', 'image_paths','ori_shape','img_process_infos']: 30 | continue 31 | elif key in ['images']: 32 | batch_dict[key] = kornia.image_to_tensor(val).float().cuda().contiguous() 33 | elif key in ['image_shape']: 34 | batch_dict[key] = torch.from_numpy(val).int().cuda() 35 | else: 36 | batch_dict[key] = torch.from_numpy(val).float().cuda() 37 | 38 | 39 | def model_fn_decorator(): 40 | ModelReturn = namedtuple('ModelReturn', ['loss', 'tb_dict', 'disp_dict']) 41 | 42 | def model_func(model, batch_dict): 43 | load_data_to_gpu(batch_dict) 44 | ret_dict, tb_dict, disp_dict = model(batch_dict) 45 | 46 | loss = ret_dict['loss'].mean() 47 | if hasattr(model, 'update_global_step'): 48 | model.update_global_step() 49 | else: 50 | model.module.update_global_step() 51 | 52 | return ModelReturn(loss, tb_dict, disp_dict) 53 | 54 | return model_func 55 | -------------------------------------------------------------------------------- /pcdet/models/detectors/pillarnet.py: -------------------------------------------------------------------------------- 1 | from .detector3d_template import Detector3DTemplate 2 | 3 | 4 | class PillarNet(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 | 36 | def post_processing(self, batch_dict): 37 | post_process_cfg = self.model_cfg.POST_PROCESSING 38 | batch_size = batch_dict['batch_size'] 39 | final_pred_dict = batch_dict['final_box_dicts'] 40 | recall_dict = {} 41 | for index in range(batch_size): 42 | pred_boxes = final_pred_dict[index]['pred_boxes'] 43 | 44 | recall_dict = self.generate_recall_record( 45 | box_preds=pred_boxes, 46 | recall_dict=recall_dict, batch_index=index, data_dict=batch_dict, 47 | thresh_list=post_process_cfg.RECALL_THRESH_LIST 48 | ) 49 | 50 | return final_pred_dict, recall_dict -------------------------------------------------------------------------------- /pcdet/models/detectors/centerpoint.py: -------------------------------------------------------------------------------- 1 | from .detector3d_template import Detector3DTemplate 2 | 3 | 4 | class CenterPoint(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 | 36 | def post_processing(self, batch_dict): 37 | post_process_cfg = self.model_cfg.POST_PROCESSING 38 | batch_size = batch_dict['batch_size'] 39 | final_pred_dict = batch_dict['final_box_dicts'] 40 | recall_dict = {} 41 | for index in range(batch_size): 42 | pred_boxes = final_pred_dict[index]['pred_boxes'] 43 | 44 | recall_dict = self.generate_recall_record( 45 | box_preds=pred_boxes, 46 | recall_dict=recall_dict, batch_index=index, data_dict=batch_dict, 47 | thresh_list=post_process_cfg.RECALL_THRESH_LIST 48 | ) 49 | 50 | return final_pred_dict, recall_dict -------------------------------------------------------------------------------- /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 .pv_rcnn_kp import PVRCNN_KP 7 | from .second_net import SECONDNet 8 | from .second_net_iou import SECONDNetIoU 9 | from .caddn import CaDDN 10 | from .voxel_rcnn import VoxelRCNN 11 | from .centerpoint import CenterPoint 12 | from .centerpoint_kp import CenterPointKP 13 | from .pv_rcnn_plusplus import PVRCNNPlusPlus 14 | from .pv_rcnn_plusplus_kp import PVRCNNPlusPlus_KP 15 | from .mppnet import MPPNet 16 | from .mppnet_e2e import MPPNetE2E 17 | from .pillarnet import PillarNet 18 | from .voxelnext import VoxelNeXt 19 | from .voxelnext_kp import VoxelNeXt_KP 20 | from .transfusion import TransFusion 21 | from .bevfusion import BevFusion 22 | 23 | __all__ = { 24 | 'Detector3DTemplate': Detector3DTemplate, 25 | 'SECONDNet': SECONDNet, 26 | 'PartA2Net': PartA2Net, 27 | 'PVRCNN': PVRCNN, 28 | 'PVRCNN_KP': PVRCNN_KP, 29 | 'PointPillar': PointPillar, 30 | 'PointRCNN': PointRCNN, 31 | 'SECONDNetIoU': SECONDNetIoU, 32 | 'CaDDN': CaDDN, 33 | 'VoxelRCNN': VoxelRCNN, 34 | 'CenterPointKP': CenterPointKP, 35 | 'CenterPoint': CenterPoint, 36 | 'PillarNet': PillarNet, 37 | 'PVRCNNPlusPlus': PVRCNNPlusPlus, 38 | 'PVRCNNPlusPlus_KP': PVRCNNPlusPlus_KP, 39 | 'MPPNet': MPPNet, 40 | 'MPPNetE2E': MPPNetE2E, 41 | 'PillarNet': PillarNet, 42 | 'VoxelNeXt': VoxelNeXt, 43 | 'VoxelNeXt_KP': VoxelNeXt_KP, 44 | 'TransFusion': TransFusion, 45 | 'BevFusion': BevFusion, 46 | } 47 | 48 | 49 | def build_detector(model_cfg, num_class, dataset): 50 | model = __all__[model_cfg.NAME]( 51 | model_cfg=model_cfg, num_class=num_class, dataset=dataset 52 | ) 53 | 54 | return model 55 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /pcdet/models/detectors/transfusion.py: -------------------------------------------------------------------------------- 1 | from .detector3d_template import Detector3DTemplate 2 | 3 | 4 | class TransFusion(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(batch_dict) 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,batch_dict): 25 | disp_dict = {} 26 | 27 | loss_trans, tb_dict = batch_dict['loss'],batch_dict['tb_dict'] 28 | tb_dict = { 29 | 'loss_trans': loss_trans.item(), 30 | **tb_dict 31 | } 32 | 33 | loss = loss_trans 34 | return loss, tb_dict, disp_dict 35 | 36 | def post_processing(self, batch_dict): 37 | post_process_cfg = self.model_cfg.POST_PROCESSING 38 | batch_size = batch_dict['batch_size'] 39 | final_pred_dict = batch_dict['final_box_dicts'] 40 | recall_dict = {} 41 | for index in range(batch_size): 42 | pred_boxes = final_pred_dict[index]['pred_boxes'] 43 | 44 | recall_dict = self.generate_recall_record( 45 | box_preds=pred_boxes, 46 | recall_dict=recall_dict, batch_index=index, data_dict=batch_dict, 47 | thresh_list=post_process_cfg.RECALL_THRESH_LIST 48 | ) 49 | 50 | return final_pred_dict, recall_dict 51 | -------------------------------------------------------------------------------- /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 | #include "voxel_query_gpu.h" 9 | #include "vector_pool_gpu.h" 10 | 11 | 12 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 13 | m.def("ball_query_wrapper", &ball_query_wrapper_stack, "ball_query_wrapper_stack"); 14 | m.def("voxel_query_wrapper", &voxel_query_wrapper_stack, "voxel_query_wrapper_stack"); 15 | 16 | m.def("farthest_point_sampling_wrapper", &farthest_point_sampling_wrapper, "farthest_point_sampling_wrapper"); 17 | m.def("stack_farthest_point_sampling_wrapper", &stack_farthest_point_sampling_wrapper, "stack_farthest_point_sampling_wrapper"); 18 | 19 | m.def("group_points_wrapper", &group_points_wrapper_stack, "group_points_wrapper_stack"); 20 | m.def("group_points_grad_wrapper", &group_points_grad_wrapper_stack, "group_points_grad_wrapper_stack"); 21 | 22 | m.def("three_nn_wrapper", &three_nn_wrapper_stack, "three_nn_wrapper_stack"); 23 | m.def("three_interpolate_wrapper", &three_interpolate_wrapper_stack, "three_interpolate_wrapper_stack"); 24 | m.def("three_interpolate_grad_wrapper", &three_interpolate_grad_wrapper_stack, "three_interpolate_grad_wrapper_stack"); 25 | 26 | m.def("query_stacked_local_neighbor_idxs_wrapper_stack", &query_stacked_local_neighbor_idxs_wrapper_stack, "query_stacked_local_neighbor_idxs_wrapper_stack"); 27 | m.def("query_three_nn_by_stacked_local_idxs_wrapper_stack", &query_three_nn_by_stacked_local_idxs_wrapper_stack, "query_three_nn_by_stacked_local_idxs_wrapper_stack"); 28 | 29 | m.def("vector_pool_wrapper", &vector_pool_wrapper_stack, "vector_pool_grad_wrapper_stack"); 30 | m.def("vector_pool_grad_wrapper", &vector_pool_grad_wrapper_stack, "vector_pool_grad_wrapper_stack"); 31 | } 32 | -------------------------------------------------------------------------------- /pcdet/models/backbones_3d/vfe/image_vfe_modules/ffn/ddn_loss/balancer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from pcdet.utils import loss_utils 5 | 6 | 7 | class Balancer(nn.Module): 8 | def __init__(self, fg_weight, bg_weight, downsample_factor=1): 9 | """ 10 | Initialize fixed foreground/background loss balancer 11 | Args: 12 | fg_weight: float, Foreground loss weight 13 | bg_weight: float, Background loss weight 14 | downsample_factor: int, Depth map downsample factor 15 | """ 16 | super().__init__() 17 | self.fg_weight = fg_weight 18 | self.bg_weight = bg_weight 19 | self.downsample_factor = downsample_factor 20 | 21 | def forward(self, loss, gt_boxes2d): 22 | """ 23 | Forward pass 24 | Args: 25 | loss: (B, H, W), Pixel-wise loss 26 | gt_boxes2d: (B, N, 4), 2D box labels for foreground/background balancing 27 | Returns: 28 | loss: (1), Total loss after foreground/background balancing 29 | tb_dict: dict[float], All losses to log in tensorboard 30 | """ 31 | # Compute masks 32 | fg_mask = loss_utils.compute_fg_mask(gt_boxes2d=gt_boxes2d, 33 | shape=loss.shape, 34 | downsample_factor=self.downsample_factor, 35 | device=loss.device) 36 | bg_mask = ~fg_mask 37 | 38 | # Compute balancing weights 39 | weights = self.fg_weight * fg_mask + self.bg_weight * bg_mask 40 | num_pixels = fg_mask.sum() + bg_mask.sum() 41 | 42 | # Compute losses 43 | loss *= weights 44 | fg_loss = loss[fg_mask].sum() / num_pixels 45 | bg_loss = loss[bg_mask].sum() / num_pixels 46 | 47 | # Get total loss 48 | loss = fg_loss + bg_loss 49 | tb_dict = {"balancer_loss": loss.item(), "fg_loss": fg_loss.item(), "bg_loss": bg_loss.item()} 50 | return loss, tb_dict 51 | -------------------------------------------------------------------------------- /pcdet/ops/pointnet2/pointnet2_stack/src/sampling.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "sampling_gpu.h" 5 | 6 | #define CHECK_CUDA(x) do { \ 7 | if (!x.type().is_cuda()) { \ 8 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 9 | exit(-1); \ 10 | } \ 11 | } while (0) 12 | #define CHECK_CONTIGUOUS(x) do { \ 13 | if (!x.is_contiguous()) { \ 14 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 15 | exit(-1); \ 16 | } \ 17 | } while (0) 18 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x) 19 | 20 | 21 | int farthest_point_sampling_wrapper(int b, int n, int m, 22 | at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor idx_tensor) { 23 | 24 | CHECK_INPUT(points_tensor); 25 | CHECK_INPUT(temp_tensor); 26 | CHECK_INPUT(idx_tensor); 27 | 28 | const float *points = points_tensor.data(); 29 | float *temp = temp_tensor.data(); 30 | int *idx = idx_tensor.data(); 31 | 32 | farthest_point_sampling_kernel_launcher(b, n, m, points, temp, idx); 33 | return 1; 34 | } 35 | 36 | 37 | int stack_farthest_point_sampling_wrapper(at::Tensor points_tensor, 38 | at::Tensor temp_tensor, at::Tensor xyz_batch_cnt_tensor, at::Tensor idx_tensor, 39 | at::Tensor num_sampled_points_tensor) { 40 | 41 | CHECK_INPUT(points_tensor); 42 | CHECK_INPUT(temp_tensor); 43 | CHECK_INPUT(idx_tensor); 44 | CHECK_INPUT(xyz_batch_cnt_tensor); 45 | CHECK_INPUT(num_sampled_points_tensor); 46 | 47 | int batch_size = xyz_batch_cnt_tensor.size(0); 48 | int N = points_tensor.size(0); 49 | const float *points = points_tensor.data(); 50 | float *temp = temp_tensor.data(); 51 | int *xyz_batch_cnt = xyz_batch_cnt_tensor.data(); 52 | int *idx = idx_tensor.data(); 53 | int *num_sampled_points = num_sampled_points_tensor.data(); 54 | 55 | stack_farthest_point_sampling_kernel_launcher(N, batch_size, points, temp, xyz_batch_cnt, idx, num_sampled_points); 56 | return 1; 57 | } -------------------------------------------------------------------------------- /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 "interpolate_gpu.h" 16 | 17 | 18 | void three_nn_wrapper_fast(int b, int n, int m, at::Tensor unknown_tensor, 19 | at::Tensor known_tensor, at::Tensor dist2_tensor, at::Tensor idx_tensor) { 20 | const float *unknown = unknown_tensor.data(); 21 | const float *known = known_tensor.data(); 22 | float *dist2 = dist2_tensor.data(); 23 | int *idx = idx_tensor.data(); 24 | 25 | three_nn_kernel_launcher_fast(b, n, m, unknown, known, dist2, idx); 26 | } 27 | 28 | 29 | void three_interpolate_wrapper_fast(int b, int c, int m, int n, 30 | at::Tensor points_tensor, 31 | at::Tensor idx_tensor, 32 | at::Tensor weight_tensor, 33 | at::Tensor out_tensor) { 34 | 35 | const float *points = points_tensor.data(); 36 | const float *weight = weight_tensor.data(); 37 | float *out = out_tensor.data(); 38 | const int *idx = idx_tensor.data(); 39 | 40 | three_interpolate_kernel_launcher_fast(b, c, m, n, points, idx, weight, out); 41 | } 42 | 43 | 44 | void three_interpolate_grad_wrapper_fast(int b, int c, int n, int m, 45 | at::Tensor grad_out_tensor, 46 | at::Tensor idx_tensor, 47 | at::Tensor weight_tensor, 48 | at::Tensor grad_points_tensor) { 49 | 50 | const float *grad_out = grad_out_tensor.data(); 51 | const float *weight = weight_tensor.data(); 52 | float *grad_points = grad_points_tensor.data(); 53 | const int *idx = idx_tensor.data(); 54 | 55 | three_interpolate_grad_kernel_launcher_fast(b, c, n, m, grad_out, idx, weight, grad_points); 56 | } 57 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /pcdet/models/detectors/pv_rcnn_plusplus.py: -------------------------------------------------------------------------------- 1 | from .detector3d_template import Detector3DTemplate 2 | 3 | 4 | class PVRCNNPlusPlus(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 | batch_dict = self.vfe(batch_dict) 11 | batch_dict = self.backbone_3d(batch_dict) 12 | batch_dict = self.map_to_bev_module(batch_dict) 13 | batch_dict = self.backbone_2d(batch_dict) 14 | batch_dict = self.dense_head(batch_dict) 15 | 16 | batch_dict = self.roi_head.proposal_layer( 17 | batch_dict, nms_config=self.roi_head.model_cfg.NMS_CONFIG['TRAIN' if self.training else 'TEST'] 18 | ) 19 | if self.training: 20 | targets_dict = self.roi_head.assign_targets(batch_dict) 21 | batch_dict['rois'] = targets_dict['rois'] 22 | batch_dict['roi_labels'] = targets_dict['roi_labels'] 23 | batch_dict['roi_targets_dict'] = targets_dict 24 | num_rois_per_scene = targets_dict['rois'].shape[1] 25 | if 'roi_valid_num' in batch_dict: 26 | batch_dict['roi_valid_num'] = [num_rois_per_scene for _ in range(batch_dict['batch_size'])] 27 | 28 | batch_dict = self.pfe(batch_dict) 29 | batch_dict = self.point_head(batch_dict) 30 | batch_dict = self.roi_head(batch_dict) 31 | 32 | if self.training: 33 | loss, tb_dict, disp_dict = self.get_training_loss() 34 | 35 | ret_dict = { 36 | 'loss': loss 37 | } 38 | return ret_dict, tb_dict, disp_dict 39 | else: 40 | pred_dicts, recall_dicts = self.post_processing(batch_dict) 41 | return pred_dicts, recall_dicts 42 | 43 | def get_training_loss(self): 44 | disp_dict = {} 45 | loss_rpn, tb_dict = self.dense_head.get_loss() 46 | if self.point_head is not None: 47 | loss_point, tb_dict = self.point_head.get_loss(tb_dict) 48 | else: 49 | loss_point = 0 50 | loss_rcnn, tb_dict = self.roi_head.get_loss(tb_dict) 51 | 52 | loss = loss_rpn + loss_point + loss_rcnn 53 | return loss, tb_dict, disp_dict 54 | -------------------------------------------------------------------------------- /pcdet/models/backbones_3d/focal_sparse_conv/SemanticSeg/basic_blocks.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | class BasicBlock1D(nn.Module): 4 | 5 | def __init__(self, in_channels, out_channels, **kwargs): 6 | """ 7 | Initializes convolutional block 8 | Args: 9 | in_channels: int, Number of input channels 10 | out_channels: int, Number of output channels 11 | **kwargs: Dict, Extra arguments for nn.Conv2d 12 | """ 13 | super().__init__() 14 | self.in_channels = in_channels 15 | self.out_channels = out_channels 16 | self.conv = nn.Conv1d(in_channels=in_channels, 17 | out_channels=out_channels, 18 | **kwargs) 19 | self.bn = nn.BatchNorm1d(out_channels) 20 | self.relu = nn.ReLU(inplace=True) 21 | 22 | def forward(self, features): 23 | """ 24 | Applies convolutional block 25 | Args: 26 | features: (B, C_in, H, W), Input features 27 | Returns: 28 | x: (B, C_out, H, W), Output features 29 | """ 30 | x = self.conv(features) 31 | x = self.bn(x) 32 | x = self.relu(x) 33 | return x 34 | 35 | class BasicBlock2D(nn.Module): 36 | 37 | def __init__(self, in_channels, out_channels, **kwargs): 38 | """ 39 | Initializes convolutional block 40 | Args: 41 | in_channels: int, Number of input channels 42 | out_channels: int, Number of output channels 43 | **kwargs: Dict, Extra arguments for nn.Conv2d 44 | """ 45 | super().__init__() 46 | self.in_channels = in_channels 47 | self.out_channels = out_channels 48 | self.conv = nn.Conv2d(in_channels=in_channels, 49 | out_channels=out_channels, 50 | **kwargs) 51 | self.bn = nn.BatchNorm2d(out_channels) 52 | self.relu = nn.ReLU(inplace=True) 53 | 54 | def forward(self, features): 55 | """ 56 | Applies convolutional block 57 | Args: 58 | features: (B, C_in, H, W), Input features 59 | Returns: 60 | x: (B, C_out, H, W), Output features 61 | """ 62 | x = self.conv(features) 63 | x = self.bn(x) 64 | x = self.relu(x) 65 | return x 66 | -------------------------------------------------------------------------------- /pcdet/ops/ingroup_inds/src/ingroup_inds_kernel.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include "cuda_fp16.h" 8 | 9 | #define CHECK_CALL(call) \ 10 | do \ 11 | { \ 12 | const cudaError_t error_code = call; \ 13 | if (error_code != cudaSuccess) \ 14 | { \ 15 | printf("CUDA Error:\n"); \ 16 | printf(" File: %s\n", __FILE__); \ 17 | printf(" Line: %d\n", __LINE__); \ 18 | printf(" Error code: %d\n", error_code); \ 19 | printf(" Error text: %s\n", \ 20 | cudaGetErrorString(error_code)); \ 21 | exit(1); \ 22 | } \ 23 | } while (0) 24 | 25 | #define THREADS_PER_BLOCK 256 26 | #define DIVUP(m, n) ((m) / (n) + ((m) % (n) > 0)) 27 | 28 | // #define DEBUG 29 | // #define ASSERTION 30 | 31 | __global__ void ingroup_inds_kernel( 32 | const long *group_inds, 33 | long *out_inds, 34 | int *ingroup_counter, 35 | int N 36 | ) { 37 | 38 | int idx = blockIdx.x * blockDim.x + threadIdx.x; 39 | if (idx >= N) return; 40 | long this_group_id = group_inds[idx]; 41 | 42 | int cnt = atomicAdd(&ingroup_counter[this_group_id], 1); 43 | out_inds[idx] = cnt; 44 | } 45 | 46 | 47 | void ingroup_inds_launcher( 48 | const long *group_inds, 49 | long *out_inds, 50 | int N, 51 | int max_group_id 52 | ) { 53 | 54 | int *ingroup_counter = NULL; 55 | CHECK_CALL(cudaMalloc(&ingroup_counter, (max_group_id + 1) * sizeof(int))); 56 | CHECK_CALL(cudaMemset(ingroup_counter, 0, (max_group_id + 1) * sizeof(int))); 57 | 58 | dim3 blocks(DIVUP(N, THREADS_PER_BLOCK)); 59 | dim3 threads(THREADS_PER_BLOCK); 60 | 61 | ingroup_inds_kernel<<>>( 62 | group_inds, 63 | out_inds, 64 | ingroup_counter, 65 | N 66 | ); 67 | 68 | cudaFree(ingroup_counter); 69 | 70 | #ifdef DEBUG 71 | CHECK_CALL(cudaGetLastError()); 72 | CHECK_CALL(cudaDeviceSynchronize()); 73 | #endif 74 | 75 | return; 76 | 77 | } -------------------------------------------------------------------------------- /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 | 34 | if self.point_encoding_config.get('filter_sweeps', False) and 'timestamp' in self.src_feature_list: 35 | max_sweeps = self.point_encoding_config.max_sweeps 36 | idx = self.src_feature_list.index('timestamp') 37 | dt = np.round(data_dict['points'][:, idx], 2) 38 | max_dt = sorted(np.unique(dt))[min(len(np.unique(dt))-1, max_sweeps-1)] 39 | data_dict['points'] = data_dict['points'][dt <= max_dt] 40 | 41 | return data_dict 42 | 43 | def absolute_coordinates_encoding(self, points=None): 44 | if points is None: 45 | num_output_features = len(self.used_feature_list) 46 | return num_output_features 47 | 48 | assert points.shape[-1] == len(self.src_feature_list) 49 | point_feature_list = [points[:, 0:3]] 50 | for x in self.used_feature_list: 51 | if x in ['x', 'y', 'z']: 52 | continue 53 | idx = self.src_feature_list.index(x) 54 | point_feature_list.append(points[:, idx:idx+1]) 55 | point_features = np.concatenate(point_feature_list, axis=1) 56 | 57 | return point_features, True 58 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /docs/DEMO.md: -------------------------------------------------------------------------------- 1 | # Quick Demo 2 | 3 | Here we provide a quick demo to test a pretrained model on the custom point cloud data and visualize the predicted results. 4 | 5 | We suppose you already followed the [INSTALL.md](INSTALL.md) to install the `OpenPCDet` repo successfully. 6 | 7 | 1. Download the provided pretrained models as shown in the [README.md](../README.md). 8 | 9 | 2. Make sure you have already installed the [`Open3D`](https://github.com/isl-org/Open3D) (faster) or `mayavi` visualization tools. 10 | If not, you could install it as follows: 11 | ``` 12 | pip install open3d 13 | # or 14 | pip install mayavi 15 | ``` 16 | 17 | 3. Prepare your custom point cloud data (skip this step if you use the original KITTI data). 18 | * You need to transform the coordinate of your custom point cloud to 19 | the unified normative coordinate of `OpenPCDet`, that is, x-axis points towards to front direction, 20 | y-axis points towards to the left direction, and z-axis points towards to the top direction. 21 | * (Optional) the z-axis origin of your point cloud coordinate should be about 1.6m above the ground surface, 22 | since currently the provided models are trained on the KITTI dataset. 23 | * Set the intensity information, and save your transformed custom data to `numpy file`: 24 | ```python 25 | # Transform your point cloud data 26 | ... 27 | 28 | # Save it to the file. 29 | # The shape of points should be (num_points, 4), that is [x, y, z, intensity] (Only for KITTI dataset). 30 | # If you doesn't have the intensity information, just set them to zeros. 31 | # If you have the intensity information, you should normalize them to [0, 1]. 32 | points[:, 3] = 0 33 | np.save(`my_data.npy`, points) 34 | ``` 35 | 36 | 4. Run the demo with a pretrained model (e.g. PV-RCNN) and your custom point cloud data as follows: 37 | ```shell 38 | python demo.py --cfg_file cfgs/kitti_models/pv_rcnn.yaml \ 39 | --ckpt pv_rcnn_8369.pth \ 40 | --data_path ${POINT_CLOUD_DATA} 41 | ``` 42 | Here `${POINT_CLOUD_DATA}` could be in any of the following format: 43 | * Your transformed custom data with a single numpy file like `my_data.npy`. 44 | * Your transformed custom data with a directory to test with multiple point cloud data. 45 | * The original KITTI `.bin` data within `data/kitti`, like `data/kitti/training/velodyne/000008.bin`. 46 | 47 | Then you could see the predicted results with visualized point cloud as follows: 48 | 49 |

50 | 51 |

52 | -------------------------------------------------------------------------------- /pcdet/models/backbones_3d/vfe/image_vfe_modules/f2v/frustum_to_voxel.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from .frustum_grid_generator import FrustumGridGenerator 5 | from .sampler import Sampler 6 | 7 | 8 | class FrustumToVoxel(nn.Module): 9 | 10 | def __init__(self, model_cfg, grid_size, pc_range, disc_cfg): 11 | """ 12 | Initializes module to transform frustum features to voxel features via 3D transformation and sampling 13 | Args: 14 | model_cfg: EasyDict, Module configuration 15 | grid_size: [X, Y, Z], Voxel grid size 16 | pc_range: [x_min, y_min, z_min, x_max, y_max, z_max], Voxelization point cloud range (m) 17 | disc_cfg: EasyDict, Depth discretiziation configuration 18 | """ 19 | super().__init__() 20 | self.model_cfg = model_cfg 21 | self.grid_size = grid_size 22 | self.pc_range = pc_range 23 | self.disc_cfg = disc_cfg 24 | self.grid_generator = FrustumGridGenerator(grid_size=grid_size, 25 | pc_range=pc_range, 26 | disc_cfg=disc_cfg) 27 | self.sampler = Sampler(**model_cfg.SAMPLER) 28 | 29 | def forward(self, batch_dict): 30 | """ 31 | Generates voxel features via 3D transformation and sampling 32 | Args: 33 | batch_dict: 34 | frustum_features: (B, C, D, H_image, W_image), Image frustum features 35 | lidar_to_cam: (B, 4, 4), LiDAR to camera frame transformation 36 | cam_to_img: (B, 3, 4), Camera projection matrix 37 | image_shape: (B, 2), Image shape [H, W] 38 | Returns: 39 | batch_dict: 40 | voxel_features: (B, C, Z, Y, X), Image voxel features 41 | """ 42 | # Generate sampling grid for frustum volume 43 | grid = self.grid_generator(lidar_to_cam=batch_dict["trans_lidar_to_cam"], 44 | cam_to_img=batch_dict["trans_cam_to_img"], 45 | image_shape=batch_dict["image_shape"]) # (B, X, Y, Z, 3) 46 | 47 | # Sample frustum volume to generate voxel volume 48 | voxel_features = self.sampler(input_features=batch_dict["frustum_features"], 49 | grid=grid) # (B, C, X, Y, Z) 50 | 51 | # (B, C, X, Y, Z) -> (B, C, Z, Y, X) 52 | voxel_features = voxel_features.permute(0, 1, 4, 3, 2) 53 | batch_dict["voxel_features"] = voxel_features 54 | return batch_dict 55 | -------------------------------------------------------------------------------- /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 "group_points_gpu.h" 13 | 14 | #define CHECK_CUDA(x) do { \ 15 | if (!x.type().is_cuda()) { \ 16 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 17 | exit(-1); \ 18 | } \ 19 | } while (0) 20 | #define CHECK_CONTIGUOUS(x) do { \ 21 | if (!x.is_contiguous()) { \ 22 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \ 23 | exit(-1); \ 24 | } \ 25 | } while (0) 26 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x) 27 | 28 | 29 | int group_points_grad_wrapper_stack(int B, int M, int C, int N, int nsample, 30 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor idx_batch_cnt_tensor, 31 | at::Tensor features_batch_cnt_tensor, at::Tensor grad_features_tensor) { 32 | 33 | CHECK_INPUT(grad_out_tensor); 34 | CHECK_INPUT(idx_tensor); 35 | CHECK_INPUT(idx_batch_cnt_tensor); 36 | CHECK_INPUT(features_batch_cnt_tensor); 37 | CHECK_INPUT(grad_features_tensor); 38 | 39 | const float *grad_out = grad_out_tensor.data(); 40 | const int *idx = idx_tensor.data(); 41 | const int *idx_batch_cnt = idx_batch_cnt_tensor.data(); 42 | const int *features_batch_cnt = features_batch_cnt_tensor.data(); 43 | float *grad_features = grad_features_tensor.data(); 44 | 45 | group_points_grad_kernel_launcher_stack(B, M, C, N, nsample, grad_out, idx, idx_batch_cnt, features_batch_cnt, grad_features); 46 | return 1; 47 | } 48 | 49 | 50 | int group_points_wrapper_stack(int B, int M, int C, int nsample, 51 | at::Tensor features_tensor, at::Tensor features_batch_cnt_tensor, 52 | at::Tensor idx_tensor, at::Tensor idx_batch_cnt_tensor, at::Tensor out_tensor) { 53 | 54 | CHECK_INPUT(features_tensor); 55 | CHECK_INPUT(features_batch_cnt_tensor); 56 | CHECK_INPUT(idx_tensor); 57 | CHECK_INPUT(idx_batch_cnt_tensor); 58 | CHECK_INPUT(out_tensor); 59 | 60 | const float *features = features_tensor.data(); 61 | const int *idx = idx_tensor.data(); 62 | const int *features_batch_cnt = features_batch_cnt_tensor.data(); 63 | const int *idx_batch_cnt = idx_batch_cnt_tensor.data(); 64 | float *out = out_tensor.data(); 65 | 66 | group_points_kernel_launcher_stack(B, M, C, nsample, features, features_batch_cnt, idx, idx_batch_cnt, out); 67 | return 1; 68 | } -------------------------------------------------------------------------------- /pcdet/models/backbones_3d/vfe/image_vfe_modules/ffn/ddn_loss/ddn_loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | from .balancer import Balancer 6 | from pcdet.utils import transform_utils 7 | 8 | try: 9 | from kornia.losses.focal import FocalLoss 10 | except: 11 | pass 12 | # print('Warning: kornia is not installed. This package is only required by CaDDN') 13 | 14 | 15 | class DDNLoss(nn.Module): 16 | 17 | def __init__(self, 18 | weight, 19 | alpha, 20 | gamma, 21 | disc_cfg, 22 | fg_weight, 23 | bg_weight, 24 | downsample_factor): 25 | """ 26 | Initializes DDNLoss module 27 | Args: 28 | weight: float, Loss function weight 29 | alpha: float, Alpha value for Focal Loss 30 | gamma: float, Gamma value for Focal Loss 31 | disc_cfg: dict, Depth discretiziation configuration 32 | fg_weight: float, Foreground loss weight 33 | bg_weight: float, Background loss weight 34 | downsample_factor: int, Depth map downsample factor 35 | """ 36 | super().__init__() 37 | self.device = torch.cuda.current_device() 38 | self.disc_cfg = disc_cfg 39 | self.balancer = Balancer(downsample_factor=downsample_factor, 40 | fg_weight=fg_weight, 41 | bg_weight=bg_weight) 42 | 43 | # Set loss function 44 | self.alpha = alpha 45 | self.gamma = gamma 46 | self.loss_func = FocalLoss(alpha=self.alpha, gamma=self.gamma, reduction="none") 47 | self.weight = weight 48 | 49 | def forward(self, depth_logits, depth_maps, gt_boxes2d): 50 | """ 51 | Gets DDN loss 52 | Args: 53 | depth_logits: (B, D+1, H, W), Predicted depth logits 54 | depth_maps: (B, H, W), Depth map [m] 55 | gt_boxes2d: torch.Tensor (B, N, 4), 2D box labels for foreground/background balancing 56 | Returns: 57 | loss: (1), Depth distribution network loss 58 | tb_dict: dict[float], All losses to log in tensorboard 59 | """ 60 | tb_dict = {} 61 | 62 | # Bin depth map to create target 63 | depth_target = transform_utils.bin_depths(depth_maps, **self.disc_cfg, target=True) 64 | 65 | # Compute loss 66 | loss = self.loss_func(depth_logits, depth_target) 67 | 68 | # Compute foreground/background balancing 69 | loss, tb_dict = self.balancer(loss=loss, gt_boxes2d=gt_boxes2d) 70 | 71 | # Final loss 72 | loss *= self.weight 73 | tb_dict.update({"ddn_loss": loss.item()}) 74 | 75 | return loss, tb_dict 76 | -------------------------------------------------------------------------------- /pcdet/datasets/kitti/kitti_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from ...utils import box_utils 3 | 4 | 5 | def transform_annotations_to_kitti_format(annos, map_name_to_kitti=None, info_with_fakelidar=False): 6 | """ 7 | Args: 8 | annos: 9 | map_name_to_kitti: dict, map name to KITTI names (Car, Pedestrian, Cyclist) 10 | info_with_fakelidar: 11 | Returns: 12 | 13 | """ 14 | for anno in annos: 15 | # For lyft and nuscenes, different anno key in info 16 | if 'name' not in anno: 17 | anno['name'] = anno['gt_names'] 18 | anno.pop('gt_names') 19 | 20 | for k in range(anno['name'].shape[0]): 21 | anno['name'][k] = map_name_to_kitti[anno['name'][k]] 22 | 23 | anno['bbox'] = np.zeros((len(anno['name']), 4)) 24 | anno['bbox'][:, 2:4] = 50 # [0, 0, 50, 50] 25 | anno['truncated'] = np.zeros(len(anno['name'])) 26 | anno['occluded'] = np.zeros(len(anno['name'])) 27 | if 'boxes_lidar' in anno: 28 | gt_boxes_lidar = anno['boxes_lidar'].copy() 29 | else: 30 | gt_boxes_lidar = anno['gt_boxes_lidar'].copy() 31 | 32 | if len(gt_boxes_lidar) > 0: 33 | if info_with_fakelidar: 34 | gt_boxes_lidar = box_utils.boxes3d_kitti_fakelidar_to_lidar(gt_boxes_lidar) 35 | 36 | gt_boxes_lidar[:, 2] -= gt_boxes_lidar[:, 5] / 2 37 | anno['location'] = np.zeros((gt_boxes_lidar.shape[0], 3)) 38 | anno['location'][:, 0] = -gt_boxes_lidar[:, 1] # x = -y_lidar 39 | anno['location'][:, 1] = -gt_boxes_lidar[:, 2] # y = -z_lidar 40 | anno['location'][:, 2] = gt_boxes_lidar[:, 0] # z = x_lidar 41 | dxdydz = gt_boxes_lidar[:, 3:6] 42 | anno['dimensions'] = dxdydz[:, [0, 2, 1]] # lwh ==> lhw 43 | anno['rotation_y'] = -gt_boxes_lidar[:, 6] - np.pi / 2.0 44 | anno['alpha'] = -np.arctan2(-gt_boxes_lidar[:, 1], gt_boxes_lidar[:, 0]) + anno['rotation_y'] 45 | else: 46 | anno['location'] = anno['dimensions'] = np.zeros((0, 3)) 47 | anno['rotation_y'] = anno['alpha'] = np.zeros(0) 48 | 49 | return annos 50 | 51 | 52 | def calib_to_matricies(calib): 53 | """ 54 | Converts calibration object to transformation matricies 55 | Args: 56 | calib: calibration.Calibration, Calibration object 57 | Returns 58 | V2R: (4, 4), Lidar to rectified camera transformation matrix 59 | P2: (3, 4), Camera projection matrix 60 | """ 61 | V2C = np.vstack((calib.V2C, np.array([0, 0, 0, 1], dtype=np.float32))) # (4, 4) 62 | R0 = np.hstack((calib.R0, np.zeros((3, 1), dtype=np.float32))) # (3, 4) 63 | R0 = np.vstack((R0, np.array([0, 0, 0, 1], dtype=np.float32))) # (4, 4) 64 | V2R = R0 @ V2C 65 | P2 = calib.P2 66 | return V2R, P2 -------------------------------------------------------------------------------- /pcdet/models/backbones_3d/vfe/image_vfe.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from .vfe_template import VFETemplate 4 | from .image_vfe_modules import ffn, f2v 5 | 6 | 7 | class ImageVFE(VFETemplate): 8 | def __init__(self, model_cfg, grid_size, point_cloud_range, depth_downsample_factor, **kwargs): 9 | super().__init__(model_cfg=model_cfg) 10 | self.grid_size = grid_size 11 | self.pc_range = point_cloud_range 12 | self.downsample_factor = depth_downsample_factor 13 | self.module_topology = [ 14 | 'ffn', 'f2v' 15 | ] 16 | self.build_modules() 17 | 18 | def build_modules(self): 19 | """ 20 | Builds modules 21 | """ 22 | for module_name in self.module_topology: 23 | module = getattr(self, 'build_%s' % module_name)() 24 | self.add_module(module_name, module) 25 | 26 | def build_ffn(self): 27 | """ 28 | Builds frustum feature network 29 | Returns: 30 | ffn_module: nn.Module, Frustum feature network 31 | """ 32 | ffn_module = ffn.__all__[self.model_cfg.FFN.NAME]( 33 | model_cfg=self.model_cfg.FFN, 34 | downsample_factor=self.downsample_factor 35 | ) 36 | self.disc_cfg = ffn_module.disc_cfg 37 | return ffn_module 38 | 39 | def build_f2v(self): 40 | """ 41 | Builds frustum to voxel transformation 42 | Returns: 43 | f2v_module: nn.Module, Frustum to voxel transformation 44 | """ 45 | f2v_module = f2v.__all__[self.model_cfg.F2V.NAME]( 46 | model_cfg=self.model_cfg.F2V, 47 | grid_size=self.grid_size, 48 | pc_range=self.pc_range, 49 | disc_cfg=self.disc_cfg 50 | ) 51 | return f2v_module 52 | 53 | def get_output_feature_dim(self): 54 | """ 55 | Gets number of output channels 56 | Returns: 57 | out_feature_dim: int, Number of output channels 58 | """ 59 | out_feature_dim = self.ffn.get_output_feature_dim() 60 | return out_feature_dim 61 | 62 | def forward(self, batch_dict, **kwargs): 63 | """ 64 | Args: 65 | batch_dict: 66 | images: (N, 3, H_in, W_in), Input images 67 | **kwargs: 68 | Returns: 69 | batch_dict: 70 | voxel_features: (B, C, Z, Y, X), Image voxel features 71 | """ 72 | batch_dict = self.ffn(batch_dict) 73 | batch_dict = self.f2v(batch_dict) 74 | return batch_dict 75 | 76 | def get_loss(self): 77 | """ 78 | Gets DDN loss 79 | Returns: 80 | loss: (1), Depth distribution network loss 81 | tb_dict: dict[float], All losses to log in tensorboard 82 | """ 83 | 84 | loss, tb_dict = self.ffn.get_loss() 85 | return loss, tb_dict 86 | -------------------------------------------------------------------------------- /tools/cfgs/dataset_configs/waymo_dataset_multiframe.yaml: -------------------------------------------------------------------------------- 1 | DATASET: 'WaymoDataset' 2 | DATA_PATH: '../data/waymo' 3 | PROCESSED_DATA_TAG: 'waymo_processed_data_v0_5_0' 4 | 5 | POINT_CLOUD_RANGE: [-75.2, -75.2, -2, 75.2, 75.2, 4] 6 | 7 | DATA_SPLIT: { 8 | 'train': train, 9 | 'test': val 10 | } 11 | 12 | SAMPLED_INTERVAL: { 13 | 'train': 5, 14 | 'test': 1 15 | } 16 | 17 | FILTER_EMPTY_BOXES_FOR_TRAIN: True 18 | DISABLE_NLZ_FLAG_ON_POINTS: True 19 | 20 | USE_SHARED_MEMORY: False # it will load the data to shared memory to speed up (DO NOT USE IT IF YOU DO NOT FULLY UNDERSTAND WHAT WILL HAPPEN) 21 | SHARED_MEMORY_FILE_LIMIT: 35000 # set it based on the size of your shared memory 22 | 23 | SEQUENCE_CONFIG: 24 | ENABLED: True 25 | SAMPLE_OFFSET: [-3, 0] 26 | 27 | TRAIN_WITH_SPEED: True 28 | 29 | 30 | DATA_AUGMENTOR: 31 | DISABLE_AUG_LIST: ['placeholder'] 32 | AUG_CONFIG_LIST: 33 | - NAME: gt_sampling 34 | USE_ROAD_PLANE: False 35 | DB_INFO_PATH: 36 | - waymo_processed_data_v0_5_0_waymo_dbinfos_train_sampled_1_multiframe_-4_to_0.pkl 37 | 38 | USE_SHARED_MEMORY: False # set it to True to speed up (it costs about 50GB? shared memory) 39 | DB_DATA_PATH: 40 | - waymo_processed_data_v0_5_0_gt_database_train_sampled_1_multiframe_-4_to_0_global.npy 41 | 42 | PREPARE: { 43 | filter_by_min_points: ['Vehicle:5', 'Pedestrian:5', 'Cyclist:5'], 44 | filter_by_difficulty: [-1], 45 | } 46 | 47 | SAMPLE_GROUPS: ['Vehicle:15', 'Pedestrian:10', 'Cyclist:10'] 48 | NUM_POINT_FEATURES: 6 49 | REMOVE_EXTRA_WIDTH: [0.0, 0.0, 0.0] 50 | LIMIT_WHOLE_SCENE: True 51 | 52 | FILTER_OBJ_POINTS_BY_TIMESTAMP: True 53 | TIME_RANGE: [0.3, 0.0] # 0.3s-0.0s indicates 4 frames 54 | 55 | - NAME: random_world_flip 56 | ALONG_AXIS_LIST: ['x', 'y'] 57 | 58 | - NAME: random_world_rotation 59 | WORLD_ROT_ANGLE: [-0.78539816, 0.78539816] 60 | 61 | - NAME: random_world_scaling 62 | WORLD_SCALE_RANGE: [0.95, 1.05] 63 | 64 | 65 | POINT_FEATURE_ENCODING: { 66 | encoding_type: absolute_coordinates_encoding, 67 | used_feature_list: ['x', 'y', 'z', 'intensity', 'elongation', 'timestamp'], 68 | src_feature_list: ['x', 'y', 'z', 'intensity', 'elongation', 'timestamp'], 69 | } 70 | 71 | 72 | DATA_PROCESSOR: 73 | - NAME: mask_points_and_boxes_outside_range 74 | REMOVE_OUTSIDE_BOXES: True 75 | USE_CENTER_TO_FILTER: True 76 | 77 | - NAME: shuffle_points 78 | SHUFFLE_ENABLED: { 79 | 'train': True, 80 | 'test': True 81 | } 82 | 83 | - NAME: transform_points_to_voxels 84 | VOXEL_SIZE: [0.1, 0.1, 0.15] 85 | MAX_POINTS_PER_VOXEL: 5 86 | MAX_NUMBER_OF_VOXELS: { 87 | 'train': 180000, 88 | 'test': 400000 89 | } 90 | -------------------------------------------------------------------------------- /tools/cfgs/dataset_configs/waymo_dataset.yaml: -------------------------------------------------------------------------------- 1 | DATASET: 'WaymoDataset' 2 | DATA_PATH: '../data/waymo' 3 | PROCESSED_DATA_TAG: 'waymo_processed_data_v0_5_0' 4 | 5 | POINT_CLOUD_RANGE: [-75.2, -75.2, -2, 75.2, 75.2, 4] 6 | 7 | DATA_SPLIT: { 8 | 'train': train, 9 | 'test': val 10 | } 11 | 12 | SAMPLED_INTERVAL: { 13 | 'train': 5, 14 | 'test': 1 15 | } 16 | 17 | FILTER_EMPTY_BOXES_FOR_TRAIN: True 18 | DISABLE_NLZ_FLAG_ON_POINTS: True 19 | 20 | USE_SHARED_MEMORY: False # it will load the data to shared memory to speed up (DO NOT USE IT IF YOU DO NOT FULLY UNDERSTAND WHAT WILL HAPPEN) 21 | SHARED_MEMORY_FILE_LIMIT: 35000 # set it based on the size of your shared memory 22 | 23 | DATA_AUGMENTOR: 24 | DISABLE_AUG_LIST: ['placeholder'] 25 | AUG_CONFIG_LIST: 26 | - NAME: gt_sampling 27 | USE_ROAD_PLANE: False 28 | DB_INFO_PATH: 29 | - waymo_processed_data_v0_5_0_waymo_dbinfos_train_sampled_1.pkl 30 | 31 | USE_SHARED_MEMORY: False # set it to True to speed up (it costs about 15GB shared memory) 32 | DB_DATA_PATH: 33 | - waymo_processed_data_v0_5_0_gt_database_train_sampled_1_global.npy 34 | 35 | BACKUP_DB_INFO: 36 | # if the above DB_INFO cannot be found, will use this backup one 37 | DB_INFO_PATH: waymo_processed_data_v0_5_0_waymo_dbinfos_train_sampled_1_multiframe_-4_to_0.pkl 38 | DB_DATA_PATH: waymo_processed_data_v0_5_0_gt_database_train_sampled_1_multiframe_-4_to_0_global.npy 39 | NUM_POINT_FEATURES: 6 40 | 41 | PREPARE: { 42 | filter_by_min_points: ['Vehicle:5', 'Pedestrian:5', 'Cyclist:5'], 43 | filter_by_difficulty: [-1], 44 | } 45 | 46 | SAMPLE_GROUPS: ['Vehicle:15', 'Pedestrian:10', 'Cyclist:10'] 47 | NUM_POINT_FEATURES: 5 48 | REMOVE_EXTRA_WIDTH: [0.0, 0.0, 0.0] 49 | LIMIT_WHOLE_SCENE: True 50 | 51 | - NAME: random_world_flip 52 | ALONG_AXIS_LIST: ['x', 'y'] 53 | 54 | - NAME: random_world_rotation 55 | WORLD_ROT_ANGLE: [-0.78539816, 0.78539816] 56 | 57 | - NAME: random_world_scaling 58 | WORLD_SCALE_RANGE: [0.95, 1.05] 59 | 60 | 61 | POINT_FEATURE_ENCODING: { 62 | encoding_type: absolute_coordinates_encoding, 63 | used_feature_list: ['x', 'y', 'z', 'intensity', 'elongation'], 64 | src_feature_list: ['x', 'y', 'z', 'intensity', 'elongation'], 65 | } 66 | 67 | 68 | DATA_PROCESSOR: 69 | - NAME: mask_points_and_boxes_outside_range 70 | REMOVE_OUTSIDE_BOXES: True 71 | 72 | - NAME: shuffle_points 73 | SHUFFLE_ENABLED: { 74 | 'train': True, 75 | 'test': True 76 | } 77 | 78 | - NAME: transform_points_to_voxels 79 | VOXEL_SIZE: [0.1, 0.1, 0.15] 80 | MAX_POINTS_PER_VOXEL: 5 81 | MAX_NUMBER_OF_VOXELS: { 82 | 'train': 150000, 83 | 'test': 150000 84 | } 85 | -------------------------------------------------------------------------------- /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, CosineAnnealing 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 in ['adam_onecycle','adam_cosineanneal']: 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 | betas = optim_cfg.get('BETAS', (0.9, 0.99)) 29 | betas = tuple(betas) 30 | optimizer_func = partial(optim.Adam, betas=betas) 31 | optimizer = OptimWrapper.create( 32 | optimizer_func, 3e-3, get_layer_groups(model), wd=optim_cfg.WEIGHT_DECAY, true_wd=True, bn_wd=True 33 | ) 34 | else: 35 | raise NotImplementedError 36 | 37 | return optimizer 38 | 39 | 40 | def build_scheduler(optimizer, total_iters_each_epoch, total_epochs, last_epoch, optim_cfg): 41 | decay_steps = [x * total_iters_each_epoch for x in optim_cfg.DECAY_STEP_LIST] 42 | def lr_lbmd(cur_epoch): 43 | cur_decay = 1 44 | for decay_step in decay_steps: 45 | if cur_epoch >= decay_step: 46 | cur_decay = cur_decay * optim_cfg.LR_DECAY 47 | return max(cur_decay, optim_cfg.LR_CLIP / optim_cfg.LR) 48 | 49 | lr_warmup_scheduler = None 50 | total_steps = total_iters_each_epoch * total_epochs 51 | if optim_cfg.OPTIMIZER == 'adam_onecycle': 52 | lr_scheduler = OneCycle( 53 | optimizer, total_steps, optim_cfg.LR, list(optim_cfg.MOMS), optim_cfg.DIV_FACTOR, optim_cfg.PCT_START 54 | ) 55 | elif optim_cfg.OPTIMIZER == 'adam_cosineanneal': 56 | lr_scheduler = CosineAnnealing( 57 | optimizer, total_steps, total_epochs, optim_cfg.LR, list(optim_cfg.MOMS), optim_cfg.PCT_START, optim_cfg.WARMUP_ITER 58 | ) 59 | else: 60 | lr_scheduler = lr_sched.LambdaLR(optimizer, lr_lbmd, last_epoch=last_epoch) 61 | 62 | if optim_cfg.LR_WARMUP: 63 | lr_warmup_scheduler = CosineWarmupLR( 64 | optimizer, T_max=optim_cfg.WARMUP_EPOCH * len(total_iters_each_epoch), 65 | eta_min=optim_cfg.LR / optim_cfg.DIV_FACTOR 66 | ) 67 | 68 | return lr_scheduler, lr_warmup_scheduler 69 | -------------------------------------------------------------------------------- /pcdet/ops/bev_pool/bev_pool.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from . import bev_pool_ext 4 | 5 | __all__ = ["bev_pool"] 6 | 7 | 8 | class QuickCumsum(torch.autograd.Function): 9 | @staticmethod 10 | def forward(ctx, x, geom_feats, ranks): 11 | x = x.cumsum(0) 12 | kept = torch.ones(x.shape[0], device=x.device, dtype=torch.bool) 13 | kept[:-1] = ranks[1:] != ranks[:-1] 14 | 15 | x, geom_feats = x[kept], geom_feats[kept] 16 | x = torch.cat((x[:1], x[1:] - x[:-1])) 17 | 18 | # save kept for backward 19 | ctx.save_for_backward(kept) 20 | 21 | # no gradient for geom_feats 22 | ctx.mark_non_differentiable(geom_feats) 23 | 24 | return x, geom_feats 25 | 26 | @staticmethod 27 | def backward(ctx, gradx, gradgeom): 28 | (kept,) = ctx.saved_tensors 29 | back = torch.cumsum(kept, 0) 30 | back[kept] -= 1 31 | 32 | val = gradx[back] 33 | 34 | return val, None, None 35 | 36 | 37 | class QuickCumsumCuda(torch.autograd.Function): 38 | @staticmethod 39 | def forward(ctx, x, geom_feats, ranks, B, D, H, W): 40 | kept = torch.ones(x.shape[0], device=x.device, dtype=torch.bool) 41 | kept[1:] = ranks[1:] != ranks[:-1] 42 | interval_starts = torch.where(kept)[0].int() 43 | interval_lengths = torch.zeros_like(interval_starts) 44 | interval_lengths[:-1] = interval_starts[1:] - interval_starts[:-1] 45 | interval_lengths[-1] = x.shape[0] - interval_starts[-1] 46 | geom_feats = geom_feats.int() 47 | 48 | out = bev_pool_ext.bev_pool_forward( 49 | x, 50 | geom_feats, 51 | interval_lengths, 52 | interval_starts, 53 | B, 54 | D, 55 | H, 56 | W, 57 | ) 58 | 59 | ctx.save_for_backward(interval_starts, interval_lengths, geom_feats) 60 | ctx.saved_shapes = B, D, H, W 61 | return out 62 | 63 | @staticmethod 64 | def backward(ctx, out_grad): 65 | interval_starts, interval_lengths, geom_feats = ctx.saved_tensors 66 | B, D, H, W = ctx.saved_shapes 67 | 68 | out_grad = out_grad.contiguous() 69 | x_grad = bev_pool_ext.bev_pool_backward( 70 | out_grad, 71 | geom_feats, 72 | interval_lengths, 73 | interval_starts, 74 | B, 75 | D, 76 | H, 77 | W, 78 | ) 79 | 80 | return x_grad, None, None, None, None, None, None 81 | 82 | 83 | def bev_pool(feats, coords, B, D, H, W): 84 | assert feats.shape[0] == coords.shape[0] 85 | 86 | ranks = ( 87 | coords[:, 0] * (W * D * B) 88 | + coords[:, 1] * (D * B) 89 | + coords[:, 2] * B 90 | + coords[:, 3] 91 | ) 92 | indices = ranks.argsort() 93 | feats, coords, ranks = feats[indices], coords[indices], ranks[indices] 94 | 95 | x = QuickCumsumCuda.apply(feats, coords, ranks, B, D, H, W) 96 | x = x.permute(0, 4, 1, 2, 3).contiguous() 97 | return x 98 | -------------------------------------------------------------------------------- /pcdet/models/backbones_3d/focal_sparse_conv/SemanticSeg/pyramid_ffn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from .basic_blocks import BasicBlock2D 4 | from .sem_deeplabv3 import SemDeepLabV3 5 | 6 | class PyramidFeat2D(nn.Module): 7 | 8 | def __init__(self, optimize, model_cfg): 9 | """ 10 | Initialize 2D feature network via pretrained model 11 | Args: 12 | model_cfg: EasyDict, Dense classification network config 13 | """ 14 | super().__init__() 15 | self.model_cfg = model_cfg 16 | self.is_optimize = optimize 17 | 18 | # Create modules 19 | self.ifn = SemDeepLabV3( 20 | num_classes=model_cfg.num_class, 21 | backbone_name=model_cfg.backbone, 22 | **model_cfg.args 23 | ) 24 | self.reduce_blocks = torch.nn.ModuleList() 25 | self.out_channels = {} 26 | for _idx, _channel in enumerate(model_cfg.channel_reduce["in_channels"]): 27 | _channel_out = model_cfg.channel_reduce["out_channels"][_idx] 28 | self.out_channels[model_cfg.args['feat_extract_layer'][_idx]] = _channel_out 29 | block_cfg = {"in_channels": _channel, 30 | "out_channels": _channel_out, 31 | "kernel_size": model_cfg.channel_reduce["kernel_size"][_idx], 32 | "stride": model_cfg.channel_reduce["stride"][_idx], 33 | "bias": model_cfg.channel_reduce["bias"][_idx]} 34 | self.reduce_blocks.append(BasicBlock2D(**block_cfg)) 35 | 36 | def get_output_feature_dim(self): 37 | return self.out_channels 38 | 39 | def forward(self, images): 40 | """ 41 | Predicts depths and creates image depth feature volume using depth distributions 42 | Args: 43 | images: (N, 3, H_in, W_in), Input images 44 | Returns: 45 | batch_dict: 46 | frustum_features: (N, C, D, H_out, W_out), Image depth features 47 | """ 48 | # Pixel-wise depth classification 49 | batch_dict = {} 50 | ifn_result = self.ifn(images) 51 | 52 | for _idx, _layer in enumerate(self.model_cfg.args['feat_extract_layer']): 53 | image_features = ifn_result[_layer] 54 | # Channel reduce 55 | if self.reduce_blocks[_idx] is not None: 56 | image_features = self.reduce_blocks[_idx](image_features) 57 | 58 | batch_dict[_layer+"_feat2d"] = image_features 59 | 60 | if self.training: 61 | # detach feature from graph if not optimize 62 | if "logits" in ifn_result: 63 | ifn_result["logits"].detach_() 64 | if not self.is_optimize: 65 | image_features.detach_() 66 | 67 | return batch_dict 68 | 69 | def get_loss(self): 70 | """ 71 | Gets loss 72 | Args: 73 | Returns: 74 | loss: (1), Network loss 75 | tb_dict: dict[float], All losses to log in tensorboard 76 | """ 77 | return None, None 78 | -------------------------------------------------------------------------------- /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('----------- %s -----------' % (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.safe_load(f, Loader=yaml.FullLoader) 56 | except: 57 | yaml_config = yaml.safe_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.safe_load(f, Loader=yaml.FullLoader) 75 | except: 76 | new_config = yaml.safe_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 | -------------------------------------------------------------------------------- /tools/test.py: -------------------------------------------------------------------------------- 1 | import sys 2 | sys.path.insert(0, '../') 3 | import glob 4 | import os 5 | import re 6 | import time 7 | from tensorboardX import SummaryWriter 8 | 9 | from eval_utils import eval_utils 10 | from pcdet.config import cfg 11 | 12 | 13 | def get_no_evaluated_ckpt(ckpt_dir, ckpt_record_file, args): 14 | ckpt_list = glob.glob(os.path.join(ckpt_dir, '*checkpoint_epoch_*.pth')) 15 | ckpt_list.sort(key=os.path.getmtime) 16 | evaluated_ckpt_list = [float(x.strip()) for x in open(ckpt_record_file, 'r').readlines()] 17 | 18 | for cur_ckpt in ckpt_list: 19 | num_list = re.findall('checkpoint_epoch_(.*).pth', cur_ckpt) 20 | if num_list.__len__() == 0: 21 | continue 22 | 23 | epoch_id = num_list[-1] 24 | if 'optim' in epoch_id: 25 | continue 26 | if float(epoch_id) not in evaluated_ckpt_list and int(float(epoch_id)) >= args.start_epoch: 27 | return epoch_id, cur_ckpt 28 | return -1, None 29 | 30 | 31 | 32 | def repeat_eval_ckpt(model, test_loader, args, eval_output_dir, logger, ckpt_dir, dist_test=False): 33 | # evaluated ckpt record 34 | ckpt_record_file = eval_output_dir / ('eval_list_%s.txt' % cfg.DATA_CONFIG.DATA_SPLIT['test']) 35 | with open(ckpt_record_file, 'a'): 36 | pass 37 | 38 | # tensorboard log 39 | if cfg.LOCAL_RANK == 0: 40 | tb_log = SummaryWriter(log_dir=str(eval_output_dir / ('tensorboard_%s' % cfg.DATA_CONFIG.DATA_SPLIT['test']))) 41 | total_time = 0 42 | first_eval = True 43 | 44 | while True: 45 | # check whether there is checkpoint which is not evaluated 46 | cur_epoch_id, cur_ckpt = get_no_evaluated_ckpt(ckpt_dir, ckpt_record_file, args) 47 | if cur_epoch_id == -1 or int(float(cur_epoch_id)) < args.start_epoch: 48 | wait_second = 30 49 | if cfg.LOCAL_RANK == 0: 50 | print('Wait %s seconds for next check (progress: %.1f / %d minutes): %s \r' 51 | % (wait_second, total_time * 1.0 / 60, args.max_waiting_mins, ckpt_dir), end='', flush=True) 52 | time.sleep(wait_second) 53 | total_time += 30 54 | if total_time > args.max_waiting_mins * 60 and (first_eval is False): 55 | break 56 | continue 57 | 58 | total_time = 0 59 | first_eval = False 60 | 61 | model.load_params_from_file(filename=cur_ckpt, logger=logger, to_cpu=dist_test) 62 | model.cuda() 63 | 64 | # start evaluation 65 | cur_result_dir = eval_output_dir / ('epoch_%s' % cur_epoch_id) / cfg.DATA_CONFIG.DATA_SPLIT['test'] 66 | tb_dict = eval_utils.eval_one_epoch( 67 | cfg, args, model, test_loader, cur_epoch_id, logger, dist_test=dist_test, 68 | result_dir=cur_result_dir 69 | ) 70 | 71 | if cfg.LOCAL_RANK == 0: 72 | for key, val in tb_dict.items(): 73 | tb_log.add_scalar(key, val, cur_epoch_id) 74 | 75 | # record this epoch which has been evaluated 76 | with open(ckpt_record_file, 'a') as f: 77 | print('%s' % cur_epoch_id, file=f) 78 | logger.info('Epoch %s has been evaluated' % cur_epoch_id) 79 | -------------------------------------------------------------------------------- /pcdet/models/backbones_image/img_neck/generalized_lss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from ...model_utils.basic_block_2d import BasicBlock2D 5 | 6 | 7 | class GeneralizedLSSFPN(nn.Module): 8 | """ 9 | This module implements FPN, which creates pyramid features built on top of some input feature maps. 10 | This code is adapted from https://github.com/open-mmlab/mmdetection/blob/main/mmdet/models/necks/fpn.py with minimal modifications. 11 | """ 12 | def __init__(self, model_cfg): 13 | super().__init__() 14 | self.model_cfg = model_cfg 15 | in_channels = self.model_cfg.IN_CHANNELS 16 | out_channels = self.model_cfg.OUT_CHANNELS 17 | num_ins = len(in_channels) 18 | num_outs = self.model_cfg.NUM_OUTS 19 | start_level = self.model_cfg.START_LEVEL 20 | end_level = self.model_cfg.END_LEVEL 21 | 22 | self.in_channels = in_channels 23 | 24 | if end_level == -1: 25 | self.backbone_end_level = num_ins - 1 26 | else: 27 | self.backbone_end_level = end_level 28 | assert end_level <= len(in_channels) 29 | assert num_outs == end_level - start_level 30 | self.start_level = start_level 31 | self.end_level = end_level 32 | 33 | self.lateral_convs = nn.ModuleList() 34 | self.fpn_convs = nn.ModuleList() 35 | 36 | for i in range(self.start_level, self.backbone_end_level): 37 | l_conv = BasicBlock2D( 38 | in_channels[i] + (in_channels[i + 1] if i == self.backbone_end_level - 1 else out_channels), 39 | out_channels, kernel_size=1, bias = False 40 | ) 41 | fpn_conv = BasicBlock2D(out_channels,out_channels, kernel_size=3, padding=1, bias = False) 42 | self.lateral_convs.append(l_conv) 43 | self.fpn_convs.append(fpn_conv) 44 | 45 | def forward(self, batch_dict): 46 | """ 47 | Args: 48 | batch_dict: 49 | image_features (list[tensor]): Multi-stage features from image backbone. 50 | Returns: 51 | batch_dict: 52 | image_fpn (list(tensor)): FPN features. 53 | """ 54 | # upsample -> cat -> conv1x1 -> conv3x3 55 | inputs = batch_dict['image_features'] 56 | assert len(inputs) == len(self.in_channels) 57 | 58 | # build laterals 59 | laterals = [inputs[i + self.start_level] for i in range(len(inputs))] 60 | 61 | # build top-down path 62 | used_backbone_levels = len(laterals) - 1 63 | for i in range(used_backbone_levels - 1, -1, -1): 64 | x = F.interpolate( 65 | laterals[i + 1], 66 | size=laterals[i].shape[2:], 67 | mode='bilinear', align_corners=False, 68 | ) 69 | laterals[i] = torch.cat([laterals[i], x], dim=1) 70 | laterals[i] = self.lateral_convs[i](laterals[i]) 71 | laterals[i] = self.fpn_convs[i](laterals[i]) 72 | 73 | # build outputs 74 | outs = [laterals[i] for i in range(used_backbone_levels)] 75 | batch_dict['image_fpn'] = tuple(outs) 76 | return batch_dict 77 | -------------------------------------------------------------------------------- /tools/cfgs/dataset_configs/waymo_dataset_kp.yaml: -------------------------------------------------------------------------------- 1 | DATASET: 'WaymoDatasetKP' 2 | DATA_PATH: '../data/waymo' 3 | PROCESSED_DATA_TAG: 'waymo_processed_data_v0_5_0' 4 | 5 | POINT_CLOUD_RANGE: [-75.2, -75.2, -2, 75.2, 75.2, 4] 6 | 7 | DATA_SPLIT: { 8 | 'train': train, 9 | 'test': val 10 | } 11 | 12 | SAMPLED_INTERVAL: { 13 | 'train': 1, 14 | 'test': 1 15 | } 16 | 17 | FILTER_EMPTY_BOXES_FOR_TRAIN: True 18 | DISABLE_NLZ_FLAG_ON_POINTS: True 19 | 20 | USE_SHARED_MEMORY: False # it will load the data to shared memory to speed up (DO NOT USE IT IF YOU DO NOT FULLY UNDERSTAND WHAT WILL HAPPEN) 21 | SHARED_MEMORY_FILE_LIMIT: 35000 # set it based on the size of your shared memory 22 | 23 | DATA_AUGMENTOR: 24 | DISABLE_AUG_LIST: ['placeholder'] 25 | AUG_CONFIG_LIST: 26 | - NAME: gt_sampling 27 | IS_KEYPOINT: True 28 | USE_ROAD_PLANE: False 29 | DB_INFO_PATH: 30 | - waymo_processed_data_v0_5_0_waymo_dbinfos_train_sampled_1.pkl 31 | 32 | USE_SHARED_MEMORY: False # set it to True to speed up (it costs about 15GB shared memory) 33 | DB_DATA_PATH: 34 | - waymo_processed_data_v0_5_0_gt_database_train_sampled_1_global.npy 35 | 36 | BACKUP_DB_INFO: 37 | # if the above DB_INFO cannot be found, will use this backup one 38 | DB_INFO_PATH: waymo_processed_data_v0_5_0_waymo_dbinfos_train_sampled_1_multiframe_-4_to_0.pkl 39 | DB_DATA_PATH: waymo_processed_data_v0_5_0_gt_database_train_sampled_1_multiframe_-4_to_0_global.npy 40 | NUM_POINT_FEATURES: 6 41 | 42 | PREPARE: { 43 | filter_by_min_points: ['Vehicle:5', 'Pedestrian:5', 'Cyclist:5'], 44 | filter_by_difficulty: [-1], 45 | } 46 | 47 | # TODO: add more Vehicle to vary the background. 48 | SAMPLE_GROUPS: ['Vehicle:15', 'Pedestrian:10', 'Cyclist:10'] 49 | NUM_POINT_FEATURES: 5 50 | REMOVE_EXTRA_WIDTH: [0.0, 0.0, 0.0] 51 | LIMIT_WHOLE_SCENE: True 52 | 53 | - NAME: random_world_flip 54 | ALONG_AXIS_LIST: ['x', 'y'] 55 | 56 | - NAME: random_world_rotation 57 | WORLD_ROT_ANGLE: [-0.78539816, 0.78539816] 58 | 59 | - NAME: random_world_scaling 60 | WORLD_SCALE_RANGE: [0.95, 1.05] 61 | 62 | 63 | REMOVE_BOXES_WITHOUT_KEYPOINTS: true 64 | LABEL_MAPPING: 65 | STRATEGY: { 66 | 'Pedestrian': 'Human', 67 | 'Cyclist': 'Human', 68 | } # merge_class in waymo_dataset_kp 69 | RESULTS: ['Human'] 70 | 71 | 72 | POINT_FEATURE_ENCODING: { 73 | encoding_type: absolute_coordinates_encoding, 74 | used_feature_list: ['x', 'y', 'z', 'intensity', 'elongation'], 75 | src_feature_list: ['x', 'y', 'z', 'intensity', 'elongation'], 76 | } 77 | 78 | DATA_PROCESSOR: 79 | - NAME: mask_points_and_boxes_outside_range 80 | REMOVE_OUTSIDE_BOXES: True 81 | 82 | - NAME: shuffle_points 83 | SHUFFLE_ENABLED: { 84 | 'train': True, 85 | 'test': True 86 | } 87 | 88 | - NAME: transform_points_to_voxels 89 | VOXEL_SIZE: [0.1, 0.1, 0.1] 90 | MAX_POINTS_PER_VOXEL: 5 91 | MAX_NUMBER_OF_VOXELS: { 92 | 'train': 150000, 93 | 'test': 150000 94 | } 95 | -------------------------------------------------------------------------------- /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, **kwargs): 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 | -------------------------------------------------------------------------------- /pcdet/models/backbones_3d/vfe/dynamic_mean_vfe.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from .vfe_template import VFETemplate 4 | 5 | try: 6 | import torch_scatter 7 | except Exception as e: 8 | # Incase someone doesn't want to use dynamic pillar vfe and hasn't installed torch_scatter 9 | pass 10 | 11 | from .vfe_template import VFETemplate 12 | 13 | 14 | class DynamicMeanVFE(VFETemplate): 15 | def __init__(self, model_cfg, num_point_features, voxel_size, grid_size, point_cloud_range, **kwargs): 16 | super().__init__(model_cfg=model_cfg) 17 | self.num_point_features = num_point_features 18 | 19 | self.grid_size = torch.tensor(grid_size).cuda() 20 | self.voxel_size = torch.tensor(voxel_size).cuda() 21 | self.point_cloud_range = torch.tensor(point_cloud_range).cuda() 22 | 23 | self.voxel_x = voxel_size[0] 24 | self.voxel_y = voxel_size[1] 25 | self.voxel_z = voxel_size[2] 26 | self.x_offset = self.voxel_x / 2 + point_cloud_range[0] 27 | self.y_offset = self.voxel_y / 2 + point_cloud_range[1] 28 | self.z_offset = self.voxel_z / 2 + point_cloud_range[2] 29 | 30 | self.scale_xyz = grid_size[0] * grid_size[1] * grid_size[2] 31 | self.scale_yz = grid_size[1] * grid_size[2] 32 | self.scale_z = grid_size[2] 33 | 34 | def get_output_feature_dim(self): 35 | return self.num_point_features 36 | 37 | @torch.no_grad() 38 | def forward(self, batch_dict, **kwargs): 39 | """ 40 | Args: 41 | batch_dict: 42 | voxels: (num_voxels, max_points_per_voxel, C) 43 | voxel_num_points: optional (num_voxels) 44 | **kwargs: 45 | 46 | Returns: 47 | vfe_features: (num_voxels, C) 48 | """ 49 | batch_size = batch_dict['batch_size'] 50 | points = batch_dict['points'] # (batch_idx, x, y, z, i, e) 51 | 52 | # # debug 53 | point_coords = torch.floor((points[:, 1:4] - self.point_cloud_range[0:3]) / self.voxel_size).int() 54 | mask = ((point_coords >= 0) & (point_coords < self.grid_size)).all(dim=1) 55 | points = points[mask] 56 | point_coords = point_coords[mask] 57 | merge_coords = points[:, 0].int() * self.scale_xyz + \ 58 | point_coords[:, 0] * self.scale_yz + \ 59 | point_coords[:, 1] * self.scale_z + \ 60 | point_coords[:, 2] 61 | points_data = points[:, 1:].contiguous() 62 | 63 | unq_coords, unq_inv, unq_cnt = torch.unique(merge_coords, return_inverse=True, return_counts=True) 64 | 65 | points_mean = torch_scatter.scatter_mean(points_data, unq_inv, dim=0) 66 | 67 | unq_coords = unq_coords.int() 68 | voxel_coords = torch.stack((unq_coords // self.scale_xyz, 69 | (unq_coords % self.scale_xyz) // self.scale_yz, 70 | (unq_coords % self.scale_yz) // self.scale_z, 71 | unq_coords % self.scale_z), dim=1) 72 | voxel_coords = voxel_coords[:, [0, 3, 2, 1]] 73 | 74 | batch_dict['voxel_features'] = points_mean.contiguous() 75 | batch_dict['voxel_coords'] = voxel_coords.contiguous() 76 | return batch_dict 77 | -------------------------------------------------------------------------------- /pcdet/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from functools import partial 3 | from torch.utils.data import DataLoader 4 | from torch.utils.data import DistributedSampler as _DistributedSampler 5 | 6 | from pcdet.utils import common_utils 7 | 8 | from .dataset import DatasetTemplate 9 | from .kitti.kitti_dataset import KittiDataset 10 | from .nuscenes.nuscenes_dataset import NuScenesDataset 11 | from .waymo.waymo_dataset import WaymoDataset 12 | from .waymo.waymo_dataset_kp import WaymoDatasetKP 13 | from .pandaset.pandaset_dataset import PandasetDataset 14 | from .lyft.lyft_dataset import LyftDataset 15 | from .once.once_dataset import ONCEDataset 16 | from .argo2.argo2_dataset import Argo2Dataset 17 | from .custom.custom_dataset import CustomDataset 18 | 19 | __all__ = { 20 | 'DatasetTemplate': DatasetTemplate, 21 | 'KittiDataset': KittiDataset, 22 | 'NuScenesDataset': NuScenesDataset, 23 | 'WaymoDataset': WaymoDataset, 24 | 'WaymoDatasetKP': WaymoDatasetKP, 25 | 'PandasetDataset': PandasetDataset, 26 | 'LyftDataset': LyftDataset, 27 | 'ONCEDataset': ONCEDataset, 28 | 'CustomDataset': CustomDataset, 29 | 'Argo2Dataset': Argo2Dataset 30 | } 31 | 32 | 33 | class DistributedSampler(_DistributedSampler): 34 | 35 | def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True): 36 | super().__init__(dataset, num_replicas=num_replicas, rank=rank) 37 | self.shuffle = shuffle 38 | 39 | def __iter__(self): 40 | if self.shuffle: 41 | g = torch.Generator() 42 | g.manual_seed(self.epoch) 43 | indices = torch.randperm(len(self.dataset), generator=g).tolist() 44 | else: 45 | indices = torch.arange(len(self.dataset)).tolist() 46 | 47 | indices += indices[:(self.total_size - len(indices))] 48 | assert len(indices) == self.total_size 49 | 50 | indices = indices[self.rank:self.total_size:self.num_replicas] 51 | assert len(indices) == self.num_samples 52 | 53 | return iter(indices) 54 | 55 | 56 | def build_dataloader(dataset_cfg, class_names, batch_size, dist, root_path=None, workers=4, seed=None, 57 | logger=None, training=True, merge_all_iters_to_one_epoch=False, total_epochs=0): 58 | 59 | dataset = __all__[dataset_cfg.DATASET]( 60 | dataset_cfg=dataset_cfg, 61 | class_names=class_names, 62 | root_path=root_path, 63 | training=training, 64 | logger=logger, 65 | ) 66 | 67 | if merge_all_iters_to_one_epoch: 68 | assert hasattr(dataset, 'merge_all_iters_to_one_epoch') 69 | dataset.merge_all_iters_to_one_epoch(merge=True, epochs=total_epochs) 70 | 71 | if dist: 72 | if training: 73 | sampler = _DistributedSampler(dataset) 74 | else: 75 | rank, world_size = common_utils.get_dist_info() 76 | sampler = DistributedSampler(dataset, world_size, rank, shuffle=False) 77 | else: 78 | sampler = None 79 | dataloader = DataLoader( 80 | dataset, batch_size=batch_size, pin_memory=False, num_workers=workers, 81 | shuffle=(sampler is None) and training, collate_fn=dataset.collate_batch, 82 | drop_last=True if training else False, sampler=sampler, timeout=0, worker_init_fn=partial(common_utils.worker_init_fn, seed=seed) 83 | ) 84 | 85 | return dataset, dataloader, sampler 86 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /pcdet/utils/transform_utils.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | 4 | try: 5 | from kornia.geometry.conversions import ( 6 | convert_points_to_homogeneous, 7 | convert_points_from_homogeneous, 8 | ) 9 | except: 10 | pass 11 | # print('Warning: kornia is not installed. This package is only required by CaDDN') 12 | 13 | 14 | def project_to_image(project, points): 15 | """ 16 | Project points to image 17 | Args: 18 | project [torch.tensor(..., 3, 4)]: Projection matrix 19 | points [torch.Tensor(..., 3)]: 3D points 20 | Returns: 21 | points_img [torch.Tensor(..., 2)]: Points in image 22 | points_depth [torch.Tensor(...)]: Depth of each point 23 | """ 24 | # Reshape tensors to expected shape 25 | points = convert_points_to_homogeneous(points) 26 | points = points.unsqueeze(dim=-1) 27 | project = project.unsqueeze(dim=1) 28 | 29 | # Transform points to image and get depths 30 | points_t = project @ points 31 | points_t = points_t.squeeze(dim=-1) 32 | points_img = convert_points_from_homogeneous(points_t) 33 | points_depth = points_t[..., -1] - project[..., 2, 3] 34 | 35 | return points_img, points_depth 36 | 37 | 38 | def normalize_coords(coords, shape): 39 | """ 40 | Normalize coordinates of a grid between [-1, 1] 41 | Args: 42 | coords: (..., 3), Coordinates in grid 43 | shape: (3), Grid shape 44 | Returns: 45 | norm_coords: (.., 3), Normalized coordinates in grid 46 | """ 47 | min_n = -1 48 | max_n = 1 49 | shape = torch.flip(shape, dims=[0]) # Reverse ordering of shape 50 | 51 | # Subtract 1 since pixel indexing from [0, shape - 1] 52 | norm_coords = coords / (shape - 1) * (max_n - min_n) + min_n 53 | return norm_coords 54 | 55 | 56 | def bin_depths(depth_map, mode, depth_min, depth_max, num_bins, target=False): 57 | """ 58 | Converts depth map into bin indices 59 | Args: 60 | depth_map: (H, W), Depth Map 61 | mode: string, Discretiziation mode (See https://arxiv.org/pdf/2005.13423.pdf for more details) 62 | UD: Uniform discretiziation 63 | LID: Linear increasing discretiziation 64 | SID: Spacing increasing discretiziation 65 | depth_min: float, Minimum depth value 66 | depth_max: float, Maximum depth value 67 | num_bins: int, Number of depth bins 68 | target: bool, Whether the depth bins indices will be used for a target tensor in loss comparison 69 | Returns: 70 | indices: (H, W), Depth bin indices 71 | """ 72 | if mode == "UD": 73 | bin_size = (depth_max - depth_min) / num_bins 74 | indices = ((depth_map - depth_min) / bin_size) 75 | elif mode == "LID": 76 | bin_size = 2 * (depth_max - depth_min) / (num_bins * (1 + num_bins)) 77 | indices = -0.5 + 0.5 * torch.sqrt(1 + 8 * (depth_map - depth_min) / bin_size) 78 | elif mode == "SID": 79 | indices = num_bins * (torch.log(1 + depth_map) - math.log(1 + depth_min)) / \ 80 | (math.log(1 + depth_max) - math.log(1 + depth_min)) 81 | else: 82 | raise NotImplementedError 83 | 84 | if target: 85 | # Remove indicies outside of bounds 86 | mask = (indices < 0) | (indices > num_bins) | (~torch.isfinite(indices)) 87 | indices[mask] = num_bins 88 | 89 | # Convert to integer 90 | indices = indices.type(torch.int64) 91 | return indices 92 | -------------------------------------------------------------------------------- /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 | 39 | 40 | class PointPillarScatter3d(nn.Module): 41 | def __init__(self, model_cfg, grid_size, **kwargs): 42 | super().__init__() 43 | 44 | self.model_cfg = model_cfg 45 | self.nx, self.ny, self.nz = self.model_cfg.INPUT_SHAPE 46 | self.num_bev_features = self.model_cfg.NUM_BEV_FEATURES 47 | self.num_bev_features_before_compression = self.model_cfg.NUM_BEV_FEATURES // self.nz 48 | 49 | def forward(self, batch_dict, **kwargs): 50 | pillar_features, coords = batch_dict['pillar_features'], batch_dict['voxel_coords'] 51 | 52 | batch_spatial_features = [] 53 | batch_size = coords[:, 0].max().int().item() + 1 54 | for batch_idx in range(batch_size): 55 | spatial_feature = torch.zeros( 56 | self.num_bev_features_before_compression, 57 | self.nz * self.nx * self.ny, 58 | dtype=pillar_features.dtype, 59 | device=pillar_features.device) 60 | 61 | batch_mask = coords[:, 0] == batch_idx 62 | this_coords = coords[batch_mask, :] 63 | indices = this_coords[:, 1] * self.ny * self.nx + this_coords[:, 2] * self.nx + this_coords[:, 3] 64 | indices = indices.type(torch.long) 65 | pillars = pillar_features[batch_mask, :] 66 | pillars = pillars.t() 67 | spatial_feature[:, indices] = pillars 68 | batch_spatial_features.append(spatial_feature) 69 | 70 | batch_spatial_features = torch.stack(batch_spatial_features, 0) 71 | batch_spatial_features = batch_spatial_features.view(batch_size, self.num_bev_features_before_compression * self.nz, self.ny, self.nx) 72 | batch_dict['spatial_features'] = batch_spatial_features 73 | return batch_dict -------------------------------------------------------------------------------- /pcdet/ops/pointnet2/pointnet2_stack/src/vector_pool_gpu.h: -------------------------------------------------------------------------------- 1 | /* 2 | Vector-pool aggregation based local feature aggregation for point cloud. 3 | PV-RCNN++: Point-Voxel Feature Set Abstraction With Local Vector Representation for 3D Object Detection 4 | https://arxiv.org/abs/2102.00463 5 | 6 | Written by Shaoshuai Shi 7 | All Rights Reserved 2020. 8 | */ 9 | 10 | 11 | #ifndef _STACK_VECTOR_POOL_GPU_H 12 | #define _STACK_VECTOR_POOL_GPU_H 13 | 14 | #include 15 | #include 16 | #include 17 | #include 18 | 19 | 20 | int query_stacked_local_neighbor_idxs_kernel_launcher_stack( 21 | const float *support_xyz, const int *xyz_batch_cnt, const float *new_xyz, const int *new_xyz_batch_cnt, 22 | int *stack_neighbor_idxs, int *start_len, int *cumsum, int avg_length_of_neighbor_idxs, 23 | float max_neighbour_distance, int batch_size, int M, int nsample, int neighbor_type); 24 | 25 | int query_stacked_local_neighbor_idxs_wrapper_stack(at::Tensor support_xyz_tensor, at::Tensor xyz_batch_cnt_tensor, 26 | at::Tensor new_xyz_tensor, at::Tensor new_xyz_batch_cnt_tensor, 27 | at::Tensor stack_neighbor_idxs_tensor, at::Tensor start_len_tensor, at::Tensor cumsum_tensor, 28 | int avg_length_of_neighbor_idxs, float max_neighbour_distance, int nsample, int neighbor_type); 29 | 30 | 31 | int query_three_nn_by_stacked_local_idxs_kernel_launcher_stack( 32 | const float *support_xyz, const float *new_xyz, const float *new_xyz_grid_centers, 33 | int *new_xyz_grid_idxs, float *new_xyz_grid_dist2, 34 | const int *stack_neighbor_idxs, const int *start_len, 35 | int M, int num_total_grids); 36 | 37 | int query_three_nn_by_stacked_local_idxs_wrapper_stack(at::Tensor support_xyz_tensor, 38 | at::Tensor new_xyz_tensor, at::Tensor new_xyz_grid_centers_tensor, 39 | at::Tensor new_xyz_grid_idxs_tensor, at::Tensor new_xyz_grid_dist2_tensor, 40 | at::Tensor stack_neighbor_idxs_tensor, at::Tensor start_len_tensor, 41 | int M, int num_total_grids); 42 | 43 | 44 | int vector_pool_wrapper_stack(at::Tensor support_xyz_tensor, at::Tensor xyz_batch_cnt_tensor, 45 | at::Tensor support_features_tensor, at::Tensor new_xyz_tensor, at::Tensor new_xyz_batch_cnt_tensor, 46 | at::Tensor new_features_tensor, at::Tensor new_local_xyz, 47 | at::Tensor point_cnt_of_grid_tensor, at::Tensor grouped_idxs_tensor, 48 | int num_grid_x, int num_grid_y, int num_grid_z, float max_neighbour_distance, int use_xyz, 49 | int num_max_sum_points, int nsample, int neighbor_type, int pooling_type); 50 | 51 | 52 | int vector_pool_kernel_launcher_stack( 53 | const float *support_xyz, const float *support_features, const int *xyz_batch_cnt, 54 | const float *new_xyz, float *new_features, float * new_local_xyz, const int *new_xyz_batch_cnt, 55 | int *point_cnt_of_grid, int *grouped_idxs, 56 | int num_grid_x, int num_grid_y, int num_grid_z, float max_neighbour_distance, 57 | int batch_size, int N, int M, int num_c_in, int num_c_out, int num_total_grids, int use_xyz, 58 | int num_max_sum_points, int nsample, int neighbor_type, int pooling_type); 59 | 60 | 61 | int vector_pool_grad_wrapper_stack(at::Tensor grad_new_features_tensor, 62 | at::Tensor point_cnt_of_grid_tensor, at::Tensor grouped_idxs_tensor, 63 | at::Tensor grad_support_features_tensor); 64 | 65 | 66 | void vector_pool_grad_kernel_launcher_stack( 67 | const float *grad_new_features, const int *point_cnt_of_grid, const int *grouped_idxs, 68 | float *grad_support_features, int N, int M, int num_c_out, int num_c_in, int num_total_grids, 69 | int num_max_sum_points); 70 | 71 | #endif 72 | -------------------------------------------------------------------------------- /tools/process_tools/create_integrated_database.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pickle as pkl 3 | from pathlib import Path 4 | import tqdm 5 | import copy 6 | import os 7 | 8 | 9 | def create_integrated_db_with_infos(args, root_path): 10 | """ 11 | Args: 12 | args: 13 | Returns: 14 | 15 | """ 16 | # prepare 17 | db_infos_path = args.src_db_info 18 | db_info_global_path = db_infos_path 19 | global_db_path = root_path / (args.new_db_name + '.npy') 20 | 21 | db_infos = pkl.load(open(db_infos_path, 'rb')) 22 | db_info_global = copy.deepcopy(db_infos) 23 | start_idx = 0 24 | global_db_list = [] 25 | 26 | for category, class_info in db_infos.items(): 27 | print('>>> Start processing %s' % category) 28 | for idx, info in tqdm.tqdm(enumerate(class_info), total=len(class_info)): 29 | obj_path = root_path / info['path'] 30 | obj_points = np.fromfile(str(obj_path), dtype=np.float32).reshape( 31 | [-1, args.num_point_features]) 32 | num_points = obj_points.shape[0] 33 | if num_points != info['num_points_in_gt']: 34 | obj_points = np.fromfile(str(obj_path), dtype=np.float64).reshape([-1, args.num_point_features]) 35 | num_points = obj_points.shape[0] 36 | obj_points = obj_points.astype(np.float32) 37 | assert num_points == info['num_points_in_gt'] 38 | 39 | db_info_global[category][idx]['global_data_offset'] = (start_idx, start_idx + num_points) 40 | start_idx += num_points 41 | global_db_list.append(obj_points) 42 | 43 | global_db = np.concatenate(global_db_list) 44 | 45 | with open(global_db_path, 'wb') as f: 46 | np.save(f, global_db) 47 | 48 | with open(db_info_global_path, 'wb') as f: 49 | pkl.dump(db_info_global, f) 50 | 51 | print(f"Successfully create integrated database at {global_db_path}") 52 | print(f"Successfully create integrated database info at {db_info_global_path}") 53 | 54 | return db_info_global, global_db 55 | 56 | 57 | def verify(info, whole_db, root_path, num_point_features): 58 | obj_path = root_path / info['path'] 59 | obj_points = np.fromfile(str(obj_path), dtype=np.float32).reshape([-1, num_point_features]) 60 | mean_origin = obj_points.mean() 61 | 62 | start_idx, end_idx = info['global_data_offset'] 63 | obj_points_new = whole_db[start_idx:end_idx] 64 | mean_new = obj_points_new.mean() 65 | 66 | assert mean_origin == mean_new 67 | 68 | print("Verification pass!") 69 | 70 | 71 | if __name__ == '__main__': 72 | import argparse 73 | 74 | parser = argparse.ArgumentParser(description='arg parser') 75 | parser.add_argument('--src_db_info', type=str, default='../../data/waymo/waymo_processed_data_v0_5_0_waymo_dbinfos_train_sampled_1_multiframe_-4_to_0_tail_parallel.pkl', help='') 76 | parser.add_argument('--new_db_name', type=str, default='waymo_processed_data_v0_5_0_gt_database_train_sampled_1_multiframe_-4_to_0_tail_parallel_global', help='') 77 | parser.add_argument('--num_point_features', type=int, default=6, help='number of feature channels for points') 78 | parser.add_argument('--class_name', type=str, default='Vehicle', help='category name for verification') 79 | 80 | args = parser.parse_args() 81 | 82 | root_path = Path(os.path.dirname(args.src_db_info)) 83 | 84 | db_infos_global, whole_db = create_integrated_db_with_infos(args, root_path) 85 | # simple verify 86 | verify(db_infos_global[args.class_name][0], whole_db, root_path, args.num_point_features) 87 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # VoxelKP: A Voxel-based Network Architecture for Human Keypoint Estimation in LiDAR Data 2 | 3 | We present VoxelKP, a novel fully sparse network architecture tailored for human keypoint estimation in LiDAR data. We achieved the state-of-the-art performances without additional data. 4 | 5 | ## Performance Demo 6 | 7 | [![Watch the video](https://img.youtube.com/vi/u-xHv_OAO0M/hqdefault.jpg)](https://youtu.be/u-xHv_OAO0M) 8 | 9 | ## Installation 10 | 11 | Please refer to the installation of `OpenPCDet`. 12 | 13 | ## Getting Started 14 | 15 | Please refer to [GETTING_STARTED.md](docs/GETTING_STARTED.md) to learn more usage about this project. 16 | 17 | This is the inference only code. 18 | 19 | Checkpoints can be downloaded from [here](https://huggingface.co/shijianjian/VoxelKP). 20 | 21 | ```bash 22 | $ cd tools 23 | $ python waymo_visualizer.py --ckpt CHECKPOINT 24 | ``` 25 | 26 | ## Train 27 | 28 | 1. Prepare dataset, referring to [OpenPCDet](https://github.com/open-mmlab/OpenPCDet/blob/master/docs/GETTING_STARTED.md). 29 | 30 | 2. Install the module for compling. 31 | ```bash 32 | $ pip install -e . 33 | ``` 34 | 35 | 3. The training script. 36 | 37 | ```bash 38 | $ cd tools 39 | $ python train.py --cfg_file ./cfgs/waymo_models/kp_effv2next4_voxelnext_iou_aug_bev_channel.yaml --epochs 20 40 | ``` 41 | 42 | if using multiple GPUs 43 | 44 | ```bash 45 | $ bash scripts/dist_train.sh 8 --cfg_file ./cfgs/waymo_models/kp_effv2next4_voxelnext_iou_aug_bev_channel.yaml --epochs 36 --workers 0 46 | ``` 47 | 48 | 49 | ## Benchmarks 50 | 51 | There is a limited number of relevant research for this task. Most of the prior works utilize additional training data beyond the 3D keypoint data within the Waymo dataset. To provide a fair comparison, we need to consider approaches that use extra data and those that rely solely on Waymo ground truth separately. 52 | 53 | 54 | 55 | 56 | ## Experiment Results 57 | 58 | We report the full spectrum of the evaluation, including MPJPE, OKS@AP, and PEM. 59 | 60 |
61 | 62 | | Part | MPJPE | OKS@KP | PEM | 63 | |-----------|:------:|:-------:|:----:| 64 | |Head | 0.0570 | 0.6393 | 0.1569 | 65 | |Shoulders | 0.0669 | 0.8917 | 0.1563 | 66 | |Elbows | 0.0948 | 0.7197 | 0.1746 | 67 | |Wrists | 0.1467 | 0.3791 | 0.1987 | 68 | |Hips | 0.0670 | 0.9533 | 0.1576 | 69 | |Knees | 0.0820 | 0.8586 | 0.1660 | 70 | |Ankles | 0.1084 | 0.7581 | 0.1765 | 71 | |All | 0.0887 | 0.7300 | 0.1695 | 72 | 73 |
74 | 75 | Our visual results show that our VoxelKP offers improved keypoint estimation with precise locations and fewer false positives. 76 | 77 | 78 | 79 | A visual demonstration of our baseline model (top) and the proposed VoxelKP (bottom). The insets are color-coded according to the legend in the figure. In the green-colored insets, a comparison with the ground truth is shown, with ground truth in red and predictions in blue. 80 | 81 | 82 | ## Architecture 83 | 84 | 85 | 86 | ## Cite our work 87 | ``` 88 | @inproceedings{shi2025voxelkp, 89 | title={VoxelKP: A Voxel-based Network Architecture for Human Keypoint Estimation in LiDAR Data}, 90 | author={Shi, Jian and Wonka, Peter}, 91 | booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision}, 92 | pages={28282--28291}, 93 | year={2025} 94 | } 95 | ``` 96 | 97 | ## Acknowledgement 98 | This repository is built on top of `OpenPCDet` and `VoxelNeXt`. 99 | 100 | We use `sptr` implementation from [here](https://github.com/dvlab-research/SparseTransformer). 101 | -------------------------------------------------------------------------------- /pcdet/ops/bev_pool/src/bev_pool.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | // CUDA function declarations 5 | void bev_pool(int b, int d, int h, int w, int n, int c, int n_intervals, const float* x, 6 | const int* geom_feats, const int* interval_starts, const int* interval_lengths, float* out); 7 | 8 | void bev_pool_grad(int b, int d, int h, int w, int n, int c, int n_intervals, const float* out_grad, 9 | const int* geom_feats, const int* interval_starts, const int* interval_lengths, float* x_grad); 10 | 11 | 12 | /* 13 | Function: pillar pooling (forward, cuda) 14 | Args: 15 | x : input features, FloatTensor[n, c] 16 | geom_feats : input coordinates, IntTensor[n, 4] 17 | interval_lengths : starting position for pooled point, IntTensor[n_intervals] 18 | interval_starts : how many points in each pooled point, IntTensor[n_intervals] 19 | Return: 20 | out : output features, FloatTensor[b, d, h, w, c] 21 | */ 22 | at::Tensor bev_pool_forward( 23 | const at::Tensor _x, 24 | const at::Tensor _geom_feats, 25 | const at::Tensor _interval_lengths, 26 | const at::Tensor _interval_starts, 27 | int b, int d, int h, int w 28 | ) { 29 | int n = _x.size(0); 30 | int c = _x.size(1); 31 | int n_intervals = _interval_lengths.size(0); 32 | const at::cuda::OptionalCUDAGuard device_guard(device_of(_x)); 33 | const float* x = _x.data_ptr(); 34 | const int* geom_feats = _geom_feats.data_ptr(); 35 | const int* interval_lengths = _interval_lengths.data_ptr(); 36 | const int* interval_starts = _interval_starts.data_ptr(); 37 | 38 | auto options = 39 | torch::TensorOptions().dtype(_x.dtype()).device(_x.device()); 40 | at::Tensor _out = torch::zeros({b, d, h, w, c}, options); 41 | float* out = _out.data_ptr(); 42 | bev_pool( 43 | b, d, h, w, n, c, n_intervals, x, 44 | geom_feats, interval_starts, interval_lengths, out 45 | ); 46 | return _out; 47 | } 48 | 49 | 50 | /* 51 | Function: pillar pooling (backward, cuda) 52 | Args: 53 | out_grad : input features, FloatTensor[b, d, h, w, c] 54 | geom_feats : input coordinates, IntTensor[n, 4] 55 | interval_lengths : starting position for pooled point, IntTensor[n_intervals] 56 | interval_starts : how many points in each pooled point, IntTensor[n_intervals] 57 | Return: 58 | x_grad : output features, FloatTensor[n, 4] 59 | */ 60 | at::Tensor bev_pool_backward( 61 | const at::Tensor _out_grad, 62 | const at::Tensor _geom_feats, 63 | const at::Tensor _interval_lengths, 64 | const at::Tensor _interval_starts, 65 | int b, int d, int h, int w 66 | ) { 67 | int n = _geom_feats.size(0); 68 | int c = _out_grad.size(4); 69 | int n_intervals = _interval_lengths.size(0); 70 | const at::cuda::OptionalCUDAGuard device_guard(device_of(_out_grad)); 71 | const float* out_grad = _out_grad.data_ptr(); 72 | const int* geom_feats = _geom_feats.data_ptr(); 73 | const int* interval_lengths = _interval_lengths.data_ptr(); 74 | const int* interval_starts = _interval_starts.data_ptr(); 75 | 76 | auto options = 77 | torch::TensorOptions().dtype(_out_grad.dtype()).device(_out_grad.device()); 78 | at::Tensor _x_grad = torch::zeros({n, c}, options); 79 | float* x_grad = _x_grad.data_ptr(); 80 | 81 | bev_pool_grad( 82 | b, d, h, w, n, c, n_intervals, out_grad, 83 | geom_feats, interval_starts, interval_lengths, x_grad 84 | ); 85 | 86 | return _x_grad; 87 | } 88 | 89 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 90 | m.def("bev_pool_forward", &bev_pool_forward, 91 | "bev_pool_forward"); 92 | m.def("bev_pool_backward", &bev_pool_backward, 93 | "bev_pool_backward"); 94 | } 95 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /pcdet/utils/object3d_custom.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_custom_obj_level() 37 | 38 | def get_custom_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_custom_format(self): 79 | custom_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 custom_str 84 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /tools/cfgs/dataset_configs/waymo_dataset_kp_augv2.yaml: -------------------------------------------------------------------------------- 1 | DATASET: 'WaymoDatasetKP' 2 | DATA_PATH: '../data/waymo' 3 | PROCESSED_DATA_TAG: 'waymo_processed_data_v0_5_0' 4 | 5 | POINT_CLOUD_RANGE: [-75.2, -75.2, -2, 75.2, 75.2, 4] 6 | 7 | DATA_SPLIT: { 8 | 'train': train, 9 | 'test': val 10 | } 11 | 12 | SAMPLED_INTERVAL: { 13 | 'train': 1, 14 | 'test': 1 15 | } 16 | 17 | FILTER_EMPTY_BOXES_FOR_TRAIN: True 18 | DISABLE_NLZ_FLAG_ON_POINTS: True 19 | 20 | USE_SHARED_MEMORY: False # it will load the data to shared memory to speed up (DO NOT USE IT IF YOU DO NOT FULLY UNDERSTAND WHAT WILL HAPPEN) 21 | SHARED_MEMORY_FILE_LIMIT: 35000 # set it based on the size of your shared memory 22 | 23 | DATA_AUGMENTOR: 24 | DISABLE_AUG_LIST: ['placeholder'] 25 | AUG_CONFIG_LIST: 26 | - NAME: gt_sampling 27 | IS_KEYPOINT: True 28 | USE_ROAD_PLANE: False 29 | DB_INFO_PATH: 30 | - waymo_processed_data_v0_5_0_waymo_dbinfos_train_sampled_1.pkl 31 | 32 | USE_SHARED_MEMORY: False # set it to True to speed up (it costs about 15GB shared memory) 33 | DB_DATA_PATH: 34 | - waymo_processed_data_v0_5_0_gt_database_train_sampled_1_global.npy 35 | 36 | BACKUP_DB_INFO: 37 | # if the above DB_INFO cannot be found, will use this backup one 38 | DB_INFO_PATH: waymo_processed_data_v0_5_0_waymo_dbinfos_train_sampled_1_multiframe_-4_to_0.pkl 39 | DB_DATA_PATH: waymo_processed_data_v0_5_0_gt_database_train_sampled_1_multiframe_-4_to_0_global.npy 40 | NUM_POINT_FEATURES: 6 41 | 42 | PREPARE: { 43 | filter_by_min_points: ['Vehicle:5', 'Pedestrian:5', 'Cyclist:5'], 44 | filter_by_difficulty: [-1], 45 | } 46 | 47 | # TODO: add more Vehicle to vary the background. 48 | SAMPLE_GROUPS: ['Vehicle:15', 'Pedestrian:10', 'Cyclist:10'] 49 | NUM_POINT_FEATURES: 5 50 | REMOVE_EXTRA_WIDTH: [0.0, 0.0, 0.0] 51 | LIMIT_WHOLE_SCENE: True 52 | 53 | - NAME: random_world_flip 54 | ALONG_AXIS_LIST: ['x', 'y'] 55 | 56 | - NAME: random_world_rotation 57 | WORLD_ROT_ANGLE: [-0.78539816, 0.78539816] 58 | 59 | - NAME: random_world_scaling 60 | WORLD_SCALE_RANGE: [0.95, 1.05] 61 | 62 | - NAME: random_local_rotation 63 | LOCAL_ROT_ANGLE: [-0.15707963267, 0.15707963267] 64 | 65 | - NAME: random_local_scaling 66 | LOCAL_SCALE_RANGE: [0.95, 1.05] 67 | 68 | - NAME: random_local_inter_object_noise 69 | PROBABILITY: 0.3 70 | POINT_RATIO: 0.1 71 | 72 | - NAME: random_local_frustum_dropout 73 | PROBABILITY: 0.05 74 | INTENSITY_RANGE: [ 0, 0.2 ] 75 | DIRECTION: ['top', 'bottom', 'left', 'right'] 76 | 77 | - NAME: random_local_noise 78 | PROBABILITY: 0.3 79 | JITTERING_RATIO: 0.2 80 | POINT_RATIO: 0.2 81 | 82 | 83 | REMOVE_BOXES_WITHOUT_KEYPOINTS: true 84 | LABEL_MAPPING: 85 | STRATEGY: { 86 | 'Pedestrian': 'Human', 87 | 'Cyclist': 'Human', 88 | } # merge_class in waymo_dataset_kp 89 | RESULTS: ['Human'] 90 | 91 | 92 | POINT_FEATURE_ENCODING: { 93 | encoding_type: absolute_coordinates_encoding, 94 | used_feature_list: ['x', 'y', 'z', 'intensity', 'elongation'], 95 | src_feature_list: ['x', 'y', 'z', 'intensity', 'elongation'], 96 | } 97 | 98 | DATA_PROCESSOR: 99 | - NAME: mask_points_and_boxes_outside_range 100 | REMOVE_OUTSIDE_BOXES: True 101 | 102 | - NAME: shuffle_points 103 | SHUFFLE_ENABLED: { 104 | 'train': True, 105 | 'test': True 106 | } 107 | 108 | - NAME: transform_points_to_voxels 109 | VOXEL_SIZE: [0.1, 0.1, 0.1] 110 | MAX_POINTS_PER_VOXEL: 5 111 | MAX_NUMBER_OF_VOXELS: { 112 | 'train': 150000, 113 | 'test': 150000 114 | } 115 | -------------------------------------------------------------------------------- /pcdet/models/detectors/bevfusion.py: -------------------------------------------------------------------------------- 1 | from .detector3d_template import Detector3DTemplate 2 | from .. import backbones_image, view_transforms 3 | from ..backbones_image import img_neck 4 | from ..backbones_2d import fuser 5 | 6 | class BevFusion(Detector3DTemplate): 7 | def __init__(self, model_cfg, num_class, dataset): 8 | super().__init__(model_cfg=model_cfg, num_class=num_class, dataset=dataset) 9 | self.module_topology = [ 10 | 'vfe', 'backbone_3d', 'map_to_bev_module', 'pfe', 11 | 'image_backbone','neck','vtransform','fuser', 12 | 'backbone_2d', 'dense_head', 'point_head', 'roi_head' 13 | ] 14 | self.module_list = self.build_networks() 15 | 16 | def build_neck(self,model_info_dict): 17 | if self.model_cfg.get('NECK', None) is None: 18 | return None, model_info_dict 19 | neck_module = img_neck.__all__[self.model_cfg.NECK.NAME]( 20 | model_cfg=self.model_cfg.NECK 21 | ) 22 | model_info_dict['module_list'].append(neck_module) 23 | 24 | return neck_module, model_info_dict 25 | 26 | def build_vtransform(self,model_info_dict): 27 | if self.model_cfg.get('VTRANSFORM', None) is None: 28 | return None, model_info_dict 29 | 30 | vtransform_module = view_transforms.__all__[self.model_cfg.VTRANSFORM.NAME]( 31 | model_cfg=self.model_cfg.VTRANSFORM 32 | ) 33 | model_info_dict['module_list'].append(vtransform_module) 34 | 35 | return vtransform_module, model_info_dict 36 | 37 | def build_image_backbone(self, model_info_dict): 38 | if self.model_cfg.get('IMAGE_BACKBONE', None) is None: 39 | return None, model_info_dict 40 | image_backbone_module = backbones_image.__all__[self.model_cfg.IMAGE_BACKBONE.NAME]( 41 | model_cfg=self.model_cfg.IMAGE_BACKBONE 42 | ) 43 | image_backbone_module.init_weights() 44 | model_info_dict['module_list'].append(image_backbone_module) 45 | 46 | return image_backbone_module, model_info_dict 47 | 48 | def build_fuser(self, model_info_dict): 49 | if self.model_cfg.get('FUSER', None) is None: 50 | return None, model_info_dict 51 | 52 | fuser_module = fuser.__all__[self.model_cfg.FUSER.NAME]( 53 | model_cfg=self.model_cfg.FUSER 54 | ) 55 | model_info_dict['module_list'].append(fuser_module) 56 | model_info_dict['num_bev_features'] = self.model_cfg.FUSER.OUT_CHANNEL 57 | return fuser_module, model_info_dict 58 | 59 | def forward(self, batch_dict): 60 | 61 | for i,cur_module in enumerate(self.module_list): 62 | batch_dict = cur_module(batch_dict) 63 | 64 | if self.training: 65 | loss, tb_dict, disp_dict = self.get_training_loss(batch_dict) 66 | 67 | ret_dict = { 68 | 'loss': loss 69 | } 70 | return ret_dict, tb_dict, disp_dict 71 | else: 72 | pred_dicts, recall_dicts = self.post_processing(batch_dict) 73 | return pred_dicts, recall_dicts 74 | 75 | def get_training_loss(self,batch_dict): 76 | disp_dict = {} 77 | 78 | loss_trans, tb_dict = batch_dict['loss'],batch_dict['tb_dict'] 79 | tb_dict = { 80 | 'loss_trans': loss_trans.item(), 81 | **tb_dict 82 | } 83 | 84 | loss = loss_trans 85 | return loss, tb_dict, disp_dict 86 | 87 | def post_processing(self, batch_dict): 88 | post_process_cfg = self.model_cfg.POST_PROCESSING 89 | batch_size = batch_dict['batch_size'] 90 | final_pred_dict = batch_dict['final_box_dicts'] 91 | recall_dict = {} 92 | for index in range(batch_size): 93 | pred_boxes = final_pred_dict[index]['pred_boxes'] 94 | 95 | recall_dict = self.generate_recall_record( 96 | box_preds=pred_boxes, 97 | recall_dict=recall_dict, batch_index=index, data_dict=batch_dict, 98 | thresh_list=post_process_cfg.RECALL_THRESH_LIST 99 | ) 100 | 101 | return final_pred_dict, recall_dict 102 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /pcdet/models/backbones_3d/vfe/image_vfe_modules/ffn/depth_ffn.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.nn.functional as F 3 | 4 | from . import ddn, ddn_loss 5 | from pcdet.models.model_utils.basic_block_2d import BasicBlock2D 6 | 7 | 8 | class DepthFFN(nn.Module): 9 | 10 | def __init__(self, model_cfg, downsample_factor): 11 | """ 12 | Initialize frustum feature network via depth distribution estimation 13 | Args: 14 | model_cfg: EasyDict, Depth classification network config 15 | downsample_factor: int, Depth map downsample factor 16 | """ 17 | super().__init__() 18 | self.model_cfg = model_cfg 19 | self.disc_cfg = model_cfg.DISCRETIZE 20 | self.downsample_factor = downsample_factor 21 | 22 | # Create modules 23 | self.ddn = ddn.__all__[model_cfg.DDN.NAME]( 24 | num_classes=self.disc_cfg["num_bins"] + 1, 25 | backbone_name=model_cfg.DDN.BACKBONE_NAME, 26 | **model_cfg.DDN.ARGS 27 | ) 28 | self.channel_reduce = BasicBlock2D(**model_cfg.CHANNEL_REDUCE) 29 | self.ddn_loss = ddn_loss.__all__[model_cfg.LOSS.NAME]( 30 | disc_cfg=self.disc_cfg, 31 | downsample_factor=downsample_factor, 32 | **model_cfg.LOSS.ARGS 33 | ) 34 | self.forward_ret_dict = {} 35 | 36 | def get_output_feature_dim(self): 37 | return self.channel_reduce.out_channels 38 | 39 | def forward(self, batch_dict): 40 | """ 41 | Predicts depths and creates image depth feature volume using depth distributions 42 | Args: 43 | batch_dict: 44 | images: (N, 3, H_in, W_in), Input images 45 | Returns: 46 | batch_dict: 47 | frustum_features: (N, C, D, H_out, W_out), Image depth features 48 | """ 49 | # Pixel-wise depth classification 50 | images = batch_dict["images"] 51 | ddn_result = self.ddn(images) 52 | image_features = ddn_result["features"] 53 | depth_logits = ddn_result["logits"] 54 | 55 | # Channel reduce 56 | if self.channel_reduce is not None: 57 | image_features = self.channel_reduce(image_features) 58 | 59 | # Create image feature plane-sweep volume 60 | frustum_features = self.create_frustum_features(image_features=image_features, 61 | depth_logits=depth_logits) 62 | batch_dict["frustum_features"] = frustum_features 63 | 64 | if self.training: 65 | self.forward_ret_dict["depth_maps"] = batch_dict["depth_maps"] 66 | self.forward_ret_dict["gt_boxes2d"] = batch_dict["gt_boxes2d"] 67 | self.forward_ret_dict["depth_logits"] = depth_logits 68 | return batch_dict 69 | 70 | def create_frustum_features(self, image_features, depth_logits): 71 | """ 72 | Create image depth feature volume by multiplying image features with depth distributions 73 | Args: 74 | image_features: (N, C, H, W), Image features 75 | depth_logits: (N, D+1, H, W), Depth classification logits 76 | Returns: 77 | frustum_features: (N, C, D, H, W), Image features 78 | """ 79 | channel_dim = 1 80 | depth_dim = 2 81 | 82 | # Resize to match dimensions 83 | image_features = image_features.unsqueeze(depth_dim) 84 | depth_logits = depth_logits.unsqueeze(channel_dim) 85 | 86 | # Apply softmax along depth axis and remove last depth category (> Max Range) 87 | depth_probs = F.softmax(depth_logits, dim=depth_dim) 88 | depth_probs = depth_probs[:, :, :-1] 89 | 90 | # Multiply to form image depth feature volume 91 | frustum_features = depth_probs * image_features 92 | return frustum_features 93 | 94 | def get_loss(self): 95 | """ 96 | Gets DDN loss 97 | Args: 98 | Returns: 99 | loss: (1), Depth distribution network loss 100 | tb_dict: dict[float], All losses to log in tensorboard 101 | """ 102 | loss, tb_dict = self.ddn_loss(**self.forward_ret_dict) 103 | return loss, tb_dict 104 | -------------------------------------------------------------------------------- /tools/cfgs/waymo_models/kp_effv2next4_voxelnext_iou_aug_bev_channel.yaml: -------------------------------------------------------------------------------- 1 | CLASS_NAMES: ['Pedestrian', 'Cyclist'] 2 | 3 | DATA_CONFIG: 4 | _BASE_CONFIG_: cfgs/dataset_configs/waymo_dataset_kp_augv2.yaml 5 | # _BASE_CONFIG_: cfgs/dataset_configs/waymo_dataset_kp.yaml 6 | SAMPLED_INTERVAL: { 7 | 'train': 1, 8 | 'test': 1 9 | } 10 | 11 | MODEL: 12 | NAME: VoxelNeXt 13 | 14 | VFE: 15 | NAME: MeanVFE 16 | 17 | BACKBONE_3D: 18 | NAME: VoxelResBackBone8xVoxelNeXtEffv2Next4 19 | SPCONV_KERNEL_SIZES: [5, 5, 3, 3] 20 | OUT_CHANNEL: 384 21 | BEV_CHANNEL: 384 22 | CHANNELS: [32, 64, 128, 256, 256] 23 | 24 | DENSE_HEAD: 25 | NAME: VoxelNeXtHeadKPMerge 26 | IOU_BRANCH: True 27 | CLASS_AGNOSTIC: False 28 | INPUT_FEATURES: 384 29 | 30 | MAP_CLASS_NAMES: { 31 | 'Pedestrian': 'Human', 32 | 'Cyclist': 'Human', 33 | } # merge_class in waymo_dataset_kp 34 | CLASS_NAMES_EACH_HEAD: [ 35 | ['Human'] 36 | ] 37 | 38 | SHARED_CONV_CHANNEL: 384 39 | USE_BIAS_BEFORE_NORM: True 40 | NUM_HM_CONV: 2 41 | SEPARATE_HEAD_CFG: 42 | HEAD_ORDER: [ 'dim', 'rot' ] # xyz in loc 43 | KP_HEAD_ORDER: [ 'loc_x', 'loc_y', 'loc_z' ] 44 | HEAD_DICT: { 45 | 'dim': {'out_channels': 3, 'num_conv': 2}, 46 | 'rot': {'out_channels': 2, 'num_conv': 2}, 47 | 'iou': {'out_channels': 1, 'num_conv': 2}, 48 | 'loc_x': { 'out_channels': 15, 'num_conv': 2 }, 49 | 'loc_y': { 'out_channels': 15, 'num_conv': 2 }, 50 | 'loc_z': { 'out_channels': 15, 'num_conv': 2 }, 51 | 'kp_vis': { 'out_channels': 14, 'num_conv': 2 }, 52 | } 53 | RECTIFIER: [0.68, 0.71, 0.65] 54 | TARGET_ASSIGNER_CONFIG: 55 | FEATURE_MAP_STRIDE: 8 56 | NUM_MAX_OBJS: 500 57 | GAUSSIAN_OVERLAP: 0.1 58 | MIN_RADIUS: 2 59 | 60 | LOSS_CONFIG: 61 | LOSS_WEIGHTS: { 62 | 'cls_weight': 1.0, 63 | 'loc_weight': 2.0, 64 | 'code_weights': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], 65 | 'kp_x_code_weights': [ 66 | 1.0, 1.0, 1.0, 1.0, 1.0, 3.0, 3.0, 67 | 1.0, 1.0, 1.0, 1.0, 3.0, 3.0, 1.0, 68 | ], 69 | 'kp_y_code_weights': [ 70 | 1.0, 1.0, 1.0, 1.0, 1.0, 3.0, 3.0, 71 | 1.0, 1.0, 1.0, 1.0, 3.0, 3.0, 1.0, 72 | ], 73 | 'kp_z_code_weights': [ 74 | 1.0, 1.0, 1.0, 1.0, 1.0, 3.0, 3.0, 75 | 1.0, 1.0, 1.0, 1.0, 3.0, 3.0, 1.0, 76 | ], 77 | # Since the sum operation of the loss 78 | 'kp_x_loc_weight': 2, 79 | 'kp_y_loc_weight': 2, 80 | 'kp_z_loc_weight': 2, 81 | 'kp_visibility_code_weights': [ 82 | 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 83 | 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 84 | ], 85 | 'kp_visibility_weights': 1, 86 | } 87 | 88 | POST_PROCESSING: 89 | SCORE_THRESH: 0.3 90 | POST_CENTER_LIMIT_RANGE: [-75.2, -75.2, -2, 75.2, 75.2, 4] 91 | MAX_OBJ_PER_SAMPLE: 500 92 | NMS_CONFIG: 93 | NMS_TYPE: nms_gpu 94 | NMS_THRESH: [0.8, 0.55, 0.55] #0.7 95 | NMS_PRE_MAXSIZE: [2048, 1024, 1024] #[4096] 96 | NMS_POST_MAXSIZE: [200, 150, 150] #500 97 | 98 | POST_PROCESSING: 99 | RECALL_THRESH_LIST: [0.3, 0.5, 0.7] 100 | 101 | EVAL_METRIC: waymo 102 | 103 | 104 | OPTIMIZATION: 105 | BATCH_SIZE_PER_GPU: 2 106 | NUM_EPOCHS: 12 107 | 108 | OPTIMIZER: adam_onecycle 109 | LR: 0.003 110 | WEIGHT_DECAY: 0.01 111 | MOMENTUM: 0.9 112 | 113 | MOMS: [0.95, 0.85] 114 | PCT_START: 0.4 115 | DIV_FACTOR: 10 116 | DECAY_STEP_LIST: [25, 35] 117 | LR_DECAY: 0.3 118 | LR_CLIP: 0.0000001 119 | 120 | LR_WARMUP: False 121 | WARMUP_EPOCH: 1 122 | 123 | GRAD_NORM_CLIP: 10 124 | -------------------------------------------------------------------------------- /pcdet/models/model_utils/transfusion_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | import torch.nn.functional as F 4 | 5 | def clip_sigmoid(x, eps=1e-4): 6 | y = torch.clamp(x.sigmoid_(), min=eps, max=1 - eps) 7 | return y 8 | 9 | 10 | class PositionEmbeddingLearned(nn.Module): 11 | """ 12 | Absolute pos embedding, learned. 13 | """ 14 | 15 | def __init__(self, input_channel, num_pos_feats=288): 16 | super().__init__() 17 | self.position_embedding_head = nn.Sequential( 18 | nn.Conv1d(input_channel, num_pos_feats, kernel_size=1), 19 | nn.BatchNorm1d(num_pos_feats), 20 | nn.ReLU(inplace=True), 21 | nn.Conv1d(num_pos_feats, num_pos_feats, kernel_size=1)) 22 | 23 | def forward(self, xyz): 24 | xyz = xyz.transpose(1, 2).contiguous() 25 | position_embedding = self.position_embedding_head(xyz) 26 | return position_embedding 27 | 28 | 29 | class TransformerDecoderLayer(nn.Module): 30 | def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu", 31 | self_posembed=None, cross_posembed=None, cross_only=False): 32 | super().__init__() 33 | self.cross_only = cross_only 34 | if not self.cross_only: 35 | self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) 36 | self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) 37 | # Implementation of Feedforward model 38 | self.linear1 = nn.Linear(d_model, dim_feedforward) 39 | self.dropout = nn.Dropout(dropout) 40 | self.linear2 = nn.Linear(dim_feedforward, d_model) 41 | 42 | self.norm1 = nn.LayerNorm(d_model) 43 | self.norm2 = nn.LayerNorm(d_model) 44 | self.norm3 = nn.LayerNorm(d_model) 45 | self.dropout1 = nn.Dropout(dropout) 46 | self.dropout2 = nn.Dropout(dropout) 47 | self.dropout3 = nn.Dropout(dropout) 48 | 49 | def _get_activation_fn(activation): 50 | """Return an activation function given a string""" 51 | if activation == "relu": 52 | return F.relu 53 | if activation == "gelu": 54 | return F.gelu 55 | if activation == "glu": 56 | return F.glu 57 | raise RuntimeError(F"activation should be relu/gelu, not {activation}.") 58 | 59 | self.activation = _get_activation_fn(activation) 60 | 61 | self.self_posembed = self_posembed 62 | self.cross_posembed = cross_posembed 63 | 64 | def with_pos_embed(self, tensor, pos_embed): 65 | return tensor if pos_embed is None else tensor + pos_embed 66 | 67 | def forward(self, query, key, query_pos, key_pos, key_padding_mask=None, attn_mask=None): 68 | # NxCxP to PxNxC 69 | if self.self_posembed is not None: 70 | query_pos_embed = self.self_posembed(query_pos).permute(2, 0, 1) 71 | else: 72 | query_pos_embed = None 73 | if self.cross_posembed is not None: 74 | key_pos_embed = self.cross_posembed(key_pos).permute(2, 0, 1) 75 | else: 76 | key_pos_embed = None 77 | 78 | query = query.permute(2, 0, 1) 79 | key = key.permute(2, 0, 1) 80 | 81 | if not self.cross_only: 82 | q = k = v = self.with_pos_embed(query, query_pos_embed) 83 | query2 = self.self_attn(q, k, value=v)[0] 84 | query = query + self.dropout1(query2) 85 | query = self.norm1(query) 86 | 87 | query2 = self.multihead_attn(query=self.with_pos_embed(query, query_pos_embed), 88 | key=self.with_pos_embed(key, key_pos_embed), 89 | value=self.with_pos_embed(key, key_pos_embed), 90 | key_padding_mask=key_padding_mask, attn_mask=attn_mask)[0] 91 | 92 | query = query + self.dropout2(query2) 93 | query = self.norm2(query) 94 | 95 | query2 = self.linear2(self.dropout(self.activation(self.linear1(query)))) 96 | query = query + self.dropout3(query2) 97 | query = self.norm3(query) 98 | 99 | # NxCxP to PxNxC 100 | query = query.permute(1, 2, 0) 101 | return query 102 | 103 | --------------------------------------------------------------------------------