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

mineru/backend/hybrid/__init__.py

@@ -0,0 +1 @@
+#  Copyright (c) Opendatalab. All rights reserved.

mineru/backend/hybrid/hybrid_analyze.py

@@ -0,0 +1,526 @@
+#  Copyright (c) Opendatalab. All rights reserved.
+import os
+import time
+from collections import defaultdict
+
+import cv2
+import numpy as np
+from loguru import logger
+from mineru_vl_utils import MinerUClient
+from mineru_vl_utils.structs import BlockType
+from tqdm import tqdm
+
+from mineru.backend.hybrid.hybrid_model_output_to_middle_json import result_to_middle_json
+from mineru.backend.pipeline.model_init import HybridModelSingleton
+from mineru.backend.vlm.vlm_analyze import ModelSingleton
+from mineru.data.data_reader_writer import DataWriter
+from mineru.utils.config_reader import get_device
+from mineru.utils.enum_class import ImageType, NotExtractType
+from mineru.utils.model_utils import crop_img, get_vram, clean_memory
+from mineru.utils.ocr_utils import get_adjusted_mfdetrec_res, get_ocr_result_list, sorted_boxes, merge_det_boxes, \
+    update_det_boxes, OcrConfidence
+from mineru.utils.pdf_classify import classify
+from mineru.utils.pdf_image_tools import load_images_from_pdf
+
+os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'  # 让mps可以fallback
+os.environ['NO_ALBUMENTATIONS_UPDATE'] = '1'  # 禁止albumentations检查更新
+
+MFR_BASE_BATCH_SIZE = 16
+OCR_DET_BASE_BATCH_SIZE = 16
+
+not_extract_list = [item.value for item in NotExtractType]
+
+def ocr_classify(pdf_bytes, parse_method: str = 'auto',) -> bool:
+    # 确定OCR设置
+    _ocr_enable = False
+    if parse_method == 'auto':
+        if classify(pdf_bytes) == 'ocr':
+            _ocr_enable = True
+    elif parse_method == 'ocr':
+        _ocr_enable = True
+    return _ocr_enable
+
+def ocr_det(
+    hybrid_pipeline_model,
+    np_images,
+    results,
+    mfd_res,
+    _ocr_enable,
+    batch_radio: int = 1,
+):
+    ocr_res_list = []
+    if not hybrid_pipeline_model.enable_ocr_det_batch:
+        # 非批处理模式 - 逐页处理
+        for np_image, page_mfd_res, page_results in tqdm(
+            zip(np_images, mfd_res, results),
+            total=len(np_images),
+            desc="OCR-det"
+        ):
+            ocr_res_list.append([])
+            img_height, img_width = np_image.shape[:2]
+            for res in page_results:
+                if res['type'] not in not_extract_list:
+                    continue
+                x0 = max(0, int(res['bbox'][0] * img_width))
+                y0 = max(0, int(res['bbox'][1] * img_height))
+                x1 = min(img_width, int(res['bbox'][2] * img_width))
+                y1 = min(img_height, int(res['bbox'][3] * img_height))
+                if x1 <= x0 or y1 <= y0:
+                    continue
+                res['poly'] = [x0, y0, x1, y0, x1, y1, x0, y1]
+                new_image, useful_list = crop_img(
+                    res, np_image, crop_paste_x=50, crop_paste_y=50
+                )
+                adjusted_mfdetrec_res = get_adjusted_mfdetrec_res(
+                    page_mfd_res, useful_list
+                )
+                bgr_image = cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR)
+                ocr_res = hybrid_pipeline_model.ocr_model.ocr(
+                    bgr_image, mfd_res=adjusted_mfdetrec_res, rec=False
+                )[0]
+                if ocr_res:
+                    ocr_result_list = get_ocr_result_list(
+                        ocr_res, useful_list, _ocr_enable, bgr_image, hybrid_pipeline_model.lang
+                    )
+
+                    ocr_res_list[-1].extend(ocr_result_list)
+    else:
+        # 批处理模式 - 按语言和分辨率分组
+        # 收集所有需要OCR检测的裁剪图像
+        all_cropped_images_info = []
+
+        for np_image, page_mfd_res, page_results in zip(
+                np_images, mfd_res, results
+        ):
+            ocr_res_list.append([])
+            img_height, img_width = np_image.shape[:2]
+            for res in page_results:
+                if res['type'] not in not_extract_list:
+                    continue
+                x0 = max(0, int(res['bbox'][0] * img_width))
+                y0 = max(0, int(res['bbox'][1] * img_height))
+                x1 = min(img_width, int(res['bbox'][2] * img_width))
+                y1 = min(img_height, int(res['bbox'][3] * img_height))
+                if x1 <= x0 or y1 <= y0:
+                    continue
+                res['poly'] = [x0, y0, x1, y0, x1, y1, x0, y1]
+                new_image, useful_list = crop_img(
+                    res, np_image, crop_paste_x=50, crop_paste_y=50
+                )
+                adjusted_mfdetrec_res = get_adjusted_mfdetrec_res(
+                    page_mfd_res, useful_list
+                )
+                bgr_image = cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR)
+                all_cropped_images_info.append((
+                    bgr_image, useful_list, adjusted_mfdetrec_res, ocr_res_list[-1]
+                ))
+
+        # 按分辨率分组并同时完成padding
+        RESOLUTION_GROUP_STRIDE = 64  # 32
+
+        resolution_groups = defaultdict(list)
+        for crop_info in all_cropped_images_info:
+            cropped_img = crop_info[0]
+            h, w = cropped_img.shape[:2]
+            # 直接计算目标尺寸并用作分组键
+            target_h = ((h + RESOLUTION_GROUP_STRIDE - 1) // RESOLUTION_GROUP_STRIDE) * RESOLUTION_GROUP_STRIDE
+            target_w = ((w + RESOLUTION_GROUP_STRIDE - 1) // RESOLUTION_GROUP_STRIDE) * RESOLUTION_GROUP_STRIDE
+            group_key = (target_h, target_w)
+            resolution_groups[group_key].append(crop_info)
+
+        # 对每个分辨率组进行批处理
+        for (target_h, target_w), group_crops in tqdm(resolution_groups.items(), desc=f"OCR-det"):
+            # 对所有图像进行padding到统一尺寸
+            batch_images = []
+            for crop_info in group_crops:
+                img = crop_info[0]
+                h, w = img.shape[:2]
+                # 创建目标尺寸的白色背景
+                padded_img = np.ones((target_h, target_w, 3), dtype=np.uint8) * 255
+                padded_img[:h, :w] = img
+                batch_images.append(padded_img)
+
+            # 批处理检测
+            det_batch_size = min(len(batch_images), batch_radio*OCR_DET_BASE_BATCH_SIZE)
+            batch_results = hybrid_pipeline_model.ocr_model.text_detector.batch_predict(batch_images, det_batch_size)
+
+            # 处理批处理结果
+            for crop_info, (dt_boxes, _) in zip(group_crops, batch_results):
+                bgr_image, useful_list, adjusted_mfdetrec_res, ocr_page_res_list = crop_info
+
+                if dt_boxes is not None and len(dt_boxes) > 0:
+                    # 处理检测框
+                    dt_boxes_sorted = sorted_boxes(dt_boxes)
+                    dt_boxes_merged = merge_det_boxes(dt_boxes_sorted) if dt_boxes_sorted else []
+
+                    # 根据公式位置更新检测框
+                    dt_boxes_final = (update_det_boxes(dt_boxes_merged, adjusted_mfdetrec_res)
+                                      if dt_boxes_merged and adjusted_mfdetrec_res
+                                      else dt_boxes_merged)
+
+                    if dt_boxes_final:
+                        ocr_res = [box.tolist() if hasattr(box, 'tolist') else box for box in dt_boxes_final]
+                        ocr_result_list = get_ocr_result_list(
+                            ocr_res, useful_list, _ocr_enable, bgr_image, hybrid_pipeline_model.lang
+                        )
+                        ocr_page_res_list.extend(ocr_result_list)
+    return ocr_res_list
+
+def mask_image_regions(np_images, results):
+    # 根据vlm返回的结果，在每一页中将image、table、equation块mask成白色背景图像
+    for np_image, vlm_page_results in zip(np_images, results):
+        img_height, img_width = np_image.shape[:2]
+        # 收集需要mask的区域
+        mask_regions = []
+        for block in vlm_page_results:
+            if block['type'] in [BlockType.IMAGE, BlockType.TABLE, BlockType.EQUATION]:
+                bbox = block['bbox']
+                # 批量转换归一化坐标到像素坐标,并进行边界检查
+                x0 = max(0, int(bbox[0] * img_width))
+                y0 = max(0, int(bbox[1] * img_height))
+                x1 = min(img_width, int(bbox[2] * img_width))
+                y1 = min(img_height, int(bbox[3] * img_height))
+                # 只添加有效区域
+                if x1 > x0 and y1 > y0:
+                    mask_regions.append((y0, y1, x0, x1))
+        # 批量应用mask
+        for y0, y1, x0, x1 in mask_regions:
+            np_image[y0:y1, x0:x1, :] = 255
+    return np_images
+
+def normalize_poly_to_bbox(item, page_width, page_height):
+    """将poly坐标归一化为bbox"""
+    poly = item['poly']
+    x0 = min(max(poly[0] / page_width, 0), 1)
+    y0 = min(max(poly[1] / page_height, 0), 1)
+    x1 = min(max(poly[4] / page_width, 0), 1)
+    y1 = min(max(poly[5] / page_height, 0), 1)
+    item['bbox'] = [round(x0, 3), round(y0, 3), round(x1, 3), round(y1, 3)]
+    item.pop('poly', None)
+
+
+def _process_ocr_and_formulas(
+    images_pil_list,
+    results,
+    language,
+    inline_formula_enable,
+    _ocr_enable,
+    batch_radio: int = 1,
+):
+    """处理OCR和公式识别"""
+
+    # 遍历results,对文本块截图交由OCR识别
+    # 根据_ocr_enable决定ocr只开det还是det+rec
+    # 根据inline_formula_enable决定是使用mfd和ocr结合的方式,还是纯ocr方式
+
+    # 将PIL图片转换为numpy数组
+    np_images = [np.asarray(pil_image).copy() for pil_image in images_pil_list]
+
+    # 获取混合模型实例
+    hybrid_model_singleton = HybridModelSingleton()
+    hybrid_pipeline_model = hybrid_model_singleton.get_model(
+        lang=language,
+        formula_enable=inline_formula_enable,
+    )
+
+    if inline_formula_enable:
+        # 在进行`行内`公式检测和识别前，先将图像中的图片、表格、`行间`公式区域mask掉
+        np_images = mask_image_regions(np_images, results)
+        # 公式检测
+        images_mfd_res = hybrid_pipeline_model.mfd_model.batch_predict(np_images, batch_size=1, conf=0.5)
+        # 公式识别
+        inline_formula_list = hybrid_pipeline_model.mfr_model.batch_predict(
+            images_mfd_res,
+            np_images,
+            batch_size=batch_radio*MFR_BASE_BATCH_SIZE,
+            interline_enable=True,
+        )
+    else:
+        inline_formula_list = [[] for _ in range(len(images_pil_list))]
+
+    mfd_res = []
+    for page_inline_formula_list in inline_formula_list:
+        page_mfd_res = []
+        for formula in page_inline_formula_list:
+            formula['category_id'] = 13
+            page_mfd_res.append({
+                "bbox": [int(formula['poly'][0]), int(formula['poly'][1]),
+                         int(formula['poly'][4]), int(formula['poly'][5])],
+            })
+        mfd_res.append(page_mfd_res)
+
+    # vlm没有执行ocr，需要ocr_det
+    ocr_res_list = ocr_det(
+        hybrid_pipeline_model,
+        np_images,
+        results,
+        mfd_res,
+        _ocr_enable,
+        batch_radio=batch_radio,
+    )
+
+    # 如果需要ocr则做ocr_rec
+    if _ocr_enable:
+        need_ocr_list = []
+        img_crop_list = []
+        for page_ocr_res_list in ocr_res_list:
+            for ocr_res in page_ocr_res_list:
+                if 'np_img' in ocr_res:
+                    need_ocr_list.append(ocr_res)
+                    img_crop_list.append(ocr_res.pop('np_img'))
+        if len(img_crop_list) > 0:
+            # Process OCR
+            ocr_result_list = hybrid_pipeline_model.ocr_model.ocr(img_crop_list, det=False, tqdm_enable=True)[0]
+
+            # Verify we have matching counts
+            assert len(ocr_result_list) == len(need_ocr_list), f'ocr_result_list: {len(ocr_result_list)}, need_ocr_list: {len(need_ocr_list)}'
+
+            # Process OCR results for this language
+            for index, need_ocr_res in enumerate(need_ocr_list):
+                ocr_text, ocr_score = ocr_result_list[index]
+                need_ocr_res['text'] = ocr_text
+                need_ocr_res['score'] = float(f"{ocr_score:.3f}")
+                if ocr_score < OcrConfidence.min_confidence:
+                    need_ocr_res['category_id'] = 16
+                else:
+                    layout_res_bbox = [need_ocr_res['poly'][0], need_ocr_res['poly'][1],
+                                       need_ocr_res['poly'][4], need_ocr_res['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',
+                                '(cid:)', '(ci:)', '(cd:1)', 'cd:)', 'c)', '(cd:)', 'c', 'id:)',
+                                ':)', '√:)', '√i:)', '−i:)', '−:' , 'i:)',
+                            ]
+                            and ocr_score < 0.8
+                            and layout_res_width < layout_res_height
+                    ):
+                        need_ocr_res['category_id'] = 16
+
+    return inline_formula_list, ocr_res_list, hybrid_pipeline_model
+
+
+def _normalize_bbox(
+    inline_formula_list,
+    ocr_res_list,
+    images_pil_list,
+):
+    """归一化坐标并生成最终结果"""
+    for page_inline_formula_list, page_ocr_res_list, page_pil_image in zip(
+            inline_formula_list, ocr_res_list, images_pil_list
+    ):
+        if page_inline_formula_list or page_ocr_res_list:
+            page_width, page_height = page_pil_image.size
+            # 处理公式列表
+            for formula in page_inline_formula_list:
+                normalize_poly_to_bbox(formula, page_width, page_height)
+            # 处理OCR结果列表
+            for ocr_res in page_ocr_res_list:
+                normalize_poly_to_bbox(ocr_res, page_width, page_height)
+
+
+def get_batch_ratio(device):
+    """
+    根据显存大小或环境变量获取 batch ratio
+    """
+    # 1. 优先尝试从环境变量获取
+    """
+    c/s架构分离部署时，建议通过设置环境变量 MINERU_HYBRID_BATCH_RATIO 来指定 batch ratio
+    建议的设置值如如下，以下配置值已考虑一定的冗余，单卡多终端部署时为了保证稳定性，可以额外保留一个client端的显存作为整体冗余
+    单个client端显存大小 | MINERU_HYBRID_BATCH_RATIO
+    ------------------|------------------------
+    <= 6   GB         | 8
+    <= 4.5 GB         | 4
+    <= 3   GB         | 2
+    <= 2.5 GB         | 1
+    例如：
+    export MINERU_HYBRID_BATCH_RATIO=4
+    """
+    env_val = os.getenv("MINERU_HYBRID_BATCH_RATIO")
+    if env_val:
+        try:
+            batch_ratio = int(env_val)
+            logger.info(f"hybrid batch ratio (from env): {batch_ratio}")
+            return batch_ratio
+        except ValueError as e:
+            logger.warning(f"Invalid MINERU_HYBRID_BATCH_RATIO value: {env_val}, switching to auto mode. Error: {e}")
+
+    # 2. 根据显存自动推断
+    """
+    根据总显存大小粗略估计 batch ratio，需要排除掉vllm等推理框架占用的显存开销
+    """
+    gpu_memory = get_vram(device)
+    if gpu_memory >= 32:
+        batch_ratio = 16
+    elif gpu_memory >= 16:
+        batch_ratio = 8
+    elif gpu_memory >= 12:
+        batch_ratio = 4
+    elif gpu_memory >= 8:
+        batch_ratio = 2
+    else:
+        batch_ratio = 1
+
+    logger.info(f"hybrid batch ratio (auto, vram={gpu_memory}GB): {batch_ratio}")
+    return batch_ratio
+
+
+def _should_enable_vlm_ocr(ocr_enable: bool, language: str, inline_formula_enable: bool) -> bool:
+    """判断是否启用VLM OCR"""
+    force_enable = os.getenv("MINERU_FORCE_VLM_OCR_ENABLE", "0").lower() in ("1", "true", "yes")
+    if force_enable:
+        return True
+
+    force_pipeline = os.getenv("MINERU_HYBRID_FORCE_PIPELINE_ENABLE", "0").lower() in ("1", "true", "yes")
+    return (
+            ocr_enable
+            and language in ["ch", "en"]
+            and inline_formula_enable
+            and not force_pipeline
+    )
+
+
+def doc_analyze(
+        pdf_bytes,
+        image_writer: DataWriter | None,
+        predictor: MinerUClient | None = None,
+        backend="transformers",
+        parse_method: str = 'auto',
+        language: str = 'ch',
+        inline_formula_enable: bool = True,
+        model_path: str | None = None,
+        server_url: str | None = None,
+        **kwargs,
+):
+    # 初始化预测器
+    if predictor is None:
+        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.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")
+
+    # 获取设备信息
+    device = get_device()
+
+    # 确定OCR配置
+    _ocr_enable = ocr_classify(pdf_bytes, parse_method=parse_method)
+    _vlm_ocr_enable = _should_enable_vlm_ocr(_ocr_enable, language, inline_formula_enable)
+
+    infer_start = time.time()
+    # VLM提取
+    if _vlm_ocr_enable:
+        results = predictor.batch_two_step_extract(images=images_pil_list)
+        hybrid_pipeline_model = None
+        inline_formula_list = [[] for _ in images_pil_list]
+        ocr_res_list = [[] for _ in images_pil_list]
+    else:
+        batch_ratio = get_batch_ratio(device)
+        results = predictor.batch_two_step_extract(
+            images=images_pil_list,
+            not_extract_list=not_extract_list
+        )
+        inline_formula_list, ocr_res_list, hybrid_pipeline_model = _process_ocr_and_formulas(
+            images_pil_list,
+            results,
+            language,
+            inline_formula_enable,
+            _ocr_enable,
+            batch_radio=batch_ratio,
+        )
+        _normalize_bbox(inline_formula_list, ocr_res_list, images_pil_list)
+    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")
+
+    # 生成中间JSON
+    middle_json = result_to_middle_json(
+        results,
+        inline_formula_list,
+        ocr_res_list,
+        images_list,
+        pdf_doc,
+        image_writer,
+        _ocr_enable,
+        _vlm_ocr_enable,
+        hybrid_pipeline_model,
+    )
+
+    clean_memory(device)
+    return middle_json, results, _vlm_ocr_enable
+
+
+async def aio_doc_analyze(
+    pdf_bytes,
+    image_writer: DataWriter | None,
+    predictor: MinerUClient | None = None,
+    backend="transformers",
+    parse_method: str = 'auto',
+    language: str = 'ch',
+    inline_formula_enable: bool = True,
+    model_path: str | None = None,
+    server_url: str | None = None,
+    **kwargs,
+):
+    # 初始化预测器
+    if predictor is None:
+        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.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")
+
+    # 获取设备信息
+    device = get_device()
+
+    # 确定OCR配置
+    _ocr_enable = ocr_classify(pdf_bytes, parse_method=parse_method)
+    _vlm_ocr_enable = _should_enable_vlm_ocr(_ocr_enable, language, inline_formula_enable)
+
+    infer_start = time.time()
+    # VLM提取
+    if _vlm_ocr_enable:
+        results = await predictor.aio_batch_two_step_extract(images=images_pil_list)
+        hybrid_pipeline_model = None
+        inline_formula_list = [[] for _ in images_pil_list]
+        ocr_res_list = [[] for _ in images_pil_list]
+    else:
+        batch_ratio = get_batch_ratio(device)
+        results = await predictor.aio_batch_two_step_extract(
+            images=images_pil_list,
+            not_extract_list=not_extract_list
+        )
+        inline_formula_list, ocr_res_list, hybrid_pipeline_model = _process_ocr_and_formulas(
+            images_pil_list,
+            results,
+            language,
+            inline_formula_enable,
+            _ocr_enable,
+            batch_radio=batch_ratio,
+        )
+        _normalize_bbox(inline_formula_list, ocr_res_list, images_pil_list)
+    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")
+
+    # 生成中间JSON
+    middle_json = result_to_middle_json(
+        results,
+        inline_formula_list,
+        ocr_res_list,
+        images_list,
+        pdf_doc,
+        image_writer,
+        _ocr_enable,
+        _vlm_ocr_enable,
+        hybrid_pipeline_model,
+    )
+
+    clean_memory(device)
+    return middle_json, results, _vlm_ocr_enable
+