diff --git a/advanced_autopilot.py b/advanced_autopilot.py new file mode 100644 index 0000000..41178f7 --- /dev/null +++ b/advanced_autopilot.py @@ -0,0 +1,418 @@ +""" +Расширенная система автопилота с улучшенным отслеживанием координат и угла БПЛА +""" + +import random +import cv2 +import numpy as np +from PIL import Image +import math +from typing import Optional, Tuple, Dict, List +from dataclasses import dataclass +import time + +random.seed(1) + +@dataclass +class DroneState: + """Структура для хранения состояния БПЛА""" + x: float = 0.0 + y: float = 0.0 + altitude: float = 10.0 + angle: float = 0.0 + velocity_x: float = 0.0 + velocity_y: float = 0.0 + angular_velocity: float = 0.0 + timestamp: float = 0.0 + confidence: float = 1.0 + +class AdvancedAutoPilot: + """ + Расширенный автопилот с улучшенным отслеживанием координат и угла БПЛА + """ + + def __init__(self, + initial_x: float = 0.0, + initial_y: float = 0.0, + initial_altitude: float = 10.0, + camera_fov: float = 60.0, + image_width: int = 640, + image_height: int = 480): + + # Состояние БПЛА + self.state = DroneState( + x=initial_x, + y=initial_y, + altitude=initial_altitude, + timestamp=time.time() + ) + + # Параметры камеры + self.camera_fov = math.radians(camera_fov) + self.image_width = image_width + self.image_height = image_height + + # Вычисляем коэффициент перевода пикселей в метры + self.update_pixel_to_meter_ratio() + + # История состояний для фильтрации + self.state_history: List[DroneState] = [] + self.max_history_size = 10 + + # Параметры фильтрации + self.velocity_filter_alpha = 0.3 # Коэффициент сглаживания скорости + self.position_filter_alpha = 0.7 # Коэффициент сглаживания позиции + + # Инициализация детекторов + self.orb_detector = cv2.ORB_create( + nfeatures=2000, + 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) + + # Предыдущее изображение + self.prev_image: Optional[np.ndarray] = None + self.prev_timestamp: float = 0.0 + + # Счетчики + self.frame_count = 0 + self.successful_tracking_frames = 0 + + # Статистика + self.tracking_stats = { + 'total_frames': 0, + 'successful_frames': 0, + 'failed_frames': 0, + 'avg_processing_time': 0.0 + } + + def update_pixel_to_meter_ratio(self): + """Обновляет коэффициент перевода пикселей в метры на основе высоты и FOV камеры""" + # Вычисляем размер пикселя на земле при текущей высоте + ground_distance = 2 * self.state.altitude * math.tan(self.camera_fov / 2) + self.pixel_to_meter_ratio = ground_distance / self.image_width + + def detect_and_match_keypoints(self, img1: np.ndarray, img2: np.ndarray) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[List], Optional[List], Optional[List]]: + """Обнаруживает и сопоставляет ключевые точки между двумя изображениями""" + # Обнаружение ключевых точек и дескрипторов + kp1, des1 = self.orb_detector.detectAndCompute(img1, None) + kp2, des2 = self.orb_detector.detectAndCompute(img2, None) + + if des1 is None or des2 is None or len(kp1) < 10 or len(kp2) < 10: + 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) < 8: # Увеличиваем минимальное количество совпадений + return None, None, 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) -> Optional[Dict]: + """Оценивает матрицу трансформации с улучшенной обработкой ошибок""" + try: + # Используем RANSAC для оценки матрицы гомографии + H, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 3.0) + + if H is None: + return None + + # Проверяем качество матрицы + if not self.is_valid_homography(H): + return None + + # Извлекаем параметры трансформации + 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 + + # Проверяем разумность масштаба + if scale < 0.5 or scale > 2.0: + return None + + return { + 'translation': (tx, ty), + 'rotation': angle, + 'scale': scale, + 'homography': H, + 'mask': mask, + 'inliers_ratio': np.sum(mask) / len(mask) if mask is not None else 0.0 + } + except Exception as e: + print(f"Ошибка при оценке трансформации: {e}") + return None + + def is_valid_homography(self, H: np.ndarray) -> bool: + """Проверяет валидность матрицы гомографии""" + if H is None: + return False + + # Проверяем, что матрица не вырождена + det = np.linalg.det(H) + if abs(det) < 1e-6: + return False + + # Проверяем, что элементы матрицы разумны + if np.any(np.abs(H) > 1000): + return False + + return True + + def update_drone_position(self, transformation_info: Dict, dt: float): + """Обновляет позицию и состояние БПЛА с учетом временного интервала""" + tx, ty = transformation_info['translation'] + rotation = transformation_info['rotation'] + scale = transformation_info['scale'] + inliers_ratio = transformation_info.get('inliers_ratio', 1.0) + + # Конвертируем смещение в пикселях в метры + dx_meters = tx * self.pixel_to_meter_ratio + dy_meters = ty * self.pixel_to_meter_ratio + + # Применяем поворот к смещению (учитываем текущий угол БПЛА) + cos_angle = math.cos(self.state.angle) + sin_angle = math.sin(self.state.angle) + + # Поворачиваем смещение в глобальные координаты + dx_global = dx_meters * cos_angle - dy_meters * sin_angle + dy_global = dx_meters * sin_angle + dy_meters * cos_angle + + # Вычисляем скорости + velocity_x = dx_global / dt if dt > 0 else 0.0 + velocity_y = dy_global / dt if dt > 0 else 0.0 + angular_velocity = rotation / dt if dt > 0 else 0.0 + + # Применяем фильтрацию + self.state.velocity_x = (self.velocity_filter_alpha * velocity_x + + (1 - self.velocity_filter_alpha) * self.state.velocity_x) + self.state.velocity_y = (self.velocity_filter_alpha * velocity_y + + (1 - self.velocity_filter_alpha) * self.state.velocity_y) + self.state.angular_velocity = (self.velocity_filter_alpha * angular_velocity + + (1 - self.velocity_filter_alpha) * self.state.angular_velocity) + + # Обновляем позицию с фильтрацией + new_x = self.state.x + dx_global + new_y = self.state.y + dy_global + new_angle = self.state.angle + rotation + + self.state.x = (self.position_filter_alpha * new_x + + (1 - self.position_filter_alpha) * self.state.x) + self.state.y = (self.position_filter_alpha * new_y + + (1 - self.position_filter_alpha) * self.state.y) + self.state.angle = (self.position_filter_alpha * new_angle + + (1 - self.position_filter_alpha) * self.state.angle) + + # Нормализуем угол + self.state.angle = math.atan2(math.sin(self.state.angle), math.cos(self.state.angle)) + + # Обновляем временную метку и уверенность + self.state.timestamp = time.time() + self.state.confidence = min(1.0, inliers_ratio * 1.5) # Увеличиваем уверенность при хороших совпадениях + + def handle(self, image: Image.Image) -> bool: + """Обрабатывает новый кадр и обновляет состояние БПЛА""" + start_time = time.time() + self.frame_count += 1 + self.tracking_stats['total_frames'] += 1 + + # Конвертируем изображение + current_image = np.array(image) + current_timestamp = time.time() + + if self.prev_image is None: + self.prev_image = current_image + self.prev_timestamp = current_timestamp + return False + + # Вычисляем временной интервал + dt = current_timestamp - self.prev_timestamp + + # Обнаруживаем и сопоставляем ключевые точки + 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, dt) + + # Добавляем состояние в историю + self.add_state_to_history() + + # Обновляем статистику + self.successful_tracking_frames += 1 + self.tracking_stats['successful_frames'] += 1 + + # Выводим информацию + self.print_tracking_info(transformation_info) + + # Визуализация + self.visualize_tracking(current_image, kp1, kp2, matches, transformation_info) + + success = True + else: + self.tracking_stats['failed_frames'] += 1 + success = False + else: + self.tracking_stats['failed_frames'] += 1 + success = False + + # Обновляем предыдущее изображение + self.prev_image = current_image + self.prev_timestamp = current_timestamp + + # Обновляем статистику времени обработки + processing_time = time.time() - start_time + self.tracking_stats['avg_processing_time'] = ( + (self.tracking_stats['avg_processing_time'] * (self.frame_count - 1) + processing_time) / self.frame_count + ) + + return success + + def add_state_to_history(self): + """Добавляет текущее состояние в историю""" + # Создаем копию текущего состояния + state_copy = DroneState( + x=self.state.x, + y=self.state.y, + altitude=self.state.altitude, + angle=self.state.angle, + velocity_x=self.state.velocity_x, + velocity_y=self.state.velocity_y, + angular_velocity=self.state.angular_velocity, + timestamp=self.state.timestamp, + confidence=self.state.confidence + ) + + self.state_history.append(state_copy) + + # Ограничиваем размер истории + if len(self.state_history) > self.max_history_size: + self.state_history.pop(0) + + def get_drone_state(self) -> Dict: + """Возвращает текущее состояние БПЛА в виде словаря""" + return { + 'x': self.state.x, + 'y': self.state.y, + 'altitude': self.state.altitude, + 'angle': self.state.angle, + 'angle_degrees': math.degrees(self.state.angle), + 'velocity_x': self.state.velocity_x, + 'velocity_y': self.state.velocity_y, + 'angular_velocity': self.state.angular_velocity, + 'confidence': self.state.confidence, + 'frame_count': self.frame_count, + 'timestamp': self.state.timestamp + } + + def get_tracking_stats(self) -> Dict: + """Возвращает статистику отслеживания""" + success_rate = (self.tracking_stats['successful_frames'] / + max(1, self.tracking_stats['total_frames'])) * 100 + + return { + **self.tracking_stats, + 'success_rate_percent': success_rate, + 'current_altitude': self.state.altitude, + 'pixel_to_meter_ratio': self.pixel_to_meter_ratio + } + + def set_altitude(self, altitude: float): + """Устанавливает высоту БПЛА и пересчитывает масштаб""" + self.state.altitude = altitude + self.update_pixel_to_meter_ratio() + + def reset_position(self, x: float = 0.0, y: float = 0.0, angle: float = 0.0): + """Сбрасывает позицию и угол БПЛА""" + self.state.x = x + self.state.y = y + self.state.angle = angle + self.state.velocity_x = 0.0 + self.state.velocity_y = 0.0 + self.state.angular_velocity = 0.0 + self.state_history.clear() + self.frame_count = 0 + self.successful_tracking_frames = 0 + + def print_tracking_info(self, transformation_info: Dict): + """Выводит информацию об отслеживании""" + state = self.get_drone_state() + stats = self.get_tracking_stats() + + print(f"Frame {self.frame_count}:") + print(f" Position: ({state['x']:.3f}, {state['y']:.3f}) m") + print(f" Angle: {state['angle_degrees']:.1f}°") + print(f" Velocity: ({state['velocity_x']:.2f}, {state['velocity_y']:.2f}) m/s") + print(f" Confidence: {state['confidence']:.2f}") + print(f" Success Rate: {stats['success_rate_percent']:.1f}%") + + def visualize_tracking(self, current_image: np.ndarray, kp1, kp2, matches, transformation_info: Dict): + """Визуализирует процесс отслеживания""" + if kp1 is None or kp2 is None or matches is None: + return + + # Рисуем сопоставления + img_matches = cv2.drawMatches(self.prev_image, kp1, current_image, kp2, matches, None, + flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS) + + # Добавляем информацию о состоянии + state = self.get_drone_state() + stats = self.get_tracking_stats() + + info_lines = [ + f"Position: ({state['x']:.2f}, {state['y']:.2f}) m", + f"Angle: {state['angle_degrees']:.1f}°", + f"Velocity: ({state['velocity_x']:.2f}, {state['velocity_y']:.2f}) m/s", + f"Altitude: {state['altitude']:.1f} m", + f"Confidence: {state['confidence']:.2f}", + f"Success Rate: {stats['success_rate_percent']:.1f}%" + ] + + for i, line in enumerate(info_lines): + cv2.putText(img_matches, line, (10, 30 + i * 25), + cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) + + cv2.imshow('Advanced Drone Tracking', img_matches) + cv2.waitKey(1) + + def act(self) -> float: + """Возвращает угол поворота для управления дроном""" + return self.state.angle + + def get_trajectory_history(self) -> Tuple[List[float], List[float], List[float]]: + """Возвращает историю траектории""" + x_history = [state.x for state in self.state_history] + y_history = [state.y for state in self.state_history] + angle_history = [math.degrees(state.angle) for state in self.state_history] + + return x_history, y_history, angle_history \ No newline at end of file diff --git a/autopilot.py b/autopilot.py index 5415ee1..85a4d46 100644 --- a/autopilot.py +++ b/autopilot.py @@ -1,9 +1,8 @@ import random - import cv2 - import numpy as np from PIL import Image +import math random.seed(1) @@ -15,12 +14,19 @@ class Pilot: 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 - def __init__(self): + def __init__(self, initial_x: float = 0.0, initial_y: float = 0.0): self.prev_image = None - self.angle = 0 + self.angle = 0.0 + self.x = initial_x + self.y = initial_y + self.frame_count = 0 + # Инициализация ORB детектора self.orb_detector = cv2.ORB_create( nfeatures=1000, @@ -32,6 +38,7 @@ class AutoPilot(Pilot): patchSize=31, fastThreshold=20 ) + # Инициализация матчера для сопоставления ключевых точек self.bf_matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True) @@ -48,7 +55,7 @@ class AutoPilot(Pilot): kp2, des2 = self.orb_detector.detectAndCompute(img2, None) if des1 is None or des2 is None: - return None, None, None + return None, None, None, None, None # Сопоставление ключевых точек matches = self.bf_matcher.match(des1, des2) @@ -63,7 +70,7 @@ class AutoPilot(Pilot): good_matches.append(match) if len(good_matches) < 4: - return None, None, None + return None, None, None, None, None # Извлечение координат сопоставленных точек src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2) @@ -109,6 +116,48 @@ class AutoPilot(Pilot): '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 = tx + dy_meters = ty + + # Применяем поворот к смещению (учитываем текущий угол БПЛА) + cos_angle = math.cos(self.angle) + sin_angle = math.sin(self.angle) + + # Поворачиваем смещение в глобальные координаты + 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): """ @@ -132,9 +181,19 @@ class AutoPilot(Pilot): 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) + # Добавляем информацию о позиции БПЛА + 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, 120), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2) + cv2.putText(img_matches, angle_text, (10, 150), 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) return @@ -144,23 +203,29 @@ class AutoPilot(Pilot): # Обнаруживаем и сопоставляем ключевые точки 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}") + print(f"Frame {self.frame_count}:") + 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'] + # Обновляем позицию и угол БПЛА + self.update_drone_position(transformation_info) + + # Выводим текущее состояние БПЛА + drone_state = self.get_drone_state() + print(f" Drone Position: ({drone_state['x']:.2f}, {drone_state['y']:.2f})") + print(f" Drone Angle: {drone_state['angle_degrees']:.1f}°") # Визуализация (опционально) img_matches = self.visualize_matches(self.prev_image, current_image, kp1, kp2, matches, transformation_info) - cv2.imshow('Matches', img_matches) + cv2.imshow('Drone Tracking', img_matches) cv2.waitKey(1) # Обновляем предыдущее изображение @@ -170,6 +235,13 @@ class AutoPilot(Pilot): """Возвращает угол поворота для управления дроном""" return self.angle + 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): def __init__(self, velocity: float = 1): diff --git a/examples.py b/examples.py new file mode 100644 index 0000000..1a3e291 --- /dev/null +++ b/examples.py @@ -0,0 +1,222 @@ +""" +Примеры использования системы отслеживания БПЛА +""" + +import numpy as np +from PIL import Image +import cv2 +from autopilot import AutoPilot +import matplotlib.pyplot as plt + +def example_basic_usage(): + """Базовый пример использования автопилота""" + print("=== Базовый пример использования ===") + + # Создаем автопилот с начальными параметрами + autopilot = AutoPilot( + initial_x=0.0, # Начальная X координата + initial_y=0.0, # Начальная Y координата + initial_altitude=15.0 # Начальная высота + ) + + # Создаем тестовые изображения + img1 = create_simple_test_image() + img2 = create_simple_test_image(offset_x=20, offset_y=10) + + # Обрабатываем изображения + autopilot.handle(img1) + autopilot.handle(img2) + + # Получаем состояние + state = autopilot.get_drone_state() + print(f"Позиция БПЛА: ({state['x']:.2f}, {state['y']:.2f})") + print(f"Угол БПЛА: {state['angle_degrees']:.1f}°") + print(f"Высота: {state['altitude']:.1f}м") + print(f"Обработано кадров: {state['frame_count']}") + +def example_altitude_change(): + """Пример изменения высоты и пересчета масштаба""" + print("\n=== Пример изменения высоты ===") + + autopilot = AutoPilot(initial_altitude=10.0) + + # Создаем изображения + img1 = create_simple_test_image() + img2 = create_simple_test_image(offset_x=10, offset_y=5) + + # Обрабатываем при высоте 10м + autopilot.handle(img1) + autopilot.handle(img2) + + state1 = autopilot.get_drone_state() + print(f"При высоте 10м: ({state1['x']:.2f}, {state1['y']:.2f})") + + # Изменяем высоту на 20м + autopilot.set_altitude(20.0) + + # Обрабатываем те же изображения при новой высоте + autopilot.reset_position() + autopilot.handle(img1) + autopilot.handle(img2) + + state2 = autopilot.get_drone_state() + print(f"При высоте 20м: ({state2['x']:.2f}, {state2['y']:.2f})") + + # Коэффициент должен измениться в 2 раза + ratio = state2['x'] / state1['x'] if state1['x'] != 0 else 0 + print(f"Коэффициент изменения: {ratio:.2f} (ожидается ~2.0)") + +def example_circular_movement(): + """Пример отслеживания кругового движения""" + print("\n=== Пример кругового движения ===") + + autopilot = AutoPilot() + + # Создаем базовое изображение + base_img = create_simple_test_image() + + # Параметры кругового движения + radius = 30 + center_x, center_y = 320, 240 + num_steps = 8 + + x_history = [] + y_history = [] + angle_history = [] + + # Симулируем круговое движение + for i in range(num_steps): + angle = 2 * np.pi * i / num_steps + + # Вычисляем смещение + dx = radius * np.cos(angle) + dy = radius * np.sin(angle) + + # Создаем смещенное изображение + current_img = create_simple_test_image( + offset_x=int(dx), + offset_y=int(dy), + rotation=angle * 0.1 + ) + + # Обрабатываем + autopilot.handle(current_img) + + # Сохраняем историю + state = autopilot.get_drone_state() + x_history.append(state['x']) + y_history.append(state['y']) + angle_history.append(state['angle_degrees']) + + print(f"Шаг {i+1}: ({state['x']:.2f}, {state['y']:.2f}), угол: {state['angle_degrees']:.1f}°") + + # Визуализируем траекторию + plot_circular_trajectory(x_history, y_history, angle_history) + +def example_position_reset(): + """Пример сброса позиции""" + print("\n=== Пример сброса позиции ===") + + autopilot = AutoPilot() + + # Создаем и обрабатываем изображения + img1 = create_simple_test_image() + img2 = create_simple_test_image(offset_x=15, offset_y=10) + + autopilot.handle(img1) + autopilot.handle(img2) + + state_before = autopilot.get_drone_state() + print(f"До сброса: ({state_before['x']:.2f}, {state_before['y']:.2f})") + + # Сбрасываем позицию + autopilot.reset_position(x=100.0, y=200.0, angle=np.pi/4) + + state_after = autopilot.get_drone_state() + print(f"После сброса: ({state_after['x']:.2f}, {state_after['y']:.2f}), угол: {state_after['angle_degrees']:.1f}°") + +def create_simple_test_image(offset_x=0, offset_y=0, rotation=0): + """Создает простое тестовое изображение с заданным смещением""" + width, height = 640, 480 + + # Создаем базовое изображение с геометрическими фигурами + img = np.zeros((height, width, 3), dtype=np.uint8) + + # Добавляем фигуры + cv2.rectangle(img, (100, 100), (200, 200), (255, 0, 0), -1) + cv2.circle(img, (400, 300), 50, (0, 255, 0), -1) + cv2.line(img, (50, 400), (550, 400), (0, 0, 255), 5) + cv2.rectangle(img, (300, 150), (350, 200), (255, 255, 0), -1) + + # Применяем трансформацию + if offset_x != 0 or offset_y != 0 or rotation != 0: + cos_rot = np.cos(rotation) + sin_rot = np.sin(rotation) + + M = np.float32([ + [cos_rot, -sin_rot, offset_x], + [sin_rot, cos_rot, offset_y] + ]) + + img = cv2.warpAffine(img, M, (width, height), + borderMode=cv2.BORDER_REPLICATE) + + return Image.fromarray(img) + +def plot_circular_trajectory(x_history, y_history, angle_history): + """Визуализирует круговую траекторию""" + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6)) + + # График траектории + ax1.plot(x_history, y_history, 'b-o', linewidth=2, markersize=8) + ax1.plot(x_history[0], y_history[0], 'go', markersize=12, label='Начало') + ax1.plot(x_history[-1], y_history[-1], 'ro', markersize=12, label='Конец') + ax1.set_xlabel('X координата (м)') + ax1.set_ylabel('Y координата (м)') + ax1.set_title('Круговая траектория БПЛА') + ax1.grid(True) + ax1.legend() + ax1.axis('equal') + + # График угла + ax2.plot(range(len(angle_history)), angle_history, 'r-o', linewidth=2, markersize=8) + ax2.set_xlabel('Номер шага') + ax2.set_ylabel('Угол поворота (градусы)') + ax2.set_title('Изменение угла поворота') + ax2.grid(True) + + plt.tight_layout() + plt.show() + +def example_error_handling(): + """Пример обработки ошибок и граничных случаев""" + print("\n=== Пример обработки ошибок ===") + + autopilot = AutoPilot() + + # Создаем изображение без текстуры (однородное) + blank_img = Image.fromarray(np.zeros((480, 640, 3), dtype=np.uint8)) + + # Обрабатываем - должно работать без ошибок + try: + autopilot.handle(blank_img) + autopilot.handle(blank_img) + print("Обработка однородного изображения прошла успешно") + except Exception as e: + print(f"Ошибка при обработке: {e}") + + # Проверяем состояние + state = autopilot.get_drone_state() + print(f"Состояние после обработки: ({state['x']:.2f}, {state['y']:.2f})") + +if __name__ == "__main__": + print("Примеры использования системы отслеживания БПЛА\n") + + # Запускаем все примеры + example_basic_usage() + example_altitude_change() + example_circular_movement() + example_position_reset() + example_error_handling() + + print("\nВсе примеры выполнены успешно!") \ No newline at end of file diff --git a/test_tracking.py b/test_tracking.py new file mode 100644 index 0000000..e011826 --- /dev/null +++ b/test_tracking.py @@ -0,0 +1,172 @@ +import numpy as np +from PIL import Image +import cv2 +from autopilot import AutoPilot +import matplotlib.pyplot as plt +import time + +def create_test_image(width=640, height=480, pattern_type="random"): + """Создает тестовое изображение с различными паттернами""" + if pattern_type == "random": + # Создаем изображение с случайными точками + img = np.random.randint(0, 255, (height, width, 3), dtype=np.uint8) + # Добавляем несколько четких объектов + cv2.rectangle(img, (100, 100), (200, 200), (255, 0, 0), -1) + cv2.circle(img, (400, 300), 50, (0, 255, 0), -1) + cv2.line(img, (50, 400), (550, 400), (0, 0, 255), 5) + elif pattern_type == "grid": + # Создаем сетку + img = np.zeros((height, width, 3), dtype=np.uint8) + for i in range(0, width, 50): + cv2.line(img, (i, 0), (i, height), (255, 255, 255), 1) + for i in range(0, height, 50): + cv2.line(img, (0, i), (width, i), (255, 255, 255), 1) + elif pattern_type == "circles": + # Создаем изображение с кругами + img = np.zeros((height, width, 3), dtype=np.uint8) + for i in range(5): + x = np.random.randint(50, width-50) + y = np.random.randint(50, height-50) + radius = np.random.randint(20, 80) + color = (np.random.randint(0, 255), np.random.randint(0, 255), np.random.randint(0, 255)) + cv2.circle(img, (x, y), radius, color, -1) + + return Image.fromarray(img) + +def simulate_drone_movement(base_image, dx, dy, rotation=0): + """Симулирует движение БПЛА, создавая смещенное изображение""" + img_array = np.array(base_image) + height, width = img_array.shape[:2] + + # Создаем матрицу трансформации + cos_rot = np.cos(rotation) + sin_rot = np.sin(rotation) + + # Матрица поворота и смещения + M = np.float32([ + [cos_rot, -sin_rot, dx], + [sin_rot, cos_rot, dy] + ]) + + # Применяем трансформацию + transformed = cv2.warpAffine(img_array, M, (width, height), + borderMode=cv2.BORDER_REPLICATE) + + return Image.fromarray(transformed) + +def test_drone_tracking(): + """Тестирует систему отслеживания БПЛА""" + print("Запуск тестирования системы отслеживания БПЛА...") + + # Создаем автопилот + autopilot = AutoPilot(initial_x=0.0, initial_y=0.0, initial_altitude=10.0) + + # Создаем базовое изображение + base_image = create_test_image(pattern_type="circles") + + # Траектория движения (dx, dy, rotation) + trajectory = [ + (10, 0, 0), # Движение вперед + (10, 5, 0.1), # Движение вперед с небольшим поворотом + (0, 10, 0), # Движение вправо + (-5, 5, -0.1), # Движение назад-вправо с поворотом + (5, -5, 0.05), # Движение вперед-влево + (0, -10, 0), # Движение влево + (-10, 0, 0), # Движение назад + ] + + # Массивы для хранения истории позиций + x_history = [] + y_history = [] + angle_history = [] + + # Обрабатываем каждое изображение в траектории + for i, (dx, dy, rotation) in enumerate(trajectory): + print(f"\n--- Кадр {i+1} ---") + + # Создаем смещенное изображение + current_image = simulate_drone_movement(base_image, dx, dy, rotation) + + # Обрабатываем изображение + autopilot.handle(current_image) + + # Получаем состояние БПЛА + state = autopilot.get_drone_state() + x_history.append(state['x']) + y_history.append(state['y']) + angle_history.append(state['angle_degrees']) + + # Обновляем базовое изображение для следующего кадра + base_image = current_image + + # Небольшая пауза для визуализации + time.sleep(0.5) + + # Закрываем окно OpenCV + cv2.destroyAllWindows() + + # Визуализируем траекторию + plot_trajectory(x_history, y_history, angle_history) + +def plot_trajectory(x_history, y_history, angle_history): + """Визуализирует траекторию движения БПЛА""" + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6)) + + # График траектории + ax1.plot(x_history, y_history, 'b-o', linewidth=2, markersize=8) + ax1.plot(x_history[0], y_history[0], 'go', markersize=12, label='Начало') + ax1.plot(x_history[-1], y_history[-1], 'ro', markersize=12, label='Конец') + ax1.set_xlabel('X координата (м)') + ax1.set_ylabel('Y координата (м)') + ax1.set_title('Траектория движения БПЛА') + ax1.grid(True) + ax1.legend() + ax1.axis('equal') + + # График угла поворота + ax2.plot(range(len(angle_history)), angle_history, 'r-o', linewidth=2, markersize=8) + ax2.set_xlabel('Номер кадра') + ax2.set_ylabel('Угол поворота (градусы)') + ax2.set_title('Изменение угла поворота БПЛА') + ax2.grid(True) + + plt.tight_layout() + plt.show() + +def test_single_frame(): + """Тестирует обработку одного кадра""" + print("Тестирование обработки одного кадра...") + + autopilot = AutoPilot() + + # Создаем два похожих изображения + img1 = create_test_image(pattern_type="circles") + img2 = simulate_drone_movement(img1, 15, 10, 0.05) + + # Обрабатываем первое изображение + autopilot.handle(img1) + + # Обрабатываем второе изображение + autopilot.handle(img2) + + # Получаем состояние + state = autopilot.get_drone_state() + print(f"Позиция БПЛА: ({state['x']:.2f}, {state['y']:.2f})") + print(f"Угол БПЛА: {state['angle_degrees']:.1f}°") + + cv2.destroyAllWindows() + +if __name__ == "__main__": + print("Выберите тест:") + print("1. Тест одного кадра") + print("2. Тест полной траектории") + + choice = input("Введите номер теста (1 или 2): ").strip() + + if choice == "1": + test_single_frame() + elif choice == "2": + test_drone_tracking() + else: + print("Неверный выбор. Запускаю тест одного кадра...") + test_single_frame() \ No newline at end of file