图片由作者在Flux-Pro平台上生成

你是否曾经想要在你的无尽图像数据集中找到一张图片，却发现过程太繁琐了？在这个教程中，我们将构建一个图像相似度搜索引擎，以便你可以轻松地通过文本查询或参考图片来查找图片。为了方便起见，本教程的完整代码已放在文章底部，以Colab notebook的形式提供。

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

这里有一个简要概述：

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

CLIP（对比语言-图像预训练）模型是由OpenAI开发的多模态视觉和语言模型，它将图像和文本映射到相同的潜在空间。由于我们将使用图像和文本查询来搜索图像，我们将使用CLIP模型嵌入我们的数据。有关CLIP的进一步阅读，你可以查阅我之前的文章。

FAISS (Facebook AI 相似度搜索) 是由 Meta 开发的一个开源库。它围绕着存储数据库嵌入向量的 Index 对象构建。FAISS 可以高效地进行稠密向量的相似度搜索和聚类，我们将使用它来索引我们的数据集并检索与查询相似的照片。

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

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

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

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

    def 生成_clip嵌入(images_path, model):  

        image_paths = glob(os.path.join(images_path, '**/*.jpg'), recursive=True)  

        嵌入 = []  
        for img_path in image_paths:  
            image = Image.open(img_path)  
            embedding = model.encode(image)  
            嵌入.append(embedding)  

        return 嵌入, image_paths  

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

    嵌入, image_paths = 生成_clip嵌入(IMAGES_PATH, model)

下一步是根据嵌入向量列表创建一个FAISS索引。FAISS提供了多种相似性搜索的距离度量，包括内积（IP）和L2（欧几里得）距离。

FAISS 还提供了多种索引选项。它可以使用近似或压缩技术来高效处理大规模数据集，同时平衡搜索速度和准确性。在本教程中，我们将使用“Flat”索引，该索引通过将查询向量与数据集中的每个向量进行比较来进行暴力搜索，确保结果的准确性，但代价是更高的计算复杂度。

    def 创建_faiss索引(embeddings, image_paths, output_path):  

        dimension = len(embeddings[0])  
        index = faiss.IndexFlatIP(dimension)  
        index = faiss.IndexIDMap(index)  

        vectors = np.array(embeddings).astype(np.float32)  

        # 将向量添加到索引中，并使用ID  
        index.add_with_ids(vectors, np.array(range(len(embeddings))))  

        # 保存索引  
        faiss.write_index(index, output_path)  
        print(f"索引创建并保存到 {output_path}")  

        # 保存图像路径  
        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 = 创建_faiss索引(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_path} 加载索引")  
        return index, image_paths  

    index, image_paths = load_faiss_index(OUTPUT_INDEX_PATH)

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

    def retrieve_similar_images(query, model, index, image_paths, top_k=3):  

        # 查询预处理:  
        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的值。我们在实现中使用的距离度量是余弦相似度。该函数返回查询和一个检索到的图像路径列表。

使用文本查询搜索:

现在我们可以开始查看搜索结果了。辅助函数 visualize_results 用于展示结果。你可以在相关的 Colab notebook 中找到它。以文本查询“ball”为例，让我们来探索检索到的最相似的3张图片。

    query = '球'  
    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)

查询和检索的图像

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

您可以在提供的Colab notebook中尝试其他查询，看看模型在不同文本和图像输入下的表现。

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

恭喜你一路看到这里。点击 👍 表示你的赞赏，提升一下算法的自尊心 🤓

想了解更多？

• 探索 我撰写的更多文章
• 订阅 以获取我发布新文章的通知
• 在 Linkedin 上关注我

Colab Notebook 链接