rs_regularization.py

# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
This code is based on https://github.com/niecongchong/RS-building-regularization
Ths copyright of niecongchong/RS-building-regularization is as follows:
Apache License [see LICENSE for details]
"""

import cv2
import matplotlib.pyplot as plt
import numpy as np
from .rdp_alg import rdp
from .cal_point import cal_ang, cal_dist, cal_azimuth
from .rotate_ang import Nrotation_angle_get_coor_coordinates, Srotation_angle_get_coor_coordinates
from .cal_line import line, intersection, par_line_dist, point_in_line


def boundary_regularization(contours, img_shape, epsilon=6):
    h, w = img_shape[0:2]
    # 轮廓定位
    contours = np.squeeze(contours)
    # 轮廓精简DP
    contours = rdp(contours, epsilon=epsilon)
    contours[:, 1] = h - contours[:, 1]
    # 轮廓规则化
    dists = []
    azis = []
    azis_index = []
    # 获取每条边的长度和方位角
    for i in range(contours.shape[0]):
        cur_index = i
        next_index = i + 1 if i < contours.shape[0] - 1 else 0
        prev_index = i - 1
        cur_point = contours[cur_index]
        nest_point = contours[next_index]
        prev_point = contours[prev_index]
        dist = cal_dist(cur_point, nest_point)
        azi = cal_azimuth(cur_point, nest_point)
        dists.append(dist)
        azis.append(azi)
        azis_index.append([cur_index, next_index])
    # 以最长的边的方向作为主方向
    longest_edge_idex = np.argmax(dists)
    main_direction = azis[longest_edge_idex]
    # 方向纠正，绕中心点旋转到与主方向垂直或者平行
    correct_points = []
    para_vetr_idxs = []  # 0平行 1垂直
    for i, (azi, (point_0_index,
                  point_1_index)) in enumerate(zip(azis, azis_index)):
        if i == longest_edge_idex:
            correct_points.append(
                [contours[point_0_index], contours[point_1_index]])
            para_vetr_idxs.append(0)
        else:
            # 确定旋转角度
            rotate_ang = main_direction - azi
            if np.abs(rotate_ang) < 180 / 4:
                rotate_ang = rotate_ang
                para_vetr_idxs.append(0)
            elif np.abs(rotate_ang) >= 90 - 180 / 4:
                rotate_ang = rotate_ang + 90
                para_vetr_idxs.append(1)
            # 执行旋转任务
            point_0 = contours[point_0_index]
            point_1 = contours[point_1_index]
            point_middle = (point_0 + point_1) / 2
            if rotate_ang > 0:
                rotate_point_0 = Srotation_angle_get_coor_coordinates(
                    point_0, point_middle, np.abs(rotate_ang))
                rotate_point_1 = Srotation_angle_get_coor_coordinates(
                    point_1, point_middle, np.abs(rotate_ang))
            elif rotate_ang < 0:
                rotate_point_0 = Nrotation_angle_get_coor_coordinates(
                    point_0, point_middle, np.abs(rotate_ang))
                rotate_point_1 = Nrotation_angle_get_coor_coordinates(
                    point_1, point_middle, np.abs(rotate_ang))
            else:
                rotate_point_0 = point_0
                rotate_point_1 = point_1
            correct_points.append([rotate_point_0, rotate_point_1])
    correct_points = np.array(correct_points)
    # 相邻边校正，垂直取交点，平行平移短边或者加线
    final_points = []
    final_points.append(correct_points[0][0])
    for i in range(correct_points.shape[0] - 1):
        cur_index = i
        next_index = i + 1 if i < correct_points.shape[0] - 1 else 0
        cur_edge_point_0 = correct_points[cur_index][0]
        cur_edge_point_1 = correct_points[cur_index][1]
        next_edge_point_0 = correct_points[next_index][0]
        next_edge_point_1 = correct_points[next_index][1]
        cur_para_vetr_idx = para_vetr_idxs[cur_index]
        next_para_vetr_idx = para_vetr_idxs[next_index]
        if cur_para_vetr_idx != next_para_vetr_idx:
            # 垂直取交点
            L1 = line(cur_edge_point_0, cur_edge_point_1)
            L2 = line(next_edge_point_0, next_edge_point_1)
            point_intersection = intersection(L1, L2)
            final_points.append(point_intersection)
        elif cur_para_vetr_idx == next_para_vetr_idx:
            # 平行分两种，一种加短线，一种平移，取决于距离阈值
            L1 = line(cur_edge_point_0, cur_edge_point_1)
            L2 = line(next_edge_point_0, next_edge_point_1)
            marg = par_line_dist(L1, L2)
            if marg < 3:
                # 平移
                point_move = point_in_line(
                    next_edge_point_0[0], next_edge_point_0[1],
                    cur_edge_point_0[0], cur_edge_point_0[1],
                    cur_edge_point_1[0], cur_edge_point_1[1])
                final_points.append(point_move)
                # 更新平移之后的下一条边
                correct_points[next_index][0] = point_move
                correct_points[next_index][1] = point_in_line(
                    next_edge_point_1[0], next_edge_point_1[1],
                    cur_edge_point_0[0], cur_edge_point_0[1],
                    cur_edge_point_1[0], cur_edge_point_1[1])
            else:
                # 加线
                add_mid_point = (cur_edge_point_1 + next_edge_point_0) / 2
                add_point_1 = point_in_line(
                    add_mid_point[0], add_mid_point[1], cur_edge_point_0[0],
                    cur_edge_point_0[1], cur_edge_point_1[0],
                    cur_edge_point_1[1])
                add_point_2 = point_in_line(
                    add_mid_point[0], add_mid_point[1], next_edge_point_0[0],
                    next_edge_point_0[1], next_edge_point_1[0],
                    next_edge_point_1[1])
                final_points.append(add_point_1)
                final_points.append(add_point_2)
    final_points.append(final_points[0])
    final_points = np.array(final_points)
    final_points[:, 1] = h - final_points[:, 1]
    final_points = final_points[np.newaxis, :].transpose((1, 0, 2))
    return final_points


# def rs_build_re(mask):
#     # 中值滤波，去噪
#     ori_img = cv2.medianBlur(mask, 5)
#     ori_img = cv2.cvtColor(ori_img, cv2.COLOR_BGR2GRAY)
#     ret, ori_img = cv2.threshold(ori_img, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
#     # 连通域分析
#     num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(ori_img, connectivity=8)
#     # 遍历联通域
#     for i in range(1, num_labels):
#         img = np.zeros_like(labels)
#         index = np.where(labels==i)
#         img[index] = 255
#         img = np.array(img, dtype=np.uint8)
#         regularization_contour = boundary_regularization(img).astype(np.int32)
#         cv2.polylines(img=mask, pts=[regularization_contour], isClosed=True, color=(255, 0, 0), thickness=5)
#         single_out = np.zeros_like(mask)
#         cv2.polylines(img=single_out, pts=[regularization_contour], isClosed=True, color=(255, 0, 0), thickness=5)
#         cv2.imwrite('single_out_{}.jpg'.format(i), single_out)