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 def __init__(self, initial_x: float = 0.0, initial_y: float = 0.0): self.prev_image = None self.angle = 0.0 self.x = initial_x self.y = initial_y self.frame_count = 0 # Инициализация 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 # Извлечение координат сопоставленных точек src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2) dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).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) # Поворачиваем смещение в глобальные координаты 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}" 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) # Добавляем информацию о позиции БПЛА 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, 120), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2) cv2.putText(img_matches, angle_text, (10, 150), 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) return # Конвертируем текущее изображение current_image = self.image_to_numpy(image) # Обнаруживаем и сопоставляем ключевые точки 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: print(f"Frame {self.frame_count}:") print(f" Translation: {transformation_info['translation']}") print(f" Rotation: {transformation_info['rotation']:.4f} rad") print(f" Scale: {transformation_info['scale']:.4f}") # Обновляем позицию и угол БПЛА self.update_drone_position(transformation_info) # Выводим текущее состояние БПЛА drone_state = self.get_drone_state() print(f" Drone Position: ({drone_state['x']:.2f}, {drone_state['y']:.2f})") print(f" Drone Angle: {drone_state['angle_degrees']:.1f}°") # Визуализация (опционально) 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: """Возвращает угол поворота для управления дроном""" return self.angle 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): def __init__(self, velocity: float = 1): pass def act(self) -> float: return 0.1 # 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()