图像搜索引擎DIY

你是否曾经想在无穷无尽的图像数据集中查找图像,但发现这太过繁琐?在本教程中,我们将构建一个图像相似性搜索引擎,以便使用文本查询或参考图像轻松查找图像。为方便起见,本文底部以 Colab 笔记本的形式提供了本教程的完整代码。

1、概述

图像的语义可以用称为嵌入的数值向量表示。比较这些低维嵌入向量(而不是原始图像)可以实现高效的相似性搜索。对于数据集中的每个图像,我们将创建一个嵌入向量并将其存储在索引中。当提供文本查询或参考图像时,将生成其嵌入并将其与索引嵌入进行比较以检索最相似的图像。

以下是简要概述:

  • 嵌入:使用 CLIP 模型提取图像的嵌入。
  • 索引:嵌入存储为 FAISS 索引。
  • 检索:使用 FAISS,查询的嵌入与索引嵌入进行比较,以检索最相似的图像。

1.1 CLIP 模型

由 OpenAI 开发的 CLIP(对比语言-图像预训练)模型是一种多模态视觉和语言模型,可将图像和文本映射到相同的潜在空间。由于我们将使用图像和文本查询来搜索图像,因此我们将使用 CLIP 模型来嵌入我们的数据。

1.2 FAISS 索引

FAISS(Facebook AI 相似性搜索)是由 Meta 开发的开源库。它围绕存储数据库嵌入向量的 Index 对象构建。FAISS 支持高效的相似性搜索和密集向量聚类,我们将使用它来索引我们的数据集并检索与查询相似的照片。

2、代码实现

为了创建本教程的图像数据集,我从 Pexels 收集了 52 张不同主题的图像。为了感受一下,让我们观察 10 张随机图像:

2.1 从图像数据集中提取 CLIP 嵌入

为了提取 CLIP 嵌入,我们首先使用 HuggingFace的 SentenceTransformer 库加载 CLIP 模型:

model = SentenceTransformer('clip-ViT-B-32')

接下来,我们将创建一个函数,使用 glob 遍历我们的数据集目录,使用 PIL的 Image.open 打开每个图像,并使用 CLIP的 model.encode 为每个图像生成一个嵌入向量。它返回嵌入向量列表和图像数据集的路径列表:

def generate_clip_embeddings(images_path, model):

    image_paths = glob(os.path.join(images_path, '**/*.jpg'), recursive=True)
    
    embeddings = []
    for img_path in image_paths:
        image = Image.open(img_path)
        embedding = model.encode(image)
        embeddings.append(embedding)
    
    return embeddings, image_paths



IMAGES_PATH = '/path/to/images/dataset'

embeddings, image_paths = generate_clip_embeddings(IMAGES_PATH, model)

2.2 生成 FAISS 索引

下一步是从嵌入向量列表创建 FAISS 索引。FAISS 为相似性搜索提供了各种距离度量,包括内积 (IP) 和 L2(欧几里得)距离。

FAISS 还提供各种索引选项。它可以使用近似或压缩技术来有效处理大型数据集,同时平衡搜索速度和准确性。在本教程中,我们将使用“平面”索引,它通过将查询向量与数据集中的每个向量进行比较来执行强力搜索,以确保以更高的计算复杂度为代价获得准确的结果。

def create_faiss_index(embeddings, image_paths, output_path):

    dimension = len(embeddings[0])
    index = faiss.IndexFlatIP(dimension)
    index = faiss.IndexIDMap(index)
    
    vectors = np.array(embeddings).astype(np.float32)

    # Add vectors to the index with IDs
    index.add_with_ids(vectors, np.array(range(len(embeddings))))
    
    # Save the index
    faiss.write_index(index, output_path)
    print(f"Index created and saved to {output_path}")
    
    # Save image paths
    with open(output_path + '.paths', 'w') as f:
        for img_path in image_paths:
            f.write(img_path + '\n')
    
    return index


OUTPUT_INDEX_PATH = "/content/vector.index"
index = create_faiss_index(embeddings, image_paths, OUTPUT_INDEX_PATH)

faiss.IndexFlatIP 初始化内积相似度索引,将其包装在 faiss.IndexIDMap 中,以将每个向量与 ID 关联。接下来, index.add_with_ids 将向量添加到具有连续 ID 的索引中,并将索引与图像路径一起保存到磁盘。

索引可以立即使用,也可以保存到磁盘以备将来使用。要加载 FAISS 索引,我们将使用以下函数:

def load_faiss_index(index_path):
    index = faiss.read_index(index_path)
    with open(index_path + '.paths', 'r') as f:
        image_paths = [line.strip() for line in f]
    print(f"Index loaded from {index_path}")
    return index, image_paths

index, image_paths = load_faiss_index(OUTPUT_INDEX_PATH)

2.3 通过文本查询或参考图像检索图像

建立 FAISS 索引后,我们现在可以使用文本查询或参考图像检索图像。如果查询是图像路径,则使用 PIL的 Image.open 打开查询。接下来,使用 CLIP的 model.encode 提取查询嵌入向量。

def retrieve_similar_images(query, model, index, image_paths, top_k=3):
    
    # query preprocess:
    if query.endswith(('.png', '.jpg', '.jpeg', '.tiff', '.bmp', '.gif')):
        query = Image.open(query)

    query_features = model.encode(query)
    query_features = query_features.astype(np.float32).reshape(1, -1)

    distances, indices = index.search(query_features, top_k)

    retrieved_images = [image_paths[int(idx)] for idx in indices[0]]

    return query, retrieved_images

检索发生在 index.search 方法上。它实现了 k-最近邻 (kNN) 搜索,以找到与查询向量最相似的 k 个向量。我们可以通过更改 top_k 参数来调整 k 的值。在我们的实现中,kNN 搜索中使用的距离度量是余弦相似度。该函数返回查询和检索图像路径的列表。

现在我们准备使用文本查询进行搜索。辅助函数 visualize_results 显示结果。你可以在关联的 Colab 笔记本中找到它。让我们探索针对文本查询“ball”检索到的最相似的 3 幅图像,例如:

query = 'ball'
query, retrieved_images = retrieve_similar_images(query, model, index, image_paths, top_k=3)
visualize_results(query, retrieved_images)
使用查询“a ball”检索图像

对于查询“animal”,我们得到:

使用查询“动物”检索图像

使用参考图像搜索:

query ='/content/drive/MyDrive/Colab Notebooks/my_medium_projects/Image_similarity_search/image_dataset/pexels-w-w-299285-889839.jpg'
query, retrieved_images = retrieve_similar_images(query, model, index, image_paths, top_k=3)
visualize_results(query, retrieved_images)
查询和检索到的图像

如我们所见,我们为现成的预训练模型获得了非常酷的结果。当我们通过眼睛绘画的参考图像进行搜索时,除了找到原始图像外,它还找到了一个匹配的眼镜和另一幅画。这展示了查询图像语义含义的不同方面。

你可以在提供的 Colab 笔记本上尝试其他查询,以查看模型在不同文本和图像输入下的表现。

3、结束语

在本教程中,我们使用 CLIP 和 FAISS 构建了一个基本的图像相似性搜索引擎。检索到的图像与查询具有相似的语义含义,表明该方法的有效性。

虽然 CLIP 对于零样本模型显示出不错的结果,但它在分布外数据、细粒度任务上可能表现出较低的性能,并继承了训练数据的自然偏差。为了克服这些限制,你可以尝试其他类似 CLIP 的预训练模型(如 OpenClip),或者在自己的自定义数据集上微调 CLIP。


原文链接:Building an Image Similarity Search Engine with FAISS and CLIP

BimAnt翻译整理,转载请标明出处