mineru 2.2.2__py3-none-any.whl → 2.5.0__py3-none-any.whl

mineru/backend/pipeline/pipeline_middle_json_mkcontent.py

@@ -245,14 +245,14 @@ def make_blocks_to_content_list(para_block, img_buket_path, page_idx, page_size)
             if block['type'] == BlockType.TABLE_FOOTNOTE:
                 para_content[BlockType.TABLE_FOOTNOTE].append(merge_para_with_text(block))
 
-    page_weight, page_height = page_size
+    page_width, page_height = page_size
     para_bbox = para_block.get('bbox')
     if para_bbox:
         x0, y0, x1, y1 = para_bbox
         para_content['bbox'] = [
-            int(x0 * 1000 / page_weight),
+            int(x0 * 1000 / page_width),
             int(y0 * 1000 / page_height),
-            int(x1 * 1000 / page_weight),
+            int(x1 * 1000 / page_width),
             int(y1 * 1000 / page_height),
         ]
 

mineru/backend/vlm/model_output_to_middle_json.py

@@ -0,0 +1,123 @@
+import os
+import time
+
+import cv2
+import numpy as np
+from loguru import logger
+
+from mineru.backend.vlm.vlm_magic_model import MagicModel
+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.pdf_image_tools import get_crop_img
+from mineru.utils.table_merge import merge_table
+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, image_dict, page, image_writer, page_index) -> dict:
+    """将blocks转换为页面信息"""
+
+    scale = image_dict["scale"]
+    # page_pil_img = image_dict["img_pil"]
+    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, width, height)
+    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()
+
+    # 如果有标题优化需求，则对title_blocks截图det
+    if heading_level_import_success:
+        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)
+
+    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, images_list, pdf_doc, image_writer):
+    middle_json = {"pdf_info": [], "_backend":"vlm", "_version_name": __version__}
+    for index, page_blocks in enumerate(model_output_blocks_list):
+        page = pdf_doc[index]
+        image_dict = images_list[index]
+        page_info = blocks_to_page_info(page_blocks, image_dict, page, image_writer, index)
+        middle_json["pdf_info"].append(page_info)
+
+    """表格跨页合并"""
+    table_enable = get_table_enable(os.getenv('MINERU_VLM_TABLE_ENABLE', 'True').lower() == 'true')
+    if table_enable:
+        merge_table(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/vlm/vlm_analyze.py

@@ -1,16 +1,20 @@
 # Copyright (c) Opendatalab. All rights reserved.
+import os
 import time
 
 from loguru import logger
 
+from .model_output_to_middle_json import result_to_middle_json
 from ...data.data_reader_writer import DataWriter
 from mineru.utils.pdf_image_tools import load_images_from_pdf
-from .base_predictor import BasePredictor
-from .predictor import get_predictor
-from .token_to_middle_json import result_to_middle_json
+from ...utils.config_reader import get_device
+
 from ...utils.enum_class import ImageType
+from ...utils.model_utils import get_vram
 from ...utils.models_download_utils import auto_download_and_get_model_root_path
 
+from mineru_vl_utils import MinerUClient
+
 
 class ModelSingleton:
     _instance = None
@@ -27,24 +31,101 @@ class ModelSingleton:
         model_path: str | None,
         server_url: str | None,
         **kwargs,
-    ) -> BasePredictor:
+    ) -> MinerUClient:
         key = (backend, model_path, server_url)
         if key not in self._models:
-            if backend in ['transformers', 'sglang-engine'] and not model_path:
+            start_time = time.time()
+            model = None
+            processor = None
+            vllm_llm = None
+            vllm_async_llm = None
+            batch_size = 0
+            if backend in ['transformers', 'vllm-engine', "vllm-async-engine"] and not model_path:
                 model_path = auto_download_and_get_model_root_path("/","vlm")
-            self._models[key] = get_predictor(
+                if backend == "transformers":
+                    try:
+                        from transformers import (
+                            AutoProcessor,
+                            Qwen2VLForConditionalGeneration,
+                        )
+                        from transformers import __version__ as transformers_version
+                    except ImportError:
+                        raise ImportError("Please install transformers to use the transformers backend.")
+
+                    from packaging import version
+                    if version.parse(transformers_version) >= version.parse("4.56.0"):
+                        dtype_key = "dtype"
+                    else:
+                        dtype_key = "torch_dtype"
+                    device = get_device()
+                    model = Qwen2VLForConditionalGeneration.from_pretrained(
+                        model_path,
+                        device_map={"": device},
+                        **{dtype_key: "auto"},  # type: ignore
+                    )
+                    processor = AutoProcessor.from_pretrained(
+                        model_path,
+                        use_fast=True,
+                    )
+                    try:
+                        vram = get_vram(device)
+                        if vram is not None:
+                            gpu_memory = int(os.getenv('MINERU_VIRTUAL_VRAM_SIZE', round(vram)))
+                            if gpu_memory >= 16:
+                                batch_size = 8
+                            elif gpu_memory >= 8:
+                                batch_size = 4
+                            else:
+                                batch_size = 1
+                            logger.info(f'gpu_memory: {gpu_memory} GB, batch_size: {batch_size}')
+                        else:
+                            # Default batch_ratio when VRAM can't be determined
+                            batch_size = 1
+                            logger.info(f'Could not determine GPU memory, using default batch_ratio: {batch_size}')
+                    except Exception as e:
+                        logger.warning(f'Error determining VRAM: {e}, using default batch_ratio: 1')
+                        batch_size = 1
+                elif backend == "vllm-engine":
+                    try:
+                        import vllm
+                    except ImportError:
+                        raise ImportError("Please install vllm to use the vllm-engine backend.")
+                    if "gpu_memory_utilization" not in kwargs:
+                        kwargs["gpu_memory_utilization"] = 0.5
+                    if "model" not in kwargs:
+                        kwargs["model"] = model_path
+                    # 使用kwargs为 vllm初始化参数
+                    vllm_llm = vllm.LLM(**kwargs)
+                elif backend == "vllm-async-engine":
+                    try:
+                        from vllm.engine.arg_utils import AsyncEngineArgs
+                        from vllm.v1.engine.async_llm import AsyncLLM
+                    except ImportError:
+                        raise ImportError("Please install vllm to use the vllm-async-engine backend.")
+                    if "gpu_memory_utilization" not in kwargs:
+                        kwargs["gpu_memory_utilization"] = 0.5
+                    if "model" not in kwargs:
+                        kwargs["model"] = model_path
+                    # 使用kwargs为 vllm初始化参数
+                    vllm_async_llm = AsyncLLM.from_engine_args(AsyncEngineArgs(**kwargs))
+            self._models[key] = MinerUClient(
                 backend=backend,
-                model_path=model_path,
+                model=model,
+                processor=processor,
+                vllm_llm=vllm_llm,
+                vllm_async_llm=vllm_async_llm,
                 server_url=server_url,
-                **kwargs,
+                batch_size=batch_size,
             )
+            elapsed = round(time.time() - start_time, 2)
+            logger.info(f"get {backend} predictor cost: {elapsed}s")
         return self._models[key]
 
 
 def doc_analyze(
     pdf_bytes,
     image_writer: DataWriter | None,
-    predictor: BasePredictor | None = None,
+    predictor: MinerUClient | None = None,
     backend="transformers",
     model_path: str | None = None,
     server_url: str | None = None,
@@ -54,13 +135,13 @@ def doc_analyze(
         predictor = ModelSingleton().get_model(backend, model_path, server_url, **kwargs)
 
     # load_images_start = time.time()
-    images_list, pdf_doc = load_images_from_pdf(pdf_bytes, image_type=ImageType.BASE64)
-    images_base64_list = [image_dict["img_base64"] for image_dict in images_list]
+    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")
 
     # infer_start = time.time()
-    results = predictor.batch_predict(images=images_base64_list)
+    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")
 
@@ -71,7 +152,7 @@ def doc_analyze(
 async def aio_doc_analyze(
     pdf_bytes,
     image_writer: DataWriter | None,
-    predictor: BasePredictor | None = None,
+    predictor: MinerUClient | None = None,
     backend="transformers",
     model_path: str | None = None,
     server_url: str | None = None,
@@ -81,13 +162,13 @@ async def aio_doc_analyze(
         predictor = ModelSingleton().get_model(backend, model_path, server_url, **kwargs)
 
     # load_images_start = time.time()
-    images_list, pdf_doc = load_images_from_pdf(pdf_bytes, image_type=ImageType.BASE64)
-    images_base64_list = [image_dict["img_base64"] for image_dict in images_list]
+    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")
 
     # infer_start = time.time()
-    results = await predictor.aio_batch_predict(images=images_base64_list)
+    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")
     middle_json = result_to_middle_json(results, images_list, pdf_doc, image_writer)