如何训练自己的数据集之——交通红绿灯数据集检测系统
【交通信号灯检测数据集】10742张，3类
names: ['green', 'red','yellow']
名称：[绿灯，红灯，黄灯]
共10742张，8:1:1比例划分，（train;8593张，val：1074张，test：1075张标注文件为YOLO适用的txt格式。可以直接用于模型训练。
map50:87.8%
其中green：80.3%；red：84%；yellow：98.9%
交通信号灯检测数据集介绍
数据集概述
该数据集专为交通信号灯检测设计，包含10742张图像，每张图像均附带有YOLO格式的.txt标签文件，并按照8:1:1的比例划分为训练集、验证集和测试集。以下是该数据集的详细信息：
- 数据量 ：共10742张图像。
- 类别 ：三类交通信号灯（绿灯、红灯、黄灯）。
- 数据格式 ：每张图像均附有
.txt标签文件（YOLO格式）。 - 数据划分 ：
- 训练集：8593张图像。
- 验证集：1074张图像。
- 测试集：1075张图像。
- 性能指标 ：
- mAP50：87.8%
- green：80.3%
- red：84%
- yellow：98.9%
数据集特点
- 专业性强 ：专为交通信号灯检测设计，适用于智能交通系统、自动驾驶等领域。
- 数据完整性 ：每张图像均附有详细的标签信息，便于模型学习。
- 格式统一 ：所有标签统一为YOLO格式，方便直接用于模型训练。
- 配置文件 ：包含一个
.yaml文件，用于描述数据集的配置信息。 - 性能优异 ：在不同的类别上都有较高的准确率，尤其是黄色信号灯达到了98.9%，表明模型在黄色信号灯的检测上表现非常出色。
数据集内容
- 图像文件 ：包含10742张JPG/PNG格式的图像文件。
- 标签文件 ：每张图像对应一个
.txt格式的标签文件，包含对象的边界框位置信息。 - 配置文件 ：一个
.yaml文件，描述了数据集的基本配置信息，如类别数量、类别名称、训练集路径等。
数据集结构示例
假设数据集的根目录为 traffic_light_detection_dataset，其结构可能如下所示：
traffic_light_detection_dataset/
├── train/
│ ├── images/
│ │ ├── image_000001.jpg
│ │ ├── image_000002.jpg
│ │ └── ...
│ ├── labels/
│ │ ├── image_000001.txt
│ │ ├── image_000002.txt
│ │ └── ...
├── valid/
│ ├── images/
│ │ ├── image_000001.jpg
│ │ ├── image_000002.jpg
│ │ └── ...
│ ├── labels/
│ │ ├── image_000001.txt
│ │ ├── image_000002.txt
│ │ └── ...
├── test/
│ ├── images/
│ │ ├── image_000001.jpg
│ │ ├── image_000002.jpg
│ │ └── ...
│ ├── labels/
│ │ ├── image_000001.txt
│ │ ├── image_000002.txt
│ │ └── ...
└── data.yaml # 数据集配置文件
数据集配置文件 data.yaml
创建一个 data.yaml 文件来描述您的数据集。这里假设数据集被放置在一个名为 traffic_light_detection_dataset 的目录中，且包含 train、valid 和 test 子目录。
# data.yaml 文件
train: ../traffic_light_detection_dataset/train/images/
val: ../traffic_light_detection_dataset/valid/images/
test: ../traffic_light_detection_dataset/test/images/
nc: 3 # number of classes
names: ['green', 'red', 'yellow'] # 类别名称
关键训练代码
安装YOLOv7
如果您还没有安装YOLOv7框架，请按照官方文档执行以下命令：
git clone https://github.com/WongKinYiu/yolov7.git # 克隆YOLOv7仓库
cd yolov7
pip install -r requirements.txt # 安装依赖项
使用YOLOv7训练
使用以下命令开始训练模型：
cd yolov7
python train.py --img 640 --batch 16 --epochs 100 --data ../traffic_light_detection_dataset/data.yaml --weights yolov7.pt --cache
自定义训练脚本
如果需要更详细的控制，可以编写一个Python脚本来执行训练过程。以下是一个简单的脚本示例：
import torch
from yolov7.utils.datasets import LoadImagesAndLabels
from yolov7.models.yolo import Model
def train_model():
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
data_yaml = '../traffic_light_detection_dataset/data.yaml'
# 加载数据集
train_loader = LoadImagesAndLabels(data_yaml, img_size=(640, 640), batch_size=16, augment=True)
# 加载预训练模型
model = Model('yolov7.yaml').to(device)
# 设置训练参数
epochs = 100
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=0.0005)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
# 开始训练
for epoch in range(epochs):
model.train()
for batch_idx, (images, targets) in enumerate(train_loader):
images = images.to(device)
targets = [t.to(device) for t in targets]
# 前向传播
output = model(images)
# 计算损失
loss = model.loss(output, targets)
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 更新学习率
scheduler.step()
# 保存模型
torch.save(model.state_dict(), 'trained_model.pt')
print('Training complete.')
if __name__ == '__main__':
train_model()
使用训练好的权重文件进行测试
如果您已经有了训练好的模型权重文件，可以直接使用它来进行测试。假设权重文件名为 best.pt，可以使用以下命令
python detect.py --weights best.pt --img 640 --conf 0.4 --source ../traffic_light_detection_dataset/test/images/
总结
这个示例展示了如何使用YOLOv7框架训练交通信号灯检测数据集。您可以根据自己的需求调整脚本中的参数和逻辑。通过使用这个数据集和相应的训练代码，您可以有效地训练出一个能够在多种条件下检测交通信号灯的模型。在实际应用中，可以根据具体需求调整模型结构和训练参数。此外，提供的训练代码可以让您快速启动训练流程，并获得良好的检测效果。
如果您希望扩展类别或者有不同的数据集划分选项（比如加入更多的训练集、验证集和测试集），可以在 data.yaml 文件中进行相应的调整，并在训练和测试过程中指定相应的路径。