生成 AI で物体検出 ! Amazon Nova でバウンディングボックスを描画しちゃおう

Detect bounding box of objects in the image, only detect {item_name} category objects with high confidence, output in a list of bounding box format.

 

Output example:

[

    {"{item_name}": [x1, y1, x2, y2]},

    ...

]

 

Result:

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/dogfood.jpg",  # 検出対象の画像
    category_list=["The darker brown feed"],  # 検出したいラベル
)