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