瀏覽代碼

fix: refine image rotation logic and replace logging with logger for better consistency

myhloli 3 月之前
父節點
當前提交
21ce40a90c
共有 2 個文件被更改,包括 6 次插入12 次删除
  1. 4 9
      mineru/model/ori_cls/paddle_ori_cls.py
  2. 2 3
      mineru/model/table/rec/unet_table/unet_table.py

+ 4 - 9
mineru/model/ori_cls/paddle_ori_cls.py

@@ -101,15 +101,10 @@ class PaddleOrientationClsModel:
                     (result,) = self.sess.run(None, {"x": x})
                     (result,) = self.sess.run(None, {"x": x})
                     label = self.labels[np.argmax(result)]
                     label = self.labels[np.argmax(result)]
 
 
-                    if label == "90":
-                        rotation = cv2.ROTATE_90_COUNTERCLOCKWISE
-                        img = cv2.rotate(np.asarray(img), rotation)
-                    # elif label == "180":
-                    #     rotation = cv2.ROTATE_180
-                    #     img = cv2.rotate(np.asarray(img), rotation)
-                    elif label == "270":
+                    if label == "270":
                         rotation = cv2.ROTATE_90_CLOCKWISE
                         rotation = cv2.ROTATE_90_CLOCKWISE
                         img = cv2.rotate(np.asarray(img), rotation)
                         img = cv2.rotate(np.asarray(img), rotation)
-                    else:
-                        img = np.array(img)
+                    else:  # 除了270度,都认为是90度
+                        rotation = cv2.ROTATE_90_COUNTERCLOCKWISE
+                        img = cv2.rotate(np.asarray(img), rotation)
         return img
         return img

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

@@ -1,5 +1,4 @@
 import html
 import html
-import logging
 import os
 import os
 import time
 import time
 import traceback
 import traceback
@@ -61,7 +60,7 @@ class UnetTableRecognition:
         img = self.load_img(img)
         img = self.load_img(img)
         polygons, rotated_polygons = self.table_structure(img, **kwargs)
         polygons, rotated_polygons = self.table_structure(img, **kwargs)
         if polygons is None:
         if polygons is None:
-            logging.warning("polygons is None.")
+            logger.warning("polygons is None.")
             return UnetTableOutput("", None, None, 0.0)
             return UnetTableOutput("", None, None, 0.0)
 
 
         try:
         try:
@@ -102,7 +101,7 @@ class UnetTableRecognition:
             elapse = time.perf_counter() - s
             elapse = time.perf_counter() - s
 
 
         except Exception:
         except Exception:
-            logging.warning(traceback.format_exc())
+            logger.warning(traceback.format_exc())
             return UnetTableOutput("", None, None, 0.0)
             return UnetTableOutput("", None, None, 0.0)
         return UnetTableOutput(pred_html, polygons, logi_points, elapse)
         return UnetTableOutput(pred_html, polygons, logi_points, elapse)