zenloger/autopilot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add shared window

Browse Source
This commit is contained in:
russian_proger
2025-06-30 01:33:09 +03:00
parent 8b6bb3e22c
commit 2d85d78def
4 changed files with 316 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files

89
autopilot.py
View File

@@ -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

9
main.py
View File

@@ -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()

115
simulator.py
View File

@@ -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("Симуляция завершена. Окно визуализации остается открытым для анализа.")

221
visualization.py Normal file
View 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)