__init__.py 2.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114
  1. """
  2. 模型适配器模块
  3. 提供统一的接口适配不同的模型后端
  4. """
  5. from .base import (
  6. BaseAdapter,
  7. BasePreprocessor,
  8. BaseLayoutDetector,
  9. BaseVLRecognizer,
  10. BaseOCRRecognizer
  11. )
  12. from .paddle_layout_detector import PaddleLayoutDetector
  13. from .paddle_vl_adapter import PaddleVLRecognizer
  14. from .docling_layout_adapter import DoclingLayoutDetector
  15. # 可选导入 MinerU 适配器
  16. try:
  17. from .mineru_adapter import (
  18. MinerUPreprocessor,
  19. MinerULayoutDetector,
  20. MinerUVLRecognizer,
  21. MinerUOCRRecognizer
  22. )
  23. MINERU_AVAILABLE = True
  24. except ImportError:
  25. MINERU_AVAILABLE = False
  26. __all__ = [
  27. # 基类
  28. 'BaseAdapter',
  29. 'BasePreprocessor',
  30. 'BaseLayoutDetector',
  31. 'BaseVLRecognizer',
  32. 'BaseOCRRecognizer',
  33. # PaddleX 适配器
  34. 'PaddleLayoutDetector',
  35. 'PaddleVLRecognizer',
  36. # Docling 适配器
  37. 'DoclingLayoutDetector',
  38. ]
  39. # 如果 MinerU 可用,添加到导出列表
  40. if MINERU_AVAILABLE:
  41. __all__.extend([
  42. 'MinerUPreprocessor',
  43. 'MinerULayoutDetector',
  44. 'MinerUVLRecognizer',
  45. 'MinerUOCRRecognizer',
  46. ])
  47. def get_layout_detector(config: dict):
  48. """
  49. 根据配置获取布局检测器
  50. Args:
  51. config: 配置字典,包含 module 和其他参数
  52. Returns:
  53. BaseLayoutDetector 实例
  54. """
  55. module = config.get('module', 'paddle')
  56. if module == 'paddle':
  57. return PaddleLayoutDetector(config)
  58. elif module == 'mineru':
  59. if not MINERU_AVAILABLE:
  60. raise ImportError("MinerU adapter not available")
  61. return MinerULayoutDetector(config)
  62. elif module == 'docling':
  63. return DoclingLayoutDetector(config)
  64. else:
  65. raise ValueError(f"Unknown layout detection module: {module}")
  66. def get_preprocessor(config: dict):
  67. """根据配置获取预处理器"""
  68. module = config.get('module', 'mineru')
  69. if module == 'mineru':
  70. if not MINERU_AVAILABLE:
  71. raise ImportError("MinerU adapter not available")
  72. return MinerUPreprocessor(config)
  73. else:
  74. raise ValueError(f"Unknown preprocessor module: {module}")
  75. def get_vl_recognizer(config: dict):
  76. """根据配置获取VL识别器"""
  77. module = config.get('module', 'mineru')
  78. if module == 'paddle':
  79. return PaddleVLRecognizer(config)
  80. elif module == 'mineru':
  81. if not MINERU_AVAILABLE:
  82. raise ImportError("MinerU adapter not available")
  83. return MinerUVLRecognizer(config)
  84. else:
  85. raise ValueError(f"Unknown VL recognizer module: {module}")
  86. def get_ocr_recognizer(config: dict):
  87. """根据配置获取OCR识别器"""
  88. module = config.get('module', 'mineru')
  89. if module == 'mineru':
  90. if not MINERU_AVAILABLE:
  91. raise ImportError("MinerU adapter not available")
  92. return MinerUOCRRecognizer(config)
  93. else:
  94. raise ValueError(f"Unknown OCR recognizer module: {module}")