import cv2 import numpy as np from 视觉代码python.config import ( HORIZONTAL_ANGLE_THRESHOLD, NEARBY_LIGHT_BAR_THRESHOLD, LIGHT_BAR_IOU_THRESHOLD, ARMOR_DISTANCE_RATIO_RANGE, ARMOR_LENGTH_DIFF_RATIO, ARMOR_ANGLE_DIFF_THRESHOLD ) class ArmorDetector: """装甲板检测类：灯条提取、配对""" def __init__(self): self.angle_threshold_rad = np.radians(HORIZONTAL_ANGLE_THRESHOLD) def _extract_initial_light_bars(self, mask): """从掩码中提取初始灯条（轮廓+最小外接矩形）""" contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) initial_light_bars = [] for contour in contours: rect = cv2.minAreaRect(contour) box = cv2.boxPoints(rect) # 获取矩形四个顶点 if len(box) == 0: # 确保顶点数有效 continue box = np.int32(box) # 转为int32（OpenCV要求的类型） (x, y), (w, h), angle = rect contour_area = cv2.contourArea(contour) rect_area = w * h # 筛选：面积足够 + 填充度高 if rect_area > 30 and (contour_area / rect_area if rect_area > 0 else 0) > 0.4: initial_light_bars.append({"rect": rect, "box": box, "area": rect_area}) return initial_light_bars def _merge_nearby_light_bars(self, initial_light_bars): """合并临近灯条（避免同一灯条被拆分）""" processed_boxes = [] visited = [False] * len(initial_light_bars) for i in range(len(initial_light_bars)): if visited[i]: continue nearby_indices = [i] (x1, y1), _, _ = initial_light_bars[i]["rect"] for j in range(i + 1, len(initial_light_bars)): if visited[j]: continue (x2, y2), _, _ = initial_light_bars[j]["rect"] distance = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2) # 距离+IOU筛选 if distance < NEARBY_LIGHT_BAR_THRESHOLD: box_i = initial_light_bars[i]["box"] box_j = initial_light_bars[j]["box"] if len(box_i) == 0 or len(box_j) == 0: continue b1 = cv2.boundingRect(box_i) b2 = cv2.boundingRect(box_j) inter_area = self._calculate_iou(b1, b2) if inter_area > LIGHT_BAR_IOU_THRESHOLD: nearby_indices.append(j) visited[j] = True # 合并临近灯条的顶点，生成新矩形 all_valid_points = [] for k in nearby_indices: box_k = initial_light_bars[k]["box"] if len(box_k) > 0: all_valid_points.append(box_k) if not all_valid_points: continue all_points = np.vstack(all_valid_points) merged_rect = cv2.minAreaRect(all_points) merged_box = np.int32(cv2.boxPoints(merged_rect)) # 转为int32 processed_boxes.append(merged_box) visited[i] = True return processed_boxes def _calculate_iou(self, b1, b2): """计算两个bounding box的IOU""" b1 = [b1[0], b1[1], b1[0] + b1[2], b1[1] + b1[3]] b2 = [b2[0], b2[1], b2[0] + b2[2], b2[1] + b2[3]] inter_x1 = max(b1[0], b2[0]) inter_y1 = max(b1[1], b2[1]) inter_x2 = min(b1[2], b2[2]) inter_y2 = min(b1[3], b2[3]) if inter_x1 >= inter_x2 or inter_y1 >= inter_y2: return 0.0 inter_area = (inter_x2 - inter_x1) * (inter_y2 - inter_y1) b1_area = (b1[2] - b1[0]) * (b1[3] - b1[1]) b2_area = (b2[2] - b2[0]) * (b2[3] - b2[1]) union_area = b1_area + b2_area - inter_area return inter_area / union_area if union_area > 0 else 0.0 def _filter_valid_light_bars(self, processed_boxes): """筛选有效灯条（排除水平灯条）""" valid_light_bars = [] for box in processed_boxes: if len(box) == 0: continue rect = cv2.minAreaRect(box) (x, y), (w, h), angle = rect length = w if w > h else h width = h if w > h else w main_angle = angle if w > h else angle + 90 main_angle = main_angle % 180 if main_angle > 90: main_angle -= 180 main_angle_rad = np.radians(main_angle) center = (x, y) end1 = ( center[0] + (length / 2) * np.cos(main_angle_rad), center[1] + (length / 2) * np.sin(main_angle_rad) ) end2 = ( center[0] - (length / 2) * np.cos(main_angle_rad), center[1] - (length / 2) * np.sin(main_angle_rad) ) rect_area = length * width contour_area = cv2.contourArea(box) area_ratio = contour_area / rect_area if rect_area > 0 else 0 # 筛选：面积足够 + 填充度高 + 非水平 if (rect_area > 40 and area_ratio > 0.4 and abs(main_angle_rad) > self.angle_threshold_rad): valid_light_bars.append({ "center": center, "size": (length, width), "angle": main_angle, "angle_rad": main_angle_rad, "area": rect_area, "center_line": [end1, end2], "center_line_length": length, "box": box }) return valid_light_bars def _pair_light_bars_to_armor(self, light_bars, target_color): """灯条配对为装甲板（按置信度排序）""" armor_plates = [] used_indices = set() for i in range(len(light_bars)): if i in used_indices: continue lb1 = light_bars[i] for j in range(i + 1, len(light_bars)): if j in used_indices: continue lb2 = light_bars[j] # 角度差异筛选 angle_diff = abs(lb1["angle_rad"] - lb2["angle_rad"]) angle_diff = min(angle_diff, 2 * np.pi - angle_diff) if angle_diff > ARMOR_ANGLE_DIFF_THRESHOLD: continue # 距离比例筛选 dx = lb2["center"][0] - lb1["center"][0] dy = lb2["center"][1] - lb1["center"][1] distance = np.sqrt(dx ** 2 + dy ** 2) avg_length = (lb1["center_line_length"] + lb2["center_line_length"]) / 2 distance_ratio = distance / avg_length if not (ARMOR_DISTANCE_RATIO_RANGE[0] < distance_ratio < ARMOR_DISTANCE_RATIO_RANGE[1]): continue # 长度差异筛选 length_diff_ratio = abs(lb1["center_line_length"] - lb2["center_line_length"]) / avg_length if length_diff_ratio > ARMOR_LENGTH_DIFF_RATIO: continue armor_center = ( (lb1["center"][0] + lb2["center"][0]) / 2, (lb1["center"][1] + lb2["center"][1]) / 2 ) confidence = (lb1["area"] + lb2["area"]) / (distance + 1) armor_plates.append({ "color": target_color, "center": armor_center, "confidence": confidence, "pair": (lb1, lb2), "corners_2d": None }) used_indices.add(i) used_indices.add(j) break return sorted(armor_plates, key=lambda x: x["confidence"], reverse=True) def detect(self, mask, target_color): """完整检测流程：灯条提取→合并→筛选→配对→3D估计""" initial_light_bars = self._extract_initial_light_bars(mask) if not initial_light_bars: return [], [] processed_boxes = self._merge_nearby_light_bars(initial_light_bars) if not processed_boxes: return [], [] valid_light_bars = self._filter_valid_light_bars(processed_boxes) if len(valid_light_bars) < 2: return valid_light_bars, [] armor_plates = self._pair_light_bars_to_armor(valid_light_bars, target_color) if not armor_plates: return valid_light_bars, [] return valid_light_bars, armor_plates