基于YOLOv8n算法实现水体污染物识别

1 选题及其背景介绍

（1）题目

利用YOLOv8n算法在M1芯片上实现水体污染物识别

（2）背景

水体污染是全球面临的一个严重问题，影响人类健康、生态系统和经济发展。传统的水质监测方法往往依赖于手工取样和实验室分析，这些方法不仅耗时耗力，而且无法实时监测大范围的水体状态。随着人工智能技术的发展，尤其是计算机视觉领域的进步，利用图像识别技术进行水体污染物的检测成为了一个创新的解决方案。

YOLOv8n是一个最新的、轻量级的目标检测模型，它能够实时处理图像，识别出图像中的各种目标。这一技术在水体污染检测中的应用，能够基于从无人机、卫星或固定摄像头采集的水体图像，自动识别污染源和污染物种类，如油污、浮渣、有害化学物质等。

通过YOLOv8n实现的水体污染自动识别系统能够大幅提高监测的频率和反应速度，及时发现污染事件，缩短处理时间。而且，自动化的监测系统减少了传统监测所需的人工和物理资源，降低了监测成本。此外，利用高分辨率图像和先进的算法可以提高检测的准确性，同时能覆盖更广的区域。

在本研究中，我们解决了以下问题：

1. 数据获取与处理：本文介绍了如何获取获取高质量的水体图像数据，并进行适当的预处理以适用于模型训练。
2. 模型的精确度和鲁棒性：本文描述了如何训练一个既精确又能适应不同水体环境变化的污染检测模型。
3. 实时性能优化：本文考虑到实时监测的需求，并且研究需要优化模型的运行效率，所以需要确保其在各种操作环境下都能快速准确地进行检测。
4. 多种污染物的识别：本文描述了如何扩展模型的能力，使其不仅能检测到常见的污染类型，还能识别多种不同的污染物。

2 数据预处理

2.1 数据的获取

为了进行水体污染物的识别，本研究首先需要获取适合的数据集，本项目中使用的数据集是从Kaggle网站获取的，而其中“Pollution Detection Dataset”数据集特别适用于水体污染物识别。具体的数据集信息和下载链接为：xxx。

2.2 原始数据集的介绍

本数据集专为水体污染物的目标检测而设计，包含5130组图像。数据集的主要目的是通过机器学习模型来自动识别并定位图像中的污染物，以便进行进一步的分析和处理。

该数据集图像主要包括各类水体污染物，如塑料瓶、包装袋、其他难降解物质等，图像采集自不同的水域环境，包括河流、湖泊和海洋，以确保模型的泛化能力，如图b-1为具体的数据集内容。

图b-1 原始数据集

2.3 数据预处理

对于目标检测任务，如水体污染物检测，合适的预处理可以增强模型对图像中物体的识别能力，特别是在多变的环境条件下，以下为本研究数据集的预处理操作：

1.像素数据的自动定向：读取图像的EXIF数据，根据其中的定向信息调整图像的方向，以消除由于设备方向变化导致的图像方向不一致的问题，具体代码如下：

from PIL import Image, ExifTags

def fix_orientation(img_path):

    img = Image.open(img_path)

    try:

        for orientation in ExifTags.TAGS.keys():

            if ExifTags.TAGS[orientation] == 'Orientation':

                break

        exif = dict(img._getexif().items())

        if exif[orientation] == 3:

            img = img.rotate(180, expand=True)

        elif exif[orientation] == 6:

            img = img.rotate(270, expand=True)

        elif exif[orientation] == 8:

            img = img.rotate(90, expand=True)

    except (AttributeError, KeyError, IndexError):

        # cases: image don't have getexif

        pass

    return img

2. 将图像尺寸调整为640x640，忽略原始宽高比，通过统一图像尺寸以符合研究模型的输入需求。

3. 随机裁剪图像0到30%的区域，以增加模型对局部特征的识别能力。

4. 在-12到+12度间随机旋转图像，同时在水平和垂直方向上进行-2°到+2°的随机剪切，通过对图像进行轻微的旋转和剪切，能够模拟和增强实验模型对真实世界变化的适应性。

3 模型训练

3.1 打标签

在开始标注之前，首先定义了15种不同的水体污染物类别。这些类别涵盖了常见的水体污染源，包括但不限于以下物品：

• Mask: 面具，可能在水中被丢弃。
• Can: 罐头或饮料罐。
• Cellphone: 手机，水中的电子垃圾。
• Electronics: 其他电子设备。
• G(bottle): 玻璃瓶。
• Glove: 手套，包括塑料和橡胶手套。
• Metal: 金属物品。
• Misc: 其他杂项物品。
• Net: 渔网或类似的网状物品。
• P(bag): 塑料袋。
• P(bottle): 塑料瓶。
• Plastic: 其他形式的塑料物品。
• Rod: 钓竿或类似的长棒状物体。
• Sunglasses: 太阳镜。
• Tire: 轮胎。

其次，本研究使用labelImg工具进行图像标注。labelImg 是一个图形化的图像标注工具，支持为本研究提供深度学习训练所需的标注数据，具体的标准过程如下：

首先，打开labelImg工具，并且加载包含水体污染物图片的文件夹。而后，在labelImg中输入预定义的类别，对图片中的每个水体污染物对象，使用鼠标绘制边界框，确保边界框紧密地围绕目标对象。最后，每标注完一张图片，保存该图片的标注信息。labelImg将标注信息保存为TXT格式的文件，每个文件对应一个图片，内容包括边界框的坐标，这种标注文件的格式适用于YOLO（You Only Look Once）目标检测模型。

3.2 拆分数据集

本研究将数据集按照7:2:1划分为训练集、验证集和测试集，具体划分数据集的代码如下：

import os

import random

import shutil

# 原始数据集

root_dir = ' /Users/bcy/resource/water_pollution'  

# 划分比例

train_ratio = 0.7

valid_ratio = 0.2

test_ratio = 0.1

random.seed(35)# 设置随机数生成器

# 拆分后的数据集目录

split_dir = '/Users/bcy/resource/water/split_dataset' # 可以自定义拆分后的目录名

os.makedirs(os.path.join(split_dir, 'train/images'), exist_ok=True)

os.makedirs(os.path.join(split_dir, 'train/labels'), exist_ok=True)

os.makedirs(os.path.join(split_dir, 'valid/images'), exist_ok=True)

os.makedirs(os.path.join(split_dir, 'valid/labels'), exist_ok=True)

os.makedirs(os.path.join(split_dir, 'test/images'), exist_ok=True)

os.makedirs(os.path.join(split_dir, 'test/labels'), exist_ok=True)

# 获取图片文件列表和标签文件列表

image_files = os.listdir(os.path.join(root_dir, 'images'))

label_files = os.listdir(os.path.join(root_dir, 'labels'))

# 打乱列表，保证图像文件与标签文件匹配

combined_files = list(zip(image_files, label_files))

random.shuffle(combined_files)

image_files_shuffled, label_files_shuffled = zip(*combined_files)

# 按照比例计算划分的边界

train_bound = int(train_ratio * len(image_files_shuffled))

valid_bound = int((train_ratio + valid_ratio) * len(image_files_shuffled))

# 将图片和标注文件按比例移动到相应的目录

for i, (image_file, label_file) in enumerate(zip(image_files_shuffled, label_files_shuffled)):

    if i < train_bound:

        shutil.move(os.path.join(root_dir, 'images', image_file), os.path.join(split_dir, 'train/images', image_file))

        shutil.move(os.path.join(root_dir, 'labels', label_file), os.path.join(split_dir, 'train/labels', label_file))

    elif i < valid_bound:

        shutil.move(os.path.join(root_dir, 'images', image_file), os.path.join(split_dir, 'valid/images', image_file))

        shutil.move(os.path.join(root_dir, 'labels', label_file), os.path.join(split_dir, 'valid/labels', label_file))

    else:

        shutil.move(os.path.join(root_dir, 'images', image_file), os.path.join(split_dir, 'test/images', image_file))

        shutil.move(os.path.join(root_dir, 'labels', label_file), os.path.join(split_dir, 'test/labels', label_file))

拆分后的数据集的文件结构如下：

• 训练集：包含3628组图像及其相应的标签，用于模型的训练和学习。
• 验证集：包含1001组图像及其标签，主要用于在训练过程中验证模型的性能，调整模型参数。
• 测试集：包含501组图像及其标签，用于在模型训练完成后评估其最终的性能。

3.3 创建数据配置文件

前文已将数据集分为训练集、验证集和测试集的图片，这些图片已经标注了不同类型的水体污染物（如塑料瓶、包装袋等），具体的数据集信息如下：

类别数量：15种。

类别名称：‘Mask’ ,‘can’, ‘cellphone’, ‘electronics’, ‘gbottle’, ‘glove’, ‘metal’, ‘misc’, ‘net’, ‘pbag’, ‘pbottle’, ‘plastic’。

本研究创建一个Python字典来存储有关数据集的关键信息。这个字典包括了每个数据集部分的路径、类别数量以及类别名称，具体的代码如下：

import os

# 设置数据集目录和类别信息

INPUT_DIR = '/Users/bcy/resource/dataset'  # 更改为你的数据集实际路径

num_classes = 15

classes = ['Mask', 'can', 'cellphone', 'electronics', 'gbottle', 'glove', 'metal', 'misc', 'net', 'pbag', 'pbottle', 'plastic']

# 构建配置文件字典

dict_file = {

    'train': os.path.join(INPUT_DIR, 'train'),

    'val': os.path.join(INPUT_DIR, 'valid'),

    'test': os.path.join(INPUT_DIR, 'test'),

    'nc': num_classes,

    'names': classes

}

而后，将字典内容写入一个YAML文件，这个文件将用于配置深度学习模型的训练过程。本文使用Python的yaml库来实现这一功能，具体的代码实现如下：

import yaml

WORK_DIR = '/Users/bcy/resource'  # 设置工作目录的路径

# 写入YAML文件

with open(os.path.join(WORK_DIR, 'data.yaml'), 'w+') as file:

    yaml.dump(dict_file, file)

以上代码会在指定的工作目录下创建一个data.yaml文件，其中包含了所有必要的数据集配置信息。具体的数据集的文件结构如图b-2所示。

图b-2  拆分数据集后的文件结构

3.4 创建模型配置文件

为实现水体污染物的识别，本文配置并加载一个适合的深度学习模型YOLOv8n模型。具体的模型配置步骤如下：

首先，本文从预训练的权重文件中导入YOLOv8n模型。这个权重文件包含了模型训练好的参数，可以直接用于识别新的图像。具体代码如下：

from yolov8 import YOLO

# 模型文件路径

model_path = 'yolov8n.pt'

# 创建YOLO模型实例

model = YOLO(model_path)

其次，本研究为模型配置模型参数，包括输入图像的大小、置信度阈值、非极大抑制（NMS）阈值等。

为了在M1芯片上优化性能和精度，我们可能需要调整输入图像的大小以匹配芯片处理能力，相关的参赛配置代码如下：

# 设置模型输入图像的大小

model.image_size = 640  # 假设我们使用640x640的输入图像

model.confidence_threshold = 0.5# 设置置信度阈值，只考虑置信度高于此阈值的检测结果

model.nms_threshold = 0.4# 设置NMS阈值，用于过滤重叠的边界框

而后，为了使模型能够正确解析输入数据，需要对输入的图像进行适当的预处理，包括归一化图像像素值、调整图像大小等步骤，相关的代码如下：

from torchvision import transforms

# 定义图像预处理转换

transform = transforms.Compose([

    transforms.Resize((640, 640)),  # 调整图像大小

    transforms.ToTensor(),          # 将图像转换为张量

    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化

])

# 应用预处理到具体图像

image_tensor = transform(image)# image 是PIL.Image格式

最后，完成模型配置和图像预处理后，将预处理后的图像输入模型进行预测。这将产生图像中每个检测到的对象的类别、置信度和边界框位置。

3.5 训练

为了训练模型以识别水体污染物，我们将使用YOLOv8n模型并通过指定一系列训练参数来进行优化。训练过程将使用以下配置：

·  task: 指定训练任务的类型。

·  imgsz: 定义模型输入的图像大小。由于YOLO模型通常需要较大的输入尺寸来准确捕捉图像细节，本研究设置为640x640像素。

·  batch: 确定每个批次的图像数量。这里设置为32，这是一个平衡训练速度和内存使用的常用值。

·  epochs: 定义训练过程将完整遍历数据集的次数，这里设置为30。增加epochs数量通常可以提高模型的性能，直到一定的饱和点。

·  data: 指向包含训练、验证和测试集信息的YAML配置文件的路径。这里使用os.path.join(WORK_DIR,'data.yaml')来确保路径的正确性。

·  mode: 设置训练模式。在这里，我们使用train模式进行模型训练。其他选项包括val用于验证和predict用于预测。

·  name: 定义训练会话的名称，这有助于保存和引用特定的训练结果。在此例中，使用yolov8n_v1_train作为标识

具体的代码配置如下：

model.train(

    data=os.path.join(WORK_DIR,'data.yaml'),

    task='detect',

    imgsz=640,

    epochs=30,

    batch=32,

    mode='train',

    name='yolov8n_v1_train'

)

模型的训练结果如图b-3所示。

图b-3 模型训练结果

具体的YOLOv8n模型实例训练生成的结果如图b-4所示。

图b-4 YOLOv8n模型实例的训练结果

4 模型评估结果

为了确保模型达到最佳性能，我们需要进行系统的模型评估。在训练期间，模型的多个版本（权重）会被保存，其中包括每个epoch产生的权重和评估指标认为最好的权重，本研究使用这些权重创建了一个YOLO模型实例，具体的配置信息如下：

from yolov8 import YOLO

# 模型权重文件的路径

best_model_path = '/User/bcy/runs/detect/yolov8n_v1_train5/weights/best.pt'

model = YOLO(best_model_path)# 加载模型

为评估模型的实际表现能力，本文使用该YOLO模型实例对测试集进行预测。下文代码片段展示了对测试集中的图像进行预测以及将结果保存到指定目录的具体操作。

test_images_dir = os.path.join(INPUT_DIR, 'test', 'images')# 测试集图片的目录

results = model.predict(source=test_images_dir, save=True)# 使用模型进行预测并保存结果

本研究评估模型的具体结果如图b-5所示：

图b-5 YOLOv8n模型实例的预测结果

5 模型推理

为直观评估模型在测试集上的表现，本研究将从预测结果中随机抽取一定数量的样本并展示它们。

本文随机选择16张图片数据进行预测，具体的预测结果如下：

图b-6 模型推理结果

如图b-6，图b-7，图b-8，，图b-9所示为该YOLOv8n模型实例的评估结果，其中，图b-6中的 F1-Confidence Curve展示在不同置信度阈值下各类别的F1分数，图中显示当置信度阈值为0.8时，多数类别的F1分数较低。而图b-7中的Precision-Confidence Curve展示在不同置信度阈值下各类别的精确度，如图，随着置信度阈值的增加，某些类别的精确度有所提高，这表明较高的置信度可以减少误报。图b-8展示的Confusion Matrix Normalized展示模型预测结果与实际标签之间的关系，帮助识别哪些类别容易被混淆，本实验中Mask类别和cellphone类别在对角线上的值较高，表示这些类别的预测准确率较高。图b-9中的Loss Curves展示模型训练过程中损失函数的变化，其中，train/box_loss, train/obj_loss, train/cls_loss表示训练集上的边界框损失、目标损失和类别损失，val/box_loss, val/obj_loss, val/cls_loss表示验证集上的相应损失，损失越低表示模型预测越准确。从图中可以看出，随着训练的进行，大部分损失都呈现下降趋势，说明模型正逐渐拟合数据。

图b-6 F1-Confidence Curve

图b-7 Precision-Confidence Curve

图b-8 Confusion Matrix Normalized

图b-9 Loss Curves