import random import cv2 import numpy as np from PIL import Image import math random.seed(1) class Pilot: def __init__(self): pass def handle(self, image: Image): pass def act(self) -> float | None: pass class AutoPilot(Pilot): prev_image: np.ndarray | None angle: float x: float # Координата X БПЛА y: float # Координата Y БПЛА orb_detector: cv2.ORB bf_matcher: cv2.BFMatcher frame_count: int image_center: tuple # Центр изображения (x, y) viz_manager: object # Менеджер визуализации (опционально) def __init__(self, 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 # Менеджер визуализации # Инициализация 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) 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, pt1[1] - center_y1]) # Координаты точки во втором изображении pt2 = kp2[match.trainIdx].pt dst_pts.append([pt2[0] - center_x2, pt2[1] - center_y2]) # Конвертируем в 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) 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 = H[0, 2], H[1, 2] # Вычисляем угол поворота 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 = tx dy_meters = ty # Применяем поворот к смещению (учитываем текущий угол БПЛА) 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): 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.viz_manager: kp, _ = self.orb_detector.detectAndCompute(self.prev_image, None) self.viz_manager.update_detection(self.prev_image, kp) return # Конвертируем текущее изображение current_image = self.image_to_numpy(image) # Обновляем центр изображения height, width = current_image.shape[:2] self.image_center = (width // 2, height // 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() 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: # Обновляем детекцию ключевых точек self.viz_manager.update_detection(current_image, kp2) # Обновляем сопоставление точек self.viz_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 def act(self) -> float | None: """ Возвращает угол поворота для управления дроном, чтобы он стремился к точке (0, 0). Если дрон находится рядом с началом координат (в радиусе 1 метра), возвращает None. """ # Расстояние до цели (0, 0) distance_to_target = math.sqrt(self.x**2 + self.y**2) # Если дрон находится рядом с целью (в радиусе 1 метра), останавливаемся if distance_to_target < 1.0: return None # Вычисляем угол к цели target_angle = math.atan2(-self.y, -self.x) # Отрицательные координаты, так как движемся к (0,0) # Вычисляем разность углов (направление поворота) 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.05, angle_diff / 2), -0.05) 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) -> float | None: self.counter += 1 if self.counter > 40: return None return 1 / (self.counter + 20) # 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()