diff --git a/mineru/backend/hybrid/hybrid_analyze.py b/mineru/backend/hybrid/hybrid_analyze.py index aaa548d0..9e364288 100644 --- a/mineru/backend/hybrid/hybrid_analyze.py +++ b/mineru/backend/hybrid/hybrid_analyze.py @@ -431,6 +431,7 @@ def _apply_flash_table_orientation_labels( rotate_labels = table_orientation_cls_model.batch_predict( table_inputs, det_batch_size=max(1, batch_ratio * OCR_DET_BASE_BATCH_SIZE), + tqdm_enable=True, ) if len(rotate_labels) != len(table_layout_refs): raise ValueError("Table orientation prediction result count mismatch") diff --git a/mineru/backend/pipeline/batch_analyze.py b/mineru/backend/pipeline/batch_analyze.py index ca18fd07..bc60237d 100644 --- a/mineru/backend/pipeline/batch_analyze.py +++ b/mineru/backend/pipeline/batch_analyze.py @@ -472,6 +472,7 @@ class BatchAnalyze: rotate_labels = table_orientation_cls_model.batch_predict( table_res_list_all_page, det_batch_size=self.batch_ratio * OCR_DET_BASE_BATCH_SIZE, + tqdm_enable=True, ) if len(rotate_labels) != len(table_res_list_all_page): raise ValueError( diff --git a/mineru/model/ocr/pytorch_paddle.py b/mineru/model/ocr/pytorch_paddle.py index 08db11ed..48743bfe 100644 --- a/mineru/model/ocr/pytorch_paddle.py +++ b/mineru/model/ocr/pytorch_paddle.py @@ -289,6 +289,7 @@ class PytorchPaddleOCR(TextSystem): mfd_res=None, tqdm_enable=False, tqdm_desc="OCR-rec Predict", + tqdm_progress_bar=None, ): assert isinstance(img, (np.ndarray, list, str, bytes)) if isinstance(img, list) and det == True: @@ -338,7 +339,12 @@ class PytorchPaddleOCR(TextSystem): if not isinstance(img, list): img = preprocess_image(img) img = [img] - rec_res, elapse = self.text_recognizer(img, tqdm_enable=tqdm_enable, tqdm_desc=tqdm_desc) + rec_res, elapse = self.text_recognizer( + img, + tqdm_enable=tqdm_enable, + tqdm_desc=tqdm_desc, + tqdm_progress_bar=tqdm_progress_bar, + ) # logger.debug("rec_res num : {}, elapsed : {}".format(len(rec_res), elapse)) ocr_res.append(rec_res) return ocr_res diff --git a/mineru/model/table/cls/mineru_table_ori_cls.py b/mineru/model/table/cls/mineru_table_ori_cls.py index 2207fd16..e31da4fb 100644 --- a/mineru/model/table/cls/mineru_table_ori_cls.py +++ b/mineru/model/table/cls/mineru_table_ori_cls.py @@ -5,6 +5,7 @@ from collections import defaultdict from typing import List, Dict import cv2 import numpy as np +from tqdm import tqdm # 旋转候选门控回到旧规则，先尽量召回疑似旋转表，再由 OCR rec 评分决定最终角度。 @@ -227,6 +228,26 @@ class MineruTableOrientationClsModel: return self._select_rotation_label_by_scores(score_by_label) + @staticmethod + def _set_progress_description(progress_bar, desc: str): + """切换复用进度条的阶段描述，并兼容测试替身和 tqdm 对象。""" + if progress_bar is None: + return + if hasattr(progress_bar, "set_description"): + progress_bar.set_description(desc) + else: + progress_bar.desc = desc + + @staticmethod + def _extend_progress_total(progress_bar, count: int): + """按新增工作量动态扩展总进度，保证最终 total 覆盖 det/score/rec。""" + if progress_bar is None or count <= 0: + return + current_total = progress_bar.total if progress_bar.total is not None else progress_bar.n + progress_bar.total = current_total + count + if hasattr(progress_bar, "refresh"): + progress_bar.refresh() + @classmethod def _collect_portrait_image_groups( cls, @@ -291,6 +312,7 @@ class MineruTableOrientationClsModel: resolution_groups: Dict[tuple[int, int], list[Dict]], det_batch_size: int, resolution_group_stride: int, + progress_bar=None, ) -> list[Dict]: """对表格批量做 OCR det，并筛选需要进入多角度评分的候选。""" rotated_imgs = [] @@ -304,11 +326,14 @@ class MineruTableOrientationClsModel: for img_info, (dt_boxes, _elapse) in zip(group_imgs, batch_results): if self._is_rotation_candidate_by_det_boxes(dt_boxes): rotated_imgs.append(img_info) + if progress_bar is not None: + progress_bar.update(len(group_imgs)) return rotated_imgs def _build_score_tasks_for_candidates( self, rotated_imgs: list[Dict], + progress_bar=None, ) -> tuple[list[tuple[Dict, list[Dict]]], list[np.ndarray]]: """为所有旋转候选构造三角度评分任务，并汇总成一次 OCR rec 输入。""" img_score_tasks = [] @@ -322,9 +347,17 @@ class MineruTableOrientationClsModel: task["crop_start"] = crop_start task["crop_end"] = crop_end img_score_tasks.append((img_info, tasks)) + if progress_bar is not None: + progress_bar.update(1) return img_score_tasks, all_crop_imgs - def _recognize_orientation_crops(self, all_crop_imgs: list[np.ndarray]) -> list: + def _recognize_orientation_crops( + self, + all_crop_imgs: list[np.ndarray], + tqdm_enable: bool = False, + tqdm_desc: str = "Table orientation", + tqdm_progress_bar=None, + ) -> list: """对所有候选角度 crop 合并执行 OCR rec，返回可按 slice 回填的结果。""" if not all_crop_imgs: return [] @@ -333,19 +366,38 @@ class MineruTableOrientationClsModel: all_crop_imgs, det=False, rec=True, + tqdm_enable=tqdm_enable and tqdm_progress_bar is None, + tqdm_desc=f"{tqdm_desc} rec", + tqdm_progress_bar=tqdm_progress_bar, ) return rec_ocr_res[0] if rec_ocr_res else [] - def _score_rotation_candidates(self, rotated_imgs: list[Dict]) -> Dict[int, str]: + def _score_rotation_candidates( + self, + rotated_imgs: list[Dict], + tqdm_enable: bool = False, + tqdm_desc: str = "Table orientation", + progress_bar=None, + ) -> Dict[int, str]: """批量评分旋转候选，并返回原始表格下标到最终角度标签的映射。""" if not rotated_imgs: return {} label_by_index = {} img_score_tasks, all_crop_imgs = self._build_score_tasks_for_candidates( - rotated_imgs + rotated_imgs, + progress_bar=progress_bar, + ) + self._extend_progress_total(progress_bar, len(all_crop_imgs)) + self._set_progress_description(progress_bar, f"{tqdm_desc} rec") + if progress_bar is not None and hasattr(progress_bar, "refresh"): + progress_bar.refresh() + rec_res = self._recognize_orientation_crops( + all_crop_imgs, + tqdm_enable=tqdm_enable, + tqdm_desc=tqdm_desc, + tqdm_progress_bar=progress_bar, ) - rec_res = self._recognize_orientation_crops(all_crop_imgs) for img_info, tasks in img_score_tasks: score_by_label = self._score_orientation_tasks_with_rec(tasks, rec_res) @@ -358,6 +410,8 @@ class MineruTableOrientationClsModel: self, imgs: List[Dict], det_batch_size: int, + tqdm_enable: bool = False, + tqdm_desc: str = "Table orientation", ) -> List[str]: """ 批量预测传入表格图片的旋转角度，只返回角度，不修改输入图片。 @@ -368,13 +422,38 @@ class MineruTableOrientationClsModel: imgs, RESOLUTION_GROUP_STRIDE, ) - rotated_imgs = self._detect_rotation_candidates( - resolution_groups, - det_batch_size, - RESOLUTION_GROUP_STRIDE, - ) - label_by_index = self._score_rotation_candidates(rotated_imgs) - for index, label in label_by_index.items(): - rotate_labels[index] = label + total_images = sum(len(group_imgs) for group_imgs in resolution_groups.values()) + progress_bar = None + if tqdm_enable: + progress_bar = tqdm( + total=total_images, + desc=f"{tqdm_desc} det", + leave=True, + ) + try: + rotated_imgs = self._detect_rotation_candidates( + resolution_groups, + det_batch_size, + RESOLUTION_GROUP_STRIDE, + progress_bar=progress_bar, + ) + self._extend_progress_total(progress_bar, len(rotated_imgs)) + self._set_progress_description(progress_bar, f"{tqdm_desc} score") + if progress_bar is not None and hasattr(progress_bar, "refresh"): + progress_bar.refresh() + label_by_index = self._score_rotation_candidates( + rotated_imgs, + tqdm_enable=tqdm_enable, + tqdm_desc=tqdm_desc, + progress_bar=progress_bar, + ) + for index, label in label_by_index.items(): + rotate_labels[index] = label + finally: + self._set_progress_description(progress_bar, tqdm_desc) + if progress_bar is not None and hasattr(progress_bar, "refresh"): + progress_bar.refresh() + if progress_bar is not None: + progress_bar.close() return rotate_labels diff --git a/mineru/model/utils/tools/infer/predict_rec.py b/mineru/model/utils/tools/infer/predict_rec.py index 8bcd8555..8ace3943 100644 --- a/mineru/model/utils/tools/infer/predict_rec.py +++ b/mineru/model/utils/tools/infer/predict_rec.py @@ -288,7 +288,13 @@ class TextRecognizer(BaseOCRV20): return img - def __call__(self, img_list, tqdm_enable=False, tqdm_desc="OCR-rec Predict"): + def __call__( + self, + img_list, + tqdm_enable=False, + tqdm_desc="OCR-rec Predict", + tqdm_progress_bar=None, + ): img_num = len(img_list) # Calculate the aspect ratio of all text bars width_list = [] @@ -301,8 +307,14 @@ class TextRecognizer(BaseOCRV20): rec_res = [['', 0.0]] * img_num batch_num = self.rec_batch_num elapse = 0 - # for beg_img_no in range(0, img_num, batch_num): - with tqdm(total=img_num, desc=tqdm_desc, disable=not tqdm_enable) as pbar: + # tqdm_progress_bar 由上层复用时，不再创建内部 OCR-rec 进度条。 + pbar = tqdm_progress_bar + should_close_pbar = False + if pbar is None: + pbar = tqdm(total=img_num, desc=tqdm_desc, disable=not tqdm_enable) + should_close_pbar = True + + try: index = 0 for beg_img_no in range(0, img_num, batch_num): end_img_no = min(img_num, beg_img_no + batch_num) @@ -428,6 +440,9 @@ class TextRecognizer(BaseOCRV20): current_batch_size = min(batch_num, img_num - index * batch_num) index += 1 pbar.update(current_batch_size) + finally: + if should_close_pbar: + pbar.close() # Fix NaN values in recognition results for i in range(len(rec_res)):