🤖

入力線画を全く変えずに画像生成AIに色塗りさせる方法

2024/05/30に公開

タイトル通り、入力線画を(基本的には)1pixelも変えずに画像生成AIに色塗りをさせる方法について解説していきます。

本題に入る前に

そもそも画像生成AIによる色塗り（着彩）について知らない方向けに、既存技術でAIに色塗りをさせるとはどのような事を指すか？またどのような課題があるか？について一度整理します。

画像生成AIによる着彩

テキストから画像を生成できるという事で一躍有名になった画像生成AIですが、テキストだけでは生成したい画像を詳細に指示することが困難であるという課題を抱えていました。
そこで、テキストと比較してより具体的に生成したい画像を指示する方法として。ControlNetを用いた画像とテキストを入力して新たな画像を生成する手法が現れます。

この応用として、色塗りがされていない線画と指示テキストを入力することで、入力した線画に対し色を塗った状態の画像を生成するという方法が存在していました。

この手法は線画への色付けができるControlnetのモデル名をとって。「Controlnet Lineart」とか、略して「Lineart」などと呼ばれています。

さて、このcontrolNet Lineartを用いた線画への着色手法ですが、一つ大きな問題を抱えていました。
前述の通り、ControlNetはあくまで画像とテキストを入力として新しい画像を生成する手法であるため、入力された線画をそのまま使用するわけではありません。
特に、入力された画像はVAEを経て潜在空間に投影する必要があるため情報の欠損がおき、画像生成時に
復元された線画と入力された線画が完全に一致しないという問題がありました。

これは、1ピクセルずれるだけで大きく印象が変わってしまうイラストの着彩に使う機能としては、かなり大きな欠陥となります。
この問題はStableDiffusion v1.5のモデルでは特に顕著で、この精度ではイラスト制作のワークフローには組み込むのは難しいという扱いになっていました。

SDXLControlNetの出現

上述の欠陥を抱えているためなかなかうまい利用方法を見出す事が出来ず、しばらく触れていなかったこの着彩技術ですが、再度注目するきっかけとなった出来事が最近ありました。
それが、SD1.5の後継となるSDXLに対応した高精度なControlnet Lineartモデルの出現です。

このモデルはかたらぎ（@redraw_0）さんによって開発されたモデルで、SD 1.5版のLineartと比較してもかなり精度が良いものとなっていました。

入力線画
SDXL Lineartによる着彩
入力した線画（青線）を重ねた画像

手法の説明

では本題に入っていきましょう。
SDXL Lineartの登場によって入力した線画にかなり忠実な着彩ができるようになりましたが、それでも多少の線画改変は起こってしまっており、入力した線画とは微妙に印象が異なる画像が生成されてしまっています。
この課題を「入力した線画をそのまま出力するControlNetモデル」と「古典的画像処理による後処理」という二つのアプローチでの解決を目指しました。

ControlNetを用いた更なる入力線画に対する忠実度向上

まず一つ目に試したのが、ControlNetを用いて入力線画に対する忠実度を更に向上する方法でした。
具体的には、入力した画像をそのまま出力するように訓練したControlNetモデル（Line2Line）を用意し、SDXL Lineartによる着彩時に低weightで混ぜ込みます。

def get_cn_pipeline():
    controlnets = [
        ControlNetModel.from_pretrained("./controlnet/lineart", torch_dtype=torch.float16, use_safetensors=True),
        ControlNetModel.from_pretrained("mattyamonaca/controlnet_line2line_xl", torch_dtype=torch.float16)
    ]

    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
    pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
        "cagliostrolab/animagine-xl-3.1", controlnet=controlnets, vae=vae, torch_dtype=torch.float16
    )

    pipe.enable_model_cpu_offload()
    return pipe

def generate(pipe, detectors, prompt, negative_prompt):
    default_pos = "1girl, bestquality, 4K, ((white background)), no background"
    default_neg = "shadow, (worst quality, low quality:1.2), (lowres:1.2), (bad anatomy:1.2), (greyscale, monochrome:1.4)"
    prompt = default_pos + prompt 
    negative_prompt = default_neg + negative_prompt 
    print(type(pipe))
    image = pipe(
                prompt=prompt,
                negative_prompt = negative_prompt,
                image=detectors,
                num_inference_steps=50,
                controlnet_conditioning_scale=[1.0, 0.2],
            ).images[0]
    return image

これにより、SDXL Lineart単体で着彩を行うよりもより線画に忠実な着彩を行うことができます

SDXL Lineart + Line2Lineによる着彩
入力した線画（青線）を重ねた画像

古典的画像処理による後処理

このLine2Lineモデルを作った事で忠実度が上がり、ニュアンスまで保持した線画着彩が行えるようになったのですが、やはり多少のズレは残ってしまいます。
これは前述した通り、入力画像が一度VAEを通して潜在空間に投影されている以上、どうしても完全な復元というのが難しいためです。
そのため、これ以上の精度を求める場合は確率的な操作を行うAIではなく、条件に基づいて確実に画像を処理できる古典的な画像処理技術を用いるのが良いと判断しました。

単純な発想として、入力した線画を完全に維持したいのであれば入力した線画を保持しておいて生成した画像の上から貼りつければ良いのではないか？というものがあります。
この手法の問題は、張り付ける側の画像にも生成された線画が存在しており、ただ上から張り付けただけでは生成された微妙にずれた線と重なってしまうという点が挙げられます。

この問題は、生成される線画と入力した線画のズレが少ない場合は古典的画像処理である程度解決が可能です。
実際に加工結果を確認しながら処理を追っていきます。
入力線画と生成結果はこちら。

この時点でもかなり入力線画への忠実度は高いのですが、重ねてみると完全には一致していないことがわかります。

まず、入力した線画を生成画像に上から張り付ける際に、以下の三つの処理を行ってから貼り付けを行います。

貼り付ける線画からαチャンネルを削除し、透明度をなくして2値化
画像を構成する主要な色を取得し、主要な色となるべく異なる色で線画を塗りつぶし
入力線画をOpencvのdilate関数を使って太くする

この太くした線画を線画①とします

処理結果としては以下のようになります。

さらに、今度は以下の二つの処理を行った後にもう一度入力線画を上から張り付けます。

貼り付ける線画からαチャンネルを削除し、透明度をなくして2値化
画像を構成する主要な色を取得し、主要な色となるべく異なる色かつ、太くした線画とは異なる色で線画を塗りつぶし

この線画を線画②とします

拡大して見てみます。

生成された線画と入力線画のズレが線画①によって塗りつぶされており、さらにその上から入力線画と1ピクセルもずれがない線画②が貼られた状態になっています。
線画①の領域はRGB値で取得できるので、このRGB値と同じ値をもつピクセルを周辺と辻褄があうように修正してやればよいというわけです。

先ほど、「生成される線画と入力した線画のズレが少ない場合」は古典的画像処理である程度解決が可能と書いたのですが、これをより正確に表すと、「生成される線画と入力した線画のズレが、線画①によって太くした幅の中に納まる場合」は解決が可能、となります。

また、線画①の領域と周辺の辻褄合わせの手順は以下のようにして行います。

線画①と同じRGB値のピクセルをすべて取得
取得したピクセルのうち、隣接するピクセルが線画①及び線画②と異なるRGB値であった場合、そのピクセルのRGB値で自身を上書き
もし、隣接する全てのピクセルのRGB値が線画①及び線画②のRGB値と同じだった場合は、何の操作も行わず処理を続行する
1-3の処理を、1で取得できるピクセルの数が変わらなくなるまで繰り返す

上記の手順をコード化したもの


def replace_color(image, color_1, color_2, alpha_np):
    # 画像データを配列に変換
    data = np.array(image)

    # RGBAモードの画像であるため、形状変更時に4チャネルを考慮
    original_shape = data.shape
    data = data.reshape(-1, 4)  # RGBAのため、4チャネルでフラット化

    # color_1のマッチングを検索する際にはRGB値のみを比較
    matches = np.all(data[:, :3] == color_1, axis=1)

    # 変更を追跡するためのフラグ
    nochange_count = 0
    idx = 0

    while np.any(matches):
        idx += 1
        new_matches = np.zeros_like(matches)
        match_num = np.sum(matches)
        for i in tqdm(range(len(data))):
            if matches[i]:
                x, y = divmod(i, original_shape[1])
                neighbors = [
                    (x-1, y), (x+1, y), (x, y-1), (x, y+1)  # 上下左右
                ]
                replacement_found = False
                for nx, ny in neighbors:
                    if 0 <= nx < original_shape[0] and 0 <= ny < original_shape[1]:
                        ni = nx * original_shape[1] + ny
                        # RGBのみ比較し、アルファは無視
                        if not np.all(data[ni, :3] == color_1, axis=0) and not np.all(data[ni, :3] == color_2, axis=0):
                            data[i, :3] = data[ni, :3]  # RGB値のみ更新
                            replacement_found = True
                            continue
                if not replacement_found:
                    new_matches[i] = True
        matches = new_matches
        if match_num == np.sum(matches):
             nochange_count += 1
        if nochange_count > 5:
            break

    # 最終的な画像をPIL形式で返す
    data = data.reshape(original_shape)
    data[:, :, 3] = 255 - alpha_np
    return Image.fromarray(data, 'RGBA')

わかりやすく処理過程をGif化したものがこちらになります。

青く塗られた部分（線画①の領域）が処理が進むごとに周辺色で塗りつぶされて行っていることがわかります。
最終的に得られる結果がこちら

青く塗られた部分（線画①の領域）が、周辺色で塗りつぶされほとんどなくなっていることがわかります。
ただし、青く塗られた部分（線画①の領域）が残ってしまっている部分もあります。
これは、線画②で囲まれた線画①の領域に、1ピクセルも他の色が存在しない場合に発生します。
色の参照が可能な隣接ピクセルが存在しないため、このような事が起こってしまう訳ですね。

この問題は、線画①を張り付ける際に、線を太くする幅を抑えることである程度抑制することができます。
しかしその場合は、生成される線と入力線画のズレもその分抑えられている必要があるため、より画像生成時の忠実度が求められる点に注意です。

現状、システムとして提供する際は、この残ってしまった線画①の領域は透過した状態で出力するようにしています。
このようにすることで、後から残ってしまった領域を埋める際に、自分の好きな色で塗りつぶしやすくなります。

最後に

今回は入力線画を全く変えずに画像生成AIに色塗りさせる方法について解説をしていきました。
この記事で紹介した手法に則って線画着彩を行うプログラムを、以下のgithub リポジトリに公開しています。
もし興味がある方がいらっしゃいましたら、ぜひ遊んでみてください！