基于YOLOv8进行路面积水识别数据集+系统代码。括数据集准备、模型训练、评估和推理路面道路积水数据集 4524，xml和txt标签都有

道路积水数据集 路面积水识别数据集 图片数量4524，xml和txt标签都有；公路积水数据集
✓类别：puddle；

✓图片数量4524，xml和txt标签都有；
✓类别：puddle；



好的，下面是一个完整的项目代码和详细配置说明，用于基于YOLOv8进行路面积水识别。该项目包括数据集准备、模型训练、评估和推理，并附带一个简单的Streamlit界面。

完整项目代码

1. 数据集准备与转换

首先，我们需要将数据集从XML格式转换为YOLO格式（如果数据集已经是YOLO格式，则跳过此步骤）。

import xml.etree.ElementTree as ET
import os
from pathlib import Path

def convert_xml_to_yolo(xml_dir, output_dir):
    for xml_file in Path(xml_dir).glob('*.xml'):
        tree = ET.parse(xml_file)
        root = tree.getroot()
        
        img_name = root.find('filename').text
        img_width = int(root.find('size/width').text)
        img_height = int(root.find('size/height').text)
        
        txt_file_path = os.path.join(output_dir, Path(img_name).stem + '.txt')
        with open(txt_file_path, 'w') as txt_file:
            for obj in root.findall('object'):
                cls = obj.find('name').text
                if cls != 'puddle':
                    continue
                
                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 / img_width
                y_center = (ymin + ymax) / 2 / img_height
                width = (xmax - xmin) / img_width
                height = (ymax - ymin) / img_height
                
                txt_file.write(f"0 {x_center} {y_center} {width} {height}\n")

# 使用示例
convert_xml_to_yolo('datasets/puddle_dataset/xml', 'datasets/puddle_dataset/yolo/labels')

2. 创建数据集配置文件 (data.yaml)

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

3. 分割数据集

为了训练和验证，我们需要将数据集分割成训练集和验证集。

import os
import random
from pathlib import Path
import shutil

def split_dataset(data_dir, train_ratio=0.8):
    images = list(Path(data_dir).glob('*.jpg'))
    random.shuffle(images)
    
    num_train = int(len(images) * train_ratio)
    train_images = images[:num_train]
    val_images = images[num_train:]
    
    train_dir = Path(data_dir).parent / 'train'
    val_dir = Path(data_dir).parent / 'val'
    
    train_img_dir = train_dir / 'images'
    train_label_dir = train_dir / 'labels'
    val_img_dir = val_dir / 'images'
    val_label_dir = val_dir / 'labels'
    
    train_img_dir.mkdir(parents=True, exist_ok=True)
    train_label_dir.mkdir(parents=True, exist_ok=True)
    val_img_dir.mkdir(parents=True, exist_ok=True)
    val_label_dir.mkdir(parents=True, exist_ok=True)
    
    for img in train_images:
        label_path = img.with_suffix('.txt')
        shutil.copy(img, train_img_dir / img.name)
        shutil.copy(label_path, train_label_dir / label_path.name)
    
    for img in val_images:
        label_path = img.with_suffix('.txt')
        shutil.copy(img, val_img_dir / img.name)
        shutil.copy(label_path, val_label_dir / label_path.name)

# 使用示例
split_dataset('./datasets/puddle_dataset/images')

4. 训练脚本

接下来是使用YOLOv8进行训练的脚本。

import torch
from ultralytics import YOLO

# 设置随机种子以保证可重复性
torch.manual_seed(42)

# 定义数据集路径
dataset_config = 'data.yaml'

# 加载预训练的YOLOv8n模型
model = YOLO('yolov8n.pt')

# 训练模型
results = model.train(
    data=dataset_config,
    epochs=100,
    imgsz=512,
    batch=16,
    name='puddle_detection',
    project='runs/train'
)

# 评估模型
metrics = model.val()

# 保存最佳模型权重
best_model_weights = 'runs/train/puddle_detection/weights/best.pt'
print(f"Best model weights saved to {best_model_weights}")

5. 推理脚本

以下是一个简单的推理脚本，用于测试模型的性能。

from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image

# 加载模型
model = YOLO('runs/train/puddle_detection/weights/best.pt')

# 图片检测函数
def detect_image(model, image_path, conf_threshold=0.5):
    results = model.predict(image_path, conf=conf_threshold)[0]
    annotated_frame = annotate_image(image_path, results, model)
    return annotated_frame

# 标注图像函数
def annotate_image(image_path, results, model):
    frame = cv2.imread(image_path)
    for result in results.boxes.cpu().numpy():
        r = result.xyxy[0].astype(int)
        cls = int(result.cls[0])
        conf = result.conf[0]
        
        label = f"{model.names[cls]} {conf:.2f}"
        color = (0, 255, 0)
        cv2.rectangle(frame, (r[0], r[1]), (r[2], r[3]), color, 2)
        cv2.putText(frame, label, (r[0], r[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
    return frame

# 测试一张图片
test_image_path = './datasets/puddle_dataset/images/test_image.jpg'
annotated_image = detect_image(model, test_image_path)
cv2.imwrite('annotated_test_image.jpg', annotated_image)

6. Streamlit 主界面

最后，我们创建一个简单的Streamlit界面来进行实时推理和显示结果。

import streamlit as st
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image
import tempfile

# 加载模型
@st.cache_resource
def load_model(weights_path):
    model = YOLO(weights_path)
    return model

# 图片检测函数
def detect_image(model, image, conf_threshold):
    results = model.predict(image, conf=conf_threshold)[0]
    annotated_frame = annotate_image(image, results, model)
    return annotated_frame

# 视频检测函数
def detect_video(model, video_path, conf_threshold):
    cap = cv2.VideoCapture(video_path)
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        results = model.predict(frame, conf=conf_threshold)[0]
        annotated_frame = annotate_image(frame, results, model)
        yield annotated_frame
    cap.release()

# 摄像头检测函数
def detect_camera(model, conf_threshold):
    cap = cv2.VideoCapture(0)
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        results = model.predict(frame, conf=conf_threshold)[0]
        annotated_frame = annotate_image(frame, results, model)
        yield annotated_frame
    cap.release()

# 标注图像函数
def annotate_image(image, results, model):
    for result in results.boxes.cpu().numpy():
        r = result.xyxy[0].astype(int)
        cls = int(result.cls[0])
        conf = result.conf[0]
        
        label = f"{model.names[cls]} {conf:.2f}"
        color = (0, 255, 0)
        cv2.rectangle(image, (r[0], r[1]), (r[2], r[3]), color, 2)
        cv2.putText(image, label, (r[0], r[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
    return image

# Streamlit 主界面
def main():
    st.title("Puddle Detection System")
    
    # 动态加载模型
    weights_options = ['best_puddle_yolov8.pt']  # 添加更多权重文件路径
    selected_weights = st.sidebar.selectbox("Select Model Weights", weights_options)
    model = load_model(selected_weights)
    
    # 动态调整置信度阈值
    conf_threshold = st.sidebar.slider("Confidence Threshold", min_value=0.0, max_value=1.0, value=0.5, step=0.01)
    
    # 输入方式选择
    input_type = st.sidebar.radio("Input Type", ["Image", "Video", "Camera"])
    
    if input_type == "Image":
        uploaded_file = st.file_uploader("Upload an image...", type=["jpg", "jpeg", "png"])
        if uploaded_file is not None:
            image = Image.open(uploaded_file)
            image_np = np.array(image)
            annotated_image = detect_image(model, image_np, conf_threshold)
            st.image(annotated_image, channels="BGR", caption="Detected Image")
            
            # 统计检测到的物体数量
            results = model.predict(image_np, conf=conf_threshold)[0]
            class_counts = {}
            for result in results.boxes.cpu().numpy():
                cls = int(result.cls[0])
                class_name = model.names[cls]
                if class_name in class_counts:
                    class_counts[class_name] += 1
                else:
                    class_counts[class_name] = 1
            
            st.subheader("Detection Summary:")
            for class_name, count in class_counts.items():
                st.write(f"{class_name}: {count}")
    
    elif input_type == "Video":
        uploaded_file = st.file_uploader("Upload a video...", type=["mp4", "avi"])
        if uploaded_file is not None:
            tfile = tempfile.NamedTemporaryFile(delete=False) 
            tfile.write(uploaded_file.read())
            tfpath = tfile.name
            cap = cv2.VideoCapture(tfpath)
            frame_placeholder = st.empty()
            for annotated_frame in detect_video(model, tfpath, conf_threshold):
                frame_placeholder.image(annotated_frame, channels="BGR", use_column_width=True)
            cap.release()
            os.remove(tfpath)
    
    elif input_type == "Camera":
        frame_placeholder = st.empty()
        for annotated_frame in detect_camera(model, conf_threshold):
            frame_placeholder.image(annotated_frame, channels="BGR", use_column_width=True)

if __name__ == "__main__":
    main()

文件结构

puddle_detection/
├── main.py
├── datasets/
│   └── puddle_dataset/
│       ├── xml/
│       │   ├── image1.xml
│       │   ├── image2.xml
│       │   └── ...
│       ├── yolo/
│       │   ├── labels/
│       │   │   ├── image1.txt
│       │   │   ├── image2.txt
│       │   │   └── ...
│       │   ├── train/
│       │   │   ├── images/
│       │   │   │   ├── image1.jpg
│       │   │   │   ├── image2.jpg
│       │   │   │   └── ...
│       │   │   └── labels/
│       │   │       ├── image1.txt
│       │   │       ├── image2.txt
│       │   │       └── ...
│       │   └── val/
│       │       ├── images/
│       │       │   ├── image1.jpg
│       │       │   ├── image2.jpg
│       │       │   └── ...
│       │       └── labels/
│       │           ├── image1.txt
│       │           ├── image2.txt
│       │           └── ...
│       └── images/
│           ├── image1.jpg
│           ├── image2.jpg
│           └── ...
├── best_puddle_yolov8.pt
└── requirements.txt

安装依赖项

首先，确保你已经安装了所有必要的依赖项。你可以通过以下命令安装：

pip install -r requirements.txt

requirements.txt 内容如下：

streamlit==1.25.0
opencv-python
torch==2.0
ultralytics

运行步骤总结

1. 克隆项目仓库（如果有的话）：

   git clone https://github.com/yourusername/puddle_detection.git
   cd puddle_detection
   

2. 安装依赖项：

   pip install -r requirements.txt
   

3. 数据集准备：

   • 如果数据集尚未转换，请运行数据集转换脚本。

     python convert_xml_to_yolo.py
     

   • 如果数据集尚未分割，请运行数据集分割脚本。

     python split_dataset.py
     

   • 确保数据集目录结构正确。

4. 训练模型：

   python train.py
   

5. 评估模型：
   在训练脚本中，模型会在训练结束后自动进行评估。

6. 推理测试：

   python inference.py
   

7. 运行Streamlit应用：

   streamlit run main.py
   

总结

以上是完整的基于YOLOv8的路面积水识别系统的项目介绍和代码实现。该项目支持图片、视频识别以及本地摄像头识别，并且可以通过UI界面动态调节模型置信度和选择不同的模型权重。希望这些详细的信息和代码能够帮助你顺利实施和优化你的路面积水识别系统。