image_converter

概要

メルアイコン変換器で用いた、UGATITのpytorch実装です。
詳しい解説はこちら。

想定環境

python 3.7.1, VRAM32GB以上
torch==1.4.0
torchvision==0.5.0
pip install -r requirements.txtでライブラリを揃えることができます。

プログラム

• UGATIT_train.pyは学習を実行し、学習の過程と学習済みモデルを出力するプログラムです。
• UGATIT_inference.pyはUGATIT_train.pyで出力した学習済みモデルを読み込み、推論(画像の変換)を実行、生成画像を出力するプログラムです。

使い方

以下では変換元ドメインをA、変換先ドメインをBと表現します。

学習の実行

1. UGATIT_train.pyのあるディレクトリに./datasetディレクトリを作成します
2. ./datasetディレクトリ内にgroup_Aディレクトリとgroup_Bディレクトリの2つを作成します。
3. ./dataset/group_Aディレクトリに、Aに属する画像を./dataset/group_A/*/*という形式で好きな数入れます(画像のファイル形式はpng)。
4. ./dataset/group_Bディレクトリに、Bに属する画像を./dataset/group_B/*/*という形式で好きな数入れます(画像のファイル形式はpng)。
5. UGATIT_train.pyの置いてあるディレクトリでpython UGATIT_train.pyを実行することで、「A⇄B」の変換ができるよう目指して学習を実行します。
   • 学習の過程が./output以下に出力されます。
   • 学習済みモデルが./trained_model/generator_A2B_trained_model_cpu.pthとして出力されます。

推論の実行

1. UGATIT_inference.pyのあるディレクトリに./conversionディレクトリを作成します
2. ./conversion内にtargetディレクトリを作成し、Aに属する画像を好きな数入れます。
3. UGATIT_inference.pyの置いてあるディレクトリでpython UGATIT_inference.pyを実行して./conversion/target内の画像をBへ変換します
   • A→Bの変換結果が./conversion/converted/以下に出力されます。
   • 注意点として、./trained_model内に学習済みモデルgenerator_A2B_trained_model_cpu.pthがなければエラーとなります

学習には環境によっては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