简介

你是否曾想过在永无止境的图像数据集中找到一张图像，但又觉得太繁琐？在本教程中，我们将建立一个图像相似性搜索引擎，使用文本查询或参考图像轻松查找图像。

管道概述

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

下面是简要概述：

1. 嵌入： 使用 CLIP 模型提取图像的嵌入。
2. 索引： 嵌入信息以 FAISS 索引的形式存储。
3. 检索： 利用 FAISS，将查询的嵌入信息与索引中的嵌入信息进行比较，以检索出最相似的图像。

CLIP 模型

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

FAISS 索引

FAISS（Facebook 人工智能相似性搜索）是由 Meta 开发的一个开源库。它围绕存储数据库嵌入向量的索引对象构建。FAISS 可以对密集向量进行高效的相似性搜索和聚类，我们将用它为数据集编制索引，并检索与查询相似的照片。

代码实现

步骤 1 - 数据集探索

为了创建本教程的图片数据集，我从 Pexels 收集了 52 张不同主题的图片。为了获得感觉，让我们随机观察 10 张图片：

步骤 2 - 从图像数据集中提取 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)

步骤 3 - 生成 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 将向量添加到索引中，索引会连同图像路径一起保存到磁盘中。

该索引可以立即使用，也可以保存到磁盘供将来使用：

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)

步骤 4 - 通过文本查询或参考图像检索图像

建立 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-Nearest Neighbors (kNN) 搜索，以找到与查询向量最相似的 k 个向量。我们可以通过更改 top_k 参数来调整 k 值。在我们的实现中，kNN 搜索使用的距离度量是余弦相似度。函数返回查询和检索图像路径列表。

使用文本查询进行搜索

现在我们可以检查搜索结果了。辅助函数 visualize_results 可以显示搜索结果。你可以在相关的 Colab 笔记本中查看。让我们以文本查询 “球 ”为例，查看检索到的最相似的 3 张图片：

query = 'ball'
query, retrieved_images = retrieve_similar_images(query, model, index, image_paths, top_k=3)
visualize_results(query, retrieved_images)

对于查询 “动物”，我们得到的结果是

使用参考图像搜索

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)

我们可以看到，对于一个现成的预训练模型来说，我们得到的结果非常棒。当我们通过一幅眼睛绘画的参考图片进行搜索时，除了找到原始图片外，它还找到了一幅与眼镜匹配的图片和一幅与不同绘画匹配的图片。这说明了查询图片语义的不同方面。

结论

在本教程中，我们使用 CLIP 和 FAISS 构建了一个基本的图像相似性搜索引擎。检索到的图像与查询的语义相似，这表明了该方法的有效性。虽然 CLIP 在零点模型中表现出了不错的效果，但它在非分布数据、细粒度任务中可能表现出较低的性能，并且会继承其所训练数据的自然偏差。为了克服这些局限性，你可以尝试 OpenClip 中其他类似 CLIP 的预训练模型，或在自己的自定义数据集上对 CLIP 进行微调。