Zenn
🍎

[rembg x AWS Lambda] 画像背景除去Slackアプリを作る

2024/04/18に公開2件
AWS
Docker
Python
Slack
Lambda
tech

はじめに

初めて記事を書いてみます。お手柔らかに。。
みなさんは画像の背景部分消したいなあって思ったことはありますか？
いつのまにかiPhoneでもなんかできるようになってることに驚き。

ECサイトなど商品写真を扱うことがあるとけっこう使うんじゃないでしょうか。
フォトショで被写体検知して切り取り。。みたいな。

今回はそれを自動化してみたいと思います。

対象読者

  • 画像から背景を除去したい方
  • AWSアカウントをお持ちの方
  • Slackアカウントをお持ちの方
  • AWS Lambdaを使ったことがある方
  • Pythonを使ったことがある方

技術スタック

  • AWS Lambda
  • Python 3.9
  • Docker
  • rembg
  • aws sam cli

rembgについて

rembgは関数一発で画像の背景を除去できるやばい便利なPythonライブラリです。
https://github.com/danielgatis/rembg

from rembg import remove
from PIL import Image

im = Image.open("画像パス")
output = remove(im) # これ

これだけです。。びっくり。

概要

  1. Slackアプリに画像を送る
  2. 画像アップロードをトリガーにAWS Lambdaを関数URLから発火
  3. rembgでアップロード画像を背景除去
  4. Slackに背景除去された写真を送信

とてもシンプルな構成ですがちょこちょこ工夫しなきゃいけない点があるので説明していきます。

コード解説(Lambda)

まずは実際に背景除去処理をするDocker ImageのLambdaを作成します。
コードはGithubに置いてありますので動かしてみたい方はcloneしてみてください。
https://github.com/Tomoaki-Moriya/slack-rembg-lambda

Slackまわりの処理

src/slack.py
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発火後にダウンロードだとレスポンスが遅いかつ効率が悪いので。

Dockerfile
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の使用を回避します。
https://numba.pydata.org/numba-doc/dev/reference/envvars.html#envvar-NUMBA_DISABLE_JIT

画像背景除去処理

src/image.py
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

src/app.py
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まわりの設定のみ簡単に記載していきます。

このあたりはたくさんわかりやすい記事がありますのでそちらを参考いただけると幸いです。
https://zenn.dev/mokomoka/articles/6d281d27aa344e
https://zenn.dev/nyancat/articles/20211219-create-slack-app

作ったBot

4月ですし、背景削除する部署に入った新卒さんをイメージしてつくりました。
ちょうど時期的にも研修が終わってきたりして帰り道はこんな感じじゃないかなと思います。ファイト！

Event Subscriptions設定

Botを作成したらEvent Subscriptionsの設定をします。

  1. Your Appsから作成したBotを選択
  2. FeaturesのEvent Subscriptionsを選択
  3. Enable EventsをOnにする
  4. Subscribe to events on behalf of usersにfile_sharedの権限を追加

OAuth & Permissions設定とBot User OAuth Token取得

次にBotの権限設定を行います

  1. OAuth & Permissionsを選択
  2. ScopesのBot Token Scopesにchat:writeとfiles:writeの権限を追加
  3. OAuth Tokens for Your Workspaceに記載のBot User OAuth Tokenを控える

Verification Token取得

Lambdaの関数URLをEvent Subscriptionsに認証させるために必要になります

  1. Basic Infomationを選択
  2. 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を使用してデプロイします。

template.yml
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

Kazuhiro SeraKazuhiro Sera

便利な連携ですね！無限ループなどは公式 SDK の bolt-python を使うと何も考えなくても防げますし、リトライの問題も lazy listener を使うと綺麗に解決できるので、次に何か作るときはぜひ試してみてください：

tomotomo

なるほど！boltって色々できるんですね。。
試してみます！ありがとうございます！