【源码真的免费】yolo道路图像分类

源代码解压在F盘可直接运行 提取码：VpSV

data提取码：ZLp2

一、数据集处理

        本文所用数据集是下载自23年mathorcup大数据赛A题的图像数据，安装下图分类成两类，文件名中包含“normal”字符表示正常道路，否则为坑洼道路，特别说明一下不需要标注文件。

        接着运行split_train_val.py按照官网要求将数据集划分为训练集与验证集，结果如下：

import os
import shutil
import random

# 源数据集目录
source_dataset_dir = 'F:/ultralytics/loaddata'
train_dir = 'F:/ultralytics/loaddata2/train'
valid_dir = 'F:/ultralytics/loaddata2/val'

# 定义训练集和验证集的比例
train_ratio = 0.8  # 80%用于训练，20%用于验证

# 创建训练集和验证集目录
os.makedirs(train_dir, exist_ok=True)
os.makedirs(valid_dir, exist_ok=True)

# 获取数据集中所有类别的子目录
categories = os.listdir(source_dataset_dir)

# 支持的图像文件扩展名
image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']

for category in categories:
    category_dir = os.path.join(source_dataset_dir, category)

    # 检查是否为目录
    if not os.path.isdir(category_dir):
        continue  # 跳过非目录项

    images = [f for f in os.listdir(category_dir) if os.path.splitext(f)[1].lower() in image_extensions]

    # 随机打乱图像列表
    random.shuffle(images)

    # 计算训练集和验证集的分割索引
    split_index = int(len(images) * train_ratio)

    # 创建类别目录（如果不存在）
    train_category_dir = os.path.join(train_dir, category)
    valid_category_dir = os.path.join(valid_dir, category)
    os.makedirs(train_category_dir, exist_ok=True)
    os.makedirs(valid_category_dir, exist_ok=True)

    # 将图像复制到相应的目录
    for i, image in enumerate(images):
        src_path = os.path.join(category_dir, image)
        if i < split_index:
            dest_dir = train_category_dir
        else:
            dest_dir = valid_category_dir
        dest_path = os.path.join(dest_dir, image)

        try:
            shutil.copy(src_path, dest_path)
            print(f"Copied {image} to {dest_dir}")  # 可选：打印复制信息
        except Exception as e:
            print(f"Failed to copy {image} to {dest_dir}: {e}")  # 打印错误信息

print("数据集划分完成。")

        运行结果如下： 

        后面就以loaddata1作为训练数据集 .

二、模型训练

        新建一个cls-train.py文件，安装官网（https://docs.ultralytics.com/tasks/classify/#models）给的示例代码进行训练

from ultralytics import YOLO

# Load a model
#model = YOLO("yolo11n-cls.yaml")  # build a new model from YAML
model = YOLO("F:/ultralytics/yolo11n-cls.pt")  # load a pretrained model (recommended for training)
#model = YOLO("yolo11n-cls.yaml").load("yolo11n-cls.pt")  # build from YAML and transfer weights

# Train the model
results = model.train(data="F:/ultralytics/loaddata1", epochs=100, imgsz=64)

        训练结果如下：

三、模型验证 

         同样安装官网新建cls-val.py进行验证

from ultralytics import YOLO

# Load a model
#model = YOLO("yolo11n-cls.pt")  # load an official model
model = YOLO("F:/ultralytics/run/runs/classify/train/weights/best.pt")  # load a custom model

# Validate the model
metrics = model.val()  # no arguments needed, dataset and settings remembered
metrics.top1  # top1 accuracy
metrics.top5  # top5 accuracy

        运行结果如下：

四、模型测试

          同样安装官网新建cls-predict.py进行测试

from ultralytics import YOLO
import cv2
# Load a model
#model = YOLO("yolo11n-cls.pt")  # load an official model
model = YOLO("F:/ultralytics/run/runs/classify/train/weights/best.pt")  # load a custom model

# Predict with the model
results = model("F:/ultralytics/testdata_V2/0a9kzasw.jpg",show=True)  # predict on an image

         运行结果如下：