从零实现一个3D目标检测算法（2）：点云数据预处理

在上一篇文章《从零实现一个3D目标检测算法（1）：3D目标检测概述》对3D目标检测研究现状和PointPillars模型进行了介绍，在本文中我们开始写代码一步步实现PointPillars，这里我们先实现如何对点云数据进行预处理。

在图像目标检测中，一般不需要对图像进行预处理操作，直接输入原始图像即可得到最终的检测结果。

但是在点云3D目标检测中，往往需要对点云进行一定的预处理，本文将介绍在PointPillars模型中如何对点云进行预处理。这里的点云数据预处理操作同样也适用其它的基于Voxels的3D检测模型中。

文章目录

1. 模型配置文件config.py
- 1.1 将模型参数保存在日志文件
- 1.2 加载模型配置文件
- 1.3 解析终端命令修改模型配置参数
2. 点云数据预处理
- 2.1 DatasetTemplate类
- 2.2 KittiDataset类
- 2.3 KITTI数据加载器

1. 模型配置文件config.py

在这里我们将首先编写在整个工程中最重要的config.py文件，该文件主要包括三个函数。

作用分别是：加载模型配置文件pointpillar.yaml、将模型参数保存在日志文件中、以及解析终端命令修改模型配置参数。

关于上述三个函数，只需要会使用即可。首先导入需要的Python库：

from easydict import EasyDict
from pathlib import Path
import yaml

1.1 将模型参数保存在日志文件

这一部分是将整个网络模型的全部参数保存到日志文件中，在模型训练过程中每一个模块的代码往往会修改很多次。有了日志文件，我们就能很方便地查看每次所修改的地方，如果有疑问的话，可以借助日志文件快速定位问题，代码如下：

def log_config_to_file(cfg, pre='cfg', logger=None):for key, val in cfg.items():if isinstance(cfg[key], EasyDict):logger.info('\n%s.%s = edict()' % (pre, key))log_config_to_file(cfg[key], pre=pre+ '.' + key, logger=logger)continuelogger.info('%s.%s: %s' % (pre, key, val))

1.2 加载模型配置文件

下面一个函数是从配置文件pointpillar.yaml中加载网络模型参数。

在Python中我们使用字典这种数据类型来存储网络的各种参数，只需要命名好参数名称即可，如测试集，训练集名称，网络各子模块名称，损失函数名称等，在修改时也只需要修改对应参数的变量值即可，这是一个很方便的调参方式。代码如下：

def cfg_from_yaml_file(cfg_file, config):with open(cfg_file, 'r') as f:try:new_config = yaml.load(f, Loader=yaml.FullLoader)except:new_config = yaml.load(f)config.update(EasyDict(new_config))return config

1.3 解析终端命令修改模型配置参数

除了对模型配置文件.yaml进行修改外，也可以在执行时通过终端来修改模型的参数。

这时就要求程序能够获取终端信息，包括参数名称以及参数值，通常是成对出现，代码如下：

def cfg_from_list(cfg_list, config):"""Set config keys via list (e.g., from command line)."""from ast import literal_evalassert len(cfg_list) % 2 == 0for k, v in zip(cfg_list[0::2], cfg_list[1::2]):key_list = k.split('.')d = configfor subkey in key_list[:-1]:assert subkey in d, 'NotFoundKey: %s' % subkeyd = d[subkey]subkey = key_list[-1]assert subkey in d, 'NotFoundKey: %s' % subkeytry: value = literal_eval(v)except:value = v if type(value) != type(d[subkey]) and isinstance(d[subkey], EasyDict):key_val_list = value.split('.')for src in key_val_list:cur_key, cur_val = src.split(':')val_type = type(d[subkey][cur_key])cur_val = val_type(cur_val)d[subkey][cur_key] = cur_valelif type(value) != type(d[subkey]) and isinstance(d[subkey], list):val_list = value.split('.')for k, x in enumerate(val_list):val_list[k] = type(d[subkey][0])(x)d[subkey] = val_listelse:assert type(value) == type(d[subkey]), \'type {} dose not match original type {}'.format(type(value), type(d[subkey]))d[subkey] = value

下面我们来定义模型参数配置变量cfg，其本身是一个字典，现在我们先定义它的根路径。

至此配置文件代码编写完毕，不妨可以调用cfg_from_yaml_file函数加载yaml文件看看模型参数加载是否正确。

cfg = EasyDict()
cfg.ROOT_DIR = (Path(__file__).resolve().parent / '../').resolve()
cfg.LOCAL_RANK = 0if __name__=='__main__':pass

2. 点云数据预处理

现在我们对KITTI数据集进行预处理，最终将其加载到PyTorch的DataLoader中。

2.1 DatasetTemplate类

首先是dataset.py文件，我们使用Python中的Class来对点云数据进行预处理，数据的预处理操作都定义为Class的成员函数。

先首先定义一个DatasetTemplate类，当做点云数据的一个基本类，后面处理其它点云数据集时可以在此基础上进行不同的操作，导入必要的Python库：

import numpy as np
from collections import defaultdict
import torch.utils.data as torch_data
import sys
sys.path.append('../')
sys.path.append('../../')
from utils import common_utils
from config import cfgclass DatasetTemplate(torch_data.Dataset):def __init__(self):super().__init__()

在DatasetTemplate中我们定义两个成员函数，一个是数据准备函数prepare_data。

输入的是点云数据帧编号（idx）（idx）（idx）和原始点云数据（N，3+C1）（N，3+C1）（N，3+C1），以字典形式传输，输出为：

Voxels
Voxels坐标
每个Voxels中点的个数
Voxels中心坐标（全局坐标）
原始点云数据

输出同样以字典形式输出。

def prepare_data(self, input_dict):""":param input_dict:sample_idx: stringpoints: (N, 3 + C1):return:voxels: (N, max_points_of_each_voxel, 3 + C2), floatnum_points: (N), intcoordinates: (N, 3), [idx_z, idx_y, idx_x]voxel_centers: (N, 3)points: (M, 3 + C)"""sample_idx = input_dict['sample_idx']points = input_dict['points']points = points[:, :cfg.DATA_CONFIG.NUM_POINT_FEATURES['use']]     # voxels, coordinates, num_pointsvoxels, coordinates, num_points = self.voxel_generator.generate(points, \max_voxels=cfg.DATA_CONFIG[self.mode].MAX_NUMBER_OF_VOXELS)    # voxel_centersvoxel_centers = (coordinates[:, ::-1] + 0.5) * self.voxel_generator.voxel_size \+ self.voxel_generator.point_cloud_range[0:3]print('voxel_centers.shape is: ', voxel_centers.shape)       # (11719, 3)if cfg.DATA_CONFIG.MASK_POINTS_BY_RANGE:points = common_utils.mask_points_by_range(points, cfg.DATA_CONFIG.POINT_CLOUD_RANGE)example = {}example.update({'voxels': voxels,'num_points': num_points,'coordinates': coordinates,'voxel_centers': voxel_centers,'points': points})return example

另一个函数是collate_batch，作用是在加载数据集时如何选取数据。

@staticmethod
def collate_batch(batch_list, _unused=False):example_merged = defaultdict(list)for example in batch_list:for k, v in example.items():example_merged[k].append(v)ret = {}for key, elems in example_merged.items():if key in ['voxels', 'num_points', 'voxel_centers', 'seg_labels', 'part_labels', 'bbox_reg_labels']:ret[key] = np.concatenate(elems, axis=0)elif key in ['coordinates', 'points']:coors = []for i, coor in enumerate(elems):coor_pad = np.pad(coor, ((0, 0), (1, 0)), mode='constant', constant_values=i)coors.append(coor_pad)ret[key] = np.concatenate(coors, axis=0)elif key in ['gt_boxes']:max_gt = 0batch_size = elems.__len__()for k in range(batch_size):max_gt = max(max_gt, elems[k].__len__())batch_gt_boxes3d = np.zeros((batch_size, max_gt, elems[0].shape[-1]), dtype=np.float32)for k in range(batch_size):batch_gt_boxes3d[k, :elems[k].__len__(), :] = elems[k]ret[key] = batch_gt_boxes3delse:ret[key] = np.stack(elems, axis=0)ret['batch_size'] = batch_list.__len__()return ret

2.2 KittiDataset类

现在我们编写kitti_dataset.py，主要目的是创造KittiDataset类，首先是导入所需库：

import os
import sys
import pickle
import copy
import numpy as np
from pathlib import Path
import torch
import sys
sys.path.append('../')
sys.path.append('../../')
from config import cfg
from spconv.utils import VoxelGenerator
from ..dataset import DatasetTemplate

在这里我们首先定义一个BaseKittiDataset类，这里初始化只有一个参数，就是点云数据的存储路径root_path。

class BaseKittiDataset(DatasetTemplate):def __init__(self, root_path):super().__init__()self.root_path = root_path

现在我们编写获取点云数据的get_lidar函数，KITTI中点云数据是以二进制格式保存的，每个点有4个信息：(x,y,z,r)(x,y,z,r)(x,y,z,r)，数据类型为float32，代码如下：

def get_lidar(self, idx):lidar_file = os.path.join(self.root_path, 'velodyne', '%06d.bin' % idx)assert os.path.exists(lidar_file)return np.fromfile(lidar_file, dtype=np.float32).reshape([-1, 4])

此外我们也可以编写函数get_infos来获取点云信息，具体为：

def get_infos(self, idx):import concurrent.futures as futuresinfo = {}pc_info = {'num_features':4, 'lidar_idx': idx}info['point_cloud'] = pc_inforeturn info

这里有一个生成最终预测结果的函数，因为模型计算时使用的是GPU，而要保存时需要转化为CPU可访问的数据。

预测信息有box尺寸box3d_lidar，分值scores，目标类型标签label_preds，以及点云编号sample_idx。

@staticmethod
def generate_prediction_dict(input_dict, index, record_dict):# finally generate predictions.sample_idx = input_dict['sample_idx'][index] if 'sample_idx' in input_dict else -1boxes3d_lidar_preds = record_dict['boxes'].cpu().numpy()if boxes3d_lidar_preds.shape[0] == 0:return {'sample_idx': sample_idx}predictions_dict ={'box3d_lidar': boxes3d_lidar_preds,'scores': record_dict['scores'].cpu.numpy(),'label_preds': record_dict['labels'].cpu().numpy(),'sample_idx': sample_idx}return predictions_dict

现在我们就可以创建KittiDataset类了，同样初始化时需要设置数据路径，这里我们需要将模式设置为TEST：

class KittiDataset(BaseKittiDataset):def __init__(self, root_path, logger=None):super().__init__(root_path=root_path)self.logger = loggerself.mode = 'TEST'self.kitti_infos = []self.include_kitti_data(self.mode, logger)self.dataset_init(logger)

在初始化时，有一个dataset_init函数，这个函数是用来生成voxel_generator的，使用的库为Spconv，在上面的prepare_data函数中会使用这个voxel_generator，代码如下：

def dataset_init(self, logger):voxel_generator_cfg = cfg.DATA_CONFIG.VOXEL_GENERATORself.voxel_generator = VoxelGenerator(voxel_size=voxel_generator_cfg.VOXEL_SIZE,point_cloud_range=cfg.DATA_CONFIG.POINT_CLOUD_RANGE,max_num_points=voxel_generator_cfg.MAX_POINTS_PER_VOXEL)

include_kitti_data函数是用来加载pkl文件的，我们会将待处理的点云信息存储在pkl文件中，这样测试模型时只需使用这一个文件就可以访问全部点云数据了：

def include_kitti_data(self, mode, logger):if cfg.LOCAL_RANK == 0 and logger is not None:logger.info('Loading KITTI dataset')kitti_infos = []for info_path in cfg.DATA_CONFIG[mode].INFO_PATH:        info_path = cfg.ROOT_DIR / info_pathwith open(info_path, 'rb') as f:infos = pickle.load(f)kitti_infos.append(infos)self.kitti_infos.extend(kitti_infos)if cfg.LOCAL_RANK == 0 and logger is not None:logger.info('Total samples for KITTI dataset: %d' % (len(kitti_infos)))

此外我们也可以对点云进行筛选，下面的代码为选取xxx在[0,70.4][0, 70.4][0,70.4]，yyy在[−40,40][-40, 40][−40,40]，zzz在[−3,1][-3, 1][−3,1]范围的点，这个一般要根据具体应用场景来设置。

@staticmethod
def get_valid_flag(pts_lidar):'''Valid points should be in the PC_AREA_SCOPE'''val_flag_x = np.logical_and(pts_lidar[:, 0]>=0, pts_lidar[:, 0]<=70.4)val_flag_y = np.logical_and(pts_lidar[:, 1]>=-40, pts_lidar[:, 1]<=40)val_flag_z = np.logical_and(pts_lidar[:, 2]>=-3, pts_lidar[:, 2]<=1)val_flag_merge = np.logical_and(val_flag_x, val_flag_y, val_flag_z)pts_valid_flag = val_flag_mergereturn pts_valid_flag

然后，就是编写__getitem__函数：

def __len__(self):return len(self.kitti_infos)def __getitem__(self, index):info = copy.deepcopy(self.kitti_infos[index])sample_idx = info['point_cloud']['lidar_idx']points = self.get_lidar(sample_idx)pts_valid_flag = self.get_valid_flag(points[:, 0:3])points = points[pts_valid_flag]input_dict = {'points': points, 'sample_idx': sample_idx}example = self.prepare_data(input_dict=input_dict)example['sample_idx'] = sample_idxreturn example

下面是create_kitti_infos函数：

def create_kitti_infos(data_path, save_path):dataset = BaseKittiDataset(root_path=data_path)val_filename = save_path / ('kitti_infos_val.pkl')print('val_filename is: ', val_filename)print('---------------Start to generate data infos---------------')kitti_infos_val = dataset.get_infos(idx)print(kitti_infos_val)with open(val_filename, 'wb') as f:pickle.dump(kitti_infos_val, f)print('Kitti info val file is saved to %s' % val_filename)

最后编写main函数，函数主要作用是获取终端信息，生成kitti_infos。

if __name__=='__main__':if sys.argv.__len__() > 1 and sys.argv[1] == 'create_kitti_infos':create_kitti_infos(data_path=cfg.ROOT_DIR / 'data',save_path=cfg.ROOT_DIR / 'data')

生成后的kitti_infos如下：

{'point_cloud': {'num_features': 4, 'lidar_idx': '000010'}}

2.3 KITTI数据加载器

现在编写__init__.py，这里的作用是通过DataLoader加载点云数据，这在PyTorch是十分常见的，代码如下：

import os
from pathlib import Path
import torch
from torch.utils.data import DataLoader
from .kitti.kitti_dataset import KittiDataset, BaseKittiDataset
from config import cfg__all__ = {'BaseKittiDataset': BaseKittiDataset,'KittiDataset': KittiDataset}def build_dataloader(data_dir, batch_size, logger=None):data_dir = Path(data_dir) if os.path.isabs(data_dir) else cfg.ROOT_DIR / data_dirdataset = __all__[cfg.DATA_CONFIG.DATASET](root_path=data_dir, logger=logger)dataloader = DataLoader(dataset, batch_size=batch_size, pin_memory=True, shuffle=False, collate_fn=dataset.collate_batch, drop_last=False)return dataset, dataloader

至此，点云数据预处理部分我们就已经完成了，预处理后的点云数据将变成如下形式：

原始points
voxels及其坐标
voxels中心位置
点的数量
点云帧编号
batch_size

下一篇文章中我们将开始实现PointPillars的网络部分。

input_dict`：{`voxels`, `num_points`, `coordinates`, `voxel_centers` ,  `points`, `sample_idx`,  `batch_size`}