基于YOLOv8模型的雨天道路目标检测系统 如何训练使用 极端天气 黑夜_白天 下雨雨天道路目标检测数据集 3600张 雨天 带标注 voc yolo

雨天道路目标检测数据集 3600张 雨天 带标注 voc yolo
白天

分类名: (图片张数， 标注个数)
car: (1134， 10528)
truck: (394， 722)
person: (387， 1585)
rider: (47， 56)
bike: (64, 86) .
bus:(187， 256)
motor :(29，32)
总数: (1150， 13265)
总类(nc): 7类

晚上
分类名: (图片张数， 标注个数)
car: (2470， 21655)
truck: (384， 499)
per son:(510， 1532)
bus: (220， 248)
bike: (76， 121)
rider: (64, 71)
motor: (43， 49)
总数: (2494， 24175)
总类(nc): 7类


构建一个基于YOLOv8模型的雨天道路目标检测系统。我们将使用PyTorch 1.10.0，并提供两个主要程序：一个是train.py，用于训练模型；另一个是predict.py，用于加载训练好的模型并进行预测。

1. 环境准备

首先，确保你已经安装了所需的依赖项。你可以使用以下命令安装这些依赖项：

pip install torch==1.10.0 torchvision matplotlib opencv-python ultralytics

2. 数据集准备

假设你的数据集已经准备好，并且分为训练集和验证集。数据集目录结构如下：

rainy_road_dataset/
├── images/
│   ├── train_day/
│   ├── val_day/
│   ├── train_night/
│   └── val_night/
├── labels_voc/
│   ├── train_day/
│   ├── val_day/
│   ├── train_night/
│   └── val_night/
├── labels_yolo/
│   ├── train_day/
│   ├── val_day/
│   ├── train_night/
│   └── val_night/
└── data.yaml

3. 数据集配置文件 (data.yaml)

创建一个data.yaml文件，配置数据集的路径和类别信息：

path: ./rainy_road_dataset  # 数据集路径
train_day: images/train_day  # 白天训练集图像路径
val_day: images/val_day  # 白天验证集图像路径
train_night: images/train_night  # 晚上训练集图像路径
val_night: images/val_night  # 晚上验证集图像路径

nc: 7  # 类别数
names: ['car', 'truck', 'person', 'rider', 'bike', 'bus', 'motor']  # 类别名称

4. 转换标注格式

假设标注文件是VOC格式的XML文件，我们需要将它们转换为YOLO格式的TXT文件。

转换脚本

import xml.etree.ElementTree as ET
import os

def convert_voc_to_yolo(voc_file, yolo_file, class_names):
    tree = ET.parse(voc_file)
    root = tree.getroot()

    width = int(root.find('size/width').text)
    height = int(root.find('size/height').text)

    with open(yolo_file, 'w') as f:
        for obj in root.findall('object'):
            class_name = obj.find('name').text
            if class_name not in class_names:
                continue
            class_id = class_names.index(class_name)

            bbox = obj.find('bndbox')
            x_min = float(bbox.find('xmin').text)
            y_min = float(bbox.find('ymin').text)
            x_max = float(bbox.find('xmax').text)
            y_max = float(bbox.find('ymax').text)

            x_center = (x_min + x_max) / 2.0 / width
            y_center = (y_min + y_max) / 2.0 / height
            w = (x_max - x_min) / width
            h = (y_max - y_min) / height

            f.write(f"{class_id} {x_center} {y_center} {w} {h}\n")

def convert_all_voc_to_yolo(voc_dir, yolo_dir, class_names):
    os.makedirs(yolo_dir, exist_ok=True)
    for filename in os.listdir(voc_dir):
        if filename.endswith('.xml'):
            voc_file = os.path.join(voc_dir, filename)
            yolo_file = os.path.join(yolo_dir, filename.replace('.xml', '.txt'))
            convert_voc_to_yolo(voc_file, yolo_file, class_names)

if __name__ == "__main__":
    class_names = ['car', 'truck', 'person', 'rider', 'bike', 'bus', 'motor']
    
    voc_train_day_dir = 'rainy_road_dataset/labels_voc/train_day'
    yolo_train_day_dir = 'rainy_road_dataset/labels_yolo/train_day'
    convert_all_voc_to_yolo(voc_train_day_dir, yolo_train_day_dir, class_names)

    voc_val_day_dir = 'rainy_road_dataset/labels_voc/val_day'
    yolo_val_day_dir = 'rainy_road_dataset/labels_yolo/val_day'
    convert_all_voc_to_yolo(voc_val_day_dir, yolo_val_day_dir, class_names)

    voc_train_night_dir = 'rainy_road_dataset/labels_voc/train_night'
    yolo_train_night_dir = 'rainy_road_dataset/labels_yolo/train_night'
    convert_all_voc_to_yolo(voc_train_night_dir, yolo_train_night_dir, class_names)

    voc_val_night_dir = 'rainy_road_dataset/labels_voc/val_night'
    yolo_val_night_dir = 'rainy_road_dataset/labels_yolo/val_night'
    convert_all_voc_to_yolo(voc_val_night_dir, yolo_val_night_dir, class_names)

5. 训练脚本 (train.py)

from ultralytics import YOLO

def train_model(data_yaml_path, model_config, epochs, batch_size, img_size, augment):
    # 加载模型
    model = YOLO(model_config)

    # 训练模型
    results = model.train(
        data=data_yaml_path,
        epochs=epochs,
        batch=batch_size,
        imgsz=img_size,
        augment=augment
    )

    # 保存模型
    model.save("runs/train/rainy_road/best.pt")

if __name__ == "__main__":
    data_yaml_path = 'rainy_road_dataset/data.yaml'
    model_config = 'yolov8n.yaml'
    epochs = 100
    batch_size = 16
    img_size = 640
    augment = True

    train_model(data_yaml_path, model_config, epochs, batch_size, img_size, augment)

6. 预测脚本 (predict.py)

import cv2
import torch
from ultralytics import YOLO

def predict_image(image_path, model_path, img_size=640):
    # 加载模型
    model = YOLO(model_path)

    # 读取图像
    image = cv2.imread(image_path)
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # 进行预测
    results = model(image_rgb, size=img_size)

    # 处理预测结果
    for result in results:
        boxes = result.boxes.xyxy.cpu().numpy()
        scores = result.boxes.conf.cpu().numpy()
        labels = result.boxes.cls.cpu().numpy().astype(int)

        for box, score, label in zip(boxes, scores, labels):
            x1, y1, x2, y2 = map(int, box)
            class_name = ['car', 'truck', 'person', 'rider', 'bike', 'bus', 'motor'][label]
            color = (0, 255, 0) if label == 0 else (0, 0, 255)
            cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
            cv2.putText(image, f'{class_name} {score:.2f}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)

    # 显示图像
    cv2.imshow('Prediction', image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

if __name__ == "__main__":
    image_path = 'path_to_your_image.jpg'
    model_path = 'runs/train/rainy_road/best.pt'
    predict_image(image_path, model_path)

7. 运行脚本

1. 转换标注格式：

   python convert_voc_to_yolo.py
   

2. 训练模型：

   python train.py
   

3. 进行预测：

   python predict.py
   

8. 详细解释

转换标注格式脚本 (convert_voc_to_yolo.py)

1. 导入依赖项：

   • import xml.etree.ElementTree as ET：导入XML解析库。
   • import os：导入操作系统接口库。

2. 定义转换函数：

   • convert_voc_to_yolo：将单个VOC格式的XML文件转换为YOLO格式的TXT文件。
   • convert_all_voc_to_yolo：将指定目录下的所有VOC格式的XML文件转换为YOLO格式的TXT文件。

3. 主函数：

   • 设置类别名称和各个子集的路径。
   • 调用convert_all_voc_to_yolo函数进行转换。

训练脚本 (train.py)

1. 导入依赖项：

   • from ultralytics import YOLO：导入YOLOv8模型。

2. 定义训练函数：

   • train_model：加载模型，设置训练参数，训练模型，并保存最佳模型。

3. 主函数：

   • 设置数据集路径、模型配置、训练参数等。
   • 调用train_model函数进行训练。

预测脚本 (predict.py)

1. 导入依赖项：

   • import cv2：导入OpenCV库。
   • import torch：导入PyTorch库。
   • from ultralytics import YOLO：导入YOLOv8模型。

2. 定义预测函数：

   • predict_image：加载模型，读取图像，进行预测，处理预测结果，并显示带有标注的图像。

3. 主函数：

   • 设置图像路径和模型路径。
   • 调用predict_image函数进行预测。

总结

通过以上步骤，你可以构建一个基于YOLOv8模型的雨天道路目标检测系统。convert_voc_to_yolo.py用于将VOC格式的标注文件转换为YOLO格式，train.py用于训练模型，predict.py用于加载训练好的模型并进行预测