110 lines
4.2 KiB
Python
110 lines
4.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 enum import Enum
|
||
|
||
class SimMode(Enum):
|
||
OPERATOR = 1
|
||
AUTONOME = 2
|
||
|
||
class Simulator:
|
||
driver: webdriver.Chrome
|
||
mode: SimMode
|
||
|
||
operatorPilot: Pilot
|
||
autonomePilot: Pilot
|
||
|
||
angle: float
|
||
|
||
def __init__(self, operatorPilot: Pilot, autonomePilot: Pilot):
|
||
self.mode = SimMode.OPERATOR
|
||
self.operatorPilot = operatorPilot
|
||
self.autonomePilot = autonomePilot
|
||
|
||
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 loop(self):
|
||
|
||
html = self.driver.find_element(By.TAG_NAME, 'html')
|
||
action = ActionChains(self.driver)
|
||
velocity = 10
|
||
|
||
for i in range(100):
|
||
# Сдвиг камеры
|
||
action = ActionChains(self.driver)
|
||
action.move_to_element_with_offset(html, 200, 200)
|
||
action.click_and_hold()
|
||
dangle = self.operatorPilot.act()
|
||
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()
|
||
|
||
# Загрузка скриншота
|
||
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")
|
||
sleep(0.25)
|
||
|
||
print("last position: ", self.driver.current_url)
|