248 lines
11 KiB
Python
248 lines
11 KiB
Python
#!/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_drone_view = 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.drone_view_image = None
|
||
|
||
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, width_ratios=[1, 1])
|
||
|
||
# Вид БПЛА (левый верхний угол)
|
||
self.ax_drone_view = self.fig.add_subplot(gs[0, 0])
|
||
self.ax_drone_view.set_title('Вид БПЛА')
|
||
self.ax_drone_view.axis('off')
|
||
|
||
# Глобальная карта (левый нижний угол)
|
||
self.ax_global_map = self.fig.add_subplot(gs[1, 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[0, 1])
|
||
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 = operator_indices[-1:] + [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_drone_view(self, image: np.ndarray):
|
||
"""Обновляет вид БПЛА"""
|
||
self.drone_view_image = image.copy()
|
||
|
||
self.ax_drone_view.clear()
|
||
self.ax_drone_view.set_title('Вид БПЛА')
|
||
|
||
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_drone_view.imshow(image_rgb)
|
||
|
||
self.ax_drone_view.axis('off')
|
||
|
||
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) |