import random import cv2 import numpy as np from PIL import Image 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 orb_detector: cv2.ORB bf_matcher: cv2.BFMatcher def __init__(self): self.prev_image = None self.angle = 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 # Сопоставление ключевых точек 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 # Извлечение координат сопоставленных точек 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 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) return img_matches def handle(self, image: Image): 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"Translation: {transformation_info['translation']}") print(f"Rotation: {transformation_info['rotation']:.4f} rad") print(f"Scale: {transformation_info['scale']:.4f}") # Обновляем угол дрона self.angle += transformation_info['rotation'] # Визуализация (опционально) img_matches = self.visualize_matches(self.prev_image, current_image, kp1, kp2, matches, transformation_info) cv2.imshow('Matches', img_matches) cv2.waitKey(1) # Обновляем предыдущее изображение self.prev_image = current_image def act(self) -> float: """Возвращает угол поворота для управления дроном""" return self.angle 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()