#!/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()