autopilot/simulator.py

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
from yandex_map import YandexMap

from PIL import Image

import os

class Simulator:
    angle: float
    yandexMap: YandexMap

    # Менеджер визуализации
    viz_manager: VisualizationManager
    current_x: float
    current_y: float

    def __init__(self, yandexMap: YandexMap = None):
        self.yandexMap = yandexMap

        # Инициализация переменных для отслеживания траектории
        self.angle = 0.0
        self.current_x = 0.0
        self.current_y = 0.0
        
        # Создаем папку для изображений, если её нет
        os.makedirs('./images', exist_ok=True)

    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 handle(self, dangle: float, velocity: float = 50) -> None:
        """ Сдвиг камеры """
        html = self.yandexMap.driver.find_element(By.TAG_NAME, 'html')

        action = ActionChains(self.yandexMap.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()

    def get_photo(self) -> Image:
        png = self.yandexMap.driver.get_screenshot_as_png()
        im = Image.open(BytesIO(png))
        return im

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