指定设备
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
工具包
import os
import cv2
from PIL import Image, ImageFont, ImageDraw
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
%matplotlib inline
import torch
from torchvision import models
import torch.nn.functional as F
from torchvision import transforms
加载预训练模型
从torchvision导入的models中包含了许多预训练模型
model = models.resnet18(pretrained=True)
model = model.eval()
model = model.to(device)
图像预处理
transforms从torvhvison导入
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])
])
导入图像
使用pillow
img_pil = Image.open(img_path) # 用 pillow 载入
执行预测
unsqueeze增加一个维度后才可传入模型。
可以看到这里的test_transform是一个对象,也就是经过上述compose后的返回值。
input_img = test_transform(img_pil).unsqueeze(0).to(device) # 预处理
在模型输出得到类别的分数,由于需要置信度,需要使用softmax归一化,结果是所有类别的和为1,并且得到了置信度。
pred_softmax = F.softmax(pred_logits, dim=1) # 对 logit 分数做 softmax 运算
可视化
关键是y传入时的调用
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轴取值范围
获取top k结果,注意类别与置信度的解析
n = 10的
top_n = torch.topk(pred_softmax, n) # 取置信度最大的 n 个结果
pred_ids = top_n[1].cpu().detach().numpy().squeeze() # 解析出类别
confs = top_n[0].cpu().detach().numpy().squeeze() # 解析出置信度