如何使用yolov8训练——电力设备漏油检测图像数据集，数据集总数为687张，标注为voc格式，也可以转yolo格式

电力设备漏油检测图像数据集，数据集总数为687张，标注为voc格式，也可以转yolo格式

如何使用YOLOv8训练电力设备漏油检测图像数据集的详细步骤和代码。我们将使用PyTorch和YOLOv8的官方实现来进行训练。

项目结构

OilLeakDetection/
├── data/
│ ├── images/
│ │ ├── train/
│ │ ├── val/
│ │ └── test/
│ ├── labels/
│ │ ├── train/
│ │ ├── val/
│ │ └── test/
├── models/
│ └── yolov8/
├── utils/
│ ├── data_loader.py
│ ├── metrics.py
│ └── plot.py
├── main.py
├── train.py
├── infer.py
└── README.md

1. 数据集准备
   数据集格式
   假设你的数据集已经标注好，并且格式为VOC格式。每个图像对应一个XML文件，包含目标框的坐标和类别标签。你可以使用工具将VOC格式转换为YOLO格式。

数据集目录结构
深色版本
data/
├── images/
│ ├── train/
│ ├── val/
│ └── test/
└── labels/
├── train/
├── val/
└── test/
2. 安装YOLOv8
首先，克隆YOLOv8仓库并安装必要的依赖项：

=(https://i-blog.csdnimg.cn/direct/a66d3e1332a34f168071de1db29ebadd.png)

cd models/
git clone https://github.com/ultralytics/ultralytics.git
cd ultralytics
pip install -r requirements.txt
3. 数据配置文件
创建一个数据配置文件data/oil_leak.yaml，指定数据集路径和类别信息：

train: …/data/images/train/
val: …/data/images/val/
test: …/data/images/test/

nc: 1 # 类别数量
names:

• oil_leak

1. 训练模型
   编写一个训练脚本train.py：

import os
from ultralytics import YOLO

def main():
# 设置训练参数
data_yaml = ‘…/data/oil_leak.yaml’
model_yaml = ‘yolov8n.yaml’ # 选择合适的YOLOv8模型
epochs = 100
batch_size = 16
img_size = 640
device = ‘cuda’ if torch.cuda.is_available() else ‘cpu’

# 加载模型
model = YOLO(model_yaml)

# 开始训练
results = model.train(
    data=data_yaml,
    epochs=epochs,
    batch=batch_size,
    imgsz=img_size,
    device=device
)

if name == ‘main’:
main()
运行训练脚本：

python train.py
5. 推理和可视化
编写一个推理脚本infer.py，用于加载训练好的模型并对新图像进行预测：
import os
import torch
import cv2
import numpy as np
from ultralytics import YOLO
from ultralytics.yolo.utils.general import non_max_suppression, scale_coords
from ultralytics.yolo.utils.plots import Annotator

def load_model(weights_path, device):
model = YOLO(weights_path)
model.to(device)
return model

def infer_image(model, img_path, img_size, device):
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, (img_size, img_size))
img = img.transpose(2, 0, 1)
img = img.astype(np.float32) / 255.0
img = torch.from_numpy(img).unsqueeze(0).to(device)

with torch.no_grad():
    pred = model(img)[0]

pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45)

return pred, img

def draw_boxes(img, pred, names):
annotator = Annotator(img, line_width=2, example=str(names))
for i, det in enumerate(pred):
if len(det):
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img.shape[2:]).round()
for *xyxy, conf, cls in reversed(det):
c = int(cls)
label = f’{names[c]} {conf:.2f}’
annotator.box_label(xyxy, label, color=(0, 255, 0))
return annotator.result()

if name == ‘main’:
weights_path = ‘runs/train/exp/weights/best.pt’
img_path = ‘…/data/images/test/000001.jpg’
img_size = 640
device = ‘cuda’ if torch.cuda.is_available() else ‘cpu’

model = load_model(weights_path, device)
pred, img = infer_image(model, img_path, img_size, device)

# 读取原始图像
original_img = cv2.imread(img_path)

# 绘制预测框
annotated_img = draw_boxes(original_img, pred, names=['oil_leak'])

# 显示结果
cv2.imshow('Result', annotated_img)
cv2.waitKey(0)
cv2.destroyAllWindows()

1. 评价指标
   编写一个脚本evaluate.py来计算模型的评价指标（如IoU、mAP等）：

import os
import torch
from ultralytics import YOLO
from ultralytics.yolo.utils.general import check_img_size
from ultralytics.yolo.utils.datasets import create_dataloader

def main():
data_yaml = ‘…/data/oil_leak.yaml’
weights_path = ‘runs/train/exp/weights/best.pt’
img_size = 640
batch_size = 16
device = ‘cuda’ if torch.cuda.is_available() else ‘cpu’

# 加载数据
dataloader = create_dataloader('../data/images/val/', img_size, batch_size, 32, single_cls=False, pad=0.5, rect=True)[0]

# 加载模型
model = YOLO(weights_path)
model.to(device)

# 验证模型
results = model.val(
    data=data_yaml,
    batch=batch_size,
    imgsz=img_size,
    device=device,
    dataloader=dataloader
)

if name == ‘main’:
main()
运行评价脚本：

python evaluate.py
7. 运行项目
确保数据集和标签文件已经准备好，并放在相应的目录中。
运行训练脚本：

python train.py
运行推理脚本：

python infer.py
运行评价脚本：

python evaluate.py
8. 代码说明
数据集准备：确保数据集已经标注好，并且格式为YOLOv8所需的格式。
训练模型：使用train.py脚本训练模型，指定数据配置文件和训练参数。
推理和可视化：使用infer.py脚本加载训练好的模型并对新图像进行预测，使用evaluate.py脚本计算模型的评价指标。
希望这些代码和说明能帮助你完成基于YOLOv8的电力设备漏油检测项目。