视觉定位 Grounding

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21import os
import base64
import cv2
from arkruntime import Ark

# 配置参数
DEFAULT_MODEL = "doubao-seed-2-1-pro-260915"  # Replace with Model ID .
IMAGE_PATH = "./ark_demo_img.png"  # 替换为您本地测试图片的路径
PROMPT = "框出中间狼卡通形象的头部的位置，输出 bounding box 的坐标"
BBOX_TAG_START = "<bbox>"
BBOX_TAG_END = "</bbox>"

# 读取API密钥
api_key = os.getenv("ARK_API_KEY")

# 创建Ark客户端
client = Ark(
    api_key=api_key,
    # The base URL for model invocation
    base_url="https://ark.cn-beijing.volces.com/api/v3",
)

1
2with open(IMAGE_PATH, "rb") as f:
    base64_image = base64.b64encode(f.read()).decode('utf-8')

1
2
3
4
5
6
7
8
9
10
11
12
13
14response = client.chat.completions.create(
    model=DEFAULT_MODEL,
    messages=[{
        "role": "user",
        "content": [{
            "type": "image_url",  # 图片输入
            "image_url": {"url": f"data:image/png;base64,{base64_image}"}
        }, {
            "type": "text",  # 文本提示
            "text": PROMPT
        }]
    }]
)
bbox_content = response.choices[0].message.content

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15# 检查结果格式是否正确
if not (bbox_content.startswith(BBOX_TAG_START) and bbox_content.endswith(BBOX_TAG_END)):
    print("错误：边界框格式不正确，缺少标签包裹")
    exit(1)

# 解析坐标值
try:
    coords_str = bbox_content[len(BBOX_TAG_START):-len(BBOX_TAG_END)]
    coords = list(map(int, coords_str.split()))
    if len(coords) != 4:  # 验证坐标数量(xmin, ymin, xmax, ymax)
        raise ValueError("坐标数量不正确，需要4个数值")
    x_min, y_min, x_max, y_max = coords
except ValueError as e:
    print(f"坐标解析失败: {str(e)}")
    exit(1)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17# 读取原图
image = cv2.imread(IMAGE_PATH)

# 获取图像尺寸并缩放坐标
h, w = image.shape[:2]
x_min_real = int(x_min * w / 1000)
y_min_real = int(y_min * h / 1000)
x_max_real = int(x_max * w / 1000)
y_max_real = int(y_max * h / 1000)

# 绘制红色边界框
cv2.rectangle(image, (x_min_real, y_min_real), (x_max_real, y_max_real), (0, 0, 255), 3)

# 保存结果图片
output_path = os.path.splitext(IMAGE_PATH)[0] + "_with_bboxes.png"
cv2.imwrite(output_path, image)
print(f"成功保存标注图片: {output_path}")