320 lines
10 KiB
Python
320 lines
10 KiB
Python
# 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") |