207 lines
9.2 KiB
Python
207 lines
9.2 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匹配(推荐,GPU约1.2s,无文字图标label=None可VLM保底) | crop=拼接crop OCR(备选,更精确)
|
||
依赖: ultralytics, rapidocr-onnxruntime, pillow, numpy；GPU另需兼容的onnxruntime-gpu及CUDA/cuDNN
|
||
OCR: 自动优先 CUDAExecutionProvider(本机3090实测1.1s, CPU约6s)，不可用回退CPU
|
||
权重: temp/weights/icon_detect/model.pt (OmniParser-v2.0 icon_detect, 40MB)
|
||
"""
|
||
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(det_use_cuda=True, rec_use_cuda=True, cls_use_cuda=True)
|
||
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()
|