YOLOv8 tutorial series: 4. Use yolov8 warehouse to train your own image classification data set (including inference prediction)

YOLOv8 tutorial series: 4. Use yolov8 warehouse to train your own image classification data set (including inference prediction)

0 Preface

Yolov8 is the latest generation of You Only Look Once object detection model, which will be developed by the Ultralytics research team in 2022. Compared with the previous Yolo version, Yolov8 has greatly improved in speed and accuracy.

On image classification tasks, Yolov8 uses the following features:

1. The network backbone uses CSPResNeXt as the backbone network. This network combines the advantages of ResNet and ResNeXt and can extract richer features.
2. The Cross Stage Partial Network is used as the connection method of the network to connect the feature maps of different stages to enhance the expressive ability of the features.
3. The attention mechanism is introduced to learn more discriminative features of the target.
4. Using a larger input resolution, namely 640x640, can learn richer information.
5. In the loss function, label modulation, IoU prediction and novel Grid Sensitive loss are combined to better improve the positioning ability of the model. More efficient matrix operations are used in reasoning, which can speed up the reasoning speed of the model.

In general, Yolov8 has made significant progress compared with the previous version through the improvement of network design, optimization of loss function and efficient inference implementation, and can achieve higher accuracy in image classification tasks.
yolov8 official document

1. Data preparation

First prepare your own data by category, the format is as follows, imagenet dataset format (folder name 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. Data division

Create a split.py file in the upper directory of datasets and run the following script:

# 工具类
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)

At the end of the run, a divided data set will be generated according to the ratio of 8:2 between the training set and the verification set, nameddataset

3. Start training

Create a new classify_train.pyfile and adjust the relevant parameters according to your own situation.

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. Inference prediction

Create a classify_infer.pynew script, pay attention to modify the following path and name list, after the operation is completed, the predicted picture will be generated in the specified directory

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)