Merge pull request #5139 from myhloli/dev

Optimize OCR batch processing and enhance detection accuracy

mineru/backend/hybrid/hybrid_analyze.py

@@ -257,63 +257,39 @@ def ocr_det(
                     bgr_image, det_image, useful_list, adjusted_mfdetrec_res, ocr_res_list[-1]
                 ))
 
-        # 按分辨率分组并同时完成padding
-        RESOLUTION_GROUP_STRIDE = 64  # 32
+        batch_images = [crop_info[1] for crop_info in all_cropped_images_info]
+        det_batch_size = min(len(batch_images), batch_ratio * OCR_DET_BASE_BATCH_SIZE)
+        batch_results = run_ocr_inference(
+            hybrid_pipeline_model.ocr_model.text_detector.batch_predict,
+            batch_images,
+            det_batch_size,
+            tqdm_enable=True,
+            tqdm_desc="OCR-det",
+        )
 
-        resolution_groups = defaultdict(list)
-        for crop_info in all_cropped_images_info:
-            cropped_img = crop_info[1]
-            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 crop_info, (dt_boxes, _) in zip(all_cropped_images_info, batch_results):
+            bgr_image, _det_image, useful_list, adjusted_mfdetrec_res, ocr_page_res_list = crop_info
 
-        # 对每个分辨率组进行批处理
-        for (target_h, target_w), group_crops in tqdm(resolution_groups.items(), desc="OCR-det"):
-            # 对所有图像进行padding到统一尺寸
-            batch_images = []
-            for crop_info in group_crops:
-                img = crop_info[1]
-                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)
+            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 []
 
-            # 批处理检测
-            det_batch_size = min(len(batch_images), batch_ratio * OCR_DET_BASE_BATCH_SIZE)
-            batch_results = run_ocr_inference(
-                hybrid_pipeline_model.ocr_model.text_detector.batch_predict,
-                batch_images,
-                det_batch_size,
-            )
+                # 根据公式位置更新检测框
+                dt_boxes_final = (update_det_boxes(dt_boxes_merged, adjusted_mfdetrec_res)
+                                  if dt_boxes_merged and adjusted_mfdetrec_res
+                                  else dt_boxes_merged)
 
-            # 处理批处理结果
-            for crop_info, (dt_boxes, _) in zip(group_crops, batch_results):
-                bgr_image, _det_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,
-                            False,
-                            bgr_image,
-                            hybrid_pipeline_model.lang,
-                        )
-                        ocr_page_res_list.extend(ocr_result_list)
+                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,
+                        False,
+                        bgr_image,
+                        hybrid_pipeline_model.lang,
+                    )
+                    ocr_page_res_list.extend(ocr_result_list)
     return ocr_res_list
 
 

mineru/backend/pipeline/batch_analyze.py

@@ -59,6 +59,7 @@ class BatchAnalyze:
         enable_ocr_det_batch: bool = True,
         table_ori_cls_batch_enabled: bool | None = None,
         text_ocr_det_batch_enabled: bool | None = None,
+        table_ocr_det_batch_enabled: bool | None = None,
         mask_inline_formula_for_ocr_det: bool = True,
     ):
         self.batch_ratio = batch_ratio
@@ -72,6 +73,9 @@ class BatchAnalyze:
         self.text_ocr_det_batch_enabled = (
             enable_ocr_det_batch if text_ocr_det_batch_enabled is None else text_ocr_det_batch_enabled
         )
+        self.table_ocr_det_batch_enabled = (
+            enable_ocr_det_batch if table_ocr_det_batch_enabled is None else table_ocr_det_batch_enabled
+        )
         self.mask_inline_formula_for_ocr_det = (
             get_ocr_det_mask_inline_formula_enable(mask_inline_formula_for_ocr_det)
         )
@@ -92,6 +96,72 @@ class BatchAnalyze:
             return bgr_image
         return self._apply_mask_boxes_to_image(bgr_image, mask_boxes)
 
+    def _build_table_ocr_det_items(self, table_res_list_all_page: list[dict]) -> list[dict]:
+        """构造表格 OCR-det 输入项，保留原图、遮罩图和后续回填所需信息。"""
+        table_det_items = []
+        for index, table_res_dict in enumerate(table_res_list_all_page):
+            bgr_image = cv2.cvtColor(table_res_dict["table_img"], cv2.COLOR_RGB2BGR)
+            table_inline_objects = (
+                table_res_dict.get("table_inline_objects", [])
+                if self._table_supports_inline_objects(table_res_dict)
+                else []
+            )
+            inline_mask_boxes = [
+                {"bbox": inline_object["table_rel_mask_bbox"]}
+                for inline_object in table_inline_objects
+            ]
+            formula_mask_boxes = [
+                {"bbox": inline_object["table_rel_mask_bbox"]}
+                for inline_object in table_inline_objects
+                if inline_object["kind"] == "formula"
+            ]
+            det_image = (
+                self._apply_mask_boxes_to_image(bgr_image, inline_mask_boxes)
+                if inline_mask_boxes
+                else bgr_image
+            )
+            table_det_items.append(
+                {
+                    "bgr_image": bgr_image,
+                    "det_image": det_image,
+                    "formula_mask_boxes": formula_mask_boxes,
+                    "lang": table_res_dict["lang"],
+                    "table_id": index,
+                }
+            )
+        return table_det_items
+
+    def _append_table_ocr_det_result(
+        self,
+        table_det_item: dict,
+        dt_boxes,
+        rec_img_lang_group: dict,
+    ) -> None:
+        """将单表 OCR-det 结果整理成 OCR-rec 输入，并保持表格回填顺序。"""
+        if dt_boxes is None or len(dt_boxes) == 0:
+            return
+
+        ocr_result = dt_boxes
+        formula_mask_boxes = table_det_item["formula_mask_boxes"]
+        if formula_mask_boxes:
+            ocr_result = update_det_boxes(ocr_result, formula_mask_boxes)
+            if not ocr_result:
+                return
+
+        ocr_result = sorted_boxes(ocr_result)
+        for dt_box in ocr_result:
+            dt_box_array = np.asarray(dt_box, dtype=np.float32)
+            rec_img_lang_group.setdefault(table_det_item["lang"], []).append(
+                {
+                    "cropped_img": get_rotate_crop_image_for_text_rec(
+                        table_det_item["bgr_image"],
+                        dt_box_array.copy(),
+                    ),
+                    "dt_box": dt_box_array.copy(),
+                    "table_id": table_det_item["table_id"],
+                }
+            )
+
     @staticmethod
     def _prune_empty_ocr_text_blocks(layout_res: list[dict], ocr_enable: bool) -> None:
         if not ocr_enable or not layout_res:
@@ -487,7 +557,7 @@ class BatchAnalyze:
                     f"Table classification failed: {e}, using default model"
                 )
 
-            # OCR det 过程，顺序执行
+            # OCR det 过程，默认使用 detector 内部分桶 batch，关闭开关时回退逐表单张路径。
             rec_img_lang_group = defaultdict(list)
             det_ocr_engine = atom_model_manager.get_atom_model(
                 atom_model_name=AtomicModel.OCR,
@@ -495,46 +565,40 @@ class BatchAnalyze:
                 det_db_unclip_ratio=1.6,
                 enable_merge_det_boxes=False,
             )
-            for index, table_res_dict in enumerate(
-                    tqdm(table_res_list_all_page, desc="Table-ocr det")
-            ):
-                bgr_image = cv2.cvtColor(table_res_dict["table_img"], cv2.COLOR_RGB2BGR)
-                table_inline_objects = (
-                    table_res_dict.get("table_inline_objects", [])
-                    if self._table_supports_inline_objects(table_res_dict)
-                    else []
-                )
-                inline_mask_boxes = [
-                    {"bbox": inline_object["table_rel_mask_bbox"]}
-                    for inline_object in table_inline_objects
-                ]
-                formula_mask_boxes = [
-                    {"bbox": inline_object["table_rel_mask_bbox"]}
-                    for inline_object in table_inline_objects
-                    if inline_object["kind"] == "formula"
-                ]
-                det_image = (
-                    self._apply_mask_boxes_to_image(bgr_image, inline_mask_boxes)
-                    if inline_mask_boxes
-                    else bgr_image
-                )
-                ocr_result = run_ocr_inference(
-                    det_ocr_engine.ocr, det_image, rec=False
-                )[0]
-                if ocr_result and formula_mask_boxes:
-                    ocr_result = update_det_boxes(ocr_result, formula_mask_boxes)
-                if ocr_result:
-                    ocr_result = sorted_boxes(ocr_result)
-                # 构造需要 OCR 识别的图片字典，包括cropped_img, dt_box, table_id，并按照语言进行分组
-                for dt_box in ocr_result:
-                    rec_img_lang_group[table_res_dict["lang"]].append(
-                        {
-                            "cropped_img": get_rotate_crop_image_for_text_rec(
-                                bgr_image, np.asarray(dt_box, dtype=np.float32)
-                            ),
-                            "dt_box": np.asarray(dt_box, dtype=np.float32),
-                            "table_id": index,
-                        }
+            table_det_items = self._build_table_ocr_det_items(table_res_list_all_page)
+            if self.table_ocr_det_batch_enabled:
+                det_images = [table_det_item["det_image"] for table_det_item in table_det_items]
+                if det_images:
+                    det_batch_size = max(
+                        1,
+                        min(len(det_images), self.batch_ratio * OCR_DET_BASE_BATCH_SIZE),
+                    )
+                    batch_results = run_ocr_inference(
+                        det_ocr_engine.text_detector.batch_predict,
+                        det_images,
+                        det_batch_size,
+                        tqdm_enable=True,
+                        tqdm_desc="Table-ocr det",
+                    )
+                    if len(batch_results) != len(table_det_items):
+                        raise ValueError("Table OCR det batch result count mismatch")
+                    for table_det_item, (dt_boxes, _) in zip(table_det_items, batch_results):
+                        self._append_table_ocr_det_result(
+                            table_det_item,
+                            dt_boxes,
+                            rec_img_lang_group,
+                        )
+            else:
+                for table_det_item in tqdm(table_det_items, desc="Table-ocr det"):
+                    ocr_result = run_ocr_inference(
+                        det_ocr_engine.ocr,
+                        table_det_item["det_image"],
+                        rec=False,
+                    )[0]
+                    self._append_table_ocr_det_result(
+                        table_det_item,
+                        ocr_result,
+                        rec_img_lang_group,
                     )
 
             # OCR rec，按照语言分批处理
@@ -687,69 +751,48 @@ class BatchAnalyze:
                     lang=lang
                 )
 
-                # 按分辨率分组并同时完成padding
-                # RESOLUTION_GROUP_STRIDE = 32
-                RESOLUTION_GROUP_STRIDE = 64
+                batch_images = [crop_info[1] for crop_info in lang_crop_list]
+                det_batch_size = min(
+                    len(batch_images), self.batch_ratio * OCR_DET_BASE_BATCH_SIZE
+                )
+                batch_results = run_ocr_inference(
+                    ocr_model.text_detector.batch_predict,
+                    batch_images,
+                    det_batch_size,
+                    tqdm_enable=True,
+                    tqdm_desc=f"OCR-det {lang}",
+                )
 
-                resolution_groups = defaultdict(list)
-                for crop_info in lang_crop_list:
-                    cropped_img = crop_info[1]
-                    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 crop_info, (dt_boxes, _) in zip(lang_crop_list, batch_results):
+                    (
+                        bgr_image,
+                        _det_image,
+                        useful_list,
+                        ocr_res_list_dict,
+                        adjusted_mfdetrec_res,
+                        _lang,
+                    ) = crop_info
 
-                # 对每个分辨率组进行批处理
-                for (target_h, target_w), group_crops in tqdm(resolution_groups.items(), desc=f"OCR-det {lang}"):
-                    # 对所有图像进行padding到统一尺寸
-                    batch_images = []
-                    for crop_info in group_crops:
-                        img = crop_info[1]
-                        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)
+                    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 []
 
-                    # 批处理检测
-                    det_batch_size = min(len(batch_images), self.batch_ratio * OCR_DET_BASE_BATCH_SIZE)
-                    batch_results = run_ocr_inference(
-                        ocr_model.text_detector.batch_predict, batch_images, det_batch_size
-                    )
+                        # 根据公式位置更新检测框
+                        dt_boxes_final = (update_det_boxes(dt_boxes_merged, adjusted_mfdetrec_res)
+                                          if dt_boxes_merged and adjusted_mfdetrec_res
+                                          else dt_boxes_merged)
 
-                    # 处理批处理结果
-                    for crop_info, (dt_boxes, _) in zip(group_crops, batch_results):
-                        (
-                            bgr_image,
-                            _det_image,
-                            useful_list,
-                            ocr_res_list_dict,
-                            adjusted_mfdetrec_res,
-                            _lang,
-                        ) = 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_res_list_dict['ocr_enable'],
-                                    bgr_image,
-                                    _lang,
-                                )
-                                ocr_res_list_dict['layout_res'].extend(ocr_result_list)
+                        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_res_list_dict['ocr_enable'],
+                                bgr_image,
+                                _lang,
+                            )
+                            ocr_res_list_dict['layout_res'].extend(ocr_result_list)
 
             # 清理显存
             clean_vram(self.model.device, vram_threshold=8)

mineru/model/ocr/pytorch_paddle.py

@@ -245,7 +245,7 @@ class PytorchPaddleOCR(TextSystem):
             logger.debug("no valid image provided")
             return None, None
 
-        ori_im = img.copy()
+        ori_im = img
         dt_boxes, elapse = self.text_detector(img)
 
         if dt_boxes is None:

mineru/model/table/cls/mineru_table_ori_cls.py

@@ -2,9 +2,11 @@
 
 from PIL import Image
 from collections import defaultdict
+import inspect
 from typing import List, Dict
 import cv2
 import numpy as np
+from loguru import logger
 from tqdm import tqdm
 
 
@@ -16,7 +18,7 @@ ROTATED_TEXT_MIN_BOXES = 3
 ORIENTATION_SCORE_MAX_SAMPLE_BOXES = 18
 ORIENTATION_SCORE_MIN_VALID_RESULTS = 5
 ORIENTATION_ZERO_SCORE_PRIORITY_THRESHOLD = 0.9
-ORIENTATION_SCORE_TIE_THRESHOLD = 0.1
+ORIENTATION_SCORE_TIE_THRESHOLD = 0.08
 ORIENTATION_SCORE_LABELS = ("0", "90", "270")
 
 
@@ -128,11 +130,14 @@ class MineruTableOrientationClsModel:
             return None
         return img[ymin:ymax, xmin:xmax].copy()
 
-    def _build_orientation_score_task(self, label: str, img_bgr: np.ndarray) -> Dict:
-        """为单个角度构造评分任务，只做 det、抽样和切图，不执行 rec。"""
-        det_ocr_res = self.ocr_engine.ocr(img_bgr, rec=False)
-        det_res = det_ocr_res[0] if det_ocr_res else None
-        sampled_boxes = self._sample_det_boxes(det_res)
+    def _build_orientation_score_task_from_det_boxes(
+        self,
+        label: str,
+        img_bgr: np.ndarray,
+        det_boxes,
+    ) -> Dict:
+        """根据已有 OCR det 框构造评分任务，复用 0 度门控结果并统一裁图逻辑。"""
+        sampled_boxes = self._sample_det_boxes(det_boxes)
 
         img_crop_list = []
         for box in sampled_boxes:
@@ -148,6 +153,16 @@ class MineruTableOrientationClsModel:
             "crop_end": len(img_crop_list),
         }
 
+    def _build_orientation_score_task(self, label: str, img_bgr: np.ndarray) -> Dict:
+        """为单个角度构造评分任务，只做 det、抽样和切图，不执行 rec。"""
+        det_ocr_res = self.ocr_engine.ocr(img_bgr, rec=False)
+        det_res = det_ocr_res[0] if det_ocr_res else None
+        return self._build_orientation_score_task_from_det_boxes(
+            label,
+            img_bgr,
+            det_res,
+        )
+
     def _build_orientation_score_tasks(self, img_bgr: np.ndarray) -> List[Dict]:
         """为一张表构造 0/90/270 三个角度的评分任务。"""
         tasks = []
@@ -181,12 +196,33 @@ class MineruTableOrientationClsModel:
         score_by_label = {}
         rec_res = rec_res or []
         for task in tasks:
+            if "score" in task:
+                score_by_label[task["label"]] = task["score"]
+                continue
             crop_start = task.get("crop_start", 0)
             crop_end = task.get("crop_end", crop_start + task.get("crop_count", 0))
             task_rec_res = rec_res[crop_start:crop_end]
             score_by_label[task["label"]] = self._score_rec_results(task_rec_res)
         return score_by_label
 
+    @staticmethod
+    def _debug_log_orientation_rec_scores(
+        table_index: int,
+        score_by_label: Dict[str, tuple[float, int, int]],
+    ) -> None:
+        """输出单张表格各旋转候选的 OCR-rec 分数，便于排查误旋转。"""
+        score_parts = []
+        for label in ORIENTATION_SCORE_LABELS:
+            score, valid_count, char_count = score_by_label.get(label, (0.0, 0, 0))
+            score_parts.append(
+                f"{label} score={float(score):.4f} "
+                f"valid_count={valid_count} char_count={char_count}"
+            )
+        logger.debug(
+            f"Table orientation rec scores table_index={table_index}: "
+            f"{'; '.join(score_parts)}"
+        )
+
     def _score_rotation_candidate_by_ocr(self, img_bgr: np.ndarray) -> tuple[float, int, int]:
         """对单个候选角度执行 OCR det+抽样 rec，返回平均置信度、有效文本数和字符数。"""
         task = self._build_orientation_score_task("", img_bgr)
@@ -226,6 +262,7 @@ class MineruTableOrientationClsModel:
             rotated_img = self._rotate_image_by_label(img_bgr, label)
             score_by_label[label] = self._score_rotation_candidate_by_ocr(rotated_img)
 
+        self._debug_log_orientation_rec_scores(-1, score_by_label)
         return self._select_rotation_label_by_scores(score_by_label)
 
     @staticmethod
@@ -286,6 +323,24 @@ class MineruTableOrientationClsModel:
             )
         return resolution_groups
 
+    @classmethod
+    def _collect_orientation_images(cls, imgs: List[Dict]) -> list[Dict]:
+        """扁平收集有效表格图，首轮 det 的分桶和 batch 交给 OCR detector 内部处理。"""
+        orientation_imgs = []
+        for index, img in enumerate(imgs):
+            bgr_img = cls._to_bgr_table_image(img)
+            img_height, img_width = bgr_img.shape[:2]
+            if img_height <= 0 or img_width <= 0:
+                continue
+
+            orientation_imgs.append(
+                {
+                    "index": index,
+                    "table_img_bgr": bgr_img,
+                }
+            )
+        return orientation_imgs
+
     @classmethod
     def _pad_group_images(
         cls,
@@ -307,46 +362,136 @@ class MineruTableOrientationClsModel:
             batch_images.append(padded_img)
         return batch_images
 
+    def _batch_detect_text_boxes(
+        self,
+        img_list: list[np.ndarray],
+        det_batch_size: int,
+        tqdm_enable: bool = False,
+        tqdm_desc: str = "OCR-det Predict",
+        progress_bar=None,
+    ):
+        """统一调用 OCR detector batch_predict，并兼容不支持进度参数的测试替身。"""
+        if not img_list:
+            return []
+
+        max_batch_size = max(1, min(len(img_list), int(det_batch_size)))
+        batch_predict = self.ocr_engine.text_detector.batch_predict
+
+        progress_kwargs = {}
+        try:
+            signature = inspect.signature(batch_predict)
+            params = signature.parameters
+        except (TypeError, ValueError):
+            params = {}
+
+        if "tqdm_enable" in params:
+            progress_kwargs["tqdm_enable"] = tqdm_enable
+        if "tqdm_desc" in params:
+            progress_kwargs["tqdm_desc"] = tqdm_desc
+        if "tqdm_progress_bar" in params:
+            progress_kwargs["tqdm_progress_bar"] = progress_bar
+
+        batch_results = batch_predict(img_list, max_batch_size, **progress_kwargs)
+
+        if progress_bar is not None and "tqdm_progress_bar" not in progress_kwargs:
+            progress_bar.update(len(img_list))
+
+        return batch_results
+
     def _detect_rotation_candidates(
         self,
-        resolution_groups: Dict[tuple[int, int], list[Dict]],
+        orientation_imgs: list[Dict],
         det_batch_size: int,
-        resolution_group_stride: int,
+        tqdm_enable: bool = False,
+        tqdm_desc: str = "Table orientation",
         progress_bar=None,
     ) -> list[Dict]:
         """对表格批量做 OCR det，并筛选需要进入多角度评分的候选。"""
         rotated_imgs = []
-        for _group_key, group_imgs in resolution_groups.items():
-            batch_images = self._pad_group_images(group_imgs, resolution_group_stride)
-            batch_results = self.ocr_engine.text_detector.batch_predict(
-                batch_images,
-                max(1, min(len(batch_images), det_batch_size)),
-            )
+        batch_images = [img_info["table_img_bgr"] for img_info in orientation_imgs]
+        batch_results = self._batch_detect_text_boxes(
+            batch_images,
+            det_batch_size,
+            tqdm_enable=tqdm_enable and progress_bar is None,
+            tqdm_desc=f"{tqdm_desc} det",
+            progress_bar=progress_bar,
+        )
 
-            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))
+        for img_info, (dt_boxes, _elapse) in zip(orientation_imgs, batch_results):
+            if not self._is_rotation_candidate_by_det_boxes(dt_boxes):
+                continue
+            candidate_info = dict(img_info)
+            candidate_info["gate_det_boxes"] = dt_boxes
+            rotated_imgs.append(candidate_info)
         return rotated_imgs
 
+    @staticmethod
+    def _add_score_task_crops(task: Dict, all_crop_imgs: list[np.ndarray]) -> None:
+        """将有效评分任务加入 OCR-rec 输入，不足阈值的任务直接记为 0 分。"""
+        if task["crop_count"] < ORIENTATION_SCORE_MIN_VALID_RESULTS:
+            task["score"] = (0.0, 0, 0)
+            task["crop_start"] = len(all_crop_imgs)
+            task["crop_end"] = len(all_crop_imgs)
+            return
+
+        crop_start = len(all_crop_imgs)
+        all_crop_imgs.extend(task["crops"])
+        crop_end = len(all_crop_imgs)
+        task["crop_start"] = crop_start
+        task["crop_end"] = crop_end
+
     def _build_score_tasks_for_candidates(
         self,
         rotated_imgs: list[Dict],
+        det_batch_size: int,
         progress_bar=None,
     ) -> tuple[list[tuple[Dict, list[Dict]]], list[np.ndarray]]:
         """为所有旋转候选构造三角度评分任务，并汇总成一次 OCR rec 输入。"""
         img_score_tasks = []
         all_crop_imgs = []
+        score_det_images = []
+        score_det_tasks = []
+
         for img_info in rotated_imgs:
-            tasks = self._build_orientation_score_tasks(img_info["table_img_bgr"])
-            for task in tasks:
-                crop_start = len(all_crop_imgs)
-                all_crop_imgs.extend(task["crops"])
-                crop_end = len(all_crop_imgs)
-                task["crop_start"] = crop_start
-                task["crop_end"] = crop_end
+            table_img_bgr = img_info["table_img_bgr"]
+            tasks = [
+                self._build_orientation_score_task_from_det_boxes(
+                    "0",
+                    table_img_bgr,
+                    img_info.get("gate_det_boxes"),
+                )
+            ]
+            for label in ("90", "270"):
+                rotated_img = self._rotate_image_by_label(table_img_bgr, label)
+                task = {
+                    "label": label,
+                    "rotated_img_bgr": rotated_img,
+                    "crops": [],
+                    "crop_count": 0,
+                    "crop_start": 0,
+                    "crop_end": 0,
+                }
+                tasks.append(task)
+                score_det_images.append(rotated_img)
+                score_det_tasks.append(task)
             img_score_tasks.append((img_info, tasks))
+
+        score_det_results = self._batch_detect_text_boxes(
+            score_det_images,
+            det_batch_size,
+        )
+        for task, (dt_boxes, _elapse) in zip(score_det_tasks, score_det_results):
+            score_task = self._build_orientation_score_task_from_det_boxes(
+                task["label"],
+                task["rotated_img_bgr"],
+                dt_boxes,
+            )
+            task.update(score_task)
+            task.pop("rotated_img_bgr", None)
+
+        for _img_info, tasks in img_score_tasks:
+            for task in tasks:
+                self._add_score_task_crops(task, all_crop_imgs)
             if progress_bar is not None:
                 progress_bar.update(1)
         return img_score_tasks, all_crop_imgs
@@ -375,6 +520,7 @@ class MineruTableOrientationClsModel:
     def _score_rotation_candidates(
         self,
         rotated_imgs: list[Dict],
+        det_batch_size: int,
         tqdm_enable: bool = False,
         tqdm_desc: str = "Table orientation",
         progress_bar=None,
@@ -386,6 +532,7 @@ class MineruTableOrientationClsModel:
         label_by_index = {}
         img_score_tasks, all_crop_imgs = self._build_score_tasks_for_candidates(
             rotated_imgs,
+            det_batch_size,
             progress_bar=progress_bar,
         )
         self._extend_progress_total(progress_bar, len(all_crop_imgs))
@@ -401,6 +548,7 @@ class MineruTableOrientationClsModel:
 
         for img_info, tasks in img_score_tasks:
             score_by_label = self._score_orientation_tasks_with_rec(tasks, rec_res)
+            self._debug_log_orientation_rec_scores(img_info["index"], score_by_label)
             label_by_index[img_info["index"]] = self._select_rotation_label_by_scores(
                 score_by_label
             )
@@ -416,13 +564,9 @@ class MineruTableOrientationClsModel:
         """
         批量预测传入表格图片的旋转角度，只返回角度，不修改输入图片。
         """
-        RESOLUTION_GROUP_STRIDE = 128
         rotate_labels = ["0"] * len(imgs)
-        resolution_groups = self._collect_orientation_image_groups(
-            imgs,
-            RESOLUTION_GROUP_STRIDE,
-        )
-        total_images = sum(len(group_imgs) for group_imgs in resolution_groups.values())
+        orientation_imgs = self._collect_orientation_images(imgs)
+        total_images = len(orientation_imgs)
         progress_bar = None
         if tqdm_enable:
             progress_bar = tqdm(
@@ -432,9 +576,10 @@ class MineruTableOrientationClsModel:
             )
         try:
             rotated_imgs = self._detect_rotation_candidates(
-                resolution_groups,
+                orientation_imgs,
                 det_batch_size,
-                RESOLUTION_GROUP_STRIDE,
+                tqdm_enable=tqdm_enable,
+                tqdm_desc=tqdm_desc,
                 progress_bar=progress_bar,
             )
             self._extend_progress_total(progress_bar, len(rotated_imgs))
@@ -443,6 +588,7 @@ class MineruTableOrientationClsModel:
                 progress_bar.refresh()
             label_by_index = self._score_rotation_candidates(
                 rotated_imgs,
+                det_batch_size,
                 tqdm_enable=tqdm_enable,
                 tqdm_desc=tqdm_desc,
                 progress_bar=progress_bar,

mineru/model/table/rec/unet_table/main.py

@@ -28,6 +28,21 @@ from .utils_table_recover import (
     gather_ocr_list_by_row,
 )
 
+BLANK_CELL_REC_DROP_TEXTS = {
+    "1",
+    "一",
+    "—",
+    "口",
+    "■",
+    "（204号",
+    "（20",
+    "（2",
+    "（2号",
+    "（20号",
+    "号",
+    "（204",
+}
+
 
 @dataclass
 class WiredTableInput:
@@ -155,20 +170,18 @@ class WiredTableRecognition:
             )
         return res
 
-    # def fill_blank_rec(
-    #     self,
-    #     img: np.ndarray,
-    #     sorted_polygons: np.ndarray,
-    #     cell_box_map: Dict[int, List[str]],
-    # ) -> Dict[int, List[Any]]:
-    #     """找到poly对应为空的框，尝试将直接将poly框直接送到识别中"""
-    #     for i in range(sorted_polygons.shape[0]):
-    #         if cell_box_map.get(i):
-    #             continue
-    #         box = sorted_polygons[i]
-    #         cell_box_map[i] = [[box, "", 1]]
-    #         continue
-    #     return cell_box_map
+    @staticmethod
+    def _should_drop_blank_cell_rec_result(text: str, score) -> bool:
+        """判断空单元格二次 OCR-rec 结果是否应作为噪声过滤。"""
+        try:
+            if float(score) < 0.6:
+                return True
+        except (TypeError, ValueError):
+            return True
+
+        normalized_text = "" if text is None else str(text).strip()
+        return not normalized_text or normalized_text in BLANK_CELL_REC_DROP_TEXTS
+
     def fill_blank_rec(
         self,
         img: np.ndarray,
@@ -209,13 +222,6 @@ class WiredTableRecognition:
         if len(img_crop_list) > 0:
             # 进行ocr识别
             ocr_result = self.ocr_engine.ocr(img_crop_list, det=False)
-            # ocr_result = [[]]
-            # for crop_img in img_crop_list:
-            #     tmp_ocr_result = self.ocr_engine.ocr(crop_img)
-            #     if tmp_ocr_result[0] and len(tmp_ocr_result[0]) > 0 and isinstance(tmp_ocr_result[0], list) and len(tmp_ocr_result[0][0]) == 2:
-            #         ocr_result[0].append(tmp_ocr_result[0][0][1])
-            #     else:
-            #         ocr_result[0].append(("", 0.0))
 
             if not ocr_result or not isinstance(ocr_result, list) or len(ocr_result) == 0:
                 logger.warning("OCR engine returned no results or invalid result for image crops.")
@@ -231,7 +237,7 @@ class WiredTableRecognition:
                 # 处理ocr结果
                 ocr_text, ocr_score = ocr_res
                 # logger.debug(f"OCR result for box {i}: {ocr_text} with score {ocr_score}")
-                if ocr_score < 0.6 or ocr_text in ['1','口','■','（204号', '（20', '（2', '（2号', '（20号', '号', '（204']:
+                if self._should_drop_blank_cell_rec_result(ocr_text, ocr_score):
                     # logger.warning(f"Low confidence OCR result for box {i}: {ocr_text} with score {ocr_score}")
                     box = sorted_polygons[i]
                     cell_box_map[i] = [[box, "", 0.1]]

mineru/model/utils/pytorchocr/base_ocr_v20.py

@@ -107,6 +107,6 @@ class BaseOCRV20:
         # print('model is loaded: {}'.format(weights_path))
 
     def inference(self, inputs):
-        with torch.no_grad():
+        with torch.inference_mode():
             infer = self.net(inputs)
         return infer

mineru/model/utils/pytorchocr/modeling/heads/rec_multi_head.py

@@ -1,5 +1,4 @@
 # Copyright (c) Opendatalab. All rights reserved.
-import torch.nn.functional as F
 from torch import nn
 
 from ..necks.rnn import EncoderWithLightSVTR, Im2Seq, SequenceEncoder
@@ -71,7 +70,8 @@ class MultiHead(nn.Module):
             ctc_encoder = ctc_encoder.squeeze(dim=2).permute(0, 2, 1)
             predicts = self.head(ctc_encoder)
             if not self.training:
-                predicts = F.softmax(predicts, dim=2)
+                # 推理阶段保留 raw logits，交给 CTC 解码端按需计算 max 概率，避免整块 softmax 矩阵搬运。
+                return {"ctc_logits": predicts, "ctc_use_raw_logits": True}
             return predicts
         ctc_encoder = self.ctc_encoder(x)
         return self.ctc_head(ctc_encoder)

mineru/model/utils/pytorchocr/postprocess/rec_postprocess.py

@@ -181,11 +181,33 @@ class CTCLabelDecode(BaseRecLabelDecode):
         super(CTCLabelDecode, self).__init__(character_dict_path,
                                              use_space_char)
 
+    def _decode_raw_logits(self, preds):
+        """从 raw logits 直接计算 CTC argmax 和 max softmax 概率，避免完整 softmax。"""
+        logits = preds["ctc_logits"]
+        if torch.is_tensor(logits):
+            preds_idx = logits.argmax(dim=2)
+            max_logits = logits.amax(dim=2)
+            preds_prob = torch.exp(max_logits - torch.logsumexp(logits, dim=2))
+            return preds_prob.float().cpu().numpy(), preds_idx.cpu().numpy()
+
+        logits = np.asarray(logits)
+        preds_idx = logits.argmax(axis=2)
+        max_logits = logits.max(axis=2)
+        stable_logits = logits - max_logits[:, :, None]
+        logsumexp = max_logits + np.log(np.exp(stable_logits).sum(axis=2))
+        preds_prob = np.exp(max_logits - logsumexp).astype(np.float32)
+        return preds_prob, preds_idx
+
     def __call__(self, preds, label=None, return_word_box=False, *args, **kwargs):
-        preds_prob, preds_idx = preds.max(axis=2)
+        if isinstance(preds, dict) and preds.get("ctc_use_raw_logits"):
+            preds_prob, preds_idx = self._decode_raw_logits(preds)
+        else:
+            preds_prob, preds_idx = preds.max(axis=2)
+            preds_idx = preds_idx.cpu().numpy()
+            preds_prob = preds_prob.float().cpu().numpy()
         text = self.decode(
-            preds_idx.cpu().numpy(),
-            preds_prob.float().cpu().numpy(),
+            preds_idx,
+            preds_prob,
             is_remove_duplicate=True,
             return_word_box=return_word_box,
         )
@@ -787,4 +809,4 @@ class CANLabelDecode(BaseRecLabelDecode):
         if label is None:
             return text
         label = self.decode(label)
-        return text, label
+        return text, label

mineru/model/utils/tools/infer/predict_cls.py

@@ -86,11 +86,10 @@ class TextClassifier(BaseOCRV20):
                 norm_img = self.resize_norm_img(img_list[indices[ino]])
                 norm_img = norm_img[np.newaxis, :]
                 norm_img_batch.append(norm_img)
-            norm_img_batch = np.concatenate(norm_img_batch)
-            norm_img_batch = norm_img_batch.copy()
+            norm_img_batch = np.ascontiguousarray(np.concatenate(norm_img_batch))
             starttime = time.time()
 
-            with torch.no_grad():
+            with torch.inference_mode():
                 inp = torch.from_numpy(norm_img_batch)
                 inp = inp.to(self.device)
                 inp = self._to_inference_dtype(inp)

mineru/model/utils/tools/infer/predict_det.py

@@ -1,9 +1,11 @@
 # Copyright (c) Opendatalab. All rights reserved.
 import sys
+from collections import defaultdict
 
 import numpy as np
 import time
 import torch
+from tqdm import tqdm
 from ...pytorchocr.base_ocr_v20 import BaseOCRV20
 from . import pytorchocr_utility as utility
 from ...pytorchocr.data import create_operators, transform
@@ -124,6 +126,90 @@ class TextDetector(BaseOCRV20):
             if hasattr(module, 'rep'):
                 module.rep()
 
+    def _preprocess_det_image(self, img):
+        """执行 OCR-det 单图预处理，并保留后处理需要的原始尺寸信息。"""
+        data = {'image': img}
+        data = transform(data, self.preprocess_op)
+        if data is None:
+            return None
+
+        img_processed, shape_list = data
+        if img_processed is None:
+            return None
+
+        return np.ascontiguousarray(img_processed), shape_list, img.shape
+
+    def _build_det_preds(self, outputs):
+        """将模型输出统一转换为后处理需要的 float32 numpy 结构。"""
+        preds = {}
+        if self.det_algorithm == "EAST":
+            preds['f_geo'] = outputs['f_geo'].float().cpu().numpy()
+            preds['f_score'] = outputs['f_score'].float().cpu().numpy()
+        elif self.det_algorithm == 'SAST':
+            preds['f_border'] = outputs['f_border'].float().cpu().numpy()
+            preds['f_score'] = outputs['f_score'].float().cpu().numpy()
+            preds['f_tco'] = outputs['f_tco'].float().cpu().numpy()
+            preds['f_tvo'] = outputs['f_tvo'].float().cpu().numpy()
+        elif self.det_algorithm in ['DB', 'PSE', 'DB++']:
+            preds['maps'] = outputs['maps'].float().cpu().numpy()
+        elif self.det_algorithm == 'FCE':
+            for i, (_k, output) in enumerate(outputs.items()):
+                preds['level_{}'.format(i)] = output.float().cpu().numpy()
+        else:
+            raise NotImplementedError
+        return preds
+
+    def _postprocess_det_batch(self, preds, batch_shapes, ori_shapes):
+        """对完整 batch 执行一次 OCR-det 后处理，再逐张裁剪过滤检测框。"""
+        post_results = self.postprocess_op(preds, batch_shapes)
+        batch_results = []
+        for post_result, ori_shape in zip(post_results, ori_shapes):
+            dt_boxes = post_result['points']
+            dt_boxes = self._filter_det_res(dt_boxes, ori_shape)
+            batch_results.append(dt_boxes)
+        return batch_results
+
+    def _batch_process_preprocessed(self, batch_items):
+        """对已经完成预处理且形状一致的图片执行批量推理。"""
+        starttime = time.time()
+        if not batch_items:
+            return [], 0
+
+        batch_data = [item[1] for item in batch_items]
+        batch_shapes = [item[2] for item in batch_items]
+        ori_shapes = [item[3] for item in batch_items]
+
+        try:
+            batch_tensor = np.ascontiguousarray(np.stack(batch_data, axis=0))
+            batch_shapes = np.stack(batch_shapes, axis=0)
+        except Exception:
+            batch_results = []
+            for _index, img_processed, shape_list, ori_shape in batch_items:
+                single_tensor = np.expand_dims(np.ascontiguousarray(img_processed), axis=0)
+                single_shape = np.expand_dims(shape_list, axis=0)
+                with torch.inference_mode():
+                    inp = torch.from_numpy(single_tensor)
+                    inp = inp.to(self.device)
+                    inp = self._to_inference_dtype(inp)
+                    outputs = self.net(inp)
+                preds = self._build_det_preds(outputs)
+                dt_boxes = self._postprocess_det_batch(preds, single_shape, [ori_shape])[0]
+                batch_results.append((dt_boxes, 0))
+            return batch_results, time.time() - starttime
+
+        with torch.inference_mode():
+            inp = torch.from_numpy(batch_tensor)
+            inp = inp.to(self.device)
+            inp = self._to_inference_dtype(inp)
+            outputs = self.net(inp)
+
+        preds = self._build_det_preds(outputs)
+        dt_boxes_batch = self._postprocess_det_batch(preds, batch_shapes, ori_shapes)
+        total_elapse = time.time() - starttime
+        batch_elapse = total_elapse / len(batch_items)
+        batch_results = [(dt_boxes, batch_elapse) for dt_boxes in dt_boxes_batch]
+        return batch_results, total_elapse
+
     def _should_only_clip_det_res(self):
         if self.det_algorithm == "SAST" and getattr(self, "det_sast_polygon", False):
             return True
@@ -149,93 +235,34 @@ class TextDetector(BaseOCRV20):
             """
         starttime = time.time()
 
-        # 预处理所有图像
-        batch_data = []
-        batch_shapes = []
-        ori_imgs = []
-
-        for img in img_list:
-            ori_im = img.copy()
-            ori_imgs.append(ori_im)
-
-            data = {'image': img}
-            data = transform(data, self.preprocess_op)
-            if data is None:
-                # 如果预处理失败，返回空结果
+        batch_items = []
+        for index, img in enumerate(img_list):
+            preprocessed = self._preprocess_det_image(img)
+            if preprocessed is None:
                 return [(None, 0) for _ in img_list], 0
+            img_processed, shape_list, ori_shape = preprocessed
+            batch_items.append((index, img_processed, shape_list, ori_shape))
 
-            img_processed, shape_list = data
-            batch_data.append(img_processed)
-            batch_shapes.append(shape_list)
+        batch_results, _elapsed = self._batch_process_preprocessed(batch_items)
+        return batch_results, time.time() - starttime
 
-        # 堆叠成批处理张量
-        try:
-            batch_tensor = np.stack(batch_data, axis=0)
-            batch_shapes = np.stack(batch_shapes, axis=0)
-        except Exception as e:
-            # 如果堆叠失败，回退到逐个处理
-            batch_results = []
-            for img in img_list:
-                dt_boxes, elapse = self.__call__(img)
-                batch_results.append((dt_boxes, elapse))
-            return batch_results, time.time() - starttime
-
-        # 批处理推理
-        with torch.no_grad():
-            inp = torch.from_numpy(batch_tensor)
-            inp = inp.to(self.device)
-            inp = self._to_inference_dtype(inp)
-            outputs = self.net(inp)
-
-        # 处理输出
-        preds = {}
-        if self.det_algorithm == "EAST":
-            preds['f_geo'] = outputs['f_geo'].float().cpu().numpy()
-            preds['f_score'] = outputs['f_score'].float().cpu().numpy()
-        elif self.det_algorithm == 'SAST':
-            preds['f_border'] = outputs['f_border'].float().cpu().numpy()
-            preds['f_score'] = outputs['f_score'].float().cpu().numpy()
-            preds['f_tco'] = outputs['f_tco'].float().cpu().numpy()
-            preds['f_tvo'] = outputs['f_tvo'].float().cpu().numpy()
-        elif self.det_algorithm in ['DB', 'PSE', 'DB++']:
-            preds['maps'] = outputs['maps'].float().cpu().numpy()
-        elif self.det_algorithm == 'FCE':
-            for i, (k, output) in enumerate(outputs.items()):
-                preds['level_{}'.format(i)] = output.float().cpu().numpy()
-        else:
-            raise NotImplementedError
-
-        # 后处理每个图像的结果
-        batch_results = []
-        total_elapse = time.time() - starttime
-
-        for i in range(len(img_list)):
-            # 提取单个图像的预测结果
-            single_preds = {}
-            for key, value in preds.items():
-                if isinstance(value, np.ndarray):
-                    single_preds[key] = value[i:i + 1]  # 保持批次维度
-                else:
-                    single_preds[key] = value
-
-            # 后处理
-            post_result = self.postprocess_op(single_preds, batch_shapes[i:i + 1])
-            dt_boxes = post_result[0]['points']
-
-            # 过滤和裁剪检测框
-            dt_boxes = self._filter_det_res(dt_boxes, ori_imgs[i].shape)
-
-            batch_results.append((dt_boxes, total_elapse / len(img_list)))
-
-        return batch_results, total_elapse
-
-    def batch_predict(self, img_list, max_batch_size=8):
+    def batch_predict(
+        self,
+        img_list,
+        max_batch_size=8,
+        tqdm_enable=False,
+        tqdm_desc="OCR-det Predict",
+        tqdm_progress_bar=None,
+    ):
         """
         批处理预测方法，支持多张图像同时检测
 
         Args:
             img_list: 图像列表
             max_batch_size: 最大批处理大小
+            tqdm_enable: 是否显示内部 OCR-det 进度条
+            tqdm_desc: 内部 OCR-det 进度条描述
+            tqdm_progress_bar: 外部复用进度条，传入时不在本方法内关闭
 
         Returns:
             batch_results: 批处理结果列表，每个元素为(dt_boxes, elapse)
@@ -243,14 +270,38 @@ class TextDetector(BaseOCRV20):
         if not img_list:
             return []
 
-        batch_results = []
+        progress_bar = tqdm_progress_bar
+        should_close_progress = False
+        if progress_bar is None:
+            progress_bar = tqdm(total=len(img_list), desc=tqdm_desc, disable=not tqdm_enable)
+            should_close_progress = True
 
-        # 分批处理
-        for i in range(0, len(img_list), max_batch_size):
-            batch_imgs = img_list[i:i + max_batch_size]
-            # assert尺寸一致
-            batch_dt_boxes, batch_elapse = self._batch_process_same_size(batch_imgs)
-            batch_results.extend(batch_dt_boxes)
+        max_batch_size = max(1, int(max_batch_size))
+        batch_results = [(None, 0)] * len(img_list)
+        grouped_items = defaultdict(list)
+
+        try:
+            for index, img in enumerate(img_list):
+                preprocessed = self._preprocess_det_image(img)
+                if preprocessed is None:
+                    progress_bar.update(1)
+                    continue
+                img_processed, shape_list, ori_shape = preprocessed
+                grouped_items[img_processed.shape].append(
+                    (index, img_processed, shape_list, ori_shape)
+                )
+
+            for group_items in grouped_items.values():
+                for i in range(0, len(group_items), max_batch_size):
+                    batch_items = group_items[i:i + max_batch_size]
+                    group_results, _batch_elapse = self._batch_process_preprocessed(batch_items)
+                    for batch_item, batch_result in zip(batch_items, group_results):
+                        original_index = batch_item[0]
+                        batch_results[original_index] = batch_result
+                    progress_bar.update(len(batch_items))
+        finally:
+            if should_close_progress:
+                progress_bar.close()
 
         return batch_results
 
@@ -310,43 +361,11 @@ class TextDetector(BaseOCRV20):
         return dt_boxes_new
 
     def __call__(self, img):
-        ori_shape = img.shape
-        data = {'image': img}
-        data = transform(data, self.preprocess_op)
-        img, shape_list = data
-        if img is None:
+        preprocessed = self._preprocess_det_image(img)
+        if preprocessed is None:
             return None, 0
-        img = np.expand_dims(img, axis=0)
-        shape_list = np.expand_dims(shape_list, axis=0)
-        img = img.copy()
-        starttime = time.time()
-
-        with torch.no_grad():
-            inp = torch.from_numpy(img)
-            inp = inp.to(self.device)
-            inp = self._to_inference_dtype(inp)
-            outputs = self.net(inp)
-
-        preds = {}
-        if self.det_algorithm == "EAST":
-            preds['f_geo'] = outputs['f_geo'].float().cpu().numpy()
-            preds['f_score'] = outputs['f_score'].float().cpu().numpy()
-        elif self.det_algorithm == 'SAST':
-            preds['f_border'] = outputs['f_border'].float().cpu().numpy()
-            preds['f_score'] = outputs['f_score'].float().cpu().numpy()
-            preds['f_tco'] = outputs['f_tco'].float().cpu().numpy()
-            preds['f_tvo'] = outputs['f_tvo'].float().cpu().numpy()
-        elif self.det_algorithm in ['DB', 'PSE', 'DB++']:
-            preds['maps'] = outputs['maps'].float().cpu().numpy()
-        elif self.det_algorithm == 'FCE':
-            for i, (k, output) in enumerate(outputs.items()):
-                preds['level_{}'.format(i)] = output.float().cpu().numpy()
-        else:
-            raise NotImplementedError
-
-        post_result = self.postprocess_op(preds, shape_list)
-        dt_boxes = post_result[0]['points']
-        dt_boxes = self._filter_det_res(dt_boxes, ori_shape)
-
-        elapse = time.time() - starttime
-        return dt_boxes, elapse
+        img_processed, shape_list, ori_shape = preprocessed
+        batch_results, _elapsed = self._batch_process_preprocessed(
+            [(0, img_processed, shape_list, ori_shape)]
+        )
+        return batch_results[0]

mineru/model/utils/tools/infer/predict_rec.py

@@ -364,8 +364,7 @@ class TextRecognizer(BaseOCRV20):
                                                         max_wh_ratio)
                         norm_img = norm_img[np.newaxis, :]
                         norm_img_batch.append(norm_img)
-                norm_img_batch = np.concatenate(norm_img_batch)
-                norm_img_batch = norm_img_batch.copy()
+                norm_img_batch = np.ascontiguousarray(np.concatenate(norm_img_batch))
 
                 if self.rec_algorithm == "SRN":
                     starttime = time.time()
@@ -376,7 +375,7 @@ class TextRecognizer(BaseOCRV20):
                     gsrm_slf_attn_bias2_list = np.concatenate(
                         gsrm_slf_attn_bias2_list)
 
-                    with torch.no_grad():
+                    with torch.inference_mode():
                         inp = torch.from_numpy(norm_img_batch)
                         encoder_word_pos_inp = torch.from_numpy(encoder_word_pos_list)
                         gsrm_word_pos_inp = torch.from_numpy(gsrm_word_pos_list)
@@ -407,7 +406,7 @@ class TextRecognizer(BaseOCRV20):
                     #     valid_ratios,
                     # ]
 
-                    with torch.no_grad():
+                    with torch.inference_mode():
                         inp = torch.from_numpy(norm_img_batch)
                         inp = inp.to(self.device)
                         inp = self._to_inference_dtype(inp)
@@ -422,7 +421,7 @@ class TextRecognizer(BaseOCRV20):
                     inp = [torch.from_numpy(e_i) for e_i in inputs]
                     inp = [e_i.to(self.device) for e_i in inp]
                     inp = [self._to_inference_dtype(e_i) for e_i in inp]
-                    with torch.no_grad():
+                    with torch.inference_mode():
                         outputs = self.net(inp)
                         outputs = [v.cpu().numpy() for k, v in enumerate(outputs)]
 
@@ -431,13 +430,13 @@ class TextRecognizer(BaseOCRV20):
                 else:
                     starttime = time.time()
 
-                    with torch.no_grad():
+                    with torch.inference_mode():
                         inp = torch.from_numpy(norm_img_batch)
                         inp = inp.to(self.device)
                         inp = self._to_inference_dtype(inp)
                         preds = self.net(inp)
 
-                with torch.no_grad():
+                with torch.inference_mode():
                     rec_result = self.postprocess_op(preds)
 
                 for rno in range(len(rec_result)):