はじめに

OpenAIのCLIP (Contrastive Language-Image Pre-training) モデルを活用し、テキストによる画像検索を試してみました。
CLIPは、画像とテキストを同じベクトル空間に埋め込むことができるマルチモーダル学習モデルのため、テキストと画像の類似度を計算し、類似度の高い画像を取得することでテキストによる画像検索が可能になります。
実装した処理の概要は以下の通りです。

事前登録
- 検索対象となる画像を画像埋め込みに変換します
- 変換した画像埋め込みをベクトルデータベース（Pinecone）に登録します
検索
- 検索文字列をテキスト埋め込みに変換します
- 変換したテキスト埋め込みを使用して、Pinecone内で類似の画像埋め込みを検索します

当初、画像検索機能をNext.jsのAPI routesで実装しましたが、APIリクエストごとにCLIPモデルの再ロードが発生しうまく動作しませんでした。対応方法を見つけることができなかったため、結局、Expressのカスタムサーバに実装し直しました。サーバー起動時に一度だけCLIPモデルをロードし、メモリに常駐させることでモデルの再ロードを発生させないようにしました。

環境

Python 3.10 (データ登録に使用）
TypeScript 4.9.5
フレームワーク: Next.js 14（Express.jsを使用したカスタムサーバ）
画像・テキスト処理: @xenova/transformers ライブラリ
ベクトルデータベース: Pinecone

画像埋め込みの登録

ベクトルデータベース(Pinecone)を使用して画像検索させるため、Pythonのバッチ処理を用いて、ディスク上の画像データを埋め込みベクトルに変換し、生成されたベクトルをPineconeに登録しました。

モデルのロード

画像を埋め込みベクトルに変換するにあたり、まずCLIPモデルとその関連コンポーネントをロードしました（下記関数内）。

from transformers import AutoProcessor,CLIPProcessor, CLIPModel, CLIPTokenizer

# Set the device
device = "cuda" if torch.cuda.is_available() else "cpu"
model_ID = "openai/clip-vit-base-patch32"

def get_model_info(model_ID, device):
	model = CLIPModel.from_pretrained(model_ID).to(device)
	processor = AutoProcessor.from_pretrained(model_ID)
	tokenizer = CLIPTokenizer.from_pretrained(model_ID)
    # Return model, processor & tokenizer
	return model, processor, tokenizer

get_model_info関数内で、openai/clip-vit-base-patch32のCLIPモデル、前処理を行うprocessor、トークンを数値IDに変換するトークナイザーを取得します。

画像の埋め込みベクトル化

ロードしたCLIPモデルと前処理用のコンポーネントを使用して画像を埋め込みベクトルに変換する関数を作成しました。

def get_single_image_embedding(my_image):
    image = processor(images=my_image , return_tensors="pt")
    embedding = model.get_image_features(**image).float()
    # convert the embeddings to numpy array
    return embedding.cpu().detach().numpy()

次の処理を実施しています。

processorで入力画像をモデルに入力可能な形式のテンソルに変換
画像の特徴を表すベクトル（埋め込みベクトル）を取得しデータ型をfloat32に変換
PytorchテンソルからNumpy配列に変換して返す

上の関数では、画像1枚づつ処理していますが、processorに画像のリストを渡すことで複数の画像を一度に処理できます。

埋め込みベクトルの登録

画像を埋め込みベクトルに変換してベクトルデータベースに登録します。

from PIL import Image
from pinecone import Pinecone

 # ... (中間の処理を省略)

# Pineconeとの接続 
pc = Pinecone(api_key=pinecorn_api_key)
index = pc.Index(pinecorn_Index)
 
 # ... (中間の処理を省略)
 
        imageFeature_np = get_single_image_embedding(Image.open(filename))
        imageEmbedding = imageFeature_np[0].tolist()
        # imageEmbedding = imageFeature_np.flatten().tolist() 
        
        metaInfo = createMetaInfo(webCamInfo)
        upsert_response = index.upsert(
            vectors=[
                {
                    'id' : webcamId,
                    'values': imageEmbedding,
                    'metadata': metaInfo,
                }
            ],
            namespace="webcamInfo"
        )

ディスク上の画像をget_single_image_embeddingで埋め込みベクトルに変換し、Pythonのリストに変換しています。次のindex.upsertの引数にセットさせるため試行錯誤し、Pythonのリストに変換しました。
index.upsert メソッドを使用して、ベクトルデータベースに情報を挿入または更新します。挿入するデータには以下を含めました。

id: 一意の識別子
values: 画像の埋め込みベクトル
metadata: 追加情報(画像のURL等の情報を含めていますが、説明は省きます）

画像検索の実装

画像のベクトル表現をPythonで作成し、ベクトルデータベースに登録しましたが、検索処理はTypeScriptで行いました。CLIPモデルを使ってテキストをベクトルに変換し、データベース内の画像ベクトルと類似度の高いものを探し出します。

モデル初期化

Embeddingクラスを作成し、サーバー起動時にコンストラクタ内で一度だけCLIPモデルをロードし、以降の検索のリクエストでインスタンスを再利用しました。

import {Processor,AutoProcessor,RawImage,AutoTokenizer,CLIPTextModelWithProjection,CLIPVisionModelWithProjection, PreTrainedTokenizer} from '@xenova/transformers';

class Embedding {
    private model: CLIPVisionModelWithProjection |null = null;
    private tokenizer: PreTrainedTokenizer | null = null;
    private textModel: CLIPTextModelWithProjection | null = null; 
    private imageProcessor: Processor | null = null;

    constructor() {
        this.initializeModels();
    }
    
    private async initializeModels() {
        const model_id = process.env.MODEL_ID; //Xenova/clip-vit-base-patch32
        console.log(`start initialize model ${model_id}`);
        
        this.model = await CLIPVisionModelWithProjection.from_pretrained(model_id as string);
        this.tokenizer = await AutoTokenizer.from_pretrained(model_id as string);
        this.textModel = await CLIPTextModelWithProjection.from_pretrained(model_id as string);
        this.imageProcessor = await AutoProcessor.from_pretrained(model_id as string);

        console.log(`end initialize model ${model_id}`);
    }

    async getTextEmbedding(text: string): Promise<number[]> {
        // 以後のコードは後述
    }
}

class内のプロパティに次のものを定義し、

model: 画像処理用のCLIPビジョンモデル
tokenizer: テキストをトークン化するためのオートトークナイザー
textModel: テキスト処理用のCLIPテキストモデル
imageProcessor: 画像の前処理を行うプロセッサー

initialize()メソッド内で、それぞれ設定しました。

CLIPVisionModelWithProjection.from_pretrained(): 画像処理用のCLIPモデルをロード
AutoTokenizer.from_pretrained(): テキストトークナイザーをロード
CLIPTextModelWithProjection.from_pretrained(): テキスト処理用のCLIPモデルをロード
AutoProcessor.from_pretrained(): 画像プロセッサーをロード
それぞれ'Xenova/clip-vit-base-patch32'という事前学習済みモデルを使用しました。

テキスト埋め込み

検索用の文字列をテキスト埋め込みのベクトルに変換します。また、次のgetTextEmbedding関数は、Embeddingクラス内に定義しています。

    async getTextEmbedding(text: string): Promise<number[]> {
      const textInputs = await this.tokenizer(text, { padding: true, truncation: true });
      const { text_embeds } = await this.textModel(textInputs);
      return Array.from(text_embeds.data);
    }

initialize()でロードしたコンポーネントを使用して、次のような処理を実行しています。

テキストのトークン化
トークンをもとにテキストの埋め込みベクトルの生成
生成したテキストの埋め込みベクトルをJavaScript配列に変換して返す

検索処理

作成したテキストの埋め込みベクトルを用いて、Pinecone(ベクトルデータベース:index名 wemcamInfo)から類似画像を検索します。

import { Index } from '@pinecone-database/pinecone';

    async searchImages(query: string, count: number) {
      const queryEmbedding = await embedding.getTextEmbedding(query);
      const response = await pineconeIndex.namespace('webcamInfo').query({
        topK: count,
        vector: queryEmbedding,
        includeValues: false,
        includeMetadata: true
      });

      # 以後の処理は省略
      
    }

上記コードのembeddingは、Embeddingクラスのインスタンス。getTextEmbedding(query)でテキストの埋め込みベクトルを作成し、ベクトルデータベースのindexに検索リクエストを実行します。検索方法は、コサイン類似度を使用していますがこれは、Pineconeで、index作成時に設定しています。
query関数の引数は次のとおりです。
・topK: count: 最も類似度の高い結果を指定された数(count)だけ返します。
・vector: queryEmbedding: 先ほど生成したクエリの埋め込みベクトルを使用して検索します。
・includeValues: false: 結果にベクトル値を含めません。
・includeMetadata: true: 結果にメタデータを含めます。

responseには、id、メタデータ、コサイン類似度のスコアのJSON配列がtopKで指定した数分含まれます。

おわりに

以前作成した画像検索システムを振り返ってみました。Clipモデルの扱いも含めて、Next.jsに関する知識が不足していたため、どのようなアーキテクチャで構築すべきだったか疑問が残りました。特にモデルの再ロードを防ぐためにNext.jsのカスタムサーバで構築したのですが、Vercelにデプロイできず代替案を検討しようと思いました。