Stable DiffusionのWeb APIを用いて画像上の人物はそのままに背景だけをAI生成で入れ替えてみた

Supershipの名畑です。まさかのてーきゅう新作動画が公開。てーきゅうの第10期が放送されるのを私はいつまでも待ち続けておりますよ。

はじめに

タイトルのままです。
Photoshop等の画像編集ソフトを使わず、コードのみで、画像上の人物をそのままに背景だけをAI生成したいというケースがあったので、そのコードを残します。Pythonです。

特に目新しいことはないですが、よくあるニーズだと思うので、どなたかの役に立つと嬉しいです。

入力画像と出力画像

今回の記事では先に結果をお伝えします。

以下の入力画像を元として、出力画像を生成します。

入力画像

出力画像

ライブラリ

API呼び出しのためにrequestsを入れておきます。

$ pip install requests

バージョンを確認しておきました。

$ pip list | grep  -e requests
requests                  2.31.0

API Key

今回、背景の生成にはStability AIのWeb APIを用います。

そのためにDeveloper PlatformでAPI Keyを発行する必要があります。価格はPricingをご参照ください。

ハードコーディングを避けるため、環境変数に保存しておきます。

私の環境での設定ファイルはzshrcです。

$ open ~/.zshrc

今回はSTABILITY_API_KEYという名称で保存しました。

export STABILITY_API_KEY =ここに取得したAPI keyを書く

コード

今回はRemove Background APIを用いて人物のみを切り抜き、切り抜いた画像をInpaint APIに渡した上で生成を行います。

プロンプトには「Large glass windows with a view of the metropolis behind」を指定しています。日本語訳すると「後ろに大都市の光景が広がっている大きなガラスの窓」です。

import os
import requests

# ファイルパス
input_path = './input.png'  # 元画像
mask_path = './mask.png'  # マスク用画像（途中生成/削除）
output_path = './output.png'  # 出力画像

# API Keyの取得確認
api_key = os.getenv("STABILITY_API_KEY")
if api_key is None:
    raise Exception("Missing Stability API key.")

headers = {
    "Accept": "image/*",
    "Authorization": f"Bearer {api_key}"
}

# Remove Background API 呼び出し
response = requests.post(
    f"https://api.stability.ai/v2beta/stable-image/edit/remove-background",
    headers=headers,
    files={
        "image": open(input_path, "rb")
    },
    data={
        "output_format": "png"
    },
)

# マスク画像保存
if response.status_code == 200:
    with open(mask_path, 'wb') as file:
        file.write(response.content)
else:
    raise Exception(str(response.json()))

# Inpaint API 呼び出し
response = requests.post(
    "https://api.stability.ai/v2beta/stable-image/edit/inpaint",
    headers=headers,
    files={
        "image": open(mask_path, "rb"),
    },
    data={
        "prompt": "Large glass windows with a view of the metropolis behind",  # プロンプト
        "output_format": "png",
        "grow_mask": 0,  # マスク周辺のぼかしをなくす
    },
)

# マスク画像削除
os.remove(mask_path)

# 出力画像保存
if response.status_code == 200:
    with open(output_path, "wb") as file:
        file.write(response.content)
else:
    raise Exception(str(response.json()))

rembgを使用した場合

画像の切り抜きにrembgを使う方法も紹介しておきます。こちらの方法であればAPIの呼び出しが一度で済むためコストを安く抑えられますが、切り抜き精度に差があります。

まずインストール。

$ pip install rembg

バージョンを確認しておきました。

$ pip list | grep -e rembg
rembg                     2.0.59

コードは以下です。

from rembg import remove
import os
import requests

# ファイルパス
input_path = './input.png'  # 入力画像のパス
mask_path = './mask.png'  # マスク画像のパス（途中生成/削除）
output_path = './output.png'  # 出力画像のパス

# 背景を除去したマスク画像生成
with open(input_path, 'rb') as i:
    with open(mask_path, 'wb') as o:
        input_image = i.read()
        mask_image = remove(input_image)
        o.write(mask_image)

# API Keyの取得確認
api_key = os.getenv("STABILITY_API_KEY")
if api_key is None:
    raise Exception("Missing Stability API key.")

# Inpaint API 呼び出し
response = requests.post(
    "https://api.stability.ai/v2beta/stable-image/edit/inpaint",
    headers={
        "Accept": "image/*",
        "Authorization": f"Bearer {api_key}"
    },
    files={
        "image": open(mask_path, "rb"),
    },
    data={
        "prompt": "Large glass windows with a view of the metropolis behind",  # プロンプト
        "output_format": "png",
        "grow_mask": 0,  # マスク周辺のぼかしをなくす
    },
)

# マスク画像削除
os.remove(mask_path)

# 出力画像保存
if response.status_code == 200:
    with open(output_path, "wb") as file:
        file.write(response.content)
else:
    raise Exception(str(response.json()))

出力画像は以下です。今回のケースだと切り抜きの精度もそれほど気にならないですかね。

Stable Diffusionの環境をローカルに構築すればAPI呼び出しコストは0にできますので、興味がある方は調べてやってみてください。

最後に

コードだけで完結させられるのが本当に楽で嬉しいです。

本当に様々な効率化が進んでいると実感する日々です。

宣伝

SupershipのQiita Organizationを合わせてご覧いただけますと嬉しいです。他のメンバーの記事も多数あります。

Supershipではプロダクト開発やサービス開発に関わる方を絶賛募集しております。
興味がある方はSupership株式会社 採用サイトよりご確認ください。