611 lines
23 KiB
Python
611 lines
23 KiB
Python
import cv2
|
||
import numpy as np
|
||
import os
|
||
# main.py
|
||
# ───────────── 导入 utils 内的工具函数 ─────────────
|
||
from utils import *
|
||
|
||
# 之后就可以直接调用，例如：
|
||
# tl, tr, br, bl = get_bounding_box(mask_path)
|
||
# top_line, right_line, bottom_line, left_line = calculate_lines(tl, tr, br, bl)
|
||
|
||
def resize_image(image, target_size=(1000, 600)):
|
||
"""
|
||
将图像缩放到指定尺寸
|
||
"""
|
||
return cv2.resize(image, target_size, interpolation=cv2.INTER_LINEAR)
|
||
|
||
def save_eroded(mask_color, out_path, kernel_size=(5, 5), erode_iter=10, dilate_iter=10):
|
||
"""
|
||
保存腐蚀10次再膨胀10次后的掩膜图像，便于可视化
|
||
返回处理后的图像
|
||
"""
|
||
gray = cv2.cvtColor(mask_color, cv2.COLOR_BGR2GRAY)
|
||
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, kernel_size)
|
||
eroded = cv2.erode(gray, kernel, iterations=erode_iter)
|
||
morphed = cv2.dilate(eroded, kernel, iterations=dilate_iter)
|
||
morphed_bgr = cv2.cvtColor(morphed, cv2.COLOR_GRAY2BGR)
|
||
cv2.imwrite(out_path, morphed_bgr)
|
||
return morphed_bgr
|
||
|
||
def detect_circles(image_input, bounding_box=None, min_radius=40, max_radius=100, min_dist=20, param1=50, param2=17):
|
||
"""
|
||
对给定的图像进行霍夫圆检测，返回圆心坐标和半径
|
||
|
||
参数
|
||
----
|
||
image_input : 图像文件路径或图像数组
|
||
bounding_box: 可选，(top_left, top_right, bottom_right, bottom_left) 形式的矩形框
|
||
"""
|
||
# 读取图像并转换为灰度图
|
||
if isinstance(image_input, str):
|
||
# 如果是文件路径
|
||
image = cv2.imread(image_input, cv2.IMREAD_GRAYSCALE)
|
||
if image is None:
|
||
raise FileNotFoundError(f"图像文件 {image_input} 未找到")
|
||
else:
|
||
# 如果是图像数组
|
||
image = cv2.cvtColor(image_input, cv2.COLOR_BGR2GRAY) if image_input.ndim == 3 else image_input.copy()
|
||
|
||
# 根据图像宽度动态调整max_radius
|
||
h, w = image.shape[:2]
|
||
if max(h, w) < 700:
|
||
max_radius = 60
|
||
min_radius = 30
|
||
|
||
# 如果提供了矩形框，则只在框内（去除左右10%）的区域检测
|
||
if bounding_box:
|
||
top_left, top_right, _, _ = bounding_box
|
||
x_min = top_left[0]
|
||
x_max = top_right[0]
|
||
box_width = x_max - x_min
|
||
|
||
# 计算要保留的区域
|
||
start_x = int(x_min + box_width * 0.1)
|
||
end_x = int(x_max - box_width * 0.1)
|
||
|
||
# 创建一个全黑的掩膜
|
||
mask = np.zeros_like(image)
|
||
# 将保留区域设为白色
|
||
mask[:, start_x:end_x] = 255
|
||
# 将原图与掩膜相乘，只保留目标区域
|
||
image = cv2.bitwise_and(image, mask)
|
||
|
||
else:
|
||
# 兼容旧逻辑：屏蔽整张图的左右10%
|
||
w = image.shape[1]
|
||
image[:, :int(w * 0.05)] = 0
|
||
image[:, int(w * 0.9):] = 0
|
||
|
||
|
||
# 图像模糊处理，减少噪声
|
||
blurred = cv2.GaussianBlur(image, (9, 9), 2)
|
||
|
||
# 霍夫圆变换检测圆
|
||
circles = cv2.HoughCircles(
|
||
blurred, cv2.HOUGH_GRADIENT, dp=1.2, minDist=min_dist,
|
||
param1=param1, param2=param2, minRadius=min_radius, maxRadius=max_radius
|
||
)
|
||
|
||
# 如果检测到圆
|
||
if circles is not None:
|
||
circles = np.uint16(np.around(circles))
|
||
centers = [(c[0], c[1], c[2]) for c in circles[0]]
|
||
# 检查前两个圆心距离
|
||
if len(centers) >= 2:
|
||
x1, y1, _ = centers[0]
|
||
x2, y2, _ = centers[1]
|
||
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
|
||
dist = abs(x1 - x2) + abs(y1 - y2)
|
||
if dist <= 200:
|
||
# 增大max_radius并重新检测
|
||
new_max_radius = max_radius + 20
|
||
circles2 = cv2.HoughCircles(
|
||
blurred, cv2.HOUGH_GRADIENT, dp=1.2, minDist=min_dist,
|
||
param1=param1, param2=param2, minRadius=min_radius, maxRadius=new_max_radius
|
||
)
|
||
if circles2 is not None:
|
||
circles2 = np.uint16(np.around(circles2))
|
||
centers2 = [(c[0], c[1], c[2]) for c in circles2[0]]
|
||
return centers2
|
||
return centers # 返回 [(x1, y1, r1), (x2, y2, r2), ...]
|
||
else:
|
||
print("[Circle] 未检测到圆")
|
||
return [] # 返回空列表表示未检测到圆
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
def draw_line_equation(image, line_info, color=(0, 255, 0), thickness=2):
|
||
"""
|
||
根据线性表达式绘制线段（而非无限直线）
|
||
"""
|
||
if line_info[0] is None:
|
||
# 垂直线 x = const
|
||
_, x, y1, y2 = line_info
|
||
pt1 = (int(x), int(min(y1, y2)))
|
||
pt2 = (int(x), int(max(y1, y2)))
|
||
elif line_info[0] == 0:
|
||
# 水平线 y = const
|
||
_, y, x1, x2 = line_info
|
||
pt1 = (int(min(x1, x2)), int(y))
|
||
pt2 = (int(max(x1, x2)), int(y))
|
||
else:
|
||
# 斜率线 y = ax + b
|
||
a, b, x1, x2 = line_info
|
||
x1, x2 = int(x1), int(x2)
|
||
pt1 = (x1, int(a * x1 + b))
|
||
pt2 = (x2, int(a * x2 + b))
|
||
|
||
cv2.line(image, pt1, pt2, color, thickness)
|
||
|
||
|
||
def draw_circles(image, circles, color=(0, 0, 255), thickness=2):
|
||
"""
|
||
根据圆心坐标和半径在图像上绘制圆心和圆，并保存图像
|
||
"""
|
||
for (x, y, r) in circles:
|
||
# 画圆边
|
||
# cv2.circle(image, (x, y), r, (0, 255, 0), 2)
|
||
# 画圆心
|
||
cv2.circle(image, (x, y), 3, (0, 0, 255), -1)
|
||
|
||
|
||
|
||
# ------------------------- 找 A,B 点 -------------------------
|
||
def find_A_B(circles):
|
||
"""
|
||
circles : [(x, y, r), ...] 来自 detect_circles()
|
||
返回 (A, B) —— 均为 (x, y) 或 None
|
||
规则：检测到的第 1 个圆心为 A，第 2 个为 B
|
||
"""
|
||
if len(circles) == 0:
|
||
return None, None
|
||
elif len(circles) == 1:
|
||
return (circles[0][0], circles[0][1]), None
|
||
else:
|
||
A = (circles[0][0], circles[0][1])
|
||
B = (circles[1][0], circles[1][1])
|
||
return A, B
|
||
|
||
|
||
# ------------------------- 找 C 点 -------------------------
|
||
def find_C(image, line_info,
|
||
thresh_val=250,
|
||
do_morph=False,
|
||
kernel_size=(5, 5),
|
||
iterations=3,
|
||
adj = -5):
|
||
"""
|
||
C 点：沿给定直线自"下 → 上"遇到的第一个白色像素
|
||
"""
|
||
# 灰度化
|
||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if image.ndim == 3 else image.copy()
|
||
|
||
# 可选形态学
|
||
if do_morph:
|
||
gray = apply_morphology(gray,
|
||
kernel_size=kernel_size,
|
||
erode_iter=iterations,
|
||
dilate_iter=iterations)
|
||
|
||
# 垂直线 x = const
|
||
if line_info[0] is None:
|
||
_, x_const, y1, y2 = line_info
|
||
x_const = int(x_const)
|
||
for y in range(int(max(y1, y2)), int(min(y1, y2)) - 1, -1): # ⬅️ 下→上
|
||
if gray[y, x_const] >= thresh_val:
|
||
return x_const, y + adj
|
||
|
||
# 水平线 y = const（极少需要）
|
||
elif line_info[0] == 0:
|
||
_, y_const, x1, x2 = line_info
|
||
y_const = int(y_const)
|
||
for y in range(gray.shape[0] - 1, -1, -1): # ⬅️ 底→顶
|
||
if gray[y, int(x1)] >= thresh_val:
|
||
return int(x1), y + adj
|
||
|
||
# 斜线 y = ax + b
|
||
else:
|
||
a, b, x1, x2 = line_info
|
||
xs = np.linspace(x1, x2, num=abs(int(x2 - x1)) + 1, dtype=int)
|
||
ys = (a * xs + b).astype(int)
|
||
for x_i, y_i in sorted(zip(xs, ys), key=lambda p: -p[1]): # ⬅️ y 降序
|
||
if 0 <= x_i < gray.shape[1] and 0 <= y_i < gray.shape[0]:
|
||
if gray[y_i, x_i] >= thresh_val:
|
||
return x_i, y_i + adj
|
||
return None
|
||
|
||
|
||
# ------------------------- 找 D 点 -------------------------
|
||
def find_D(image, line_info,
|
||
thresh_val=250,
|
||
do_morph=False,
|
||
kernel_size=(3, 3),
|
||
iterations=9,
|
||
adj = - 5):
|
||
"""
|
||
最右侧直线（或给定 line_info）与图像区域的第一个白色像素交点
|
||
—— 采用"自下而上"扫描策略。
|
||
|
||
参数
|
||
----
|
||
image : BGR/灰度图
|
||
line_info : 与 get_line_equation 返回格式一致
|
||
thresh_val : 像素灰度值 ≥ thresh_val 判定为白
|
||
do_morph : 是否做形态学去噪
|
||
kernel_size : 形态学核尺寸
|
||
iterations : 腐蚀/膨胀迭代次数
|
||
"""
|
||
# 1) 灰度化
|
||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if image.ndim == 3 else image.copy()
|
||
|
||
|
||
# 2) 形态学（可选）
|
||
if do_morph:
|
||
gray = apply_morphology(gray,
|
||
kernel_size=kernel_size,
|
||
erode_iter=iterations,
|
||
dilate_iter=iterations+1)
|
||
|
||
# 3) 自下而上扫描
|
||
if line_info[0] is None: # 垂直线 x = const
|
||
_, x_const, y1, y2 = line_info
|
||
x_const = int(x_const)
|
||
for y in range(int(max(y1, y2)), int(min(y1, y2)) - 1, -1):
|
||
if gray[y, x_const] >= thresh_val:
|
||
return x_const, y + adj
|
||
|
||
elif line_info[0] == 0: # 水平线 y = const（少见）
|
||
_, y_const, x1, x2 = line_info
|
||
y_const = int(y_const)
|
||
for y in range(gray.shape[0] - 1, -1, -1):
|
||
if gray[y, int(x1)] >= thresh_val:
|
||
return int(x1), y + adj
|
||
|
||
else: # 斜率线 y = a x + b
|
||
a, b, x1, x2 = line_info
|
||
xs = np.linspace(x1, x2, num=abs(int(x2 - x1)) + 1, dtype=int)
|
||
ys = (a * xs + b).astype(int)
|
||
for x_i, y_i in sorted(zip(xs, ys), key=lambda p: -p[1]): # y 由大到小
|
||
if 0 <= x_i < gray.shape[1] and 0 <= y_i < gray.shape[0]:
|
||
if gray[y_i, x_i] >= thresh_val:
|
||
return x_i, y_i + adj
|
||
return None
|
||
|
||
|
||
def find_EFH(A, B, G, bottom_line):
|
||
"""
|
||
参数
|
||
----
|
||
A, B, G : 三个已知点坐标 (x, y)，或 None
|
||
bottom_line : 外接矩形底边的 line_info（来自 calculate_lines）
|
||
返回
|
||
----
|
||
(E, F, H) : 三个垂足坐标，若对应输入为 None 则返回 None
|
||
"""
|
||
E = perpendicular_foot(A, bottom_line) if A else None
|
||
F = perpendicular_foot(B, bottom_line) if B else None
|
||
H = perpendicular_foot(G, bottom_line) if G else None
|
||
return E, F, H
|
||
# ------------------------- 找 G 点 -------------------------
|
||
def find_G(image, line_info,
|
||
kernel_size=(3, 3),
|
||
erode_iter=3,
|
||
dilate_iter=3,
|
||
adj=3):
|
||
"""
|
||
侧视图汽车最高点 G：沿直线从左到右（或上到下）找到遇到的第一个白色像素。
|
||
"""
|
||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||
|
||
# 形态学处理（固定参数版本）
|
||
# gray = apply_morphology(gray,
|
||
# kernel_size=kernel_size,
|
||
# erode_iter=erode_iter,
|
||
# dilate_iter=dilate_iter)
|
||
|
||
# 【调试】保存形态学处理后的图像，检查车顶是否被腐蚀
|
||
cv2.imwrite("debug_G_morphed_mask.png", gray)
|
||
|
||
|
||
if line_info[0] is None: # 垂直线 x = const
|
||
_, x, y1, y2 = line_info
|
||
for y in range(int(min(y1, y2)), int(max(y1, y2))):
|
||
if gray[y, int(x)] >= 150:
|
||
return int(x), y + adj
|
||
|
||
elif line_info[0] == 0: # 水平线 y = const
|
||
_, y, x1, x2 = line_info
|
||
for x in range(int(min(x1, x2)), int(max(x1, x2))):
|
||
if gray[int(y), x] >= 150:
|
||
return x, int(y) + adj
|
||
|
||
else: # 斜率线 y = ax + b
|
||
a, b, x1, x2 = line_info
|
||
for x in range(int(x1), int(x2) + 1):
|
||
y = int(a * x + b)
|
||
if 0 <= y < gray.shape[0] and gray[y, x] >= 150:
|
||
return x, y + adj
|
||
return None
|
||
|
||
|
||
def find_bottom_gap_midpoint(mask_color, bottom_line, gap=20, direction='left'):
|
||
"""
|
||
遍历底边采样点，往上找gap距离内的白色点，
|
||
第一个dy<=gap的点为左端点，第一个dy>gap的点的前一个为右端点，
|
||
L/M点为区间中心。
|
||
"""
|
||
gray = cv2.cvtColor(mask_color, cv2.COLOR_BGR2GRAY)
|
||
h, w = gray.shape
|
||
|
||
# 采样底边线段上的点
|
||
if bottom_line[0] is None: # 垂直线
|
||
_, x_const, y1, y2 = bottom_line
|
||
ys = np.linspace(y1, y2, abs(int(y2 - y1)) + 1, dtype=int)
|
||
xs = np.full_like(ys, int(x_const))
|
||
elif bottom_line[0] == 0: # 水平线
|
||
_, y_const, x1, x2 = bottom_line
|
||
xs = np.linspace(x1, x2, abs(int(x2 - x1)) + 1, dtype=int)
|
||
ys = np.full_like(xs, int(y_const))
|
||
else: # 斜率线
|
||
a, b, x1, x2 = bottom_line
|
||
xs = np.linspace(x1, x2, abs(int(x2 - x1)) + 1, dtype=int)
|
||
ys = (a * xs + b).astype(int)
|
||
|
||
# 按方向决定遍历顺序
|
||
if direction == 'left':
|
||
idx_range = range(len(xs))
|
||
else:
|
||
idx_range = range(len(xs) - 1, -1, -1)
|
||
|
||
dy_list = [] # 存储(x, y, dy)
|
||
for i in idx_range:
|
||
x, y = xs[i], ys[i]
|
||
dy = None
|
||
for d in range(1, gap + 2):
|
||
yy = y - d
|
||
if 0 <= x < w and 0 <= yy < h and gray[yy, x] >= 200:
|
||
dy = d
|
||
break
|
||
dy_list.append((x, y, dy))
|
||
|
||
# 找左端点（第一个dy<=gap的点）
|
||
left_idx = None
|
||
for i, (x, y, dy) in enumerate(dy_list):
|
||
if dy is not None and dy <= gap:
|
||
left_idx = i
|
||
break
|
||
if left_idx is None:
|
||
return None
|
||
# 找右端点（第一个dy>gap的点的前一个）
|
||
right_idx = None
|
||
for i in range(left_idx, len(dy_list)):
|
||
dy = dy_list[i][2]
|
||
if dy is None or dy > gap:
|
||
right_idx = i - 1
|
||
break
|
||
if right_idx is None:
|
||
right_idx = len(dy_list) - 1
|
||
# 区间中心
|
||
mid_idx = (left_idx + right_idx) // 2
|
||
print(f"[find_bottom_gap_midpoint] direction={direction}, 左端点=({dy_list[left_idx][0]}, {dy_list[left_idx][1]}), 右端点=({dy_list[right_idx][0]}, {dy_list[right_idx][1]})")
|
||
return (int(dy_list[mid_idx][0]), int(dy_list[mid_idx][1]))
|
||
|
||
def preprocess_mask(mask_color, kernel_size=(3, 3), erode_iter=2, dilate_iter=2):
|
||
gray = cv2.cvtColor(mask_color, cv2.COLOR_BGR2GRAY)
|
||
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, kernel_size)
|
||
eroded = cv2.erode(gray, kernel, iterations=erode_iter)
|
||
morphed = cv2.dilate(eroded, kernel, iterations=dilate_iter)
|
||
return cv2.cvtColor(morphed, cv2.COLOR_GRAY2BGR)
|
||
|
||
|
||
def process_side(mask_path, rgb_path, out_dir):
|
||
"""
|
||
处理侧视图，检测并标注 A-H 点
|
||
"""
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
mask_color = cv2.imread(mask_path, cv2.IMREAD_COLOR)
|
||
rgb_color = cv2.imread(rgb_path, cv2.IMREAD_COLOR)
|
||
if mask_color is None or rgb_color is None:
|
||
raise FileNotFoundError("无法读取 mask 或 RGB 图像")
|
||
|
||
# 先对mask做轻微腐蚀膨胀预处理
|
||
mask_color = preprocess_mask(mask_color, kernel_size=(3, 3), erode_iter=2, dilate_iter=2)
|
||
|
||
# 1. 先获取外接矩形
|
||
top_left, top_right, bottom_right, bottom_left = get_bounding_box(mask_color)
|
||
|
||
# 2. 将矩形框传入霍夫圆检测，只在框内特定区域检测
|
||
bounding_box_coords = (top_left, top_right, bottom_right, bottom_left)
|
||
circles = detect_circles(mask_color, bounding_box=bounding_box_coords)
|
||
|
||
# 3. 计算边线
|
||
top_line, right_line, bottom_line, left_line = calculate_lines(
|
||
top_left, top_right, bottom_right, bottom_left
|
||
)
|
||
A, B = find_A_B(circles)
|
||
C = find_C(mask_color, left_line)
|
||
D = find_D(mask_color, right_line)
|
||
G = find_G(mask_color, top_line)
|
||
E, F, H = find_EFH(A, B, G, bottom_line)
|
||
|
||
points = {'A': A, 'B': B, 'C': C, 'D': D, 'E': E, 'F': F, 'G': G, 'H': H}
|
||
# 用原始mask_color和rgb_color做可视化
|
||
for canvas in (mask_color, rgb_color):
|
||
draw_line_equation(canvas, top_line)
|
||
draw_line_equation(canvas, right_line)
|
||
draw_line_equation(canvas, bottom_line)
|
||
draw_line_equation(canvas, left_line)
|
||
for label, pt in points.items():
|
||
if pt is not None:
|
||
cv2.circle(canvas, pt, 3, (0, 0, 255), -1)
|
||
if label == 'H':
|
||
cv2.putText(canvas, label, (pt[0]+5, pt[1]+20),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 1, cv2.LINE_AA)
|
||
else:
|
||
cv2.putText(canvas, label, (pt[0]+5, pt[1]-5),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 1, cv2.LINE_AA)
|
||
cv2.imwrite(os.path.join(out_dir, 'side_output_mask.png'), mask_color)
|
||
cv2.imwrite(os.path.join(out_dir, 'side_output_rgb.png'), rgb_color)
|
||
print("A =", A, " B =", B,"C =", C, " D =", D," G =", G , "E =", E ,"F =", F ,"H =", H )
|
||
print(f"[side] 结果已保存到 {out_dir}")
|
||
|
||
def process_rear(mask_path, rgb_path, out_dir):
|
||
"""
|
||
处理后视图，检测并标注 P, N, O, Q, R 点
|
||
"""
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
mask_color = cv2.imread(mask_path, cv2.IMREAD_COLOR)
|
||
rgb_color = cv2.imread(rgb_path, cv2.IMREAD_COLOR)
|
||
if mask_color is None or rgb_color is None:
|
||
raise FileNotFoundError("无法读取 mask 或 RGB 图像")
|
||
|
||
# 先验知识的外接矩形（与正视图一致）
|
||
top_left, top_right, bottom_right, bottom_left = get_front_bounding_box_with_prior(mask_color)
|
||
top_line, right_line, bottom_line, left_line = calculate_lines(
|
||
top_left, top_right, bottom_right, bottom_left
|
||
)
|
||
|
||
# 保存绘制外接矩形的图片
|
||
mask_with_rect = mask_color.copy()
|
||
draw_line_equation(mask_with_rect, top_line)
|
||
draw_line_equation(mask_with_rect, right_line)
|
||
draw_line_equation(mask_with_rect, bottom_line)
|
||
draw_line_equation(mask_with_rect, left_line)
|
||
cv2.imwrite(os.path.join(out_dir, 'rear_with_rect.png'), mask_with_rect)
|
||
|
||
# P: 上边线段中点
|
||
P = ((top_left[0] + top_right[0]) // 2, (top_left[1] + top_right[1]) // 2)
|
||
# N: 左边线段的中点
|
||
N = ((top_left[0] + bottom_left[0]) // 2, (top_left[1] + bottom_left[1]) // 2)
|
||
# O: 右边线段的中点
|
||
O = ((top_right[0] + bottom_right[0]) // 2, (top_right[1] + bottom_right[1]) // 2)
|
||
|
||
R = find_bottom_gap_midpoint(mask_color, bottom_line, gap=8, direction='left')
|
||
Q = find_bottom_gap_midpoint(mask_color, bottom_line, gap=8, direction='right')
|
||
|
||
points = {'P': P, 'N': N, 'O': O, 'Q': Q, 'R': R}
|
||
for canvas in (mask_color, rgb_color):
|
||
draw_line_equation(canvas, top_line)
|
||
draw_line_equation(canvas, right_line)
|
||
draw_line_equation(canvas, bottom_line)
|
||
draw_line_equation(canvas, left_line)
|
||
for label, pt in points.items():
|
||
if pt is not None:
|
||
cv2.circle(canvas, pt, 3, (0, 0, 255), -1)
|
||
cv2.putText(canvas, label, (pt[0]+5, pt[1]-5),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 1, cv2.LINE_AA)
|
||
cv2.imwrite(os.path.join(out_dir, 'rear_output_mask.png'), mask_color)
|
||
cv2.imwrite(os.path.join(out_dir, 'rear_output_rgb.png'), rgb_color)
|
||
print("P =", P, "N =", N, "O =", O, "Q =", Q, "R =", R)
|
||
print(f"[rear] 结果已保存到 {out_dir}")
|
||
|
||
def get_front_bounding_box_with_prior(mask):
|
||
"""
|
||
先将mask上1/2区域设为黑色，仅用下1/3区域的白色像素计算左右边界，上下边界用原始mask。
|
||
返回 (top_left, top_right, bottom_right, bottom_left)
|
||
"""
|
||
h, w = mask.shape[:2]
|
||
# 1. 生成处理后的mask（上2/3黑，下1/3保留）
|
||
mask_proc = mask.copy()
|
||
y_cut = h * 1 // 2
|
||
mask_proc[:y_cut, :] = 0
|
||
# 2. 计算左右边界（下1/3区域）
|
||
# 只考虑白色像素
|
||
gray_proc = cv2.cvtColor(mask_proc, cv2.COLOR_BGR2GRAY) if mask_proc.ndim == 3 else mask_proc
|
||
coords = cv2.findNonZero((gray_proc > 127).astype(np.uint8))
|
||
if coords is None:
|
||
raise ValueError("处理后mask没有白色区域，无法计算左右边界")
|
||
xs = coords[:, 0, 0]
|
||
ys = coords[:, 0, 1]
|
||
left_x = np.min(xs)
|
||
right_x = np.max(xs)
|
||
# 3. 计算上下边界（原始mask）
|
||
gray = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY) if mask.ndim == 3 else mask
|
||
coords_full = cv2.findNonZero((gray > 127).astype(np.uint8))
|
||
if coords_full is None:
|
||
raise ValueError("原始mask没有白色区域，无法计算上下边界")
|
||
top_y = np.min(coords_full[:, 0, 1])
|
||
bottom_y = np.max(coords_full[:, 0, 1])
|
||
# 4. 返回四个点
|
||
top_left = (left_x, top_y)
|
||
top_right = (right_x, top_y)
|
||
bottom_right = (right_x, bottom_y)
|
||
bottom_left = (left_x, bottom_y)
|
||
return top_left, top_right, bottom_right, bottom_left
|
||
|
||
def process_front(mask_path, rgb_path, out_dir):
|
||
"""
|
||
处理前视图，检测并标注 K, I, J, L, M 点
|
||
"""
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
mask_color = cv2.imread(mask_path, cv2.IMREAD_COLOR)
|
||
rgb_color = cv2.imread(rgb_path, cv2.IMREAD_COLOR)
|
||
if mask_color is None or rgb_color is None:
|
||
raise FileNotFoundError("无法读取 mask 或 RGB 图像")
|
||
|
||
# 先进行腐蚀膨胀处理
|
||
#mask_processed = save_eroded(mask_color, os.path.join(out_dir, 'front_eroded_mask.png'),
|
||
# kernel_size=(5, 5), erode_iter=8, dilate_iter=5)
|
||
mask_processed = mask_color
|
||
# 使用先验知识的外接矩形
|
||
top_left, top_right, bottom_right, bottom_left = get_front_bounding_box_with_prior(mask_processed)
|
||
top_line, right_line, bottom_line, left_line = calculate_lines(
|
||
top_left, top_right, bottom_right, bottom_left
|
||
)
|
||
|
||
# 保存绘制外接矩形的图片
|
||
mask_with_rect = mask_processed.copy()
|
||
draw_line_equation(mask_with_rect, top_line)
|
||
draw_line_equation(mask_with_rect, right_line)
|
||
draw_line_equation(mask_with_rect, bottom_line)
|
||
draw_line_equation(mask_with_rect, left_line)
|
||
cv2.imwrite(os.path.join(out_dir, 'front_with_rect.png'), mask_with_rect)
|
||
|
||
# K: 上边线段中点
|
||
K = ((top_left[0] + top_right[0]) // 2, (top_left[1] + top_right[1]) // 2)
|
||
# I: 左边线段的高度一半处
|
||
I = (left_line[1], (left_line[2] + left_line[3]) // 2) if left_line[0] is None else None
|
||
# J: 右边线段的高度一半处
|
||
J = (right_line[1], (right_line[2] + right_line[3]) // 2) if right_line[0] is None else None
|
||
M = find_bottom_gap_midpoint(mask_color, bottom_line, gap=8, direction='left')
|
||
L = find_bottom_gap_midpoint(mask_color, bottom_line, gap=8, direction='right')
|
||
|
||
points = {'K': K, 'I': I, 'J': J, 'L': L, 'M': M}
|
||
for canvas in (mask_color, rgb_color):
|
||
draw_line_equation(canvas, top_line)
|
||
draw_line_equation(canvas, right_line)
|
||
draw_line_equation(canvas, bottom_line)
|
||
draw_line_equation(canvas, left_line)
|
||
for label, pt in points.items():
|
||
if pt is not None:
|
||
cv2.circle(canvas, pt, 3, (0, 0, 255), -1)
|
||
cv2.putText(canvas, label, (pt[0]+5, pt[1]-5),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 1, cv2.LINE_AA)
|
||
cv2.imwrite(os.path.join(out_dir, 'front_output_mask.png'), mask_color)
|
||
cv2.imwrite(os.path.join(out_dir, 'front_output_rgb.png'), rgb_color)
|
||
print("K =", K, "I =", I, "J =", J, "L =", L, "M =", M)
|
||
print(f"[front] 结果已保存到 {out_dir}")
|
||
|
||
if __name__ == '__main__':
|
||
side_mask = '../segment-anything-main/output/5_mask.png'
|
||
side_rgb = '../ultralytics-main/input/5.jpg'
|
||
|
||
front_mask = '../segment-anything-main/output/00213_mask.png'
|
||
front_rgb = '../ultralytics-main/input/00213.jpg'
|
||
|
||
rear_mask = '../segment-anything-main/output/3_mask.png'
|
||
rear_rgb = '../ultralytics-main/input/3.jpg'
|
||
|
||
out_dir = './result'
|
||
|
||
process_side(side_mask, side_rgb, out_dir)
|
||
# process_front(front_mask, front_rgb, out_dir)
|
||
# process_rear(rear_mask, rear_rgb, out_dir)
|