图像特征提取

图像特征提取是根据给定图像提取语义上有意义的特征的任务。这在图像相似度匹配和图像检索等许多用例中都有应用。此外，大多数计算机视觉模型都可以用于图像特征提取，其中可以移除特定任务的头部（图像分类、目标检测等）并获取特征。这些特征在更高级别上非常有用：边缘检测、角点检测等。它们还可能包含有关现实世界的信息（例如猫的外观），具体取决于模型的深度。因此，这些输出可以用于在特定数据集上训练新的分类器。

在本指南中，你将：

学习如何基于 image-feature-extraction 管道构建一个简单的图像相似度系统。
通过裸模型推理完成相同的任务。

使用 image-feature-extraction 管道进行图像相似度匹配

我们有两张猫坐在渔网上的图像，其中一张是生成的。

from PIL import Image
import requests

img_urls = ["https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/cats.png", "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/cats.jpeg"]
image_real = Image.open(requests.get(img_urls[0], stream=True).raw).convert("RGB")
image_gen = Image.open(requests.get(img_urls[1], stream=True).raw).convert("RGB")

让我们看看管道的实际操作。首先，初始化管道。如果你不向其传递任何模型，管道将自动使用 google/vit-base-patch16-224 进行初始化。如果你想计算相似度，将 pool 设置为 True。

import torch
from transformers import pipeline
from accelerate.test_utils.testing import get_backend
# automatically detects the underlying device type (CUDA, CPU, XPU, MPS, etc.)
DEVICE, _, _ = get_backend()
pipe = pipeline(task="image-feature-extraction", model_name="google/vit-base-patch16-384", device=DEVICE, pool=True)

要使用 pipe 进行推理，请将两张图像都传递给它。

outputs = pipe([image_real, image_gen])

输出包含这两张图像的池化嵌入。

# get the length of a single output
print(len(outputs[0][0]))
# show outputs
print(outputs)

# 768
# [[[-0.03909236937761307, 0.43381670117378235, -0.06913255900144577,

要获得相似度分数，我们需要将它们传递给一个相似度函数。

from torch.nn.functional import cosine_similarity

similarity_score = cosine_similarity(torch.Tensor(outputs[0]),
                                     torch.Tensor(outputs[1]), dim=1)

print(similarity_score)

# tensor([0.6043])

如果你想在池化之前获取最后的隐藏状态，请避免为 pool 参数传递任何值，因为它默认为 False。这些隐藏状态对于基于模型的特征训练新分类器或模型很有用。

pipe = pipeline(task="image-feature-extraction", model_name="google/vit-base-patch16-224", device=DEVICE)
outputs = pipe(image_real)

由于输出未池化，我们得到最后的隐藏状态，其中第一维是批量大小，最后两维是嵌入形状。

import numpy as np
print(np.array(outputs).shape)
# (1, 197, 768)

使用 AutoModel 获取特征和相似度

我们还可以使用 transformers 的 AutoModel 类来获取特征。AutoModel 加载任何不带任务特定头部的 transformers 模型，我们可以使用它来获取特征。

from transformers import AutoImageProcessor, AutoModel

processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
model = AutoModel.from_pretrained("google/vit-base-patch16-224").to(DEVICE)

让我们编写一个简单的推理函数。我们将首先将输入传递给 processor，然后将其输出传递给 model。

def infer(image):
  inputs = processor(image, return_tensors="pt").to(DEVICE)
  outputs = model(**inputs)
  return outputs.pooler_output

我们可以直接将图像传递给此函数并获取嵌入。

embed_real = infer(image_real)
embed_gen = infer(image_gen)

我们可以再次对嵌入计算相似度。

from torch.nn.functional import cosine_similarity

similarity_score = cosine_similarity(embed_real, embed_gen, dim=1)
print(similarity_score)

# tensor([0.6061], device='cuda:0', grad_fn=<SumBackward1>)

< > 在 GitHub 上更新

Transformers

图像特征提取

使用 image-feature-extraction 管道进行图像相似度匹配

使用 AutoModel 获取特征和相似度