cedro-blog

CLIPを使って、大量の画像の中から自分が探したい画像をテキストで検索する

1.はじめに

 2021年1月OpenAIは、画像と画像を説明するテキストのペア4億組を学習させた、汎用画像分類モデルCLIPを発表しました。今回は、このCLIPを使って、大量の画像の中から自分が探したい画像をテキストで検索するシステムを作ってみます。

2.CLIPとは?

 まず、CLIPに行っている事前学習の内容を見てみましょう。

 CLIPに入力されるのは、「画像」「その内容を説明するテキスト」のペアがN個です。Image Encoderは、N個の画像からその特徴を表すベクトル I_1〜I_N を取り出そうとします。TextEncoderは、N個のテキストからその特徴を表すベクトル T_1〜T_N を取り出そうとします。

 ベクトルには、2つのベクトルの内積(COS類似度)が大きいほど類似度が高く、内積が小さいほど類似度が低い、という性質があります。CLIPは、この性質を利用して、ペアであるベクトルの内積、I_1・T_1, I_2・T_2, I_3・T_3 , … , I_N・T_N を最大化し、ペアではないベクトルの内積を最小化するように、Image EncoderとText Encoderのパラメータを学習します。

 CLIPは、この事前学習によって、画像からもテキストからも、その特徴を適切に表すベクトルを取得できるようになります。次に、事前学習したCLIPを使って画像分類をする段階を見てみましょう。

 ラベルによる画像分類を考えます。ラベルは普通「単語」なので、「テキスト」を学習して来たCLIPには上手く解釈できません。なので、例えばラベルがdogであれば「A photo of a dog」という文にして、N個のラベルの文をText Encoderへ入力しラベルの特徴ベクトル T_1〜T_N を取り出します。

 一方で、分類したい画像からImage Encoderで画像の特徴ベクトル I_1 を取り出します。そして、画像の特徴ベクトルとラベルの特徴ベクトルの掛け算を行った結果をみて、一番大きい値となったところを分類すべきラベルと判断すれば良いわけです。

 ということで、CLIPは膨大な量の事前学習を行っているので、全く学習を行わずに分類タスクを解けるのです。

3.CLIPを画像検索に利用する

 それでは、大量の画像の中からテキストの内容に最もふさわしい画像を検索するには、どうしたら良いでしょうか。実は、CLIPを使うと簡単です。

 大量の画像から画像の特徴ベクトルI_1〜I_Nを抽出し、テキストからテキストの特徴ベクトル Tを抽出し、内積(COS類似度)が一番大きくなった画像を選べば良いわけです。それでは、早速コードを動かしてみましょう。

4.コード

 コードはGoogle Colabで動かす形にしてGithubに上げてありますので、それに沿って説明して行きます。自分で動かしてみたい方は、この「リンク」をクリックし表示されたノートブックの先頭にある「Colab on Web」ボタンをクリックすると動かせます。

 まず、セットアップを行います。

 CLIPのモデル化によって、clip.tokenize() でテキストをトークン(数字)にし、model.encode_text() でトークンからテキストの特徴ベクトルを抽出できるようになります。また、model.encode_image() で画像から画像の特徴ベクトルを抽出できるようになります。

 検索に使う画像は、有名な顔画像のデータセット CelebAから先頭の5,000枚を取り出したものをグーグルドライブからダウンロードして使います。画像サイズは、178×218です。

 次に、5,000枚の画像を読み込みます。

 img フォルダーからpng画像を1枚づつ読み、サイズを224×224に調整しアペンドします。そしてテンソルに変換し、正規化(輝度の平均を引き、輝度のバラツキで割る)します。結果、image_input のシェイプは(5000,3,224,224)となります。

 次に、検索するテキストを入力します。text = She is a charming woman with blonde hair and blue eyes(彼女はブロンドの髪で青い目をしたチャーミングな女性です)でやってみます。

 clip.tokenize() で、テキストをトークン(数字)に変換します。トークン化した text_input は余分な部分をゼロで埋め、シェイプは (1,77) になっています。

 それでは、image_inputtext_input を元に、COS類似度を計算します。

 model.encode_image() で画像の特徴ベクトル image_features を抽出し、model.encode_text() でテキストの特徴ベクトル text_features を抽出します。そして、torch.cosine_similarity() でその2つのベクトルのCOS類似度 text_probs を計算します。

 それでは、検索結果を見てみましょう。text_probs には5,000枚の画像のCOS類似度計算の結果が格納されているので、高い順にソートしてTOP3に該当するインデックスの画像ファイルを表示します。

 おっ!CLIP 中々やりますね。結構イケてる感じです。画像の下に表示されているのがCOS類似度の値です。COS類似度は数字が大きいほど類似度が高く、+1〜 −1の値をとります。

 もう1つやってみましょう。画像は読み込み済みなので、テキスト入力のところから再度実行すればOKです。text = He is a dandy middle-aged man wearing glass(彼はメガネを掛けた中年のダンディな男性)でやってみると、

 3番目の男性が少し若そうですが、5,000枚の中ではそこそこ類似度が高いのかもしれません。

 それにしても、CLIP便利です。色々な応用が考えられそうですね。では、また。

(オリジナルgithub)https://github.com/openai/CLIP