pytorch调用已训练图片分类模型

import os
import shutil
import time
import torch
import torch.nn as nn
from torchvision import transforms
import torch.nn.functional as F
from PIL import Image

#指定被操作的文件夹绝对路径
mubiao_path = os.path.join('D:\\', 'mubiao')

#在目标路径下指定待分类文件夹，要分类的图都放在这里
weifen_path = os.path.join(mubiao_path, 'weifen')

# 检查并创建已分类目录
yifen_path = mubiao_path + '\yifen'
if not os.path.exists(yifen_path):
    os.mkdir(yifen_path)

# 根据模型已有的类别序号创建文件夹,这些分类文件都放在yifen_path里
class_dirs = ['1', '100']
for class_dir in class_dirs:
    class_path = os.path.join(yifen_path, class_dir)
    if not os.path.exists(class_path):
        os.mkdir(class_path)


# 定义模型结构
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)  # 第一层卷积层，输入通道数为3，输出通道数为32，卷积核大小为3x3，边缘填充1
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)  # 第二层卷积层，输入通道数为32，输出通道数为64，卷积核大小为3x3，边缘填充1
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)  # 最大池化层，池化核大小为2x2，步长为2
        self.fc1 = nn.Linear(64 * 56 * 56, 128)  # 第一层全连接层，输入大小为64*56*56，输出大小为128
        self.fc2 = nn.Linear(128, 2)  # 第二层全连接层，输入大小为128，输出大小为2

    def forward(self, x):
        # 第一层卷积
        x = self.pool(torch.relu(self.conv1(x)))
        # 第二层卷积
        x = self.pool(torch.relu(self.conv2(x)))
        # 全连接层
        x = x.view(-1, 64 * 56 * 56)  # 将特征图展开成一维向量
        x = torch.relu(self.fc1(x))   # 激活函数ReLU
        x = F.softmax(self.fc2(x), dim=1)  # 添加可信度，相似度
        return x

# 加载模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  #如果支持cuda
model = Net().to(device)  #模型转移到cuda
model.load_state_dict(torch.load('fenlei_rmb_1_10.pth', map_location=device))  #加载模型参数，并将其映射到指定cuda
model.eval()  #将模型设置为评估模式，这意味着模型将不再进行梯度计算,所以在实际运用中都要讲模型设置为该模式

# 定义图像预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),   # 将图像大小重置为 224x224 像素
    transforms.ToTensor(),  # 将图像转换为 PyTorch 的张量格式
    transforms.Normalize(mean=[0.5], std=[0.5])   # 将图像像素值归一化到 [-1,1] 范围内
])

#插入一个时间登记
date_start = time.process_time()

# 开始批量处理图像
for filename in os.listdir(weifen_path):
    # 构造图像路径
    filepath = os.path.join(weifen_path, filename)

    # 判断文件是否为图像文件，这里指定要分类的格式
    if not filename.endswith(('.jpg', '.jpeg', '.png', '.bmp','.webp','.gif')):
        continue

    # 打开图像
    image = Image.open(filepath)

    # 将图像转换为张量，图像只有三个维度，因为要符合运算规范所以要转成四维
    image_tensor = transform(image).unsqueeze(0).to(device)  #unsqueeze(0)的作用是在张量的第一个维度前再插入一个空维度

    # 进行分类
    with torch.no_grad():  #禁用梯度计算上下文，该上下文下的计算不会影响到梯度
        outputs = model(image_tensor)  #将得到的图像张量进行前向传递计算，得到模型的输出
        _, predicted = torch.max(outputs.data, 1)  #这里predicted就已经是模型里面的结果，也即是识别的类型序号，但获得的是一个张量类型

    predicted = predicted.item()  #将前分类张量转为序号
    probabilities = torch.nn.functional.softmax(outputs[0], dim=0)  #获得分类的可信度

    if probabilities[predicted] > 0.7:   #如果可信度大于0.7则执行下面的操作
        # 移动图像到分类后的目录
        class_dir = class_dirs[predicted]  # 获得分类序号
        output_path = os.path.join(yifen_path, class_dir, filename)  # 根据分类序号获得绝对当前文件绝对路径
        shutil.move(filepath, output_path)  # 将当前文件路径移动到目标路径


#插入一个结束时间登记
end_date = time.process_time()

#输出总用时时间
print(f'总用时:{end_date - date_start}秒')