166 lines
7.0 KiB
Python
166 lines
7.0 KiB
Python
import numpy as np
|
||
import cv2
|
||
|
||
def _are_lines_mergeable(line1_pts, line2_pts,
|
||
max_angle_diff_deg,
|
||
max_ep_gap_abs, max_ep_gap_factor,
|
||
max_p_dist_abs, max_p_dist_factor):
|
||
"""
|
||
判断两条线段是否可以合并。
|
||
lineX_pts: [x1, y1, x2, y2]
|
||
max_angle_diff_deg: 最大角度差 (度)
|
||
max_ep_gap_abs: 端点间最大绝对间隙
|
||
max_ep_gap_factor: 端点间最大相对间隙 (基于较短线段长度的因子)
|
||
max_p_dist_abs: 最大绝对垂直距离
|
||
max_p_dist_factor: 最大相对垂直距离 (基于较短线段长度的因子)
|
||
"""
|
||
x1a, y1a, x2a, y2a = line1_pts
|
||
x1b, y1b, x2b, y2b = line2_pts
|
||
|
||
p1a, p2a = np.array([x1a, y1a]), np.array([x2a, y2a])
|
||
p1b, p2b = np.array([x1b, y1b]), np.array([x2b, y2b])
|
||
|
||
len_a = np.linalg.norm(p2a - p1a)
|
||
len_b = np.linalg.norm(p2b - p1b)
|
||
|
||
if len_a < 1e-3 or len_b < 1e-3: return False # 避免零长度线段
|
||
shorter_len = min(len_a, len_b)
|
||
|
||
# 1. 角度相似性
|
||
angle_a = np.arctan2(y2a - y1a, x2a - x1a)
|
||
angle_b = np.arctan2(y2b - y1b, x2b - x1b)
|
||
angle_diff = np.degrees(angle_a - angle_b)
|
||
|
||
while angle_diff <= -180: angle_diff += 360
|
||
while angle_diff > 180: angle_diff -= 360
|
||
|
||
is_angle_similar = abs(angle_diff) < max_angle_diff_deg or \
|
||
abs(abs(angle_diff) - 180) < max_angle_diff_deg
|
||
if not is_angle_similar:
|
||
return False
|
||
|
||
# 2. 共线性: 检查一个线段的端点到另一个线段（无限延长）的垂直距离
|
||
current_max_allowed_perp_dist = max(max_p_dist_abs, max_p_dist_factor * shorter_len)
|
||
|
||
def get_perp_dist_point_to_inf_line(p, line_p1, line_p2, line_len):
|
||
if line_len < 1e-3: return np.linalg.norm(p - line_p1) # 线是点
|
||
return abs((line_p2[1]-line_p1[1])*p[0] - (line_p2[0]-line_p1[0])*p[1] + \
|
||
line_p2[0]*line_p1[1] - line_p2[1]*line_p1[0]) / line_len
|
||
|
||
dist_b1_to_a_inf = get_perp_dist_point_to_inf_line(p1b, p1a, p2a, len_a)
|
||
dist_b2_to_a_inf = get_perp_dist_point_to_inf_line(p2b, p1a, p2a, len_a)
|
||
dist_a1_to_b_inf = get_perp_dist_point_to_inf_line(p1a, p1b, p2b, len_b)
|
||
dist_a2_to_b_inf = get_perp_dist_point_to_inf_line(p2a, p1b, p2b, len_b)
|
||
|
||
# 如果两条线都很短，那么允许的垂直距离也应该小
|
||
# 如果一条线的端点离另一条无限延长的线太远，则不合并
|
||
if min(dist_b1_to_a_inf, dist_b2_to_a_inf) > current_max_allowed_perp_dist and \
|
||
min(dist_a1_to_b_inf, dist_a2_to_b_inf) > current_max_allowed_perp_dist:
|
||
return False
|
||
|
||
# 3. 邻近性/重叠性: 端点之间的最小距离
|
||
min_ep_dist = min(np.linalg.norm(p1a-p1b), np.linalg.norm(p1a-p2b),
|
||
np.linalg.norm(p2a-p1b), np.linalg.norm(p2a-p2b))
|
||
|
||
current_max_allowed_ep_gap = max(max_ep_gap_abs, max_ep_gap_factor * shorter_len)
|
||
if min_ep_dist > current_max_allowed_ep_gap:
|
||
return False
|
||
|
||
# 附加检查：确保线段确实是"连接的"或"重叠的"
|
||
# 计算一个端点到另一个 *线段* (非无限线) 的距离
|
||
def dist_point_to_segment(p, seg_a, seg_b):
|
||
ap = p - seg_a
|
||
ab = seg_b - seg_a
|
||
ab_len_sq = np.dot(ab, ab)
|
||
if ab_len_sq < 1e-6: return np.linalg.norm(ap)
|
||
|
||
t = np.dot(ap, ab) / ab_len_sq
|
||
if t < 0.0: closest_point = seg_a
|
||
elif t > 1.0: closest_point = seg_b
|
||
else: closest_point = seg_a + t * ab
|
||
return np.linalg.norm(p - closest_point)
|
||
|
||
dist_p1a_to_seg_b = dist_point_to_segment(p1a, p1b, p2b)
|
||
dist_p2a_to_seg_b = dist_point_to_segment(p2a, p1b, p2b)
|
||
dist_p1b_to_seg_a = dist_point_to_segment(p1b, p1a, p2a)
|
||
dist_p2b_to_seg_a = dist_point_to_segment(p2b, p1a, p2a)
|
||
|
||
if min(dist_p1a_to_seg_b, dist_p2a_to_seg_b, dist_p1b_to_seg_a, dist_p2b_to_seg_a) > current_max_allowed_ep_gap:
|
||
return False
|
||
|
||
return True
|
||
|
||
def _merge_two_lines(line1_pts, line2_pts):
|
||
points = np.array([
|
||
line1_pts[:2], line1_pts[2:],
|
||
line2_pts[:2], line2_pts[2:]
|
||
], dtype=np.float32)
|
||
|
||
vx, vy, x0, y0 = cv2.fitLine(points.reshape(-1,1,2), cv2.DIST_L2, 0, 0.01, 0.01)
|
||
vx, vy, x0, y0 = vx[0], vy[0], x0[0], y0[0]
|
||
|
||
projected_params = []
|
||
line_origin = np.array([x0, y0])
|
||
line_direction = np.array([vx, vy])
|
||
|
||
for i in range(4):
|
||
pt_vec = points[i] - line_origin
|
||
param = np.dot(pt_vec, line_direction)
|
||
projected_params.append(param)
|
||
|
||
min_param = min(projected_params)
|
||
max_param = max(projected_params)
|
||
|
||
p_start = line_origin + min_param * line_direction
|
||
p_end = line_origin + max_param * line_direction
|
||
|
||
return [int(round(p_start[0])), int(round(p_start[1])),
|
||
int(round(p_end[0])), int(round(p_end[1]))]
|
||
|
||
def _merge_collinear_lines_iterative(lines,
|
||
min_initial_len,
|
||
max_angle_diff_deg,
|
||
max_ep_gap_abs, max_ep_gap_factor,
|
||
max_p_dist_abs, max_p_dist_factor):
|
||
"""
|
||
"""
|
||
if lines is None:
|
||
return None
|
||
|
||
# 过滤掉初始太短的线
|
||
current_lines = [line for line in lines if np.linalg.norm(np.array(line[:2]) - np.array(line[2:])) >= min_initial_len]
|
||
|
||
if not current_lines or len(current_lines) < 2:
|
||
return np.array([[line] for line in current_lines]) if current_lines else np.array([])
|
||
|
||
made_change_in_pass = True
|
||
while made_change_in_pass:
|
||
made_change_in_pass = False
|
||
i = 0
|
||
while i < len(current_lines):
|
||
j = i + 1
|
||
merged_line_i = False
|
||
while j < len(current_lines):
|
||
if _are_lines_mergeable(current_lines[i], current_lines[j],
|
||
max_angle_diff_deg,
|
||
max_ep_gap_abs, max_ep_gap_factor,
|
||
max_p_dist_abs, max_p_dist_factor):
|
||
|
||
merged_line = _merge_two_lines(current_lines[i], current_lines[j])
|
||
current_lines[i] = merged_line # 更新第i条线
|
||
current_lines.pop(j) # 移除第j条线
|
||
made_change_in_pass = True
|
||
merged_line_i = True
|
||
# 由于current_lines[i]已改变，且列表长度已改变，
|
||
# 最安全的是从头开始或至少重新评估current_lines[i]
|
||
# 但对于这种贪心策略，继续尝试将（新的）current_lines[i]与后续线合并是可以的
|
||
# j的位置不需要改变，因为它指向被删除元素后的下一个元素
|
||
continue # j 现在指向下一条未检查的线，相对于更新后的i
|
||
j += 1
|
||
|
||
# 如果 line_i 在内部循环中被修改过，则外部的 i 也应该重新评估
|
||
# 但由于我们是迭代地进行整个过程(while made_change_in_pass), 最终会收敛
|
||
i += 1
|
||
|
||
return np.array([[line] for line in current_lines]) if current_lines else np.array([])
|