207 lines
7.7 KiB
Python
207 lines
7.7 KiB
Python
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() |