Browse Source

feat: add initial project structure and CLI entry points for mineru

myhloli 3 months ago
parent
commit
fc5179ce5e
1 changed files with 4 additions and 61 deletions
  1. 4 61
      mineru/model/table/rec/unet_table/table_structure_unet.py

+ 4 - 61
mineru/model/table/rec/unet_table/table_structure_unet.py

@@ -94,8 +94,7 @@ class TSRUnet:
         extend_line = (
             kwargs.get("extend_line", enhance_box_line) if kwargs else enhance_box_line
         )  # 是否进行线段延长使得端点连接
-        # 是否进行旋转修正
-        rotated_fix = kwargs.get("rotated_fix") if kwargs else True
+
         ori_shape = img.shape
         pred = np.uint8(pred)
         hpred = copy.deepcopy(pred)  # 横线
@@ -131,16 +130,9 @@ class TSRUnet:
             rowboxes, colboxes = final_adjust_lines(rowboxes, colboxes)
         line_img = np.zeros(img.shape[:2], dtype="uint8")
         line_img = draw_lines(line_img, rowboxes + colboxes, color=255, lineW=2)
-        rotated_angle = self.cal_rotate_angle(line_img)
-        if rotated_fix and abs(rotated_angle) > 0.3:
-            rotated_line_img = self.rotate_image(line_img, rotated_angle)
-            rotated_polygons = self.cal_region_boxes(rotated_line_img)
-            polygons = self.unrotate_polygons(
-                rotated_polygons, rotated_angle, line_img.shape
-            )
-        else:
-            polygons = self.cal_region_boxes(line_img)
-            rotated_polygons = polygons.copy()
+
+        polygons = self.cal_region_boxes(line_img)
+        rotated_polygons = polygons.copy()
         return polygons, rotated_polygons
 
     def cal_region_boxes(self, tmp):
@@ -155,52 +147,3 @@ class TSRUnet:
             adjust_box=False,
         )  # 最后一个参数改为False
         return np.array(ceilboxes)
-
-    def cal_rotate_angle(self, tmp):
-        # 计算最外侧的旋转框
-        contours, _ = cv2.findContours(tmp, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        if not contours:
-            return 0
-        largest_contour = max(contours, key=cv2.contourArea)
-        rect = cv2.minAreaRect(largest_contour)
-        # 计算旋转角度
-        angle = rect[2]
-        if angle < -45:
-            angle += 90
-        elif angle > 45:
-            angle -= 90
-        return angle
-
-    def rotate_image(self, image, angle):
-        # 获取图像的中心点
-        (h, w) = image.shape[:2]
-        center = (w // 2, h // 2)
-
-        # 计算旋转矩阵
-        M = cv2.getRotationMatrix2D(center, angle, 1.0)
-
-        # 进行旋转
-        rotated_image = cv2.warpAffine(
-            image, M, (w, h), flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_REPLICATE
-        )
-
-        return rotated_image
-
-    def unrotate_polygons(
-        self, polygons: np.ndarray, angle: float, img_shape: tuple
-    ) -> np.ndarray:
-        # 将多边形旋转回原始位置
-        (h, w) = img_shape
-        center = (w // 2, h // 2)
-        M_inv = cv2.getRotationMatrix2D(center, -angle, 1.0)
-
-        # 将 (N, 8) 转换为 (N, 4, 2)
-        polygons_reshaped = polygons.reshape(-1, 4, 2)
-
-        # 批量逆旋转
-        unrotated_polygons = cv2.transform(polygons_reshaped, M_inv)
-
-        # 将 (N, 4, 2) 转换回 (N, 8)
-        unrotated_polygons = unrotated_polygons.reshape(-1, 8)
-
-        return unrotated_polygons