174 lines
7.2 KiB
Python
174 lines
7.2 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 visualization import VisualizationManager, SimMode
|
||
|
||
import os
|
||
|
||
class Simulator:
|
||
driver: webdriver.Chrome
|
||
mode: SimMode
|
||
|
||
operatorPilot: Pilot
|
||
autonomePilot: Pilot
|
||
|
||
angle: float
|
||
|
||
# Менеджер визуализации
|
||
viz_manager: VisualizationManager
|
||
current_x: float
|
||
current_y: float
|
||
|
||
def __init__(self, operatorPilot: Pilot, autonomePilot: Pilot, viz_manager: VisualizationManager = None):
|
||
self.mode = SimMode.OPERATOR
|
||
self.operatorPilot = operatorPilot
|
||
self.autonomePilot = autonomePilot
|
||
|
||
# Инициализация переменных для траектории
|
||
self.current_x = 0.0
|
||
self.current_y = 0.0
|
||
|
||
# Создаем менеджер визуализации
|
||
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)
|
||
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()
|
||
|
||
# Режим спутника
|
||
sleep(1)
|
||
action.move_to_element_with_offset(self.driver.find_element(By.XPATH, "//div[@class='rounded-controls']/div[@class='rounded-controls__child'][5]//button"), -500, -500)
|
||
action.click()
|
||
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
|
||
|
||
def update_map(self):
|
||
"""
|
||
Обновляет карту траектории полета
|
||
"""
|
||
self.viz_manager.update_global_map(self.current_x, self.current_y, self.mode)
|
||
|
||
def loop(self):
|
||
|
||
html = self.driver.find_element(By.TAG_NAME, 'html')
|
||
action = ActionChains(self.driver)
|
||
|
||
# Добавляем начальную точку в траекторию
|
||
self.update_trajectory(0, 0)
|
||
self.viz_manager.update_global_map(self.current_x, self.current_y, self.mode)
|
||
|
||
for i in range(1000):
|
||
signal = None
|
||
if self.mode == SimMode.OPERATOR:
|
||
signal = self.operatorPilot.act()
|
||
if signal is None:
|
||
self.mode = SimMode.AUTONOME
|
||
print("Режим возвращения домой!")
|
||
|
||
if self.mode == SimMode.AUTONOME:
|
||
signal = self.autonomePilot.act()
|
||
|
||
if signal is None:
|
||
break
|
||
|
||
dangle, velocity = signal
|
||
drone_x, drone_y = self.autonomePilot.get_position()
|
||
self.viz_manager.update_error_plot(i, drone_x, drone_y, self.current_x, self.current_y)
|
||
|
||
# Сдвиг камеры
|
||
action = ActionChains(self.driver)
|
||
action.move_to_element_with_offset(html, 200, 200)
|
||
action.click_and_hold()
|
||
|
||
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, math.pi / 2 - self.angle)
|
||
|
||
# Передаем изображение в AutoPilot для анализа
|
||
self.autonomePilot.handle(rotated_im)
|
||
|
||
# Обновляем визуализацию каждые несколько итераций для производительности
|
||
self.update_map()
|
||
self.viz_manager.update_display()
|
||
|
||
# Финальное обновление карты
|
||
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})")
|
||
|
||
# Показываем карту до закрытия, но не поднимаем на передний план
|
||
self.viz_manager.show_final()
|
||
print("Симуляция завершена. Окно визуализации остается открытым для анализа.")
|