[Stable Diffusion] AIで指定範囲をテキストで指示して画像修復 [inpainting]

2022年8月30日火曜日

Artificial Intelligence

本記事では、機械学習手法Stable Diffusionを用いて画像の特定の範囲をテキストに応じた画像に変更する画像修復を実行する方法をご紹介します。

アイキャッチ

Stable Diffusion

概要

Stable Diffusionは、拡散モデルによる画像合成モデルです。
自然言語で入力されたテキスト(prompt)から画像を生成するText to Imageタスクを実現します。

Latent Diffusionをベースとした本モデルは、非常に大規模なデータセットであるLAION-5Bを用いてトレーニングされています。

GoogleのImagenやOpenAIのGLIDE, DALLE-2など最新の画像生成モデルは社会に与える影響が大きいことなどを理由に学習済みモデルの公開を控えるケースが目立っていましたが、
Stable Diffusionでは、AIの平等な使用機会を与えることを目的に学習済みモデルの公開に踏み切っています。

なお、以下の記事では有償にはなりますが、Lattent DiffusionやGLIDEの技術解説やデモを紹介しています。よろしければご参照ください。
技術解説までは、無料で参照可能です。

本記事では上記手法を用いて、inpaitingタスクをGoogle Colaboratoryで実行していきます。

なお、テキストから画像を生成するtext2imgタスク、画像とテキストを入力し新たに画像を生成するimg2textは以下でご紹介しています。

デモ(Colaboratory)

それでは、実際に動かしながらinpaintingタスクを実行していきます。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

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

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

環境セットアップ

初めにHuggingFaceのアカウントを作成します。
HuggingFaceにアクセスし画面右上のSignUpよりアカウントを作成します。
登録したメールアドレスに認証メールが届くのでメールに記載されたリンクにアクセスしてアカウント登録は完了です。

続いて、こちらのCompVis/stable-diffusion-v1-4にアクセスし記載の内容を確認の上承認します。

承認

承認後、画面右上のアカウントのアイコンから[Settings]->[Access Tokens]に移動しNewTokenを発行し、メモしておきます。

Settings

以降の作業はGoogle Colaboratoryで行います。

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

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

まずGithubからコードをcloneします。

%cd /content

!git clone https://github.com/huggingface/diffusers.git
!sed -E -i "s/sample\(generator=generator\)/sample\(\)/" /content/diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py

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

%cd /content

!pip install diffusers transformers scipy ftfy

最後にライブラリをインポートします。

import matplotlib.pyplot as plt
from PIL import Image, ImageDraw

import torch
from torch import autocast

%cd /content/diffusers

from src.diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint import StableDiffusionInpaintPipeline

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

学習済みモデルのセットアップ

続いて、先ほど発行したアクセストークンを使用して、モデルをダウンロードします。

%cd /content/diffusers
access_tokens="先ほどメモしたAccessToken" # @param {type:"string"}

# load pretrain model

%cd /content/diffusers

device = "cuda"
pipe = StableDiffusionInpaintingPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    revision="fp16", 
    torch_dtype=torch.float16,
    use_auth_token=access_tokens
).to(device)

マスク画像生成

それでは、promptと画像を設定し指定範囲を画像修復していきます。
画像修復は、マスク画像で指定された白色領域のみに対して画像生成することで実現します。

まず、入力画像の取得と、マスク画像の生成を行います。

img_w, img_h = 768, 512

%cd /content/diffusers

init_image = Image.open(filename)
if '.png' in filename:
  init_image = init_image.convert('RGB')
init_image = init_image.resize((img_w, img_h))

plt.imshow(init_image)
plt.axis('off')
plt.show()

以下の画像を使用します。

input

続いて、画像修復範囲を矩形で指定します。

x1 = 400  #@param {type:"integer"}
y1 = 0    #@param {type:"integer"}
x2 = 650  #@param {type:"integer"}
y2 = 300  #@param {type:"integer"}

%cd /content/diffusers

temp = Image.open(filename)
if '.png' in filename:
  temp = temp.convert('RGB')
mask_img = temp.resize((img_w, img_h))
draw = ImageDraw.Draw(mask_img)

# 矩形の設定
rectcolor = (255, 0, 0)
linewidth = 4
draw.rectangle(
    [(x1, y1), (x2, y2)], # 左上x,y 右下x,y
    outline=rectcolor, width=linewidth,
    )
mask_img

出力結果は以下の通りです。

矩形設定後画像

指定された矩形に従ってマスク画像を生成します。

# mask画像生成

# 一度すべて黒く塗りつぶし
rectcolor = (0, 0, 0)
draw.rectangle(
    [(0, 0), (img_w, img_h)],
    outline=rectcolor, fill=(0, 0, 0)
    )
# 先ほど設定した矩形箇所白く塗りつぶし
rectcolor = (255, 255, 255)
draw.rectangle(
   [(x1, y1), (x2, y2)],
    outline=rectcolor, fill=(255, 255, 255)
    )
mask_img

以下のマスク画像が生成されます。

マスク画像

白色の領域が修復範囲です。

Inpainting

入力画像、マスク画像がそろったのであとはモデルに入力するのみです。

promptganesha faceを設定して実行してみます。

prompt = "ganesha face" #@param {type:"string"}

generator = torch.Generator(device).manual_seed(12)

with autocast("cuda"):
  images = pipe(
      prompt=prompt,
      init_image=init_image,
      mask_image=mask_img,
      strength=0.75,
      guidance_scale=7.5,
      generator=generator,
  )["sample"]

images[0].save("inpainting_01.png")

出力結果は以下の通りです。

result

矩形範囲のみガネーシャの顔に修復することができました。

まとめ

本記事では、Stable Diffusionを用いてinpaintingを行う方法をご紹介しました。
かなり高精度なモデルのため悪用厳禁であることは言うまでもありません。

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


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

参考文献

1.  論文 - High-Resolution Image Synthesis with Latent Diffusion Models

2. GitHub - CompVis/stable-diffusion

AIで副業ならココから!

まずは無料会員登録

プロフィール

メーカーで研究開発を行う現役エンジニア
組み込み機器開発や機会学習モデル開発に従事しています

本ブログでは最新AI技術を中心にソースコード付きでご紹介します


Twitter

カテゴリ

このブログを検索

ブログ アーカイブ

TeDokology