简介基于U-Net的车道线检测实战资源面向图像分割入门者及自动驾驶感知方向的开发者。内容围绕TuSimple数据集完成模型训练与预测系统讲解U-Net收缩路径与扩展路径的编码-解码结构、跳跃连接原理并给出数据预处理、损失函数选择、训练调参与效果评估的方法。压缩包共15个文件约7.89MB以7个Python训练/测试/数据处理脚本为核心辅以2个说明文档、2个日志文本以及4个avi/mp4演示视频便于复现完整的训练与预测流程。目前已有581人学习适合希望结合开源数据集快速上手车道线检测项目、理解像素级分割任务并迁移到其他目标检测场景的读者。通过该包可获得一整套可运行的U-Net车道线检测代码、路面实线与虚线的测试效果录屏以及针对不同光照和积水路况的预测演示帮助快速评估模型鲁棒性并拓展实际部署思路。1. 用一个下午跑通U-Net车道线检测这份PyTorch工程到底该怎么下、怎么改、怎么避坑车道线检测在自动驾驶里是一个典型的图像分割任务但很多人第一次上手时总下意识把它当成目标检测来做——画框、给类别、出置信度结果在弯道和实线虚线交替的场景里频繁翻车。真正实用的做法是把每条车道线当成像素级目标用U-Net这种编码器-解码器结构去做语义分割输入一张图输出一张和原图同尺寸的mask车道线像素标1背景标0。这份zip包是一个基于PyTorch的U-Net车道线检测工程已经在TuSimple数据集上完成了训练和推理验证代码里带了train.py、dataset.py、model.py、predict.py、test_onvideo.py这几个核心脚本还有训练好的checkpoints和验证视频适合想快速看到U-Net跑车道线效果的人也适合想在此基础上改自己的数据集的从业者。本文不吹不黑直接拆解这个工程的代码结构、训练参数和实际踩过的坑。2. TuSimple数据集的落地处理从json标注到二值mask的完整转换2.1 TuSimple的标注格式为什么不能直接拿来训练不少人是先把TuSimple下载下来解压后直接写Dataloader结果发现训练出来的模型输出一团糊。原因很简单TuSimple的标注不是像素级mask而是json格式的曲线点坐标。每张图对应一个json文件里面有一个lanes字段存放着每条车道线在图像上的离散点位置还有一个tions字段存放的是样本划分关系。这种格式的优点是标注成本低、人工只需要标关键点但缺点是U-Net训练时需要的是和输入图像尺寸一致的二值图。所以拿到这份工程后第一步要做的是把json标注转成mask图。转换的逻辑并不复杂对每一条车道线取它的所有点坐标用cv2.polylines把点连成线再用cv2.fillPoly把线加宽成一条有厚度的像素带。# process_label.py 核心逻辑简化版 import cv2 import numpy as np import json def json_to_mask(json_path, img_shape(720, 1280)): # img_shape 为 (H, W)TuSimple 原始图像是 720x1280 mask np.zeros((img_shape[0], img_shape[1]), dtypenp.uint8) with open(json_path, r, encodingutf-8) as f: anno json.load(f) for lane in anno[lanes]: # lane 是一个长度为 N 的列表元素是 x 坐标 # 对应的 y 坐标在 anno[h_samples] 中两者一一对应 pts [] for x, y in zip(lane, anno[h_samples]): if x 0: pts.append([x, y]) if len(pts) 2: continue pts np.array(pts, dtypenp.int32).reshape(-1, 1, 2) # 先把点连成线再fillPoly加宽度避免断线 cv2.polylines(mask, [pts], isClosedFalse, color1, thickness5) cv2.fillPoly(mask, [pts], color1) return mask这里有一个容易被忽略的细节TuSimple标注里x坐标为-1的点表示该位置没有车道线转换时一定要过滤掉否则cv2.polylines会画出莫名其妙的斜线。另外thickness5是为了让车道线在mask里有一定宽度这个值需要根据训练效果调整太窄了模型学不到车道线的连通性太宽了会让实线和虚线的边界模糊。2.2 dataset.py的设计训练时做在线增强评估时保持干净这份工程里的dataset.py没有把预处理结果提前存盘而是采用在线处理的方式。# dataset.py 关键结构 from torch.utils.data import Dataset import torchvision.transforms as T class LaneDataset(Dataset): def __init__(self, data_list, img_h256, img_w512, is_trainTrue): self.data_list data_list # 每一行是 图像路径,json路径 self.img_h img_h self.img_w img_w self.is_train is_train def __len__(self): return len(self.data_list) def __getitem__(self, idx): img_path, json_path self.data_list[idx].split(,) # 读原图转RGBTuSimple原图就是RGB三通道 img cv2.imread(img_path) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask json_to_mask(json_path, img_shapeimg.shape[:2]) # 统一resize到模型输入尺寸 img cv2.resize(img, (self.img_w, self.img_h)) mask cv2.resize(mask, (self.img_w, self.img_h), interpolationcv2.INTER_NEAREST) if self.is_train: # 训练时加轻度增强随机水平翻转 随机亮度抖动 if np.random.rand() 0.5: img cv2.flip(img, 1) mask cv2.flip(mask, 1) img T.ToTensor()(img) img T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225])(img) mask torch.from_numpy(mask.astype(np.float32)) return img, mask在线增强的收益在这个任务上很大。TuSimple训练集有三千多张图像如果不做增强U-Net很容易过拟合到路面的固定纹理上尤其是高速公路场景下路面高度相似模型会直接记忆“哪条路有线”而不是学习“线长什么样”。随机水平翻转对车道线检测尤其有效因为车道线本身具有左右对称性。亮度抖动则是为了让模型在阴影、逆光路段更鲁棒。注意mask的resize必须用INTER_NEAREST如果用了双线性插值边界会出现灰值训练时会被当成介于0和1之间的模糊标签导致预测结果出现一圈灰边。2.3 图像尺寸的选择256x512和720x1280之间的取舍工程默认把训练图像resize到256x512这个选择在精度和显存之间取了平衡。U-Net的每一层下采样会把空间尺寸减半输入256x512经过四次下采样后特征图是16x32这个尺寸足够模型捕捉到“线在路面上”的全局结构同时也保留了足够的定位细节。如果把输入放大到720x1280显存占用会暴涨在普通消费级显卡上很容易出OOM。但在实际训练中我发现256x512的输入会在远处车道线位置造成严重的精度损失。TuSimple的标注线常常延伸到图像上边缘附近原图上那里的车道线宽度只有几个像素缩放到256x512后这些细线在mask里几乎消失。这也是为什么工程里的checkpoints在近处车道线上效果很好但在图像上边缘区域偶尔出现断线。如果你的场景需要重点检测远处车道线建议把输入尺寸改成320x640显存充裕的话可以试试384x768解码器上采样到原始分辨率后细节保留会明显改善。3. model.py中的U-Net结构收缩路径、扩展路径和跳跃连接的实现细节3.1 为什么这个任务选U-Net而不是DeepLab或FCN很多人在车道线检测里第一反应是某某分割模型但U-Net在这个场景里有一个不可替代的优点跳跃连接。车道线在图像里只占极少量像素背景占比常常超过95%这意味着模型很容易把注意力全放在背景上学到“所有东西都是背景”这个平庸解。DeepLab系列靠空洞卷积扩大感受野但对细线结构的保留能力弱FCN的浅层特征融合方式太粗糙恢复出来的边界不够锐利。U-Net的扩展路径每一层都会把上采样结果和收缩路径对应层的特征做concat这样浅层的边缘信息和深层的语义信息在同一个特征图上相遇。对于车道线这种既有明确局部方向、又有全局连续性约束的目标这种结构天然合适。这份工程里的model.py就是标准的U-Net实现没有做魔改入门者拿它当模板理解结构非常合适。3.2 模型代码逐段拆解通道数、上采样方式、激活函数的选择# model.py 核心结构 import torch import torch.nn as nn class DoubleConv(nn.Module): # 每次下采样/上采样之前先做两次卷积 def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x) class UNet(nn.Module): def __init__(self, in_channels3, out_channels1, features[64, 128, 256, 512]): super().__init__() # 收缩路径 self.enc1 DoubleConv(in_channels, features[0]) self.enc2 DoubleConv(features[0], features[1]) self.enc3 DoubleConv(features[1], features[2]) self.enc4 DoubleConv(features[2], features[3]) self.pool nn.MaxPool2d(2) # 瓶颈层 self.bottleneck DoubleConv(features[3], features[3] * 2) # 扩展路径 self.up4 nn.ConvTranspose2d(features[3] * 2, features[3], 2, stride2) self.dec4 DoubleConv(features[3] features[3], features[3]) self.up3 nn.ConvTranspose2d(features[3], features[2], 2, stride2) self.dec3 DoubleConv(features[2] features[2], features[2]) self.up2 nn.ConvTranspose2d(features[2], features[1], 2, stride2) self.dec2 DoubleConv(features[1] features[1], features[1]) self.up1 nn.ConvTranspose2d(features[1], features[0], 2, stride2) self.dec1 DoubleConv(features[0] features[0], features[0]) # 输出层把通道数压到1接sigmoid self.outc nn.Conv2d(features[0], out_channels, 1) def forward(self, x): e1 self.enc1(x) e2 self.enc2(self.pool(e1)) e3 self.enc3(self.pool(e2)) e4 self.enc4(self.pool(e3)) b self.bottleneck(self.pool(e4)) d4 self.dec4(torch.cat([self.up4(b), e4], dim1)) d3 self.dec3(torch.cat([self.up3(d4), e3], dim1)) d2 self.dec2(torch.cat([self.up2(d3), e2], dim1)) d1 self.dec1(torch.cat([self.up1(d2), e1], dim1)) return torch.sigmoid(self.outc(d1))代码里features列表[64, 128, 256, 512]控制的是每一层卷积核数量这个参数直接决定了模型的参数量和显存占用。需要说明的是它比原始U-Net论文里的通道数论文第一层是64设计得更收敛原始论文的第一层通道数是64但下采样次数更多。这份工程的Unet用ConvTranspose2d做上采样相比双线性插值后接卷积的做法它的优势是上采样核是可以学习的模型能根据任务自动决定如何在空间上填充特征缺点是如果kernel_size设置不当会产生棋盘格伪影这里用kernel_size2, stride2避免了重叠区域。跳跃连接时直接torch.cat通道数翻倍后输入dec层。这种做法会让解码器的前几层计算量变大但信息保留最完整。另一种做法是像某些变体那样先降通道再concat省显存但效果略差。对于车道线这种边缘信息至关重要的任务保持原样concat是大概率正确的选择。3.3 输出层设计为什么是sigmoid而不是softmax这个模型输出通道数是1用sigmoid做像素级二分类前背景等价于对每个像素独立判断“是不是车道线”。如果你把输出改成2通道再用softmax理论上也能做但实际效果差不多而且显存占用更大。关键点在于sigmoid配合BCE损失对每个像素的梯度是独立的不会像softmax那样在不同类别之间做竞争。车道线检测中前景像素数量远少于背景两个类别天然不平衡softmax归一化会让小类别的梯度被大类别的概率变化所压制。工程用sigmoid是更合理的选择。4. train.py训练流程拆解损失函数、优化器、学习率策略与验证逻辑4.1 训练脚本的整体结构与参数定义# train.py 关键参数 import argparse parser argparse.ArgumentParser() parser.add_argument(--epochs, typeint, default100) parser.add_argument(--batch_size, typeint, default8) parser.add_argument(--lr, typefloat, default1e-3) parser.add_argument(--img_h, typeint, default256) parser.add_argument(--img_w, typeint, default512) parser.add_argument(--dataset_path, typestr, default./data) parser.add_argument(--checkpoint_dir, typestr, default./checkpoints) parser.add_argument(--val_interval, typeint, default5, help每N个epoch做一次验证)训练时建议batch_size不要超过8因为U-Net在256x512输入下一张图的中间特征图显存占用已经不小。如果显存只有6GBbatch_size2也可以跑只是收敛速度会慢一些。epochs设100并不是说一定要跑满这个工程的验证逻辑会在每val_interval个epoch评估一次然后保存验证IoU最高的权重所以更重要的是训练中期能不能看到验证指标在上升。学习率用1e-3配合Adam是一个安全起点但注意Adam虽然收敛快后期容易出现验证loss平台期如果发现验证IoU在50轮之后纹丝不动可以考虑在第60轮手动把学习率降到1e-4通常能再往上走一截。4.2 损失函数的选择BCE与Dice的组合是车道线的关键# train.py 损失函数部分 import torch.nn.functional as F def dice_loss(pred, target, smooth1.0): pred pred.contiguous().view(-1) target target.contiguous().view(-1) intersection (pred * target).sum() return 1 - (2. * intersection smooth) / (pred.sum() target.sum() smooth) def bce_dice_loss(pred, target): bce F.binary_cross_entropy(pred, target, reductionmean) dice dice_loss(pred, target) return bce dice工程里同时计算BCE和Dice损失然后相加这是像素级分割任务里被验证过多次的组合方式。单独用BCE时由于背景像素占比极大模型会倾向于把一切都预测成背景因为这样loss已经很低了。Dice系数直接衡量预测和标注的重叠程度对前景区域的变化非常敏感即使只漏掉一小段车道线Dice也会明显下降。组合起来之后BCE保证每个像素都能被优化到Dice则把优化重心拉回到前景区域。这个组合在这份工程里是现成的不建议删掉其中一个。4.3 验证指标IoU的计算方式与可视化输出# train.py 验证部分 def compute_iou(pred, target, threshold0.5): pred (pred threshold).float() intersection (pred * target).sum() union pred.sum() target.sum() - intersection iou intersection / (union 1e-6) return iouIoU计算的核心是取阈值0.5把连续的概率图转成二值mask然后统计预测和标注的交集/并集。这里的细节点是union的计算方式——如果直接用torch.logical_or再sum数据类型处理不好容易报错工程里用pred.sum() target.sum() - intersection是更笨但绝对不出错的方式。注意使用threshold0.5评估模型部署效果如果某张图上预测结果中间的裂缝宽度明显大于真实车道线但形状对得上不代表模型学坏了更可能是阈值设定问题——调成0.3或者0.7再测一次就能看出模型输出的概率集中在哪个区间。5. 避坑指南TuSimple训练中四个最容易翻车的问题5.1 数据集划分陷阱误把训练视频帧当独立样本现象训练loss下降很快但验证IoU很低而且在验证集上预测的车道线位置明显偏移。原因TuSimple发布的数据集是视频帧序列相邻帧之间内容高度相似。如果把一组从同一段视频里截取的连续帧同时放进训练集和验证集模型其实已经“见过”验证集的画面了但实际场景中遇到的是全新的路况所以验证评估失真训练出来的模型泛化能力差。解决划分数据时按clip视频片段分组而不是按单张图像随机打散确保同一个视频片段的帧只出现在训练集或验证集中的一个。简单做法是给每个json文件名加上视频ID前缀按前缀做stratified split。5.2 车道线标注点稀疏导致的mask断裂现象用cv2.polylines转换后mask里有些弯道路段的车道线像虚线一样断开训练后模型预测的线也不连续。原因TuSimple的h_samples是固定y采样在弯道曲率大的位置x坐标变化快相邻标注点之间的像素距离拉大如果直接用直线连接会在弧线部分产生明显的切线与真实弧线的偏差某些情况下偏差过大导致相邻线段在视觉上断裂。解决转换标注时不要直接polyline连线而是在相邻标注点之间做一次线性插值或三次样条插值每两个点之间补出5到10个中间点再连线。# 插值补点示例 import numpy as np def densify_points(pts, n8): # pts: Nx2 数组按车道线方向排列 dense [] for i in range(len(pts) - 1): x1, y1 pts[i] x2, y2 pts[i 1] t np.linspace(0, 1, n) xs x1 (x2 - x1) * t ys y1 (y2 - y1) * t dense.extend(zip(xs, ys)) return np.array(dense)这段插值代码的作用是在两个标注点之间均匀插入n个采样点使得后续的polyline连线在弯道处更贴合真实车道线。注意只对相邻两个点距离超过一定阈值比如30像素的线段做插值点已经足够密的路段再插值意义不大反而增加计算量。5.3 训练过程中的显存OOM现象batch_size8时训练到第10轮突然报CUDA out of memory但前面一直正常。原因PyTorch在训练过程中每个step的中间激活值会保留到backward结束才释放。U-Net的跳跃连接让每一层解码器都保存了编码器的特征图这些特征图在batch_size8时占用已经逼近显存峰值当某些批次图像的纹理复杂度较高导致中间特征非零值更多时峰值会短暂超过显存容量。解决不要单一调低batch_size可以优先在train.py里设置torch.cuda.empty_cache()在每轮结束后调用并把torch.utils.checkpoint用在最深的编码器块上。使用gradient checkpointing会牺牲约20%的训练速度但能省出超过一半的激活值显存。# 调用梯度检查点的方法 from torch.utils.checkpoint import checkpoint_sequential # 假设enc4是一个包含多个子模块的Sequential e4 checkpoint_sequential(self.enc4_modules, chunks2, xself.pool(e3))关键点是checkpoint_sequential在forward时不会保存中间激活值而是在backward时重新做一次forward来算出梯度所以省显存但费算力。这个方法只适合放在网络最深、特征图通道数最大的层浅层特征图空间尺寸大但通道数少放在那里收益小。5.4 预测结果里出现整片白色三角形伪影现象predict.py生成的预测mask里图像左右两侧边缘区域出现大块白色三角形看起来像光照变化但方向固定。原因训练时数据增强用了随机水平翻转但验证和预测时没有做翻转。U-Net的收缩路径对边缘的感受野约束不对称模型可能在翻转后的样本上把“图像的左边”学成“车道线一定出现的区域”于是翻转增强在训练集上有效但在推理时反而带来边缘幻觉。解决一个更稳妥的方案是去掉水平翻转增强改用随机亮度、随机对比度和轻微随机旋转旋转角度控制在±2度以内。车道线方向本身是全局一致的水平翻转会让模型学到“左右车道线互换”的对称性但对左右边缘的风格纹理也做了互换容易在单侧阴影路面上产生伪影。如果一定要保留翻转增强推理时要同时预测原图和水平翻转后的图把两次预测结果按对应位置做平均。6. 从预测到视频验证test_onvideo.py里的推理加速与可视化保存技巧6.1 视频推理中的预处理一致性使用test_onvideo.py时最容易忽略的是推理预处理必须和训练完全一致。工程里训练时先对图像做resize再转Tensor并做ImageNet归一化推理时如果直接读视频帧就去预测少了归一化那一步输出的概率值会整体偏移本来0.6的置信度可能降到0.3导致大量漏检。我通常会在predict函数里把推理预处理单独封装成一个函数确保训练和测试共用同一套参数。# test_onvideo.py 推理帧预处理 def preprocess_frame(frame, img_h256, img_w512): frame cv2.resize(frame, (img_w, img_h)) frame cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame torch.from_numpy(frame).float().permute(2, 0, 1) / 255.0 mean torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1) std torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1) frame (frame - mean) / std return frame.unsqueeze(0)6.2 大图的滑窗预测与拼接策略工程里predict.py是单图直接推理如果想把模型用在原始720x1280分辨率的视频上又不重新训练最稳妥的技巧是滑窗预测。把原图切成左半和右半两个512x720的窗口分别送进模型推理再把两个mask拼接回720x1280。这个方案能明显提升远处车道线的检测正确率因为每个窗口内的目标相对尺寸变大了。拼接时有一个中间区域的权重过渡问题不能忽略——直接硬拼接会在画面中央出现一条可见的分界缝因为模型的感受野在左右窗口的中央区域最充分边缘区域由于padding的影响预测压力较大正确率偏低。正确做法是设置一个重叠加宽约32像素重叠区域内两边的预测结果按线性权重混合。def sliding_predict(model, full_img, window_w512, overlap32, img_h256): H, W full_img.shape[:2] # 计算滑窗个数 n_windows int(np.ceil((W - window_w) / (window_w - overlap))) 1 full_mask np.zeros((H, W), dtypenp.float32) weight_map np.zeros((H, W), dtypenp.float32) for i in range(n_windows): start_x i * (window_w - overlap) end_x start_x window_w if end_x W: start_x W - window_w end_x W window full_img[:, start_x:end_x] # 推理得到mask mask predict_window(model, window) full_mask[:, start_x:end_x] mask weight_map[:, start_x:end_x] 1.0 full_mask full_mask / np.maximum(weight_map, 1e-6) return full_mask参数上window_w取512是一个常用值如果窗口太宽模型对窗口边缘的目标感知依然不足太窄则滑窗数量太多推理时间成倍增长。overlap取32像素时窗口间重叠区域占比约6%既能平滑过渡又不至于让推理耗时增加太多。6.3 输出掩码后处理去除小孔洞与连通域筛选在很多视频帧里模型会把路肩边缘、破损路面误判成车道线。单纯调高阈值效果有限因为真车道的置信度可能只在0.4附近而误判区域的置信度有时反而高。这时候我会在保存视频之前加一个连通域筛选的后处理步骤使用OpenCV的connectedComponentsWithStats找出所有独立连通域只保留面积在某个区间内的区域。# 提取车道线主连通域 import cv2 import numpy as np def filter_lane_mask(binary_mask, min_area100, max_area20000): num_labels, labels, stats, _ cv2.connectedComponentsWithStats(binary_mask, connectivity8) out np.zeros_like(binary_mask) for i in range(1, num_labels): area stats[i, cv2.CC_STAT_AREA] if min_area area max_area: out[labels i] 1 return out这里min_area和max_area不是固定值需要根据你使用的分辨率来换算。例如256x512推理输出下真实车道线在画面宽度上通常横跨至少100像素、面积在500到5000之间小于这个值的基本是噪点或反光碎片大于这个值的很可能是把整段路肩都包进去了。把视频流每一帧过一遍这个滤波基本可以消除大部分误检。后处理这一步是那种看起来不起眼但效果极好的操作我用相同的checkpoints推理加了这个滤波后视频里可见的误检帧数减少了大概六成。从那以后我每次做车道线视频演示都会强制走一遍“预处理一致性检查 → 滑窗或单图推理 → 连通域筛选 → 原图画线叠加”这四步跑出来的视频才敢拿给业务那边看而不是只丢一个裸mask过去。希望帮到你。本文还有配套的精品资源点击获取
