如何用Yolo训练自己的数据集和权重 ——光伏板热斑缺陷数据集
发布时间
阅读量：
阅读量
包含光伏板热斑缺陷的大量图像数据集，并用于目标检测的任务，可按照以下步骤执行
数据集描述
该数据集包含以下三种类别：
- 热斑（hotspot） ：表示光伏板上的热斑问题。
- 热点（hotpoint） ：表示光伏板上的热点问题。
- 遮挡（occlusion） ：表示光伏板上的覆盖现象。
数据集的规模较大，具体划分如下：
- 训练集 ：被用来进行训练的图片占大多数。
- 验证集 ：被用来评估模型性能的图片较少。
- 测试集 ：被用来进行全面评估的图片数量有限。
数据集组织
假设你的数据集目录结构如下：
pv_hotspot_dataset/
├── train/
│ ├── images/
│ └── labels/
├── valid/
│ ├── images/
│ └── labels/
├── test/
│ ├── images/
│ └── labels/
└── README.txt # 数据说明其中：
- 该存储目录包含了训练数据集中的图像及对应标注信息。 *
【* 该存储目录也包含了用于验证的数据及其标注信息。 *
【* 该存储目录同样包含了用于测试的数据及其标注信息。 *
数据预处理
因为数据量庞大而推荐使用脚本自动化处理来完成分类任务。如随机选取部分样本用于验证和测试。
import os
import random
import shutil
def split_dataset(src_img_dir, src_ann_dir, dest_train_img_dir, dest_train_ann_dir, dest_valid_img_dir, dest_valid_ann_dir, dest_test_img_dir, dest_test_ann_dir, train_ratio=0.7, valid_ratio=0.2, test_ratio=0.1):
"""
将数据集分为训练集、验证集和测试集。
"""
filenames = os.listdir(src_img_dir)
random.shuffle(filenames)
total = len(filenames)
train_end = int(total * train_ratio)
valid_end = int(total * (train_ratio + valid_ratio))
train_files = filenames[:train_end]
valid_files = filenames[train_end:valid_end]
test_files = filenames[valid_end:]
for files, dest_img_dir, dest_ann_dir in [
(train_files, dest_train_img_dir, dest_train_ann_dir),
(valid_files, dest_valid_img_dir, dest_valid_ann_dir),
(test_files, dest_test_img_dir, dest_test_ann_dir)
]:
os.makedirs(dest_img_dir, exist_ok=True)
os.makedirs(dest_ann_dir, exist_ok=True)
for filename in files:
shutil.copy(f'{src_img_dir}/{filename}', dest_img_dir)
shutil.copy(f'{src_ann_dir}/{filename.replace(".jpg", ".txt")}', dest_ann_dir)
# 示例调用
split_dataset(
src_img_dir='path/to/pv_hotspot_dataset/images/',
src_ann_dir='path/to/pv_hotspot_dataset/labels/',
dest_train_img_dir='path/to/pv_hotspot_dataset/train/images/',
dest_train_ann_dir='path/to/pv_hotspot_dataset/train/labels/',
dest_valid_img_dir='path/to/pv_hotspot_dataset/valid/images/',
dest_valid_ann_dir='path/to/pv_hotspot_dataset/valid/labels/',
dest_test_img_dir='path/to/pv_hotspot_dataset/test/images/',
dest_test_ann_dir='path/to/pv_hotspot_dataset/test/labels/'
)数据加载器
我们需要定义一个数据加载器来读取图像和标签文件：
import torch
from torch.utils.data import Dataset, DataLoader
import os
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
class PVHotspotDataset(Dataset):
def __init__(self, images_dir, labels_dir, transform=None):
self.images_dir = images_dir
self.labels_dir = labels_dir
self.transform = transform
self.image_filenames = [os.path.join(images_dir, f) for f in os.listdir(images_dir)]
self.label_filenames = [os.path.join(labels_dir, f.replace('.jpg', '.txt')) for f in os.listdir(images_dir)]
def __len__(self):
return len(self.image_filenames)
def __getitem__(self, idx):
if torch.is_tensor(idx):
idx = idx.tolist()
img_name = self.image_filenames[idx]
label_name = self.label_filenames[idx]
image = Image.open(img_name).convert("RGB")
with open(label_name, 'r') as f:
labels = []
for line in f.readlines():
label, cx, cy, w, h = line.strip().split()
labels.append([int(label), float(cx), float(cy), float(w), float(h)])
labels = np.array(labels)
if self.transform:
image = self.transform(image)
return image, labels
# 示例转换
transform = transforms.Compose([
transforms.Resize((448, 448)),
transforms.ToTensor(),
])
# 创建数据集实例
train_dataset = PVHotspotDataset(
images_dir='path/to/pv_hotspot_dataset/train/images/',
labels_dir='path/to/pv_hotspot_dataset/train/labels/',
transform=transform
)
valid_dataset = PVHotspotDataset(
images_dir='path/to/pv_hotspot_dataset/valid/images/',
labels_dir='path/to/pv_hotspot_dataset/valid/labels/',
transform=transform
)
test_dataset = PVHotspotDataset(
images_dir='path/to/pv_hotspot_dataset/test/images/',
labels_dir='path/to/pv_hotspot_dataset/test/labels/',
transform=transform
)
# 创建数据加载器
train_dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=4)
valid_dataloader = DataLoader(valid_dataset, batch_size=4, shuffle=False, num_workers=4)
test_dataloader = DataLoader(test_dataset, batch_size=4, shuffle=False, num_workers=4)模型定义
我们可以使用YOLOv5作为基础模型来进行目标检测：
# 假设YOLOv5已经安装并且可用
from yolov5.models.common import DetectMultiBackend
from yolov5.utils.torch_utils import select_device
device = select_device('')
model = DetectMultiBackend('yolov5s.pt', device=device, dnn=False) # load FP32 model
model.nc = 3 # number of classes
model.names = ['hotspot', 'hotpoint', 'occlusion'] # class names模型训练
接下来定义训练循环：
from yolov5.train import train
# 调用训练函数
train(
data='path/to/data.yaml',
imgsz=448,
batch_size=4,
epochs=100,
workers=4,
device=device,
project='runs/train',
name='pv_hotspot_detection',
exist_ok=True
)模型评估
在训练完成后，我们需要评估模型的性能：
from yolov5.val import val
# 调用评估函数
val(
data='path/to/data.yaml',
weights='path/to/weights/best.pt',
imgsz=448,
batch_size=4,
workers=4,
device=device,
project='runs/val',
name='pv_hotspot_detection',
exist_ok=True
)模型预测
下面是一个使用训练好的模型进行预测的Python脚本示例：
from yolov5.detect import detect
# 调用预测函数
detect(
weights='path/to/weights/best.pt',
source='path/to/your/image.jpg',
imgsz=448,
conf_thres=0.5,
iou_thres=0.45,
device=device,
project='runs/detect',
name='pv_hotspot_detection',
exist_ok=True
)数据配置文件
创建一个data.yaml文件来指定数据集路径、类别数量等信息：
path: path/to/pv_hotspot_dataset/
train: train/images/
val: valid/images/
test: test/images/ # test set
nc: 3 # number of classes
names: ['hotspot', 'hotpoint', 'occlusion'] # class names完整的训练和预测流程
克隆YOLOv5仓库 ：
1. git clone https://github.com/ultralytics/yolov5.git
2. cd yolov5
3. pip install -r requirements.txt数据预处理 ：
split_dataset.py运行训练脚本 ：
python train.py --img 448 --batch 4 --epochs 100 --data path/to/data.yaml --cfg models/yolov5s.yaml --weights '' --name pv_hotspot_detection --cache运行预测脚本 ：
python detect.py --weights path/to/weights/best.pt --source path/to/your/image.jpg --img 448 --conf 0.5 --iou 0.45注意事项
- 数据集质量 ：保证数据集的质量要素及其相关属性。
- 模型选择 ：建议采用更为先进的模型架构（如YOLOv5l或YOLOv5x）以提升性能水平。
- 超参数调整 ：根据具体情况优化超参数设置，如批量大小（
--batch）、图像尺寸（--img）等。 - 监控性能 ：通过监控损失函数值与mAP指标来评估模型表现状态。
按照以下所述的步骤,你可能能够利用YOLOv5来建立一个光伏板热斑缺陷检测的数据集,并利用训练好的模型来进行预测
全部评论 (0)
还没有任何评论哟~