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 = 0.8274 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, 5.0) RH, rmask = cv2.findHomography(dst_pts, src_pts, cv2.RANSAC, 5.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()