parent
37677385af
commit
7390b9e6ad
5 changed files with 429 additions and 130 deletions
@ -0,0 +1,278 @@ |
||||
# Copyright (c) 2022 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. |
||||
|
||||
import itertools |
||||
import warnings |
||||
|
||||
import cv2 |
||||
import numpy as np |
||||
from skimage import morphology |
||||
from scipy import ndimage, optimize |
||||
|
||||
with warnings.catch_warnings(): |
||||
warnings.filterwarnings("ignore", category=DeprecationWarning) |
||||
from sklearn import metrics |
||||
from sklearn.cluster import KMeans |
||||
|
||||
from .utils import prepro_mask, calc_distance |
||||
|
||||
|
||||
def cut_road_connection(mask: np.ndarray, line_width: int=6) -> np.ndarray: |
||||
""" |
||||
Connecting cut road lines. |
||||
|
||||
The original article refers to |
||||
Wang B, Chen Z, et al. "Road extraction of high-resolution satellite remote sensing images in U-Net network with consideration of connectivity." |
||||
(http://hgs.publish.founderss.cn/thesisDetails?columnId=4759509). |
||||
|
||||
This algorithm has no public code. |
||||
The implementation procedure refers to original article, |
||||
and it is not fully consistent with the article: |
||||
1. The way to determine the optimal number of clusters k used in k-means clustering is not described in the original article. In this implementation, we use the k that reports the highest silhouette score. |
||||
2. We unmark the breakpoints if the angle between the two road extensions is less than 90°. |
||||
|
||||
Args: |
||||
mask (np.ndarray): Mask of road. |
||||
line_width (int, optional): Width of the line used for patching. |
||||
. Default is 6. |
||||
|
||||
Returns: |
||||
np.ndarray: Mask of road after connecting cut road lines. |
||||
""" |
||||
mask = prepro_mask(mask) |
||||
skeleton = morphology.skeletonize(mask).astype("uint8") |
||||
break_points = _find_breakpoint(skeleton) |
||||
labels = _k_means(break_points) |
||||
match_points = _get_match_points(break_points, labels) |
||||
res = _draw_curve(mask, skeleton, match_points, line_width) |
||||
return res |
||||
|
||||
|
||||
def _find_breakpoint(skeleton): |
||||
kernel_3x3 = np.ones((3, 3), dtype="uint8") |
||||
k3 = ndimage.convolve(skeleton, kernel_3x3) |
||||
point_map = np.zeros_like(k3) |
||||
point_map[k3 == 2] = 1 |
||||
point_map *= skeleton * 255 |
||||
# boundary filtering |
||||
filter_w = 5 |
||||
cropped = point_map[filter_w:-filter_w, filter_w:-filter_w] |
||||
padded = np.pad(cropped, (filter_w, filter_w), mode="constant") |
||||
breakpoints = np.column_stack(np.where(padded == 255)) |
||||
return breakpoints |
||||
|
||||
|
||||
def _k_means(data): |
||||
silhouette_int = -1 # threshold |
||||
labels = None |
||||
for k in range(2, data.shape[0]): |
||||
kms = KMeans(k, random_state=66) |
||||
labels_tmp = kms.fit_predict(data) # train |
||||
silhouette = metrics.silhouette_score(data, labels_tmp) |
||||
if silhouette > silhouette_int: # better |
||||
silhouette_int = silhouette |
||||
labels = labels_tmp |
||||
return labels |
||||
|
||||
|
||||
def _get_match_points(break_points, labels): |
||||
match_points = {} |
||||
for point, lab in zip(break_points, labels): |
||||
if lab in match_points.keys(): |
||||
match_points[lab].append(point) |
||||
else: |
||||
match_points[lab] = [point] |
||||
return match_points |
||||
|
||||
|
||||
def _draw_curve(mask, skeleton, match_points, line_width): |
||||
result = mask * 255 |
||||
for v in match_points.values(): |
||||
p_num = len(v) |
||||
if p_num == 2: |
||||
points_list = _curve_backtracking(v, skeleton) |
||||
if points_list is not None: |
||||
result = _broken_wire_repair(result, points_list, line_width) |
||||
elif p_num == 3: |
||||
sim_v = list(itertools.combinations(v, 2)) |
||||
min_di = 1e6 |
||||
for vij in sim_v: |
||||
di = calc_distance(vij[0][np.newaxis], vij[1][np.newaxis]) |
||||
if di < min_di: |
||||
vv = vij |
||||
min_di = di |
||||
points_list = _curve_backtracking(vv, skeleton) |
||||
if points_list is not None: |
||||
result = _broken_wire_repair(result, points_list, line_width) |
||||
return result |
||||
|
||||
|
||||
def _curve_backtracking(add_lines, skeleton): |
||||
points_list = [] |
||||
p1 = add_lines[0] |
||||
p2 = add_lines[1] |
||||
bpk1, ps1 = _calc_angle_by_road(p1, skeleton) |
||||
bpk2, ps2 = _calc_angle_by_road(p2, skeleton) |
||||
if _check_angle(bpk1, bpk2): |
||||
points_list.append(( |
||||
np.array( |
||||
ps1, dtype="int64"), |
||||
add_lines[0], |
||||
add_lines[1], |
||||
np.array( |
||||
ps2, dtype="int64"), )) |
||||
return points_list |
||||
else: |
||||
return None |
||||
|
||||
|
||||
def _broken_wire_repair(mask, points_list, line_width): |
||||
d_mask = mask.copy() |
||||
for points in points_list: |
||||
nx, ny = _line_cubic(points) |
||||
for i in range(len(nx) - 1): |
||||
loc_p1 = (int(ny[i]), int(nx[i])) |
||||
loc_p2 = (int(ny[i + 1]), int(nx[i + 1])) |
||||
cv2.line(d_mask, loc_p1, loc_p2, [255], line_width) |
||||
return d_mask |
||||
|
||||
|
||||
def _calc_angle_by_road(p, skeleton, num_circle=10): |
||||
def _not_in(p1, ps): |
||||
for p in ps: |
||||
if p1[0] == p[0] and p1[1] == p[1]: |
||||
return False |
||||
return True |
||||
|
||||
h, w = skeleton.shape |
||||
tmp_p = p.tolist() if isinstance(p, np.ndarray) else p |
||||
tmp_p = [int(tmp_p[0]), int(tmp_p[1])] |
||||
ps = [] |
||||
ps.append(tmp_p) |
||||
for _ in range(num_circle): |
||||
t_x = 0 if tmp_p[0] - 1 < 0 else tmp_p[0] - 1 |
||||
t_y = 0 if tmp_p[1] - 1 < 0 else tmp_p[1] - 1 |
||||
b_x = w if tmp_p[0] + 1 >= w else tmp_p[0] + 1 |
||||
b_y = h if tmp_p[1] + 1 >= h else tmp_p[1] + 1 |
||||
if int(np.sum(skeleton[t_x:b_x + 1, t_y:b_y + 1])) <= 3: |
||||
for i in range(t_x, b_x + 1): |
||||
for j in range(t_y, b_y + 1): |
||||
if skeleton[i, j] == 1: |
||||
pp = [int(i), int(j)] |
||||
if _not_in(pp, ps): |
||||
tmp_p = pp |
||||
ps.append(tmp_p) |
||||
# calc angle |
||||
theta = _angle_regression(ps) |
||||
dx, dy = np.cos(theta), np.sin(theta) |
||||
# calc direction |
||||
start = ps[-1] |
||||
end = ps[0] |
||||
if end[1] < start[1] or (end[1] == start[1] and end[0] < start[0]): |
||||
dx *= -1 |
||||
dy *= -1 |
||||
return [dx, dy], start |
||||
|
||||
|
||||
def _angle_regression(datas): |
||||
def _linear(x: float, k: float, b: float) -> float: |
||||
return k * x + b |
||||
|
||||
xs = [] |
||||
ys = [] |
||||
for data in datas: |
||||
xs.append(data[0]) |
||||
ys.append(data[1]) |
||||
xs_arr = np.array(xs) |
||||
ys_arr = np.array(ys) |
||||
# horizontal |
||||
if len(np.unique(xs_arr)) == 1: |
||||
theta = np.pi / 2 |
||||
# vertical |
||||
elif len(np.unique(ys_arr)) == 1: |
||||
theta = 0 |
||||
# cross calc |
||||
else: |
||||
k1, b1 = optimize.curve_fit(_linear, xs_arr, ys_arr)[0] |
||||
k2, b2 = optimize.curve_fit(_linear, ys_arr, xs_arr)[0] |
||||
err1 = 0 |
||||
err2 = 0 |
||||
for x, y in zip(xs_arr, ys_arr): |
||||
err1 += abs(_linear(x, k1, b1) - y) / np.sqrt(k1**2 + 1) |
||||
err2 += abs(_linear(y, k2, b2) - x) / np.sqrt(k2**2 + 1) |
||||
if err1 <= err2: |
||||
theta = (np.arctan(k1) + 2 * np.pi) % (2 * np.pi) |
||||
else: |
||||
theta = (np.pi / 2.0 - np.arctan(k2) + 2 * np.pi) % (2 * np.pi) |
||||
# [0, 180) |
||||
theta = theta * 180 / np.pi + 90 |
||||
while theta >= 180: |
||||
theta -= 180 |
||||
theta -= 90 |
||||
if theta < 0: |
||||
theta += 180 |
||||
return theta * np.pi / 180 |
||||
|
||||
|
||||
def _cubic(x, y): |
||||
def _func(x, a, b, c, d): |
||||
return a * x**3 + b * x**2 + c * x + d |
||||
|
||||
arr_x = np.array(x).reshape((4, )) |
||||
arr_y = np.array(y).reshape((4, )) |
||||
popt1 = np.polyfit(arr_x, arr_y, 3) |
||||
popt2 = np.polyfit(arr_y, arr_x, 3) |
||||
x_min = np.min(arr_x) |
||||
x_max = np.max(arr_x) |
||||
y_min = np.min(arr_y) |
||||
y_max = np.max(arr_y) |
||||
nx = np.arange(x_min, x_max + 1, 1) |
||||
y_estimate = [_func(i, popt1[0], popt1[1], popt1[2], popt1[3]) for i in nx] |
||||
ny = np.arange(y_min, y_max + 1, 1) |
||||
x_estimate = [_func(i, popt2[0], popt2[1], popt2[2], popt2[3]) for i in ny] |
||||
if np.max(y_estimate) - np.min(y_estimate) <= np.max(x_estimate) - np.min( |
||||
x_estimate): |
||||
return nx, y_estimate |
||||
else: |
||||
return x_estimate, ny |
||||
|
||||
|
||||
def _line_cubic(points): |
||||
xs = [] |
||||
ys = [] |
||||
for p in points: |
||||
x, y = p |
||||
xs.append(x) |
||||
ys.append(y) |
||||
nx, ny = _cubic(xs, ys) |
||||
return nx, ny |
||||
|
||||
|
||||
def _get_theta(dy, dx): |
||||
theta = np.arctan2(dy, dx) * 180 / np.pi |
||||
if theta < 0.0: |
||||
theta = 360.0 - abs(theta) |
||||
return float(theta) |
||||
|
||||
|
||||
def _check_angle(bpk1, bpk2, ang_threshold=90): |
||||
af1 = _get_theta(bpk1[0], bpk1[1]) |
||||
af2 = _get_theta(bpk2[0], bpk2[1]) |
||||
ang_diff = abs(af1 - af2) |
||||
if ang_diff > 180: |
||||
ang_diff = 360 - ang_diff |
||||
if ang_diff > ang_threshold: |
||||
return True |
||||
else: |
||||
return False |
Loading…
Reference in new issue