如何使用深度学习方法进行目标检测，训练自己的数据集——太阳能电池板缺陷数据集
发布时间
阅读量：
阅读量
基于深度学习的目标检测技术的具体流程如下。该数据集涵盖四个缺陷类别，并经过划分成训练集、验证集和测试集。
数据集描述
该数据集包含以下四种类别：
- 鸟类排泄物（bird_feces） ：指代太阳能电池板表面所出现的鸟类排泄物污渍。
- 干净状态（clean_state） ：指代没有缺陷且保持良好功能状态的太阳能电池板表面。
- 裂痕（crack） ：指代太阳能电池板表面因机械或物理因素造成的裂纹或破损痕迹。
- 尘埃（dust_particle） ：指代漂浮于太阳能电池板表面的微小尘埃颗粒物。
数据集的分布如下：
- 训练集 ：3482张图片
- 验证集 ：1087张图片
- 测试集 ：584张图片
该标注文件采用YOLO兼容的.txt格式（亦可选XML或JSON格式，具体由提供数据决定）
数据集组织
假设你的数据集目录结构如下：
solar_panel_defects/
├── train/
│ ├── images/
│ └── labels/
├── valid/
│ ├── images/
│ └── labels/
├── test/
│ ├── images/
│ └── labels/
└── README.txt # 数据说明
AI生成项目
其中：
train/存储了训练数据的图片集合及其对应的标注信息。valid/存储了验证数据的图片数据与标注信息文件。test/存储了测试数据的图片集合及其对应的标注信息。
数据预处理
如果标注文件不符合YOLO适用的.txt格式，则必须将其他格式（如XML或JSON）转码为YOLO格式，并附带一个简单的脚本示例说明如何完成转码过程。假设标注文件被视为XML格式处理。
import xml.etree.ElementTree as ET
import os
def convert_xml_to_yolo(xml_file, img_size=(448, 448)):
tree = ET.parse(xml_file)
root = tree.getroot()
size = root.find('size')
width = int(size.find('width').text)
height = int(size.find('height').text)
labels = []
for obj in root.iter('object'):
cls = obj.find('name').text
if cls not in ['bird_drop', 'clean', 'cracked', 'dust']:
continue
cls_id = ['bird_drop', 'clean', 'cracked', 'dust'].index(cls)
xmlbox = obj.find('bndbox')
b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text),
float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
bb = convert_coordinates(b, width, height, img_size)
labels.append([cls_id, *bb])
return labels
def convert_coordinates(box, orig_w, orig_h, target_size):
dw, dh = target_size
x_center = ((box[1] + box[0]) / 2) / orig_w
y_center = ((box[3] + box[2]) / 2) / orig_h
w = (box[1] - box[0]) / orig_w
h = (box[3] - box[2]) / orig_h
x_center *= dw
y_center *= dh
w *= dw
h *= dh
return x_center, y_center, w, h
def convert_annotations(src_xml_dir, src_img_dir, dest_label_dir):
os.makedirs(dest_label_dir, exist_ok=True)
for xml_file in os.listdir(src_xml_dir):
if xml_file.endswith('.xml'):
labels = convert_xml_to_yolo(os.path.join(src_xml_dir, xml_file))
img_file = xml_file.replace('.xml', '.jpg')
img_path = os.path.join(src_img_dir, img_file)
img = Image.open(img_path)
img_size = img.size
with open(os.path.join(dest_label_dir, img_file.replace('.jpg', '.txt')), 'w') as f:
for label in labels:
f.write('%d %.6f %.6f %.6f %.6f\n' % tuple(label))
# 示例调用
convert_annotations(
src_xml_dir='path/to/solar_panel_defects/train/labels/',
src_img_dir='path/to/solar_panel_defects/train/images/',
dest_label_dir='path/to/solar_panel_defects/train/labels_converted/'
)
AI生成项目
数据加载器
我们需要定义一个数据加载器来读取图像和标签文件：
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 SolarPanelDefectDataset(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 = SolarPanelDefectDataset(
images_dir='path/to/solar_panel_defects/train/images/',
labels_dir='path/to/solar_panel_defects/train/labels_converted/',
transform=transform
)
valid_dataset = SolarPanelDefectDataset(
images_dir='path/to/solar_panel_defects/valid/images/',
labels_dir='path/to/solar_panel_defects/valid/labels_converted/',
transform=transform
)
test_dataset = SolarPanelDefectDataset(
images_dir='path/to/solar_panel_defects/test/images/',
labels_dir='path/to/solar_panel_defects/test/labels_converted/',
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)
AI生成项目
模型定义
我们可以使用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 = 4 # number of classes
model.names = ['bird_drop', 'clean', 'cracked', 'dust'] # class names
AI生成项目模型训练
接下来定义训练循环：
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='solar_panel_defects',
exist_ok=True
)
AI生成项目
模型评估
在训练完成后，我们需要评估模型的性能：
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='solar_panel_defects',
exist_ok=True
)
AI生成项目
模型预测
下面是一个使用训练好的模型进行预测的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='solar_panel_defects',
exist_ok=True
)
AI生成项目
数据配置文件
创建一个data.yaml文件来指定数据集路径、类别数量等信息：
path: path/to/solar_panel_defects/
train: train/images/
val: valid/images/
test: test/images/ # test set
nc: 4 # number of classes
names: ['bird_drop', 'clean', 'cracked', 'dust'] # class names
AI生成项目完整的训练和预测流程
克隆YOLOv5仓库 ：
1. git clone https://github.com/ultralytics/yolov5.git
2. cd yolov5
3. pip install -r requirements.txt
AI生成项目数据预处理 ：
python train.py --img 448 --batch 4 --epochs 100 --data path/to/data.yaml --cfg models/yolov5s.yaml --weights '' --name solar_panel_defects --cache
AI生成项目预测脚本 ：
python detect.py --weights path/to/weights/best.pt --source path/to/your/image.jpg --img 448 --conf 0.5 --iou 0.45
AI生成项目注意事项
- 数据集质量：确保数据集的质量标准得到满足（包括清晰度、标注准确性等）。
- 模型选择：优化配置：建议采用大模型版本以提升性能表现。
- 超参数调整：根据实际需求进行参数优化配置（如批量大小
--batch、图像尺寸--img等）。 - 监控性能：实时跟踪关键指标：在训练过程中实时跟踪损失函数值和平均精度（mAP）等关键指标。
按照以下流程建议你利用YOLOv5训练一组太阳能电池板缺陷检测数据集，并同时利用训练好的模型来进行预测。
全部评论 (0)
还没有任何评论哟~