MinerU/zhch/universal_doc_parser/models/adapters/__init__.py

"""
模型适配器模块
提供统一的接口适配不同的模型后端
"""

from .base import (
    BaseAdapter,
    BasePreprocessor,
    BaseLayoutDetector,
    BaseVLRecognizer,
    BaseOCRRecognizer
)

from .paddle_layout_detector import PaddleLayoutDetector

# 可选导入 MinerU 适配器
try:
    from .mineru_adapter import (
        MinerUPreprocessor,
        MinerULayoutDetector,
        MinerUVLRecognizer,
        MinerUOCRRecognizer
    )
    MINERU_AVAILABLE = True
except ImportError:
    MINERU_AVAILABLE = False

__all__ = [
    # 基类
    'BaseAdapter',
    'BasePreprocessor',
    'BaseLayoutDetector',
    'BaseVLRecognizer',
    'BaseOCRRecognizer',
    
    # PaddleX 适配器
    'PaddleLayoutDetector',
]

# 如果 MinerU 可用,添加到导出列表
if MINERU_AVAILABLE:
    __all__.extend([
        'MinerUPreprocessor',
        'MinerULayoutDetector',
        'MinerUVLRecognizer',
        'MinerUOCRRecognizer',
    ])


def get_layout_detector(config: dict):
    """
    根据配置获取布局检测器
    
    Args:
        config: 配置字典,包含 module 和其他参数
        
    Returns:
        BaseLayoutDetector 实例
    """
    module = config.get('module', 'paddle')
    
    if module == 'paddle':
        return PaddleLayoutDetector(config)
    elif module == 'mineru':
        if not MINERU_AVAILABLE:
            raise ImportError("MinerU adapter not available")
        return MinerULayoutDetector(config)
    else:
        raise ValueError(f"Unknown layout detection module: {module}")


def get_preprocessor(config: dict):
    """根据配置获取预处理器"""
    module = config.get('module', 'mineru')
    
    if module == 'mineru':
        if not MINERU_AVAILABLE:
            raise ImportError("MinerU adapter not available")
        return MinerUPreprocessor(config)
    else:
        raise ValueError(f"Unknown preprocessor module: {module}")


def get_vl_recognizer(config: dict):
    """根据配置获取VL识别器"""
    module = config.get('module', 'mineru')
    
    if module == 'mineru':
        if not MINERU_AVAILABLE:
            raise ImportError("MinerU adapter not available")
        return MinerUVLRecognizer(config)
    else:
        raise ValueError(f"Unknown VL recognizer module: {module}")


def get_ocr_recognizer(config: dict):
    """根据配置获取OCR识别器"""
    module = config.get('module', 'mineru')
    
    if module == 'mineru':
        if not MINERU_AVAILABLE:
            raise ImportError("MinerU adapter not available")
        return MinerUOCRRecognizer(config)
    else:
        raise ValueError(f"Unknown OCR recognizer module: {module}")