Procházet zdrojové kódy

fix: improve image orientation handling in batch analysis and streamline table classification process

myhloli před 3 měsíci
rodič
revize
11c4a0c6b6
1 změnil soubory, kde provedl 16 přidání a 15 odebrání
  1. 16 15
      mineru/backend/pipeline/batch_analyze.py

+ 16 - 15
mineru/backend/pipeline/batch_analyze.py

@@ -251,17 +251,32 @@ class BatchAnalyze:
             for table_res_dict in tqdm(table_res_list_all_page, desc="Table Predict"):
                 _lang = table_res_dict['lang']
 
+                # 调整图片方向
+                img_orientation_cls_model = atom_model_manager.get_atom_model(
+                    atom_model_name=AtomicModel.ImgOrientationCls,
+                )
+                try:
+                    table_img = img_orientation_cls_model.predict(
+                        table_res_dict["table_img"]
+                    )
+                except Exception as e:
+                    logger.warning(
+                        f"Image orientation classification failed: {e}, using original image"
+                    )
+                    table_img = table_res_dict["table_img"]
+
                 # 有线表/无线表分类
                 table_cls_model = atom_model_manager.get_atom_model(
                     atom_model_name=AtomicModel.TableCls,
                 )
                 try:
-                    table_label = table_cls_model.predict(table_res_dict["table_img"])
+                    table_label = table_cls_model.predict(table_img)
                 except Exception as e:
                     table_label = AtomicModel.WirelessTable
                     logger.warning(
                         f"Table classification failed: {e}, using default model {table_label}"
                     )
+                # table_label = AtomicModel.WirelessTable
                 # logger.debug(f"Table classification result: {table_label}")
                 if table_label not in [AtomicModel.WirelessTable, AtomicModel.WiredTable]:
                     raise ValueError(
@@ -274,20 +289,6 @@ class BatchAnalyze:
                     lang=_lang,
                 )
 
-                # 调整图片方向
-                img_orientation_cls_model = atom_model_manager.get_atom_model(
-                    atom_model_name=AtomicModel.ImgOrientationCls,
-                )
-                try:
-                    table_img = img_orientation_cls_model.predict(
-                        table_res_dict["table_img"]
-                    )
-                except Exception as e:
-                    logger.warning(
-                        f"Image orientation classification failed: {e}, using original image"
-                    )
-                    table_img = table_res_dict["table_img"]
-
                 html_code, table_cell_bboxes, logic_points, elapse = table_model.predict(table_img)
                 # 判断是否返回正常
                 if html_code: