最佳经纬度谷歌地图 API免费版推荐与使用指南
生成式 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)
如我们所见,我们为现成的预训练模型获得了非常酷的结果。当我们通过眼睛绘画的参考图像进行搜索时,除了找到原始图像外,它还找到了一个匹配的眼镜和另一幅画。这展示了查询图像语义含义的不同方面。
3、结束语
在本教程中,我们使用 CLIP 和 FAISS 构建了一个基本的图像相似性搜索引擎。检索到的图像与查询具有相似的语义含义,表明该方法的有效性。
虽然 CLIP 对于零样本模型显示出不错的结果,但它在分布外数据、细粒度任务上可能表现出较低的性能,并继承了训练数据的自然偏差。为了克服这些限制,你可以尝试其他类似 CLIP 的预训练模型(如 OpenClip),或者在自己的自定义数据集上微调 CLIP。