基于图像的搜索引擎

引言

在视觉数据领域，高效准确的图像检索是关键。这篇博客文章提供了一步步的指南，教你如何使用开源工具构建一个基于图像的搜索引擎。学完本文，你将掌握创建强大、可定制的搜索引擎以实现精确图像检索的技能。这篇博客将助你以全新且激动人心的方式探索视觉数据。准备好投入图像检索的世界，使用开源工具解锁高效视觉数据搜索的秘密！

嵌入数据

嵌入对于构建基于图像的搜索引擎至关重要，因为它将图像转换为易于比较和处理的格式。此步骤通过将图像表示为高维空间中的唯一向量来实现高效图像检索。

我们首先下载所需的库。至于 `loadimg`，它是我开发的一个 Python 库，用于轻松读取和转换图像。如果你愿意，可以跳过这个库。如果你对它感兴趣并希望为其发展做出贡献，可以查看我的 GitHub 仓库。

pip install -qU datasets accelerate loadimg faiss-cpu

然后我们可以继续加载我们的数据集。

from datasets import load_dataset
data = load_dataset("not-lain/pokemon",split="train")

data

>>> Dataset({
    features: ['image', 'text'],
    num_rows: 898
})

之后，我们来加载模型，我这里用的是 CLIP，但你可以用任何其他类似的模型。

import torch
from transformers import AutoProcessor, AutoModelForZeroShotImageClassification # or you can use CLIPProcessor, CLIPModel

device = 'cuda' if torch.cuda.is_available() else 'cpu'

processor = AutoProcessor.from_pretrained("openai/clip-vit-large-patch14")
model = AutoModelForZeroShotImageClassification.from_pretrained("openai/clip-vit-large-patch14", device_map = device)

现在到了最重要的部分，我们将嵌入数据集并把嵌入数据添加到新列`embeddings`中。建议你在此处使用GPU或任何其他加速器，因为这是一个非常缓慢的过程。

def embed(batch):
    pixel_values = processor(images = batch["image"], return_tensors="pt")['pixel_values']
    pixel_values = pixel_values.to(device)
    img_emb = model.get_image_features(pixel_values)
    batch["embeddings"] = img_emb
    return batch

embedded_dataset = dataset.map(embed, batched=True, batch_size=16)

建议将嵌入的数据存储在数据库中，无论是本地的、HF、Pinecone、ChromaDB 还是任何其他替代方案，以避免再次嵌入数据集。

💡提示

尽管这与我们这里的工作无关，但你也可以使用相同的模型来创建文本类型数据的嵌入。首先通过 `tokens = processor(text = "some text here", padding=True, return_tensors="pt").to(device)` 处理输入，然后将其传递给你的模型，通过 `text_emb = model.get_text_features(**tokens)` 创建嵌入。

embedded_dataset.push_to_hub("not-lain/embedded-pokemon")

检索图像

一旦你的数据集嵌入并存储在数据库中，我们现在可以开始定义检索逻辑。我们首先需要在此处加载与上一节中使用的相似模型。

import torch
from transformers import AutoProcessor, AutoModelForZeroShotImageClassification # or you can use CLIPProcessor, CLIPModel

device = 'cuda' if torch.cuda.is_available() else 'cpu'

processor = AutoProcessor.from_pretrained("openai/clip-vit-large-patch14")
model = AutoModelForZeroShotImageClassification.from_pretrained("openai/clip-vit-large-patch14", device_map = device)

我们还需要加载我们的嵌入式数据集。

from datasets import load_dataset

dataset = load_dataset("not-lain/embedded-pokemon", split="train")

你需要给嵌入列添加一个 Faiss 索引，以便将其设置为相似性搜索索引。

Faiss 是一个用于高效相似性搜索和密集向量聚类的库，它对于大规模图像检索任务特别有用。通过向嵌入列添加 Faiss 索引，可以实现快速准确的最近邻搜索，从而有效地根据图像嵌入检索相似图像。这一步显著提高了图像搜索引擎的性能，尤其是在处理大量图像时。

dataset = dataset.add_faiss_index("embeddings")

现在要检索最相似的图像，你需要首先创建新图像的嵌入，然后从数据集中检索最相似的条目。

import numpy as np

def search(query: str, k: int = 4 ):
    """a function that embeds a new image and returns the most probable results"""

    pixel_values = processor(images = query, return_tensors="pt")['pixel_values'] # embed new image
    pixel_values = pixel_values.to(device)
    img_emb = model.get_image_features(pixel_values)[0] # because it's a single element
    img_emb = img_emb.cpu().detach().numpy() # convert to numpy because the datasets library does not support torch vectors

    scores, retrieved_examples = dataset.get_nearest_examples( # retrieve results
        "embeddings", img_emb, # compare our new embedded image with the dataset embeddings
        k=k # get only top k results
    )

    return retrieved_examples

我们来测试一下我们的算法，你可以从加载一张图片开始。

from loadimg import load_img
image = load_img("https://img.pokemondb.net/artwork/large/charmander.jpg")
image

之后，你可以检索最相似的条目。

这些条目按降序排列，第一个条目与我们的输入图像最相似。

retrieved_examples = search(image)

我们来可视化一下结果。

import matplotlib.pyplot as plt
f, axarr = plt.subplots(2,2)
for index in range(4):
    i,j = index//2, index%2
    axarr[i,j].set_title(retrieved_examples["text"][index])
    axarr[i,j].imshow(retrieved_examples["image"][index])
    axarr[i,j].axis('off')
plt.show()

演示

现在让我们把所有东西整合到一个应用程序中，看看我们的工作效果如何。你可以将以下应用程序视为一个很好的参考：https://huggingface.co/spaces/not-lain/image-retriever

致谢

我在此感谢 Pinecone 的 文档，它们在开发本博客中使用的脚本方面给予了我极大的帮助 🌲

我要感谢 HuggingFace Discord 服务器中所有支持我工作的人，特别是 lunarflu、christopher 和 tomaarsen ❤️

如果你喜欢这篇博客文章，请考虑点赞，这将有助于我展示我的工作 🤗

最后，如果你想请求其他博客文章，可以通过 Twitter、电子邮件 或 LinkedIn 与我联系 ✉️