AutoAndroidController/py/string-reg-location.py

"""
文字识别和定位模块
功能：在截图中查找指定文字，并返回文字在截图中的坐标
使用 OnnxOCR 进行文字识别
"""

import sys
import os
import cv2
from pathlib import Path
from typing import Optional, Tuple, Dict, Any

# 添加 OnnxOCR 路径到 sys.path
current_dir = Path(__file__).parent
onnxocr_path = current_dir / 'OnnxOCR'
if str(onnxocr_path) not in sys.path:
    sys.path.insert(0, str(onnxocr_path))

from onnxocr.onnx_paddleocr import ONNXPaddleOcr

# 设置环境变量，跳过模型源连接检查（避免首次运行时超时）
os.environ['DISABLE_MODEL_SOURCE_CHECK'] = 'True'


def find_text_location(
    screenshot_path: str,
    target_text: str,
    device_width: int = None,
    device_height: int = None,
    use_angle_cls: bool = True,
    lang: str = 'ch'
) -> Optional[Dict[str, Any]]:
    """
    在截图中查找目标文字并返回坐标
    
    Args:
        screenshot_path: 截图文件路径
        target_text: 要查找的文字
        device_width: 设备实际宽度（像素），如果提供则会将坐标转换到设备分辨率
        device_height: 设备实际高度（像素），如果提供则会将坐标转换到设备分辨率
        use_angle_cls: 是否使用角度分类器，默认True (对应 OnnxOCR 的 use_angle_cls)
        lang: 语言类型，'ch'表示中英文混合，'en'表示英文，默认'ch' (OnnxOCR 内部处理)
    
    Returns:
        如果找到文字，返回包含坐标信息的字典：
        {
            "found": True,
            "x": 中心点x坐标,
            "y": 中心点y坐标,
            "width": 文字框宽度,
            "height": 文字框高度,
            "bbox": [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]  # 文字框的四个角点
        }
        如果未找到，返回 {"found": False}
    """
    # 检查文件是否存在
    screenshot = Path(screenshot_path)
    if not screenshot.exists():
        raise FileNotFoundError(f"截图文件不存在: {screenshot_path}")
    
    # 检查目标文字是否为空
    if not target_text or not target_text.strip():
        raise ValueError("目标文字不能为空")
    
    try:
        # 读取截图以获取实际尺寸
        img = cv2.imread(str(screenshot))
        if img is None:
            raise ValueError(f"无法读取截图文件: {screenshot_path}")
        
        screenshot_height, screenshot_width = img.shape[:2]
        
        # 计算缩放比例（如果提供了设备分辨率）
        scale_x = 1.0
        scale_y = 1.0
        if device_width is not None and device_height is not None:
            scale_x = device_width / screenshot_width
            scale_y = device_height / screenshot_height
        
        # 初始化 OnnxOCR（首次调用可能需要一些时间）
        # OnnxOCR 需要传入图片对象，而不是路径
        ocr = ONNXPaddleOcr(use_angle_cls=use_angle_cls, use_gpu=False)
        
        # 识别截图中的所有文字（OnnxOCR 需要传入 cv2 读取的图片对象）
        result = ocr.ocr(img, det=True, rec=True, cls=use_angle_cls)
        
        # 如果识别结果为空
        if not result or not result[0]:
            return {"found": False}
        
        # 在识别结果中查找目标文字
        # result[0] 是一个列表，每个元素是一行文字的识别结果
        # 格式: [[[x1,y1], [x2,y2], [x3,y3], [x4,y4]], (文字内容, 置信度)]
        for line in result[0]:
            if not line:
                continue
            
            # line[0] 是四个角点坐标
            # line[1] 是 (文字内容, 置信度)
            bbox = line[0]  # 四个角点: [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
            text_info = line[1]
            recognized_text = text_info[0]  # 识别的文字内容
            confidence = text_info[1]  # 置信度
            
            # 检查识别的文字是否包含目标文字（支持部分匹配和完全匹配）
            # 使用 in 操作符支持部分匹配，如果需要完全匹配可以使用 ==
            if target_text in recognized_text or recognized_text in target_text:
                # 找到匹配，计算文字框的中心点和尺寸
                # bbox 是四个角点的列表: [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
                x_coords = [point[0] for point in bbox]
                y_coords = [point[1] for point in bbox]
                
                # 计算边界框（基于截图尺寸）
                min_x = int(min(x_coords))
                max_x = int(max(x_coords))
                min_y = int(min(y_coords))
                max_y = int(max(y_coords))
                
                # 将坐标转换到设备分辨率（如果提供了设备分辨率）
                min_x = int(min_x * scale_x)
                max_x = int(max_x * scale_x)
                min_y = int(min_y * scale_y)
                max_y = int(max_y * scale_y)
                
                # 转换 bbox 坐标
                scaled_bbox = [[int(p[0] * scale_x), int(p[1] * scale_y)] for p in bbox]
                
                # 计算中心点（基于设备分辨率）
                center_x = int((min_x + max_x) / 2)
                center_y = int((min_y + max_y) / 2)
                
                # 计算宽度和高度（基于设备分辨率）
                width = max_x - min_x
                height = max_y - min_y
                
                return {
                    "found": True,
                    "x": center_x,
                    "y": center_y,
                    "width": width,
                    "height": height,
                    "bbox": scaled_bbox,  # 已转换到设备分辨率的坐标
                    "text": recognized_text,  # 实际识别的文字
                    "confidence": float(confidence)  # 置信度
                }
        
        # 未找到匹配的文字
        return {"found": False}
        
    except Exception as e:
        raise RuntimeError(f"OCR 识别过程中出错: {str(e)}")


def find_text_location_multiple(
    screenshot_path: str,
    target_text: str,
    device_width: int = None,
    device_height: int = None,
    use_angle_cls: bool = True,
    lang: str = 'ch'
) -> list:
    """
    在截图中查找目标文字的所有出现位置（可能有多处匹配）
    
    Args:
        screenshot_path: 截图文件路径
        target_text: 要查找的文字
        device_width: 设备实际宽度（像素）
        device_height: 设备实际高度（像素）
        use_angle_cls: 是否使用角度分类器
        lang: 语言类型
    
    Returns:
        返回所有匹配位置的列表，每个元素为包含坐标信息的字典
    """
    screenshot = Path(screenshot_path)
    if not screenshot.exists():
        raise FileNotFoundError(f"截图文件不存在: {screenshot_path}")
    
    if not target_text or not target_text.strip():
        raise ValueError("目标文字不能为空")
    
    try:
        # 读取截图以获取实际尺寸
        img = cv2.imread(str(screenshot))
        if img is None:
            raise ValueError(f"无法读取截图文件: {screenshot_path}")
        
        screenshot_height, screenshot_width = img.shape[:2]
        
        # 计算缩放比例（如果提供了设备分辨率）
        scale_x = 1.0
        scale_y = 1.0
        if device_width is not None and device_height is not None:
            scale_x = device_width / screenshot_width
            scale_y = device_height / screenshot_height
        
        # 初始化 OnnxOCR
        ocr = ONNXPaddleOcr(use_angle_cls=use_angle_cls, use_gpu=False)
        # OnnxOCR 需要传入图片对象
        result = ocr.ocr(img, det=True, rec=True, cls=use_angle_cls)
        
        if not result or not result[0]:
            return []
        
        matches = []
        for line in result[0]:
            if not line:
                continue
            
            bbox = line[0]
            text_info = line[1]
            recognized_text = text_info[0]
            confidence = text_info[1]
            
            # 检查是否匹配
            if target_text in recognized_text or recognized_text in target_text:
                x_coords = [point[0] for point in bbox]
                y_coords = [point[1] for point in bbox]
                
                # 计算边界框（基于截图尺寸）
                min_x = int(min(x_coords))
                max_x = int(max(x_coords))
                min_y = int(min(y_coords))
                max_y = int(max(y_coords))
                
                # 将坐标转换到设备分辨率（如果提供了设备分辨率）
                min_x = int(min_x * scale_x)
                max_x = int(max_x * scale_x)
                min_y = int(min_y * scale_y)
                max_y = int(max_y * scale_y)
                
                # 转换 bbox 坐标
                scaled_bbox = [[int(p[0] * scale_x), int(p[1] * scale_y)] for p in bbox]
                
                center_x = int((min_x + max_x) / 2)
                center_y = int((min_y + max_y) / 2)
                width = max_x - min_x
                height = max_y - min_y
                
                matches.append({
                    "found": True,
                    "x": center_x,
                    "y": center_y,
                    "width": width,
                    "height": height,
                    "bbox": scaled_bbox,
                    "text": recognized_text,
                    "confidence": float(confidence)
                })
        
        return matches
        
    except Exception as e:
        raise RuntimeError(f"OCR 识别过程中出错: {str(e)}")


if __name__ == "__main__":
    # 测试示例
    if len(sys.argv) < 3:
        print("用法: python string-reg-location.py <截图路径> <要查找的文字> [设备宽度] [设备高度]")
        print("示例: python string-reg-location.py screenshot.png \"你好\" 1080 2400")
        sys.exit(1)
    
    screenshot_path = sys.argv[1]
    target_text = sys.argv[2]
    
    device_width = None
    device_height = None
    if len(sys.argv) >= 5:
        try:
            device_width = int(sys.argv[3])
            device_height = int(sys.argv[4])
        except ValueError:
            print("警告: 无法解析设备分辨率，将使用截图原始尺寸")
    
    try:
        result = find_text_location(screenshot_path, target_text, device_width, device_height)
        if result.get("found"):
            x = result["x"]
            y = result["y"]
            w = result["width"]
            h = result["height"]
            print(f"找到文字！坐标: x={x}, y={y}, 宽度={w}, 高度={h}")
            print(f"识别的文字: {result.get('text', '')}")
            print(f"置信度: {result.get('confidence', 0):.2f}")
            print(f"JSON格式: {{\"x\": {x}, \"y\": {y}, \"width\": {w}, \"height\": {h}}}")
        else:
            print("未找到匹配的文字")
    except Exception as e:
        print(f"错误: {e}")
        sys.exit(1)