Amazon Kendra の Custom Document Enrichment と Amazon Bedrock で画像検索に対応する

こんにちは、機械学習チーム YAMALEX の駿です。
YAMALEX は Acroquest 社内で発足した、会社の未来の技術を創る、機械学習がメインテーマのデータサイエンスチームです。
（詳細はリンク先をご覧ください。）

皆さんは、「前のプレゼン資料に使った、犬の画像はどこいったかな？あの画像が欲しいので、探してくれないかな？」と無茶振りされたことはありませんか？

そんな時でも、「舌を出して喜んでいる」と検索すれば画像がヒットし、こんな無茶振りにも応えることができるシステムを Amazon Kendra （以下、 Kendra ）で構築しました。

舌を出して喜んでいる犬

ちょっと待って

Kendra は機械学習を利用した検索サービスで、ウェブサイトや S3 に保存したドキュメントなどをもとに、適切な検索結果を返します。

しかし、 Kendra で検索できるのはテキストだけで、画像を S3 に保存しても Kendra には取り込まれず無視されてしまう、という問題がありそうです。

そこで今回は、 Kendra の Custom Data Enrichment と、 Amazon Bedrock で利用できる Claude 3 を利用して、画像の内容を文字列化することで、 Kendra に取り込んで検索をできるようにしてみました。

概要
- Amazon Kendra Custom Data Enrichment とは？
- Amazon Bedrock とは？
構成
- 1. 取り込み
- 2. クエリ
事前構築
結果
まとめ

概要

Amazon Kendra Custom Data Enrichment とは？

Kendra にはドキュメントの取り込みプロセス中にコンテンツとドキュメントのメタデータを変更する Custom Data Enrichment （以下、 CDE ）という機能があります。

ドキュメントのパスや更新日時を使って簡単な条件式をもとにメタデータを追加するほかに、 AWS Lambda の関数を呼び出して Kendra に取り込まれる前のデータを加工したり、読み込み後にメタデータを追加したりすることが可能です。
例えば画像に対して OCR を行ってテキストを抽出する、テキストを翻訳する、といった処理を行うことで、必要なドキュメントが検索しやすくなります。

今回はテキスト情報が含まれない画像を検索対象にするため OCR は使わずに、画像の説明文を LLM に生成させ、これをコンテントとして Kendra に取り込みます。

docs.aws.amazon.com

Amazon Bedrock とは？

Claude 3 は Bedrock 上で利用可能な Anthropic が開発した生成 AI モデルです。

今回はその中でも高速かつ軽量な Claude 3 Haiku を利用します。
Claude 3 モデルファミリーは、テキストだけではなく画像認識が可能な言語モデルですが、最も軽量な Haiku でも、ちゃんと画像認識をしてくれます。

aws.amazon.com

構成

1. 取り込み

S3 に置いた画像を Kendra の CDE を使って文字列に変換し、 Kendra のインデックスに取り込みます（ Sync ）。
CDE は Lambda 関数を呼び出し、その中で Claude 3 が画像を文字列に変換しています。

Kendra の Sync 時の流れ

① [User] Kendra の Sync を実行する
② [Kendra] S3 に置かれたインデックス対象のドキュメントを走査し、S3 上の一時置き場にコピーする
（一時置き場は Kendra の CDE 設定で指定した S3 バケットに、 Kendra が自動で作成します）
③ [Kendra] ドキュメント毎に Lambda を実行する
④ [Lambda] S3 の一時置き場からファイルをダウンロードする
⑤ [Lambda] Bedrock の Claude 3 を呼び出し、画像を文字列に変換する
⑥ [Lambda] 変換後のファイルを S3 に保存する
⑦ [Kendra] 変換後のファイルを Kendra のインデックスに取り込む

2. クエリ

FastAPI 使って、文字列を入力として Kendra のインデックスを検索し、ヒットしたドキュメントの画像を画面に表示するアプリを作成しました。

画像本体は Kendra のインデックスに保存されず、 S3 のパスのみ取得できるため、画面側で S3 に取りに行くようにしています。

クエリ時の流れ

① [アプリ] ユーザから検索テキストを受け取る
② [Kendra] 検索テキストを使って Kendra のインデックスを検索し画像パスを返す
③ [アプリ] 画像を S3 からダウンロード

事前構築

1. 検索対象の画像を S3 に保存

検索対象のデータとして、下のデータセットから犬、猫、ハムスターを各 70 枚ほどずつに絞って、あらかじめ S3 に保存します。

S3 内のディレクトリ構造は次のようにします。

s3_bucket/
└── images/
     ├── cat/
     │   ├── xxxx.jpg
     │   └── ...
     ├── dog/
     │   ├── yyyy.jpg
     │   └── ...
     └── ...

www.kaggle.com

2. Kendra の CDE で用いる Lambda 関数をデプロイ

S3 から画像をダウンロードし、 Claude 3 Haiku で説明文を生成、 S3 にテキストをアップロードする関数を作ります。
Kendra の CDE に設定して取り込み前に実行されるものです。

例えば次の画像はこのような説明文が生成され、 Kendra に取り込まれます。

この画像には、黒い犬が写っています。犬は口の中に大きなピンクのドーナツ型のおもちゃを咥えており、楽しそうに遊んでいます。犬の表情は喜びに満ちており、おもちゃを噛んで遊ぶ様子が捉えられています。

犬は緑の芝生の上に立っており、背景には木製のフェンスが見えます。フェンスの向こうには庭の家具が置かれた庭園が広がっています。全体的に自然豊かな雰囲気の中で、犬が楽しく遊んでいる様子が伝わってきます。

犬

おもちゃ

庭園

注意点

下に示すコードは内部で pillow を使っているため、 pillow を含めた媒体をデプロイする必要があります。
タイムアウトは10秒～30秒にしておく必要があります。

800x800 の画像を1枚処理するのに8秒ほどかかっていました。
IAM ロールには次の権限がついていることを確認してください。
- bedrock:InvokeModel
- s3:GetObject
- s3:PutObject

import base64
import io
import json
import textwrap
import xml.etree.ElementTree as ET

import boto3
from PIL import Image

s3 = boto3.client("s3")
bedrock = boto3.client("bedrock-runtime")

MODEL_ID = "anthropic.claude-3-haiku-20240307-v1:0"


def invoke(img_b64: str, text: str) -> str:
    """Claudeを呼び出す"""
    user_content = [
        {
            "type": "image",
            "source": {"type": "base64", "media_type": "image/jpeg", "data": img_b64},
        },
        {"type": "text", "text": text},
    ]
    body = {
        "anthropic_version": "bedrock-2023-05-31",
        "messages": [{
            "role": "user",
            "content": user_content,
        }],
        "temperature": 0,
        "max_tokens": 1000,
    }

    response = bedrock.invoke_model(
        body=json.dumps({k: v for k, v in body.items() if v is not None}),
        modelId=MODEL_ID,
    )

    body = json.loads(response.get("body").read())
    contents = body.get("content", [])
    texts = []
    for content in contents:
        res_text = content.get("text")
        if res_text:
            texts.append(res_text)
    result = " ".join(texts)

    return result


def lambda_handler(event, _):
    """画像を読み込みテキスト化する"""
    # ソースのbucket,keyからpreExtraction用に一時的にコピーされたオブジェクトのbucket,keyが取得できます。
    bucket = event.get("s3Bucket")
    key = event.get("s3ObjectKey")

    if not key.endswith(".jpg"):
        # JPGでない場合、何もせずに取り込む
        return {
            "version": event.get("version", "v0"),
            "s3ObjectKey": key,
        }

    img = Image.open(s3.get_object(Bucket=bucket, Key=key)["Body"])
    # 下2つの関数はpillowを使って画像をリサイズ、Base64化する関数です。詳細割愛
    img = resize(img, max_size=800)
    img_b64 = image_to_base64(img)

    # 画像を Claude 3 を使って説明文に変換する
    prompt = textwrap.dedent("""\
        この画像について説明してください。
        出力フォーマットは次の通りです。

        ```
        <explanation>
            <foreground>ここで前景について説明してください</foreground>
            <background>ここで背景について説明してください</background>
            <keywords>
            <keyword>画像を良く説明する３つのキーワードを書いてください</keyword>
            </keywords>
        </explanation>
        ```
    """)
    response = invoke(img_b64, prompt)

    # Claude用にXMLで出力させているので、整形
    root = ET.fromstring(response)
    foreground_text = root.find("foreground").text
    background_text = root.find("background").text
    keywords = root.findall(".//keyword")
    keywords_text = "\n".join(f"- {keyword.text}" for keyword in keywords)

    output = f"{foreground_text}\n\n{background_text}\n\n{keywords_text}"

    # テキスト化後のコンテンツをS3に保存
    updated_key = key.replace(".jpg", ".txt")
    s3.put_object(Bucket=bucket, Key=updated_key, Body=output)

    category = key.split("/")[-2]

    # ファイルの場所とメタデータの更新内容を返す
    result = {
        "version": event.get("version", "v0"),
        "s3ObjectKey": updated_key,
        "metadataUpdates": [
            {"name": "_category", "value": {"stringValue": category}},
        ],
    }
    return result

3. Kendra で CDE を設定

ここではインデックスとデータソースは作成済みで Sync 前の状態を仮定します。
インデックス、データソースの作成方法はこちらの記事を参考にしてください。

acro-engineer.hatenablog.com

サイドメニューの Enrichments>Document enrichments を選択し「 Add document enrichment 」を押下
最初に basic operations を設定する画面ではデータソースのみ選択して、「 Next 」を押下
Lambda 関数の設定をする画面で Lambda 関数の ARN などを指定します
- pre-extraction に Lambda 関数の ARN と S3 Bucket を指定
  
  ここで指定する S3 バケットはデータソースに使っているものと一致している必要はありませんが、今回は同じバケットを選択しました。
  この後 Sync を実行した際に、このバケットに pre-extraction/ というディレクトリが作られ、取り込み前の一時保存場所として使われます。
- IAM ロールは自動で新しいロールを作成するオプションを選択
  
  S3 の Read/Write 権限や Lambda 関数を実行する権限が設定されます。
- 他の値は空のままで「 Next 」を押下
設定内容を確認して、「 Add document enrichment 」を押下

以上で、取り込み（ Sync ）を実行する準備が整いました。

CDE を追加する

Lambda 関数 ARN と S3 Bucket を指定する

結果

いろいろな文章で検索を実行してみました。

想像していたよりもクエリ通りの画像が表示されています。

おもちゃを口にくわえた

芝生に座っている

舌を出して喜んでいる

まとめ

Amazon Kendra の Custom Data Enrichment と Amazon Bedrock で使える Claude 3 Haiku を使って、テキストで画像を検索することができました。
取り込みの前後で Lambda を動かせると、なんでもできてしまうので、活用の可能性が広がりますね。

個人的には、「舌を出して喜んでいる」の右下の子が好きです。

Acroquest Technologyでは、キャリア採用を行っています。

ディープラーニング等を使った自然言語／画像／音声／動画解析の研究開発

Elasticsearch等を使ったデータ収集／分析／可視化

マイクロサービス、DevOps、最新のOSSやクラウドサービスを利用する開発プロジェクト

書籍・雑誌等の執筆や、社内外での技術の発信・共有によるエンジニアとしての成長

　少しでも上記に興味を持たれた方は、是非以下のページをご覧ください。 www.wantedly.com

Taste of Tech Topics

Acroquest Technology株式会社のエンジニアが書く技術ブログ