108 lines
4.1 KiB
Python
108 lines
4.1 KiB
Python
import math
|
||
import os
|
||
from io import BytesIO
|
||
from time import sleep
|
||
from PIL import Image
|
||
import cv2
|
||
import numpy as np
|
||
from position import Position
|
||
from vision_chunk import VisionChunk
|
||
from yandex_map import YandexMap
|
||
import constants
|
||
import utility
|
||
|
||
class Simulator:
|
||
def __init__(self, yandex_map: YandexMap = None):
|
||
self.yandex_map = yandex_map
|
||
# Используем новый конструктор с yaw, pitch, roll
|
||
self.pos = Position(x=0, y=0, z=1, yaw=0, pitch=0, roll=0)
|
||
|
||
os.makedirs('./images', exist_ok=True)
|
||
|
||
def set_pitch(self, pitch_deg: float):
|
||
"""Установить тангаж камеры (градусы, -10 до 10)"""
|
||
clamped_pitch = max(-10, min(10, pitch_deg))
|
||
self.pos.pitch = math.radians(clamped_pitch)
|
||
|
||
def set_roll(self, roll_deg: float):
|
||
"""Установить крен камеры (градусы, -10 до 10)"""
|
||
clamped_roll = max(-10, min(10, roll_deg))
|
||
self.pos.roll = math.radians(clamped_roll)
|
||
|
||
def _apply_perspective_transform(self, image: Image.Image) -> Image.Image:
|
||
"""
|
||
Применяет перспективную трансформацию к изображению.
|
||
Возвращает квадратное изображение 700x700.
|
||
"""
|
||
img_array = np.array(image)
|
||
print(img_array.shape)
|
||
h, w, _ = img_array.shape
|
||
|
||
# Применяем трансформацию
|
||
pos = self.pos.copy()
|
||
pos.x = 0
|
||
pos.y = 0
|
||
K = utility.calc_camera_matrix(w, h)
|
||
K = constants.K
|
||
img_array = img_array[:constants.CHUNK_WIDTH, :constants.CHUNK_WIDTH]
|
||
transformed = cv2.warpPerspective(img_array, pos.get_homography_matrix(K), (constants.CHUNK_WIDTH, constants.CHUNK_WIDTH))
|
||
|
||
return Image.fromarray(transformed)
|
||
|
||
def update_trajectory(self, dx: float, dy: float):
|
||
"""Обновляет координаты дрона"""
|
||
self.pos.x += dx * self.pos.z
|
||
self.pos.y += dy * self.pos.z
|
||
|
||
def handle(self, dangle: float, velocity: float = 50) -> None:
|
||
"""
|
||
Управление движением дрона.
|
||
dangle - изменение угла курса (радианы)
|
||
velocity - скорость движения
|
||
"""
|
||
from selenium.webdriver.common.by import By
|
||
from selenium.webdriver.common.action_chains import ActionChains
|
||
|
||
html = self.yandex_map.driver.find_element(By.TAG_NAME, 'html')
|
||
action = ActionChains(self.yandex_map.driver)
|
||
action.move_to_element_with_offset(html, 200, 200)
|
||
action.click_and_hold()
|
||
|
||
# Обновляем yaw в объекте Position
|
||
self.pos.yaw += dangle
|
||
velocity = max(velocity, 10)
|
||
|
||
# Вычисляем смещение на основе текущего yaw
|
||
dx = math.cos(math.pi / 2 + self.pos.yaw) * velocity / self.pos.z
|
||
dy = math.sin(math.pi / 2 + self.pos.yaw) * velocity / self.pos.z
|
||
|
||
self.update_trajectory(dx, dy)
|
||
|
||
action.move_by_offset(-dx, dy)
|
||
action.release()
|
||
action.perform()
|
||
|
||
def set_zoom(self, zoom_level: float):
|
||
"""Программное изменение масштаба"""
|
||
self.pos.z = zoom_level
|
||
|
||
def get_chunk(self) -> VisionChunk:
|
||
"""Получить текущий снимок с камеры дрона"""
|
||
png = self.yandex_map.driver.get_screenshot_as_png()
|
||
im = Image.open(BytesIO(png))
|
||
|
||
# Применяем перспективную трансформацию
|
||
transformed_im = self._apply_perspective_transform(im)
|
||
|
||
return VisionChunk(transformed_im)
|
||
|
||
def get_orientation_info(self) -> dict:
|
||
"""Получить полную информацию об ориентации дрона"""
|
||
return {
|
||
'position': (self.pos.x, self.pos.y, self.pos.z),
|
||
'yaw_deg': math.degrees(self.pos.yaw),
|
||
'pitch_deg': math.degrees(self.pos.pitch),
|
||
'roll_deg': math.degrees(self.pos.roll),
|
||
'rotation_matrix': self.pos.get_rotation_matrix()
|
||
}
|