diff --git a/autopilot.py b/autopilot.py index 8992e6e..83b7e6c 100644 --- a/autopilot.py +++ b/autopilot.py @@ -19,13 +19,17 @@ class AutoPilot(Pilot): orb_detector: cv2.ORB bf_matcher: cv2.BFMatcher frame_count: int + image_center: tuple # Центр изображения (x, y) + viz_manager: object # Менеджер визуализации (опционально) - def __init__(self, initial_x: float = 0.0, initial_y: float = 0.0): + def __init__(self, viz_manager=None): self.prev_image = None self.angle = 0.0 - self.x = initial_x - self.y = initial_y + self.x = 0.0 + self.y = 0.0 self.frame_count = 0 + self.image_center = (0, 0) # Будет обновлено при получении первого изображения + self.viz_manager = viz_manager # Менеджер визуализации # Инициализация ORB детектора self.orb_detector = cv2.ORB_create( @@ -72,15 +76,35 @@ class AutoPilot(Pilot): if len(good_matches) < 4: return None, None, None, None, None - # Извлечение координат сопоставленных точек - src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2) - dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2) + # Получаем центр изображения + height1, width1 = img1.shape[:2] + height2, width2 = img2.shape[:2] + center_x1, center_y1 = width1 // 2, height1 // 2 + center_x2, center_y2 = width2 // 2, height2 // 2 + + # Извлекаем координаты сопоставленных точек и отцентрируем их + src_pts = [] + dst_pts = [] + + for match in good_matches: + # Координаты точки в первом изображении + pt1 = kp1[match.queryIdx].pt + src_pts.append([pt1[0] - center_x1, pt1[1] - center_y1]) + + # Координаты точки во втором изображении + pt2 = kp2[match.trainIdx].pt + dst_pts.append([pt2[0] - center_x2, pt2[1] - center_y2]) + + # Конвертируем в numpy массивы + src_pts = np.float32(src_pts).reshape(-1, 1, 2) + dst_pts = np.float32(dst_pts).reshape(-1, 1, 2) return src_pts, dst_pts, good_matches, kp1, kp2 def estimate_transformation_matrix(self, src_pts: np.ndarray, dst_pts: np.ndarray): """ Оценивает матрицу трансформации на основе сопоставленных ключевых точек + Точки уже отцентрированы относительно центра изображения """ # Используем RANSAC для оценки матрицы гомографии H, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0) @@ -97,7 +121,7 @@ class AutoPilot(Pilot): a11, a12 = H[0, 0], H[0, 1] a21, a22 = H[1, 0], H[1, 1] - # Смещение + # Смещение (уже отцентрировано) tx, ty = H[0, 2], H[1, 2] # Вычисляем угол поворота @@ -109,7 +133,7 @@ class AutoPilot(Pilot): scale = (scale_x + scale_y) / 2 return { - 'translation': (tx, ty), + 'translation': (tx, ty), # Уже отцентрировано 'rotation': angle, 'scale': scale, 'homography': H, @@ -119,12 +143,13 @@ class AutoPilot(Pilot): def update_drone_position(self, transformation_info: dict): """ Обновляет позицию и угол БПЛА на основе трансформации изображения + Координаты уже отцентрированы относительно центра изображения """ tx, ty = transformation_info['translation'] rotation = transformation_info['rotation'] scale = transformation_info['scale'] - # Конвертируем смещение в пикселях в метры + # Координаты уже отцентрированы, поэтому используем их напрямую dx_meters = tx dy_meters = ty @@ -133,8 +158,9 @@ class AutoPilot(Pilot): sin_angle = math.sin(self.angle) # Поворачиваем смещение в глобальные координаты - dx_global = dx_meters * cos_angle - dy_meters * sin_angle - dy_global = dx_meters * sin_angle + dy_meters * cos_angle + # Обратите внимание: dy_meters инвертирован, так как в изображениях Y направлен вниз + dx_global = dx_meters * cos_angle - (-dy_meters) * sin_angle + dy_global = dx_meters * sin_angle + (-dy_meters) * cos_angle # Обновляем координаты БПЛА self.x += dx_global @@ -176,18 +202,20 @@ class AutoPilot(Pilot): 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}" + info_text4 = f"Image Center: {self.image_center}" cv2.putText(img_matches, info_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) cv2.putText(img_matches, info_text2, (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) cv2.putText(img_matches, info_text3, (10, 90), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) + cv2.putText(img_matches, info_text4, (10, 120), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) # Добавляем информацию о позиции БПЛА drone_state = self.get_drone_state() pos_text = f"Drone Pos: ({drone_state['x']:.2f}, {drone_state['y']:.2f})" angle_text = f"Drone Angle: {drone_state['angle_degrees']:.1f}°" - cv2.putText(img_matches, pos_text, (10, 120), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2) - cv2.putText(img_matches, angle_text, (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2) + cv2.putText(img_matches, pos_text, (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2) + cv2.putText(img_matches, angle_text, (10, 180), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2) return img_matches @@ -196,11 +224,24 @@ class AutoPilot(Pilot): if self.prev_image is None: self.prev_image = self.image_to_numpy(image) + # Вычисляем центр изображения + height, width = self.prev_image.shape[:2] + self.image_center = (width // 2, height // 2) + + # Обновляем визуализацию детекции + if self.viz_manager: + kp, _ = self.orb_detector.detectAndCompute(self.prev_image, None) + self.viz_manager.update_detection(self.prev_image, kp) + return # Конвертируем текущее изображение current_image = self.image_to_numpy(image) + # Обновляем центр изображения + height, width = current_image.shape[:2] + self.image_center = (width // 2, height // 2) + # Обнаруживаем и сопоставляем ключевые точки src_pts, dst_pts, matches, kp1, kp2 = self.detect_and_match_keypoints(self.prev_image, current_image) @@ -209,11 +250,6 @@ class AutoPilot(Pilot): transformation_info = self.estimate_transformation_matrix(src_pts, dst_pts) if transformation_info: - # print(f"Frame {self.frame_count}:") - # print(f" Translation: {transformation_info['translation']}") - # print(f" Rotation: {transformation_info['rotation']:.4f} rad") - # print(f" Scale: {transformation_info['scale']:.4f}") - # Обновляем позицию и угол БПЛА self.update_drone_position(transformation_info) @@ -222,11 +258,20 @@ class AutoPilot(Pilot): print(f" [Pilot] Drone Position: ({drone_state['x']:.2f}, {drone_state['y']:.2f})") print(f" [Pilot] Angle: {drone_state['angle_degrees']:.1f}°") - # Визуализация (опционально) - img_matches = self.visualize_matches(self.prev_image, current_image, + # Обновляем визуализацию + if self.viz_manager: + # Обновляем детекцию ключевых точек + self.viz_manager.update_detection(current_image, kp2) + # Обновляем сопоставление точек + self.viz_manager.update_matches(self.prev_image, current_image, kp1, kp2, matches, transformation_info) - cv2.imshow('Drone Tracking', img_matches) - cv2.waitKey(1) + + # Визуализация (опционально, если нет менеджера визуализации) + if not self.viz_manager: + img_matches = self.visualize_matches(self.prev_image, current_image, + kp1, kp2, matches, transformation_info) + cv2.imshow('Drone Tracking', img_matches) + cv2.waitKey(1) # Обновляем предыдущее изображение self.prev_image = current_image diff --git a/main.py b/main.py index ecb1a09..035720b 100644 --- a/main.py +++ b/main.py @@ -1,6 +1,13 @@ from autopilot import AutoPilot, RandomPilot from simulator import Simulator +from visualization import VisualizationManager + +# Создаем менеджер визуализации +viz_manager = VisualizationManager("Drone Autopilot - Global Map & Detection") # Создаем симулятор с AutoPilot для обработки изображений -simulator = Simulator(RandomPilot(), AutoPilot()) +# Передаем менеджер визуализации в автопилот +simulator = Simulator(RandomPilot(), AutoPilot(viz_manager=viz_manager), viz_manager=viz_manager) + +# Запускаем симуляцию simulator.loop() diff --git a/simulator.py b/simulator.py index c251fd1..5d1e9d1 100644 --- a/simulator.py +++ b/simulator.py @@ -8,16 +8,10 @@ from selenium.webdriver.common.by import By from selenium.webdriver.common.action_chains import ActionChains from autopilot import Pilot +from visualization import VisualizationManager, SimMode -from enum import Enum -import matplotlib.pyplot as plt -import numpy as np import os -class SimMode(Enum): - OPERATOR = 1 - AUTONOME = 2 - class Simulator: driver: webdriver.Chrome mode: SimMode @@ -27,42 +21,32 @@ class Simulator: angle: float - # Переменные для отслеживания траектории - trajectory_x: list - trajectory_y: list - trajectory_modes: list # Для отслеживания режимов полета + # Менеджер визуализации + viz_manager: VisualizationManager current_x: float current_y: float - fig: plt.Figure - ax: plt.Axes - def __init__(self, operatorPilot: Pilot, autonomePilot: Pilot): + def __init__(self, operatorPilot: Pilot, autonomePilot: Pilot, viz_manager: VisualizationManager = None): self.mode = SimMode.OPERATOR self.operatorPilot = operatorPilot self.autonomePilot = autonomePilot # Инициализация переменных для траектории - self.trajectory_x = [] - self.trajectory_y = [] - self.trajectory_modes = [] self.current_x = 0.0 self.current_y = 0.0 - # Создание окна для карты - plt.ion() # Включаем интерактивный режим - self.fig, self.ax = plt.subplots(figsize=(10, 8)) - self.ax.set_title('Global Map - Траектория полета беспилотника') - self.ax.set_xlabel('X координата') - self.ax.set_ylabel('Y координата') - self.ax.grid(True, alpha=0.3) - self.ax.axhline(y=0, color='k', linestyle='-', alpha=0.3) - self.ax.axvline(x=0, color='k', linestyle='-', alpha=0.3) + # Создаем менеджер визуализации + self.viz_manager = viz_manager + + # Передаем менеджер визуализации в автопилот, если он поддерживает это + if hasattr(self.autonomePilot, 'viz_manager'): + self.autonomePilot.viz_manager = self.viz_manager # Создаем папку для изображений, если её нет os.makedirs('./images', exist_ok=True) options = webdriver.ChromeOptions() - options.add_experimental_option("detach", True) + # options.add_experimental_option("detach", True) self.driver = webdriver.Chrome(options) self.driver.get("https://yandex.ru/maps/43/kazan/?ll=49.103814%2C55.794258&z=14") sleep(2) @@ -112,69 +96,12 @@ class Simulator: # Обновляем текущие координаты self.current_x += dx self.current_y += dy - - # Добавляем точку в траекторию - self.trajectory_x.append(self.current_x) - self.trajectory_y.append(self.current_y) - self.trajectory_modes.append(self.mode) def update_map(self): """ Обновляет карту траектории полета """ - self.ax.clear() - self.ax.set_title('Global Map - Траектория полета беспилотника') - self.ax.set_xlabel('X координата') - self.ax.set_ylabel('Y координата') - self.ax.grid(True, alpha=0.3) - self.ax.axhline(y=0, color='k', linestyle='-', alpha=0.3) - self.ax.axvline(x=0, color='k', linestyle='-', alpha=0.3) - - if len(self.trajectory_x) > 1: - # Разделяем траекторию по режимам - operator_indices = [i for i, mode in enumerate(self.trajectory_modes) if mode == SimMode.OPERATOR] - autonome_indices = [i for i, mode in enumerate(self.trajectory_modes) if mode == 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.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.plot(autonome_x, autonome_y, 'r-', linewidth=2, label='Автономный режим') - - # Рисуем начальную точку (зеленая) - self.ax.plot(self.trajectory_x[0], self.trajectory_y[0], 'go', markersize=8, label='Начальная точка') - - # Рисуем текущую позицию (черная) - self.ax.plot(self.current_x, self.current_y, 'ko', markersize=6, label='Текущая позиция') - - # Рисуем целевую точку (0, 0) - желтая - self.ax.plot(0, 0, 'yo', markersize=8, label='Цель (0, 0)') - - self.ax.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) - - # Учитываем целевую точку (0, 0) - 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.set_xlim(x_min - margin, x_max + margin) - self.ax.set_ylim(y_min - margin, y_max + margin) - - plt.draw() - plt.pause(0.25) # Небольшая пауза для обновления окна + self.viz_manager.update_global_map(self.current_x, self.current_y, self.mode) def loop(self): @@ -184,7 +111,7 @@ class Simulator: # Добавляем начальную точку в траекторию self.update_trajectory(0, 0) - sleep(2) + self.viz_manager.update_global_map(self.current_x, self.current_y, self.mode) for i in range(1000): dangle = None @@ -200,15 +127,11 @@ class Simulator: if dangle is None: break - # Сдвиг камеры action = ActionChains(self.driver) action.move_to_element_with_offset(html, 200, 200) action.click_and_hold() - print(f" [Simulator] Drone Position: ({self.current_x:.2f}, {self.current_y:.2f})") - print(f" [Simulator] Angle: {self.angle * 180 / math.pi:.1f}°") - self.angle += dangle dx = math.cos(self.angle) * velocity dy = math.sin(self.angle) * velocity @@ -232,17 +155,19 @@ class Simulator: rotated_im.save(f"./images/{i}.png") - # Обновляем карту каждые несколько итераций для производительности + # Обновляем визуализацию каждые несколько итераций для производительности if i % 5 == 0: self.update_map() + self.viz_manager.update_display() - + self.viz_manager.pause(0.25) # Финальное обновление карты self.update_map() + self.viz_manager.update_display() print("last position: ", self.driver.current_url) print(f"Финальная позиция: ({self.current_x:.2f}, {self.current_y:.2f})") - # Показываем карту до закрытия - plt.ioff() - plt.show() + # Показываем карту до закрытия, но не поднимаем на передний план + self.viz_manager.show_final() + print("Симуляция завершена. Окно визуализации остается открытым для анализа.") diff --git a/visualization.py b/visualization.py new file mode 100644 index 0000000..f5f44e0 --- /dev/null +++ b/visualization.py @@ -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) \ No newline at end of file