YOLOv8教程系列：四、使用yolov8仓库训练自己的图像分类数据集（含推理预测）

0.引言

Yolov8是最新一代的You Only Look Once目标检测模型,它由Ultralytics研究团队在2022年开发。相比于之前的Yolo版本,Yolov8在速度和精度上都有很大的提升。

在图像分类任务上,Yolov8使用了以下几个特点:

网络backbone采用了CSPResNeXt作为主干网络,此网络结合了ResNet和ResNeXt的优点,能够提取更丰富的特征。
使用了Cross Stage Partial Network作为网络的连接方式,将不同阶段的特征图连接起来,增强特征的表达能力。
引入了注意力机制,可以学习到目标的更加 discriminative 的特征。
使用了更大的输入分辨率,即640x640,可以学习到更丰富的信息。
损失函数上,融合了标签调制、IoU预测以及新颖的Grid Sensitive损失,可以更好地改善模型的定位能力。在推理上使用了更高效的矩阵运算,可以加速模型的推理速度。

总的来说,Yolov8通过网络设计的改进,损失函数的优化以及高效的推理实现,相比之前的版本取得了显著的进步,在图像分类任务上能够达到更高的精度。
yolov8官方文档
**加粗样式**

1.数据准备

先把自己的数据按类别准备好，格式如下，imagenet数据集格式（文件夹名为datasets）：

.
├── ./datasets
│ ├── ./datasets/OK
│ │ ├── ./datasets/OK/1.jpg
│ │ ├── ./datasets/OK/2.jpg
│ │ ├── ./datasets/OK/3.jpg
│ │ ├── …
│ ├── ./datasets/NG
│ │ ├── ./datasets/NG/1.jpg
│ │ ├── ./datasets/NG/1.jpg
│ │ ├── ./datasets/NG/1.jpg
│ │ ├── …

2.数据划分

在datasets上层目录下新建一个split.py文件，运行下面脚本：

# 工具类
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 = "datasets"
    target_data_folder = "dataset/"
    data_set_split(src_data_folder, target_data_folder)

运行结束，会按训练集和验证集8:2的比例生成一个划分后的数据集，名为dataset

3.开始训练

新建一个classify_train.py文件，根据自己情况，调整相关参数即可

from ultralytics import YOLO

model = YOLO("yolo-cls/yolov8s-cls.pt")
model.train(data='/home/lzj/03.AlgoDemo/yolov8/dataset/', epochs=100, batch=2, imgsz=1280)

4.推理预测

新建一个classify_infer.py的脚步，注意修改下面的路径和名称列表，运行结束后，会在指定目录下生成预测的图片

import cv2
import os
from ultralytics import YOLO
from tqdm import tqdm

def read_path(file_pathname, model, name_dict, save_folder):
    file_dir = os.listdir(file_pathname)
    for k,v in name_dict.items():
        name_foler = os.path.join(save_folder, v)
        os.makedirs(name_foler)
    #遍历该目录下的所有图片文件
    for filename in tqdm(file_dir):
        print(filename)
        img = cv2.imread(file_pathname+'/'+filename)
        results = model.predict(source=img)

        for result in results:
            # print(result.names)
            name_dict = result.names
            print(name_dict)
            probs = result.probs.cpu().numpy()
            top1_index = result.probs.top1
            class_name = name_dict[top1_index]
            print(class_name)
            save_img_path = os.path.join(save_folder, class_name, filename)
            cv2.imwrite(save_img_path, img)
        print('---------------------------')



if __name__ == '__main__':
    name_dict = {0: 'NG', 1: 'OK'}
    save_folder = 'classify_infer_folder'
    load_img_folder = '/home/lzj/Downloads/pb'
    model = YOLO('runs/classify/train46/weights/best.pt')
    read_path(load_img_folder, model, name_dict, save_folder)