一个可以用单张图片测试网络分类性能的函数

import torch
from torch import nn
from PIL import Image
from torchvision import transforms, datasets
import cv2
import os
import shutil

class_names = ['0', '1']  # 这个顺序很重要，要和训练时候的类名顺序一致
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
def label_predict(model_dir,picture_dir):
    ##载入模型并读取权重
    model = torch.load(model_dir,map_location='cuda:0')
    model.to(device)
    model.eval()
    img_path = picture_dir
    img=cv2.imread(img_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    transform = transforms.Compose(
        [
         transforms.ToTensor(),
         transforms.Normalize(mean, std)
         ])
    img = transform(img).cuda()
    img = img.unsqueeze(0)
    outputs = model(img)
    print(outputs)
    # 输出概率最大的类别
    _, indices = torch.max(outputs, 1)
    result = class_names[indices]
    return result
class_num=2
src_folder=r'../data_G/data-new4-U-D-smaller-224/test'
dir_folder=r'../data_G/data-new4-U-D-smaller-224/after-train-test'
model_dir=r'./train_model/test-224/test-U-D/net_better.pkl'
for i in range(class_num):
    if os.path.exists(os.path.join(dir_folder,str(i))) is not True:
        os.makedirs(os.path.join(dir_folder,str(i)))
    if os.path.exists(os.path.join(dir_folder,str(i))) is not True:
        os.makedirs(os.path.join(dir_folder,str(i)))
for fi in os.listdir(src_folder):
    for filename in os.listdir(os.path.join(src_folder,fi)):
        if filename.split('.')[-1].upper() in ("JPG", "JPEG", "PNG", "BMP", "GIF"):
            print(os.path.join(src_folder,fi,filename))
            img_dir=os.path.join(src_folder,fi,filename)
            label=label_predict(model_dir,img_dir)
            shutil.copyfile(os.path.join(src_folder,fi,filename), os.path.join(dir_folder,fi,label+'_'+filename+'.jpg'))
            #os.rename(os.path.join(src_folder,fi,filename), os.path.join(dir_folder,fi,label+'.jpg'))
            #cv2.imwrite(os.path.join(dir_folder,fi,file,filename), result)