Amazon Nova でバウンディングボックスを描画しちゃおう

[ {"white dog": [114, 29, 395, 312]}, {"Toy carrot": [88, 132, 321, 233]} ]  

pip install boto3 Pillow

import base64
import boto3
import io
import json
import re # 正規表現を扱うライブラリ

from PIL import Image, ImageDraw, ImageFont # Pillow ライブラリから必要なモジュールをインポート

modelId = "us.amazon.nova-pro-v1:0" # 使用する Nova モデルの ID
accept = "application/json"
contentType = "application/json"
bedrock_rt = boto3.client("bedrock-runtime", region_name="us-east-1") # Bedrock Runtimeクライアントの作成 (リージョンは適宜変更)

def safe_json_load(json_string):
    try:
        json_string = re.sub(r"\s", "", json_string) # 空白文字の除去
        json_string = re.sub(r"\(", "[", json_string) # 丸括弧を角括弧に置換
        json_string = re.sub(r"\)", "]", json_string) # 丸括弧を角括弧に置換
        bbox_set = {}
        # モデルの出力が途中で切れている場合や、不正な形式になっている場合の補正処理
        for b in re.finditer(r"\[\d+,\d+,\d+,\d+\]", json_string):
            if b.group(0) in bbox_set:
                json_string = json_string[:bbox_set[b.group(0)][1]] + "}]"
                break
            bbox_set[b.group(0)] = (b.start(), b.end())
        else:
            json_string = json_string[:bbox_set[b.group(0)][1]] + "}]"
        json_string = re.sub(r"\]\},\]$", "]}]", json_string) # 末尾の余分な文字の修正
        json_string = re.sub(r"\]\],\[\"", "]},{\"", json_string) # 区切り文字の修正
        json_string = re.sub(r"\]\],\[\{\"", "]},{\"", json_string) # 区切り文字の修正
        return json.loads(json_string)
    except Exception as e:
        print(e)
        return []

def get_font(height):
    try:
        # 画像の高さに基づいてフォントサイズを調整
        font = ImageFont.truetype("Arial", size=height // 20)
    except (OSError, IOError):
        # Arialフォントが見つからない場合はデフォルトフォントを試みる
        try:
            font = ImageFont.load_default()
            print("Using default font as Arial font not found")
        except Exception as e:
            print(f"Error loading default font: {e}")
            # デフォルトフォントも読み込めない場合はエラーを発生させる
            raise

def detection(image_path, category_list, image_short_size=360):
    image_pil = Image.open(image_path)
    width, height = image_pil.size

    ratio = image_short_size/ min(width, height) # 画像の短辺が image_short_size になるようにリサイズ比率を計算
    width = round(ratio * width)
    height = round(ratio * height)

    image_pil = image_pil.resize((width, height), resample=Image.Resampling.LANCZOS)
    
    buffer = io.BytesIO()
    image_pil.save(buffer, format="webp", quality=90) # リサイズされた画像をメモリ上のバッファに webp 形式で保存
    image_data = buffer.getvalue()

    category_str = ",".join([f'"{category}"' for category in category_list])# プロンプト内で使用するために、カテゴリリストをカンマ区切りの文字列に変換
    
    # Amazon Novaモデルへの指示 (プロンプト) を作成
    prompts = """
Detect bounding box of objects in the image, only detect %s category objects with high confidence, output in a list of bounding box format.
Output example:
[
    {"%s": [x1, y1, x2, y2]},
    ...
]
""" % (category_str, category_list[0]) # category_list の最初の要素を例として使用

    # モデルの出力形式を誘導するための prefill を定義して、JSON 形式の出力に誘導
    prefill="""[
    {"
"""

    print("Input Prompt:\n", prompts)

    messages = [
        {
            "role": "user",
            "content": [
                {
                    "image": {
                        "format": 'webp',
                        "source": {
                            "bytes": image_data,
                        }
                    }
                },
                {
                    "text": prompts
                },
            ],
        },
        {
            "role": "assistant",
            "content": [
                {
                    "text": prefill
                },
            ],
        }
    ]

    # Amazon Bedrock Runtime クライアントを使用して、Amazon Nova モデルにリクエストを送信
    response = bedrock_rt.converse(
        modelId=modelId,       
        messages=messages,     
        inferenceConfig={      
            "temperature": 0.0,  # 出力のランダム性を制御
            "maxTokens": 1024,  
        },
    )

    output = prefill + response.get('output')["message"]["content"][0]["text"]

    result = safe_json_load(output)

    color_list = [
        'blue',
        'green',
        'yellow',
        'red',
        'orange',
        'pink',
        'purple',
    ]

    font = get_font(height) 


    print("Result:\n")

    for idx, item in enumerate(result):
        label = next(iter(item))
        bbox = item[label]
        x1, y1, x2, y2 = bbox
        if x1 >= x2 or y1 >= y2:# 座標値が不正 (例: x1 が x2 以上) でないか簡単なチェック
            print(f"Skipping invalid bbox for {label}: {bbox}")
            continue

        # Amazon Nova が出力する座標は[0, 1000] のスケールなので、リサイズされた画像 (image_pil) の寸法に合わせてリスケールする。
        w, h = image_pil.size 
        x1 = x1 / 1000 * w
        x2 = x2 / 1000 * w
        y1 = y1 / 1000 * h
        y2 = y2 / 1000 * h
        bbox = (x1, y1, x2, y2)
        bbox = list(map(round, bbox))
        print(f"Detect <{label}> in {bbox}")
        # draw bounding box
        draw = ImageDraw.Draw(image_pil)
        color = color_list[idx % len(color_list)]
        draw.rectangle(bbox, outline=color, width=2)
        draw.text((x1 + 4, y1 + 2), label, fill=color, font=font)
    return image_pil

detection(
        "./media/pometaro_01.png", # 検出対象の画像
        category_list=["toy"], # 検出したいラベル
    )

detection(
    "./media/pometaro_02.png", # 検出対象の画像
    category_list=["tempura shrimp","bear", "Bed"], # 検出したいラベル
)

detection(
    "./media/pometaro_03.png", # 検出対象の画像
    category_list=["brown dog","white dog"], # 検出したいラベル
)

detection(
    "./media/dogfood.jpg",  # 検出対象の画像
    category_list=["The darker brown feed"],  # 検出したいラベル
)