python
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}秒')