利用AlexNet模型进行图片分类

目录

1.加载预训练好的AlexNet模型

2.定义图片文件路径列表

3.加载ImageNet的类别标签文件

4.图像预处理

5.利用预训练好的AlexNet模型进行图像分类

6.获取预测结果的类别名称  

---

1.加载预训练好的AlexNet模型

# 加载预训练好的AlexNet模型
alexnet = models.alexnet(pretrained=True)

2.定义图片文件路径列表

# 定义图片文件路径列表
image_files = ['image1.jpg', 'image2.jpg', 'image3.jpg']

3.加载ImageNet的类别标签文件

# 加载ImageNet的类别标签文件
labels_url = "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt"
class_labels = requests.get(labels_url).text.split('\n')

class_labels是切分好的种类名称

4.图像预处理

    transform = 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])
    ])
    input_image = transform(image).unsqueeze(0)

• transform = transforms.Compose([...])：transforms.Compose是一个将多个转换操作组合起来的类，可以将多个数据转换操作依次执行。在这里，我们定义了一个包含多个数据转换操作的列表。

• transforms.Resize(256)：这个操作将输入的图像调整大小为 256x256 像素。这里指定了一个整数值，表示调整后的图像的宽度和高度都将被调整为 256 像素。

• transforms.CenterCrop(224)：这个操作在图像中心裁剪一个 224x224 大小的区域。通过这一步，我们将图像裁剪为模型所需的输入大小。

• transforms.ToTensor()：这个操作将 PIL Image 或者 ndarray 数据类型的图片转换为 Tensor。在神经网络中，一般需要将图像转换成 Tensor 格式，以便进行后续的计算。

• transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])：这个操作对图像进行标准化处理，将图像的每个通道数值减去均值（mean）再除以标准差（std）。这有助于将输入数据归一化到一个较小的范围，有利于模型的训练和收敛。

• input_image = transform(image).unsqueeze(0):

  在深度学习中，对于图像数据，通常会使用三维张量来表示，即（通道数，高度，宽度）。这里的通道数是指图像的颜色通道，一般为 RGB 彩色图像时通道数为 3，灰度图像时通道数为 1。高度和宽度分别代表图像的像素高度和宽度。

  当我们将图像进行预处理后，得到的结果是一个三维张量，但是在输入到深度学习模型中时，一般会要求添加一个额外的维度，即批处理维度（batch dimension）。这是因为在训练或推理过程中，往往会同时处理多个输入样本，而且模型的输入通常要求包含一个批处理维度。这个批处理维度通常放在最前面，即第 0 维度。

  因此，在这段代码中，input_image = transform(image).unsqueeze(0) 的作用是对经过预处理的图像张量在第 0 维度上增加一个维度，从而将其变成一个四维张量，符合模型输入要求。这样处理之后，我们就可以将 input_image 作为模型的输入进行后续的计算了。

5.利用预训练好的AlexNet模型进行图像分类

# 利用预训练好的AlexNet模型进行图像分类
    output = alexnet(input_image)
    _, predicted_idx = torch.max(output, 1)

• output = alexnet(input_image)：首先，将经过预处理的图片张量 input_image 输入到预训练好的 AlexNet 模型中进行推理。模型会对输入的图像进行前向传播计算，得到一个输出张量 output 。这个输出张量包含了关于图像属于每个类别的预测概率值。

• _, predicted_idx = torch.max(output, 1)：接着，我们使用 PyTorch 的 torch.max函数来找到输出张量output 中每行的最大值和对应的索引。torch.max(output, 1)会返回每行最大值的张量以及对应的索引张量。通过这个操作，我们可以得到预测概率最高的类别的索引 predicted_idx。

• 下划线 _ 是一个不需要的变量。在这段代码中，_, predicted_idx = torch.max(output, 1) 中的下划线 _ 实际上是一个占位符，用来接收torch.max函数返回的最大值张量，但在后续的代码中并没有使用到这个值。

6.获取预测结果的类别名称  

# 获取预测结果的类别名称
    predicted_label = class_labels[predicted_idx.item()]

• predicted_idx.item()： 针对预测得到的类别索引predicted_idx，首先使用.item()方法将其转换为 Python 中的标量值，即将PyTorch张量中的数值提取出来。

• predicted_label = class_labels[predicted_idx.item()]：接着，将获取到的预测类别索引作为索引值，从类别标签列表class_labels中取出对应的类别名称，并将其赋值给predicted_label变量。这样就实现了根据神经网络模型预测的类别索引，获取对应的类别名称。

完整代码：

import torch
import torchvision.models as models
from torchvision import transforms
from PIL import Image
import requests
import ssl
ssl._create_default_https_context = ssl._create_unverified_context

# 加载预训练好的AlexNet模型
alexnet = models.alexnet(pretrained=True)

# 打印网络结构
print(alexnet)

# 定义图片文件路径列表
image_files = ['image1.jpg', 'image2.jpg', 'image3.jpg']

# 加载ImageNet的类别标签文件
labels_url = "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt"
class_labels = requests.get(labels_url).text.split('\n')

# 遍历文件路径列表，读取并预测每张图片
for image_file in image_files:
    # 读入图像
    image = Image.open(image_file)

    # 对图像进行预处理
    transform = 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])
    ])
    input_image = transform(image).unsqueeze(0)

    # 利用预训练好的AlexNet模型进行图像分类
    output = alexnet(input_image)
    _, predicted_idx = torch.max(output, 1)

    # 获取预测结果的类别名称
    predicted_label = class_labels[predicted_idx.item()]

    # 打印预测结果
    print(f"Predicted class for {image_file}: {predicted_label}")