Zenn
🧊

アニメ風イラスト特化の画像検索ツールを作ってみた - Vision Transformer x LSI x StreamLit

2024/10/18に公開
ディープラーニング
画像検索
イラスト
visiontransformer
潜在的意味インデクシング
tech

どうも、 ryo_grid です。
　
手元に大量の画像ファイルがあって、うまいこと検索できないかなあと思い、ちょいとコードを書いてみました。

作ったものの概要

  • 深層学習系の技術も活用した画像検索アプリケーション
    • => Vision Transformer（ViT）を内部的に利用したTagger
  • Googleフォトのようなサービスでいまいち検索が効かない2次元イラスト画像に特化
    • ただし、今回の実装では英語のタグ（キーワード）での検索のみ
    • 日本語非対応
  • ローカルにあるファイルの検索のためのアプリケーションであり、どこかにデータをアップロードするといったことが必要ない
  • 潜在的意味インデクシング（LSI）とかもやってみた

使用技術

画像についたタグ群とクエリ内のタグ群を、インデックス対象の画像群についたタグ群の集合（正確には行列）からLSIで生成した埋め込みモデルを利用してベクトル化し、それらの間の類似度をとることで検索を行います。

  • 画像へのタグ付け: Visual Transformer (ViT) Tagger
  • 埋め込みモデル: LSIで作成。実質的にはタグ付けの精度をカバーするために使用

LSIを適用することで、検索をかけたいニッチなタグAがあったとして、それがTaggerによって認識できない場合が多くとも、共起頻度の高いタグ群があれば、それらがついている画像が、Aで検索した場合でもヒットすることが期待できます。
なお、Web UIはStreamLitを使用して実装しました。

環境

少なくとも以下の環境では動作実績があります

  • （OS: Windows 11 Pro x64 (23H2)）
  • （Processor: AMD Ryzen 7 5700X 8-Core Processor 4.50 GHz）
  • Python: 3.10.4
  • pip: 22.0.4

インストール and 利用の手順

  • リポジトリをcloneしてきて中に移動します

    $ git clone https://github.com/ryogrid/anime-illust-image-searcher.git; cd anime-illust-image-searcher

  • 必要なライブラリをインストール

    $ pip install -r requirements.txt

  • 画像ファイルのあるディレクトリを指定して各画像にタグ付けを行います

    $ python tagging.py --dir "画像ファイルが含まれるディレクトリのパス"
    • ディレクトリは再帰的に検索されます
    • ファイル数によっては結構時間がかかるかもです
      • 前記のプロセッサスペックにて1ファイルあたり約0.5秒 (GPU等によるアクセラレーションは無し）
    • 実行が終わると画像ファイルに対応するタグの情報が tags-wd-tagger.txt に保存されます

  • 後でLSIをやる際に必要なので、タグ数をカウントします

    $ python counttag.py
    • 標準出力にユニークなタグ数の集計結果が表示されます

  • LSI適用後の次元数を指定してモデルと検索用のインデックスを生成します

    $ python genmodel.py --dim モデル次元数
    • LSIの内部処理: 特異値分解による次元削減
    • 次元数（削減後の次元数）はcounttag.pyで確認した値 x 0.8 程度で良いかなあと考えています
    • 結構時間がかかると思います（汗）
      • 約1000ファイルを対象に、800次元から700次元に落とす場合で数分程度
      • 約34万ファイルを対象に、7500次元から5000次元に落とす場合で約3.4時間
      • 計算量のオーダーとしては、計算手法にもよるのでこれで間違いないというものが見つからなかったのですが、元の次元か、削減後の次元のいずれかの2乗には比例すると思われます...

  • インデックスができたのでアプリを起動します（Web UI）

    $ streamlit run webui.py
    • pythonコマンドではなくstreamlitのコマンドを叩く必要があるので注意
    • streamLitによってローカルで起動されたアプリケーションサーバへアクセスする形で、自動的にブラウザでアプリが開きます

補足

  • 自身の画像ファイル群にタグをつけた結果現れたタグでしか検索はできません
    • 検索で使えないタグをクエリに含めるとエラーになります
      • アプリが落ちたりはしないのでご安心を
    • その場合、ちょっと手間ですが、タグ付けした結果が保存される tags-wd-tagger.txt をgrepなり秀丸やらの巨大なファイルが開けるエディタで検索して確認してみて下さい
  • tagging.py は onnxruntime で GPU や Apple Silicon の Neural Engine 等々のアクセラレータが使われるようにスクリプトを修正すれば高速化できるかもです
  • genmodel.py も gensim が GPU なんかを使ってくれるようにすれば速くなるかも
    • gensimは GPU 非対応でした...
  • Taggerには SmilingWolf/wd-vit-tagger-v3 を使用しています
    • ViT な Tagger のモデルはそれなりに存在しますが、このモデルはアニメ風イラストにおおむね特化して学習が行われています

便利機能

  • クエリの各タグに重みづけができます
    • 整数のみ
    • 負の重みも指定可（ネガティブプロンプトみたいなもの）
    • 例
      • "girl:3 dragon"
      • "girl:2 boy:3"
      • "girl dragon:2 boy:-3"
  • スライドショー機能
  • 検索結果のエクスポート機能
    • 検索結果をテキストファイルに出力させられます
      • webui.py が実行されたシェルでのカレントパスにクエリの内容とタイムスタンプから成るファイル名で作成されます
    • 文字コードはWindows環境だとsjis、それ以外だとutf-8
    • Irfan View などのソフトウェアでは、テキストファイルなファイルリストを読み込んでごにょごにょしたり、スライドショーしたりできるので、餅は餅屋で、ビューワ機能はその手のものに求めるた方がいい気がします

Python環境不要のバイナリパッケージ

  • Windows環境向けのものだけどうにか用意しました
    • v1.0.0
    • 使用方法はPythonスクリプトを利用する場合と基本的に同じです
    • 詳細はリポジトリに記述した説明をご参照下さい

TODO

  • インクリメンタルなインデックスの更新機能の追加
  • 画像ファイルを指定して類似画像検索を行う機能の追加

スクリーンショット

約1000枚のイラストを使用したデモでのスクリーンショットです。
画像ファイルは いらすとや さんのものを使わせてもらいました。
一部UIが最新のものと異なりますがその点はご容赦下さい。

  • 「standing」での検索結果
    standing 検索結果

  • 「standing animal」での検索結果
    standing animal 検索結果

  • 各画像の情報も見られます

    • image

  • スライドショー機能

    • こんな感じの表示で5秒間隔でスライドショーさせられます。終わりまで行くと最初に戻ります
    • image

Enjoy!

Discussion