☰
基于UNet与Tusimple数据集的车道线分割实战:从标注转换到模型调优
2026/10/6 8:32:54 网站建设 项目流程

简介:这份资源面向车道线检测与自动驾驶视觉方向的初学者及进阶开发者,提供基于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_height=720, img_width=1280, line_width=5): """ 将 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_ok=True) for item in data: # 初始化全黑掩码 mask = np.zeros((img_height, img_width), dtype=np.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, key=lambda p: p[1]) for i in range(len(points) - 1): cv2.line(mask, points[i], points[i+1], 255, thickness=line_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_width=5 )

这段代码的逻辑很直接:逐条读取 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, interpolation=T.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_classes=1, pretrained=True): super().__init__() # 编码器:ResNet34 前 4 个 stage resnet = models.resnet34(pretrained=pretrained) 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, stride=2) self.dec4 = self._conv_block(256 + 256, 256) self.up3 = nn.ConvTranspose2d(256, 128, 2, stride=2) self.dec3 = self._conv_block(128 + 128, 128) self.up2 = nn.ConvTranspose2d(128, 64, 2, stride=2) self.dec2 = self._conv_block(64 + 64, 64) self.up1 = nn.ConvTranspose2d(64, 32, 2, stride=2) 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, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True) ) 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], dim=1)) d3 = self.up3(d4) d3 = self.dec3(torch.cat([d3, e2], dim=1)) d2 = self.up2(d3) d2 = self.dec2(torch.cat([d2, e1], dim=1)) d1 = self.up1(d2) d1 = self.dec1(torch.cat([d1, e0], dim=1)) 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.5,Dice 权重 1.0,让 Dice 主导梯度。

class DiceLoss(nn.Module): def __init__(self, smooth=1.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。学习率调度用 CosineAnnealingLR,T_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_decay=1e-4) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=50, eta_min=1e-6 )

训练 epoch 数一般设 50 到 80,Tusimple 训练集不大,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_workers=4 和 pin_memory=True,训练循环里每个 batch 结束后 del 掉中间变量。如果还不行,把图片 resize 到 256x512 而不是 512x1024,显存占用能降一半。

6. 进阶技巧:用形态学后处理把 IoU 再提两个点

训练完模型直接输出二值掩码,边缘往往有毛刺,细车道线可能断成几段。我一般会在推理后加一套形态学后处理,顺序是:先闭运算连接断裂,再开运算去掉孤立噪点,最后按连通域面积过滤掉太小的区域。这套操作在验证集上通常能把 IoU 提 1 到 2 个点,F1 提 2 到 3 个点,而且不增加任何训练成本。

import cv2 import numpy as np def postprocess(mask, kernel_size=5, min_area=100): """ 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, connectivity=8) 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 result

kernel_size 一般设 5,弯道多的场景可以设 7 让闭运算连接能力更强。min_area 根据图片分辨率调,256x512 的图上设 100 左右比较合适,太小去不掉噪点,太大会把短车道线片段也滤掉。注意闭运算和开运算的顺序不能反,先闭后开才能既连接断裂又去噪,反过来会把该连的也断开。

验证后处理效果时,不要只看整体 IoU,要单独统计断裂车道线的数量变化。我一般会写个小脚本,对每张验证图算连通域数量,后处理前后对比,如果连通域数量从 5 降到 3 而车道线实际是 3 条,说明后处理有效。另外后处理参数不要只在验证集上调,要留一部分测试集做最终确认,避免过拟合到验证集。

这套流程跑下来,Tusimple 验证集上 IoU 能到 0.75 到 0.82,F1 到 0.85 左右,单卡训练 4 小时以内。如果追求更高精度,可以换更深的编码器或者加注意力模块,但推理速度会下降,实际部署时要权衡。我自己踩过最大的坑是掩码线宽设了 2,训了两天 loss 都不降,后来改成 6 才正常。希望帮到你。

本文还有配套的精品资源,点击获取

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询