‹ Back

使用FAISS和CLIP構建圖像相似性搜索引擎的實踐指南*

 

圖靈學院編譯/2024年8月27日

 

    在當今數位時代，隨著圖像數據集的不斷增長，如何有效地搜索和管理這些海量圖像成為了一個重要的課題。本文將介紹如何使用FAISS和CLIP模型來構建一個高效的圖像相似性搜索引擎，讓您能夠通過文本或圖片查詢來輕鬆找到所需的圖像。

 

一、簡介

 

您是否曾經在龐大的圖像數據集中尋找某張圖片，卻因數量龐大而感到繁瑣？本文將引導您一步步構建一個圖像相似性搜索引擎，讓您可以使用文本查詢或參考圖像輕鬆地找到所需的圖片。本文提供了完整的程式碼，您可以在複製進行實踐。

 

二、系統概覽

 

圖像的語義訊息可以通過數字向量（即嵌入）來表示。與其比較原始圖像，將這些低維度的嵌入向量進行比較可以更高效地進行相似性搜索。對於數據集中的每個圖像，我們將生成其嵌入向量並存儲在索引中。當提供文本查詢或參考圖像時，我們將生成其嵌入並與索引中的嵌入進行比較，以檢索最相似的圖像。

 

三、CLIP模型介紹

 

CLIP（Contrastive Language-Image Pre-training）是由OpenAI開發的一種多模態視覺與語言模型，它將圖像和文本映射到相同的潛在空間中。由於我們將使用文本和圖像查詢來搜索圖像，因此選擇使用CLIP模型來嵌入數據。如果您對CLIP模型有更深入的興趣，可以參考原文作者之前的文章。

 

四、FAISS索引介紹

 

FAISS（Facebook AI Similarity Search）是Meta開發的一個開源庫，它圍繞索引物件來存儲數據庫嵌入向量。FAISS支持密集向量的高效相似性搜索和聚類，我們將使用它來為數據集建立索引並檢索與查詢最相似的圖像。

 

 

五、程式碼實現步驟

 

Step 1. 數據集探索

 

為了構建本教程的圖像數據集，作者從Pexels上收集了52張不同主題的圖像。首先，我們來觀察隨機選取的10張圖像，以了解數據集的整體風格。

Step 2. 提取圖像嵌入

 

要提取CLIP嵌入，我們首先需要使用HuggingFace的SentenceTransformer庫加載CLIP模型：

 

“python
model = SentenceTransformer('clip-ViT-B-32')
”

 

接下來，我們將創建一個函數，該函數遍歷數據集目錄，打開每個圖像並使用CLIP模型進行編碼，生成嵌入向量。該函數返回嵌入向量列表和圖像路徑列表。

 

“python
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_e

mbeddings(IMAGES_PATH, model)
”

Step 3. 建立FAISS索引

 

接下來的步驟是使用嵌入向量列表創建FAISS索引。FAISS提供多種距離度量選項，包括內積（IP）和L2（歐幾里得）距離。在本教程中，我們將使用‘Flat’索引，這是一種逐一比較數據集向量的暴力搜索方法，保證精確結果但計算複雜度較高。

 

“python
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)
    index.add_with_ids(vectors, np.array(range(len(embeddings))))
    faiss.write_index(index, output_path)
    print(f"Index created and saved to {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 = create_faiss_index(embeddings, image_paths, OUTPUT_INDEX_PATH)
”

 

FAISS索引創建後，可以直接使用或保存到磁碟以便未來使用。要加載FAISS索引，我們可以使用以下函數：

 

“python
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)
”

Step 4. 通過文本查詢或參考圖像檢索圖像

 

當FAISS索引建立完成後，我們就可以通過文本查詢或參考圖像來檢索相似圖像。如果查詢是一張圖像，則首先使用PIL的`Image.open`打開該圖像，然後使用CLIP模型進行編碼。

 

“python
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-最近鄰（kNN）搜索，找到與查詢向量最相似的k個向量。kNN搜索使用的距離度量為餘弦相似度。函數返回查詢結果和檢索到的圖像路徑列表。

 

文本查詢範例：

 

現在，我們可以檢查搜索結果。您可以在提供的Colab筆記本中找到顯示結果的輔助函數`visualize_results`。例如，對於文本查詢「ball」，我們可以檢索到最相似的3張圖片。

 

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

 

參考圖像查詢範例：

 

使用圖像作為查詢，我們也能取得不錯的結果。例如，當我們使用一張眼睛繪畫的參考圖像進行搜索時，系統除了找到原始圖像，還找到了一張眼鏡和另一張不同繪畫的匹配圖像，展示了查詢圖像的不同語義層面。

 

"python
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在Zero Shot模型中展示了不錯的結果，但在處理分布外數據、細粒度任務時可能表現不佳，而且它可能繼承了訓練數據中的自然偏見。為了克服這些限制，您可以嘗試其他類似CLIP的預訓練模型，如OpenClip，或者在自定義數據集上對CLIP進行微調。

 

 

*本文譯自Lihi Gur Arie博士在Towards Data Science上發表的《Building an Image Similarity Search Engine with FAISS and CLIP》。

原文連結