做一个图像分类任务（二）预测单张图像-英文_单张图片测试给出类别代码-CSDN博客

本文链接：https://blog.csdn.net/codeplayme/article/details/138002230

文章目录

建立目录

用来存放测试的图片与结果

# 建立目录
if os.path.exists('test_img'):
    print('{} is exited'.format('test_img'))
else:
    # 存放测试图片
    os.mkdir('test_img')

if os.path.exists('output'):
    print('{} is exited'.format('output'))
else:
    # 存放结果文件
    os.mkdir('output')

代码

使用pytorch自带的 resnet18 模型来预测单张图片的类别，其流程如下：

载入数据集
对图像进行预处理
模型预测
softmax输出结果
对结果进行处理

# coding:gbk
import os
from tkinter import N
from tkinter.tix import InputOnly
from tkinter.ttk import LabeledScale
from tokenize import Ignore
import cv2
import pandas as pd
import numpy as np
import torch
import matplotlib.pyplot as plt
from torchvision import models
from torchvision import transforms
from PIL import Image
import torch.nn.functional as F
import matplotlib
from IPython.display import display
# 预训练图像分类模型预测单张图像-英文
def predict_single_pic_en(img_path,output_path,csv_path):
    # 使用GPU
    device=torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    print('device:',device)
    model=models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
    # dir(models)# 可以看到所有的模型
    model=model.eval()
    model=model.to(device)
    # 图像预处理：测试集图像预处理-RCTN：缩放裁剪、转Tensor、归一化 只接收pillow格式的图像
    test_transform=transforms.Compose([transforms.Resize(256),
                                    transforms.CenterCrop(224),
                                    transforms.ToTensor(),
                                    transforms.Normalize(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225]) ])
    # 载入测试集
    img_pil=Image.open(img_path)
    # print(img_pil)
    # print(np.array(img_pil).shape)

    # 对图像预处理
    input_img=test_transform(img_pil)
    print(input_img.shape)
    input_img=input_img.unsqueeze(0).to(device)
    print(input_img.shape)
    # 前向传播，得到所有类别的logit预测分数
    pred_logits=model(input_img)
    # print(pred_logits)
    pred_softmax=F.softmax(pred_logits,dim=1)
    # print(pred_softmax)
    # 对预测结果进行分析
    # 绘制置信图
    confidence_map(pred_softmax,img_path)
    # 取预测结果前10的
    n=10
    show_result_on_pic(pred_softmax,img_path,output_path,csv_path,n)
    # 图片和柱状图一起显示
    show_result_with_bar(output_path,csv_path,pred_softmax)
    
# 测试代码
img_path='test_img/02.webp'
output_path='output/img_pred.jpg'
csv_path='imagenet_class_index.csv'
predict_single_pic_en(img_path,output_path,csv_path)

对预测结果绘制置信图

# 绘制置信图
def confidence_map(pred_softmax,img_path):
    # 中文字体
    matplotlib.rc('font',family='SimHei')
    # 用来正常显示负号
    plt.rcParams['axes.unicode_minus']=False
    x=range(1000)
    y=pred_softmax.cpu().detach().numpy()[0]
    # 绘制各个类别置信度的柱状图
    plt.figure(figsize=(8,4))
    plt.bar(x,y,alpha=0.5,width=0.3,color='yellow',edgecolor='red',lw=3)
    # y轴取值范围
    plt.ylim([0,1.0])
    plt.xlabel('Class',fontsize=20)
    plt.ylabel('Confidence',fontsize=20)
    # 坐标文字大小
    plt.tick_params(labelsize=16)
    plt.title(img_path,fontsize=25)
    plt.show()

输出结果：
在这里插入图片描述

把预测结果写在原图和表格上

top_n=torch.topk(pred_softmax,n)：取tensor的topk元素（降序后的前n个元素值以及索引）
pred_ids=top_n[1].cpu().detach().numpy().squeeze()：
cpu()：将数据迁移到CPU中；
detach()：阻断反向传播；
numpy()：将tensor转为numpy数据；
squeeze()：将数组的形状中删除单维度条目，即把shape中维度为1的删掉。

# 把结果写在图像上
def show_result_on_pic(pred_softmax,img_path,output_path,csv_path,n=10):
    # 取置信度最大的n个结果
    top_n=torch.topk(pred_softmax,n)
    # 解析出类别
    pred_ids=top_n[1].cpu().detach().numpy().squeeze()
    # 解析置信度
    confs=top_n[0].cpu().detach().numpy().squeeze()
    # 载入imageNet1000图像分类标签
    df=pd.read_csv(csv_path)
    idx_to_labels={}
    for _,row in df.iterrows():
        idx_to_labels[row['ID']]=[row['wordnet'],row['class']]

    # 写到表格中
    pred_excel(idx_to_labels,pred_ids,confs,n)

    # 把分类的结果写在原图上
    # img_bar=cv2.imread(img_path)
    # 解决中文路径的问题
    img_bgr = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8),-1)
    for i in range(n):
        # 获取类别名称
        class_name=idx_to_labels[pred_ids[i]][1]
        # 获取置信度
        confidence=confs[i]*100
        text='{:<15}{:>.4f}'.format(class_name,confidence)
        # print(text)
        # 图片，文字，左上角坐标，字体，bgr颜色，线宽
        img_bgr=cv2.putText(img_bgr,text,(24,50+40*i),cv2.FONT_HERSHEY_SIMPLEX,1.25,(0,0,255),2)

    cv2.imwrite(output_path,img_bgr)
    # 读取
    img=cv2.imread(output_path)
    cv2.imshow('Image',img)
    cv2.waitKey(5000)
    cv2.destroyAllWindows()

# 图像和柱状图一起显示
def show_result_with_bar(output_path,csv_path,pred_softmax):
    # 中文字体
    matplotlib.rc('font',family='SimHei')
    # 用来正常显示负号
    plt.rcParams['axes.unicode_minus']=False
    fig=plt.figure(figsize=(18,6))
    # 绘制左图-预测图
    ax1=plt.subplot(1,2,1)
    # 解决中文路径的问题
    img_pred_bgr = cv2.imdecode(np.fromfile(output_path, dtype=np.uint8),-1)
    img_pred=cv2.cvtColor(img_pred_bgr,cv2.COLOR_RGB2BGR)
    ax1.imshow(img_pred)
    ax1.axis('off')

    # 绘制右图-柱状图
    ax2=plt.subplot(1,2,2)
    # 载入imageNet1000图像分类标签
    df=pd.read_csv(csv_path)
    x=df['ID']
    y=pred_softmax.cpu().detach().numpy()[0]
    ax2.bar(x,y,alpha=0.5,width=0.3,color='yellow',edgecolor='red',lw=3)
    # 设置y轴的取值范围
    plt.ylim([0,1.0])
    plt.title('{} Classification'.format(img_path),fontsize=30)
    plt.xlabel('Class',fontsize=20)
    plt.ylabel('Confidence',fontsize=20)
    # 设置坐标文字大小
    ax2.tick_params(labelsize=16)
    # 布局调整器，能够自动调整子图的参数，使得
    plt.tight_layout()
    plt.show()
    fig.savefig('output/预测图+柱状图.jpg')

# 预测结果表格输出
def pred_excel(idx_to_labels,pred_ids,confs,n=10):
    # 预测结果表格
    pred_df=pd.DataFrame()
    for i in range(n):
        # 类别名称
        class_name=idx_to_labels[pred_ids[i]][1]
        # 类别号
        label_idx=int(pred_ids[i])
        # WordNet
        wordnet=idx_to_labels[pred_ids[i]][0]
        # 置信度
        confidence=confs[i]*100
        # 添加一行
        pred_df=pred_df.append({'Class':class_name,'Class_ID':label_idx,'Confidence(%)':confidence,'WordNet':wordnet},ignore_index=True)
    # 展示预测表格
    display(pred_df)

输出结果：
在这里插入图片描述

在这里插入图片描述

这里的代码与github上可能有一丢丢不同，大家想看全面点还是要看github上原版的代码以及原作者的视频讲解，博主这里只是博主按需学习的一些经验哦！
声明：这里只是做一个图像分类任务，不作为商业用途，不涉及任何利益交易，甚至不作为毕业设计。侵权可以删哦！

代码出处：github地址
视频讲解：b站