PyTorch图像分类全流程实战--预训练模型预测图像分类02

主要内容

今天的任务是学习预训练模型的使用，模型是Resnet18，使用的torchvision包由流行的数据集、模型体系结构和通用的计算机视觉图像转换组成。简单地说就是常用数据集+常见模型+常见图像增强方法。步骤包括：载入预训练模型，图像预处理（缩放裁剪、转 Tensor、归一化），执行前向预测，预测结果分析（得到各类别的置信度，解析类别，载入标签，可视化分类结果）。

载入需要的库

import os

import cv2

import pandas as pd
import numpy as np

import torch

import matplotlib.pyplot as plt
%matplotlib inline
#图像预处理
from torchvision import transforms
# 用 pillow 载入
from PIL import Image
import torch.nn.functional as F

GPU

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

载入预训练图像分类模型

model = models.resnet18(pretrained=True)
model = model.eval()#评估模式
model = model.to(device)

model.eval()
不启用 BatchNormalization 和 Dropout，保证BN和dropout不发生变化，pytorch框架会自动把BN和Dropout固定住，不会取平均，而是用训练好的值
【pytorch系列】model.eval()用法详解

图像预处理

# 测试集图像预处理-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])
                                    ])

Transforms,数据在读取到pytorch之后通常都需要对数据进行预处理，包括尺寸缩放、转换张量、数据中心化或标准化等等，这些操作都是通过transforms进行的数据在读取到pytorch之后通常都需要对数据进行预处理，包括尺寸缩放、转换张量、数据中心化或标准化等等，if 这些操作都是通过transforms进行的。
python transforms_2.2 图像预处理——transforms（笔记）

#载入一张测试图像
img_path = 'test_img/7.jpeg'
# 用 pillow 载入
from PIL import Image
img_pil = Image.open(img_path)
print(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)

torch.unsqueeze(input, dim, out=None) 增维，dim=0：最粗粒度的方向，在第1维插入一个维度
Python-squeeze()、unsqueeze()函数的理解
所以input_img.shape由torch.Size([3, 224, 224])->torch.Size([1, 3, 224, 224])

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

torch.nn.functional
torch.nn.functional.x 为函数,与torch.nn不同, torch.nn.x中包含了初始化需要的参数等 attributes 而torch.nn.functional.x则需要把相应的weights 作为输入参数传递,才能完成运算, 所以用torch.nn.functional创建模型时需要创建并初始化相应参数.
softmax
pytorch中nn.functional()学习总结

#预测结果分析
#各类别置信度柱状图
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)

# 解析出类别
pred_ids = top_n[1].cpu().detach().numpy().squeeze()
# 解析出置信度
confs = top_n[0].cpu().detach().numpy().squeeze()

#载入ImageNet 1000图像分类标签
df = pd.read_csv('imagenet_class_index.csv')
idx_to_labels = {}
for idx, row in df.iterrows():
    idx_to_labels[row['ID']] = [row['wordnet'], row['class']]

df.iterrows DataFrame行作为(index, Series)对

# 用 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) # 展示预测结果表格

参考文献：

【1】pytorch——torchvision预训练模型与数据集全览
【2】【pytorch系列】model.eval()用法详解
【3】python transforms_2.2 图像预处理——transforms（笔记）
【4】Python-squeeze()、unsqueeze()函数的理解
【5】pytorch中nn.functional()学习总结