feat: compute key points

autopilot.py

@@ -1,25 +1,182 @@
-import math
 import random
-from time import sleep
+
+import cv2
 
 import numpy as np
-
-# from scipy import 
-from matplotlib import pyplot as plt
+from PIL import Image
 
 random.seed(1)
 
 class Pilot:
     def __init__(self): pass
-    def handle(self, image: np.ndarray): 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 (random.random() - 0.5) * 0.1
+        return 0.1
 
 # def _test():
 #     randomPilot = RandomPilot()