181 lines
5.6 KiB
Python
181 lines
5.6 KiB
Python
import cv2
|
|
import numpy as np
|
|
|
|
def detect_arrow_direction(image):
|
|
"""
|
|
从图像中提取绿色箭头并判断其指向方向（左或右）
|
|
|
|
参数:
|
|
image: 输入图像，可以是文件路径或者已加载的图像数组
|
|
|
|
返回:
|
|
direction: 字符串，"left"表示左箭头，"right"表示右箭头，"unknown"表示无法确定
|
|
"""
|
|
# 如果输入是字符串（文件路径），则加载图像
|
|
if isinstance(image, str):
|
|
img = cv2.imread(image)
|
|
else:
|
|
img = image.copy()
|
|
|
|
if img is None:
|
|
print("无法加载图像")
|
|
return "unknown"
|
|
|
|
# 转换到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)
|
|
|
|
# 将掩码转为灰度图
|
|
gray = mask.copy()
|
|
|
|
# 查找轮廓
|
|
contours, _ = cv2.findContours(gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
|
|
# 如果没有找到轮廓，返回未知
|
|
if not contours:
|
|
return "unknown"
|
|
|
|
# 找到最大的轮廓（假设是箭头）
|
|
max_contour = max(contours, key=cv2.contourArea)
|
|
|
|
# 获取轮廓的最小外接矩形
|
|
x, y, w, h = cv2.boundingRect(max_contour)
|
|
|
|
# 计算轮廓的矩，用于确定箭头方向
|
|
M = cv2.moments(max_contour)
|
|
|
|
# 避免除以零
|
|
if M["m00"] != 0:
|
|
cx = int(M["m10"] / M["m00"])
|
|
cy = int(M["m01"] / M["m00"])
|
|
else:
|
|
cx, cy = x + w//2, y + h//2
|
|
|
|
# 将图像分为左右两部分
|
|
left_region = gray[y:y+h, x:cx]
|
|
right_region = gray[y:y+h, cx:x+w]
|
|
|
|
# 计算每个区域中白色像素的数量（箭头部分）
|
|
left_pixels = cv2.countNonZero(left_region)
|
|
right_pixels = cv2.countNonZero(right_region)
|
|
|
|
# 计算每个区域中白色像素的密度
|
|
left_density = left_pixels / (left_region.size + 1e-10)
|
|
right_density = right_pixels / (right_region.size + 1e-10)
|
|
|
|
# 根据左右区域的像素密度确定箭头方向
|
|
# 如果箭头指向右侧，右侧区域的箭头头部应该有更多的像素密度
|
|
# 如果箭头指向左侧，左侧区域的箭头头部应该有更多的像素密度
|
|
if right_density > left_density:
|
|
return "right"
|
|
else:
|
|
return "left"
|
|
|
|
def visualize_arrow_detection(image, save_path=None):
|
|
"""
|
|
可视化箭头检测过程，显示中间结果
|
|
|
|
参数:
|
|
image: 输入图像，可以是文件路径或者已加载的图像数组
|
|
save_path: 保存结果图像的路径（可选）
|
|
"""
|
|
# 如果输入是字符串（文件路径），则加载图像
|
|
if isinstance(image, str):
|
|
img = cv2.imread(image)
|
|
else:
|
|
img = image.copy()
|
|
|
|
if img is None:
|
|
print("无法加载图像")
|
|
return
|
|
|
|
# 转换到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)
|
|
|
|
# 获取轮廓的最小外接矩形
|
|
x, y, w, h = cv2.boundingRect(max_contour)
|
|
cv2.rectangle(output, (x, y), (x + w, y + h), (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)
|
|
|
|
# 绘制分割线
|
|
cv2.line(output, (cx, y), (cx, y + h), (0, 255, 255), 2)
|
|
|
|
# 获取箭头方向
|
|
direction = detect_arrow_direction(img)
|
|
|
|
# 在图像上添加方向文本
|
|
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)
|
|
|
|
# 创建一个包含所有图像的窗口
|
|
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"
|
|
|
|
# 检测箭头方向
|
|
direction = detect_arrow_direction(image_path)
|
|
print(f"检测到的箭头方向: {direction}")
|
|
|
|
# 可视化检测过程
|
|
visualize_arrow_detection(image_path)
|