如何使用YOLOv5来训练——建筑工地安全图像数据集,并附上详细的训练代码和步调。这个数据集包罗10个类别,标注为YOLO格式。安全帽面罩安全锥等数据集进行检测
建筑工地安全活动图像数据集
yolo格式
0:“安全帽”,1:“面罩”,2:“无安全帽”、3:“无面罩”、4:“无安全背心”、5:“人”、6:“安全锥”、7:“安全背心”,8:“机器”,9:“车辆”
train(2605)val(114)test(82)
如何使用YOLOv5来训练一个包罗建筑工地安全图像的数据集,并附上详细的训练代码和步调。这个数据集包罗10个类别,标注为YOLO格式。
数据集描述
- 数据量:
- 训练集:2605张图像
- 验证集:114张图像
- 测试集:82张图像
- 类别:
- 0: 安全帽(Safety Helmet)
- 1: 面罩(Face Mask)
- 2: 无安全帽(No Safety Helmet)
- 3: 无面罩(No Face Mask)
- 4: 无安全背心(No Safety Vest)
- 5: 人(Person)
- 6: 安全锥(Safety Cone)
- 7: 安全背心(Safety Vest)
- 8: 机器(Machine)
- 9: 车辆(Vehicle)
- 标注格式:YOLO格式
数据集构造
假设你的数据集目录结构如下:
- construction_site_safety_dataset/
- ├── train/
- │ ├── images/
- │ └── labels/
- ├── valid/
- │ ├── images/
- │ └── labels/
- ├── test/
- │ ├── images/
- │ └── labels/
- └── data.yaml # 数据配置文件
复制代码 此中:
- train/ 包罗训练集的图像和TXT标签文件。
- valid/ 包罗验证集的图像和TXT标签文件。
- test/ 包罗测试集的图像和TXT标签文件。
- data.yaml 包罗数据设置信息。
数据设置文件
创建或确认data.yaml文件是否精确设置了数据集路径和类别信息:
- train: ./train/images/ # 训练集图像路径
- val: ./valid/images/ # 验证集图像路径
- test: ./test/images/ # 测试集图像路径
- # Classes
- nc: 10 # 类别数量
- names:
- - safety_helmet
- - face_mask
- - no_safety_helmet
- - no_face_mask
- - no_safety_vest
- - person
- - safety_cone
- - safety_vest
- - machine
- - vehicle # 类别名称列表
复制代码 安装YOLOv5
如果你还没有安装YOLOv5,可以使用以下命令安装:
训练模型
使用YOLOv5训练模型的命令非常简单,你可以直接使用以下命令开始训练:
- [/code] [code]cd path/to/construction_site_safety_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/目录下,你可以查看训练过程中的损失、精度等信息。
数据增强
为了进一步提高模型性能,可以使用数据增强技术。以下是一个简单的数据增强示例:
- 安装albumentations库:
- pip install -U albumentations
复制代码 - 在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来训练一个建筑工地安全图像数据集,并使用训练好的模型进行猜测
免责声明:如果侵犯了您的权益,请联系站长,我们会及时删除侵权内容,谢谢合作!更多信息从访问主页:qidao123.com:ToB企服之家,中国第一个企服评测及商务社交产业平台。 |