mi-task/utils/detect_bar.py

117 lines
4.6 KiB
Python
Raw Normal View History

import cv2
import numpy as np
def detect_bar_top_distance(image, observe=False, delay=800):
"""
识别图片中灰色栏杆到图像顶部的距离像素
:param image: 输入BGR格式图片
:param observe: 是否显示中间处理结果
:param delay: 显示结果的延迟时间(ms)
:return: 距离像素未检测到返回None
"""
# 转为HSV便于分离灰色
hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
# 调整灰色范围 - 更精确地针对多种灰色
lower_gray = np.array([0, 0, 30]) # 降低亮度下限以捕获更暗的灰色
upper_gray = np.array([180, 40, 120]) # 增加饱和度和亮度上限以捕获更多灰色变体
# 创建掩码
mask = cv2.inRange(hsv, lower_gray, upper_gray)
# 尝试使用边缘检测作为辅助方法
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
edges = cv2.Canny(gray, 50, 150)
# 仅保留图像上半部分的边缘,因为栏杆通常在上部
height = edges.shape[0]
edges[int(height*0.6):, :] = 0
# 使用霍夫变换检测水平线条
lines = cv2.HoughLinesP(edges, 1, np.pi/180, threshold=50, minLineLength=100, maxLineGap=20)
# 如果找到线条,创建线条掩码
line_mask = np.zeros_like(gray)
if lines is not None:
for line in lines:
x1, y1, x2, y2 = line[0]
# 仅保留接近水平的线条
if abs(y2 - y1) < 20: # 水平线的y差异很小
cv2.line(line_mask, (x1, y1), (x2, y2), 255, 5)
# 合并HSV掩码和线条掩码
combined_mask = cv2.bitwise_or(mask, line_mask)
# 增强形态学操作以更好地提取横杆
# 先进行闭操作连接相近区域
kernel_close = np.ones((5, 30), np.uint8)
combined_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_CLOSE, kernel_close)
# 使用开操作去除小噪点
kernel_open = np.ones((3, 3), np.uint8)
combined_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_OPEN, kernel_open)
# 再次进行水平方向的膨胀,增强横杆连通性
kernel_dilate = np.ones((1, 20), np.uint8)
combined_mask = cv2.dilate(combined_mask, kernel_dilate, iterations=1)
if observe:
cv2.imshow("HSV Mask", mask)
cv2.imshow("Edge Detection", edges)
cv2.imshow("Combined Mask", combined_mask)
cv2.waitKey(delay)
# 查找轮廓
contours, _ = cv2.findContours(combined_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if not contours:
# 如果没有在组合掩码中找到轮廓,尝试在原始掩码中查找
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if not contours:
return None
# 按面积排序找到最大的几个轮廓
contours = sorted(contours, key=cv2.contourArea, reverse=True)
# 从最大的轮廓中筛选出形状类似横杆的轮廓
for contour in contours[:5]: # 只检查最大的5个轮廓
x, y, w, h = cv2.boundingRect(contour)
# 横杆特征:宽度远大于高度,且宽度占据图像较大比例
aspect_ratio = float(w) / h if h > 0 else 0
width_ratio = float(w) / image.shape[1]
# 放宽轮廓筛选条件
if aspect_ratio > 3 and width_ratio > 0.2 and h < image.shape[0] * 0.3:
if observe:
debug_img = image.copy()
cv2.rectangle(debug_img, (x, y), (x+w, y+h), (0, 255, 0), 2)
cv2.imshow("Detected Bar", debug_img)
cv2.waitKey(delay)
# 返回栏杆顶部到图像顶部的距离
return y
# 如果基于轮廓没有找到合适的栏杆,尝试使用霍夫线检测结果
if lines is not None:
horizontal_lines = []
for line in lines:
x1, y1, x2, y2 = line[0]
# 筛选水平线条
if abs(y2 - y1) < 10 and abs(x2 - x1) > image.shape[1] * 0.3:
horizontal_lines.append((min(y1, y2), abs(x2 - x1)))
# 按线条长度排序
horizontal_lines.sort(key=lambda x: x[1], reverse=True)
if horizontal_lines and len(horizontal_lines) > 0:
# 返回最长的水平线的y坐标
if observe:
debug_img = image.copy()
y = horizontal_lines[0][0]
cv2.line(debug_img, (0, y), (image.shape[1], y), (0, 255, 0), 2)
cv2.imshow("Detected Bar (Hough)", debug_img)
cv2.waitKey(delay)
return horizontal_lines[0][0]
return None