# This file is created by Minyi Liu (GitHub ID: MiniMinyi) # The hash algorithm is copied from: # https://github.com/hjaurum/DHash/blob/master/dHash.py def load_image_from_path(img_path): """ Load an image from path :param img_path: The path to the image :return: """ import cv2 import numpy as np # Fix for reading images with Chinese paths on Windows # Use IMREAD_COLOR to ensure 3 channels (BGR) as cv2.imread does by default return cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR) def load_image_from_buf(img_bytes): """ Load an image from a byte array :param img_bytes: The byte array of an image :return: """ import cv2 import numpy img_bytes = numpy.frombuffer(img_bytes, dtype=numpy.uint8) return cv2.imdecode(img_bytes, cv2.IMREAD_UNCHANGED) # ------------------------------------------------------------- # 辅助函数：将 OpenCV 图像对象编码为 Base64 # ------------------------------------------------------------- import cv2 import numpy as np import base64 import os from typing import Optional, List, Tuple, Dict, Any # 全局变量用于缓存模型实例 _model_handler = None _model_config = None class OmniParserModelManager: """ OmniParser模型管理器，使用单例模式确保模型只初始化一次 """ _instance = None _handler = None def __new__(cls): if cls._instance is None: cls._instance = super(OmniParserModelManager, cls).__new__(cls) return cls._instance def __init__(self): if self._handler is None: self._initialize_model() def _initialize_model(self): """初始化OmniParser模型""" from .handler import EndpointHandler # 默认配置 config = { "model_dir": os.path.join(os.path.dirname(__file__), "."), "mode": "custom", "ocr_languages": ["ch_sim", "en"], "image_size": 1280, "bbox_threshold": 0.2, "iou_threshold": 0.5, } print("初始化OmniParser模型...") self._handler = EndpointHandler( model_dir=config["model_dir"], enable_yolo=True, enable_ocr=True, enable_caption=False, ocr_languages=config["ocr_languages"], ) self._config = config print("OmniParser模型初始化完成") def get_handler(self): """获取模型处理器实例""" return self._handler def get_config(self): """获取模型配置""" return self._config def get_model_manager() -> OmniParserModelManager: """获取模型管理器单例实例""" return OmniParserModelManager() def load_image_from_buf(img_bytes): """ Load image from bytes :param img_bytes: bytes of image :return: numpy.ndarray, representing an image in opencv """ import cv2 import numpy img_bytes = numpy.frombuffer(img_bytes, dtype=numpy.uint8) return cv2.imdecode(img_bytes, cv2.IMREAD_UNCHANGED) def _encode_cv_image(img): """ 将 OpenCV 图像 (numpy.ndarray) 编码为 Base64 字符串。 """ import cv2 import base64 # 将图像编码为 PNG 格式的字节流 success, encoded_image = cv2.imencode(".png", img) if not success: raise ValueError("无法将图像编码为 Base64") # 转换为 Base64 字符串 (不带 MIME 头) raw_b64 = base64.b64encode(encoded_image).decode("ascii") # 加上 MIME 头 return f"data:image/png;base64,{raw_b64}" def enhance_contrast_gamma(image, gamma=1.0): """ Apply Power Law Transformation to the image. Output = (Input/255) ^ gamma * 255 Gamma > 1.0 will darken the image (crush shadows). Gamma < 1.0 will brighten the image. """ if gamma == 1.0: return image import numpy as np import cv2 # Direct power law transformation table = np.array([((i / 255.0) ** gamma) * 255 for i in np.arange(0, 256)]).astype("uint8") return cv2.LUT(image, table) def find_views(img) -> List[Dict[str, Any]]: """ 使用 OmniParser 模型查找给定 UI 截图中的矩形视图。 默认应用 Gamma=4.0 的对比度增强以压暗背景（处理弹窗场景）。 :param img: numpy.ndarray, representing an image in opencv :return: List[ViewDict] - 符合统一 ViewDict 格式的视图列表 """ import cv2 # 默认应用 Gamma 校正 (Gamma=4.0) 以压暗背景 # 用户需求：默认执行此预处理，无需外部传参 print(f"[CV] Applying Default Gamma Correction (gamma=4.0)...") img = enhance_contrast_gamma(img, gamma=4.0) # 保存处理后的图像到当前目录，方便查看 import os import cv2 # output_path = os.path.join(".", "gamma_enhanced.png") # cv2.imwrite(output_path, img) # print(f"[CV] 已保存 Gamma 校正后的图像到: {os.path.abspath(output_path)}") # 获取图像的原始尺寸 height, width = img.shape[:2] # 获取模型管理器实例（单例模式，确保只初始化一次） model_manager = get_model_manager() handler = model_manager.get_handler() config = model_manager.get_config() # 构造推理请求 image_b64 = _encode_cv_image(img) payload = { "inputs": { "image": image_b64, # 模型推理尺寸可以设为一个标准值，例如 1280x1280 "image_size": {"w": config["image_size"], "h": config["image_size"]}, "bbox_threshold": config["bbox_threshold"], "iou_threshold": config["iou_threshold"], } } # 执行推理 result = handler(payload) # 解析和转换结果为统一的 ViewDict 格式 views = [] if "bboxes" in result: # 打印模型识别出的元素总数（过滤前） total_detected = len(result["bboxes"]) print(f"[CV] 模型识别出的元素总数（过滤前）: {total_detected}") # 打印各类型元素数量 icon_count = sum(1 for box in result["bboxes"] if box.get("type") == "icon") text_count = sum(1 for box in result["bboxes"] if box.get("type") == "text") print(f"[CV] - 图标元素数量: {icon_count}") print(f"[CV] - 文本元素数量: {text_count}") for idx, box_data in enumerate(result["bboxes"]): # 模型返回的归一化坐标 [x1, y1, x2, y2] bbox_norm = box_data.get("bbox") if not bbox_norm or len(bbox_norm) < 4: continue # 归一化坐标 x1_norm, y1_norm, x2_norm, y2_norm = bbox_norm # 转换为绝对像素坐标 x1 = int(x1_norm * width) y1 = int(y1_norm * height) x2 = int(x2_norm * width) y2 = int(y2_norm * height) # 过滤掉不合理的尺寸 if x2 <= x1 or y2 <= y1: continue # 获取文本信息（如果有） text = box_data.get("content", "") element_type = box_data.get("type", "cv_element") is_interactive = box_data.get("interactivity", True) # 按照 _get_view_signature 格式生成 signature，供 input_policy 使用 view_text_sig = text if text and len(text) <= 50 else "None" signature = "[class]%s[resource_id]%s[text]%s[%s,,]" % ( element_type or "None", "", # CV 视图无 resource_id view_text_sig or "None", "enabled", ) # 构建统一的 ViewDict 格式 view = { # === 必需字段 === "bounds": [[x1, y1], [x2, y2]], "text": text, "visible": True, "enabled": True, "clickable": is_interactive, "editable": False, "scrollable": False, "children": [], "view_str": f"cv_{idx}_{x1}_{y1}_{x2}_{y2}", "signature": signature, # === 可选字段 === "content_description": "", "resource_id": "", "class_name": element_type, "temp_id": idx, "parent": -1, "source": "cv", } views.append(view) return views def calculate_dhash(img): """ Calculate the dhash value of an image. :param img: numpy.ndarray, representing an image in opencv :return: """ difference = _calculate_pixel_difference(img) # convert to hex decimal_value = 0 hash_string = "" for index, value in enumerate(difference): if value: decimal_value += value * (2 ** (index % 8)) if index % 8 == 7: # every eight binary bit to one hex number hash_string += str(hex(decimal_value)[2:-1].rjust(2, "0")) # 0xf=>0x0f decimal_value = 0 return hash_string def _calculate_pixel_difference(img): """ Calculate difference between pixels :param img: numpy.ndarray, representing an image in opencv """ import cv2 resize_width = 18 resize_height = 16 # 1. resize to 18*16 smaller_image = cv2.resize(img, (resize_width, resize_height)) # 2. calculate grayscale grayscale_image = cv2.cvtColor(smaller_image, cv2.COLOR_BGR2GRAY) # 3. calculate difference between pixels difference = [] for row in range(resize_height): for col in range(resize_width - 1): difference.append(grayscale_image[row][col] > grayscale_image[row][col + 1]) return difference def dhash_hamming_distance(dhash1, dhash2): """ Calculate the hamming distance between two dhash values :param dhash1: str, the dhash of an image returned by `calculate_dhash` :param dhash2: str, the dhash of an image returned by `calculate_dhash` :return: int, the hamming distance between two dhash values """ difference = (int(dhash1, 16)) ^ (int(dhash2, 16)) return bin(difference).count("1")