├── lidardet ├── __init__.py ├── losses │ ├── __init__.py │ ├── lovasz_softmax.py │ └── loss_utils.py ├── ops │ ├── __init__.py │ └── lidar_bev │ │ ├── __init__.py │ │ └── build.sh ├── utils │ ├── __init__.py │ ├── buildtools │ │ ├── __init__.py │ │ ├── pybind11_build.py │ │ └── command.py │ ├── input │ │ ├── point_cloud │ │ │ ├── __init__.py │ │ │ ├── box_ops.cc │ │ │ ├── point_cloud_ops.cc │ │ │ ├── point_cloud_ops.h │ │ │ ├── box_ops.h │ │ │ ├── bev_ops.py │ │ │ └── point_cloud_ops.py │ │ ├── __init__.py │ │ └── voxel_generator.py │ ├── registry.py │ ├── loader.py │ ├── precision_recall.py │ ├── config.py │ ├── range_image.py │ ├── box_coder_utils.py │ ├── dist.py │ ├── scan_unfolding.py │ ├── find.py │ ├── projection.py │ ├── laserscanvis.py │ ├── common.py │ └── geometry.py ├── datasets │ ├── augmentor │ │ ├── __init__.py │ │ ├── augmentor_utils.py │ │ ├── data_augmentor.py │ │ └── database_sampler.py │ ├── processor │ │ ├── __init__.py │ │ └── data_processor.py │ ├── __init__.py │ ├── registry.py │ ├── trajectory │ │ ├── __init__.py │ │ └── utils.py │ ├── builder.py │ └── base.py └── models │ ├── trajectory_predictor │ ├── __init__.py │ ├── conv_traj.py │ ├── covernet.py │ ├── header.py │ ├── predictor_base.py │ ├── conv_header.py │ └── backbone.py │ ├── registry.py │ ├── builder.py │ └── __init__.py ├── datasets └── .gitignore ├── data ├── .gitignore └── kitti │ ├── trajset.npy │ └── traj_pred_kitti10.gif ├── tools ├── scripts │ ├── dist_test.sh │ └── dist_train.sh ├── eval_utils │ └── trajectory_prediction.py ├── test.py ├── train_utils │ ├── optimization │ │ ├── __init__.py │ │ └── learning_schedules_fastai.py │ └── train_utils.py └── train.py ├── setup.py ├── README.md ├── .gitignore └── config └── trajectory_prediction └── transformer.yaml /lidardet/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /lidardet/losses/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /lidardet/ops/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /lidardet/utils/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /lidardet/ops/lidar_bev/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /lidardet/utils/buildtools/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /datasets/.gitignore: -------------------------------------------------------------------------------- 1 | * 2 | !.gitignore 3 | -------------------------------------------------------------------------------- /data/.gitignore: -------------------------------------------------------------------------------- 1 | *.pkl 2 | /kitti/gt_database 3 | -------------------------------------------------------------------------------- /lidardet/utils/input/point_cloud/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /lidardet/utils/input/__init__.py: -------------------------------------------------------------------------------- 1 | from . import voxel_generator 2 | -------------------------------------------------------------------------------- /lidardet/datasets/augmentor/__init__.py: -------------------------------------------------------------------------------- 1 | from .data_augmentor import DataAugmentor 2 | -------------------------------------------------------------------------------- /lidardet/datasets/processor/__init__.py: -------------------------------------------------------------------------------- 1 | from .data_processor import DataProcessor 2 | -------------------------------------------------------------------------------- /lidardet/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | from .trajectory import MTPDataset, CoverNetDataset, HeatMapDataset 2 | -------------------------------------------------------------------------------- /data/kitti/trajset.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jiaolong/trajectory-prediction/HEAD/data/kitti/trajset.npy -------------------------------------------------------------------------------- /data/kitti/traj_pred_kitti10.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jiaolong/trajectory-prediction/HEAD/data/kitti/traj_pred_kitti10.gif -------------------------------------------------------------------------------- /lidardet/models/trajectory_predictor/__init__.py: -------------------------------------------------------------------------------- 1 | from .predictor_base import PredictorBase 2 | from .conv_traj import ConvTrajPredictor 3 | -------------------------------------------------------------------------------- /lidardet/datasets/registry.py: -------------------------------------------------------------------------------- 1 | from lidardet.utils.registry import Registry 2 | 3 | DATASETS = Registry('datasets') 4 | AUGMENTORS = Registry('augmentors') 5 | -------------------------------------------------------------------------------- /lidardet/models/registry.py: -------------------------------------------------------------------------------- 1 | from ..utils.registry import Registry 2 | 3 | # trajectory predictor 4 | TRAJECTORY_PREDICTOR = Registry('trajectory_predictor') 5 | -------------------------------------------------------------------------------- /lidardet/datasets/trajectory/__init__.py: -------------------------------------------------------------------------------- 1 | from .mtp_dataset import MTPDataset 2 | from .heatmap_dataset import HeatMapDataset 3 | from .covernet_dataset import CoverNetDataset 4 | -------------------------------------------------------------------------------- /lidardet/ops/lidar_bev/build.sh: -------------------------------------------------------------------------------- 1 | c++ -O3 -Wall -shared -std=c++11 -fPIC -I /usr/include/eigen3 `python -m pybind11 --includes` bev.cpp -o bev`python3.6-config --extension-suffix` 2 | -------------------------------------------------------------------------------- /tools/scripts/dist_test.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -x 4 | NGPUS=$1 5 | PY_ARGS=${@:2} 6 | 7 | python3 -m torch.distributed.launch --nproc_per_node=${NGPUS} tools/test.py --launcher pytorch ${PY_ARGS} 8 | 9 | -------------------------------------------------------------------------------- /tools/scripts/dist_train.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -x 4 | NGPUS=$1 5 | PY_ARGS=${@:2} 6 | 7 | python3 -m torch.distributed.launch --nproc_per_node=${NGPUS} tools/train.py --launcher pytorch ${PY_ARGS} 8 | -------------------------------------------------------------------------------- /lidardet/utils/input/point_cloud/box_ops.cc: -------------------------------------------------------------------------------- 1 | #include "box_ops.h" 2 | 3 | PYBIND11_MODULE(box_ops_cc, m) { 4 | m.doc() = "box ops written by c++"; 5 | m.def("rbbox_iou", &rbbox_iou, 6 | py::return_value_policy::reference_internal, "rbbox iou", 7 | "box_corners"_a = 1, "qbox_corners"_a = 2, "standup_iou"_a = 3, 8 | "standup_thresh"_a = 4); 9 | m.def("rbbox_iou", &rbbox_iou, 10 | py::return_value_policy::reference_internal, "rbbox iou", 11 | "box_corners"_a = 1, "qbox_corners"_a = 2, "standup_iou"_a = 3, 12 | "standup_thresh"_a = 4); 13 | } 14 | -------------------------------------------------------------------------------- /lidardet/utils/input/point_cloud/point_cloud_ops.cc: -------------------------------------------------------------------------------- 1 | 2 | #include "point_cloud_ops.h" 3 | 4 | PYBIND11_MODULE(point_cloud_ops_cc, m) { 5 | m.doc() = "pybind11 example plugin"; 6 | m.def("points_to_voxel", &points_to_voxel, "matrix tensor_square", 7 | "points"_a = 1, "voxels"_a = 2, "coors"_a = 3, 8 | "num_points_per_voxel"_a = 4, "voxel_size"_a = 5, "coors_range"_a = 6, 9 | "max_points"_a = 7, "max_voxels"_a = 8); 10 | m.def("points_to_voxel", &points_to_voxel, "matrix tensor_square", 11 | "points"_a = 1, "voxels"_a = 2, "coors"_a = 3, 12 | "num_points_per_voxel"_a = 4, "voxel_size"_a = 5, "coors_range"_a = 6, 13 | "max_points"_a = 7, "max_voxels"_a = 8); 14 | } -------------------------------------------------------------------------------- /lidardet/models/builder.py: -------------------------------------------------------------------------------- 1 | from ..utils.registry import build_from_cfg 2 | from torch import nn 3 | 4 | from .registry import ( 5 | TRAJECTORY_PREDICTOR 6 | ) 7 | 8 | def build(cfg, registry, default_args=None): 9 | if isinstance(cfg, list): 10 | modules = [build_from_cfg(cfg_, registry, default_args) for cfg_ in cfg] 11 | return nn.Sequential(*modules) 12 | else: 13 | return build_from_cfg(cfg, registry, default_args) 14 | 15 | def build_trajectory_predictor(cfg): 16 | return build(cfg, TRAJECTORY_PREDICTOR) 17 | 18 | def build_model(cfg, train_cfg=None, test_cfg=None): 19 | if cfg.task == 'trajectory_prediction': 20 | return build_trajectory_predictor(cfg) 21 | else: 22 | raise ValueError("Unknown task {}".format(cfg.task)) 23 | -------------------------------------------------------------------------------- /lidardet/utils/registry.py: -------------------------------------------------------------------------------- 1 | import inspect 2 | 3 | class Registry(object): 4 | def __init__(self, name): 5 | self.name = name 6 | self.obj_dict = {} 7 | 8 | def get(self, key): 9 | if not key in self.obj_dict: 10 | raise ValueError("{} is not a registered class.".format(key)) 11 | return self.obj_dict.get(key, None) 12 | 13 | def register(self, cls): 14 | if not inspect.isclass(cls): 15 | raise TypeError('module must be a class, but got {}'.format(cls)) 16 | cls_name = cls.__name__ 17 | if cls_name in self.obj_dict: 18 | raise KeyError("{} is already registered in {}".format(cls_name, self.name)) 19 | 20 | self.obj_dict[cls_name] = cls 21 | return cls 22 | 23 | def build_from_cfg(cfg, registry, default_args=None): 24 | """ 25 | Build an object from config dict. 26 | """ 27 | obj_type = cfg.pop('type') 28 | obj_class = registry.get(obj_type) 29 | if default_args is not None: 30 | return obj_class(cfg=cfg, **default_args) 31 | else: 32 | return obj_class(cfg=cfg) 33 | -------------------------------------------------------------------------------- /lidardet/models/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from collections import namedtuple 4 | 5 | # trajectory predictor 6 | from .trajectory_predictor import * 7 | 8 | from .builder import ( 9 | build_model 10 | ) 11 | 12 | def load_data_to_gpu(batch_dict): 13 | for key, val in batch_dict.items(): 14 | if not isinstance(val, np.ndarray): 15 | continue 16 | if key in ['frame_id', 'metadata', 'calib', 'image_shape']: 17 | continue 18 | batch_dict[key] = torch.from_numpy(val).float().cuda() 19 | 20 | def model_fn_decorator(): 21 | ModelReturn = namedtuple('ModelReturn', ['loss', 'tb_dict', 'disp_dict']) 22 | 23 | def model_func(model, batch_dict): 24 | load_data_to_gpu(batch_dict) 25 | 26 | ret_dict, tb_dict, disp_dict = model(batch_dict) 27 | 28 | loss = ret_dict['loss'].mean() 29 | if hasattr(model, 'update_global_step'): 30 | model.update_global_step() 31 | else: 32 | model.module.update_global_step() 33 | 34 | return ModelReturn(loss, tb_dict, disp_dict) 35 | 36 | return model_func 37 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | from setuptools import Extension 4 | from setuptools import setup, find_packages 5 | 6 | __version__ = '0.0.0' 7 | 8 | py_dir = '/usr/local/lib/python{}/dist-packages/'.format(sys.version[:3]) 9 | 10 | # lidar BEV cpp extension 11 | lidar_bev_ext_module = Extension( 12 | name = 'lidardet.ops.lidar_bev.bev', 13 | sources = ['lidardet/ops/lidar_bev/bev.cpp'], 14 | include_dirs = ['/usr/include/eigen3', py_dir + 'pybind11/include'], 15 | language='c++' 16 | ) 17 | 18 | if __name__ == '__main__': 19 | 20 | setup( 21 | name='lidardet', 22 | version="0.0.0", 23 | description='LidarDet is a general codebase for learning based perception task', 24 | install_requires=[ 25 | 'numpy', 26 | 'torch>=1.1', 27 | 'numba', 28 | 'tensorboardX', 29 | 'easydict', 30 | 'pyyaml' 31 | ], 32 | author='jiaolong', 33 | author_email='jiaolongxu@gmail.com', 34 | license='Apache License 2.0', 35 | packages=find_packages(exclude=['tools', 'data', 'output', 'cache', 'ros', 'docker', 'config']), 36 | ext_modules=[lidar_bev_ext_module] 37 | ) 38 | -------------------------------------------------------------------------------- /lidardet/utils/input/voxel_generator.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from .point_cloud.point_cloud_ops import points_to_voxel 3 | 4 | 5 | class VoxelGenerator: 6 | def __init__(self, voxel_size, point_cloud_range, max_num_points, max_voxels=20000): 7 | point_cloud_range = np.array(point_cloud_range, dtype=np.float32) 8 | # [0, -40, -3, 70.4, 40, 1] 9 | voxel_size = np.array(voxel_size, dtype=np.float32) 10 | grid_size = (point_cloud_range[3:] - point_cloud_range[:3]) / voxel_size 11 | grid_size = np.round(grid_size).astype(np.int64) 12 | 13 | self._voxel_size = voxel_size 14 | self._point_cloud_range = point_cloud_range 15 | self._max_num_points = max_num_points 16 | self._max_voxels = max_voxels 17 | self._grid_size = grid_size 18 | 19 | def generate(self, points, max_voxels=20000): 20 | return points_to_voxel( 21 | points, 22 | self._voxel_size, 23 | self._point_cloud_range, 24 | self._max_num_points, 25 | True, 26 | self._max_voxels, 27 | ) 28 | 29 | @property 30 | def voxel_size(self): 31 | return self._voxel_size 32 | 33 | @property 34 | def max_num_points_per_voxel(self): 35 | return self._max_num_points 36 | 37 | @property 38 | def point_cloud_range(self): 39 | return self._point_cloud_range 40 | 41 | @property 42 | def grid_size(self): 43 | return self._grid_size 44 | -------------------------------------------------------------------------------- /lidardet/models/trajectory_predictor/conv_traj.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import torch.nn as nn 4 | import numpy as np 5 | 6 | from .conv_header import ConvHeader 7 | from .backbone import Bottleneck, BackBone 8 | from .predictor_base import PredictorBase 9 | from ..registry import TRAJECTORY_PREDICTOR 10 | 11 | @TRAJECTORY_PREDICTOR.register 12 | class ConvTrajPredictor(PredictorBase): 13 | def __init__(self, cfg): 14 | super().__init__(cfg=cfg) 15 | 16 | self.backbone = BackBone(Bottleneck, cfg.backbone) 17 | self.header = ConvHeader(cfg.header) 18 | 19 | for m in self.modules(): 20 | if isinstance(m, nn.Conv2d): 21 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 22 | m.weight.data.normal_(0, math.sqrt(2. / n)) 23 | elif isinstance(m, nn.BatchNorm2d): 24 | m.weight.data.fill_(1) 25 | m.bias.data.zero_() 26 | 27 | def forward(self, batch_dict): 28 | x = batch_dict['lidar_bev'] 29 | img_hmi = batch_dict['img_hmi'] 30 | x = torch.cat([x, img_hmi], dim=1) 31 | 32 | batch_dict['input'] = x 33 | batch_dict = self.backbone(batch_dict) 34 | batch_dict = self.header(batch_dict) 35 | 36 | if self.training: 37 | loss, tb_dict, disp_dict = self.get_training_loss() 38 | 39 | ret_dict = { 40 | 'loss': loss 41 | } 42 | return ret_dict, tb_dict, disp_dict 43 | else: 44 | pred_dicts, recall_dicts = self.post_processing(batch_dict) 45 | return pred_dicts, recall_dicts 46 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Trajectory Prediction for Autonomous Driving with Topometric Map 2 | 3 | ![image](https://github.com/Jiaolong/trajectory-prediction/tree/main/data/kitti/traj_pred_kitti10.gif) 4 | 5 | Repository for the paper ["Trajectory Prediction for Autonomous Driving with Topometric Map"](https://arxiv.org/abs/2105.03869). 6 | ``` 7 | @inproceedings{traj-pred:2022, 8 | title={Trajectory Prediction for Autonomous Driving with Topometric Map}, 9 | author={J. Xu, L. Xiao, D. Zhao etal}, 10 | booktitle={ICRA}, 11 | year={2022} 12 | } 13 | ``` 14 | 15 | ## Requirements 16 | 17 | - python 3.6+ 18 | 19 | - pytorch 1.4+ 20 | 21 | ## Install build requirements 22 | 23 | ```shell 24 | pip install -v -e . # or "python setup.py develop" 25 | ``` 26 | 27 | ## Pretrained models 28 | 29 | Pre-trained weights can be downloaded [here](https://pan.baidu.com/s/1Ns7qjW352rMXJhleGJN2TQ)(code: uf9g) 30 | 31 | ## Dataset 32 | 33 | ``` 34 | ├── datasets 35 | │   └── KITTI_RAW 36 |    └── trajectory_prediction 37 |    ├── 07 38 |    └── 10 39 | ``` 40 | 41 | Testing dataset kitti-10 can be downloaded [here](https://pan.baidu.com/s/1DrPRNWfMOy7JMc_TOzdV7w)(code: kbuf) 42 | 43 | ## Train & Test 44 | 45 | ### Train 46 | 47 | * Train with multiple GPUs: 48 | ```shell script 49 | sh tools/scripts/dist_train.sh ${NUM_GPUS} -c ${CONFIG_FILE} 50 | ``` 51 | 52 | * Train with a single GPU: 53 | ```shell script 54 | python tools/train.py --cfg config/trajectory_prediction/transformer.yaml 55 | ``` 56 | 57 | ### Test 58 | 59 | * Test with a pretrained model: 60 | ```shell script 61 | python tools/test.py --cfg config/trajectory_prediction/transformer.yaml --ckpt cache/transformer_epoch_120.pth 62 | ``` 63 | -------------------------------------------------------------------------------- /tools/eval_utils/trajectory_prediction.py: -------------------------------------------------------------------------------- 1 | import tqdm 2 | import time 3 | import pickle 4 | import numpy as np 5 | import torch 6 | from lidardet.models import load_data_to_gpu 7 | 8 | def eval_trajectory_prediction(cfg, model, dataloader, logger, dist_test=False, save_to_file=False, result_dir=None): 9 | result_dir.mkdir(parents=True, exist_ok=True) 10 | 11 | dataset = dataloader.dataset 12 | 13 | if dist_test: 14 | num_gpus = torch.cuda.device_count() 15 | local_rank = cfg.local_rank % num_gpus 16 | model = torch.nn.parallel.DistributedDataParallel( 17 | model, 18 | device_ids=[local_rank], 19 | broadcast_buffers=False 20 | ) 21 | model.eval() 22 | 23 | if cfg.local_rank == 0: 24 | progress_bar = tqdm.tqdm(total=len(dataloader), leave=True, desc='eval', dynamic_ncols=True) 25 | start_time = time.time() 26 | result_dicts_list = [] 27 | for i, batch_dict in enumerate(dataloader): 28 | load_data_to_gpu(batch_dict) 29 | with torch.no_grad(): 30 | pred_dicts, ret_dict = model(batch_dict) 31 | 32 | result_dicts = dataset.generate_prediction_dicts(batch_dict, pred_dicts, 33 | output_path = result_dir if save_to_file else None) 34 | result_dicts_list.append(result_dicts) 35 | 36 | if cfg.local_rank == 0: 37 | # progress_bar.set_postfix(disp_dict) 38 | progress_bar.update() 39 | 40 | if cfg.local_rank == 0: 41 | progress_bar.close() 42 | 43 | sec_per_example = (time.time() - start_time) / len(dataloader.dataset) 44 | logger.info('Test finished(sec_per_example: %.4f second).' % sec_per_example) 45 | 46 | if cfg.local_rank != 0: 47 | return {} 48 | 49 | dataset.evaluation(result_dicts_list) 50 | 51 | logger.info('Result is save to %s' % result_dir) 52 | logger.info('****************Evaluation done.*****************') 53 | -------------------------------------------------------------------------------- /lidardet/datasets/builder.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.utils.data import DataLoader 3 | from torch.utils.data import DistributedSampler as _DistributedSampler 4 | 5 | from lidardet.utils.dist import get_dist_info 6 | from lidardet.utils.registry import build_from_cfg 7 | from .registry import DATASETS 8 | 9 | class DistributedSampler(_DistributedSampler): 10 | 11 | def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True): 12 | super().__init__(dataset, num_replicas=num_replicas, rank=rank) 13 | self.shuffle = shuffle 14 | 15 | def __iter__(self): 16 | if self.shuffle: 17 | g = torch.Generator() 18 | g.manual_seed(self.epoch) 19 | indices = torch.randperm(len(self.dataset), generator=g).tolist() 20 | else: 21 | indices = torch.arange(len(self.dataset)).tolist() 22 | 23 | indices += indices[:(self.total_size - len(indices))] 24 | assert len(indices) == self.total_size 25 | 26 | indices = indices[self.rank:self.total_size:self.num_replicas] 27 | assert len(indices) == self.num_samples 28 | 29 | return iter(indices) 30 | 31 | def build_dataset(cfg, logger=None): 32 | dataset = build_from_cfg(cfg, DATASETS, dict(logger=logger)) 33 | return dataset 34 | 35 | def build_dataloader(cfg, dist=False, training=False, logger=None): 36 | dataset = build_dataset(cfg, logger) 37 | 38 | if dist: 39 | if training: 40 | sampler = torch.utils.data.distributed.DistributedSampler(dataset) 41 | else: 42 | rank, world_size = get_dist_info() 43 | sampler = DistributedSampler(dataset, world_size, rank, shuffle=False) 44 | else: 45 | sampler = None 46 | 47 | dataloader = DataLoader( 48 | dataset, batch_size=cfg.batch_size, pin_memory=True, num_workers=cfg.num_workers, 49 | shuffle=(sampler is None) and training, collate_fn=dataset.collate_batch, 50 | drop_last=False, sampler=sampler, timeout=0 51 | ) 52 | 53 | return dataset, dataloader, sampler 54 | -------------------------------------------------------------------------------- /lidardet/datasets/augmentor/augmentor_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from lidardet.utils import geometry 3 | 4 | 5 | def random_flip_along_x(gt_boxes, points): 6 | """ 7 | Args: 8 | gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading] 9 | points: (M, 3 + C) 10 | Returns: 11 | """ 12 | enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5]) 13 | if enable: 14 | gt_boxes[:, 1] = -gt_boxes[:, 1] 15 | gt_boxes[:, 6] = -gt_boxes[:, 6] 16 | points[:, 1] = -points[:, 1] 17 | return gt_boxes, points 18 | 19 | 20 | def random_flip_along_y(gt_boxes, points): 21 | """ 22 | Args: 23 | gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading] 24 | points: (M, 3 + C) 25 | Returns: 26 | """ 27 | enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5]) 28 | if enable: 29 | gt_boxes[:, 0] = -gt_boxes[:, 0] 30 | gt_boxes[:, 6] = -(gt_boxes[:, 6] + np.pi) 31 | points[:, 0] = -points[:, 0] 32 | return gt_boxes, points 33 | 34 | 35 | def global_rotation(gt_boxes, points, rot_range): 36 | """ 37 | Args: 38 | gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading] 39 | points: (M, 3 + C), 40 | rot_range: [min, max] 41 | Returns: 42 | """ 43 | noise_rotation = np.random.uniform(rot_range[0], rot_range[1]) 44 | points = geometry.rotate_points_along_z(points[np.newaxis, :, :], np.array([noise_rotation]))[0] 45 | gt_boxes[:, 0:3] = geometry.rotate_points_along_z(gt_boxes[np.newaxis, :, 0:3], np.array([noise_rotation]))[0] 46 | gt_boxes[:, 6] += noise_rotation 47 | return gt_boxes, points 48 | 49 | 50 | def global_scaling(gt_boxes, points, scale_range): 51 | """ 52 | Args: 53 | gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading] 54 | points: (M, 3 + C), 55 | scale_range: [min, max] 56 | Returns: 57 | """ 58 | if scale_range[1] - scale_range[0] < 1e-3: 59 | return gt_boxes, points 60 | noise_scale = np.random.uniform(scale_range[0], scale_range[1]) 61 | points[:, :3] *= noise_scale 62 | gt_boxes[:, :6] *= noise_scale 63 | return gt_boxes, points 64 | -------------------------------------------------------------------------------- /tools/test.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import pprint 4 | from lidardet.utils import dist 5 | from lidardet.utils.config import get_config 6 | from lidardet.utils.common import get_logger, set_random_seed 7 | 8 | from lidardet.models.builder import build_model 9 | from lidardet.datasets.builder import build_dataloader 10 | 11 | from eval_utils.trajectory_prediction import eval_trajectory_prediction 12 | 13 | def eval_single_model(model, test_loader, cfg, ckpt_file, logger, dist_test): 14 | logger.info('Evaluting {:s}'.format(str(ckpt_file))) 15 | # load checkpoint 16 | model.load_params_from_file(filename=ckpt_file, logger=logger, to_cpu=dist_test) 17 | model.cuda() 18 | 19 | eval_trajectory_prediction(cfg, model, test_loader, logger, dist_test=dist_test, 20 | result_dir=cfg.output_dir, save_to_file=cfg.get('save_to_file', True)) 21 | 22 | def main(): 23 | 24 | cfg = get_config() 25 | logger = get_logger(cfg.log_dir, cfg.tag) 26 | # log to file 27 | logger.info(pprint.pformat(cfg)) 28 | 29 | if cfg.launcher == 'none': 30 | dist_test = False 31 | else: 32 | logger.info('Start distributed testing ...') 33 | cfg.batch_size, cfg.local_rank = dist.init_dist_pytorch( 34 | cfg.batch_size, cfg.local_rank, backend='nccl' 35 | ) 36 | cfg.data.val.batch_size = cfg.batch_size 37 | dist_test = True 38 | 39 | 40 | if dist_test: 41 | total_gpus = dist.get_world_size() 42 | logger.info('total_batch_size: %d' % (total_gpus * cfg.batch_size)) 43 | 44 | test_set, test_loader, sampler = build_dataloader( 45 | cfg.data.val, dist=dist_test, training=False, logger=logger) 46 | 47 | model = build_model(cfg.model) 48 | # logger.info(model) 49 | 50 | with torch.no_grad(): 51 | if cfg.eval_all: 52 | ckpt_files = list(cfg.model_dir.glob('*.pth')) 53 | ckpt_files = sorted(ckpt_files, key=os.path.getmtime, reverse=True) 54 | for ckpt_file in ckpt_files: 55 | eval_single_model(model, test_loader, cfg, ckpt_file=ckpt_file, logger=logger, dist_test=dist_test) 56 | else: 57 | eval_single_model(model, test_loader, cfg, ckpt_file=cfg.ckpt, logger=logger, dist_test=dist_test) 58 | 59 | if __name__ == '__main__': 60 | main() 61 | 62 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | pyntcloud_plot.* 2 | pyntcloud_plot_assets/ 3 | /cache 4 | # Byte-compiled / optimized / DLL files 5 | __pycache__/ 6 | *.py[cod] 7 | *$py.class 8 | 9 | # C extensions 10 | *.o 11 | *.so 12 | 13 | # Distribution / packaging 14 | .Python 15 | build/ 16 | develop-eggs/ 17 | dist/ 18 | downloads/ 19 | eggs/ 20 | .eggs/ 21 | lib/ 22 | lib64/ 23 | parts/ 24 | sdist/ 25 | var/ 26 | wheels/ 27 | pip-wheel-metadata/ 28 | share/python-wheels/ 29 | *.egg-info/ 30 | .installed.cfg 31 | *.egg 32 | MANIFEST 33 | 34 | # PyInstaller 35 | # Usually these files are written by a python script from a template 36 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 37 | *.manifest 38 | *.spec 39 | 40 | # Installer logs 41 | pip-log.txt 42 | pip-delete-this-directory.txt 43 | 44 | # Unit test / coverage reports 45 | htmlcov/ 46 | .tox/ 47 | .nox/ 48 | .coverage 49 | .coverage.* 50 | .cache 51 | nosetests.xml 52 | coverage.xml 53 | *.cover 54 | *.py,cover 55 | .hypothesis/ 56 | .pytest_cache/ 57 | 58 | # Translations 59 | *.mo 60 | *.pot 61 | 62 | # Django stuff: 63 | *.log 64 | local_settings.py 65 | db.sqlite3 66 | db.sqlite3-journal 67 | 68 | # Flask stuff: 69 | instance/ 70 | .webassets-cache 71 | 72 | # Scrapy stuff: 73 | .scrapy 74 | 75 | # Sphinx documentation 76 | docs/_build/ 77 | 78 | # PyBuilder 79 | target/ 80 | 81 | # Jupyter Notebook 82 | .ipynb_checkpoints 83 | 84 | # IPython 85 | profile_default/ 86 | ipython_config.py 87 | 88 | # pyenv 89 | .python-version 90 | 91 | # pipenv 92 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 93 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 94 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 95 | # install all needed dependencies. 96 | #Pipfile.lock 97 | 98 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 99 | __pypackages__/ 100 | 101 | # Celery stuff 102 | celerybeat-schedule 103 | celerybeat.pid 104 | 105 | # SageMath parsed files 106 | *.sage.py 107 | 108 | # Environments 109 | .env 110 | .venv 111 | env/ 112 | venv/ 113 | ENV/ 114 | env.bak/ 115 | venv.bak/ 116 | 117 | # Spyder project settings 118 | .spyderproject 119 | .spyproject 120 | 121 | # Rope project settings 122 | .ropeproject 123 | 124 | # mkdocs documentation 125 | /site 126 | 127 | # mypy 128 | .mypy_cache/ 129 | .dmypy.json 130 | dmypy.json 131 | 132 | # Pyre type checker 133 | .pyre/ 134 | -------------------------------------------------------------------------------- /lidardet/utils/input/point_cloud/point_cloud_ops.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | namespace py = pybind11; 10 | using namespace pybind11::literals; 11 | 12 | template 13 | int points_to_voxel(py::array_t points, py::array_t voxels, 14 | py::array_t coors, 15 | py::array_t num_points_per_voxel, 16 | py::array_t coor_to_voxelidx, 17 | std::vector voxel_size, 18 | std::vector coors_range, int max_points, 19 | int max_voxels) { 20 | auto points_rw = points.template mutable_unchecked<2>(); 21 | auto voxels_rw = voxels.template mutable_unchecked<3>(); 22 | auto coors_rw = coors.mutable_unchecked<2>(); 23 | auto num_points_per_voxel_rw = num_points_per_voxel.mutable_unchecked<1>(); 24 | auto coor_to_voxelidx_rw = coor_to_voxelidx.mutable_unchecked(); 25 | auto N = points_rw.shape(0); 26 | auto num_features = points_rw.shape(1); 27 | // auto ndim = points_rw.shape(1) - 1; 28 | constexpr int ndim_minus_1 = NDim - 1; 29 | int voxel_num = 0; 30 | bool failed = false; 31 | int coor[NDim]; 32 | int c; 33 | int grid_size[NDim]; 34 | for (int i = 0; i < NDim; ++i) { 35 | grid_size[i] = 36 | round((coors_range[NDim + i] - coors_range[i]) / voxel_size[i]); 37 | } 38 | int voxelidx, num; 39 | for (int i = 0; i < N; ++i) { 40 | failed = false; 41 | for (int j = 0; j < NDim; ++j) { 42 | c = floor((points_rw(i, j) - coors_range[j]) / voxel_size[j]); 43 | if ((c < 0 || c >= grid_size[j])) { 44 | failed = true; 45 | break; 46 | } 47 | coor[ndim_minus_1 - j] = c; 48 | } 49 | if (failed) 50 | continue; 51 | voxelidx = coor_to_voxelidx_rw(coor[0], coor[1], coor[2]); 52 | if (voxelidx == -1) { 53 | voxelidx = voxel_num; 54 | if (voxel_num >= max_voxels) 55 | break; 56 | voxel_num += 1; 57 | coor_to_voxelidx_rw(coor[0], coor[1], coor[2]) = voxelidx; 58 | for (int k = 0; k < NDim; ++k) { 59 | coors_rw(voxelidx, k) = coor[k]; 60 | } 61 | } 62 | num = num_points_per_voxel_rw(voxelidx); 63 | if (num < max_points) { 64 | for (int k = 0; k < num_features; ++k) { 65 | voxels_rw(voxelidx, num, k) = points_rw(i, k); 66 | } 67 | num_points_per_voxel_rw(voxelidx) += 1; 68 | } 69 | } 70 | return voxel_num; 71 | } 72 | -------------------------------------------------------------------------------- /tools/train_utils/optimization/__init__.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.optim as optim 3 | import torch.optim.lr_scheduler as lr_sched 4 | from functools import partial 5 | from .fastai_optim import OptimWrapper 6 | from .learning_schedules_fastai import OneCycle, CosineWarmupLR 7 | 8 | 9 | def build_optimizer(model, optim_cfg): 10 | if optim_cfg.optimizer == 'adam': 11 | optimizer = optim.Adam(model.parameters(), lr=optim_cfg.lr, weight_decay=optim_cfg.weight_decay) 12 | elif optim_cfg.optimizer == 'sgd': 13 | optimizer = optim.SGD( 14 | model.parameters(), lr=optim_cfg.lr, weight_decay=optim_cfg.weight_decay, 15 | momentum=optim_cfg.momentum 16 | ) 17 | elif optim_cfg.optimizer == 'adam_onecycle': 18 | def children(m: nn.Module): 19 | return list(m.children()) 20 | 21 | def num_children(m: nn.Module) -> int: 22 | return len(children(m)) 23 | 24 | flatten_model = lambda m: sum(map(flatten_model, m.children()), []) if num_children(m) else [m] 25 | get_layer_groups = lambda m: [nn.Sequential(*flatten_model(m))] 26 | 27 | optimizer_func = partial(optim.Adam, betas=(0.9, 0.99)) 28 | optimizer = OptimWrapper.create( 29 | optimizer_func, 3e-3, get_layer_groups(model), wd=optim_cfg.weight_decay, true_wd=True, bn_wd=True 30 | ) 31 | else: 32 | raise NotImplementedError 33 | 34 | return optimizer 35 | 36 | 37 | def build_scheduler(optimizer, total_iters_each_epoch, total_epochs, last_epoch, optim_cfg): 38 | decay_steps = [x * total_iters_each_epoch for x in optim_cfg.decay_step_list] 39 | def lr_lbmd(cur_epoch): 40 | cur_decay = 1 41 | for decay_step in decay_steps: 42 | if cur_epoch >= decay_step: 43 | cur_decay = cur_decay * optim_cfg.lr_decay 44 | return max(cur_decay, optim_cfg.lr_clip / optim_cfg.lr) 45 | 46 | lr_warmup_scheduler = None 47 | total_steps = total_iters_each_epoch * total_epochs 48 | if optim_cfg.optimizer == 'adam_onecycle': 49 | lr_scheduler = OneCycle( 50 | optimizer, total_steps, optim_cfg.lr, list(optim_cfg.moms), optim_cfg.div_factor, optim_cfg.pct_start 51 | ) 52 | else: 53 | lr_scheduler = lr_sched.LambdaLR(optimizer, lr_lbmd, last_epoch=last_epoch) 54 | 55 | if optim_cfg.lr_warmup: 56 | lr_warmup_scheduler = CosineWarmupLR( 57 | #optimizer, T_max=optim_cfg.warmup_epoch * len(total_iters_each_epoch), 58 | optimizer, T_max=optim_cfg.warmup_epoch * total_iters_each_epoch, 59 | eta_min=optim_cfg.lr / optim_cfg.div_factor 60 | ) 61 | 62 | return lr_scheduler, lr_warmup_scheduler 63 | -------------------------------------------------------------------------------- /lidardet/utils/loader.py: -------------------------------------------------------------------------------- 1 | import importlib 2 | import logging 3 | import os 4 | import sys 5 | from pathlib import Path 6 | 7 | logger = logging.getLogger("lidardet.utils.loader") 8 | 9 | CUSTOM_LOADED_MODULES = {} 10 | 11 | 12 | def _get_possible_module_path(paths): 13 | ret = [] 14 | for p in paths: 15 | p = Path(p) 16 | for path in p.glob("*"): 17 | if path.suffix in ["py", ".so"] or (path.is_dir()): 18 | if path.stem.isidentifier(): 19 | ret.append(path) 20 | return ret 21 | 22 | 23 | def _get_regular_import_name(path, module_paths): 24 | path = Path(path) 25 | for mp in module_paths: 26 | mp = Path(mp) 27 | if mp == path: 28 | return path.stem 29 | try: 30 | relative_path = path.relative_to(Path(mp)) 31 | parts = list((relative_path.parent / relative_path.stem).parts) 32 | module_name = ".".join([mp.stem] + parts) 33 | return module_name 34 | except Exception: 35 | pass 36 | return None 37 | 38 | 39 | def import_file(path, name: str = None, add_to_sys=True, disable_warning=False): 40 | global CUSTOM_LOADED_MODULES 41 | path = Path(path) 42 | module_name = path.stem 43 | try: 44 | user_paths = os.environ["PYTHONPATH"].split(os.pathsep) 45 | except KeyError: 46 | user_paths = [] 47 | possible_paths = _get_possible_module_path(user_paths) 48 | model_import_name = _get_regular_import_name(path, possible_paths) 49 | if model_import_name is not None: 50 | return import_name(model_import_name) 51 | if name is not None: 52 | module_name = name 53 | spec = importlib.util.spec_from_file_location(module_name, path) 54 | module = importlib.util.module_from_spec(spec) 55 | spec.loader.exec_module(module) 56 | if not disable_warning: 57 | logger.warning( 58 | ( 59 | f"Failed to perform regular import for file {path}. " 60 | "this means this file isn't in any folder in PYTHONPATH " 61 | "or don't have __init__.py in that project. " 62 | "directly file import may fail and some reflecting features are " 63 | "disabled even if import succeed. please add your project to PYTHONPATH " 64 | "or add __init__.py to ensure this file can be regularly imported. " 65 | ) 66 | ) 67 | 68 | if add_to_sys: # this will enable find objects defined in a file. 69 | # avoid replace system modules. 70 | if module_name in sys.modules and module_name not in CUSTOM_LOADED_MODULES: 71 | raise ValueError(f"{module_name} exists in system.") 72 | CUSTOM_LOADED_MODULES[module_name] = module 73 | sys.modules[module_name] = module 74 | return module 75 | 76 | 77 | def import_name(name, package=None): 78 | module = importlib.import_module(name, package) 79 | return module 80 | -------------------------------------------------------------------------------- /config/trajectory_prediction/transformer.yaml: -------------------------------------------------------------------------------- 1 | # train 2 | batch_size: &batch_size 64 3 | epochs: 120 4 | start_epoch: 0 5 | launcher: 'none' 6 | ckpt_save_interval: 2 7 | max_ckpt_save_num: 20 8 | 9 | # global settings 10 | point_cloud_range: &point_cloud_range [-32.0, -16.0, -2.5, 32.0, 32.0, 3.5] 11 | voxel_size: &voxel_size [0.16, 0.16, 4] 12 | num_points_per_trajectory: &num_points_per_trajectory 10 # ~ 20m length 13 | 14 | # dataset settings 15 | dataset_type: &dataset_type "HeatMapDataset" 16 | data_root: &data_root "datasets/KITTI_RAW/trajectory_prediction/" 17 | 18 | data: 19 | train: 20 | type: *dataset_type 21 | root_path: *data_root 22 | mode: 'train' 23 | subset: ['00', '02', '05', '07'] 24 | 25 | voxel_size: *voxel_size 26 | batch_size: *batch_size 27 | point_cloud_range: *point_cloud_range 28 | num_points_per_trajectory: *num_points_per_trajectory 29 | 30 | num_workers: 4 31 | road_width: 2.0 32 | sensor_height: 1.73 33 | lidar_intensity_max: 1.0 34 | lidar_bev_type: 'rgb_map' # 'height_map', 'rgb_map', 'rgb_traversability_map' 35 | 36 | random_drop: false 37 | random_shift: true 38 | max_random_shift: 0.25 39 | random_flip: true 40 | random_rotate: false 41 | use_lidar_points: false 42 | 43 | val: 44 | type: *dataset_type 45 | root_path: *data_root 46 | mode: 'val' 47 | subset: ['10'] 48 | 49 | voxel_size: *voxel_size 50 | batch_size: *batch_size 51 | point_cloud_range: *point_cloud_range 52 | num_points_per_trajectory: *num_points_per_trajectory 53 | 54 | num_workers: 4 55 | road_width: 2.0 56 | sensor_height: 1.73 57 | lidar_intensity_max: 1.0 58 | lidar_bev_type: 'rgb_map' # 'height_map', 'rgb_map', 'rgb_traversability_map' 59 | use_lidar_points: false 60 | 61 | # model settings 62 | model: 63 | type: ConvTrajPredictor 64 | task: 'trajectory_prediction' 65 | 66 | backbone: 67 | use_bn: true 68 | input_dim: 4 69 | num_blocks: [3, 6, 6, 3] 70 | 71 | header: 72 | use_bn: true 73 | use_road_seg: true 74 | use_transformer: true 75 | weight_loss_road: 1.0 76 | weight_loss_heatmap: 1.0 77 | weight_loss_waypoint: 1.0 78 | num_points_per_trajectory: *num_points_per_trajectory 79 | 80 | transformer: 81 | use_position_encoder: true 82 | use_transformer_encoder: true 83 | use_transformer_decoder: true 84 | num_points_per_trajectory: *num_points_per_trajectory 85 | feature_encoder_layers: [256, 128] 86 | position_encoder_layers: [256, 128] 87 | input_dim: 1900 # 7500 88 | d_model: 64 89 | dim_feedforward: 128 90 | num_layers: 6 91 | nhead: 8 92 | activation: "relu" 93 | dropout: 0.1 94 | loss_type: 'mse' # 'smooth_l1', 'mse' 95 | 96 | optimization: 97 | optimizer: adam_onecycle 98 | weight_decay: 0.01 99 | 100 | # scheduler 101 | lr: 0.003 102 | moms: [0.95, 0.85] 103 | div_factor: 10 104 | pct_start: 0.4 105 | 106 | grad_norm_clip: 10 107 | -------------------------------------------------------------------------------- /lidardet/utils/precision_recall.py: -------------------------------------------------------------------------------- 1 | """Universal procedure of calculating precision and recall.""" 2 | import bisect 3 | 4 | 5 | def match_gt_with_preds(ground_truth, predictions, match_labels): 6 | """Match a ground truth with every predictions and return matched index.""" 7 | max_confidence = 0. 8 | matched_idx = -1 9 | for i, pred in enumerate(predictions): 10 | if match_labels(ground_truth, pred[1]) and max_confidence < pred[0]: 11 | max_confidence = pred[0] 12 | matched_idx = i 13 | return matched_idx 14 | 15 | 16 | def get_confidence_list(ground_truths_list, predictions_list, match_labels): 17 | """Generate a list of confidence of true positives and false positives.""" 18 | assert len(ground_truths_list) == len(predictions_list) 19 | true_positive_list = [] 20 | false_positive_list = [] 21 | num_samples = len(ground_truths_list) 22 | for i in range(num_samples): 23 | ground_truths = ground_truths_list[i] 24 | predictions = predictions_list[i] 25 | prediction_matched = [False] * len(predictions) 26 | for ground_truth in ground_truths: 27 | idx = match_gt_with_preds(ground_truth, predictions, match_labels) 28 | if idx >= 0: 29 | prediction_matched[idx] = True 30 | true_positive_list.append(predictions[idx][0]) 31 | else: 32 | true_positive_list.append(.0) 33 | for idx, pred_matched in enumerate(prediction_matched): 34 | if not pred_matched: 35 | false_positive_list.append(predictions[idx][0]) 36 | return true_positive_list, false_positive_list 37 | 38 | 39 | def calc_precision_recall(ground_truths_list, predictions_list, match_labels): 40 | """Adjust threshold to get mutiple precision recall sample.""" 41 | true_positive_list, false_positive_list = get_confidence_list( 42 | ground_truths_list, predictions_list, match_labels) 43 | true_positive_list = sorted(true_positive_list) 44 | false_positive_list = sorted(false_positive_list) 45 | thresholds = sorted(list(set(true_positive_list))) 46 | recalls = [0.] 47 | precisions = [0.] 48 | for thresh in reversed(thresholds): 49 | if thresh == 0.: 50 | recalls.append(1.) 51 | precisions.append(0.) 52 | break 53 | false_negatives = bisect.bisect_left(true_positive_list, thresh) 54 | true_positives = len(true_positive_list) - false_negatives 55 | true_negatives = bisect.bisect_left(false_positive_list, thresh) 56 | false_positives = len(false_positive_list) - true_negatives 57 | recalls.append(true_positives / (true_positives+false_negatives)) 58 | precisions.append(true_positives / (true_positives + false_positives)) 59 | return precisions, recalls 60 | 61 | 62 | def calc_average_precision(precisions, recalls): 63 | """Calculate average precision defined in VOC contest.""" 64 | total_precision = 0. 65 | for i in range(11): 66 | index = next(conf[0] for conf in enumerate(recalls) if conf[1] >= i/10) 67 | total_precision += max(precisions[index:]) 68 | return total_precision / 11 69 | -------------------------------------------------------------------------------- /lidardet/datasets/base.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from pathlib import Path 3 | from collections import defaultdict 4 | from torch.utils.data import Dataset 5 | 6 | from .registry import DATASETS 7 | from .augmentor import DataAugmentor 8 | from .processor import DataProcessor 9 | 10 | @DATASETS.register 11 | class PointCloudDataset(Dataset): 12 | def __init__(self, cfg, logger=None): 13 | self.cfg = cfg 14 | self.logger = logger 15 | self.class_names = cfg.class_names 16 | self.root_path = Path(cfg.root_path) 17 | 18 | if self.cfg.get('augmentor', None): 19 | self.data_augmentor = DataAugmentor(self.root_path, cfg.augmentor, self.class_names, logger) 20 | 21 | if self.cfg.get('pre_processor', None): 22 | self.pre_processor = DataProcessor(cfg.pre_processor) 23 | 24 | def __len__(self): 25 | raise NotImplementedError 26 | 27 | def forward(self, index): 28 | raise NotImplementedError 29 | 30 | def augment_data(self, data_dict): 31 | if data_dict.get('gt_names', None) is not None: 32 | gt_boxes_mask = np.array([n in self.class_names for n in data_dict['gt_names']], dtype=np.bool_) 33 | 34 | data_dict = self.data_augmentor.forward( 35 | data_dict={ 36 | **data_dict, 37 | 'gt_boxes_mask': gt_boxes_mask 38 | } 39 | ) 40 | else: 41 | data_dict = self.data_augmentor.forward( 42 | data_dict={**data_dict}) 43 | 44 | if data_dict.get('gt_boxes', None) is not None: 45 | if len(data_dict['gt_boxes']) == 0: 46 | new_index = np.random.randint(self.__len__()) 47 | return self.__getitem__(new_index) 48 | 49 | return data_dict 50 | 51 | def pre_process(self, data_dict): 52 | data_dict = self.pre_processor.forward(data_dict) 53 | return data_dict 54 | 55 | @staticmethod 56 | def collate_batch(batch_list, _unused=False): 57 | data_dict = defaultdict(list) 58 | for cur_sample in batch_list: 59 | for key, val in cur_sample.items(): 60 | data_dict[key].append(val) 61 | batch_size = len(batch_list) 62 | ret = {} 63 | 64 | for key, val in data_dict.items(): 65 | if key in ['voxels', 'voxel_num_points']: 66 | ret[key] = np.concatenate(val, axis=0) 67 | elif key in ['points', 'voxel_coords']: 68 | coors = [] 69 | for i, coor in enumerate(val): 70 | coor_pad = np.pad(coor, ((0, 0), (1, 0)), mode='constant', constant_values=i) 71 | coors.append(coor_pad) 72 | ret[key] = np.concatenate(coors, axis=0) 73 | elif key in ['gt_boxes']: 74 | max_gt = max([len(x) for x in val]) 75 | batch_gt_boxes3d = np.zeros((batch_size, max_gt, val[0].shape[-1]), dtype=np.float32) 76 | for k in range(batch_size): 77 | batch_gt_boxes3d[k, :val[k].__len__(), :] = val[k] 78 | ret[key] = batch_gt_boxes3d 79 | else: 80 | ret[key] = np.stack(val, axis=0) 81 | 82 | ret['batch_size'] = batch_size 83 | return ret 84 | -------------------------------------------------------------------------------- /lidardet/utils/input/point_cloud/box_ops.h: -------------------------------------------------------------------------------- 1 | #ifndef BOX_OPS_H 2 | #define BOX_OPS_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | namespace py = pybind11; 9 | using namespace pybind11::literals; 10 | 11 | template 12 | inline py::array_t constant(ShapeContainer shape, DType value){ 13 | // create ROWMAJOR array. 14 | py::array_t array(shape); 15 | std::fill(array.mutable_data(), array.mutable_data() + array.size(), value); 16 | return array; 17 | } 18 | 19 | template 20 | inline py::array_t zeros(std::vector shape){ 21 | return constant>(shape, 0); 22 | } 23 | template 24 | py::array_t 25 | rbbox_iou(py::array_t box_corners, py::array_t qbox_corners, 26 | py::array_t standup_iou, DType standup_thresh) { 27 | namespace bg = boost::geometry; 28 | typedef bg::model::point point_t; 29 | typedef bg::model::polygon polygon_t; 30 | polygon_t poly, qpoly; 31 | std::vector poly_inter, poly_union; 32 | DType inter_area, union_area; 33 | auto box_corners_r = box_corners.template unchecked<3>(); 34 | auto qbox_corners_r = qbox_corners.template unchecked<3>(); 35 | auto standup_iou_r = standup_iou.template unchecked<2>(); 36 | auto N = box_corners_r.shape(0); 37 | auto K = qbox_corners_r.shape(0); 38 | py::array_t overlaps = zeros({N, K}); 39 | auto overlaps_rw = overlaps.template mutable_unchecked<2>(); 40 | if (N == 0 || K == 0) { 41 | return overlaps; 42 | } 43 | for (int k = 0; k < K; ++k) { 44 | for (int n = 0; n < N; ++n) { 45 | if (standup_iou_r(n, k) <= standup_thresh) 46 | continue; 47 | bg::append(poly, point_t(box_corners_r(n, 0, 0), box_corners_r(n, 0, 1))); 48 | bg::append(poly, point_t(box_corners_r(n, 1, 0), box_corners_r(n, 1, 1))); 49 | bg::append(poly, point_t(box_corners_r(n, 2, 0), box_corners_r(n, 2, 1))); 50 | bg::append(poly, point_t(box_corners_r(n, 3, 0), box_corners_r(n, 3, 1))); 51 | bg::append(poly, point_t(box_corners_r(n, 0, 0), box_corners_r(n, 0, 1))); 52 | bg::append(qpoly, 53 | point_t(qbox_corners_r(k, 0, 0), qbox_corners_r(k, 0, 1))); 54 | bg::append(qpoly, 55 | point_t(qbox_corners_r(k, 1, 0), qbox_corners_r(k, 1, 1))); 56 | bg::append(qpoly, 57 | point_t(qbox_corners_r(k, 2, 0), qbox_corners_r(k, 2, 1))); 58 | bg::append(qpoly, 59 | point_t(qbox_corners_r(k, 3, 0), qbox_corners_r(k, 3, 1))); 60 | bg::append(qpoly, 61 | point_t(qbox_corners_r(k, 0, 0), qbox_corners_r(k, 0, 1))); 62 | 63 | bg::intersection(poly, qpoly, poly_inter); 64 | 65 | if (!poly_inter.empty()) { 66 | inter_area = bg::area(poly_inter.front()); 67 | bg::union_(poly, qpoly, poly_union); 68 | if (!poly_union.empty()) { 69 | union_area = bg::area(poly_union.front()); 70 | overlaps_rw(n, k) = inter_area / union_area; 71 | } 72 | poly_union.clear(); 73 | } 74 | poly.clear(); 75 | qpoly.clear(); 76 | poly_inter.clear(); 77 | } 78 | } 79 | return overlaps; 80 | } 81 | 82 | 83 | #endif -------------------------------------------------------------------------------- /tools/train.py: -------------------------------------------------------------------------------- 1 | import pprint 2 | import torch 3 | from tensorboardX import SummaryWriter 4 | from lidardet.utils.config import get_config 5 | from lidardet.utils.common import get_logger, set_random_seed 6 | 7 | from lidardet.utils import dist 8 | from lidardet.models import model_fn_decorator 9 | from lidardet.models.builder import build_model 10 | from lidardet.datasets.builder import build_dataloader 11 | 12 | from train_utils.optimization import build_optimizer, build_scheduler 13 | from train_utils.train_utils import train_model 14 | 15 | def main(): 16 | 17 | cfg = get_config() 18 | logger = get_logger(cfg.log_dir, cfg.tag) 19 | logger.info(pprint.pformat(cfg)) 20 | 21 | if cfg.launcher == 'none': 22 | dist_train = False 23 | else: 24 | cfg.batch_size, cfg.local_rank = dist.init_dist_pytorch( 25 | cfg.batch_size, cfg.local_rank, backend='nccl' 26 | ) 27 | cfg.data.train.batch_size = cfg.batch_size 28 | dist_train = True 29 | 30 | set_random_seed(cfg['random_seed']) 31 | 32 | if dist_train: 33 | total_gpus = dist.get_world_size() 34 | logger.info('total_batch_size: %d' % (total_gpus * cfg.batch_size)) 35 | 36 | tb_log = SummaryWriter(log_dir=str(cfg.log_dir)) if cfg.local_rank == 0 else None 37 | 38 | train_set, train_loader, train_sampler = build_dataloader( 39 | cfg.data.train, dist=dist_train, training=True, logger=logger) 40 | 41 | #data = train_set[0] 42 | model = build_model(cfg.model) 43 | if cfg.get('sync_bn', False): 44 | model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) 45 | model.cuda() 46 | 47 | optimizer = build_optimizer(model, cfg.optimization) 48 | 49 | start_epoch = it = 0 50 | last_epoch = -1 51 | if cfg.ckpt is not None: 52 | #it, start_epoch = 53 | model.load_params_with_optimizer(cfg.ckpt, to_cpu=dist, optimizer=optimizer, logger=logger) 54 | #last_epoch = start_epoch + 1 55 | 56 | model.train() 57 | if dist_train: 58 | model = torch.nn.parallel.DistributedDataParallel( 59 | model, device_ids=[cfg.local_rank % torch.cuda.device_count()]) 60 | logger.info(model) 61 | 62 | total_iters_each_epoch = len(train_loader) 63 | lr_scheduler, lr_warmup_scheduler = build_scheduler( 64 | optimizer, total_iters_each_epoch=total_iters_each_epoch, total_epochs=cfg.epochs, 65 | last_epoch=last_epoch, optim_cfg=cfg.optimization 66 | ) 67 | 68 | # -----------------------start training--------------------------- 69 | logger.info('*'*20 + 'Start training' +'*'*20) 70 | 71 | train_model( 72 | model, 73 | optimizer, 74 | train_loader, 75 | model_func=model_fn_decorator(), 76 | lr_scheduler=lr_scheduler, 77 | optim_cfg=cfg.optimization, 78 | start_epoch=start_epoch, 79 | total_epochs=cfg.epochs, 80 | start_iter=it, 81 | rank=cfg.local_rank, 82 | tb_log=tb_log, 83 | ckpt_save_dir=cfg.model_dir, 84 | train_sampler=train_sampler, 85 | lr_warmup_scheduler=lr_warmup_scheduler, 86 | ckpt_save_interval=cfg.ckpt_save_interval, 87 | max_ckpt_save_num=cfg.max_ckpt_save_num, 88 | merge_all_iters_to_one_epoch=False 89 | ) 90 | logger.info('*'*20 + 'End training' +'*'*20) 91 | 92 | 93 | if __name__ == '__main__': 94 | main() 95 | -------------------------------------------------------------------------------- /lidardet/utils/buildtools/pybind11_build.py: -------------------------------------------------------------------------------- 1 | import shutil 2 | import subprocess 3 | import tempfile 4 | from pathlib import Path 5 | 6 | from ..find import find_cuda_device_arch 7 | from ..loader import import_file 8 | 9 | from .command import CUDALink, Gpp, Nvcc, compile_libraries, out 10 | 11 | 12 | class Pybind11Link(Gpp): 13 | def __init__( 14 | self, 15 | sources, 16 | target, 17 | std="c++11", 18 | includes: list = None, 19 | defines: dict = None, 20 | cflags: str = None, 21 | libraries: dict = None, 22 | lflags: str = None, 23 | extra_cflags: str = None, 24 | extra_lflags: str = None, 25 | build_directory: str = None, 26 | ): 27 | pb11_includes = ( 28 | subprocess.check_output("python3 -m pybind11 --includes", shell=True) 29 | .decode("utf8") 30 | .strip("\n") 31 | ) 32 | cflags = cflags or "-fPIC -O3 " 33 | cflags += pb11_includes 34 | super().__init__( 35 | sources, 36 | target, 37 | std, 38 | includes, 39 | defines, 40 | cflags, 41 | link=True, 42 | libraries=libraries, 43 | lflags=lflags, 44 | extra_cflags=extra_cflags, 45 | extra_lflags=extra_lflags, 46 | build_directory=build_directory, 47 | ) 48 | 49 | 50 | class Pybind11CUDALink(CUDALink): 51 | def __init__( 52 | self, 53 | sources, 54 | target, 55 | std="c++11", 56 | includes: list = None, 57 | defines: dict = None, 58 | cflags: str = None, 59 | libraries: dict = None, 60 | lflags: str = None, 61 | extra_cflags: str = None, 62 | extra_lflags: str = None, 63 | build_directory: str = None, 64 | ): 65 | pb11_includes = ( 66 | subprocess.check_output("python3 -m pybind11 --includes", shell=True) 67 | .decode("utf8") 68 | .strip("\n") 69 | ) 70 | cflags = cflags or "-fPIC -O3 " 71 | cflags += pb11_includes 72 | super().__init__( 73 | sources, 74 | target, 75 | std, 76 | includes, 77 | defines, 78 | cflags, 79 | libraries=libraries, 80 | lflags=lflags, 81 | extra_cflags=extra_cflags, 82 | extra_lflags=extra_lflags, 83 | build_directory=build_directory, 84 | ) 85 | 86 | 87 | def load_pb11( 88 | sources, 89 | target, 90 | cwd=".", 91 | cuda=False, 92 | arch=None, 93 | num_workers=4, 94 | includes: list = None, 95 | build_directory=None, 96 | compiler="g++", 97 | ): 98 | cmd_groups = [] 99 | cmds = [] 100 | outs = [] 101 | main_sources = [] 102 | if arch is None: 103 | arch = find_cuda_device_arch() 104 | 105 | for s in sources: 106 | s = str(s) 107 | if ".cu" in s or ".cu.cc" in s: 108 | assert cuda is True, "cuda must be true if contain cuda file" 109 | cmds.append(Nvcc(s, out(s), arch)) 110 | outs.append(out(s)) 111 | else: 112 | main_sources.append(s) 113 | 114 | if cuda is True and arch is None: 115 | raise ValueError("you must specify arch if sources contains" " cuda files") 116 | cmd_groups.append(cmds) 117 | if cuda: 118 | cmd_groups.append( 119 | [Pybind11CUDALink(outs + main_sources, target, includes=includes)] 120 | ) 121 | else: 122 | cmd_groups.append( 123 | [Pybind11Link(outs + main_sources, target, includes=includes)] 124 | ) 125 | for cmds in cmd_groups: 126 | compile_libraries(cmds, cwd, num_workers=num_workers, compiler=compiler) 127 | 128 | return import_file(target, add_to_sys=False, disable_warning=True) 129 | -------------------------------------------------------------------------------- /lidardet/models/trajectory_predictor/covernet.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from torch import nn 4 | from torch.nn import functional as F 5 | 6 | class CoverNet(nn.Module): 7 | """ Implementation of CoverNet https://arxiv.org/pdf/1911.10298.pdf """ 8 | 9 | def __init__(self, cfg): 10 | """ 11 | Inits Covernet. 12 | """ 13 | 14 | super().__init__() 15 | 16 | n_hidden_layers = cfg.get('n_hidden_layers', [4096]) 17 | num_modes = cfg.num_modes 18 | n_hidden_layers = [cfg.input_dim] + n_hidden_layers + [num_modes] 19 | 20 | linear_layers = [nn.Linear(in_dim, out_dim) 21 | for in_dim, out_dim in zip(n_hidden_layers[:-1], n_hidden_layers[1:])] 22 | 23 | self.head = nn.ModuleList(linear_layers) 24 | 25 | self.loss_function = nn.CrossEntropyLoss() 26 | 27 | def forward(self, batch_dict): 28 | 29 | logits = batch_dict['reg_features'] 30 | for linear in self.head: 31 | logits = linear(logits) 32 | 33 | 34 | if self.training: 35 | self.logits = logits 36 | self.label = batch_dict['traj_label'] 37 | else: 38 | logits = F.softmax(logits, dim=-1) 39 | 40 | batch_dict['logits'] = logits 41 | return batch_dict 42 | 43 | def get_loss(self): 44 | loss = self.loss_function(self.logits, self.label.long()) 45 | return loss 46 | 47 | def mean_pointwise_l2_distance(lattice, ground_truth): 48 | """ 49 | Computes the index of the closest trajectory in the lattice as measured by l1 distance. 50 | :param lattice: Lattice of pre-generated trajectories. Shape [num_modes, n_timesteps, state_dim] 51 | :param ground_truth: Ground truth trajectory of agent. Shape [1, n_timesteps, state_dim]. 52 | :return: Index of closest mode in the lattice. 53 | """ 54 | stacked_ground_truth = ground_truth.repeat(lattice.shape[0], 1, 1) 55 | return torch.pow(lattice - stacked_ground_truth, 2).sum(dim=2).sqrt().mean(dim=1).argmin() 56 | 57 | 58 | class ConstantLatticeLoss: 59 | """ 60 | Computes the loss for a constant lattice CoverNet model. 61 | """ 62 | 63 | def __init__(self, lattice, similarity_function = mean_pointwise_l2_distance): 64 | """ 65 | Inits the loss. 66 | :param lattice: numpy array of shape [n_modes, n_timesteps, state_dim] 67 | :param similarity_function: Function that computes the index of the closest trajectory in the lattice 68 | to the actual ground truth trajectory of the agent. 69 | """ 70 | 71 | self.lattice = torch.Tensor(lattice) 72 | self.similarity_func = similarity_function 73 | 74 | def __call__(self, batch_logits, batch_ground_truth_trajectory): 75 | """ 76 | Computes the loss on a batch. 77 | :param batch_logits: Tensor of shape [batch_size, n_modes]. Output of a linear layer since this class 78 | uses nn.functional.cross_entropy. 79 | :param batch_ground_truth_trajectory: Tensor of shape [batch_size, n_timesteps, state_dim] 80 | :return: Average element-wise loss on the batch. 81 | """ 82 | 83 | # If using GPU, need to copy the lattice to the GPU if haven't done so already 84 | # This ensures we only copy it once 85 | if self.lattice.device != batch_logits.device: 86 | self.lattice = self.lattice.to(batch_logits.device) 87 | 88 | batch_losses = torch.Tensor().requires_grad_(True).to(batch_logits.device) 89 | 90 | for logit, ground_truth in zip(batch_logits, batch_ground_truth_trajectory): 91 | 92 | closest_lattice_trajectory = self.similarity_func(self.lattice, ground_truth.unsqueeze(0)) 93 | label = torch.LongTensor([closest_lattice_trajectory]).to(batch_logits.device) 94 | classification_loss = F.cross_entropy(logit, label) 95 | 96 | batch_losses = torch.cat((batch_losses, classification_loss.unsqueeze(0)), 0) 97 | 98 | return batch_losses.mean() 99 | -------------------------------------------------------------------------------- /lidardet/utils/config.py: -------------------------------------------------------------------------------- 1 | import json 2 | import yaml 3 | import argparse 4 | from pathlib import Path 5 | #from easydict import EasyDict 6 | from permissive_dict import PermissiveDict as Dict 7 | 8 | def get_config_from_json(json_file): 9 | """ 10 | Get the config from a json file 11 | Input: 12 | - json_file: json configuration file 13 | Return: 14 | - config: namespace 15 | - config_dict: dictionary 16 | """ 17 | # parse the configurations from the config json file provided 18 | with open(json_file, 'r') as config_file: 19 | config_dict = json.load(config_file) 20 | 21 | # convert the dictionary to a namespace using bunch lib 22 | config = Dict(config_dict) 23 | 24 | return config, config_dict 25 | 26 | def get_config_from_yaml(yaml_file): 27 | """ 28 | Get the config from yaml file 29 | Input: 30 | - yaml_file: yaml configuration file 31 | Return: 32 | - config: namespace 33 | - config_dict: dictionary 34 | """ 35 | 36 | with open(yaml_file) as fp: 37 | if hasattr(yaml, 'FullLoader'): 38 | config_dict = yaml.load(fp, Loader=yaml.FullLoader) 39 | else: 40 | config_dict = yaml.load(fp) 41 | 42 | # convert the dictionary to a namespace using bunch lib 43 | config = Dict(config_dict) 44 | return config, config_dict 45 | 46 | def get_args(): 47 | argparser = argparse.ArgumentParser(description=__doc__) 48 | argparser.add_argument( '-c', '--cfg', metavar='C', required=True, help='The Configuration file') 49 | argparser.add_argument( '-s', '--seed', default=100, type=int, help='The random seed') 50 | argparser.add_argument( '-m', '--ckpt', type=str, help='The model path') 51 | argparser.add_argument( '--local_rank', type=int, default=0, help='local rank for distributed training') 52 | argparser.add_argument( '--launcher', type=str, default='none', help='launcher for distributed training') 53 | argparser.add_argument('--eval_all', action='store_true', default=False, help='whether to evaluate all checkpoints') 54 | args = argparser.parse_args() 55 | return args 56 | 57 | def merge_new_config(config, new_config): 58 | if '_base_' in new_config: 59 | with open(new_config['_base_'], 'r') as f: 60 | if hasattr(yaml, 'FullLoader'): 61 | yaml_config = yaml.load(f, Loader=yaml.FullLoader) 62 | else: 63 | yaml_config = yaml.load(f) 64 | print(yaml_config) 65 | config.update(Dict(yaml_config)) 66 | 67 | for key, val in new_config.items(): 68 | if not isinstance(val, dict): 69 | config[key] = val 70 | continue 71 | if key not in config: 72 | config[key] = Dict() 73 | merge_new_config(config[key], val) 74 | 75 | return config 76 | 77 | def cfg_from_file(config_file): 78 | if config_file.endswith('json'): 79 | new_config, _ = get_config_from_json(config_file) 80 | elif config_file.endswith('yaml'): 81 | new_config, _ = get_config_from_yaml(config_file) 82 | else: 83 | raise Exception("Only .json and .yaml are supported!") 84 | 85 | config = Dict() 86 | merge_new_config(config=config, new_config=new_config) 87 | return config 88 | 89 | def get_config(): 90 | args = get_args() 91 | config_file = args.cfg 92 | random_seed = args.seed 93 | 94 | config = cfg_from_file(config_file) 95 | 96 | config.eval_all = args.eval_all 97 | config.local_rank = args.local_rank 98 | config.ckpt = args.ckpt 99 | config.launcher = args.launcher 100 | 101 | config.random_seed = random_seed 102 | config.tag = Path(config_file).stem 103 | config.cache_dir = Path('cache') / config.tag / str(config.random_seed) 104 | config.model_dir = config.cache_dir / 'models' 105 | config.log_dir = config.cache_dir / 'logs' 106 | config.output_dir = config.cache_dir / 'output' 107 | 108 | # create the experiments dirs 109 | config.cache_dir.resolve().mkdir(parents=True, exist_ok=True) 110 | config.model_dir.resolve().mkdir(exist_ok=True) 111 | config.log_dir.resolve().mkdir(exist_ok=True) 112 | config.output_dir.resolve().mkdir(exist_ok=True) 113 | 114 | cfg = Dict() 115 | merge_new_config(config=cfg, new_config=config) 116 | return cfg 117 | -------------------------------------------------------------------------------- /lidardet/utils/range_image.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | class RangeImage(object): 4 | 5 | def __init__(self, cfg = {}): 6 | self.cfg = cfg 7 | self.max_range = cfg.get('max_range', 80.0) 8 | self.min_range = cfg.get('min_range', 2.0) 9 | self.fov_up = cfg.get('fov_up', 3.0) / 180.0 * np.pi 10 | self.fov_down = cfg.get('fov_down', -25.0) / 180.0 * np.pi 11 | 12 | # get vertical field of view total in radians 13 | self.fov = abs(self.fov_down) + abs(self.fov_up) 14 | 15 | self.cols = cfg.get('cols', 1024) # 1800 16 | self.rows = cfg.get('rows', 64) 17 | 18 | def from_points(self, points_in, normalize=True): 19 | """ 20 | Project points to range image 21 | """ 22 | points = np.copy(points_in) 23 | # get depth of all points 24 | depth = np.linalg.norm(points[:, :3], 2, axis=1) 25 | # filter by range limit 26 | points = points[(depth > self.min_range) & (depth < self.max_range)] 27 | depth = depth[(depth > self.min_range) & (depth < self.max_range)] 28 | 29 | # extract x, y, z and intensity values 30 | x = points[:,0] 31 | y = points[:,1] 32 | z = points[:,2] 33 | t = points[:,3] 34 | 35 | # get horizontal and vertical angles [radian] 36 | yaw = -np.arctan2(y, x) # [-pi, pi] 37 | pitch = np.arcsin(z / depth) 38 | 39 | # get projections in image coords 40 | proj_x = 0.5 * (yaw / np.pi + 1.0) # in [0.0, 1.0] 41 | proj_y = 1.0 - (pitch + abs(self.fov_down)) / self.fov # in [0.0, 1.0] 42 | 43 | # scale to image size using angular resolution 44 | proj_x *= self.cols # in [0.0, cols] 45 | proj_y *= self.rows # in [0.0, rows] 46 | 47 | # round and clamp for use as index 48 | proj_x = np.floor(proj_x) 49 | proj_x = np.minimum(self.cols - 1, proj_x) 50 | proj_x = np.maximum(0, proj_x).astype(np.int32) # in [0,W-1] 51 | 52 | proj_y = np.floor(proj_y) 53 | proj_y = np.minimum(self.rows - 1, proj_y) 54 | proj_y = np.maximum(0, proj_y).astype(np.int32) # in [0,H-1] 55 | 56 | # sort depth in ascending order to keep more far distance points 57 | indices = np.arange(depth.shape[0]) 58 | order = np.argsort(depth) 59 | depth = depth[order] 60 | indices = indices[order] 61 | points = points[order] 62 | proj_y = proj_y[order] 63 | proj_x = proj_x[order] 64 | 65 | proj_range = np.zeros((self.rows, self.cols, 5), dtype=np.float32) # [H,W] range (0 is no data) 66 | 67 | if normalize: 68 | depth /= self.max_range 69 | 70 | proj_range[proj_y, proj_x, 0] = depth 71 | proj_range[proj_y, proj_x, 1] = points[:,3] 72 | proj_range[proj_y, proj_x, 2] = points[:,0] 73 | proj_range[proj_y, proj_x, 3] = points[:,1] 74 | proj_range[proj_y, proj_x, 4] = points[:,2] 75 | return proj_range 76 | 77 | def to_points(self, img_in, denormalize=True): 78 | img = np.copy(img_in) 79 | proj_y = np.float32(np.arange(self.rows)) / self.rows 80 | proj_x = np.float32(np.arange(self.cols)) / self.cols 81 | 82 | v_angles = (1.0 - proj_y) * self.fov - abs(self.fov_down) 83 | h_angles = (2 * proj_x - 1.0) * np.pi 84 | 85 | points = [] 86 | coordinates = [] 87 | 88 | if denormalize: 89 | img[:,:,0] *= self.max_range 90 | 91 | for i in range(self.rows): 92 | for j in range(self.cols): 93 | depth = img[i, j, 0] 94 | intensity = 0 95 | if img.shape[2] >= 2: 96 | intensity = img[i, j, 1] 97 | 98 | if depth < self.min_range: 99 | continue 100 | 101 | h_angle = h_angles[j] 102 | v_angle = v_angles[i] 103 | x = np.sin(h_angle) * np.cos(v_angle) * depth 104 | y = np.cos(h_angle) * np.cos(v_angle) * depth 105 | z = np.sin(v_angle) * depth 106 | point = np.array([x, y, z, intensity]).astype(np.float32) 107 | points.append(point) 108 | coordinates.append(np.array([i, j]).astype(np.int32)) 109 | 110 | return np.array(points), np.array(coordinates) 111 | -------------------------------------------------------------------------------- /lidardet/models/trajectory_predictor/header.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | 4 | from .mtp import MTP 5 | from .covernet import CoverNet 6 | from .backbone import conv3x3 7 | 8 | class Header(nn.Module): 9 | 10 | def __init__(self, cfg): 11 | super(Header, self).__init__() 12 | self.cfg = cfg 13 | self.use_bn = cfg.use_bn 14 | self.with_attention = cfg.with_attention 15 | bias = not self.use_bn 16 | self.conv1 = conv3x3(96, 96, bias=bias) 17 | self.bn1 = nn.BatchNorm2d(96) 18 | self.conv2 = conv3x3(96, 96, bias=bias) 19 | self.bn2 = nn.BatchNorm2d(96) 20 | self.conv3 = conv3x3(96, 96, bias=bias) 21 | self.bn3 = nn.BatchNorm2d(96) 22 | self.conv4 = conv3x3(96, 96, bias=bias) 23 | self.bn4 = nn.BatchNorm2d(96) 24 | 25 | dim_conv5 = 384 26 | if self.with_attention: 27 | self.conv7 = nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1, bias=bias) 28 | self.bn7 = nn.BatchNorm2d(16) 29 | self.conv8 = nn.Conv2d(16, 32, kernel_size=3, stride=2, padding=1, bias=bias) 30 | self.bn8 = nn.BatchNorm2d(32) 31 | #self.pool = nn.MaxPool2d(kernel_size=4, padding=(1,0)) 32 | dim_conv5 += 32 33 | 34 | self.conv5 = nn.Conv2d(dim_conv5, 512, kernel_size=3, stride=2, padding=0, bias=bias) 35 | self.bn5 = nn.BatchNorm2d(512) 36 | self.conv6 = nn.Conv2d(512, 1024, kernel_size=3, stride=2, padding=0, bias=bias) 37 | self.bn6 = nn.BatchNorm2d(1024) 38 | self.seg_head = conv3x3(96, 1, bias=True) 39 | 40 | if self.cfg.get('mtp'): 41 | self.pred_head = MTP(cfg.mtp) 42 | elif self.cfg.get('covernet'): 43 | self.pred_head = CoverNet(cfg.covernet) 44 | 45 | self.seg_loss_func = nn.BCELoss() 46 | 47 | def forward(self, batch_dict): 48 | x = batch_dict['seg_features'] 49 | 50 | x = self.conv1(x) 51 | if self.use_bn: 52 | x = self.bn1(x) 53 | x = self.conv2(x) 54 | if self.use_bn: 55 | x = self.bn2(x) 56 | x = self.conv3(x) 57 | if self.use_bn: 58 | x = self.bn3(x) 59 | x = self.conv4(x) 60 | if self.use_bn: 61 | x = self.bn4(x) 62 | 63 | seg = torch.sigmoid(self.seg_head(x)) # [b, 1, 75, 100] 64 | 65 | y = batch_dict['reg_features'] # [b, 384, 19, 25] 66 | 67 | if self.with_attention: 68 | z = self.conv7(seg) 69 | z = self.bn7(z) 70 | z = self.conv8(z) 71 | z = self.bn8(z) 72 | y = torch.cat([y, z], dim=1) 73 | 74 | y = self.conv5(y) 75 | if self.use_bn: 76 | y = self.bn5(y) 77 | y = self.conv6(y) 78 | if self.use_bn: 79 | y = self.bn6(y) 80 | y = y.mean([2, 3]) 81 | batch_dict['reg_features'] = y 82 | 83 | batch_dict = self.pred_head(batch_dict) 84 | 85 | if self.training: 86 | self.seg_pred = seg.squeeze(1) 87 | self.seg_target = batch_dict['img_ins'] 88 | 89 | batch_dict['pred_seg'] = seg.squeeze(1) 90 | return batch_dict 91 | 92 | def get_loss(self): 93 | loss_pred = self.pred_head.get_loss() 94 | loss_seg = self.seg_loss_func(self.seg_pred, self.seg_target) 95 | 96 | loss = self.cfg.weight_loss_seg * loss_seg + loss_pred 97 | tb_dict = {'loss_pred': loss_pred, 'loss_seg': loss_seg} 98 | return loss, tb_dict 99 | 100 | def get_prediction(self, batch_dict): 101 | pred_traj_list = [] 102 | if self.cfg.get('mtp'): 103 | pred_traj_batch = batch_dict['trajectory_predictions'] 104 | mode_prob_batch = batch_dict['mode_probabilities'] 105 | for i, mode_prob in enumerate(mode_prob_batch): 106 | order = torch.argsort(mode_prob, dim=0, descending=True) 107 | traj_sorted = pred_traj_batch[i][order].cpu().numpy() 108 | pred_traj_list.append(traj_sorted) 109 | elif self.cfg.get('covernet'): 110 | logits = batch_dict['logits'] 111 | for i, mode_prob in enumerate(logits): 112 | order = torch.argsort(mode_prob, dim=0, descending=True) 113 | pred_traj_list.append(order.cpu().numpy()) 114 | 115 | return pred_traj_list 116 | -------------------------------------------------------------------------------- /lidardet/utils/box_coder_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | class ResidualCoder(object): 5 | def __init__(self, code_size=7, **kwargs): 6 | super().__init__() 7 | self.code_size = code_size 8 | 9 | @staticmethod 10 | def encode_torch(boxes, anchors): 11 | """ 12 | Args: 13 | boxes: (N, 7 + C) [x, y, z, dx, dy, dz, heading, ...] 14 | anchors: (N, 7 + C) [x, y, z, dx, dy, dz, heading, ...] 15 | 16 | Returns: 17 | 18 | """ 19 | anchors[:, 3:6] = torch.clamp_min(anchors[:, 3:6], min=1e-5) 20 | boxes[:, 3:6] = torch.clamp_min(boxes[:, 3:6], min=1e-5) 21 | 22 | xa, ya, za, dxa, dya, dza, ra, *cas = torch.split(anchors, 1, dim=-1) 23 | xg, yg, zg, dxg, dyg, dzg, rg, *cgs = torch.split(boxes, 1, dim=-1) 24 | 25 | diagonal = torch.sqrt(dxa ** 2 + dya ** 2) 26 | xt = (xg - xa) / diagonal 27 | yt = (yg - ya) / diagonal 28 | zt = (zg - za) / dza 29 | dxt = torch.log(dxg / dxa) 30 | dyt = torch.log(dyg / dya) 31 | dzt = torch.log(dzg / dza) 32 | rt = rg - ra 33 | 34 | cts = [g - a for g, a in zip(cgs, cas)] 35 | return torch.cat([xt, yt, zt, dxt, dyt, dzt, rt, *cts], dim=-1) 36 | 37 | @staticmethod 38 | def decode_torch(box_encodings, anchors): 39 | """ 40 | Args: 41 | box_encodings: (B, N, 7 + C) or (N, 7 + C) [x, y, z, dx, dy, dz, heading, ...] 42 | anchors: (B, N, 7 + C) or (N, 7 + C) [x, y, z, dx, dy, dz, heading, ...] 43 | 44 | Returns: 45 | 46 | """ 47 | xa, ya, za, dxa, dya, dza, ra, *cas = torch.split(anchors, 1, dim=-1) 48 | xt, yt, zt, dxt, dyt, dzt, rt, *cts = torch.split(box_encodings, 1, dim=-1) 49 | 50 | diagonal = torch.sqrt(dxa ** 2 + dya ** 2) 51 | xg = xt * diagonal + xa 52 | yg = yt * diagonal + ya 53 | zg = zt * dza + za 54 | 55 | dxg = torch.exp(dxt) * dxa 56 | dyg = torch.exp(dyt) * dya 57 | dzg = torch.exp(dzt) * dza 58 | rg = rt + ra 59 | 60 | cgs = [t + a for t, a in zip(cts, cas)] 61 | return torch.cat([xg, yg, zg, dxg, dyg, dzg, rg, *cgs], dim=-1) 62 | 63 | 64 | class PreviousResidualDecoder(object): 65 | def __init__(self, code_size=7, **kwargs): 66 | super().__init__() 67 | self.code_size = code_size 68 | 69 | @staticmethod 70 | def decode_torch(box_encodings, anchors): 71 | """ 72 | Args: 73 | box_encodings: (B, N, 7 + ?) x, y, z, w, l, h, r, custom values 74 | anchors: (B, N, 7 + C) or (N, 7 + C) [x, y, z, dx, dy, dz, heading, ...] 75 | 76 | Returns: 77 | 78 | """ 79 | xa, ya, za, dxa, dya, dza, ra, *cas = torch.split(anchors, 1, dim=-1) 80 | xt, yt, zt, wt, lt, ht, rt, *cts = torch.split(box_encodings, 1, dim=-1) 81 | 82 | diagonal = torch.sqrt(dxa ** 2 + dya ** 2) 83 | xg = xt * diagonal + xa 84 | yg = yt * diagonal + ya 85 | zg = zt * dza + za 86 | 87 | dxg = torch.exp(lt) * dxa 88 | dyg = torch.exp(wt) * dya 89 | dzg = torch.exp(ht) * dza 90 | rg = rt + ra 91 | 92 | cgs = [t + a for t, a in zip(cts, cas)] 93 | return torch.cat([xg, yg, zg, dxg, dyg, dzg, rg, *cgs], dim=-1) 94 | 95 | 96 | class PreviousResidualRoIDecoder(object): 97 | def __init__(self, code_size=7, **kwargs): 98 | super().__init__() 99 | self.code_size = code_size 100 | 101 | @staticmethod 102 | def decode_torch(box_encodings, anchors): 103 | """ 104 | Args: 105 | box_encodings: (B, N, 7 + ?) x, y, z, w, l, h, r, custom values 106 | anchors: (B, N, 7 + C) or (N, 7 + C) [x, y, z, dx, dy, dz, heading, ...] 107 | 108 | Returns: 109 | 110 | """ 111 | xa, ya, za, dxa, dya, dza, ra, *cas = torch.split(anchors, 1, dim=-1) 112 | xt, yt, zt, wt, lt, ht, rt, *cts = torch.split(box_encodings, 1, dim=-1) 113 | 114 | diagonal = torch.sqrt(dxa ** 2 + dya ** 2) 115 | xg = xt * diagonal + xa 116 | yg = yt * diagonal + ya 117 | zg = zt * dza + za 118 | 119 | dxg = torch.exp(lt) * dxa 120 | dyg = torch.exp(wt) * dya 121 | dzg = torch.exp(ht) * dza 122 | rg = ra - rt 123 | 124 | cgs = [t + a for t, a in zip(cts, cas)] 125 | return torch.cat([xg, yg, zg, dxg, dyg, dzg, rg, *cgs], dim=-1) 126 | -------------------------------------------------------------------------------- /lidardet/utils/input/point_cloud/bev_ops.py: -------------------------------------------------------------------------------- 1 | import math 2 | 3 | import numba 4 | import numpy as np 5 | 6 | 7 | @numba.jit(nopython=True) 8 | def _points_to_bevmap_reverse_kernel( 9 | points, 10 | voxel_size, 11 | coors_range, 12 | coor_to_voxelidx, 13 | # coors_2d, 14 | bev_map, 15 | height_lowers, 16 | # density_norm_num=16, 17 | with_reflectivity=False, 18 | max_voxels=40000, 19 | ): 20 | # put all computations to one loop. 21 | # we shouldn't create large array in main jit code, otherwise 22 | # reduce performance 23 | N = points.shape[0] 24 | ndim = points.shape[1] - 1 25 | # ndim = 3 26 | ndim_minus_1 = ndim - 1 27 | grid_size = (coors_range[3:] - coors_range[:3]) / voxel_size 28 | # np.round(grid_size) 29 | # grid_size = np.round(grid_size).astype(np.int64)(np.int32) 30 | grid_size = np.round(grid_size, 0, grid_size).astype(np.int32) 31 | height_slice_size = voxel_size[-1] 32 | coor = np.zeros(shape=(3,), dtype=np.int32) # DHW 33 | voxel_num = 0 34 | failed = False 35 | for i in range(N): 36 | failed = False 37 | for j in range(ndim): 38 | c = np.floor((points[i, j] - coors_range[j]) / voxel_size[j]) 39 | if c < 0 or c >= grid_size[j]: 40 | failed = True 41 | break 42 | coor[ndim_minus_1 - j] = c 43 | if failed: 44 | continue 45 | voxelidx = coor_to_voxelidx[coor[0], coor[1], coor[2]] 46 | if voxelidx == -1: 47 | voxelidx = voxel_num 48 | if voxel_num >= max_voxels: 49 | break 50 | voxel_num += 1 51 | coor_to_voxelidx[coor[0], coor[1], coor[2]] = voxelidx 52 | # coors_2d[voxelidx] = coor[1:] 53 | bev_map[-1, coor[1], coor[2]] += 1 54 | height_norm = bev_map[coor[0], coor[1], coor[2]] 55 | incomimg_height_norm = ( 56 | points[i, 2] - height_lowers[coor[0]] 57 | ) / height_slice_size 58 | if incomimg_height_norm > height_norm: 59 | bev_map[coor[0], coor[1], coor[2]] = incomimg_height_norm 60 | if with_reflectivity: 61 | bev_map[-2, coor[1], coor[2]] = points[i, 3] 62 | # return voxel_num 63 | 64 | 65 | def points_to_bev( 66 | points, 67 | voxel_size, 68 | coors_range, 69 | with_reflectivity=False, 70 | density_norm_num=16, 71 | max_voxels=40000, 72 | ): 73 | """convert kitti points(N, 4) to a bev map. return [C, H, W] map. 74 | this function based on algorithm in points_to_voxel. 75 | takes 5ms in a reduced pointcloud with voxel_size=[0.1, 0.1, 0.8] 76 | 77 | Args: 78 | points: [N, ndim] float tensor. points[:, :3] contain xyz points and 79 | points[:, 3] contain reflectivity. 80 | voxel_size: [3] list/tuple or array, float. xyz, indicate voxel size 81 | coors_range: [6] list/tuple or array, float. indicate voxel range. 82 | format: xyzxyz, minmax 83 | with_reflectivity: bool. if True, will add a intensity map to bev map. 84 | Returns: 85 | bev_map: [num_height_maps + 1(2), H, W] float tensor. 86 | `WARNING`: bev_map[-1] is num_points map, NOT density map, 87 | because calculate density map need more time in cpu rather than gpu. 88 | if with_reflectivity is True, bev_map[-2] is intensity map. 89 | """ 90 | if not isinstance(voxel_size, np.ndarray): 91 | voxel_size = np.array(voxel_size, dtype=points.dtype) 92 | if not isinstance(coors_range, np.ndarray): 93 | coors_range = np.array(coors_range, dtype=points.dtype) 94 | voxelmap_shape = (coors_range[3:] - coors_range[:3]) / voxel_size 95 | voxelmap_shape = tuple(np.round(voxelmap_shape).astype(np.int32).tolist()) 96 | voxelmap_shape = voxelmap_shape[::-1] # DHW format 97 | coor_to_voxelidx = -np.ones(shape=voxelmap_shape, dtype=np.int32) 98 | # coors_2d = np.zeros(shape=(max_voxels, 2), dtype=np.int32) 99 | bev_map_shape = list(voxelmap_shape) 100 | bev_map_shape[0] += 1 101 | height_lowers = np.linspace( 102 | coors_range[2], coors_range[5], voxelmap_shape[0], endpoint=False 103 | ) 104 | if with_reflectivity: 105 | bev_map_shape[0] += 1 106 | bev_map = np.zeros(shape=bev_map_shape, dtype=points.dtype) 107 | _points_to_bevmap_reverse_kernel( 108 | points, 109 | voxel_size, 110 | coors_range, 111 | coor_to_voxelidx, 112 | bev_map, 113 | height_lowers, 114 | with_reflectivity, 115 | max_voxels, 116 | ) 117 | return bev_map 118 | -------------------------------------------------------------------------------- /tools/train_utils/optimization/learning_schedules_fastai.py: -------------------------------------------------------------------------------- 1 | # This file is modified from https://github.com/traveller59/second.pytorch 2 | 3 | import numpy as np 4 | import math 5 | from functools import partial 6 | import torch.optim.lr_scheduler as lr_sched 7 | from .fastai_optim import OptimWrapper 8 | 9 | 10 | class LRSchedulerStep(object): 11 | def __init__(self, fai_optimizer: OptimWrapper, total_step, lr_phases, 12 | mom_phases): 13 | # if not isinstance(fai_optimizer, OptimWrapper): 14 | # raise TypeError('{} is not a fastai OptimWrapper'.format( 15 | # type(fai_optimizer).__name__)) 16 | self.optimizer = fai_optimizer 17 | self.total_step = total_step 18 | self.lr_phases = [] 19 | 20 | for i, (start, lambda_func) in enumerate(lr_phases): 21 | if len(self.lr_phases) != 0: 22 | assert self.lr_phases[-1][0] < start 23 | if isinstance(lambda_func, str): 24 | lambda_func = eval(lambda_func) 25 | if i < len(lr_phases) - 1: 26 | self.lr_phases.append((int(start * total_step), int(lr_phases[i + 1][0] * total_step), lambda_func)) 27 | else: 28 | self.lr_phases.append((int(start * total_step), total_step, lambda_func)) 29 | assert self.lr_phases[0][0] == 0 30 | self.mom_phases = [] 31 | for i, (start, lambda_func) in enumerate(mom_phases): 32 | if len(self.mom_phases) != 0: 33 | assert self.mom_phases[-1][0] < start 34 | if isinstance(lambda_func, str): 35 | lambda_func = eval(lambda_func) 36 | if i < len(mom_phases) - 1: 37 | self.mom_phases.append((int(start * total_step), int(mom_phases[i + 1][0] * total_step), lambda_func)) 38 | else: 39 | self.mom_phases.append((int(start * total_step), total_step, lambda_func)) 40 | assert self.mom_phases[0][0] == 0 41 | 42 | def step(self, step): 43 | for start, end, func in self.lr_phases: 44 | if step >= start: 45 | self.optimizer.lr = func((step - start) / (end - start)) 46 | for start, end, func in self.mom_phases: 47 | if step >= start: 48 | self.optimizer.mom = func((step - start) / (end - start)) 49 | 50 | 51 | def annealing_cos(start, end, pct): 52 | # print(pct, start, end) 53 | "Cosine anneal from `start` to `end` as pct goes from 0.0 to 1.0." 54 | cos_out = np.cos(np.pi * pct) + 1 55 | return end + (start - end) / 2 * cos_out 56 | 57 | 58 | class OneCycle(LRSchedulerStep): 59 | def __init__(self, fai_optimizer, total_step, lr_max, moms, div_factor, 60 | pct_start): 61 | self.lr_max = lr_max 62 | self.moms = moms 63 | self.div_factor = div_factor 64 | self.pct_start = pct_start 65 | a1 = int(total_step * self.pct_start) 66 | a2 = total_step - a1 67 | low_lr = self.lr_max / self.div_factor 68 | lr_phases = ((0, partial(annealing_cos, low_lr, self.lr_max)), 69 | (self.pct_start, 70 | partial(annealing_cos, self.lr_max, low_lr / 1e4))) 71 | mom_phases = ((0, partial(annealing_cos, *self.moms)), 72 | (self.pct_start, partial(annealing_cos, 73 | *self.moms[::-1]))) 74 | fai_optimizer.lr, fai_optimizer.mom = low_lr, self.moms[0] 75 | super().__init__(fai_optimizer, total_step, lr_phases, mom_phases) 76 | 77 | 78 | class CosineWarmupLR(lr_sched._LRScheduler): 79 | def __init__(self, optimizer, T_max, eta_min=0, last_epoch=-1): 80 | self.T_max = T_max 81 | self.eta_min = eta_min 82 | super(CosineWarmupLR, self).__init__(optimizer, last_epoch) 83 | 84 | def get_lr(self): 85 | return [self.eta_min + (base_lr - self.eta_min) * 86 | (1 - math.cos(math.pi * self.last_epoch / self.T_max)) / 2 87 | for base_lr in self.base_lrs] 88 | 89 | 90 | class FakeOptim: 91 | def __init__(self): 92 | self.lr = 0 93 | self.mom = 0 94 | 95 | 96 | if __name__ == "__main__": 97 | import matplotlib.pyplot as plt 98 | 99 | opt = FakeOptim() # 3e-3, wd=0.4, div_factor=10 100 | schd = OneCycle(opt, 100, 3e-3, (0.95, 0.85), 10.0, 0.1) 101 | 102 | lrs = [] 103 | moms = [] 104 | for i in range(100): 105 | schd.step(i) 106 | lrs.append(opt.lr) 107 | moms.append(opt.mom) 108 | plt.plot(lrs) 109 | # plt.plot(moms) 110 | plt.show() 111 | plt.plot(moms) 112 | plt.show() 113 | -------------------------------------------------------------------------------- /lidardet/models/trajectory_predictor/predictor_base.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import os 3 | import torch.nn as nn 4 | 5 | from ..registry import TRAJECTORY_PREDICTOR 6 | 7 | @TRAJECTORY_PREDICTOR.register 8 | class PredictorBase(nn.Module): 9 | def __init__(self, cfg): 10 | super().__init__() 11 | self.model_cfg = cfg 12 | self.register_buffer('global_step', torch.LongTensor(1).zero_()) 13 | 14 | @property 15 | def mode(self): 16 | return 'TRAIN' if self.training else 'TEST' 17 | 18 | def update_global_step(self): 19 | self.global_step += 1 20 | 21 | def forward(self, **kwargs): 22 | raise NotImplementedError 23 | 24 | def post_processing(self, batch_dict): 25 | pred_dicts = {} 26 | ret_dicts = {} 27 | pred_dicts['pred_seg'] = batch_dict['pred_seg'].cpu().numpy() 28 | if 'pred_heatmap' in batch_dict: 29 | pred_dicts['pred_heatmap'] = batch_dict['pred_heatmap'].cpu().numpy() 30 | pred_dicts['pred_traj'] = self.header.get_prediction(batch_dict) 31 | return pred_dicts, ret_dicts 32 | 33 | def get_training_loss(self): 34 | disp_dict = {} 35 | 36 | loss, tb_dict = self.header.get_loss() 37 | 38 | tb_dict = { 39 | 'loss': loss.item(), 40 | **tb_dict 41 | } 42 | 43 | return loss, tb_dict, disp_dict 44 | 45 | def load_params_from_file(self, filename, logger=None, to_cpu=False): 46 | if not os.path.isfile(filename): 47 | raise FileNotFoundError 48 | 49 | if logger: 50 | logger.info('==> Loading parameters from checkpoint %s to %s' % (filename, 'CPU' if to_cpu else 'GPU')) 51 | loc_type = torch.device('cpu') if to_cpu else None 52 | checkpoint = torch.load(filename, map_location=loc_type) 53 | model_state_disk = checkpoint['model_state'] 54 | 55 | if logger and 'version' in checkpoint: 56 | logger.info('==> Checkpoint trained from version: %s' % checkpoint['version']) 57 | 58 | update_model_state = {} 59 | for key, val in model_state_disk.items(): 60 | if key in self.state_dict(): 61 | if self.state_dict()[key].shape == model_state_disk[key].shape: 62 | update_model_state[key] = val 63 | # logger.info('Update weight %s: %s' % (key, str(val.shape))) 64 | #else: 65 | # logger.info('Shape not matched %s: self --> %s vs disk --> %s ' % (key, str(self.state_dict()[key].shape), str(val.shape))) 66 | 67 | 68 | state_dict = self.state_dict() 69 | state_dict.update(update_model_state) 70 | self.load_state_dict(state_dict) 71 | 72 | for key in state_dict: 73 | if key not in update_model_state and logger: 74 | logger.info('Not updated weight %s: %s' % (key, str(state_dict[key].shape))) 75 | 76 | if logger: 77 | logger.info('==> Done (loaded %d/%d)' % (len(update_model_state), len(self.state_dict()))) 78 | else: 79 | print('==> Done (loaded %d/%d)' % (len(update_model_state), len(self.state_dict()))) 80 | 81 | def load_params_with_optimizer(self, filename, to_cpu=False, optimizer=None, logger=None): 82 | if not os.path.isfile(filename): 83 | raise FileNotFoundError 84 | 85 | logger.info('==> Loading parameters from checkpoint %s to %s' % (filename, 'CPU' if to_cpu else 'GPU')) 86 | loc_type = torch.device('cpu') if to_cpu else None 87 | checkpoint = torch.load(filename, map_location=loc_type) 88 | epoch = checkpoint.get('epoch', -1) 89 | it = checkpoint.get('it', 0.0) 90 | self.load_state_dict(checkpoint['model_state']) 91 | 92 | if optimizer is not None: 93 | if 'optimizer_state' in checkpoint and checkpoint['optimizer_state'] is not None: 94 | logger.info('==> Loading optimizer parameters from checkpoint %s to %s' 95 | % (filename, 'CPU' if to_cpu else 'GPU')) 96 | optimizer.load_state_dict(checkpoint['optimizer_state']) 97 | else: 98 | assert filename[-4] == '.', filename 99 | src_file, ext = filename[:-4], filename[-3:] 100 | optimizer_filename = '%s_optim.%s' % (src_file, ext) 101 | if os.path.exists(optimizer_filename): 102 | optimizer_ckpt = torch.load(optimizer_filename, map_location=loc_type) 103 | optimizer.load_state_dict(optimizer_ckpt['optimizer_state']) 104 | 105 | if 'version' in checkpoint: 106 | print('==> Checkpoint trained from version: %s' % checkpoint['version']) 107 | logger.info('==> Done') 108 | 109 | return it, epoch 110 | -------------------------------------------------------------------------------- /lidardet/datasets/processor/data_processor.py: -------------------------------------------------------------------------------- 1 | from functools import partial 2 | import numpy as np 3 | #from ...utils import box_utils, geometry 4 | #from ...utils.common import scan_downsample, scan_upsample, scan_to_range 5 | 6 | class DataProcessor(object): 7 | def __init__(self, processor_configs): 8 | self.grid_size = self.voxel_size = None 9 | self.data_processor_queue = [] 10 | for cur_cfg in processor_configs: 11 | cur_processor = getattr(self, cur_cfg['name'])(config=cur_cfg) 12 | self.data_processor_queue.append(cur_processor) 13 | 14 | def remove_points_and_boxes_outside_range(self, data_dict=None, config=None): 15 | 16 | if data_dict is None: 17 | return partial(self.remove_points_and_boxes_outside_range, config=config) 18 | 19 | point_cloud_range = np.array(config['point_cloud_range'], dtype=np.float32) 20 | mask = geometry.mask_points_by_range(data_dict['points'], point_cloud_range) 21 | data_dict['points'] = data_dict['points'][mask] 22 | if data_dict.get('gt_boxes', None) is not None and config['remove_outside_boxes']: 23 | mask = box_utils.mask_boxes_outside_range_numpy( 24 | data_dict['gt_boxes'], point_cloud_range, min_num_corners=config.get('min_num_corners', 1) 25 | ) 26 | data_dict['gt_boxes'] = data_dict['gt_boxes'][mask] 27 | return data_dict 28 | 29 | def normalization(self, data_dict=None, config=None): 30 | if data_dict is None: 31 | return partial(self.normalization, config=config) 32 | 33 | img = data_dict['range_image_in'] 34 | for i in range(img.shape[0]): 35 | img[i,...] = (img[i,...] - config['mean'][i]) / config['std'][i] 36 | 37 | data_dict['range_image_in'] = img 38 | return data_dict 39 | 40 | def shuffle_points(self, data_dict=None, config=None): 41 | if data_dict is None: 42 | return partial(self.shuffle_points, config=config) 43 | 44 | points = data_dict['points'] 45 | shuffle_idx = np.random.permutation(points.shape[0]) 46 | points = points[shuffle_idx] 47 | data_dict['points'] = points 48 | 49 | return data_dict 50 | 51 | def get_range_image(self, data_dict=None, config=None): 52 | if data_dict is None: 53 | return partial(self.get_range_image, config=config) 54 | points = data_dict['points'] 55 | range_image, _, _ = scan_to_range(points, normalize=True) 56 | range_image = range_image[:,:,:2] 57 | data_dict['range_image_even'] = range_image[::2, :, :].transpose((2, 0, 1)) 58 | data_dict['range_image_odd'] = range_image[1::2, :, :].transpose((2, 0, 1)) 59 | 60 | if 'points' in data_dict: 61 | data_dict.pop('points') 62 | if 'gt_boxes' in data_dict: 63 | data_dict.pop('gt_boxes') 64 | return data_dict 65 | 66 | def scan_downsample(self, data_dict=None, config=None): 67 | if data_dict is None: 68 | return partial(self.scan_downsample, config=config) 69 | points = data_dict['points'] 70 | points_lr = scan_downsample(points) 71 | data_dict['points'] = points_lr 72 | return data_dict 73 | 74 | def scan_upsample(self, data_dict=None, config=None): 75 | if data_dict is None: 76 | return partial(self.scan_upsample, config=config) 77 | points = data_dict['points'] 78 | points_dense = scan_upsample(points) 79 | data_dict['points'] = points_dense 80 | return data_dict 81 | 82 | def voxelization(self, data_dict=None, config=None, voxel_generator=None): 83 | if data_dict is None: 84 | from spconv.utils import VoxelGenerator 85 | point_cloud_range = np.array(config['point_cloud_range'], dtype=np.float32) 86 | voxel_generator = VoxelGenerator( 87 | voxel_size=config['voxel_size'], 88 | point_cloud_range=point_cloud_range, 89 | max_num_points=config['max_points_per_voxel'], 90 | max_voxels=config['max_num_voxels'], full_mean=False 91 | ) 92 | grid_size = (point_cloud_range[3:6] - point_cloud_range[0:3]) / np.array(config['voxel_size']) 93 | self.grid_size = np.round(grid_size).astype(np.int64) 94 | self.voxel_size = config['voxel_size'] 95 | return partial(self.voxelization, voxel_generator=voxel_generator) 96 | points = data_dict['points'] 97 | voxels, coordinates, num_points = voxel_generator.generate(points) 98 | data_dict['use_lead_xyz'] = True 99 | #if not data_dict['use_lead_xyz']: 100 | # voxels = voxels[..., 3:] # remove xyz in voxels(N, 3) 101 | 102 | data_dict['voxels'] = voxels 103 | data_dict['voxel_coords'] = coordinates 104 | data_dict['voxel_num_points'] = num_points 105 | return data_dict 106 | 107 | def forward(self, data_dict): 108 | """ 109 | Args: 110 | data_dict: 111 | points: (N, 3 + C_in) 112 | gt_boxes: optional, (N, 7 + C) [x, y, z, dx, dy, dz, heading, ...] 113 | gt_names: optional, (N), string 114 | ... 115 | 116 | Returns: 117 | """ 118 | 119 | for cur_processor in self.data_processor_queue: 120 | data_dict = cur_processor(data_dict=data_dict) 121 | 122 | return data_dict 123 | -------------------------------------------------------------------------------- /lidardet/losses/lovasz_softmax.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | MIT License 4 | 5 | Copyright (c) 2018 Maxim Berman 6 | Copyright (c) 2020 Tiago Cortinhal, George Tzelepis and Eren Erdal Aksoy 7 | 8 | 9 | Permission is hereby granted, free of charge, to any person obtaining a copy 10 | of this software and associated documentation files (the "Software"), to deal 11 | in the Software without restriction, including without limitation the rights 12 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 13 | copies of the Software, and to permit persons to whom the Software is 14 | furnished to do so, subject to the following conditions: 15 | 16 | The above copyright notice and this permission notice shall be included in all 17 | copies or substantial portions of the Software. 18 | 19 | """ 20 | import torch 21 | import torch.nn as nn 22 | from torch.autograd import Variable 23 | 24 | 25 | try: 26 | from itertools import ifilterfalse 27 | except ImportError: 28 | from itertools import filterfalse as ifilterfalse 29 | 30 | 31 | def isnan(x): 32 | return x != x 33 | 34 | 35 | def mean(l, ignore_nan=False, empty=0): 36 | """ 37 | nanmean compatible with generators. 38 | """ 39 | l = iter(l) 40 | if ignore_nan: 41 | l = ifilterfalse(isnan, l) 42 | try: 43 | n = 1 44 | acc = next(l) 45 | except StopIteration: 46 | if empty == 'raise': 47 | raise ValueError('Empty mean') 48 | return empty 49 | for n, v in enumerate(l, 2): 50 | acc += v 51 | if n == 1: 52 | return acc 53 | return acc / n 54 | 55 | 56 | def lovasz_grad(gt_sorted): 57 | """ 58 | Computes gradient of the Lovasz extension w.r.t sorted errors 59 | See Alg. 1 in paper 60 | """ 61 | p = len(gt_sorted) 62 | gts = gt_sorted.sum() 63 | intersection = gts - gt_sorted.float().cumsum(0) 64 | union = gts + (1 - gt_sorted).float().cumsum(0) 65 | jaccard = 1. - intersection / union 66 | if p > 1: # cover 1-pixel case 67 | jaccard[1:p] = jaccard[1:p] - jaccard[0:-1] 68 | return jaccard 69 | 70 | 71 | def lovasz_softmax(probas, labels, classes='present', per_image=False, ignore=None): 72 | """ 73 | Multi-class Lovasz-Softmax loss 74 | probas: [B, C, H, W] Variable, class probabilities at each prediction (between 0 and 1). 75 | Interpreted as binary (sigmoid) output with outputs of size [B, H, W]. 76 | labels: [B, H, W] Tensor, ground truth labels (between 0 and C - 1) 77 | classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average. 78 | per_image: compute the loss per image instead of per batch 79 | ignore: void class labels 80 | """ 81 | if per_image: 82 | loss = mean(lovasz_softmax_flat(*flatten_probas(prob.unsqueeze(0), lab.unsqueeze(0), ignore), classes=classes) 83 | for prob, lab in zip(probas, labels)) 84 | else: 85 | loss = lovasz_softmax_flat(*flatten_probas(probas, labels, ignore), classes=classes) 86 | return loss 87 | 88 | 89 | def lovasz_softmax_flat(probas, labels, classes='present'): 90 | """ 91 | Multi-class Lovasz-Softmax loss 92 | probas: [P, C] Variable, class probabilities at each prediction (between 0 and 1) 93 | labels: [P] Tensor, ground truth labels (between 0 and C - 1) 94 | classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average. 95 | """ 96 | if probas.numel() == 0: 97 | # only void pixels, the gradients should be 0 98 | return probas * 0. 99 | C = probas.size(1) 100 | losses = [] 101 | class_to_sum = list(range(C)) if classes in ['all', 'present'] else classes 102 | for c in class_to_sum: 103 | fg = (labels == c).float() # foreground for class c 104 | if (classes is 'present' and fg.sum() == 0): 105 | continue 106 | if C == 1: 107 | if len(classes) > 1: 108 | raise ValueError('Sigmoid output possible only with 1 class') 109 | class_pred = probas[:, 0] 110 | else: 111 | class_pred = probas[:, c] 112 | errors = (Variable(fg) - class_pred).abs() 113 | errors_sorted, perm = torch.sort(errors, 0, descending=True) 114 | perm = perm.data 115 | fg_sorted = fg[perm] 116 | losses.append(torch.dot(errors_sorted, Variable(lovasz_grad(fg_sorted)))) 117 | return mean(losses) 118 | 119 | 120 | def flatten_probas(probas, labels, ignore=None): 121 | """ 122 | Flattens predictions in the batch 123 | """ 124 | if probas.dim() == 3: 125 | # assumes output of a sigmoid layer 126 | B, H, W = probas.size() 127 | probas = probas.view(B, 1, H, W) 128 | B, C, H, W = probas.size() 129 | probas = probas.permute(0, 2, 3, 1).contiguous().view(-1, C) # B * H * W, C = P, C 130 | labels = labels.view(-1) 131 | if ignore is None: 132 | return probas, labels 133 | valid = (labels != ignore) 134 | vprobas = probas[valid.nonzero().squeeze()] 135 | vlabels = labels[valid] 136 | return vprobas, vlabels 137 | 138 | 139 | class Lovasz_softmax(nn.Module): 140 | def __init__(self, classes='present', per_image=False, ignore=None): 141 | super(Lovasz_softmax, self).__init__() 142 | self.classes = classes 143 | self.per_image = per_image 144 | self.ignore = ignore 145 | 146 | def forward(self, probas, labels): 147 | return lovasz_softmax(probas, labels, self.classes, self.per_image, self.ignore) 148 | -------------------------------------------------------------------------------- /lidardet/utils/dist.py: -------------------------------------------------------------------------------- 1 | """ 2 | This file contains primitives for multi-gpu communication. 3 | This is useful when doing distributed training. 4 | """ 5 | 6 | import functools 7 | import pickle 8 | import sys 9 | import time 10 | from getpass import getuser 11 | from socket import gethostname 12 | 13 | import torch 14 | import torch.distributed as dist 15 | import torch.multiprocessing as mp 16 | 17 | 18 | def init_dist_pytorch(batch_size, local_rank, tcp_port=18888, backend='nccl'): 19 | if mp.get_start_method(allow_none=True) is None: 20 | mp.set_start_method('spawn') 21 | 22 | num_gpus = torch.cuda.device_count() 23 | torch.cuda.set_device(local_rank % num_gpus) 24 | dist.init_process_group( 25 | backend=backend, 26 | init_method='tcp://127.0.0.1:%d' % tcp_port, 27 | rank=local_rank, 28 | world_size=num_gpus 29 | ) 30 | assert batch_size % num_gpus == 0, 'Batch size should be matched with GPUS: (%d, %d)' % ( 31 | batch_size, num_gpus) 32 | batch_size_each_gpu = batch_size // num_gpus 33 | rank = dist.get_rank() 34 | return batch_size_each_gpu, rank 35 | 36 | 37 | def get_host_info(): 38 | return "{}@{}".format(getuser(), gethostname()) 39 | 40 | 41 | def get_dist_info(): 42 | if torch.__version__ < "1.0": 43 | initialized = dist._initialized 44 | else: 45 | initialized = dist.is_initialized() 46 | if initialized: 47 | rank = dist.get_rank() 48 | world_size = dist.get_world_size() 49 | else: 50 | rank = 0 51 | world_size = 1 52 | return rank, world_size 53 | 54 | 55 | def master_only(func): 56 | @functools.wraps(func) 57 | def wrapper(*args, **kwargs): 58 | rank, _ = get_dist_info() 59 | if rank == 0: 60 | return func(*args, **kwargs) 61 | 62 | return wrapper 63 | 64 | 65 | def get_time_str(): 66 | return time.strftime("%Y%m%d_%H%M%S", time.localtime()) 67 | 68 | 69 | def get_world_size(): 70 | if not dist.is_available(): 71 | return 1 72 | if not dist.is_initialized(): 73 | return 1 74 | return dist.get_world_size() 75 | 76 | 77 | def get_rank(): 78 | if not dist.is_available(): 79 | return 0 80 | if not dist.is_initialized(): 81 | return 0 82 | return dist.get_rank() 83 | 84 | 85 | def is_main_process(): 86 | return get_rank() == 0 87 | 88 | 89 | def synchronize(): 90 | """ 91 | Helper function to synchronize (barrier) among all processes when 92 | using distributed training 93 | """ 94 | if not dist.is_available(): 95 | return 96 | if not dist.is_initialized(): 97 | return 98 | world_size = dist.get_world_size() 99 | if world_size == 1: 100 | return 101 | dist.barrier() 102 | 103 | 104 | def all_gather(data): 105 | """ 106 | Run all_gather on arbitrary picklable data (not necessarily tensors) 107 | Args: 108 | data: any picklable object 109 | Returns: 110 | list[data]: list of data gathered from each rank 111 | """ 112 | world_size = get_world_size() 113 | if world_size == 1: 114 | return [data] 115 | 116 | # serialized to a Tensor 117 | buffer = pickle.dumps(data) 118 | storage = torch.ByteStorage.from_buffer(buffer) 119 | tensor = torch.ByteTensor(storage).to("cuda") 120 | 121 | # obtain Tensor size of each rank 122 | local_size = torch.IntTensor([tensor.numel()]).to("cuda") 123 | size_list = [torch.IntTensor([0]).to("cuda") for _ in range(world_size)] 124 | dist.all_gather(size_list, local_size) 125 | size_list = [int(size.item()) for size in size_list] 126 | max_size = max(size_list) 127 | 128 | # receiving Tensor from all ranks 129 | # we pad the tensor because torch all_gather does not support 130 | # gathering tensors of different shapes 131 | tensor_list = [] 132 | for _ in size_list: 133 | tensor_list.append(torch.ByteTensor(size=(max_size,)).to("cuda")) 134 | if local_size != max_size: 135 | padding = torch.ByteTensor(size=(max_size - local_size,)).to("cuda") 136 | tensor = torch.cat((tensor, padding), dim=0) 137 | dist.all_gather(tensor_list, tensor) 138 | 139 | data_list = [] 140 | for size, tensor in zip(size_list, tensor_list): 141 | buffer = tensor.cpu().numpy().tobytes()[:size] 142 | data_list.append(pickle.loads(buffer)) 143 | 144 | return data_list 145 | 146 | 147 | def reduce_dict(input_dict, average=True): 148 | """ 149 | Args: 150 | input_dict (dict): all the values will be reduced 151 | average (bool): whether to do average or sum 152 | Reduce the values in the dictionary from all processes so that process with rank 153 | 0 has the averaged results. Returns a dict with the same fields as 154 | input_dict, after reduction. 155 | """ 156 | world_size = get_world_size() 157 | if world_size < 2: 158 | return input_dict 159 | with torch.no_grad(): 160 | names = [] 161 | values = [] 162 | # sort the keys so that they are consistent across processes 163 | for k in sorted(input_dict.keys()): 164 | names.append(k) 165 | values.append(input_dict[k]) 166 | values = torch.stack(values, dim=0) 167 | dist.reduce(values, dst=0) 168 | if dist.get_rank() == 0 and average: 169 | # only main process gets accumulated, so only divide by 170 | # world_size in this case 171 | values /= world_size 172 | reduced_dict = {k: v for k, v in zip(names, values)} 173 | return reduced_dict 174 | -------------------------------------------------------------------------------- /lidardet/models/trajectory_predictor/conv_header.py: -------------------------------------------------------------------------------- 1 | import math 2 | import numpy as np 3 | import torch 4 | from torch import nn 5 | 6 | from .backbone import conv3x3 7 | from .transformer import Transformer 8 | 9 | class ConvHeader(nn.Module): 10 | 11 | def __init__(self, cfg): 12 | super(ConvHeader, self).__init__() 13 | self.cfg = cfg 14 | self.use_bn = cfg.use_bn 15 | self.use_transformer = cfg.use_transformer 16 | self.use_road_seg = cfg.use_road_seg 17 | 18 | bias = not self.use_bn 19 | self.conv1 = conv3x3(96, 96, bias=bias) 20 | self.bn1 = nn.BatchNorm2d(96) 21 | self.conv2 = conv3x3(96, 96, bias=bias) 22 | self.bn2 = nn.BatchNorm2d(96) 23 | self.conv3 = conv3x3(96, 96, bias=bias) 24 | self.bn3 = nn.BatchNorm2d(96) 25 | self.conv4 = conv3x3(96, 96, bias=bias) 26 | self.bn4 = nn.BatchNorm2d(96) 27 | 28 | self.heatmap_head = conv3x3(96, self.cfg.num_points_per_trajectory, bias=True) 29 | self.softmax = nn.Softmax(dim=-1) 30 | 31 | npoints = self.cfg.num_points_per_trajectory 32 | if self.use_road_seg: 33 | self.road_head = conv3x3(96, 1, bias=True) 34 | self.conv7 = nn.Conv2d(1, npoints, 3, 2, padding=1, bias=True) 35 | 36 | if self.cfg.use_transformer: 37 | self.conv5 = nn.Conv2d(96, npoints, 3, 2, padding=1, bias=True) 38 | # depth-wise convolution 39 | if self.cfg.transformer.use_position_encoder: 40 | self.conv6 = nn.Conv2d(npoints, npoints, 3, 2, groups=npoints, padding=1, bias=True) 41 | self.waypoint_predictor = Transformer(self.cfg.transformer) 42 | 43 | self.road_loss_func = nn.BCELoss().cuda() 44 | self.heatmap_loss_func = nn.MSELoss().cuda() 45 | 46 | def spatial_softmax(self, x): 47 | b, c, h, w = x.shape 48 | x = self.softmax(x.view(b, c, -1)) 49 | x = x.view(b, c, h, w) 50 | 51 | argmax = torch.argmax(x.view(b * c, -1), dim=1) 52 | argmax_x, argmax_y = torch.remainder(argmax, w).float(), torch.floor(torch.div(argmax.float(), float(w))) 53 | argmax_x = argmax_x.view((b, c, -1)) / float(w) 54 | argmax_y = argmax_y.view((b, c, -1)) / float(h) 55 | pos_pred = torch.cat([argmax_x, argmax_y], dim=2) 56 | 57 | return x, pos_pred 58 | 59 | def forward(self, batch_dict): 60 | feature_map = batch_dict['seg_features'] 61 | 62 | x = self.conv1(feature_map) 63 | if self.use_bn: 64 | x = self.bn1(x) 65 | x = self.conv2(x) 66 | if self.use_bn: 67 | x = self.bn2(x) 68 | x = self.conv3(x) 69 | if self.use_bn: 70 | x = self.bn3(x) 71 | x = self.conv4(x) 72 | if self.use_bn: 73 | x = self.bn4(x) 74 | 75 | heatmap = self.heatmap_head(x) # [b, num_points_per_trajectory, 75, 100] 76 | heatmap, pred_pos = self.spatial_softmax(heatmap) 77 | 78 | if self.use_road_seg: 79 | road = torch.sigmoid(self.road_head(x)) # [b, 1, 75, 100] 80 | batch_dict['pred_seg'] = road.squeeze(1) 81 | else: 82 | batch_dict['pred_seg'] = heatmap.sum(dim=1) 83 | 84 | batch_dict['pred_heatmap'] = heatmap.sum(dim=1) 85 | 86 | if self.cfg.use_transformer: 87 | waypoints_feature = self.conv5(feature_map) 88 | b, c, h, w = waypoints_feature.shape 89 | if self.cfg.transformer.use_position_encoder: 90 | x2 = self.conv6(heatmap) 91 | batch_dict['pos_feature'] = x2.view(b, c, -1) 92 | 93 | if self.use_road_seg: 94 | x3 = self.conv7(road) 95 | waypoints_feature += x3 96 | 97 | waypoints_feature = waypoints_feature.view(b, c, -1) 98 | 99 | batch_dict['waypoints_feature'] = waypoints_feature 100 | batch_dict = self.waypoint_predictor(batch_dict) 101 | else: 102 | batch_dict['waypoints_pred'] = pred_pos 103 | 104 | if self.training: 105 | 106 | if self.use_road_seg: 107 | self.road_pred = road.squeeze(1) 108 | self.road_target = batch_dict['img_ins'] 109 | 110 | self.heatmap_pred = heatmap 111 | self.heatmap_targets = batch_dict['heatmap'] 112 | 113 | return batch_dict 114 | 115 | def get_loss(self): 116 | 117 | loss_heatmap = self.heatmap_loss_func(self.heatmap_pred, self.heatmap_targets) 118 | loss_heatmap *= self.cfg.weight_loss_heatmap 119 | loss = loss_heatmap 120 | tb_dict = {'loss_heatmap': loss_heatmap} 121 | 122 | if self.cfg.use_transformer: 123 | loss_waypoint = self.waypoint_predictor.get_loss() 124 | loss_waypoint *= self.cfg.weight_loss_waypoint 125 | loss += loss_waypoint 126 | tb_dict['loss_waypoint'] = loss_waypoint 127 | 128 | if self.use_road_seg: 129 | loss_road = self.road_loss_func(self.road_pred, self.road_target) 130 | loss_road *= self.cfg.weight_loss_road 131 | loss += loss_road 132 | tb_dict['loss_road'] = loss_road 133 | 134 | return loss, tb_dict 135 | 136 | def get_prediction(self, batch_dict): 137 | pred_points = batch_dict['waypoints_pred'] 138 | return pred_points.detach().cpu().numpy() 139 | 140 | def parse_predicted_waypoints(self, prediction): 141 | assert isinstance(prediction, torch.Tensor) 142 | predicted_points = [] 143 | prediction = prediction.detach().cpu().numpy() 144 | c, h, w = prediction.shape[0:3] 145 | num_points = c 146 | assert num_points == self.cfg.num_points_per_trajectory 147 | for n in range(num_points): 148 | tmp = prediction[n] 149 | y, x = np.unravel_index(np.argmax(tmp), tmp.shape) 150 | point = np.array([x, y]) 151 | predicted_points.append(point) 152 | return np.array(predicted_points).astype(np.float32) 153 | -------------------------------------------------------------------------------- /lidardet/datasets/trajectory/utils.py: -------------------------------------------------------------------------------- 1 | import math 2 | import random 3 | import numpy as np 4 | from scipy.spatial.distance import directed_hausdorff 5 | 6 | from lidardet.utils import geometry 7 | 8 | def hausdorff_dist(traj1, traj2): 9 | """ 10 | Double side hausdorff distance 11 | """ 12 | dist1 = directed_hausdorff(traj1, traj2)[0] 13 | dist2 = directed_hausdorff(traj2, traj1)[0] 14 | return max(dist1, dist2) 15 | 16 | def l2_distance_to_trajectory_set(traj, traj_set): 17 | num_set, num_points = traj_set.shape[:2] 18 | traj = traj[:num_points,:2] 19 | min_dist = 1e5 20 | min_idx = 0 21 | for i, traj_ref in enumerate(traj_set): 22 | _,_,metric = interp_l2_distance(traj, traj_ref) 23 | d = metric['ADE'] 24 | if d < min_dist: 25 | min_dist = d 26 | min_idx = i 27 | return min_idx, min_dist 28 | 29 | def angle_to_trajectory_set(traj, traj_set): 30 | num_set, num_points = traj_set.shape[:2] 31 | traj = traj[:num_points,:2] 32 | angles = [] 33 | for traj_ref in traj_set: 34 | angle = 0 35 | for i in range(num_points): 36 | angle += angle_between(traj[i,:], traj_ref[i,:]) 37 | angles.append(angle / num_points) 38 | idx = np.argmin(np.array(angles)) 39 | return idx, angles[idx] 40 | 41 | def angle_between(v1, v2): 42 | norm = np.linalg.norm(v1) * np.linalg.norm(v2) 43 | if math.isclose(norm, 0): 44 | return 0 45 | 46 | dot_product = v1.dot(v2) 47 | angle = math.degrees(math.acos(max(min(dot_product / norm, 1), -1))) 48 | if angle == 180: 49 | return angle - 1e-5 50 | 51 | return angle 52 | 53 | def l2_distance(traj, traj_ref): 54 | metric = {} 55 | d = np.linalg.norm(traj - traj_ref, axis=1) 56 | metric['ADE'] = np.mean(d) 57 | metric['FDE'] = d[-1] 58 | metric['MAX'] = np.max(d) 59 | return metric 60 | 61 | def interp_l2_distance(traj, traj_ref, samples=20): 62 | x1, y1 = traj_ref[-1,:2] 63 | if abs(y1) > abs(x1): 64 | ymin = np.min(traj_ref[:,1]) 65 | ymax = np.max(traj_ref[:,1]) 66 | yvals = np.linspace(ymin, ymax, samples) 67 | xinterp = np.interp(yvals, traj[:,1], traj[:,0]) 68 | xinterp_ref = np.interp(yvals, traj_ref[:,1], traj_ref[:,0]) 69 | traj_interp = np.vstack([xinterp, yvals]).T 70 | traj_interp_ref = np.vstack([xinterp_ref, yvals]).T 71 | else: 72 | xmin = np.min(traj_ref[:,0]) 73 | xmax = np.max(traj_ref[:,0]) 74 | xvals = np.linspace(xmin, xmax, samples) 75 | yinterp = np.interp(xvals, traj[:,0], traj[:,1]) 76 | yinterp_ref = np.interp(xvals, traj_ref[:,0], traj_ref[:,1]) 77 | traj_interp = np.vstack([xvals, yinterp]).T 78 | traj_interp_ref = np.vstack([xvals, yinterp_ref]).T 79 | 80 | metric = l2_distance(traj_interp, traj_interp_ref) 81 | return traj_interp, traj_interp_ref, metric 82 | 83 | def random_shift(traj_list, res, max_dist=1.0): 84 | rx = max_dist * 2 * (random.random() - 0.5) # [-1.0, 1.0) 85 | #print('rx1', rx) 86 | rx = rx / res 87 | for i in range(len(traj_list)): 88 | traj_list[i][:,0] += rx 89 | #print(rx) 90 | return traj_list 91 | 92 | def random_flip(points, traj_list): 93 | """ 94 | Random flip along y axis 95 | Args: 96 | points: (M, 4) 97 | traj_list: list of trajectories with shape of (N, 2) 98 | Returns: 99 | """ 100 | enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5]) 101 | if enable: 102 | points[:, 0] = -points[:, 0] 103 | for i in range(len(traj_list)): 104 | traj_list[i][:, 0] = - traj_list[i][:, 0] 105 | return points, traj_list 106 | 107 | def random_rotate(points, traj_list, rot_range=[0, np.pi/4]): 108 | """ 109 | Random rotate along z axis 110 | Args: 111 | points: (M, 4) 112 | traj_list: list of trajectories with shape of (N, 2) 113 | rot_range: [min, max] 114 | Returns: 115 | """ 116 | noise_rotation = np.random.uniform(rot_range[0], rot_range[1]) 117 | points = geometry.rotate_points_along_z(points[np.newaxis, :, :], np.array([noise_rotation]))[0] 118 | for i, traj in enumerate(traj_list): 119 | traj_points = np.zeros((1, traj.shape[0], 3)) 120 | traj_points[0,:,:2] = traj[:,:2] 121 | traj_list[i][:,:2] = geometry.rotate_points_along_z(traj_points, np.array([noise_rotation]))[0][:,:2] 122 | return points, traj_list 123 | 124 | def random_drop(points, drop_ratio=0.5): 125 | enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5]) 126 | if enable: 127 | np.random.shuffle(points) 128 | points = points[:int(points.shape[0] * drop_ratio), :] 129 | return points 130 | 131 | class HeatMapGenerator(): 132 | def __init__(self, img_w, img_h, scale = 1.0): 133 | self.img_w = int(img_w * scale) 134 | self.img_h = int(img_h * scale) 135 | sigma = min(self.img_h, self.img_w) / 64.0 136 | size = 6 * sigma + 3 137 | x = np.arange(0, size, 1, float) 138 | y = x[:, np.newaxis] 139 | x0 = y0 = 3 * sigma + 1 140 | self.g = np.exp(-((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2)) 141 | self.sigma = sigma 142 | 143 | def __call__(self, points): 144 | hms = np.zeros((points.shape[0], self.img_h, self.img_w), dtype=np.float32) 145 | sigma = self.sigma 146 | points[:,0] *= self.img_w 147 | points[:,1] *= self.img_h 148 | for idx, p in enumerate(points): 149 | x, y = int(p[0]), int(p[1]) 150 | ul = int(x - 3*sigma - 1), int(y - 3*sigma - 1) 151 | br = int(x + 3*sigma + 2), int(y + 3*sigma + 2) 152 | 153 | c,d = max(0, -ul[0]), min(br[0], self.img_w) - ul[0] 154 | a,b = max(0, -ul[1]), min(br[1], self.img_h) - ul[1] 155 | 156 | cc,dd = max(0, ul[0]), min(br[0], self.img_w) 157 | aa,bb = max(0, ul[1]), min(br[1], self.img_h) 158 | hms[idx, aa:bb,cc:dd] = np.maximum(hms[idx, aa:bb,cc:dd], self.g[a:b,c:d]) 159 | return hms 160 | -------------------------------------------------------------------------------- /tools/train_utils/train_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import os 3 | import glob 4 | import tqdm 5 | from torch.nn.utils import clip_grad_norm_ 6 | 7 | 8 | def train_one_epoch(model, optimizer, train_loader, model_func, lr_scheduler, accumulated_iter, optim_cfg, 9 | rank, tbar, total_it_each_epoch, dataloader_iter, tb_log=None, leave_pbar=False): 10 | if total_it_each_epoch == len(train_loader): 11 | dataloader_iter = iter(train_loader) 12 | 13 | if rank == 0: 14 | pbar = tqdm.tqdm(total=total_it_each_epoch, leave=leave_pbar, desc='train', dynamic_ncols=True) 15 | 16 | for cur_it in range(total_it_each_epoch): 17 | try: 18 | batch = next(dataloader_iter) 19 | except StopIteration: 20 | dataloader_iter = iter(train_loader) 21 | batch = next(dataloader_iter) 22 | print('new iters') 23 | 24 | #lr_scheduler.step(accumulated_iter) 25 | 26 | try: 27 | cur_lr = float(optimizer.lr) 28 | except: 29 | cur_lr = optimizer.param_groups[0]['lr'] 30 | 31 | if tb_log is not None: 32 | tb_log.add_scalar('learning_rate', cur_lr, accumulated_iter) 33 | 34 | model.train() 35 | optimizer.zero_grad() 36 | 37 | loss, tb_dict, disp_dict = model_func(model, batch) 38 | 39 | loss.backward() 40 | if optim_cfg.get('grad_norm_clip', None): 41 | clip_grad_norm_(model.parameters(), optim_cfg.grad_norm_clip) 42 | optimizer.step() 43 | 44 | lr_scheduler.step(accumulated_iter) 45 | 46 | accumulated_iter += 1 47 | disp_dict.update({'loss': loss.item(), 'lr': cur_lr}) 48 | 49 | # log to console and tensorboard 50 | if rank == 0: 51 | pbar.update() 52 | pbar.set_postfix(dict(total_it=accumulated_iter)) 53 | tbar.set_postfix(disp_dict) 54 | tbar.refresh() 55 | 56 | if tb_log is not None: 57 | tb_log.add_scalar('train_loss', loss, accumulated_iter) 58 | tb_log.add_scalar('learning_rate', cur_lr, accumulated_iter) 59 | for key, val in tb_dict.items(): 60 | tb_log.add_scalar('train_' + key, val, accumulated_iter) 61 | if rank == 0: 62 | pbar.close() 63 | return accumulated_iter 64 | 65 | 66 | def train_model(model, optimizer, train_loader, model_func, lr_scheduler, optim_cfg, 67 | start_epoch, total_epochs, start_iter, rank, tb_log, ckpt_save_dir, train_sampler=None, 68 | lr_warmup_scheduler=None, ckpt_save_interval=1, max_ckpt_save_num=50, 69 | merge_all_iters_to_one_epoch=False): 70 | accumulated_iter = start_iter 71 | with tqdm.trange(start_epoch, total_epochs, desc='epochs', dynamic_ncols=True, leave=(rank == 0)) as tbar: 72 | total_it_each_epoch = len(train_loader) 73 | if merge_all_iters_to_one_epoch: 74 | assert hasattr(train_loader.dataset, 'merge_all_iters_to_one_epoch') 75 | train_loader.dataset.merge_all_iters_to_one_epoch(merge=True, epochs=total_epochs) 76 | total_it_each_epoch = len(train_loader) // max(total_epochs, 1) 77 | 78 | dataloader_iter = iter(train_loader) 79 | for cur_epoch in tbar: 80 | if train_sampler is not None: 81 | train_sampler.set_epoch(cur_epoch) 82 | 83 | # train one epoch 84 | if lr_warmup_scheduler is not None and cur_epoch < optim_cfg.warmup_epoch: 85 | cur_scheduler = lr_warmup_scheduler 86 | else: 87 | cur_scheduler = lr_scheduler 88 | accumulated_iter = train_one_epoch( 89 | model, optimizer, train_loader, model_func, 90 | lr_scheduler=cur_scheduler, 91 | accumulated_iter=accumulated_iter, optim_cfg=optim_cfg, 92 | rank=rank, tbar=tbar, tb_log=tb_log, 93 | leave_pbar=(cur_epoch + 1 == total_epochs), 94 | total_it_each_epoch=total_it_each_epoch, 95 | dataloader_iter=dataloader_iter 96 | ) 97 | 98 | # save trained model 99 | trained_epoch = cur_epoch + 1 100 | if trained_epoch % ckpt_save_interval == 0 and rank == 0: 101 | 102 | ckpt_list = glob.glob(str(ckpt_save_dir / 'checkpoint_epoch_*.pth')) 103 | ckpt_list.sort(key=os.path.getmtime) 104 | 105 | if ckpt_list.__len__() >= max_ckpt_save_num: 106 | for cur_file_idx in range(0, len(ckpt_list) - max_ckpt_save_num + 1): 107 | os.remove(ckpt_list[cur_file_idx]) 108 | 109 | ckpt_name = ckpt_save_dir / ('checkpoint_epoch_%d' % trained_epoch) 110 | save_checkpoint( 111 | checkpoint_state(model, optimizer, trained_epoch, accumulated_iter), filename=ckpt_name, 112 | ) 113 | 114 | 115 | def model_state_to_cpu(model_state): 116 | model_state_cpu = type(model_state)() # ordered dict 117 | for key, val in model_state.items(): 118 | model_state_cpu[key] = val.cpu() 119 | return model_state_cpu 120 | 121 | 122 | def checkpoint_state(model=None, optimizer=None, epoch=None, it=None): 123 | optim_state = optimizer.state_dict() if optimizer is not None else None 124 | if model is not None: 125 | if isinstance(model, torch.nn.parallel.DistributedDataParallel): 126 | model_state = model_state_to_cpu(model.module.state_dict()) 127 | else: 128 | model_state = model.state_dict() 129 | else: 130 | model_state = None 131 | 132 | try: 133 | import pcdet 134 | version = 'pcdet+' + pcdet.__version__ 135 | except: 136 | version = 'none' 137 | 138 | return {'epoch': epoch, 'it': it, 'model_state': model_state, 'optimizer_state': optim_state, 'version': version} 139 | 140 | 141 | def save_checkpoint(state, filename='checkpoint'): 142 | if False and 'optimizer_state' in state: 143 | optimizer_state = state['optimizer_state'] 144 | state.pop('optimizer_state', None) 145 | optimizer_filename = '{}_optim.pth'.format(filename) 146 | torch.save({'optimizer_state': optimizer_state}, optimizer_filename) 147 | 148 | filename = '{}.pth'.format(filename) 149 | torch.save(state, filename) 150 | -------------------------------------------------------------------------------- /lidardet/utils/scan_unfolding.py: -------------------------------------------------------------------------------- 1 | """ 2 | https://github.com/ltriess/kitti_scan_unfolding 3 | """ 4 | import numpy as np 5 | 6 | def get_kitti_columns(points: np.array, number_of_columns: int = 2000) -> np.array: 7 | """ Returns the column indices for unfolding one or more raw KITTI scans """ 8 | azi = np.arctan2(points[..., 1], points[..., 0]) 9 | cols = number_of_columns * (np.pi - azi) / (2 * np.pi) 10 | # In rare cases it happens that the index is exactly one too large. 11 | cols = np.minimum(cols, number_of_columns - 1) 12 | return np.int32(cols) 13 | 14 | 15 | def get_kitti_rows(points: np.array, threshold: float = -0.005) -> np.array: 16 | """ Returns the row indices for unfolding one or more raw KITTI scans """ 17 | azimuth_flipped = -np.arctan2(points[..., 1], -points[..., 0]) 18 | azi_diffs = azimuth_flipped[..., 1:] - azimuth_flipped[..., :-1] 19 | jump_mask = np.greater(threshold, azi_diffs) 20 | ind = np.add(np.where(jump_mask), 1) 21 | rows = np.zeros_like(points[..., 0]) 22 | rows[..., ind] += 1 23 | return np.int32(np.cumsum(rows, axis=-1)) 24 | 25 | 26 | def projection(points: np.array, *channels, image_size: tuple = (64, 2000)): 27 | """ Scan unfolding of raw KITTI point cloud of shape [N, (x,y,z)] 28 | This functions performs a cylindrical projection of a 3D point cloud given in 29 | cartesian coordinates. The function aims to recover the internal data structure 30 | of the KITTI HDL-64 sensor from its distinct scan pattern that can be approximated 31 | from the original KITTI Raw Dataset. 32 | The resulting projected image (H, W) contains missing values, these positions 33 | are filled with `-1`. If you set `return_inverse` to True, you may receive 34 | a mask for the valid measurements of the image with `valid_image = inverse >= 0`. 35 | This function has less information loss than techniques which compute the 36 | azimuth and elevation angle and sort the points accordingly into the image. 37 | Therefore, most entries in `valid_image` shall be True, except for locations 38 | that hit the ego vehicle or are pointed towards the sky. 39 | The measurements are ordered in decreasing depth before being assigned to image 40 | locations. In case more than one value is scattered to this position, the one 41 | closest to the ego vehicle is selected. To know which points actually made it into 42 | the projection, set `return_active` to True. 43 | Arguments: 44 | points : np.array(shape=(N, 3), dtype=np.float) 45 | The raw KITTI point cloud (does not work with ego-motion corrected data!). 46 | image_size : tuple of 2 integers 47 | The image size of the projection in (H, W). (Default is (64, 2000). 48 | channels : * times np.array(shape=(N, D)) 49 | Additional channels to project in the same way as `points`. 50 | Returns: 51 | Dict( 52 | points=np.array(shape=(H, W, 3), dtype=np.float), # projected point cloud 53 | depth=np.array(shape=(N,), dtype=np.float) # projected depth 54 | channels=List(*), # list of projected additional channels 55 | indices=np.array(shape=(N, 2), dtype=np.int) # image indices from the list 56 | inverse=np.array(shape=(H, W), dtype=np.int) # list indices from the image 57 | active=np.array(shape=(N,), dtype=np.bool) # active array 58 | Raises: 59 | IndexError if projection is not possible, e.g. if you do not use a KITTI 60 | point cloud, or the KITTI point cloud is ego motion corrected 61 | """ 62 | 63 | output = {} 64 | 65 | assert points.shape[1] == 3, "Points must contain N xyz coordinates." 66 | if len(channels) > 0: 67 | assert all( 68 | isinstance(x, np.ndarray) for x in channels 69 | ), "All channels must be numpy arrays." 70 | assert all( 71 | x.shape[0] == points.shape[0] for x in channels 72 | ), "All channels must have the same first dimension as `points`." 73 | 74 | # Get depth of all points for ordering. 75 | depth_list = np.linalg.norm(points, 2, axis=1) 76 | 77 | # Get the indices of the rows and columns to project to. 78 | proj_column = get_kitti_columns(points, number_of_columns=image_size[1]) 79 | proj_row = get_kitti_rows(points) 80 | 81 | if np.any(proj_row >= image_size[0]) or np.any(proj_column >= image_size[1]): 82 | raise IndexError( 83 | "Cannot find valid image indices for this point cloud and image size. " 84 | "Are you sure you entered the correct image size? This function only works " 85 | "with raw KITTI HDL-64 point clouds (no ego motion corrected data allowed)!" 86 | ) 87 | 88 | # Store a copy in original order. 89 | output["indices"] = np.stack([np.copy(proj_row), np.copy(proj_column)], axis=-1) 90 | 91 | # Get the indices in order of decreasing depth. 92 | indices = np.arange(depth_list.shape[0]) 93 | order = np.argsort(depth_list)[::-1] 94 | 95 | indices = indices[order] 96 | proj_column = proj_column[order] 97 | proj_row = proj_row[order] 98 | 99 | # Project the points. 100 | points_img = np.full(shape=(*image_size, 3), fill_value=-1, dtype=np.float32) 101 | points_img[proj_row, proj_column] = points[order] 102 | output["points"] = points_img 103 | 104 | # The depth projection. 105 | depth_img = np.full(shape=image_size, fill_value=-1, dtype=np.float32) 106 | depth_img[proj_row, proj_column] = depth_list[order] 107 | output["depth"] = depth_img 108 | 109 | # Convert all channels. 110 | projected_channels = [] 111 | for channel in channels: 112 | # Set the shape. 113 | _shape = ( 114 | (*image_size, channel.shape[1]) 115 | if len(channel.shape) > 1 116 | else (*image_size,) 117 | ) 118 | 119 | # Initialize the image. 120 | _image = np.full(shape=_shape, fill_value=-1, dtype=channel.dtype) 121 | 122 | # Assign the values. 123 | _image[proj_row, proj_column] = channel[order] 124 | projected_channels.append(_image) 125 | output["channels"] = projected_channels 126 | 127 | # Get the inverse indices mapping. 128 | list_indices_img = np.full(image_size, -1, dtype=np.int32) 129 | list_indices_img[proj_row, proj_column] = indices 130 | output["inverse"] = list_indices_img 131 | 132 | # Set which points are used in the projection. 133 | active_list = np.full(depth_list.shape, fill_value=0, dtype=np.int32) 134 | active_list[list_indices_img] = 1 135 | output["active"] = active_list.astype(np.bool) 136 | 137 | return output 138 | -------------------------------------------------------------------------------- /lidardet/datasets/augmentor/data_augmentor.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from functools import partial 4 | 5 | from lidardet.utils.geometry import limit_period 6 | from lidardet.utils.common import check_numpy_to_torch, scan_downsample, scan_upsample 7 | from . import augmentor_utils, database_sampler 8 | 9 | class DataAugmentor(object): 10 | def __init__(self, root_path, augmentor_configs, class_names, logger=None): 11 | self.root_path = root_path 12 | self.class_names = class_names 13 | self.logger = logger 14 | 15 | self.data_augmentor_queue = [] 16 | for cur_cfg in augmentor_configs: 17 | cur_augmentor = getattr(self, cur_cfg['name'])(config=cur_cfg) 18 | self.data_augmentor_queue.append(cur_augmentor) 19 | 20 | def gt_sampling(self, config=None): 21 | db_sampler = database_sampler.DataBaseSampler( 22 | root_path=self.root_path, 23 | sampler_cfg=config, 24 | class_names=self.class_names, 25 | logger=self.logger 26 | ) 27 | return db_sampler 28 | 29 | def __getstate__(self): 30 | d = dict(self.__dict__) 31 | del d['logger'] 32 | return d 33 | 34 | def __setstate__(self, d): 35 | self.__dict__.update(d) 36 | 37 | def scan_down_up(self, data_dict=None, config=None): 38 | if data_dict is None: 39 | return partial(self.scan_down_up, config=config) 40 | points = data_dict['points'] 41 | if config.random: 42 | enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5]) 43 | if enable: 44 | points_lr = scan_downsample(points, output_rings = 'even_or_odd') 45 | points_hr = scan_upsample(points_lr) 46 | data_dict['points'] = points_hr 47 | else: 48 | points_lr = scan_downsample(points, output_rings = 'even_or_odd') 49 | points_hr = scan_upsample(points_lr) 50 | data_dict['points'] = points_hr 51 | 52 | return data_dict 53 | 54 | def scan_downsample(self, data_dict=None, config=None): 55 | if data_dict is None: 56 | return partial(self.scan_downsample, config=config) 57 | points = data_dict['points'] 58 | if config.random: 59 | enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5]) 60 | if enable: 61 | points_lr = scan_downsample(points, output_rings = config.output_rings) 62 | data_dict['points'] = points_lr 63 | else: 64 | points_lr = scan_downsample(points, output_rings = config.output_rings) 65 | data_dict['points'] = points_lr 66 | 67 | return data_dict 68 | 69 | def random_image_flip(self, data_dict=None, config=None): 70 | if data_dict is None: 71 | return partial(self.random_image_flip, config=config) 72 | 73 | enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5]) 74 | if enable: 75 | img = data_dict['range_image_in'][:,:,::-1] 76 | gt = data_dict['range_image_gt'][:,:,::-1] 77 | data_dict['range_image_in'] = img 78 | data_dict['range_image_gt'] = gt 79 | 80 | return data_dict 81 | 82 | def random_noise(self, data_dict=None, config=None): 83 | if data_dict is None: 84 | return partial(self.random_noise, config=config) 85 | 86 | enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5]) 87 | if enable: 88 | img = data_dict['range_image_in'] 89 | noise = np.random.normal(0, config.sigma, img.shape) # mu, sigma, size 90 | for i in range(img.shape[0]): 91 | noise[i, img[0] == 0] = 0 92 | 93 | img += noise 94 | data_dict['range_image_in'] = img 95 | 96 | return data_dict 97 | 98 | def random_world_flip(self, data_dict=None, config=None): 99 | if data_dict is None: 100 | return partial(self.random_world_flip, config=config) 101 | gt_boxes, points = data_dict['gt_boxes'], data_dict['points'] 102 | for cur_axis in config['axis']: 103 | assert cur_axis in ['x', 'y'] 104 | gt_boxes, points = getattr(augmentor_utils, 'random_flip_along_%s' % cur_axis)( 105 | gt_boxes, points, 106 | ) 107 | 108 | data_dict['gt_boxes'] = gt_boxes 109 | data_dict['points'] = points 110 | return data_dict 111 | 112 | 113 | def random_world_rotation(self, data_dict=None, config=None): 114 | if data_dict is None: 115 | return partial(self.random_world_rotation, config=config) 116 | rot_range = config['angle'] 117 | if not isinstance(rot_range, list): 118 | rot_range = [-rot_range, rot_range] 119 | gt_boxes, points = augmentor_utils.global_rotation( 120 | data_dict['gt_boxes'], data_dict['points'], rot_range=rot_range 121 | ) 122 | 123 | data_dict['gt_boxes'] = gt_boxes 124 | data_dict['points'] = points 125 | return data_dict 126 | 127 | 128 | def random_world_scaling(self, data_dict=None, config=None): 129 | if data_dict is None: 130 | return partial(self.random_world_scaling, config=config) 131 | gt_boxes, points = augmentor_utils.global_scaling( 132 | data_dict['gt_boxes'], data_dict['points'], config['scale'] 133 | ) 134 | data_dict['gt_boxes'] = gt_boxes 135 | data_dict['points'] = points 136 | return data_dict 137 | 138 | def forward(self, data_dict): 139 | """ 140 | Args: 141 | data_dict: 142 | points: (N, 3 + C_in) 143 | gt_boxes: optional, (N, 7) [x, y, z, dx, dy, dz, heading] 144 | gt_names: optional, (N), string 145 | ... 146 | 147 | Returns: 148 | """ 149 | for cur_augmentor in self.data_augmentor_queue: 150 | data_dict = cur_augmentor(data_dict=data_dict) 151 | 152 | if 'gt_boxes' in data_dict: 153 | data_dict['gt_boxes'][:, 6] = limit_period( 154 | data_dict['gt_boxes'][:, 6], offset=0.5, period=2 * np.pi 155 | ) 156 | if 'calib' in data_dict: 157 | data_dict.pop('calib') 158 | if 'road_plane' in data_dict: 159 | data_dict.pop('road_plane') 160 | if 'gt_boxes_mask' in data_dict: 161 | gt_boxes_mask = data_dict['gt_boxes_mask'] 162 | data_dict['gt_boxes'] = data_dict['gt_boxes'][gt_boxes_mask] 163 | data_dict['gt_names'] = data_dict['gt_names'][gt_boxes_mask] 164 | data_dict.pop('gt_boxes_mask') 165 | return data_dict 166 | -------------------------------------------------------------------------------- /lidardet/models/trajectory_predictor/backbone.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | 4 | def conv3x3(in_planes, out_planes, stride=1, bias=False): 5 | """3x3 convolution with padding""" 6 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 7 | padding=1, bias=bias) 8 | 9 | class BasicBlock(nn.Module): 10 | expansion = 1 11 | 12 | def __init__(self, in_planes, planes, stride=1, downsample=None): 13 | super(BasicBlock, self).__init__() 14 | self.conv1 = conv3x3(in_planes, planes, stride, bias=True) 15 | self.bn1 = nn.BatchNorm2d(planes) 16 | self.relu = nn.ReLU(inplace=True) 17 | self.conv2 = conv3x3(planes, planes, bias=True) 18 | self.bn2 = nn.BatchNorm2d(planes) 19 | self.downsample = downsample 20 | self.stride = stride 21 | 22 | def forward(self, x): 23 | residual = x 24 | 25 | out = self.conv1(x) 26 | #out = self.bn1(out) 27 | out = self.relu(out) 28 | 29 | out = self.conv2(out) 30 | #out = self.bn2(out) 31 | 32 | if self.downsample is not None: 33 | residual = self.downsample(x) 34 | 35 | out += residual 36 | out = self.relu(out) 37 | 38 | return out 39 | 40 | class Bottleneck(nn.Module): 41 | expansion = 4 42 | 43 | def __init__(self, in_planes, planes, stride=1, downsample=None, use_bn=True): 44 | super(Bottleneck, self).__init__() 45 | bias = not use_bn 46 | self.use_bn = use_bn 47 | self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=bias) 48 | self.bn1 = nn.BatchNorm2d(planes) 49 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=bias) 50 | self.bn2 = nn.BatchNorm2d(planes) 51 | self.conv3 = nn.Conv2d(planes, self.expansion*planes, kernel_size=1, bias=bias) 52 | self.bn3 = nn.BatchNorm2d(self.expansion*planes) 53 | self.downsample = downsample 54 | self.relu = nn.ReLU(inplace=True) 55 | 56 | def forward(self, x): 57 | residual = x 58 | out = self.conv1(x) 59 | if self.use_bn: 60 | out = self.bn1(out) 61 | out = self.relu(out) 62 | 63 | out = self.conv2(out) 64 | if self.use_bn: 65 | out = self.bn2(out) 66 | out = self.relu(out) 67 | out = self.conv3(out) 68 | if self.use_bn: 69 | out = self.bn3(out) 70 | 71 | if self.downsample is not None: 72 | residual = self.downsample(x) 73 | out = self.relu(residual + out) 74 | return out 75 | 76 | class BackBone(nn.Module): 77 | 78 | def __init__(self, block, cfg): 79 | super(BackBone, self).__init__() 80 | self.use_bn = cfg.use_bn 81 | num_block = cfg.num_blocks 82 | 83 | # Block 1 84 | self.conv1 = conv3x3(cfg.input_dim, 32) 85 | self.conv2 = conv3x3(32, 32) 86 | self.bn1 = nn.BatchNorm2d(32) 87 | self.bn2 = nn.BatchNorm2d(32) 88 | self.relu = nn.ReLU(inplace=True) 89 | 90 | 91 | # Block 2-5 92 | self.in_planes = 32 93 | self.block2 = self._make_layer(block, 24, num_blocks=num_block[0]) 94 | self.block3 = self._make_layer(block, 48, num_blocks=num_block[1]) 95 | self.block4 = self._make_layer(block, 64, num_blocks=num_block[2]) 96 | self.block5 = self._make_layer(block, 96, num_blocks=num_block[3]) 97 | 98 | # Lateral layers 99 | self.latlayer1 = nn.Conv2d(384, 196, kernel_size=1, stride=1, padding=0) 100 | self.latlayer2 = nn.Conv2d(256, 128, kernel_size=1, stride=1, padding=0) 101 | self.latlayer3 = nn.Conv2d(192, 96, kernel_size=1, stride=1, padding=0) 102 | 103 | # Top-down layers 104 | self.deconv1 = nn.ConvTranspose2d(196, 128, kernel_size=3, stride=2, padding=1, output_padding=1) 105 | self.deconv2 = nn.ConvTranspose2d(128, 96, kernel_size=3, stride=2, padding=1, output_padding=(0, 1)) 106 | 107 | def save_grad(self, grad): 108 | self.grad = grad 109 | 110 | def forward(self, batch_dict): 111 | x = batch_dict['input'] 112 | x = self.conv1(x) 113 | if self.use_bn: 114 | x = self.bn1(x) 115 | x = self.relu(x) 116 | 117 | x = self.conv2(x) 118 | if self.use_bn: 119 | x = self.bn2(x) 120 | c1 = self.relu(x) 121 | 122 | # bottom up layers 123 | c2 = self.block2(c1) 124 | c3 = self.block3(c2) 125 | c4 = self.block4(c3) 126 | c5 = self.block5(c4) 127 | 128 | # for Grad-CAM visualization 129 | #c5.register_hook(self.save_grad) 130 | #self.activations = c5 131 | 132 | l5 = self.latlayer1(c5) 133 | l4 = self.latlayer2(c4) 134 | p5 = l4 + self.deconv1(l5) 135 | l3 = self.latlayer3(c3) 136 | p5d = self.deconv2(p5) 137 | p4 = l3 + self.deconv2(p5) 138 | 139 | batch_dict['seg_features'] = p4 140 | batch_dict['reg_features'] = c5 141 | return batch_dict 142 | 143 | def _make_layer(self, block, planes, num_blocks): 144 | if self.use_bn: 145 | downsample = nn.Sequential( 146 | nn.Conv2d(self.in_planes, planes * block.expansion, 147 | kernel_size=1, stride=2, bias=False), 148 | nn.BatchNorm2d(planes * block.expansion) 149 | ) 150 | else: 151 | downsample = nn.Conv2d(self.in_planes, planes * block.expansion, 152 | kernel_size=1, stride=2, bias=True) 153 | 154 | layers = [] 155 | layers.append(block(self.in_planes, planes, stride=2, downsample=downsample)) 156 | self.in_planes = planes * block.expansion 157 | for i in range(1, num_blocks): 158 | layers.append(block(self.in_planes, planes, stride=1)) 159 | self.in_planes = planes * block.expansion 160 | return nn.Sequential(*layers) 161 | 162 | def _upsample_add(self, x, y): 163 | '''Upsample and add two feature maps. 164 | Args: 165 | x: (Variable) top feature map to be upsampled. 166 | y: (Variable) lateral feature map. 167 | Returns: 168 | (Variable) added feature map. 169 | Note in PyTorch, when input size is odd, the upsampled feature map 170 | with `F.upsample(..., scale_factor=2, mode='nearest')` 171 | maybe not equal to the lateral feature map size. 172 | e.g. 173 | original input size: [N,_,15,15] -> 174 | conv2d feature map size: [N,_,8,8] -> 175 | upsampled feature map size: [N,_,16,16] 176 | So we choose bilinear upsample which supports arbitrary output sizes. 177 | ''' 178 | _, _, H, W = y.size() 179 | return F.upsample(x, size=(H, W), mode='bilinear') + y 180 | -------------------------------------------------------------------------------- /lidardet/losses/loss_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from lidardet.utils import box_utils 6 | 7 | 8 | class SigmoidFocalClassificationLoss(nn.Module): 9 | """ 10 | Sigmoid focal cross entropy loss. 11 | """ 12 | 13 | def __init__(self, gamma: float = 2.0, alpha: float = 0.25): 14 | """ 15 | Args: 16 | gamma: Weighting parameter to balance loss for hard and easy examples. 17 | alpha: Weighting parameter to balance loss for positive and negative examples. 18 | """ 19 | super(SigmoidFocalClassificationLoss, self).__init__() 20 | self.alpha = alpha 21 | self.gamma = gamma 22 | 23 | @staticmethod 24 | def sigmoid_cross_entropy_with_logits(input: torch.Tensor, target: torch.Tensor): 25 | """ PyTorch Implementation for tf.nn.sigmoid_cross_entropy_with_logits: 26 | max(x, 0) - x * z + log(1 + exp(-abs(x))) in 27 | https://www.tensorflow.org/api_docs/python/tf/nn/sigmoid_cross_entropy_with_logits 28 | 29 | Args: 30 | input: (B, #anchors, #classes) float tensor. 31 | Predicted logits for each class 32 | target: (B, #anchors, #classes) float tensor. 33 | One-hot encoded classification targets 34 | 35 | Returns: 36 | loss: (B, #anchors, #classes) float tensor. 37 | Sigmoid cross entropy loss without reduction 38 | """ 39 | loss = torch.clamp(input, min=0) - input * target + \ 40 | torch.log1p(torch.exp(-torch.abs(input))) 41 | return loss 42 | 43 | def forward(self, input: torch.Tensor, target: torch.Tensor, weights: torch.Tensor): 44 | """ 45 | Args: 46 | input: (B, #anchors, #classes) float tensor. 47 | Predicted logits for each class 48 | target: (B, #anchors, #classes) float tensor. 49 | One-hot encoded classification targets 50 | weights: (B, #anchors) float tensor. 51 | Anchor-wise weights. 52 | 53 | Returns: 54 | weighted_loss: (B, #anchors, #classes) float tensor after weighting. 55 | """ 56 | pred_sigmoid = torch.sigmoid(input) 57 | alpha_weight = target * self.alpha + (1 - target) * (1 - self.alpha) 58 | pt = target * (1.0 - pred_sigmoid) + (1.0 - target) * pred_sigmoid 59 | focal_weight = alpha_weight * torch.pow(pt, self.gamma) 60 | 61 | bce_loss = self.sigmoid_cross_entropy_with_logits(input, target) 62 | 63 | loss = focal_weight * bce_loss 64 | 65 | if weights.shape.__len__() == 2 or \ 66 | (weights.shape.__len__() == 1 and target.shape.__len__() == 2): 67 | weights = weights.unsqueeze(-1) 68 | 69 | assert weights.shape.__len__() == loss.shape.__len__() 70 | 71 | return loss * weights 72 | 73 | 74 | class WeightedSmoothL1Loss(nn.Module): 75 | """ 76 | Code-wise Weighted Smooth L1 Loss modified based on fvcore.nn.smooth_l1_loss 77 | https://github.com/facebookresearch/fvcore/blob/master/fvcore/nn/smooth_l1_loss.py 78 | | 0.5 * x ** 2 / beta if abs(x) < beta 79 | smoothl1(x) = | 80 | | abs(x) - 0.5 * beta otherwise, 81 | where x = input - target. 82 | """ 83 | def __init__(self, beta: float = 1.0 / 9.0, code_weights: list = None): 84 | """ 85 | Args: 86 | beta: Scalar float. 87 | L1 to L2 change point. 88 | For beta values < 1e-5, L1 loss is computed. 89 | code_weights: (#codes) float list if not None. 90 | Code-wise weights. 91 | """ 92 | super(WeightedSmoothL1Loss, self).__init__() 93 | self.beta = beta 94 | if code_weights is not None: 95 | self.code_weights = np.array(code_weights, dtype=np.float32) 96 | self.code_weights = torch.from_numpy(self.code_weights).cuda() 97 | 98 | @staticmethod 99 | def smooth_l1_loss(diff, beta): 100 | if beta < 1e-5: 101 | loss = torch.abs(diff) 102 | else: 103 | n = torch.abs(diff) 104 | loss = torch.where(n < beta, 0.5 * n ** 2 / beta, n - 0.5 * beta) 105 | 106 | return loss 107 | 108 | def forward(self, input: torch.Tensor, target: torch.Tensor, weights: torch.Tensor = None): 109 | """ 110 | Args: 111 | input: (B, #anchors, #codes) float tensor. 112 | Ecoded predicted locations of objects. 113 | target: (B, #anchors, #codes) float tensor. 114 | Regression targets. 115 | weights: (B, #anchors) float tensor if not None. 116 | 117 | Returns: 118 | loss: (B, #anchors) float tensor. 119 | Weighted smooth l1 loss without reduction. 120 | """ 121 | diff = input - target 122 | # code-wise weighting 123 | if self.code_weights is not None: 124 | diff = diff * self.code_weights.view(1, 1, -1) 125 | 126 | loss = self.smooth_l1_loss(diff, self.beta) 127 | 128 | # anchor-wise weighting 129 | if weights is not None: 130 | assert weights.shape[0] == loss.shape[0] and weights.shape[1] == loss.shape[1] 131 | loss = loss * weights.unsqueeze(-1) 132 | 133 | return loss 134 | 135 | 136 | class WeightedCrossEntropyLoss(nn.Module): 137 | """ 138 | Transform input to fit the fomation of PyTorch offical cross entropy loss 139 | with anchor-wise weighting. 140 | """ 141 | def __init__(self): 142 | super(WeightedCrossEntropyLoss, self).__init__() 143 | 144 | def forward(self, input: torch.Tensor, target: torch.Tensor, weights: torch.Tensor): 145 | """ 146 | Args: 147 | input: (B, #anchors, #classes) float tensor. 148 | Predited logits for each class. 149 | target: (B, #anchors, #classes) float tensor. 150 | One-hot classification targets. 151 | weights: (B, #anchors) float tensor. 152 | Anchor-wise weights. 153 | 154 | Returns: 155 | loss: (B, #anchors) float tensor. 156 | Weighted cross entropy loss without reduction 157 | """ 158 | input = input.permute(0, 2, 1) 159 | target = target.argmax(dim=-1) 160 | loss = F.cross_entropy(input, target, reduction='none') * weights 161 | return loss 162 | 163 | 164 | def get_corner_loss_lidar(pred_bbox3d: torch.Tensor, gt_bbox3d: torch.Tensor): 165 | """ 166 | Args: 167 | pred_bbox3d: (N, 7) float Tensor. 168 | gt_bbox3d: (N, 7) float Tensor. 169 | 170 | Returns: 171 | corner_loss: (N) float Tensor. 172 | """ 173 | assert pred_bbox3d.shape[0] == gt_bbox3d.shape[0] 174 | 175 | pred_box_corners = box_utils.boxes_to_corners_3d(pred_bbox3d) 176 | gt_box_corners = box_utils.boxes_to_corners_3d(gt_bbox3d) 177 | 178 | gt_bbox3d_flip = gt_bbox3d.clone() 179 | gt_bbox3d_flip[:, 6] += np.pi 180 | gt_box_corners_flip = box_utils.boxes_to_corners_3d(gt_bbox3d_flip) 181 | # (N, 8) 182 | corner_dist = torch.min(torch.norm(pred_box_corners - gt_box_corners, dim=2), 183 | torch.norm(pred_box_corners - gt_box_corners_flip, dim=2)) 184 | # (N, 8) 185 | corner_loss = WeightedSmoothL1Loss.smooth_l1_loss(corner_dist, beta=1.0) 186 | 187 | return corner_loss.mean(dim=1) 188 | -------------------------------------------------------------------------------- /lidardet/utils/find.py: -------------------------------------------------------------------------------- 1 | import glob 2 | import json 3 | import os 4 | import subprocess 5 | import sys 6 | import tempfile 7 | from pathlib import Path 8 | 9 | def _get_info_from_anaconda_info(info, split=":"): 10 | info = info.strip("\n").replace(" ", "") 11 | info_dict = {} 12 | latest_key = "" 13 | for line in info.splitlines(): 14 | if split in line: 15 | pair = line.split(split) 16 | info_dict[pair[0]] = pair[1] 17 | latest_key = pair[0] 18 | else: 19 | if not isinstance(info_dict[latest_key], list): 20 | info_dict[latest_key] = [info_dict[latest_key]] 21 | info_dict[latest_key].append(line) 22 | return info_dict 23 | 24 | 25 | def find_anaconda(): 26 | # try find in default path 27 | path = Path.home() / "anaconda3" 28 | if path.exists(): 29 | return path 30 | # try conda in cmd 31 | try: 32 | info = subprocess.check_output("conda info", shell=True).decode("utf-8") 33 | info_dict = _get_info_from_anaconda_info(info) 34 | return info_dict["activeenvlocation"] 35 | except subprocess.CalledProcessError: 36 | raise RuntimeError("find anadonda failed") 37 | 38 | 39 | def find_cuda(): 40 | """Finds the CUDA install path.""" 41 | # Guess #1 42 | cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") 43 | if cuda_home is None: 44 | # Guess #2 45 | if sys.platform == "win32": 46 | cuda_homes = glob.glob( 47 | "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v*.*" 48 | ) 49 | if len(cuda_homes) == 0: 50 | cuda_home = "" 51 | else: 52 | cuda_home = cuda_homes[0] 53 | else: 54 | cuda_home = "/usr/local/cuda" 55 | if not os.path.exists(cuda_home): 56 | # Guess #3 57 | try: 58 | which = "where" if sys.platform == "win32" else "which" 59 | nvcc = subprocess.check_output([which, "nvcc"]).decode().rstrip("\r\n") 60 | cuda_home = os.path.dirname(os.path.dirname(nvcc)) 61 | except Exception: 62 | cuda_home = None 63 | if cuda_home is None: 64 | raise RuntimeError( 65 | "No CUDA runtime is found, using CUDA_HOME='{}'".format(cuda_home) 66 | ) 67 | return cuda_home 68 | 69 | 70 | def find_cuda_device_arch(): 71 | if sys.platform == "win32": 72 | # TODO: add windows support 73 | return None 74 | cuda_home = find_cuda() 75 | if cuda_home is None: 76 | return None 77 | cuda_home = Path(cuda_home) 78 | try: 79 | device_query_path = cuda_home / "extras/demo_suite/deviceQuery" 80 | if not device_query_path.exists(): 81 | source = """ 82 | #include 83 | #include 84 | int main(){ 85 | int nDevices; 86 | cudaGetDeviceCount(&nDevices); 87 | for (int i = 0; i < nDevices; i++) { 88 | cudaDeviceProp prop; 89 | cudaGetDeviceProperties(&prop, i); 90 | std::cout << prop.major << "." << prop.minor << std::endl; 91 | } 92 | return 0; 93 | } 94 | """ 95 | with tempfile.NamedTemporaryFile("w", suffix=".cc") as f: 96 | f_path = Path(f.name) 97 | f.write(source) 98 | f.flush() 99 | try: 100 | # TODO: add windows support 101 | cmd = ( 102 | f"g++ {f.name} -o {f_path.stem}" 103 | f" -I{cuda_home / 'include'} -L{cuda_home / 'lib64'} -lcudart" 104 | ) 105 | print(cmd) 106 | subprocess.check_output(cmd, shell=True, cwd=f_path.parent) 107 | cmd = f"./{f_path.stem}" 108 | arches = ( 109 | subprocess.check_output(cmd, shell=True, cwd=f_path.parent) 110 | .decode() 111 | .rstrip("\r\n") 112 | .split("\n") 113 | ) 114 | if len(arches) < 1: 115 | return None 116 | arch = arches[0] 117 | except Exception: 118 | return None 119 | else: 120 | cmd = f"{str(device_query_path)} | grep 'CUDA Capability'" 121 | arch = ( 122 | subprocess.check_output(cmd, shell=True) 123 | .decode() 124 | .rstrip("\r\n") 125 | .split(" ")[-1] 126 | ) 127 | # assert len(arch) == 2 128 | arch_list = [int(s) for s in arch.split(".")] 129 | arch_int = arch_list[0] * 10 + arch_list[1] 130 | find_work_arch = False 131 | while arch_int > 10: 132 | try: 133 | res = subprocess.check_output( 134 | "nvcc -arch=sm_{}".format(arch_int), 135 | shell=True, 136 | stderr=subprocess.STDOUT, 137 | ) 138 | except subprocess.CalledProcessError as e: 139 | if "No input files specified" in e.output.decode(): 140 | find_work_arch = True 141 | break 142 | elif ( 143 | "is not defined for option 'gpu-architecture'" in e.output.decode() 144 | ): 145 | arch_int -= 1 146 | else: 147 | raise RuntimeError("unknown error") 148 | if find_work_arch: 149 | arch = f"sm_{arch_int}" 150 | else: 151 | arch = None 152 | 153 | except Exception: 154 | arch = None 155 | return arch 156 | 157 | 158 | def get_gpu_memory_usage(): 159 | if sys.platform == "win32": 160 | # TODO: add windows support 161 | return None 162 | cuda_home = find_cuda() 163 | if cuda_home is None: 164 | return None 165 | cuda_home = Path(cuda_home) 166 | source = """ 167 | #include 168 | #include 169 | int main(){ 170 | int nDevices; 171 | cudaGetDeviceCount(&nDevices); 172 | size_t free_m, total_m; 173 | // output json format. 174 | std::cout << "["; 175 | for (int i = 0; i < nDevices; i++) { 176 | cudaSetDevice(i); 177 | cudaMemGetInfo(&free_m, &total_m); 178 | std::cout << "[" << free_m << "," << total_m << "]"; 179 | if (i != nDevices - 1) 180 | std::cout << "," << std::endl; 181 | } 182 | std::cout << "]" << std::endl; 183 | return 0; 184 | } 185 | """ 186 | with tempfile.NamedTemporaryFile("w", suffix=".cc") as f: 187 | f_path = Path(f.name) 188 | f.write(source) 189 | f.flush() 190 | try: 191 | # TODO: add windows support 192 | cmd = ( 193 | f"g++ {f.name} -o {f_path.stem} -std=c++11" 194 | f" -I{cuda_home / 'include'} -L{cuda_home / 'lib64'} -lcudart" 195 | ) 196 | print(cmd) 197 | subprocess.check_output(cmd, shell=True, cwd=f_path.parent) 198 | cmd = f"./{f_path.stem}" 199 | usages = subprocess.check_output( 200 | cmd, shell=True, cwd=f_path.parent 201 | ).decode() 202 | usages = json.loads(usages) 203 | return usages 204 | except Exception: 205 | return None 206 | return None 207 | 208 | 209 | if __name__ == "__main__": 210 | print(find_cuda_device_arch()) 211 | -------------------------------------------------------------------------------- /lidardet/utils/projection.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | class PointProjection(object): 4 | 5 | def __init__(self, cfg): 6 | self.cfg = cfg 7 | 8 | @staticmethod 9 | def xyz_to_cylinder(points_xyz): 10 | points_x, points_y, points_z = np.split(points_xyz, 3, axis=-1) 11 | points_rho = np.sqrt(points_x**2 + points_y**2) 12 | points_phi = np.arctan2(points_y, points_x) 13 | points_cylinder = np.stack([points_phi, points_z, points_rho], axis=-1) 14 | return points_cylinder 15 | 16 | @staticmethod 17 | def cylinder_to_xyz(points_cylinder): 18 | points_phi, points_z, points_rho = np.split(points_cylinder, 3, axis=-1) 19 | points_x = points_rho * np.cos(points_phi) 20 | points_y = points_rho * np.sin(points_phi) 21 | points_xyz = np.stack([points_x, points_y, points_z], axis=-1) 22 | return points_xyz 23 | 24 | def get_kitti_columns(points: np.array, number_of_columns: int = 2000) -> np.array: 25 | """ Returns the column indices for unfolding one or more raw KITTI scans """ 26 | azi = np.arctan2(points[..., 1], points[..., 0]) 27 | cols = number_of_columns * (np.pi - azi) / (2 * np.pi) 28 | # In rare cases it happens that the index is exactly one too large. 29 | cols = np.minimum(cols, number_of_columns - 1) 30 | return np.int32(cols) 31 | 32 | 33 | def get_kitti_rows(points: np.array, threshold: float = -0.005) -> np.array: 34 | """ Returns the row indices for unfolding one or more raw KITTI scans """ 35 | azimuth_flipped = -np.arctan2(points[..., 1], -points[..., 0]) 36 | azi_diffs = azimuth_flipped[..., 1:] - azimuth_flipped[..., :-1] 37 | jump_mask = np.greater(threshold, azi_diffs) 38 | ind = np.add(np.where(jump_mask), 1) 39 | rows = np.zeros_like(points[..., 0]) 40 | rows[..., ind] += 1 41 | return np.int32(np.cumsum(rows, axis=-1)) 42 | 43 | 44 | def projection(points: np.array, *channels, image_size: tuple = (64, 2000)): 45 | """ 46 | Scan unfolding of raw KITTI point cloud of shape [N, (x,y,z)] 47 | This functions performs a cylindrical projection of a 3D point cloud given in 48 | cartesian coordinates. The function aims to recover the internal data structure 49 | of the KITTI HDL-64 sensor from its distinct scan pattern that can be approximated 50 | from the original KITTI Raw Dataset. 51 | The resulting projected image (H, W) contains missing values, these positions 52 | are filled with `-1`. If you set `return_inverse` to True, you may receive 53 | a mask for the valid measurements of the image with `valid_image = inverse >= 0`. 54 | This function has less information loss than techniques which compute the 55 | azimuth and elevation angle and sort the points accordingly into the image. 56 | Therefore, most entries in `valid_image` shall be True, except for locations 57 | that hit the ego vehicle or are pointed towards the sky. 58 | The measurements are ordered in decreasing depth before being assigned to image 59 | locations. In case more than one value is scattered to this position, the one 60 | closest to the ego vehicle is selected. To know which points actually made it into 61 | the projection, set `return_active` to True. 62 | Arguments: 63 | points : np.array(shape=(N, 3), dtype=np.float) 64 | The raw KITTI point cloud (does not work with ego-motion corrected data!). 65 | image_size : tuple of 2 integers 66 | The image size of the projection in (H, W). (Default is (64, 2000). 67 | channels : * times np.array(shape=(N, D)) 68 | Additional channels to project in the same way as `points`. 69 | Returns: 70 | Dict( 71 | points=np.array(shape=(H, W, 3), dtype=np.float), # projected point cloud 72 | depth=np.array(shape=(N,), dtype=np.float) # projected depth 73 | channels=List(*), # list of projected additional channels 74 | indices=np.array(shape=(N, 2), dtype=np.int) # image indices from the list 75 | inverse=np.array(shape=(H, W), dtype=np.int) # list indices from the image 76 | active=np.array(shape=(N,), dtype=np.bool) # active array 77 | Raises: 78 | IndexError if projection is not possible, e.g. if you do not use a KITTI 79 | point cloud, or the KITTI point cloud is ego motion corrected 80 | """ 81 | 82 | output = {} 83 | 84 | assert points.shape[1] == 3, "Points must contain N xyz coordinates." 85 | if len(channels) > 0: 86 | assert all( 87 | isinstance(x, np.ndarray) for x in channels 88 | ), "All channels must be numpy arrays." 89 | assert all( 90 | x.shape[0] == points.shape[0] for x in channels 91 | ), "All channels must have the same first dimension as `points`." 92 | 93 | # Get depth of all points for ordering. 94 | depth_list = np.linalg.norm(points, 2, axis=1) 95 | 96 | # Get the indices of the rows and columns to project to. 97 | proj_column = get_kitti_columns(points, number_of_columns=image_size[1]) 98 | proj_row = get_kitti_rows(points) 99 | 100 | if np.any(proj_row >= image_size[0]) or np.any(proj_column >= image_size[1]): 101 | raise IndexError( 102 | "Cannot find valid image indices for this point cloud and image size. " 103 | "Are you sure you entered the correct image size? This function only works " 104 | "with raw KITTI HDL-64 point clouds (no ego motion corrected data allowed)!" 105 | ) 106 | 107 | # Store a copy in original order. 108 | output["indices"] = np.stack([np.copy(proj_row), np.copy(proj_column)], axis=-1) 109 | 110 | # Get the indices in order of decreasing depth. 111 | indices = np.arange(depth_list.shape[0]) 112 | order = np.argsort(depth_list)[::-1] 113 | 114 | indices = indices[order] 115 | proj_column = proj_column[order] 116 | proj_row = proj_row[order] 117 | 118 | # Project the points. 119 | points_img = np.full(shape=(*image_size, 3), fill_value=-1, dtype=np.float32) 120 | points_img[proj_row, proj_column] = points[order] 121 | output["points"] = points_img 122 | 123 | # The depth projection. 124 | depth_img = np.full(shape=image_size, fill_value=-1, dtype=np.float32) 125 | depth_img[proj_row, proj_column] = depth_list[order] 126 | output["depth"] = depth_img 127 | 128 | # Convert all channels. 129 | projected_channels = [] 130 | for channel in channels: 131 | # Set the shape. 132 | _shape = ( 133 | (*image_size, channel.shape[1]) 134 | if len(channel.shape) > 1 135 | else (*image_size,) 136 | ) 137 | 138 | # Initialize the image. 139 | _image = np.full(shape=_shape, fill_value=-1, dtype=channel.dtype) 140 | 141 | # Assign the values. 142 | _image[proj_row, proj_column] = channel[order] 143 | projected_channels.append(_image) 144 | output["channels"] = projected_channels 145 | 146 | # Get the inverse indices mapping. 147 | list_indices_img = np.full(image_size, -1, dtype=np.int32) 148 | list_indices_img[proj_row, proj_column] = indices 149 | output["inverse"] = list_indices_img 150 | 151 | # Set which points are used in the projection. 152 | active_list = np.full(depth_list.shape, fill_value=0, dtype=np.int32) 153 | active_list[list_indices_img] = 1 154 | output["active"] = active_list.astype(np.bool) 155 | 156 | return output 157 | -------------------------------------------------------------------------------- /lidardet/utils/input/point_cloud/point_cloud_ops.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | import numba 4 | import numpy as np 5 | 6 | 7 | @numba.jit(nopython=True) 8 | def _points_to_voxel_reverse_kernel( 9 | points, 10 | voxel_size, 11 | coors_range, 12 | num_points_per_voxel, 13 | coor_to_voxelidx, 14 | voxels, 15 | coors, 16 | max_points=35, 17 | max_voxels=20000, 18 | ): 19 | # put all computations to one loop. 20 | # we shouldn't create large array in main jit code, otherwise 21 | # reduce performance 22 | N = points.shape[0] 23 | # ndim = points.shape[1] - 1 24 | ndim = 3 25 | ndim_minus_1 = ndim - 1 26 | grid_size = (coors_range[3:] - coors_range[:3]) / voxel_size 27 | # np.round(grid_size) 28 | # grid_size = np.round(grid_size).astype(np.int64)(np.int32) 29 | grid_size = np.round(grid_size, 0, grid_size).astype(np.int32) 30 | coor = np.zeros(shape=(3,), dtype=np.int32) 31 | voxel_num = 0 32 | failed = False 33 | for i in range(N): 34 | failed = False 35 | for j in range(ndim): 36 | c = np.floor((points[i, j] - coors_range[j]) / voxel_size[j]) 37 | if c < 0 or c >= grid_size[j]: 38 | failed = True 39 | break 40 | coor[ndim_minus_1 - j] = c 41 | if failed: 42 | continue 43 | voxelidx = coor_to_voxelidx[coor[0], coor[1], coor[2]] 44 | if voxelidx == -1: 45 | voxelidx = voxel_num 46 | if voxel_num >= max_voxels: 47 | break 48 | voxel_num += 1 49 | coor_to_voxelidx[coor[0], coor[1], coor[2]] = voxelidx 50 | coors[voxelidx] = coor 51 | num = num_points_per_voxel[voxelidx] 52 | if num < max_points: 53 | voxels[voxelidx, num] = points[i] 54 | num_points_per_voxel[voxelidx] += 1 55 | return voxel_num 56 | 57 | 58 | @numba.jit(nopython=True) 59 | def _points_to_voxel_kernel( 60 | points, 61 | voxel_size, 62 | coors_range, 63 | num_points_per_voxel, 64 | coor_to_voxelidx, 65 | voxels, 66 | coors, 67 | max_points=35, 68 | max_voxels=20000, 69 | ): 70 | # need mutex if write in cuda, but numba.cuda don't support mutex. 71 | # in addition, pytorch don't support cuda in dataloader(tensorflow support this). 72 | # put all computations to one loop. 73 | # we shouldn't create large array in main jit code, otherwise 74 | # decrease performance 75 | N = points.shape[0] 76 | # ndim = points.shape[1] - 1 77 | ndim = 3 78 | grid_size = (coors_range[3:] - coors_range[:3]) / voxel_size 79 | # grid_size = np.round(grid_size).astype(np.int64)(np.int32) 80 | grid_size = np.round(grid_size, 0, grid_size).astype(np.int32) 81 | 82 | lower_bound = coors_range[:3] 83 | upper_bound = coors_range[3:] 84 | coor = np.zeros(shape=(3,), dtype=np.int32) 85 | voxel_num = 0 86 | failed = False 87 | for i in range(N): 88 | failed = False 89 | for j in range(ndim): 90 | c = np.floor((points[i, j] - coors_range[j]) / voxel_size[j]) 91 | if c < 0 or c >= grid_size[j]: 92 | failed = True 93 | break 94 | coor[j] = c 95 | if failed: 96 | continue 97 | voxelidx = coor_to_voxelidx[coor[0], coor[1], coor[2]] 98 | if voxelidx == -1: 99 | voxelidx = voxel_num 100 | if voxel_num >= max_voxels: 101 | break 102 | voxel_num += 1 103 | coor_to_voxelidx[coor[0], coor[1], coor[2]] = voxelidx 104 | coors[voxelidx] = coor 105 | num = num_points_per_voxel[voxelidx] 106 | if num < max_points: 107 | voxels[voxelidx, num] = points[i] 108 | num_points_per_voxel[voxelidx] += 1 109 | return voxel_num 110 | 111 | 112 | def points_to_voxel( 113 | points, voxel_size, coors_range, max_points=35, reverse_index=True, max_voxels=20000 114 | ): 115 | """convert kitti points(N, >=3) to voxels. This version calculate 116 | everything in one loop. now it takes only 4.2ms(complete point cloud) 117 | with jit and 3.2ghz cpu.(don't calculate other features) 118 | Note: this function in ubuntu seems faster than windows 10. 119 | 120 | Args: 121 | points: [N, ndim] float tensor. points[:, :3] contain xyz points and 122 | points[:, 3:] contain other information such as reflectivity. 123 | voxel_size: [3] list/tuple or array, float. xyz, indicate voxel size 124 | coors_range: [6] list/tuple or array, float. indicate voxel range. 125 | format: xyzxyz, minmax 126 | max_points: int. indicate maximum points contained in a voxel. 127 | reverse_index: boolean. indicate whether return reversed coordinates. 128 | if points has xyz format and reverse_index is True, output 129 | coordinates will be zyx format, but points in features always 130 | xyz format. 131 | max_voxels: int. indicate maximum voxels this function create. 132 | for second, 20000 is a good choice. you should shuffle points 133 | before call this function because max_voxels may drop some points. 134 | 135 | Returns: 136 | voxels: [M, max_points, ndim] float tensor. only contain points. 137 | coordinates: [M, 3] int32 tensor. 138 | num_points_per_voxel: [M] int32 tensor. 139 | """ 140 | if not isinstance(voxel_size, np.ndarray): 141 | voxel_size = np.array(voxel_size, dtype=points.dtype) 142 | if not isinstance(coors_range, np.ndarray): 143 | coors_range = np.array(coors_range, dtype=points.dtype) 144 | voxelmap_shape = (coors_range[3:] - coors_range[:3]) / voxel_size 145 | voxelmap_shape = tuple(np.round(voxelmap_shape).astype(np.int32).tolist()) 146 | if reverse_index: 147 | voxelmap_shape = voxelmap_shape[::-1] 148 | # don't create large array in jit(nopython=True) code. 149 | num_points_per_voxel = np.zeros(shape=(max_voxels,), dtype=np.int32) 150 | coor_to_voxelidx = -np.ones(shape=voxelmap_shape, dtype=np.int32) 151 | voxels = np.zeros( 152 | shape=(max_voxels, max_points, points.shape[-1]), dtype=points.dtype 153 | ) 154 | coors = np.zeros(shape=(max_voxels, 3), dtype=np.int32) 155 | if reverse_index: 156 | voxel_num = _points_to_voxel_reverse_kernel( 157 | points, 158 | voxel_size, 159 | coors_range, 160 | num_points_per_voxel, 161 | coor_to_voxelidx, 162 | voxels, 163 | coors, 164 | max_points, 165 | max_voxels, 166 | ) 167 | 168 | else: 169 | voxel_num = _points_to_voxel_kernel( 170 | points, 171 | voxel_size, 172 | coors_range, 173 | num_points_per_voxel, 174 | coor_to_voxelidx, 175 | voxels, 176 | coors, 177 | max_points, 178 | max_voxels, 179 | ) 180 | 181 | coors = coors[:voxel_num] 182 | voxels = voxels[:voxel_num] 183 | num_points_per_voxel = num_points_per_voxel[:voxel_num] 184 | return voxels, coors, num_points_per_voxel 185 | 186 | 187 | @numba.jit(nopython=True) 188 | def bound_points_jit(points, upper_bound, lower_bound): 189 | # to use nopython=True, np.bool is not supported. so you need 190 | # convert result to np.bool after this function. 191 | N = points.shape[0] 192 | ndim = points.shape[1] 193 | keep_indices = np.zeros((N,), dtype=np.int32) 194 | success = 0 195 | for i in range(N): 196 | success = 1 197 | for j in range(ndim): 198 | if points[i, j] < lower_bound[j] or points[i, j] >= upper_bound[j]: 199 | success = 0 200 | break 201 | keep_indices[i] = success 202 | return keep_indices 203 | -------------------------------------------------------------------------------- /lidardet/utils/laserscanvis.py: -------------------------------------------------------------------------------- 1 | import vispy 2 | from vispy.scene import visuals, SceneCanvas 3 | import numpy as np 4 | from matplotlib import pyplot as plt 5 | 6 | 7 | class LaserScanVis: 8 | """Class that creates and handles a visualizer for a pointcloud""" 9 | 10 | def __init__(self, scan, scan_names, label_names, offset=0, 11 | semantics=True, instances=False): 12 | self.scan = scan 13 | self.scan_names = scan_names 14 | self.label_names = label_names 15 | self.offset = offset 16 | self.semantics = semantics 17 | self.instances = instances 18 | # sanity check 19 | if not self.semantics and self.instances: 20 | print("Instances are only allowed in when semantics=True") 21 | raise ValueError 22 | 23 | self.reset() 24 | self.update_scan() 25 | 26 | def reset(self): 27 | """ Reset. """ 28 | # last key press (it should have a mutex, but visualization is not 29 | # safety critical, so let's do things wrong) 30 | self.action = "no" # no, next, back, quit are the possibilities 31 | 32 | # new canvas prepared for visualizing data 33 | self.canvas = SceneCanvas(keys='interactive', show=True) 34 | # interface (n next, b back, q quit, very simple) 35 | self.canvas.events.key_press.connect(self.key_press) 36 | self.canvas.events.draw.connect(self.draw) 37 | # grid 38 | self.grid = self.canvas.central_widget.add_grid() 39 | 40 | # laserscan part 41 | self.scan_view = vispy.scene.widgets.ViewBox( 42 | border_color='white', parent=self.canvas.scene) 43 | self.grid.add_widget(self.scan_view, 0, 0) 44 | self.scan_vis = visuals.Markers() 45 | self.scan_view.camera = 'turntable' 46 | self.scan_view.add(self.scan_vis) 47 | visuals.XYZAxis(parent=self.scan_view.scene) 48 | # add semantics 49 | if self.semantics: 50 | print("Using semantics in visualizer") 51 | self.sem_view = vispy.scene.widgets.ViewBox( 52 | border_color='white', parent=self.canvas.scene) 53 | self.grid.add_widget(self.sem_view, 0, 1) 54 | self.sem_vis = visuals.Markers() 55 | self.sem_view.camera = 'turntable' 56 | self.sem_view.add(self.sem_vis) 57 | visuals.XYZAxis(parent=self.sem_view.scene) 58 | # self.sem_view.camera.link(self.scan_view.camera) 59 | 60 | if self.instances: 61 | print("Using instances in visualizer") 62 | self.inst_view = vispy.scene.widgets.ViewBox( 63 | border_color='white', parent=self.canvas.scene) 64 | self.grid.add_widget(self.inst_view, 0, 2) 65 | self.inst_vis = visuals.Markers() 66 | self.inst_view.camera = 'turntable' 67 | self.inst_view.add(self.inst_vis) 68 | visuals.XYZAxis(parent=self.inst_view.scene) 69 | # self.inst_view.camera.link(self.scan_view.camera) 70 | 71 | # img canvas size 72 | self.multiplier = 1 73 | self.canvas_W = 1024 74 | self.canvas_H = 64 75 | if self.semantics: 76 | self.multiplier += 1 77 | if self.instances: 78 | self.multiplier += 1 79 | 80 | # new canvas for img 81 | self.img_canvas = SceneCanvas(keys='interactive', show=True, 82 | size=(self.canvas_W, self.canvas_H * self.multiplier)) 83 | # grid 84 | self.img_grid = self.img_canvas.central_widget.add_grid() 85 | # interface (n next, b back, q quit, very simple) 86 | self.img_canvas.events.key_press.connect(self.key_press) 87 | self.img_canvas.events.draw.connect(self.draw) 88 | 89 | # add a view for the depth 90 | self.img_view = vispy.scene.widgets.ViewBox( 91 | border_color='white', parent=self.img_canvas.scene) 92 | self.img_grid.add_widget(self.img_view, 0, 0) 93 | self.img_vis = visuals.Image(cmap='viridis') 94 | self.img_view.add(self.img_vis) 95 | 96 | # add semantics 97 | if self.semantics: 98 | self.sem_img_view = vispy.scene.widgets.ViewBox( 99 | border_color='white', parent=self.img_canvas.scene) 100 | self.img_grid.add_widget(self.sem_img_view, 1, 0) 101 | self.sem_img_vis = visuals.Image(cmap='viridis') 102 | self.sem_img_view.add(self.sem_img_vis) 103 | 104 | # add instances 105 | if self.instances: 106 | self.inst_img_view = vispy.scene.widgets.ViewBox( 107 | border_color='white', parent=self.img_canvas.scene) 108 | self.img_grid.add_widget(self.inst_img_view, 2, 0) 109 | self.inst_img_vis = visuals.Image(cmap='viridis') 110 | self.inst_img_view.add(self.inst_img_vis) 111 | 112 | def get_mpl_colormap(self, cmap_name): 113 | cmap = plt.get_cmap(cmap_name) 114 | 115 | # Initialize the matplotlib color map 116 | sm = plt.cm.ScalarMappable(cmap=cmap) 117 | 118 | # Obtain linear color range 119 | color_range = sm.to_rgba(np.linspace(0, 1, 256), bytes=True)[:, 2::-1] 120 | 121 | return color_range.reshape(256, 3).astype(np.float32) / 255.0 122 | 123 | def update_scan(self): 124 | # first open data 125 | self.scan.open_scan(self.scan_names[self.offset]) 126 | if self.semantics: 127 | self.scan.open_label(self.label_names[self.offset]) 128 | self.scan.colorize() 129 | 130 | # then change names 131 | title = "scan " + str(self.offset) + " of " + str(len(self.scan_names)) 132 | self.canvas.title = title 133 | self.img_canvas.title = title 134 | 135 | # then do all the point cloud stuff 136 | 137 | # plot scan 138 | power = 16 139 | # print() 140 | range_data = np.copy(self.scan.unproj_range) 141 | # print(range_data.max(), range_data.min()) 142 | range_data = range_data**(1 / power) 143 | # print(range_data.max(), range_data.min()) 144 | viridis_range = ((range_data - range_data.min()) / 145 | (range_data.max() - range_data.min()) * 146 | 255).astype(np.uint8) 147 | viridis_map = self.get_mpl_colormap("viridis") 148 | viridis_colors = viridis_map[viridis_range] 149 | self.scan_vis.set_data(self.scan.points, 150 | face_color=viridis_colors[..., ::-1], 151 | edge_color=viridis_colors[..., ::-1], 152 | size=1) 153 | 154 | # plot semantics 155 | if self.semantics: 156 | self.sem_vis.set_data(self.scan.points, 157 | face_color=self.scan.sem_label_color[..., ::-1], 158 | edge_color=self.scan.sem_label_color[..., ::-1], 159 | size=1) 160 | 161 | # plot instances 162 | if self.instances: 163 | self.inst_vis.set_data(self.scan.points, 164 | face_color=self.scan.inst_label_color[..., ::-1], 165 | edge_color=self.scan.inst_label_color[..., ::-1], 166 | size=1) 167 | 168 | # now do all the range image stuff 169 | # plot range image 170 | data = np.copy(self.scan.proj_range) 171 | # print(data[data > 0].max(), data[data > 0].min()) 172 | data[data > 0] = data[data > 0]**(1 / power) 173 | data[data < 0] = data[data > 0].min() 174 | # print(data.max(), data.min()) 175 | data = (data - data[data > 0].min()) / \ 176 | (data.max() - data[data > 0].min()) 177 | # print(data.max(), data.min()) 178 | self.img_vis.set_data(data) 179 | self.img_vis.update() 180 | 181 | if self.semantics: 182 | self.sem_img_vis.set_data(self.scan.proj_sem_color[..., ::-1]) 183 | self.sem_img_vis.update() 184 | 185 | if self.instances: 186 | self.inst_img_vis.set_data(self.scan.proj_inst_color[..., ::-1]) 187 | self.inst_img_vis.update() 188 | 189 | # interface 190 | def key_press(self, event): 191 | self.canvas.events.key_press.block() 192 | self.img_canvas.events.key_press.block() 193 | if event.key == 'N': 194 | self.offset += 1 195 | self.update_scan() 196 | elif event.key == 'B': 197 | self.offset -= 1 198 | self.update_scan() 199 | elif event.key == 'Q' or event.key == 'Escape': 200 | self.destroy() 201 | 202 | def draw(self, event): 203 | if self.canvas.events.key_press.blocked(): 204 | self.canvas.events.key_press.unblock() 205 | if self.img_canvas.events.key_press.blocked(): 206 | self.img_canvas.events.key_press.unblock() 207 | 208 | def destroy(self): 209 | # destroy the visualization 210 | self.canvas.close() 211 | self.img_canvas.close() 212 | vispy.app.quit() 213 | 214 | def run(self): 215 | vispy.app.run() 216 | -------------------------------------------------------------------------------- /lidardet/utils/common.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import torch 4 | import pickle 5 | import shutil 6 | import logging 7 | import datetime 8 | import numpy as np 9 | import subprocess 10 | from pathlib import Path 11 | import random as pyrandom 12 | import torch.distributed as dist 13 | 14 | from .dist import get_dist_info 15 | 16 | def scan_upsample(points_array, input_rings=32, vertical_fov=26.8, bottom_angle=-24.8): 17 | im0, inds_e, inds_o = scan_to_range(points_array, input_rings, vertical_fov, bottom_angle) 18 | h, w, c = im0.shape 19 | points_new = [] 20 | for i in range(h - 1): 21 | for j in range(w): 22 | d1, t1, v_angle1, h_angle1 = im0[i, j, :] 23 | d2, t2, v_angle2, h_angle2 = im0[i + 1, j, :] 24 | if d1 != 0 and d2 != 0: 25 | t = (t1 + t2) * 0.5 26 | d = (d1 + d2) * 0.5 27 | v_angle = (v_angle1 + v_angle2) * 0.5 28 | h_angle = (h_angle1 + h_angle2) * 0.5 29 | x = np.sin(h_angle) * np.cos(v_angle) * d 30 | y = np.cos(h_angle) * np.cos(v_angle) * d 31 | z = np.sin(v_angle) * d 32 | point = np.array([x, y, z, t]) 33 | points_new.append(point) 34 | points_new = np.array(points_new) 35 | points_hr = np.vstack((points_array, points_new)) 36 | return points_hr 37 | 38 | def scan_downsample(points_array, input_rings=64, vertical_fov=26.8, bottom_angle=-24.8, output_rings='even'): 39 | range_image, inds_e, inds_o = scan_to_range(points_array, input_rings, vertical_fov, bottom_angle) 40 | if output_rings == 'even': 41 | return points_array[inds_e,:4] 42 | elif output_rings == 'odd': 43 | return points_array[inds_o,:4] 44 | elif output_rings == 'even_or_odd': 45 | even = np.random.choice([False, True], replace=False, p=[0.5, 0.5]) 46 | if even: 47 | return points_array[inds_e,:4] 48 | return points_array[inds_o,:4] 49 | elif output_rings == 'random': 50 | inds = inds_e + inds_o 51 | np.random.shuffle(inds) 52 | inds = inds[:int(len(inds) * 0.5)] 53 | return points_array[inds,:4] 54 | else: 55 | raise ValueError('Unknown output_rings value: %s', output_rings) 56 | 57 | def range_to_scan(range_image, num_rings=64, vertical_fov=26.8, bottom_angle=-24.8): 58 | max_range = 80.0 59 | min_range = 2.0 60 | image_cols = 1024 61 | 62 | ang_res_x = 360.0 / float(image_cols) # horizontal resolution 63 | ang_res_y = vertical_fov / float(num_rings - 1) # vertical resolution 64 | row_ids = np.arange(num_rings) 65 | col_ids = np.arange(image_cols) 66 | v_angles = np.float32(row_ids * ang_res_y) + bottom_angle 67 | h_angles = np.float32(col_ids + 1 - image_cols / 2) * ang_res_x + 90 68 | v_angles = v_angles / 180.0 * np.pi 69 | h_angles = h_angles / 180.0 * np.pi 70 | 71 | range_image[:,:,0] *= max_range 72 | 73 | h, w, c = range_image.shape 74 | points = [] 75 | inds_even = [] 76 | inds_odd = [] 77 | for i in range(h): 78 | for j in range(w): 79 | depth, intensity = range_image[i, j, :] 80 | if depth < min_range: 81 | continue 82 | h_angle = h_angles[j] 83 | v_angle = v_angles[i] 84 | x = np.sin(h_angle) * np.cos(v_angle) * depth 85 | y = np.cos(h_angle) * np.cos(v_angle) * depth 86 | z = np.sin(v_angle) * depth 87 | point = np.array([x, y, z, int(intensity)]) 88 | points.append(point) 89 | idx = len(points) - 1 90 | if i % 2 == 0: 91 | inds_even.append(idx) 92 | else: 93 | inds_odd.append(idx) 94 | return np.array(points), inds_even, inds_odd 95 | 96 | def scan_to_range(points_array, input_rings=64, vertical_fov=26.8, bottom_angle=-24.8, normalize=False): 97 | # range image size, depends on your sensor, i.e., VLP-16: 16x1800, OS1-64: 64x1024 98 | image_rows_full = input_rings 99 | max_range = 80.0 100 | min_range = 2.0 101 | image_cols = 1024 102 | 103 | ang_res_x = 360.0 / float(image_cols) # horizontal resolution 104 | ang_res_y = vertical_fov / float(image_rows_full - 1) # vertical resolution 105 | ang_start_y = bottom_angle 106 | 107 | # project points to range image 108 | # channels: range, intensity, horizon_angle, vertical_angle 109 | range_image = np.zeros((image_rows_full, image_cols, 4), dtype=np.float32) 110 | x = points_array[:,0] 111 | y = points_array[:,1] 112 | z = points_array[:,2] 113 | t = points_array[:,3] 114 | # find row id 115 | vertical_angle = np.arctan2(z, np.sqrt(x * x + y * y)) * 180.0 / np.pi 116 | 117 | relative_vertical_angle = vertical_angle - ang_start_y 118 | rowId = np.int_(np.round_(relative_vertical_angle / ang_res_y)) 119 | 120 | # find column id 121 | horitontal_angle = np.arctan2(x, y) * 180.0 / np.pi 122 | colId = -np.int_((horitontal_angle - 90.0) / ang_res_x) + image_cols / 2; 123 | shift_ids = np.where(colId>=image_cols) 124 | colId[shift_ids] = colId[shift_ids] - image_cols 125 | # filter range 126 | thisRange = np.sqrt(x * x + y * y + z * z) 127 | thisRange[thisRange > max_range] = 0 128 | thisRange[thisRange < min_range] = 0 129 | 130 | if normalize: 131 | thisRange /= max_range 132 | 133 | # save range info to range image 134 | inds = [] 135 | inds_odd_row = [] 136 | inds_even_row = [] 137 | for i in range(len(thisRange)): 138 | if rowId[i] < 0 or rowId[i] >= image_rows_full or colId[i] < 0 or colId[i] >= image_cols: 139 | continue 140 | range_image[int(rowId[i]), int(colId[i]), 0] = thisRange[i] 141 | range_image[int(rowId[i]), int(colId[i]), 1] = t[i] 142 | range_image[int(rowId[i]), int(colId[i]), 2] = vertical_angle[i] * np.pi / 180.0 143 | range_image[int(rowId[i]), int(colId[i]), 3] = horitontal_angle[i] * np.pi / 180.0 144 | 145 | if thisRange[i] > 0: 146 | inds.append(i) 147 | if rowId[i] % 2 == 0: 148 | inds_even_row.append(i) 149 | else: 150 | inds_odd_row.append(i) 151 | return range_image, inds_even_row, inds_odd_row 152 | 153 | def check_numpy_to_torch(x): 154 | if isinstance(x, np.ndarray): 155 | return torch.from_numpy(x).float(), True 156 | return x, False 157 | 158 | def get_voxel_centers(voxel_coords, downsample_times, voxel_size, point_cloud_range): 159 | """ 160 | Args: 161 | voxel_coords: (N, 3) 162 | downsample_times: 163 | voxel_size: 164 | point_cloud_range: 165 | 166 | Returns: 167 | 168 | """ 169 | assert voxel_coords.shape[1] == 3 170 | voxel_centers = voxel_coords[:, [2, 1, 0]].float() # (xyz) 171 | voxel_size = torch.tensor(voxel_size, device=voxel_centers.device).float() * downsample_times 172 | pc_range = torch.tensor(point_cloud_range[0:3], device=voxel_centers.device).float() 173 | voxel_centers = (voxel_centers + 0.5) * voxel_size + pc_range 174 | return voxel_centers 175 | 176 | def keep_arrays_by_name(gt_names, used_classes): 177 | inds = [i for i, x in enumerate(gt_names) if x in used_classes] 178 | inds = np.array(inds, dtype=np.int64) 179 | return inds 180 | 181 | def set_random_seed(seed=3): 182 | pyrandom.seed(seed) 183 | np.random.seed(seed) 184 | torch.manual_seed(seed) 185 | torch.cuda.manual_seed_all(seed) 186 | 187 | def get_logger(logdir, name): 188 | logger = logging.getLogger(name) 189 | logger.setLevel(logging.INFO) 190 | formatter = logging.Formatter("%(asctime)s %(levelname)s %(message)s") 191 | 192 | ts = str(datetime.datetime.now()).split(".")[0].replace(" ", "_") 193 | ts = ts.replace(":", "_").replace("-", "_") 194 | file_path = Path(logdir) / "run_{}.log".format(ts) 195 | file_hdlr = logging.FileHandler(str(file_path)) 196 | file_hdlr.setFormatter(formatter) 197 | 198 | strm_hdlr = logging.StreamHandler(sys.stdout) 199 | strm_hdlr.setFormatter(formatter) 200 | 201 | logger.addHandler(file_hdlr) 202 | logger.addHandler(strm_hdlr) 203 | return logger 204 | 205 | def merge_results_dist(result_part, size, tmpdir): 206 | rank, world_size = get_dist_info() 207 | os.makedirs(tmpdir, exist_ok=True) 208 | 209 | dist.barrier() 210 | pickle.dump(result_part, open(os.path.join(tmpdir, 'result_part_{}.pkl'.format(rank)), 'wb')) 211 | dist.barrier() 212 | 213 | if rank != 0: 214 | return None 215 | 216 | part_list = [] 217 | for i in range(world_size): 218 | part_file = os.path.join(tmpdir, 'result_part_{}.pkl'.format(i)) 219 | part_list.append(pickle.load(open(part_file, 'rb'))) 220 | 221 | ordered_results = [] 222 | for res in zip(*part_list): 223 | ordered_results.extend(list(res)) 224 | ordered_results = ordered_results[:size] 225 | shutil.rmtree(tmpdir) 226 | return ordered_results 227 | -------------------------------------------------------------------------------- /lidardet/datasets/augmentor/database_sampler.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pickle 3 | #from ...utils import box_utils 4 | #from ...ops.iou3d_nms import iou3d_nms_utils 5 | 6 | 7 | class DataBaseSampler(object): 8 | def __init__(self, root_path, sampler_cfg, class_names, logger=None): 9 | self.root_path = root_path 10 | self.class_names = class_names 11 | self.sampler_cfg = sampler_cfg 12 | self.logger = logger 13 | self.db_infos = {} 14 | for class_name in class_names: 15 | self.db_infos[class_name] = [] 16 | 17 | for db_info_path in sampler_cfg['db_info_path']: 18 | db_info_path = self.root_path.resolve() / db_info_path 19 | with open(str(db_info_path), 'rb') as f: 20 | infos = pickle.load(f) 21 | [self.db_infos[cur_class].extend(infos[cur_class]) for cur_class in class_names] 22 | 23 | for func_name, val in sampler_cfg['prepare'].items(): 24 | self.db_infos = getattr(self, func_name)(self.db_infos, val) 25 | 26 | self.sample_groups = {} 27 | self.sample_class_num = {} 28 | self.limit_whole_scene = sampler_cfg.get('limit_whole_scene', False) 29 | for x in sampler_cfg['sample_groups']: 30 | class_name, sample_num = x.split(':') 31 | if class_name not in class_names: 32 | continue 33 | self.sample_class_num[class_name] = sample_num 34 | self.sample_groups[class_name] = { 35 | 'sample_num': sample_num, 36 | 'pointer': len(self.db_infos[class_name]), 37 | 'indices': np.arange(len(self.db_infos[class_name])) 38 | } 39 | 40 | def __getstate__(self): 41 | d = dict(self.__dict__) 42 | del d['logger'] 43 | return d 44 | 45 | def __setstate__(self, d): 46 | self.__dict__.update(d) 47 | 48 | def filter_by_difficulty(self, db_infos, removed_difficulty): 49 | new_db_infos = {} 50 | for key, dinfos in db_infos.items(): 51 | pre_len = len(dinfos) 52 | new_db_infos[key] = [ 53 | info for info in dinfos 54 | if info['difficulty'] not in removed_difficulty 55 | ] 56 | if self.logger is not None: 57 | self.logger.info('Database filter by difficulty %s: %d => %d' % (key, pre_len, len(new_db_infos[key]))) 58 | return new_db_infos 59 | 60 | def filter_by_min_points(self, db_infos, min_gt_points_list): 61 | for name_num in min_gt_points_list: 62 | name, min_num = name_num.split(':') 63 | min_num = int(min_num) 64 | if min_num > 0 and name in db_infos.keys(): 65 | filtered_infos = [] 66 | for info in db_infos[name]: 67 | if info['num_points_in_gt'] >= min_num: 68 | filtered_infos.append(info) 69 | 70 | if self.logger is not None: 71 | self.logger.info('Database filter by min points %s: %d => %d' % 72 | (name, len(db_infos[name]), len(filtered_infos))) 73 | db_infos[name] = filtered_infos 74 | 75 | return db_infos 76 | 77 | def sample_with_fixed_number(self, class_name, sample_group): 78 | """ 79 | Args: 80 | class_name: 81 | sample_group: 82 | Returns: 83 | 84 | """ 85 | sample_num, pointer, indices = int(sample_group['sample_num']), sample_group['pointer'], sample_group['indices'] 86 | if pointer >= len(self.db_infos[class_name]): 87 | indices = np.random.permutation(len(self.db_infos[class_name])) 88 | pointer = 0 89 | 90 | sampled_dict = [self.db_infos[class_name][idx] for idx in indices[pointer: pointer + sample_num]] 91 | pointer += sample_num 92 | sample_group['pointer'] = pointer 93 | sample_group['indices'] = indices 94 | return sampled_dict 95 | 96 | @staticmethod 97 | def put_boxes_on_road_planes(gt_boxes, road_planes, calib): 98 | """ 99 | Only validate in KITTIDataset 100 | Args: 101 | gt_boxes: (N, 7 + C) [x, y, z, dx, dy, dz, heading, ...] 102 | road_planes: [a, b, c, d] 103 | calib: 104 | 105 | Returns: 106 | """ 107 | a, b, c, d = road_planes 108 | center_cam = calib.lidar_to_rect(gt_boxes[:, 0:3]) 109 | cur_height_cam = (-d - a * center_cam[:, 0] - c * center_cam[:, 2]) / b 110 | center_cam[:, 1] = cur_height_cam 111 | cur_lidar_height = calib.rect_to_lidar(center_cam)[:, 2] 112 | mv_height = gt_boxes[:, 2] - gt_boxes[:, 5] / 2 - cur_lidar_height 113 | gt_boxes[:, 2] -= mv_height # lidar view 114 | return gt_boxes, mv_height 115 | 116 | def add_sampled_boxes_to_scene(self, data_dict, sampled_gt_boxes, total_valid_sampled_dict): 117 | gt_boxes_mask = data_dict['gt_boxes_mask'] 118 | gt_boxes = data_dict['gt_boxes'][gt_boxes_mask] 119 | gt_names = data_dict['gt_names'][gt_boxes_mask] 120 | points = data_dict['points'] 121 | if self.sampler_cfg.get('use_road_plane', False): 122 | sampled_gt_boxes, mv_height = self.put_boxes_on_road_planes( 123 | sampled_gt_boxes, data_dict['road_plane'], data_dict['calib'] 124 | ) 125 | data_dict.pop('calib') 126 | data_dict.pop('road_plane') 127 | 128 | obj_points_list = [] 129 | for idx, info in enumerate(total_valid_sampled_dict): 130 | file_path = self.root_path / info['path'] 131 | obj_points = np.fromfile(str(file_path), dtype=np.float32).reshape( 132 | [-1, self.sampler_cfg['num_point_features']]) 133 | 134 | obj_points[:, :3] += info['box3d_lidar'][:3] 135 | 136 | if self.sampler_cfg.get('use_road_plane', False): 137 | # mv height 138 | obj_points[:, 2] -= mv_height[idx] 139 | 140 | obj_points_list.append(obj_points) 141 | 142 | obj_points = np.concatenate(obj_points_list, axis=0) 143 | sampled_gt_names = np.array([x['name'] for x in total_valid_sampled_dict]) 144 | 145 | large_sampled_gt_boxes = box_utils.enlarge_box3d( 146 | sampled_gt_boxes[:, 0:7], extra_width=self.sampler_cfg['remove_extra_width'] 147 | ) 148 | points = box_utils.remove_points_in_boxes3d(points, large_sampled_gt_boxes) 149 | points = np.concatenate([obj_points, points], axis=0) 150 | gt_names = np.concatenate([gt_names, sampled_gt_names], axis=0) 151 | gt_boxes = np.concatenate([gt_boxes, sampled_gt_boxes], axis=0) 152 | data_dict['gt_boxes'] = gt_boxes 153 | data_dict['gt_names'] = gt_names 154 | data_dict['points'] = points 155 | return data_dict 156 | 157 | def __call__(self, data_dict): 158 | """ 159 | Args: 160 | data_dict: 161 | gt_boxes: (N, 7 + C) [x, y, z, dx, dy, dz, heading, ...] 162 | 163 | Returns: 164 | 165 | """ 166 | gt_boxes = data_dict['gt_boxes'] 167 | gt_names = data_dict['gt_names'] 168 | existed_boxes = gt_boxes 169 | total_valid_sampled_dict = [] 170 | for class_name, sample_group in self.sample_groups.items(): 171 | if self.limit_whole_scene: 172 | num_gt = np.sum(class_name == gt_names) 173 | sample_group['sample_num'] = str(int(self.sample_class_num[class_name]) - num_gt) 174 | if int(sample_group['sample_num']) > 0: 175 | sampled_dict = self.sample_with_fixed_number(class_name, sample_group) 176 | 177 | sampled_boxes = np.stack([x['box3d_lidar'] for x in sampled_dict], axis=0).astype(np.float32) 178 | 179 | if self.sampler_cfg.get('database_with_fakelidar', False): 180 | sampled_boxes = box_utils.boxes3d_kitti_fakelidar_to_lidar(sampled_boxes) 181 | 182 | iou1 = iou3d_nms_utils.boxes_bev_iou_cpu(sampled_boxes, existed_boxes) 183 | iou2 = iou3d_nms_utils.boxes_bev_iou_cpu(sampled_boxes, sampled_boxes) 184 | iou2[range(sampled_boxes.shape[0]), range(sampled_boxes.shape[0])] = 0 185 | valid_mask = ((iou1.max(axis=1) + iou2.max(axis=1)) == 0).nonzero()[0] 186 | valid_sampled_dict = [sampled_dict[x] for x in valid_mask] 187 | valid_sampled_boxes = sampled_boxes[valid_mask] 188 | 189 | existed_boxes = np.concatenate((existed_boxes, valid_sampled_boxes), axis=0) 190 | total_valid_sampled_dict.extend(valid_sampled_dict) 191 | 192 | sampled_gt_boxes = existed_boxes[gt_boxes.shape[0]:, :] 193 | if total_valid_sampled_dict.__len__() > 0: 194 | data_dict = self.add_sampled_boxes_to_scene(data_dict, sampled_gt_boxes, total_valid_sampled_dict) 195 | 196 | data_dict.pop('gt_boxes_mask') 197 | return data_dict 198 | -------------------------------------------------------------------------------- /lidardet/utils/buildtools/command.py: -------------------------------------------------------------------------------- 1 | import multiprocessing 2 | import os 3 | import re 4 | import subprocess 5 | from concurrent.futures import ProcessPoolExecutor 6 | from enum import Enum 7 | from functools import partial 8 | from pathlib import Path 9 | 10 | from lidardet.utils.find import find_cuda, find_cuda_device_arch 11 | 12 | 13 | class Gpp: 14 | def __init__( 15 | self, 16 | sources, 17 | target, 18 | std="c++11", 19 | includes: list = None, 20 | defines: dict = None, 21 | cflags: str = None, 22 | compiler="g++", 23 | link=False, 24 | libraries: dict = None, 25 | lflags: str = None, 26 | extra_cflags: str = None, 27 | extra_lflags: str = None, 28 | build_directory: str = None, 29 | ): 30 | if not isinstance(sources, (list, tuple)): 31 | sources = [sources] 32 | if build_directory is not None: 33 | build_directory = Path(build_directory) 34 | new_sources = [] 35 | for p in sources: 36 | if not Path(p).is_absolute(): 37 | new_sources.append(str(build_directory / p)) 38 | else: 39 | new_sources.append(p) 40 | sources = new_sources 41 | target = Path(target) 42 | if not target.is_absolute(): 43 | target = str(build_directory / target) 44 | self.sources = [str(p) for p in sources] 45 | self.target = str(target) 46 | self.std = std 47 | self.includes = includes or [] 48 | self.cflags = cflags or "-fPIC -O3" 49 | self.defines = defines or {} 50 | self.compiler = compiler 51 | self.link = link 52 | self.libraries = libraries or {} 53 | self.lflags = lflags or "" 54 | self.extra_cflags = extra_cflags or "" 55 | self.extra_lflags = extra_lflags or "" 56 | 57 | def shell(self, target: str = None, compiler: str = None): 58 | defines = [f"-D {n}={v}" for n, v in self.defines.items()] 59 | includes = [f"-I{inc}" for inc in self.includes] 60 | libraries = [ 61 | f"-L{n} {' '.join(['-l' + l for l in v])}" 62 | for n, v in self.libraries.items() 63 | ] 64 | compiler = compiler or self.compiler 65 | string = f"{compiler} -std={self.std} " 66 | if self.link: 67 | string += " -shared " 68 | else: 69 | string += " -c " 70 | target = target or self.target 71 | string += ( 72 | f"-o {target} {' '.join(self.sources)} " 73 | f"{' '.join(defines)} " 74 | f"{' '.join(includes)} " 75 | f"{self.cflags} {self.extra_cflags}" 76 | f"{' '.join(libraries)} " 77 | f"{self.lflags} {self.extra_lflags}" 78 | ) 79 | return re.sub(r" +", r" ", string) 80 | 81 | 82 | class Link: 83 | def __init__(self, outs, target, compiler="ld", build_directory: str = None): 84 | if not isinstance(outs, (list, tuple)): 85 | outs = [outs] 86 | if build_directory is not None: 87 | build_directory = Path(build_directory) 88 | new_outs = [] 89 | for p in outs: 90 | if not Path(p).is_absolute(): 91 | new_outs.append(str(build_directory / p)) 92 | else: 93 | new_outs.append(p) 94 | outs = new_outs 95 | target = Path(target) 96 | if target.is_absolute(): 97 | target = str(build_directory / target) 98 | self.outs = [str(p) for p in outs] 99 | self.target = str(target) 100 | self.compiler = compiler 101 | 102 | def shell(self, target: str = None): 103 | string = f"{self.compiler} -r " 104 | if target is None: 105 | target = self.target 106 | string += f"-o {target} {' '.join(self.outs)} " 107 | return string 108 | 109 | 110 | class Nvcc(Gpp): 111 | def __init__( 112 | self, 113 | sources, 114 | target, 115 | arch=None, 116 | std="c++11", 117 | includes: list = None, 118 | defines: dict = None, 119 | cflags: str = None, 120 | extra_cflags: str = None, 121 | extra_lflags: str = None, 122 | build_directory: str = None, 123 | ): 124 | if arch is None: 125 | arch = find_cuda_device_arch() 126 | if arch is None: 127 | raise ValueError("you must specify arch if use cuda.") 128 | 129 | cflags = ( 130 | cflags or f"-x cu -Xcompiler -fPIC -arch={arch} --expt-relaxed-constexpr" 131 | ) 132 | try: 133 | cuda_home = find_cuda() 134 | except: 135 | cuda_home = None 136 | if cuda_home is not None: 137 | cuda_include = Path(cuda_home) / "include" 138 | includes = includes or [] 139 | includes += [str(cuda_include)] 140 | super().__init__( 141 | sources, 142 | target, 143 | std, 144 | includes, 145 | defines, 146 | cflags, 147 | compiler="nvcc", 148 | extra_cflags=extra_cflags, 149 | extra_lflags=extra_lflags, 150 | build_directory=build_directory, 151 | ) 152 | 153 | 154 | class CUDALink(Gpp): 155 | def __init__( 156 | self, 157 | sources, 158 | target, 159 | std="c++11", 160 | includes: list = None, 161 | defines: dict = None, 162 | cflags: str = None, 163 | libraries: dict = None, 164 | lflags: str = None, 165 | extra_cflags: str = None, 166 | extra_lflags: str = None, 167 | build_directory: str = None, 168 | ): 169 | includes = includes or [] 170 | defines = defines or {} 171 | libraries = libraries or {} 172 | cflags = cflags or "-fPIC -O3" 173 | try: 174 | cuda_home = find_cuda() 175 | except: 176 | cuda_home = None 177 | if cuda_home is not None: 178 | cuda_include = Path(cuda_home) / "include" 179 | includes += [str(cuda_include)] 180 | cuda_lib_path = Path(cuda_home) / "lib64" 181 | cuda_libs = {str(cuda_lib_path): ["cublas", "cudart"]} 182 | libraries = {**libraries, **cuda_libs} 183 | super().__init__( 184 | sources, 185 | target, 186 | std, 187 | includes, 188 | defines, 189 | cflags, 190 | link=True, 191 | libraries=libraries, 192 | lflags=lflags, 193 | extra_cflags=extra_cflags, 194 | extra_lflags=extra_lflags, 195 | build_directory=build_directory, 196 | ) 197 | 198 | 199 | class NodeState(Enum): 200 | Evaled = "evaled" 201 | Normal = "normal" 202 | Error = "error" 203 | 204 | 205 | class Node: 206 | def __init__(self, name=None): 207 | self.name = name 208 | self.prev = [] 209 | self.next = [] 210 | self.state = NodeState.Normal 211 | 212 | def __call__(self, *nodes): 213 | for node in nodes: 214 | self.prev.append(node) 215 | node.next.append(self) 216 | return self 217 | 218 | def _eval(self, *args, **kw): 219 | return True 220 | 221 | def eval(self, *args, **kw): 222 | for p in self.prev: 223 | if not p.eval(*args, **kw): 224 | self.state = NodeState.Error 225 | return False 226 | if self.state == NodeState.Normal: 227 | if self._eval(*args, **kw): 228 | self.state = NodeState.Evaled 229 | else: 230 | self.state = NodeState.Error 231 | return True 232 | return True 233 | 234 | def reset(self): 235 | self.state = NodeState.Normal 236 | self.prev = [] 237 | self.next = [] 238 | for node in self.prev: 239 | node.reset() 240 | 241 | 242 | class TargetNode(Node): 243 | def __init__(self, srcs, hdrs, deps, copts, name=None): 244 | super().__init__(name) 245 | self.srcs = srcs 246 | self.hdrs = hdrs 247 | self.deps = deps 248 | self.copts = copts 249 | 250 | def _eval(self, executor): 251 | pass 252 | 253 | 254 | def compile_func(cmd, code_folder, compiler): 255 | if not isinstance(cmd, (Link, Nvcc)): 256 | shell = cmd.shell(compiler=compiler) 257 | else: 258 | shell = cmd.shell() 259 | print(shell) 260 | cwd = None 261 | if code_folder is not None: 262 | cwd = str(code_folder) 263 | 264 | ret = subprocess.run(shell, shell=True, cwd=cwd) 265 | if ret.returncode != 0: 266 | raise RuntimeError("compile failed with retcode", ret.returncode) 267 | return ret 268 | 269 | 270 | def compile_libraries(cmds, code_folder=None, compiler: str = None, num_workers=-1): 271 | if num_workers == -1: 272 | num_workers = min(len(cmds), multiprocessing.cpu_count()) 273 | # for cmd in cmds: 274 | # print(cmd.shell()) 275 | if num_workers == 0: 276 | rets = map( 277 | partial(compile_func, code_folder=code_folder, compiler=compiler), cmds 278 | ) 279 | else: 280 | with ProcessPoolExecutor(num_workers) as pool: 281 | func = partial(compile_func, code_folder=code_folder, compiler=compiler) 282 | rets = pool.map(func, cmds) 283 | 284 | if any([r.returncode != 0 for r in rets]): 285 | cmds.clear() 286 | return False 287 | cmds.clear() 288 | return True 289 | 290 | 291 | def out(path): 292 | return Path(path).parent / (Path(path).stem + ".o") 293 | -------------------------------------------------------------------------------- /lidardet/utils/geometry.py: -------------------------------------------------------------------------------- 1 | import numba 2 | import torch 3 | import numpy as np 4 | from .common import check_numpy_to_torch 5 | 6 | def limit_period(val, offset=0.5, period=np.pi): 7 | val, is_numpy = check_numpy_to_torch(val) 8 | ans = val - torch.floor(val / period + offset) * period 9 | return ans.numpy() if is_numpy else ans 10 | 11 | def rotate_points_along_z(points, angle): 12 | """ 13 | Args: 14 | points: (B, N, 3 + C) 15 | angle: (B), angle along z-axis, angle increases x ==> y 16 | Returns: 17 | 18 | """ 19 | points, is_numpy = check_numpy_to_torch(points) 20 | angle, _ = check_numpy_to_torch(angle) 21 | 22 | cosa = torch.cos(angle) 23 | sina = torch.sin(angle) 24 | zeros = angle.new_zeros(points.shape[0]) 25 | ones = angle.new_ones(points.shape[0]) 26 | rot_matrix = torch.stack(( 27 | cosa, sina, zeros, 28 | -sina, cosa, zeros, 29 | zeros, zeros, ones 30 | ), dim=1).view(-1, 3, 3).float() 31 | points_rot = torch.matmul(points[:, :, 0:3], rot_matrix) 32 | points_rot = torch.cat((points_rot, points[:, :, 3:]), dim=-1) 33 | return points_rot.numpy() if is_numpy else points_rot 34 | 35 | def camera_to_lidar(points, r_rect, velo2cam): 36 | pts = np.concatenate( 37 | [points[:, :3], np.ones([points.shape[0], 1])], axis=1) 38 | pts = pts @ np.linalg.inv((r_rect @ velo2cam).T) 39 | points[:, :3] = pts[:, :3] 40 | return points 41 | 42 | 43 | def lidar_to_camera(points, r_rect, velo2cam): 44 | pts = np.concatenate( 45 | [points[:, :3], np.ones([points.shape[0], 1])], axis=1) 46 | pts = pts @ (r_rect @ velo2cam).T 47 | points[:, :3] = pts[:, :3] 48 | return points 49 | 50 | 51 | def box_camera_to_lidar(data, r_rect, velo2cam): 52 | xyz = data[:, 0:3] 53 | l, h, w = data[:, 3:4], data[:, 4:5], data[:, 5:6] 54 | r = data[:, 6:7] 55 | xyz_lidar = camera_to_lidar(xyz, r_rect, velo2cam) 56 | return np.concatenate([xyz_lidar, w, l, h, r], axis=1) 57 | 58 | 59 | def box_lidar_to_camera(data, r_rect, velo2cam): 60 | xyz_lidar = data[:, 0:3] 61 | w, l, h = data[:, 3:4], data[:, 4:5], data[:, 5:6] 62 | r = data[:, 6:7] 63 | xyz = lidar_to_camera(xyz_lidar, r_rect, velo2cam) 64 | return np.concatenate([xyz, l, h, w, r], axis=1) 65 | 66 | 67 | def projection_matrix_to_CRT_kitti(P): 68 | """ 69 | 将投影矩阵P利用QR分解分解出摄像机内外参数 70 | 输入： 71 | P：投影矩阵，3*4 72 | 输出： 73 | K：内参数矩阵，3*3 74 | R：旋转矩阵，3*3 75 | T：平移向量，3*1 76 | """ 77 | # P = K @ [R|T] 78 | # K is upper triangular matrix, so we need to inverse CR and use QR 79 | # stable for all kitti camera projection matrix 80 | CR = P[0:3, 0:3] 81 | CT = P[0:3, 3] 82 | RinvCinv = np.linalg.inv(CR) 83 | Rinv, Cinv = np.linalg.qr(RinvCinv) 84 | K = np.linalg.inv(Cinv) 85 | R = np.linalg.inv(Rinv) 86 | T = Cinv @ CT 87 | return K, R, T 88 | 89 | 90 | def get_frustum(bbox_image, C, near_clip=0.001, far_clip=100): 91 | fku = C[0, 0] 92 | fkv = -C[1, 1] 93 | u0v0 = C[0:2, 2] 94 | z_points = np.array([near_clip] * 4 + [far_clip] * 95 | 4, dtype=C.dtype)[:, np.newaxis] 96 | b = bbox_image 97 | box_corners = np.array( 98 | [[b[0], b[1]], [b[0], b[3]], [b[2], b[3]], [b[2], b[1]]], dtype=C.dtype 99 | ) 100 | near_box_corners = (box_corners - u0v0) / np.array( 101 | [fku / near_clip, -fkv / near_clip], dtype=C.dtype 102 | ) 103 | far_box_corners = (box_corners - u0v0) / np.array( 104 | [fku / far_clip, -fkv / far_clip], dtype=C.dtype 105 | ) 106 | ret_xy = np.concatenate( 107 | [near_box_corners, far_box_corners], axis=0) # [8, 2] 108 | ret_xyz = np.concatenate([ret_xy, z_points], axis=1) 109 | return ret_xyz 110 | 111 | 112 | @numba.jit(nopython=True) 113 | def corner_to_surfaces_3d_jit(corners): 114 | """convert 3d box corners from corner function above 115 | to surfaces that normal vectors all direct to internal. 116 | 117 | Args: 118 | corners (float array, [N, 8, 3]): 3d box corners. 119 | Returns: 120 | surfaces (float array, [N, 6, 4, 3]): 121 | """ 122 | # box_corners: [N, 8, 3], must from corner functions in this module 123 | num_boxes = corners.shape[0] 124 | surfaces = np.zeros((num_boxes, 6, 4, 3), dtype=corners.dtype) 125 | corner_idxes = np.array( 126 | [0, 1, 2, 3, 7, 6, 5, 4, 0, 3, 7, 4, 1, 5, 6, 2, 0, 4, 5, 1, 3, 2, 6, 7] 127 | ).reshape(6, 4) 128 | for i in range(num_boxes): 129 | for j in range(6): 130 | for k in range(4): 131 | surfaces[i, j, k] = corners[i, corner_idxes[j, k]] 132 | return surfaces 133 | 134 | 135 | def points_in_convex_polygon_3d_jit(points, polygon_surfaces, num_surfaces=None): 136 | """check points is in 3d convex polygons. 137 | Args: 138 | points: [num_points, 3] array. 139 | polygon_surfaces: [num_polygon, max_num_surfaces, 140 | max_num_points_of_surface, 3] 141 | array. all surfaces' normal vector must direct to internal. 142 | max_num_points_of_surface must at least 3. 143 | num_surfaces: [num_polygon] array. indicate how many surfaces 144 | a polygon contain 145 | Returns: 146 | [num_points, num_polygon] bool array. 147 | """ 148 | max_num_surfaces, max_num_points_of_surface = polygon_surfaces.shape[1:3] 149 | num_points = points.shape[0] 150 | num_polygons = polygon_surfaces.shape[0] 151 | if num_surfaces is None: 152 | num_surfaces = np.full((num_polygons,), 9999999, dtype=np.int64) 153 | normal_vec, d = surface_equ_3d_jitv2(polygon_surfaces[:, :, :3, :]) 154 | # normal_vec: [num_polygon, max_num_surfaces, 3] 155 | # d: [num_polygon, max_num_surfaces] 156 | return _points_in_convex_polygon_3d_jit( 157 | points, polygon_surfaces, normal_vec, d, num_surfaces 158 | ) 159 | 160 | 161 | @numba.njit 162 | def surface_equ_3d_jitv2(surfaces): 163 | # polygon_surfaces: [num_polygon, num_surfaces, num_points_of_polygon, 3] 164 | num_polygon = surfaces.shape[0] 165 | max_num_surfaces = surfaces.shape[1] 166 | normal_vec = np.zeros((num_polygon, max_num_surfaces, 3), dtype=surfaces.dtype) 167 | d = np.zeros((num_polygon, max_num_surfaces), dtype=surfaces.dtype) 168 | sv0 = surfaces[0, 0, 0] - surfaces[0, 0, 1] 169 | sv1 = surfaces[0, 0, 0] - surfaces[0, 0, 1] 170 | for i in range(num_polygon): 171 | for j in range(max_num_surfaces): 172 | sv0[0] = surfaces[i, j, 0, 0] - surfaces[i, j, 1, 0] 173 | sv0[1] = surfaces[i, j, 0, 1] - surfaces[i, j, 1, 1] 174 | sv0[2] = surfaces[i, j, 0, 2] - surfaces[i, j, 1, 2] 175 | sv1[0] = surfaces[i, j, 1, 0] - surfaces[i, j, 2, 0] 176 | sv1[1] = surfaces[i, j, 1, 1] - surfaces[i, j, 2, 1] 177 | sv1[2] = surfaces[i, j, 1, 2] - surfaces[i, j, 2, 2] 178 | normal_vec[i, j, 0] = sv0[1] * sv1[2] - sv0[2] * sv1[1] 179 | normal_vec[i, j, 1] = sv0[2] * sv1[0] - sv0[0] * sv1[2] 180 | normal_vec[i, j, 2] = sv0[0] * sv1[1] - sv0[1] * sv1[0] 181 | 182 | d[i, j] = ( 183 | -surfaces[i, j, 0, 0] * normal_vec[i, j, 0] 184 | - surfaces[i, j, 0, 1] * normal_vec[i, j, 1] 185 | - surfaces[i, j, 0, 2] * normal_vec[i, j, 2] 186 | ) 187 | return normal_vec, d 188 | 189 | @numba.njit 190 | def _points_in_convex_polygon_3d_jit( 191 | points, polygon_surfaces, normal_vec, d, num_surfaces=None 192 | ): 193 | """check points is in 3d convex polygons. 194 | Args: 195 | points: [num_points, 3] array. 196 | polygon_surfaces: [num_polygon, max_num_surfaces, 197 | max_num_points_of_surface, 3] 198 | array. all surfaces' normal vector must direct to internal. 199 | max_num_points_of_surface must at least 3. 200 | num_surfaces: [num_polygon] array. indicate how many surfaces 201 | a polygon contain 202 | Returns: 203 | [num_points, num_polygon] bool array. 204 | """ 205 | max_num_surfaces, max_num_points_of_surface = polygon_surfaces.shape[1:3] 206 | num_points = points.shape[0] 207 | num_polygons = polygon_surfaces.shape[0] 208 | ret = np.ones((num_points, num_polygons), dtype=np.bool_) 209 | sign = 0.0 210 | for i in range(num_points): 211 | for j in range(num_polygons): 212 | for k in range(max_num_surfaces): 213 | if k > num_surfaces[j]: 214 | break 215 | sign = ( 216 | points[i, 0] * normal_vec[j, k, 0] 217 | + points[i, 1] * normal_vec[j, k, 1] 218 | + points[i, 2] * normal_vec[j, k, 2] 219 | + d[j, k] 220 | ) 221 | if sign >= 0: 222 | ret[i, j] = False 223 | break 224 | return ret 225 | 226 | def mask_points_by_range(points, limit_range): 227 | mask = (points[:, 0] >= limit_range[0]) & (points[:, 0] <= limit_range[3]) \ 228 | & (points[:, 1] >= limit_range[1]) & (points[:, 1] <= limit_range[4]) 229 | return mask 230 | 231 | def remove_outside_points(points, rect, Trv2c, P2, image_shape): 232 | # 5x faster than remove_outside_points_v1(2ms vs 10ms) 233 | C, R, T = projection_matrix_to_CRT_kitti(P2) 234 | image_bbox = [0, 0, image_shape[1], image_shape[0]] 235 | frustum = get_frustum(image_bbox, C) 236 | frustum -= T 237 | frustum = np.linalg.inv(R) @ frustum.T 238 | frustum = camera_to_lidar(frustum.T, rect, Trv2c) 239 | frustum_surfaces = corner_to_surfaces_3d_jit(frustum[np.newaxis, ...]) 240 | indices = points_in_convex_polygon_3d_jit(points[:, :3], frustum_surfaces) 241 | points = points[indices.reshape([-1])] 242 | return points 243 | --------------------------------------------------------------------------------