mineru 2.6.7__py3-none-any.whl → 2.7.0__py3-none-any.whl

mineru/backend/hybrid/hybrid_model_output_to_middle_json.py

@@ -0,0 +1,212 @@
+#  Copyright (c) Opendatalab. All rights reserved.
+
+import os
+import time
+
+import cv2
+import numpy as np
+from loguru import logger
+
+from mineru.backend.hybrid.hybrid_magic_model import MagicModel
+from mineru.backend.utils import cross_page_table_merge
+from mineru.utils.config_reader import get_table_enable, get_llm_aided_config
+from mineru.utils.cut_image import cut_image_and_table
+from mineru.utils.enum_class import ContentType
+from mineru.utils.hash_utils import bytes_md5
+from mineru.utils.ocr_utils import OcrConfidence
+from mineru.utils.pdf_image_tools import get_crop_img
+from mineru.version import __version__
+
+
+heading_level_import_success = False
+llm_aided_config = get_llm_aided_config()
+if llm_aided_config:
+    title_aided_config = llm_aided_config.get('title_aided', {})
+    if title_aided_config.get('enable', False):
+        try:
+            from mineru.utils.llm_aided import llm_aided_title
+            from mineru.backend.pipeline.model_init import AtomModelSingleton
+            heading_level_import_success = True
+        except Exception as e:
+            logger.warning("The heading level feature cannot be used. If you need to use the heading level feature, "
+                            "please execute `pip install mineru[core]` to install the required packages.")
+
+
+def blocks_to_page_info(
+        page_blocks,
+        page_inline_formula,
+        page_ocr_res,
+        image_dict,
+        page,
+        image_writer,
+        page_index,
+        _ocr_enable,
+        _vlm_ocr_enable,
+) -> dict:
+    """将blocks转换为页面信息"""
+
+    scale = image_dict["scale"]
+    page_pil_img = image_dict["img_pil"]
+    page_img_md5 = bytes_md5(page_pil_img.tobytes())
+    width, height = map(int, page.get_size())
+
+    magic_model = MagicModel(
+        page_blocks,
+        page_inline_formula,
+        page_ocr_res,
+        page,
+        scale,
+        page_pil_img,
+        width,
+        height,
+        _ocr_enable,
+        _vlm_ocr_enable,
+    )
+    image_blocks = magic_model.get_image_blocks()
+    table_blocks = magic_model.get_table_blocks()
+    title_blocks = magic_model.get_title_blocks()
+    discarded_blocks = magic_model.get_discarded_blocks()
+    code_blocks = magic_model.get_code_blocks()
+    ref_text_blocks = magic_model.get_ref_text_blocks()
+    phonetic_blocks = magic_model.get_phonetic_blocks()
+    list_blocks = magic_model.get_list_blocks()
+
+    # 如果有标题优化需求，计算标题的平均行高
+    if heading_level_import_success:
+        if _vlm_ocr_enable:  # vlm_ocr导致没有line信息，需要重新det获取平均行高
+            atom_model_manager = AtomModelSingleton()
+            ocr_model = atom_model_manager.get_atom_model(
+                atom_model_name='ocr',
+                ocr_show_log=False,
+                det_db_box_thresh=0.3,
+                lang='ch_lite'
+            )
+            for title_block in title_blocks:
+                title_pil_img = get_crop_img(title_block['bbox'], page_pil_img, scale)
+                title_np_img = np.array(title_pil_img)
+                # 给title_pil_img添加上下左右各50像素白边padding
+                title_np_img = cv2.copyMakeBorder(
+                    title_np_img, 50, 50, 50, 50, cv2.BORDER_CONSTANT, value=[255, 255, 255]
+                )
+                title_img = cv2.cvtColor(title_np_img, cv2.COLOR_RGB2BGR)
+                ocr_det_res = ocr_model.ocr(title_img, rec=False)[0]
+                if len(ocr_det_res) > 0:
+                    # 计算所有res的平均高度
+                    avg_height = np.mean([box[2][1] - box[0][1] for box in ocr_det_res])
+                    title_block['line_avg_height'] = round(avg_height/scale)
+        else:  # 有line信息，直接计算平均行高
+            for title_block in title_blocks:
+                lines = title_block.get('lines', [])
+                if lines:
+                    # 使用列表推导式和内置函数,一次性计算平均高度
+                    avg_height = sum(line['bbox'][3] - line['bbox'][1] for line in lines) / len(lines)
+                    title_block['line_avg_height'] = round(avg_height)
+                else:
+                    title_block['line_avg_height'] = title_block['bbox'][3] - title_block['bbox'][1]
+
+    text_blocks = magic_model.get_text_blocks()
+    interline_equation_blocks = magic_model.get_interline_equation_blocks()
+
+    all_spans = magic_model.get_all_spans()
+    # 对image/table/interline_equation的span截图
+    for span in all_spans:
+        if span["type"] in [ContentType.IMAGE, ContentType.TABLE, ContentType.INTERLINE_EQUATION]:
+            span = cut_image_and_table(span, page_pil_img, page_img_md5, page_index, image_writer, scale=scale)
+
+    page_blocks = []
+    page_blocks.extend([
+        *image_blocks,
+        *table_blocks,
+        *code_blocks,
+        *ref_text_blocks,
+        *phonetic_blocks,
+        *title_blocks,
+        *text_blocks,
+        *interline_equation_blocks,
+        *list_blocks,
+    ])
+    # 对page_blocks根据index的值进行排序
+    page_blocks.sort(key=lambda x: x["index"])
+
+    page_info = {"para_blocks": page_blocks, "discarded_blocks": discarded_blocks, "page_size": [width, height], "page_idx": page_index}
+    return page_info
+
+
+def result_to_middle_json(
+        model_output_blocks_list,
+        inline_formula_list,
+        ocr_res_list,
+        images_list,
+        pdf_doc,
+        image_writer,
+        _ocr_enable,
+        _vlm_ocr_enable,
+        hybrid_pipeline_model,
+):
+    middle_json = {
+        "pdf_info": [],
+        "_backend": "hybrid",
+        "_ocr_enable": _ocr_enable,
+        "_vlm_ocr_enable": _vlm_ocr_enable,
+        "_version_name": __version__
+    }
+
+    for index, (page_blocks, page_inline_formula, page_ocr_res) in enumerate(zip(model_output_blocks_list, inline_formula_list, ocr_res_list)):
+        page = pdf_doc[index]
+        image_dict = images_list[index]
+        page_info = blocks_to_page_info(
+            page_blocks, page_inline_formula, page_ocr_res,
+            image_dict, page, image_writer, index,
+            _ocr_enable, _vlm_ocr_enable
+        )
+        middle_json["pdf_info"].append(page_info)
+
+    if not (_vlm_ocr_enable or _ocr_enable):
+        """后置ocr处理"""
+        need_ocr_list = []
+        img_crop_list = []
+        text_block_list = []
+        for page_info in middle_json["pdf_info"]:
+            for block in page_info['para_blocks']:
+                if block['type'] in ['table', 'image', 'list', 'code']:
+                    for sub_block in block['blocks']:
+                        if not sub_block['type'].endswith('body'):
+                            text_block_list.append(sub_block)
+                elif block['type'] in ['text', 'title', 'ref_text']:
+                    text_block_list.append(block)
+            for block in page_info['discarded_blocks']:
+                text_block_list.append(block)
+        for block in text_block_list:
+            for line in block['lines']:
+                for span in line['spans']:
+                    if 'np_img' in span:
+                        need_ocr_list.append(span)
+                        img_crop_list.append(span['np_img'])
+                        span.pop('np_img')
+        if len(img_crop_list) > 0:
+            ocr_res_list = hybrid_pipeline_model.ocr_model.ocr(img_crop_list, det=False, tqdm_enable=True)[0]
+            assert len(ocr_res_list) == len(
+                need_ocr_list), f'ocr_res_list: {len(ocr_res_list)}, need_ocr_list: {len(need_ocr_list)}'
+            for index, span in enumerate(need_ocr_list):
+                ocr_text, ocr_score = ocr_res_list[index]
+                if ocr_score > OcrConfidence.min_confidence:
+                    span['content'] = ocr_text
+                    span['score'] = float(f"{ocr_score:.3f}")
+                else:
+                    span['content'] = ''
+                    span['score'] = 0.0
+
+    """表格跨页合并"""
+    table_enable = get_table_enable(os.getenv('MINERU_VLM_TABLE_ENABLE', 'True').lower() == 'true')
+    if table_enable:
+        cross_page_table_merge(middle_json["pdf_info"])
+
+    """llm优化标题分级"""
+    if heading_level_import_success:
+        llm_aided_title_start_time = time.time()
+        llm_aided_title(middle_json["pdf_info"], title_aided_config)
+        logger.info(f'llm aided title time: {round(time.time() - llm_aided_title_start_time, 2)}')
+
+    # 关闭pdf文档
+    pdf_doc.close()
+    return middle_json

mineru/backend/pipeline/batch_analyze.py

@@ -420,7 +420,15 @@ class BatchAnalyze:
                                                layout_res_item['poly'][4], layout_res_item['poly'][5]]
                             layout_res_width = layout_res_bbox[2] - layout_res_bbox[0]
                             layout_res_height = layout_res_bbox[3] - layout_res_bbox[1]
-                            if ocr_text in ['（204号', '（20', '（2', '（2号', '（20号', '号', '（204'] and ocr_score < 0.8 and layout_res_width < layout_res_height:
+                            if (
+                                    ocr_text in [
+                                        '（204号', '（20', '（2', '（2号', '（20号', '号', '（204',
+                                        '(cid:)', '(ci:)', '(cd:1)', 'cd:)', 'c)', '(cd:)', 'c', 'id:)',
+                                        ':)', '√:)', '√i:)', '−i:)', '−:', 'i:)',
+                                    ]
+                                    and ocr_score < 0.8
+                                    and layout_res_width < layout_res_height
+                            ):
                                 layout_res_item['category_id'] = 16
 
                     total_processed += len(img_crop_list)

mineru/backend/pipeline/model_init.py

@@ -14,6 +14,7 @@ from ...model.table.cls.paddle_table_cls import PaddleTableClsModel
 # from ...model.table.rec.RapidTable import RapidTableModel
 from ...model.table.rec.slanet_plus.main import RapidTableModel
 from ...model.table.rec.unet_table.main import UnetTableModel
+from ...utils.config_reader import get_device
 from ...utils.enum_class import ModelPath
 from ...utils.models_download_utils import auto_download_and_get_model_root_path
 
@@ -267,4 +268,98 @@ class MineruPipelineModel:
                 lang=self.lang,
             )
 
-        logger.info('DocAnalysis init done!')
+        logger.info('DocAnalysis init done!')
+
+
+class HybridModelSingleton:
+    _instance = None
+    _models = {}
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def get_model(
+        self,
+        lang=None,
+        formula_enable=None,
+    ):
+        key = (lang, formula_enable)
+        if key not in self._models:
+            self._models[key] = MineruHybridModel(
+                lang=lang,
+                formula_enable=formula_enable,
+            )
+        return self._models[key]
+
+def ocr_det_batch_setting(device):
+    # 检测torch的版本号
+    import torch
+    from packaging import version
+    if version.parse(torch.__version__) >= version.parse("2.8.0") or str(device).startswith('mps'):
+        enable_ocr_det_batch = False
+    else:
+        enable_ocr_det_batch = True
+    return enable_ocr_det_batch
+
+class MineruHybridModel:
+    def __init__(
+            self,
+            device=None,
+            lang=None,
+            formula_enable=True,
+    ):
+        if device is not None:
+            self.device = device
+        else:
+            self.device = get_device()
+
+        self.lang = lang
+
+        self.enable_ocr_det_batch = ocr_det_batch_setting(self.device)
+
+        if str(self.device).startswith('npu'):
+            try:
+                import torch_npu
+                if torch_npu.npu.is_available():
+                    torch_npu.npu.set_compile_mode(jit_compile=False)
+            except Exception as e:
+                raise RuntimeError(
+                    "NPU is selected as device, but torch_npu is not available. "
+                    "Please ensure that the torch_npu package is installed correctly."
+                ) from e
+
+        self.atom_model_manager = AtomModelSingleton()
+
+        # 初始化OCR模型
+        self.ocr_model = self.atom_model_manager.get_atom_model(
+            atom_model_name=AtomicModel.OCR,
+            det_db_box_thresh=0.3,
+            lang=self.lang
+        )
+
+        if formula_enable:
+            # 初始化公式检测模型
+            self.mfd_model = self.atom_model_manager.get_atom_model(
+                atom_model_name=AtomicModel.MFD,
+                mfd_weights=str(
+                    os.path.join(auto_download_and_get_model_root_path(ModelPath.yolo_v8_mfd), ModelPath.yolo_v8_mfd)
+                ),
+                device=self.device,
+            )
+
+            # 初始化公式解析模型
+            if MFR_MODEL == "unimernet_small":
+                mfr_model_path = ModelPath.unimernet_small
+            elif MFR_MODEL == "pp_formulanet_plus_m":
+                mfr_model_path = ModelPath.pp_formulanet_plus_m
+            else:
+                logger.error('MFR model name not allow')
+                exit(1)
+
+            self.mfr_model = self.atom_model_manager.get_atom_model(
+                atom_model_name=AtomicModel.MFR,
+                mfr_weight_dir=str(os.path.join(auto_download_and_get_model_root_path(mfr_model_path), mfr_model_path)),
+                device=self.device,
+            )

mineru/backend/pipeline/pipeline_analyze.py

@@ -86,6 +86,7 @@ def doc_analyze(
     all_image_lists = []
     all_pdf_docs = []
     ocr_enabled_list = []
+    load_images_start = time.time()
     for pdf_idx, pdf_bytes in enumerate(pdf_bytes_list):
         # 确定OCR设置
         _ocr_enable = False
@@ -99,10 +100,7 @@ def doc_analyze(
         _lang = lang_list[pdf_idx]
 
         # 收集每个数据集中的页面
-        # load_images_start = time.time()
         images_list, pdf_doc = load_images_from_pdf(pdf_bytes, image_type=ImageType.PIL)
-        # load_images_time = round(time.time() - load_images_start, 2)
-        # logger.debug(f"load images cost: {load_images_time}, speed: {round(len(images_list) / load_images_time, 3)} images/s")
         all_image_lists.append(images_list)
         all_pdf_docs.append(pdf_doc)
         for page_idx in range(len(images_list)):
@@ -111,6 +109,8 @@ def doc_analyze(
                 pdf_idx, page_idx,
                 img_dict['img_pil'], _ocr_enable, _lang,
             ))
+    load_images_time = round(time.time() - load_images_start, 2)
+    logger.debug(f"load images cost: {load_images_time}, speed: {round(len(all_pages_info) / load_images_time, 3)} images/s")
 
     # 准备批处理
     images_with_extra_info = [(info[2], info[3], info[4]) for info in all_pages_info]
@@ -123,6 +123,7 @@ def doc_analyze(
     # 执行批处理
     results = []
     processed_images_count = 0
+    infer_start = time.time()
     for index, batch_image in enumerate(batch_images):
         processed_images_count += len(batch_image)
         logger.info(
@@ -131,6 +132,8 @@ def doc_analyze(
         )
         batch_results = batch_image_analyze(batch_image, formula_enable, table_enable)
         results.extend(batch_results)
+    infer_time = round(time.time() - infer_start, 2)
+    logger.debug(f"infer finished, cost: {infer_time}, speed: {round(len(results) / infer_time, 3)} page/s")
 
     # 构建返回结果
     infer_results = []
@@ -185,7 +188,6 @@ def batch_image_analyze(
         batch_ratio = 1
     logger.info(
             f'GPU Memory: {gpu_memory} GB, Batch Ratio: {batch_ratio}. '
-            f'You can set MINERU_VIRTUAL_VRAM_SIZE environment variable to adjust GPU memory allocation.'
     )
 
     # 检测torch的版本号

mineru/backend/pipeline/pipeline_middle_json_mkcontent.py

@@ -1,25 +1,12 @@
-import re
 from loguru import logger
 
+from mineru.utils.char_utils import full_to_half_exclude_marks, is_hyphen_at_line_end
 from mineru.utils.config_reader import get_latex_delimiter_config
 from mineru.backend.pipeline.para_split import ListLineTag
 from mineru.utils.enum_class import BlockType, ContentType, MakeMode
 from mineru.utils.language import detect_lang
 
 
-def __is_hyphen_at_line_end(line):
-    """Check if a line ends with one or more letters followed by a hyphen.
-
-    Args:
-    line (str): The line of text to check.
-
-    Returns:
-    bool: True if the line ends with one or more letters followed by a hyphen, False otherwise.
-    """
-    # Use regex to check if the line ends with one or more letters followed by a hyphen
-    return bool(re.search(r'[A-Za-z]+-\s*$', line))
-
-
 def make_blocks_to_markdown(paras_of_layout,
                                       mode,
                                       img_buket_path='',
@@ -102,25 +89,6 @@ def make_blocks_to_markdown(paras_of_layout,
     return page_markdown
 
 
-def full_to_half(text: str) -> str:
-    """Convert full-width characters to half-width characters using code point manipulation.
-
-    Args:
-        text: String containing full-width characters
-
-    Returns:
-        String with full-width characters converted to half-width
-    """
-    result = []
-    for char in text:
-        code = ord(char)
-        # Full-width letters and numbers (FF21-FF3A for A-Z, FF41-FF5A for a-z, FF10-FF19 for 0-9)
-        if (0xFF21 <= code <= 0xFF3A) or (0xFF41 <= code <= 0xFF5A) or (0xFF10 <= code <= 0xFF19):
-            result.append(chr(code - 0xFEE0))  # Shift to ASCII range
-        else:
-            result.append(char)
-    return ''.join(result)
-
 latex_delimiters_config = get_latex_delimiter_config()
 
 default_delimiters = {
@@ -140,7 +108,7 @@ def merge_para_with_text(para_block):
     for line in para_block['lines']:
         for span in line['spans']:
             if span['type'] in [ContentType.TEXT]:
-                span['content'] = full_to_half(span['content'])
+                span['content'] = full_to_half_exclude_marks(span['content'])
                 block_text += span['content']
     block_lang = detect_lang(block_text)
 
@@ -166,22 +134,45 @@ def merge_para_with_text(para_block):
             content = content.strip()
 
             if content:
-                langs = ['zh', 'ja', 'ko']
+
+                if span_type == ContentType.INTERLINE_EQUATION:
+                    para_text += content
+                    continue
+
+                # 定义CJK语言集合(中日韩)
+                cjk_langs = {'zh', 'ja', 'ko'}
                 # logger.info(f'block_lang: {block_lang}, content: {content}')
-                if block_lang in langs: # 中文/日语/韩文语境下，换行不需要空格分隔,但是如果是行内公式结尾，还是要加空格
-                    if j == len(line['spans']) - 1 and span_type not in [ContentType.INLINE_EQUATION]:
+
+                # 判断是否为行末span
+                is_last_span = j == len(line['spans']) - 1
+
+                if block_lang in cjk_langs: # 中文/日语/韩文语境下，换行不需要空格分隔,但是如果是行内公式结尾，还是要加空格
+                    if is_last_span and span_type not in [ContentType.INLINE_EQUATION]:
                         para_text += content
                     else:
                         para_text += f'{content} '
                 else:
+                    # 西方文本语境下 每行的最后一个span判断是否要去除连字符
                     if span_type in [ContentType.TEXT, ContentType.INLINE_EQUATION]:
                         # 如果span是line的最后一个且末尾带有-连字符，那么末尾不应该加空格,同时应该把-删除
-                        if j == len(line['spans'])-1 and span_type == ContentType.TEXT and __is_hyphen_at_line_end(content):
-                            para_text += content[:-1]
+                        if (
+                                is_last_span
+                                and span_type == ContentType.TEXT
+                                and is_hyphen_at_line_end(content)
+                        ):
+                            # 如果下一行的第一个span是小写字母开头，删除连字符
+                            if (
+                                    i + 1 < len(para_block['lines'])
+                                    and para_block['lines'][i + 1].get('spans')
+                                    and para_block['lines'][i + 1]['spans'][0].get('type') == ContentType.TEXT
+                                    and para_block['lines'][i + 1]['spans'][0].get('content', '')
+                                    and para_block['lines'][i + 1]['spans'][0]['content'][0].islower()
+                            ):
+                                para_text += content[:-1]
+                            else:  # 如果没有下一行，或者下一行的第一个span不是小写字母开头，则保留连字符但不加空格
+                                para_text += content
                         else:  # 西方文本语境下 content间需要空格分隔
                             para_text += f'{content} '
-                    elif span_type == ContentType.INTERLINE_EQUATION:
-                        para_text += content
             else:
                 continue
 

mineru/backend/vlm/utils.py

@@ -72,7 +72,9 @@ def set_lmdeploy_backend(device_type: str) -> str:
 
 def set_default_gpu_memory_utilization() -> float:
     from vllm import __version__ as vllm_version
-    if version.parse(vllm_version) >= version.parse("0.11.0"):
+    device = get_device()
+    gpu_memory = get_vram(device)
+    if version.parse(vllm_version) >= version.parse("0.11.0") and gpu_memory <= 8:
         return 0.7
     else:
         return 0.5

mineru/backend/vlm/vlm_analyze.py

@@ -202,16 +202,16 @@ def doc_analyze(
     if predictor is None:
         predictor = ModelSingleton().get_model(backend, model_path, server_url, **kwargs)
 
-    # load_images_start = time.time()
+    load_images_start = time.time()
     images_list, pdf_doc = load_images_from_pdf(pdf_bytes, image_type=ImageType.PIL)
     images_pil_list = [image_dict["img_pil"] for image_dict in images_list]
-    # load_images_time = round(time.time() - load_images_start, 2)
-    # logger.info(f"load images cost: {load_images_time}, speed: {round(len(images_base64_list)/load_images_time, 3)} images/s")
+    load_images_time = round(time.time() - load_images_start, 2)
+    logger.debug(f"load images cost: {load_images_time}, speed: {round(len(images_pil_list)/load_images_time, 3)} images/s")
 
-    # infer_start = time.time()
+    infer_start = time.time()
     results = predictor.batch_two_step_extract(images=images_pil_list)
-    # infer_time = round(time.time() - infer_start, 2)
-    # logger.info(f"infer finished, cost: {infer_time}, speed: {round(len(results)/infer_time, 3)} page/s")
+    infer_time = round(time.time() - infer_start, 2)
+    logger.debug(f"infer finished, cost: {infer_time}, speed: {round(len(results)/infer_time, 3)} page/s")
 
     middle_json = result_to_middle_json(results, images_list, pdf_doc, image_writer)
     return middle_json, results
@@ -229,15 +229,15 @@ async def aio_doc_analyze(
     if predictor is None:
         predictor = ModelSingleton().get_model(backend, model_path, server_url, **kwargs)
 
-    # load_images_start = time.time()
+    load_images_start = time.time()
     images_list, pdf_doc = load_images_from_pdf(pdf_bytes, image_type=ImageType.PIL)
     images_pil_list = [image_dict["img_pil"] for image_dict in images_list]
-    # load_images_time = round(time.time() - load_images_start, 2)
-    # logger.debug(f"load images cost: {load_images_time}, speed: {round(len(images_pil_list)/load_images_time, 3)} images/s")
+    load_images_time = round(time.time() - load_images_start, 2)
+    logger.debug(f"load images cost: {load_images_time}, speed: {round(len(images_pil_list)/load_images_time, 3)} images/s")
 
-    # infer_start = time.time()
+    infer_start = time.time()
     results = await predictor.aio_batch_two_step_extract(images=images_pil_list)
-    # infer_time = round(time.time() - infer_start, 2)
-    # logger.info(f"infer finished, cost: {infer_time}, speed: {round(len(results)/infer_time, 3)} page/s")
+    infer_time = round(time.time() - infer_start, 2)
+    logger.debug(f"infer finished, cost: {infer_time}, speed: {round(len(results)/infer_time, 3)} page/s")
     middle_json = result_to_middle_json(results, images_list, pdf_doc, image_writer)
     return middle_json, results