├── .gitignore
├── LICENSE
├── README.md
├── data
├── kitti
│ └── ImageSets
│ │ ├── test.txt
│ │ ├── train.txt
│ │ └── val.txt
└── waymo
│ └── ImageSets
│ ├── train.txt
│ └── val.txt
├── docs
├── DEMO.md
├── GETTING_STARTED.md
├── INSTALL.md
├── dataset_vs_model.png
├── demo.png
├── model_framework.png
├── multiple_models_demo.png
└── open_mmlab.png
├── pcdet
├── __init__.py
├── config.py
├── datasets
│ ├── __init__.py
│ ├── augmentor
│ │ ├── augmentor_utils.py
│ │ ├── data_augmentor.py
│ │ └── database_sampler.py
│ ├── dataset.py
│ ├── kitti
│ │ ├── kitti_dataset.py
│ │ ├── kitti_object_eval_python
│ │ │ ├── LICENSE
│ │ │ ├── README.md
│ │ │ ├── eval.py
│ │ │ ├── evaluate.py
│ │ │ ├── kitti_common.py
│ │ │ └── rotate_iou.py
│ │ └── kitti_utils.py
│ ├── nuscenes
│ │ ├── nuscenes_dataset.py
│ │ └── nuscenes_utils.py
│ ├── processor
│ │ ├── data_processor.py
│ │ └── point_feature_encoder.py
│ └── waymo
│ │ ├── waymo_dataset.py
│ │ ├── waymo_eval.py
│ │ └── waymo_utils.py
├── models
│ ├── __init__.py
│ ├── backbones_2d
│ │ ├── __init__.py
│ │ ├── base_bev_backbone.py
│ │ └── map_to_bev
│ │ │ ├── __init__.py
│ │ │ ├── conv2d_collapse.py
│ │ │ ├── height_compression.py
│ │ │ └── pointpillar_scatter.py
│ ├── backbones_3d
│ │ ├── __init__.py
│ │ ├── pfe
│ │ │ ├── __init__.py
│ │ │ └── voxel_set_abstraction.py
│ │ ├── pointnet2_backbone.py
│ │ ├── spconv_backbone.py
│ │ ├── spconv_unet.py
│ │ └── vfe
│ │ │ ├── __init__.py
│ │ │ ├── image_vfe.py
│ │ │ ├── image_vfe_modules
│ │ │ ├── f2v
│ │ │ │ ├── __init__.py
│ │ │ │ ├── frustum_grid_generator.py
│ │ │ │ ├── frustum_to_voxel.py
│ │ │ │ └── sampler.py
│ │ │ └── ffn
│ │ │ │ ├── __init__.py
│ │ │ │ ├── ddn
│ │ │ │ ├── __init__.py
│ │ │ │ ├── ddn_deeplabv3.py
│ │ │ │ └── ddn_template.py
│ │ │ │ ├── ddn_loss
│ │ │ │ ├── __init__.py
│ │ │ │ ├── balancer.py
│ │ │ │ └── ddn_loss.py
│ │ │ │ └── depth_ffn.py
│ │ │ ├── mean_vfe.py
│ │ │ ├── pillar_vfe.py
│ │ │ └── vfe_template.py
│ ├── dense_heads
│ │ ├── __init__.py
│ │ ├── anchor_head_multi.py
│ │ ├── anchor_head_single.py
│ │ ├── anchor_head_template.py
│ │ ├── point_head_box.py
│ │ ├── point_head_simple.py
│ │ ├── point_head_template.py
│ │ ├── point_intra_part_head.py
│ │ └── target_assigner
│ │ │ ├── anchor_generator.py
│ │ │ ├── atss_target_assigner.py
│ │ │ └── axis_aligned_target_assigner.py
│ ├── detectors
│ │ ├── PartA2_net.py
│ │ ├── __init__.py
│ │ ├── caddn.py
│ │ ├── detector3d_template.py
│ │ ├── point_rcnn.py
│ │ ├── pointpillar.py
│ │ ├── pv_rcnn.py
│ │ ├── second_net.py
│ │ ├── second_net_iou.py
│ │ └── voxel_rcnn.py
│ ├── model_utils
│ │ ├── basic_block_2d.py
│ │ └── model_nms_utils.py
│ └── roi_heads
│ │ ├── __init__.py
│ │ ├── partA2_head.py
│ │ ├── pointrcnn_head.py
│ │ ├── pvrcnn_head.py
│ │ ├── roi_head_template.py
│ │ ├── second_head.py
│ │ ├── target_assigner
│ │ └── proposal_target_layer.py
│ │ └── voxelrcnn_head.py
├── ops
│ ├── iou3d_nms
│ │ ├── iou3d_nms_utils.py
│ │ └── src
│ │ │ ├── iou3d_cpu.cpp
│ │ │ ├── iou3d_cpu.h
│ │ │ ├── iou3d_nms.cpp
│ │ │ ├── iou3d_nms.h
│ │ │ ├── iou3d_nms_api.cpp
│ │ │ └── iou3d_nms_kernel.cu
│ ├── pointnet2
│ │ ├── pointnet2_batch
│ │ │ ├── pointnet2_modules.py
│ │ │ ├── pointnet2_utils.py
│ │ │ └── src
│ │ │ │ ├── ball_query.cpp
│ │ │ │ ├── ball_query_gpu.cu
│ │ │ │ ├── ball_query_gpu.h
│ │ │ │ ├── cuda_utils.h
│ │ │ │ ├── group_points.cpp
│ │ │ │ ├── group_points_gpu.cu
│ │ │ │ ├── group_points_gpu.h
│ │ │ │ ├── interpolate.cpp
│ │ │ │ ├── interpolate_gpu.cu
│ │ │ │ ├── interpolate_gpu.h
│ │ │ │ ├── pointnet2_api.cpp
│ │ │ │ ├── sampling.cpp
│ │ │ │ ├── sampling_gpu.cu
│ │ │ │ └── sampling_gpu.h
│ │ └── pointnet2_stack
│ │ │ ├── pointnet2_modules.py
│ │ │ ├── pointnet2_utils.py
│ │ │ ├── src
│ │ │ ├── ball_query.cpp
│ │ │ ├── ball_query_gpu.cu
│ │ │ ├── ball_query_gpu.h
│ │ │ ├── cuda_utils.h
│ │ │ ├── group_points.cpp
│ │ │ ├── group_points_gpu.cu
│ │ │ ├── group_points_gpu.h
│ │ │ ├── interpolate.cpp
│ │ │ ├── interpolate_gpu.cu
│ │ │ ├── interpolate_gpu.h
│ │ │ ├── pointnet2_api.cpp
│ │ │ ├── sampling.cpp
│ │ │ ├── sampling_gpu.cu
│ │ │ ├── sampling_gpu.h
│ │ │ ├── voxel_query.cpp
│ │ │ ├── voxel_query_gpu.cu
│ │ │ └── voxel_query_gpu.h
│ │ │ ├── voxel_pool_modules.py
│ │ │ └── voxel_query_utils.py
│ ├── roiaware_pool3d
│ │ ├── roiaware_pool3d_utils.py
│ │ └── src
│ │ │ ├── roiaware_pool3d.cpp
│ │ │ └── roiaware_pool3d_kernel.cu
│ ├── roipoint_pool3d
│ │ ├── roipoint_pool3d_utils.py
│ │ └── src
│ │ │ ├── roipoint_pool3d.cpp
│ │ │ └── roipoint_pool3d_kernel.cu
│ ├── spconv
│ │ ├── __init__.py
│ │ ├── conv.py
│ │ ├── functional.py
│ │ ├── include
│ │ │ ├── paramsgrid.h
│ │ │ ├── prettyprint.h
│ │ │ ├── pybind11_utils.h
│ │ │ ├── spconv
│ │ │ │ ├── fused_spconv_ops.h
│ │ │ │ ├── geometry.h
│ │ │ │ ├── indice.cu.h
│ │ │ │ ├── indice.h
│ │ │ │ ├── maxpool.h
│ │ │ │ ├── mp_helper.h
│ │ │ │ ├── point2voxel.h
│ │ │ │ ├── pool_ops.h
│ │ │ │ ├── reordering.cu.h
│ │ │ │ ├── reordering.h
│ │ │ │ └── spconv_ops.h
│ │ │ ├── tensorview
│ │ │ │ ├── helper_kernel.cu.h
│ │ │ │ ├── helper_launch.h
│ │ │ │ └── tensorview.h
│ │ │ ├── torch_utils.h
│ │ │ └── utility
│ │ │ │ └── timer.h
│ │ ├── modules.py
│ │ ├── ops.py
│ │ ├── pool.py
│ │ ├── src
│ │ │ ├── all.cc
│ │ │ ├── indice.cc
│ │ │ ├── indice_cuda.cu
│ │ │ ├── maxpool.cc
│ │ │ ├── maxpool_cuda.cu
│ │ │ ├── reordering.cc
│ │ │ └── reordering_cuda.cu
│ │ ├── structure.py
│ │ └── test_utils.py
│ └── voxel
│ │ ├── __init__.py
│ │ ├── scatter_points.py
│ │ ├── src
│ │ ├── scatter_points_cpu.cpp
│ │ ├── scatter_points_cuda.cu
│ │ ├── voxelization.cpp
│ │ ├── voxelization.h
│ │ ├── voxelization_cpu.cpp
│ │ └── voxelization_cuda.cu
│ │ └── voxelize.py
└── utils
│ ├── box_coder_utils.py
│ ├── box_utils.py
│ ├── calibration_kitti.py
│ ├── common_utils.py
│ ├── loss_utils.py
│ ├── object3d_kitti.py
│ └── transform_utils.py
├── requirements.txt
├── setup.py
├── tags
└── tools
├── .ipynb_checkpoints
├── DataLoader-checkpoint.ipynb
└── Playground-checkpoint.ipynb
├── DataLoader.ipynb
├── Playground.ipynb
├── cfgs
├── dataset_configs
│ ├── kitti_dataset.yaml
│ ├── nuscenes_dataset.yaml
│ └── waymo_dataset.yaml
├── kitti_models
│ ├── CaDDN.yaml
│ ├── PartA2.yaml
│ ├── PartA2_free.yaml
│ ├── pointpillar.yaml
│ ├── pointrcnn.yaml
│ ├── pointrcnn_iou.yaml
│ ├── pv_rcnn.yaml
│ ├── second.yaml
│ ├── second_iou.yaml
│ ├── second_multihead.yaml
│ ├── spg.yaml
│ └── voxel_rcnn_car.yaml
├── nuscenes_models
│ ├── cbgs_pp_multihead.yaml
│ └── cbgs_second_multihead.yaml
└── waymo_models
│ ├── PartA2.yaml
│ ├── pv_rcnn.yaml
│ └── second.yaml
├── demo.py
├── eval_utils
└── eval_utils.py
├── plotvoxel.py
├── scripts
├── dist_test.sh
├── dist_train.sh
├── slurm_test_mgpu.sh
├── slurm_test_single.sh
└── slurm_train.sh
├── spg_model.py
├── spg_pointpillars.py
├── spg_test.py
├── spg_train.py
├── test.py
├── train.py
├── train_utils
├── optimization
│ ├── __init__.py
│ ├── fastai_optim.py
│ └── learning_schedules_fastai.py
└── train_utils.py
└── visual_utils
└── visualize_utils.py
/.gitignore:
--------------------------------------------------------------------------------
1 | **__pycache__**
2 | **build**
3 | **egg-info**
4 | **dist**
5 | data/
6 | *.pyc
7 | venv/
8 | *.idea/
9 | *.so
10 | *.yaml
11 | *.sh
12 | *.pth
13 | *.pkl
14 | *.zip
15 | *.bin
16 | output
17 | version.py
18 | tools/*.png
19 | tools/*.csv
20 | tools/*.p
21 | tools/*.npy
22 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # Semantic Point Generation
2 |
3 |
4 | This repo consists of code for running the classification head of the paper [SPG: Unsupervised Domain Adaptation for 3D Object Detection via Semantic Point Generation](https://arxiv.org/abs/2108.06709).
5 |
6 | All scripts can be found in tools prefixed with spg.
7 |
8 | # Spconv-OpenPCDet
9 | OpenPCDet with spconv package **already included** for **one-step** installation. Uses spconv & voxel CUDA ops from mmdetection3d repository.
10 |
11 | ## Installation
12 | Here is how I setup for my environment:
13 | ```
14 | conda create -n spconv-openpcdet python=3.7
15 | conda activate spconv-openpcdet
16 | conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=11.1 -c pytorch -c conda-forge
17 | pip install -r requirements.txt
18 | python setup.py develop
19 | ```
20 | I imagine that earlier versions of pytorch would work as well (since OpenPCDet supports 1.1, 1.3, 1.5 and mmdetection3d works for earlier versions).
21 |
22 | Note that mmcv, and the mm* packages do not need to be installed.
23 |
24 | ## Acknowledgement
25 | This repository is basically a copy of OpenPCDet, with some elements of mmdetection3d's usage of spconv within it.
26 | Thanks to [divadi](https://github.com/Divadi/Spconv-OpenPCDet) for helping us with installation.
27 |
--------------------------------------------------------------------------------
/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 `mayavi` visualization tools. If not, you could install it as follows:
10 | ```
11 | pip install mayavi
12 | ```
13 |
14 | 3. Prepare your custom point cloud data (skip this step if you use the original KITTI data).
15 | * You need to transform the coordinate of your custom point cloud to
16 | the unified normative coordinate of `OpenPCDet`, that is, x-axis points towards to front direction,
17 | y-axis points towards to the left direction, and z-axis points towards to the top direction.
18 | * (Optional) the z-axis origin of your point cloud coordinate should be about 1.6m above the ground surface,
19 | since currently the provided models are trained on the KITTI dataset.
20 | * Set the intensity information, and save your transformed custom data to `numpy file`:
21 | ```python
22 | # Transform your point cloud data
23 | ...
24 |
25 | # Save it to the file.
26 | # The shape of points should be (num_points, 4), that is [x, y, z, intensity] (Only for KITTI dataset).
27 | # If you doesn't have the intensity information, just set them to zeros.
28 | # If you have the intensity information, you should normalize them to [0, 1].
29 | points[:, 3] = 0
30 | np.save(`my_data.npy`, points)
31 | ```
32 |
33 | 4. Run the demo with a pretrained model (e.g. PV-RCNN) and your custom point cloud data as follows:
34 | ```shell
35 | python demo.py --cfg_file cfgs/kitti_models/pv_rcnn.yaml \
36 | --ckpt pv_rcnn_8369.pth \
37 | --data_path ${POINT_CLOUD_DATA}
38 | ```
39 | Here `${POINT_CLOUD_DATA}` could be in any of the following format:
40 | * Your transformed custom data with a single numpy file like `my_data.npy`.
41 | * Your transformed custom data with a directory to test with multiple point cloud data.
42 | * The original KITTI `.bin` data within `data/kitti`, like `data/kitti/training/velodyne/000008.bin`.
43 |
44 | Then you could see the predicted results with visualized point cloud as follows:
45 |
46 |
47 | 
48 |
49 |
--------------------------------------------------------------------------------
/docs/INSTALL.md:
--------------------------------------------------------------------------------
1 | # Installation
2 |
3 | ### Requirements
4 | All the codes are tested in the following environment:
5 | * Linux (tested on Ubuntu 14.04/16.04)
6 | * Python 3.6+
7 | * PyTorch 1.1 or higher (tested on PyTorch 1.1, 1,3, 1,5)
8 | * CUDA 9.0 or higher (PyTorch 1.3+ needs CUDA 9.2+)
9 | * [`spconv v1.0 (commit 8da6f96)`](https://github.com/traveller59/spconv/tree/8da6f967fb9a054d8870c3515b1b44eca2103634) or [`spconv v1.2`](https://github.com/traveller59/spconv)
10 |
11 |
12 | ### Install `pcdet v0.3`
13 | NOTE: Please re-install `pcdet v0.3` by running `python setup.py develop` even if you have already installed previous version.
14 |
15 | a. Clone this repository.
16 | ```shell
17 | git clone https://github.com/open-mmlab/OpenPCDet.git
18 | ```
19 |
20 | b. Install the dependent libraries as follows:
21 |
22 | * Install the dependent python libraries:
23 | ```
24 | pip install -r requirements.txt
25 | ```
26 |
27 | * Install the SparseConv library, we use the implementation from [`[spconv]`](https://github.com/traveller59/spconv).
28 | * If you use PyTorch 1.1, then make sure you install the `spconv v1.0` with ([commit 8da6f96](https://github.com/traveller59/spconv/tree/8da6f967fb9a054d8870c3515b1b44eca2103634)) instead of the latest one.
29 | * If you use PyTorch 1.3+, then you need to install the `spconv v1.2`. As mentioned by the author of [`spconv`](https://github.com/traveller59/spconv), you need to use their docker if you use PyTorch 1.4+.
30 |
31 | c. Install this `pcdet` library by running the following command:
32 | ```shell
33 | python setup.py develop
34 | ```
35 |
--------------------------------------------------------------------------------
/docs/dataset_vs_model.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/prithusuresh/semantic-point-generation/e9dd107b86e6068bd1de38b6dc177108a54deee1/docs/dataset_vs_model.png
--------------------------------------------------------------------------------
/docs/demo.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/prithusuresh/semantic-point-generation/e9dd107b86e6068bd1de38b6dc177108a54deee1/docs/demo.png
--------------------------------------------------------------------------------
/docs/model_framework.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/prithusuresh/semantic-point-generation/e9dd107b86e6068bd1de38b6dc177108a54deee1/docs/model_framework.png
--------------------------------------------------------------------------------
/docs/multiple_models_demo.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/prithusuresh/semantic-point-generation/e9dd107b86e6068bd1de38b6dc177108a54deee1/docs/multiple_models_demo.png
--------------------------------------------------------------------------------
/docs/open_mmlab.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/prithusuresh/semantic-point-generation/e9dd107b86e6068bd1de38b6dc177108a54deee1/docs/open_mmlab.png
--------------------------------------------------------------------------------
/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/config.py:
--------------------------------------------------------------------------------
1 | from pathlib import Path
2 |
3 | import yaml
4 | from easydict import EasyDict
5 |
6 |
7 | def log_config_to_file(cfg, pre='cfg', logger=None):
8 | for key, val in cfg.items():
9 | if isinstance(cfg[key], EasyDict):
10 | logger.info('\n%s.%s = edict()' % (pre, key))
11 | log_config_to_file(cfg[key], pre=pre + '.' + key, logger=logger)
12 | continue
13 | logger.info('%s.%s: %s' % (pre, key, val))
14 |
15 |
16 | def cfg_from_list(cfg_list, config):
17 | """Set config keys via list (e.g., from command line)."""
18 | from ast import literal_eval
19 | assert len(cfg_list) % 2 == 0
20 | for k, v in zip(cfg_list[0::2], cfg_list[1::2]):
21 | key_list = k.split('.')
22 | d = config
23 | for subkey in key_list[:-1]:
24 | assert subkey in d, 'NotFoundKey: %s' % subkey
25 | d = d[subkey]
26 | subkey = key_list[-1]
27 | assert subkey in d, 'NotFoundKey: %s' % subkey
28 | try:
29 | value = literal_eval(v)
30 | except:
31 | value = v
32 |
33 | if type(value) != type(d[subkey]) and isinstance(d[subkey], EasyDict):
34 | key_val_list = value.split(',')
35 | for src in key_val_list:
36 | cur_key, cur_val = src.split(':')
37 | val_type = type(d[subkey][cur_key])
38 | cur_val = val_type(cur_val)
39 | d[subkey][cur_key] = cur_val
40 | elif type(value) != type(d[subkey]) and isinstance(d[subkey], list):
41 | val_list = value.split(',')
42 | for k, x in enumerate(val_list):
43 | val_list[k] = type(d[subkey][0])(x)
44 | d[subkey] = val_list
45 | else:
46 | assert type(value) == type(d[subkey]), \
47 | 'type {} does not match original type {}'.format(type(value), type(d[subkey]))
48 | d[subkey] = value
49 |
50 |
51 | def merge_new_config(config, new_config):
52 | if '_BASE_CONFIG_' in new_config:
53 | with open(new_config['_BASE_CONFIG_'], 'r') as f:
54 | try:
55 | yaml_config = yaml.load(f, Loader=yaml.FullLoader)
56 | except:
57 | yaml_config = yaml.load(f)
58 | config.update(EasyDict(yaml_config))
59 |
60 | for key, val in new_config.items():
61 | if not isinstance(val, dict):
62 | config[key] = val
63 | continue
64 | if key not in config:
65 | config[key] = EasyDict()
66 | merge_new_config(config[key], val)
67 |
68 | return config
69 |
70 |
71 | def cfg_from_yaml_file(cfg_file, config):
72 | with open(cfg_file, 'r') as f:
73 | try:
74 | new_config = yaml.load(f, Loader=yaml.FullLoader)
75 | except:
76 | new_config = yaml.load(f)
77 |
78 | merge_new_config(config=config, new_config=new_config)
79 |
80 | return config
81 |
82 |
83 | cfg = EasyDict()
84 | cfg.ROOT_DIR = (Path(__file__).resolve().parent / '../').resolve()
85 | cfg.LOCAL_RANK = 0
86 |
--------------------------------------------------------------------------------
/pcdet/datasets/__init__.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from torch.utils.data import DataLoader
3 | from torch.utils.data import DistributedSampler as _DistributedSampler
4 |
5 | from pcdet.utils import common_utils
6 |
7 | from .dataset import DatasetTemplate
8 | from .kitti.kitti_dataset import KittiDataset
9 | from .nuscenes.nuscenes_dataset import NuScenesDataset
10 | from .waymo.waymo_dataset import WaymoDataset
11 |
12 | __all__ = {
13 | 'DatasetTemplate': DatasetTemplate,
14 | 'KittiDataset': KittiDataset,
15 | 'NuScenesDataset': NuScenesDataset,
16 | 'WaymoDataset': WaymoDataset
17 | }
18 |
19 |
20 | class DistributedSampler(_DistributedSampler):
21 |
22 | def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True):
23 | super().__init__(dataset, num_replicas=num_replicas, rank=rank)
24 | self.shuffle = shuffle
25 |
26 | def __iter__(self):
27 | if self.shuffle:
28 | g = torch.Generator()
29 | g.manual_seed(self.epoch)
30 | indices = torch.randperm(len(self.dataset), generator=g).tolist()
31 | else:
32 | indices = torch.arange(len(self.dataset)).tolist()
33 |
34 | indices += indices[:(self.total_size - len(indices))]
35 | assert len(indices) == self.total_size
36 |
37 | indices = indices[self.rank:self.total_size:self.num_replicas]
38 | assert len(indices) == self.num_samples
39 |
40 | return iter(indices)
41 |
42 |
43 | def build_dataloader(dataset_cfg, class_names, batch_size, dist, root_path=None, workers=4,
44 | logger=None, training=True, merge_all_iters_to_one_epoch=False, total_epochs=0):
45 |
46 | dataset = __all__[dataset_cfg.DATASET](
47 | dataset_cfg=dataset_cfg,
48 | class_names=class_names,
49 | root_path=root_path,
50 | training=training,
51 | logger=logger,
52 | )
53 |
54 | if merge_all_iters_to_one_epoch:
55 | assert hasattr(dataset, 'merge_all_iters_to_one_epoch')
56 | dataset.merge_all_iters_to_one_epoch(merge=True, epochs=total_epochs)
57 |
58 | if dist:
59 | if training:
60 | sampler = torch.utils.data.distributed.DistributedSampler(dataset)
61 | else:
62 | rank, world_size = common_utils.get_dist_info()
63 | sampler = DistributedSampler(dataset, world_size, rank, shuffle=False)
64 | else:
65 | sampler = None
66 | dataloader = DataLoader(
67 | dataset, batch_size=batch_size, pin_memory=True, num_workers=workers,
68 | shuffle=(sampler is None) and training, collate_fn=dataset.collate_batch,
69 | drop_last=False, sampler=sampler, timeout=0
70 | )
71 |
72 | return dataset, dataloader, sampler
73 |
--------------------------------------------------------------------------------
/pcdet/datasets/augmentor/augmentor_utils.py:
--------------------------------------------------------------------------------
1 | import copy
2 | import numpy as np
3 |
4 | from ...utils import common_utils
5 |
6 |
7 | def random_flip_along_x(gt_boxes, points):
8 | """
9 | Args:
10 | gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
11 | points: (M, 3 + C)
12 | Returns:
13 | """
14 | enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5])
15 | if enable:
16 | gt_boxes[:, 1] = -gt_boxes[:, 1]
17 | gt_boxes[:, 6] = -gt_boxes[:, 6]
18 | points[:, 1] = -points[:, 1]
19 |
20 | if gt_boxes.shape[1] > 7:
21 | gt_boxes[:, 8] = -gt_boxes[:, 8]
22 |
23 | return gt_boxes, points
24 |
25 |
26 | def random_flip_along_y(gt_boxes, points):
27 | """
28 | Args:
29 | gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
30 | points: (M, 3 + C)
31 | Returns:
32 | """
33 | enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5])
34 | if enable:
35 | gt_boxes[:, 0] = -gt_boxes[:, 0]
36 | gt_boxes[:, 6] = -(gt_boxes[:, 6] + np.pi)
37 | points[:, 0] = -points[:, 0]
38 |
39 | if gt_boxes.shape[1] > 7:
40 | gt_boxes[:, 7] = -gt_boxes[:, 7]
41 |
42 | return gt_boxes, points
43 |
44 |
45 | def global_rotation(gt_boxes, points, rot_range):
46 | """
47 | Args:
48 | gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
49 | points: (M, 3 + C),
50 | rot_range: [min, max]
51 | Returns:
52 | """
53 | noise_rotation = np.random.uniform(rot_range[0], rot_range[1])
54 | points = common_utils.rotate_points_along_z(points[np.newaxis, :, :], np.array([noise_rotation]))[0]
55 | gt_boxes[:, 0:3] = common_utils.rotate_points_along_z(gt_boxes[np.newaxis, :, 0:3], np.array([noise_rotation]))[0]
56 | gt_boxes[:, 6] += noise_rotation
57 | if gt_boxes.shape[1] > 7:
58 | gt_boxes[:, 7:9] = common_utils.rotate_points_along_z(
59 | np.hstack((gt_boxes[:, 7:9], np.zeros((gt_boxes.shape[0], 1))))[np.newaxis, :, :],
60 | np.array([noise_rotation])
61 | )[0][:, 0:2]
62 |
63 | return gt_boxes, points
64 |
65 |
66 | def global_scaling(gt_boxes, points, scale_range):
67 | """
68 | Args:
69 | gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading]
70 | points: (M, 3 + C),
71 | scale_range: [min, max]
72 | Returns:
73 | """
74 | if scale_range[1] - scale_range[0] < 1e-3:
75 | return gt_boxes, points
76 | noise_scale = np.random.uniform(scale_range[0], scale_range[1])
77 | points[:, :3] *= noise_scale
78 | gt_boxes[:, :6] *= noise_scale
79 | return gt_boxes, points
80 |
81 | def random_image_flip_horizontal(image, depth_map, gt_boxes, calib):
82 | """
83 | Performs random horizontal flip augmentation
84 | Args:
85 | image: (H_image, W_image, 3), Image
86 | depth_map: (H_depth, W_depth), Depth map
87 | gt_boxes: (N, 7), 3D box labels in LiDAR coordinates [x, y, z, w, l, h, ry]
88 | calib: calibration.Calibration, Calibration object
89 | Returns:
90 | aug_image: (H_image, W_image, 3), Augmented image
91 | aug_depth_map: (H_depth, W_depth), Augmented depth map
92 | aug_gt_boxes: (N, 7), Augmented 3D box labels in LiDAR coordinates [x, y, z, w, l, h, ry]
93 | """
94 | # Randomly augment with 50% chance
95 | enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5])
96 |
97 | if enable:
98 | # Flip images
99 | aug_image = np.fliplr(image)
100 | aug_depth_map = np.fliplr(depth_map)
101 |
102 | # Flip 3D gt_boxes by flipping the centroids in image space
103 | aug_gt_boxes = copy.copy(gt_boxes)
104 | locations = aug_gt_boxes[:, :3]
105 | img_pts, img_depth = calib.lidar_to_img(locations)
106 | W = image.shape[1]
107 | img_pts[:, 0] = W - img_pts[:, 0]
108 | pts_rect = calib.img_to_rect(u=img_pts[:, 0], v=img_pts[:, 1], depth_rect=img_depth)
109 | pts_lidar = calib.rect_to_lidar(pts_rect)
110 | aug_gt_boxes[:, :3] = pts_lidar
111 | aug_gt_boxes[:, 6] = -1 * aug_gt_boxes[:, 6]
112 |
113 | else:
114 | aug_image = image
115 | aug_depth_map = depth_map
116 | aug_gt_boxes = gt_boxes
117 |
118 | return aug_image, aug_depth_map, aug_gt_boxes
--------------------------------------------------------------------------------
/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/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/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/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 k in range(anno['name'].shape[0]):
16 | anno['name'][k] = map_name_to_kitti[anno['name'][k]]
17 |
18 | anno['bbox'] = np.zeros((len(anno['name']), 4))
19 | anno['bbox'][:, 2:4] = 50 # [0, 0, 50, 50]
20 | anno['truncated'] = np.zeros(len(anno['name']))
21 | anno['occluded'] = np.zeros(len(anno['name']))
22 | if 'boxes_lidar' in anno:
23 | gt_boxes_lidar = anno['boxes_lidar'].copy()
24 | else:
25 | gt_boxes_lidar = anno['gt_boxes_lidar'].copy()
26 |
27 | if len(gt_boxes_lidar) > 0:
28 | if info_with_fakelidar:
29 | gt_boxes_lidar = box_utils.boxes3d_kitti_fakelidar_to_lidar(gt_boxes_lidar)
30 |
31 | gt_boxes_lidar[:, 2] -= gt_boxes_lidar[:, 5] / 2
32 | anno['location'] = np.zeros((gt_boxes_lidar.shape[0], 3))
33 | anno['location'][:, 0] = -gt_boxes_lidar[:, 1] # x = -y_lidar
34 | anno['location'][:, 1] = -gt_boxes_lidar[:, 2] # y = -z_lidar
35 | anno['location'][:, 2] = gt_boxes_lidar[:, 0] # z = x_lidar
36 | dxdydz = gt_boxes_lidar[:, 3:6]
37 | anno['dimensions'] = dxdydz[:, [0, 2, 1]] # lwh ==> lhw
38 | anno['rotation_y'] = -gt_boxes_lidar[:, 6] - np.pi / 2.0
39 | anno['alpha'] = -np.arctan2(-gt_boxes_lidar[:, 1], gt_boxes_lidar[:, 0]) + anno['rotation_y']
40 | else:
41 | anno['location'] = anno['dimensions'] = np.zeros((0, 3))
42 | anno['rotation_y'] = anno['alpha'] = np.zeros(0)
43 |
44 | return annos
45 |
46 |
47 | def calib_to_matricies(calib):
48 | """
49 | Converts calibration object to transformation matricies
50 | Args:
51 | calib: calibration.Calibration, Calibration object
52 | Returns
53 | V2R: (4, 4), Lidar to rectified camera transformation matrix
54 | P2: (3, 4), Camera projection matrix
55 | """
56 | V2C = np.vstack((calib.V2C, np.array([0, 0, 0, 1], dtype=np.float32))) # (4, 4)
57 | R0 = np.hstack((calib.R0, np.zeros((3, 1), dtype=np.float32))) # (3, 4)
58 | R0 = np.vstack((R0, np.array([0, 0, 0, 1], dtype=np.float32))) # (4, 4)
59 | V2R = R0 @ V2C
60 | P2 = calib.P2
61 | return V2R, P2
--------------------------------------------------------------------------------
/pcdet/datasets/processor/point_feature_encoder.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 |
3 |
4 | class PointFeatureEncoder(object):
5 | def __init__(self, config, point_cloud_range=None):
6 | super().__init__()
7 | self.point_encoding_config = config
8 | assert list(self.point_encoding_config.src_feature_list[0:3]) == ['x', 'y', 'z']
9 | self.used_feature_list = self.point_encoding_config.used_feature_list
10 | self.src_feature_list = self.point_encoding_config.src_feature_list
11 | self.point_cloud_range = point_cloud_range
12 |
13 | @property
14 | def num_point_features(self):
15 | return getattr(self, self.point_encoding_config.encoding_type)(points=None)
16 |
17 | def forward(self, data_dict):
18 | """
19 | Args:
20 | data_dict:
21 | points: (N, 3 + C_in)
22 | ...
23 | Returns:
24 | data_dict:
25 | points: (N, 3 + C_out),
26 | use_lead_xyz: whether to use xyz as point-wise features
27 | ...
28 | """
29 | data_dict['points'], use_lead_xyz = getattr(self, self.point_encoding_config.encoding_type)(
30 | data_dict['points']
31 | )
32 | data_dict['use_lead_xyz'] = use_lead_xyz
33 | return data_dict
34 |
35 | def absolute_coordinates_encoding(self, points=None):
36 | if points is None:
37 | num_output_features = len(self.used_feature_list)
38 | return num_output_features
39 |
40 | point_feature_list = [points[:, 0:3]]
41 | for x in self.used_feature_list:
42 | if x in ['x', 'y', 'z']:
43 | continue
44 | idx = self.src_feature_list.index(x)
45 | point_feature_list.append(points[:, idx:idx+1])
46 | point_features = np.concatenate(point_feature_list, axis=1)
47 | return point_features, True
48 |
--------------------------------------------------------------------------------
/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 not isinstance(val, np.ndarray):
26 | continue
27 | elif key in ['frame_id', 'metadata', 'calib']:
28 | continue
29 | elif key in ['images']:
30 | batch_dict[key] = image_to_tensor(val).float().cuda().contiguous()
31 | elif key in ['image_shape']:
32 | batch_dict[key] = torch.from_numpy(val).int().cuda()
33 | else:
34 |
35 | if isinstance(batch_dict[key], np.ndarray):
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/backbones_2d/__init__.py:
--------------------------------------------------------------------------------
1 | from .base_bev_backbone import BaseBEVBackbone
2 |
3 | __all__ = {
4 | 'BaseBEVBackbone': BaseBEVBackbone
5 | }
6 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_2d/map_to_bev/__init__.py:
--------------------------------------------------------------------------------
1 | from .height_compression import HeightCompression
2 | from .pointpillar_scatter import PointPillarScatter
3 | from .conv2d_collapse import Conv2DCollapse
4 |
5 | __all__ = {
6 | 'HeightCompression': HeightCompression,
7 | 'PointPillarScatter': PointPillarScatter,
8 | 'Conv2DCollapse': Conv2DCollapse
9 | }
10 |
--------------------------------------------------------------------------------
/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/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/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 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_3d/__init__.py:
--------------------------------------------------------------------------------
1 | from .pointnet2_backbone import PointNet2Backbone, PointNet2MSG
2 | from .spconv_backbone import VoxelBackBone8x, VoxelResBackBone8x
3 | from .spconv_unet import UNetV2
4 |
5 | __all__ = {
6 | 'VoxelBackBone8x': VoxelBackBone8x,
7 | 'UNetV2': UNetV2,
8 | 'PointNet2Backbone': PointNet2Backbone,
9 | 'PointNet2MSG': PointNet2MSG,
10 | 'VoxelResBackBone8x': VoxelResBackBone8x,
11 | }
12 |
--------------------------------------------------------------------------------
/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/__init__.py:
--------------------------------------------------------------------------------
1 | from .mean_vfe import MeanVFE
2 | from .pillar_vfe import PillarVFE
3 | from .image_vfe import ImageVFE
4 | from .vfe_template import VFETemplate
5 |
6 | __all__ = {
7 | 'VFETemplate': VFETemplate,
8 | 'MeanVFE': MeanVFE,
9 | 'PillarVFE': PillarVFE,
10 | 'ImageVFE': ImageVFE
11 | }
12 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/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/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/models/backbones_3d/vfe/image_vfe_modules/f2v/sampler.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | import torch.nn.functional as F
4 |
5 |
6 | class Sampler(nn.Module):
7 |
8 | def __init__(self, mode="bilinear", padding_mode="zeros"):
9 | """
10 | Initializes module
11 | Args:
12 | mode: string, Sampling mode [bilinear/nearest]
13 | padding_mode: string, Padding mode for outside grid values [zeros/border/reflection]
14 | """
15 | super().__init__()
16 | self.mode = mode
17 | self.padding_mode = padding_mode
18 |
19 | def forward(self, input_features, grid):
20 | """
21 | Samples input using sampling grid
22 | Args:
23 | input_features: (B, C, D, H, W), Input frustum features
24 | grid: (B, X, Y, Z, 3), Sampling grids for input features
25 | Returns
26 | output_features: (B, C, X, Y, Z) Output voxel features
27 | """
28 | # Sample from grid
29 | output = F.grid_sample(input=input_features, grid=grid, mode=self.mode, padding_mode=self.padding_mode)
30 | return output
31 |
--------------------------------------------------------------------------------
/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/__init__.py:
--------------------------------------------------------------------------------
1 | from .ddn_deeplabv3 import DDNDeepLabV3
2 |
3 | __all__ = {
4 | 'DDNDeepLabV3': DDNDeepLabV3
5 | }
6 |
--------------------------------------------------------------------------------
/pcdet/models/backbones_3d/vfe/image_vfe_modules/ffn/ddn/ddn_deeplabv3.py:
--------------------------------------------------------------------------------
1 | import torchvision
2 |
3 | from .ddn_template import DDNTemplate
4 |
5 |
6 | class DDNDeepLabV3(DDNTemplate):
7 |
8 | def __init__(self, backbone_name, **kwargs):
9 | """
10 | Initializes DDNDeepLabV3 model
11 | Args:
12 | backbone_name: string, ResNet Backbone Name [ResNet50/ResNet101]
13 | """
14 | if backbone_name == "ResNet50":
15 | constructor = torchvision.models.segmentation.deeplabv3_resnet50
16 | elif backbone_name == "ResNet101":
17 | constructor = torchvision.models.segmentation.deeplabv3_resnet101
18 | else:
19 | raise NotImplementedError
20 |
21 | super().__init__(constructor=constructor, **kwargs)
22 |
--------------------------------------------------------------------------------
/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_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/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/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 |
--------------------------------------------------------------------------------
/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/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/models/dense_heads/__init__.py:
--------------------------------------------------------------------------------
1 | from .anchor_head_multi import AnchorHeadMulti
2 | from .anchor_head_single import AnchorHeadSingle
3 | from .anchor_head_template import AnchorHeadTemplate
4 | from .point_head_box import PointHeadBox
5 | from .point_head_simple import PointHeadSimple
6 | from .point_intra_part_head import PointIntraPartOffsetHead
7 |
8 | __all__ = {
9 | 'AnchorHeadTemplate': AnchorHeadTemplate,
10 | 'AnchorHeadSingle': AnchorHeadSingle,
11 | 'PointIntraPartOffsetHead': PointIntraPartOffsetHead,
12 | 'PointHeadSimple': PointHeadSimple,
13 | 'PointHeadBox': PointHeadBox,
14 | 'AnchorHeadMulti': AnchorHeadMulti,
15 | }
16 |
--------------------------------------------------------------------------------
/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/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 |
--------------------------------------------------------------------------------
/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/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/detectors/__init__.py:
--------------------------------------------------------------------------------
1 | from .detector3d_template import Detector3DTemplate
2 | from .PartA2_net import PartA2Net
3 | from .point_rcnn import PointRCNN
4 | from .pointpillar import PointPillar
5 | from .pv_rcnn import PVRCNN
6 | from .second_net import SECONDNet
7 | from .second_net_iou import SECONDNetIoU
8 | from .caddn import CaDDN
9 | from .voxel_rcnn import VoxelRCNN
10 |
11 | __all__ = {
12 | 'Detector3DTemplate': Detector3DTemplate,
13 | 'SECONDNet': SECONDNet,
14 | 'PartA2Net': PartA2Net,
15 | 'PVRCNN': PVRCNN,
16 | 'PointPillar': PointPillar,
17 | 'PointRCNN': PointRCNN,
18 | 'SECONDNetIoU': SECONDNetIoU,
19 | 'CaDDN': CaDDN,
20 | 'VoxelRCNN': VoxelRCNN
21 | }
22 |
23 |
24 | def build_detector(model_cfg, num_class, dataset):
25 | model = __all__[model_cfg.NAME](
26 | model_cfg=model_cfg, num_class=num_class, dataset=dataset
27 | )
28 |
29 | return model
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/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/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 | if self.training:
13 | loss, tb_dict, disp_dict = self.get_training_loss()
14 |
15 | ret_dict = {
16 | 'loss': loss
17 | }
18 | return ret_dict, tb_dict, disp_dict
19 | else:
20 | pred_dicts, recall_dicts = self.post_processing(batch_dict)
21 | return pred_dicts, recall_dicts
22 |
23 | def get_training_loss(self):
24 | disp_dict = {}
25 |
26 | loss_rpn, tb_dict = self.dense_head.get_loss()
27 | tb_dict = {
28 | 'loss_rpn': loss_rpn.item(),
29 | **tb_dict
30 | }
31 |
32 | loss = loss_rpn
33 | return loss, tb_dict, disp_dict
34 |
--------------------------------------------------------------------------------
/pcdet/models/detectors/pv_rcnn.py:
--------------------------------------------------------------------------------
1 | from .detector3d_template import Detector3DTemplate
2 |
3 |
4 | class PVRCNN(Detector3DTemplate):
5 | def __init__(self, model_cfg, num_class, dataset):
6 | super().__init__(model_cfg=model_cfg, num_class=num_class, dataset=dataset)
7 | self.module_list = self.build_networks()
8 |
9 | def forward(self, batch_dict):
10 | for cur_module in self.module_list:
11 | batch_dict = cur_module(batch_dict)
12 |
13 | if self.training:
14 | loss, tb_dict, disp_dict = self.get_training_loss()
15 |
16 | ret_dict = {
17 | 'loss': loss
18 | }
19 | return ret_dict, tb_dict, disp_dict
20 | else:
21 | pred_dicts, recall_dicts = self.post_processing(batch_dict)
22 | return pred_dicts, recall_dicts
23 |
24 | def get_training_loss(self):
25 | disp_dict = {}
26 | loss_rpn, tb_dict = self.dense_head.get_loss()
27 | loss_point, tb_dict = self.point_head.get_loss(tb_dict)
28 | loss_rcnn, tb_dict = self.roi_head.get_loss(tb_dict)
29 |
30 | loss = loss_rpn + loss_point + loss_rcnn
31 | return loss, tb_dict, disp_dict
32 |
--------------------------------------------------------------------------------
/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/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 | return loss, tb_dict, disp_dict
33 |
--------------------------------------------------------------------------------
/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/model_utils/model_nms_utils.py:
--------------------------------------------------------------------------------
1 | import torch
2 |
3 | from ...ops.iou3d_nms import iou3d_nms_utils
4 |
5 |
6 | def class_agnostic_nms(box_scores, box_preds, nms_config, score_thresh=None):
7 | src_box_scores = box_scores
8 | if score_thresh is not None:
9 | scores_mask = (box_scores >= score_thresh)
10 | box_scores = box_scores[scores_mask]
11 | box_preds = box_preds[scores_mask]
12 |
13 | selected = []
14 | if box_scores.shape[0] > 0:
15 | box_scores_nms, indices = torch.topk(box_scores, k=min(nms_config.NMS_PRE_MAXSIZE, box_scores.shape[0]))
16 | boxes_for_nms = box_preds[indices]
17 | keep_idx, selected_scores = getattr(iou3d_nms_utils, nms_config.NMS_TYPE)(
18 | boxes_for_nms[:, 0:7], box_scores_nms, nms_config.NMS_THRESH, **nms_config
19 | )
20 | selected = indices[keep_idx[:nms_config.NMS_POST_MAXSIZE]]
21 |
22 | if score_thresh is not None:
23 | original_idxs = scores_mask.nonzero().view(-1)
24 | selected = original_idxs[selected]
25 | return selected, src_box_scores[selected]
26 |
27 |
28 | def multi_classes_nms(cls_scores, box_preds, nms_config, score_thresh=None):
29 | """
30 | Args:
31 | cls_scores: (N, num_class)
32 | box_preds: (N, 7 + C)
33 | nms_config:
34 | score_thresh:
35 |
36 | Returns:
37 |
38 | """
39 | pred_scores, pred_labels, pred_boxes = [], [], []
40 | for k in range(cls_scores.shape[1]):
41 | if score_thresh is not None:
42 | scores_mask = (cls_scores[:, k] >= score_thresh)
43 | box_scores = cls_scores[scores_mask, k]
44 | cur_box_preds = box_preds[scores_mask]
45 | else:
46 | box_scores = cls_scores[:, k]
47 | cur_box_preds = box_preds
48 |
49 | selected = []
50 | if box_scores.shape[0] > 0:
51 | box_scores_nms, indices = torch.topk(box_scores, k=min(nms_config.NMS_PRE_MAXSIZE, box_scores.shape[0]))
52 | boxes_for_nms = cur_box_preds[indices]
53 | keep_idx, selected_scores = getattr(iou3d_nms_utils, nms_config.NMS_TYPE)(
54 | boxes_for_nms[:, 0:7], box_scores_nms, nms_config.NMS_THRESH, **nms_config
55 | )
56 | selected = indices[keep_idx[:nms_config.NMS_POST_MAXSIZE]]
57 |
58 | pred_scores.append(box_scores[selected])
59 | pred_labels.append(box_scores.new_ones(len(selected)).long() * k)
60 | pred_boxes.append(cur_box_preds[selected])
61 |
62 | pred_scores = torch.cat(pred_scores, dim=0)
63 | pred_labels = torch.cat(pred_labels, dim=0)
64 | pred_boxes = torch.cat(pred_boxes, dim=0)
65 |
66 | return pred_scores, pred_labels, pred_boxes
67 |
--------------------------------------------------------------------------------
/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 .second_head import SECONDHead
5 | from .voxelrcnn_head import VoxelRCNNHead
6 | from .roi_head_template import RoIHeadTemplate
7 |
8 |
9 | __all__ = {
10 | 'RoIHeadTemplate': RoIHeadTemplate,
11 | 'PartA2FCHead': PartA2FCHead,
12 | 'PVRCNNHead': PVRCNNHead,
13 | 'SECONDHead': SECONDHead,
14 | 'PointRCNNHead': PointRCNNHead,
15 | 'VoxelRCNNHead': VoxelRCNNHead
16 | }
17 |
--------------------------------------------------------------------------------
/pcdet/ops/iou3d_nms/iou3d_nms_utils.py:
--------------------------------------------------------------------------------
1 | """
2 | 3D IoU Calculation and Rotated NMS
3 | Written by Shaoshuai Shi
4 | All Rights Reserved 2019-2020.
5 | """
6 | import torch
7 |
8 | from ...utils import common_utils
9 | from . import iou3d_nms_cuda
10 |
11 |
12 | def boxes_bev_iou_cpu(boxes_a, boxes_b):
13 | """
14 | Args:
15 | boxes_a: (N, 7) [x, y, z, dx, dy, dz, heading]
16 | boxes_b: (M, 7) [x, y, z, dx, dy, dz, heading]
17 |
18 | Returns:
19 | ans_iou: (N, M)
20 | """
21 | boxes_a, is_numpy = common_utils.check_numpy_to_torch(boxes_a)
22 | boxes_b, is_numpy = common_utils.check_numpy_to_torch(boxes_b)
23 | assert not (boxes_a.is_cuda or boxes_b.is_cuda), 'Only support CPU tensors'
24 | assert boxes_a.shape[1] == 7 and boxes_b.shape[1] == 7
25 | ans_iou = boxes_a.new_zeros(torch.Size((boxes_a.shape[0], boxes_b.shape[0])))
26 | iou3d_nms_cuda.boxes_iou_bev_cpu(boxes_a.contiguous(), boxes_b.contiguous(), ans_iou)
27 |
28 | return ans_iou.numpy() if is_numpy else ans_iou
29 |
30 |
31 | def boxes_iou_bev(boxes_a, boxes_b):
32 | """
33 | Args:
34 | boxes_a: (N, 7) [x, y, z, dx, dy, dz, heading]
35 | boxes_b: (M, 7) [x, y, z, dx, dy, dz, heading]
36 |
37 | Returns:
38 | ans_iou: (N, M)
39 | """
40 | assert boxes_a.shape[1] == boxes_b.shape[1] == 7
41 | ans_iou = torch.cuda.FloatTensor(torch.Size((boxes_a.shape[0], boxes_b.shape[0]))).zero_()
42 |
43 | iou3d_nms_cuda.boxes_iou_bev_gpu(boxes_a.contiguous(), boxes_b.contiguous(), ans_iou)
44 |
45 | return ans_iou
46 |
47 |
48 | def boxes_iou3d_gpu(boxes_a, boxes_b):
49 | """
50 | Args:
51 | boxes_a: (N, 7) [x, y, z, dx, dy, dz, heading]
52 | boxes_b: (M, 7) [x, y, z, dx, dy, dz, heading]
53 |
54 | Returns:
55 | ans_iou: (N, M)
56 | """
57 | assert boxes_a.shape[1] == boxes_b.shape[1] == 7
58 |
59 | # height overlap
60 | boxes_a_height_max = (boxes_a[:, 2] + boxes_a[:, 5] / 2).view(-1, 1)
61 | boxes_a_height_min = (boxes_a[:, 2] - boxes_a[:, 5] / 2).view(-1, 1)
62 | boxes_b_height_max = (boxes_b[:, 2] + boxes_b[:, 5] / 2).view(1, -1)
63 | boxes_b_height_min = (boxes_b[:, 2] - boxes_b[:, 5] / 2).view(1, -1)
64 |
65 | # bev overlap
66 | overlaps_bev = torch.cuda.FloatTensor(torch.Size((boxes_a.shape[0], boxes_b.shape[0]))).zero_() # (N, M)
67 | iou3d_nms_cuda.boxes_overlap_bev_gpu(boxes_a.contiguous(), boxes_b.contiguous(), overlaps_bev)
68 |
69 | max_of_min = torch.max(boxes_a_height_min, boxes_b_height_min)
70 | min_of_max = torch.min(boxes_a_height_max, boxes_b_height_max)
71 | overlaps_h = torch.clamp(min_of_max - max_of_min, min=0)
72 |
73 | # 3d iou
74 | overlaps_3d = overlaps_bev * overlaps_h
75 |
76 | vol_a = (boxes_a[:, 3] * boxes_a[:, 4] * boxes_a[:, 5]).view(-1, 1)
77 | vol_b = (boxes_b[:, 3] * boxes_b[:, 4] * boxes_b[:, 5]).view(1, -1)
78 |
79 | iou3d = overlaps_3d / torch.clamp(vol_a + vol_b - overlaps_3d, min=1e-6)
80 |
81 | return iou3d
82 |
83 |
84 | def nms_gpu(boxes, scores, thresh, pre_maxsize=None, **kwargs):
85 | """
86 | :param boxes: (N, 7) [x, y, z, dx, dy, dz, heading]
87 | :param scores: (N)
88 | :param thresh:
89 | :return:
90 | """
91 | assert boxes.shape[1] == 7
92 | order = scores.sort(0, descending=True)[1]
93 | if pre_maxsize is not None:
94 | order = order[:pre_maxsize]
95 |
96 | boxes = boxes[order].contiguous()
97 | keep = torch.LongTensor(boxes.size(0))
98 | num_out = iou3d_nms_cuda.nms_gpu(boxes, keep, thresh)
99 | return order[keep[:num_out].cuda()].contiguous(), None
100 |
101 |
102 | def nms_normal_gpu(boxes, scores, thresh, **kwargs):
103 | """
104 | :param boxes: (N, 7) [x, y, z, dx, dy, dz, heading]
105 | :param scores: (N)
106 | :param thresh:
107 | :return:
108 | """
109 | assert boxes.shape[1] == 7
110 | order = scores.sort(0, descending=True)[1]
111 |
112 | boxes = boxes[order].contiguous()
113 |
114 | keep = torch.LongTensor(boxes.size(0))
115 | num_out = iou3d_nms_cuda.nms_normal_gpu(boxes, keep, thresh)
116 | return order[keep[:num_out].cuda()].contiguous(), None
117 |
--------------------------------------------------------------------------------
/pcdet/ops/iou3d_nms/src/iou3d_cpu.h:
--------------------------------------------------------------------------------
1 | #ifndef IOU3D_CPU_H
2 | #define IOU3D_CPU_H
3 |
4 | #include
5 | #include
6 | #include
7 | #include
8 |
9 | int boxes_iou_bev_cpu(at::Tensor boxes_a_tensor, at::Tensor boxes_b_tensor, at::Tensor ans_iou_tensor);
10 |
11 | #endif
12 |
--------------------------------------------------------------------------------
/pcdet/ops/iou3d_nms/src/iou3d_nms.h:
--------------------------------------------------------------------------------
1 | #ifndef IOU3D_NMS_H
2 | #define IOU3D_NMS_H
3 |
4 | #include
5 | #include
6 | #include
7 | #include
8 |
9 | int boxes_overlap_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b, at::Tensor ans_overlap);
10 | int boxes_iou_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b, at::Tensor ans_iou);
11 | int nms_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh);
12 | int nms_normal_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh);
13 |
14 | #endif
15 |
--------------------------------------------------------------------------------
/pcdet/ops/iou3d_nms/src/iou3d_nms_api.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 | #include
4 | #include
5 | #include
6 |
7 | #include "iou3d_cpu.h"
8 | #include "iou3d_nms.h"
9 |
10 |
11 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
12 | m.def("boxes_overlap_bev_gpu", &boxes_overlap_bev_gpu, "oriented boxes overlap");
13 | m.def("boxes_iou_bev_gpu", &boxes_iou_bev_gpu, "oriented boxes iou");
14 | m.def("nms_gpu", &nms_gpu, "oriented nms gpu");
15 | m.def("nms_normal_gpu", &nms_normal_gpu, "nms gpu");
16 | m.def("boxes_iou_bev_cpu", &boxes_iou_bev_cpu, "oriented boxes iou");
17 | }
18 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_batch/src/ball_query.cpp:
--------------------------------------------------------------------------------
1 | /*
2 | batch version of ball query, modified from the original implementation of official PointNet++ codes.
3 | Written by Shaoshuai Shi
4 | All Rights Reserved 2018.
5 | */
6 |
7 |
8 | #include
9 | #include
10 | #include
11 | #include
12 | #include
13 | #include "ball_query_gpu.h"
14 |
15 | extern THCState *state;
16 |
17 | #define CHECK_CUDA(x) do { \
18 | if (!x.type().is_cuda()) { \
19 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \
20 | exit(-1); \
21 | } \
22 | } while (0)
23 | #define CHECK_CONTIGUOUS(x) do { \
24 | if (!x.is_contiguous()) { \
25 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \
26 | exit(-1); \
27 | } \
28 | } while (0)
29 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x)
30 |
31 |
32 | int ball_query_wrapper_fast(int b, int n, int m, float radius, int nsample,
33 | at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor idx_tensor) {
34 | CHECK_INPUT(new_xyz_tensor);
35 | CHECK_INPUT(xyz_tensor);
36 | const float *new_xyz = new_xyz_tensor.data();
37 | const float *xyz = xyz_tensor.data();
38 | int *idx = idx_tensor.data();
39 |
40 | ball_query_kernel_launcher_fast(b, n, m, radius, nsample, new_xyz, xyz, idx);
41 | return 1;
42 | }
43 |
--------------------------------------------------------------------------------
/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/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/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/group_points.cpp:
--------------------------------------------------------------------------------
1 | /*
2 | batch version of point grouping, modified from the original implementation of official PointNet++ codes.
3 | Written by Shaoshuai Shi
4 | All Rights Reserved 2018.
5 | */
6 |
7 |
8 | #include
9 | #include
10 | #include
11 | #include
12 | #include
13 | #include "group_points_gpu.h"
14 |
15 | extern THCState *state;
16 |
17 |
18 | int group_points_grad_wrapper_fast(int b, int c, int n, int npoints, int nsample,
19 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor) {
20 |
21 | float *grad_points = grad_points_tensor.data();
22 | const int *idx = idx_tensor.data();
23 | const float *grad_out = grad_out_tensor.data();
24 |
25 | group_points_grad_kernel_launcher_fast(b, c, n, npoints, nsample, grad_out, idx, grad_points);
26 | return 1;
27 | }
28 |
29 |
30 | int group_points_wrapper_fast(int b, int c, int n, int npoints, int nsample,
31 | at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor) {
32 |
33 | const float *points = points_tensor.data();
34 | const int *idx = idx_tensor.data();
35 | float *out = out_tensor.data();
36 |
37 | group_points_kernel_launcher_fast(b, c, n, npoints, nsample, points, idx, out);
38 | return 1;
39 | }
40 |
--------------------------------------------------------------------------------
/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/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/pointnet2/pointnet2_batch/src/interpolate.cpp:
--------------------------------------------------------------------------------
1 | /*
2 | batch version of point interpolation, modified from the original implementation of official PointNet++ codes.
3 | Written by Shaoshuai Shi
4 | All Rights Reserved 2018.
5 | */
6 |
7 |
8 | #include
9 | #include
10 | #include
11 | #include
12 | #include
13 | #include
14 | #include
15 | #include
16 | #include "interpolate_gpu.h"
17 |
18 | extern THCState *state;
19 |
20 |
21 | void three_nn_wrapper_fast(int b, int n, int m, at::Tensor unknown_tensor,
22 | at::Tensor known_tensor, at::Tensor dist2_tensor, at::Tensor idx_tensor) {
23 | const float *unknown = unknown_tensor.data();
24 | const float *known = known_tensor.data();
25 | float *dist2 = dist2_tensor.data();
26 | int *idx = idx_tensor.data();
27 |
28 | three_nn_kernel_launcher_fast(b, n, m, unknown, known, dist2, idx);
29 | }
30 |
31 |
32 | void three_interpolate_wrapper_fast(int b, int c, int m, int n,
33 | at::Tensor points_tensor,
34 | at::Tensor idx_tensor,
35 | at::Tensor weight_tensor,
36 | at::Tensor out_tensor) {
37 |
38 | const float *points = points_tensor.data();
39 | const float *weight = weight_tensor.data();
40 | float *out = out_tensor.data();
41 | const int *idx = idx_tensor.data();
42 |
43 | three_interpolate_kernel_launcher_fast(b, c, m, n, points, idx, weight, out);
44 | }
45 |
46 | void three_interpolate_grad_wrapper_fast(int b, int c, int n, int m,
47 | at::Tensor grad_out_tensor,
48 | at::Tensor idx_tensor,
49 | at::Tensor weight_tensor,
50 | at::Tensor grad_points_tensor) {
51 |
52 | const float *grad_out = grad_out_tensor.data();
53 | const float *weight = weight_tensor.data();
54 | float *grad_points = grad_points_tensor.data();
55 | const int *idx = idx_tensor.data();
56 |
57 | three_interpolate_grad_kernel_launcher_fast(b, c, n, m, grad_out, idx, weight, grad_points);
58 | }
59 |
--------------------------------------------------------------------------------
/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("furthest_point_sampling_wrapper", &furthest_point_sampling_wrapper, "furthest_point_sampling_wrapper");
20 |
21 | m.def("three_nn_wrapper", &three_nn_wrapper_fast, "three_nn_wrapper_fast");
22 | m.def("three_interpolate_wrapper", &three_interpolate_wrapper_fast, "three_interpolate_wrapper_fast");
23 | m.def("three_interpolate_grad_wrapper", &three_interpolate_grad_wrapper_fast, "three_interpolate_grad_wrapper_fast");
24 | }
25 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_batch/src/sampling.cpp:
--------------------------------------------------------------------------------
1 | /*
2 | batch version of point sampling and gathering, modified from the original implementation of official PointNet++ codes.
3 | Written by Shaoshuai Shi
4 | All Rights Reserved 2018.
5 | */
6 |
7 |
8 | #include
9 | #include
10 | #include
11 | #include
12 |
13 | #include "sampling_gpu.h"
14 |
15 | extern THCState *state;
16 |
17 |
18 | int gather_points_wrapper_fast(int b, int c, int n, int npoints,
19 | at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor){
20 | const float *points = points_tensor.data();
21 | const int *idx = idx_tensor.data();
22 | float *out = out_tensor.data();
23 |
24 | gather_points_kernel_launcher_fast(b, c, n, npoints, points, idx, out);
25 | return 1;
26 | }
27 |
28 |
29 | int gather_points_grad_wrapper_fast(int b, int c, int n, int npoints,
30 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor) {
31 |
32 | const float *grad_out = grad_out_tensor.data();
33 | const int *idx = idx_tensor.data();
34 | float *grad_points = grad_points_tensor.data();
35 |
36 | gather_points_grad_kernel_launcher_fast(b, c, n, npoints, grad_out, idx, grad_points);
37 | return 1;
38 | }
39 |
40 |
41 | int furthest_point_sampling_wrapper(int b, int n, int m,
42 | at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor idx_tensor) {
43 |
44 | const float *points = points_tensor.data();
45 | float *temp = temp_tensor.data();
46 | int *idx = idx_tensor.data();
47 |
48 | furthest_point_sampling_kernel_launcher(b, n, m, points, temp, idx);
49 | return 1;
50 | }
51 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_batch/src/sampling_gpu.h:
--------------------------------------------------------------------------------
1 | #ifndef _SAMPLING_GPU_H
2 | #define _SAMPLING_GPU_H
3 |
4 | #include
5 | #include
6 | #include
7 |
8 |
9 | int gather_points_wrapper_fast(int b, int c, int n, int npoints,
10 | at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor);
11 |
12 | void gather_points_kernel_launcher_fast(int b, int c, int n, int npoints,
13 | const float *points, const int *idx, float *out);
14 |
15 |
16 | int gather_points_grad_wrapper_fast(int b, int c, int n, int npoints,
17 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor);
18 |
19 | void gather_points_grad_kernel_launcher_fast(int b, int c, int n, int npoints,
20 | const float *grad_out, const int *idx, float *grad_points);
21 |
22 |
23 | int furthest_point_sampling_wrapper(int b, int n, int m,
24 | at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor idx_tensor);
25 |
26 | void furthest_point_sampling_kernel_launcher(int b, int n, int m,
27 | const float *dataset, float *temp, int *idxs);
28 |
29 | #endif
30 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_stack/src/ball_query.cpp:
--------------------------------------------------------------------------------
1 | /*
2 | Stacked-batch-data version of ball query, modified from the original implementation of official PointNet++ codes.
3 | Written by Shaoshuai Shi
4 | All Rights Reserved 2019-2020.
5 | */
6 |
7 |
8 | #include
9 | #include
10 | #include
11 | #include
12 | #include
13 | #include "ball_query_gpu.h"
14 |
15 | extern THCState *state;
16 |
17 | #define CHECK_CUDA(x) do { \
18 | if (!x.type().is_cuda()) { \
19 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \
20 | exit(-1); \
21 | } \
22 | } while (0)
23 | #define CHECK_CONTIGUOUS(x) do { \
24 | if (!x.is_contiguous()) { \
25 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \
26 | exit(-1); \
27 | } \
28 | } while (0)
29 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x)
30 |
31 | int ball_query_wrapper_stack(int B, int M, float radius, int nsample,
32 | at::Tensor new_xyz_tensor, at::Tensor new_xyz_batch_cnt_tensor,
33 | at::Tensor xyz_tensor, at::Tensor xyz_batch_cnt_tensor, at::Tensor idx_tensor) {
34 | CHECK_INPUT(new_xyz_tensor);
35 | CHECK_INPUT(xyz_tensor);
36 | CHECK_INPUT(new_xyz_batch_cnt_tensor);
37 | CHECK_INPUT(xyz_batch_cnt_tensor);
38 |
39 | const float *new_xyz = new_xyz_tensor.data();
40 | const float *xyz = xyz_tensor.data();
41 | const int *new_xyz_batch_cnt = new_xyz_batch_cnt_tensor.data();
42 | const int *xyz_batch_cnt = xyz_batch_cnt_tensor.data();
43 | int *idx = idx_tensor.data();
44 |
45 | ball_query_kernel_launcher_stack(B, M, radius, nsample, new_xyz, new_xyz_batch_cnt, xyz, xyz_batch_cnt, idx);
46 | return 1;
47 | }
48 |
--------------------------------------------------------------------------------
/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/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/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 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_stack/src/group_points.cpp:
--------------------------------------------------------------------------------
1 | /*
2 | Stacked-batch-data version of point grouping, modified from the original implementation of official PointNet++ codes.
3 | Written by Shaoshuai Shi
4 | All Rights Reserved 2019-2020.
5 | */
6 |
7 |
8 | #include
9 | #include
10 | #include
11 | #include
12 | #include
13 | #include "group_points_gpu.h"
14 |
15 | extern THCState *state;
16 | #define CHECK_CUDA(x) do { \
17 | if (!x.type().is_cuda()) { \
18 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \
19 | exit(-1); \
20 | } \
21 | } while (0)
22 | #define CHECK_CONTIGUOUS(x) do { \
23 | if (!x.is_contiguous()) { \
24 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \
25 | exit(-1); \
26 | } \
27 | } while (0)
28 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x)
29 |
30 |
31 | int group_points_grad_wrapper_stack(int B, int M, int C, int N, int nsample,
32 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor idx_batch_cnt_tensor,
33 | at::Tensor features_batch_cnt_tensor, at::Tensor grad_features_tensor) {
34 |
35 | CHECK_INPUT(grad_out_tensor);
36 | CHECK_INPUT(idx_tensor);
37 | CHECK_INPUT(idx_batch_cnt_tensor);
38 | CHECK_INPUT(features_batch_cnt_tensor);
39 | CHECK_INPUT(grad_features_tensor);
40 |
41 | const float *grad_out = grad_out_tensor.data();
42 | const int *idx = idx_tensor.data();
43 | const int *idx_batch_cnt = idx_batch_cnt_tensor.data();
44 | const int *features_batch_cnt = features_batch_cnt_tensor.data();
45 | float *grad_features = grad_features_tensor.data();
46 |
47 | group_points_grad_kernel_launcher_stack(B, M, C, N, nsample, grad_out, idx, idx_batch_cnt, features_batch_cnt, grad_features);
48 | return 1;
49 | }
50 |
51 |
52 | int group_points_wrapper_stack(int B, int M, int C, int nsample,
53 | at::Tensor features_tensor, at::Tensor features_batch_cnt_tensor,
54 | at::Tensor idx_tensor, at::Tensor idx_batch_cnt_tensor, at::Tensor out_tensor) {
55 |
56 | CHECK_INPUT(features_tensor);
57 | CHECK_INPUT(features_batch_cnt_tensor);
58 | CHECK_INPUT(idx_tensor);
59 | CHECK_INPUT(idx_batch_cnt_tensor);
60 | CHECK_INPUT(out_tensor);
61 |
62 | const float *features = features_tensor.data();
63 | const int *idx = idx_tensor.data();
64 | const int *features_batch_cnt = features_batch_cnt_tensor.data();
65 | const int *idx_batch_cnt = idx_batch_cnt_tensor.data();
66 | float *out = out_tensor.data();
67 |
68 | group_points_kernel_launcher_stack(B, M, C, nsample, features, features_batch_cnt, idx, idx_batch_cnt, out);
69 | return 1;
70 | }
--------------------------------------------------------------------------------
/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_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 |
10 |
11 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
12 | m.def("ball_query_wrapper", &ball_query_wrapper_stack, "ball_query_wrapper_stack");
13 | m.def("voxel_query_wrapper", &voxel_query_wrapper_stack, "voxel_query_wrapper_stack");
14 |
15 | m.def("furthest_point_sampling_wrapper", &furthest_point_sampling_wrapper, "furthest_point_sampling_wrapper");
16 |
17 | m.def("group_points_wrapper", &group_points_wrapper_stack, "group_points_wrapper_stack");
18 | m.def("group_points_grad_wrapper", &group_points_grad_wrapper_stack, "group_points_grad_wrapper_stack");
19 |
20 | m.def("three_nn_wrapper", &three_nn_wrapper_stack, "three_nn_wrapper_stack");
21 | m.def("three_interpolate_wrapper", &three_interpolate_wrapper_stack, "three_interpolate_wrapper_stack");
22 | m.def("three_interpolate_grad_wrapper", &three_interpolate_grad_wrapper_stack, "three_interpolate_grad_wrapper_stack");
23 | }
24 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_stack/src/sampling.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 | #include
4 | #include
5 |
6 | #include "sampling_gpu.h"
7 |
8 | extern THCState *state;
9 | #define CHECK_CUDA(x) do { \
10 | if (!x.type().is_cuda()) { \
11 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \
12 | exit(-1); \
13 | } \
14 | } while (0)
15 | #define CHECK_CONTIGUOUS(x) do { \
16 | if (!x.is_contiguous()) { \
17 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \
18 | exit(-1); \
19 | } \
20 | } while (0)
21 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x)
22 |
23 |
24 | int furthest_point_sampling_wrapper(int b, int n, int m,
25 | at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor idx_tensor) {
26 |
27 | CHECK_INPUT(points_tensor);
28 | CHECK_INPUT(temp_tensor);
29 | CHECK_INPUT(idx_tensor);
30 |
31 | const float *points = points_tensor.data();
32 | float *temp = temp_tensor.data();
33 | int *idx = idx_tensor.data();
34 |
35 | furthest_point_sampling_kernel_launcher(b, n, m, points, temp, idx);
36 | return 1;
37 | }
38 |
--------------------------------------------------------------------------------
/pcdet/ops/pointnet2/pointnet2_stack/src/sampling_gpu.h:
--------------------------------------------------------------------------------
1 | #ifndef _SAMPLING_GPU_H
2 | #define _SAMPLING_GPU_H
3 |
4 | #include
5 | #include
6 | #include
7 |
8 |
9 | int furthest_point_sampling_wrapper(int b, int n, int m,
10 | at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor idx_tensor);
11 |
12 | void furthest_point_sampling_kernel_launcher(int b, int n, int m,
13 | const float *dataset, float *temp, int *idxs);
14 |
15 | #endif
16 |
--------------------------------------------------------------------------------
/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
9 | #include "voxel_query_gpu.h"
10 |
11 | extern THCState *state;
12 |
13 | #define CHECK_CUDA(x) do { \
14 | if (!x.type().is_cuda()) { \
15 | fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \
16 | exit(-1); \
17 | } \
18 | } while (0)
19 | #define CHECK_CONTIGUOUS(x) do { \
20 | if (!x.is_contiguous()) { \
21 | fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \
22 | exit(-1); \
23 | } \
24 | } while (0)
25 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x)
26 |
27 |
28 | int voxel_query_wrapper_stack(int M, int R1, int R2, int R3, int nsample, float radius,
29 | int z_range, int y_range, int x_range, at::Tensor new_xyz_tensor, at::Tensor xyz_tensor,
30 | at::Tensor new_coords_tensor, at::Tensor point_indices_tensor, at::Tensor idx_tensor) {
31 | CHECK_INPUT(new_coords_tensor);
32 | CHECK_INPUT(point_indices_tensor);
33 | CHECK_INPUT(new_xyz_tensor);
34 | CHECK_INPUT(xyz_tensor);
35 |
36 | const float *new_xyz = new_xyz_tensor.data();
37 | const float *xyz = xyz_tensor.data();
38 | const int *new_coords = new_coords_tensor.data();
39 | const int *point_indices = point_indices_tensor.data();
40 | int *idx = idx_tensor.data();
41 |
42 | 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);
43 | return 1;
44 | }
45 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/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/ops/spconv/__init__.py:
--------------------------------------------------------------------------------
1 | # Copyright 2019 Yan Yan
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 | from .conv import (SparseConv2d, SparseConv3d, SparseConvTranspose2d,
16 | SparseConvTranspose3d, SparseInverseConv2d,
17 | SparseInverseConv3d, SubMConv2d, SubMConv3d)
18 | from .modules import SparseModule, SparseSequential
19 | from .pool import SparseMaxPool2d, SparseMaxPool3d
20 | from .structure import SparseConvTensor, scatter_nd
21 |
22 | __all__ = [
23 | 'SparseConv2d',
24 | 'SparseConv3d',
25 | 'SubMConv2d',
26 | 'SubMConv3d',
27 | 'SparseConvTranspose2d',
28 | 'SparseConvTranspose3d',
29 | 'SparseInverseConv2d',
30 | 'SparseInverseConv3d',
31 | 'SparseModule',
32 | 'SparseSequential',
33 | 'SparseMaxPool2d',
34 | 'SparseMaxPool3d',
35 | 'SparseConvTensor',
36 | 'scatter_nd',
37 | ]
38 |
--------------------------------------------------------------------------------
/pcdet/ops/spconv/functional.py:
--------------------------------------------------------------------------------
1 | # Copyright 2019 Yan Yan
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 | from torch.autograd import Function
16 |
17 | from . import ops as ops
18 |
19 |
20 | class SparseConvFunction(Function):
21 |
22 | @staticmethod
23 | def forward(ctx, features, filters, indice_pairs, indice_pair_num,
24 | num_activate_out):
25 | ctx.save_for_backward(indice_pairs, indice_pair_num, features, filters)
26 | return ops.indice_conv(features, filters, indice_pairs,
27 | indice_pair_num, num_activate_out, False)
28 |
29 | @staticmethod
30 | def backward(ctx, grad_output):
31 | indice_pairs, indice_pair_num, features, filters = ctx.saved_tensors
32 | input_bp, filters_bp = ops.indice_conv_backward(
33 | features, filters, grad_output, indice_pairs, indice_pair_num,
34 | False)
35 |
36 | return input_bp, filters_bp, None, None, None
37 |
38 |
39 | class SparseInverseConvFunction(Function):
40 |
41 | @staticmethod
42 | def forward(ctx, features, filters, indice_pairs, indice_pair_num,
43 | num_activate_out):
44 | ctx.save_for_backward(indice_pairs, indice_pair_num, features, filters)
45 | return ops.indice_conv(features, filters, indice_pairs,
46 | indice_pair_num, num_activate_out, True, False)
47 |
48 | @staticmethod
49 | def backward(ctx, grad_output):
50 | indice_pairs, indice_pair_num, features, filters = ctx.saved_tensors
51 | input_bp, filters_bp = ops.indice_conv_backward(
52 | features, filters, grad_output, indice_pairs, indice_pair_num,
53 | True, False)
54 |
55 | return input_bp, filters_bp, None, None, None
56 |
57 |
58 | class SubMConvFunction(Function):
59 |
60 | @staticmethod
61 | def forward(ctx, features, filters, indice_pairs, indice_pair_num,
62 | num_activate_out):
63 | ctx.save_for_backward(indice_pairs, indice_pair_num, features, filters)
64 | return ops.indice_conv(features, filters, indice_pairs,
65 | indice_pair_num, num_activate_out, False, True)
66 |
67 | @staticmethod
68 | def backward(ctx, grad_output):
69 | indice_pairs, indice_pair_num, features, filters = ctx.saved_tensors
70 | input_bp, filters_bp = ops.indice_conv_backward(
71 | features, filters, grad_output, indice_pairs, indice_pair_num,
72 | False, True)
73 |
74 | return input_bp, filters_bp, None, None, None
75 |
76 |
77 | class SparseMaxPoolFunction(Function):
78 |
79 | @staticmethod
80 | def forward(ctx, features, indice_pairs, indice_pair_num,
81 | num_activate_out):
82 | out = ops.indice_maxpool(features, indice_pairs, indice_pair_num,
83 | num_activate_out)
84 | ctx.save_for_backward(indice_pairs, indice_pair_num, features, out)
85 | return out
86 |
87 | @staticmethod
88 | def backward(ctx, grad_output):
89 | indice_pairs, indice_pair_num, features, out = ctx.saved_tensors
90 | input_bp = ops.indice_maxpool_backward(features, out, grad_output,
91 | indice_pairs, indice_pair_num)
92 | return input_bp, None, None, None
93 |
94 |
95 | indice_conv = SparseConvFunction.apply
96 | indice_inverse_conv = SparseInverseConvFunction.apply
97 | indice_subm_conv = SubMConvFunction.apply
98 | indice_maxpool = SparseMaxPoolFunction.apply
99 |
--------------------------------------------------------------------------------
/pcdet/ops/spconv/include/paramsgrid.h:
--------------------------------------------------------------------------------
1 | // Copyright 2019 Yan Yan
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | #ifndef PARAMS_GRID_H_
16 | #define PARAMS_GRID_H_
17 | #include
18 | #include
19 |
20 | namespace detail {
21 | template
22 | int getTotalSize(std::vector arg) {
23 | return arg.size();
24 | }
25 |
26 | template
27 | int getTotalSize(std::vector arg, std::vector... args) {
28 | return arg.size() * getTotalSize(args...);
29 | }
30 | template
31 | int getSize(std::vector arg) {
32 | return arg.size();
33 | }
34 |
35 | template
36 | void assigner(TT &src, std::vector counter, std::vector &arg) {
37 | std::get(src) = arg[counter[Idx]];
38 | }
39 |
40 | template
41 | void assigner(TT &src, std::vector counter, std::vector &arg,
42 | std::vector &... args) {
43 | std::get(src) = arg[counter[Idx]];
44 | assigner(src, counter, args...);
45 | }
46 | } // namespace detail
47 | template
48 | std::vector> paramsGrid(std::vector... args) {
49 | int length = detail::getTotalSize(args...);
50 | std::vector sizes = {detail::getSize(args)...};
51 | int size = sizes.size();
52 |
53 | std::vector> params(length);
54 | std::vector counter(size);
55 | for (int i = 0; i < length; ++i) {
56 | detail::assigner<0>(params[i], counter, args...);
57 | counter[size - 1] += 1;
58 | for (int c = size - 1; c >= 0; --c) {
59 | if (counter[c] == sizes[c] && c > 0) {
60 | counter[c - 1] += 1;
61 | counter[c] = 0;
62 | }
63 | }
64 | }
65 | return params;
66 | }
67 |
68 | #endif
69 |
--------------------------------------------------------------------------------
/pcdet/ops/spconv/include/pybind11_utils.h:
--------------------------------------------------------------------------------
1 | // Copyright 2019 Yan Yan
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | #pragma once
16 | #include
17 | #include
18 | #include // everything needed for embedding
19 | #include
20 | #include
21 | #include
22 | #include
23 |
24 | #include
25 |
26 | namespace py = pybind11;
27 |
28 | template
29 | std::vector array2Vector(TPyObject arr){
30 | py::array arr_np = arr;
31 | size_t size = arr.attr("size").template cast();
32 | py::array_t arr_cc = arr_np;
33 | std::vector data(arr_cc.data(), arr_cc.data() + size);
34 | return data;
35 | }
36 |
37 | template
38 | std::vector arrayT2Vector(py::array_t arr)
39 | {
40 | std::vector data(arr.data(), arr.data() + arr.size());
41 | return data;
42 | }
43 |
44 | template
45 | tv::TensorView array2TensorView(TPyObject arr){
46 | py::array arr_np = arr;
47 | py::array_t arr_cc = arr_np;
48 | tv::Shape shape;
49 | for (int i = 0; i < arr_cc.ndim(); ++i){
50 | shape.push_back(arr_cc.shape(i));
51 | }
52 | return tv::TensorView(arr_cc.mutable_data(), shape);
53 | }
54 | template
55 | tv::TensorView arrayT2TensorView(py::array_t arr){
56 | tv::Shape shape;
57 | for (int i = 0; i < arr.ndim(); ++i){
58 | shape.push_back(arr.shape(i));
59 | }
60 | return tv::TensorView(arr.mutable_data(), shape);
61 | }
62 |
--------------------------------------------------------------------------------
/pcdet/ops/spconv/include/spconv/indice.h:
--------------------------------------------------------------------------------
1 | // Copyright 2019 Yan Yan
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | #ifndef SPARSE_CONV_INDICE_FUNCTOR_H_
16 | #define SPARSE_CONV_INDICE_FUNCTOR_H_
17 | #include
18 |
19 | namespace spconv {
20 | namespace functor {
21 | template
22 | struct CreateConvIndicePairFunctorP1 {
23 | Index operator()(const Device& d, tv::TensorView indicesIn,
24 | tv::TensorView indicesOut,
25 | tv::TensorView gridsOut,
26 | tv::TensorView indicePairs,
27 | tv::TensorView indiceNum,
28 | tv::TensorView indicePairUnique,
29 | const tv::SimpleVector kernelSize,
30 | const tv::SimpleVector stride,
31 | const tv::SimpleVector padding,
32 | const tv::SimpleVector dilation,
33 | const tv::SimpleVector outSpatialShape,
34 | bool transpose);
35 | };
36 |
37 | template
38 | struct CreateConvIndicePairFunctorP2 {
39 | Index operator()(const Device& d, tv::TensorView indicesIn,
40 | tv::TensorView indicesOut,
41 | tv::TensorView gridsOut,
42 | tv::TensorView indicePairs,
43 | tv::TensorView indiceNum,
44 | tv::TensorView indicePairUnique,
45 | const tv::SimpleVector outSpatialShape,
46 | bool transpose, bool resetGrid = false);
47 | };
48 |
49 | template
50 | struct CreateConvIndicePairFunctor {
51 | Index operator()(const Device& d, tv::TensorView indicesIn,
52 | tv::TensorView indicesOut,
53 | tv::TensorView gridsOut,
54 | tv::TensorView indicePairs,
55 | tv::TensorView indiceNum,
56 | const tv::SimpleVector kernelSize,
57 | const tv::SimpleVector stride,
58 | const tv::SimpleVector padding,
59 | const tv::SimpleVector dilation,
60 | const tv::SimpleVector outSpatialShape,
61 | bool transpose, bool resetGrid = false);
62 | };
63 |
64 | template
65 | struct CreateSubMIndicePairFunctor {
66 | Index operator()(const Device& d, tv::TensorView