terrainClassHandover/main_regularization_selfExpand.py

import os
import cv2
import matplotlib.pyplot as plt
import numpy as np
from rdp_alg import rdp
from cal_dist_ang import cal_ang, cal_dist, azimuthAngle
from rotate_ang import Nrotation_angle_get_coor_coordinates, Srotation_angle_get_coor_coordinates
from line_intersection import line, intersection, par_line_dist, point_in_line


def boundary_regularization(img, epsilon=6):
    h, w = img.shape[0:2]

    # 轮廓定位
    contours, hierarchy = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)  # 检索所有轮廓
    # contours, hierarchy = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)  # 只检索最外面的轮廓
    contours = np.squeeze(contours[0])  # [[x1,y1], [x2, y2],...]
    # print("line17", contours)
    # 轮廓精简（DP）
    contours = rdp(contours, epsilon=epsilon)
    # print("line20", contours[:, 1], h)  # [ 409, 415, 539, 573, 610], 27710
    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 = azimuthAngle(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]
    return final_points


imgPath = "./input"
imgList = os.listdir(imgPath)
for i in range(len(imgList)):
    img = cv2.imread(imgPath + os.sep + imgList[i])  # 读取彩色的分割图像
    imgGray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)


    imgDB = imgGray.copy()
    imgDB[imgDB == 38] = 0  # 删除建筑物  filterBuilding.png
    # imgDB = cv2.cvtColor(imgDB, cv2.COLOR_BGR2GRAY)

    imgGray[imgGray != 38] = 0
    ori_img1 = cv2.cvtColor(imgGray, cv2.COLOR_GRAY2BGR)  # rgb值相同的24位图像

    h, w = ori_img1.shape[0], ori_img1.shape[1]
    # 中值滤波，去噪
    ori_img = cv2.medianBlur(ori_img1, 5)  # 滤波核大小为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)  # 参数8表示8连通。返回值：所有连通域的数目，图像上每一像素的标记，每一个标记的统计信息，连通域的中心点

    # 遍历连通域
    allCnt = []
    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=ori_img1, pts=[regularization_contour], isClosed=True, color=(255, 0, 0), thickness=5)  # 原始
        # print("line153", type(regularization_contour))  # [[999, 666], [222, 111],... ]

        # single_out = np.zeros_like(ori_img1)
        # cv2.polylines(img=single_out, pts=[regularization_contour], isClosed=True, color=(255, 0, 0), thickness=5)
        # cv2.imwrite('./middle/' + 'single_out_{}.jpg'.format(i), single_out)

        rows = regularization_contour.shape[0]
        regularization_contour = regularization_contour.reshape(rows, 1, 2)
        regularization_contour = regularization_contour.astype(int)
        allCnt.append(regularization_contour)
        # print("line162", regularization_contour)
        # print("line162", regularization_contour.shape)

    buildingMask = np.zeros((h, w), dtype='uint8')
    cv2.fillPoly(buildingMask, allCnt, color=38)
    img2 = buildingMask.copy()
    cv2.imwrite("./output/building.png", img2)


    buildingMask[buildingMask == 0] = 255
    buildingMask[buildingMask == 38] = 0  # step2.png
    img3 = cv2.bitwise_and(imgDB, imgDB, mask=buildingMask)  # 在去掉建筑物区域的图像中，再去掉“优化后的建筑物”边界范围内的区域
    finalResult = cv2.bitwise_or(img2, img3)

    cv2.imwrite('./output/finalResult.png', finalResult)





























# imgPath = "./input"
# imgList = os.listdir(imgPath)
# for i in range(len(imgList)):
#     img = cv2.imread(imgPath + os.sep + imgList[i])  # 读取彩色的分割图像
#     imgGray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
#     img = cv2.cvtColor(imgGray, cv2.COLOR_GRAY2BGR)  # rgb值相同的24位图像
#
#
# ori_img1 = cv2.imread('./input/1.png')
# h, w = ori_img1.shape[0], ori_img1.shape[1]
# # 中值滤波，去噪
# ori_img = cv2.medianBlur(ori_img1, 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)
#
#
# # 遍历联通域
# allCnt = []
# 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=ori_img1, pts=[regularization_contour], isClosed=True, color=(255, 0, 0), thickness=5)  # 原始
#     # print("line153", type(regularization_contour))  # [[999, 666], [222, 111],... ]
#
#     # single_out = np.zeros_like(ori_img1)
#     # cv2.polylines(img=single_out, pts=[regularization_contour], isClosed=True, color=(255, 0, 0), thickness=5)
#     # cv2.imwrite('./middle/' + 'single_out_{}.jpg'.format(i), single_out)
#
#     rows = regularization_contour.shape[0]
#     regularization_contour = regularization_contour.reshape(rows, 1, 2)
#     regularization_contour = regularization_contour.astype(int)
#     allCnt.append(regularization_contour)
#     # print("line162", regularization_contour)
#     # print("line162", regularization_contour.shape)
#
# mask = np.zeros((h, w), dtype='uint8')
# cv2.fillPoly(mask, allCnt, color=38)
# cv2.imwrite("./output/new.png", mask)
#
# # cv2.imwrite('./output/result.png', ori_img1)