diff --git a/mineru/backend/pipeline/batch_analyze.py b/mineru/backend/pipeline/batch_analyze.py index 7cd6ddef..ca18fd07 100644 --- a/mineru/backend/pipeline/batch_analyze.py +++ b/mineru/backend/pipeline/batch_analyze.py @@ -57,16 +57,11 @@ class BatchAnalyze: table_ori_cls_batch_enabled: bool | None = None, text_ocr_det_batch_enabled: bool | None = None, mask_inline_formula_for_ocr_det: bool = True, - formula_recognition_scope: str = "all", - ocr_rec_enable: bool = True, - seal_ocr_rec_enable: bool = True, ): self.batch_ratio = batch_ratio self.formula_enable = get_formula_enable(formula_enable) self.table_enable = get_table_enable(table_enable) self.model_manager = model_manager - self.ocr_rec_enable = ocr_rec_enable - self.seal_ocr_rec_enable = seal_ocr_rec_enable self.enable_ocr_det_batch = enable_ocr_det_batch self.table_ori_cls_batch_enabled = ( enable_ocr_det_batch if table_ori_cls_batch_enabled is None else table_ori_cls_batch_enabled @@ -77,10 +72,6 @@ class BatchAnalyze: self.mask_inline_formula_for_ocr_det = ( get_ocr_det_mask_inline_formula_enable(mask_inline_formula_for_ocr_det) ) - if formula_recognition_scope not in {"all", "inline_only", "none"}: - raise ValueError(f"Unsupported formula_recognition_scope: {formula_recognition_scope}") - # 控制公式识别范围，默认保持pipeline原行为；hybrid-flash可只保留公式det而跳过MFR。 - self.formula_recognition_scope = formula_recognition_scope @staticmethod def _apply_mask_boxes_to_image( @@ -268,9 +259,6 @@ class BatchAnalyze: return match.expand(replacement) return text - def _formula_recognition_enabled(self) -> bool: - return self.formula_enable and self.formula_recognition_scope != "none" - @classmethod def _extract_table_inline_objects( cls, @@ -369,7 +357,7 @@ class BatchAnalyze: self.model = self.model_manager.get_model( lang=None, - formula_enable=self._formula_recognition_enabled(), + formula_enable=self.formula_enable, table_enable=self.table_enable, ) atom_model_manager = AtomModelSingleton() @@ -388,52 +376,38 @@ class BatchAnalyze: clean_vram(self.model.device, vram_threshold=8) if self.formula_enable: - all_formula_labels = ["display_formula", "inline_formula"] - formula_labels = all_formula_labels - if self.formula_recognition_scope == "inline_only": - formula_labels = ["inline_formula"] images_mfd_res = [] for layout_res in images_layout_res: page_formula_res = [] for res in layout_res: - if res.get("label") in all_formula_labels: + if res.get("label") in ["display_formula", "inline_formula"]: res.setdefault("latex", "") - if res.get("label") in formula_labels: page_formula_res.append(res) images_mfd_res.append(page_formula_res) - if self.formula_recognition_scope != "none": - # 公式识别 - images_formula_list = run_mfr_inference( - self.model.mfr_model.batch_predict, - images_mfd_res, - np_images, - batch_size=self.batch_ratio * MFR_BASE_BATCH_SIZE, - ) - mfr_count = 0 - for image_index in range(len(np_images)): - mfr_count += len(images_formula_list[image_index]) - for formula_res, formula_with_latex in zip( - images_mfd_res[image_index], images_formula_list[image_index] - ): - formula_res["latex"] = formula_with_latex.get("latex", "") + # 公式识别 + images_formula_list = run_mfr_inference( + self.model.mfr_model.batch_predict, + images_mfd_res, + np_images, + batch_size=self.batch_ratio * MFR_BASE_BATCH_SIZE, + ) + mfr_count = 0 + for image_index in range(len(np_images)): + mfr_count += len(images_formula_list[image_index]) + for formula_res, formula_with_latex in zip( + images_mfd_res[image_index], images_formula_list[image_index] + ): + formula_res["latex"] = formula_with_latex.get("latex", "") - # 清理显存 - clean_vram(self.model.device, vram_threshold=8) - else: - for page_formula_res in images_mfd_res: - for formula_res in page_formula_res: - formula_res["latex"] = "" + # 清理显存 + clean_vram(self.model.device, vram_threshold=8) else: for layout_res in images_layout_res: # 移除所有的"inline_formula" layout_res[:] = [res for res in layout_res if res.get("label") != "inline_formula"] - - - ocr_should_recognize_text = bool(self.ocr_rec_enable) - ocr_res_list_all_page = [] table_res_list_all_page = [] for index in range(len(np_images)): @@ -784,7 +758,7 @@ class BatchAnalyze: ocr_result_list = get_ocr_result_list( ocr_res, useful_list, - ocr_res_list_dict['ocr_enable'] and ocr_should_recognize_text, + ocr_res_list_dict['ocr_enable'], bgr_image, _lang, ) @@ -828,7 +802,7 @@ class BatchAnalyze: ocr_result_list = get_ocr_result_list( ocr_res, useful_list, - ocr_res_list_dict['ocr_enable'] and ocr_should_recognize_text, + ocr_res_list_dict['ocr_enable'], bgr_image, _lang, ) @@ -917,11 +891,10 @@ class BatchAnalyze: total_processed += len(img_crop_list) seal_ocr_items = [] - if self.seal_ocr_rec_enable: - for ocr_res_list_dict in ocr_res_list_all_page: - for layout_res_item in ocr_res_list_dict['layout_res']: - if layout_res_item.get("label") == "seal": - seal_ocr_items.append((ocr_res_list_dict, layout_res_item)) + for ocr_res_list_dict in ocr_res_list_all_page: + for layout_res_item in ocr_res_list_dict['layout_res']: + if layout_res_item.get("label") == "seal": + seal_ocr_items.append((ocr_res_list_dict, layout_res_item)) seal_ocr_model = None for ocr_res_list_dict, layout_res_item in tqdm(seal_ocr_items, desc="Seal Predict"): @@ -969,7 +942,7 @@ class BatchAnalyze: for ocr_res_list_dict in ocr_res_list_all_page: self._prune_empty_ocr_text_blocks( ocr_res_list_dict["layout_res"], - ocr_res_list_dict["ocr_enable"] and self.ocr_rec_enable, + ocr_res_list_dict["ocr_enable"], ) return images_layout_res diff --git a/mineru/backend/pipeline/pipeline_analyze.py b/mineru/backend/pipeline/pipeline_analyze.py index e8670d18..b288712b 100644 --- a/mineru/backend/pipeline/pipeline_analyze.py +++ b/mineru/backend/pipeline/pipeline_analyze.py @@ -331,10 +331,7 @@ def doc_analyze_streaming( def batch_image_analyze( images_with_extra_info: List[Tuple[Image.Image, bool, str]], formula_enable=True, - table_enable=True, - formula_recognition_scope="all", - ocr_rec_enable=True, - seal_ocr_rec_enable=True): + table_enable=True): from .batch_analyze import BatchAnalyze @@ -385,9 +382,6 @@ def batch_image_analyze( formula_enable, table_enable, enable_ocr_det_batch, - formula_recognition_scope=formula_recognition_scope, - ocr_rec_enable=ocr_rec_enable, - seal_ocr_rec_enable=seal_ocr_rec_enable, ) results = batch_model(images_with_extra_info)