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 copy2def 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):passelse: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):passelse: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_scalecurrent_idx = 0train_num = 0val_num = 0for 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 + 1else:copy2(src_img_path, val_folder)val_num = val_num + 1current_idx = current_idx + 1print("*********************************{}*************************************".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 YOLOmodel = 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 tqdmdef 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.namesprint(name_dict)probs = result.probs.cpu().numpy()top1_index = result.probs.top1class_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)