249 lines
10 KiB
Python
249 lines
10 KiB
Python
import math
|
||
from io import BytesIO
|
||
from time import sleep
|
||
|
||
from PIL import Image
|
||
from selenium import webdriver
|
||
from selenium.webdriver.common.by import By
|
||
from selenium.webdriver.common.action_chains import ActionChains
|
||
|
||
from autopilot import Pilot
|
||
|
||
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
|
||
|
||
operatorPilot: Pilot
|
||
autonomePilot: Pilot
|
||
|
||
angle: float
|
||
|
||
# Переменные для отслеживания траектории
|
||
trajectory_x: list
|
||
trajectory_y: list
|
||
trajectory_modes: list # Для отслеживания режимов полета
|
||
current_x: float
|
||
current_y: float
|
||
fig: plt.Figure
|
||
ax: plt.Axes
|
||
|
||
def __init__(self, operatorPilot: Pilot, autonomePilot: Pilot):
|
||
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)
|
||
|
||
# Создаем папку для изображений, если её нет
|
||
os.makedirs('./images', exist_ok=True)
|
||
|
||
options = webdriver.ChromeOptions()
|
||
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)
|
||
|
||
action = ActionChains(self.driver)
|
||
|
||
# Закрытие левой панели
|
||
action.click(self.driver.find_element(By.CLASS_NAME, 'sidebar-toggle-button'))
|
||
action.move_to_element_with_offset(self.driver.find_element(By.XPATH, "//div[@class='rounded-controls']/div[@class='rounded-controls__child'][5]//button"), 5, 5)
|
||
action.perform()
|
||
|
||
# Режим спутника
|
||
action.click(self.driver.find_element(By.CLASS_NAME, '_key_satellite'))
|
||
action.perform()
|
||
|
||
self.angle = 0
|
||
|
||
def rotate_image_like_drone(self, image: Image.Image, angle: float) -> Image.Image:
|
||
"""
|
||
Поворачивает картинку как будто съемка ведется с летящего дрона.
|
||
Выделяет концентрический квадрат, поворачивает его и извлекает результат.
|
||
"""
|
||
# Получаем размеры изображения
|
||
width, height = image.size
|
||
square_size = min(width, height)
|
||
cropped_image = image.crop((0, 0, square_size, square_size))
|
||
cropped_image = cropped_image.rotate(angle / math.pi * 180, expand=True)
|
||
|
||
# Определяем размер концентрического квадрата (80% от минимальной стороны)
|
||
local_square_size = int(square_size / 2 ** 0.5)
|
||
|
||
# Вычисляем координаты для центрирования квадрата
|
||
left = (cropped_image.width - local_square_size) // 2
|
||
top = (cropped_image.height - local_square_size) // 2
|
||
right = left + local_square_size
|
||
bottom = top + local_square_size
|
||
|
||
# Вырезаем концентрический квадрат
|
||
final_image = cropped_image.crop((left, top, right, bottom))
|
||
|
||
return final_image
|
||
|
||
def update_trajectory(self, dx: float, dy: float):
|
||
"""
|
||
Обновляет траекторию полета беспилотника
|
||
"""
|
||
# Обновляем текущие координаты
|
||
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) # Небольшая пауза для обновления окна
|
||
|
||
def loop(self):
|
||
|
||
html = self.driver.find_element(By.TAG_NAME, 'html')
|
||
action = ActionChains(self.driver)
|
||
velocity = 10
|
||
|
||
# Добавляем начальную точку в траекторию
|
||
self.update_trajectory(0, 0)
|
||
sleep(2)
|
||
|
||
for i in range(1000):
|
||
dangle = None
|
||
if self.mode == SimMode.OPERATOR:
|
||
dangle = self.operatorPilot.act()
|
||
if dangle is None:
|
||
self.mode = SimMode.AUTONOME
|
||
print("Режим возвращения домой!")
|
||
|
||
if self.mode == SimMode.AUTONOME:
|
||
dangle = self.autonomePilot.act()
|
||
|
||
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
|
||
action.move_by_offset(-dx, dy)
|
||
action.release()
|
||
action.perform()
|
||
|
||
# Обновляем траекторию
|
||
self.update_trajectory(dx, dy)
|
||
|
||
# Загрузка скриншота
|
||
png = self.driver.get_screenshot_as_png()
|
||
im = Image.open(BytesIO(png))
|
||
im = im.crop([0, 80, im.width-80, im.height-60])
|
||
|
||
# Применяем поворот как будто съемка с дрона
|
||
rotated_im = self.rotate_image_like_drone(im, self.angle + math.pi / 2)
|
||
|
||
# Передаем изображение в AutoPilot для анализа
|
||
self.autonomePilot.handle(rotated_im)
|
||
|
||
rotated_im.save(f"./images/{i}.png")
|
||
|
||
# Обновляем карту каждые несколько итераций для производительности
|
||
if i % 5 == 0:
|
||
self.update_map()
|
||
|
||
|
||
|
||
# Финальное обновление карты
|
||
self.update_map()
|
||
print("last position: ", self.driver.current_url)
|
||
print(f"Финальная позиция: ({self.current_x:.2f}, {self.current_y:.2f})")
|
||
|
||
# Показываем карту до закрытия
|
||
plt.ioff()
|
||
plt.show()
|