feat: add rth map

.gitignore

@@ -3,4 +3,5 @@ __pycache__
 trajectories
 z
 chunks/*
-images/*
+images/*
+test_runs/*

main.py

@@ -1,4 +1,5 @@
 from google_map import GoogleMap
+from datetime import datetime
 from pathlib import Path
 from position import Position
 from simulator import Simulator
@@ -17,6 +18,7 @@ import matplotlib.pyplot as plt
 import numpy as np
 import pickle
 import random
+import shutil
 import utility
 
 def get_map(map_name: str = 'google', lat=49.103814, lon=55.794258, zoom=18):
@@ -36,6 +38,28 @@ def get_trajectory_points(bg_img: str) -> list[(float, float)]:
     points = list(map(lambda p: [p[0] / trajectoryDrawer.img.shape[1], p[1] / trajectoryDrawer.img.shape[0]], trajectoryDrawer.points))
     return points
 
+def make_map_extent(width: float, height: float, origin_point: list[float], pixel_ratio: float) -> tuple[float, float, float, float]:
+    """Возвращает границы снимка карты в координатах маршрута."""
+    origin_x = origin_point[0] * width
+    origin_y = origin_point[1] * height
+    left = (0 - origin_x) * pixel_ratio
+    right = (width - origin_x) * pixel_ratio
+    bottom = (origin_y - height) * pixel_ratio
+    top = origin_y * pixel_ratio
+    return (left, right, bottom, top)
+
+def make_test_run_dir(name: str) -> Path:
+    run_dir = Path('test_runs') / f"{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}_{name}"
+    run_dir.mkdir(parents=True, exist_ok=True)
+    return run_dir
+
+def move_map_safely(online_map: YandexMap | GoogleMap, dx: float, dy: float, step: int = 300):
+    """Двигает карту короткими drag-шагами, чтобы Selenium не выходил за viewport."""
+    steps = max(1, int(np.ceil(max(abs(dx), abs(dy)) / step)))
+    for _ in range(steps):
+        online_map.move(dx / steps, dy / steps)
+        sleep(0.05)
+
 def build(name: str, map_name: str, lat: float, lon: float):
 
     # Создание папки с информацией о маршруте
@@ -48,7 +72,10 @@ def build(name: str, map_name: str, lat: float, lon: float):
     dir_chunks.mkdir()
 
     make_global_photo('map.jpg', map_name, lat, lon, 15)
-    points = get_trajectory_points('map.jpg')
+    map_path = dir / 'map.jpg'
+    shutil.copyfile('map.jpg', map_path)
+
+    points = get_trajectory_points(str(map_path))
     online_map: YandexMap | GoogleMap = get_map(map_name, lat, lon, 15)
 
 
@@ -106,7 +133,11 @@ def build(name: str, map_name: str, lat: float, lon: float):
     data = {
         'points': points_coords,
         'chunk_positions': positions,
-        'initial_geolocation': geo
+        'initial_geolocation': geo,
+        'map_image': 'map.jpg',
+        'map_size': (width, height),
+        'map_zero_point': points[0],
+        'map_extent': make_map_extent(width, height, points[0], online_map.pixel_ratio),
     }
 
     print(points_coords)
@@ -160,10 +191,26 @@ def run(
         )
     print("R: ", r)
 
-    points = data['points'] / online_map.pixel_ratio
-    print("READ POINTS:", points)
+    forward_points = data['points']
+    return_points = forward_points[::-1].copy()
+    points = return_points / online_map.pixel_ratio
+    print("READ RTH POINTS:", points)
 
     vis_manager = VisualizationManager()
+    map_path = dir / data.get('map_image', 'map.jpg')
+    if map_path.exists() and 'map_extent' in data:
+        turn_point = data['chunk_positions'][-1] if data.get('chunk_positions') else forward_points[-1]
+        build_trajectory = np.array([[pos.x, pos.y] for pos in data.get('chunk_positions', [])])
+        if len(build_trajectory) == 0:
+            build_trajectory = forward_points
+        vis_manager.set_route_map(
+            map_path,
+            data['map_extent'],
+            build_trajectory,
+            turn_point,
+            forward_points[0],
+        )
+
     pilot = autopilot.AutoPilot(
         points,
         chunks,
@@ -175,17 +222,27 @@ def run(
     simulator = Simulator(online_map)
     pilot.target_idx = 0
 
+    turn_position = data['chunk_positions'][-1].copy() if data.get('chunk_positions') else Position(forward_points[-1][0], forward_points[-1][1])
+    start_position = turn_position / online_map.pixel_ratio
+    start_position.yaw += np.pi
+    simulator.pos = start_position.copy()
+    pilot.pos = start_position.copy()
+    move_map_safely(online_map, start_position.x, start_position.y)
+    vis_manager.update_rth_trajectory(turn_position.x, turn_position.y)
+
     chunk = simulator.get_chunk()
     command = pilot.handle(chunk)
 
     vis_manager.update_display()
     vis_manager.pause(1)
 
-    vis_manager.set_target_points(data['points'])
+    vis_manager.set_target_points(return_points)
 
     proc_time = np.array([])
 
     errors = []
+    start_point = forward_points[0]
+    final_return_error = None
     
     sleep(1)
 
@@ -213,6 +270,7 @@ def run(
         vis_manager.set_target_index(pilot.target_idx)
         vis_manager.update_drone_trajectory(pilot.pos.x * online_map.pixel_ratio, pilot.pos.y * online_map.pixel_ratio)
         vis_manager.update_global_map(simulator.pos.x * online_map.pixel_ratio, simulator.pos.y * online_map.pixel_ratio)
+        vis_manager.update_rth_trajectory(simulator.pos.x * online_map.pixel_ratio, simulator.pos.y * online_map.pixel_ratio)
         vis_manager.update_error_plot(i, pilot.pos.x * online_map.pixel_ratio, pilot.pos.y * online_map.pixel_ratio, simulator.pos.x * online_map.pixel_ratio, simulator.pos.y * online_map.pixel_ratio)
 
         errors.append(np.hypot((pilot.pos.x - simulator.pos.x) * online_map.pixel_ratio, (pilot.pos.y - simulator.pos.y) * online_map.pixel_ratio))
@@ -227,14 +285,27 @@ def run(
         print("Simulator coords:", simulator.pos)
         sleep(0.5)
         simulator.handle(command.dangle, command.velocity)
-        if i == 0 and map_name == 'google':
-            simulator.pos.x = 0
-            simulator.pos.y = 0
 
+    final_point = np.array([simulator.pos.x * online_map.pixel_ratio, simulator.pos.y * online_map.pixel_ratio])
+    final_return_error = float(np.hypot(final_point[0] - start_point[0], final_point[1] - start_point[1]))
+    vis_manager.set_final_point(final_point[0], final_point[1], final_return_error)
     print("Errors:", errors)
     print("MSE:", (np.array(errors) ** 2).mean())
     print("RMSE:", (np.array(errors) ** 2).mean() ** 0.5)
+    print("Return error:", final_return_error)
     print("Average FPS:", 1 / proc_time.mean())
+    test_run_dir = make_test_run_dir(name)
+    vis_manager.save_plots(test_run_dir)
+    with (test_run_dir / 'metrics.pkl').open('wb') as file:
+        pickle.dump({
+            'position_errors': errors,
+            'mse': float((np.array(errors) ** 2).mean()) if errors else None,
+            'rmse': float((np.array(errors) ** 2).mean() ** 0.5) if errors else None,
+            'return_error': final_return_error,
+            'start_point': start_point,
+            'turn_point': np.array([turn_position.x, turn_position.y]),
+            'final_point': final_point,
+        }, file)
     vis_manager.show_final()
 
 def parse_args():

visualization.py

@@ -3,9 +3,10 @@
 Модуль для управления общим окном визуализации
 """
 
-from PIL import Image
-from enum import Enum
-from scipy.interpolate import make_interp_spline
+from PIL import Image
+from enum import Enum
+from pathlib import Path
+from scipy.interpolate import make_interp_spline
 
 import cv2
 import matplotlib
@@ -31,8 +32,9 @@ class VisualizationManager:
         self.window_title = window_title
         self.fig = None
         self.ax_error_plot = None  # График погрешности позиции
-        self.ax_global_map = None
-        self.ax_detection = None
+        self.ax_global_map = None
+        self.ax_route_map = None
+        self.ax_detection = None
         self.ax_matches = None
         self.ax_chunk_matches = None
         self.ax_motion_vectors = None
@@ -48,7 +50,18 @@ class VisualizationManager:
         
         # Данные для траектории БПЛА (его собственное видение)
         self.drone_trajectory_x = []
-        self.drone_trajectory_y = []
+        self.drone_trajectory_y = []
+
+        # Данные для RTH-карты
+        self.map_image = None
+        self.map_extent = None
+        self.forward_route = None
+        self.turn_point = None
+        self.start_point = None
+        self.final_point = None
+        self.final_return_error = None
+        self.rth_trajectory_x = []
+        self.rth_trajectory_y = []
         
         # Данные для графика погрешности
         self.error_times = []
@@ -95,9 +108,15 @@ class VisualizationManager:
         self.ax_matches.axis('off')
         
         # Сопоставление точек (средний средний угол)
-        self.ax_chunk_matches = self.fig.add_subplot(gs[1, 1:3])
-        self.ax_chunk_matches.set_title('Chunk Matching')
-        self.ax_chunk_matches.axis('off')
+        self.ax_route_map = self.fig.add_subplot(gs[1, 1])
+        self.ax_route_map.set_title('RTH Map - маршрут на карте')
+        self.ax_route_map.set_xlabel('X координата')
+        self.ax_route_map.set_ylabel('Y координата')
+        self.ax_route_map.grid(True, alpha=0.3)
+
+        self.ax_chunk_matches = self.fig.add_subplot(gs[1, 2])
+        self.ax_chunk_matches.set_title('Chunk Matching')
+        self.ax_chunk_matches.axis('off')
         
         # Визуализация движения ключевых точек (левый нижний угол)
         # self.ax_motion_vectors = self.fig.add_subplot(gs[1, 1])
@@ -112,16 +131,104 @@ class VisualizationManager:
         # Настройки окна
         self.fig.canvas.manager.window.attributes('-topmost', False)
         
-        plt.tight_layout()
-        plt.show(block=False)
+        self.fig.subplots_adjust(left=0.06, right=0.98, bottom=0.06, top=0.94, hspace=0.3, wspace=0.3)
+        plt.show(block=False)
 
-    def set_target_points(self, target_pts):
-        """ Обновление списка координат целевых точек """
-        self.target_pts = target_pts
+    def set_target_points(self, target_pts):
+        """ Обновление списка координат целевых точек """
+        self.target_pts = target_pts
 
-    def set_target_index(self, target_idx):
-        """ Обновление номера целевой точки """
-        self.target_idx = target_idx
+    def set_target_index(self, target_idx):
+        """ Обновление номера целевой точки """
+        self.target_idx = target_idx
+
+    def set_route_map(self, map_path, map_extent, forward_route, turn_point, start_point):
+        """Настраивает отдельный график маршрута на фоне карты."""
+        self.map_image = np.array(Image.open(map_path))
+        self.map_extent = map_extent
+        self.forward_route = np.array(forward_route)
+        self.turn_point = np.array([turn_point.x, turn_point.y]) if hasattr(turn_point, 'x') else np.array(turn_point)
+        self.start_point = np.array(start_point)
+        self._draw_route_map()
+
+    def update_rth_trajectory(self, x: float, y: float):
+        """Добавляет точку обратного полета на карту RTH."""
+        self.rth_trajectory_x.append(x)
+        self.rth_trajectory_y.append(y)
+        self._draw_route_map()
+
+    def set_final_point(self, x: float, y: float, return_error: float | None = None):
+        """Фиксирует последнюю точку и итоговую ошибку возврата."""
+        self.final_point = np.array([x, y])
+        self.final_return_error = return_error
+        self._draw_route_map()
+
+    def _draw_route_map(self):
+        if self.ax_route_map is None:
+            return
+
+        self.ax_route_map.clear()
+        self.ax_route_map.set_title('RTH Map - маршрут на карте')
+        self.ax_route_map.set_xlabel('X координата')
+        self.ax_route_map.set_ylabel('Y координата')
+        self.ax_route_map.grid(True, alpha=0.3)
+
+        if self.map_image is not None and self.map_extent is not None:
+            self.ax_route_map.imshow(self.map_image, extent=self.map_extent, origin='upper', alpha=0.85)
+
+        all_points = []
+
+        if self.forward_route is not None and len(self.forward_route) > 0:
+            self.ax_route_map.plot(
+                self.forward_route[:, 0],
+                self.forward_route[:, 1],
+                color='red',
+                linewidth=2,
+                label='Маршрут туда',
+            )
+            all_points.append(self.forward_route)
+
+        if len(self.rth_trajectory_x) > 0:
+            rth_points = np.column_stack([self.rth_trajectory_x, self.rth_trajectory_y])
+            self.ax_route_map.plot(
+                self.rth_trajectory_x,
+                self.rth_trajectory_y,
+                color='gold',
+                linewidth=2,
+                label='Маршрут обратно',
+            )
+            all_points.append(rth_points)
+
+        if self.start_point is not None:
+            self.ax_route_map.plot(self.start_point[0], self.start_point[1], 'go', markersize=8, label='Старт')
+            all_points.append(self.start_point.reshape(1, 2))
+
+        if self.turn_point is not None:
+            self.ax_route_map.plot(self.turn_point[0], self.turn_point[1], 'rx', markersize=10, markeredgewidth=2.5, label='RTH')
+            all_points.append(self.turn_point.reshape(1, 2))
+
+        if self.final_point is not None:
+            self.ax_route_map.plot(self.final_point[0], self.final_point[1], 'bo', markersize=8, label='Финиш')
+            all_points.append(self.final_point.reshape(1, 2))
+
+        if self.final_return_error is not None:
+            self.ax_route_map.text(
+                0.02,
+                0.98,
+                f"Ошибка возврата: {self.final_return_error:.2f}",
+                transform=self.ax_route_map.transAxes,
+                va='top',
+                fontsize=8,
+                bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8),
+            )
+
+        if all_points:
+            points = np.vstack(all_points)
+            margin = 50
+            self.ax_route_map.set_xlim(points[:, 0].min() - margin, points[:, 0].max() + margin)
+            self.ax_route_map.set_ylim(points[:, 1].min() - margin, points[:, 1].max() + margin)
+
+        self.ax_route_map.legend(loc='best')
     
     def update_global_map(self, x: float, y: float):
         """Обновляет глобальную карту"""
@@ -573,11 +680,33 @@ class VisualizationManager:
         """Закрывает окно"""
         plt.close(self.fig)
     
-    def show_final(self):
-        """Показывает финальное состояние окна"""
-        plt.ioff()
-        print("Симуляция завершена. Окно визуализации остается открытым для анализа.")
-        plt.pause(100000)
-
-    def pause(self, duration: float):
-        plt.pause(duration) 
+    def show_final(self):
+        """Показывает финальное состояние окна"""
+        plt.ioff()
+        print("Симуляция завершена. Окно визуализации остается открытым для анализа.")
+        plt.pause(100000)
+
+    def pause(self, duration: float):
+        plt.pause(duration) 
+
+    def save_plots(self, output_dir: Path | str):
+        """Сохраняет итоговое окно и отдельные графики в папку тестового запуска."""
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+        self.fig.canvas.draw()
+        self.fig.savefig(output_dir / 'visualization.png', dpi=150, bbox_inches='tight')
+
+        axes = {
+            'error_plot.png': self.ax_error_plot,
+            'global_map.png': self.ax_global_map,
+            'rth_map.png': self.ax_route_map,
+            'keypoint_detection.png': self.ax_motion_gomography,
+            'feature_matching.png': self.ax_matches,
+            'chunk_matching.png': self.ax_chunk_matches,
+        }
+        renderer = self.fig.canvas.get_renderer()
+        for filename, axes_item in axes.items():
+            if axes_item is None:
+                continue
+            bbox = axes_item.get_tightbbox(renderer).expanded(1.08, 1.15)
+            self.fig.savefig(output_dir / filename, dpi=150, bbox_inches=bbox)