如何训练自己的数据集之，河道垃圾数据集

计算机c9硕士算法工程师

已于 2024-10-28 17:20:55 修改

阅读量962

点赞数 9

分类专栏：数据集垃圾类别数据集文章标签：无人机 YOLO

于 2024-10-27 05:10:12 首次发布

本文链接：https://blog.csdn.net/2401_88441190/article/details/143263461

版权

数据集同时被 2 个专栏收录

318 篇文章

订阅专栏

垃圾类别数据集

7 篇文章

订阅专栏

河道垃圾水污染数据集（无人机视角数据集共3000张图片，可直接用于河道垃圾、水污染功能检测；标注数据人工检查过一遍，保证标注质量；已标注yolo格式、voc格式，可直接训练；继续介绍标签类别及标签个数：水污染：488，漂浮物：5495，废弃船：1215，捕鱼养殖：710，废弃物：156

河道垃圾与水污染数据集介绍

数据集概述

该数据集专注于河道垃圾和水污染的检测，特别设计用于无人机视角下的功能检测。数据集包含3000张图片，已经经过人工检查，确保标注的质量可靠。数据集中的标注已经转换为YOLO格式和VOC格式，可以直接用于模型训练。

数据集特点

高质量标注：所有标注数据均经过人工检查，确保标注的准确性。
多样化的标签：数据集包含五个标签类别，分别是水污染、漂浮物、废弃船、捕鱼养殖和废弃物，涵盖了河道中常见的污染物类型。
多格式支持：数据集提供YOLO格式和VOC格式的标注文件，方便使用不同框架的研究者和开发者。

数据集内容

图像文件：共有3000张JPG/PNG格式的图像文件。
标注文件：每张图像都配有YOLO格式和VOC格式的标注文件。

数据集结构示例

假设数据集的根目录为 river_pollution_dataset，其结构可能如下所示：

river_pollution_dataset/
├── images/
│   ├── image_0001.jpg
│   ├── image_0002.jpg
│   └── ...
├── labels_yolo/
│   ├── image_0001.txt
│   ├── image_0002.txt
│   └── ...
├── labels_voc/
│   ├── image_0001.xml
│   ├── image_0002.xml
│   └── ...
└── README.md  # 数据集说明

标签类别及数量

水污染（Water Pollution）：488个实例
漂浮物（Floating Debris）：5495个实例
废弃船（Abandoned Ships）：1215个实例
捕鱼养殖（Fishing and Aquaculture）：710个实例
废弃物（Waste）：156个实例

这些标签反映了河道中存在的主要污染类型，可用于训练模型以检测和分类不同的污染情况。

数据集使用方法

YOLO格式标注示例

每张图像的标签文件（.txt）包含一行或多行数据，每一行代表一个物体的标签，格式如下：

<class_index> <x_center> <y_center> <width> <height>

其中：

<class_index> 是物体类别的索引（0表示水污染，1表示漂浮物，2表示废弃船，3表示捕鱼养殖，4表示废弃物）；
<x_center> 是物体中心点相对于图像宽度的比例位置；
<y_center> 是物体中心点相对于图像高度的比例位置；
<width> 是物体宽度相对于图像宽度的比例；
<height> 是物体高度相对于图像高度的比例。

VOC格式标注示例

每个图像都有一个XML文件，包含图像的基本信息（如尺寸）以及每个物体的位置信息（边界框坐标）和类别标签。例如：

<annotation>
    <folder>river_pollution_dataset</folder>
    <filename>image_0001.jpg</filename>
    <size>
        <width>640</width>
        <height>480</height>
        <depth>3</depth>
    </size>
    <object>
        <name>Floating Debris</name>
        <pose>Unspecified</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>120</xmin>
            <ymin>100</ymin>
            <xmax>180</xmax>
            <ymax>150</ymax>
        </bndbox>
    </object>
    <!-- 更多物体标签... -->
</annotation>

使用建议

数据预处理：根据模型的需求，可能需要将图像统一到相同的分辨率，或者采用数据增强技术以增加模型的鲁棒性。
模型训练：选择适合的检测框架（如YOLO、Faster R-CNN等），并根据框架的要求调整标注文件格式。
性能评估：在训练过程中，应该定期评估模型在验证集上的表现，以便及时调整模型参数或训练策略。

应用领域

此数据集特别适用于河道管理和环境保护的应用，如河道垃圾监测、水质评估、环境治理等。通过对河道中各类污染源的准确识别，可以辅助相关部门制定有效的管理和治理措施。

总结

此河道垃圾与水污染数据集因其高质量的标注和多样化的标签类别，为无人机视角下的河道污染检测提供了坚实的数据基础。对于从事相关研究和应用的科学家和工程师来说，这是一个非常有用的数据资源。

准备工作

安装YOLOv5：确保您已经安装了YOLOv5及其依赖项。
准备数据集：确保您的数据集已经按照YOLOv5所需的格式准备好，并且有相应的data.yaml文件来描述数据集的路径和其他相关信息。

示例代码

安装YOLOv5

如果您还没有安装YOLOv5，请按照官方文档执行以下命令：

git clone https://github.com/ultralytics/yolov5.git  # clone repo
cd yolov5
pip install -r requirements.txt  # install dependencies

准备data.yaml

创建一个data.yaml文件来描述您的数据集。这里假设数据集被放置在一个名为river_pollution_dataset的目录中，且包含images和labels子目录。

# data.yaml 文件
train: ../river_pollution_dataset/images/train/
val: ../river_pollution_dataset/images/val/
test: ../river_pollution_dataset/images/test/  # optional

nc: 5  # number of classes
names: ['water_pollution', 'floating_debris', 'abandoned_ships', 'fishing_aquaculture', 'waste']  # class names

训练模型

使用以下命令开始训练模型：

cd yolov5
python train.py --img 640 --batch 16 --epochs 300 --data ../river_pollution_dataset/data.yaml --weights yolov5s.pt --cache

自定义训练脚本

如果需要更详细的控制，可以编写一个Python脚本来执行训练过程。以下是一个简单的脚本示例：

import torch
from utils.datasets import LoadImagesAndLabels  # 导入数据加载器
from models.experimental import attempt_load  # 导入模型加载器
from utils.torch_utils import select_device  # 导入选用设备的函数
from utils.general import check_dataset  # 导入检查数据集的函数

def main():
    device = select_device('')  # 选择设备，自动选择GPU/CPU
    data_yaml = '../river_pollution_dataset/data.yaml'
    train_images_folder = '../river_pollution_dataset/images/train/'
    train_labels_folder = '../river_pollution_dataset/labels/train/'
    val_images_folder = '../river_pollution_dataset/images/val/'
    val_labels_folder = '../river_pollution_dataset/labels/val/'

    # 加载数据集
    train_set = LoadImagesAndLabels(train_images_folder, train_labels_folder)
    val_set = LoadImagesAndLabels(val_images_folder, val_labels_folder)

    # 创建数据加载器
    train_loader = torch.utils.data.DataLoader(train_set, batch_size=16, shuffle=True, num_workers=4)
    val_loader = torch.utils.data.DataLoader(val_set, batch_size=16, shuffle=False, num_workers=4)

    # 加载预训练模型
    model = attempt_load('yolov5s.pt', map_location=device)  # 加载预训练权重
    model.to(device)

    # 设置损失函数和优化器
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    # 训练循环
    epochs = 300
    for epoch in range(epochs):
        model.train()
        for images, targets, _, _ in train_loader:
            images = images.to(device)
            targets = [t.to(device) for t in targets]

            # 前向传播
            outputs = model(images)

            # 计算损失
            loss = criterion(outputs, targets)

            # 反向传播和优化
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item()}')

    torch.save(model.state_dict(), 'trained_model.pt')
    print('Training complete.')

if __name__ == '__main__':
    main()

注意事项

确保data.yaml文件中的路径是正确的，并且数据集的结构与上面描述的一致。
调整批量大小、学习率、迭代次数等超参数以适应您的计算资源和任务需求。
如果数据集很大，您可能需要更多的计算资源和时间来完成训练。
这个脚本仅作为一个起点，您可能需要根据实际情况做进一步的修改。

使用YOLOv5命令行训练

如果您更倾向于使用命令行来训练模型，可以使用以下命令：

cd yolov5
python train.py --img 640 --batch 16 --epochs 300 --data ../river_pollution_dataset/data.yaml --weights yolov5s.pt --cache

测试模型

在训练完成后，您可以通过以下命令测试模型的性能：

python detect.py --weights runs/train/exp/weights/best.pt --img 640 --conf 0.4 --source ../river_pollution_dataset/images/test/

总结

这个示例展示了如何使用YOLOv5框架训练一个基于无人机视角下的河道垃圾与水污染检测模型。您可以根据自己的需求调整脚本中的参数和逻辑。通过使用这个数据集和相应的训练代码，您可以有效地训练出一个能够识别河道中不同污染类型的目标检测模型。