Răsfoiți Sursa

feat: 添加 DiT 适配器的可选导入和布局检测支持

- 在适配器模块中添加对 DiT (Document Image Transformer) 的可选导入,确保在缺少依赖时不会导致导入错误。
- 更新 `get_layout_detector` 函数,支持 DiT 布局检测,并在不可用时抛出明确的导入错误。
- 将 DiT 适配器添加到导出列表中,以便在可用时使用。
zhch158_admin 1 săptămână în urmă
părinte
comite
1cdd879991
1 a modificat fișierele cu 15 adăugiri și 0 ștergeri
  1. 15 0
      ocr_tools/universal_doc_parser/models/adapters/__init__.py

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

@@ -15,6 +15,13 @@ from .paddle_layout_detector import PaddleLayoutDetector
 from .paddle_vl_adapter import PaddleVLRecognizer
 
 from .docling_layout_adapter import DoclingLayoutDetector
+
+# 可选导入 DiT 适配器
+try:
+    from .dit_layout_adapter import DitLayoutDetector
+    DIT_AVAILABLE = True
+except ImportError:
+    DIT_AVAILABLE = False
 # 可选导入 MinerU 适配器
 try:
     from .mineru_adapter import (
@@ -44,6 +51,10 @@ __all__ = [
     'DoclingLayoutDetector',
 ]
 
+# 如果 DiT 可用,添加到导出列表
+if DIT_AVAILABLE:
+    __all__.append('DitLayoutDetector')
+
 # 如果 MinerU 可用,添加到导出列表
 if MINERU_AVAILABLE:
     __all__.extend([
@@ -75,6 +86,10 @@ def get_layout_detector(config: dict):
         return MinerULayoutDetector(config)
     elif module == 'docling':
         return DoclingLayoutDetector(config)
+    elif module == 'dit':
+        if not DIT_AVAILABLE:
+            raise ImportError("DiT adapter not available. Please ensure detectron2 and ditod are installed.")
+        return DitLayoutDetector(config)
     else:
         raise ValueError(f"Unknown layout detection module: {module}")