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

[Supership](https://supership.jp/)の名畑です。まさかの[てーきゅう新作動画](https://www.nicovideo.jp/watch/sm44272728)が公開。[てーきゅう](http://te-kyu.com/)の第10期が放送されるのを私はいつまでも待ち続けておりますよ。

## はじめに

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

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

## 入力画像と出力画像

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

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

### 入力画像

![input.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/869934/1f5318df-68a2-bba0-73ad-14d010653b82.png)

### 出力画像

![output.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/869934/810e939f-a3f9-9faf-a58b-f12d0efd9e82.png)

## ライブラリ

API呼び出しのために[requests](https://pypi.org/project/requests/)を入れておきます。

```shell-session
$ pip install requests
```

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

```shell-session
$ pip list | grep  -e requests
requests                  2.31.0
```

## API Key

今回、背景の生成には[Stability AI](https://ja.stability.ai/)のWeb APIを用います。

そのために[Developer Platform](https://platform.stability.ai/)でAPI Keyを発行する必要があります。価格は[Pricing](https://platform.stability.ai/pricing)をご参照ください。

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

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

```shell-session
$ open ~/.zshrc
```

今回は**STABILITY\_API\_KEY**という名称で保存しました。

```shell-session
export STABILITY_API_KEY =ここに取得したAPI keyを書く
```

## コード

今回は[Remove Background API](https://platform.stability.ai/docs/api-reference#tag/Edit/paths/~1v2beta~1stable-image~1edit~1remove-background/post)を用いて人物のみを切り抜き、切り抜いた画像を[Inpaint API](https://platform.stability.ai/docs/api-reference#tag/Edit/paths/~1v2beta~1stable-image~1edit~1inpaint/post)に渡した上で生成を行います。

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

```python
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](https://github.com/danielgatis/rembg)を使う方法も紹介しておきます。こちらの方法であればAPIの呼び出しが一度で済むためコストを安く抑えられますが、切り抜き精度に差があります。

まずインストール。

```shell-session
$ pip install rembg
```

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

```shell-session
$ pip list | grep -e rembg
rembg                     2.0.59
```

コードは以下です。

```python
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()))
```

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

![output_2.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/869934/2a11c0f4-808a-19cb-f3f3-4ea65a034941.png)

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

## 最後に

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

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


## 宣伝

[SupershipのQiita Organization](https://qiita.com/organizations/supership-inc)を合わせてご覧いただけますと嬉しいです。他のメンバーの記事も多数あります。

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