使用Hugging Face Datasets和Transformers进行图像相似度比较

使用Hugging Face Datasets和Transformers比较图像相似度

在本文中,您将学习使用🤗 Transformers构建图像相似性系统。查找查询图像与候选图像之间的相似性是信息检索系统的重要用例,例如反向图像搜索。系统试图回答的问题是:给定一个查询图像和一组候选图像,哪些图像与查询图像最相似。

我们将利用🤗 datasets库,因为它可以无缝支持并行处理,这在构建此系统时非常方便。

尽管本文使用了一个基于ViT的模型(nateraw/vit-base-beans)和一个特定的数据集(Beans),但可以扩展为使用其他支持视觉模态和其他图像数据集的模型。您可以尝试一些著名的模型:

  • Swin Transformer
  • ConvNeXT
  • RegNet

此外,本文中介绍的方法也可以扩展到其他模态。

要学习完整工作的图像相似性系统,您可以参考开头链接的Colab笔记本。

我们如何定义相似性?

要构建此系统,我们首先需要定义我们想要计算两个图像之间的相似性的方式。一种广泛流行的做法是计算给定图像的密集表示(嵌入),然后使用余弦相似度度量来确定两个图像的相似程度。

在本文中,我们将使用“嵌入”来表示向量空间中的图像。这使我们能够以有意义的方式压缩图像的高维像素空间(例如224 x 224 x 3)到一个更低维的空间(例如768)。这样做的主要优点是减少了后续步骤中的计算时间。

计算嵌入

为了从图像中计算嵌入,我们将使用一个具有对输入图像在向量空间中表示方式的理解的视觉模型。这种类型的模型通常也被称为图像编码器。

为了加载模型,我们利用了AutoModel类。它提供了一个接口,用于从Hugging Face Hub加载任何兼容的模型检查点。除了模型之外,我们还加载了与模型关联的处理器,用于数据预处理。

from transformers import AutoFeatureExtractor, AutoModel


model_ckpt = "nateraw/vit-base-beans"
extractor = AutoFeatureExtractor.from_pretrained(model_ckpt)
model = AutoModel.from_pretrained(model_ckpt)

在这种情况下,检查点是通过在beans数据集上微调Vision Transformer模型获得的。

这里可能会有一些问题:

Q1:为什么我们没有使用AutoModelForImageClassification

这是因为我们想获取图像的密集表示,而不是离散的类别,而AutoModelForImageClassification会提供这些离散的类别。

Q2:为什么选择这个特定的检查点?

如前所述,我们正在使用特定的数据集构建系统。因此,与使用通用模型(例如在ImageNet-1k数据集上训练的模型)相比,最好使用在使用的数据集上进行了微调的模型。这样,底层模型更好地理解输入图像。

请注意,您还可以使用通过自监督预训练获得的检查点。检查点不一定需要来自有监督的学习。实际上,如果自监督模型训练得好,可以获得令人印象深刻的检索性能。

现在我们有了一个用于计算嵌入的模型,我们需要一些候选图像进行查询。

加载候选图像数据集

在一段时间后,我们将构建哈希表来将候选图像映射到哈希值。在查询时,我们将使用这些哈希表。我们将在相应的部分中更多地讨论哈希表,但现在为了有一组候选图像,我们将使用beans数据集的train拆分。

from datasets import load_dataset


dataset = load_dataset("beans")

这是训练集中的一个样本:

数据集有三个特征:

dataset["train"].features
>>> {'image_file_path': Value(dtype='string', id=None),
 'image': Image(decode=True, id=None),
 'labels': ClassLabel(names=['angular_leaf_spot', 'bean_rust', 'healthy'], id=None)}

为了演示图像相似性系统,我们将使用候选图像数据集中的100个样本,以保持整体运行时间较短。

num_samples = 100
seed = 42
candidate_subset = dataset["train"].shuffle(seed=seed).select(range(num_samples))

寻找相似图像的过程

下面是获取相似图像的过程的图示概述。

对上面的图进行分解,我们有:

  1. 从候选图像(candidate_subset)中提取嵌入,将它们存储在一个矩阵中。
  2. 选择一个查询图像并提取其嵌入。
  3. 迭代计算嵌入矩阵(在步骤1中计算)中的每个候选嵌入与查询嵌入之间的相似度得分。通常我们使用类似字典的映射来维护候选图像的某个标识符与相似度得分之间的对应关系。
  4. 根据相似度得分对映射结构进行排序,并返回相应的标识符。我们使用这些标识符来获取候选样本。

我们可以编写一个简单的实用程序,并将其映射到候选图像数据集以高效地计算嵌入。

import torch 

def extract_embeddings(model: torch.nn.Module):
    """计算嵌入的实用程序。"""
    device = model.device

    def pp(batch):
        images = batch["image"]
        # `transformation_chain` 是我们将输入图像应用于模型之前的预处理转换的组合。
        # 有关更多详细信息,请查看附带的 Colab 笔记本。
        image_batch_transformed = torch.stack(
            [transformation_chain(image) for image in images]
        )
        new_batch = {"pixel_values": image_batch_transformed.to(device)}
        with torch.no_grad():
            embeddings = model(**new_batch).last_hidden_state[:, 0].cpu()
        return {"embeddings": embeddings}

    return pp

我们可以这样映射 extract_embeddings()

device = "cuda" if torch.cuda.is_available() else "cpu"
extract_fn = extract_embeddings(model.to(device))
candidate_subset_emb = candidate_subset.map(extract_fn, batched=True, batch_size=batch_size)

接下来,为了方便起见,我们创建一个包含候选图像标识符的列表。

candidate_ids = []

for id in tqdm(range(len(candidate_subset_emb))):
    label = candidate_subset_emb[id]["labels"]

    # 创建一个唯一的标识符。
    entry = str(id) + "_" + str(label)

    candidate_ids.append(entry)

我们将使用所有候选图像的嵌入矩阵来计算与查询图像的相似度得分。我们已经计算了候选图像的嵌入。在下一个单元格中,我们只需将它们合并成一个矩阵。

all_candidate_embeddings = np.array(candidate_subset_emb["embeddings"])
all_candidate_embeddings = torch.from_numpy(all_candidate_embeddings)

我们将使用余弦相似度来计算两个嵌入向量之间的相似度得分。然后,我们将使用它来获取给定查询样本的相似候选样本。

def compute_scores(emb_one, emb_two):
    """计算两个向量之间的余弦相似度。"""
    scores = torch.nn.functional.cosine_similarity(emb_one, emb_two)
    return scores.numpy().tolist()


def fetch_similar(image, top_k=5):
    """使用`image`作为查询,获取`top_k`个相似图像。"""
    # 准备用于计算嵌入的输入查询图像。
    image_transformed = transformation_chain(image).unsqueeze(0)
    new_batch = {"pixel_values": image_transformed.to(device)}

    # 计算嵌入。
    with torch.no_grad():
        query_embeddings = model(**new_batch).last_hidden_state[:, 0].cpu()

    # 一次性计算与所有候选图像的相似度得分。
    # 我们还创建了一个候选图像标识符与其与查询图像的相似度得分之间的映射。
    sim_scores = compute_scores(all_candidate_embeddings, query_embeddings)
    similarity_mapping = dict(zip(candidate_ids, sim_scores))
 
    # 对映射字典进行排序,并返回`top_k`个候选项。
    similarity_mapping_sorted = dict(
        sorted(similarity_mapping.items(), key=lambda x: x[1], reverse=True)
    )
    id_entries = list(similarity_mapping_sorted.keys())[:top_k]

    ids = list(map(lambda x: int(x.split("_")[0]), id_entries))
    labels = list(map(lambda x: int(x.split("_")[-1]), id_entries))
    return ids, labels

执行查询

在拥有所有工具的情况下,我们可以进行相似性搜索。让我们从beans数据集的test拆分中选择一个查询图像:

test_idx = np.random.choice(len(dataset["test"]))
test_sample = dataset["test"][test_idx]["image"]
test_label = dataset["test"][test_idx]["labels"]

sim_ids, sim_labels = fetch_similar(test_sample)
print(f"查询标签:{test_label}")
print(f"前5个候选标签:{sim_labels}")

结果为:

查询标签:0
前5个候选标签:[0, 0, 0, 0, 0]

看起来我们的系统得到了正确的一组相似图像。可视化后的结果如下:

进一步扩展和结论

现在我们有一个可用的图像相似性系统。但实际上,您将处理更多候选图像。考虑到这一点,我们当前的方法有多个缺点:

  • 如果我们按原样存储嵌入,内存需求会迅速增加,尤其是当处理数百万个候选图像时。在我们的情况下,嵌入是768维,这在大规模情况下仍然相对较高。
  • 具有高维嵌入会直接影响检索部分涉及的后续计算。

如果我们能够在不破坏它们的含义的情况下降低嵌入的维度,我们仍然可以在速度和检索质量之间保持良好的平衡。本文的Colab笔记本实现并展示了使用随机投影和局部敏感哈希实现此目的的实用程序。

🤗 Datasets提供了与FAISS的直接集成,进一步简化了构建相似性系统的过程。假设您已经提取了候选图像(beans数据集)的嵌入,并将它们存储在名为embeddings的特征中。现在,您可以轻松使用数据集的add_faiss_index()方法构建一个密集索引:

dataset_with_embeddings.add_faiss_index(column="embeddings")

构建索引后,可以使用dataset_with_embeddings来检索给定查询嵌入的最近示例,使用get_nearest_examples()方法:

scores, retrieved_examples = dataset_with_embeddings.get_nearest_examples(
    "embeddings", qi_embedding, k=top_k
)

该方法返回得分和相应的候选示例。要了解更多信息,您可以查阅官方文档和此笔记本。

最后,您可以尝试以下构建迷你图像相似性应用程序的空间:

在本文中,我们快速介绍了构建图像相似性系统的入门步骤。如果您觉得这篇文章有趣,我们强烈建议您在我们在这里讨论的概念基础上进行更多的构建,以便更加熟悉内部工作原理。

还想继续学习吗?以下是一些可能对您有用的其他资源:

  • Faiss:高效相似性搜索库
  • ScaNN:高效向量相似性搜索
  • 将图像搜索器集成到移动应用程序中