pytorch_图片分类_net.zip

import torch as t import torchvision as tv from PIL import Image import matplotlib.pyplot as plt from torch.autograd import Variable import numpy as np from net import Net bCuda = t.cuda.is_available() # 是否开启 GPU bCuda = False # 不启用GPU 我的电脑不支持 device = t.device("cuda:0" if bCuda else "cpu") img_size = 32 # 图片大小，可以改 # 对Tensor进行变换 颜色转换 mean=给定均值：(R,G,B) std=方差：（R，G，B） normalize = tv.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform = tv.transforms.Compose( [tv.transforms.Resize([img_size, img_size]), tv.transforms.CenterCrop([img_size, img_size]), tv.transforms.ToTensor(), normalize]) # 分类数组 classes = ['凹下标志-0', '凸上标志-1', '打滑标志-2', '左弯标志-3', '右弯标志-4', '连续转弯标志-5', '00020-6', '00021-7', '00022-8', '00023-9'] # 显示图片方法 def imshow(img): plt.imshow(img) plt.show() # 单张图片调用 def prediect(model, img_path, imgType, isShowSoftmax=False, isShowImg=False): t.no_grad() image_PIL = Image.open(img_path) # imshow(image_PIL) image_tensor = transform(image_PIL) # 以下语句等效于 img = torch.unsqueeze(image_tensor, 0) image_tensor.unsqueeze_(0) # 没有这句话会报错 image_tensor = image_tensor.to(device) out = model(image_tensor) # 得到预测结果，并且从大到小排序 _, indices = t.sort(out, descending=True) # 返回每个预测值的百分数 percentage = t.nn.functional.softmax(out, dim=1)[0] * 100 # 是否显示每个分类的预测值 item = indices[0] if isShowSoftmax: for idx in item: ss = percentage[idx] value = ss.item(); name = classes[idx] print('名称：', name, '预测值：', value) # 预测最大值 _, predicted = t.max(out.data, 1) maxPredicted = classes[predicted.item()] maxAccuracy = percentage[item[0]].item() if imgType == maxPredicted: print('预测正确，预测结果：', maxPredicted, '预测值：', maxAccuracy) else: print('预测错误，正确结果：', imgType, '，预测结果：', maxPredicted, '预测值：', maxAccuracy, '图片：', img_path) if isShowImg: plt.imshow(image_PIL) plt.show() # 测试集 def loadtestdata(): path = "./imageData/test/" testset = tv.datasets.ImageFolder(path, transform=transform) testloader = t.utils.data.DataLoader(testset, batch_size=40, shuffle=True, num_workers=6) return testloader # 测试全部 def testAll(model): testloader = loadtestdata() dataiter = iter(testloader) images, labels = dataiter.next() print(labels) print('真实值: ' , " ".join('%5s' % classes[labels[j]] for j in range(25))) # 打印前25个GT（test集里图片的标签） outputs = model(Variable(images)) _, predicted = t.max(outputs.data, 1) print('预测值: ', " ".join('%5s' % classes[predicted[j]] for j in range(25))) # 打印前25个预测值 imshow2(tv.utils.make_grid(images, nrow=5)) # nrow是每行显示的图片数量，缺省值为8 def imshow2(img): img = img / 2 + 0.5 # unnormalize npimg = img.numpy() plt.imshow(np.transpose(npimg, (1, 2, 0))) plt.show() if __name__ == '__main__': # 直接加载 model = t.load('model/10_net.pkl') # 加载2 ，看官方的解释 # model = Net() # 10 分类数量 # load_weights = t.load('model/10_net_params.pkl', map_location='cpu') # model.load_state_dict(load_weights) model = model.to(device) # GPU model.eval() # 运行模式 # 测试全部图片 #testAll(model) # 测试一张图片 # # 凹下标志-0 # prediect(model,'imageData/test/00000/01160_00000.png', classes[0], False, False) # prediect(model,'imageData/test/00000/01160_00001.png', classes[0], False, False) # prediect(model,'imageData/test/00000/01160_00002.png', classes[0], False, False) # prediect(model,'imageData/test/00000/01798_00000.png', classes[0], False, False) # prediect(model,'imageData/test/00000/01798_00001.png', classes[0], False, False) # prediect(model,'imageData/test/00000/01798_00002.png', classes[0], False, False) # # # 凸上标志-1 # prediect(model,'imageData/test/00001/00029_00000.png', classes[1], False, False) # prediect(model,'imageData/test/00001/00029_00001.png', classes[1], False, False) # prediect(model,'imageData/test/00001/00029_00002.png', classes[1], False, False) # prediect(model,'imageData/test/00001/00079_00000.png', classes[1], False, False) # prediect(model,'imageData/test/00001/00079_00002.png', classes[1], False, False) # prediect(model,'imageData/test/00001/00079_00001.png', classes[1], False, False) # # # 打滑标志-2 # prediect(model,'imageData/test/00002/01503_00000.png', classes[2], False, False) # prediect(model,'imageData/test/00002/01503_00001.png', classes[2], False, False) # prediect(model,'imageData/test/00002/01503_00002.png', classes[2], False, False) # prediect(model,'imageData/test/00002/01515_00000.png', classes[2], False, False) # prediect(model,'imageData/test/00002/01515_00001.png', classes[2], False, False) # prediect(model,'imageData/test/00002/01515_00002.png', classes[2], False, False) # # # 左弯标志-3 # prediect(model,'imageData/test/00003/00207_00000.png', classes[3], False, False) # prediect(model,'imageData/test/00003/00207_00001.png', classes[3], False, False) # prediect(model,'imageData/test/00003/00207_00002.png', classes[3], False, False) # prediect(model,'imageData/test/00003/00211_00000.png', classes[3], False, False) # prediect(model,'imageData/test/00003/00211_00001.png', classes[3], False, False) # prediect(model,'imageData/test/00003/00211_00002.png', classes[3], False, False) # prediect(model,'imageData/test/00003/02664_00000.png', classes[3], False, False) # prediect(model,'imageData/test/00003/02664_00001.png', classes[3], False, False) # prediect(model,'imageData/test/00003/02664_00002.png', classes[3], False, False) # # # 右弯标志-4 # prediect(model,'imageData/test/00004/00214_00000.png', classes[4], False, False) # prediect(model,'imageData/test/00004/00214_00001.png', classes[4], False, False) # prediect(model,'imageData/test/00004/00214_00002.png', classes[4], False, False) # prediect(model,'imageData/test/00004/00282_00000.png', classes[4], False, False) # prediect(model,'imageData/test/00004/00282_00001.png', classes[4], False, False) # prediect(model,'imageData/test/00004/00282_00002.png', classes[4], False, False) # prediect(model,'imageData/test/00004/02567_00000.png', classes[4], False, False) # prediect(model,'imageData/test/00004/02567_00001.png', classes[4], False, False) # prediect(model,'imageData/test/00004/02567_00002.png', classes[4], False, False) # prediect(model,'imageData/test/00004/02660_00000.png', classes[4], False, False) # prediect(model,'imageData/test/00004/02660_00001.png', classes[4], False, False) # prediect(model,'imageData/test/00004/02660_00002.png', classes[4], False, False) # # # 连续转弯标志-5 # prediect(model,'imageData/test/00005/00575_00000.png', classes[5], False, False) # prediect(model,'imageData/test/00005/00575_00001.png', classes[5], False, False) # prediect(model,'imageData/test/00005/00575_00002.png', classes[5], False, False) # prediect(model,'imageData/test/00005/01893_00000.png', classes[5], False, False) # prediect(model,'imageData/test/00005/01893_00001.png', classes[5], False, False) # prediect(model,'imageData/te