AI计算机视觉笔记三十一:基于UNetMultiLane的多车道线等识别

简介: 该项目基于开源数据集 VIL100 实现了 UNetMultiLane,用于多车道线及车道线类型的识别。数据集中标注了六个车道的车道线及其类型。项目详细记录了从环境搭建到模型训练与测试的全过程,并提供了在 CPU 上进行训练和 ONNX 转换的代码示例。训练过程约需 4 小时完成 50 个 epoch。此外,还实现了视频检测功能,可在视频中实时识别车道线及其类型。

UNetMultiLane 多车道线、车道线类型识别。

数据是基于开源数据集 VIL100。其中数据标注了所在的六个车道的车道线和车道线的类型。

8条车道线(六个车道),对应的顺序是:7,5,3,1,2,4,6,8。其中1,2对应的自车所在的车道,从左往右标记。

车道线的类别(10个类别):单条白色实线、单条白色虚线、单条黄色实线、单条黄色虚线、双条白色实线、双条黄色实线、双条黄色虚线、双条白色实虚线、双条白色黄色实线、双条白色虚实线。

从环境开发到部署记录全过程。

由于没有使用CUDA,所以训练和转换ONNX在CPU上进行,山水无移大佬的是使用CUDA处理,这里修改了一点他的代码。

一、环境创建

1、创建虚拟环境

conda create -n UNet_env python=3.9

2、激活

conda activate UNet_env

3、安装轮子

pip install torch -i https://pypi.tuna.tsinghua.edu.cn/simple

pip install numpy -i https://pypi.tuna.tsinghua.edu.cn/simple

pip install Pillow  -i https://pypi.tuna.tsinghua.edu.cn/simple

pip install torchvision -i https://pypi.tuna.tsinghua.edu.cn/simple

pip install tqdm -i https://pypi.tuna.tsinghua.edu.cn/simple

pip install opencv-python -i https://pypi.tuna.tsinghua.edu.cn/simple

pip install onnx -i https://pypi.tuna.tsinghua.edu.cn/simple

pip install onnxruntime -i https://pypi.tuna.tsinghua.edu.cn/simple

二、训练

cqu20160901/UNetMultiLane: UNetMultiLane 多车道线、车道线类型识别,自我学习使用,没有经过大量测试,不免有问题,不喜不扰。 (github.com)

代码比较简单,不多,由于使用的是CPU所以修改了train.py,直接附上。

train.py

from getdata import GetLaneDataset
from torch.utils.data import DataLoader as DataLoader
from unet import UNetMultiLane as Net
import torch
from torch.autograd import Variable
from loss import softmax_focal_loss
from tqdm import tqdm
import sys


train_image_txt = './data/VIL100/train.txt'
test_image_txt = './data/VIL100/test.txt'

data_main_path = './data/VIL100'
weights_save_path = './weights'

input_height = 480
input_width = 640
lane_num = 9   # 8 条车道线 + 1
type_num = 11  # 10 种线的类型 + 1

epoch_num = 50
batch_size = 2
learn_rate = 0.001
num_workers = 1

width_mult = 0.25


def test_eval(model, epoch):
    log_txt = open(weights_save_path + './log.txt', 'a')
    dataset = GetLaneDataset(image_txt=test_image_txt, data_main_path=data_main_path, input_height=input_height, input_width=input_width, train_mode=False)
    images_num = len(dataset)
    print('eval images num is:', images_num)
    model.eval()
    criterion = softmax_focal_loss

    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, drop_last=True)
    seg_loss_total = 0
    cls_loss_total = 0

    for image, label_mask, label_type in tqdm(dataloader):

        image, label_mask, label_type = Variable(image), Variable(label_mask), Variable(label_type)
        if torch.cuda.is_available():
            image, label_mask, label_type = image.cuda(), label_mask.cuda(), label_type.cuda()
            print("Model is running on GPU")
        else:
            print("Model is running on CPU")

        pred_out = model(image)
        seg_loss, cls_loss = criterion(pred_out, label_mask, label_type.squeeze())

        seg_loss_total += seg_loss.item()
        cls_loss_total += cls_loss.item()

    print(f"{'eval_seg_loss':>13} {'eval_cls_loss':>13} {'eval_total_loss':>15}")
    print(f"{seg_loss_total:>13} {cls_loss_total:>13} {(seg_loss_total + seg_loss_total):>15}")

    save_line = 'epoch:' + str(epoch) + ',seg_loss_total:' + str(seg_loss_total) + ', cls_loss_total:' + str(cls_loss_total) + ', total_loss:' + str(seg_loss_total + seg_loss_total) + '\n'
    log_txt.write(save_line)
    log_txt.close()

    return seg_loss_total + seg_loss_total


def train():
    dataset = GetLaneDataset(image_txt=train_image_txt, data_main_path=data_main_path, input_height=input_height, input_width=input_width, train_mode=True)
    images_num = len(dataset)
    print('train images num is:', images_num)

    model = Net(in_channels=3, lane_num=lane_num, type_num=type_num, width_mult=width_mult, is_deconv=True, is_batchnorm=True, is_ds=True)
    if torch.cuda.is_available():
        model = model.cuda()
        print("Model is running on GPU")
    else:
        print("Model is running on CPU")

    model.train()

    criterion = softmax_focal_loss

    optimizer = torch.optim.Adam(model.parameters(), lr=learn_rate)
    print(f"{'epoch'}/{'epoch_num'} | {'seg_loss':>8} | {'cls_loss':>8} | {'total_loss':>10}")

    eval_loss_total = 1e7
    best_epoch = 0
    for epoch in range(epoch_num):
        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, drop_last=True)

        seg_loss_total = 0
        cls_loss_total = 0

        for image, label_mask, label_type in tqdm(dataloader):

            image, label_mask, label_type = Variable(image), Variable(label_mask), Variable(label_type)
            if torch.cuda.is_available():
                image, label_mask, label_type = image.cuda(), label_mask.cuda(), label_type.cuda()
                print("Model is running on GPU")
            else:
                print("Model is running on CPU")

            pred_out = model(image)
            seg_loss, cls_loss = criterion(pred_out, label_mask, label_type.squeeze())

            total_loss = seg_loss + cls_loss
            total_loss.backward()
            optimizer.step()
            optimizer.zero_grad()

            seg_loss1 = "%.4f" % seg_loss.item()
            cls_loss1 = "%.4f" % cls_loss.item()
            total_loss1 = "%.4f" % total_loss.item()

            text = f"{epoch}/{epoch_num} {seg_loss1:>8} {cls_loss1:>8} {total_loss1:>8}"
            sys.stdout.write(text)
            sys.stdout.flush()

            seg_loss_total += seg_loss.item()
            cls_loss_total += cls_loss.item()

        seg_loss_total1 = "%.4f" % seg_loss_total
        cls_loss_total1 = "%.4f" % cls_loss_total
        total_loss_total1 = "%.4f" % (seg_loss_total + cls_loss_total)

        print()
        print(f"{'epoch':<5} {'epoch_num':<9} {'seg_loss_total':>14} {'cls_loss_total':>14} {'total_loss_total':>16}")
        print(f"{epoch:<5} {epoch_num:<9} {seg_loss_total1:>14} {cls_loss_total1:>14} {total_loss_total1:>16}")

        torch.save(model.state_dict(), weights_save_path + '/epoch_{0}.pth'.format(epoch + 1))

        eval_loss = test_eval(model, epoch)
        model.train()
        if eval_loss < eval_loss_total:
            eval_loss_total = eval_loss
            torch.save(model.state_dict(), weights_save_path + '/best.pth')
            best_epoch = epoch
    log_txt = open(weights_save_path + './log.txt', 'a')
    save_line = 'eval best epoch is:' + str(best_epoch) + '\n'
    log_txt.write(save_line)
    log_txt.close()


if __name__ == '__main__':
    print('start train ...')
    train()

执行python train.py开始训练,由于是cpu,50轮大概4小时。
image.png

三、测试

原代码测试使用的是GPU,这里也修改成cpu,并增加了视频检测功能。

from unet import UNetMultiLane as Net
import torch
import numpy as np
import cv2
from torch.autograd import Variable


width_mult = 0.25

input_height = 480
input_width = 640
lane_num = 9   # 8 条车道线 + 1
type_num = 11  # 10 种线的类型 + 1


color_list = [(100, 149, 237), (0, 0, 255), (173, 255, 47), (240, 255, 255), (0, 100, 0),
              (47, 79, 79), (255, 228, 196), (138, 43, 226), (165, 42, 42), (222, 184, 135)]

lane_Id_type = [7, 5, 3, 1, 2, 4, 6, 8]


line_type = ['No lane markings',
             'Single white solid line',
             'Single white dashed line',
             'Single solid yellow line',
             'Single yellow dashed line',
             'Double solid white lines',
             'Double solid yellow lines',
             'Double yellow dashed lines',
             'Double white yellow solid lines',
             'Double white dashed lines',
             'Double white solid dashed lines']


def precess_image(img_src, resize_w, resize_h):
    image = cv2.resize(img_src, (resize_w, resize_h), interpolation=cv2.INTER_LINEAR)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = image.astype(np.float32)
    image /= 255
    image = image.transpose((2, 0, 1))
    return image


def softmax(x, axis):
    x -= np.max(x, axis=axis, keepdims=True)
    value = np.exp(x) / np.sum(np.exp(x), axis=1, keepdims=True)
    return value


def test_image():
    weights_path = './weights/best.pth'
    image_path = './images/00078.jpg'

    model = Net(in_channels=3, lane_num=lane_num, type_num=type_num, width_mult=width_mult, is_deconv=True, is_batchnorm=True, is_ds=True)

    if torch.cuda.is_available():
        model = model.cuda()
        print("Model moved to GPU.")
    else:
        print("No CUDA device available, model will run on CPU.")

    model.load_state_dict(torch.load(weights_path))
    model.eval()

    origin_image = cv2.imread(image_path)
    image_height, image_width = origin_image.shape[:2]

    input_image = precess_image(origin_image, input_width, input_height)

    input_image = torch.from_numpy(input_image)
    input_image = Variable(input_image.unsqueeze(0))
    # 判断 input_image 是否在 CUDA 设备上
    if torch.cuda.is_available():
        input_image = input_image.cuda()
        print("输入图像在CUDA设备上")
    else:
        print("输入图像在CPU上")

    output = model(input_image)

    seg_output = softmax(output[0].cpu().detach().numpy(), axis=1)[0]
    cls_output = softmax(output[1].cpu().detach().numpy(), axis=2)[0]

    cls_output = np.argmax(cls_output, axis=1)

    mask = np.zeros(shape=(input_height, input_width, 3))

    lane_id = []
    write_pos = []

    for i in range(mask.shape[0] - 1, 0, -1):
        for j in range(mask.shape[1] - 1, 0, -1):
            max_index = np.argmax(seg_output[:, i, j])
            if max_index not in lane_id:
                lane_id.append(max_index)
                if i > input_height - 20 or j > input_width - 20:
                    write_pos.append([j - 20, i - 20])
                else:
                    write_pos.append([j, i])
            if max_index != 0 and seg_output[max_index, i, j] > 0.5:
                mask[i, j, :] = color_list[max_index]

    mask = cv2.resize(mask, (image_width, image_height))

    for i in range(len(lane_id)):
        if lane_id[i] == 0:
            continue

        lane_type = cls_output[lane_Id_type.index(lane_id[i])]

        px = int(write_pos[i][0] / input_width * image_width)
        py = int(write_pos[i][1] / input_height * image_height)

        cv2.putText(origin_image, str(lane_id[i]), (px, py), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2, cv2.LINE_AA)
        cv2.putText(origin_image, str(lane_type), (px, py + 20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2, cv2.LINE_AA)

    cv2.putText(origin_image, 'lane_id: 7-5-3-1-2-4-6-8', (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2, cv2.LINE_AA)

    cv2.putText(origin_image, 'line type:', (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2, cv2.LINE_AA)
    for i in range(len(line_type)):
        cv2.putText(origin_image, str(i) + ': ' + str(line_type[i]), (10, 80 + i * 30), cv2.FONT_HERSHEY_SIMPLEX, 0.6,(0, 0, 0), 2, cv2.LINE_AA)

    opencv_image = np.clip(np.array(origin_image) + np.array(mask) * 0.4, a_min=0, a_max=255)
    opencv_image = opencv_image.astype("uint8")
    cv2.imwrite('./images/result.jpg', opencv_image)


def test_video():
    weights_path = './weights/best.pth'
    video_path = './images/test.mp4'

    model = Net(in_channels=3, lane_num=lane_num, type_num=type_num, width_mult=width_mult, is_deconv=True, is_batchnorm=True, is_ds=True)

    if torch.cuda.is_available():
        model = model.cuda()
        print("Model moved to GPU.")
    else:
        print("No CUDA device available, model will run on CPU.")

    model.load_state_dict(torch.load(weights_path))
    model.eval()

    cap = cv2.VideoCapture(video_path)

    # 获取视频的宽度和高度
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # 设置新的宽度和高度
    new_width = int(width * 0.7)
    new_height = int(height * 0.7)

    # 创建一个新的VideoWriter对象,用于保存缩放后的视频
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter('./images/output.mp4', fourcc, cap.get(cv2.CAP_PROP_FPS), (new_width, new_height))
    k = 0
    while True:
        ret, origin_image = cap.read()
        if not ret:
            break
        k = k+1
        print("k = ", k)
        image_height, image_width = origin_image.shape[:2]

        input_image = precess_image(origin_image, input_width, input_height)

        input_image = torch.from_numpy(input_image)
        input_image = Variable(input_image.unsqueeze(0))
        # 判断 input_image 是否在 CUDA 设备上
        if torch.cuda.is_available():
            input_image = input_image.cuda()
            print("输入图像在CUDA设备上")
        else:
            print("输入图像在CPU上")

        output = model(input_image)

        seg_output = softmax(output[0].cpu().detach().numpy(), axis=1)[0]
        cls_output = softmax(output[1].cpu().detach().numpy(), axis=2)[0]

        cls_output = np.argmax(cls_output, axis=1)

        mask = np.zeros(shape=(input_height, input_width, 3))

        lane_id = []
        write_pos = []

        for i in range(mask.shape[0] - 1, 0, -1):
            for j in range(mask.shape[1] - 1, 0, -1):
                max_index = np.argmax(seg_output[:, i, j])
                if max_index not in lane_id:
                    lane_id.append(max_index)
                    if i > input_height - 20 or j > input_width - 20:
                        write_pos.append([j - 20, i - 20])
                    else:
                        write_pos.append([j, i])
                if max_index != 0 and seg_output[max_index, i, j] > 0.5:
                    mask[i, j, :] = color_list[max_index]

        mask = cv2.resize(mask, (image_width, image_height))

        for i in range(len(lane_id)):
            if lane_id[i] == 0:
                continue

            lane_type = cls_output[lane_Id_type.index(lane_id[i])]

            px = int(write_pos[i][0] / input_width * image_width)
            py = int(write_pos[i][1] / input_height * image_height)

            cv2.putText(origin_image, str(lane_id[i]), (px, py), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2, cv2.LINE_AA)
            cv2.putText(origin_image, str(lane_type), (px, py + 20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2, cv2.LINE_AA)

        cv2.putText(origin_image, 'lane_id: 7-5-3-1-2-4-6-8', (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2, cv2.LINE_AA)

        cv2.putText(origin_image, 'line type:', (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2, cv2.LINE_AA)
        for i in range(len(line_type)):
            cv2.putText(origin_image, str(i) + ': ' + str(line_type[i]), (10, 80 + i * 30), cv2.FONT_HERSHEY_SIMPLEX, 0.6,(0, 0, 0), 2, cv2.LINE_AA)

        opencv_image = np.clip(np.array(origin_image) + np.array(mask) * 0.4, a_min=0, a_max=255)
        opencv_image = opencv_image.astype("uint8")

        # 缩放帧
        resized_frame = cv2.resize(opencv_image, (new_width, new_height))

        # 写入新的VideoWriter对象
        out.write(resized_frame)

        #cv2.imshow("Image",resized_frame)

        #if cv2.waitKey(1) & 0xFF == ord("q"):
        #    break



if __name__ == '__main__':
    print('test image ...')
    # test_image()
    test_video()

测试结果
image.png

相关实践学习
部署Stable Diffusion玩转AI绘画(GPU云服务器)
本实验通过在ECS上从零开始部署Stable Diffusion来进行AI绘画创作,开启AIGC盲盒。
相关文章
|
8天前
|
人工智能 测试技术 API
AI计算机视觉笔记二十 九:yolov10竹签模型,自动数竹签
本文介绍了如何在AutoDL平台上搭建YOLOv10环境并进行竹签检测与计数。首先从官网下载YOLOv10源码并创建虚拟环境,安装依赖库。接着通过官方模型测试环境是否正常工作。然后下载自定义数据集并配置`mycoco128.yaml`文件,使用`yolo detect train`命令或Python代码进行训练。最后,通过命令行或API调用测试训练结果,并展示竹签计数功能。如需转载,请注明原文出处。
|
8天前
|
机器学习/深度学习 人工智能 PyTorch
AI计算机视觉笔记三十二:LPRNet车牌识别
LPRNet是一种基于Pytorch的高性能、轻量级车牌识别框架,适用于中国及其他国家的车牌识别。该网络无需对字符进行预分割,采用端到端的轻量化设计,结合了squeezenet和inception的思想。其创新点在于去除了RNN,仅使用CNN与CTC Loss,并通过特定的卷积模块提取上下文信息。环境配置包括使用CPU开发板和Autodl训练环境。训练和测试过程需搭建虚拟环境并安装相关依赖,执行训练和测试脚本时可能遇到若干错误,需相应调整代码以确保正确运行。使用官方模型可获得较高的识别准确率,自行训练时建议增加训练轮数以提升效果。
|
8天前
|
人工智能 开发工具 计算机视觉
AI计算机视觉笔记三十:yolov8_obb旋转框训练
本文介绍了如何使用AUTODL环境搭建YOLOv8-obb的训练流程。首先创建虚拟环境并激活,然后通过指定清华源安装ultralytics库。接着下载YOLOv8源码,并使用指定命令开始训练,过程中可能会下载yolov8n.pt文件。训练完成后,可使用相应命令进行预测测试。
|
8天前
|
人工智能 监控 算法
AI计算机视觉笔记二十 八:基于YOLOv8实例分割的DeepSORT多目标跟踪
本文介绍了YOLOv8实例分割与DeepSORT视觉跟踪算法的结合应用,通过YOLOv8进行目标检测分割,并利用DeepSORT实现特征跟踪,在复杂环境中保持目标跟踪的准确性与稳定性。该技术广泛应用于安全监控、无人驾驶等领域。文章提供了环境搭建、代码下载及测试步骤,并附有详细代码示例。
|
4月前
|
机器学习/深度学习 计算机视觉
AIGC核心技术——计算机视觉(CV)预训练大模型
【1月更文挑战第13天】AIGC核心技术——计算机视觉(CV)预训练大模型
540 3
AIGC核心技术——计算机视觉(CV)预训练大模型
|
9月前
|
机器学习/深度学习 PyTorch 算法框架/工具
Azure 机器学习 - 使用 ONNX 对来自 AutoML 的计算机视觉模型进行预测
Azure 机器学习 - 使用 ONNX 对来自 AutoML 的计算机视觉模型进行预测
107 0
|
8天前
|
人工智能 测试技术 PyTorch
AI计算机视觉笔记二十四:YOLOP 训练+测试+模型评估
本文介绍了通过正点原子的ATK-3568了解并实现YOLOP(You Only Look Once for Panoptic Driving Perception)的过程,包括训练、测试、转换为ONNX格式及在ONNX Runtime上的部署。YOLOP由华中科技大学团队于2021年发布,可在Jetson TX2上达到23FPS,实现了目标检测、可行驶区域分割和车道线检测的多任务学习。文章详细记录了环境搭建、训练数据准备、模型转换和测试等步骤,并解决了ONNX转换过程中的问题。
|
2月前
|
自然语言处理 监控 自动驾驶
大模型在自然语言处理(NLP)、计算机视觉(CV)和多模态模型等领域应用最广
【7月更文挑战第26天】大模型在自然语言处理(NLP)、计算机视觉(CV)和多模态模型等领域应用最广
60 11
|
3月前
|
编解码 机器人 测试技术
2024年6月计算机视觉论文推荐:扩散模型、视觉语言模型、视频生成等
6月还有一周就要结束了,我们今天来总结2024年6月上半月发表的最重要的论文,重点介绍了计算机视觉领域的最新研究和进展。
109 8
|
2月前
|
机器学习/深度学习 人工智能 自然语言处理
计算机视觉借助深度学习实现了革命性进步,从图像分类到复杂场景理解,深度学习模型如CNN、RNN重塑了领域边界。
【7月更文挑战第2天】计算机视觉借助深度学习实现了革命性进步,从图像分类到复杂场景理解,深度学习模型如CNN、RNN重塑了领域边界。AlexNet开启新时代,后续模型不断优化,推动对象检测、语义分割、图像生成等领域发展。尽管面临数据隐私、模型解释性等挑战,深度学习已广泛应用于安防、医疗、零售和农业,预示着更智能、高效的未来,同时也强调了技术创新、伦理考量的重要性。
43 1

热门文章

最新文章