wd-taggerをC++から使う方法(onnxruntime事始め)

はじめに

久しぶりのプログラムの話です。
いまやもはや猫も杓子も生成AIということで、我々プログラマも生成AIをプログラムでどう使っていくか…というのを知っておかないといけない時代になってきました。
あのInterface誌でも2024年8月号は「プログラミングで体験 生成AI」という特集が大々的に組まれており、プログラミングの面からAIに対してどうやって行くべきかなどの特集が組まれたりする昨今です。

AIといっても色々ありますが、今回は比較的とっつきやすく面白く非常にわかりやすい題材として、wd-taggerでいろんな画像を判定してみようという話を進めてみます。

実はwd-taggerのサンプルなどはpythonでは既にそこら中に実装があるのでpythonを使ったほうがはるかに圧倒的に楽ではあるのですが、コードを読んだ所そんなに難しいことはしていないので、今回は敢えてC++で書くとどうなるか…といった感じで、C++でのサンプルと簡単な解説をここで載せてみます。

なお、今回ボクが作ったコードサンプルは以下です
https://github.com/chromabox/wd-tagger-cpp-sample/

今回はあくまでサンプルということでGUI系ライブラリは使わずで、ターミナル(コマンドプロンプト)から動くものを作りました
環境としてはUbuntu 22.04なのですが、さほど特殊なことはしていないので、CMakeなどの設定をちょっと変更すればWindowsでもビルドできるかなと思います。

使い方としては…

$ ./wdtagger 画像ファイル名

という感じです。

例として、ボクがGithubで使っているアバターの以下の絵

を読ませると次のような結果になります

$ ./wdtagger 846461.png 
read label file...
read label ok.
loading wd-tagger model... 
load wd-tagger model success
Running predict...
predict success.
result -----  
ratings: 
   general: 0.0733535
   sensitive: 0.920688
   questionable: 0.00349146
   explicit: 0.000387698
tags: 
   1girl: 0.991317
   solo: 0.98317
   instrument: 0.866745
   guitar: 0.774588
   brown_hair: 0.746188
   purple_eyes: 0.68452
   animal_ears: 0.629489
   short_hair: 0.615854
   bow: 0.60101
   chibi: 0.375334
charas: 
   ushiromiya_maria: 0.667614
   ryuuguu_rena: 0.0163996
   chen: 0.00753877
   hirasawa_yui: 0.00524783
   yakumo_ran: 0.00188139
   houjou_satoko: 0.00141031
   misty_(pokemon): 0.00120723
   souryuu_asuka_langley: 0.0010246
   may_(pokemon): 0.00101718
   alice_margatroid: 0.000895023

画像を読ませるとなにやら文字列と数値が出てきました。
次からはもうちょっとこのあたりについて詳しく解説します

wd-taggerとは

そもそもwd-tagger とはなんぞや？ということですが、wd-taggerとは waifudiffusion tagger の略称で、主に二次元のイラストを入れるとそのイラストを構成する要素を単語(タグ)で答えをつらつらと返してくれるというものです。
先程の例もつらつら英単語が出てきましたが、あれは画像を構成する要素はおおよそこれだろうという答えを出してくれているということになります。

wd-tagger を作った作者（SmilingWolf氏）が用意してくれているウェブ上から動くサンプルは以下になります
https://huggingface.co/spaces/SmilingWolf/wd-tagger
(全くの余談ですが、海外オタクの用語で「二次元の俺の嫁」は「wife」ではなく「waifu」というらしいですはい...ということでwaifu diffusionなわけですね)

たとえば、このずんだもんの画像(元ネタ https://zunko.jp/sozai/zundamon/zunmon026.png) を突っ込むと以下のような結果を返します。
(SmilingWolf/wd-v1-4-vit-tagger-v2の場合)

1girl, solo, green hair, hair ornament, japanese clothes, flower, hair flower, short hair, kimono, full body, sandals, yellow eyes, white background, short kimono, personification, looking at viewer, simple background, standing, sash, leaf, obi, blush, wide sleeves

一人の女の子、ソロ、緑の髪、髪飾り、日本の衣装、花、髪の毛に花、ショートヘア、着物…ということでこのずんだもんの絵の要素を認識していることがわかります。

これの何処にAIを使っているのかといえば、ずばり画像の要素判定に使っています。
wd-taggerのモデルは、予めいろんな画像で「こういう画像がきたらこの単語を返す」というような学習をされているのでこのようなことが出来る訳です。
当然のことながらこれはAIモデルであるので、学習していない画像がやってきても適切な単語を返します。
つまり柔軟性も多少はあるよということです。

wd-taggerにもいろいろなモデルがありますが、wd-v1-4-vit-tagger-v2の仕組みを示した模式図はこのようになっています。
(煩雑すぎるためNN間の線は省略していますが全部つながってると思ってください)

入力層が縦横Pixel(448x448)と、RGB値の3に対応しています。で、真ん中に学習済みのニューラルネットモデルがあって、出力層が9083個あります。
出力層の配列に入っている値は、その要素がどれだけ含まれているかの強さが０〜１までの浮動小数点で入っています。

なおこの9083の層はそのままモデルと一緒に配布されている、このcsvファイルの行数+1と対応しています
https://huggingface.co/SmilingWolf/wd-v1-4-vit-tagger-v2/raw/main/selected_tags.csv

csvファイルを抜粋するとこんな感じになっています

tag_id,name,category,count
9999999,general,9,807858
9999998,sensitive,9,3771700
9999997,questionable,9,769899
9999996,explicit,9,560281
470575,1girl,0,4225150
212816,solo,0,3515897
13197,long_hair,0,2982517
8601,breasts,0,2323580
469576,looking_at_viewer,0,2089971
3389,blush,0,2040471
(略)

つまり出力層[4]は1girlの強さを表している…ということになります。
出力層[6]だとlong_hair…つまり「長い髪の毛が絵に含まれているかどうか」を示します。
出力層[6]値が0だと絵に髪の毛の要素は無いかも知れない、この値が0.9999とかだと絵に含まれている髪の毛は長い…ということになります

この仕組みはよく見ると、機械学習の典型的な例題でよく出てくるMNIST(0〜9の手書き文字をAIに判定させ、0〜9か答えを得る)と似ています。
というわけで、処理の流れは以下になります。

１．機械学習モデルをロードする
２．CSVを読み込む
３．画像を縦横448x448、チャンネルはRGBに直す
４．ロードしたモデルに変換した画像をwd-taggerモデルへ流す
５．結果をCSVと照らし合わせる

こんな感じで、一見すると小難しそうなAIで画像判定もそんなに難しくはないことがわかるかと思います。

なお、ボクがC++で作ったサンプルは以下になります。
https://github.com/chromabox/wd-tagger-cpp-sample
このコードを見ながら、以下読み進めていただくとわかりやすいかと思います。

onnxruntimeでのモデルロードなど

wd-taggerはOnnx形式のモデルであるため、onnxruntimeを使う必要があります。
onnxruntimeはpythonの他にC/C++でも利用可能なライブラリも提供されているのでこれを使用します。
ライブラリはここのRelaeseページからダウンロードができます。
https://github.com/microsoft/onnxruntime/releases
Windows環境だとDLLなどですが、Linux(ubuntu)だと共有ライブラリ(.so)とincludeファイルが同封されています。

これをビルド時にリンクすればいいわけですが、CMakeだと以下の記述になるかなと思います。

set(INC_DIRS
  ${CMAKE_CURRENT_SOURCE_DIR}/onnxlib/onnxruntime-linux-x64-1.18.0/include
)

target_include_directories(${PROJECT_NAME} PUBLIC ${INC_DIRS})

target_link_directories(${PROJECT_NAME} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/onnxlib/onnxruntime-linux-x64-1.18.0/lib)
target_link_libraries(${PROJECT_NAME} PRIVATE onnxruntime)

サンプルの完全なCMakeのCMakeLists.txtは以下になります
https://github.com/chromabox/wd-tagger-cpp-sample/blob/master/CMakeLists.txt

リンクをせずに、プラグインみたいな形で動的にonnxruntimeそのものをロードすることも出来ますが、少しややこしいのでここでは触れません。
ただプラグインみたいな形式でロードできると、onnxruntimeはCUDA向けとか他のGPU向けとかあるのでそのあたりで選択肢が増やせるのかなと思います。

モデルのロードは簡単で、以下のようになります

#define WD_MODEL_ONNX	"../models/wd-vit-tagger-v2/model.onnx"
...

	Ort::Session ortsession{nullptr};
	Ort::SessionOptions sessionOptions;
	Ort::Env ortenv(ORT_LOGGING_LEVEL_WARNING,"wdtagger");

	// Onnxモデルのロード
	std::cout << "loading wd-tagger model... " << std::endl;
	try{
		ortsession = Ort::Session(ortenv, WD_MODEL_ONNX, sessionOptions);
	}catch(Ort::Exception& e) {
		std::cerr << "ort::session Error:" << e.what() << std::endl;
	}
	std::cout << "load wd-tagger model success" << std::endl;

Ort::Sessionがモデルをロードしたセッションで、これに対して画像を入力して推論をさせたりします。
オプションや環境(Ort::Env)は特に指定なしです。

モデルへ渡したり受け取ったりするデータ形状はテンソルにする必要があります
モデルをロードするとそのモデルが必要としているテンソル形状などがわかるためそれを取得して、テンソルを作り、その中にデータを入れる必要があります
テンソルと配列の違いについては以下がわかりやすいです
https://ml-wiki.sys.affrc.go.jp/ai-workshop/_media/%E3%83%86%E3%83%B3%E3%82%BD%E3%83%AB.pdf
表現の上では多次元配列とそんなに変わりませんが、数学的意味合いが違うということのようです(ボクもこの辺はあまりよく理解していません...)

ということで、必要なテンソル形状を取得するには以下のようにします。

	// モデルが要求する形状などをとってくる
	std::vector<int64_t> input_shapes = ortsession.GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
	std::vector<int64_t> output_shapes = ortsession.GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();

onnxモデルは複数の入出力形状を取ることが出来るため、本来はGetInputTypeInfo(0)と決め打ちをせずにまずは入力形状がいくつ存在するかを見る必要がありますが、今回は題材がわかっているので0を決め打ちしています。
帰ってくるvectorには配列要素数が入ってきます。
wd-vit-tagger-v2のモデルの場合、input_shapesには「1,448,448,3」が入っています。
これは「バッチ数,横画像サイズ,縦画像サイズ,色チャンネル数」ということで、448x448x3の配列を用意すれば良いということになります。

output_shapesには「1,9083」ということで、そのまま9083の要素をもった配列を用意すればOKです。

csvファイルを読み込む

モデルに付随しているcsvファイルを読み込む必要があります。
出力層の各配列が何の画像要素を示しているのかということを表しています。

この辺りの処理はPythonだとリストやcsv用外部モジュール類などを用いて楽に出来ますが、C++の場合は自力で行う必要がある作業です。
具体的なコードはloadlabelあたりを見てもらうということにして、csvファイルの中身は次のような感じで区切られています。

タグID,要素名,カテゴリ,出現カウント

このうち、「要素名」と「カテゴリ」が必要になるのでこれを取得して記憶しておきます。
カテゴリはこのタグ要素が一般のタグなのか、レーティングを示したタグなのか、キャラ名を示したタグなのかが番号で記されています。
よって「カテゴリ」内で各種ソート処理をする必要があるのなら、カテゴリごとに分類もしておくこともあるでしょう。

実装方法は色々あると思いますが、ボクはこのようなクラスを作ってマスターのVectorにとりあえず全部記憶させておき、カテゴリごとに分けてクラスのポインタだけを記録するVectorをカテゴリごとに作っています。

class TaggerLabel
{
public:
	std::string		name;				// ラベル名
	int				category;			// カテゴリ
	float			score;				// 推論した結果を入れるためのスコア

public:
	TaggerLabel(const std::string &sname, const std::string scategory_str){
		name = sname;
		category = ::atoi(scategory_str.c_str());
		score = 0;
	}
};
typedef std::vector<TaggerLabel>		TaggerLabelVec;
typedef std::vector<TaggerLabel*>		TaggerLabelPtrVec;

画像を縦横448x448、チャンネルはRGBに直す

次は入力された画像をモデルが読める形に変換する必要があるのですが、画像ファイルをロードして…とかこれを純粋C++で作るのはなかなかしんどい作業になるので、ここは流石にOpenCVを使って楽します。

	// イメージをファイルから読む
	cv::Mat src_image;
	src_image = cv::imread(argv[1], 1);
	if (! src_image.data){
		std::cerr << "cv::imread Error: can not read image" << std::endl;
		return -1;
	}
	// TODO: チャンネルが3以外は今回パス
	if (src_image.channels() != 3) {
		std::cerr << "sorry this program color channel 3 image only..." << std::endl;
		return -1;
	}
    ...
	cv::Size model_insize;
	model_insize.width = (int)input_shapes[1];		// batch,width,height,channelの順番
	model_insize.height = (int)input_shapes[2];
    ...
	cv::Mat tgt_image;
	cv::resize(src_image,tgt_image, model_insize, 1, 1, cv::INTER_CUBIC);
	if(! tgt_image.isContinuous()){		// 念の為
		std::cerr << "Error : cv::mat is not memory continuous..." << std::endl;
		return -1;
	}

本来の実装では無理やりResizeをするのではなく元画像の縦横比率を維持してリサイズ後、余白は白で埋めるという作業が必要ですが簡略化のためこうしています。
チャンネルをRGBのみにしてAは抜く必要がありますが、今回は例なのでチャンネル数3以外は見ないということで...

これで、モデルが必要とする448x448x3の画像を用意できましたが、最終的に浮動小数点のテンソルにしないといけないため以下のようにします。

	// モデルへの入力データの確保
	std::vector<float> model_input_data(1 * model_insize.width * model_insize.height *3);    
	// 浮動小数点型へ直す
	{
		unsigned char* isrc = tgt_image.data;
		for (auto it = model_input_data.begin(); it != model_input_data.end(); ++it){
			*it = static_cast<float>(*isrc);
			isrc++;
		}
	}

OpenCVのMatは内部的にはBGRの並びで画素を記憶しており、wd-taggerのモデルもBGRの並びで画像を学習しているので、上のようなベタなコピーでも許されます。
しかし、OpenCV以外のライブラリを使って実現する場合は浮動小数点側の配列がBGRになるように、並び順に注意してコピーしなければなりません。

モデルに変換した画像を流し、推論を実行する

これで推論に必要なものは揃ったので、いよいよ推論…ですが、前にも「モデルにはテンソル形式で渡さないといけない」と述べた通り、C++のVector型をそのままの形では受けてもらえません。
入出力用のテンソルを作成する必要があります。
onnxruntimeでは、Ort::Value型でやり取りする必要があり以下のようにします。

	int	model_outsize = (int)output_shapes[1];		// batch,出力サイズの順番
	// モデルへの出力データの確保
	std::vector<float> model_output_data(1 * model_outsize);
	//入出力用のテンソルを作成する
	Ort::MemoryInfo ortmem(Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtDeviceAllocator, OrtMemType::OrtMemTypeCPU));
	Ort::Value input_tensor = Ort::Value::CreateTensor<float>(ortmem, model_input_data.data(), model_input_data.size(), input_shapes.data(), input_shapes.size());
	Ort::Value output_tensor = Ort::Value::CreateTensor<float>(ortmem, model_output_data.data(), model_output_data.size(), output_shapes.data(), output_shapes.size());

これを行って、それぞれのfloat型のVectorにそれぞれOrt::Valueを結びつけたあとで以下で推論実行します。

	try {
		std::vector<char const *> input_node_names;
		std::vector<char const *> output_node_names;
		std::string istr = ortsession.GetInputNameAllocated(0, ortallocator).get();
		std::string ostr = ortsession.GetOutputNameAllocated(0, ortallocator).get();
		input_node_names.push_back(istr.c_str());
		output_node_names.push_back(ostr.c_str());
		
		ortsession.Run(
			Ort::RunOptions{nullptr},
			input_node_names.data(),
			&input_tensor,
			1,
			output_node_names.data(),
			&output_tensor,
			1
		);
	}catch(Ort::Exception& e) {
		std::cerr << "ort::session::Run Error:" << e.what() << std::endl;
		return -1;
	}catch(...){
		std::cerr << "ort::session::Run : unknown Error" << std::endl;
		return -1;
	}

Ort::SessionのRunには入力するテンソルの名前を設定する必要があるわけですが、名前はGetInputNameAllocatedで取得しています。
予め固定の名前がわかっているのであればこれをする必要があまりないのですが、一応こうしています。
出力名についても同じです。
また、なんらかの推論エラーが発生すると例外を返してくるのでtry-catchで例外をキャッチするのを忘れないようにしたほうが良いでしょう。

例外を出されず無事に帰ってきた場合は、何らかの推論が成功したということになります。

結果を確認する

モデルから出力されたデータはoutput_tensorに結びつけた

std::vector<float> model_output_data(1 * model_outsize);

に入っているので、model_output_dataの要素数とCSVの行数+1の対応をそれぞれ照らし合わせます。
model_output_dataの各要素に含まれている値は、値0〜1までの浮動小数点値になっており、1に近いほどその要素が強いということになります。
一般タグだと、0.3辺りから有効ということらしいです。

サンプルでは単純にソートをしています。この辺は色々実装者で工夫出来るかなと思います。

終わりに

ということで、AIをC++でというとなかなかに大層な感じで尻込みしてしまいがちですが、実はこと推論に関しては「AIモデルがわかるデータへ前処理→AIモデルに投げる→出力を人間がわかるように直す」というだけです。
そこまで大げさな処理でもなく、簡単な画像の判定くらいなら300行くらいで収まる範囲内で出来る話というのがわかると思います。

画像生成AIであるStableDiffusionや文章生成のLLMはこのような単純な構造ではありません。
これらは今回のように単一のAIモデルの結果のみを取得…というわけには行かずいろんなAIモデルが複雑に組み合わさって一つの機能を実現しています。

という感じで今回のように一筋縄では行かないですが…結局やってることは入力データへ変換→AIモデルに投げる→出力データを加工というのが複数あるということで、とりあえず今回で基礎は理解出来たのかなと…

あと、onnxruntimeの他にC++で使用可能なAI系ライブラリとしては、ggml(よくLLama.cppで使われています)などありますが、pythonでは非常に有名なPytorchも実はC++で使用可能なライブラリがあるそうなのでそれでもやってみるとか色々ありそうです。