Serie de tutoriales de YOLOv8: 4. Utilice el almacén yolov8 para entrenar su propio conjunto de datos de clasificación de imágenes (incluida la predicción de inferencia)

Serie de tutoriales de YOLOv8: 4. Utilice el almacén yolov8 para entrenar su propio conjunto de datos de clasificación de imágenes (incluida la predicción de inferencia)

0 Prefacio

Yolov8 es la última generación del modelo de detección de objetos You Only Look Once, que será desarrollado por el equipo de investigación de Ultralytics en 2022. En comparación con la versión anterior de Yolo, Yolov8 ha mejorado enormemente en velocidad y precisión.

En tareas de clasificación de imágenes, Yolov8 utiliza las siguientes funciones:

1. La red troncal utiliza CSPResNeXt como red troncal, que combina las ventajas de ResNet y ResNeXt y puede extraer funciones más ricas.
2. La red parcial entre etapas se utiliza como método de conexión de la red para conectar los mapas de características de diferentes etapas para mejorar la capacidad expresiva de las características.
3. El mecanismo de atención se introduce para aprender características más discriminativas del objetivo.
4. Utilizando una resolución de entrada mayor, concretamente 640x640, se puede obtener información más rica.
5. En la función de pérdida, la modulación de etiquetas, la predicción de IoU y la nueva pérdida sensible a la red se combinan para mejorar la capacidad de posicionamiento del modelo. En el razonamiento se utilizan operaciones matriciales más eficientes, que pueden acelerar la velocidad de razonamiento del modelo.

En general, Yolov8 ha logrado avances significativos en comparación con la versión anterior mediante la mejora del diseño de la red, la optimización de la función de pérdida y la implementación eficiente de la inferencia, y puede lograr una mayor precisión en las tareas de clasificación de imágenes.
documento oficial yolov8

1. Preparación de datos

Primero prepare sus propios datos por categoría, el formato es el siguiente, formato de conjunto de datos de imagenet (nombre de carpeta 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. División de datos

Cree un archivo split.py en el directorio superior de conjuntos de datos y ejecute el siguiente 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)

Al final de la ejecución, se generará un conjunto de datos dividido según la proporción de 8:2 entre el conjunto de entrenamiento y el conjunto de verificación, denominadodataset

3. Empieza a entrenar

Cree un nuevo classify_train.pyarchivo y ajuste los parámetros relevantes según su propia situación.

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. Predicción de inferencia

Cree un classify_infer.pynuevo script, preste atención para modificar la siguiente ruta y lista de nombres. Una vez completada la operación, la imagen predicha se generará en el directorio especificado.

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)