まいだいありー

機械学習、技術系、日記など勉強したことのメモを書けたらなと思います。

アニメキャラの顔を色指定して自動彩色( pix2pix with color hint )

はじめに

AnimeFace で pix2pix をやってみたで pix2pix モデルを用いてアニメキャラの顔の自動彩色を実験しました. その際は, 単に線画と元の画像のペアが似てるかどうかを学習させ, Generator が線画に彩色したものを生成させるというものでした.
今回はその続きとして, ユーザがカラーヒントを与え, そのカラーヒントの通りに自動彩色させるような Generator を作成してみたいと思います.

最初に結論を述べておくと, 割とうまく行きました...!


pix2pix with color hint

f:id:kenzo1122:20210603005518p:plain:w400:h250

図1: 学習モデルの概要


学習する方法は pix2pixと何も変わりませんが(ここを参照), 細かいところで異なる部分があるのでそこだけ説明します.


カラーヒントとして使う画像

色を指定して自動彩色させるには, どのようなカラーヒントを与えるかが重要です.

今回は, 複数個のサークルを元画像に貼り付けてマスクする Circle Mask (勝手に命名) を用いました. Circle Mask を画像に適用することで良い感じにマスクされるところとされないところが出来上がり, これをカラーヒントとして用いました.

f:id:kenzo1122:20210602235144p:plain

図2: Circle Mask

また, ロバスト性を向上させるため, データセットをバッチごとに切り出す際にランダムで Circle Mask を適用させています. なので epoch 数を重ねるほど多様なカラーヒントと出会うので色指定による彩色の精度は上がるはずです.


カラーヒントを Generator に付与する

カラーヒントの情報を Generator に与える方法は単純で, 入力のチャンネル方向にカラーヒントの画像と線画を結合しました. ( PyTorchでは, (batch_size, 4, 256, 256 ) となる )

Generator のネットワーク

f:id:kenzo1122:20210603004042p:plain

図3: Generator ( SLEBlock を用いた U-Net)


今回は, U-Net に Skip-layer excitation Block (SLEBlock) を組み合わせたネットワークで実験しました.
SLEBlockは 小さい特徴量マップを一本のベクトルに潰し, 大きい特徴量マップと積をとった出力をする (Attention のような役割もありそう? )モジュールで, LightweightGANのGeneratorで用いられていました.

f:id:kenzo1122:20210602231133p:plain:w250:h400

図4: Skip-layer excitation Block (引用[1])


LightweightGANのGenerator ではガウス分布からサンプリングしたベクトルをアップサンプリングする構造ですが(つまり Decoderのみ), 以下の図5の緑の矢印のように途中のデコード部分を逆にみることで U-Net 的な構造でも使えるのではという発想です.

f:id:kenzo1122:20210602234529p:plain

図5: LightweightGANのGenerator (引用[1])


バニラ U-Net では Encoder の各特徴量マップの大きさが一致するものだけを Decoder に skip-connection で結合していましたが, SLEBlock を用いることでskip-connection によって伝わる情報が局所的な特徴と全体の特徴を考慮したものとなるので, 取り入れてみました.

実際に用いたレイヤーの詳細は, Encoder のレイヤーでは, Conv -> BatchNorm -> LeakyReLU -> AvgPool のダウンサンプリング, Decoder のレイヤーは, 図5のレイヤーと同じです.

Discriminator

Discrminatorの入力は, [線画,元画像, Circlemaskのカラーヒント] と [線画, 生成画像, Circlemaskのカラーヒント] であり, PyTorch表記での次元は (batch_size, 7, 256,256)です また用いたレイヤーは, Conv -> BatchNorm -> LeakyReLU -> AvgPool のダウンサンプリングです.

実験環境

Google Colab で実験しました.

コード

github.com

実験結果

最初に学習データとテストデータを混合したデータで, 線画とCircle Mask によるカラーヒントをGeneratorに入力した時の自動彩色結果です.
当然ですが綺麗に彩色されています.

f:id:kenzo1122:20210603155602p:plain
図6:Circle Mask と線画を入力に用いた彩色


次に学習データとテストデータを混合したデータで, 線画とそれに対応関係でない元画像のCircleMask(つまりシャッフルした) によるカラーヒントをGeneratorに入力した時の自動彩色結果です.
これは彩色というよりかはスタイル変換のような結果として示します.
本来ランダムにカラーヒントを与えて彩色することが目的ではないので, よくあるスタイル変換という文脈とした用途では向いてませんね.
さらに学習データ内の各パーツにない色がカラーヒントとして与えられると, 汚い彩色になりがちです. (例えば肌に彩度の低いオレンジ系ではないカラーヒントが与えられた時など)

f:id:kenzo1122:20210606010839p:plainf:id:kenzo1122:20210606010750p:plainf:id:kenzo1122:20210606010649p:plain
図6 : CircleMaskをシャッフルした結果


次に学習データとテストデータから数枚サンプルし, Clip Studioでカラーヒントとして用いる色をつけた画像(以下の実験結果の真ん中) を用意しました. そのカラーヒントの画像とそれに対応する線画をGenerator に入力した時の自動彩色結果を以下に示します.


学習データの線画と人間によるカラーヒントを用いた自動彩色結果

f:id:kenzo1122:20210603154449p:plainf:id:kenzo1122:20210603154454p:plainf:id:kenzo1122:20210603154458p:plainf:id:kenzo1122:20210603154504p:plainf:id:kenzo1122:20210603154509p:plain
図7 : 人間によるカラーヒントと線画を入力に用いた彩色(学習データ)


テストデータの線画と人間によるカラーヒントを用いた自動彩色結果

f:id:kenzo1122:20210603153141p:plainf:id:kenzo1122:20210603153136p:plainf:id:kenzo1122:20210603153132p:plainf:id:kenzo1122:20210603153127p:plainf:id:kenzo1122:20210603153053p:plainf:id:kenzo1122:20210603153035p:plainf:id:kenzo1122:20210603153146p:plain
図8 : 人間によるカラーヒントと線画を入力に用いた彩色(テストデータ)

以上が実験結果ですが, 彩色結果から一応色指定による自動彩色は成功しています. 学習データとテストデータの彩色結果の差は見受けられませんので, 汎用性は割とあるのかもれません.

ただ, Circle Maskによるカラーヒントと人間がつけるカラーヒントとでは彩色結果に大きな差があることがわかります. これはおそらくCircle Mask によって切り取られるカラー情報の形や, 切り取られたカラーが1色ではないことなどに依存してしまってることが原因だと思います. 対処としては, マスクする方法を変える,ブラーを入れるなどバリエーションを持たせることが必要です.

また, 実験途中で線画を作る手法を変更したら彩色結果が酷く劣化したので, 線画の種類を増やすため全体が細い線画と輪郭が太い線画を確率的選択をし, 入力として用いています ( 今の時点であらゆる線画に対応させる意味はないですが )

次は高解像度の画像で, 今回問題となった部分の対処法を見つけて同じことをやってみたいと思います, 以上!

参考