📚 课程目录第4章车道线检测数据集标注格式转换

4. 车道线检测数据集标注格式转换

任务说明

本案例处理车道线检测数据集的标注格式统一问题。以TuSimple数据集为例，其原始标注为JSON格式（每条车道线由一组离散点构成），而主流训练框架（如基于分割或Anchor的车道线模型）通常需要掩码（Mask）或固定数量点的结构化标签。目标是将TuSimple的JSON标注转换为两种通用格式：二值分割掩码（PNG图像，便于训练分割模型）和固定点序列（JSON，便于训练点回归模型），并输出可视化验证图。

实现方案

代码读取TuSimple的label.json，解析每条样本的车道线点集，通过多项式拟合与重采样将每条线统一为固定数量的点（本案例取56个点），同时将点集绘制为二值掩码。输出文件按统一命名规则存放，并生成一张叠加了标注点的可视化图像用于人工校验。

代码实现

import json
import os
import numpy as np
import cv2
from PIL import Image, ImageDraw

def fit_and_resample(points, num_points=56):
    """将一条车道线的点集拟合为y方向均匀采样的固定点数"""
    pts = np.array(points, dtype=np.float64)
    if len(pts) < 2:
        return None
    # 按y排序，保证从图像底部向上处理
    pts = pts[np.argsort(pts[:, 1])]
    y_vals = pts[:, 1]
    x_vals = pts[:, 0]
    # 使用2阶多项式拟合x=f(y)
    coeffs = np.polyfit(y_vals, x_vals, 2)
    y_new = np.linspace(y_vals.min(), y_vals.max(), num_points)
    x_new = np.polyval(coeffs, y_new)
    return np.stack([x_new, y_new], axis=1)

def convert_tusimple(json_path, output_dir, img_h=720, img_w=1280):
    os.makedirs(output_dir, exist_ok=True)
    mask_dir = os.path.join(output_dir, 'masks')
    pts_dir = os.path.join(output_dir, 'points')
    vis_dir = os.path.join(output_dir, 'visualize')
    for d in [mask_dir, pts_dir, vis_dir]:
        os.makedirs(d, exist_ok=True)

    with open(json_path, 'r') as f:
        lines = f.readlines()

    for idx, line in enumerate(lines):
        data = json.loads(line)
        img_name = os.path.basename(data['raw_file'])
        base_name = os.path.splitext(img_name)[0]

        # 1. 生成二值分割掩码
        mask = Image.new('L', (img_w, img_h), 0)
        draw = ImageDraw.Draw(mask)
        lanes = data['lanes']  # 每条线x坐标列表
        h_samples = data['h_samples']  # 对应的y坐标列表
        all_points = []
        for lane_x in lanes:
            pts = [(x, y) for x, y in zip(lane_x, h_samples) if x > 0]
            if len(pts) < 2:
                continue
            # 绘制宽度为8像素的线
            draw.line(pts, fill=255, width=8)
            # 保存原始点用于后续重采样
            all_points.append(pts)

        mask_path = os.path.join(mask_dir, base_name + '.png')
        mask.save(mask_path)

        # 2. 生成固定点序列
        fixed_points = []
        for pts in all_points:
            resampled = fit_and_resample(pts)
            if resampled is not None:
                fixed_points.append(resampled.tolist())

        pts_path = os.path.join(pts_dir, base_name + '.json')
        with open(pts_path, 'w') as f:
            json.dump({'image': img_name, 'lanes': fixed_points}, f)

        # 3. 可视化验证图
        img = cv2.imread(os.path.join(os.path.dirname(json_path), data['raw_file']))
        if img is None:
            continue
        for pts in fixed_points:
            pts = np.array(pts, dtype=np.int32)
            for i in range(len(pts) - 1):
                cv2.line(img, tuple(pts[i]), tuple(pts[i+1]), (0, 0, 255), 3)
        vis_path = os.path.join(vis_dir, base_name + '.jpg')
        cv2.imwrite(vis_path, img)

        if idx >= 9:  # 本案例处理前10条样本
            break

    print(f'转换完成，共处理 {idx+1} 条样本')
    print(f'掩码输出: {mask_dir}')
    print(f'点序列输出: {pts_dir}')
    print(f'可视化输出: {vis_dir}')

if __name__ == '__main__':
    # 使用示例（请根据实际路径修改）
    convert_tusimple('label.json', './output')

输入数据格式说明

TuSimple的label.json每行是一个JSON对象，关键字段为：raw_file（图像相对路径）、lanes（每条车道线的x坐标列表，无效点为-2）、h_samples（固定的y坐标采样列表）。本案例代码自动过滤x<0的无效点。

落地与运行

安装依赖：pip install numpy opencv-python pillow
将TuSimple数据集中的label.json与图像目录放在同一父目录下
修改__main__中的路径后运行：python convert_tusimple.py

验证方法与判定标准

运行后检查输出目录：

掩码验证：打开masks/下的PNG，白色像素应形成连续的车道线形状，无断裂或明显噪声；
点序列验证：读取points/下的JSON，每条车道线应包含56个点，且y坐标单调递增；
可视化验证：查看visualize/下的JPG，红色线条应准确覆盖原图中的车道线位置，无偏移或错位。

满足以上三点即判定转换成功。若掩码出现断线，可适当增大绘制宽度；若拟合点偏离原线，可调整多项式阶数。