斑马线+行人+交通灯(红绿灯)多类别YoLo数据集

置顶 naanaaan

已于 2025-04-16 13:56:28 修改

阅读量4.6k

点赞数 14

文章标签： YOLO 人工智能目标检测自动驾驶

于 2023-08-24 14:07:21 首次发布

本文链接：https://blog.csdn.net/na_a_na_a_an/article/details/132472844

版权

1. 数据集介绍

● 数据集原为本人毕设所用，数据集图片全部从滴滴数据集 D²-City 中抽帧获得，本人做了二次标注
● 数据集详细说明文档 📥下载 📥

① 数据集基本情况

Alt

原始三类别数据集

Alt

细分灯色六类别数据集

数据集共 12554 张图片，全为行车记录仪视角，标注效果👆；
数据集默认提供 txt 标签，适合一切目标检测模型训练，包括 YOLO 全系列、DETR…
数据集按照 交通灯是否分灯色标注 有以下两个版本：

注：3类别数据集训练的模型只能判断画面中是否有交通灯；6类别数据集训练的模型才能直接识别灯色

原始 3 类别细分灯色 6 类别
行人行人
斑马线斑马线
交通灯红灯
绿灯
黄灯
黑灯(坏灯或闪烁中的灯)

原始 3 类别	细分灯色 6 类别
行人	行人
斑马线	斑马线
交通灯	红灯
	绿灯
	黄灯
	黑灯(坏灯或闪烁中的灯)

② 数据集标注情况

标注效果视频：点这观看

Alt

原始三类别数据分布

Alt

灯色部分数据分布

上面统计了各类实例标注框的数量、位置与大小的情况，每一个矩形色块内含有以下内容：

柱形图：指明该类具有多少实例(即标注数量)，数值显示在柱子顶端；
X-Y 分布图：x 、y 分别为按帧宽高归一化后的标注框中心点横坐标与纵坐标，描述了该类实例标注框的位置分布。图中深色处即为该类实例在图中主要分布的位置。
width-height 分布图：width 、height 分别为按帧宽高归一化后的标注框宽度与高度，描述了该类实例标注框的大小分布。图中深色处即为该类实例标注框的主要宽度与高度。

2. 数据集使用

① 准备数据

请严格走完以下 123 步

视频下载：链接（此步骤必需 ! ! !）
标注下载：（以下各项按自己需要，选择一项或多项）

按以下步骤，分帧得数据集图片：
A. 安装opencv: 终端激活环境后，输入pip install opencv-python；
B.按注释配置以下脚本的 label_dir、save_dir、video_path；
(注意：若下载了3类别的，需要保证 save_dir 有 18G 的空间；单交通灯的需要 8G；单斑马线的需要 10G；单行人的需要 16G)
C.运行脚本，确保 save_dir 中最终的图片数 = label_dir 中的txt文件数

import os
import re
import cv2 
from glob import iglob

def validate_input(video_path, label_dir):
    assert os.path.exists(video_path), f'无法找到视频文件: {video_path}，请检查该路径是否存在!'
    assert video_path.endswith(".mp4"), 'video_path 需设置为所下载的mp4视频文件路径!'

    def label_check(file_path):
        file_name = os.path.basename(file_path)
        if file_name.startswith('frame'):
            try:
                idx = re.search(r'\d{5}', file_name).group()
            except AttributeError:
                raise ValueError(f'错误的文件名 {file_name}，请勿修改txt文件名，若已修改，请删除并重新下载!')
            return int(idx)
        else:
            raise ValueError(f'错误的文件名 {file_name}，请勿修改txt文件名或向标签所在文件夹增加自己的文件!')

    valid_idx = list(map( label_check, iglob(os.path.join(label_dir, '*.txt'), recursive=False) ))
    if not valid_idx:
        raise ValueError(f'无法在 {label_dir} 目录下找到txt标签文件，请检查该路径是否存在，或txt标签文件是否直接存在于该目录下!')

    return valid_idx

def main(video_path, label_dir, save_dir):
    valid_idx = validate_input(video_path, label_dir)

    cap = cv2.VideoCapture(video_path)
    frame_idx = 0
    count = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            cap.release()
            break
        if frame_idx in valid_idx:
            cv2.imwrite(os.path.join(save_dir, f'frame{frame_idx:05d}.png'), frame)
            count += 1
            print(f'\rsaving {count}...',end='')

        frame_idx += 1

if __name__ == '__main__':
    label_dir = r'' # 含有所有txt文件的文件夹路径, 不用删掉前面的r, 路径中最好不要有中文
    save_dir = r'' # 图片将要保存路径, 不用删掉前面的r, 路径中最好不要有中文
    video_path = r'' # 下载的视频路径, 不用删掉前面的r, 路径中最好不要有中文
    
    video_path = os.path.abspath(video_path)    
    label_dir = os.path.abspath(label_dir)    
    save_dir = os.path.abspath(save_dir)

    os.makedirs(save_dir,exist_ok=True)
    main(video_path, label_dir, save_dir)

② 划分数据集

按照自定比例(如8:2)，以随机抽样的方式抽取训练集与测试集（可自行搜索YOLO数据集划分代码）

③ 开始训练

数据集可用于一切目标检测模型，包括 YOLO全系列、DETR等，这里以 YOLOV5 为例
原三类别数据集上 YOLOV5 训练结果 | 📥下载 📥 (仅供借鉴，不含权重)

去YoLoV5官网下载预训练权重，如 yolov5m.pt
在 train.py 里完成超参数配置即可开始训练。本人利用YoLoV5的m6权重，在三类别数据集上以imgsz 640训练300轮，mAP 0.5 有 0.956，mAP 0.5~0.95 有 0.7299。