feat: enhance progress tracking in batch analysis and OCR processes with tqdm integration

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")

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(

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

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

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)):