どうも、 ryo_grid です。
手元に大量の2次元イラスト画像があって、うまいこと検索できないかなあと思い、ちょいとコードを書いてみました。
- Anime Style Illustration Specific Image Search App with ViT Tagger x BM25/Doc2Vec
作ったものの概要
- 深層学習系の技術も活用した画像検索アプリケーション
- => Vision Transformer（ViT）を内部的に利用したTaggerとDoc2Vec
- Googleフォトのようなサービスでいまいち検索が効かない2次元イラスト画像に特化
-
Taggerをアニメ風イラストに特化したものとしている
- （根本的に実現方法が違うと思われるので、Taggerどうこうの問題ではおそらくないですが）
- ただし、英語のタグしかつけてくれないので、必然的に英語でしか検索はできない
-
Taggerをアニメ風イラストに特化したものとしている
- ローカルにあるファイルの検索のためのアプリケーションであり、どこかにデータをアップロードするといったことが必要ない
- Doc2Vecで作成したモデルを用いて、タグ文字列のリストからベクトルを生成する
使用技術
画像についたタグ群とクエリ内のタグ群を、インデックス対象の画像群についたタグ群の集合（正確には行列）からLSIで生成した埋め込みモデルを利用してベクトル化し、それらの間の類似度をとることで検索を行います。
- 画像へのタグ付け: Visual Transformer (ViT) Tagger
-
検索処理: 複数メトリクスのアンサンブルとリランキング
- 利用するメトリクス
- 埋め込みモデルを用いた画像のタグと検索クエリのベクトル化とそれらの間の類似度（コサイン類似度）
- 埋め込みモデルはDoc2Vecで作成。実質的にはタグ付けの精度をカバーするために使用
- bm25で求めた類似度
- 埋め込みモデルを用いた画像のタグと検索クエリのベクトル化とそれらの間の類似度（コサイン類似度）
- 利用するメトリクス
Doc2Vecを用いることで、検索をかけたいニッチなタグAがあったとして、それがTaggerによって認識できない場合が多くとも、共起頻度の高いタグ群があれば、それらがついている画像が、Aで検索した場合でもヒットすることが期待できます。
リランキング（後述）の仕組みも同様の効果が期待できるはずです。
- 結果出力までのステップ
- 1） 上の2つでのスコアを1:1でアンサンブルした結果をスコアとして一旦内部的にランキングを作成
- 2） 上位10個の検索結果のベクトルを類似度で重みづけした上で平均をとったものを、クエリから得られた埋め込みベクトルの代わりとして扱って類似度を算出
- ユーザがクエリをあれこれトライ&エラーした結果として上位の結果が期待通りであるとなった場合、それらの画像についているタグはクエリの内容を精度の観点で補完してくれるものであろうという仮定
- 上位が期待通りの結果になっていない時点で行うリランキングのステップは無駄な処理とも言えるが、UI的に、ユーザがこれでいける、と思った時だけリランキングを指示させるというのは分かりにくいので毎回行うように実装した
- 3） ①でのスコアと②でのスコアを7:3の割合でアンサンブルして最終的なスコアとし、ソートしてランキング表示（リランキング）
- 表示するトップ10個は変更しない。理由はリランキングの前のスコアで選ばれたトップ10の結果がユーザの期待通りの画像であることがリランキングによる結果改善の前提だが、ユーザがリランキング前の結果を見られないとユーザが前提を成立させる方向にクエリを改善できなくなってしまうため
なお、Web UIはStreamLitを使用して実装しました。
環境
少なくとも以下の環境では動作実績があります
- （OS: Windows 11 Pro x64 (23H2)）
- （Processor: AMD Ryzen 7 5700X 8-Core Processor 4.50 GHz）
- （Memory: 64GB）
- 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ファイルあたり約1.7秒 (GPU等によるアクセラレーションは無し）
- 実行が終わると画像ファイルに対応するタグの情報が
tags-wd-tagger.txtに保存されます
-
Doc2Vecのモデルと検索用のインデックスを生成します
$ python genmodel.py
-
インデックスができたのでアプリを起動します（Web UI）
$ streamlit run webui.py- pythonコマンドではなくstreamlitのコマンドを叩く必要があるので注意
- StreamLitによってローカルで起動されたアプリケーションサーバへアクセスする形で、自動的にブラウザでアプリが開きます
補足
- 自身の画像ファイル群にタグをつけた結果現れたタグでしか検索はできません
- 検索で使えないタグをクエリに含めるとエラーになります
- アプリが落ちたりはしないのでご安心を
- その場合、ちょっと手間ですが、タグ付けした結果が保存される tags-wd-tagger.txt をgrepなり秀丸やらの巨大なファイルが開けるエディタで検索して確認してみて下さい
- 検索で使えないタグをクエリに含めるとエラーになります
- tagging.py は pytorch v2.4.1 でGPUが叩けるようCUDAライブラリとcuDNNライブラリをセットアップすれば速くなるかもしれません
- CUDAインタフェースが通るようであれば使うようには実装してあります
-
genmodel.py も gensim が GPU なんかを使ってくれるようにすれば速くなるかも- gensimは GPU 非対応でした...
- Taggerには "WD EVA02-Large Tagger v3" model を使用しています
- ViT な Tagger のモデルはそれなりに存在しますが、このモデルはアニメ風イラストにおおむね特化して学習が行われています
便利機能
- クエリの各タグに重みづけができます
- 整数のみ
- 例
- "girl:3 dragon"
- "girl:2 boy:3"
- 検索対象とする画像について、指定したタグを必須もしくは排除する
- 必須とする時: 重みの数値の前に'+'をつける
- 例: "girl:+2 dragon:2 boy"
- 上のようにした場合 'girl' とタグのついていない画像は検索結果に現れなくなります
- 排除する時: 重みの数値の前に'-'をつける
- 例: "girl:2 dragon:2 boy:-1"
- 上のようにした場合 'boy' とタグのついた画像は検索結果に現れなくなります
- 数値の値は意味を持たない形になりますが、実装上省略はできません...
- 必須とする時: 重みの数値の前に'+'をつける
- スライドショー機能
- 検索結果のエクスポート機能
- 検索結果をテキストファイルに出力させられます
- webui.py が実行されたシェルでのカレントパスにクエリの内容とタイムスタンプから成るファイル名で作成されます
- 文字コードはWindows環境だとsjis、それ以外だとutf-8
- Irfan View などのソフトウェアでは、テキストファイルなファイルリストを読み込んでごにょごにょしたり、スライドショーしたりできるので、餅は餅屋で、ビューワ機能はその手のものに求めた方がいい気もしていたりします
- 検索結果をテキストファイルに出力させられます
Python環境不要のバイナリパッケージ
残TODO
インクリメンタルなインデックスの更新機能の追加- 画像ファイルを指定して類似画像検索を行う機能の追加
スクリーンショット
約1000枚のイラストを使用したデモでのスクリーンショットです。
画像ファイルは いらすとや さんのものを使わせてもらいました。
一部UIが最新のものと異なりますがその点はご容赦下さい。
-
「standing」での検索結果
-
「standing animal」での検索結果
-
各画像の情報も見られます
-
スライドショー機能
- こんな感じの表示で5秒間隔でスライドショーさせられます。終わりまで行くと最初に戻ります
Enjoy!