345 lines
12 KiB
Python
345 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
点阵灯盘检测算法
|
||
|
||
核心功能：
|
||
1. 预处理与离散点聚合
|
||
2. 连通域分析与圆拟合
|
||
3. 多灯盘分离
|
||
4. 后处理过滤
|
||
"""
|
||
import cv2
|
||
import numpy as np
|
||
from sklearn.cluster import DBSCAN
|
||
|
||
class DotMatrixDetector:
|
||
def __init__(self):
|
||
"""初始化点阵灯盘检测器"""
|
||
# 默认HSV参数（针对高亮绿色LED）
|
||
self.h_low = 35 # H通道最小值（绿色起始）
|
||
self.h_high = 85 # H通道最大值（绿色结束）
|
||
self.s_low = 60 # S通道最小值（饱和度阈值，降低以适应LED）
|
||
self.s_high = 255 # S通道最大值
|
||
self.v_low = 200 # V通道最小值（亮度阈值，提高以过滤环境光）
|
||
self.v_high = 255 # V通道最大值
|
||
|
||
# 形态学参数
|
||
self.morph_kernel_size = 20 # 闭运算核大小（略大于LED点阵间距）
|
||
self.morph_open_kernel = 5 # 开运算核大小（去除噪点）
|
||
|
||
# 几何参数
|
||
self.min_disk_area = 500 # 最小灯盘面积（像素）
|
||
self.min_led_count = 10 # 最小LED数量
|
||
self.circularity_threshold = 0.4 # 圆度阈值（点阵灯盘较宽松）
|
||
|
||
# 检测模式
|
||
self.detection_mode = 'least_squares' # least_squares, min_enclosing, dbscan
|
||
|
||
# DBSCAN参数
|
||
self.dbscan_eps = 30 # 空间距离阈值
|
||
self.dbscan_min_samples = 5 # 最小样本数
|
||
|
||
def set_hsv_range(self, h_low, h_high, s_low, s_high, v_low, v_high):
|
||
"""设置HSV颜色范围"""
|
||
self.h_low = h_low
|
||
self.h_high = h_high
|
||
self.s_low = s_low
|
||
self.s_high = s_high
|
||
self.v_low = v_low
|
||
self.v_high = v_high
|
||
|
||
def set_morph_parameters(self, kernel_size, open_kernel):
|
||
"""设置形态学参数"""
|
||
self.morph_kernel_size = kernel_size
|
||
self.morph_open_kernel = open_kernel
|
||
|
||
def set_geometry_parameters(self, min_area, min_leds, circularity):
|
||
"""设置几何参数"""
|
||
self.min_disk_area = min_area
|
||
self.min_led_count = min_leds
|
||
self.circularity_threshold = circularity
|
||
|
||
def set_detection_mode(self, mode):
|
||
"""设置检测模式"""
|
||
self.detection_mode = mode
|
||
|
||
def detect(self, frame):
|
||
"""
|
||
检测点阵灯盘
|
||
|
||
Args:
|
||
frame: BGR格式的图像
|
||
|
||
Returns:
|
||
results: 检测结果列表，每个元素包含圆心坐标、半径、置信度等
|
||
mask: 二值掩码图像
|
||
raw_leds: 原始LED点集
|
||
aggregated_mask: 形态学聚合后的掩码
|
||
"""
|
||
# 1. 预处理：颜色空间转换
|
||
hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
|
||
|
||
# 2. 颜色阈值分割
|
||
lower_green = np.array([self.h_low, self.s_low, self.v_low])
|
||
upper_green = np.array([self.h_high, self.s_high, self.v_high])
|
||
mask = cv2.inRange(hsv, lower_green, upper_green)
|
||
|
||
# 3. 形态学操作：闭运算（聚合离散点）
|
||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (self.morph_kernel_size, self.morph_kernel_size))
|
||
aggregated_mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
|
||
|
||
# 4. 形态学操作：开运算（去除噪点）
|
||
open_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (self.morph_open_kernel, self.morph_open_kernel))
|
||
aggregated_mask = cv2.morphologyEx(aggregated_mask, cv2.MORPH_OPEN, open_kernel)
|
||
|
||
# 5. 提取原始LED点
|
||
raw_leds = self._extract_raw_leds(mask)
|
||
|
||
# 6. 检测灯盘
|
||
results = []
|
||
|
||
if self.detection_mode == 'dbscan' and len(raw_leds) >= self.min_led_count:
|
||
# 使用DBSCAN聚类检测多灯盘
|
||
results = self._detect_with_dbscan(raw_leds)
|
||
else:
|
||
# 使用连通域分析
|
||
results = self._detect_with_contours(aggregated_mask, mask)
|
||
|
||
# 7. 后处理：非极大值抑制
|
||
results = self._non_maximum_suppression(results)
|
||
|
||
return results, mask, raw_leds, aggregated_mask
|
||
|
||
def _extract_raw_leds(self, mask):
|
||
"""提取原始LED点"""
|
||
# 查找连通域
|
||
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
|
||
leds = []
|
||
for contour in contours:
|
||
# 计算中心点
|
||
M = cv2.moments(contour)
|
||
if M['m00'] > 0:
|
||
cx = int(M['m10'] / M['m00'])
|
||
cy = int(M['m01'] / M['m00'])
|
||
leds.append((cx, cy))
|
||
|
||
return np.array(leds)
|
||
|
||
def _detect_with_contours(self, aggregated_mask, raw_mask):
|
||
"""基于连通域的灯盘检测"""
|
||
results = []
|
||
|
||
# 查找轮廓
|
||
contours, _ = cv2.findContours(aggregated_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
|
||
for i, contour in enumerate(contours):
|
||
# 计算轮廓面积
|
||
area = cv2.contourArea(contour)
|
||
if area < self.min_disk_area:
|
||
continue
|
||
|
||
# 在原始掩码中提取该连通域内的LED点
|
||
led_points = self._extract_leds_in_contour(contour, raw_mask)
|
||
if len(led_points) < self.min_led_count:
|
||
continue
|
||
|
||
# 根据检测模式拟合圆
|
||
if self.detection_mode == 'min_enclosing':
|
||
center, radius = cv2.minEnclosingCircle(led_points)
|
||
cx, cy = center
|
||
else: # least_squares
|
||
center, radius = self._fit_circle_least_squares(led_points)
|
||
if center is None:
|
||
continue
|
||
cx, cy = center
|
||
|
||
# 计算圆度
|
||
circularity = self._calculate_circularity(contour, radius)
|
||
if circularity < self.circularity_threshold:
|
||
continue
|
||
|
||
# 计算置信度
|
||
confidence = self._calculate_confidence(len(led_points), circularity, radius)
|
||
|
||
# 添加结果
|
||
results.append({
|
||
'id': i,
|
||
'center': (cx, cy),
|
||
'radius': radius,
|
||
'confidence': confidence,
|
||
'led_count': len(led_points),
|
||
'circularity': circularity,
|
||
'contour': contour,
|
||
'led_points': led_points
|
||
})
|
||
|
||
return results
|
||
|
||
def _detect_with_dbscan(self, led_points):
|
||
"""使用DBSCAN聚类检测多灯盘"""
|
||
results = []
|
||
|
||
if len(led_points) < self.min_led_count:
|
||
return results
|
||
|
||
# DBSCAN聚类
|
||
db = DBSCAN(eps=self.dbscan_eps, min_samples=self.dbscan_min_samples).fit(led_points)
|
||
labels = db.labels_
|
||
|
||
# 处理每个簇
|
||
unique_labels = set(labels)
|
||
for i, label in enumerate(unique_labels):
|
||
if label == -1: # 噪声点
|
||
continue
|
||
|
||
# 提取簇中的点
|
||
cluster_points = led_points[labels == label]
|
||
if len(cluster_points) < self.min_led_count:
|
||
continue
|
||
|
||
# 拟合圆
|
||
center, radius = self._fit_circle_least_squares(cluster_points)
|
||
if center is None:
|
||
continue
|
||
cx, cy = center
|
||
|
||
# 计算置信度
|
||
circularity = self._calculate_cluster_circularity(cluster_points, center, radius)
|
||
confidence = self._calculate_confidence(len(cluster_points), circularity, radius)
|
||
|
||
# 添加结果
|
||
results.append({
|
||
'id': i,
|
||
'center': (cx, cy),
|
||
'radius': radius,
|
||
'confidence': confidence,
|
||
'led_count': len(cluster_points),
|
||
'circularity': circularity,
|
||
'led_points': cluster_points
|
||
})
|
||
|
||
return results
|
||
|
||
def _extract_leds_in_contour(self, contour, mask):
|
||
"""提取轮廓内的LED点"""
|
||
# 创建掩码
|
||
contour_mask = np.zeros_like(mask)
|
||
cv2.drawContours(contour_mask, [contour], -1, 255, -1)
|
||
|
||
# 与原始掩码相交
|
||
combined_mask = cv2.bitwise_and(mask, contour_mask)
|
||
|
||
# 提取点
|
||
return self._extract_raw_leds(combined_mask)
|
||
|
||
def _fit_circle_least_squares(self, points):
|
||
"""最小二乘拟合圆"""
|
||
if len(points) < 3:
|
||
return None, None
|
||
|
||
# 最小二乘拟合
|
||
x = points[:, 0]
|
||
y = points[:, 1]
|
||
|
||
# 构建矩阵
|
||
A = np.column_stack((x, y, np.ones_like(x)))
|
||
b = x**2 + y**2
|
||
|
||
# 求解线性方程组
|
||
try:
|
||
x_opt = np.linalg.lstsq(A, b, rcond=None)[0]
|
||
cx = x_opt[0] / 2
|
||
cy = x_opt[1] / 2
|
||
radius = np.sqrt(x_opt[2] + cx**2 + cy**2)
|
||
return (cx, cy), radius
|
||
except:
|
||
return None, None
|
||
|
||
def _calculate_circularity(self, contour, radius):
|
||
"""计算圆度"""
|
||
area = cv2.contourArea(contour)
|
||
expected_area = np.pi * radius**2
|
||
if expected_area == 0:
|
||
return 0
|
||
return area / expected_area
|
||
|
||
def _calculate_cluster_circularity(self, points, center, radius):
|
||
"""计算聚类点集的圆度"""
|
||
if len(points) < 3:
|
||
return 0
|
||
|
||
# 计算点到圆心的距离
|
||
distances = np.sqrt((points[:, 0] - center[0])**2 + (points[:, 1] - center[1])**2)
|
||
|
||
# 计算距离的标准差
|
||
std_distance = np.std(distances)
|
||
if radius == 0:
|
||
return 0
|
||
|
||
# 圆度：标准差越小，圆度越高
|
||
circularity = 1.0 / (1.0 + std_distance / radius)
|
||
return circularity
|
||
|
||
def _calculate_confidence(self, led_count, circularity, radius):
|
||
"""计算置信度"""
|
||
# LED数量得分
|
||
led_score = min(led_count / 50, 1.0)
|
||
|
||
# 圆度得分
|
||
circ_score = min(circularity / 0.7, 1.0)
|
||
|
||
# 综合得分
|
||
confidence = (led_score * 0.6 + circ_score * 0.4) * 100
|
||
return confidence
|
||
|
||
def _non_maximum_suppression(self, results, iou_threshold=0.5):
|
||
"""非极大值抑制"""
|
||
if not results:
|
||
return []
|
||
|
||
# 按置信度排序
|
||
results.sort(key=lambda x: x['confidence'], reverse=True)
|
||
|
||
filtered_results = []
|
||
while results:
|
||
current = results.pop(0)
|
||
filtered_results.append(current)
|
||
|
||
# 过滤重叠的结果
|
||
results = [r for r in results if not self._is_overlapping(current, r, iou_threshold)]
|
||
|
||
return filtered_results
|
||
|
||
def _is_overlapping(self, result1, result2, threshold):
|
||
"""判断两个圆是否重叠"""
|
||
cx1, cy1 = result1['center']
|
||
r1 = result1['radius']
|
||
cx2, cy2 = result2['center']
|
||
r2 = result2['radius']
|
||
|
||
# 计算圆心距
|
||
distance = np.sqrt((cx1 - cx2)**2 + (cy1 - cy2)**2)
|
||
|
||
# 计算重叠面积
|
||
if distance >= r1 + r2:
|
||
return False
|
||
|
||
if distance <= abs(r1 - r2):
|
||
return True
|
||
|
||
# 计算重叠区域
|
||
a = (r1**2 - r2**2 + distance**2) / (2 * distance)
|
||
b = distance - a
|
||
h = np.sqrt(r1**2 - a**2)
|
||
|
||
area1 = r1**2 * np.arccos(a / r1) - a * h
|
||
area2 = r2**2 * np.arccos(b / r2) - b * h
|
||
overlap_area = area1 + area2
|
||
|
||
# 计算IoU
|
||
total_area = np.pi * (r1**2 + r2**2) - overlap_area
|
||
iou = overlap_area / total_area
|
||
|
||
return iou > threshold
|