资讯动态

基于UNet与Tusimple数据集的车道线分割实战:从标注转换到模型调优

发布时间:2026/10/6 8:32:58 来源:尧图企业网站定制
简介这份资源面向车道线检测与自动驾驶视觉方向的初学者及进阶开发者提供基于U-Net模型在TuSimple数据集上完成训练与预测的完整工程。包内共15个文件以7个Python脚本为核心涵盖模型定义、数据集加载、标签处理、训练与视频推理等环节另含2个txt说明、2个md文档及avi、mp4演示视频压缩包约7.89MB结构紧凑便于快速复现。已有583人学习下载。读者可借此理解U-Net收缩与扩展路径、跳跃连接的设计思路掌握TuSimple数据预处理、损失函数与优化器选择、训练验证及调参流程并通过IoU、Precision、Recall等指标评估效果。附带的实线、虚线及路面有水等场景视频直观展示模型在弯曲车道、遮挡与光照变化下的表现为后续数据增强、模型融合与超参数调优提供参考。1. 车道线分割实战从 Tusimple 数据集到 UNet 预测效果的完整路径拿到「使用 unet 模型结构在 Tusimple 数据集上训练得到预测车道线的效果」这个题目很多人第一反应是找现成代码跑一遍结果发现预测出来的车道线要么断断续续要么在弯道处直接糊成一片。Tusimple 数据集本身是自动驾驶领域车道线检测的经典基准标注形式是点序列而非像素级掩码这意味着你不能直接把标注丢进 UNet 当分割标签用。UNet 的优势在于编码器-解码器结构配合跳跃连接能同时保留车道线的语义信息和精细边界但前提是你得先把 Tusimple 的 json 标注转成二值分割掩码。这套流程适合有一定 PyTorch 基础、想从零复现车道线分割效果的从业者也适合已经跑通过其他分割任务、想迁移到车道线场景的工程师。下面按数据准备、模型搭建、训练调参、效果验证的顺序拆开讲中间会重点说清楚标注转换和损失函数选择这两个最容易翻车的地方。2. Tusimple 标注转分割掩码json 到二值图的四个关键步骤2.1 先搞清楚 Tusimple 的标注长什么样Tusimple 数据集的训练集包含约 3600 个视频片段每个片段取最后一帧作为标注帧标注文件是 json 格式。每条标注记录里有三个关键字段raw_file 指向图片路径lanes 是车道线点序列h_samples 是这些点对应的 y 坐标。lanes 里每个元素是一条车道线的 x 坐标列表-2 表示该位置没有车道线点。这个结构和语义分割常用的掩码图完全不同你不能直接拿 lanes 当标签训练 UNet。常见做法是先把点序列按顺序连成线再在空白掩码上画线并做膨胀处理。这里有个细节Tusimple 的标注点比较稀疏相邻点之间间隔约 10 个像素如果只画 1 像素宽的线UNet 训练时正负样本极度不平衡模型会倾向于全部预测为背景。我一般会把线宽设为 5 到 8 像素既保证正样本比例合理又不会让不同车道线粘连。2.2 转换脚本从 json 到掩码图import json import cv2 import numpy as np import os def tusimple_to_mask(json_path, output_dir, img_height720, img_width1280, line_width5): 将 Tusimple json 标注转换为二值分割掩码 json_path: 标注文件路径 output_dir: 掩码保存目录 line_width: 车道线绘制宽度建议 5-8 with open(json_path, r) as f: data [json.loads(line) for line in f.readlines()] os.makedirs(output_dir, exist_okTrue) for item in data: # 初始化全黑掩码 mask np.zeros((img_height, img_width), dtypenp.uint8) for lane in item[lanes]: points [] for x, y in zip(lane, item[h_samples]): if x ! -2: # -2 表示该点不存在 points.append((x, y)) if len(points) 2: continue # 按 y 坐标排序后连线 points sorted(points, keylambda p: p[1]) for i in range(len(points) - 1): cv2.line(mask, points[i], points[i1], 255, thicknessline_width) # 保存掩码文件名与图片对应 base_name os.path.basename(item[raw_file]).replace(.jpg, .png) cv2.imwrite(os.path.join(output_dir, base_name), mask) print(f转换完成共处理 {len(data)} 条标注) # 调用示例 tusimple_to_mask( json_path./tusimple/train_set/label_data_0313.json, output_dir./tusimple/train_masks, line_width5 )这段代码的逻辑很直接逐条读取 json 记录对每条车道线提取有效点按 y 坐标排序后相邻点连线。参数 line_width 控制线宽太小会导致正样本不足太大则会让相邻车道线在掩码上粘连一般 5 到 8 之间比较稳妥。img_height 和 img_width 要和原图一致Tusimple 图片是 720x1280。转换完成后建议随机抽几张掩码叠加到原图上目视检查确认车道线位置和宽度合理。2.3 数据集划分与 DataLoader 封装Tusimple 官方没有给出验证集划分常见做法是从训练集里按 8:2 随机划分。注意划分时要按视频片段划分而不是按帧划分否则同一段视频的相邻帧会同时出现在训练集和验证集里导致验证指标虚高。我一般会先提取所有 raw_file 的片段前缀按前缀分组后再划分。from torch.utils.data import Dataset, DataLoader from PIL import Image import torchvision.transforms as T class LaneDataset(Dataset): def __init__(self, img_dir, mask_dir, file_list, img_size(256, 512)): self.img_dir img_dir self.mask_dir mask_dir self.file_list file_list self.img_size img_size self.img_transform T.Compose([ T.Resize(img_size), T.ToTensor(), T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) self.mask_transform T.Compose([ T.Resize(img_size, interpolationT.InterpolationMode.NEAREST), T.ToTensor() ]) def __len__(self): return len(self.file_list) def __getitem__(self, idx): name self.file_list[idx] img Image.open(os.path.join(self.img_dir, name)).convert(RGB) mask Image.open(os.path.join(self.mask_dir, name.replace(.jpg, .png))).convert(L) img self.img_transform(img) mask self.mask_transform(mask) mask (mask 0.5).float() # 二值化 return img, mask掩码的 Resize 必须用 NEAREST 插值用双线性插值会把二值掩码变成灰度图边界处出现 0 到 1 之间的值训练时损失计算会出问题。图片归一化用了 ImageNet 的均值和标准差因为 UNet 编码器通常加载预训练权重保持输入分布一致很重要。3. UNet 模型搭建编码器选型与跳跃连接的两个改动点3.1 为什么车道线分割更适合轻量编码器标准 UNet 的编码器是 5 层下采样每层通道数翻倍从 64 到 1024。这个结构在医学影像分割上表现很好但车道线场景有两个不同一是车道线是细长结构不需要那么大的感受野二是 Tusimple 图片分辨率高5 层下采样后特征图只有 45x80细车道线信息丢失严重。我一般会把编码器换成 ResNet34 的前 4 个 stage输出 stride 为 16 的特征图这样既保留了足够的空间分辨率又能利用 ImageNet 预训练权重加速收敛。import torch import torch.nn as nn import torchvision.models as models class UNetLane(nn.Module): def __init__(self, num_classes1, pretrainedTrue): super().__init__() # 编码器ResNet34 前 4 个 stage resnet models.resnet34(pretrainedpretrained) self.encoder0 nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu) # stride 2 self.encoder1 nn.Sequential(resnet.maxpool, resnet.layer1) # stride 4 self.encoder2 resnet.layer2 # stride 8 self.encoder3 resnet.layer3 # stride 16 self.encoder4 resnet.layer4 # stride 32 # 解码器 self.up4 nn.ConvTranspose2d(512, 256, 2, stride2) self.dec4 self._conv_block(256 256, 256) self.up3 nn.ConvTranspose2d(256, 128, 2, stride2) self.dec3 self._conv_block(128 128, 128) self.up2 nn.ConvTranspose2d(128, 64, 2, stride2) self.dec2 self._conv_block(64 64, 64) self.up1 nn.ConvTranspose2d(64, 32, 2, stride2) self.dec1 self._conv_block(32 32, 32) self.final nn.Conv2d(32, num_classes, 1) def _conv_block(self, in_ch, out_ch): return 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): # 编码 e0 self.encoder0(x) # 1/2 e1 self.encoder1(e0) # 1/4 e2 self.encoder2(e1) # 1/8 e3 self.encoder3(e2) # 1/16 e4 self.encoder4(e3) # 1/32 # 解码 跳跃连接 d4 self.up4(e4) d4 self.dec4(torch.cat([d4, e3], dim1)) d3 self.up3(d4) d3 self.dec3(torch.cat([d3, e2], dim1)) d2 self.up2(d3) d2 self.dec2(torch.cat([d2, e1], dim1)) d1 self.up1(d2) d1 self.dec1(torch.cat([d1, e0], dim1)) return self.final(d1)这个结构比标准 UNet 少了一次下采样最终输出是原图的 1/2 分辨率。训练时标签也要对应下采样到 1/2推理时再把输出上采样回原图。这样做的好处是车道线这种细结构在 1/2 分辨率下仍有 2 到 3 像素宽不会因为下采样太狠而消失。跳跃连接保留了编码器各阶段的特征解码器逐级融合边界定位会更准。3.2 损失函数Dice Loss 和 BCE 怎么配车道线分割是典型的类别不平衡问题背景像素占 95% 以上。只用 BCE Loss 的话模型很快学会全预测背景就能拿到很低的 loss但车道线一条都出不来。常见做法是 BCE 和 Dice Loss 按权重相加我一般设 BCE 权重 0.5Dice 权重 1.0让 Dice 主导梯度。class DiceLoss(nn.Module): def __init__(self, smooth1.0): super().__init__() self.smooth smooth def forward(self, pred, target): pred torch.sigmoid(pred) pred pred.view(-1) target target.view(-1) intersection (pred * target).sum() dice (2. * intersection self.smooth) / (pred.sum() target.sum() self.smooth) return 1 - dice # 组合损失 bce_loss nn.BCEWithLogitsLoss() dice_loss DiceLoss() def criterion(pred, target): return 0.5 * bce_loss(pred, target) 1.0 * dice_loss(pred, target)Dice Loss 的 smooth 参数防止分母为零一般设 1.0。注意 pred 要先过 sigmoid 再算 Dice而 BCEWithLogitsLoss 内部已经包含 sigmoid所以两者输入的都是 logits不要重复做 sigmoid。训练初期 Dice Loss 波动比较大可以先用 BCE 单独训 5 个 epoch 再切到组合损失收敛会更稳。4. 训练调参与效果验证学习率、批次大小和评估指标4.1 学习率策略和批次大小怎么定编码器加载了预训练权重学习率不能设太大否则预训练特征会被破坏。我一般把编码器学习率设为 1e-4解码器设为 1e-3用 PyTorch 的参数组分开设置。优化器选 AdamW权重衰减 1e-4。学习率调度用 CosineAnnealingLRT_max 设为总 epoch 数eta_min 设为 1e-6。批次大小受显存限制Tusimple 图片 resize 到 256x512 后单张 1080Ti 上 batch size 可以设到 8。如果显存不够可以降到 4 并配合梯度累积累积步数设为 2等效 batch size 还是 8。注意 BatchNorm 在 batch size 小于 4 时统计量不稳定这时候要么冻结 BN 层要么改用 GroupNorm。# 优化器参数组 encoder_params list(model.encoder0.parameters()) list(model.encoder1.parameters()) \ list(model.encoder2.parameters()) list(model.encoder3.parameters()) \ list(model.encoder4.parameters()) decoder_params [p for n, p in model.named_parameters() if encoder not in n] optimizer torch.optim.AdamW([ {params: encoder_params, lr: 1e-4}, {params: decoder_params, lr: 1e-3} ], weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max50, eta_min1e-6 )训练 epoch 数一般设 50 到 80Tusimple 训练集不大50 个 epoch 在单卡上大约 3 到 4 小时。每个 epoch 结束后在验证集上算 IoU 和 F1保存 IoU 最高的权重。如果验证 IoU 连续 10 个 epoch 不提升就提前停止避免过拟合。4.2 评估指标IoU 和 F1 之外还要看什么车道线分割的常规指标是 IoU 和 F1但这两个指标对细长结构不够敏感。一条车道线断成两截IoU 可能只掉几个点但实际使用中断裂的车道线没法做后续的拟合和跟踪。我一般会额外算两个指标一是连通域数量理想情况下每条车道线应该是一个连通域如果连通域数量明显多于车道线数量说明有断裂二是车道线像素的召回率单独看正样本的召回比整体 IoU 更能反映车道线是否完整。def evaluate(model, dataloader, device): model.eval() iou_sum, f1_sum, recall_sum 0, 0, 0 with torch.no_grad(): for img, mask in dataloader: img, mask img.to(device), mask.to(device) pred torch.sigmoid(model(img)) pred (pred 0.5).float() intersection (pred * mask).sum() union pred.sum() mask.sum() - intersection iou (intersection 1e-6) / (union 1e-6) tp intersection fp pred.sum() - tp fn mask.sum() - tp precision tp / (tp fp 1e-6) recall tp / (tp fn 1e-6) f1 2 * precision * recall / (precision recall 1e-6) iou_sum iou.item() f1_sum f1.item() recall_sum recall.item() n len(dataloader) return iou_sum / n, f1_sum / n, recall_sum / n推理时阈值设 0.5 是默认做法但实际调参时可以在验证集上扫一遍 0.3 到 0.7选 F1 最高的阈值。有时候 0.4 比 0.5 能多召回一些断裂的车道线片段后续用形态学闭运算连起来整体效果反而更好。5. 避坑与排查Tusimple 训练 UNet 的五个血泪教训5.1 掩码全黑或全白现象训练几个 epoch 后 loss 几乎不变可视化预测结果发现输出全黑或全白。原因通常是标注转换时线宽设得太小正样本比例低于 1%模型直接摆烂全预测背景或者掩码保存时用了 jpg 格式压缩导致二值掩码出现灰度值二值化后全变成 0。解决线宽调到 5 以上掩码必须存 png 格式转换后随机抽 10 张检查正样本比例正常应该在 3% 到 8% 之间。5.2 验证集 IoU 虚高现象训练集和验证集 IoU 都到 0.8 以上但拿实际视频跑预测车道线抖动严重。原因是数据集划分时按帧随机划分同一段视频的相邻帧同时进了训练集和验证集模型记住了这段视频的车道线位置。解决按视频片段前缀分组后再划分确保验证集的视频片段在训练集中完全没出现过。5.3 弯道处车道线断裂现象直道预测很完整一到弯道就断成几截。原因是 UNet 的下采样倍数太高弯道处车道线曲率大在低分辨率特征图上相邻点被合并了。解决把编码器下采样倍数从 32 降到 16或者保持 32 但在解码器最后加一层上采样到原图分辨率再算损失。另一个办法是训练时对弯道样本过采样Tusimple 里弯道片段占比不高过采样能让模型多关注这些难例。5.4 相邻车道线粘连现象预测结果里两条相邻车道线糊成一条粗线。原因是掩码转换时线宽太大或者不同车道线的点序列在连线时交叉了。解决线宽降到 5转换时检查每条车道线的点是否按 y 坐标严格递增如果出现 y 坐标相同但 x 不同的点说明标注有噪声需要手动过滤。推理后处理可以用分水岭算法或者简单的连通域分析把粘连的线分开。5.5 显存溢出现象训练到一半报 CUDA out of memory。原因通常是 batch size 设太大或者数据加载时没有及时释放中间变量。解决batch size 降到 4 并开梯度累积数据加载用 num_workers4 和 pin_memoryTrue训练循环里每个 batch 结束后 del 掉中间变量。如果还不行把图片 resize 到 256x512 而不是 512x1024显存占用能降一半。6. 进阶技巧用形态学后处理把 IoU 再提两个点训练完模型直接输出二值掩码边缘往往有毛刺细车道线可能断成几段。我一般会在推理后加一套形态学后处理顺序是先闭运算连接断裂再开运算去掉孤立噪点最后按连通域面积过滤掉太小的区域。这套操作在验证集上通常能把 IoU 提 1 到 2 个点F1 提 2 到 3 个点而且不增加任何训练成本。import cv2 import numpy as np def postprocess(mask, kernel_size5, min_area100): mask: 二值掩码0 或 255 kernel_size: 形态学核大小 min_area: 最小连通域面积小于此值的区域被过滤 kernel cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) # 闭运算连接断裂 closed cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel) # 开运算去噪 opened cv2.morphologyEx(closed, cv2.MORPH_OPEN, kernel) # 连通域过滤 num_labels, labels, stats, _ cv2.connectedComponentsWithStats(opened, connectivity8) result np.zeros_like(opened) for i in range(1, num_labels): if stats[i, cv2.CC_STAT_AREA] min_area: result[labels i] 255 return resultkernel_size 一般设 5弯道多的场景可以设 7 让闭运算连接能力更强。min_area 根据图片分辨率调256x512 的图上设 100 左右比较合适太小去不掉噪点太大会把短车道线片段也滤掉。注意闭运算和开运算的顺序不能反先闭后开才能既连接断裂又去噪反过来会把该连的也断开。验证后处理效果时不要只看整体 IoU要单独统计断裂车道线的数量变化。我一般会写个小脚本对每张验证图算连通域数量后处理前后对比如果连通域数量从 5 降到 3 而车道线实际是 3 条说明后处理有效。另外后处理参数不要只在验证集上调要留一部分测试集做最终确认避免过拟合到验证集。这套流程跑下来Tusimple 验证集上 IoU 能到 0.75 到 0.82F1 到 0.85 左右单卡训练 4 小时以内。如果追求更高精度可以换更深的编码器或者加注意力模块但推理速度会下降实际部署时要权衡。我自己踩过最大的坑是掩码线宽设了 2训了两天 loss 都不降后来改成 6 才正常。希望帮到你。本文还有配套的精品资源点击获取

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑