目录
一、github下载模型
二、环境部署
三、 数据集制作
3.1 labelme下载安装:
3.2 数据集格式介绍:
3.3数据标注过程:
3.4 数据转换:
四、 训练:
五、测试效果
一、github下载模型
模型地址:https://github.com/AnshChoudhary/Lane-Detection-UNet
二、环境部署
文件下载后压缩,推荐创建一个conda虚拟环境(miniconda的下载和安装不多赘述):
conda create -n unet python=3.8 conda activate unet pip install -r requirements.txt (特别要注意torch和torchvision的版本以及自己的cuda版本是否一致)如果对torch和torchvision环境配置不清除可以在pytorch官网查询:(博主的环境为torch==2.1.0,torchvision=0.16.0,cuda=11.x) https://pytorch.org/get-started/previous-versions/
三、 数据集制作
3.1 labelme下载安装:
车道线标注使用labelme特别方便,建议下载labelme工具(最新版要求python3.9版本及以上,如果你的python版本低于3.9,建议新建conda虚拟环境)
conda create -n labelme python=3.10 conda activate labelme pip install labelme这里介绍一个key,如果你想单独快速下载包,可以在pip install命令中添加-i参数指定镜像地址。以清华源为例,你可以使用以下命令:
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple labelme下载完成之后在虚拟环境终端输入:
labelme即可出现labelme的UI界面,你就安装成功啦!
3.2 数据集格式介绍:
由于要自己准备数据集,所以我们得先了解模型需要的数据集格式,模型中train.py文件中的数据集是这样要求的:
train_dataset = BDD100KDataset( images_dir='/home/orin/Downloads/data/jpg', masks_dir='/home/orin/Downloads/data/mask', transform=transform ) val_dataset = BDD100KDataset( images_dir='/home/orin/Downloads/data/val/jpg', masks_dir='/home/orin/Downloads/data/val/mask', transform=transform )所以我们的文件格式可以这样创建:
/home/orin/Downloads/data/ ├── jpg/ # 训练集图像目录 │ ├── image1.jpg │ ├── image2.jpg │ └── ... ├── mask/ # 训练集标注目录(如分割mask或车道线多边形) │ ├── image1_mask.png │ ├── image2_mask.png │ └── ... └── val/ # 验证集目录 ├── jpg/ # 验证集图像目录 │ ├── val_image1.jpg │ ├── val_image2.jpg │ └── ... └── mask/ # 验证集标注目录 ├── val_image1_mask.png ├── val_image2_mask.png └── ...3.3数据标注过程:
首先打开图片目录(右上角还可以点击自动保存和更改输出路径加快标注速度):
单击右键后点击创建折线:
标注:
保存后得到json文件:
├── jpg/ # 训练集图像目录│ └── ...├── mask/ # 训练集标注目录(如分割mask或车道线多边形)│ └── ...└── val/ # 验证集目录├── jpg/ # 验证集图像目录│ └── ...└── mask/ # 验证集标注目录└── ...
3.4 数据转换:
由于数据集需要图片掩码,这里需要新建python程序将json数据格式转换为图片掩码:
import os import json import cv2 import numpy as np def generate_mask(image_path, annotation, mask_path): # 读取图片尺寸 img = cv2.imread(image_path) height, width = img.shape[:2] # 创建空白掩码 mask = np.zeros((height, width), dtype=np.uint8) # 处理json中每个shape,绘制折线 for shape in annotation.get('shapes', []): points = shape['points'] pts = np.array(points, np.int32) # 绘制线条 # 你可以调整线宽,比如为2或3 cv2.polylines(mask, [pts], isClosed=False, color=255, thickness=3) # 保存掩码 cv2.imwrite(mask_path, mask) # 你的路径 images_dir = r'D:\AI\Lane-Detection-UNet-main\data_process\data\jpg' masks_dir = r'D:\AI\Lane-Detection-UNet-main\data_process\data\mask' json_dir = r"D:\AI\Lane-Detection-UNet-main\data_process\data\json" if not os.path.exists(masks_dir): os.makedirs(masks_dir) for image_filename in os.listdir(images_dir): if not image_filename.endswith('.jpg'): continue image_path = os.path.join(images_dir, image_filename) json_filename = os.path.splitext(image_filename)[0] + '.json' json_path = os.path.join(json_dir, json_filename) if not os.path.exists(json_path): print(f"对应的json文件不存在:{json_path}") continue with open(json_path, 'r', encoding='utf-8') as f: annotation_data = json.load(f) # 生成掩码路径 mask_path = os.path.join(masks_dir, image_filename.replace('.jpg', '.jpg')) # 调用绘制线条生成掩码 generate_mask(image_path, annotation_data, mask_path)即可得到图片掩码:
到此,数据集准备完毕!!!
四、 训练
进入模型目录,输入训练代码(注意改动标注图片路径)即可开始训练:
CUDA_VISIBLE_DEVICES=0 python train.py训练完成之后会得到一个best_unet_lane_detection.pth权重文件:
五、测试效果
博主自己改动了一下inference.py代码,增加了一个二值化的图:
CUDA_VISIBLE_DEVICES=0 python inference.py (注意改动测试图片路径)训练过程有什么问题的欢迎和博主讨论:
博主qq:2904385943
训练成功,完结撒花!!!