416 lines
16 KiB
Python
416 lines
16 KiB
Python
import subprocess
|
||
import os
|
||
import xml.etree.ElementTree as ET
|
||
import re
|
||
import time
|
||
import sys
|
||
import json
|
||
|
||
try:
|
||
from PIL import Image, ImageDraw, ImageFont
|
||
except ImportError:
|
||
print("❌ 缺少依赖库，请先运行: pip install pillow")
|
||
sys.exit(1)
|
||
|
||
# ==============================================================================
|
||
# 1. 全局配置
|
||
# ==============================================================================
|
||
class Config:
|
||
ADB_PATH = "adb"
|
||
|
||
# 输出文件路径 (Agent 的输入来源)
|
||
LOCAL_XML_PATH = "./tree/ui_dump.xml"
|
||
JSON_PATH = "./tree/ui_elements.json" # 最终的结构化数据
|
||
DEBUG_JSON_PATH = "./tree/ui_elements_debug.json" # 调试用的完整节点数据
|
||
RAW_SCREEN_PATH = "./tree/screen_raw.png" # 原始截图
|
||
LABELED_SCREEN_PATH = "./tree/screen_labeled.png" # 标注后的截图 (SoM)
|
||
|
||
# 逻辑开关
|
||
ENABLE_CHILD_TEXT_MERGE = True # 合并子元素文本
|
||
EXPORT_ALL_VISIBLE_NODES = True # 导出所有可见节点，而不只 clickable + 纯文本
|
||
ONLY_VISIBLE_TO_USER = True # 仅导出 visible-to-user=true 的节点
|
||
SHRINK_STEP_PX = 5 # 视觉内缩像素 (防遮挡)
|
||
|
||
# ==============================================================================
|
||
# 2. 几何与数学工具
|
||
# ==============================================================================
|
||
def parse_bounds_rect(bounds_str):
|
||
"""解析 bounds 字符串为列表 [x1, y1, x2, y2]"""
|
||
if not bounds_str: return None
|
||
try:
|
||
coords = list(map(int, re.findall(r'\d+', bounds_str)))
|
||
if len(coords) == 4: return coords
|
||
except: pass
|
||
return None
|
||
|
||
def get_center(rect):
|
||
"""计算矩形几何中心"""
|
||
return (rect[0] + rect[2]) // 2, (rect[1] + rect[3]) // 2
|
||
|
||
def get_area(rect):
|
||
return (rect[2] - rect[0]) * (rect[3] - rect[1])
|
||
|
||
def resource_id_to_label(resource_id):
|
||
"""把 resource-id 转成更易读的调试标签"""
|
||
if not resource_id:
|
||
return ""
|
||
tail = resource_id.split("/")[-1]
|
||
tail = tail.replace("_", " ").replace("-", " ").strip()
|
||
return tail
|
||
|
||
def is_inside(inner_rect, outer_rect):
|
||
"""判断 inner 是否完全在 outer 内部"""
|
||
return (inner_rect[0] >= outer_rect[0] and
|
||
inner_rect[1] >= outer_rect[1] and
|
||
inner_rect[2] <= outer_rect[2] and
|
||
inner_rect[3] <= outer_rect[3])
|
||
|
||
def is_point_in_rect(point, rect):
|
||
"""判断点是否在矩形内"""
|
||
x, y = point
|
||
return rect[0] <= x <= rect[2] and rect[1] <= y <= rect[3]
|
||
|
||
def clip_rect(rect, screen_w, screen_h):
|
||
"""屏幕边缘裁剪 (处理只露出一半的控件)"""
|
||
x1, y1, x2, y2 = rect
|
||
new_x1 = max(0, x1)
|
||
new_y1 = max(0, y1)
|
||
new_x2 = min(screen_w, x2)
|
||
new_y2 = min(screen_h, y2)
|
||
if new_x2 <= new_x1 or new_y2 <= new_y1:
|
||
return None
|
||
return [new_x1, new_y1, new_x2, new_y2]
|
||
|
||
def find_safe_point(parent_rect, children_rects):
|
||
"""
|
||
智能避让算法：在父控件中寻找一个不与子控件重叠的坐标
|
||
"""
|
||
p_x1, p_y1, p_x2, p_y2 = parent_rect
|
||
w, h = p_x2 - p_x1, p_y2 - p_y1
|
||
|
||
# 搜索网格: 中心 -> 十字线 -> 四角
|
||
search_grid = [
|
||
(0.5, 0.5), (0.5, 0.25), (0.5, 0.75),
|
||
(0.25, 0.5), (0.75, 0.5),
|
||
(0.2, 0.2), (0.8, 0.2), (0.2, 0.8), (0.8, 0.8)
|
||
]
|
||
|
||
for rx, ry in search_grid:
|
||
cx = int(p_x1 + w * rx)
|
||
cy = int(p_y1 + h * ry)
|
||
|
||
# 检查该点是否撞到了子控件
|
||
hit = False
|
||
for child in children_rects:
|
||
if is_point_in_rect((cx, cy), child):
|
||
hit = True
|
||
break
|
||
if not hit: return (cx, cy) # 找到安全点
|
||
|
||
return get_center(parent_rect) # 兜底
|
||
|
||
# ==============================================================================
|
||
# 3. ADB 基础操作
|
||
# ==============================================================================
|
||
def run_cmd(cmd):
|
||
try:
|
||
if Config.ADB_PATH != "adb" and cmd.startswith("adb"):
|
||
cmd = cmd.replace("adb", f'"{Config.ADB_PATH}"', 1)
|
||
res = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, encoding="utf-8")
|
||
return res.stdout.strip()
|
||
except Exception: return None
|
||
|
||
def get_devices():
|
||
out = run_cmd("adb devices")
|
||
if not out: return []
|
||
devs = []
|
||
for line in out.split('\n'):
|
||
if "\t" in line and "List of" not in line:
|
||
parts = line.split("\t")
|
||
devs.append({"serial": parts[0].strip(), "status": parts[1].strip()})
|
||
return devs
|
||
|
||
def get_screen_size(serial):
|
||
prefix = f"adb -s {serial}" if serial else "adb"
|
||
out = run_cmd(f"{prefix} shell wm size")
|
||
if out:
|
||
m = re.search(r'(\d+)x(\d+)', out)
|
||
if m: return int(m.group(1)), int(m.group(2))
|
||
return 1080, 2400
|
||
|
||
def select_device():
|
||
devs = [d for d in get_devices() if d['status'] == 'device']
|
||
if not devs: return None
|
||
if len(devs) == 1: return devs[0]['serial']
|
||
# 只有多台设备时才交互，否则自动选
|
||
print("多台设备在线，请选择:")
|
||
for i, d in enumerate(devs): print(f"{i+1}. {d['serial']}")
|
||
try:
|
||
idx = int(input("序号: ")) - 1
|
||
return devs[idx]['serial'] if 0 <= idx < len(devs) else None
|
||
except: return None
|
||
|
||
# ==============================================================================
|
||
# 4. 数据解析核心 (Parsing Layer)
|
||
# ==============================================================================
|
||
def parse_ui_tree(device_serial):
|
||
"""解析 XML，执行合并、裁剪、避让，返回清洗后的列表"""
|
||
prefix = f"adb -s {device_serial}"
|
||
screen_w, screen_h = get_screen_size(device_serial)
|
||
|
||
# 0. 清除旧缓存
|
||
run_cmd(f"{prefix} shell rm -f /sdcard/window_dump.xml")
|
||
|
||
# 1. Dump XML
|
||
run_cmd(f"{prefix} shell uiautomator dump /sdcard/window_dump.xml")
|
||
run_cmd(f"{prefix} pull /sdcard/window_dump.xml {Config.LOCAL_XML_PATH}")
|
||
|
||
if not os.path.exists(Config.LOCAL_XML_PATH): return []
|
||
|
||
tree = ET.parse(Config.LOCAL_XML_PATH)
|
||
root = tree.getroot()
|
||
|
||
all_clickables = []
|
||
all_non_clickables = []
|
||
non_clickable_texts = []
|
||
|
||
# 2. 初步提取与屏幕裁剪
|
||
for node in root.iter():
|
||
text = node.attrib.get("text", "")
|
||
desc = node.attrib.get("content-desc", "")
|
||
content = text if text else desc
|
||
resource_id = node.attrib.get("resource-id", "")
|
||
visible_to_user = node.attrib.get("visible-to-user", "true") != "false"
|
||
enabled = node.attrib.get("enabled", "true") == "true"
|
||
if Config.ONLY_VISIBLE_TO_USER and not visible_to_user:
|
||
continue
|
||
|
||
raw_rect = parse_bounds_rect(node.attrib.get("bounds", ""))
|
||
if not raw_rect: continue
|
||
|
||
# 裁剪：确保只处理屏幕内的元素
|
||
visible_rect = clip_rect(raw_rect, screen_w, screen_h)
|
||
if not visible_rect: continue
|
||
|
||
item = {
|
||
"id": id(node),
|
||
"type": node.attrib.get("class", ""),
|
||
"content": content,
|
||
"text": text,
|
||
"desc": desc,
|
||
"resource_id": resource_id,
|
||
"resource_label": resource_id_to_label(resource_id),
|
||
"package": node.attrib.get("package", ""),
|
||
"enabled": enabled,
|
||
"visible_to_user": visible_to_user,
|
||
"rect": visible_rect,
|
||
"area": get_area(visible_rect),
|
||
"center": get_center(visible_rect),
|
||
"temp_children_text": []
|
||
}
|
||
|
||
if node.attrib.get("clickable") == "true":
|
||
all_clickables.append(item)
|
||
else:
|
||
all_non_clickables.append(item)
|
||
if content:
|
||
non_clickable_texts.append(item)
|
||
|
||
# 3. 文本合并 (子元素文本归属给父按钮)
|
||
if Config.ENABLE_CHILD_TEXT_MERGE:
|
||
for txt in non_clickable_texts:
|
||
# 找到包含该文本的所有空白父按钮
|
||
parents = [btn for btn in all_clickables if not btn["content"] and is_inside(txt["rect"], btn["rect"])]
|
||
if parents:
|
||
# 归属给面积最小(最内层)的父按钮
|
||
best = sorted(parents, key=lambda x: x["area"])[0]
|
||
best["temp_children_text"].append(txt["content"])
|
||
|
||
# 应用合并
|
||
for btn in all_clickables:
|
||
if not btn["content"] and btn["temp_children_text"]:
|
||
btn["content"] = " ".join(btn["temp_children_text"])
|
||
|
||
# 4. 智能避让 (点击穿透处理)
|
||
for parent in all_clickables:
|
||
# 找到当前 parent 内部的所有子点击区域
|
||
children = [c['rect'] for c in all_clickables if c['id'] != parent['id'] and is_inside(c['rect'], parent['rect'])]
|
||
if children:
|
||
# 检查当前中心点是否安全，不安全则寻找新点
|
||
if any(is_point_in_rect(parent['center'], c) for c in children):
|
||
parent['center'] = find_safe_point(parent['rect'], children)
|
||
|
||
# 5. 选择需要导出的非 clickable 节点
|
||
# 全量模式：导出所有可见非 clickable 节点
|
||
# 兼容模式：仅导出不在 clickable 内部的纯文本节点
|
||
if Config.EXPORT_ALL_VISIBLE_NODES:
|
||
export_non_clickables = all_non_clickables
|
||
else:
|
||
export_non_clickables = []
|
||
for txt in non_clickable_texts:
|
||
inside_clickable = any(is_inside(txt["rect"], btn["rect"]) for btn in all_clickables)
|
||
if not inside_clickable:
|
||
export_non_clickables.append(txt)
|
||
|
||
# 6. 生成最终数据结构 (格式化为 Agent 需要的 JSON)
|
||
final_data = []
|
||
export_items = []
|
||
for item in all_clickables:
|
||
export_items.append({
|
||
**item,
|
||
"clickable": True,
|
||
})
|
||
for item in export_non_clickables:
|
||
export_items.append({
|
||
**item,
|
||
"clickable": False,
|
||
})
|
||
|
||
for i, item in enumerate(export_items):
|
||
final_data.append({
|
||
"id": i, # 映射给 Agent 的短 ID
|
||
"context": item["content"], # 文本内容
|
||
"center": item["center"], # 安全点击坐标 (x, y)
|
||
"type": item["type"], # 控件类型 (辅助信息)
|
||
"clickable": item["clickable"], # 是否可点击
|
||
"bbox": item["rect"], # 边界框 (辅助信息)
|
||
"text": item["text"], # 原始 text
|
||
"desc": item["desc"], # 原始 content-desc
|
||
"resource_id": item["resource_id"],
|
||
"resource_label": item["resource_label"],
|
||
"package": item["package"],
|
||
"enabled": item["enabled"],
|
||
"visible_to_user": item["visible_to_user"],
|
||
|
||
# 以下字段仅用于内部画图计算层级，保存后可忽略
|
||
"_area": item["area"],
|
||
"_raw_id": item["id"]
|
||
})
|
||
|
||
return final_data
|
||
|
||
# ==============================================================================
|
||
# 5. 视觉处理核心 (Visual Layer)
|
||
# ==============================================================================
|
||
def draw_labels(device_serial, ui_elements):
|
||
"""截图并绘制 SoM 标记 (仅绘制有文本的节点，带层级内缩)"""
|
||
prefix = f"adb -s {device_serial}" if device_serial else "adb"
|
||
run_cmd(f"{prefix} shell screencap -p /sdcard/screen.png")
|
||
run_cmd(f"{prefix} pull /sdcard/screen.png {Config.RAW_SCREEN_PATH}")
|
||
|
||
if not os.path.exists(Config.RAW_SCREEN_PATH): return None
|
||
|
||
image = Image.open(Config.RAW_SCREEN_PATH).convert("RGB")
|
||
draw = ImageDraw.Draw(image)
|
||
try: font = ImageFont.truetype("arial.ttf", 24)
|
||
except: font = ImageFont.load_default()
|
||
|
||
drawable_items = [item for item in ui_elements if str(item.get("context", "")).strip()]
|
||
if not drawable_items:
|
||
image.save(Config.LABELED_SCREEN_PATH)
|
||
return Config.LABELED_SCREEN_PATH
|
||
|
||
# 计算每个元素的层级 (用于视觉内缩)
|
||
# 复杂度 O(N^2)，但在 UI 元素数量级下很快
|
||
for item in drawable_items:
|
||
level = 0
|
||
for other in drawable_items:
|
||
if item['id'] == other['id']: continue
|
||
# 如果 item 在 other 内部，层级+1
|
||
if is_inside(item['bbox'], other['bbox']):
|
||
# 处理相同大小的重叠情况
|
||
if other['_area'] > item['_area']: level += 1
|
||
elif other['_area'] == item['_area'] and item['id'] > other['id']: level += 1
|
||
item['_shrink_level'] = level
|
||
|
||
# 绘制
|
||
for item in drawable_items:
|
||
rect = item['bbox']
|
||
idx = item['id']
|
||
level = item.get('_shrink_level', 0)
|
||
|
||
# 计算内缩
|
||
gap = min(level * Config.SHRINK_STEP_PX, min(rect[2]-rect[0], rect[3]-rect[1])//2 - 2)
|
||
gap = max(0, gap)
|
||
|
||
vis_rect = [rect[0]+gap, rect[1]+gap, rect[2]-gap, rect[3]-gap]
|
||
|
||
# 画框
|
||
draw.rectangle(vis_rect, outline="red", width=3)
|
||
|
||
# 画标签
|
||
tag = str(idx)
|
||
bbox = font.getbbox(tag)
|
||
w, h = bbox[2]-bbox[0], bbox[3]-bbox[1]
|
||
|
||
tx, ty = vis_rect[0], vis_rect[1] - h - 4
|
||
if ty < 0: ty = vis_rect[1] # 顶端边界处理
|
||
|
||
draw.rectangle([tx, ty, tx+w+8, ty+h+4], fill="red", outline="red")
|
||
draw.text((tx+4, ty), tag, fill="white", font=font)
|
||
|
||
image.save(Config.LABELED_SCREEN_PATH)
|
||
return Config.LABELED_SCREEN_PATH
|
||
|
||
# ==============================================================================
|
||
# 主程序入口
|
||
# ==============================================================================
|
||
if __name__ == "__main__":
|
||
device = select_device()
|
||
if not device:
|
||
print("❌ 未连接设备")
|
||
sys.exit(1)
|
||
|
||
os.makedirs("./tree", exist_ok=True)
|
||
|
||
print(f"🚀 开始处理设备: {device}")
|
||
start_time = time.time()
|
||
|
||
# 1. 获取并清洗数据
|
||
ui_data = parse_ui_tree(device)
|
||
|
||
if ui_data:
|
||
# 2. 生成视觉反馈图
|
||
draw_labels(device, ui_data)
|
||
|
||
# 3. 清理不需要的内部字段，只保留对外导出字段
|
||
final_json_data = []
|
||
debug_json_data = []
|
||
for item in ui_data:
|
||
final_json_data.append({
|
||
"id": item["id"],
|
||
"context": item["context"], # 文本描述
|
||
"center": item["center"], # 点击坐标
|
||
"type": item["type"], # 控件类型
|
||
"clickable": item["clickable"], # 是否可点击
|
||
# "bbox": item["bbox"] # 可选：如果 Agent 需要知道大小
|
||
})
|
||
debug_json_data.append({
|
||
"id": item["id"],
|
||
"context": item["context"],
|
||
"text": item["text"],
|
||
"desc": item["desc"],
|
||
"resource_id": item["resource_id"],
|
||
"resource_label": item["resource_label"],
|
||
"center": item["center"],
|
||
"bbox": item["bbox"],
|
||
"type": item["type"],
|
||
"package": item["package"],
|
||
"clickable": item["clickable"],
|
||
"enabled": item["enabled"],
|
||
"visible_to_user": item["visible_to_user"],
|
||
})
|
||
|
||
# 4. 保存 JSON
|
||
with open(Config.JSON_PATH, "w", encoding="utf-8") as f:
|
||
json.dump(final_json_data, f, ensure_ascii=False, indent=2)
|
||
with open(Config.DEBUG_JSON_PATH, "w", encoding="utf-8") as f:
|
||
json.dump(debug_json_data, f, ensure_ascii=False, indent=2)
|
||
|
||
print(f"✅ 处理完成 (耗时 {time.time()-start_time:.2f}s)")
|
||
print(f"📂 数据已保存: {Config.JSON_PATH}")
|
||
print(f"🧪 调试数据已保存: {Config.DEBUG_JSON_PATH}")
|
||
print(f"🖼️ 图片已保存: {Config.LABELED_SCREEN_PATH}")
|
||
else:
|
||
print("⚠️ 未发现可见 UI 元素或解析失败")
|