使用 Vision Transformer 进行图像分类

介绍

图像分类是计算机视觉中的一项核心任务。构建更好的分类器来识别图片中存在的物体是一个活跃的研究领域，因为它在从面部识别到制造业质量控制等领域都有应用。

最先进的图像分类器基于 transformers 架构，该架构最初在 NLP 任务中流行起来。这种架构通常被称为 Vision Transformer (ViT)。这些模型非常适合与 Gradio 的 image 输入组件一起使用，因此在本教程中，我们将构建一个 Web 演示来使用 Gradio 对图像进行分类。我们只需一行 Python 代码即可构建整个 Web 应用程序，它将看起来像页面底部的演示。

让我们开始吧！

先决条件

请确保您已经安装了 gradio Python 包。

步骤 1 — 选择 Vision 图像分类模型

首先，我们需要一个图像分类模型。在本教程中，我们将使用 Hugging Face Model Hub 中的一个模型。该 Hub 包含数千个模型，涵盖数十种不同的机器学习任务。

展开左侧边栏的“Tasks”类别，然后选择“Image Classification”作为我们感兴趣的任务。然后，您将看到 Hub 上所有旨在对图像进行分类的模型。

在撰写本文时，最受欢迎的模型是 google/vit-base-patch16-224，它已在分辨率为 224x224 像素的 ImageNet 图像上进行了训练。我们将使用此模型进行演示。

步骤 2 — 使用 Gradio 加载 Vision Transformer 模型

当使用 Hugging Face Hub 中的模型时，我们不需要为演示定义输入或输出组件。同样，我们也不需要关心预处理或后处理的细节。所有这些都将从模型标签中自动推断出来。

除了 import 语句之外，只需一行 Python 代码即可加载和启动演示。

我们使用 gr.Interface.load() 方法并传入模型的路径，包括 huggingface/ 以指定它来自 Hugging Face Hub。

import gradio as gr

gr.Interface.load(
             "huggingface/google/vit-base-patch16-224",
             examples=["alligator.jpg", "laptop.jpg"]).launch()

请注意，我们添加了一个额外的参数 examples，它允许我们使用一些预定义的示例来预填充我们的界面。

这将生成以下界面，您可以在浏览器中直接尝试。当您输入图像时，它会自动进行预处理并发送到 Hugging Face Hub API，在那里它会通过模型并作为人类可解释的预测返回。尝试上传您自己的图像！