CycleGANで浮世絵風の画像を生成してみる

AI実装

CycleGANはGANをつかったアルゴリズムの一種ですが、Pix2Pixと異なり入力と出力のペア画像を用意する必要がないImage to image translationモデルとなっています。それらしいデータを大量にあつめてクセを学習させることで画像変換するAIをつくることができます。今回はCycleGANで浮世絵風の画像を生成してみます。用意するデータとしては「浮世絵の風景っぽい画像」と「風景画像」のデータセットになります。ペアである必要はないので、これらは別の風景モノでOKです。

…っと書きましたが実際に学習させようとするとCycleGANは莫大なメモリを必要とするため学習にものすごい時間がかかります。ここがCycleGANの弱点かも😅…今回はあらかじめ学習した重みデータを用いて画像変換する方法を説明します。

動画もつくりました!こちらも併せてご覧ください。

CycleGANについて

CycleGANの仕組みを簡単に説明します。今ドメイン「X」と「Y」を↓の様に定義します。

  • X:風景写真の画像
  • Y:浮世絵の風景画像

CycleGANは↓のようにGeneratorが2つ(G,F)、Discriminatorが2つ(Dx,Dy)、計 4つのCNNが連動しています。

X→Yへ画像変換するG(x)が本物か偽物かを判断するのがDy、逆にF(y)を判断するのがDxのような関係になっています。つまり↓の感じです。

  • 風景写真を浮世絵風にする:G(x)
  • 生成した浮世絵風画像が本物か判断する:Dy
  • 浮世絵風画像を風景写真へ変換する:F(y)
  • 生成した風景写真が本物か判断する:Dx

↑を繰り返すだけではうまくいかないためCycleGANにはもう一つの面白い特徴があります。Cycle Consistency Lossという損失関数があります。これは「風景画像を浮世絵風に変換して、変換した浮世絵風画像を風景画像へ変換すれば元の画像へ戻るはず」ということをニューラルネットワークで表現するための損失関数です。

CycleGANでは Cycle Consistency Loss とAdversarial Loss(本物と偽物の損失関数)、Identity Loss※の3つの損失関数を最小にするように学習を進めていく仕組みとなっています。CNNが4つ、損失関数も3つあるので学習にはめっちゃ時間かかるんですね。

※ Identity Loss :「変換するターゲットの画像をGやFへ入力したら、そのままの画像が出力されるはず」 ということを表現する損失関数です。つまり浮世絵画像をG(x)へ入力し、出力画像を元の浮世絵画像と比較することです。

CycleGANの浮世絵画像生成モデルを実装してみる

さてCycleGANを動かしていきます。今回もローカル環境で実装していきます。

ソースコードは↓のGitHubを用いています。git cloneもしくはデスクトップ等へダウンロード→展開してください。

GitHub - junyanz/pytorch-CycleGAN-and-pix2pix: Image-to-Image Translation in PyTorch
Image-to-Image Translation in PyTorch. Contribute to junyanz/pytorch-CycleGAN-and-pix2pix development by creating an account on GitHub.

学習環境は↓

  • Windows 10
  • CPU : AMD Ryzen 7 5800
  • GPU : RTX3070 8GB

ですが、今回は学習済の重みデータを使用するだけなので、スペックは低くても大丈夫だと思います。

必要なライブリは↓です。AnacondaPromptなどでインストールしてください。PyTorchはCUDAのバージョンに注意してください。この辺りが不安な方は前の記事で詳しく説明しているので参考にしてください。

pip install torch==1.8.0+cu111 torchvision==0.9.0+cu111 torchaudio==0.8.0 -f https://download.pytorch.org/whl/torch_stable.html
pip install dominate
pip install visdom

次にデータセットと学習済の重みデータをダウンロードしていくのですが、Linux環境で下記のコードを入力すると一発で変換してくれるので便利です。Windows10でLinuxを動かす場合はwsl2(Windows Subsystem for Linux 2)をインストールし、Ubuntu上で実行することができます。ここを参考に設定することができます。

Linuxの準備が整いましたらpytorch-CycleGAN-and-pix2pixディレクトリまで移動し↓のコードを実行してください。

bash ./scripts/download_cyclegan_model.sh style_ukiyoe
bash ./datasets/download_cyclegan_dataset.sh ukiyoe2photo

実行すると、checkpointsディレクトリ下へstyle_ukiyoe_pretrainedディレクトリが生成され、その中にlatest_net_G.pthというデータファイルがダウンロードされていることを確認してください。また、datasetsディレクトリ下へukiyoe2photoディレクトリが生成され、その中にtestA、testB、trainA、trainBディレクトリがあり、浮世絵画像と風景画像がダウンロードされていることを確認してください。

確認できましたら、↓のコードを実行すると浮世絵風の画像が生成されます。

python test.py --dataroot datasets/ukiyoe2photo/testB --name style_ukiyoe_pretrained --model test --no_dropout

実行すると、resultディレクトリへstyle_ukiyoe_pretrainedディレクトリが生成され画像変換されていると思います。

自前の画像を浮世絵風に変換する

まずは自前画像を用意しdatasetsディレクトリ下へimagesディレクトリを作成し、そこへ保存してください。

/datasets
     -> images
          1.jpg
          2.jpg
          3.jpg
          4.jpg
          ……

画像のサイズはバラバラでも大丈夫ですが、そのまま変換すると256×256の画像へ変換されてしまうので↓のコードで実行してください。

python test.py --dataroot datasets/images --name style_ukiyoe_pretrained --model test --no_dropout --preprocess scale_width --load_size 1080

追加した引数の説明を簡単にすると↓です。

  • –preprocess scale_width:256×256に切り抜かれるのを防ぎます
  • –load_size 1080:1080pix幅になるように画像をリサイズします

結果↓の感じになりました。

思ったより浮世絵風になりました。面白いのでお時間あれば是非挑戦してみてください。

コメント

タイトルとURLをコピーしました