Python + CycleGanで茶毛のウマをシマウマに変換する

画像生成系のCycleGanを実装します．Python + CycleGanで茶毛のウマをシマウマに変換します．

有料枠設定にしていますが，下記のサイトで他の記事もみれます．youtubeの投げ銭的な物として，お考えください．

Python + CycleGanで茶毛のウマをシマウマに変換する

本稿の成果物

作業ディレクトリのファイル構成

プロジェクトディレクトリはdiablos_ganとしています．度々，省略しています．

diablos_gan
├── /pytorch-CycleGAN-and-pix2pix
│   ├── /checkpoints <- 学習モデルの保存先
│   │   └── /horse2zebra
│   │       └── latest_net_G.pth <- 学習モデル
│   │   └── /horse2zebra_pretrained
│   │       └── latest_net_G.pth <- 学習済みモデル
│   ├── /results <- 出力される変換画像
│   │   └── /horse2zebra
│   │   └── /horse2zebra_pretrained
│   ├── /datasets <- 学習データ
│   │   └── /horse2zebra
│   ├── /scripts <- モデルのダウンロード
│   │   └── download_cyclegan_model.sh
│   ├── /datasets <- データセットのダウンロード
│   │   └── download_cyclegan_dataset.sh
│   ├── requirements.txt
│   ├── train.py
│   ├── test.py
│   └── (省略)
└── diablos_gan.ipynb <- 実行用ノートブック

pytorch-CycleGAN-and-pix2pix

pytorch-CycleGAN-and-pix2pixのダウンロード

Google ColabとGoogle Driveを連携させて，gitからjunyanz/pytorch-CycleGAN-and-pix2pix (opens new window)をダウンロードします．

# Google ColabとGoogle Driveを連携
from google.colab import drive
drive.mount('/content/drive')

%%bash
# ディレクトリの移動
cd /content/drive/My\ Drive/diablos_gan
# gitのダウンロード
git clone https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix.git

CycleGANのデモ

起動チェックのため，茶毛ウマ->シマウマ，シマウマ->茶毛ウマへ色を変換する学習済みモデルを実行します

学習済みモデルをダウンロード

以下のコマンドで，学習済みモデルをダウンロードします．
学習済みモデルは， ./checkpoints/horse2zebra_pretrained/latest_net_G.pthとして格納されています．

# ディレクトリの移動
%cd /content/drive/My\ Drive/diablos_gan/pytorch-CycleGAN-and-pix2pix
# 学習済みモデルをダウンロード
!bash ./scripts/download_cyclegan_model.sh horse2zebra

データセットをダウンロード

以下のコマンドで，シマウマと茶毛ウマのデータセットをダウンロードします．
データセットは，./datasets/horse2zebraに格納されます．
trainA・testAは茶毛ウマ，trainB・testBはシマウマの画像となっています．

# 現在ディレクトリ
# /content/drive/My\ Drive/diablos_gan/pytorch-CycleGAN-and-pix2pix

# データセットをダウンロード
!bash ./datasets/download_cyclegan_dataset.sh horse2zebra
#モジュールのダウンロード
以下のコマンドで，モジュールをダウンロードします．

# 現在ディレクトリ
# /content/drive/My\ Drive/diablos_gan/pytorch-CycleGAN-and-pix2pix
# モジュールのダウンロード
!pip install -r requirements.txt

デモの実行

以下のコマンドで，学習済みモデルからテストを実行します．
--datarootはCucleganで変換するデータセット，--nameは/checkpointsの中で使用したいモデルが格納されているディレクトリ名です．(/checkpoints/ --nameの引数 /latest_net_G.pth)

# 現在ディレクトリ
# /content/drive/My\ Drive/diablos_gan/pytorch-CycleGAN-and-pix2pix

# デモの実行
!python test.py --dataroot datasets/horse2zebra/testA \
               --name horse2zebra_pretrained \
               --model test --no_dropout

学習済みモデルでの変換前画像

学習済みモデルでの変換後画像

CycleGANの学習

先ほどダウンロードしたシマウマとウマのデータセットから，CycleGANの学習します．
google colabで実行する場合は，--display_id 0と設定しておかないと学習が途中で止まります．

# 現在ディレクトリ
# /content/drive/My\ Drive/diablos_gan/pytorch-CycleGAN-and-pix2pix

# 学習実行
!python train.py --dataroot ./datasets/horse2zebra \
                --n_epochs 50 \
                --name horse2zebra \
                --model cycle_gan \
                --display_id 0 \
                --gpu_ids 0

学習モデルのテスト

以下のコマンドで学習のテストができます．
--nameでは/checkpointsの中にある学習モデルのディレクトリ名を記述します．
学習モデルの名前はlatest_net_G.pthとしてください．
Aのデータ -> Bのデータに変換するモデルは＊＊＊_G_A.pthで，Bのデータ -> Aのデータに変換するのモデルは＊＊＊_G_B.pthです．
※trainA・testAはウマ，trainB・testBはシマウマの画像となっています．

# 現在ディレクトリ
# /content/drive/My\ Drive/diablos_gan/pytorch-CycleGAN-and-pix2pix

# 学習モデルの名前変更
!cp checkpoints/horse2zebra/latest_net_G_A.pth checkpoints/horse2zebra/latest_net_G.pth

# 学習モデルのテスト実行
!python test.py --dataroot datasets/horse2zebra/testA \
               --name horse2zebra \
               --results_dir ./results/ \
               --model test --no_dropout

学習モデル(50エポック)での変換前画像

学習モデル(50エポック)での変換後画像

まとめ

画像生成系のCycleGanを実装し，茶毛のウマをシマウマに変換しました．
今度は，ディアブロス亜種->ディアブロスに変換するモデルを実装します．

参考サイト

Python + CycleGanで茶毛のウマをシマウマに変換する