PPLanedet:用于车道线检测的工具包

★★★ 本文源自AI Studio社区精品项目，【点击此处】查看更多精品内容 >>>

---

PPLanedet: 基于PaddlePaddle的车道线检测工具包

github传送门

1、引言

     目前Aistudio中对于车道线检测的方法比较散，各位大佬复现的风格各不相同，为了能够方便科研人员或开发者能够通过一个框架来实现不同车道线检测算法，我们开源了PPLanedet工具包，PPLanedet属于个人开发项目，并不属于PaddlePaddle官方的框架(例如Paddledetection、Paddleseg等)，因此我们希望大家能够加入到我们去不断完善该工具包。

2、PPLanedet特点

1、将车道线检测框架解耦成不同的模块组件，通过组合不同的模块组件，用户可以便捷地构建自定义的车道线检测模型；

2、目前PPLanedet复现了SCNN、ERFNet、UFLD三种算法，覆盖了基于分割模型和基于关键点的两类车道线检测模型；

3、训练速度较快，SCNN训练只需要2个小时，UFLD也只需要1个多小时；

4、预训练的backbone与Paddleseg对接，Paddleseg中backbone预训练权重PPLanedet中均可使用。

3、安装PPLanedet

# git clone PPLanedet
!git clone https://github.com/zkyseu/PPlanedet

%cd PPlanedet

#安装依赖库
!pip install -r requirements.txt --user

!python setup.py build install

4、模块组建介绍

     1、首先是数据集代码，代码在PPlanedet/pplanedet/datasets文件夹下面。目前PPLanedet提供了两种在车道线检测中应用最为广泛的数据集: Tusimple和CULane。这里我们以Tusimple讲述一下如何准备训练数据集。

#1、准备数据集

#首先需要下载Tusimple数据集，aistudio上面已经有人上传了，因此我们直接挂载数据集即可，如果本地的话则需要下载数据集
#下载地址为https://github.com/TuSimple/tusimple-benchmark/issues/3

#2、解压数据集

#这里可以解压到任何你想要解压的文件夹当中，aistudio里面我们就解压到data文件夹下面，解压到data文件夹下面可以避免服务器保存数据过大，导致下一次加载过慢
!unzip -d /home/aistudio/data /home/aistudio/data/data119539/train_set.zip
!unzip -d /home/aistudio/data /home/aistudio/data/data119539/test_set.zip

#标签需要和解压数据在同一层目录，因此这里需要copy一份标签
!cp data/data119539/test_label.json data/

#具体的数据集结构可以如下所示，大家对照一下看看自己本地数据集格式和预期的是否一样
"""
$TUSIMPLEROOT/clips # data folders
$TUSIMPLEROOT/lable_data_xxxx.json # label json file x4
$TUSIMPLEROOT/test_tasks_0627.json # test tasks json file
$TUSIMPLEROOT/test_label.json # test label json file

如果在aistudio TUSIMPLEROOT=/home/aistudio/data/
"""

3、解压完数据后再配置文件里面定义一下数据集信息，具体如下，配置文件路径在 /home/aistudio/PPlanedet/configs 下面

dataset_path = '/home/aistudio/data/' #这里就是TUSIMPLEROOT路径

dataset = dict(
    train=dict(
        name='TuSimple',
        data_root=dataset_path,
        split='trainval',
        processes=train_transform, #这里是数据增强下面会讲到
    ),
    val=dict(
        name='TuSimple',
        data_root=dataset_path,
        split='test',
        processes=val_transform,
    ),
    test=dict(
        name='TuSimple',
        data_root=dataset_path,
        split='test',
        processes=val_transform,
    )
)

test_json_file='/home/fyj/zky/tusimple/test_label.json' #测试集标签位置

好啦，以上就是配置Tusimple数据集的方式，是不是很简单。可能有同学有疑问数据集怎么调用，这里PPLanedet中采用模块注册机制来调用各个已经定义好的模块。注册代码在 /home/aistudio/PPlanedet/pplanedet/utils/registry.py 里面。简单来说，大家可以这么想，就是我们通过注册一个个模块让他们在一个系统里面可以通过配置文件方式方便地进行调用。

2、接下来和大家介绍一下数据集定义，总的来说就是先定义一个BaseDateset用于最后dataloader数据读取，数据增强等分成具体数据集(例如Tusimple和CUlane)。BaseDataset代码如下。

# @DATASETS.register() 实际运行该代码需要解除注释
class BaseDataset(Dataset):
    def __init__(self, data_root, split, processes=None,
            cfg=None):
        self.cfg = cfg
        self.logger = logging.getLogger(__name__)
        self.data_root = data_root
        self.training = 'train' in split 
        self.processes = Compose(processes, cfg)


    def view(self, predictions, img_metas):
        """
        这个函数作用是验证时候可视化测试集
        """
        img_metas = [item for img_meta in img_metas.data for item in img_meta]
        for lanes, img_meta in zip(predictions, img_metas):
            img_name = img_meta['img_name']
            img = cv2.imread(osp.join(self.data_root, img_name))
            out_file = osp.join(self.cfg.work_dir, 'visualization',
                                img_name.replace('/', '_'))
            lanes = [lane.to_array(self.cfg) for lane in lanes]
            imshow_lanes(img, lanes, out_file=out_file)

    def __len__(self):
        return len(self.data_infos)

    def __getitem__(self, idx):
        """
        具体流程为:1、读取数据路径; 2、读取标签; 3、数据预处理
        """
        data_info = self.data_infos[idx] # 数据集路径具体定义见Tusimple和CULane
        if not osp.isfile(data_info['img_path']):
            raise FileNotFoundError('cannot find file: {}'.format(data_info['img_path']))

        img = cv2.imread(data_info['img_path'])

        img = img[self.cfg.cut_height:, :, :]
        sample = data_info.copy()
        sample.update({'img': img})

        if self.training:  #训练时候读取标签
            label = cv2.imread(sample['mask_path'], cv2.IMREAD_UNCHANGED)
            if len(label.shape) > 2:
                label = label[:, :, 0]
            label = label.squeeze()
            label = label[self.cfg.cut_height:, :]
            sample.update({'mask': label})

        sample = self.processes(sample) #数据预处理

        return sample 

这个部分我们介绍如何增加自己的自定义数据集。我们以Tusimple数据集为例子

1、定义自己数据集类并继承BaseDataset(如果觉得BaseDataset类不符合自己要求可以直接定义为paddle.io.dataset)

2、完善自己数据集类中数据集路径的定义，数据预处理，评价指标等。这里就完成了数据集类的定义。

3、在数据集类定义上加上@DATASETS.register() 表明要注册自己数据集。

4、不能忘记一步，需要在datasets这个文件夹里的__init__.py文件里import自己定义的数据集类

看完步骤我们来看代码

# @DATASETS.register() #表明自己要定义下面这个数据集函数，这一行要紧跟要定义数据集类上一行。
class TuSimple(BaseDataset): #继承BaseDataset
    def __init__(self, data_root, split, processes=None, cfg=None):
        super().__init__(data_root, split, processes, cfg)
        self.anno_files = SPLIT_FILES[split] 
        self.load_annotations()
        self.h_samples = list(range(160, 720, 10))

    def load_annotations(self): #定义数据集路径
        self.logger.info('Loading TuSimple annotations...')
        self.data_infos = []
        max_lanes = 0
        for anno_file in self.anno_files:
            anno_file = osp.join(self.data_root, anno_file)
            with open(anno_file, 'r') as anno_obj:
                lines = anno_obj.readlines()
            for line in lines:
                data = json.loads(line)
                y_samples = data['h_samples']
                gt_lanes = data['lanes']
                mask_path = data['raw_file'].replace('clips', 'seg_label')[:-3] + 'png'
                lanes = [[(x, y) for (x, y) in zip(lane, y_samples) if x >= 0] for lane in gt_lanes]
                lanes = [lane for lane in lanes if len(lane) > 0]
                max_lanes = max(max_lanes, len(lanes))
                self.data_infos.append({
                    'img_path': osp.join(self.data_root, data['raw_file']),
                    'img_name': data['raw_file'],
                    'mask_path': osp.join(self.data_root, mask_path),
                    'lanes': lanes,
                })

        if self.training:
            random.shuffle(self.data_infos)
        self.max_lanes = max_lanes
    #=============下面就是一些评价指标=====================
    def pred2lanes(self, pred):
        ys = np.array(self.h_samples) / self.cfg.ori_img_h
        lanes = []
        for lane in pred:
            xs = lane(ys)
            invalid_mask = xs < 0
            lane = (xs * self.cfg.ori_img_w).astype(int)
            lane[invalid_mask] = -2
            lanes.append(lane.tolist())

        return lanes

    def pred2tusimpleformat(self, idx, pred, runtime):
        runtime *= 1000.  # s to ms
        img_name = self.data_infos[idx]['img_name']
        lanes = self.pred2lanes(pred)
        output = {'raw_file': img_name, 'lanes': lanes, 'run_time': runtime}
        return json.dumps(output)

    def save_tusimple_predictions(self, predictions, filename, runtimes=None):
        if runtimes is None:
            runtimes = np.ones(len(predictions)) * 1.e-3
        lines = []
        for idx, (prediction, runtime) in enumerate(zip(predictions, runtimes)):
            line = self.pred2tusimpleformat(idx, prediction, runtime)
            lines.append(line)
        with open(filename, 'w') as output_file:
            output_file.write('\n'.join(lines))

    def evaluate(self, predictions, output_basedir, runtimes=None):
        if not os.path.exists(output_basedir):
            os.mkdir(output_basedir)
        pred_filename = os.path.join(output_basedir, 'tusimple_predictions.json')
        self.save_tusimple_predictions(predictions, pred_filename, runtimes)
        acc = 0
        try:
            json_pred = [json.loads(line)for line in open(pred_filename).readlines()]
            if len(json_pred) == 0:
                acc = -1
        except:
            acc = -1
        if acc == -1:
            return acc
        
        result, acc = LaneEval.bench_one_submit(pred_filename, self.cfg.test_json_file)
        self.logger.info(result)
        return acc

#在 __init__.py里面import自己定义的数据集类
from .tusimple import TuSimple

看完这个部分，大家可以去尝试定义自己数据级，并在配置文件里面完成数据集调用。其余的关于模型的自定义等都和这个类似，之后就不赘述了。

接着，我们来讲讲怎么来看模型，首先我们来看一下配置文件，以SCNN为例。SCNN的配置文件在 /home/aistudio/PPlanedet/configs/scnn/resnet50tusimple.py 里面。PPLanedet将模型分成了architecture、backbone、aggregators、head四个部分，architecture是指模型整体结构，如果不是特殊的Detecor类就足够了。Backbone就是指特征提取网络(例如ResNet)，aggregators是指一些上下文特征提取模块例如SCNN、ASPP等，head是指模型的decoder以及预测head。

model = dict(
    name='Detector',
)

backbone = dict(
    name='ResNet',
    output_stride=8,
    multi_grid=[1, 1, 1],
    return_idx=[0,1,2],
    pretrained='https://bj.bcebos.com/paddleseg/dygraph/resnet50_vd_ssld_v2.tar.gz',
)

featuremap_out_channel = 128

aggregator = dict(
    name='SCNN',
)

sample_y=range(710, 150, -10)


heads = dict(
    name='LaneSeg',
    decoder=dict(name='PlainDecoder'),
    thr=0.6,
    seg_loss = dict(name = 'MultiClassFocalLoss',
                   num_class = 6+1,
                   loss_weight = 1.),
    sample_y=sample_y,
)

并且，PPLanedet中将各个模块分文件夹放置，也体现了PPLanedet模块化管理的特点。这样可以使得代码整洁，可扩展性更好。

最后便是训练上的配置，包括训练轮数，优化器，数据增强，模型输出路径等，大家可参照给的配置文件例子来配置自己的模型。

epochs = 100
batch_size = 10
total_iter = (3616 // batch_size + 1) * epochs 

lr_scheduler = dict(
    name = 'PolynomialDecay',
    learning_rate = 0.025,
    decay_steps = total_iter
)

optimizer = dict(
  name = 'Momentum',
  weight_decay = 1e-4,
  momentum = 0.9
)

img_height = 368
img_width = 640
cut_height = 160
ori_img_h = 720
ori_img_w = 1280

img_norm = dict(
    mean=[0.5, 0.5, 0.5],
    std=[0.5, 0.5, 0.5]
)

train_transform = [
    dict(name='RandomRotation'),
    dict(name='RandomHorizontalFlip'),
    dict(name='Resize', size=(img_width, img_height)),
    dict(name='Normalize', img_norm=img_norm),
    dict(name='ToTensor'),
] 

val_transform = [
    dict(name='Resize', size=(img_width, img_height)),
    dict(name='Normalize', img_norm=img_norm),
    dict(name='ToTensor', keys=['img']),
] 

log_config = dict(
    name = 'LogHook',
    interval = 50
    )

custom_config = [dict(
    name = 'EvaluateHook'
    )]

device = 'gpu'
seed =  0
save_inference_dir = './inference'
output_dir = './output_dir'
best_dir = './output_dir/best_dir'
pred_save_dir = './pred_save'
num_workers = 4
num_classes = 6 + 1
view = False
ignore_label = 255
test_json_file='/home/fyj/zky/tusimple/test_label.json'

5、模型的训练、验证、demo

这个部分中我们介绍如何训练我们的模型，具体代码如下

#单卡训练
python tools/train.py -c configs/scnn/resnet50_tusimple.py

#多卡训练
# export CUDA_VISIBLE_DEVICES=0,1,2,3
# python -m paddle.distributed.launch tools/train.py -c configs/scnn/resnet50_tusimple.py

#使用预训练模型训练
# python tools/train.py -c configs/scnn/resnet50_tusimple.py --load pretrain weight path 

#恢复训练
#python tools/train.py -c configs/scnn/resnet50_tusimple.py --resume weight path #

# 不用训练，直接验证模型
python tools/train.py -c configs/scnn/resnet50_tusimple.py \
                      --load model_path \
                      --evaluate-only 

# demo 测试单张图片效果
python tools/detect.py configs/scnn/resnet50_tusimple.py --img images\
          --load_from model.pth --savedir ./vis

# 如果是分割模型，想测试分割效果的
# 首先在配置文件中增加 'seg = True'
# python tools/detect.py configs/scnn/resnet50_tusimple.py --img images\
#           --load_from model.pth --savedir ./vis

6、PPLanedet相关的项目

如果大家对PPLanedet流程还有疑问可以看看我们之前公开的项目

车道线检测经典论文UFLD复现

结束语

好啦，以上便是PPLanedet全部流程，此文档还有很多不足之处，我们欢迎各位开发者帮助我们完善，我们也希望大家能加入到我们当中完善这个项目，最后别走开，喜欢就点个赞吧，也帮github点个star！谢谢大家！