|
|
@@ -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}")
|
|
|
|