Parcourir la source

fix: adjust confidence thresholds and refine table detection logic for improved accuracy

myhloli il y a 3 mois
Parent
commit
a01c4bbe66

+ 3 - 1
mineru/model/layout/doclayout_yolo.py

@@ -60,10 +60,12 @@ class DocLayoutYOLOModel:
         with tqdm(total=len(images), desc="Layout Predict") as pbar:
             for idx in range(0, len(images), batch_size):
                 batch = images[idx: idx + batch_size]
+                if batch_size == 1:
+                    conf = 0.9 * self.conf
                 predictions = self.model.predict(
                     batch,
                     imgsz=self.imgsz,
-                    conf=self.conf,
+                    conf=conf,
                     iou=self.iou,
                     verbose=False,
                 )

+ 1 - 1
mineru/model/table/cls/paddle_table_cls.py

@@ -67,6 +67,6 @@ class PaddleTableClsModel:
         idx = np.argmax(result)
         conf = float(np.max(result))
         # logger.debug(f"Table classification result: {self.labels[idx]} with confidence {conf:.4f}")
-        if idx == 0 and conf < 0.85:
+        if idx == 0 and conf < 0.8:
             idx = 1
         return self.labels[idx], conf

+ 2 - 2
mineru/model/table/rec/unet_table/unet_table.py

@@ -62,7 +62,7 @@ class UnetTableRecognition:
         img = self.load_img(img)
         polygons, rotated_polygons = self.table_structure(img, **kwargs)
         if polygons is None:
-            logger.warning("polygons is None.")
+            # logger.warning("polygons is None.")
             return UnetTableOutput("", None, None, 0.0)
 
         try:
@@ -243,7 +243,7 @@ class UnetTableModel:
                 # 判断是否使用无线表格模型的结果
                 if (
                     wired_len <= round(wireless_len * 0.5)  # 有线模型检测到的单元格数太少(低于无线模型的50%)
-                    or (wireless_len < wired_len < (2 * wireless_len) and table_cls_score <= 0.949)  # 有线模型检测到的单元格数反而更多
+                    or ((wireless_len < wired_len) and (wired_len < (2 * wireless_len)) and table_cls_score <= 0.949)  # 有线模型检测到的单元格数反而更多
                     or (0 <= gap_of_len <= 5 and wired_len <= round(wireless_len * 0.75))  # 两者相差不大但有线模型结果较少
                     or (gap_of_len == 0 and wired_len <= 4)  # 单元格数量完全相等且总量小于等于4
                 ):