YOLOv8做图片分类 太强了吧！！！

做完验证码识别之后，本人深刻意识到YOLOv8的强大。验证码任务完成之后，接到了一个实验装置水位检测的任务，即水位低于某一水平，认定为低水位；高于某一水平，认定为高水位；其余的可以认定为正常水位。emmmm，刚开始接到这个任务，想的是标注然后用YOLOv8做目标检测，然后发现标注很麻烦，识别效果也很差。然后突然想到YOLOv8可以做图像分类、检测和分割，就试了一下，不得不说啊，真的是方便省事效果还好。那就记录一下吧，省的下次忘了怎么做。

虚拟环境

新建虚拟环境

conda create -n your_env_name python=x.x

注意，YOLOv8要求：Python>=3.7，PyTorch>=1.7.

查看虚拟环境是否新建成功，有哪些虚拟环境，可以通过以下命令：

conda env list 
conda info -e

激活虚拟环境

conda activate your_env_name

查看虚拟环境中有哪些包，可以使用下面的命令

pip list
conda list

删除虚拟环境

conda remove -n your_env_name --all

安装

pip install ultralytics

为了加速安装，可以添加镜像（附一个镜像汇总链接：https://zhuanlan.zhihu.com/p/298372291），比如使用清华镜像：

pip install ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple

剩下的就不用管了，只需要等待安装完成即可。

准备数据集

第一步，分类放置

在准备数据集方面，相比于目标检测，这个简直太方便了。将图片分类放到不同的文件夹里面，比如我分成了low、normal、high三类，就是三个文件夹，如图

第二步，划分训练集和验证集

还有一个问题是需要划分训练集和验证集，训练集和验证集中都需要包含所有类别图片，这个依靠代码就可以完成啦。把代码放到下面（这个我也是在网上找的，但实在找不到原文了，故没有放链接，sorry）：

import os
import random
from shutil import copy2

def data_set_split(src_data_folder, target_data_folder, train_scale=0.8, val_scale=0.2):
    '''
    读取源数据文件夹，生成划分好的文件夹，分为train、val两个文件夹进行
    :param src_data_folder: 源文件夹
    :param target_data_folder: 目标文件夹
    :param train_scale: 训练集比例
    :param val_scale: 验证集比例
    :return:
    '''
    print("开始数据集划分")
    class_names = os.listdir(src_data_folder)
    # 在目标目录下创建文件夹
    split_names = ['train', 'val']
    for split_name in split_names:
        split_path = os.path.join(target_data_folder, split_name)
        if os.path.isdir(split_path):
            pass
        else:
            os.makedirs(split_path)
        # 然后在split_path的目录下创建类别文件夹
        for class_name in class_names:
            class_split_path = os.path.join(split_path, class_name)
            if os.path.isdir(class_split_path):
                pass
            else:
                os.makedirs(class_split_path)

    # 按照比例划分数据集，并进行数据图片的复制
    # 首先进行分类遍历
    for class_name in class_names:
        current_class_data_path = os.path.join(src_data_folder, class_name)
        current_all_data = os.listdir(current_class_data_path)
        current_data_length = len(current_all_data)
        current_data_index_list = list(range(current_data_length))
        random.shuffle(current_data_index_list)

        train_folder = os.path.join(os.path.join(target_data_folder, 'train'), class_name)
        val_folder = os.path.join(os.path.join(target_data_folder, 'val'), class_name)
        train_stop_flag = current_data_length * train_scale
        current_idx = 0
        train_num = 0
        val_num = 0
        for i in current_data_index_list:
            src_img_path = os.path.join(current_class_data_path, current_all_data[i])
            if current_idx <= train_stop_flag:
                copy2(src_img_path, train_folder)
                train_num = train_num + 1
            else:
                copy2(src_img_path, val_folder)
                val_num = val_num + 1

            current_idx = current_idx + 1

        print("*********************************{}*************************************".format(class_name))
        print("{}类按照{}：{}的比例划分完成，一共{}张图片".format(class_name, train_scale, val_scale, current_data_length))
        print("训练集{}：{}张".format(train_folder, train_num))
        print("验证集{}：{}张".format(val_folder, val_num))


if __name__ == '__main__':
    src_data_folder = "C:/Users/lenovo/Desktop/liquid"
    target_data_folder = "C:/Users/lenovo/Desktop/liquid_split"
    data_set_split(src_data_folder, target_data_folder)

划分完毕，最终数据集格式如下：

下载预训练权重

下载地址：https://docs.ultralytics.com/zh/tasks/classify/#models。官方给了五个预训练权重，从小到大依次是yolov8n-cls.pt、yolov8s-cls.pt、yolov8m-cls.pt、yolov8l-cls.pt和yolov8x-cls.pt，因为我这边数据集比较小，我用的是yolov8n-cls.pt。

训练

准备好了数据集和预训练权重之后，就可以开始训练了。

python代码

from ultralytics import YOLO

# Load a model
model = YOLO('yolov8n-cls.pt')  # 加载预训练模型

# Train the model
results = model.train(data='liquid_split', epochs=100, imgsz=64)

命令行

yolo classify train data=liquid_split model=yolov8n-cls.pt epochs=100 imgsz=64

结果

训练结果出乎意料的好啊哈哈哈哈，结果详细解读可以点击这里，我太懒了不想码字了实在是。

预测

训练模型保存在runs/classify/weights/best.pt

from ultralytics import YOLO

# 加载训练好的模型
model = YOLO('best.pt')


results = model('IMG_20240204_160532.jpg') #image 1/1 D:\PyProject\IMG_20240204_160532.jpg: 64x64 high 1.00, normal 0.00, low 0.00, 72.4ms
print(results[0].names) #{0: 'high', 1: 'low', 2: 'normal'}
print(results[0].probs.top1) #0
print(results[0].names[results[0].probs.top1])#high

参数说明

最后，关于训练参数、验证参数说明、预测参数说明都可以点击这里跳转。