点击下方卡片，关注“小白玩转Python”公众号主要工作：阐述如何利用文本或图像查询来检索您的图像数据库，借助CLIP模型的嵌入技术与FAISS索引系统。
介绍
您是否曾面临在庞大的图像数据库中寻找特定图片的难题？本教程将指导您构建一个图像相似性搜索引擎，使您能够通过文本描述或提供参考图片来快速定位所需图像。为方便操作，本教程的完整代码已在文章末尾以Colab笔记本形式提供。
流程概述
图像的语义信息可以通过一个数值向量——嵌入向量——来表示。通过比较这些嵌入向量而非原始图像，可以高效地进行相似性搜索。我们将为数据库中的每张图像生成一个嵌入向量，并存储于索引中。当用户提交文本查询或参考图像时，系统将生成相应的嵌入向量，并与索引中的向量进行匹配，以找出最相似的图像。 具体步骤如下：
嵌入：使用CLIP模型提取图像的嵌入。
索引：将嵌入存储为FAISS索引。
检索：使用FAISS，查询的嵌入与索引中的嵌入进行比较，以检索最相似的图像。
CLIP模型
CLIP（Contrastive Language-Image Pre-training）模型是由OpenAI开发的多模态视觉和语言预训练模型，它能够将图像和文本映射到统一的潜在空间。由于我们将使用图像和文本查询来搜索图像，CLIP模型将被用于生成我们数据的嵌入向量。若您希望了解更多关于CLIP的信息，可以参阅我之前的相关文章。
FAISS索引
FAISS（Facebook AI Similarity Search）是Meta开发的开源库，专门用于高效相似性搜索和聚类。FAISS围绕Index对象构建，该对象负责存储数据库的嵌入向量。我们将利用FAISS来索引我们的图像数据集，并检索与用户查询相似的图片。
代码实现
步骤1 - 数据集探索
为了创建本教程的图像数据集，我从Pexels收集了52张不同主题的图像。为了感受一下，让我们观察10张随机图像：
介绍
您是否曾经想在无尽的图像数据集中找到一张图片，却发现这太繁琐了？在这个教程中，我们将构建一个图像相似性搜索引擎，以便使用文本查询或参考图像轻松找到图像。为了方便您，本教程的完整代码在文章底部以Colab笔记本的形式提供。
流程概述
图像的语义含义可以通过一个数值向量表示，称为嵌入。比较这些低维嵌入向量，而不是原始图像，可以进行高效的相似性搜索。对于数据集中的每张图像，我们将创建一个嵌入向量并将其存储在索引中。当提供文本查询或参考图像时，将生成其嵌入并将其与索引中的嵌入进行比较，以检索最相似的图像。
以下是简要概述：
嵌入：使用CLIP模型提取图像的嵌入。
索引：将嵌入存储为FAISS索引。
检索：使用FAISS，查询的嵌入与索引中的嵌入进行比较，以检索最相似的图像。
CLIP模型
由OpenAI开发的CLIP（Contrastive Language-Image Pre-training）模型是一个多模态视觉和语言模型，它将图像和文本映射到相同的潜在空间。由于我们将使用图像和文本查询来搜索图像，我们将使用CLIP模型来嵌入我们的数据。有关CLIP的更多阅读，您可以查看我之前在这里的文章。
FAISS索引
FAISS（Facebook AI Similarity Search）是由Meta开发的开源库。它围绕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
· END ·
🌟 想要变身计算机视觉小能手？快来「小白玩转Python」公众号！
回复“Python视觉实战项目”，解锁31个超有趣的视觉项目大礼包！🎁
本文仅供学习交流使用，如有侵权请联系作者删除
扫一扫
1500
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
