Procházet zdrojové kódy

feat(mineru_adapter): add MinerUVLLayoutDetector for enhanced layout detection

- Introduced the MinerUVLLayoutDetector class to improve layout detection capabilities using the MinerU VL model.
- Updated the __init__.py file to include the new detector in the module's exports.
- Enhanced block type mapping and initialization processes for better integration with existing OCR tools.
zhch158_admin před 2 dny
rodič
revize
c2d08236cf

+ 2 - 0
ocr_tools/universal_doc_parser/models/adapters/__init__.py

@@ -27,6 +27,7 @@ try:
     from .mineru_adapter import (
         MinerUPreprocessor,
         MinerULayoutDetector,
+        MinerUVLLayoutDetector,
         MinerUVLRecognizer,
         MinerUOCRRecognizer
     )
@@ -60,6 +61,7 @@ if MINERU_AVAILABLE:
     __all__.extend([
         'MinerUPreprocessor',
         'MinerULayoutDetector',
+        'MinerUVLLayoutDetector',
         'MinerUVLRecognizer',
         'MinerUOCRRecognizer',
         'MinerUWiredTableRecognizer',

+ 121 - 1
ocr_tools/universal_doc_parser/models/adapters/mineru_adapter.py

@@ -160,6 +160,124 @@ class MinerULayoutDetector(BaseLayoutDetector):
             print(f"❌ Layout detection failed: {e}")
             return []
 
+class MinerUVLLayoutDetector(BaseLayoutDetector):
+    """MinerU-VL 版式检测适配器(基于 mineru_vl_utils 的 layout_detect)"""
+
+    # 类别映射:MinerU-VL BlockType → EnhancedDocPipeline 类别体系
+    # 参考:models/adapters/docling_layout_adapter.py 中 CATEGORY_MAP 的用法
+    CATEGORY_MAP = {
+        # 直接可用/同名
+        "text": "text",
+        "title": "title",
+        "header": "header",
+        "footer": "footer",
+        "page_number": "page_number",
+        "page_footnote": "page_footnote",
+        "aside_text": "aside_text",
+        "ref_text": "ref_text",
+        "code": "code",
+        "algorithm": "algorithm",
+        "list": "text",  # 列表块按文本处理(pipeline TEXT_CATEGORIES)
+        "phonetic": "text",
+
+        # 表格/图片/公式等映射到 pipeline 体系
+        "table": "table_body",
+        "table_caption": "table_caption",
+        "table_footnote": "table_footnote",
+        "image": "image_body",
+        "image_caption": "image_caption",
+        "image_footnote": "image_footnote",
+        "equation": "interline_equation",
+        "equation_block": "interline_equation",
+
+        # 兜底
+        "unknown": "text",
+    }
+
+    def __init__(self, config: Dict[str, Any]):
+        super().__init__(config)
+        if not MINERU_AVAILABLE:
+            raise ImportError("MinerU components not available")
+
+        self.vlm_model = None
+
+    def initialize(self):
+        """初始化 MinerU-VL 模型(用于 layout_detect)"""
+        try:
+            backend = self.config.get("backend", "http-client")
+            server_url = self.config.get("server_url")
+            model_params = self.config.get("model_params", {})
+
+            self.vlm_model = VLMModelSingleton().get_model(
+                backend=backend,
+                model_path=None,
+                server_url=server_url,
+                **model_params,
+            )
+            print(f"✅ MinerU-VL layout detector initialized: {backend}")
+        except Exception as e:
+            print(f"❌ Failed to initialize MinerU-VL layout detector: {e}")
+            raise
+
+    def cleanup(self):
+        pass
+
+    @staticmethod
+    def _normalize_block_type(block_type: str) -> str:
+        return (block_type or "").strip().lower()
+
+    def _map_block_type(self, block_type: str) -> str:
+        """将 MinerU-VL ContentBlock.type 统一映射到 pipeline category"""
+        t = self._normalize_block_type(block_type)
+        return self.CATEGORY_MAP.get(t, "text")
+
+    def _detect_raw(
+        self,
+        image: Union[np.ndarray, Image.Image],
+        ocr_spans: Optional[List[Dict[str, Any]]] = None,
+    ) -> List[Dict[str, Any]]:
+        """版式检测(原始检测,不包含后处理)"""
+        if self.vlm_model is None:
+            raise RuntimeError("MinerU-VL model not initialized")
+
+        # 转换为PIL图像
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+
+        w, h = image.size
+        try:
+            blocks = self.vlm_model.layout_detect(image=image)
+        except TypeError:
+            # 兼容不同签名(有的实现不使用关键字参数)
+            blocks = self.vlm_model.layout_detect(image)
+
+        formatted_results: List[Dict[str, Any]] = []
+        for block in blocks or []:
+            try:
+                block_type = block.get("type") if isinstance(block, dict) else getattr(block, "type", None)
+                bbox_norm = block.get("bbox") if isinstance(block, dict) else getattr(block, "bbox", None)
+                angle = block.get("angle") if isinstance(block, dict) else getattr(block, "angle", None)
+
+                if not bbox_norm or len(bbox_norm) < 4:
+                    continue
+
+                x1n, y1n, x2n, y2n = bbox_norm[:4]
+                bbox = [round((x1n) * w, 0), round((y1n) * h, 0), round((x2n) * w, 0), round((y2n) * h, 0)]
+
+                formatted_results.append(
+                    {
+                        "category": self._map_block_type(str(block_type) if block_type is not None else ""),
+                        "bbox": bbox,
+                        "confidence": 1.0,
+                        "angle": angle,
+                        "raw": block,
+                    }
+                )
+            except Exception as e:
+                logger.debug(f"Skip invalid layout block: {e}")
+
+        return formatted_results
+
 class MinerUVLRecognizer(BaseVLRecognizer):
     """MinerU VL识别适配器"""
     
@@ -343,7 +461,8 @@ class MinerUVLRecognizer(BaseVLRecognizer):
                     results.append({
                         'html': content,
                         'markdown': self._html_to_markdown(content),
-                        'cells': self._extract_cells_from_html(content) if kwargs.get('return_cells_coordinate', False) else []
+                        # MinerU-VL 当前返回的是 HTML/结构化文本,单元格坐标由下游 matcher 负责或由 wired-table 分支产出
+                        'cells': []
                     })
                 else:
                     results.append({'html': '', 'markdown': '', 'cells': []})
@@ -549,6 +668,7 @@ class MinerUOCRRecognizer(BaseOCRRecognizer):
 __all__ = [
     'MinerUPreprocessor',
     'MinerULayoutDetector', 
+    'MinerUVLLayoutDetector',
     'MinerUVLRecognizer',
     'MinerUOCRRecognizer'
 ]