autotrain/scripts/vision_fallback.py

#!/usr/bin/env python3
"""视觉 fallback: 截屏 + 调豆包视觉 API,做"当前是哪一页 / 找哪个按钮"判断."""

from __future__ import annotations

import argparse
import base64
import json
import subprocess
import sys
import tempfile
import time
import urllib.error
import urllib.request
from pathlib import Path

import device_runtime

PROJECT_ROOT = Path(__file__).resolve().parent.parent
CONFIG_LOCAL = PROJECT_ROOT / "config.local.json"
WORK_DIR = PROJECT_ROOT / "work"


def load_vision_config() -> dict:
    if not CONFIG_LOCAL.exists():
        raise RuntimeError(f"配置文件不存在: {CONFIG_LOCAL}")
    cfg = json.loads(CONFIG_LOCAL.read_text(encoding="utf-8"))
    vision = cfg.get("vision")
    if not vision:
        raise RuntimeError("config.local.json 缺 vision 段")
    for key in ("base_url", "api_key", "model"):
        if not vision.get(key):
            raise RuntimeError(f"vision.{key} 未配置")
    return vision


def adb_screenshot(output: Path) -> Path:
    output.parent.mkdir(parents=True, exist_ok=True)
    result = device_runtime.run_adb_bytes(["exec-out", "screencap", "-p"], check=False)
    if result.returncode != 0:
        raise RuntimeError(f"adb screencap 失败: {result.stderr.decode('utf-8', 'ignore')}")
    output.write_bytes(result.stdout)
    return output


def call_vision(image_path: Path, prompt: str, vision_cfg: dict, timeout: int = 60) -> dict:
    image_b64 = base64.b64encode(image_path.read_bytes()).decode("ascii")
    payload = {
        "model": vision_cfg["model"],
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_b64}"}},
                    {"type": "text", "text": prompt},
                ],
            }
        ],
        "temperature": 0.1,
    }
    headers = {
        "Authorization": f"Bearer {vision_cfg['api_key']}",
        "Content-Type": "application/json",
    }
    url = vision_cfg["base_url"].rstrip("/") + "/chat/completions"
    start = time.time()
    request = urllib.request.Request(
        url,
        data=json.dumps(payload).encode("utf-8"),
        headers=headers,
        method="POST",
    )
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            body = response.read().decode("utf-8", "ignore")
            status = response.status
    except urllib.error.HTTPError as exc:
        body = exc.read().decode("utf-8", "ignore")
        raise RuntimeError(f"vision_api status={exc.code} body={body[:500]}") from exc
    elapsed = time.time() - start
    if status >= 300:
        raise RuntimeError(f"vision_api status={status} body={body[:500]}")
    data = json.loads(body)
    content = data.get("choices", [{}])[0].get("message", {}).get("content", "")
    usage = data.get("usage", {})
    return {
        "content": content,
        "usage": usage,
        "elapsed_seconds": round(elapsed, 2),
        "raw": data,
    }


PAGE_DETECT_PROMPT = """这是手机 12306 App 的一张截图。请只用一个 JSON 对象回答(不要 markdown 代码块),字段如下:

{
  "page": "<one of: list / verify / invoice_info / invoice_confirm / success / email_confirm / mail_sent / no_ticket / launcher / other>",
  "summary": "<简短描述当前看到的内容,不超过 30 字>",
  "actionable_buttons": ["<可点击的关键按钮文字,如 开具/确认/提交/继续开票/确定/发送至邮箱>"],
  "scroll_position": "<top / middle / bottom / unknown>"
}

页面对应关系:
- list = 扫码开票单列表(显示多个 X*Y 脱敏姓名 + 开具按钮)
- verify = 请输入乘车人 + 证件号码后8位的核验弹窗
- invoice_info = 发票信息页(有抬头/税号/金额/提交按钮)
- invoice_confirm = 发票信息确认弹窗
- success = 发票开具成功页(有 继续开票 / 发送至邮箱)
- email_confirm = 邮箱地址确认
- mail_sent = 邮件发送完成提示
- no_ticket = 发票管理页 + 您没有可开具电子发票的客票
- launcher = 桌面 / 不在 12306 内
- other = 其他识别不出的页面
"""


LIST_ROWS_PROMPT_TEMPLATE = """这是手机 12306 App 扫码开票单列表页的一张截图。
请自上而下识别屏幕上所有"乘客行"(每行有一个脱敏姓名如 张*岩、王*平,行尾右侧有可点击按钮,按钮文案可能是"开具"、"已开票"或"查看")。
**坐标系约定**:所有 bbox 用 0-1000 归一化坐标返回(x 用 [0,1000] 横向占比,y 用 [0,1000] 纵向占比),不要返回原始像素。

只用一个 JSON 对象回答(不要 markdown 代码块、不要解释),字段如下:
{{
  "rows": [
    {{
      "masked_name": "<脱敏姓名,如 张*岩;只用一个 ASCII * 号,不要 ＊>",
      "id_first": "<身份证首位,看不见就空串>",
      "id_last": "<身份证末位,看不见就空串>",
      "button_bbox": [x1, y1, x2, y2]  // 该行可点击按钮(开具/已开票/查看)的 0-1000 归一化 bbox,整数
    }}
  ],
  "scroll_position": "<top / middle / bottom>",  // bottom = 看到列表底部边缘或"没有更多了"提示
  "page_ok": <true/false>  // 当前是不是扫码开票单列表页;不是请置 false 并把 rows 留空
}}

注意:
- 只识别完整可见的行;只露出半截的行不要加进 rows
- masked_name 用半角 * 号(不要全角 ＊)
- 没看到"开具"、"已开票"或"查看"这类行尾可点击按钮的行不算"乘客行",跳过
- bbox 数值范围 [0, 1000],整数
"""


def screenshot_size(image_path: Path) -> tuple[int, int]:
    data = image_path.read_bytes()
    if len(data) >= 24 and data[:8] == b"\x89PNG\r\n\x1a\n":
        width = int.from_bytes(data[16:20], "big")
        height = int.from_bytes(data[20:24], "big")
        return width, height
    return 0, 0


def parse_json_lenient(content: str) -> dict:
    text = content.strip()
    if text.startswith("```"):
        text = text.strip("`")
        if text.startswith("json"):
            text = text[4:]
    start = text.find("{")
    end = text.rfind("}")
    if start == -1 or end == -1:
        raise RuntimeError(f"无法从模型输出里抽出 JSON: {content[:300]}")
    return json.loads(text[start:end + 1])


def denormalize_bboxes(parsed: dict, width: int, height: int) -> dict:
    """把模型返回的 0-1000 归一化 bbox 换算成截图绝对像素。

    判定规则:
      - max(bbox) 明显小于 min(image_w, image_h) 的一半,几乎不可能是绝对像素 -> 归一化
      - max(bbox) <= 1100(允许 Qwen-VL rounding 飘到 1000-1030)且图比 1100 大 -> 归一化
      - 否则按绝对像素原样保留
    **就地**修改 parsed 并返回。
    """
    rows = parsed.get("rows") or []
    if not rows or width <= 0 or height <= 0:
        return parsed
    flat = []
    for row in rows:
        bbox = row.get("button_bbox") or []
        if len(bbox) == 4:
            flat.extend(bbox)
    if not flat:
        return parsed
    max_v = max(flat)
    img_min = min(width, height)
    is_normalized = max_v <= 1100 and img_min > 1100
    if not is_normalized:
        return parsed
    for row in rows:
        bbox = row.get("button_bbox") or []
        if len(bbox) != 4:
            continue
        x1, y1, x2, y2 = bbox
        row["button_bbox"] = [
            int(round(x1 * width / 1000)),
            int(round(y1 * height / 1000)),
            int(round(x2 * width / 1000)),
            int(round(y2 * height / 1000)),
        ]
        cx = (row["button_bbox"][0] + row["button_bbox"][2]) // 2
        cy = (row["button_bbox"][1] + row["button_bbox"][3]) // 2
        row["tap_point"] = [cx, cy]
    parsed["bbox_denormalized"] = True
    return parsed


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="vision fallback 测试入口")
    sub = parser.add_subparsers(dest="cmd", required=True)

    p_ping = sub.add_parser("ping", help="不截屏,只发一张本地图测试 API 通路")
    p_ping.add_argument("--image", required=True, help="本地图片路径")
    p_ping.add_argument("--prompt", default="一句话描述这张图")

    p_detect = sub.add_parser("detect-page", help="adb 截屏 + 让豆包判断当前是哪一页")
    p_detect.add_argument("--save", default="work/last_screenshot.png", help="截屏保存路径")

    p_ask = sub.add_parser("ask", help="adb 截屏 + 自定义 prompt")
    p_ask.add_argument("--prompt", required=True)
    p_ask.add_argument("--save", default="work/last_screenshot.png")

    p_rows = sub.add_parser("list-rows", help="adb 截屏 + 让豆包识别列表页所有可见乘客行")
    p_rows.add_argument("--save", default="work/last_screenshot.png", help="截屏保存路径")
    p_rows.add_argument("--image", default=None, help="指定已有图片,跳过 adb 截屏(调试用)")

    args = parser.parse_args(argv)
    vision_cfg = load_vision_config()

    if args.cmd == "ping":
        image = Path(args.image).expanduser()
        if not image.is_absolute():
            image = PROJECT_ROOT / image
        result = call_vision(image, args.prompt, vision_cfg)
    elif args.cmd == "list-rows":
        if args.image:
            save_path = Path(args.image)
            if not save_path.is_absolute():
                save_path = PROJECT_ROOT / save_path
        else:
            save_path = Path(args.save)
            if not save_path.is_absolute():
                save_path = PROJECT_ROOT / save_path
            adb_screenshot(save_path)
        width, height = screenshot_size(save_path)
        prompt = LIST_ROWS_PROMPT_TEMPLATE.format(width=width, height=height)
        result = call_vision(save_path, prompt, vision_cfg)
        result["screenshot"] = str(save_path)
        result["image_size"] = {"width": width, "height": height}
        try:
            parsed = parse_json_lenient(result["content"])
            denormalize_bboxes(parsed, width, height)
            result["parsed"] = parsed
        except Exception as exc:
            result["parse_error"] = str(exc)
    else:
        save_path = Path(args.save)
        if not save_path.is_absolute():
            save_path = PROJECT_ROOT / save_path
        adb_screenshot(save_path)
        prompt = PAGE_DETECT_PROMPT if args.cmd == "detect-page" else args.prompt
        result = call_vision(save_path, prompt, vision_cfg)
        result["screenshot"] = str(save_path)

    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    sys.exit(main())