Transformers documentation

图像特征提取

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.2).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

图像特征提取

图像特征提取是指从给定图像中提取具有语义意义的特征。该任务有许多应用场景,包括图像相似性计算和图像检索。此外,大多数计算机视觉模型都可以用于图像特征提取:只需移除任务特定的头部(如图像分类、目标检测等),即可获得特征。这些特征在更高层面上非常有用,例如边缘检测、角点检测等。根据模型的深度不同,这些特征还可能包含关于真实世界的信息(例如猫长什么样)。因此,这些输出可用于在特定数据集上训练新的分类器。

在本指南中,您将:

  • 学习如何在 image-feature-extraction 流水线之上构建一个简单的图像相似度系统。
  • 使用裸模型推理完成相同的任务。

使用 image-feature-extraction Pipeline 进行图像相似度计算

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

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")

让我们看看 pipeline 的实际运行效果。首先,初始化 pipeline。如果不传入任何模型,pipeline 将自动使用 google/vit-base-patch16-224 进行初始化。如果需要计算相似度,请将 pool 设置为 True。

import torch
from transformers import pipeline
from accelerate import Accelerator
# automatically detects the underlying device type (CUDA, CPU, XPU, MPS, etc.)
device = Accelerator().device
pipe = pipeline(task="image-feature-extraction", model="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="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>)
Update on GitHub