CycleGAN で画像変換する。 #機械学習

始めまして！
adventカレンダー16日目を担当させていただきます。
DMM WEBCAMPでメンターをしている野村です！

よろしくお願いします！

馬の体をシマシマに変えたいな。。

突然ですが、そこのあなた！
「馬の体をシマシマにして、シマウマにしてやりたい！」と思ったことはありませんか？！
CycleGANというものを使うとそれを実現することができます。
今回は自分の勉強も兼ねてハンズオン形式で行いますので、ご自身の環境でも実装していただけると幸いです！
間違っている部分、追加の説明などがありましたらコメントお願いいたします！！！

※主にこの記事を参考に自分でも考えながら作業してみました。(ありがとうございます）
※CycleGANの仕組みについて詳しく知りたい人はこちら

Pythonを動かすには

Pythonを使うのに、ローカルでやっても、仮想環境作ってもいいですが、Googleが無料で使わせてくれるので使いましょう。
Google Colabratoryを使えば、無料でGPUを貸してくれるので低スペックパソコンでも機械学習をガンガン進めることができます。(Googleさんありがとう)

CycleGANとは

画像から新たな画像を生成する方法はCycleGANの他にもpix2pixと言う手法もあり、このpix2pixで画像をあたらに生成するにはペアとなる画像を用意する必要があります。
ペアの画像とは、形状は全く同じだけど色や模様だけが異なる画像であり、以下のような組み合わせがその一例です。

※馬の形状自体は同じだが、体の模様が異なる

しかし、このようなペアの画像が揃っている場合は非常に稀で、学習をする以前にペア画像を集めることができません。
そこでCycleGANを用いてこのペア画像の生成を行おう！ということです。
CycleGANを用いることでペア画像を作成することができるので、pix2pixを用いた機械学習を進めることができます。

では実際に次のセクションからCycleGANを用いた学習を初めて行きます！

1.Google ColabratoryをGoogle Driveにマウントする

Google Colabratoryにアクセスして、ノートブックを新規作成しましょう！
まずはGoogle ColabratoryとGoogle Driveを繋げて(マウント)することで、Google Driveの画像などのデータをそのまま使うことができます。
Google Driveも同時に開きながら作業するとフォルダの関係性がわかりやすいです！
※時間を置いて作業をしていた際に、コードを全くいじっていなくても「no such file or directory」で怒られることがありますが、マウントの作業をやり直すことで解消できると思います。
※詳しいマウント方法についてはこちら

from google.colab import drive
drive.mount('/content/drive')

2. 初回作業

2.1新規フォルダ作成

# GoogleDriveの中の、Colab Notebooksフォルダに移動
%cd /content/drive/My Drive/Colab Notebooks/

# CycleGAN_Projectというフォルダを作成
!mkdir CycleGAN_Project

# CycleGAN_Projectフォルダに移動
%cd /content/drive/My Drive/Colab Notebooks/CycleGAN_Project

2.2 GitHubからコードを入手する

とある人がすでに学習済みのCycleGANを公開してくれているのでそちらを利用します。
requirements.txtに記載されている内容にしたがってライブラリをインストールすることで作業環境を合わせています。

# GitHubから本家CycleGANをクローン
!git clone https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix.git

# GitHubからクローンしてきたフォルダに移動
%cd /content/drive/My Drive/Colab Notebooks/CycleGAN_Project/pytorch-CycleGAN-and-pix2pix

# requirements.txtの内容にしたがって必要なライブラリをインストールする
!pip install -r requirements.txt

2.3 visdomはpip installでは使えないみたいです。

こちらの記事を読む限り、普通にインストールするのではうまくいかないみたいなのでGitHubからクローンしてインストールします。

# GitHubからvisdomをクローン
!git clone https://github.com/facebookresearch/visdom.git

#　インストール
%cd visdom
!pip install -e .
!pip install visdom dominate

3. 2回目以降

2回目以降は、セクション1のマウントを行った後に以下のコードを実行します。
※2回目以降は、セクション2は飛ばします。

# クローンした本家CycleGANのフォルダに移動
%cd /content/drive/My Drive/Colab Notebooks/CycleGAN_Project/pytorch-CycleGAN-and-pix2pix

# インストール
!pip install -r requirements.txt
%cd visdom
!pip install -e .
!pip install visdom dominate
%cd /content/drive/My Drive/Colab Notebooks/CycleGAN_Project/pytorch-CycleGAN-and-pix2pix

4. 学習の準備(学習するデータを持ってくる)

とりあえずCycleGANのサンプルデータで学習させる。

# 馬とシマウマのデータセット
!bash ./datasets/download_cyclegan_dataset.sh horse2zebra

5. 学習させる(馬のデータセットで)

「xxxxxx」の部分をデータセットのフォルダ名にする。(今回なら「horse2zebra」)
「yyyyyy」の部分をモデルの名前にする。(データセットのフォルダ名と同じでいい)

# サンプルコード
!python train.py --dataroot ./datasets/xxxxxx --name yyyyyy_cyclegan --model cycle_gan

#実際のコード
%cd /content/drive/My Drive/Colab Notebooks/CycleGAN_Project/pytorch-CycleGAN-and-pix2pix

# 学習開始！
!python train.py --dataroot ./datasets/horse2zebra --name horse2zebra_cyclegan --model cycle_gan

※何にもいじっていない状態だと以下のようにGPUでエラーが発生するので、
「ランタイム>ランタイムのタイプを変更>GPU」とすることで解消できます。

学習が始まると、
%cd /content/drive/My Drive/Colab Notebooks/CycleGAN_Project/pytorch-CycleGAN-and-pix2pix/checkpoints/yyyyyy_cyclegan/web/images
の中に1エポックごとに画像が生成されていきます。
エポックとは学習させる回数のことで、今回は200回行います。多すぎても少なすぎてもダメみたいで、ここの調整は今後自分でできるようになる必要がありますね。
200エポックやり切るのには、かなり時間がかかるので一旦休憩しましょう。。。
※目安ですが1エポックに30分ちょっとかかりました。

※2021/11/27追記

当記事のコメントにもある通り、Google Colab上で割り当てられるGPUによって学習時間に違いが見られるようです。
v100というGPUでは1エポック約6分、k80というGOUでは1エポック約36分という計算になります。（参照：https://teratail.com/questions/367471）

6. 一旦停止する

学習を始めたはいいけど、やっぱり時間がかかりすぎる！ってときには一旦中断することもできます。
↓実行ボタンがクルクルしているので、停止しましょう。

7. 学習を再開する

Google Driveを見ると、5エポックごとに学習結果が保存されているのがわかります。

中断した学習を再開する場合は「--continue_train」を追加します。
また、実行ボタンを押すとエポックのカウントが1に初期化されてしまうので、再開するエポックを指定する必要があります。
今回は60エポックが完了した所で一旦停止したので「--epoch_count 61」を追加しました。
※countの後に半角スペースがあることに注意！

%cd /content/drive/My Drive/Colab Notebooks/CycleGAN_Project/pytorch-CycleGAN-and-pix2pix

# 学習再開！
!python train.py --dataroot ./datasets/horse2zebra --name horse2zebra_cyclegan --model cycle_gan --continue_train --epoch_count 60

8. 学習完了

200epochまで学習ができたら自動でプロセスが終了します。(ちょこちょこ進めて3日くらいかかりました）
※GoogleDriveの容量は4GBほど使用しました。

以下のコマンドで、学習したモデルを使用して馬の画像をシマウマの画像に変換しましょう。
下のコードではtestAに入っている画像を変換しています。

!bash ./scripts/download_cyclegan_model.sh horse2zebra

#変換する
!python test.py --dataroot datasets/horse2zebra/testA --name horse2zebra_pretrained --model test --no_dropout

以下のコマンドで元画像とシマウマ変換画像を出力しましょう。

# 元画像を出力
import matplotlib.pyplot as plt

# ****に出力したい番号を入れる。
img = plt.imread('./results/horse2zebra_pretrained/test_latest/images/n02381460_****_real.png')
plt.imshow(img)

# シマウマ変換画像を出力
import matplotlib.pyplot as plt

# ****に出力したい番号を入れる。
img = plt.imread('./results/horse2zebra_pretrained/test_latest/images/n02381460_****_fake.png')
plt.imshow(img)

左が変換前の元画像、右側が学習させたことによって生成したシマウマ変換画像です。

9. まとめ

今回は元々画像が用意されている中でボタンを押していくだけで機械学習ができましたが、今後は自分で画像を用意し、モデルなどもいじれるようになるといいですね。。(頑張ります)
最後まで見ていただきありがとうございました！

10. 自前の画像を変換する(2021/2/24追記)

testCというフォルダを ./datasets/horse2zebra の配下に新たに作成し、ネットから適当に拾ってきた馬の画像を入れておきます

変換のときに使用したコードのパスをtestCに書き換えます。

!bash ./scripts/download_cyclegan_model.sh horse2zebra

#変換する
!python test.py --dataroot datasets/horse2zebra/testC --name horse2zebra_pretrained --model test --no_dropout

変換した画像を出力してみると、、

できてます！いい感じにシマシマになってますね！！