feat: main loop

This commit is contained in:
2025-09-22 22:29:51 +03:00
parent afd54c55f0
commit 7a1a4050c8
5 changed files with 214 additions and 82 deletions

View File

@@ -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
View File

@@ -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()

View File

@@ -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')

View File

@@ -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):
"""Закрывает окно"""

View File

@@ -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