398 lines
16 KiB
Python
398 lines
16 KiB
Python
from enum import Enum
|
||
from pathlib import Path
|
||
import math
|
||
import random
|
||
|
||
import cv2
|
||
import numpy as np
|
||
from PIL import Image
|
||
|
||
from visualization import VisualizationManager
|
||
|
||
random.seed(1)
|
||
|
||
class Pilot:
|
||
def __init__(self): pass
|
||
def handle(self, image: Image): pass
|
||
def act(self) -> tuple[float, float] | None: pass
|
||
def get_position(self) -> tuple[float, float]: pass
|
||
|
||
class PilotCommand:
|
||
velocity: float
|
||
dangle: float
|
||
stop: bool
|
||
|
||
def __init__(self, dangle: float = 0, velocity: float = 100, stop: bool = False):
|
||
self.dangle = dangle
|
||
self.stop = stop
|
||
self.velocity = velocity
|
||
|
||
MOVE_KOF: float = 1 / 1.2
|
||
|
||
class AutoPilot(Pilot):
|
||
|
||
# Состояние одометрии
|
||
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)
|
||
vis_manager: VisualizationManager # Менеджер визуализации (опционально)
|
||
|
||
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.vis_manager = viz_manager # Менеджер визуализации
|
||
|
||
# Инициализация ORB детектора
|
||
self.orb_detector = cv2.ORB_create(
|
||
nfeatures=1000,
|
||
scaleFactor=1.2,
|
||
nlevels=8,
|
||
edgeThreshold=31,
|
||
firstLevel=0,
|
||
WTA_K=2,
|
||
patchSize=31,
|
||
fastThreshold=20
|
||
)
|
||
|
||
# Инициализация матчера для сопоставления ключевых точек
|
||
self.bf_matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
|
||
|
||
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:
|
||
"""Конвертирует PIL Image в numpy array для OpenCV"""
|
||
return np.array(image)
|
||
|
||
def detect_and_match_keypoints(self, img1: np.ndarray, img2: np.ndarray):
|
||
"""
|
||
Обнаруживает ключевые точки ORB и сопоставляет их между двумя изображениями
|
||
"""
|
||
# Обнаружение ключевых точек и дескрипторов
|
||
kp1, des1 = self.orb_detector.detectAndCompute(img1, None)
|
||
kp2, des2 = self.orb_detector.detectAndCompute(img2, None)
|
||
|
||
if des1 is None or des2 is None:
|
||
return None, None, None, None, None
|
||
|
||
# Сопоставление ключевых точек
|
||
matches = self.bf_matcher.match(des1, des2)
|
||
|
||
# Сортировка совпадений по расстоянию
|
||
matches = sorted(matches, key=lambda x: x.distance)
|
||
|
||
# Фильтрация хороших совпадений (расстояние меньше порога)
|
||
good_matches = []
|
||
for match in matches:
|
||
if match.distance < 50: # Порог расстояния
|
||
good_matches.append(match)
|
||
|
||
if len(good_matches) < 4:
|
||
return None, None, None, None, None
|
||
|
||
# Получаем центр изображения
|
||
height1, width1 = img1.shape[:2]
|
||
height2, width2 = img2.shape[:2]
|
||
center_x1, center_y1 = width1 // 2, height1 // 2
|
||
center_x2, center_y2 = width2 // 2, height2 // 2
|
||
|
||
# Извлекаем координаты сопоставленных точек и отцентрируем их
|
||
src_pts = []
|
||
dst_pts = []
|
||
|
||
for match in good_matches:
|
||
# Координаты точки в первом изображении
|
||
pt1 = kp1[match.queryIdx].pt
|
||
src_pts.append([pt1[0] - center_x1, center_y1 - pt1[1]])
|
||
|
||
# Координаты точки во втором изображении
|
||
pt2 = kp2[match.trainIdx].pt
|
||
dst_pts.append([pt2[0] - center_x2, center_y2 - pt2[1]])
|
||
|
||
# Конвертируем в numpy массивы
|
||
src_pts = np.float32(src_pts).reshape(-1, 1, 2)
|
||
dst_pts = np.float32(dst_pts).reshape(-1, 1, 2)
|
||
|
||
return src_pts, dst_pts, good_matches, kp1, kp2
|
||
|
||
def estimate_transformation_matrix(self, src_pts: np.ndarray, dst_pts: np.ndarray):
|
||
"""
|
||
Оценивает матрицу трансформации на основе сопоставленных ключевых точек
|
||
Точки уже отцентрированы относительно центра изображения
|
||
"""
|
||
# Используем RANSAC для оценки матрицы гомографии
|
||
H, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 1.0)
|
||
RH, rmask = cv2.findHomography(dst_pts, src_pts, cv2.RANSAC, 1.0)
|
||
|
||
if H is None:
|
||
return None
|
||
|
||
# Извлекаем параметры трансформации из матрицы гомографии
|
||
# H = [a11 a12 tx]
|
||
# [a21 a22 ty]
|
||
# [0 0 1 ]
|
||
|
||
# Масштаб и поворот
|
||
a11, a12 = H[0, 0], H[0, 1]
|
||
a21, a22 = H[1, 0], H[1, 1]
|
||
|
||
# Смещение (уже отцентрировано)
|
||
tx, ty = RH[0, 2] * MOVE_KOF, RH[1, 2] * MOVE_KOF
|
||
|
||
print(" [Pilot] translate:", tx, ty)
|
||
|
||
# Вычисляем угол поворота
|
||
angle = -np.arctan2(a21, a11)
|
||
|
||
# Вычисляем масштаб
|
||
scale_x = np.sqrt(a11**2 + a21**2)
|
||
scale_y = np.sqrt(a12**2 + a22**2)
|
||
scale = (scale_x + scale_y) / 2
|
||
|
||
return {
|
||
'translation': (tx, ty), # Уже отцентрировано
|
||
'rotation': angle,
|
||
'scale': scale,
|
||
'homography': H,
|
||
'mask': mask
|
||
}
|
||
|
||
def update_drone_position(self, transformation_info: dict):
|
||
"""
|
||
Обновляет позицию и угол БПЛА на основе трансформации изображения
|
||
Координаты уже отцентрированы относительно центра изображения
|
||
"""
|
||
tx, ty = transformation_info['translation']
|
||
rotation = transformation_info['rotation']
|
||
scale = transformation_info['scale']
|
||
|
||
# Координаты уже отцентрированы, поэтому используем их напрямую
|
||
dx_meters = ty
|
||
dy_meters = tx
|
||
|
||
# Применяем поворот к смещению (учитываем текущий угол БПЛА)
|
||
cos_angle = math.cos(self.angle)
|
||
sin_angle = math.sin(self.angle)
|
||
|
||
# Поворачиваем смещение в глобальные координаты
|
||
# Обратите внимание: dy_meters инвертирован, так как в изображениях Y направлен вниз
|
||
dx_global = dx_meters * cos_angle - dy_meters * sin_angle
|
||
dy_global = dx_meters * sin_angle + dy_meters * cos_angle
|
||
|
||
# Обновляем координаты БПЛА
|
||
self.x += dx_global
|
||
self.y += dy_global
|
||
|
||
# Обновляем угол БПЛА
|
||
self.angle += rotation
|
||
|
||
# Нормализуем угол в диапазоне [-π, π]
|
||
self.angle = math.atan2(math.sin(self.angle), math.cos(self.angle))
|
||
|
||
def get_drone_state(self) -> dict:
|
||
"""
|
||
Возвращает текущее состояние БПЛА
|
||
"""
|
||
return {
|
||
'x': self.x,
|
||
'y': self.y,
|
||
'angle': self.angle,
|
||
'angle_degrees': math.degrees(self.angle),
|
||
'frame_count': self.frame_count
|
||
}
|
||
|
||
def visualize_matches(self, img1: np.ndarray, img2: np.ndarray,
|
||
kp1, kp2, matches, transformation_info):
|
||
"""
|
||
Визуализирует сопоставленные ключевые точки и трансформацию
|
||
"""
|
||
# Рисуем сопоставления
|
||
img_matches = cv2.drawMatches(img1, kp1, img2, kp2, matches, None,
|
||
flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS)
|
||
|
||
# Добавляем информацию о трансформации
|
||
if transformation_info:
|
||
tx, ty = transformation_info['translation']
|
||
angle = transformation_info['rotation']
|
||
scale = transformation_info['scale']
|
||
|
||
info_text = f"Translation: ({tx:.2f}, {ty:.2f})"
|
||
info_text2 = f"Rotation: {angle:.2f} rad ({np.degrees(angle):.1f}°)"
|
||
info_text3 = f"Scale: {scale:.2f}"
|
||
info_text4 = f"Image Center: {self.image_center}"
|
||
|
||
cv2.putText(img_matches, info_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
|
||
cv2.putText(img_matches, info_text2, (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
|
||
cv2.putText(img_matches, info_text3, (10, 90), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
|
||
cv2.putText(img_matches, info_text4, (10, 120), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
|
||
|
||
# Добавляем информацию о позиции БПЛА
|
||
drone_state = self.get_drone_state()
|
||
pos_text = f"Drone Pos: ({drone_state['x']:.2f}, {drone_state['y']:.2f})"
|
||
angle_text = f"Drone Angle: {drone_state['angle_degrees']:.1f}°"
|
||
|
||
cv2.putText(img_matches, pos_text, (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)
|
||
cv2.putText(img_matches, angle_text, (10, 180), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)
|
||
|
||
return img_matches
|
||
|
||
def handle(self, image: Image) -> PilotCommand:
|
||
self.frame_count += 1
|
||
|
||
if self.prev_image is None:
|
||
self.prev_image = self.image_to_numpy(image)
|
||
# Вычисляем центр изображения
|
||
height, width = self.prev_image.shape[:2]
|
||
self.image_center = (width // 2, height // 2)
|
||
|
||
# Обновляем визуализацию детекции
|
||
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)
|
||
|
||
return PilotCommand()
|
||
|
||
# Конвертируем текущее изображение
|
||
current_image = self.image_to_numpy(image)
|
||
|
||
# Обновляем центр изображения
|
||
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)
|
||
|
||
if transformation_info:
|
||
# Обновляем позицию и угол БПЛА
|
||
self.update_drone_position(transformation_info)
|
||
|
||
# Выводим текущее состояние БПЛА
|
||
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)
|
||
|
||
# Пытаемся найти ориентир на картинке:
|
||
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.points[self.target_idx][1] - self.y, self.points[self.target_idx][0] - self.x)
|
||
|
||
# Вычисляем разность углов (направление поворота)
|
||
angle_diff = target_angle - self.angle
|
||
|
||
# Нормализуем разность углов в диапазон [-π, π]
|
||
angle_diff %= 2 * math.pi
|
||
if angle_diff >= math.pi:
|
||
angle_diff -= 2 * math.pi
|
||
|
||
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):
|
||
"""Сбрасывает позицию и угол БПЛА"""
|
||
self.x = x
|
||
self.y = y
|
||
self.angle = angle
|
||
self.frame_count = 0
|
||
|
||
|
||
class RandomPilot(Pilot):
|
||
counter: int
|
||
|
||
def __init__(self, velocity: float = 1):
|
||
self.counter = 0
|
||
|
||
def act(self) -> tuple[float, float] | None:
|
||
self.counter += 1
|
||
if self.counter > 300:
|
||
return None
|
||
|
||
return 1 / (self.counter + 20), 10.0
|
||
|
||
# def _test():
|
||
# randomPilot = RandomPilot()
|
||
# point = [0, 0]
|
||
# iter_count = 100
|
||
# points = [point.copy()]
|
||
|
||
# for i in range(iter_count):
|
||
# dx, dy = randomPilot.step()
|
||
# prev_point = point.copy()
|
||
# point[0] += dx
|
||
# point[1] += dy
|
||
# points.append(point.copy())
|
||
|
||
# coords = list(zip(*points))
|
||
# padding = 5
|
||
# plt.axis([
|
||
# min(coords[0]) - padding, max(coords[0]) + padding,
|
||
# min(coords[1]) - padding, max(coords[1]) + padding])
|
||
|
||
# for i in range(iter_count):
|
||
# plt.plot(coords[0][i:i+2], coords[1][i:i+2], color='#5e5')
|
||
# plt.pause(0.05)
|
||
|
||
# sleep(1)
|
||
|
||
# if __name__ == '__main__':
|
||
# _test() |