Bläddra i källkod

fix: enhance batch analysis with table orientation classification and prediction improvements

myhloli 3 månader sedan
förälder
incheckning
f1fb900ea5

+ 117 - 114
mineru/backend/pipeline/batch_analyze.py

@@ -18,7 +18,8 @@ YOLO_LAYOUT_BASE_BATCH_SIZE = 1
 MFD_BASE_BATCH_SIZE = 1
 MFR_BASE_BATCH_SIZE = 16
 OCR_DET_BASE_BATCH_SIZE = 16
-ORI_TAB_CLS_BATCH_SIZE = 16
+TABLE_ORI_CLS_BATCH_SIZE = 16
+TABLE_Wired_Wireless_CLS_BATCH_SIZE = 16
 
 
 class BatchAnalyze:
@@ -105,7 +106,120 @@ class BatchAnalyze:
                                                 'table_img':table_img,
                                               })
 
-        # OCR检测处理
+        # 表格识别 table recognition
+        if self.table_enable:
+
+            # 图片旋转批量处理
+            img_orientation_cls_model = atom_model_manager.get_atom_model(
+                atom_model_name=AtomicModel.ImgOrientationCls,
+            )
+            try:
+                img_orientation_cls_model.batch_predict(table_res_list_all_page,
+                                                        det_batch_size=self.batch_ratio * OCR_DET_BASE_BATCH_SIZE,
+                                                        batch_size=TABLE_ORI_CLS_BATCH_SIZE)
+            except Exception as e:
+                logger.warning(
+                    f"Image orientation classification failed: {e}, using original image"
+                )
+
+            # 表格分类
+            table_cls_model = atom_model_manager.get_atom_model(
+                atom_model_name=AtomicModel.TableCls,
+            )
+            try:
+                table_cls_model.batch_predict(table_res_list_all_page,
+                                              batch_size=TABLE_Wired_Wireless_CLS_BATCH_SIZE)
+            except Exception as e:
+                logger.warning(
+                    f"Table classification failed: {e}, using default model"
+                )
+
+            # OCR det 过程,顺序执行
+            rec_img_lang_group = defaultdict(list)
+            for index, table_res_dict in enumerate(
+                    tqdm(table_res_list_all_page, desc="Table-ocr det")
+            ):
+                ocr_engine = atom_model_manager.get_atom_model(
+                    atom_model_name=AtomicModel.OCR,
+                    det_db_box_thresh=0.5,
+                    det_db_unclip_ratio=1.6,
+                    # lang= table_res_dict["lang"],
+                    enable_merge_det_boxes=False,
+                )
+                bgr_image = cv2.cvtColor(table_res_dict["table_img"], cv2.COLOR_RGB2BGR)
+                ocr_result = ocr_engine.ocr(bgr_image, rec=False)[0]
+                # 构造需要 OCR 识别的图片字典,包括cropped_img, dt_box, table_id,并按照语言进行分组
+                for dt_box in ocr_result:
+                    rec_img_lang_group[_lang].append(
+                        {
+                            "cropped_img": get_rotate_crop_image(
+                                bgr_image, np.asarray(dt_box, dtype=np.float32)
+                            ),
+                            "dt_box": np.asarray(dt_box, dtype=np.float32),
+                            "table_id": index,
+                        }
+                    )
+
+            # OCR rec,按照语言分批处理
+            for _lang, rec_img_list in rec_img_lang_group.items():
+                ocr_engine = atom_model_manager.get_atom_model(
+                    atom_model_name=AtomicModel.OCR,
+                    det_db_box_thresh=0.5,
+                    det_db_unclip_ratio=1.6,
+                    lang=_lang,
+                    enable_merge_det_boxes=False,
+                )
+                cropped_img_list = [item["cropped_img"] for item in rec_img_list]
+                ocr_res_list = \
+                ocr_engine.ocr(cropped_img_list, det=False, tqdm_enable=True, tqdm_desc="Table-ocr rec")[0]
+                # 按照 table_id 将识别结果进行回填
+                for img_dict, ocr_res in zip(rec_img_list, ocr_res_list):
+                    if table_res_list_all_page[img_dict["table_id"]].get("ocr_result"):
+                        table_res_list_all_page[img_dict["table_id"]]["ocr_result"].append(
+                            [img_dict["dt_box"], html.escape(ocr_res[0]), ocr_res[1]]
+                        )
+                    else:
+                        table_res_list_all_page[img_dict["table_id"]]["ocr_result"] = [
+                            [img_dict["dt_box"], html.escape(ocr_res[0]), ocr_res[1]]
+                        ]
+
+            # 先对所有表格使用无线表格模型,然后对分类为有线的表格使用有线表格模型
+            wireless_table_model = atom_model_manager.get_atom_model(
+                atom_model_name=AtomicModel.WirelessTable,
+            )
+            wireless_table_model.batch_predict(table_res_list_all_page)
+
+            # 单独拿出有线表格进行预测
+            wired_table_res_list = []
+            for table_res_dict in table_res_list_all_page:
+                if table_res_dict["table_res"]["cls_label"] == AtomicModel.WiredTable:
+                    wired_table_res_list.append(table_res_dict)
+            if wired_table_res_list:
+                for table_res_dict in tqdm(
+                        wired_table_res_list, desc="Table-wired Predict"
+                ):
+                    wired_table_model = atom_model_manager.get_atom_model(
+                        atom_model_name=AtomicModel.WiredTable,
+                        lang=table_res_dict["lang"],
+                    )
+                    table_res_dict["table_res"]["html"] = wired_table_model.predict(
+                        table_res_dict["table_img"],
+                        table_res_dict["ocr_result"],
+                        table_res_dict["table_res"].get("html", None)
+                    )
+
+            # 表格格式清理
+            for table_res_dict in table_res_list_all_page:
+                html_code = table_res_dict["table_res"].get("html", "")
+
+                # 检查html_code是否包含'<table>'和'</table>'
+                if "<table>" in html_code and "</table>" in html_code:
+                    # 选用<table>到</table>的内容,放入table_res_dict['table_res']['html']
+                    start_index = html_code.find("<table>")
+                    end_index = html_code.rfind("</table>") + len("</table>")
+                    table_res_dict["table_res"]["html"] = html_code[start_index:end_index]
+
+        # OCR det
         if self.enable_ocr_det_batch:
             # 批处理模式 - 按语言和分辨率分组
             # 收集所有需要OCR检测的裁剪图像
@@ -256,118 +370,7 @@ class BatchAnalyze:
 
                         ocr_res_list_dict['layout_res'].extend(ocr_result_list)
 
-        # 表格识别 table recognition
-        if self.table_enable:
-            # 图片旋转批量处理
-
-            img_orientation_cls_model = atom_model_manager.get_atom_model(
-                atom_model_name=AtomicModel.ImgOrientationCls,
-            )
-            try:
-                img_orientation_cls_model.batch_predict(table_res_list_all_page, batch_size=self.batch_ratio * OCR_DET_BASE_BATCH_SIZE)
-            except Exception as e:
-                logger.warning(
-                    f"Image orientation classification failed: {e}, using original image"
-                )
-            # 表格分类
-            table_cls_model = atom_model_manager.get_atom_model(
-                atom_model_name=AtomicModel.TableCls,
-            )
-            try:
-                table_cls_model.batch_predict(table_res_list_all_page)
-            except Exception as e:
-                logger.warning(
-                    f"Table classification failed: {e}, using default model"
-                )
-            rec_img_lang_group = defaultdict(list)
-            # OCR det 过程,顺序执行
-            for index, table_res_dict in enumerate(
-                tqdm(table_res_list_all_page, desc="Table OCR det")
-            ):
-                _lang = table_res_dict["lang"]
-                ocr_engine = atom_model_manager.get_atom_model(
-                    atom_model_name=AtomicModel.OCR,
-                    det_db_box_thresh=0.5,
-                    det_db_unclip_ratio=1.6,
-                    lang=_lang,
-                    enable_merge_det_boxes=False,
-                )
-                bgr_image = cv2.cvtColor(
-                    np.asarray(table_res_dict["table_img"]), cv2.COLOR_RGB2BGR
-                )
-                ocr_result = ocr_engine.ocr(bgr_image, det=True, rec=False)[0]
-                # 构造需要 OCR 识别的图片字典,包括cropped_img, dt_box, table_id,并按照语言进行分组
-                for dt_box in ocr_result:
-                    rec_img_lang_group[_lang].append(
-                        {
-                            "cropped_img": get_rotate_crop_image(
-                                bgr_image, np.asarray(dt_box, dtype=np.float32)
-                            ),
-                            "dt_box": np.asarray(dt_box, dtype=np.float32),
-                            "table_id": index,
-                        }
-                    )
-            # OCR rec,按照语言分批处理
-            for _lang, rec_img_list in rec_img_lang_group.items():
-                ocr_engine = atom_model_manager.get_atom_model(
-                    atom_model_name=AtomicModel.OCR,
-                    det_db_box_thresh=0.5,
-                    det_db_unclip_ratio=1.6,
-                    lang=_lang,
-                    enable_merge_det_boxes=False,
-                )
-                cropped_img_list = [item["cropped_img"] for item in rec_img_list]
-                ocr_res_list = ocr_engine.ocr(
-                    cropped_img_list, det=False, rec=True, tqdm_enable=True
-                )[0]
-                # 按照 table_id 将识别结果进行回填
-                for img_dict, ocr_res in zip(rec_img_list, ocr_res_list):
-                    if table_res_list_all_page[img_dict["table_id"]].get("ocr_result"):
-                        table_res_list_all_page[img_dict["table_id"]]["ocr_result"].append(
-                            [img_dict["dt_box"], html.escape(ocr_res[0]), ocr_res[1]]
-                        )
-                    else:
-                        table_res_list_all_page[img_dict["table_id"]]["ocr_result"] = [
-                            [img_dict["dt_box"], html.escape(ocr_res[0]), ocr_res[1]]
-                        ]
-
-            # 先对所有表格使用无线表格模型,然后对分类为有线的表格使用有线表格模型
-            wireless_table_model = atom_model_manager.get_atom_model(
-                atom_model_name=AtomicModel.WirelessTable,
-            )
-            wireless_table_model.batch_predict(table_res_list_all_page)
-
-            # 单独拿出有线表格进行预测
-            wired_table_res_list = []
-            for table_res_dict in table_res_list_all_page:
-                if table_res_dict["table_res"]["cls_label"] == AtomicModel.WiredTable:
-                    wired_table_res_list.append(table_res_dict)
-            for table_res_dict in tqdm(
-                wired_table_res_list, desc="Wired Table Predict"
-            ):
-                if table_res_dict["table_res"]["cls_label"] == AtomicModel.WiredTable:
-                    wired_table_model = atom_model_manager.get_atom_model(
-                        atom_model_name=AtomicModel.WiredTable,
-                        lang=table_res_dict["lang"],
-                    )
-                    html_code = wired_table_model.predict(
-                        table_res_dict["table_img"],
-                        table_res_dict["ocr_result"],
-                        table_res_dict["table_res"].get("html", None)
-                    )
-                    # 检查html_code是否包含'<table>'和'</table>'
-                    if "<table>" in html_code and "</table>" in html_code:
-                        # 选用<table>到</table>的内容,放入table_res_dict['table_res']['html']
-                        start_index = html_code.find("<table>")
-                        end_index = html_code.rfind("</table>") + len("</table>")
-                        table_res_dict["table_res"]["html"] = html_code[
-                            start_index:end_index
-                        ]
-                    else:
-                        logger.warning(
-                            "wired table recognition processing fails, not found expected HTML table end"
-                        )
-
+        # OCR rec
         # Create dictionaries to store items by language
         need_ocr_lists_by_lang = {}  # Dict of lists for each language
         img_crop_lists_by_lang = {}  # Dict of lists for each language

+ 11 - 6
mineru/backend/pipeline/model_init.py

@@ -114,13 +114,18 @@ class AtomModelSingleton:
         lang = kwargs.get('lang', None)
 
         if atom_model_name in [AtomicModel.WiredTable, AtomicModel.WirelessTable]:
-            key = (atom_model_name, lang)
+            key = (
+                atom_model_name,
+                lang
+            )
         elif atom_model_name in [AtomicModel.OCR]:
-            key = (atom_model_name,
-                   kwargs.get('det_db_box_thresh', 0.3),
-                   lang, kwargs.get('det_db_unclip_ratio', 1.8),
-                   kwargs.get('enable_merge_det_boxes', True)
-                   )
+            key = (
+                atom_model_name,
+                kwargs.get('det_db_box_thresh', 0.3),
+                lang,
+                kwargs.get('det_db_unclip_ratio', 1.8),
+                kwargs.get('enable_merge_det_boxes', True)
+            )
         else:
             key = atom_model_name
 

+ 2 - 1
mineru/model/ocr/paddleocr2pytorch/pytorch_paddle.py

@@ -105,6 +105,7 @@ class PytorchPaddleOCR(TextSystem):
             rec=True,
             mfd_res=None,
             tqdm_enable=False,
+            tqdm_desc="OCR-rec Predict",
             ):
         assert isinstance(img, (np.ndarray, list, str, bytes))
         if isinstance(img, list) and det == True:
@@ -149,7 +150,7 @@ 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)
+                    rec_res, elapse = self.text_recognizer(img, tqdm_enable=tqdm_enable, tqdm_desc=tqdm_desc)
                     # logger.debug("rec_res num  : {}, elapsed : {}".format(len(rec_res), elapse))
                     ocr_res.append(rec_res)
                 return ocr_res

+ 2 - 2
mineru/model/ocr/paddleocr2pytorch/tools/infer/predict_rec.py

@@ -288,7 +288,7 @@ class TextRecognizer(BaseOCRV20):
 
         return img
 
-    def __call__(self, img_list, tqdm_enable=False):
+    def __call__(self, img_list, tqdm_enable=False, tqdm_desc="OCR-rec Predict"):
         img_num = len(img_list)
         # Calculate the aspect ratio of all text bars
         width_list = []
@@ -302,7 +302,7 @@ class TextRecognizer(BaseOCRV20):
         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='OCR-rec Predict', disable=not tqdm_enable) as pbar:
+        with tqdm(total=img_num, desc=tqdm_desc, disable=not tqdm_enable) as pbar:
             index = 0
             for beg_img_no in range(0, img_num, batch_num):
                 end_img_no = min(img_num, beg_img_no + batch_num)

+ 42 - 39
mineru/model/ori_cls/paddle_ori_cls.py

@@ -134,9 +134,7 @@ class PaddleOrientationClsModel:
     def batch_preprocess(self, imgs):
         res_imgs = []
         for img_info in imgs:
-            # PIL图像转cv2
-            img = cv2.cvtColor(np.asarray(img_info["table_img"]), cv2.COLOR_RGB2BGR)
-            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            img = np.asarray(img_info["table_img"])
             # 放大图片,使其最短边长为256
             h, w = img.shape[:2]
             scale = 256 / min(h, w)
@@ -174,33 +172,30 @@ class PaddleOrientationClsModel:
         return x
 
     def batch_predict(
-        self, imgs: List[Dict], batch_size: int
+        self, imgs: List[Dict], det_batch_size: int, batch_size: int = 16
     ) -> None:
         """
         批量预测传入的包含图片信息列表的旋转信息,并且将旋转过的图片正确地旋转回来
         """
-        RESOLUTION_GROUP_STRIDE = 64
+        RESOLUTION_GROUP_STRIDE = 128
         # 跳过长宽比小于1.2的图片
         resolution_groups = defaultdict(list)
         for img in imgs:
             # RGB图像转换BGR
-            table_img: np.ndarray = cv2.cvtColor(img["table_img"], cv2.COLOR_RGB2BGR)
-            img["table_img_bgr"] = table_img
-            img_height, img_width = table_img.shape[:2]
+            bgr_img: np.ndarray = cv2.cvtColor(np.asarray(img["table_img"]), cv2.COLOR_RGB2BGR)
+            img["table_img_bgr"] = bgr_img
+            img_height, img_width = bgr_img.shape[:2]
             img_aspect_ratio = img_height / img_width if img_width > 0 else 1.0
-            img_is_portrait = img_aspect_ratio > 1.2
-            if img_is_portrait:
-                h, w = img["table_img_bgr"].shape[:2]
-                normalized_h = ((h + RESOLUTION_GROUP_STRIDE) // RESOLUTION_GROUP_STRIDE) * RESOLUTION_GROUP_STRIDE  # 向上取整到RESOLUTION_GROUP_STRIDE的倍数
-                normalized_w = ((w + RESOLUTION_GROUP_STRIDE) // RESOLUTION_GROUP_STRIDE) * RESOLUTION_GROUP_STRIDE
+            if img_aspect_ratio > 1.2:
+                # 归一化尺寸到RESOLUTION_GROUP_STRIDE的倍数
+                normalized_h = ((img_height + RESOLUTION_GROUP_STRIDE) // RESOLUTION_GROUP_STRIDE) * RESOLUTION_GROUP_STRIDE  # 向上取整到RESOLUTION_GROUP_STRIDE的倍数
+                normalized_w = ((img_width + RESOLUTION_GROUP_STRIDE) // RESOLUTION_GROUP_STRIDE) * RESOLUTION_GROUP_STRIDE
                 group_key = (normalized_h, normalized_w)
                 resolution_groups[group_key].append(img)
 
-            # 对每个分辨率组进行批处理
-        for group_key, group_imgs in tqdm(
-            resolution_groups.items(), desc=f"ORI CLS OCR-det"
-        ):
-
+        # 对每个分辨率组进行批处理
+        rotated_imgs = []
+        for group_key, group_imgs in tqdm(resolution_groups.items(), desc="Table-ori cls stage1 predict"):
             # 计算目标尺寸(组内最大尺寸,向上取整到RESOLUTION_GROUP_STRIDE的倍数)
             max_h = max(img["table_img_bgr"].shape[0] for img in group_imgs)
             max_w = max(img["table_img_bgr"].shape[1] for img in group_imgs)
@@ -210,22 +205,21 @@ class PaddleOrientationClsModel:
             # 对所有图像进行padding到统一尺寸
             batch_images = []
             for img in group_imgs:
-                table_img_ndarray = img["table_img_bgr"]
-                h, w = table_img_ndarray.shape[:2]
+                bgr_img = img["table_img_bgr"]
+                h, w = bgr_img.shape[:2]
                 # 创建目标尺寸的白色背景
                 padded_img = np.ones((target_h, target_w, 3), dtype=np.uint8) * 255
                 # 将原图像粘贴到左上角
-                padded_img[:h, :w] = table_img_ndarray
+                padded_img[:h, :w] = bgr_img
                 batch_images.append(padded_img)
 
             # 批处理检测
-            det_batch_size = min(len(batch_images), batch_size)  # 增加批处理大小
             batch_results = self.ocr_engine.text_detector.batch_predict(
-                batch_images, det_batch_size
+                batch_images, min(len(batch_images), det_batch_size)
             )
 
-            rotated_imgs = []
             # 根据批处理结果检测图像是否旋转,旋转的图像放入列表中,继续进行旋转角度的预测
+
             for index, (img_info, (dt_boxes, elapse)) in enumerate(
                 zip(group_imgs, batch_results)
             ):
@@ -245,18 +239,27 @@ class PaddleOrientationClsModel:
 
                 if vertical_count >= len(dt_boxes) * 0.28 and vertical_count >= 3:
                     rotated_imgs.append(img_info)
-            if len(rotated_imgs) > 0:
-                x = self.batch_preprocess(rotated_imgs)
-                results = self.sess.run(None, {"x": x})
-                for img_info, res in zip(rotated_imgs, results[0]):
-                    label = self.labels[np.argmax(res)]
-                    if label == "270":
-                        img_info["table_img"] = cv2.rotate(
-                            np.asarray(img_info["table_img"]),
-                            cv2.ROTATE_90_CLOCKWISE,
-                        )
-                    elif label == "90":
-                        img_info["table_img"] = cv2.rotate(
-                            np.asarray(img_info["table_img"]),
-                            cv2.ROTATE_90_COUNTERCLOCKWISE,
-                        )
+
+        # 对旋转的图片进行旋转角度预测
+        if len(rotated_imgs) > 0:
+            imgs = self.list_2_batch(rotated_imgs, batch_size=batch_size)
+            with tqdm(total=len(rotated_imgs), desc="Table-ori cls stage2 predict") as pbar:
+                for img_batch in imgs:
+                    x = self.batch_preprocess(img_batch)
+                    results = self.sess.run(None, {"x": x})
+                    for img_info, res in zip(rotated_imgs, results[0]):
+                        label = self.labels[np.argmax(res)]
+                        if label == "270":
+                            img_info["table_img"] = cv2.rotate(
+                                np.asarray(img_info["table_img"]),
+                                cv2.ROTATE_90_CLOCKWISE,
+                            )
+                        elif label == "90":
+                            img_info["table_img"] = cv2.rotate(
+                                np.asarray(img_info["table_img"]),
+                                cv2.ROTATE_90_COUNTERCLOCKWISE,
+                            )
+                        else:
+                            # 180度和0度不做处理
+                            pass
+                        pbar.update(1)

+ 17 - 16
mineru/model/table/cls/paddle_table_cls.py

@@ -6,6 +6,7 @@ import cv2
 import numpy as np
 import onnxruntime
 from loguru import logger
+from tqdm import tqdm
 
 from mineru.backend.pipeline.model_list import AtomicModel
 from mineru.utils.enum_class import ModelPath
@@ -97,9 +98,7 @@ class PaddleTableClsModel:
     def batch_preprocess(self, imgs):
         res_imgs = []
         for img in imgs:
-            # PIL图像转cv2
-            img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
-            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            img = np.asarray(img)
             # 放大图片,使其最短边长为256
             h, w = img.shape[:2]
             scale = 256 / min(h, w)
@@ -139,16 +138,18 @@ class PaddleTableClsModel:
         imgs = [item["table_img"] for item in img_info_list]
         imgs = self.list_2_batch(imgs, batch_size=batch_size)
         label_res = []
-        for img_batch in imgs:
-            x = self.batch_preprocess(img_batch)
-            result = self.sess.run(None, {"x": x})
-            for img_res in result[0]:
-                idx = np.argmax(img_res)
-                conf = float(np.max(img_res))
-                # logger.debug(f"Table classification result: {self.labels[idx]} with confidence {conf:.4f}")
-                if idx == 0 and conf < 0.8:
-                    idx = 1
-                label_res.append((self.labels[idx],conf))
-        for img_info, (label, conf) in zip(img_info_list, label_res):
-            img_info['table_res']["cls_label"] = label
-            img_info['table_res']["cls_score"] = conf
+        with tqdm(total=len(img_info_list), desc="Table-wired/wireless cls predict") as pbar:
+            for img_batch in imgs:
+                x = self.batch_preprocess(img_batch)
+                result = self.sess.run(None, {"x": x})
+                for img_res in result[0]:
+                    idx = np.argmax(img_res)
+                    conf = float(np.max(img_res))
+                    # logger.debug(f"Table classification result: {self.labels[idx]} with confidence {conf:.4f}")
+                    if idx == 0 and conf < 0.8:
+                        idx = 1
+                    label_res.append((self.labels[idx],conf))
+                pbar.update(len(img_batch))
+            for img_info, (label, conf) in zip(img_info_list, label_res):
+                img_info['table_res']["cls_label"] = label
+                img_info['table_res']["cls_score"] = round(conf, 3)

+ 19 - 31
mineru/model/table/rec/slanet_plus/main.py

@@ -245,34 +245,22 @@ class RapidTableModel(object):
 
     def batch_predict(self, table_res_list: List[Dict], batch_size: int = 4) -> None:
         """对传入的字典列表进行批量预测,无返回值"""
-        for index in tqdm(
-            range(0, len(table_res_list), batch_size),
-            desc=f"Wireless Table Batch Predict, total={len(table_res_list)}, batch_size={batch_size}",
-        ):
-            batch_imgs = [
-                cv2.cvtColor(np.asarray(table_res_list[i]["table_img"]), cv2.COLOR_RGB2BGR)
-                for i in range(index, min(index + batch_size, len(table_res_list)))
-            ]
-            batch_ocrs = [
-                table_res_list[i]["ocr_result"]
-                for i in range(index, min(index + batch_size, len(table_res_list)))
-            ]
-            results = self.table_model.batch_predict(
-                batch_imgs, batch_ocrs, batch_size=batch_size
-            )
-            for i, result in enumerate(results):
-                if result.pred_html:
-                    # 检查html_code是否包含'<table>'和'</table>'
-                    if '<table>' in result.pred_html and '</table>' in result.pred_html:
-                        # 选用<table>到</table>的内容,放入table_res_dict['table_res']['html']
-                        start_index = result.pred_html.find('<table>')
-                        end_index = result.pred_html.rfind('</table>') + len('</table>')
-                        table_res_list[index + i]['table_res']['html'] = result.pred_html[start_index:end_index]
-                    else:
-                        logger.warning(
-                            'wireless table recognition processing fails, not found expected HTML table end'
-                        )
-                else:
-                    logger.warning(
-                        "wireless table recognition processing fails, not get html return"
-                    )
+        with tqdm(total=len(table_res_list), desc="Table-wireless Predict") as pbar:
+            for index in range(0, len(table_res_list), batch_size):
+                batch_imgs = [
+                    cv2.cvtColor(np.asarray(table_res_list[i]["table_img"]), cv2.COLOR_RGB2BGR)
+                    for i in range(index, min(index + batch_size, len(table_res_list)))
+                ]
+                batch_ocrs = [
+                    table_res_list[i]["ocr_result"]
+                    for i in range(index, min(index + batch_size, len(table_res_list)))
+                ]
+                results = self.table_model.batch_predict(
+                    batch_imgs, batch_ocrs, batch_size=batch_size
+                )
+                for i, result in enumerate(results):
+                    if result.pred_html:
+                        table_res_list[index + i]['table_res']['html'] = result.pred_html
+
+                # 更新进度条
+                pbar.update(len(results))

+ 12 - 3
mineru/model/table/rec/unet_table/main.py

@@ -16,7 +16,7 @@ from .table_structure_unet import TSRUnet
 from mineru.utils.enum_class import ModelPath
 from mineru.utils.models_download_utils import auto_download_and_get_model_root_path
 from .table_recover import TableRecover
-from .utils import InputType, LoadImage
+from .utils import InputType, LoadImage, VisTable
 from .utils_table_recover import (
     match_ocr_cell,
     plot_html_table,
@@ -187,7 +187,7 @@ class WiredTableRecognition:
                 continue
             # 判断长宽比
             if (x2 - x1) / (y2 - y1) > 20 or (y2 - y1) / (x2 - x1) > 20:
-                logger.warning(f"Box {i} has invalid aspect ratio: {x1, y1, x2, y2}")
+                # logger.warning(f"Box {i} has invalid aspect ratio: {x1, y1, x2, y2}")
                 continue
             img_crop = bgr_img[int(y1):int(y2), int(x1):int(x2)]
             img_crop_list.append(img_crop)
@@ -264,6 +264,16 @@ class UnetTableModel:
         try:
             wired_table_results = self.wired_table_model(np_img, ocr_result)
 
+            # viser = VisTable()
+            # save_html_path = f"outputs/output.html"
+            # save_drawed_path = f"outputs/output_table_vis.jpg"
+            # save_logic_path = (
+            #     f"outputs/output_table_vis_logic.jpg"
+            # )
+            # vis_imged = viser(
+            #     np_img, wired_table_results, save_html_path, save_drawed_path, save_logic_path
+            # )
+
             wired_html_code = wired_table_results.pred_html
 
             wired_len = count_table_cells_physical(wired_html_code)
@@ -275,7 +285,6 @@ class UnetTableModel:
             # 判断是否使用无线表格模型的结果
             if (
                 wired_len <= int(wireless_len * 0.55)+1  # 有线模型检测到的单元格数太少(低于无线模型的50%)
-                # or ((round(wireless_len*1.2) < wired_len) and (wired_len < (2 * wireless_len)) and table_cls_score <= 0.94)  # 有线模型检测到的单元格数反而更多
                 or (0 <= gap_of_len <= 5 and wired_len <= round(wireless_len * 0.75))  # 两者相差不大但有线模型结果较少
                 or (gap_of_len == 0 and wired_len <= 4)  # 单元格数量完全相等且总量小于等于4
             ):

+ 1 - 1
mineru/model/table/rec/unet_table/utils.py

@@ -429,7 +429,7 @@ class VisTable:
         :return:
         """
         # 读取原图
-        img = cv2.imread(img_path)
+        img = img_path
         img = cv2.copyMakeBorder(
             img, 0, 0, 0, 100, cv2.BORDER_CONSTANT, value=[255, 255, 255]
         )