This commit is contained in:
HuangHai
2026-01-26 10:50:11 +08:00
parent 4868198143
commit a662c33ecf
6 changed files with 263 additions and 924 deletions

View File

@@ -14,11 +14,9 @@ if project_root not in sys.path:
sys.path.append(project_root)
import json
from Util.VLMKit import VLMKit
from Util.EasyOcrKit import EasyOcrKit
# 初始化 VLMKit 和 EasyOcrKit
vlm_kit = VLMKit()
# 初始化 EasyOcrKit
ocr_kit = EasyOcrKit()
# 配置日志
@@ -72,302 +70,216 @@ def connect_device():
logger.error(f"设备连接失败: {e}")
return None
async def get_vlm_json(image_path, prompt):
"""
通用 VLM 分析函数,返回 JSON 数据 (自动处理归一化坐标的反归一化)
"""
try:
# 调用 VLM
response = await vlm_kit.analyze_image(image_path, prompt)
json_str = vlm_kit.extract_json(response)
result_data = json.loads(json_str)
# 获取图片尺寸进行坐标反归一化
try:
from PIL import Image
with Image.open(image_path) as img:
width, height = img.size
# 定义反归一化函数
def denormalize(point):
if not point or len(point) != 2:
return point
return [int(point[0] / 1000 * width), int(point[1] / 1000 * height)]
# 递归遍历字典进行反归一化 (仅针对常见坐标字段 center, input_box)
def recursive_denormalize(data):
if isinstance(data, dict):
for key, value in data.items():
if key in ["center", "input_box", "coordinates"] and isinstance(value, list) and len(value) == 2:
data[key] = denormalize(value)
elif isinstance(value, (dict, list)):
recursive_denormalize(value)
elif isinstance(data, list):
for item in data:
recursive_denormalize(item)
recursive_denormalize(result_data)
except Exception as e:
logger.warning(f"坐标反归一化失败: {e},将使用原始坐标")
return result_data
except Exception as e:
logger.error(f"VLM Analysis Failed: {e}", exc_info=True)
return None
async def get_vlm_analysis(image_path):
"""
仅调用 VLM 分析图片,返回原始 JSON 数据 (dict)
"""
logger.info(f"正在使用 VLM 分析图片: {image_path}")
# 构造 Prompt
prompt = """
请分析这张微信聊天截图,提取所有对话消息。
【核心规则 - 优先级最高】
1. 🚀 **从下往上扫描**:必须确保屏幕最底部的消息被识别。很多时候最底部的消息是最重要的。
2. 🔴 **未读红点 (Unread)**:极度关注语音气泡右上角的红点。如果有红点,`is_unread` 必须为 true。
3. 📦 **完整性**:识别图中【所有】可见的消息气泡,包括文本消息、语音消息、系统提示(如“昨天 10:36”、“你撤回了一条消息”)。
【消息类型判别】
- **发送者 (Sender)**:左侧头像为“对方”(Other),右侧头像为“我”(Me)。
- **语音 (Voice)**:
- 气泡内只有时长(如 5")和声波图标。
- **重点**:如果语音气泡右侧有灰色的“转文字”字样或红点,且下方没有对应的文本翻译气泡,说明它【尚未转换】。
- `status` 判断:只有当语音气泡【正下方】紧跟着一个相同发送者的文本气泡(内容是翻译结果),`status` 才为 "converted"。否则为 "unconverted"。
- **文本 (Text)**:气泡内包含具体的文字内容。
【坐标系统】
- 使用 [0-1000] 归一化坐标。返回气泡的几何中心点 `center`。
- 识别底部输入框的位置 `input_box`。
【输出格式】
请返回纯 JSON 格式:
{
"is_chat_interface": true,
"input_box": [x, y],
"messages": [
{
"type": "voice" | "text" | "system",
"sender": "对方" | "我" | "系统",
"status": "converted" | "unconverted",
"is_unread": true | false,
"center": [x, y],
"content": "消息内容或时长"
},
...
]
}
"""
2. <EFBFBD> **红点 (Unread)**:极度关注语音气泡右上角的红点。如果有红点,`is_unread` 必须为 true。
3. 📦 **完整性**:识别图中【所有】可见的消息气泡。不要遗漏任何一个,特别是连续的语音消息。
【消息类型判别】
- **发送者 (Sender)**:左侧头像为“对方”(Other),右侧头像为“我”(Me)。
- **语音 (Voice)**:气泡内只有时长(如 5")和声波图标。
- 语音气泡右侧可能有“转文字”或“取消”等灰色小字,请忽略这些文字,气泡依然是 Voice。
- `status` 判断:如果语音气泡下方【紧接着】有一个属于同一人的文本气泡,且内容看起来像翻译结果,则 `status` 为 "converted",否则为 "unconverted"。
- **文本 (Text)**:气泡内包含具体的文字内容。
【坐标系统】
- 使用 [0-1000] 归一化坐标。返回气泡的几何中心点 `center`。
- 识别底部输入框的位置 `input_box`。
【输出格式】
请返回纯 JSON 格式:
{
"is_chat_interface": true,
"input_box": [x, y],
"messages": [
{
"type": "voice",
"sender": "对方" | "我",
"status": "converted" | "unconverted",
"is_unread": true | false,
"center": [x, y],
"content": "8\""
},
...
]
}
"""
try:
# 调用 VLM
response = await vlm_kit.analyze_image(image_path, prompt)
logger.info(f"VLM Raw Response: {response}") # 打印原始响应以便调试
json_str = vlm_kit.extract_json(response)
result_data = json.loads(json_str)
# 获取图片尺寸进行坐标反归一化
try:
from PIL import Image
with Image.open(image_path) as img:
width, height = img.size
# 定义反归一化函数
def denormalize(point):
if not point or len(point) != 2:
return point
return [int(point[0] / 1000 * width), int(point[1] / 1000 * height)]
# 反归一化 input_box
if result_data.get("input_box"):
result_data["input_box"] = denormalize(result_data["input_box"])
# 反归一化 messages
if result_data.get("messages"):
for msg in result_data["messages"]:
if msg.get("center"):
msg["center"] = denormalize(msg["center"])
if msg.get("coordinates"): # 兼容旧字段
msg["coordinates"] = denormalize(msg["coordinates"])
except Exception as e:
logger.warning(f"坐标反归一化失败: {e},将使用原始坐标")
return result_data
except Exception as e:
logger.error(f"VLM Analysis Failed: {e}", exc_info=True)
return None
async def analyze_chat_image(image_path, output_path, device=None, target_name="对方"):
"""
使用 VLM 识别微信聊天截图中的对话内容、语音消息状态以及输入框位置
替代原本的 CV/OCR 方案
全面采用 CV + OCR 识别微信聊天截图中的最后一条消息
不再使用 VLM
"""
# 语音识别标志
should_trigger_convert = False
try:
result_data = await get_vlm_analysis(image_path)
# 1. 初始化
d = device if device else connect_device()
if not d:
return [], None
# 2. 读取图片
img = cv2.imread(image_path)
if img is None:
logger.error(f"无法读取图片: {image_path}")
return [], None
h, w = img.shape[:2]
if not result_data:
return [], None
# 3. 模板匹配寻找语音图标和红点
audio_template = os.path.join(TEMPLATE_DIR, "audio.jpg")
red_point_template = os.path.join(TEMPLATE_DIR, "red_point.jpg")
audio_matches = find_all_template_matches(image_path, audio_template, threshold=0.8)
red_points = find_all_template_matches(image_path, red_point_template, threshold=0.8)
# 4. OCR 识别所有文本
logger.info("正在执行 OCR 识别...")
ocr_results = ocr_kit.read_text(image_path)
# 5. 整合所有消息
messages = []
debug_img = img.copy() # 初始化调试图
try:
# 检查是否为聊天界面
is_chat = result_data.get("is_chat_interface", False)
if not is_chat:
logger.warning("VLM 判断当前不是微信聊天界面")
return None, None
# A. 添加语音消息
for ax, ay in audio_matches:
# 过滤掉顶部和底部的非聊天区域 (经验值: 顶部150, 底部250)
if ay < 150 or ay > h - 250:
continue
if isinstance(result_data, list):
# 兼容旧格式
messages = result_data
input_field_coordinates = None
else:
messages = result_data.get("messages", [])
input_field_coordinates = result_data.get("input_box") # input_box
sender = "对方" if ax < w / 2 else "我"
is_unread = False
for rx, ry in red_points:
# 红点通常在语音图标右侧且 Y 轴相近
if abs(ry - ay) < 50 and rx > ax:
is_unread = True
break
# 根据已读/未读画框:未读红框,已读绿框
color = (0, 0, 255) if is_unread else (0, 255, 0)
cv2.rectangle(debug_img, (ax-30, ay-30), (ax+30, ay+30), color, 2)
except Exception as e:
logger.error(f"解析 VLM 结果失败: {e}")
# --- 新增:判断是否已转文字 ---
is_converted = False
for bbox, text, conf in ocr_results:
c_x = int((bbox[0][0] + bbox[2][0]) / 2)
c_y = int((bbox[0][1] + bbox[2][1]) / 2)
# 转换后的文字通常在语音图标下方 30-300 像素内,且水平位置相近
if 30 < c_y - ay < 300 and abs(c_x - ax) < 200:
is_converted = True
break
label = "YES" if is_converted else "NO"
# 在框的右侧标注 YES 或 NO
cv2.putText(debug_img, label, (ax + 40, ay + 10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2)
# --- 结束 ---
messages.append({
"type": "voice",
"sender": sender,
"center": (ax, ay),
"y": ay,
"is_unread": is_unread,
"is_converted": is_converted
})
# B. 添加文本消息
# 简单策略:排除掉明显是系统时间、输入框或顶部标题的文字
for bbox, text, conf in ocr_results:
c_x = int((bbox[0][0] + bbox[2][0]) / 2)
c_y = int((bbox[0][1] + bbox[2][1]) / 2)
# 过滤区域
if c_y < 150 or c_y > h - 250:
continue
# 过滤掉单字(可能是头像旁边的文字或杂质)和某些系统词
if len(text) < 1 and "昨天" not in text and "今天" not in text:
continue
sender = "对方" if c_x < w / 2 else "我"
messages.append({
"type": "text",
"sender": sender,
"content": text,
"center": (c_x, c_y),
"y": c_y
})
# 6. 排序并找出最后一条消息
if not messages:
logger.warning("未发现任何消息")
if output_path:
cv2.imwrite(output_path, debug_img)
return [], None
# 按 Y 坐标从上到下排序
messages.sort(key=lambda x: x['y'])
last_msg = messages[-1]
if output_path:
cv2.imwrite(output_path, debug_img)
logger.info(f"调试图已保存: {output_path}")
dialogue_log = []
unconverted_voices = []
# 处理识别结果
for msg in messages:
sender = msg.get('sender', '未知')
msg_type = msg.get('type', 'other')
content = msg.get('content', '')
coords = msg.get('center', [0, 0]) # center
status = msg.get('status', 'unconverted')
is_unread = msg.get('is_unread', False)
is_converted = (status == "converted")
unread_mark = "[未读]" if is_unread else ""
# 记录对话日志
if msg_type == 'voice':
if is_converted:
dialogue_log.append(f"{sender}: {unread_mark}[语音] {content} (已转换)")
else:
dialogue_log.append(f"{sender}: {unread_mark}[语音] (待转换)")
# 将 center 转换为 coordinates 供后续使用
msg['coordinates'] = coords
unconverted_voices.append(msg)
elif msg_type == 'text':
dialogue_log.append(f"{sender}: {content}")
logger.info(f"VLM 识别: {sender} [{msg_type}] {content} (Converted: {is_converted}, Unread: {is_unread})")
# 处理未转换的语音消息
input_field_coordinates = (w // 2, int(h * 0.9)) # 默认输入框位置
# 7. 自动处理所有“红框 + NO”的语音消息
unconverted_voices = [m for m in messages if m['type'] == 'voice' and m.get('is_unread') and not m.get('is_converted')]
if unconverted_voices:
# 优先级:1. 有红点的最后一条 2. 没红点的最后一条
unread_voices = [v for v in unconverted_voices if v.get('is_unread')]
if unread_voices:
logger.info(f"发现 {len(unread_voices)} 条未读语音消息,优先处理最后一条...")
voice_to_process = unread_voices[-1]
else:
logger.info(f"发现 {len(unconverted_voices)} 条未转换语音消息,处理最后一条...")
voice_to_process = unconverted_voices[-1]
logger.info(f"发现 {len(unconverted_voices)} 条未转换的未读语音,开始处理...")
for v_msg in unconverted_voices:
vx, vy = int(v_msg['center'][0]), int(v_msg['center'][1])
logger.info(f"--- 正在处理语音消息 ({vx}, {vy}) ---")
# 仅保留选中的一条进行处理
unconverted_voices = [voice_to_process]
# A. 长按语音
logger.info(f"正在长按语音消息 ({vx}, {vy})...")
d.long_click(vx, vy, 1.5)
time.sleep(1.5)
# 使用传入的 device 或创建新连接
d = device if device else connect_device()
# B. 截图寻找“转文字”按钮
menu_shot = os.path.join(OUTPUT_DIR, f"voice_menu_{vy}.jpg")
d.screenshot(menu_shot)
zhuan_template = os.path.join(TEMPLATE_DIR, "zhun_wen_zi.jpg")
for voice in unconverted_voices:
vx, vy = voice['coordinates']
logger.info(f"长按语音消息: ({vx}, {vy})")
# 降低阈值到 0.7 以增加匹配成功率
btn_pos = find_template_match(menu_shot, zhuan_template, threshold=0.7)
if btn_pos:
btn_x, btn_y = int(btn_pos[0]), int(btn_pos[1])
logger.info(f"✅ 找到'转文字'按钮: ({btn_x}, {btn_y}),点击中...")
d.click(btn_x, btn_y)
# 等待转换完成 (根据语音长度,通常 3-5 秒足够)
logger.info("等待语音转文字完成...")
time.sleep(5.0)
# C. 再次截图 OCR 获取转换后的文字
after_convert_shot = os.path.join(OUTPUT_DIR, f"after_auto_{vy}.jpg")
d.screenshot(after_convert_shot)
convert_ocr = ocr_kit.read_text(after_convert_shot)
# 提取转换文字:寻找在语音图标下方的文字块
converted_text = ""
for c_bbox, c_text, c_conf in convert_ocr:
cc_x = (c_bbox[0][0] + c_bbox[2][0]) / 2
cc_y = (c_bbox[0][1] + c_bbox[2][1]) / 2
# 转换后的文字通常在语音图标下方 30-300 像素内,且水平位置相近
if 30 < cc_y - vy < 300 and abs(cc_x - vx) < 250:
converted_text = c_text
break
if converted_text:
logger.info(f"✨ OCR 识别成功!")
print(f"\n[语音转文字结果]: {converted_text}\n")
# 同步到消息对象
v_msg['content'] = converted_text
v_msg['is_converted'] = True
# 如果这条消息也是最后一条消息,更新 dialogue_log 需要的内容
if v_msg == last_msg:
last_msg['content'] = converted_text
else:
logger.warning("❌ OCR 未能提取到转换后的文字内容")
# D. 长按并点击“取消转文字”恢复界面
logger.info("正在恢复界面状态 (点击'取消转文字')...")
d.long_click(vx, vy, 1.5)
time.sleep(1.0)
cancel_shot = os.path.join(OUTPUT_DIR, f"cancel_menu_{vy}.jpg")
d.screenshot(cancel_shot)
cancel_template = os.path.join(TEMPLATE_DIR, "cancel_zhuan_wen_zi.jpg")
cancel_btn = find_template_match(cancel_shot, cancel_template, threshold=0.7)
# 查找“转文字” (使用 OCR)
menu_shot_path = os.path.join(os.path.dirname(image_path), "temp_menu_shot.jpg")
d.screenshot(menu_shot_path)
# OCR 识别
ocr_results = ocr_kit.read_text(menu_shot_path)
convert_btn_center = None
for bbox, text, conf in ocr_results:
if "转文字" in text or "转换为文字" in text:
# bbox is [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
# Calculate center
c_x = int((bbox[0][0] + bbox[2][0]) / 2)
c_y = int((bbox[0][1] + bbox[2][1]) / 2)
convert_btn_center = (c_x, c_y)
break
if convert_btn_center:
logger.info(f"OCR 找到 '转文字' 按钮: {convert_btn_center}")
d.click(convert_btn_center[0], convert_btn_center[1])
should_trigger_convert = True
# 动态等待: 60s语音约需10s转换,比例约 1/6
duration_str = voice.get('content', '0').replace('"', '').strip()
try:
duration = int(duration_str)
except:
duration = 10 # 默认值
wait_seconds = max(2, duration / 5.0) # 稍微多等一点,用 /5.0
logger.info(f"语音时长 {duration}s,预计等待转换 {wait_seconds:.1f}s...")
time.sleep(wait_seconds)
if cancel_btn:
c_btn_x, c_btn_y = int(cancel_btn[0]), int(cancel_btn[1])
d.click(c_btn_x, c_btn_y)
logger.info(f"✅ 已点击'取消转文字' ({c_btn_x}, {c_btn_y}),界面已恢复")
else:
logger.warning("OCR 未找到 '转文字' 菜单项")
# 点击空白处关闭菜单,避免遮挡
d.click(vx + 200, vy)
if should_trigger_convert:
# 转换完成后稍微多等一下,确保 UI 刷新
time.sleep(1.0)
# 即使触发了转换,我们也返回当前的对话日志,但在日志末尾注明正在转换
dialogue_log.append("系统: [正在转换语音...]")
return dialogue_log, input_field_coordinates
# 兜底:点击语音图标右侧空白处尝试关闭菜单
logger.warning("⚠️ 未找到'取消转文字'按钮,尝试点击空白处关闭菜单")
d.click(vx + 300, vy)
else:
logger.warning("❌ 未能找到'转文字'按钮,可能长按失败或模板不匹配")
# 尝试点击空白处退出菜单
d.click(vx + 300, vy)
# 8. 整合对话日志 (仅针对最后一条消息进行反馈)
dialogue_log = []
if last_msg['type'] == 'voice':
# 优先使用刚才转文字得到的内容
content = last_msg.get('content') or "[语音]"
dialogue_log.append(f"{last_msg['sender']}: {content}")
else:
dialogue_log.append(f"{last_msg['sender']}: {last_msg['content']}")
return dialogue_log, input_field_coordinates
except Exception as e:
logger.error(f"VLM 分析失败: {e}", exc_info=True)
logger.error(f"analyze_chat_image 失败: {e}", exc_info=True)
return [], None