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盲盒。
相关文章
|
1月前
|
人工智能 测试技术 API
AI计算机视觉笔记二十 九:yolov10竹签模型,自动数竹签
本文介绍了如何在AutoDL平台上搭建YOLOv10环境并进行竹签检测与计数。首先从官网下载YOLOv10源码并创建虚拟环境,安装依赖库。接着通过官方模型测试环境是否正常工作。然后下载自定义数据集并配置`mycoco128.yaml`文件,使用`yolo detect train`命令或Python代码进行训练。最后,通过命令行或API调用测试训练结果,并展示竹签计数功能。如需转载,请注明原文出处。
|
1月前
|
机器学习/深度学习 人工智能 PyTorch
AI计算机视觉笔记三十二:LPRNet车牌识别
LPRNet是一种基于Pytorch的高性能、轻量级车牌识别框架,适用于中国及其他国家的车牌识别。该网络无需对字符进行预分割,采用端到端的轻量化设计,结合了squeezenet和inception的思想。其创新点在于去除了RNN,仅使用CNN与CTC Loss,并通过特定的卷积模块提取上下文信息。环境配置包括使用CPU开发板和Autodl训练环境。训练和测试过程需搭建虚拟环境并安装相关依赖,执行训练和测试脚本时可能遇到若干错误,需相应调整代码以确保正确运行。使用官方模型可获得较高的识别准确率,自行训练时建议增加训练轮数以提升效果。
|
1月前
|
人工智能 开发工具 计算机视觉
AI计算机视觉笔记三十:yolov8_obb旋转框训练
本文介绍了如何使用AUTODL环境搭建YOLOv8-obb的训练流程。首先创建虚拟环境并激活,然后通过指定清华源安装ultralytics库。接着下载YOLOv8源码,并使用指定命令开始训练,过程中可能会下载yolov8n.pt文件。训练完成后,可使用相应命令进行预测测试。
|
1月前
|
人工智能 监控 算法
AI计算机视觉笔记二十 八:基于YOLOv8实例分割的DeepSORT多目标跟踪
本文介绍了YOLOv8实例分割与DeepSORT视觉跟踪算法的结合应用,通过YOLOv8进行目标检测分割,并利用DeepSORT实现特征跟踪,在复杂环境中保持目标跟踪的准确性与稳定性。该技术广泛应用于安全监控、无人驾驶等领域。文章提供了环境搭建、代码下载及测试步骤,并附有详细代码示例。
|
13天前
|
机器学习/深度学习 数据采集 人工智能
探索AI技术在文本生成中的应用与挑战
【9月更文挑战第26天】本文深入探讨了AI技术在文本生成领域的应用,并分析了其面临的挑战。通过介绍AI文本生成的基本原理、应用场景以及未来发展趋势,帮助读者全面了解该技术的潜力和局限性。同时,文章还提供了代码示例,展示了如何使用Python和相关库实现简单的文本生成模型。
41 9
|
7天前
|
人工智能 自然语言处理 搜索推荐
AI技术在智能客服系统中的应用与挑战
【9月更文挑战第32天】本文将探讨AI技术在智能客服系统中的应用及其面临的挑战。我们将分析AI技术如何改变传统客服模式,提高服务质量和效率,并讨论在实际应用中可能遇到的问题和解决方案。
111 65
|
3天前
|
机器学习/深度学习 人工智能 自然语言处理
探索AI技术在医疗领域的应用
【8月更文挑战第67天】随着人工智能技术的不断发展,其在医疗领域的应用也日益广泛。本文将探讨AI技术在医疗领域的应用,包括辅助诊断、个性化治疗方案和智能医疗设备等方面。通过实例分析,我们将了解AI技术如何改变传统医疗模式,提高医疗服务的质量和效率。
38 16
|
2天前
|
机器学习/深度学习 人工智能 算法
AI在医疗健康领域的应用
随着人工智能技术的不断发展,其在医疗健康领域的应用也日益广泛。从辅助诊断、个性化治疗方案的制定,到疾病预防和健康管理,AI技术都在发挥着重要作用。本文将探讨AI在医疗健康领域的应用,包括其在医学影像分析、基因编辑、药物研发等方面的应用,以及其对医疗行业未来发展的影响。
|
9天前
|
机器学习/深度学习 人工智能 供应链
精准农业:AI在农业生产中的应用
【10月更文挑战第1天】随着科技的发展,人工智能(AI)逐渐渗透到农业领域,通过精准监控和管理提升了农业生产效率和质量。AI在精准农业中的应用包括:精准农田管理,如个性化灌溉和施肥;作物病虫害识别与预测,及时发现并预防病虫害;智能农机自动化作业,提高作业效率;农产品质量检测与分类,确保品质;农业供应链优化,预测需求和价格。尽管面临数据收集、技术接受度等挑战,AI在精准农业中的未来前景广阔,有望实现全程自动化作业、数据驱动决策及智能预警系统,推动农业可持续发展。
43 11
|
3天前
|
机器学习/深度学习 人工智能 搜索推荐
AI技术在医疗领域的应用及挑战
【9月更文挑战第36天】随着人工智能技术的不断发展,其在医疗领域的应用也日益广泛。本文将从AI技术在医疗领域的应用场景、优势以及面临的挑战等方面进行探讨。通过分析AI技术在医疗领域的应用,我们可以更好地了解其发展趋势和未来前景。
17 3

热门文章

最新文章