企业官网建设流程全解析

使用深度学习框架目标检测算法Yolo训练_河道河流水污染水质污染程度检测数据集检测水域中 清洁程度 浑浊 正常污染等

河道河流水污染水质检测数据集 2400张
1

1

处理数据集，包括安装依赖、准备数据集、训练模型、评估模型，并最终构建一个GUI应用程序来展示模型的预测结果。使用YOLOv8作为示例模型。

数据集概述

• 类别标签：clean, turbid, pollute, normal
• 图片数量：总共2487张图片
• 标注框数量：总共2489个标注框

实现步骤

一、安装依赖

确保已安装Python（推荐版本3.8+），并安装以下依赖：

pipinstalltorch torchvision ultralytics opencv-python matplotlib lxml pyqt5

• torch和torchvision：深度学习框架。
• ultralytics：YOLOv8的官方库。
• opencv-python：用于图像处理。
• matplotlib：用于可视化。
• lxml：用于解析XML文件（如果需要处理VOC格式）。
• pyqt5：用于构建GUI应用程序。

二、准备数据集

1. 数据集目录结构

将水污染水质检测数据集按照以下格式组织：

water-quality-dataset/ ├── images/ │ ├── train/ # 训练集图片 │ ├── val/ # 验证集图片 │ └── test/ # 测试集图片 ├── labels/ │ ├── train/ # 训练集标注文件（txt格式） │ ├── val/ # 验证集标注文件（txt格式） │ └── test/ # 测试集标注文件（txt格式） └── data.yaml # 数据集配置文件

如果你的数据集是以VOC格式标注的（通常包含.xml文件），你需要将其转换为YOLO所需的.txt格式。

2. 创建data.yaml

在water-quality-dataset/目录下创建data.yaml文件，内容如下：

train:./images/trainval:./images/valtest:./images/testnc:4# 类别数量names:['clean','turbid','pollute','normal']# 类别名称

3. 转换标注格式（如果需要）

如果你的数据集是VOC格式，可以编写脚本将VOC格式转换为YOLO所需的.txt格式。

importosimportxml.etree.ElementTreeasETdefconvert_voc_to_yolo(xml_file,label_dir):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_file=os.path.join(label_dir,os.path.splitext(os.path.basename(xml_file))[0]+'.txt')withopen(label_file,'w')asf:forobjinroot.findall('object'):class_name=obj.find('name').textifclass_name=='clean':class_id=0elifclass_name=='turbid':class_id=1elifclass_name=='pollute':class_id=2elifclass_name=='normal':class_id=3else:continuebbox=obj.find('bndbox')xmin=int(bbox.find('xmin').text)ymin=int(bbox.find('ymin').text)xmax=int(bbox.find('xmax').text)ymax=int(bbox.find('ymax').text)x_center=(xmin+xmax)/2.0/image_width y_center=(ymin+ymax)/2.0/image_height width=(xmax-xmin)/float(image_width)height=(ymax-ymin)/float(image_height)f.write(f"{class_id}{x_center}{y_center}{width}{height}\n")# 示例调用xml_files=['path/to/voc/file1.xml','path/to/voc/file2.xml']label_dir='path/to/labels/'forxml_fileinxml_files:convert_voc_to_yolo(xml_file,label_dir)

请确保所有.xml文件都被转换为.txt文件，并放置在正确的目录中。

三、训练模型

1. 开始训练

运行以下代码开始训练模型：

fromultralyticsimportYOLO# 加载预训练模型model=YOLO('yolov8n.pt')# 使用 YOLOv8 Nano 模型# 训练模型results=model.train(data='path/to/water-quality-dataset/data.yaml',# 数据集配置文件路径epochs=50,# 训练轮数imgsz=640,# 输入图片尺寸batch=16,# 批量大小device='cuda'iftorch.cuda.is_available()else'cpu'# 使用 GPU 或 CPU)# 保存训练结果model.export(format='onnx')# 导出为 ONNX 格式（可选）

四、评估模型

1. 测试模型性能

运行以下代码进行模型评估：

fromultralyticsimportYOLO# 加载训练好的模型model=YOLO('runs/detect/train/weights/best.pt')# 测试模型metrics=model.val(data='path/to/water-quality-dataset/data.yaml')print(f"mAP50:{metrics.box.map50}")# 输出 mAP50 指标

五、构建 GUI 应用程序

1. 创建 PyQt5 GUI

以下是一个简单的 GUI 示例，用于加载模型并检测图片中的水质情况。

importsysimportcv2fromPyQt5.QtWidgetsimportQApplication,QMainWindow,QLabel,QPushButton,QVBoxLayout,QWidget,QFileDialogfromPyQt5.QtGuiimportQPixmap,QImagefromultralyticsimportYOLOclassWaterQualityDetectorApp(QMainWindow):def__init__(self):super().__init__()self.initUI()self.model=YOLO('runs/detect/train/weights/best.pt')# 加载训练好的模型definitUI(self):self.setWindowTitle("水质检测")self.setGeometry(100,100,800,600)self.label=QLabel(self)self.label.setText("选择图片以检测水质")self.label.setAlignment(Qt.AlignCenter)self.button=QPushButton("选择图片",self)self.button.clicked.connect(self.load_image)layout=QVBoxLayout()layout.addWidget(self.label)layout.addWidget(self.button)container=QWidget()container.setLayout(layout)self.setCentralWidget(container)defload_image(self):options=QFileDialog.Options()file_path,_=QFileDialog.getOpenFileName(self,"选择图片","","Images (*.png *.jpg *.jpeg)",options=options)iffile_path:image=cv2.imread(file_path)results=self.model(image)# 进行目标检测annotated_image=results[0].plot()# 绘制检测结果# 将 OpenCV 图像转换为 PyQt 图像height,width,channel=annotated_image.shape bytes_per_line=3*width q_img=QImage(annotated_image.data,width,height,bytes_per_line,QImage.Format_RGB888).rgbSwapped()self.label.setPixmap(QPixmap.fromImage(q_img))if__name__=="__main__":app=QApplication(sys.argv)window=WaterQualityDetectorApp()window.show()sys.exit(app.exec_())

六、运行应用程序

1. 将上述代码保存为app.py。
2. 确保已安装所有依赖库。
3. 在终端中运行以下命令启动 GUI 应用程序：

python app.py

通过以上步骤，同学ni 就可以 完成水污染水质检测的数据集准备、模型训练、评估以及构建一个简单的GUI应用程序来展示模型的检测效果

以上文字及代码仅供参考。