feat: edit simulator
This commit is contained in:
214
simulator.py
214
simulator.py
@@ -1,113 +1,171 @@
|
||||
import math
|
||||
import os
|
||||
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
|
||||
import cv2
|
||||
import numpy as np
|
||||
from geolocation import Geolocation
|
||||
from vision_chunk import VisionChunk
|
||||
from visualization import VisualizationManager, SimMode
|
||||
from yandex_map import YandexMap
|
||||
|
||||
from PIL import Image
|
||||
|
||||
import os
|
||||
|
||||
class Simulator:
|
||||
geo: Geolocation
|
||||
yandexMap: YandexMap
|
||||
|
||||
# Менеджер визуализации
|
||||
viz_manager: VisualizationManager
|
||||
|
||||
def __init__(self, yandexMap: YandexMap = None):
|
||||
self.yandexMap = yandexMap
|
||||
|
||||
# Инициализация переменных для отслеживания траектории
|
||||
def __init__(self, yandex_map: YandexMap = None):
|
||||
self.yandex_map = yandex_map
|
||||
self.geo = Geolocation(0, 0, 1, 0)
|
||||
|
||||
# Создаем папку для изображений, если её нет
|
||||
# Параметры камеры (в градусах)
|
||||
self.pitch = 0.0 # тангаж (-10 до 10)
|
||||
self.roll = 0.0 # крен (-10 до 10)
|
||||
|
||||
os.makedirs('./images', exist_ok=True)
|
||||
|
||||
def rotate_image_like_drone(self, image: Image.Image, angle: float) -> Image.Image:
|
||||
|
||||
def set_pitch(self, pitch: float):
|
||||
"""Установить тангаж камеры (градусы, -10 до 10)"""
|
||||
self.pitch = max(-10, min(10, pitch))
|
||||
|
||||
def set_roll(self, roll: float):
|
||||
"""Установить крен камеры (градусы, -10 до 10)"""
|
||||
self.roll = max(-10, min(10, roll))
|
||||
|
||||
def _calculate_camera_angles(self, velocity: float, dangle: float):
|
||||
"""Автоматический расчёт тангажа и крена на основе скорости и поворота"""
|
||||
# Тангаж: чем больше скорость, тем больше тангаж (до 10°)
|
||||
self.pitch = min(10, velocity / 10)
|
||||
|
||||
# Крен: чем больше угол поворота, тем больше крен
|
||||
self.roll = max(-10, min(10, math.degrees(dangle) * 2))
|
||||
|
||||
def _get_perspective_points(self, image_width: int, image_height: int,
|
||||
yaw_deg: float) -> tuple:
|
||||
"""
|
||||
Поворачивает картинку как будто съемка ведется с летящего дрона.
|
||||
Выделяет концентрический квадрат, поворачивает его и извлекает результат.
|
||||
Вычисляет 4 точки для перспективной трансформации.
|
||||
Учитывает тангаж, крен, поворот и масштаб.
|
||||
"""
|
||||
# Получаем размеры изображения
|
||||
width, height = image.size
|
||||
square_size = min(width, height)
|
||||
off_x = (width - square_size) / 2
|
||||
off_y = (height - square_size) / 2
|
||||
|
||||
cropped_image = image.crop((off_x, off_y, off_x + square_size, off_y + square_size))
|
||||
cropped_image = cropped_image.rotate(angle / math.pi * 180, expand=True)
|
||||
center_x = image_width / 2
|
||||
center_y = image_height / 2
|
||||
|
||||
# Определяем размер концентрического квадрата (80% от минимальной стороны)
|
||||
local_square_size = int(square_size / 2 ** 0.5)
|
||||
# Базовый размер области захвата (квадрат)
|
||||
base_size = min(image_width, image_height) * 0.7
|
||||
half_size = base_size / 2
|
||||
|
||||
# Вычисляем координаты для центрирования квадрата
|
||||
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
|
||||
# Исходные углы квадрата (до применения перспективы)
|
||||
corners = np.float32([
|
||||
[center_x - half_size, center_y - half_size], # top-left
|
||||
[center_x + half_size, center_y - half_size], # top-right
|
||||
[center_x + half_size, center_y + half_size], # bottom-right
|
||||
[center_x - half_size, center_y + half_size], # bottom-left
|
||||
])
|
||||
|
||||
# Вырезаем концентрический квадрат
|
||||
final_image = cropped_image.crop((left, top, right, bottom))
|
||||
|
||||
return final_image
|
||||
|
||||
# Применяем смещения для имитации тангажа (pitch)
|
||||
# Положительный тангаж - дрон наклонён вперёд, камера смотрит дальше
|
||||
pitch_offset = self.pitch * 3 # коэффициент для видимого эффекта
|
||||
corners[0][1] -= pitch_offset # top-left y
|
||||
corners[1][1] -= pitch_offset # top-right y
|
||||
corners[2][1] += pitch_offset # bottom-right y
|
||||
corners[3][1] += pitch_offset # bottom-left y
|
||||
|
||||
# Применяем смещения для имитации крена (roll)
|
||||
# Положительный крен - дрон наклонён вправо
|
||||
roll_offset = self.roll * 3
|
||||
corners[0][0] += roll_offset # top-left x
|
||||
corners[1][0] -= roll_offset # top-right x
|
||||
corners[2][0] -= roll_offset # bottom-right x
|
||||
corners[3][0] += roll_offset # bottom-left x
|
||||
|
||||
# Поворот вокруг центра (yaw)
|
||||
angle_rad = math.radians(yaw_deg)
|
||||
cos_a = math.cos(angle_rad)
|
||||
sin_a = math.sin(angle_rad)
|
||||
|
||||
rotated_corners = []
|
||||
for corner in corners:
|
||||
# Смещаем к началу координат
|
||||
x = corner[0] - center_x
|
||||
y = corner[1] - center_y
|
||||
|
||||
# Поворачиваем
|
||||
new_x = x * cos_a - y * sin_a
|
||||
new_y = x * sin_a + y * cos_a
|
||||
|
||||
# Возвращаем обратно
|
||||
rotated_corners.append([new_x + center_x, new_y + center_y])
|
||||
|
||||
return np.float32(rotated_corners)
|
||||
|
||||
def _apply_perspective_transform(self, image: Image.Image) -> Image.Image:
|
||||
"""
|
||||
Применяет перспективную трансформацию к изображению.
|
||||
Возвращает квадратное изображение 700x700.
|
||||
"""
|
||||
img_array = np.array(image)
|
||||
h, w = img_array.shape[:2]
|
||||
|
||||
# Получаем исходные точки с учётом перспективы
|
||||
yaw_deg = -math.degrees(self.geo.angle)
|
||||
src_points = self._get_perspective_points(w, h, yaw_deg)
|
||||
|
||||
# Целевые точки - квадрат 700x700
|
||||
dst_points = np.float32([
|
||||
[0, 0],
|
||||
[700, 0],
|
||||
[700, 700],
|
||||
[0, 700]
|
||||
])
|
||||
|
||||
# Вычисляем матрицу трансформации
|
||||
matrix = cv2.getPerspectiveTransform(src_points, dst_points)
|
||||
|
||||
# Применяем трансформацию
|
||||
transformed = cv2.warpPerspective(img_array, matrix, (700, 700))
|
||||
|
||||
return Image.fromarray(transformed)
|
||||
|
||||
def update_trajectory(self, dx: float, dy: float):
|
||||
"""
|
||||
Обновляет траекторию полета беспилотника
|
||||
"""
|
||||
# Обновляем текущие координаты
|
||||
"""Обновляет координаты дрона"""
|
||||
self.geo.x += dx * self.geo.z
|
||||
self.geo.y += dy * self.geo.z
|
||||
|
||||
def update_map(self):
|
||||
"""
|
||||
Обновляет карту траектории полета
|
||||
"""
|
||||
self.viz_manager.update_global_map(self.geo.x, self.geo.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)
|
||||
"""
|
||||
Управление движением дрона.
|
||||
dangle - изменение угла курса (радианы)
|
||||
velocity - скорость движения
|
||||
"""
|
||||
from selenium.webdriver.common.by import By
|
||||
from selenium.webdriver.common.action_chains import ActionChains
|
||||
|
||||
# Автоматический расчёт углов камеры
|
||||
self._calculate_camera_angles(velocity, dangle)
|
||||
|
||||
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()
|
||||
|
||||
|
||||
self.geo.angle += dangle
|
||||
# print(f" [Simulator] angle: {self.angle / math.pi * 180:.1f}°")
|
||||
velocity = max(velocity, 10)
|
||||
|
||||
dx = math.cos(math.pi / 2 + self.geo.angle) * velocity / self.geo.z
|
||||
dy = math.sin(math.pi / 2 + self.geo.angle) * velocity / self.geo.z
|
||||
# print(" [Simulator] dx, dy:", [dx, dy])
|
||||
|
||||
self.update_trajectory(dx, dy)
|
||||
|
||||
action.move_by_offset(-dx, dy)
|
||||
action.release()
|
||||
action.perform()
|
||||
# print(f" [Simulator] Position: {self.current_x}, {self.current_y}")
|
||||
|
||||
def change_zoom(self, direction: str = 'down'):
|
||||
""" Изменение масштаба """
|
||||
self.yandexMap.scroll(0.5, 0.5, 2, direction == 'down')
|
||||
sleep(0.4)
|
||||
if direction == 'down':
|
||||
self.geo.z /= 2
|
||||
else:
|
||||
self.geo.z *= 2
|
||||
|
||||
def set_zoom(self, zoom_level: float):
|
||||
"""Программное изменение масштаба"""
|
||||
self.geo.z = zoom_level
|
||||
|
||||
def get_chunk(self) -> VisionChunk:
|
||||
png = self.yandexMap.driver.get_screenshot_as_png()
|
||||
"""Получить текущий снимок с камеры дрона"""
|
||||
png = self.yandex_map.driver.get_screenshot_as_png()
|
||||
im = Image.open(BytesIO(png))
|
||||
|
||||
# Применяем поворот как будто съемка с дрона
|
||||
rotated_im = self.rotate_image_like_drone(im, -self.geo.angle)
|
||||
return VisionChunk(rotated_im)
|
||||
|
||||
# Применяем перспективную трансформацию
|
||||
transformed_im = self._apply_perspective_transform(im)
|
||||
|
||||
return VisionChunk(transformed_im)
|
||||
|
||||
Reference in New Issue
Block a user