MMGenerationのセットアップと学習済みモデルの使い方

本記事では、GANに特化したツールセットであるMMGenerationのセットアップ方法や学習済みモデルを用いた様々なタスクの実行方法をご紹介します。

1.MMGeneration
- 1.1.概要
2.デモ(Colaboratory)
3.Interpolation
4.まとめ
5.参考文献

MMGeneration

概要

MMGenerationはOpenMMLabが開発するPytorchをベースとしたGANに特化したツールセットです。

Pytorch1.5以降で動作する本ツールセットは、StyleGANv1, v2, v3などのUnconditional GANからBIGGANやSNGANなどのConditional GANなど様々なモデルをサポートしており、統一されたAPIを使用して容易に切り替えることが可能です。

また、トレーニングや評価等の基本的なパイプラインの構築はもちろんのこと、GAN Interpolation, GAN Projection, GAN Manipulationなど複数のタスクがツールセットに組み込まれており、モデルの構築から使用まで幅広くサポートされています。

詳細はこちらのドキュメントをご参照ください。

本記事では上記手法を用いて、画像合成やGAN Interpolationを動かしていきます。

デモ(Colaboratory)

それでは、実際に動かしながらMMGenerationを用いた画像合成や画像補間を行っていきます。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

また、下記から直接Google Colaboratoryで開くこともできます。

また、このデモはPythonで実装しています。
Pythonの実装に不安がある方、Pythonを使った機械学習について詳しく勉強したい方は、以下の書籍やオンライン講座などがおすすめです。

おすすめの書籍

おすすめのオンライン講座

環境セットアップ

それではセットアップしていきます。 Colaboratoryを開いたら下記を設定しGPUを使用するようにしてください。

「ランタイムのタイプを変更」→「ハードウェアアクセラレータ」をGPUに変更

初めにGithubからソースコードを取得します。

%cd /content
 
!git clone https://github.com/open-mmlab/mmgeneration.git

次にライブラリをインストールします。

%cd /content/mmgeneration
 
!!pip install moviepy
!pip3 install openmim
!mim install mmcv-full
 
!pip install -e .[all]

以上で環境セットアップは完了です。

画像合成

ここでは、FFHQでトレーニングされたStyleGANv2モデルを使用してランダムに人の顔画像を生成します。
MMGenerationを使用するとコンフィグファイルとモデルのパスを指定するのみでモデルが使用できます。

import mmcv
from mmgen.apis import init_model, sample_unconditional_model
 
# Specify the path to model config and checkpoint file
config_file = 'configs/styleganv2/stylegan2_c2_ffhq_1024_b4x8.py'
# you can download this checkpoint in advance and use a local file path.
checkpoint_file = 'https://download.openmmlab.com/mmgen/stylegan2/stylegan2_c2_ffhq_1024_b4x8_20210407_150045-618c9024.pth'
 
device = 'cuda:0'
# init a generatvie
model = init_model(config_file, checkpoint_file, device=device)
# sample images
fake_imgs = sample_unconditional_model(model, 4)

生成された画像を表示します。

from torchvision import utils
import matplotlib.pyplot as plt
import os
from IPython.display import Image, display
 
results = (fake_imgs[:, [2, 1, 0]] + 1.) / 2.
save_path = './outputs/unconditional/results.jpg'
 
# save images
mmcv.mkdir_or_exist(os.path.dirname(save_path))
utils.save_image(
    results, save_path, ncol=1, padding=0)
 
# display
display(Image(save_path))

実在している人だと言われても違和感がないほどの画像が生成されています。

ラベル付き画像合成

ここでは、ImageNetを用いてトレーニングされたBigGANを用いて、ラベルで指定した画像を合成してみます。

from mmgen.apis import init_model, sample_conditional_model
 
# Specify the path to model config and checkpoint file
config_file = 'configs/sagan/sagan_128_woReLUinplace_noaug_bigGAN_Glr-1e-4_Dlr-4e-4_ndisc1_imagenet1k_b32x8.py'
# you can download this checkpoint in advance and use a local file path.
checkpoint_file = 'https://download.openmmlab.com/mmgen/sagan/sagan_128_woReLUinplace_noaug_bigGAN_imagenet1k_b32x8_Glr1e-4_Dlr-4e-4_ndisc1_20210818_210232-3f5686af.pth'
 
device = 'cuda:0'
# init a generatvie
model = init_model(config_file, checkpoint_file, device=device)
 
# sample images with specific labels
# label https://gist.github.com/yrevar/942d3a0ac09ec9e5eb3a
fake_imgs = sample_conditional_model(model, 4, label=[0, 1, 2, 3])

合成結果を表示します。

results = (fake_imgs[:, [2, 1, 0]] + 1.) / 2.
save_path = './outputs/conditional/results.jpg'
 
# save images
mmcv.mkdir_or_exist(os.path.dirname(save_path))
utils.save_image(
    results, save_path, ncol=1, padding=0)
 
# display
display(Image(save_path))

StyleGAN使用時と比較して、コンフィグと学習済みモデルのパスを変更した以外ほとんど実装内容を変えずにモデルを切り替えることができています。

Interpolation

最後に、StyleGANv2を用いて、画像間を補間しながらモーフィングのような動画を生成していきます。

!python apps/interpolate_sample.py \
  configs/styleganv2/stylegan2_c2_ffhq_256_b4x8_800k.py \
  https://download.openmmlab.com/mmgen/stylegan2/stylegan2_c2_ffhq_256_b4x8_20210407_160709-7890ae1f.pth \
  --export-video \
  --samples-path outputs/interpolation \
  --endpoint 6 \
  --interval 60 \
  --space z \
  --seed 12 \
  --sample-cfg truncation=0.8

生成された動画は以下の通りです。

from moviepy.video.fx.resize import resize
from moviepy.editor import VideoFileClip
 
# display
clip = VideoFileClip('/content/mmgeneration/outputs/interpolation/lerp.mp4')
clip = resize(clip, height=420)
clip.ipython_display()