wehub-resource-sync
205 lines
8.9 KiB
Python
205 lines
8.9 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
UI元素检测 - 基于 OmniParser-v2.0 的 icon_detect YOLO 模型 + RapidOCR
|
|
|
|
首次使用前必须准备对应权重；这是 OmniParser-v2.0 的 icon_detect YOLO 模型。
|
|
若缺失模型文件，新用户/AI 应搜索并下载 OmniParser-v2.0 icon_detect YOLO 权重。
|
|
|
|
用法:
|
|
from ui_detect import detect
|
|
elements = detect("screenshot.png") # 默认match模式
|
|
elements = detect(pil_image) # 支持PIL.Image
|
|
elements = detect(img, mode='crop') # crop备选
|
|
返回: [{'bbox':[x1,y1,x2,y2], 'type':'icon'|'text', 'label':str|None, 'confidence':float}]
|
|
模式: match=YOLO+全图OCR IoU匹配(推荐,1.2s,无文字图标label=None可VLM保底) | crop=拼接crop OCR(备选,更精确,2.3s)
|
|
依赖: ultralytics, rapidocr-onnxruntime, pillow, numpy
|
|
"""
|
|
from pathlib import Path
|
|
from PIL import Image, ImageDraw
|
|
import numpy as np
|
|
import json, urllib.request, subprocess, sys, time
|
|
|
|
#print('[UI DETECT] 截图分析后必须使用物理坐标，ljqCtrl也使用物理坐标！')
|
|
|
|
DEFAULT_MODEL = str(Path(__file__).resolve().parent.parent / 'temp' / 'weights' / 'icon_detect' / 'model.pt')
|
|
|
|
try:
|
|
from rapidocr_onnxruntime import RapidOCR
|
|
_ocr = RapidOCR()
|
|
except ImportError: _ocr = None
|
|
|
|
_YOLO = None
|
|
_YOLO_PORT = 31876
|
|
|
|
def _yolo_local(image_path, conf=0.25):
|
|
global _YOLO
|
|
if _YOLO is None:
|
|
from ultralytics import YOLO
|
|
_YOLO = YOLO(DEFAULT_MODEL)
|
|
res = _YOLO(image_path, conf=conf, verbose=False)
|
|
boxes = []
|
|
for r in res:
|
|
for b in r.boxes:
|
|
x1, y1, x2, y2 = map(int, b.xyxy[0].cpu().numpy())
|
|
boxes.append([x1, y1, x2, y2, float(b.conf[0])])
|
|
return boxes
|
|
|
|
|
|
def _ping_yolo_daemon():
|
|
try: return urllib.request.urlopen(f'http://127.0.0.1:{_YOLO_PORT}/ping', timeout=0.1).read() == b'ui_detect_yolo'
|
|
except Exception: return False
|
|
|
|
def _yolo(image_path, conf=0.25):
|
|
"""YOLO检测 → list of [x1,y1,x2,y2,conf]；默认模型走跨进程daemon cache，失败回退本地"""
|
|
if not _ping_yolo_daemon():
|
|
kw = {'creationflags': getattr(subprocess, 'CREATE_NO_WINDOW', 0)} if sys.platform == 'win32' else {}
|
|
subprocess.Popen([sys.executable, __file__, '--yolo-daemon'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, **kw)
|
|
for _ in range(15):
|
|
if _ping_yolo_daemon(): break
|
|
time.sleep(0.5)
|
|
try:
|
|
data = json.dumps({'path': str(image_path), 'conf': conf}).encode('utf-8')
|
|
req = urllib.request.Request(f'http://127.0.0.1:{_YOLO_PORT}/yolo', data=data, headers={'Content-Type': 'application/json'})
|
|
return json.loads(urllib.request.urlopen(req, timeout=3).read().decode('utf-8'))['boxes']
|
|
except Exception: return _yolo_local(image_path, conf)
|
|
|
|
def _ocr_full(image_path):
|
|
"""全图OCR → list of [x1,y1,x2,y2,text,conf]"""
|
|
if not _ocr: return []
|
|
result, _ = _ocr(image_path)
|
|
if not result: return []
|
|
out = []
|
|
for bbox, text, conf in result:
|
|
xs = [p[0] for p in bbox]; ys = [p[1] for p in bbox]
|
|
out.append([int(min(xs)), int(min(ys)), int(max(xs)), int(max(ys)), text, conf])
|
|
return out
|
|
|
|
def _ocr_crops_batch(img, yolo_boxes):
|
|
"""批量OCR：将所有YOLO框crop垂直拼接为一张图，一次OCR，按y坐标映射回各box → {box_idx: text}"""
|
|
if not _ocr or not yolo_boxes: return {}
|
|
|
|
crops, offsets = [], [] # offsets: [(y_off, orig_x1, orig_y1, box_idx)]
|
|
max_w, y_cursor = 0, 0
|
|
for idx, (x1, y1, x2, y2, _) in enumerate(yolo_boxes):
|
|
crop = img.crop((x1, y1, x2, y2))
|
|
w, h = crop.size
|
|
max_w = max(max_w, w)
|
|
crops.append(crop)
|
|
offsets.append((y_cursor, x1, y1, idx))
|
|
y_cursor += h
|
|
if max_w == 0: return {}
|
|
# 垂直拼接
|
|
stitched = Image.new('RGB', (max_w, y_cursor), (255, 255, 255))
|
|
for i, crop in enumerate(crops):
|
|
stitched.paste(crop, (0, offsets[i][0]))
|
|
|
|
result, _ = _ocr(np.array(stitched))
|
|
if not result: return {}
|
|
# 映射：OCR框中心y → 归属的crop
|
|
labels = {}
|
|
for bbox, text, _ in result:
|
|
cy = sum(p[1] for p in bbox) / len(bbox)
|
|
for y_off, ox1, oy1, idx in offsets:
|
|
h = yolo_boxes[idx][3] - yolo_boxes[idx][1]
|
|
if y_off <= cy < y_off + h:
|
|
old = labels.get(idx)
|
|
labels[idx] = (old + ' ' + text) if old else text
|
|
break
|
|
return labels
|
|
|
|
def _iou(a, b):
|
|
"""计算两个bbox的交集占b面积的比例(包含率)"""
|
|
x1, y1, x2, y2 = max(a[0],b[0]), max(a[1],b[1]), min(a[2],b[2]), min(a[3],b[3])
|
|
inter = max(0, x2-x1) * max(0, y2-y1)
|
|
area_b = (b[2]-b[0]) * (b[3]-b[1])
|
|
return inter / area_b if area_b > 0 else 0
|
|
|
|
def detect(image_path, mode='match', conf=0.25, iou_thresh=0.5):
|
|
"""
|
|
统一检测入口，返回元素列表:
|
|
[{'bbox':[x1,y1,x2,y2], 'type':'icon'|'text', 'label':str|None, 'confidence':float}]
|
|
mode: 'match' = YOLO+全图OCR空间匹配(推荐, 快) | 'crop' = YOLO+拼接OCR(备选, 更精确)
|
|
支持 image_path: str 路径 或 PIL.Image 对象
|
|
"""
|
|
# 归一化：PIL Image → 临时文件
|
|
if isinstance(image_path, Image.Image):
|
|
import tempfile, os
|
|
tmp = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
|
|
image_path.save(tmp.name)
|
|
image_path = tmp.name
|
|
img = Image.open(image_path)
|
|
|
|
yolo_boxes = _yolo(image_path, conf)
|
|
elements = []
|
|
|
|
if mode == 'crop':
|
|
# YOLO元素批量OCR（拼接一次推理）
|
|
labels_map = _ocr_crops_batch(img, yolo_boxes)
|
|
for idx, (x1, y1, x2, y2, c) in enumerate(yolo_boxes):
|
|
elements.append({'bbox': [x1,y1,x2,y2], 'type': 'icon', 'label': labels_map.get(idx), 'confidence': c})
|
|
# 补充：全图OCR找未被覆盖的纯文本
|
|
for ox1, oy1, ox2, oy2, text, oc in _ocr_full(image_path):
|
|
covered = any(_iou([x1,y1,x2,y2,_,__], [ox1,oy1,ox2,oy2]) > iou_thresh
|
|
for x1,y1,x2,y2,_,__ in [(b[0],b[1],b[2],b[3],0,0) for b in yolo_boxes])
|
|
if not covered:
|
|
elements.append({'bbox': [ox1,oy1,ox2,oy2], 'type': 'text', 'label': text, 'confidence': oc})
|
|
|
|
elif mode == 'match':
|
|
ocr_items = _ocr_full(image_path)
|
|
matched_ocr = set()
|
|
for x1, y1, x2, y2, c in yolo_boxes:
|
|
label = None
|
|
for i, (ox1, oy1, ox2, oy2, text, oc) in enumerate(ocr_items):
|
|
if _iou([x1,y1,x2,y2], [ox1,oy1,ox2,oy2]) > iou_thresh:
|
|
label = text; matched_ocr.add(i); break
|
|
elements.append({'bbox': [x1,y1,x2,y2], 'type': 'icon', 'label': label, 'confidence': c})
|
|
# 未匹配的OCR作为独立text元素
|
|
for i, (ox1, oy1, ox2, oy2, text, oc) in enumerate(ocr_items):
|
|
if i not in matched_ocr:
|
|
elements.append({'bbox': [ox1,oy1,ox2,oy2], 'type': 'text', 'label': text, 'confidence': oc})
|
|
#if [x for x in elements if x['label'] is None]: print('[TIPS] crop grid + VLM to identify target no text icon if needed')
|
|
print('[TIPS] UI DETECT contains OCR, no need to run OCR again!')
|
|
return elements
|
|
|
|
def visualize_for_debug(image_path, elements, output_path=None):
|
|
"""Only use when user wants to DEBUG!"""
|
|
from PIL import ImageFont
|
|
img = Image.open(image_path)
|
|
draw = ImageDraw.Draw(img)
|
|
try:
|
|
font = ImageFont.truetype("msyh.ttc", 14)
|
|
except:
|
|
font = ImageFont.load_default()
|
|
for el in elements:
|
|
x1, y1, x2, y2 = el['bbox']
|
|
color = 'red' if el['type'] == 'icon' else 'blue'
|
|
draw.rectangle([x1, y1, x2, y2], outline=color, width=2)
|
|
tag = el.get('label') or f"{el['confidence']:.2f}"
|
|
draw.text((x1, y1-16), tag[:15], fill=color, font=font)
|
|
if output_path: img.save(output_path)
|
|
return img
|
|
|
|
def _serve_yolo_daemon():
|
|
from http.server import BaseHTTPRequestHandler, HTTPServer
|
|
class H(BaseHTTPRequestHandler):
|
|
def log_message(self, *args): pass
|
|
def handle_one_request(self): self.server.last=time.time(); return super().handle_one_request()
|
|
def do_GET(self):
|
|
if self.path == '/ping':
|
|
self.send_response(200); self.end_headers(); self.wfile.write(b'ui_detect_yolo')
|
|
else:
|
|
self.send_response(404); self.end_headers()
|
|
def do_POST(self):
|
|
try:
|
|
d = json.loads(self.rfile.read(int(self.headers.get('Content-Length', 0))))
|
|
body = json.dumps({'boxes': _yolo_local(d['path'], d.get('conf', 0.25))}).encode('utf-8')
|
|
self.send_response(200); self.end_headers(); self.wfile.write(body)
|
|
except Exception as e:
|
|
body = json.dumps({'error': repr(e)}).encode('utf-8')
|
|
self.send_response(500); self.end_headers(); self.wfile.write(body)
|
|
s=HTTPServer(('127.0.0.1', _YOLO_PORT), H); s.timeout=60; s.last=time.time()
|
|
while time.time()-s.last < 3600: s.handle_request()
|
|
|
|
if __name__ == '__main__' and '--yolo-daemon' in sys.argv:
|
|
_serve_yolo_daemon()
|