MinerU/zhch/universal_doc_parser/core/config_manager.py

"""配置管理器 - 加载和验证配置文件"""
import yaml
from pathlib import Path
from typing import Dict, Any, Optional

class ConfigManager:
    """配置管理器"""
    
    _config_cache = {}
    
    @classmethod
    def load_config(cls, config_path: str) -> Dict[str, Any]:
        """加载配置文件"""
        config_path = Path(config_path)
        
        # 缓存机制
        cache_key = str(config_path.absolute())
        if cache_key in cls._config_cache:
            return cls._config_cache[cache_key]
        
        if not config_path.exists():
            raise FileNotFoundError(f"Config file not found: {config_path}")
        
        with open(config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        
        # 配置验证和默认值设置
        config = cls._validate_and_set_defaults(config)
        
        # 缓存配置
        cls._config_cache[cache_key] = config
        
        return config
    
    @classmethod
    def _validate_and_set_defaults(cls, config: Dict[str, Any]) -> Dict[str, Any]:
        """验证配置并设置默认值"""
        # 设置默认场景名称
        if 'scene_name' not in config:
            config['scene_name'] = 'unknown'
        
        # 验证必需的配置项
        required_sections = ['preprocessor', 'layout_detection', 'vl_recognition', 'ocr_recognition']
        for section in required_sections:
            if section not in config:
                config[section] = {'module': 'mineru'}
        
        # 设置预处理器默认配置
        preprocessor_defaults = {
            'module': 'mineru',
            'orientation_classifier': {'enabled': True},
            'unwarping': {'enabled': False}
        }
        config['preprocessor'] = cls._merge_defaults(
            config.get('preprocessor', {}), preprocessor_defaults
        )
        
        # 设置版式检测默认配置
        layout_defaults = {
            'module': 'mineru',
            'model_name': 'RT-DETR-H_layout_17cls',
            'device': 'cpu',
            'batch_size': 1,
            'conf': 0.25,
            'iou': 0.45
        }
        config['layout_detection'] = cls._merge_defaults(
            config.get('layout_detection', {}), layout_defaults
        )
        
        # 设置VL识别默认配置
        vl_defaults = {
            'module': 'mineru',
            'backend': 'vllm-http-client',
            'server_url': 'http://localhost:8111/v1',
            'device': 'cpu',
            'batch_size': 1,
            'model_params': {'max_concurrency': 10, 'http_timeout': 600}
        }
        config['vl_recognition'] = cls._merge_defaults(
            config.get('vl_recognition', {}), vl_defaults
        )
        
        # 设置OCR默认配置
        ocr_defaults = {
            'module': 'mineru',
            'language': 'ch',
            'det_threshold': 0.3,
            'unclip_ratio': 1.8,
            'batch_size': 8,
            'device': 'cpu'
        }
        config['ocr_recognition'] = cls._merge_defaults(
            config.get('ocr_recognition', {}), ocr_defaults
        )
        
        # 设置输出默认配置
        output_defaults = {
            'format': 'enhanced_json',
            'save_json': True,
            'save_markdown': True,
            'save_html': True,
            'save_images': {'layout': True, 'ocr': True, 'table_cells': True},
            'coordinate_precision': 2
        }
        config['output'] = cls._merge_defaults(
            config.get('output', {}), output_defaults
        )
        
        return config
    
    @classmethod
    def _merge_defaults(cls, user_config: Dict[str, Any], defaults: Dict[str, Any]) -> Dict[str, Any]:
        """合并用户配置和默认配置"""
        result = defaults.copy()
        for key, value in user_config.items():
            if isinstance(value, dict) and key in result and isinstance(result[key], dict):
                result[key] = cls._merge_defaults(value, result[key])
            else:
                result[key] = value
        return result
    
    @classmethod
    def save_config(cls, config: Dict[str, Any], config_path: str):
        """保存配置文件"""
        config_path = Path(config_path)
        config_path.parent.mkdir(parents=True, exist_ok=True)
        
        with open(config_path, 'w', encoding='utf-8') as f:
            yaml.dump(config, f, default_flow_style=False, allow_unicode=True)
    
    @classmethod
    def clear_cache(cls):
        """清空配置缓存"""
        cls._config_cache.clear()