ref: replace project files
This commit is contained in:
360
autopilot.py
360
autopilot.py
@@ -1,360 +0,0 @@
|
||||
import random
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import math
|
||||
from visualization import VisualizationManager
|
||||
|
||||
random.seed(1)
|
||||
|
||||
class Pilot:
|
||||
def __init__(self): pass
|
||||
def handle(self, image: Image): pass
|
||||
def act(self) -> tuple[float, float] | None: pass
|
||||
def get_position(self) -> tuple[float, float]: pass
|
||||
|
||||
class AutoPilot(Pilot):
|
||||
prev_image: np.ndarray | None
|
||||
angle: float
|
||||
x: float # Координата X БПЛА
|
||||
y: float # Координата Y БПЛА
|
||||
orb_detector: cv2.ORB
|
||||
bf_matcher: cv2.BFMatcher
|
||||
frame_count: int
|
||||
image_center: tuple # Центр изображения (x, y)
|
||||
viz_manager: VisualizationManager # Менеджер визуализации (опционально)
|
||||
|
||||
def __init__(self, viz_manager=None):
|
||||
self.prev_image = None
|
||||
self.angle = 0.0
|
||||
self.x = 0.0
|
||||
self.y = 0.0
|
||||
self.frame_count = 0
|
||||
self.image_center = (0, 0) # Будет обновлено при получении первого изображения
|
||||
self.viz_manager = viz_manager # Менеджер визуализации
|
||||
|
||||
# Инициализация 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 get_position(self):
|
||||
return self.x, self.y
|
||||
|
||||
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, 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, None, None
|
||||
|
||||
# Получаем центр изображения
|
||||
height1, width1 = img1.shape[:2]
|
||||
height2, width2 = img2.shape[:2]
|
||||
center_x1, center_y1 = width1 // 2, height1 // 2
|
||||
center_x2, center_y2 = width2 // 2, height2 // 2
|
||||
|
||||
# Извлекаем координаты сопоставленных точек и отцентрируем их
|
||||
src_pts = []
|
||||
dst_pts = []
|
||||
|
||||
for match in good_matches:
|
||||
# Координаты точки в первом изображении
|
||||
pt1 = kp1[match.queryIdx].pt
|
||||
src_pts.append([center_x1 - pt1[0], pt1[1] - center_y1])
|
||||
|
||||
# Координаты точки во втором изображении
|
||||
pt2 = kp2[match.trainIdx].pt
|
||||
dst_pts.append([center_x2 - pt2[0], pt2[1] - center_y2])
|
||||
|
||||
# Конвертируем в numpy массивы
|
||||
src_pts = np.float32(src_pts).reshape(-1, 1, 2)
|
||||
dst_pts = np.float32(dst_pts).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 update_drone_position(self, transformation_info: dict):
|
||||
"""
|
||||
Обновляет позицию и угол БПЛА на основе трансформации изображения
|
||||
Координаты уже отцентрированы относительно центра изображения
|
||||
"""
|
||||
tx, ty = transformation_info['translation']
|
||||
rotation = transformation_info['rotation']
|
||||
scale = transformation_info['scale']
|
||||
|
||||
# Координаты уже отцентрированы, поэтому используем их напрямую
|
||||
dx_meters = ty
|
||||
dy_meters = tx
|
||||
|
||||
# Применяем поворот к смещению (учитываем текущий угол БПЛА)
|
||||
cos_angle = math.cos(self.angle)
|
||||
sin_angle = math.sin(self.angle)
|
||||
|
||||
# Поворачиваем смещение в глобальные координаты
|
||||
# Обратите внимание: dy_meters инвертирован, так как в изображениях Y направлен вниз
|
||||
dx_global = dx_meters * cos_angle - dy_meters * sin_angle
|
||||
dy_global = dx_meters * sin_angle + dy_meters * cos_angle
|
||||
|
||||
# Обновляем координаты БПЛА
|
||||
self.x += dx_global
|
||||
self.y += dy_global
|
||||
|
||||
# Обновляем угол БПЛА
|
||||
self.angle += rotation
|
||||
|
||||
# Нормализуем угол в диапазоне [-π, π]
|
||||
self.angle = math.atan2(math.sin(self.angle), math.cos(self.angle))
|
||||
|
||||
def get_drone_state(self) -> dict:
|
||||
"""
|
||||
Возвращает текущее состояние БПЛА
|
||||
"""
|
||||
return {
|
||||
'x': self.x,
|
||||
'y': self.y,
|
||||
'angle': self.angle,
|
||||
'angle_degrees': math.degrees(self.angle),
|
||||
'frame_count': self.frame_count
|
||||
}
|
||||
|
||||
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}"
|
||||
info_text4 = f"Image Center: {self.image_center}"
|
||||
|
||||
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)
|
||||
cv2.putText(img_matches, info_text4, (10, 120), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
|
||||
|
||||
# Добавляем информацию о позиции БПЛА
|
||||
drone_state = self.get_drone_state()
|
||||
pos_text = f"Drone Pos: ({drone_state['x']:.2f}, {drone_state['y']:.2f})"
|
||||
angle_text = f"Drone Angle: {drone_state['angle_degrees']:.1f}°"
|
||||
|
||||
cv2.putText(img_matches, pos_text, (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)
|
||||
cv2.putText(img_matches, angle_text, (10, 180), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)
|
||||
|
||||
return img_matches
|
||||
|
||||
def handle(self, image: Image):
|
||||
self.frame_count += 1
|
||||
|
||||
if self.prev_image is None:
|
||||
self.prev_image = self.image_to_numpy(image)
|
||||
# Вычисляем центр изображения
|
||||
height, width = self.prev_image.shape[:2]
|
||||
self.image_center = (width // 2, height // 2)
|
||||
|
||||
# Обновляем визуализацию детекции
|
||||
if self.viz_manager:
|
||||
kp, _ = self.orb_detector.detectAndCompute(self.prev_image, None)
|
||||
self.viz_manager.update_detection(self.prev_image, kp)
|
||||
|
||||
return
|
||||
|
||||
# Конвертируем текущее изображение
|
||||
current_image = self.image_to_numpy(image)
|
||||
|
||||
# Обновляем центр изображения
|
||||
height, width = current_image.shape[:2]
|
||||
self.image_center = (width // 2, height // 2)
|
||||
|
||||
# Обнаруживаем и сопоставляем ключевые точки
|
||||
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:
|
||||
# Обновляем позицию и угол БПЛА
|
||||
self.update_drone_position(transformation_info)
|
||||
|
||||
# Выводим текущее состояние БПЛА
|
||||
drone_state = self.get_drone_state()
|
||||
self.viz_manager.update_drone_trajectory(drone_state['x'], drone_state['y'])
|
||||
print(f" [Pilot] Drone Position: ({drone_state['x']:.2f}, {drone_state['y']:.2f})")
|
||||
print(f" [Pilot] Angle: {drone_state['angle_degrees']:.1f}°")
|
||||
|
||||
# Обновляем визуализацию
|
||||
if self.viz_manager:
|
||||
# Обновляем детекцию ключевых точек
|
||||
self.viz_manager.update_detection(current_image, kp2)
|
||||
# Обновляем сопоставление точек
|
||||
self.viz_manager.update_matches(self.prev_image, current_image,
|
||||
kp1, kp2, matches, transformation_info)
|
||||
|
||||
# Визуализация (опционально, если нет менеджера визуализации)
|
||||
if not self.viz_manager:
|
||||
img_matches = self.visualize_matches(self.prev_image, current_image,
|
||||
kp1, kp2, matches, transformation_info)
|
||||
cv2.imshow('Drone Tracking', img_matches)
|
||||
cv2.waitKey(1)
|
||||
|
||||
# Обновляем предыдущее изображение
|
||||
self.prev_image = current_image
|
||||
|
||||
def act(self) -> tuple[float, float] | None:
|
||||
"""
|
||||
Возвращает угол поворота для управления дроном, чтобы он стремился к точке (0, 0).
|
||||
Если дрон находится рядом с началом координат (в радиусе 1 метра), возвращает None.
|
||||
"""
|
||||
# Расстояние до цели (0, 0)
|
||||
distance_to_target = math.sqrt(self.x**2 + self.y**2)
|
||||
|
||||
# Если дрон находится рядом с целью, останавливаемся
|
||||
if distance_to_target < 30.0:
|
||||
return None
|
||||
|
||||
# Вычисляем угол к цели
|
||||
target_angle = math.atan2(-self.y, -self.x) # Отрицательные координаты, так как движемся к (0,0)
|
||||
|
||||
# Вычисляем разность углов (направление поворота)
|
||||
angle_diff = target_angle - self.angle
|
||||
|
||||
print(self.angle, target_angle, angle_diff)
|
||||
|
||||
# Нормализуем разность углов в диапазон [-π, π]
|
||||
angle_diff %= 2 * math.pi
|
||||
if angle_diff >= math.pi:
|
||||
angle_diff -= 2 * math.pi
|
||||
|
||||
# Возвращаем угол поворота (положительный - поворот влево, отрицательный - вправо)
|
||||
return max(min(0.1, angle_diff), -0.1), min(10., distance_to_target / 2)
|
||||
|
||||
def reset_position(self, x: float = 0.0, y: float = 0.0, angle: float = 0.0):
|
||||
"""Сбрасывает позицию и угол БПЛА"""
|
||||
self.x = x
|
||||
self.y = y
|
||||
self.angle = angle
|
||||
self.frame_count = 0
|
||||
|
||||
|
||||
class RandomPilot(Pilot):
|
||||
counter: int
|
||||
|
||||
def __init__(self, velocity: float = 1):
|
||||
self.counter = 0
|
||||
|
||||
def act(self) -> tuple[float, float] | None:
|
||||
self.counter += 1
|
||||
if self.counter > 300:
|
||||
return None
|
||||
|
||||
return 1 / (self.counter + 20), 10.0
|
||||
|
||||
# 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()
|
||||
Reference in New Issue
Block a user