Browse Source

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

myhloli 3 months ago
parent
commit
21ce40a90c

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

@@ -101,15 +101,10 @@ class PaddleOrientationClsModel:
                     (result,) = self.sess.run(None, {"x": x})
                     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
                         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

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

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