HandsomeGroup/作业/20200404136张文虎/Day8 20200404136张文虎/reco2.py

import torch
import torchvision.transforms as transforms
from PIL import Image
from lenet5 import Lenet5
import cv2
# 加载预训练的模型
model = Lenet5()
model.load_state_dict(torch.load('lenet5.pt'))
model.eval()


class Lenet5Recognizier:
    def __init__(self, model_file="./lenet5.pt"):
        super(Lenet5Recognizier, self).__init__()
        # 创建模型
        self.net = Lenet5()
        # cuda
        self.state = torch.load(model_file, map_location='cpu')
        self.net.load_state_dict(self.state)



    def recognizie(imgs):
        model = Lenet5()
        model.load_state_dict(torch.load('lenet5.pt'))
        model.eval()

        # 图像预处理
        transform = transforms.Compose([
            transforms.Resize((150, 150)),
            transforms.Grayscale(),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485], std=[0.229])
        ])

        # 加载要预测的图像
        image_path = imgs
        img = cv2.imread(image_path)

        # img = cv2.resize(img, (300, 200))

        img = Image.fromarray(img)
        img = transform(img)



        # 将图片转换为适合模型输入的张量形状
        img = img.unsqueeze(0)

        # 使用GPU进行预测（如果可用）
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = model.to(device)
        img = img.to(device)

        # 进行预测
        with torch.no_grad():
            output = model(img)
            probabilities = torch.softmax(output, dim=1)  # 对预测结果进行softmax操作

        # 获取预测结果
        _, predicted = torch.max(output, 1)
        class_names = ['布', '石头', '剪刀']
        predicted_class = class_names[predicted.item()]
        predicted_probabilities = probabilities[0].tolist()  # 将概率张量转换为Python列表
        return predicted_class,predicted_probabilities

# def recognizie(imgs):
#     model = Lenet5()
#     model.load_state_dict(torch.load('lenet5.pt'))
#     model.eval()
#     # 图像预处理
#     transform = transforms.Compose([
#         transforms.Resize((150, 150)),
#         transforms.Grayscale(),
#         transforms.ToTensor(),
#         transforms.Normalize(mean=[0.485], std=[0.229])
#     ])
#
#     # 加载要预测的图像
#     image_path = imgs
#     img = cv2.imread(image_path)
#
#     # img = cv2.resize(img, (300, 200))
#
#     img = Image.fromarray(img)
#     img = transform(img)
#
#
#
#     # 将图片转换为适合模型输入的张量形状
#     img = img.unsqueeze(0)
#
#     # 使用GPU进行预测（如果可用）
#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#     model = model.to(device)
#     img = img.to(device)
#
#     # 进行预测
#     with torch.no_grad():
#         output = model(img)
#         probabilities = torch.softmax(output, dim=1)  # 对预测结果进行softmax操作
#
#     # 获取预测结果
#     _, predicted = torch.max(output, 1)
#     class_names = ['布', '石头', '剪刀']
#     predicted_class = class_names[predicted.item()]
#     predicted_probabilities = probabilities[0].tolist()  # 将概率张量转换为Python列表
#     return predicted_class,predicted_probabilities
#
# print(recognizie('1.png'))