80 lines
2.9 KiB
Python
80 lines
2.9 KiB
Python
import cv2
|
||
import numpy as np
|
||
|
||
class Visualizer:
|
||
"""可视化：绘制灯条、装甲板、预测位置等"""
|
||
def __init__(self):
|
||
self.colors = {
|
||
"red": (0, 0, 255),
|
||
"blue": (255, 0, 0),
|
||
"light_bar": (0, 255, 255),
|
||
"armor": (0, 255, 0),
|
||
"prediction": (0, 165, 255),
|
||
"text": (255, 255, 255)
|
||
}
|
||
|
||
def draw_light_bars(self, frame, light_bars, target_color):
|
||
"""绘制灯条（矩形+中心线+标签）"""
|
||
bar_color = self.colors[target_color]
|
||
line_color = self.colors["light_bar"]
|
||
|
||
for lb in light_bars:
|
||
cv2.drawContours(frame, [lb["box"]], 0, bar_color, 2)
|
||
end1 = tuple(map(int, lb["center_line"][0]))
|
||
end2 = tuple(map(int, lb["center_line"][1]))
|
||
cv2.line(frame, end1, end2, line_color, 2)
|
||
center = tuple(map(int, lb["center"]))
|
||
cv2.putText(
|
||
frame, target_color, (center[0], center[1] - 8),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.4, bar_color, 1
|
||
)
|
||
|
||
return frame
|
||
|
||
def draw_armor_plate(self, frame, armor):
|
||
"""绘制装甲板（2D框+3D信息）"""
|
||
corners = armor.get("corners_2d")
|
||
if corners is not None:
|
||
# 关键修改：np.int0 替换为 np.int32
|
||
cv2.polylines(frame, [np.int32(corners)], True, self.colors["armor"], 2)
|
||
|
||
center = tuple(map(int, armor["center"]))
|
||
cv2.circle(frame, center, 3, self.colors["armor"], -1)
|
||
cv2.putText(
|
||
frame, "Armor", (center[0] + 5, center[1] - 5),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.5, self.colors["armor"], 2
|
||
)
|
||
|
||
return frame
|
||
|
||
def draw_prediction(self, frame, predicted_pos, label="Predicted"):
|
||
"""绘制预测位置（橙色圆点+标签）"""
|
||
if predicted_pos is None:
|
||
return frame
|
||
|
||
pos = tuple(map(int, predicted_pos))
|
||
cv2.circle(frame, pos, 5, self.colors["prediction"], -1)
|
||
cv2.putText(
|
||
frame, label, (pos[0] + 10, pos[1] - 10),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.5, self.colors["prediction"], 2
|
||
)
|
||
|
||
return frame
|
||
|
||
def draw_offset_text(self, frame, display_center, target_color, is_predicted):
|
||
"""绘制装甲板相对图像中心的偏移量"""
|
||
if display_center is None:
|
||
text = "Armor: unknown | Offset: (0, 0)"
|
||
else:
|
||
img_center = (frame.shape[1] // 2, frame.shape[0] // 2)
|
||
offset_x = int(display_center[0] - img_center[0])
|
||
offset_y = int(display_center[1] - img_center[1])
|
||
color_text = f"{target_color} (predicted)" if is_predicted else target_color
|
||
text = f"Armor: {color_text} | Offset: ({offset_x}, {offset_y})"
|
||
|
||
cv2.putText(
|
||
frame, text, (10, 30),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.8, self.colors["text"], 2
|
||
)
|
||
|
||
return frame |