ImageNet预训练图像分类模型预测单张图像

导入基础工具包

import os

import cv2

import pandas as pd
import numpy as np

import torch

import matplotlib.pyplot as plt
%matplotlib inline

计算设备确定

# 有 GPU 就用 GPU，没有就用 CPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

载入预训练模型

from torchvision import models
# 载入预训练图像分类模型

model = models.resnet18(pretrained=True) 

# model = models.resnet152(pretrained=True)
model = model.eval() #将模型设为eval
model = model.to(device)

图像预处理，比较固定的四个部分，其他分类任务也可以用。

四步：

1. 缩放裁剪
2. 中心获取
3. 转为Tensor
4. 归一化处理：更近似于正态分布，易于神经网络处理。mean、std这六个数也是通用的。

from torchvision import transforms

# 测试集图像预处理-RCTN：缩放裁剪、转 Tensor、归一化
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_path = 'test_img/banana1.jpg'
# img_path = 'test_img/husky1.jpeg'
img_path = 'test_img/basketball_shoe.jpeg'

# img_path = 'test_img/cat_dog.jpg'


# 用 pillow 载入
from PIL import Image
img_pil = Image.open(img_path)

执行图像分类预测：

input_img = test_transform(img_pil) # 预处理，将图片传入图片与处理的函数

 转换模型所需要的维度：

input_img = input_img.unsqueeze(0).to(device)
input_img.shape

运行后为：

torch.Size([1, 3, 224, 224])，即一张3通道224*224的图片

执行前向预测：
 

# 执行前向预测，得到所有类别的 logit 预测分数
pred_logits = model(input_img) 
pred_logits.shape

结果为：

torch.Size([1, 1000])

利用softmax对分数大小进行比较：

import torch.nn.functional as F
pred_softmax = F.softmax(pred_logits, dim=1) # 对 logit 分数做 softmax 运算
pred_softmax.shape

预测结果分析

对softmax结果画一个柱状图：

plt.figure(figsize=(8,4))

x = range(1000)
y = pred_softmax.cpu().detach().numpy()[0]

ax = plt.bar(x, y, alpha=0.5, width=0.3, color='yellow', edgecolor='red', lw=3)
plt.ylim([0, 1.0]) # y轴取值范围
# plt.bar_label(ax, fmt='%.2f', fontsize=15) # 置信度数值

plt.xlabel('Class', fontsize=20)
plt.ylabel('Confidence', fontsize=20)
plt.tick_params(labelsize=16) # 坐标文字大小
plt.title(img_path, fontsize=25)

plt.show()

取置信度最大的n个结果：

n = 10
top_n = torch.topk(pred_softmax, n)
top_n

out：

torch.return_types.topk(
values=tensor([[0.5988, 0.3556, 0.0064, 0.0047, 0.0041, 0.0041, 0.0037, 0.0025, 0.0022,
         0.0022]], device='cuda:0', grad_fn=<TopkBackward0>),
indices=tensor([[430, 514, 522, 630, 502, 770, 427, 768, 805,  35]], device='cuda:0'))

解析出类别：

# 解析出类别
pred_ids = top_n[1].cpu().detach().numpy().squeeze()
pred_ids

out：

array([430, 514, 522, 630, 502, 770, 427, 768, 805,  35])

如何知道430、514是哪一类？

df = pd.read_csv('imagenet_class_index.csv')

将分类结果写在原图上：

# 用 opencv 载入原图
img_bgr = cv2.imread(img_path)


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, (25, 50 + 40 * i), cv2.FONT_HERSHEY_SIMPLEX, 1.25, (0, 0, 255), 3)

# 保存图像
cv2.imwrite('output/img_pred.jpg', img_bgr)

# 载入预测结果图像
img_pred = Image.open('output/img_pred.jpg')
img_pred

 

预测结果用表格输出：

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 = idx_to_labels[pred_ids[i]][0] # 获取 WordNet
    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) # 展示预测结果表格