Class OpenDocONNX

tools/infer_doc_onnx.py:525–1233

完整的文档OCR ONNX Pipeline

# ==================== OpenDoc ONNX Pipeline ====================
class OpenDocONNX:
    """完整的文档OCR ONNX Pipeline"""

    def __init__(
        self,
        layout_model_path: Optional[str] = None,
        unirec_encoder_path: Optional[str] = None,
        unirec_decoder_path: Optional[str] = None,
        tokenizer_mapping_path: Optional[str] = None,
        use_gpu: Optional[bool] = None,
        layout_threshold: float = 0.5,
        use_layout_detection: bool = True,
        use_chart_recognition: bool = True,
        auto_download: bool = True,
        max_parallel_blocks: int = 4,
    ):
        """
        初始化OpenDoc ONNX Pipeline

        Args:
            layout_model_path: 版面检测ONNX模型路径. If None, use default cache directory.
            unirec_encoder_path: UniRec编码器ONNX模型路径. If None, use default cache directory.
            unirec_decoder_path: UniRec解码器ONNX模型路径. If None, use default cache directory.
            tokenizer_mapping_path: Tokenizer映射文件路径. If None, use default cache directory.
            use_gpu: Whether to use GPU. If None, auto-detect. If True, force GPU. If False, force CPU.
            layout_threshold: 版面检测阈值
            use_layout_detection: 是否使用版面检测
            use_chart_recognition: 是否识别图表
            auto_download: If True, automatically download missing models
            max_parallel_blocks: Maximum number of blocks to process in parallel for VLM recognition (default: 4)
        """
        self.use_layout_detection = use_layout_detection
        self.use_chart_recognition = use_chart_recognition
        self.max_parallel_blocks = max(1, max_parallel_blocks)

        # Set default paths if not provided
        if layout_model_path is None:
            cache_dir = Path.home() / '.cache' / 'openocr'
            layout_model_path = str(cache_dir / 'PP_DoclayoutV2_onnx' / 'PP-DoclayoutV2.onnx')

        # Markdown忽略的标签
        self.markdown_ignore_labels = [
            'number', 'footnote', 'header', 'footer', 'aside_text', 'footer_image', 'header_image', 'chart'
        ]

        # 为所有25种标签类型定义不同的颜色 (BGR格式)
        self.colors = {
            'abstract': (255, 128, 0),        # 橙色
            'algorithm': (128, 0, 255),       # 紫色
            'aside_text': (128, 128, 128),    # 灰色
            'chart': (0, 255, 255),           # 青色
            'content': (0, 255, 0),           # 绿色
            'display_formula': (255, 0, 255), # 品红
            'doc_title': (255, 0, 0),         # 红色
            'figure_title': (255, 128, 128),  # 浅红
            'footer': (64, 64, 64),           # 深灰
            'footer_image': (128, 64, 0),     # 棕色
            'footnote': (192, 192, 192),      # 浅灰