|
|
@@ -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
|