204 lines
8.6 KiB
Python
204 lines
8.6 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Тестовый скрипт для проверки системы отслеживания траектории беспилотника
|
||
"""
|
||
|
||
import numpy as np
|
||
from PIL import Image
|
||
import cv2
|
||
from autopilot import AutoPilot, RandomPilot
|
||
from simulator import Simulator, SimMode
|
||
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()
|
||
|
||
def test_trajectory_tracking():
|
||
"""Тестирует систему отслеживания траектории"""
|
||
print("Запуск теста системы отслеживания траектории...")
|
||
|
||
# Создаем пилотов
|
||
operator_pilot = RandomPilot(velocity=1.0)
|
||
autonome_pilot = AutoPilot(initial_x=0.0, initial_y=0.0)
|
||
|
||
# Создаем симулятор
|
||
simulator = Simulator(operator_pilot, autonome_pilot)
|
||
|
||
try:
|
||
# Запускаем симуляцию
|
||
simulator.loop()
|
||
|
||
print("Симуляция завершена!")
|
||
print(f"Всего точек траектории: {len(simulator.trajectory_x)}")
|
||
print(f"Режим оператора: {sum(1 for mode in simulator.trajectory_modes if mode == SimMode.OPERATOR)} точек")
|
||
print(f"Автономный режим: {sum(1 for mode in simulator.trajectory_modes if mode == SimMode.AUTONOME)} точек")
|
||
|
||
except Exception as e:
|
||
print(f"Ошибка во время симуляции: {e}")
|
||
finally:
|
||
# Закрываем браузер
|
||
simulator.driver.quit()
|
||
|
||
if __name__ == "__main__":
|
||
print("Выберите тест:")
|
||
print("1. Тест одного кадра")
|
||
print("2. Тест полной траектории")
|
||
|
||
choice = input("Введите номер теста (1 или 2): ").strip()
|
||
|
||
if choice == "1":
|
||
test_single_frame()
|
||
elif choice == "2":
|
||
test_trajectory_tracking()
|
||
else:
|
||
print("Неверный выбор. Запускаю тест одного кадра...")
|
||
test_single_frame() |