136 lines
4.1 KiB
Python
136 lines
4.1 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
极简UI元素检测脚本 - 基于OmniParser的YOLO模型
|
||
依赖: ultralytics, rapidocr-onnxruntime, pillow, numpy
|
||
"""
|
||
import sys
|
||
from pathlib import Path
|
||
from ultralytics import YOLO
|
||
from PIL import Image, ImageDraw
|
||
import numpy as np
|
||
|
||
DEFAULT_MODEL = str(Path(__file__).resolve().parent.parent / 'temp' / 'weights' / 'icon_detect' / 'model.pt')
|
||
|
||
# 可选：使用rapidocr做OCR
|
||
try:
|
||
from rapidocr_onnxruntime import RapidOCR
|
||
ocr_engine = RapidOCR()
|
||
HAS_OCR = True
|
||
except ImportError:
|
||
HAS_OCR = False
|
||
print("警告: rapidocr未安装，跳过OCR功能")
|
||
|
||
def detect_ui_elements(image_path, model_path=None, conf_threshold=0.25):
|
||
model_path = model_path or DEFAULT_MODEL
|
||
"""检测UI元素并返回边界框"""
|
||
# 加载模型
|
||
model = YOLO(model_path)
|
||
|
||
# 推理
|
||
results = model(image_path, conf=conf_threshold, verbose=False)
|
||
|
||
# 提取检测结果
|
||
detections = []
|
||
for result in results:
|
||
boxes = result.boxes
|
||
for box in boxes:
|
||
x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
|
||
conf = float(box.conf[0])
|
||
cls = int(box.cls[0])
|
||
detections.append({
|
||
'bbox': [int(x1), int(y1), int(x2), int(y2)],
|
||
'confidence': conf,
|
||
'class': cls
|
||
})
|
||
|
||
return detections
|
||
|
||
def ocr_text(image_path):
|
||
"""OCR识别文本"""
|
||
if not HAS_OCR:
|
||
return []
|
||
|
||
result, _ = ocr_engine(image_path)
|
||
if not result:
|
||
return []
|
||
|
||
texts = []
|
||
for item in result:
|
||
bbox, text, conf = item
|
||
texts.append({
|
||
'text': text,
|
||
'bbox': bbox,
|
||
'confidence': conf
|
||
})
|
||
return texts
|
||
|
||
def visualize(image_path, detections, ocr_results=None, output_path=None):
|
||
"""可视化检测结果"""
|
||
img = Image.open(image_path)
|
||
draw = ImageDraw.Draw(img)
|
||
|
||
# 画UI元素框（红色）
|
||
for det in detections:
|
||
x1, y1, x2, y2 = det['bbox']
|
||
draw.rectangle([x1, y1, x2, y2], outline='red', width=2)
|
||
draw.text((x1, y1-10), f"{det['confidence']:.2f}", fill='red')
|
||
|
||
# 画OCR文本框（蓝色）
|
||
if ocr_results:
|
||
for ocr in ocr_results:
|
||
bbox = ocr['bbox']
|
||
points = [(bbox[i][0], bbox[i][1]) for i in range(4)]
|
||
draw.polygon(points, outline='blue')
|
||
draw.text((points[0][0], points[0][1]-10), ocr['text'][:10], fill='blue')
|
||
|
||
if output_path:
|
||
img.save(output_path)
|
||
return img
|
||
|
||
def main():
|
||
if len(sys.argv) < 2:
|
||
print("用法: python ui_detect.py <图片路径> <模型路径> [输出路径]")
|
||
print("示例: python ui_detect.py screenshot.png weights/icon_detect/model.pt output.png")
|
||
sys.exit(1)
|
||
|
||
image_path = sys.argv[1]
|
||
model_path = sys.argv[2] if len(sys.argv) > 2 else DEFAULT_MODEL
|
||
output_path = sys.argv[3] if len(sys.argv) > 3 else "output.png"
|
||
|
||
print(f"检测图片: {image_path}")
|
||
print(f"使用模型: {model_path}")
|
||
|
||
# UI元素检测
|
||
print("\n[1/2] YOLO检测UI元素...")
|
||
detections = detect_ui_elements(image_path, model_path)
|
||
print(f"检测到 {len(detections)} 个UI元素")
|
||
for i, det in enumerate(detections, 1):
|
||
print(f" {i}. bbox={det['bbox']}, conf={det['confidence']:.3f}")
|
||
|
||
# OCR文本识别
|
||
ocr_results = None
|
||
if HAS_OCR:
|
||
print("\n[2/2] OCR识别文本...")
|
||
ocr_results = ocr_text(image_path)
|
||
print(f"识别到 {len(ocr_results)} 个文本区域")
|
||
for i, ocr in enumerate(ocr_results, 1):
|
||
print(f" {i}. text='{ocr['text']}', conf={ocr['confidence']:.3f}")
|
||
|
||
# 可视化
|
||
print(f"\n保存结果到: {output_path}")
|
||
visualize(image_path, detections, ocr_results, output_path)
|
||
|
||
# 输出JSON格式结果
|
||
import json
|
||
result = {
|
||
'ui_elements': detections,
|
||
'ocr_texts': ocr_results or []
|
||
}
|
||
json_path = output_path.replace('.png', '.json')
|
||
with open(json_path, 'w', encoding='utf-8') as f:
|
||
json.dump(result, f, ensure_ascii=False, indent=2)
|
||
print(f"JSON结果: {json_path}")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|