image_converter
image_converter copied to clipboard
UGATITによる画像変換器
image_converter
概要
メルアイコン変換器で用いた、UGATITのpytorch実装です。
詳しい解説はこちら。
想定環境
python 3.7.1, VRAM32GB以上
pip install -r requirements.txtでライブラリを揃えることができます。
プログラム
UGATIT_train.pyは学習を実行し、学習の過程と学習済みモデルを出力するプログラムです。UGATIT_inference.pyはUGATIT_train.pyで出力した学習済みモデルを読み込み、推論(画像の変換)を実行、生成画像を出力するプログラムです。
使い方
以下では変換元ドメインをA、変換先ドメインをBと表現します。
学習の実行
UGATIT_train.pyのあるディレクトリに./datasetディレクトリを作成します./datasetディレクトリ内にgroup_Aディレクトリとgroup_Bディレクトリの2つを作成します。./dataset/group_Aディレクトリに、Aに属する画像を./dataset/group_A/*/*という形式で好きな数入れます(画像のファイル形式はpng)。./dataset/group_Bディレクトリに、Bに属する画像を./dataset/group_B/*/*という形式で好きな数入れます(画像のファイル形式はpng)。UGATIT_train.pyの置いてあるディレクトリでpython UGATIT_train.pyを実行することで、「A⇄B」の変換ができるよう目指して学習を実行します。- 学習の過程が
./output以下に出力されます。 - 学習済みモデルが
./trained_model/generator_A2B_trained_model_cpu.pthとして出力されます。
- 学習の過程が
推論の実行
UGATIT_inference.pyのあるディレクトリに./conversionディレクトリを作成します./conversion内にtargetディレクトリを作成し、Aに属する画像を好きな数入れます。UGATIT_inference.pyの置いてあるディレクトリでpython UGATIT_inference.pyを実行して./conversion/target内の画像をBへ変換します- A→Bの変換結果が
./conversion/converted/以下に出力されます。 - 注意点として、
./trained_model内に学習済みモデルgenerator_A2B_trained_model_cpu.pthがなければエラーとなります
- A→Bの変換結果が
学習には環境によっては12時間以上要する場合があります。
入力された画像は256×256にリサイズされた上で学習に使われます。出力画像も256×256です。
参考
U-GAT-IT: Unsupervised Generative Attentional Networks with Adaptive Layer-Instance Normalization for Image-to-Image Translation
U-GAT-IT — Official PyTorch Implementation