当前位置：首页 > article >正文

YOLOv8 classify介绍

article 2024/11/16 17:40:11

图像分类器(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 torch

def 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 args

def 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 method
	elif task == "segment":
		model = YOLO("yolov8n-seg.pt") # load a pretrained model, should be a *.pt PyTorch model to run this method
	elif task == "classify":
		model = YOLO("yolov8n-cls.pt") # n/s/m/l/x
	else:
		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 EarlyStopping
	results = 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.yaml

	metrics = model.val() # It'll automatically evaluate the data you trained, no arguments needed, dataset and settings remembered
	if task == "classify":
		print("Top-1 Accuracy:", metrics.top1) # top1 accuracy
		print("Top-5 Accuracy:", metrics.top5) # top5 accuracy

	model.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 calibration


if __name__ == "__main__":
	# python test_yolov8_train.py --yaml datasets/melon_new_detect/melon_new_detect.yaml --epochs 1000 --task detect --imgsz 640
	colorama.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 torch

import 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.1

def 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 args

def get_images(dir):
	# supported image formats
	img_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)
					break

	return images

def 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 support

	images = 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 classify
	colorama.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%