100 lines
3.7 KiB
Python
100 lines
3.7 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
|
||
from google_map import GoogleMap
|
||
import constants
|
||
import utility
|
||
|
||
class Simulator:
|
||
def __init__(self, online_map: YandexMap | GoogleMap = None):
|
||
self.online_map = online_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)
|
||
h, w, _ = img_array.shape
|
||
|
||
# Применяем трансформацию
|
||
pos = self.pos.copy()
|
||
pos.x = 0
|
||
pos.y = 0
|
||
|
||
K_in = utility.calc_camera_matrix(w, h)
|
||
K_out = constants.K
|
||
H = pos.get_homography_matrix(K_in, K_out)
|
||
|
||
shape = (constants.CHUNK_WIDTH, constants.CHUNK_WIDTH)
|
||
transformed = cv2.warpPerspective(img_array, H, shape)
|
||
|
||
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 - скорость движения
|
||
"""
|
||
|
||
# Обновляем 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)
|
||
self.online_map.move(dx, dy)
|
||
|
||
def set_zoom(self, zoom_level: float):
|
||
"""Программное изменение масштаба"""
|
||
self.pos.z = zoom_level
|
||
|
||
def get_chunk(self) -> VisionChunk:
|
||
"""Получить текущий снимок с камеры дрона"""
|
||
im = self.online_map.make_screenshot()
|
||
|
||
# Применяем перспективную трансформацию
|
||
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()
|
||
}
|