diff --git a/autopilot.py b/autopilot.py index 3948f73..5415ee1 100644 --- a/autopilot.py +++ b/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() diff --git a/main.py b/main.py index cfb4fd7..ecb1a09 100644 --- a/main.py +++ b/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() diff --git a/requirements.txt b/requirements.txt index e607ee6..9ea02a8 100644 --- a/requirements.txt +++ b/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 diff --git a/simulator.py b/simulator.py index 8477837..63e4c8c 100644 --- a/simulator.py +++ b/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)