import os
import json
import torch
from torchvision import transforms
from PIL import Image
# 加载运算设备
device=torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# 数据处理
data_transform = transforms.Compose(
[transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
# 因为是同时预测多个图片,所以要将路径添加到列表中
img_path_list=['13290033_ebd7c7abba_n.jpg','302782756_d35cb3e468.jpg','40411019_526f3fc8d9_m.jpg']
img_list=[]
# 将转化后的数据添加到list中
for img_path in img_path_list:
img=Image.open(img_path)
img=data_transform(img)
img_list.append(img)
# 将列表转化成tensor,因为列表中已经存在3个样本,所以直接就升维了,不用unsqueeze
batch_img=torch.stack(img_list,dim=0)
# 加载预测结果与真实分类的映射
json_path='./class_indices.json'
json_file=open(json_path,'r')
class_indict=json.load(json_file)
# 构建网络
model=resnet34(num_classes=5).to(device)
# 加载训练模型参数
weights_path='./resNet34.pth'
model.load_state_dict(torch.load(weights_path,map_location=device))
# 开启验证模式
model.eval()
# 不需要进行求导
with torch.no_grad():
output=model(batch_img.to(device)).cpu()
# 这里的dim维度就是1,因为此时预测的输出结果为二维,dim=1代表,每一行所有列进行计算
# 输出维度(3,5)
predict=torch.softmax(output,dim=1)
probs,classes=torch.max(predict,dim=1)
for idx, (pro, cla) in enumerate(zip(probs, classes)):
print("image: {} class: {} prob: {:.3}".format(img_path_list[idx],
class_indict[str(cla.item())],
pro.item()))