image_to_pixle_params_yoloSAM/main/utils.py

101 lines
3.1 KiB
Python
Raw Normal View History

2025-07-14 17:36:53 +08:00
import cv2
import numpy as np
def get_bounding_box(image_input):
"""
从阈值化图像中提取最小外接矩形返回四个角点坐标顺时针
参数
----
image_input : 图像文件路径或图像数组
"""
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()
coords = np.argwhere(image > 0)
if coords.size == 0:
raise ValueError("图像中没有非零像素,无法计算外接矩形")
ymin, xmin = coords.min(axis=0)
ymax, xmax = coords.max(axis=0)
# 返回四个角点(顺时针)
return (xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)
def get_line_equation(p1, p2):
"""
给定两点 (p1, p2)返回直线方程
- 垂直线: (None, x_const, y1, y2)
- 水平线: (0, y_const, x1, x2)
- 斜率线: (a, b, x1, x2) 对应 y = ax + b范围 [x1, x2]
"""
if p2[0] == p1[0]: # 垂直线
return None, p1[0], p1[1], p2[1]
elif p2[1] == p1[1]: # 水平线
return 0, p1[1], p1[0], p2[0]
else:
a = (p2[1] - p1[1]) / (p2[0] - p1[0])
b = p1[1] - a * p1[0]
return a, b, p1[0], p2[0]
def calculate_lines(top_left, top_right, bottom_right, bottom_left):
"""
返回外接矩形四条边的线性表达式
"""
top_line = get_line_equation(top_left, top_right)
right_line = get_line_equation(top_right, bottom_right)
bottom_line = get_line_equation(bottom_right, bottom_left)
left_line = get_line_equation(bottom_left, top_left)
return top_line, right_line, bottom_line, left_line
def apply_morphology(gray,
kernel_size=(5, 5),
erode_iter=3,
dilate_iter=3):
"""
对灰度图执行"腐蚀 → 膨胀"去噪声并返回处理后的结果
参数
----
gray : np.ndarray灰度图 (H, W)
kernel_size : structuring element 尺寸
erode_iter : 腐蚀迭代次数
dilate_iter : 膨胀迭代次数
"""
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, kernel_size)
out = cv2.erode(gray, kernel, iterations=erode_iter)
out = cv2.dilate(out, kernel, iterations=dilate_iter)
return out
def perpendicular_foot(pt, line_info):
"""
计算点 pt 到指定直线 line_info 的垂足坐标
pt : (x0, y0)
line_info : get_line_equation() 的返回值
"""
x0, y0 = pt
if line_info[0] is None: # 垂直线 x = const
return int(line_info[1]), int(y0)
if line_info[0] == 0: # 水平线 y = const
return int(x0), int(line_info[1])
a, b, *_ = line_info # 斜率线 y = ax + b
xf = (x0 + a * (y0 - b)) / (1 + a**2)
yf = a * xf + b
return int(round(xf)), int(round(yf))