feat: add chunks to autopilot, test for estimation

This commit is contained in:
2025-10-01 20:47:06 +03:00
parent 7a1a4050c8
commit 4d3e0d0e59
6 changed files with 297 additions and 55 deletions

View File

@@ -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()
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.vis_manager.update_matches(
self.prev_image,
current_image,
kp1, kp2, matches,
transformation_info)
# Обновляем предыдущее изображение
self.prev_image = current_image
return PilotCommand()
# Пытаемся найти ориентир на картинке:
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)
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 < 50:
self.target_idx += 1
# Если дрон находится рядом с целью, останавливаемся
if distance_to_target < 30.0:
return None
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):
"""Сбрасывает позицию и угол БПЛА"""

64
main.py
View File

@@ -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,11 +66,13 @@ 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])
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__":

View File

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

130
test_autopilot.ipynb Normal file
View File

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

View File

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

View File

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