From 4d3e0d0e59d0704f488b947b57700e5cdc667a1a Mon Sep 17 00:00:00 2001 From: russian_proger Date: Wed, 1 Oct 2025 20:47:06 +0300 Subject: [PATCH] feat: add chunks to autopilot, test for estimation --- autopilot.py | 80 ++++++++++++++++---------- main.py | 68 +++++++++++++++------- simulator.py | 17 +++++- test_autopilot.ipynb | 130 +++++++++++++++++++++++++++++++++++++++++++ visualization.py | 43 +++++++++++++- yandex_map.py | 14 +++++ 6 files changed, 297 insertions(+), 55 deletions(-) create mode 100644 test_autopilot.ipynb diff --git a/autopilot.py b/autopilot.py index 32cecdd..20ae94d 100644 --- a/autopilot.py +++ b/autopilot.py @@ -1,4 +1,5 @@ from enum import Enum +from pathlib import Path import math import random @@ -17,12 +18,16 @@ class Pilot: def get_position(self) -> tuple[float, float]: pass class PilotCommand: + velocity: float dangle: float stop: bool - def __init__(self, dangle: float = 0, stop: bool = False): + def __init__(self, dangle: float = 0, velocity: float = 100, stop: bool = False): self.dangle = dangle self.stop = stop + self.velocity = velocity + +MOVE_KOF: float = 0.8274 class AutoPilot(Pilot): @@ -118,11 +123,11 @@ class AutoPilot(Pilot): for match in good_matches: # Координаты точки в первом изображении pt1 = kp1[match.queryIdx].pt - src_pts.append([center_x1 - pt1[0], pt1[1] - center_y1]) + src_pts.append([pt1[0] - center_x1, center_y1 - pt1[1]]) # Координаты точки во втором изображении pt2 = kp2[match.trainIdx].pt - dst_pts.append([center_x2 - pt2[0], pt2[1] - center_y2]) + dst_pts.append([pt2[0] - center_x2, center_y2 - pt2[1]]) # Конвертируем в numpy массивы src_pts = np.float32(src_pts).reshape(-1, 1, 2) @@ -137,6 +142,7 @@ class AutoPilot(Pilot): """ # Используем RANSAC для оценки матрицы гомографии H, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0) + RH, rmask = cv2.findHomography(dst_pts, src_pts, cv2.RANSAC, 5.0) if H is None: return None @@ -151,7 +157,9 @@ class AutoPilot(Pilot): a21, a22 = H[1, 0], H[1, 1] # Смещение (уже отцентрировано) - tx, ty = H[0, 2], H[1, 2] + tx, ty = RH[0, 2] * MOVE_KOF, RH[1, 2] * MOVE_KOF + + print(" [Pilot] translate:", tx, ty) # Вычисляем угол поворота angle = -np.arctan2(a21, a11) @@ -258,7 +266,7 @@ class AutoPilot(Pilot): self.image_center = (width // 2, height // 2) # Обновляем визуализацию детекции - if self.vis_manager: + if self.vis_manager is not None: kp, _ = self.orb_detector.detectAndCompute(self.prev_image, None) self.vis_manager.update_detection(self.prev_image, kp) @@ -271,9 +279,16 @@ class AutoPilot(Pilot): height, width = current_image.shape[:2] self.image_center = (width // 2, height // 2) + # Расстояние до цели + distance_to_target = math.sqrt( + (self.points[self.target_idx][0] - self.x) ** 2 + + (self.points[self.target_idx][1] - self.y) ** 2 + ) + # Обнаруживаем и сопоставляем ключевые точки src_pts, dst_pts, matches, kp1, kp2 = self.detect_and_match_keypoints(self.prev_image, current_image) + # Оцениваем смещение БПЛА if src_pts is not None and dst_pts is not None: # Оцениваем матрицу трансформации transformation_info = self.estimate_transformation_matrix(src_pts, dst_pts) @@ -284,49 +299,54 @@ class AutoPilot(Pilot): # Выводим текущее состояние БПЛА drone_state = self.get_drone_state() - self.vis_manager.update_drone_trajectory(drone_state['x'], drone_state['y']) + if self.vis_manager: + 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}°") + print(f" [Pilot] Target Index: {self.target_idx}") + print(f" [Pilot] Target Position: {self.points[self.target_idx]}") + print(f" [Pilot] Distance: {distance_to_target}") # Обновляем визуализацию if self.vis_manager: # Обновляем детекцию ключевых точек self.vis_manager.update_detection(current_image, kp2) - # Обновляем сопоставление точек - self.vis_manager.update_matches(self.prev_image, current_image, - kp1, kp2, matches, transformation_info) - - # Обновляем предыдущее изображение - self.prev_image = current_image - return PilotCommand() - def act(self) -> tuple[float, float] | None: - """ - Возвращает угол поворота для управления дроном, чтобы он стремился к точке (0, 0). - Если дрон находится рядом с началом координат (в радиусе 1 метра), возвращает None. - """ - # Расстояние до цели (0, 0) - distance_to_target = math.sqrt(self.x**2 + self.y**2) - - # Если дрон находится рядом с целью, останавливаемся - if distance_to_target < 30.0: - return None - + # Обновляем сопоставление точек + self.vis_manager.update_matches( + self.prev_image, + current_image, + kp1, kp2, matches, + transformation_info) + + # Пытаемся найти ориентир на картинке: + landmark_image = cv2.imread(Path('chunks') / f'chunk_{self.target_idx}.png', cv2.IMREAD_COLOR_RGB) + src_pts, dst_pts, matches, kp1, kp2 = self.detect_and_match_keypoints(current_image, landmark_image) + # if src_pts is not None and dst_pts is not None: + # # Оцениваем матрицу трансформации + # transformation_info = self.estimate_transformation_matrix(src_pts, dst_pts) + # if self.vis_manager: + # self.vis_manager.update_chunk_matches(current_image, landmark_image, kp1, kp2, matches, transformation_info) + + if distance_to_target < 50: + self.target_idx += 1 + + if self.target_idx == len(self.points): + return PilotCommand(stop=True) + # Вычисляем угол к цели - target_angle = math.atan2(-self.y, -self.x) # Отрицательные координаты, так как движемся к (0,0) + target_angle = math.atan2(self.points[self.target_idx][1] - self.y, self.points[self.target_idx][0] - self.x) # Вычисляем разность углов (направление поворота) angle_diff = target_angle - self.angle - - print(self.angle, target_angle, angle_diff) # Нормализуем разность углов в диапазон [-π, π] angle_diff %= 2 * math.pi if angle_diff >= math.pi: angle_diff -= 2 * math.pi - # Возвращаем угол поворота (положительный - поворот влево, отрицательный - вправо) - return max(min(0.1, angle_diff), -0.1), min(10., distance_to_target / 2) + self.prev_image = current_image + return PilotCommand(max(min(0.1, angle_diff), -0.1), min(50., distance_to_target / 2)) def reset_position(self, x: float = 0.0, y: float = 0.0, angle: float = 0.0): """Сбрасывает позицию и угол БПЛА""" diff --git a/main.py b/main.py index 1ea3cb2..ac3e046 100644 --- a/main.py +++ b/main.py @@ -29,27 +29,27 @@ 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)] + 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) + 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() @@ -66,12 +66,14 @@ np.float64(0.5699957136274587)] fastThreshold=20 ) img = cv2.imread(Path('chunks') / f'chunk_{i}.png') + kp: list[cv2.KeyPoint] 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') + if len(kp_coords) > 0: + plt.scatter(kp_coords[:, 0], kp_coords[:, 1], c='red', s=20, alpha=0.7, marker='o') plt.pause(0.2) plt.ioff() @@ -79,21 +81,45 @@ np.float64(0.5699957136274587)] # Начнём симуляцию полёта с первой точки yandexMap.scroll(points[0][0], points[0][1], 5, True) + sleep(1) + yandexMap.make_as_center(*points[0]) + sleep(5) vis_manager = VisualizationManager() - pilot = autopilot.AutoPilot(points, keypoints, vis_manager) + width, height = yandexMap.get_size() + points_coords = np.array(list(map(lambda p: [ + (p[0] - points[0][0]) * width, (points[0][1] - p[1]) * height + ], points))) + points_coords *= 2 ** 4 + pilot = autopilot.AutoPilot(points_coords, keypoints, vis_manager) simulator = Simulator(yandexMap) + pilot.target_idx = 0 - while True: + photo = simulator.get_photo() + command = pilot.handle(photo) + print("DEBUG SIZE:", yandexMap.get_size()) + print("DEBUG SIZE:", photo.size) + + vis_manager.update_display() + vis_manager.pause(10) + + for i in range(10000000000): photo = simulator.get_photo() command = pilot.handle(photo) + + # Save Image + photo.save(Path('images') / f"photo_{i}.png") + + vis_manager.update_display() + vis_manager.pause(1) if command.stop: break - simulator.handle(command.dangle) + simulator.handle(command.dangle, command.velocity) vis_manager.update_display() - + vis_manager.pause(0.2) + vis_manager.pause(50) sleep(30) if __name__ == "__main__": diff --git a/simulator.py b/simulator.py index 0850aa2..c4251db 100644 --- a/simulator.py +++ b/simulator.py @@ -43,7 +43,10 @@ class Simulator: # Получаем размеры изображения width, height = image.size square_size = min(width, height) - cropped_image = image.crop((0, 0, square_size, square_size)) + 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) # Определяем размер концентрического квадрата (80% от минимальной стороны) @@ -83,16 +86,24 @@ class Simulator: action.click_and_hold() self.angle += dangle + print(f" [Simulator] angle: {self.angle / math.pi * 180:.1f}°") + velocity = max(velocity, 10) dx = math.cos(self.angle) * velocity dy = math.sin(self.angle) * velocity + 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 get_photo(self) -> Image: + def get_photo(self) -> Image.Image: png = self.yandexMap.driver.get_screenshot_as_png() im = Image.open(BytesIO(png)) - return im + + # Применяем поворот как будто съемка с дрона + rotated_im = self.rotate_image_like_drone(im, math.pi / 2 - self.angle) + return rotated_im def loop(self): diff --git a/test_autopilot.ipynb b/test_autopilot.ipynb new file mode 100644 index 0000000..8cb53bc --- /dev/null +++ b/test_autopilot.ipynb @@ -0,0 +1,130 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + " [Pilot] translate: -1.1016611435961554 -11.000654444967525\n", + " [Pilot] Drone Position: (11.00, 1.10)\n", + " [Pilot] Angle: 5.7°\n", + " [Pilot] Target Index: 0\n", + " [Pilot] Target Position: (1500, 1000)\n", + " [Pilot] Distance: 1802.7756377319947\n", + "\n", + " [Pilot] translate: -1.3733859493278928 -62.101186076671574\n", + " [Pilot] Drone Position: (72.66, 8.66)\n", + " [Pilot] Angle: 7.0°\n", + " [Pilot] Target Index: 0\n", + " [Pilot] Target Position: (1500, 1000)\n", + " [Pilot] Distance: 1793.018946479205\n", + "\n", + " [Pilot] translate: -0.286107956133153 -61.700187952027996\n", + " [Pilot] Drone Position: (133.86, 16.44)\n", + " [Pilot] Angle: 7.3°\n", + " [Pilot] Target Index: 0\n", + " [Pilot] Target Position: (1500, 1000)\n", + " [Pilot] Distance: 1737.8358192949638\n", + "\n", + " [Pilot] translate: -1.217279449147543 -61.59798129394066\n", + " [Pilot] Drone Position: (194.81, 25.43)\n", + " [Pilot] Angle: 8.4°\n", + " [Pilot] Target Index: 0\n", + " [Pilot] Target Position: (1500, 1000)\n", + " [Pilot] Distance: 1683.363020511812\n", + "\n", + " [Pilot] translate: -0.5202028232620485 -61.25377666064388\n", + " [Pilot] Drone Position: (255.34, 34.87)\n", + " [Pilot] Angle: 8.9°\n", + " [Pilot] Target Index: 0\n", + " [Pilot] Target Position: (1500, 1000)\n", + " [Pilot] Distance: 1628.8954458349883\n", + "\n", + " [Pilot] translate: 0.030834225184026082 -62.02734738952944\n", + " [Pilot] Drone Position: (316.63, 44.41)\n", + " [Pilot] Angle: 8.8°\n", + " [Pilot] Target Index: 0\n", + " [Pilot] Target Position: (1500, 1000)\n", + " [Pilot] Distance: 1575.0091545970413\n", + "\n", + " [Pilot] translate: -0.4160221189622985 -62.497522722978935\n", + " [Pilot] Drone Position: (378.32, 54.42)\n", + " [Pilot] Angle: 9.2°\n", + " [Pilot] Target Index: 0\n", + " [Pilot] Target Position: (1500, 1000)\n", + " [Pilot] Distance: 1521.0283267047403\n", + "\n", + " [Pilot] translate: 0.5176443048484498 -60.8909129531913\n", + " [Pilot] Drone Position: (438.51, 63.67)\n", + " [Pilot] Angle: 8.7°\n", + " [Pilot] Target Index: 0\n", + " [Pilot] Target Position: (1500, 1000)\n", + " [Pilot] Distance: 1467.065793978559\n", + "\n", + " [Pilot] translate: -0.22161188590699135 -62.49031949249053\n", + " [Pilot] Drone Position: (500.24, 73.38)\n", + " [Pilot] Angle: 8.9°\n", + " [Pilot] Target Index: 0\n", + " [Pilot] Target Position: (1500, 1000)\n", + " [Pilot] Distance: 1415.4431252709192\n" + ] + } + ], + "source": [ + "from pathlib import Path\n", + "from PIL import Image\n", + "import numpy as np\n", + "from autopilot import AutoPilot\n", + "from visualization import VisualizationManager\n", + "\n", + "autopilot = AutoPilot([(1500, 1000)], [])\n", + "imgs = [Image.open(Path('images') / f'photo_{i}.png') for i in range(10)]\n", + "autopilot.handle(imgs[0])\n", + "for i in range(1, 10):\n", + " print()\n", + " autopilot.handle(imgs[i])\n" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "None\n" + ] + } + ], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.0" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/visualization.py b/visualization.py index 5de7055..9e04e7d 100644 --- a/visualization.py +++ b/visualization.py @@ -31,6 +31,7 @@ class VisualizationManager: self.ax_global_map = None self.ax_detection = None self.ax_matches = None + self.ax_chunk_matches = None # Данные для глобальной карты self.trajectory_x = [] @@ -64,7 +65,7 @@ class VisualizationManager: self.fig.canvas.manager.window.state('zoomed') # Создаем сетку 2x2 с разными размерами колонок - gs = self.fig.add_gridspec(2, 2, hspace=0.3, wspace=0.3, width_ratios=[1, 1]) + gs = self.fig.add_gridspec(2, 3, hspace=0.3, wspace=0.3, width_ratios=[.5, 1, 1]) # График погрешности позиции (левый верхний угол) self.ax_error_plot = self.fig.add_subplot(gs[0, 0]) @@ -92,10 +93,16 @@ class VisualizationManager: self.ax_matches.set_title('Feature Matching') self.ax_matches.axis('off') + # Сопоставление точек (правый нижний угол) + self.ax_chunk_matches = self.fig.add_subplot(gs[0, 2]) + self.ax_chunk_matches.set_title('Chunk Matching') + self.ax_chunk_matches.axis('off') + # Настройки окна self.fig.canvas.manager.window.attributes('-topmost', False) plt.tight_layout() + plt.show(block=False) def update_global_map(self, x: float, y: float, mode: SimMode): """Обновляет глобальную карту""" @@ -257,6 +264,40 @@ class VisualizationManager: self.ax_matches.axis('off') + def update_chunk_matches(self, img1: np.ndarray, img2: np.ndarray, + kp1, kp2, matches, transformation_info=None): + """Обновляет визуализацию сопоставления точек""" + self.ax_chunk_matches.clear() + self.ax_chunk_matches.set_title('Chunk Matching') + + if img1 is not None and img2 is not None and matches: + # Рисуем сопоставления + img_matches = cv2.drawMatches(img1, kp1, img2, kp2, matches, None, + flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS) + + # Конвертируем BGR в RGB + if len(img_matches.shape) == 3 and img_matches.shape[2] == 3: + img_matches_rgb = cv2.cvtColor(img_matches, cv2.COLOR_BGR2RGB) + else: + img_matches_rgb = img_matches + + self.ax_chunk_matches.imshow(img_matches_rgb) + + # Добавляем информацию о трансформации + if transformation_info: + tx, ty = transformation_info['translation'] + angle = transformation_info['rotation'] + + info_text = f"Translation: ({tx:.2f}, {ty:.2f})" + info_text2 = f"Rotation: {angle:.2f} rad ({np.degrees(angle):.1f}°)" + + self.ax_chunk_matches.text(10, 30, info_text, fontsize=8, color='green', + bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8)) + self.ax_chunk_matches.text(10, 90, info_text2, fontsize=8, color='green', + bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8)) + + self.ax_chunk_matches.axis('off') + def update_display(self): """Обновляет отображение всех областей""" self.fig.canvas.draw() diff --git a/yandex_map.py b/yandex_map.py index 59c0b81..3b5a913 100644 --- a/yandex_map.py +++ b/yandex_map.py @@ -59,6 +59,20 @@ class YandexMap: if i != count - 1: sleep(0.25) + def make_as_center(self, x: float, y: float): + html = self.driver.find_element(By.TAG_NAME, 'html') + + action = ActionChains(self.driver) + action.move_to_element_with_offset(html, 0, 0) + action.click_and_hold() + + dx = (x - 0.5) * self.get_size()[0] + dy = (0.5 - y) * self.get_size()[1] + print(dx, dy) + action.move_by_offset(-dx, dy) + action.release() + action.perform() + def make_screenshot(self, x: float, y: float, width: float, height: float) -> cv2.typing.MatLike: # Сохраняем скриншот self.save_photo("temp.png")