feat: add shared window
This commit is contained in:
221
visualization.py
Normal file
221
visualization.py
Normal file
@@ -0,0 +1,221 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Модуль для управления общим окном визуализации
|
||||
"""
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib.patches as patches
|
||||
import numpy as np
|
||||
from enum import Enum
|
||||
import cv2
|
||||
from PIL import Image
|
||||
import matplotlib
|
||||
|
||||
# Настройки matplotlib
|
||||
matplotlib.use('TkAgg')
|
||||
plt.rcParams['figure.raise_window'] = False
|
||||
|
||||
class SimMode(Enum):
|
||||
OPERATOR = 1
|
||||
AUTONOME = 2
|
||||
|
||||
class VisualizationManager:
|
||||
"""
|
||||
Менеджер для управления общим окном визуализации
|
||||
"""
|
||||
|
||||
def __init__(self, window_title="Drone Autopilot Visualization"):
|
||||
self.window_title = window_title
|
||||
self.fig = None
|
||||
self.ax_global_map = None
|
||||
self.ax_detection = None
|
||||
self.ax_matches = None
|
||||
|
||||
# Данные для глобальной карты
|
||||
self.trajectory_x = []
|
||||
self.trajectory_y = []
|
||||
self.trajectory_modes = []
|
||||
self.current_x = 0.0
|
||||
self.current_y = 0.0
|
||||
|
||||
# Данные для детекции
|
||||
self.current_frame = None
|
||||
self.keypoints = []
|
||||
self.matches = []
|
||||
|
||||
self._setup_window()
|
||||
|
||||
def _setup_window(self):
|
||||
"""Настраивает общее окно с несколькими областями"""
|
||||
plt.ion()
|
||||
self.fig = plt.figure(figsize=(16, 10))
|
||||
self.fig.canvas.manager.window.title(self.window_title)
|
||||
|
||||
# Открываем окно на полный экран
|
||||
self.fig.canvas.manager.window.state('zoomed')
|
||||
|
||||
# Создаем сетку 2x2
|
||||
gs = self.fig.add_gridspec(2, 2, hspace=0.3, wspace=0.3)
|
||||
|
||||
# Глобальная карта (левый верхний угол, занимает 2x1)
|
||||
self.ax_global_map = self.fig.add_subplot(gs[0, :])
|
||||
self.ax_global_map.set_title('Global Map - Траектория полета беспилотника')
|
||||
self.ax_global_map.set_xlabel('X координата')
|
||||
self.ax_global_map.set_ylabel('Y координата')
|
||||
self.ax_global_map.grid(True, alpha=0.3)
|
||||
self.ax_global_map.axhline(y=0, color='k', linestyle='-', alpha=0.3)
|
||||
self.ax_global_map.axvline(x=0, color='k', linestyle='-', alpha=0.3)
|
||||
|
||||
# Детекция ключевых точек (правый нижний угол)
|
||||
self.ax_detection = self.fig.add_subplot(gs[1, 0])
|
||||
self.ax_detection.set_title('Keypoint Detection')
|
||||
self.ax_detection.axis('off')
|
||||
|
||||
# Сопоставление точек (правый нижний угол)
|
||||
self.ax_matches = self.fig.add_subplot(gs[1, 1])
|
||||
self.ax_matches.set_title('Feature Matching')
|
||||
self.ax_matches.axis('off')
|
||||
|
||||
# Настройки окна
|
||||
self.fig.canvas.manager.window.attributes('-topmost', False)
|
||||
|
||||
plt.tight_layout()
|
||||
|
||||
def update_global_map(self, x: float, y: float, mode: SimMode):
|
||||
"""Обновляет глобальную карту"""
|
||||
self.current_x = x
|
||||
self.current_y = y
|
||||
self.trajectory_x.append(x)
|
||||
self.trajectory_y.append(y)
|
||||
self.trajectory_modes.append(mode)
|
||||
|
||||
self.ax_global_map.clear()
|
||||
self.ax_global_map.set_title('Global Map - Траектория полета беспилотника')
|
||||
self.ax_global_map.set_xlabel('X координата')
|
||||
self.ax_global_map.set_ylabel('Y координата')
|
||||
self.ax_global_map.grid(True, alpha=0.3)
|
||||
self.ax_global_map.axhline(y=0, color='k', linestyle='-', alpha=0.3)
|
||||
self.ax_global_map.axvline(x=0, color='k', linestyle='-', alpha=0.3)
|
||||
|
||||
if len(self.trajectory_x) > 1:
|
||||
# Разделяем траекторию по режимам
|
||||
operator_indices = [i for i, m in enumerate(self.trajectory_modes) if m == SimMode.OPERATOR]
|
||||
autonome_indices = [i for i, m in enumerate(self.trajectory_modes) if m == SimMode.AUTONOME]
|
||||
|
||||
# Рисуем траекторию оператора (синий цвет)
|
||||
if len(operator_indices) > 1:
|
||||
operator_x = [self.trajectory_x[i] for i in operator_indices]
|
||||
operator_y = [self.trajectory_y[i] for i in operator_indices]
|
||||
self.ax_global_map.plot(operator_x, operator_y, 'b-', linewidth=2, label='Режим оператора')
|
||||
|
||||
# Рисуем траекторию автономного режима (красный цвет)
|
||||
if len(autonome_indices) > 1:
|
||||
autonome_x = [self.trajectory_x[i] for i in autonome_indices]
|
||||
autonome_y = [self.trajectory_y[i] for i in autonome_indices]
|
||||
self.ax_global_map.plot(autonome_x, autonome_y, 'r-', linewidth=2, label='Автономный режим')
|
||||
|
||||
# Рисуем начальную точку (зеленая)
|
||||
self.ax_global_map.plot(self.trajectory_x[0], self.trajectory_y[0], 'go', markersize=8, label='Начальная точка')
|
||||
|
||||
# Рисуем текущую позицию (черная)
|
||||
self.ax_global_map.plot(self.current_x, self.current_y, 'ko', markersize=6, label='Текущая позиция')
|
||||
|
||||
# Рисуем целевую точку (0, 0) - желтая
|
||||
self.ax_global_map.plot(0, 0, 'yo', markersize=8, label='Цель (0, 0)')
|
||||
|
||||
self.ax_global_map.legend()
|
||||
|
||||
# Автоматически масштабируем оси
|
||||
if len(self.trajectory_x) > 0:
|
||||
margin = 50
|
||||
x_min, x_max = min(self.trajectory_x), max(self.trajectory_x)
|
||||
y_min, y_max = min(self.trajectory_y), max(self.trajectory_y)
|
||||
|
||||
x_min = min(x_min, 0)
|
||||
x_max = max(x_max, 0)
|
||||
y_min = min(y_min, 0)
|
||||
y_max = max(y_max, 0)
|
||||
|
||||
self.ax_global_map.set_xlim(x_min - margin, x_max + margin)
|
||||
self.ax_global_map.set_ylim(y_min - margin, y_max + margin)
|
||||
|
||||
def update_detection(self, image: np.ndarray, keypoints):
|
||||
"""Обновляет визуализацию детекции ключевых точек"""
|
||||
self.current_frame = image.copy()
|
||||
self.keypoints = keypoints
|
||||
|
||||
self.ax_detection.clear()
|
||||
self.ax_detection.set_title('Keypoint Detection')
|
||||
|
||||
if image is not None:
|
||||
# Конвертируем BGR в RGB для matplotlib
|
||||
if len(image.shape) == 3 and image.shape[2] == 3:
|
||||
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
else:
|
||||
image_rgb = image
|
||||
|
||||
self.ax_detection.imshow(image_rgb)
|
||||
|
||||
# Рисуем ключевые точки
|
||||
if keypoints:
|
||||
kp_coords = np.array([kp.pt for kp in keypoints])
|
||||
self.ax_detection.scatter(kp_coords[:, 0], kp_coords[:, 1],
|
||||
c='red', s=20, alpha=0.7, marker='o')
|
||||
|
||||
self.ax_detection.axis('off')
|
||||
|
||||
def update_matches(self, img1: np.ndarray, img2: np.ndarray,
|
||||
kp1, kp2, matches, transformation_info=None):
|
||||
"""Обновляет визуализацию сопоставления точек"""
|
||||
self.ax_matches.clear()
|
||||
self.ax_matches.set_title('Feature Matching')
|
||||
|
||||
if img1 is not None and img2 is not None and matches:
|
||||
# Рисуем сопоставления
|
||||
img_matches = cv2.drawMatches(img1, kp1, img2, kp2, matches, None,
|
||||
flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS)
|
||||
|
||||
# Конвертируем BGR в RGB
|
||||
if len(img_matches.shape) == 3 and img_matches.shape[2] == 3:
|
||||
img_matches_rgb = cv2.cvtColor(img_matches, cv2.COLOR_BGR2RGB)
|
||||
else:
|
||||
img_matches_rgb = img_matches
|
||||
|
||||
self.ax_matches.imshow(img_matches_rgb)
|
||||
|
||||
# Добавляем информацию о трансформации
|
||||
if transformation_info:
|
||||
tx, ty = transformation_info['translation']
|
||||
angle = transformation_info['rotation']
|
||||
scale = transformation_info['scale']
|
||||
|
||||
info_text = f"Translation: ({tx:.2f}, {ty:.2f})"
|
||||
info_text2 = f"Rotation: {angle:.2f} rad ({np.degrees(angle):.1f}°)"
|
||||
info_text3 = f"Scale: {scale:.2f}"
|
||||
|
||||
self.ax_matches.text(10, 30, info_text, fontsize=8, color='green',
|
||||
bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8))
|
||||
self.ax_matches.text(10, 90, info_text2, fontsize=8, color='green',
|
||||
bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8))
|
||||
self.ax_matches.text(10, 150, info_text3, fontsize=8, color='green',
|
||||
bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8))
|
||||
|
||||
self.ax_matches.axis('off')
|
||||
|
||||
def update_display(self):
|
||||
"""Обновляет отображение всех областей"""
|
||||
self.fig.canvas.draw()
|
||||
self.fig.canvas.flush_events()
|
||||
plt.pause(0.01)
|
||||
|
||||
def close(self):
|
||||
"""Закрывает окно"""
|
||||
plt.close(self.fig)
|
||||
|
||||
def show_final(self):
|
||||
"""Показывает финальное состояние окна"""
|
||||
plt.ioff()
|
||||
print("Симуляция завершена. Окно визуализации остается открытым для анализа.")
|
||||
|
||||
def pause(self, duration: float):
|
||||
plt.pause(duration)
|
||||
Reference in New Issue
Block a user