[rembg x AWS Lambda] 画像背景除去Slackアプリを作る
はじめに
初めて記事を書いてみます。お手柔らかに。。
みなさんは画像の背景部分消したいなあって思ったことはありますか？
いつのまにかiPhoneでもなんかできるようになってることに驚き。
ECサイトなど商品写真を扱うことがあるとけっこう使うんじゃないでしょうか。
フォトショで被写体検知して切り取り。。みたいな。
今回はそれを自動化してみたいと思います。
対象読者
- 画像から背景を除去したい方
- AWSアカウントをお持ちの方
- Slackアカウントをお持ちの方
- AWS Lambdaを使ったことがある方
- Pythonを使ったことがある方
技術スタック
- AWS Lambda
- Python 3.9
- Docker
- rembg
- aws sam cli
rembgについて
rembgは関数一発で画像の背景を除去できるやばい便利なPythonライブラリです。
from rembg import remove
from PIL import Image
im = Image.open("画像パス")
output = remove(im) # これ
これだけです。。びっくり。
概要
- Slackアプリに画像を送る
- 画像アップロードをトリガーにAWS Lambdaを関数URLから発火
- rembgでアップロード画像を背景除去
- Slackに背景除去された写真を送信
とてもシンプルな構成ですがちょこちょこ工夫しなきゃいけない点があるので説明していきます。
コード解説(Lambda)
まずは実際に背景除去処理をするDocker ImageのLambdaを作成します。
コードはGithubに置いてありますので動かしてみたい方はcloneしてみてください。
Slackまわりの処理
import os
from typing import Any, Final, Optional
import uuid
import requests
class SlackService:
def __init__(self) -> None:
self._verification_token: Final = os.environ["VERIFICATION_TOKEN"]
self._api_token: Final = os.environ["API_TOKEN"]
self._bot_user_id: Final = os.environ["BOT_USER_ID"]
def auth(self, body: dict) -> Optional[dict[str, Any]]:
"""
SlackのEvent SubscriptionsのURL認証を行う。
認証後のリクエストでは、ボット自身のメッセージは無視する。
無視しないと、ボットが自身のメッセージにLambdaが反応して無限ループに陥る。
"""
event_type = body.get("type")
if event_type == "url_verification" and body["token"] == self._verification_token:
return {
"challenge": body["challenge"]
}
if body["event"]["user_id"] == self._bot_user_id:
return {
"statusCode": 200
}
return None
def download_file(self, file_id: str, download_dir: str) -> str:
"""
file_idでSlackにアップロードされたファイルをダウンロードする。
"""
url_private_download, original_filename = self._get_url_private_download(
file_id)
headers = {"Authorization": f"Bearer {self._api_token}"}
file_response = requests.get(
url_private_download, headers=headers, stream=True)
if file_response.status_code == 200:
download_path = os.path.join(
download_dir, f"{uuid.uuid4()}.{original_filename.split('.')[-1]}")
with open(download_path, "wb") as f:
for chunk in file_response.iter_content(chunk_size=1024):
if chunk:
f.write(chunk)
return download_path
else:
raise ValueError(f"Failed to download file with id {file_id}")
def upload_file(self, channel_id: str, file_path: str, message: Optional[str] = None) -> None:
"""
指定したチャンネルにファイルをアップロードする。
"""
url = "https://slack.com/api/files.getUploadURLExternal"
headers = {"Authorization": f"Bearer {self._api_token}"}
with open(file_path, 'rb') as f:
payload = {
"filename": os.path.basename(file_path),
"length": os.path.getsize(file_path),
}
res = requests.post(url, headers=headers, data=payload)
if not res.ok:
raise ValueError(
f"Failed to get upload URL for channel {channel_id}")
res_body = res.json()
files = {'file': f}
upload_res = requests.post(
res_body["upload_url"], headers=headers, files=files)
if not upload_res.ok:
raise ValueError(
f"Failed to upload file to channel {channel_id}")
complete_url = "https://slack.com/api/files.completeUploadExternal"
payload = {
"files": [
{"id": res_body["file_id"],
"title": os.path.basename(file_path)}
],
"channel_id": channel_id,
}
if message:
payload["initial_comment"] = message
complete_res = requests.post(
complete_url, headers=headers, json=payload)
if not complete_res.ok:
raise ValueError(
f"Failed to complete upload for channel {channel_id}")
def _get_url_private_download(self, file_id: str) -> tuple[str, str]:
url = f"https://slack.com/api/files.info"
send_data = {
"file": file_id,
"token": self._api_token
}
res = requests.post(url, send_data)
if not res.ok:
raise ValueError(
f"Failed to get url_private_download for file {file_id}")
body = res.json()
url_private_download = body["file"]["url_private_download"]
name = body["file"]["name"]
return url_private_download, name
rembgまわり
環境設定
rembgをLambdaで使うにはちょっと工夫が必要です。
rembgはu2net.onnxという学習モデルを使用して背景除去処理をします。
最初にモデルをこちらからダウンロードしておき、Dockerイメージに含ませておきましょう。(Dockerfileの中でやってもよかったね。)
処理実行時に上記の学習モデルがない場合はダウンロードしてから処理を開始してくれますが、Lambda発火後にダウンロードだとレスポンスが遅いかつ効率が悪いので。
FROM public.ecr.aws/lambda/python:3.9
ENV NUMBA_DISABLE_JIT=1
WORKDIR ${LAMBDA_TASK_ROOT}
COPY ./src ./
COPY ./u2net.onnx ./
RUN python3.9 -m pip install -r requirements.txt -t .
CMD ["app.lambda_handler"]
ENV NUMBA_DISABLE_JIT=1という環境変数がちょっと工夫しないといけない点その2です。
rembgは内部的にnumbaを使用しており、 multiprocessingが使われています。
並行処理を行う際にdev/shmディレクトリをメモリとして使用しますが、Lambdaはご存知の通りtmpディレクトリのみに書き込み権限があるため、失敗してしまいます。
そこで上記の環境変数を使ってjitを無効化し、multiprocessingの使用を回避します。
画像背景除去処理
import os
from typing import Final
from rembg import remove
from PIL import Image
from PIL.Image import Image as PILImage
import onnxruntime as ort
from rembg.session_simple import SimpleSession
class ImageService:
def __init__(self, u2net_home: str) -> None:
"""
u2net.onnxの配置パスを引数に、rembgが使う学習モデルのセッションを作成する。
rembgが内部で自動でセッションをつくってくれるが、ディレクトリ作成する部分があり、
Lambdaではエラーが発生するため、事前にセッションを作成しておく。
"""
sess_opts = ort.SessionOptions()
self._session: Final = SimpleSession("u2net", ort.InferenceSession(
u2net_home,
providers=ort.get_available_providers(),
sess_options=sess_opts,
),
)
def remove_background(self, input_path: str) -> str:
"""
指定された画像パスの背景を除去する
"""
filename = os.path.basename(input_path)
p = input_path.replace(filename, f"removed_{filename}")
im = Image.open(input_path)
output: PILImage = remove(im, session=self._session) # type: ignore
if output.mode != "RGB":
output = output.convert("RGB")
output.save(p)
return p
lambda_handler
import json
import os
import tempfile
from image import ImageService
from slack import SlackService
slack_service = SlackService()
u2net_home = os.path.join(os.path.dirname(__file__), "u2net.onnx")
image_service = ImageService(u2net_home)
def lambda_handler(event, _):
"""
Slackのヘッダーにリトライの情報がある場合は、無視する。
それ以外の場合は、Slackから送られてきた画像を背景除去して、アップロードする。
"""
if "headers" in event and "x-slack-retry-reason" in event["headers"]:
return {
"statusCode": 200
}
body = json.loads(event["body"])
res = slack_service.auth(body)
if res:
return res
file_id = body["event"]["file_id"]
with tempfile.TemporaryDirectory() as temp_dir:
download_path = slack_service.download_file(file_id, temp_dir)
output = image_service.remove_background(download_path)
channel_id = body["event"]["channel_id"]
slack_service.upload_file(channel_id, output, "お待たせしました。。。")
Slack App作成
長くなるので作成の大枠についてはこの記事では割愛させていただき、
Event Subscriptionsまわりの設定のみ簡単に記載していきます。
このあたりはたくさんわかりやすい記事がありますのでそちらを参考いただけると幸いです。
作ったBot
4月ですし、背景削除する部署に入った新卒さんをイメージしてつくりました。
ちょうど時期的にも研修が終わってきたりして帰り道はこんな感じじゃないかなと思います。ファイト！
Event Subscriptions設定
Botを作成したらEvent Subscriptionsの設定をします。
- Your Appsから作成したBotを選択
- FeaturesのEvent Subscriptionsを選択
- Enable EventsをOnにする
- Subscribe to events on behalf of usersに
file_sharedの権限を追加
OAuth & Permissions設定とBot User OAuth Token取得
次にBotの権限設定を行います
- OAuth & Permissionsを選択
- ScopesのBot Token Scopesに
chat:writeとfiles:writeの権限を追加 - OAuth Tokens for Your Workspaceに記載のBot User OAuth Tokenを控える
Verification Token取得
Lambdaの関数URLをEvent Subscriptionsに認証させるために必要になります
- Basic Infomationを選択
- App Credentialsに記載のVerification Tokenを控える
作成したBotのメンバーIDを取得
SlackにいるBotを右クリックしてアプリの詳細を表示するをクリックで表示されます。
前段で若干触れていますが、BotからのEventかどうかをLambda側で判断するために使用します。
Event Subscriptionsで設定したfile_sharedの権限はファイルがチャンネルで共有されたときに発火します。
なのでBotからの送信かどうかを判断して無視するようにしないと、
Slackに画像ファイルアップロード -> Event発火 -> LambdaからSlackに画像ファイルアップロード -> Event発火のLambda無限ループになりますので、人生終えたい方以外はしっかり確認しておきましょう。
Lambdaデプロイ
SAM CLIを使用してデプロイします。
AWSTemplateFormatVersion: "2010-09-09"
Transform: AWS::Serverless-2016-10-31
Parameters:
VerificationToken:
Type: String
ApiToken:
Type: String
BotUserId:
Type: String
Resources:
Function:
Type: AWS::Serverless::Function
Properties:
FunctionName: "slack-rembg-lambda"
PackageType: Image
MemorySize: 5120
EphemeralStorage:
Size: 512
Timeout: 360
Environment:
Variables:
VERIFICATION_TOKEN: !Sub "${VerificationToken}"
API_TOKEN: !Sub "${ApiToken}"
BOT_USER_ID: !Sub "${BotUserId}"
Metadata:
Dockerfile: Dockerfile
DockerTag: slack-rembg-lambda
DockerContext: ./
Permission:
Type: AWS::Lambda::Permission
Properties:
Action: lambda:InvokeFunctionUrl
FunctionName: !GetAtt Function.Arn
FunctionUrlAuthType: NONE
Principal: "*"
EventInvokeConfig:
Type: AWS::Lambda::EventInvokeConfig
Properties:
FunctionName: !Ref Function
MaximumRetryAttempts: 0
Qualifier: "$LATEST"
LogGroup:
Type: AWS::Logs::LogGroup
Properties:
LogGroupName: !Sub "/aws/lambda/${Function}"
Url:
Type: AWS::Lambda::Url
Properties:
AuthType: NONE
TargetFunctionArn: !GetAtt Function.Arn
Outputs:
UrlOutput:
Value: !GetAtt Url.FunctionUrl
Export:
Name: FunctionUrl
ビルドします
sam build
LambdaのDocker Imageの置き場所としてECRにリポジトリを作成します
aws ecr create-repository --repository-name <お好きなリポジトリ名> --region <リージョン> --profile <プロファイル名>
Lambdaのソースの置き場所としてS3にバケットを作成します
aws s3api create-bucket --bucket <お好きなバケット名> --region <リージョン> --create-bucket-configuration LocationConstraint=<リージョン> --profile <プロファイル名>
Slackから取得したトークンやメンバーIDも合わせ値をセットし、デプロイします
sam deploy \
--parameter-overrides VerificationToken=<Verification Token> ApiToken=<Bot User Auth Token> BotUserId=<メンバーID> \
--stack-name slack-rembg-lambda-stack \
--region <リージョン> \
--s3-bucket <作成したS3バケット名> \
--capabilities CAPABILITY_NAMED_IAM \
--image-repository <AWSアカウントID>.dkr.ecr.<リージョン>.amazonaws.com/<作成したリポジトリ名> \
--profile <プロファイル名>
Lambdaの関数URLがOutputsに表示されるため、控えておきましょう。
Key UrlOutput
Description -
Value https://xxxx.lambda-url.region.on.aws/ <-これ
Botを動かしてみる
デプロイが成功したらさっそくBotを動かしてみましょう。
Event SubscriptionsのURL認証
実際に動かす前に、Event Subscriptionsは初回実行時にはまずトリガー発火後にリクエストするURLを認証する必要があります。
Enable EventsにあるRequest URLに先程のLambdaの関数URLを入力し、認証しましょう。
以下のようにVerifiedにチェックが入れば成功です。
※Lambdaのコールドスタート時はレスポンスが遅くなり、認証が成功しないことがあるので、失敗したらすぐリトライしましょう。
画像をBotに送信する
今回はこのりんごの画像を使うことにします。
画像を送信してみると、背景が除去されたりんごが返ってきました！
「お待たせしました。。」が疲れてる感じがでていいですね。
最後に
いかがだったでしょうか。今回は画像１枚の簡単なサンプルなのでユースケースとしては少ないかもしれませんが、Dockerベースなので、LambdaをECSやApp Runnerで応用して一括処理したりなどけっこう幅は広がるかなという印象です。
Event Subscriptionsの無限ループだけお気をつけください。<- 軽くやらかしたのはナイショ。
Discussion
便利な連携ですね！無限ループなどは公式 SDK の bolt-python を使うと何も考えなくても防げますし、リトライの問題も lazy listener を使うと綺麗に解決できるので、次に何か作るときはぜひ試してみてください：
なるほど！boltって色々できるんですね。。
試してみます！ありがとうございます！