見出し画像

【技術解説】Lightweight GANに浮世絵を学習させてグニョグニョ動かしてみた【東海道五十三次】

デザイニウム

こんにちは！デザイニウムのBBOY/エンジニアの平澤（@eatora22）です。今回はLightweight GANという機械学習モデルを使って浮世絵の画像を学習させてみたいと思います。

はじめに

まずは以下の動画をご覧ください。

こちらは、Lightweight GANに歌川広重の浮世絵作品「東海道五十三次」を学習させモーフィング動画として出力させたものです。画像データは以下URLより使わせていただいております。

GANのGeneratorによる出力は基本静止画ですが、入力に使用する潜在変数を少しずつ変化させることで出力される画像も合わせて変化し、それらを繋げることで絵が滑らかに動くモーフィング動画を作成することができます。この辺りの話は他のGANを使用した記事でも解説していますので、そちらをご参照ください。

技術解説

今回は数あるGANの中でもLightweight GANというモデルを使用します。実装としては、以下のlucidrainsさんによるものが有名なのですが、元となった論文の著者による公式実装も出ているようです（FastGAN）。

公式実装の方で動作確認してくれている方の記事もあったのですが、リポジトリ自体があまり更新されていないようです。一方でlucidrainsさんによる実装は最近も更新されており動作実績も多そうなので、とりあえず試したいのであればこちらの方が良さそう。

Lightweight GANの特徴としては、学習コストを抑えて比較的高品質な画像生成が可能である点です。元になった論文によると、100枚以下の解像度1024×1024pxの訓練画像を単一のGPU（RTX-2080）で数時間学習させるだけでも高いパフォーマンスを発揮できるそうです。過去にもPGGANやStyleGANを使った学習を行ってきましたが、数千~数万枚の画像が必要であったり学習時間も数日~数十日かかったりするイメージなので、それらと比べると非常にお手軽なモデルとなります。

導入もREADMEにあるようにpip installだけで済むので非常に楽なのですが、その他ハマった箇所もあるので本記事でいくつか補足しておきます。

Google Colabで動かすのが楽だと思うのですが、ローカル環境で動かしたい場合に「pip install lightweight-gan」を実行すると既にインストール済みのCUDAと合わないPyTorchが入ってしまいました。こちらはCUDAのバージョンと合うものを手動で入れた方が安全そうです。
上記に関連して、Colab側とローカル側でlightweight-ganのバージョンが異なると保存したモデルを読み込めない場合があります。また、Colabで実行する場合は毎回pip installすることになるかと思うのですが、元リポジトリのバージョン更新が多いため学習に時間をかけているとモデルが読み込めなくなっていることがありました。どちらもパッケージをダウングレードすることになると思います。
学習データの水増し（Augmentation）用のオプション引数としてaug-typesが用意されていますが、その中でcutoutという手法（画像の一部を黒い四角で塗りつぶす）がデフォルトで含まれています。学習データとの相性もありそうですがこれを入れていると私の場合は高頻度で生成画像にも一部黒い塗りつぶしが発生してしまったので、cutoutはできるだけ使わないようにしていました。

今回の学習対象は東海道五十三次の画像55枚を512×512pxサイズまで落としたものを使用します（53の宿場+始点の日本橋+終点の三条大橋で55枚になります）。実際に学習を行ったのが数か月前で記憶が定かではないのですが、無料版のGoogle Colabで2~3日ほど学習させて生成したのが記事冒頭のモーフィング動画です。

画像1

まずこれだけ少ない画像枚数をたったこれだけの日数&単一のGPUで学習させてここまでのクオリティで生成ができるという点で普通に凄いなと思いました。ただ、学習データが少ないためなのかモーフィングの中間出力の生成画像が曖昧な感じがします（ノイズっぽい）。また、1024×1024pxサイズでの学習も試してみたのですがあまりうまくいきませんでした。

ちなみにモーフィング動画の出力のためにgenerate-interpolationというオプション引数が用意されていますがGIF出力にしか対応していませんでした。取り急ぎですがOpenCVでmp4出力するサンプルコードも用意してみたので使ってみてください（昔書いたものですが多分動く）。※下記の関数をlightweight_gan.pyに追加しcli.py経由でgenerate_interpolationと同じ要領で呼び出して使ってください。

import cv2
import numpy as np

@torch.no_grad()
def generate_video(self, num = 0, num_image_tiles = 8, types=['default', 'ema'], num_steps = 100, fps = 30.0, linspace_end = 4):
   self.GAN.eval()
   ext = self.image_extension
   num_rows = num_image_tiles

   latent_dim = self.GAN.latent_dim
   image_size = self.GAN.image_size

   # latents and noise

   latents_low = torch.randn(num_rows ** 2, latent_dim).cuda(self.rank)
   latents_high = torch.randn(num_rows ** 2, latent_dim).cuda(self.rank)

   ratios = torch.linspace(0., linspace_end, num_steps)

   frames = []
   first_interp_latents = 0
   for i, ratio in enumerate(tqdm(ratios)):
       interp_latents = slerp(ratio, latents_low, latents_high)
       if i == 0:
           first_interp_latents = interp_latents
       elif i == (num_steps - 1):
           interp_latents = first_interp_latents

       # regular
       if 'default' in types:
           generated_images = self.generate_(self.GAN.G, interp_latents)
       # moving averages
       if 'ema' in types:
           generated_images = self.generate_(self.GAN.GE, interp_latents)
       images_grid = torchvision.utils.make_grid(generated_images, nrow = num_rows)
       
       frames.append(images_grid)

   # convert to mp4
   fourcc = cv2.VideoWriter_fourcc('m','p','4','v')
   frame_size = frames[0].shape[1]
   video = cv2.VideoWriter(str(self.results_dir / self.name / f'{str(num)}.mp4'), fourcc, fps, (frame_size, frame_size))
   for i in range(len(frames)):
       device2 = torch.device('cpu')
       frame = frames[i].to(device2)
       img = frame.numpy()
       img = np.transpose(img, [1,2,0])
       img = 255*img
       img = np.round(img).clip(0, 255).astype(np.uint8)
       img = img[:, :, [2,1,0]]
       video.write(img)
   # release video 
   video.release()


生成画像を使ってゴニョゴニョしてみた

更に面白い絵にするために、出力したモーフィング動画にひと手間加えてみようと思います。その結果生まれた映像がこちら。

モーフィングしていく東海道五十三次の浮世絵が音楽に合わせて四方へグニョグニョ動き出すようになりました。ピクセルソートされた時の色合いが良い感じですね。

こちらの加工にはTouchDesignerを使用しています。参考にさせていただいたチュートリアルはこちら。

さいごに

以上、Lightweight GANとそのモーフィング動画を使った応用例をご紹介しました。学習コストも下がり使いやすいGANが増えている中、逆に手間をかけまくることでありえないぐらい高品質な出力が可能になるモデルも登場しないのかな。

編集後記

広報のマリコです。いつも様々な機械学習を通して面白い映像を作り出しているBBOYエンジニア平澤が今回の題材に選んだのは、歌川広重の「東海道五十三次」✨個人的に浮世絵が大好きなので、とても興味深い作品でした❗ちなみに東海道五十三次のファミコンゲームも大好きだったのですが知ってる人いますか？（笑）こうやって見てみると、他の浮世絵作品や絵画のバージョンも見てみたくなりました。今後の作品も楽しみです❗

フォローやいいね！お待ちしてます😊


The Designium.inc
・オフィシャルサイト
・Interactive website
・Designium Dance
・Twitter
・Facebook

この記事が気に入ったらサポートをしてみませんか？