424 lines
14 KiB
Python
424 lines
14 KiB
Python
import cv2
|
||
import numpy as np
|
||
|
||
def detect_arrow_direction(image, observe=False, delay=500):
|
||
"""
|
||
从图像中提取绿色箭头并判断其指向方向（左或右）
|
||
|
||
参数:
|
||
image: 输入图像，可以是文件路径或者已加载的图像数组
|
||
observe: 是否输出中间状态信息和可视化结果，默认为False
|
||
delay: 展示每个步骤的等待时间(毫秒)，默认为500ms
|
||
|
||
返回:
|
||
direction: 字符串，"left"表示左箭头，"right"表示右箭头，"unknown"表示无法确定
|
||
"""
|
||
# 如果输入是字符串（文件路径），则加载图像
|
||
if isinstance(image, str):
|
||
img = cv2.imread(image)
|
||
else:
|
||
img = image.copy()
|
||
|
||
if img is None:
|
||
print("无法加载图像")
|
||
return "unknown"
|
||
|
||
if observe:
|
||
print("步骤1: 原始图像已加载")
|
||
cv2.imshow("原始图像", img)
|
||
cv2.waitKey(delay)
|
||
|
||
# 转换到HSV颜色空间以便更容易提取绿色
|
||
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
|
||
|
||
if observe:
|
||
print("步骤2: 转换到HSV颜色空间")
|
||
cv2.imshow("HSV图像", hsv)
|
||
cv2.waitKey(delay)
|
||
|
||
# 绿色的HSV范围
|
||
# 调整这些值以匹配图像中绿色的具体色调··
|
||
lower_green = np.array([40, 50, 50])
|
||
upper_green = np.array([80, 255, 255])
|
||
|
||
# 创建绿色的掩码
|
||
mask = cv2.inRange(hsv, lower_green, upper_green)
|
||
|
||
if observe:
|
||
print("步骤3: 创建绿色掩码")
|
||
cv2.imshow("绿色掩码", mask)
|
||
cv2.waitKey(delay)
|
||
|
||
# 应用掩码，只保留绿色部分
|
||
green_only = cv2.bitwise_and(img, img, mask=mask)
|
||
|
||
if observe:
|
||
print("步骤4: 提取绿色部分")
|
||
cv2.imshow("只保留绿色", green_only)
|
||
cv2.waitKey(delay)
|
||
|
||
# 将掩码转为灰度图
|
||
gray = mask.copy()
|
||
|
||
# 查找轮廓
|
||
contours, _ = cv2.findContours(gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
|
||
# 如果没有找到轮廓，返回未知
|
||
if not contours:
|
||
if observe:
|
||
print("未找到轮廓")
|
||
return "unknown"
|
||
|
||
if observe:
|
||
print(f"步骤5: 找到 {len(contours)} 个轮廓")
|
||
contour_img = img.copy()
|
||
cv2.drawContours(contour_img, contours, -1, (0, 255, 0), 2)
|
||
cv2.imshow("所有轮廓", contour_img)
|
||
cv2.waitKey(delay)
|
||
|
||
# 找到最大的轮廓（假设是箭头）
|
||
max_contour = max(contours, key=cv2.contourArea)
|
||
|
||
if observe:
|
||
print("步骤6: 提取最大轮廓")
|
||
max_contour_img = img.copy()
|
||
cv2.drawContours(max_contour_img, [max_contour], -1, (0, 0, 255), 2)
|
||
cv2.imshow("最大轮廓", max_contour_img)
|
||
cv2.waitKey(delay)
|
||
|
||
# 使用多边形近似轮廓
|
||
epsilon = 0.02 * cv2.arcLength(max_contour, True)
|
||
approx = cv2.approxPolyDP(max_contour, epsilon, True)
|
||
|
||
if observe:
|
||
print(f"步骤7: 多边形近似，顶点数: {len(approx)}")
|
||
approx_img = img.copy()
|
||
cv2.drawContours(approx_img, [approx], -1, (255, 0, 0), 2)
|
||
cv2.imshow("多边形近似", approx_img)
|
||
cv2.waitKey(delay)
|
||
|
||
# 计算凸包
|
||
hull = cv2.convexHull(max_contour)
|
||
|
||
if observe:
|
||
print("步骤8: 计算凸包")
|
||
hull_img = img.copy()
|
||
cv2.drawContours(hull_img, [hull], -1, (0, 255, 0), 2)
|
||
cv2.imshow("凸包", hull_img)
|
||
cv2.waitKey(delay)
|
||
|
||
# 计算凸缺陷
|
||
if len(max_contour) > 3:
|
||
defects = cv2.convexityDefects(max_contour, cv2.convexHull(max_contour, returnPoints=False))
|
||
else:
|
||
if observe:
|
||
print("轮廓点太少，无法准确判断")
|
||
return "unknown" # 轮廓点太少，无法准确判断
|
||
|
||
if defects is None:
|
||
if observe:
|
||
print("未找到凸缺陷")
|
||
return "unknown"
|
||
|
||
if observe:
|
||
print(f"步骤9: 计算凸缺陷，找到 {defects.shape[0]} 个缺陷点")
|
||
defect_img = img.copy()
|
||
for i in range(defects.shape[0]):
|
||
s, e, f, d = defects[i, 0]
|
||
start = tuple(max_contour[s][0])
|
||
end = tuple(max_contour[e][0])
|
||
far = tuple(max_contour[f][0])
|
||
cv2.circle(defect_img, far, 5, (0, 0, 255), -1)
|
||
cv2.imshow("凸缺陷", defect_img)
|
||
cv2.waitKey(delay)
|
||
|
||
# 获取轮廓的最小外接矩形
|
||
rect = cv2.minAreaRect(max_contour)
|
||
box = cv2.boxPoints(rect)
|
||
box = np.int0(box)
|
||
|
||
if observe:
|
||
print("步骤10: 获取最小外接矩形")
|
||
rect_img = img.copy()
|
||
cv2.drawContours(rect_img, [box], 0, (255, 0, 0), 2)
|
||
cv2.imshow("最小外接矩形", rect_img)
|
||
cv2.waitKey(delay)
|
||
|
||
# 获取中心点和角度
|
||
center = rect[0]
|
||
angle = rect[2]
|
||
|
||
if observe:
|
||
print(f"矩形中心: {center}, 角度: {angle}")
|
||
|
||
# 计算轮廓的矩，用于确定箭头方向
|
||
M = cv2.moments(max_contour)
|
||
|
||
# 避免除以零
|
||
if M["m00"] != 0:
|
||
cx = int(M["m10"] / M["m00"])
|
||
cy = int(M["m01"] / M["m00"])
|
||
else:
|
||
cx, cy = int(center[0]), int(center[1])
|
||
|
||
if observe:
|
||
print(f"步骤11: 计算质心 - 坐标: ({cx}, {cy})")
|
||
center_img = img.copy()
|
||
cv2.circle(center_img, (cx, cy), 5, (255, 0, 0), -1)
|
||
cv2.imshow("质心", center_img)
|
||
cv2.waitKey(delay)
|
||
|
||
# 改进的箭头尖端检测算法
|
||
# 1. 找到所有凸缺陷点
|
||
defect_points = []
|
||
for i in range(defects.shape[0]):
|
||
s, e, f, d = defects[i, 0]
|
||
start = tuple(max_contour[s][0])
|
||
end = tuple(max_contour[e][0])
|
||
far = tuple(max_contour[f][0])
|
||
|
||
# 计算缺陷点到中心的距离
|
||
dist = np.sqrt((far[0] - cx) ** 2 + (far[1] - cy) ** 2)
|
||
|
||
if observe:
|
||
print(f"缺陷点 {i}: 位置 {far}, 到中心距离: {dist:.2f}")
|
||
|
||
# 记录凸缺陷点及其距离
|
||
defect_points.append({
|
||
'point': far,
|
||
'distance': dist,
|
||
'start': start,
|
||
'end': end
|
||
})
|
||
|
||
# 没有缺陷点，使用矩形判断
|
||
if not defect_points:
|
||
# 判断逻辑将在后面处理
|
||
arrow_tip = None
|
||
else:
|
||
# 2. 按距离对缺陷点排序
|
||
defect_points.sort(key=lambda x: x['distance'], reverse=True)
|
||
|
||
# 3. 获取最远的几个缺陷点（可能的尖端候选）
|
||
top_n = min(3, len(defect_points))
|
||
candidates = defect_points[:top_n]
|
||
|
||
# 4. 分析这些点的位置分布
|
||
left_candidates = [p for p in candidates if p['point'][0] < cx]
|
||
right_candidates = [p for p in candidates if p['point'][0] >= cx]
|
||
|
||
# 5. 根据候选点的分布判断箭头朝向
|
||
if len(left_candidates) > len(right_candidates):
|
||
# 左侧候选点更多，箭头可能指向左边
|
||
arrow_tip = max(left_candidates, key=lambda x: x['distance'])['point']
|
||
arrow_direction = "left"
|
||
elif len(right_candidates) > len(left_candidates):
|
||
# 右侧候选点更多，箭头可能指向右边
|
||
arrow_tip = max(right_candidates, key=lambda x: x['distance'])['point']
|
||
arrow_direction = "right"
|
||
else:
|
||
# 候选点分布均衡，使用最远的点
|
||
arrow_tip = candidates[0]['point']
|
||
# 根据最远点的位置判断
|
||
if arrow_tip[0] < cx:
|
||
arrow_direction = "left"
|
||
else:
|
||
arrow_direction = "right"
|
||
|
||
if observe and arrow_tip:
|
||
print(f"步骤12: 找到可能的箭头尖端 - 位置: {arrow_tip}")
|
||
tip_img = img.copy()
|
||
cv2.circle(tip_img, arrow_tip, 8, (0, 255, 255), -1)
|
||
cv2.line(tip_img, (cx, cy), arrow_tip, (255, 0, 255), 2)
|
||
cv2.imshow("箭头尖端", tip_img)
|
||
cv2.waitKey(delay)
|
||
|
||
# 使用多种特征综合判断箭头方向
|
||
if arrow_tip is None:
|
||
# 如果没有找到明显的箭头尖端，使用最小外接矩形来判断
|
||
width = rect[1][0]
|
||
height = rect[1][1]
|
||
|
||
if observe:
|
||
print(f"未找到明显尖端，使用矩形判断 - 宽: {width}, 高: {height}")
|
||
|
||
# 矩形的方向向量
|
||
if width > height: # 水平箭头
|
||
# 考虑角度的定义
|
||
if angle > 45:
|
||
arrow_direction = "left"
|
||
else:
|
||
arrow_direction = "right"
|
||
else: # 垂直箭头，不在我们的考虑范围内
|
||
if observe:
|
||
print("垂直箭头，不在考虑范围")
|
||
return "unknown"
|
||
else:
|
||
# 已经在上面的代码中确定了方向
|
||
pass
|
||
|
||
# 附加检查：计算轮廓边界上的点，确认箭头的形状特征
|
||
x, y, w, h = cv2.boundingRect(max_contour)
|
||
aspect_ratio = float(w) / h
|
||
|
||
# 计算轮廓周长和面积
|
||
perimeter = cv2.arcLength(max_contour, True)
|
||
area = cv2.contourArea(max_contour)
|
||
|
||
# 计算轮廓的形状复杂度
|
||
complexity = perimeter / (4 * np.sqrt(area))
|
||
|
||
if observe:
|
||
print(f"轮廓分析 - 宽高比: {aspect_ratio:.2f}, 复杂度: {complexity:.2f}")
|
||
|
||
# 箭头形状特征检查
|
||
# 一般来说，箭头的复杂度会在一定范围内
|
||
if 1.1 < complexity < 3.0:
|
||
# 根据形状特征进一步确认方向判断
|
||
# 使用质心位置相对于轮廓边界的偏移
|
||
relative_cx = (cx - x) / w
|
||
|
||
if observe:
|
||
print(f"质心相对位置: {relative_cx:.2f}")
|
||
|
||
# 如果质心偏向左侧，箭头可能指向右侧
|
||
# 如果质心偏向右侧，箭头可能指向左侧
|
||
# 这是一个启发式规则，可能需要根据具体箭头形状调整
|
||
if relative_cx < 0.4:
|
||
# 质心在左侧，可能是右箭头
|
||
if arrow_direction == "left":
|
||
# 当前方法判断为左，但质心位置特征表明可能是右
|
||
# 增加额外的检查
|
||
if aspect_ratio > 1.5: # 宽大于高
|
||
arrow_direction = "right"
|
||
elif relative_cx > 0.6:
|
||
# 质心在右侧，可能是左箭头
|
||
if arrow_direction == "right":
|
||
# 当前方法判断为右，但质心位置特征表明可能是左
|
||
if aspect_ratio > 1.5: # 宽大于高
|
||
arrow_direction = "left"
|
||
|
||
if observe:
|
||
print(f"基于综合特征判断: {arrow_direction}箭头")
|
||
|
||
return arrow_direction
|
||
|
||
def visualize_arrow_detection(image, save_path=None, observe=False, delay=500):
|
||
"""
|
||
可视化箭头检测过程，显示中间结果
|
||
|
||
参数:
|
||
image: 输入图像，可以是文件路径或者已加载的图像数组
|
||
save_path: 保存结果图像的路径（可选）
|
||
observe: 是否输出中间状态信息和可视化结果，默认为False
|
||
delay: 展示每个步骤的等待时间(毫秒)，默认为500ms
|
||
"""
|
||
# 如果输入是字符串（文件路径），则加载图像
|
||
if isinstance(image, str):
|
||
img = cv2.imread(image)
|
||
else:
|
||
img = image.copy()
|
||
|
||
if img is None:
|
||
print("无法加载图像")
|
||
return
|
||
|
||
if observe:
|
||
print("\n开始可视化箭头检测过程")
|
||
|
||
# 转换到HSV颜色空间
|
||
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
|
||
|
||
# 绿色的HSV范围
|
||
lower_green = np.array([40, 50, 50])
|
||
upper_green = np.array([80, 255, 255])
|
||
|
||
# 创建绿色的掩码
|
||
mask = cv2.inRange(hsv, lower_green, upper_green)
|
||
|
||
# 应用掩码，只保留绿色部分
|
||
green_only = cv2.bitwise_and(img, img, mask=mask)
|
||
|
||
# 查找轮廓
|
||
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
|
||
# 创建输出图像
|
||
output = img.copy()
|
||
|
||
# 如果找到轮廓，绘制最大轮廓
|
||
if contours:
|
||
max_contour = max(contours, key=cv2.contourArea)
|
||
cv2.drawContours(output, [max_contour], -1, (0, 0, 255), 2)
|
||
|
||
# 获取轮廓的最小外接矩形
|
||
rect = cv2.minAreaRect(max_contour)
|
||
box = cv2.boxPoints(rect)
|
||
box = np.int0(box)
|
||
cv2.drawContours(output, [box], 0, (255, 0, 0), 2)
|
||
|
||
# 计算轮廓的矩
|
||
M = cv2.moments(max_contour)
|
||
|
||
# 避免除以零
|
||
if M["m00"] != 0:
|
||
cx = int(M["m10"] / M["m00"])
|
||
cy = int(M["m01"] / M["m00"])
|
||
|
||
# 绘制中心点
|
||
cv2.circle(output, (cx, cy), 5, (255, 0, 0), -1)
|
||
|
||
# 绘制凸包
|
||
hull = cv2.convexHull(max_contour)
|
||
cv2.drawContours(output, [hull], 0, (0, 255, 0), 2)
|
||
|
||
# 绘制凸缺陷
|
||
if len(max_contour) > 3:
|
||
defects = cv2.convexityDefects(max_contour, cv2.convexHull(max_contour, returnPoints=False))
|
||
if defects is not None:
|
||
for i in range(defects.shape[0]):
|
||
s, e, f, d = defects[i, 0]
|
||
start = tuple(max_contour[s][0])
|
||
end = tuple(max_contour[e][0])
|
||
far = tuple(max_contour[f][0])
|
||
cv2.circle(output, far, 5, (0, 0, 255), -1) # 在凸缺陷点绘制红色圆点
|
||
|
||
# 获取箭头方向
|
||
direction = detect_arrow_direction(img, observe=observe, delay=delay)
|
||
|
||
# 在图像上添加方向文本
|
||
cv2.putText(output, f"Direction: {direction}", (10, 30),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
|
||
|
||
# 如果提供了保存路径，保存结果图像
|
||
if save_path:
|
||
cv2.imwrite(save_path, output)
|
||
if observe:
|
||
print(f"结果已保存到: {save_path}")
|
||
|
||
# 创建一个包含所有图像的窗口
|
||
result = np.hstack((img, green_only, output))
|
||
|
||
# 调整大小以便查看
|
||
scale_percent = 50 # 缩放到原来的50%
|
||
width = int(result.shape[1] * scale_percent / 100)
|
||
height = int(result.shape[0] * scale_percent / 100)
|
||
dim = (width, height)
|
||
resized = cv2.resize(result, dim, interpolation=cv2.INTER_AREA)
|
||
|
||
# 显示结果
|
||
cv2.imshow('Arrow Detection Process', resized)
|
||
cv2.waitKey(0)
|
||
cv2.destroyAllWindows()
|
||
|
||
# 用法示例
|
||
if __name__ == "__main__":
|
||
# 替换为实际图像路径
|
||
image_path = "path/to/arrow/image.png"
|
||
|
||
# 检测箭头方向，使用较长的延迟时间（1500毫秒）
|
||
direction = detect_arrow_direction(image_path, observe=True, delay=1500)
|
||
print(f"检测到的箭头方向: {direction}")
|
||
|
||
# 可视化检测过程，使用较长的延迟时间
|
||
visualize_arrow_detection(image_path, observe=True, delay=1500)
|