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盲盒。
相关文章
|
2月前
|
人工智能 测试技术 API
AI计算机视觉笔记二十 九:yolov10竹签模型,自动数竹签
本文介绍了如何在AutoDL平台上搭建YOLOv10环境并进行竹签检测与计数。首先从官网下载YOLOv10源码并创建虚拟环境,安装依赖库。接着通过官方模型测试环境是否正常工作。然后下载自定义数据集并配置`mycoco128.yaml`文件,使用`yolo detect train`命令或Python代码进行训练。最后,通过命令行或API调用测试训练结果,并展示竹签计数功能。如需转载,请注明原文出处。
|
2月前
|
机器学习/深度学习 人工智能 PyTorch
AI计算机视觉笔记三十二:LPRNet车牌识别
LPRNet是一种基于Pytorch的高性能、轻量级车牌识别框架,适用于中国及其他国家的车牌识别。该网络无需对字符进行预分割,采用端到端的轻量化设计,结合了squeezenet和inception的思想。其创新点在于去除了RNN,仅使用CNN与CTC Loss,并通过特定的卷积模块提取上下文信息。环境配置包括使用CPU开发板和Autodl训练环境。训练和测试过程需搭建虚拟环境并安装相关依赖,执行训练和测试脚本时可能遇到若干错误,需相应调整代码以确保正确运行。使用官方模型可获得较高的识别准确率,自行训练时建议增加训练轮数以提升效果。
|
2月前
|
人工智能 开发工具 计算机视觉
AI计算机视觉笔记三十:yolov8_obb旋转框训练
本文介绍了如何使用AUTODL环境搭建YOLOv8-obb的训练流程。首先创建虚拟环境并激活,然后通过指定清华源安装ultralytics库。接着下载YOLOv8源码,并使用指定命令开始训练,过程中可能会下载yolov8n.pt文件。训练完成后,可使用相应命令进行预测测试。
|
2月前
|
人工智能 监控 算法
AI计算机视觉笔记二十 八:基于YOLOv8实例分割的DeepSORT多目标跟踪
本文介绍了YOLOv8实例分割与DeepSORT视觉跟踪算法的结合应用,通过YOLOv8进行目标检测分割,并利用DeepSORT实现特征跟踪,在复杂环境中保持目标跟踪的准确性与稳定性。该技术广泛应用于安全监控、无人驾驶等领域。文章提供了环境搭建、代码下载及测试步骤,并附有详细代码示例。
|
机器学习/深度学习 人工智能 自然语言处理
欧洲科学院院士:中国领先计算机视觉和机器人领域,但AI研究还不足以支撑垂直领域解决方案
如何构建一家伟大的人工智能研究机构?深度学习在自然语言处理方面有怎样的发展?中美欧人工智能市场和技术发展差异?德国人工智能研究中心(DFKI)科学董事,北京深知无限人工智能研究院(AITC)院长兼首席科学家汉斯·乌思克尔特(Hans Uszkoreit)教授,在新智元AI技术峰会上,分享了他的精彩见解。
1567 0
|
5天前
|
机器学习/深度学习 人工智能 自然语言处理
当前AI大模型在软件开发中的创新应用与挑战
2024年,AI大模型在软件开发领域的应用正重塑传统流程,从自动化编码、智能协作到代码审查和测试,显著提升了开发效率和代码质量。然而,技术挑战、伦理安全及模型可解释性等问题仍需解决。未来,AI将继续推动软件开发向更高效、智能化方向发展。
|
9天前
|
机器学习/深度学习 人工智能 自然语言处理
AI在医疗领域的应用及其挑战
【10月更文挑战第34天】本文将探讨人工智能(AI)在医疗领域的应用及其面临的挑战。我们将从AI技术的基本概念入手,然后详细介绍其在医疗领域的各种应用,如疾病诊断、药物研发、患者护理等。最后,我们将讨论AI在医疗领域面临的主要挑战,包括数据隐私、算法偏见、法规合规等问题。
28 1
|
7天前
|
机器学习/深度学习 人工智能 算法
AI在医疗领域的应用与挑战
本文探讨了人工智能(AI)在医疗领域的应用,包括其在疾病诊断、治疗方案制定、患者管理等方面的优势和潜力。同时,也分析了AI在医疗领域面临的挑战,如数据隐私、伦理问题以及技术局限性等。通过对这些内容的深入分析,旨在为读者提供一个全面了解AI在医疗领域现状和未来发展的视角。
36 10
|
7天前
|
机器学习/深度学习 人工智能 监控
探索AI在医疗领域的应用与挑战
本文深入探讨了人工智能(AI)在医疗领域中的应用现状和面临的挑战。通过分析AI技术如何助力疾病诊断、治疗方案优化、患者管理等方面的创新实践,揭示了AI技术为医疗行业带来的变革潜力。同时,文章也指出了数据隐私、算法透明度、跨学科合作等关键问题,并对未来的发展趋势进行了展望。
|
11天前
|
机器学习/深度学习 人工智能 自然语言处理
当前AI大模型在软件开发中的创新应用与挑战
【10月更文挑战第31天】2024年,AI大模型在软件开发领域的应用取得了显著进展,从自动化代码生成、智能代码审查到智能化测试,极大地提升了开发效率和代码质量。然而,技术挑战、伦理与安全问题以及模型可解释性仍是亟待解决的关键问题。开发者需不断学习和适应,以充分利用AI的优势。

热门文章

最新文章