使用FAISS和CLIP构建图像相似性搜索引擎

使用FAISS和CLIP构建图像相似性搜索引擎

引用

CSDN

1.

https://blog.csdn.net/weixin_38739735/article/details/141653205

本文将指导您构建一个图像相似性搜索引擎，使您能够通过文本描述或提供参考图片来快速定位所需图像。

流程概述

图像的语义信息可以通过一个数值向量——嵌入向量——来表示。通过比较这些嵌入向量而非原始图像，可以高效地进行相似性搜索。我们将为数据库中的每张图像生成一个嵌入向量，并存储于索引中。当用户提交文本查询或参考图像时，系统将生成相应的嵌入向量，并与索引中的向量进行匹配，以找出最相似的图像。 具体步骤如下：

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

CLIP模型

CLIP（Contrastive Language-Image Pre-training）模型是由OpenAI开发的多模态视觉和语言预训练模型，它能够将图像和文本映射到统一的潜在空间。由于我们将使用图像和文本查询来搜索图像，CLIP模型将被用于生成我们数据的嵌入向量。

FAISS索引

FAISS（Facebook AI Similarity Search）是Meta开发的开源库，专门用于高效相似性搜索和聚类。FAISS围绕Index对象构建，该对象负责存储数据库的嵌入向量。我们将利用FAISS来索引我们的图像数据集，并检索与用户查询相似的图片。

代码实现

步骤1 - 数据集探索

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

步骤2 - 从图像数据集中提取CLIP嵌入

为了提取CLIP嵌入，我们将首先使用HuggingFace SentenceTransformer库加载CLIP模型：

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

接下来，我们将创建一个函数，使用glob遍历我们的数据集目录，使用PIL Image.open打开每张图像，并使用CLIP模型.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还提供了各种索引选项。它可以在平衡搜索速度和准确性的同时，使用近似或压缩技术高效处理大型数据集。在这个教程中，我们将使用一个“Flat”索引，它通过将查询向量与数据集中的每个向量进行比较，执行暴力搜索，确保结果的准确性，但计算复杂度更高。

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)

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

有了我们的FAISS索引，我们现在可以使用文本查询或参考图像检索图像。如果查询是图像路径，使用PIL Image.open打开查询。接下来，使用CLIP模型.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)

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

您可以尝试在提供的Colab笔记本上进行其他查询，以查看模型对不同文本和图像输入的表现。

结束语

在这个教程中，我们使用CLIP和FAISS构建了一个基本的图像相似性搜索引擎。检索到的图像与查询共享相似的语义含义，表明了方法的有效性。尽管CLIP对于Zero Shot模型来说显示出了不错的结果，但它可能在分布外数据、细粒度任务上表现不佳，并继承了它所训练数据的自然偏见。为了克服这些限制，您可以尝试其他类似CLIP的预训练模型，如OpenClip，或者在您自己的自定义数据集上微调CLIP。

源码参考：https://gist.github.com/Lihi-Gur-Arie/7cac63dbffde55449d2444e402d87bfc