本記事では、JoJoGANを使って、顔写真をJoJo風のイラストに変換する方法を紹介します。
また、自作のモデルをトレーニングする方法も合わせて紹介します。

JoJoGANとは

JoJoGANとはOne shot Face Stylizationを実現するGANです。

従来手法は、複数枚の画像を学習し、学習によって獲得したスタイルを入力画像に適用しています。
しかし従来手法では、目の形、線の太さなど詳細なスタイルをモデルが獲得することができませんでした。

JoJoGANでは、一枚の画像(One shot)でモデルをトレーニングすることが可能でありながら
前述の目の形、線の太さなどの詳細なスタイルを獲得することを目指している手法です。

さらに、学習したスタイルは任意の入力画像に適用できると共に、学習したスタイルの色を適用するか、 入力画像の色を適用するか選択することも可能となっています。

出典: mchong6/JoJoGAN

上図のReferencesがトレーニングに利用したOne shot画像で、Inputsがスタイルを適用する入力画像です。
赤枠内がスタイルの色を適用した出力結果で、青枠内が入力画像の色を適用した出力結果です。
髪色などに選択の結果の差が表れています。

また、髪の質感や目の形状など細部に渡りスタイルが適用されていることが分かります。

トレーニングに必要なデータが非常に少ないため、様々なスタイルのモデルを作成することができます。

JoJoGANの導入手順

セットアップ: conda環境構築

それでは早速、開発環境にJoJoGANをセットアップしていきます。
動作確認は下記の環境で行っています。

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

それでは、JoJoGANをインストールしていきます。他の機械学習環境に影響を与えないためにMinicondaの仮想環境上に構築していきます。Minicondaのインストール手順は公式ドキュメントをご参照ください。

conda環境構築手順を下記に記載します。

# python version 3.7のconda環境を作成
$ conda create -n jojo_gan python=3.7 -y

# 作成した環境をアクティベート
$ conda activate jojo_gan
$ cd jojo_gan

# Proxy配下の場合は事前に下記を設定
$ export http_proxy="http://"username":"password"@proxy:port"
$ export https_proxy="http://"username":"password"@proxy:port"

# githubからcode clone
$ git clone https://github.com/mchong6/JoJoGAN.git

# 関連するライブラリインストール
$ pip3 install tqdm gdown scikit-learn==0.22 scipy lpips dlib opencv-python wandb PyDrive matplotlib
$ pip3 install ninja

# Pytorchのインストール
# https://pytorch.org/get-started/previous-versions/ からcudaバージョンが一致するものをインストール
$ conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
# torchの動作確認
$ python -c "import torch; print(torch.__version__); print(torch.cuda.is_available()); print(torch.cuda.device_count());"
1.10.1
True
2

# modelをダウンロード
$ mkdir models
$ cd models
$ wget http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2
$ bzip2 -dk shape_predictor_68_face_landmarks.dat.bz2

GPUを積んでいる環境であれば、torchの動作確認で事前にPytorchがGPUを認識できることを確認するようにしてください。

以上でconda環境作成は完了です。

JoJoGANデモファイルの解説

それでは早速JoJoGANを試してみます。

作成したconda環境で動作するデモ用のファイルを下記に作成しました。
https://github.com/kaz12tech/JoJoGAN/blob/master/demo.py

こちらのdemo.pyを実行すると、事前学習済みモデルを使ったスタイルの適用と、
--img_dirに格納された画像からモデルをトレーニングして獲得した新たなスタイルの適用を行います。

以下に、demo.pyの解説を記載します。

def get_args():
    parser = argparse.ArgumentParser()
    # スタイルの適用に利用する入力画像を指定
    parser.add_argument("--img_path", type=str, default='./test_input/sample_input.jpg')
    
    # 事前学習済みモデルを選択(art, arcane_multi, supergirl, arcane_jinx, arcane_caitlyn, jojo_yasuho, jojo, disney)
    parser.add_argument("--pretrained", type=str, default='arcane_multi', 
        help='please select art, arcane_multi, supergirl, arcane_jinx, arcane_caitlyn, jojo_yasuho, jojo, disney')
    
    # 事前学習済みモデル、及び、学習時にStyleの色を適用する場合はTrue
    parser.add_argument("--preserve_color", type=bool, default=True)
    
    # ランダムサンプルにスタイルを適用。サンプル数を指定 
    parser.add_argument("--n_sample", type=int, default=5)

    # seed 乱数の初期値
    parser.add_argument("--seed", type=int, default=1234)

    # styleを学習させる画像を格納したディレクトリを指定
    parser.add_argument("--img_dir", type=str, default='./train_style_images/')

    # 学習時のiteration
    parser.add_argument("--num_iter", type=int, default=200)

    args = parser.parse_args()

    return (args)

もし、Pythonの実装に不安がある場合は、こちらの書籍などがおすすめです。

それぞれの引数の意味はコメントに記載の通りですが、重要な引数は下記3点です。

1. --pretrained:学習済みモデルの選択
   (art, arcane_multi, supergirl, arcane_jinx, arcane_caitlyn, jojo_yasuho, jojo, disney)
2. --img_dir:スタイルを学習させる対象画像を格納したディレクトリ
3. --img_path:Githubより提供された学習済みモデルと、上記学習によって生成したモデルのスタイルを適用する入力画像

demo.pyの実行時は以下のようなコマンドで動作します。

$ python3 demo.py \
    --img_path ./test_input/sample_input.jpg \
    --pretrained arcane_multi \
    --preserve_color True \
    --n_sample 5 \
    --seed 3000 \
    --img_dir ./train_style_images/ \
    --num_iter 500

以降では、事前学習済みモデルのスタイル適用結果、及び、自作スタイルの適用結果をご紹介します。

事前学習済みモデルのスタイル適用結果

それでは事前学習済みモデルのスタイルを適用してみます。
今回はこちらの画像を入力してみます。

以下が、事前学習済みモデルJoJoのスタイルを適用した結果です。

以下はランダムサンプルにJoJoスタイルを適用した結果です。

目の形状や、髪質、影の付け方などone shotから獲得したスタイルであるにも関わらず
なかなか詳細にスタイルを適用できていますね。

自作スタイルの適用結果

次に、学習データを用意して自作モデルをトレーニングさせてみます。
今回はこちらのゴッホの絵を学習データとして用意しました。

1点学習データを用意する際の注意点として、瞳がはっきりと描かれており正面を向いている顔画像である必要があります。

学習時に、学習データ画像から顔部分のみをcropして使用するのですが、
いらすとやにあるような目が黒塗りの点のような顔画像ですと正常に顔部分の切り抜きが完了しません。

さて、自作モデルであるゴッホの絵から獲得したスタイルを適用した結果は以下の通りです。

油絵を塗り重ねたような質感が見事に適用されています。
初回実行時のみ学習済みモデルのダウンロード時間がかかりますが、学習時間はおよそ1分～2分程で完了しました。 手軽さを鑑みると見事な出力結果です。

またスラムダンクの桜木花道を使って学習させた結果は以下の通りでした。(学習に使った画像はこちら)

1枚目の右端の子供の写真はうまくスタイルが適用されているように思います。
全員三白眼になるので目のスタイルが強く反映される傾向がありそうです。

まとめ

本記事では、JoJoGANを使って、入力画像に事前学習済みモデルのスタイルを適用する方法や、自作のスタイルを学習したモデルのトレーニング方法を紹介しました。

学習データの収集は、多くのタスクにおいてコストの問題に直面しがちです。
本記事でご紹介した手法のように既に確立されたモデルから転移学習を行うことによって、学習データの必要数を大きく減らせると 導入コストを大きく下げられそうです。

余談ですが、JoJoGANという名称からしてもそうですが、まじめな研究しつつ見せ方を少しふざけている感じが個人的にたまりません。

また本記事では、機械学習を動かすことにフォーカスしてご紹介しました。
もう少し学術的に体系立てて学びたいという方には以下の書籍などがお勧めです。ぜひご一読下さい。

また動かせるだけから理解して応用できるエンジニアの足掛かりに下記のUdemyなどもお勧めです。

参考文献

1.  論文 - JoJoGAN: One Shot Face Stylization

2. GitHub - mchong6/JoJoGAN