python与机器学习（七）下——torchvision预训练模型测试真实图像分类

任务要求：利用torchvision中的预训练CNN模型来对真实的图像进行分类，预测每张图片的top5类别。
数据: real_image, class_index.json

导入：

import torch
from torchvision import models, datasets, transforms
from torch.utils.data import DataLoader, Dataset
from PIL import Image
import numpy as np
import torchvision

import os
import json
import time
import matplotlib.pyplot as plt
%matplotlib inline

1. 类别索引：

构建类别索引词典

f = open('./data/class_index.json')
class_index = json.load(f)
print('class num:', len(class_index))
class_dict = {int(k): v[1] for k, v in class_index.items()}
print(class_dict)

2. 预训练模型：

加载预训练CNN模型

import ssl
alexnet = models.alexnet(pretrained=True)
# import ssl
# resnet = models.resnet50(pretrained=True)

3. 图像预处理：

图像缩放、裁剪、转Tensor、归一化

image_transforms = transforms.Compose([
    transforms.Resize(256),                    
    transforms.CenterCrop(224),                
    transforms.ToTensor(),                     
    transforms.Normalize(                      
    mean=[0.485, 0.456, 0.406],                
    std=[0.229, 0.224, 0.225]                  
    )  
])

4. 测试数据集加载：

构建测试数据集，迭代返回预处理后的Tensor格式图像和原始图像

class TestDataset():
    def __init__(self, root, transforms=None):
        imgs = os.listdir(root)
        self.imgs = [os.path.join(root, img) for img in imgs]
        self.transforms = transforms
        
    def __getitem__(self, index):
        img_path = self.imgs[index]
        img_pil = Image.open(img_path)
        label = None
        img_np = np.asarray(img_pil)
        data = self.transforms(img_pil)
        return data, img_np
    
    def __len__(self):
        return len(self.imgs)

test_dir = './data/real_image/'
test_dataset = TestDataset(test_dir, image_transforms)
print('test image num:', test_dataset.__len__())

运行结果如下：

test image num: 20

5. 模型预测图像类别：

在测试模式下，对于每张图片显示原始图像，并输出模型预测的top5类别及top1类别

alexnet.eval()
for data, img_np in test_dataset:
    img = torch.unsqueeze(data, 0)
    output = alexnet(img)
    _, index = torch.max(output, 1)
    index=index.numpy()
    percentage = torch.nn.functional.softmax(output, dim=1)[0] * 100
    plt.imshow(img_np, aspect='auto')
    plt.show()
    print('top1类别:')
    print(class_dict[index[0]], percentage[index[0]].item())
    _, indices = torch.sort(output, descending=True)
    indices=indices.numpy()
    print('top5类别:')
    for idx in indices[0][:5]:
        print((class_dict[idx], percentage[idx].item()))
    print()

运行结果如下（部分）：

……