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 PIL import Image
-# from scipy import 
-from matplotlib import pyplot as plt
 
 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()

main.py

@@ -1,5 +1,6 @@
-from autopilot import RandomPilot
+from autopilot import AutoPilot, RandomPilot
 from simulator import Simulator
 
-simulator = Simulator(RandomPilot(), RandomPilot())
+# Создаем симулятор с AutoPilot для обработки изображений
+simulator = Simulator(RandomPilot(), AutoPilot())
 simulator.loop()

requirements.txt

@@ -9,6 +9,7 @@ idna==3.10
 kiwisolver==1.4.8
 matplotlib==3.10.1
 numpy==2.2.4
+opencv-python==4.9.0.80
 outcome==1.3.0.post0
 packaging==24.2
 pillow==11.2.1

simulator.py

@@ -48,6 +48,31 @@ class Simulator:
 
         self.angle = 0
 
+    def rotate_image_like_drone(self, image: Image.Image, angle: float) -> Image.Image:
+        """
+        Поворачивает картинку как будто съемка ведется с летящего дрона.
+        Выделяет концентрический квадрат, поворачивает его и извлекает результат.
+        """
+        # Получаем размеры изображения
+        width, height = image.size
+        square_size = min(width, height)
+        cropped_image = image.crop((0, 0, square_size, square_size))
+        cropped_image = cropped_image.rotate(angle / math.pi * 180, expand=True)
+        
+        # Определяем размер концентрического квадрата (80% от минимальной стороны)
+        local_square_size = int(square_size / 2 ** 0.5)
+        
+        # Вычисляем координаты для центрирования квадрата
+        left = (cropped_image.width - local_square_size) // 2
+        top = (cropped_image.height - local_square_size) // 2
+        right = left + local_square_size
+        bottom = top + local_square_size
+        
+        # Вырезаем концентрический квадрат
+        final_image = cropped_image.crop((left, top, right, bottom))
+
+        return final_image
+
     def loop(self):
 
         html = self.driver.find_element(By.TAG_NAME, 'html')
@@ -71,7 +96,14 @@ class Simulator:
             png = self.driver.get_screenshot_as_png()
             im = Image.open(BytesIO(png))
             im = im.crop([0, 80, im.width-80, im.height-60])
-            im.save(f"./images/{i}.png")
+            
+            # Применяем поворот как будто съемка с дрона
+            rotated_im = self.rotate_image_like_drone(im, self.angle + math.pi / 2)
+            
+            # Передаем изображение в AutoPilot для анализа
+            self.autonomePilot.handle(rotated_im)
+            
+            rotated_im.save(f"./images/{i}.png")
             sleep(0.25)
 
         print("last position: ", self.driver.current_url)