「あの写真どこいった？」を画像認識で解決してみる

はじめに

はじめまして、こんにちは。
フォルシアの基盤技術部でエンジニアをしている太田と申します。

突然ですが、みなさまはスマホ、デジカメ等々で写真を撮ることは多いでしょうか？
写真をよく撮られる方にあるあるのお悩みの一つに「あの写真どこ？」となることがあるかと思います。
今回はそんなお悩みが解決するような「文章による画像検索機能」を作成してみました。

キーワード画像検索機能について

今回の機能を作成するにあたって、基本機能やいくつかの要件を設けました。

• 機能概要
  • ブラウザ上で入力された検索キーワードの内容に対し、感覚的に「近い」画像を表示する機能。
    • 単語ではなく「○○の××」のような文章による検索が可能な、実際の検索時に直感的に使いやすい検索機能。
• 要件
  • Pythonで全処理ならびにブラウザへの表示まで含む機能実装を完結させる。
  • 今回は機械学習用データセットのうちの画像約3万枚分を検索対象とする。
    • データセットの実態としては、画像のリンクが大量に記載されたテキストデータになっているが、これらのリンク先画像を検索対象とする。
    • ローカルのストレージ上の画像検索を意図しているが、本質的な機能に相違はない認識。
  • ～5秒程度で検索部分の処理が完了すること。
  • 検索の処理にはGPUを用いない。

また今回の機能実装にあたり下記の記事を参考にさせていただきました。この場をお借りしてお礼申し上げます。

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

実施結果

百聞は一見に如かずということで検索結果画像をいくつか添付します。
主観的には想定通りの検索結果が表示されていると思います。
また、正味の検索にかかる時間も3秒弱であり、上記の要件を満たしています。

text = animal
検索時間 : 0:00:02.766672

text = car
検索時間 : 0:00:02.732000

詳細

ということで具体的な説明に入ります。
はじめに開発環境ですが、以下のとおりです。

Python 3.8.5
Ubuntu 22.04.2 LTS

いくつか方法はあると思いますが、素直に「画像を文章で検索する」という機能を要素分解すると下記のようになります。

＜手順＞

1. 画像を文章で表現する。
2. 画像を表す文章をベクトル化する。
   1. 以降、このベクトルを検索対象文ベクトルと表記します。
3. 検索文章（クエリ文）も同様にベクトル化する。
   1. 同様に以降、このベクトルをクエリ文ベクトルと表記します。
4. 検索対象文ベクトルとクエリ文ベクトルの「近さ」を取り、近い順に抽出する。
5. 結果をブラウザに表示する。

今回、手順3までの要素の実現に、CLIPというライブラリを利用しました。
CLIPとはOpenAIが開発した汎用画像分類モデル（を容易に利用可能にしているライブラリ）です。
特徴として、利用時に全く学習を行わずに画像分類が可能な、ゼロショット画像分類ができるという点があります。
また、自然言語処理的な技術も取り入れられているため、文章のベクトル化や、単語ではなく文章による検索も実現できます。

まず事前処理についてですが、上記手順2までを事前に実施しておくことにしています。
CLIPを用い、検索対象文ベクトルを約3万枚分の画像についてそれぞれ512次元ベクトルとして取得しています。
この処理にはGPU搭載のマシンを利用しており、3万枚分の画像のベクトル化には数時間程度かかっています。
この3万個のベクトルをローカルにnpyファイルとして、1ファイル当たりおよそ1万ベクトル分記載して保存しています。

make_caption_list.py

<略>

for image_file in image_list:

    image_ids = []
    image_vectors = []

    image_urls = np.load(f'/path/to/{image_file}')
    image_urls = image_urls.tolist()
    image_urls = [item.replace('-', '/') for item in image_urls]

    for i, image_url in enumerate(tqdm(image_urls)):
        try:
            with urllib.request.urlopen(image_url) as url:
                f = io.BytesIO(url.read())
            image = PIL.Image.open(f)
            image = preprocess(image.convert("RGB")).unsqueeze(0)

        # 画像をベクトル化
            with torch.no_grad():
                image_features = model.encode_image(image)
                image_features = image_features[0].to('cpu').detach().numpy().copy()

        # ベクトルをリストに追加
            image_ids.append(image_url.replace('/', '-'))
            image_vectors.append(image_features)

    ids_path = "/path/to/{:04d}_ids.npy".format(count)
    features_path = "/path/to/{:04d}_features.npy".format(count)

    np.save(image_ids, ids_path)
    np.save(image_vectors, features_path)

    count+=1

ここからは実際の検索処理です。
まず、検索文章もベクトル化する必要があるので、手順2までと同様にベクトル化を行い、512次元のクエリ文ベクトルを生成します。
ベクトル同士の計算についてですが、データベースを利用せずPythonのNumpyを利用してベクトル計算や類似度によるソート等を行っています。Numpyの利用により、検索時間の短縮を見込んでいます。
クエリ文ベクトルと検索対象文ベクトルそれぞれに対してcos類似度を計算し、類似度の高い上位3つの画像のURLをnpyファイルから抽出します。

keywordimagesearch.py

<略>

object_ids_path = FEATURE_KW_DIR / "{:04d}_ids.npy".format(blob_index)
features_path = FEATURE_KW_DIR / "{:04d}_features.npy".format(blob_index)

if not object_ids_path.exists():
    print("No features found")
    break

object_ids = np.load(object_ids_path, allow_pickle=True)
object_ids = np.array(object_ids)
features = np.load(features_path, allow_pickle=True)

assert len(object_ids) == len(features)

query_features = np.tile(text_features, (len(features), 1))

# コサイン類似度の計算
cos_similarities = (dot(query_features, features.T) / (norm(query_features) * norm(features, axis=1)))[0]

# 類似度が高い順にソート
similarity_indexes = np.argsort(cos_similarities)[::-1][:limit]

results.extend(zip(object_ids[similarity_indexes], cos_similarities[similarity_indexes]))

後はこれをブラウザへ表示することで、今回の「文章による画像検索機能」を実現しています。

実施結果を確認してみます。
まずは"animal", "car"といった単語で検索をしてみます。
"animal"に関しては剥製が表示されてしまっていますが、いずれも主観的には正しそうな検索結果であると思います。

text = animal
検索時間 : 0:00:02.766672

text = car
検索時間 : 0:00:02.732000

次に、形容詞を追加しクエリを少し複雑にして結果の違いを確認してみます。
"building"に対し、"buildings in a big city"ではより一般的なイメージに近い検索結果が表示されているかと思います。
また、検索時間は有意には変わっていないことも確認できました。

text = building
検索時間 : 0:00:02.336895

text = buildings in a big city
検索時間 : 0:00:02.387510

さいごに

今回開発した機能では、はじめに定めた要件は達成できました。
しかし、日本語対応していない、検索速度に改善の余地があるなど改善したい部分も明らかになったというのが所感です。
この改善点は今後継続的に改善し、適宜続きのブログを書いていければと思っています！

---

フォルシアでは現在、AI画像認識システムの開発ならびに導入支援サービスを展開しております。
少しでも気になった点がございましたら、是非お気軽にフォームよりお問い合わせください。

サービスご紹介ページ
お問い合わせフォーム

---

この記事を書いた人

太田 圭祐
2023年キャリア入社 エンジニア
このような記事を書いておいて恐縮ですが、私は日常生活でほとんど写真を撮らないタイプの人種です。