YOLOv8 classify介绍

本文主要是介绍YOLOv8 classify介绍，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

图像分类器(image classifier)的输出是单个类别标签和置信度分数。当你只需要知道图像属于哪个类别，而不需要知道该类别的目标位于何处或它们的确切形状时，图像分类非常有用。

YOLOv8支持的预训练分类模型包括：YOLOv8n-cls、YOLOv8s-cls、YOLOv8m-cls、YOLOv8l-cls、YOLOv8x-cls。分类模型在ImageNet数据集上进行预训练，而检测、分割和Pose模型是在COCO数据集上进行预训练。

数据集格式：

(1).对于Ultralytics YOLO分类任务，数据集必须在根目录下以特定的拆分目录(split-directory)结构进行组织，以便进行正确的训练、测试和验证(可选的)过程。此结构包括用于训练(train)和测试(test)阶段的单独目录，以及用于验证(val)的可选目录。

(2).每个目录都应包含数据集中每个类别的一个子目录。子目录以相应的类别命名，包含该类别的所有图像。确保每个图像文件都具有唯一的名称，并以JPEG或PNG等通用格式存储。

(3).例如CIFAR-10数据集的目录结构如下所示：

这里使用 https://blog.csdn.net/fengbingchun/article/details/141635132 中的数据集，通过YOLOv8 classify进行train和predict：

train代码如下：

import argparse
import colorama
from ultralytics import YOLO
import torchdef parse_args():parser = argparse.ArgumentParser(description="YOLOv8 train")parser.add_argument("--yaml", required=True, type=str, help="yaml file or datasets path(classify)")parser.add_argument("--epochs", required=True, type=int, help="number of training")parser.add_argument("--task", required=True, type=str, choices=["detect", "segment", "classify"], help="specify what kind of task")parser.add_argument("--imgsz", type=int, default=640, help="input net image size")args = parser.parse_args()return argsdef train(task, yaml, epochs, imgsz):if task == "detect":model = YOLO("yolov8n.pt") # load a pretrained model, should be a *.pt PyTorch model to run this methodelif task == "segment":model = YOLO("yolov8n-seg.pt") # load a pretrained model, should be a *.pt PyTorch model to run this methodelif task == "classify":model = YOLO("yolov8n-cls.pt") # n/s/m/l/xelse:raise ValueError(colorama.Fore.RED + f"Error: unsupported task: {task}")# petience: Training stopped early as no improvement observed in last patience epochs, use patience=0 to disable EarlyStoppingresults = model.train(data=yaml, epochs=epochs, imgsz=imgsz, patience=150, augment=True) # train the model, supported parameter reference, for example: runs/segment(detect)/train3/args.yamlmetrics = model.val() # It'll automatically evaluate the data you trained, no arguments needed, dataset and settings rememberedif task == "classify":print("Top-1 Accuracy:", metrics.top1) # top1 accuracyprint("Top-5 Accuracy:", metrics.top5) # top5 accuracymodel.export(format="onnx", opset=12, simplify=True, dynamic=False, imgsz=imgsz) # onnx, export the model, cannot specify dynamic=True, opencv does not support# model.export(format="torchscript", imgsz=imgsz) # libtorch# model.export(format="engine", imgsz=imgsz, dynamic=False, verbose=False, batch=1, workspace=2) # tensorrt fp32# model.export(format="engine", imgsz=imgsz, dynamic=False, verbose=False, batch=1, workspace=2, half=True) # tensorrt fp16# model.export(format="engine", imgsz=imgsz, dynamic=False, verbose=False, batch=1, workspace=2, int8=True, data=yaml) # tensorrt int8# model.export(format="openvino", imgsz=imgsz) # openvino fp32# model.export(format="openvino", imgsz=imgsz, half=True) # openvino fp16# model.export(format="openvino", imgsz=imgsz, int8=True, data=yaml) # openvino int8, INT8 export requires 'data' arg for calibrationif __name__ == "__main__":# python test_yolov8_train.py --yaml datasets/melon_new_detect/melon_new_detect.yaml --epochs 1000 --task detect --imgsz 640colorama.init(autoreset=True)args = parse_args()print("Runging on GPU") if torch.cuda.is_available() else print("Runting on CPU")train(args.task, args.yaml, args.epochs, args.imgsz)print(colorama.Fore.GREEN + "====== execution completed ======")

执行结果如下图所示：因数据集很小，训练速度很快

predict代码如下：

import colorama
import argparse
from ultralytics import YOLO
import os
import torchimport numpy as np
np.bool = np.bool_ # Fix Error: AttributeError: module 'numpy' has no attribute 'bool'. OR: downgrade numpy: pip unistall numpy; pip install numpy==1.23.1def parse_args():parser = argparse.ArgumentParser(description="YOLOv8 predict")parser.add_argument("--model", required=True, type=str, help="model file")parser.add_argument("--task", required=True, type=str, choices=["detect", "segment", "classify"], help="specify what kind of task")parser.add_argument("--dir_images", required=True, type=str, help="directory of test images")parser.add_argument("--dir_result", type=str, default="", help="directory where the image results are saved")args = parser.parse_args()return argsdef get_images(dir):# supported image formatsimg_formats = (".bmp", ".jpeg", ".jpg", ".png", ".webp")images = []for file in os.listdir(dir):if os.path.isfile(os.path.join(dir, file)):# print(file)_, extension = os.path.splitext(file)for format in img_formats:if format == extension.lower():images.append(file)breakreturn imagesdef predict(model, task, dir_images, dir_result):model = YOLO(model) # load an model, support format: *.pt, *.onnx, *.torchscript, *.engine, openvino_model# model.info() # display model information # only *.pt format supportimages = get_images(dir_images)# print("images:", images)if task == "detect" or task =="segment":os.makedirs(dir_result) #, exist_ok=True)for image in images:device = "cuda" if torch.cuda.is_available() else "cpu"results = model.predict(dir_images+"/"+image, verbose=True, device=device)# print("results:", results)if task == "detect" or task =="segment":for result in results:result.save(dir_result+"/"+image)else:print(f"class names:{results[0].names}: top5: {results[0].probs.top5}; conf:{results[0].probs.top5conf}")if __name__ == "__main__":# python test_yolov8_predict.py --model runs/detect/train10/weights/best_int8.engine --dir_images datasets/melon_new_detect/images/test --dir_result result_detect_engine_int8 --task classifycolorama.init(autoreset=True)args = parse_args()print("Runging on GPU") if torch.cuda.is_available() else print("Runting on CPU")predict(args.model, args.task, args.dir_images, args.dir_result)print(colorama.Fore.GREEN + "====== execution completed ======")

执行结果如下图所示：top1识别率100%