223 lines
8.4 KiB
Python
223 lines
8.4 KiB
Python
import cv2
|
||
import numpy as np
|
||
from 新建文件夹.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 |