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()))