生成式 AI 图像搜索引擎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)

对于查询“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)

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