浏览代码

feat: add support for Docling layout detector in ModelFactory, enhancing layout detection options and improving error handling during initialization

zhch158_admin 1 周之前
父节点
当前提交
22f3d967cf
共有 1 个文件被更改,包括 11 次插入4 次删除
  1. 11 4
      zhch/universal_doc_parser/core/model_factory.py

+ 11 - 4
zhch/universal_doc_parser/core/model_factory.py

@@ -36,14 +36,20 @@ class ModelFactory:
         if module_name == 'paddle':
             from models.adapters import PaddleLayoutDetector
             detector = PaddleLayoutDetector(config)
+        elif module_name == 'docling':
+            from models.adapters import DoclingLayoutDetector
+            detector = DoclingLayoutDetector(config)
         elif module_name == 'mineru':
             from models.adapters import MinerULayoutDetector
             detector = MinerULayoutDetector(config)
         else:
             raise ValueError(f"Unknown layout detector module: {module_name}")
-
-        detector.initialize()
-        return detector
+        try:
+            detector.initialize()
+            return detector
+        except Exception as e:
+            print(f"Failed to initialize layout detector: {e}")
+            raise
     
     @classmethod
     def create_vl_recognizer(cls, config: Dict[str, Any]) -> BaseVLRecognizer:
@@ -88,4 +94,5 @@ ModelFactory.register_adapter('layout_detector', 'mineru_adapter', 'MinerULayout
 ModelFactory.register_adapter('vl_recognizer', 'mineru_adapter', 'MinerUVLRecognizer')
 ModelFactory.register_adapter('ocr_recognizer', 'mineru_adapter', 'MinerUOCRRecognizer')
 
-ModelFactory.register_adapter('layout_detector', 'paddle_adapter', 'PaddleLayoutDetector')
+ModelFactory.register_adapter('layout_detector', 'paddle_adapter', 'PaddleLayoutDetector')
+ModelFactory.register_adapter('layout_detector', 'docling_adapter', 'DoclingLayoutDetector')