# 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