反光衣安全帽检测数据集，用于铁路工人目标检测

铁路施工场景，数据集3220+，vest,helmet,worker共三个类别

图片与标签txt一一对应，直接用yolov5-10训练

如何使用YOLOv5或YOLOv7来训练一个包含3220多张图像的反光衣安全帽检测数据集，并附上详细的训练代码和步骤。这个数据集用于铁路工人目标检测，包含三个类别：反光衣（vest）、安全帽（helmet）和工人（worker）。

数据集描述

• 数据量：3220多张图像
• 类别：
  • 0: 反光衣（vest）
  • 1: 安全帽（helmet）
  • 2: 工人（worker）
• 标注格式：YOLO格式
• 应用场景：铁路施工场景

数据集组织

假设你的数据集目录结构如下：

railway_worker_detection_dataset/
├── train/
│   ├── images/
│   └── labels/
├── valid/
│   ├── images/
│   └── labels/
└── data.yaml  # 数据配置文件

其中：

• train/ 包含训练集的图像和TXT标签文件。
• valid/ 包含验证集的图像和TXT标签文件。
• data.yaml 包含数据配置信息。

数据配置文件

创建或确认data.yaml文件是否正确配置了数据集路径和类别信息：

train: ./train/images/  # 训练集图像路径
val: ./valid/images/  # 验证集图像路径

# Classes
nc: 3  # 类别数量
names: 
  - vest
  - helmet
  - worker  # 类别名称列表

安装YOLOv5

如果你还没有安装YOLOv5，可以使用以下命令安装：

pip install yolov5

训练模型

使用YOLOv5训练模型的命令非常简单，你可以直接使用以下命令开始训练：

cd path/to/railway_worker_detection_dataset/

# 克隆YOLOv5仓库
git clone https://github.com/ultralytics/yolov5.git
cd yolov5

# 下载预训练权重
wget https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.pt

# 开始训练
python train.py --img 640 --batch 16 --epochs 100 --data ../data.yaml --weights yolov5s.pt --cache

在这个命令中：

• --img 640：输入图像的大小。
• --batch 16：批量大小。
• --epochs 100：训练轮数。
• --data ../data.yaml：指定数据配置文件。
• --weights yolov5s.pt：指定预训练权重。
• --cache：使用缓存以加快训练速度。

模型评估

训练完成后，可以使用以下命令评估模型在验证集上的表现：

python val.py --data ../data.yaml --weights runs/train/exp/weights/best.pt --img 640

这里的runs/train/exp/weights/best.pt是训练过程中产生的最佳模型权重文件。

模型预测

你可以使用训练好的模型对新图像进行预测：

python detect.py --source path/to/your/image.jpg --weights runs/train/exp/weights/best.pt --img 640 --conf 0.4 --iou 0.5

查看训练结果

训练过程中的日志和结果会保存在runs/train/目录下，你可以查看训练过程中的损失、精度等信息。

升级到YOLOv7

如果你想要使用YOLOv7进行训练，可以按照以下步骤操作：

1. 安装YOLOv7：首先需要确认YOLOv7是否已经发布，并且是否有可用的开源实现。目前YOLOv7是最新的公开版本之一，可以按照类似的方式安装和使用深色版本

2. cd path/to/railway_worker_detection_dataset/
   git clone https://github.com/WongKinYiu/yolov7.git
   cd yolov7

3. 下载预训练权重：

   wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7.pt

4. 开始训练：

   python train.py --img 640 --batch 16 --epochs 100 --data ../data.yaml --weights yolov7.pt --cache

在这个命令中：

• --img 640：输入图像的大小。
• --batch 16：批量大小。
• --epochs 100：训练轮数。
• --data ../data.yaml：指定数据配置文件。
• --weights yolov7.pt：指定预训练权重。
• --cache：使用缓存以加快训练速度。

数据增强

为了进一步提高模型性能，可以使用数据增强技术。以下是一个简单的数据增强示例：

1. 安装albumentations库：

   pip install -U albumentations

2. 在train.py中添加数据增强：

   import albumentations as A
   from albumentations.pytorch import ToTensorV2
   import cv2
   
   # 定义数据增强
   transform = A.Compose([
       A.RandomSizedBBoxSafeCrop(width=640, height=640, erosion_rate=0.2),
       A.HorizontalFlip(p=0.5),
       A.VerticalFlip(p=0.5),
       A.Rotate(limit=10, p=0.5, border_mode=cv2.BORDER_CONSTANT),
       A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2, p=0.5),
       A.GaussNoise(var_limit=(10.0, 50.0), p=0.5),
       A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
       ToTensorV2()
   ], bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))
   
   # 在数据加载器中应用数据增强
   def collate_fn(batch):
       images, targets = zip(*batch)
       transformed_images = []
       transformed_targets = []
   
       for img, target in zip(images, targets):
           bboxes = target['bboxes']
           class_labels = target['labels']
           augmented = transform(image=img, bboxes=bboxes, class_labels=class_labels)
           transformed_images.append(augmented['image'])
           transformed_targets.append({
               'bboxes': augmented['bboxes'],
               'labels': augmented['class_labels']
           })
   
       return torch.stack(transformed_images), transformed_targets

注意事项

• 数据集质量：确保数据集的质量，包括清晰度、标注准确性等。
• 模型选择：可以选择更强大的模型版本（如YOLOv5m、YOLOv5l等）以提高性能。
• 超参数调整：根据实际情况调整超参数，如批量大小（batch-size）、图像大小（imgsz）等。
• 监控性能：训练过程中监控损失函数和mAP指标，确保模型收敛。

通过上述步骤，你可以使用YOLOv5或YOLOv7来训练一个铁路工人目标检测数据集，并使用训练好的模型进行预测。希望这对你有帮助