chore: import upstream snapshot with attribution

memory/ui_detect.py

@@ -0,0 +1,204 @@
+#!/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()