👀

FAISSとCLIPを使用した画像類似性検索エンジンの構築

Li Yachao

2024/11/05に公開

 FAISSとCLIPを使用した画像類似性検索エンジンの構築
 概要本記事では、FAISSとCLIPを使用して、テキストまたは画像をクエリとして利用し、画像データベースを検索する方法について説明します。FAISSの高速検索能力とCLIPのマルチモーダル理解を組み合わせて、効率的な画像検索エンジンの構築手順を紹介します。

 目次はじめに
FAISSとCLIPの概要
画像類似性検索の手法
コードの実装
おわりに

 1. はじめに
 背景大規模な画像データベースを効率的に検索するためには、FAISS（Facebook AIが開発した高速類似性検索ライブラリ）と、CLIP（OpenAIが開発したマルチモーダル学習モデル）の組み合わせが効果的です。FAISSは高次元データに対する高速な検索を実現し、CLIPはテキストと画像の両方を理解することができるモデルです。本記事では、この2つのツールを活用して、画像検索エンジンを構築する手順を解説します。

 2. FAISSとCLIPの概要
 FAISSの特徴FAISS（Facebook AI Similarity Search）は、大規模データセットの類似性検索を高速に行うためのライブラリです。特に高次元データに対して効率的に検索を行うことができ、GPUを使用することでさらに高速化が可能です。

 CLIPの特徴CLIP（Contrastive Language-Image Pretraining）は、テキストと画像の両方を理解し、それらを関連付けることができるマルチモーダルモデルです。テキストから画像、または画像からテキストへのクエリを可能にし、幅広いアプリケーションで利用されています。

 3. 画像類似性検索の手法画像類似性検索では、入力画像またはテキストを基に、データベース内の類似した画像を見つけることを目指します。以下は、FAISSとCLIPを用いて画像類似性検索を行う手順です。
特徴ベクトルの抽出
CLIPを使用して、画像とテキストから特徴ベクトル（高次元ベクトル）を抽出します。
データのインデックス化
FAISSを使用して特徴ベクトルをインデックス化し、高速な検索を可能にします。
クエリ処理
検索したい画像またはテキストをCLIPで特徴ベクトルに変換し、FAISSで類似性検索を実行します。

 4. コードの実装データセット作成
データセットを作成するために、Pexelsから52枚の異なるテーマの画像を収集しました。下図にはランダムに10枚画像を示します。



２.  特徴抽出
次に、各画像からCLIPを使用して特徴ベクトルを抽出します。

FAISSでインデックスを作成します。これにより、検索が高速化されます。
import faiss
from transformers import SentenceTransformer
import torch

# CLIPモデル
model = SentenceTransformer('clip-ViT-B-32')

# データベース画像の特徴ベクトルを抽出
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)

インデックス生成
FAISSでインデックスを生成します。
FAISSは、内積（IP）とL2（ユークリッド）距離を含む、さまざまな類似性距離計測を提供しています。
FAISSはさらに、さまざまな索引オプションを提供しています。検索速度と精度のバランスを取るために、近似または圧縮の技術を使用します。今回は"Flat"インデックスを使用します。Flatとは、クエリベクトルとデータセット内の各ベクトルの比較し、力まかせ検索方法です。精度が高いですが、計算複雑性も高いです。
def create_faiss_index(embeddings, image_paths, output_path):
    dimension = len(embeddings[0])
    index = faiss.IndexFlatIP(dimension) #内積類似とFlatインデックスの初期化
    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"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 = "/vector.index"
index = create_faiss_index(embeddings, image_paths, OUTPUT_INDEX_PATH)

#インデックスをロード
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)

クエリ実行
まず、CLIPモデルで、入力データの特徴ベクトルを抽出します。

CLIPは画像とテキスト２種類の入力データをサポートしています。

画像の場合はPILのImage.openで開け、テキストの場合はテキストのままで特徴ベクトルを抽出します。
  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

画像データで検索
      query ='./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)

事前学習済みモデルだけで、かなりいい結果を得ています。

参考画像としての目の絵を検索すると、元の画像を見つけるだけではなく、それに合致の眼鏡や別の異なる絵も見つかります。
テキストで検索
   query = 'ball'
   query, retrieved_images = retrieve_similar_images(query, model, index, image_paths, top_k=3)
   visualize_results(query, retrieved_images)
「animal」クエリとして検索してみたら：


 5. おわりに本記事で、CLIPとFAISSを使い、基本的画像類似性検索エンジンを構築しました。上記の検索結果により、メソッドの有效性を示しています。
今回事前学習済みモデルを基づき、もし自分のデータを使い、fine-tuningしてみたら、より良いモデルが得られるかもしれません。

 参照https://mp.weixin.qq.com/s/DNgpgVBBxVGNrjldr7MV_w
https://huggingface.co/docs/transformers/main/ja/model_doc/clip
https://github.com/facebookresearch/faiss/wiki