feat: main loop
This commit is contained in:
62
autopilot.py
62
autopilot.py
@@ -1,8 +1,11 @@
|
||||
from enum import Enum
|
||||
import math
|
||||
import random
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import math
|
||||
|
||||
from visualization import VisualizationManager
|
||||
|
||||
random.seed(1)
|
||||
@@ -13,25 +16,42 @@ class Pilot:
|
||||
def act(self) -> tuple[float, float] | None: pass
|
||||
def get_position(self) -> tuple[float, float]: pass
|
||||
|
||||
class PilotCommand:
|
||||
dangle: float
|
||||
stop: bool
|
||||
|
||||
def __init__(self, dangle: float = 0, stop: bool = False):
|
||||
self.dangle = dangle
|
||||
self.stop = stop
|
||||
|
||||
class AutoPilot(Pilot):
|
||||
prev_image: np.ndarray | None
|
||||
|
||||
# Состояние одометрии
|
||||
angle: float
|
||||
x: float # Координата X БПЛА
|
||||
y: float # Координата Y БПЛА
|
||||
|
||||
# Ориентиры
|
||||
points: list[tuple[float,float]] # Координаты
|
||||
keypoints: list[any] # Ключевые точки ориентиров
|
||||
target_idx: int # Текущая целевая метка
|
||||
|
||||
# Всякие вспомогательные штуки
|
||||
prev_image: np.ndarray | None
|
||||
orb_detector: cv2.ORB
|
||||
bf_matcher: cv2.BFMatcher
|
||||
frame_count: int
|
||||
image_center: tuple # Центр изображения (x, y)
|
||||
viz_manager: VisualizationManager # Менеджер визуализации (опционально)
|
||||
vis_manager: VisualizationManager # Менеджер визуализации (опционально)
|
||||
|
||||
def __init__(self, viz_manager=None):
|
||||
def __init__(self, points = [], keypoints = [], viz_manager=None):
|
||||
self.prev_image = None
|
||||
self.angle = 0.0
|
||||
self.x = 0.0
|
||||
self.y = 0.0
|
||||
self.frame_count = 0
|
||||
self.image_center = (0, 0) # Будет обновлено при получении первого изображения
|
||||
self.viz_manager = viz_manager # Менеджер визуализации
|
||||
self.vis_manager = viz_manager # Менеджер визуализации
|
||||
|
||||
# Инициализация ORB детектора
|
||||
self.orb_detector = cv2.ORB_create(
|
||||
@@ -48,7 +68,11 @@ class AutoPilot(Pilot):
|
||||
# Инициализация матчера для сопоставления ключевых точек
|
||||
self.bf_matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
|
||||
|
||||
def get_position(self):
|
||||
self.points = points
|
||||
self.keypoints = keypoints
|
||||
self.target_idx = 0
|
||||
|
||||
def get_position(self) -> tuple[float, float]:
|
||||
return self.x, self.y
|
||||
|
||||
def image_to_numpy(self, image: Image.Image) -> np.ndarray:
|
||||
@@ -224,9 +248,9 @@ class AutoPilot(Pilot):
|
||||
|
||||
return img_matches
|
||||
|
||||
def handle(self, image: Image):
|
||||
def handle(self, image: Image) -> PilotCommand:
|
||||
self.frame_count += 1
|
||||
|
||||
|
||||
if self.prev_image is None:
|
||||
self.prev_image = self.image_to_numpy(image)
|
||||
# Вычисляем центр изображения
|
||||
@@ -234,11 +258,11 @@ class AutoPilot(Pilot):
|
||||
self.image_center = (width // 2, height // 2)
|
||||
|
||||
# Обновляем визуализацию детекции
|
||||
if self.viz_manager:
|
||||
if self.vis_manager:
|
||||
kp, _ = self.orb_detector.detectAndCompute(self.prev_image, None)
|
||||
self.viz_manager.update_detection(self.prev_image, kp)
|
||||
self.vis_manager.update_detection(self.prev_image, kp)
|
||||
|
||||
return
|
||||
return PilotCommand()
|
||||
|
||||
# Конвертируем текущее изображение
|
||||
current_image = self.image_to_numpy(image)
|
||||
@@ -260,27 +284,21 @@ class AutoPilot(Pilot):
|
||||
|
||||
# Выводим текущее состояние БПЛА
|
||||
drone_state = self.get_drone_state()
|
||||
self.viz_manager.update_drone_trajectory(drone_state['x'], drone_state['y'])
|
||||
self.vis_manager.update_drone_trajectory(drone_state['x'], drone_state['y'])
|
||||
print(f" [Pilot] Drone Position: ({drone_state['x']:.2f}, {drone_state['y']:.2f})")
|
||||
print(f" [Pilot] Angle: {drone_state['angle_degrees']:.1f}°")
|
||||
|
||||
# Обновляем визуализацию
|
||||
if self.viz_manager:
|
||||
if self.vis_manager:
|
||||
# Обновляем детекцию ключевых точек
|
||||
self.viz_manager.update_detection(current_image, kp2)
|
||||
self.vis_manager.update_detection(current_image, kp2)
|
||||
# Обновляем сопоставление точек
|
||||
self.viz_manager.update_matches(self.prev_image, current_image,
|
||||
self.vis_manager.update_matches(self.prev_image, current_image,
|
||||
kp1, kp2, matches, transformation_info)
|
||||
|
||||
# Визуализация (опционально, если нет менеджера визуализации)
|
||||
if not self.viz_manager:
|
||||
img_matches = self.visualize_matches(self.prev_image, current_image,
|
||||
kp1, kp2, matches, transformation_info)
|
||||
cv2.imshow('Drone Tracking', img_matches)
|
||||
cv2.waitKey(1)
|
||||
|
||||
# Обновляем предыдущее изображение
|
||||
self.prev_image = current_image
|
||||
return PilotCommand()
|
||||
|
||||
def act(self) -> tuple[float, float] | None:
|
||||
"""
|
||||
|
||||
89
main.py
89
main.py
@@ -3,25 +3,98 @@ from simulator import Simulator
|
||||
from visualization import VisualizationManager
|
||||
from trajectory_drawer import TrajectoryDrawer
|
||||
from yandex_map import YandexMap
|
||||
from time import sleep
|
||||
from pathlib import Path
|
||||
|
||||
def main():
|
||||
global trajectoryDrawer
|
||||
import cv2
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
import autopilot
|
||||
|
||||
def make_global_photo(filename):
|
||||
yandexMap = YandexMap()
|
||||
yandexMap.savePhoto('map.jpg')
|
||||
yandexMap.save_photo()
|
||||
yandexMap.destroy()
|
||||
|
||||
|
||||
trajectoryDrawer = TrajectoryDrawer('map.jpg')
|
||||
trajectoryDrawer.on_complete_trajectory = onCompleteTrajectory
|
||||
def get_trajectory_points(bg_img: str) -> list[(float, float)]:
|
||||
trajectoryDrawer = TrajectoryDrawer(bg_img)
|
||||
trajectoryDrawer.show()
|
||||
trajectoryDrawer.wait()
|
||||
onCompleteTrajectory(trajectoryDrawer.points)
|
||||
points = list(map(lambda p: [p[0] / trajectoryDrawer.img.shape[1], p[1] / trajectoryDrawer.img.shape[0]], trajectoryDrawer.points))
|
||||
return points
|
||||
|
||||
def onCompleteTrajectory(points):
|
||||
def main():
|
||||
# Скриншот местности
|
||||
# make_global_photo('map.jpg')
|
||||
|
||||
# Получаем траекторию от пользователя
|
||||
# points = get_trajectory_points('map.jpg')
|
||||
# print(points)
|
||||
points = [np.float64(0.5443937502226799), np.float64(0.4030424838774785)], [np.float64(0.18517133490120316), np.float64(0.4586935604608052)], [np.float64(0.1641887171838272), np.float64(0.5586383510594329)], [np.float64(0.587198290366127),
|
||||
np.float64(0.5699957136274587)]
|
||||
|
||||
# Для каждой точки сделаем приближенный снимок
|
||||
yandexMap = YandexMap()
|
||||
keypoints: list[(any, any)] = []
|
||||
|
||||
plt.ion()
|
||||
# for i in range(len(points)):
|
||||
# point = points[i]
|
||||
# yandexMap.scroll(point[0], point[1], 5, True)
|
||||
# sleep(1)
|
||||
# img = yandexMap.make_screenshot(point[0], point[1], 0.2, 0.2)
|
||||
# Path('chunks').mkdir(exist_ok=True)
|
||||
# cv2.imwrite(Path('.') / 'chunks' / f'chunk_{i}.png', img)
|
||||
# plt.subplot(1, len(points), i+1)
|
||||
# plt.imshow(img)
|
||||
# plt.pause(0.25)
|
||||
# yandexMap.scroll(point[0], point[1], 5, False)
|
||||
|
||||
plt.tight_layout()
|
||||
|
||||
# Выделим на каждой картинке ключевые точки
|
||||
for i in range(len(points)):
|
||||
orb = cv2.ORB_create(
|
||||
nfeatures=1000,
|
||||
scaleFactor=1.2,
|
||||
nlevels=8,
|
||||
edgeThreshold=31,
|
||||
firstLevel=0,
|
||||
WTA_K=2,
|
||||
patchSize=31,
|
||||
fastThreshold=20
|
||||
)
|
||||
img = cv2.imread(Path('chunks') / f'chunk_{i}.png')
|
||||
kp, des = orb.detectAndCompute(img, None)
|
||||
keypoints.append((kp, des))
|
||||
|
||||
plt.subplot(1, len(points), i+1)
|
||||
kp_coords = np.array([j.pt for j in kp])
|
||||
plt.scatter(kp_coords[:, 0], kp_coords[:, 1], c='red', s=20, alpha=0.7, marker='o')
|
||||
plt.pause(0.2)
|
||||
plt.ioff()
|
||||
|
||||
plt.show(block=True)
|
||||
|
||||
# Начнём симуляцию полёта с первой точки
|
||||
yandexMap.scroll(points[0][0], points[0][1], 5, True)
|
||||
|
||||
vis_manager = VisualizationManager()
|
||||
pilot = autopilot.AutoPilot(points, keypoints, vis_manager)
|
||||
simulator = Simulator(yandexMap)
|
||||
|
||||
while True:
|
||||
photo = simulator.get_photo()
|
||||
command = pilot.handle(photo)
|
||||
|
||||
if command.stop:
|
||||
break
|
||||
|
||||
simulator.handle(command.dangle)
|
||||
vis_manager.update_display()
|
||||
|
||||
sleep(30)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
69
simulator.py
69
simulator.py
@@ -9,63 +9,32 @@ 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:
|
||||
driver: webdriver.Chrome
|
||||
mode: SimMode
|
||||
|
||||
operatorPilot: Pilot
|
||||
autonomePilot: Pilot
|
||||
|
||||
angle: float
|
||||
|
||||
yandexMap: YandexMap
|
||||
|
||||
# Менеджер визуализации
|
||||
viz_manager: VisualizationManager
|
||||
current_x: float
|
||||
current_y: float
|
||||
|
||||
def __init__(self, operatorPilot: Pilot, autonomePilot: Pilot, viz_manager: VisualizationManager = None):
|
||||
self.mode = SimMode.OPERATOR
|
||||
self.operatorPilot = operatorPilot
|
||||
self.autonomePilot = autonomePilot
|
||||
def __init__(self, yandexMap: YandexMap = None):
|
||||
self.yandexMap = yandexMap
|
||||
|
||||
# Инициализация переменных для траектории
|
||||
# Инициализация переменных для отслеживания траектории
|
||||
self.angle = 0.0
|
||||
self.current_x = 0.0
|
||||
self.current_y = 0.0
|
||||
|
||||
# Создаем менеджер визуализации
|
||||
self.viz_manager = viz_manager
|
||||
|
||||
# Передаем менеджер визуализации в автопилот, если он поддерживает это
|
||||
if hasattr(self.autonomePilot, 'viz_manager'):
|
||||
self.autonomePilot.viz_manager = self.viz_manager
|
||||
|
||||
# Создаем папку для изображений, если её нет
|
||||
os.makedirs('./images', exist_ok=True)
|
||||
|
||||
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.driver.get_screenshot_as_png()
|
||||
|
||||
self.angle = 0
|
||||
|
||||
def rotate_image_like_drone(self, image: Image.Image, angle: float) -> Image.Image:
|
||||
"""
|
||||
Поворачивает картинку как будто съемка ведется с летящего дрона.
|
||||
@@ -105,6 +74,26 @@ class Simulator:
|
||||
"""
|
||||
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')
|
||||
|
||||
@@ -261,7 +261,7 @@ class VisualizationManager:
|
||||
"""Обновляет отображение всех областей"""
|
||||
self.fig.canvas.draw()
|
||||
self.fig.canvas.flush_events()
|
||||
plt.pause(0.01)
|
||||
plt.pause(0.2)
|
||||
|
||||
def close(self):
|
||||
"""Закрывает окно"""
|
||||
|
||||
@@ -3,13 +3,17 @@ from io import BytesIO
|
||||
from time import sleep
|
||||
import os
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from PIL import Image
|
||||
from selenium.webdriver.common.actions.wheel_input import ScrollOrigin
|
||||
from selenium import webdriver
|
||||
from selenium.webdriver.common.by import By
|
||||
from selenium.webdriver.common.action_chains import ActionChains
|
||||
|
||||
def generateURL(lat=49.103814, lon=55.794258, zoom=10):
|
||||
return f"https://yandex.ru/maps/43/kazan/?ll={lat}%2C{lon}&z={zoom}"
|
||||
return f"https://yandex.ru/maps/43/kazan/?l=sat&ll={lat}%2C{lon}&z={zoom}"
|
||||
|
||||
class YandexMap:
|
||||
def __init__(self):
|
||||
@@ -23,21 +27,69 @@ class YandexMap:
|
||||
|
||||
# Закрытие левой панели
|
||||
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.driver.execute_script("arguments[0].remove();", self.driver.find_element(By.XPATH, "//div[@data-chunk='header']"))
|
||||
self.driver.execute_script("arguments[0].remove();", self.driver.find_element(By.XPATH, "//div[@class='sidebar-toggle-button _collapsed _name_home']"))
|
||||
self.driver.execute_script("arguments[0].remove();", self.driver.find_element(By.XPATH, "//nav[@class='map-controls']"))
|
||||
self.driver.execute_script("arguments[0].remove();", self.driver.find_element(By.XPATH, "//footer"))
|
||||
|
||||
html = self.driver.find_element(By.TAG_NAME, 'html')
|
||||
action.move_to_element_with_offset(html, 200, 200)
|
||||
action.perform()
|
||||
sleep(0.2)
|
||||
|
||||
sleep(1)
|
||||
|
||||
def savePhoto(self, filename: str) -> bytes:
|
||||
def save_photo(self, filename: str) -> bytes:
|
||||
return self.driver.save_screenshot(filename)
|
||||
|
||||
def destroy(self):
|
||||
self.driver.close()
|
||||
|
||||
def get_size(self) -> tuple[int, int]:
|
||||
html = self.driver.find_element(By.TAG_NAME, 'html')
|
||||
return (html.size['width'], html.size['height'])
|
||||
|
||||
def scroll(self, x: float, y: float, count: int = 1, inner_zoom: bool = True):
|
||||
html = self.driver.find_element(By.TAG_NAME, 'html')
|
||||
|
||||
x_offset = (x - 0.5) * html.size['width']
|
||||
y_offset = (y - 0.5) * html.size['height']
|
||||
action = ActionChains(self.driver)
|
||||
|
||||
for i in range(count-1):
|
||||
action.scroll_from_origin(ScrollOrigin(html, int(x_offset), int(y_offset)), 0, -100 if inner_zoom else 100)
|
||||
action.perform()
|
||||
if i != count - 1:
|
||||
sleep(0.25)
|
||||
|
||||
def make_screenshot(self, x: float, y: float, width: float, height: float) -> cv2.typing.MatLike:
|
||||
# Сохраняем скриншот
|
||||
self.save_photo("temp.png")
|
||||
|
||||
# Загружаем изображение
|
||||
image = cv2.imread("temp.png")
|
||||
|
||||
if image is None:
|
||||
raise ValueError("Не удалось загрузить изображение temp.png")
|
||||
|
||||
# Получаем размеры исходного изображения
|
||||
img_height, img_width = image.shape[:2]
|
||||
|
||||
# Преобразуем относительные координаты в абсолютные пиксели
|
||||
center_x = int(x * img_width)
|
||||
center_y = int(y * img_height)
|
||||
crop_width = int(width * img_width)
|
||||
crop_height = int(height * img_height)
|
||||
|
||||
# Вычисляем координаты прямоугольника для кадрирования
|
||||
x1 = max(0, center_x - crop_width // 2)
|
||||
y1 = max(0, center_y - crop_height // 2)
|
||||
x2 = min(img_width, center_x + crop_width // 2)
|
||||
y2 = min(img_height, center_y + crop_height // 2)
|
||||
|
||||
# Проверяем, что прямоугольник имеет положительные размеры
|
||||
if x2 <= x1 or y2 <= y1:
|
||||
raise ValueError("Некорректные размеры для кадрирования")
|
||||
|
||||
# Кадрируем изображение
|
||||
cropped_image = image[y1:y2, x1:x2]
|
||||
|
||||
# Если нужно вернуть изображение как результат функции:
|
||||
return cropped_image
|
||||
Reference in New Issue
Block a user