使用深度学习框架目标检测yolov8训练使用——
巡检图像变电站可见光缺陷数据集
360-变电站可见光缺陷数据集
数据集包含8376张巡检图像,带xml标签,共包含17类巡检标签 缺陷分类见下图!
缺陷分类:
表计破损,绝缘子破损,渗漏油,呼吸器异常,异物
8376张图片
各类标签的数量分别为:
bjdsyc: 789
bj_ wkps: 523
yw_ nc: 883
xmbhyc: 383
kgg_ ybh: 362
gbps: 654
yw_ gkxfw: 729
hxq_ gjbs: 1174
bj bpmh: 869
jyz_pl: 410
bj. bpps: 723
slydmyw: 833
wcaqm: 567
wcgz: 815
ywzt yfyc: 331
hxq_ gjitps: 106
xy: 607
标签类别如下:
[‘bjdsyc’, ‘bj wkps’, ‘yw_ nc’, ‘xmbhyc’, ‘kgg. ybh’, ‘gbps’, ‘yw. gkxfw’, ‘hxq. gjbs’, ‘bj bpmh’, "iyz. pl’, ‘bj. bpps’,
‘sly_ dmyw’, ‘wcaqm’, ‘wcgz’, ‘ywzt yfyc’, "hxq. gitps’, ‘xy’]
以下所有内容仅供参考!
使用YOLOv8来进行巡检图像变电站可见光缺陷检测任务。以下是详细的步骤:
环境准备:
- 安装必要的库。
- 下载并组织数据集。
数据预处理:
- 将XML格式的标注文件转换为YOLOv8所需的TXT格式。
- 确认数据集的划分和标签格式。
模型定义与训练:
- 使用YOLOv8进行目标检测。
评估与可视化:
- 评估模型性能。
- 可视化结果。
环境准备
首先,我们需要安装必要的库。您可以使用以下命令来设置环境:
pipinstallultralytics lxml opencv-python-headless pandas scikit-learn数据预处理
我们需要将XML格式的标注文件转换为YOLOv8所需的TXT格式。以下是数据转换脚本convert_xml_to_yolo.py:
[<title="Convert XML Annotations to YOLO TXT Format">]importosimportxml.etree.ElementTreeasETfrompathlibimportPath# Define pathsbase_path=Path('datasets/substation_inspection')annotations_dir=base_path/'Annotations'images_dir=base_path/'JPEGImages'output_labels_dir=base_path/'labels'# Create output directory if it doesn't existos.makedirs(output_labels_dir,exist_ok=True)# Class names and their corresponding IDsclass_names={'bjdsyc':0,'bj_wkps':1,'yw_nc':2,'xmbhyc':3,'kgg_ybh':4,'gbps':5,'yw_gkxfw':6,'hxq_gjbs':7,'bj_bpmh':8,'jyz_pl':9,'bj_bpps':10,'sly_dmyw':11,'wcaqm':12,'wcgz':13,'ywzt_yfyc':14,'hxq_gitps':15,'xy':16}defconvert_annotation(xml_file):tree=ET.parse(xml_file)root=tree.getroot()image_width=int(root.find('size/width').text)image_height=int(root.find('size/height').text)label_lines=[]forobjinroot.findall('object'):class_name=obj.find('name').text.replace('.','_')# Replace dots with underscores to match keysbbox=obj.find('bndbox')xmin=float(bbox.find('xmin').text)ymin=float(bbox.find('ymin').text)xmax=float(bbox.find('xmax').text)ymax=float(bbox.find('ymax').text)# Convert bounding box to YOLO format (center_x, center_y, width, height)center_x=(xmin+xmax)/2.0/image_width center_y=(ymin+ymax)/2.0/image_height width=(xmax-xmin)/image_width height=(ymax-ymin)/image_height class_id=class_names[class_name]label_line=f"{class_id}{center_x}{center_y}{width}{height}\n"label_lines.append(label_line)returnlabel_linesforannotation_fileinannotations_dir.glob('*.xml'):image_name=annotation_file.stem+'.jpg'# Assuming images are in JPEG formatlabel_file=output_labels_dir/(annotation_file.stem+'.txt')label_lines=convert_annotation(annotation_file)withopen(label_file,'w')asf:f.writelines(label_lines)print("Conversion completed.")创建YAML配置文件
确保您有一个正确的YAML配置文件substation_inspection.yaml来描述数据集。以下是示例配置文件:
[<title="YOLOv8 Configuration File for Substation Inspection Object Detection">]train:../datasets/substation_inspection/train/imagesval:../datasets/substation_inspection/val/imagesnc:17names:['bjdsyc','bj_wkps','yw_nc','xmbhyc','kgg_ybh','gbps','yw_gkxfw','hxq_gjbs','bj_bpmh','jyz_pl','bj_bpps','sly_dmyw','wcaqm','wcgz','ywzt_yfyc','hxq_gitps','xy']将上述内容保存为datasets/substation_inspection/substation_inspection.yaml。
划分数据集
假设您的数据集已经分为训练集、验证集和测试集,并且每个集合都有对应的图像和标签文件。如果还没有划分数据集,可以使用以下脚本来随机划分数据集:
[<title="Split Dataset into Train, Val, Test Sets">]importosimportrandomfromsklearn.model_selectionimporttrain_test_splitimportpandasaspdfrompathlibimportPath# Define pathsbase_path=Path('datasets/substation_inspection')images_dir=base_path/'JPEGImages'annotations_dir=base_path/'labels'train_images_dir=base_path/'train/images'train_labels_dir=base_path/'train/labels'val_images_dir=base_path/'val/images'val_labels_dir=base_path/'val/labels'test_images_dir=base_path/'test/images'test_labels_dir=base_path/'test/labels'# Create directories if they don't existos.makedirs(train_images_dir,exist_ok=True)os.makedirs(train_labels_dir,exist_ok=True)os.makedirs(val_images_dir,exist_ok=True)os.makedirs(val_labels_dir,exist_ok=True)os.makedirs(test_images_dir,exist_ok=True)os.makedirs(test_labels_dir,exist_ok=True)# List all image filesimage_files=list(images_dir.glob('*.jpg'))# Adjust extension if necessary# Shuffle the image filesrandom.shuffle(image_files)# Split ratiostrain_ratio=0.7val_ratio=0.15test_ratio=0.15# Calculate split indicesnum_images=len(image_files)train_split=int(num_images*train_ratio)val_split=int(num_images*(train_ratio+val_ratio))# Split images and labelstrain_images=image_files[:train_split]val_images=image_files[train_split:val_split]test_images=image_files[val_split:]defcopy_files(source_images,dest_images_dir,dest_labels_dir):forimg_fileinsource_images:label_file=annotations_dir/(img_file.stem+'.txt')iflabel_file.exists():os.symlink(img_file,dest_images_dir/img_file.name)os.symlink(label_file,dest_labels_dir/label_file.name)copy_files(train_images,train_images_dir,train_labels_dir)copy_files(val_images,val_images_dir,val_labels_dir)copy_files(test_images,test_images_dir,test_labels_dir)print("Dataset splitting completed.")模型定义与训练
我们将使用YOLOv8进行目标检测。以下是训练脚本train_detection.py:
[<title="Training Script for Substation Inspection Object Detection using YOLOv8">]fromultralyticsimportYOLO# Load a modelmodel=YOLO('yolov8n.pt')# load a pretrained model (recommended for training)# Train the modelresults=model.train(data='../datasets/substation_inspection/substation_inspection.yaml',epochs=50,imgsz=128,batch=16,project='../runs/train',name='substation_inspection_detection')# Evaluate the modelmetrics=model.val()results=model.export(format='onnx')# export the trained model to ONNX format评估与可视化
使用YOLOv8自带的评估脚本来评估目标检测模型。
[<title="Evaluation Script for Substation Inspection Object Detection using YOLOv8">]fromultralyticsimportYOLO# Load the best modelbest_model=YOLO('../runs/train/substation_inspection_detection/weights/best.pt')# Evaluate the model on the validation datasetmetrics=best_model.val(data='../datasets/substation_inspection/substation_inspection.yaml',conf=0.5,iou=0.45)print(metrics)用户界面
我们将使用 PyQt5 创建一个简单的 GUI 来加载和运行模型进行实时预测。以下是用户界面脚本ui.py:
[<title="PyQt5 Main Window for Substation Inspection Object Detection">]importsysimportcv2importnumpyasnpfromPyQt5.QtWidgetsimportQApplication,QMainWindow,QLabel,QPushButton,QVBoxLayout,QWidget,QFileDialogfromPyQt5.QtGuiimportQImage,QPixmapfromPyQt5.QtCoreimportQt,QTimerfromultralyticsimportYOLO# Load modeldetection_model=YOLO('../runs/train/substation_inspection_detection/weights/best.pt')classMainWindow(QMainWindow):def__init__(self):super().__init__()self.setWindowTitle("变电站巡检图像缺陷检测系统")self.setGeometry(100,100,800,600)self.initUI()definitUI(self):self.central_widget=QWidget()self.setCentralWidget(self.central_widget)self.layout=QVBoxLayout()self.image_label=QLabel(self)self.image_label.setAlignment(Qt.AlignCenter)self.layout.addWidget(self.image_label)self.load_image_button=QPushButton("加载图像",self)self.load_image_button.clicked.connect(self.load_image)self.layout.addWidget(self.load_image_button)self.start_prediction_button=QPushButton("开始预测",self)self.start_prediction_button.clicked.connect(self.start_prediction)self.layout.addWidget(self.start_prediction_button)self.stop_prediction_button=QPushButton("停止预测",self)self.stop_prediction_button.clicked.connect(self.stop_prediction)self.layout.addWidget(self.stop_prediction_button)self.central_widget.setLayout(self.layout)self.image_path=Noneself.timer=QTimer()self.timer.timeout.connect(self.update_frame)defload_image(self):options=QFileDialog.Options()file_name,_=QFileDialog.getOpenFileName(self,"选择图像文件","","Images (*.png *.jpg *.jpeg);;All Files (*)",options=options)iffile_name:self.image_path=file_name self.display_image(file_name)defdisplay_image(self,path):pixmap=QPixmap(path)scaled_pixmap=pixmap.scaled(self.image_label.width(),self.image_label.height(),Qt.KeepAspectRatio)self.image_label.setPixmap(scaled_pixmap)defstart_prediction(self):ifself.image_pathisnotNoneandnotself.timer.isActive():self.timer.start(30)# Update frame every 30 msdefstop_prediction(self):ifself.timer.isActive():self.timer.stop()self.image_label.clear()defupdate_frame(self):original_image=cv2.imread(self.image_path)image_rgb=cv2.cvtColor(original_image,cv2.COLOR_BGR2RGB)# Detectionresults=detection_model.predict(image_rgb,size=128,conf=0.5,iou=0.45)[0]forboxinresults.boxes.cpu().numpy():r=box.xyxy[0].astype(int)cls=int(box.cls[0])conf=box.conf[0]# Map class ID to nameclass_names=['表计破损','绝缘子破损','渗漏油','呼吸器异常','异物','其他','其他','其他','其他','其他','其他','其他','其他','其他','其他','其他','其他']class_name=class_names[cls]# Draw bounding boxcv2.rectangle(image_rgb,(r[0],r[1]),(r[2],r[3]),(0,255,0),2)# Put textfont=cv2.FONT_HERSHEY_SIMPLEX cv2.putText(image_rgb,f'{class_name}({conf:.2f})',(r[0],r[1]-10),font,0.9,(0,255,0),2)h,w,ch=image_rgb.shape bytes_per_line=ch*w qt_image=QImage(image_rgb.data,w,h,bytes_per_line,QImage.Format_RGB888)pixmap=QPixmap.fromImage(qt_image)scaled_pixmap=pixmap.scaled(self.image_label.width(),self.image_label.height(),Qt.KeepAspectRatio)self.image_label.setPixmap(scaled_pixmap)if__name__=="__main__":app=QApplication(sys.argv)window=MainWindow()window.show()sys.exit(app.exec_())请确保将路径替换为您实际的路径。
使用说明
配置路径:
- 确保
datasets/substation_inspection目录结构正确,并且包含train,val, 和test子目录。 - 确保
runs/train/substation_inspection_detection/weights/best.pt是训练好的 YOLOv8 模型权重路径。
- 确保
运行脚本:
- 在终端中运行
convert_xml_to_yolo.py脚本来将XML格式的标注文件转换为YOLOv8所需的TXT格式。 - 如果尚未划分数据集,在终端中运行
split_dataset.py脚本来划分数据集。 - 在终端中运行
check_dataset.py脚本来检查数据集的有效性。 - 在终端中运行
train_detection.py脚本来训练目标检测模型。 - 在终端中运行
evaluate_detection.py来评估目标检测模型性能。 - 在终端中运行
ui.py来启动 GUI 应用程序。
- 在终端中运行
注意事项:
- 确保所有必要的工具箱已安装,特别是
ultralytics和PyQt5。 - 根据需要调整参数,如
epochs和batch_size。
- 确保所有必要的工具箱已安装,特别是
示例
假设您的数据文件夹结构如下:
datasets/ └── substation_inspection/ ├── train/ │ ├── images/ │ └── labels/ ├── val/ │ ├── images/ │ └── labels/ └── test/ ├── images/ └── labels/并且每个数据集中包含相应的图像和标签文件。运行ui.py后,您可以点击按钮来加载图像并进行变电站巡检图像缺陷检测。
总结
通过上述步骤,我们可以构建一个全面的基于深度学习的变电站巡检图像缺陷检测系统,包括数据集准备、环境部署、数据预处理、模型定义、训练、评估、结果分析、可视化以及用户界面开发。以下是所有相关的代码文件: