使用YOLOv8来训练——包含250张图像的工地钢筋计数数据集；数据集已标注为YOLO格式和VOC格式

计算机c9硕士算法工程师

1119人浏览 · 2024-11-02 06:32:43

计算机c9硕士算法工程师 · 2024-11-02 06:32:43 发布

可以直接用于智慧工地物料盘点，标签类别为rebar，共有30942个标注。

工地类——钢筋计数数据集数据集共有250张图片数据，可直接用于智慧工地物料盘点；保证标注质量；已标注yolo格式、VOC格式，可直接训练；标签类别及标签个数：rebar：30942；

使用YOLOv8来训练一个包含250张图像的钢筋计数数据集。这个数据集已标注为YOLO格式和VOC格式，可以直接用于智慧工地物料盘点，标签类别为rebar，共有30942个标注。

数据集描述

数据量：250张图像
标注格式：YOLO格式和VOC格式
应用场景：智慧工地物料盘点
标签类别及标签个数：
- rebar：30942个标注

数据集组织

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

rebar_counting_dataset/
├── images/
│   ├── train/
│   │   ├── 000001.jpg
│   │   ├── 000002.jpg
│   │   └── ...
│   └── valid/
│       ├── 000001.jpg
│       ├── 000002.jpg
│       └── ...
├── labels/
│   ├── train/
│   │   ├── 000001.txt
│   │   ├── 000002.txt
│   │   └── ...
│   └── valid/
│       ├── 000001.txt
│       ├── 000002.txt
│       └── ...
└── data.yaml  # 数据配置文件

数据配置文件

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

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

# Classes
nc: 1  # 类别数量
names: 
  - rebar  # 类别名称列表

安装YOLOv8

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

pip install ultralytics

训练模型

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

cd path/to/rebar_counting_dataset/

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

# 开始训练
python yolo.py detect train data=../data.yaml model=yolov8n.pt epochs=100 imgsz=640 batch=16

在这个命令中：

data=../data.yaml：指定数据配置文件。
model=yolov8n.pt：指定预训练权重，这里使用的是YOLOv8的小模型。
epochs=100：训练轮数。
imgsz=640：输入图像的大小。
batch=16：批量大小。

模型评估

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

python yolo.py detect val data=../data.yaml model=runs/detect/train/weights/best.pt imgsz=640

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

模型预测

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

python yolo.py detect predict source=path/to/your/image.jpg model=runs/detect/train/weights/best.pt imgsz=640 conf=0.4 iou=0.5

查看训练结果

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

数据增强

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

安装albumentations库：
```
pip install -U albumentations
```

在yolo.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