import base64 import io import os from typing import Any, Dict, List, Literal, Optional, Tuple, Union # 在任何 matplotlib 导入之前强制设置非交互式后端 # 避免在 macOS 自动化版本中弹出 Quartz 窗口导致阴塞 import matplotlib matplotlib.use('Agg') import cv2 import easyocr import numpy as np import torch from PIL import Image from PIL.Image import Image as ImageType from supervision.detection.core import Detections from supervision.draw.color import Color, ColorPalette from torchvision.ops import box_convert from torchvision.transforms import ToPILImage from transformers import AutoModelForCausalLM, AutoProcessor from transformers.image_utils import load_image from ultralytics import YOLO class EndpointHandler: """ OmniParser 推理处理器 Args: model_dir: 模型目录路径 enable_ocr: 是否启用 OCR 文字识别（默认 True） enable_caption: 是否启用图标描述生成（默认 True） enable_yolo: 是否启用 YOLO 图标检测（默认 True） ocr_languages: OCR 语言列表（默认 ["ch_sim", "en"] 中英文） """ def __init__( self, model_dir: str = os.path.dirname(os.path.abspath(__file__)), enable_ocr: bool = True, enable_caption: bool = True, enable_yolo: bool = True, ocr_languages: Optional[List[str]] = None, ) -> None: # 保存开关状态 self.enable_ocr = enable_ocr self.enable_caption = enable_caption self.enable_yolo = enable_yolo # 选择计算设备 if torch.cuda.is_available(): self.device = torch.device("cuda") device_name = torch.cuda.get_device_name(0) print(f"🚀 使用设备: GPU ({device_name})") elif torch.backends.mps.is_available(): self.device = torch.device("mps") print("🚀 使用设备: GPU (Apple MPS)") else: self.device = torch.device("cpu") print("⚠️ 使用设备: CPU (推理较慢，建议安装 CUDA 版 PyTorch)") # 显示启用的模块 modules = [] if enable_yolo: modules.append("YOLO检测") if enable_ocr: modules.append("OCR文字") if enable_caption: modules.append("图标描述") print(f"📦 启用模块: {', '.join(modules) if modules else '无'}") # YOLO 图标检测模型 self.yolo = None if enable_yolo: self.yolo = YOLO(f"{model_dir}/icon_detect/model.pt") self.yolo.to(self.device) # Florence-2 图标描述模型 self.processor = None self.model = None if enable_caption: self.processor = AutoProcessor.from_pretrained( "microsoft/Florence-2-base", trust_remote_code=True ) self.model = AutoModelForCausalLM.from_pretrained( f"{model_dir}/icon_caption", torch_dtype=torch.float16 if self.device.type != "cpu" else torch.float32, trust_remote_code=True, ).to(self.device) # EasyOCR 文字识别模型 self.ocr = None if enable_ocr: langs = ocr_languages or ["ch_sim", "en"] model_storage_directory = os.path.join(model_dir, "easyocr_models") print(f"📄 EasyOCR 模型目录: {model_storage_directory}") # EasyOCR recognition.py 内部 DataLoader 将 pin_memory 写死为 True # 但 MPS 设备不支持 pin_memory，通过 monkey-patch 消除警告 if self.device.type == "mps": try: import easyocr.recognition as _easyocr_recognition import torch.utils.data as _torch_data _original_DataLoader = _torch_data.DataLoader def _patched_DataLoader(*args, **kwargs): kwargs['pin_memory'] = False return _original_DataLoader(*args, **kwargs) _easyocr_recognition.DataLoader = _patched_DataLoader except Exception: pass # 补丁失败不影响功能 self.ocr = easyocr.Reader( langs, gpu=self.device.type == "cuda", model_storage_directory=model_storage_directory, download_enabled=True ) # box annotator self.annotator = BoxAnnotator() def __call__(self, data: Dict[str, Any]) -> Any: data = data.pop("inputs") # read image image = load_image(data["image"]) # 1. OCR Step（如果启用） ocr_texts, ocr_bboxes = [], [] if self.enable_ocr and self.ocr is not None: ocr_texts, ocr_bboxes = self.check_ocr_bboxes( image, out_format="xyxy", ocr_kwargs={"text_threshold": 0.4}, ) # 2. YOLO + Caption Step annotated_image, filtered_bboxes_out = self.get_som_labeled_img( image, image_size=data.get("image_size", None), ocr_texts=ocr_texts, ocr_bboxes=ocr_bboxes, bbox_threshold=data.get("bbox_threshold", 0.05), iou_threshold=data.get("iou_threshold", 0.5), ) return { "image": annotated_image, "bboxes": filtered_bboxes_out, } def check_ocr_bboxes( self, image: ImageType, out_format: Literal["xywh", "xyxy"] = "xywh", ocr_kwargs: Optional[Dict[str, Any]] = {}, ) -> Tuple[List[str], List[List[int]]]: # 🔧 修复点 1：RBGA -> RGBA if image.mode == "RGBA": image = image.convert("RGB") result = self.ocr.readtext(np.array(image), **ocr_kwargs) # type: ignore texts = [str(item[1]) for item in result] bboxes = [ self.coordinates_to_bbox(item[0], format=out_format) for item in result ] return (texts, bboxes) @staticmethod def coordinates_to_bbox( coordinates: np.ndarray, format: Literal["xywh", "xyxy"] = "xywh" ) -> List[int]: if format == "xywh": return [ int(coordinates[0][0]), int(coordinates[0][1]), int(coordinates[2][0] - coordinates[0][0]), int(coordinates[2][1] - coordinates[0][1]), ] elif format == "xyxy": return [ int(coordinates[0][0]), int(coordinates[0][1]), int(coordinates[2][0]), int(coordinates[2][1]), ] else: raise ValueError(f"Unsupported format: {format}") @staticmethod def bbox_area(bbox: List[int], w: int, h: int) -> int: bbox = [bbox[0] * w, bbox[1] * h, bbox[2] * w, bbox[3] * h] return (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]) @staticmethod def remove_bbox_overlap( xyxy_bboxes: List[Dict[str, Any]], ocr_bboxes: Optional[List[Dict[str, Any]]] = None, iou_threshold: Optional[float] = 0.7, ) -> List[Dict[str, Any]]: filtered_bboxes = [] if ocr_bboxes is not None: filtered_bboxes.extend(ocr_bboxes) for i, bbox_outter in enumerate(xyxy_bboxes): bbox_left = bbox_outter["bbox"] valid_bbox = True for j, bbox_inner in enumerate(xyxy_bboxes): if i == j: continue bbox_right = bbox_inner["bbox"] if ( intersection_over_union( bbox_left, bbox_right, ) > iou_threshold # type: ignore ) and (area(bbox_left) > area(bbox_right)): valid_bbox = False break if valid_bbox is False: continue if ocr_bboxes is None: filtered_bboxes.append(bbox_outter) continue box_added = False ocr_labels = [] for ocr_bbox in ocr_bboxes: if not box_added: bbox_right = ocr_bbox["bbox"] # 检查是否重叠（合并 Icon 和 OCR 文字） if overlap(bbox_right, bbox_left): try: ocr_labels.append(ocr_bbox["content"]) if ocr_bbox in filtered_bboxes: filtered_bboxes.remove(ocr_bbox) except Exception as e: self.logger.error(f"Error removing bbox overlap: {e}") continue elif overlap(bbox_left, bbox_right): box_added = True break if not box_added: filtered_bboxes.append( { "type": "icon", "bbox": bbox_outter["bbox"], "interactivity": True, "content": " ".join(ocr_labels) if ocr_labels else None, } ) return filtered_bboxes def get_som_labeled_img( self, image: ImageType, image_size: Optional[Dict[Literal["w", "h"], int]] = None, ocr_texts: Optional[List[str]] = None, ocr_bboxes: Optional[List[List[int]]] = None, bbox_threshold: float = 0.05, iou_threshold: Optional[float] = None, caption_prompt: Optional[str] = None, caption_batch_size: int = 64, ) -> Tuple[str, List[Dict[str, Any]]]: if image.mode == "RGBA": image = image.convert("RGB") w, h = image.size image_np = np.asarray(image) # YOLO 检测（如果启用） xyxy_bboxes_raw = torch.tensor([]) # 默认空 if self.enable_yolo and self.yolo is not None: if image_size is None: imgsz = [h, w] else: imgsz = [image_size.get("h", h), image_size.get("w", w)] out = self.yolo.predict( image, imgsz=imgsz, conf=bbox_threshold, iou=iou_threshold or 0.7, verbose=False, )[0] if out.boxes is not None: xyxy_bboxes_raw = out.boxes.xyxy xyxy_bboxes_raw = xyxy_bboxes_raw / torch.Tensor([w, h, w, h]).to(xyxy_bboxes_raw.device) # 处理 OCR 检测框 ocr_bboxes_normalized = [] if ocr_bboxes: ocr_bboxes_tensor = torch.tensor(ocr_bboxes) / torch.Tensor([w, h, w, h]) ocr_bboxes_normalized = ocr_bboxes_tensor.tolist() ocr_bbox_dicts = [ { "type": "text", "bbox": bbox, "interactivity": False, "content": text, "source": "box_ocr_content_ocr", } for bbox, text in zip(ocr_bboxes_normalized, ocr_texts or []) if self.bbox_area(bbox, w, h) > 0 ] # 处理 YOLO 检测框 yolo_bbox_dicts = [ { "type": "icon", "bbox": bbox, "interactivity": True, "content": None, "source": "box_yolo_content_yolo", } for bbox in xyxy_bboxes_raw.tolist() if self.bbox_area(bbox, w, h) > 0 ] filtered_bboxes = self.remove_bbox_overlap( xyxy_bboxes=yolo_bbox_dicts, ocr_bboxes=ocr_bbox_dicts if ocr_bbox_dicts else None, iou_threshold=iou_threshold or 0.7, ) filtered_bboxes_out = sorted( filtered_bboxes, key=lambda x: x["content"] is None ) starting_idx = next( ( idx for idx, bbox in enumerate(filtered_bboxes_out) if bbox["content"] is None ), -1, ) filtered_bboxes = torch.tensor([box["bbox"] for box in filtered_bboxes_out]) # 如果启用了 Caption 且有需要描述的图标 if starting_idx != -1 and self.enable_caption and self.model is not None: non_ocr_bboxes = filtered_bboxes[starting_idx:] bbox_images = [] for _, coordinates in enumerate(non_ocr_bboxes): try: xmin, xmax = ( int(coordinates[0] * image_np.shape[1]), int(coordinates[2] * image_np.shape[1]), ) ymin, ymax = ( int(coordinates[1] * image_np.shape[0]), int(coordinates[3] * image_np.shape[0]), ) cropped_image = image_np[ymin:ymax, xmin:xmax, :] cropped_image = cv2.resize(cropped_image, (64, 64)) bbox_images.append(ToPILImage()(cropped_image)) except Exception as e: self.logger.error(f"Error cropping bbox: {e}") continue if caption_prompt is None: caption_prompt = "" captions = [] for idx in range(0, len(bbox_images), caption_batch_size): # type: ignore batch = bbox_images[idx : idx + caption_batch_size] # type: ignore if not batch: break inputs = self.processor( images=batch, text=[caption_prompt] * len(batch), return_tensors="pt", do_resize=False, ) if self.device.type in {"cuda", "mps"}: inputs = inputs.to(device=self.device, dtype=torch.float16) with torch.inference_mode(): generated_ids = self.model.generate( input_ids=inputs["input_ids"], pixel_values=inputs["pixel_values"], max_new_tokens=20, num_beams=1, do_sample=False, early_stopping=False, ) generated_texts = self.processor.batch_decode( generated_ids, skip_special_tokens=True ) captions.extend([text.strip() for text in generated_texts]) ocr_texts = [f"Text Box ID {idx}: {text}" for idx, text in enumerate(ocr_texts)] # type: ignore for _, bbox in enumerate(filtered_bboxes_out): if bbox["content"] is None and captions: bbox["content"] = captions.pop(0) filtered_bboxes = box_convert( boxes=filtered_bboxes, in_fmt="xyxy", out_fmt="cxcywh" ) annotated_image = image_np.copy() bboxes_annotate = filtered_bboxes * torch.Tensor([w, h, w, h]) xyxy_annotate = box_convert( bboxes_annotate, in_fmt="cxcywh", out_fmt="xyxy" ).numpy() detections = Detections(xyxy=xyxy_annotate) labels = [str(idx) for idx in range(bboxes_annotate.shape[0])] annotated_image = self.annotator.annotate( scene=annotated_image, detections=detections, labels=labels, image_size=(w, h), ) assert w == annotated_image.shape[1] and h == annotated_image.shape[0] out_image = Image.fromarray(annotated_image) out_buffer = io.BytesIO() out_image.save(out_buffer, format="PNG") encoded_image = base64.b64encode(out_buffer.getvalue()).decode("ascii") return encoded_image, filtered_bboxes_out def area(bbox: List[int]) -> int: return (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]) def intersection_area(bbox_left: List[int], bbox_right: List[int]) -> int: # 计算两个 bbox 的交集面积 # 交集的左边界取两者的最大值，右边界取两者的最小值 x_overlap = max(0, min(bbox_left[2], bbox_right[2]) - max(bbox_left[0], bbox_right[0])) y_overlap = max(0, min(bbox_left[3], bbox_right[3]) - max(bbox_left[1], bbox_right[1])) return x_overlap * y_overlap def intersection_over_union(bbox_left: List[int], bbox_right: List[int]) -> float: intersection = intersection_area(bbox_left, bbox_right) bbox_left_area = area(bbox_left) bbox_right_area = area(bbox_right) union = bbox_left_area + bbox_right_area - intersection + 1e-6 ratio_left, ratio_right = 0, 0 if bbox_left_area > 0 and bbox_right_area > 0: ratio_left = intersection / bbox_left_area ratio_right = intersection / bbox_right_area return max(intersection / union, ratio_left, ratio_right) def overlap(bbox_left: List[int], bbox_right: List[int]) -> bool: intersection = intersection_area(bbox_left, bbox_right) ratio_left = intersection / area(bbox_left) # 🔧 修复点 3：从 0.80 降低到 0.50 # 作用：只要 OCR 文字框有一半在 YOLO 图标框内，就认为它们是一体的， # 避免因为框对齐不准导致 Icon 框被丢弃或产生两个框 return ratio_left > 0.50 class BoxAnnotator: def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.DEFAULT, # type: ignore thickness: int = 3, text_color: Color = Color.BLACK, # type: ignore text_scale: float = 0.5, text_thickness: int = 2, text_padding: int = 10, avoid_overlap: bool = True, ): self.color: Union[Color, ColorPalette] = color self.thickness: int = thickness self.text_color: Color = text_color self.text_scale: float = text_scale self.text_thickness: int = text_thickness self.text_padding: int = text_padding self.avoid_overlap: bool = avoid_overlap def annotate( self, scene: np.ndarray, detections: Detections, labels: Optional[List[str]] = None, skip_label: bool = False, image_size: Optional[Tuple[int, int]] = None, ) -> np.ndarray: font = cv2.FONT_HERSHEY_SIMPLEX for i in range(len(detections)): x1, y1, x2, y2 = detections.xyxy[i].astype(int) class_id = ( detections.class_id[i] if detections.class_id is not None else None ) idx = class_id if class_id is not None else i color = ( self.color.by_idx(idx) if isinstance(self.color, ColorPalette) else self.color ) cv2.rectangle( img=scene, pt1=(x1, y1), pt2=(x2, y2), color=color.as_bgr(), thickness=self.thickness, ) if skip_label: continue text = ( f"{class_id}" if (labels is None or len(detections) != len(labels)) else labels[i] ) text_width, text_height = cv2.getTextSize( text=text, fontFace=font, fontScale=self.text_scale, thickness=self.text_thickness, )[0] if not self.avoid_overlap: text_x = x1 + self.text_padding text_y = y1 - self.text_padding text_background_x1 = x1 text_background_y1 = y1 - 2 * self.text_padding - text_height text_background_x2 = x1 + 2 * self.text_padding + text_width text_background_y2 = y1 else: ( text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2, ) = self.get_optimal_label_pos( self.text_padding, text_width, text_height, x1, y1, x2, y2, detections, image_size, ) cv2.rectangle( img=scene, pt1=(text_background_x1, text_background_y1), pt2=(text_background_x2, text_background_y2), color=color.as_bgr(), thickness=cv2.FILLED, ) box_color = color.as_rgb() luminance = ( 0.299 * box_color[0] + 0.587 * box_color[1] + 0.114 * box_color[2] ) text_color = (0, 0, 0) if luminance > 160 else (255, 255, 255) cv2.putText( img=scene, text=text, org=(text_x, text_y), fontFace=font, fontScale=self.text_scale, color=text_color, thickness=self.text_thickness, lineType=cv2.LINE_AA, ) return scene @staticmethod def get_optimal_label_pos( text_padding, text_width, text_height, x1, y1, x2, y2, detections, image_size ): def get_is_overlap( detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size, ): is_overlap = False for i in range(len(detections)): detection = detections.xyxy[i].astype(int) if ( intersection_over_union( [ text_background_x1, text_background_y1, text_background_x2, text_background_y2, ], detection, ) > 0.3 ): is_overlap = True break if ( text_background_x1 < 0 or text_background_x2 > image_size[0] or text_background_y1 < 0 or text_background_y2 > image_size[1] ): is_overlap = True return is_overlap text_x = x1 + text_padding text_y = y1 - text_padding text_background_x1 = x1 text_background_y1 = y1 - 2 * text_padding - text_height text_background_x2 = x1 + 2 * text_padding + text_width text_background_y2 = y1 is_overlap = get_is_overlap( detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size, ) if not is_overlap: return ( text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2, ) text_x = x1 - text_padding - text_width text_y = y1 + text_padding + text_height text_background_x1 = x1 - 2 * text_padding - text_width text_background_y1 = y1 text_background_x2 = x1 text_background_y2 = y1 + 2 * text_padding + text_height is_overlap = get_is_overlap( detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size, ) if not is_overlap: return ( text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2, ) text_x = x2 + text_padding text_y = y1 + text_padding + text_height text_background_x1 = x2 text_background_y1 = y1 text_background_x2 = x2 + 2 * text_padding + text_width text_background_y2 = y1 + 2 * text_padding + text_height is_overlap = get_is_overlap( detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size, ) if not is_overlap: return ( text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2, ) text_x = x2 - text_padding - text_width text_y = y1 - text_padding text_background_x1 = x2 - 2 * text_padding - text_width text_background_y1 = y1 - 2 * text_padding - text_height text_background_x2 = x2 text_background_y2 = y1 is_overlap = get_is_overlap( detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size, ) if not is_overlap: return ( text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2, ) return ( text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2, )