YOLOv8分类任务测试代码分享

运用于yolov8做分类任务的 test 代码：

功能介绍：对模型进行 test 并输出其测试结果，其次用柱状图实现各类预测精度值的可视化处理。

一、导入的test数据集格式为：

>test

>>>class_name1

>>>>>>1.jpg

>>>>>>2.jpg

>>>class_name2

>>>>>>1.jpg

>>>>>>2.jpg

...

二、需导的包有

import cv2
import os
from ultralytics import YOLO
from tqdm import tqdm
from matplotlib import pyplot as plt

三、以下完整代码需要修改的路径

根据自己文件存放位置修改其对应指定路径

# test数据集的路劲
file_pathname= './test'

# 模型的路径
model_path='./best.pt'

# 柱状图保存路径
output_path = './test_result_image/img_test.png'

三、完整代码

import cv2
import os
from ultralytics import YOLO
from tqdm import tqdm
from matplotlib import pyplot as plt


'''
    对做分类任务模型的测试test文件:
            传入数据集的路径 file_pathname= /dataset/val  格式为:
                >test
                >>>class_name1
                >>>class_name2
                .
                .
                .

'''

# test数据集的路劲
file_pathname= './dataset/test'

# 模型的路径
model_path='./best.pt'

#加载模型
model = YOLO(model_path)

file_dir = os.listdir(file_pathname)  # 记录各类名 
# file_dir = ['Brown rust', 'Healthy', 'Septoria', 'Yellow rust'] # 记录各类名


# 初始化记录每个标签的百分比字典 （字典推导式）
label_percentages = {label: 0.0 for label in file_dir}


print('===================')
print(file_dir)

#遍历该目录下的所有图片文件
for filename in tqdm(file_dir):

    # 初始化每个标签的数量（字典推导式）
    label_counts = {label: 0 for label in file_dir}
    print(label_counts)

    class_folder_path = os.path.join(file_pathname, filename)   # 记录各class文件夹路径
    image_filename = os.listdir(class_folder_path)  # 读取各分类的文件中的 图片文件名

    print('------------------------')

    # 遍历每一张图片 
    for img_path in image_filename:
        print('img_path: '+img_path)
        img_p = os.path.join(class_folder_path, img_path)  # 图片路径
        img = cv2.imread(img_p)  # 读取图片文件，并将图片数据转换为 OpenCV 格式

        # 进行测试
        results = model.predict(source=img, device=0, save_txt=True)

        for result in results:
            # print(result.names)
            name_dict = result.names  # 识别出的类别名称的字典。
            probs = result.probs.cpu().numpy()  # .numpy()方法通常用于将数据转换为NumPy数组

            print('name_dict: ', name_dict)
            top1_index = result.probs.top1  # 第1类的索引 （即为识别出置信度最高的类别的索引）
            top1conf_index = result.probs.top1conf  # 第1类的索引置信度
            class_name = name_dict[top1_index]  # 第1识别类型名
            print('class_name:', class_name)

            # 更新标签数量
            label_counts[class_name] += 1
           
    count = label_counts[filename]
    # 计算总图片数   len(test_path)计算当前目录下文件个数
    total_images = len(image_filename)  

    # 计算每个标签的百分比       字典推到式 {   key: value      for       key, value in iterable   }  
    # label_percentage = {label: count / total_images * 100 for label, count in label_counts.items()}
    label_percentage = count / total_images * 100
    
    # 赋值其对应标签的百分比view
    label_percentages[filename] = label_percentage

    print('---------------------------')

print(label_percentages)
# 绘制柱状图
plt.bar(label_percentages.keys(), label_percentages.values(), color=['red', 'green', 'blue', 'orange'])
# 在柱状图上显示类别名称和概率值
for i, value in enumerate(label_percentages.values()):
    plt.text(i, value, f"{value:.2f}%", ha='center', va='bottom')  # ha='center' 表示水平居中显示,  va='bottom' 表示底部居中显示。

plt.xlabel('Labels')
plt.ylabel('Percentage')
plt.title('Percentage of Each Label in Predictions')
output_path = r'./test_result_image/img_test.png'
plt.savefig(output_path)
plt.show()