企业官网建设流程全解析

使用YOLOv5构建一个用于配网绝缘导线烧伤检测的深度学习模型——

_配网绝缘导线烧伤检测图像数据集，总共1500多张图像，标注为voc (xml)格式
构建一个用于配网绝缘导线烧伤检测的深度学习模型。我们将使用YOLOv5（You Only Look Once version 5）进行目标检测。以下是详细的步骤：

1. 环境准备：

   • 安装必要的库。
   • 下载并组织数据集。

2. 数据预处理：

   • 将数据集分为训练集、验证集和测试集。
   • 格式化标签文件以便于YOLOv5使用。

3. 模型定义与训练：

   • 使用YOLOv5进行目标检测。

4. 评估与可视化：

   • 评估模型性能。
   • 可视化结果。
     

环境准备

安装YOLOv5和一些其他的依赖项。您可以使用以下命令来设置环境：

pipinstalltorch torchvision torchaudio pyyaml opencv-python-headless seaborn pandas matplotlib scikit-learn lxmlgitclone https://github.com/ultralytics/yolov5cdyolov5 pipinstall-rrequirements.txt

数据预处理

假设您的数据集已经下载并存储在datasets/insulation_wire_burn目录中，其中包含images和annotations子目录。annotations子目录中的每个.xml文件对应一个图像文件，格式为VOC XML格式。

接下来，我们将编写一个脚本来检查数据集的有效性，并将VOC XML标签转换为YOLOv5所需的格式。

[<title="Data Preparation Script for Insulation Wire Burn Detection">]importosimportxml.etree.ElementTreeasETfrompathlibimportPathimportshutilimportrandomimportnumpyasnp# Define pathsbase_path=Path('datasets/insulation_wire_burn')images_dir=base_path/'images'annotations_dir=base_path/'annotations'# Create train, val, test directoriestrain_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'os.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)# Function to convert VOC XML to YOLO formatdefconvert_voc_to_yolo(xml_file,image_width,image_height):tree=ET.parse(xml_file)root=tree.getroot()yolo_lines=[]forobjinroot.findall('object'):label=obj.find('name').text bbox=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)x_center=(xmin+xmax)/2.0/image_width y_center=(ymin+ymax)/2.0/image_height width=(xmax-xmin)/image_width height=(ymax-ymin)/image_height class_id=0# Assuming only one class "burn"yolo_line=f"{class_id}{x_center}{y_center}{width}{height}\n"yolo_lines.append(yolo_line)return''.join(yolo_lines)# Split dataset into train, val, testimage_files=list(images_dir.glob('*.jpg'))# Adjust extension if necessaryrandom.shuffle(image_files)split_index_1=int(0.8*len(image_files))split_index_2=int(0.9*len(image_files))train_files=image_files[:split_index_1]val_files=image_files[split_index_1:split_index_2]test_files=image_files[split_index_2:]defcopy_and_convert(files,src_img_dir,src_ann_dir,dst_img_dir,dst_ann_dir):forimg_fileinfiles:ann_file=src_ann_dir/(img_file.stem+'.xml')ifnotann_file.exists():print(f"Missing annotation file for{img_file.name}")continue# Copy imageshutil.copy(img_file,dst_img_dir)# Convert and save labeltree=ET.parse(ann_file)root=tree.getroot()size=root.find('size')image_width=int(size.find('width').text)image_height=int(size.find('height').text)yolo_content=convert_voc_to_yolo(ann_file,image_width,image_height)withopen(dst_ann_dir/(img_file.stem+'.txt'),'w')asf:f.write(yolo_content)copy_and_convert(train_files,images_dir,annotations_dir,train_images_dir,train_labels_dir)copy_and_convert(val_files,images_dir,annotations_dir,val_images_dir,val_labels_dir)copy_and_convert(test_files,images_dir,annotations_dir,test_images_dir,test_labels_dir)# Create YOLOv5 configuration fileconfig={"train":str(train_images_dir),"val":str(val_images_dir),"nc":1,"names":["burn"]}withopen(base_path/'insulation_wire_burn.yaml','w')asf:importyaml yaml.dump(config,f)print("Dataset prepared and YOLOv5 config created.")

模型定义与训练

我们将使用YOLOv5进行目标检测。以下是训练脚本train_detection.py:

[<title="Training Script for Insulation Wire Burn Detection using YOLOv5">]importsubprocess# Run YOLOv5 training commandcommand=['python','train.py','--img','640',# Image size'--batch','16',# Batch size'--epochs','50',# Number of epochs'--data','../datasets/insulation_wire_burn/insulation_wire_burn.yaml',# Path to dataset YAML file'--cfg','./models/yolov5s.yaml',# Model configuration'--weights','yolov5s.pt',# Pretrained weights'--project','../runs/train',# Project directory'--name','wire_burn_detection'# Experiment name]subprocess.run(command)

评估与可视化

使用YOLOv5自带的评估脚本来评估目标检测模型。

[<title="Evaluation Script for Insulation Wire Burn Detection using YOLOv5">]importsubprocess# Run YOLOv5 evaluation commandcommand=['python','val.py','--data','../datasets/insulation_wire_burn/insulation_wire_burn.yaml',# Path to dataset YAML file'--weights','../runs/train/wire_burn_detection/weights/best.pt',# Path to best model weights'--img','640',# Image size'--conf','0.5',# Confidence threshold'--iou-thres','0.45',# IoU threshold'--task','val',# Task type'--save-json',# Save JSON results'--project','../runs/val',# Project directory'--name','wire_burn_detection_val'# Experiment name]subprocess.run(command)

用户界面

我们将使用 PyQt5 创建一个简单的 GUI 来加载和运行模型进行实时预测。以下是用户界面脚本ui.py:

[<title="PyQt5 Main Window for Insulation Wire Burn Detection">]importsysimportcv2importnumpyasnpfromPyQt5.QtWidgetsimportQApplication,QMainWindow,QLabel,QPushButton,QVBoxLayout,QWidget,QFileDialogfromPyQt5.QtGuiimportQImage,QPixmapfromPyQt5.QtCoreimportQt,QTimerimportultralytics.yolo.engine.model# Load modeldetection_model=ultralytics.yolo.engine.model.Model(cfg='./models/yolov5s.yaml').load('../runs/train/wire_burn_detection/weights/best.pt')classMainWindow(QMainWindow):def__init__(self):super().__init__()self.setWindowTitle("Insulation Wire Burn Detection System")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("Load Image",self)self.load_image_button.clicked.connect(self.load_image)self.layout.addWidget(self.load_image_button)self.start_prediction_button=QPushButton("Start Prediction",self)self.start_prediction_button.clicked.connect(self.start_prediction)self.layout.addWidget(self.start_prediction_button)self.stop_prediction_button=QPushButton("Stop Prediction",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,"QFileDialog.getOpenFileName()","","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)# Detectiondetection_results=detection_model(image_rgb,size=640)[0].boxes.data.cpu().numpy()forresultindetection_results:x1,y1,x2,y2,conf,cls=result x1,y1,x2,y2=int(x1),int(y1),int(x2),int(y2)cls=int(cls)# Draw bounding boxcv2.rectangle(image_rgb,(x1,y1),(x2,y2),(0,255,0),2)# Put textfont=cv2.FONT_HERSHEY_SIMPLEX cv2.putText(image_rgb,f'Burn ({conf:.2f})',(x1,y1-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_())

请确保将路径替换为您实际的路径。

使用说明

1. 配置路径：

   • 确保datasets/insulation_wire_burn目录结构正确，并且包含images和annotations子目录。
   • 确保runs/train/wire_burn_detection/weights/best.pt是训练好的 YOLOv5 模型权重路径。

2. 运行脚本：

   • 在终端中运行data_preparation.py脚本来检查数据集的有效性并创建 YOLOv5 配置文件。
   • 在终端中运行train_detection.py脚本来训练目标检测模型。
   • 在终端中运行evaluate_detection.py来评估目标检测模型性能。
   • 在终端中运行ui.py来启动 GUI 应用程序。

3. 注意事项：

   • 确保所有必要的工具箱已安装，特别是 TensorFlow 和 PyQt5。
   • 根据需要调整参数，如epochs和batch_size。

示例

假设您的数据文件夹结构如下：

datasets/ └── insulation_wire_burn/ ├── images/ └── annotations/

并且每个数据集中包含相应的图像和标签文件。运行ui.py后，您可以点击按钮来加载图像并进行绝缘导线烧伤检测。