From 7a1a4050c820a8c3c492a86fc0e1f133d05ae4e6 Mon Sep 17 00:00:00 2001 From: russian_proger Date: Mon, 22 Sep 2025 22:29:51 +0300 Subject: [PATCH] feat: main loop --- autopilot.py | 62 +++++++++++++++++++++------------ main.py | 89 +++++++++++++++++++++++++++++++++++++++++++----- simulator.py | 69 ++++++++++++++++--------------------- visualization.py | 2 +- yandex_map.py | 74 ++++++++++++++++++++++++++++++++++------ 5 files changed, 214 insertions(+), 82 deletions(-) diff --git a/autopilot.py b/autopilot.py index f4577b9..32cecdd 100644 --- a/autopilot.py +++ b/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: """ diff --git a/main.py b/main.py index 5cec3a2..1ea3cb2 100644 --- a/main.py +++ b/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() \ No newline at end of file diff --git a/simulator.py b/simulator.py index dcea4ce..0850aa2 100644 --- a/simulator.py +++ b/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') diff --git a/visualization.py b/visualization.py index de882b0..5de7055 100644 --- a/visualization.py +++ b/visualization.py @@ -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): """Закрывает окно""" diff --git a/yandex_map.py b/yandex_map.py index 03ef96c..59c0b81 100644 --- a/yandex_map.py +++ b/yandex_map.py @@ -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 \ No newline at end of file