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='13290033_ebd7c7abba_n.jpg'
# 加载图片
img=Image.open(img_path)
# 将图片进行处理，返回的是tensor
img=data_transform(img)
# 将数据升维，图片是三维，而训练时是4维，因为训练时第一个维度为每个批次的训练数据大小
# 本次预测一个图片，就要升维，变成1，表明该批次图片有1个
img=torch.unsqueeze(img,dim=0)
# 读取预测结果和真实分类的映射
json_path='./class_indices.json'
json_file=open(json_path,'r')
# 加载成字典 0:'A',1:'B',2:'C'
class_indict=json.load(json_file)
# 创建网络
model=resnet34(num_classes=5).to(device)
# 加载模型训练好的参数
weitcht_path='./resNet34.pth'
model.load_state_dict(torch.load(weitcht_path,map_location=device))
# 开启验证模式
model.eval()
# 不需要求导
with torch.no_grad():
    # 每个数据对应输出，(1,5)维度，将其降维，直接是(5,)
    output=torch.squeeze(model(img.to(device))).cpu()
    # 如果是预测，完全可以用touch.max,返回最大值索引，但是下面为了输出预测的概率，就要将其标准化，概率和为1
    # 按道理说应该是dim=1，每行的所有列，但是现在不是二维，所以要用0
    predict=torch.softmax(output,dim=0)
    # 返回最大值索引，dim=0和上面的同理
    predict_cla=torch.argmax(predict,dim=0).item()
print_res = "class: {}   prob: {:.3}".format(class_indict[str(predict_cla)],
                                                 predict[predict_cla].numpy())
print(print_res)