perception_fixed.py

"""GhostBridge 感知层示例实现（SoM - Set-of-Mark）。

本模块提供两个核心函数：
1) add_tags_to_image(image_path): 对输入截图自动加数字标签。
2) identify_element(tagged_image, user_instruction): 构建提示词并调用 Qwen API。

依赖：Pillow (PIL), requests
"""

from __future__ import annotations
import requests
import json
import os
import base64
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Tuple

from PIL import Image, ImageDraw, ImageFilter, ImageFont


@dataclass
class TagInfo:
    """用于记录单个标签的几何信息。"""

    id: int
    bbox: Tuple[int, int, int, int]  # (left, top, right, bottom)
    center: Tuple[int, int]  # (x, y)


def add_tags_to_image(image_path: str, grid_cols: int = 6, grid_rows: int = 4, top_k: int = 10) -> str:
    """为图片添加显眼的数字标签，并写出标签元数据 JSON。

    思路（新手友好版）：
    1) 读取图片并计算"边缘强度"（边缘越多，信息越丰富）。
    2) 把图片划分成网格，统计每个格子的平均边缘强度。
    3) 选取得分最高的若干格子，在中心位置绘制数字标签。

    参数：
        image_path: 原始截图路径。
        grid_cols, grid_rows: 网格划分数量（默认 6x4）。
        top_k: 选取前 K 个区域打标签。

    返回：
        tagged_image_path: 生成的带标签图片路径。

    备注：
        会额外生成一个 JSON 文件，记录每个标签的坐标信息，
        以便 identify_element() 使用。
    """

    image_path = str(image_path)
    image = Image.open(image_path).convert("RGB")
    width, height = image.size

    # 1) 生成边缘图（灰度 + 边缘检测）
    gray = image.convert("L")
    edges = gray.filter(ImageFilter.FIND_EDGES)

    # 2) 计算网格得分（平均边缘强度）
    cell_w = max(1, width // grid_cols)
    cell_h = max(1, height // grid_rows)

    scores = []
    for r in range(grid_rows):
        for c in range(grid_cols):
            left = c * cell_w
            top = r * cell_h
            right = min(width, left + cell_w)
            bottom = min(height, top + cell_h)

            # 裁剪该格子区域并计算平均亮度（边缘强度）
            patch = edges.crop((left, top, right, bottom))
            # 快速计算平均像素值
            pixels = list(patch.getdata())
            avg = sum(pixels) / max(1, len(pixels))
            scores.append(((left, top, right, bottom), avg))

    # 3) 选取得分最高的 top_k 个网格区域
    scores.sort(key=lambda x: x[1], reverse=True)
    selected = scores[: max(1, min(top_k, len(scores)))]

    # 4) 准备绘制标签
    draw = ImageDraw.Draw(image)
    try:
        font = ImageFont.truetype("arial.ttf", 20)
    except OSError:
        # 如果系统没有 Arial 字体，则使用默认字体
        font = ImageFont.load_default()

    tag_infos: List[TagInfo] = []

    for idx, (bbox, _) in enumerate(selected, start=1):
        left, top, right, bottom = bbox
        center_x = (left + right) // 2
        center_y = (top + bottom) // 2

        # 绘制标签背景（半透明黄色方块）
        text = str(idx)
        text_bbox = draw.textbbox((0, 0), text, font=font)
        text_w = text_bbox[2] - text_bbox[0]
        text_h = text_bbox[3] - text_bbox[1]
        pad = 4
        box_left = center_x - text_w // 2 - pad
        box_top = center_y - text_h // 2 - pad
        box_right = center_x + text_w // 2 + pad
        box_bottom = center_y + text_h // 2 + pad

        draw.rectangle([box_left, box_top, box_right, box_bottom], fill=(255, 230, 0))
        draw.text((center_x - text_w // 2, center_y - text_h // 2), text, fill=(0, 0, 0), font=font)

        tag_infos.append(TagInfo(id=idx, bbox=bbox, center=(center_x, center_y)))

    # 5) 保存带标签图片 + 标签元数据
    input_path = Path(image_path)
    tagged_image_path = input_path.with_name(input_path.stem + "_tagged.png")
    image.save(tagged_image_path)

    meta_path = input_path.with_name(input_path.stem + "_tags.json")
    meta_payload = {
        "source_image": str(input_path.name),
        "tagged_image": str(tagged_image_path.name),
        "tags": [
            {"id": t.id, "bbox": list(t.bbox), "center": list(t.center)} for t in tag_infos
        ],
    }
    meta_path.write_text(json.dumps(meta_payload, ensure_ascii=False, indent=2), encoding="utf-8")

    return str(tagged_image_path)


def identify_element(tagged_image: str, user_instruction: str) -> Dict[str, object]:
    """构建提示词并调用 Qwen API，返回标签 ID 与坐标。

    参数：
        tagged_image: 带标签的图片路径（add_tags_to_image 的输出）。
        user_instruction: 用户自然语言指令（例如"提交按钮"）。

    返回：
        dict, 示例：
        {
            "tag_id": 3,
            "center": [512, 384],
            "prompt": "...",
            "raw_model_response": "..."
        }
    """

    # 1) 生成 Prompt
    prompt = (
        f"用户想要点击 {user_instruction}。"
        "请看图，告诉我对应的数字标签 ID 是多少？"
    )

    # 2) 调用 Qwen API
    raw_response = _call_qwen_api(prompt=prompt, image_path=tagged_image)

    # 3) 解析返回的 tag_id
    tag_id = int(raw_response.get("tag_id", 1))

    # 4) 读取标签元数据，以便得到坐标
    meta_path = _guess_meta_path(tagged_image)
    center = _lookup_center_by_id(meta_path, tag_id)

    return {
        "tag_id": tag_id,
        "center": center,
        "prompt": prompt,
        "raw_model_response": raw_response,
    }


def _call_qwen_api(prompt: str, image_path: str) -> Dict[str, object]:
    """调用阿里云 Qwen2.5-VL API 进行图像识别。

    Args:
        prompt: 用户提示词
        image_path: 图像文件路径

    Returns:
        dict: 包含 tag_id 和其他信息的字典
    """

    import re

    # 1) 读取 API Key 和模型配置
    api_key = os.getenv("VLM_API_KEY", "")
    model = os.getenv("VLM_MODEL", "qwen3-vl-plus")

    if not api_key:
        # 没有配置 API Key 时，使用默认返回
        return {"tag_id": 1, "note": "Mock response (no API key configured)."}

    # 2) 读取并编码图像
    try:
        with open(image_path, "rb") as f:
            image_base64 = base64.b64encode(f.read()).decode("utf-8")
    except Exception as e:
        return {"tag_id": 1, "error": f"Failed to read image: {str(e)}"}

    # 3) 构建请求
    url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }

    payload = {
        "model": model,
        "input": {
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "image": f"data:image/png;base64,{image_base64}"
                        },
                        {
                            "text": prompt
                        }
                    ]
                }
            ]
        },
        "parameters": {
            "max_tokens": 2048,
            "temperature": 0.1
        }
    }

    # 4) 发送请求
    try:
        response = requests.post(url, headers=headers, json=payload, timeout=30)
        response.raise_for_status()
        result = response.json()

        # 5) 解析响应
        # 检查是否有 output 字段（成功响应）
        if "output" in result:
            content = result["output"]["choices"][0]["message"]["content"][0]["text"]
            # 尝试从响应中提取数字标签
            match = re.search(r"\d+", content)
            if match:
                tag_id = int(match.group())
            else:
                # 如果没有找到数字，返回默认值
                tag_id = 1

            return {
                "tag_id": tag_id,
                "content": content,
                "model": model,
                "note": "Real API response"
            }
        else:
            # 检查是否有错误信息
            if "code" in result:
                error_msg = result.get("message", "Unknown error")
                return {"tag_id": 1, "error": f"API error: {error_msg}"}
            else:
                return {"tag_id": 1, "error": "Unknown error"}

    except Exception as e:
        return {"tag_id": 1, "error": f"Request failed: {str(e)}"}


def _guess_meta_path(tagged_image_path: str) -> Path:
    """根据 tagged_image 的文件名推断对应的 JSON 元数据路径。"""

    p = Path(tagged_image_path)
    # 约定：xxx_tagged.png 对应 xxx_tags.json
    # 如果没有 _tagged，则尝试直接换后缀
    if p.stem.endswith("_tagged"):
        base = p.stem.replace("_tagged", "")
    else:
        base = p.stem
    return p.with_name(base + "_tags.json")


def _lookup_center_by_id(meta_path: Path, tag_id: int) -> List[int]:
    """从元数据 JSON 中查找标签 ID 的中心坐标。"""

    if not meta_path.exists():
        return []

    payload = json.loads(meta_path.read_text(encoding="utf-8"))
    for tag in payload.get("tags", []):
        if int(tag.get("id", -1)) == tag_id:
            return list(tag.get("center", []))
    return []