自動着色データセット作成のためのイラストの線画+ベタ塗り化

機械学習 画像処理 イラスト

最近自由研究でイラストの自動着色システムを開発しようと試行錯誤しています。 その過程で訓練データ用の線画やベタ塗りを抽出するコードを実装したので、メモ書きがてらご紹介したいと思います。

以下の前者のイラストを後者のような線画＋ベタ塗りのイラストに変換することが目的です。

線画抽出

まずはイラストから線画部分を抽出します。

自動着色に関する既存研究[Ci et al.,2018; Hati et al.,2019]では、データセット作成のためにXDoG [Winnemöller et al.,2012]と呼ばれるエッジ抽出フィルターを用いて着色済みのイラストから線画を人工的に生成していました。 これに倣い、今回の取り組みでもこのフィルターを使って線画を抽出しました。

XDoGフィルター

詳細な説明は論文や解説記事を見ていただくとして、ここでは論文や解説記事を参考にしたPython実装を共有します。デフォルト値は線画着色の論文で利用されていた値です。sigmaに関しては0.3, 0.4, 0.5の中からランダムに選ぶようです。

import cv2
import numpy as np

def XDoG(img, sigma, tau=0.95, k=4.5, phi=1e9, epsilon=0):
    """
    OpenCVの画像に対してXDoGフィルターを適用する。
    """
    img = np.repeat(np.mean(img, axis=2, keepdims=True), 3, axis=2)
    g_sigma = cv2.GaussianBlur(img, (0, 0), sigma)
    g_ksigma = cv2.GaussianBlur(img, (0, 0), k * sigma)
    d = g_sigma - tau * g_ksigma
    print(d.shape)
    t = np.where(
        d >= epsilon,
        np.array([255]),
        255*(1+np.tanh(phi*(d-epsilon)))
    )
    print(t.shape)
    t = t.clip(0.0, 255.0)
    return t.astype(np.uint8)

実際にこのフィルターをイラストに適用してみると、きちんと以下のように線画らしいものが取り出せていることが確認できます。

f:id:ronwall1701:20210503203111p:plain

ベタ塗り抽出

続いて完成品のイラストからベタ塗り状態のイラストを抽出する必要があります。シンプルにイラストの各領域ごとに色を平均してしまえば良さそうですが、この領域をどうやって抽出するのかという問題があります。

実はこのイラストからの領域抽出にはDanbooRegion [Zhang et al.,2020]という既存研究があり、以下のレポジトリにデータセット、ソースコード、事前学習済みモデルなどが公開されています。この手法はイラストの領域をスケルトンマップの形で表現し、イラスト→スケルトンマップの変換を学習することによりイラスト上の領域を抽出します。

github.com

pix2pixHDの学習

私の環境ではなぜか事前学習済みモデルがきちんと動かなかったので、image-to-image translationモデルであるpix2pixHD [Wang et al.,2018]を用いてこのスケルトンマップへの変換を学習しました。

PyTorch実装がgithubに公開されているので、そちらを利用させていただくことにします。‘

github.com

Anacondaを用いた環境構築は以下のように実行しました。

conda create -n pix2pixhd python=3.7
conda activate pix2pixhd
conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch -c nvidia
pip install dominate scipy==1.2.0

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

学習のためには自作データセット用に新しくディレクトリを作成し、その中にtrain_A, train_Bというそれぞれ変換前、変換後の画像を含むディレクトリを作成する必要があります。画像の名前は並べ替えたときにペアとなる画像が同じ順序になるようにすれば良さそうでした。今回はDanbooRegionデータセットのイラストを変換前画像に、スケルトンマップを変換後画像にします。

以上の準備の後、以下を実行すると学習が実行されます。DanbooRegionはデータセット用に作ったディレクトリです。特にvalidationなどはしてくれないのですが、今回のケースでは1-3エポック程度でそれらしい結果になったのでえいやで学習を打ち切ってしまいました。

python -m torch.distributed.launch train.py --name train_danboo_region --label_nc 0 --dataroot DanbooRegion --tf_log --no_instance --fp16 --resize_or_crop none

イラスト→スケルトンマップ→ベタ塗り　の変換

上の手順で学習したpix2pixHDモデルを使ってスケルトンマップを作るラッパークラスと、スケルトンマップを用いてイラストをベタ塗り化するコードを以下で公開しました。

pix2pixHDモデルのラッパークラス: https://gist.github.com/kosuke1701/4e2ac722bb0c4af9cbf9acfec3d91c3f

イラストベタ塗り化コード: https://gist.github.com/kosuke1701/6874ec7e10090ea9a59ba655b5905bc6

以上のコードを使って実際にイラストからベタ塗りを抽出してみた結果が以下になります。多少おかしいところもありますが、おおむね良くベタ塗りが抽出できていることがわかります。

f:id:ronwall1701:20210506080030p:plain
左からオリジナル画像、スケルトンマップ、ベタ塗り化された画像

参考文献

[Ci et al.,2018] Ci, Yuanzheng, et al. "User-guided deep anime line art colorization with conditional adversarial networks." Proceedings of the 26th ACM international conference on Multimedia. 2018.

[Hati et al.,2019] Hati, Yliess, et al. "PaintsTorch: a User-Guided Anime Line Art Colorization Tool with Double Generator Conditional Adversarial Network." European Conference on Visual Media Production. 2019.

[Wang et al.,2018] Wang, Ting-Chun, et al. "High-Resolution Image Synthesis and Semantic Manipulation with Conditional GANs." Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2018.

[Winnemöller et al.,2012] Winnemöller, Holger, Jan Eric Kyprianidis, and Sven C. Olsen. "Xdog: an extended difference-of-gaussians compendium including advanced image stylization." Computers & Graphics 36.6 (2012): 740-753.

[Zhang et al.,2020] Zhang, Lvmin, et al. "DanbooRegion: An Illustration Region Dataset." European Conference on Computer Vision (ECCV). 2020.