PyTorch 中的图像分类

介绍

图像分类是计算机视觉中的一项核心任务。构建更好的分类器来识别图片中存在的物体是一个活跃的研究领域，其应用范围从自动驾驶汽车到医学成像。

这类模型非常适合与 Gradio 的 图像 输入组件配合使用，因此在本教程中，我们将构建一个使用 Gradio 进行图像分类的网页演示。我们将能够用 Python 构建整个 Web 应用程序，它将像页面底部的演示一样。

让我们开始吧！

前提条件

确保您已安装 `gradio` Python 包。我们将使用预训练的图像分类模型，因此您也应该安装 `torch`。

步骤 1 — 设置图像分类模型

首先，我们需要一个图像分类模型。在本教程中，我们将使用预训练的 Resnet-18 模型，因为它很容易从 PyTorch Hub 下载。您可以使用不同的预训练模型或训练您自己的模型。

import torch

model = torch.hub.load('pytorch/vision:v0.6.0', 'resnet18', pretrained=True).eval()

因为我们将使用模型进行推理，所以我们调用了 .eval() 方法。

步骤 2 — 定义一个 predict 函数

接下来，我们需要定义一个函数，它接受用户输入（在本例中为图像），并返回预测结果。预测结果应以字典形式返回，其中键为类别名称，值为置信度概率。我们将从这个文本文件中加载类别名称。

对于我们的预训练模型，它看起来像这样

import requests
from PIL import Image
from torchvision import transforms

# Download human-readable labels for ImageNet.
response = requests.get("https://git.io/JJkYN")
labels = response.text.split("\n")

def predict(inp):
  inp = transforms.ToTensor()(inp).unsqueeze(0)
  with torch.no_grad():
    prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
    confidences = {labels[i]: float(prediction[i]) for i in range(1000)}
  return confidences

让我们来分解一下。该函数接受一个参数

• inp: 作为 PIL 图像的输入图像

然后，该函数将图像转换为 PIL 图像，最终转换为 PyTorch tensor，将其通过模型，并返回

• confidences: 预测结果，以字典形式表示，其中键是类别标签，值是置信度概率

步骤 3 — 创建 Gradio 界面

现在我们已经设置好了预测函数，我们可以在其周围创建一个 Gradio 界面。

在本例中，输入组件是一个拖放图像组件。为了创建此输入，我们使用 Image(type="pil")，它会创建该组件并处理预处理以将其转换为 PIL 图像。

输出组件将是一个 Label，它以美观的形式显示顶部标签。由于我们不想显示所有 1,000 个类别标签，我们将通过将其构造为 Label(num_top_classes=3) 来定制它，使其仅显示前 3 个图像。

最后，我们将添加另一个参数，即 examples，它允许我们使用一些预定义的示例来预填充我们的界面。Gradio 的代码如下所示

import gradio as gr

gr.Interface(fn=predict,
             inputs=gr.Image(type="pil"),
             outputs=gr.Label(num_top_classes=3),
             examples=["lion.jpg", "cheetah.jpg"]).launch()

这会生成以下界面，您可以在浏览器中直接尝试（尝试上传您自己的示例！）