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

最新推荐文章于 2024-07-25 18:10:33 发布

自动驾驶小学生

最新推荐文章于 2024-07-25 18:10:33 发布

阅读量6.5k

点赞数 15

分类专栏：论文笔记

本文链接：https://blog.csdn.net/cg129054036/article/details/107523376

版权

论文笔记专栏收录该内容

65 篇文章 170 订阅

订阅专栏

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

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

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

文章目录

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)
            continue
        logger.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_eval
    assert len(cfg_list) % 2 == 0
    for k, v in zip(cfg_list[0::2], cfg_list[1::2]):
        key_list = k.split('.')
        d = config
        for subkey in key_list[:-1]:
            assert subkey in d, 'NotFoundKey: %s' % subkey
            d = d[subkey]
        subkey = key_list[-1]
        assert subkey in d, 'NotFoundKey: %s' % subkey
        try: 
            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_val
        elif 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_list
        else:
            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 = 0

if __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 cfg

class DatasetTemplate(torch_data.Dataset):
	def __init__(self):
        super().__init__()

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

输入的是点云数据帧编号 $（ i d x ）$ 和原始点云数据 $（ N ， 3 + C 1 ）$ ，以字典形式传输，输出为：

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

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

def prepare_data(self, input_dict):
	"""
    :param input_dict:
    	sample_idx: string
        points: (N, 3 + C1)
    :return:
        voxels: (N, max_points_of_each_voxel, 3 + C2), float
        num_points: (N), int
        coordinates: (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_points
    voxels, coordinates, num_points = self.voxel_generator.generate(points, \
                                      max_voxels=cfg.DATA_CONFIG[self.mode].MAX_NUMBER_OF_VOXELS)    
       
    # voxel_centers
    voxel_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 = 0
            batch_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_boxes3d
        else:
            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)$ ，数据类型为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 futures

    info = {}
    pc_info = {'num_features':4, 'lidar_idx': idx}
    info['point_cloud'] = pc_info
    return 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 -1
    boxes3d_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 = logger
        self.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_GENERATOR
        
    self.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_path
            with 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)))

此外我们也可以对点云进行筛选，下面的代码为选取 $x$ 在 $[0, 70.4]$ ， $y$ 在 $[- 40, 40]$ ， $z$ 在 $[- 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_merge
	 return 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_idx
          
    return 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_dir

    dataset = __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`}

自动驾驶小学生

关注

15
点赞
踩
70

收藏

觉得还不错? 一键收藏
7
评论
复制链接

分享到 QQ

分享到新浪微博

扫一扫

专栏目录